feature: expose custom html based splitter

Author: cdxker

server/src/bin/crawl-worker.rs

@@ -452,7 +452,11 @@ async fn get_chunks_with_firecrawl(
             }
         }
 
-        let chunked_html = chunk_html(&page_html.clone(), &scrape_request.crawl_options);
+        let chunked_html = chunk_html(
+            &page_html.clone(),
+            scrape_request.crawl_options.heading_remove_strings.clone(),
+            scrape_request.crawl_options.body_remove_strings.clone(),
+        );
 
         for chunk in chunked_html {
             let heading = chunk.0.last().unwrap_or(&String::new()).clone();

server/src/handlers/chunk_handler.rs

@@ -14,6 +14,7 @@ use crate::middleware::api_version::APIVersion;
 use crate::operators::chunk_operator::get_metadata_from_id_query;
 use crate::operators::chunk_operator::*;
 use crate::operators::clickhouse_operator::{get_latency_from_header, ClickHouseEvent, EventQueue};
+use crate::operators::crawl_operator;
 use crate::operators::dataset_operator::{
     get_dataset_usage_query, ChunkDeleteMessage, DeleteMessage,
 };
@@ -2820,6 +2821,97 @@ pub async fn generate_off_chunks(
         .streaming(completion_stream))
 }
 
