feature: chunk on headings

Author: densumesh

server/src/bin/scrape-worker.rs

@@ -8,7 +8,6 @@ use std::sync::{
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
 
 use actix_web::web;
-use trieve_server::handlers::chunk_handler::ChunkReqPayload;
 use trieve_server::operators::chunk_operator::create_chunk_metadata;
 use trieve_server::operators::dataset_operator::get_dataset_by_id_query;
 use trieve_server::{
@@ -21,6 +20,9 @@ use trieve_server::{
     establish_connection, get_env,
     operators::crawl_operator::{get_chunk_html, get_images, get_tags, update_crawl_status},
 };
+use trieve_server::{
+    handlers::chunk_handler::ChunkReqPayload, operators::crawl_operator::chunk_markdown,
+};
 use ureq::json;
 
 async fn crawl(
@@ -92,25 +94,31 @@ async fn crawl(
         let page_markdown = page.markdown.clone().unwrap_or_default();
         let page_tags = get_tags(page_link.clone());
 
-        let chunk = ChunkReqPayload {
-            chunk_html: Some(get_chunk_html(
-                page_markdown.clone(),
-                page_title.clone(),
-                "".to_string(),
-                0,
-                None,
-            )),
-            link: Some(page_link.clone()),
-            tag_set: Some(page_tags),
-            image_urls: Some(get_images(&page_markdown.clone())),
-            metadata: Some(json!({
-                "title": page_title.clone(),
-                "description": page_description.clone(),
-                "url": page_link.clone(),
-            })),
-            ..Default::default()
-        };
-        chunks.push(chunk);
+        let chunk_html = get_chunk_html(
+            page_markdown.clone(),
+            page_title.clone(),
+            "".to_string(),
+            0,
+            None,
+        );
+
+        let chunked_markdown = chunk_markdown(&chunk_html.clone());
+
+        for chunk in chunked_markdown {
+            let chunk = ChunkReqPayload {
+                chunk_html: Some(chunk.clone()),
+                link: Some(page_link.clone()),
+                tag_set: Some(page_tags.clone()),
+                image_urls: Some(get_images(&chunk.clone())),
+                metadata: Some(json!({
+                    "title": page_title.clone(),
+                    "description": page_description.clone(),
+                    "url": page_link.clone(),
+                })),
+                ..Default::default()
+            };
+            chunks.push(chunk);
+        }
     }
 
     let dataset = get_dataset_by_id_query(

server/src/operators/crawl_operator.rs

@@ -402,10 +402,36 @@ impl Cleaners {
 }
 
 pub fn get_images(markdown_content: &str) -> Vec<String> {
-    let image_pattern = Regex::new(r"!\[.*?\]\((.*?\.(?:png|webp))\)").unwrap();
+    let image_pattern = Regex::new(r"!\[.*?\]\((.*?\.(?:png|webp|jpeg|jpg))\)").unwrap();
     image_pattern
         .captures_iter(markdown_content)
         .filter_map(|cap| cap.get(1))
         .map(|m| m.as_str().to_string())
         .collect()
 }
+
+pub fn chunk_markdown(markdown: &str) -> Vec<String> {
+    let re = Regex::new(r"(?m)^(#{1,6}\s.+)$").unwrap();
+    let mut chunks = Vec::new();
+    let mut current_chunk = String::new();
+
+    for line in markdown.lines() {
+        if re.is_match(line) {
+            if !current_chunk.is_empty() {
+                chunks.push(current_chunk.trim().to_string());
+                current_chunk = String::new();
+            }
+            current_chunk.push_str(line);
+            current_chunk.push('\n');
+        } else {
+            current_chunk.push_str(line);
+            current_chunk.push('\n');
+        }
+    }
+
+    if !current_chunk.is_empty() {
+        chunks.push(current_chunk.trim().to_string());
+    }
+
+    chunks
+}