+#[derive(Debug, Serialize, Deserialize, ToSchema)]
+#[schema(example = json!({
+    "chunk_html": "",
+    "heading_remove_strings": ["###", "##", "#"],
+    "body_remove_strings": ["Warning:", "Note:"]
+}))]
+pub struct ChunkHtmlContentReqPayload {
+    /// The HTML content to be split into chunks
+    pub chunk_html: String,
+    /// Text strings to remove from headings when creating chunks for each page
+    pub heading_remove_strings: Option<Vec<String>>,
+    /// Text strings to remove from body when creating chunks for each page
+    pub body_remove_strings: Option<Vec<String>>,
+}
+
+#[derive(Debug, Serialize, Deserialize, ToSchema)]
+#[schema(example = json!({
+    "chunks": [
+        {
+            "headings": ["Title Heading", "Sub Heading 1", "Sub Sub Heading 1"],
+            "body": "This is the body of the content"
+        },
+        {
+            "headings": ["Title Heading", "Sub Heading 1", "Sub Sub Heading 2"],
+            "body": "This is the body of the content"
+        }
+        // ...
+    ]
+}))]
+pub struct SplitHtmlResponse {
+    pub chunks: Vec<ChunkedContent>,
+}
+
+#[derive(Debug, Serialize, Deserialize, ToSchema)]
+#[schema(example = json!({
+    "headings": ["Title Heading", "Sub Heading 1", "Last SubHeading"],
+    "body": "This is the body of the content"
+}))]
+pub struct ChunkedContent {
+    /// The headings of the content in order of when they appear
+    pub headings: Vec<String>,
+    /// The body of the content
+    pub body: String,
+}
+
+/// Split HTML Content into Chunks
+///
+/// This endpoint receives a single html string and splits it into chunks based on the headings and
+/// body content. The headings are split based on headding html tags. chunk_html has a maximum size
+/// of 256Kb.
+#[utoipa::path(
+    post,
+    path = "/chunk/split",
+    context_path = "/api",
+    tag = "Chunk",
+    request_body(content = ChunkHtmlContentReqPayload, description = "JSON request payload to perform RAG on some chunks (chunks)", content_type = "application/json"),
+    responses(
+        (
+            status = 200, description = "This will be a JSON response of the chunks split from the HTML content with the headings and body",
+            body = SplitHtmlResponse,
+        ),
+        (
+            status = 413, description = "Payload too large, if the HTML contnet is greater than 256Kb",
+            body = ErrorResponseBody,
+        ),
+    ),
+)]
+#[tracing::instrument]
+pub async fn split_html_content(
+    body: web::Json<ChunkHtmlContentReqPayload>,
+) -> Result<HttpResponse, ServiceError> {
+    if body.chunk_html.bytes().len() >= 262_144 {
+        return Err(ServiceError::PayloadTooLarge(
+            "The HTML content is too large".to_string(),
+        ));
+    }
+
+    let chunked_content = crawl_operator::chunk_html(
+        &body.chunk_html,
+        body.heading_remove_strings.clone(),
+        body.body_remove_strings.clone(),
+    );
+
+    Ok(HttpResponse::Ok().json(SplitHtmlResponse {
+        chunks: chunked_content
+            .into_iter()
+            .map(|(headings, body)| ChunkedContent { headings, body })
+            .collect(),
+    }))
+}
+
 pub fn check_completion_param_validity(
     temperature: Option<f32>,
     frequency_penalty: Option<f32>,

server/src/lib.rs

@@ -175,6 +175,7 @@ impl Modify for SecurityAddon {
         handlers::chunk_handler::create_chunk,
         handlers::chunk_handler::update_chunk,
         handlers::chunk_handler::delete_chunk,
+        handlers::chunk_handler::split_html_content,
         handlers::chunk_handler::get_recommended_chunks,
         handlers::chunk_handler::update_chunk_by_tracking_id,
         handlers::chunk_handler::search_chunks,
@@ -274,6 +275,9 @@ impl Modify for SecurityAddon {
             handlers::chunk_handler::SearchResponseTypes,
             handlers::chunk_handler::CreateBatchChunkReqPayload,
             handlers::chunk_handler::SingleQueuedChunkResponse,
+            handlers::chunk_handler::ChunkHtmlContentReqPayload,
+            handlers::chunk_handler::SplitHtmlResponse,
+            handlers::chunk_handler::ChunkedContent,
             handlers::chunk_handler::BatchQueuedChunkResponse,
             handlers::chunk_handler::ReturnQueuedChunk,
             handlers::chunk_handler::RecommendChunksResponseBody,
@@ -929,6 +933,11 @@ pub fn main() -> std::io::Result<()> {
                                         .route(web::put().to(handlers::chunk_handler::update_chunk))
                                         .route(web::delete().to(handlers::chunk_handler::bulk_delete_chunk)),
                                 )
+                                .service(
+                                    web::resource("split").route(
+                                        web::post().to(handlers::chunk_handler::split_html_content),
+                                    ),
+                                )
                                 .service(web::resource("/recommend").route(
                                     web::post().to(handlers::chunk_handler::get_recommended_chunks),
                                 )

server/src/operators/crawl_operator.rs

@@ -602,7 +602,11 @@ pub fn get_tags(url: String) -> Vec<String> {
     Vec::new()
 }
 
-pub fn chunk_html(html: &str, crawl_options: &CrawlOptions) -> Vec<(Vec<String>, String)> {
+pub fn chunk_html(
+    html: &str,
+    heading_remove_strings: Option<Vec<String>>,
+    body_remove_strings: Option<Vec<String>>,
+) -> Vec<(Vec<String>, String)> {
     let re = Regex::new(r"(?i)<h[1-6].*?>").unwrap();
     let mut chunks = Vec::new();
     let mut current_chunk = String::new();
@@ -656,14 +660,14 @@ pub fn chunk_html(html: &str, crawl_options: &CrawlOptions) -> Vec<(Vec<String>,
             let mut headings_text = headings.last().unwrap_or(&String::new()).clone();
             let mut content = content.clone();
 
-            if let Some(heading_remove_strings) = &crawl_options.heading_remove_strings {
+            if let Some(heading_remove_strings) = &heading_remove_strings {
                 heading_remove_strings.iter().for_each(|remove_string| {
                     headings_text = headings_text.replace(remove_string, "");
                 });
                 headings.pop();
                 headings.push(headings_text);
             }
-            if let Some(body_remove_strings) = &crawl_options.body_remove_strings {
+            if let Some(body_remove_strings) = &body_remove_strings {
                 body_remove_strings.iter().for_each(|remove_string| {
                     content = content.replace(remove_string, "");
                 });