feature: incrementally add pages

devflowinc · Dec 14, 2024 · ee2734a · ee2734a
1 parent c589fbe
commit ee2734a
Show file tree

Hide file tree

Showing 8 changed files with 264 additions and 181 deletions.
diff --git a/frontends/search/src/components/UploadFile.tsx b/frontends/search/src/components/UploadFile.tsx
@@ -17,7 +17,11 @@ interface RequestBody {
   group_tracking_id?: string;
   metadata: any;
   time_stamp?: string;
-  use_pdf2md_ocr?: boolean;
+  pdf2md_options?: {
+    use_pdf2md_ocr: boolean;
+    system_prompt?: string;
+    split_headings?: boolean;
+  };
 }
 
 export const UploadFile = () => {
@@ -145,7 +149,7 @@ export const UploadFile = () => {
         split_delimiters: splitDelimiters(),
         target_splits_per_chunk: targetSplitsPerChunk(),
         rebalance_chunks: rebalanceChunks(),
-        use_pdf2md_ocr: useGptChunking(),
+        pdf2md_options: { use_pdf2md_ocr: useGptChunking() },
         group_tracking_id:
           groupTrackingId() === "" ? undefined : groupTrackingId(),
         // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment

diff --git a/pdf2md/server/src/models.rs b/pdf2md/server/src/models.rs
@@ -251,7 +251,7 @@ impl From<TaskResponse> for Vec<Chunk> {
 
 #[derive(Debug, serde::Serialize, serde::Deserialize, Clone)]
 pub struct GetTaskRequest {
-    pub pagination_token: Option<uuid::Uuid>,
+    pub pagination_token: Option<u32>,
     pub limit: Option<u32>,
 }
 
@@ -265,7 +265,7 @@ pub struct GetTaskResponse {
     pub status: String,
     pub created_at: String,
     pub pages: Option<Vec<Chunk>>,
-    pub pagination_token: Option<String>,
+    pub pagination_token: Option<u32>,
 }
 
 impl GetTaskResponse {
@@ -296,7 +296,7 @@ impl GetTaskResponse {
             pages_processed: task.pages_processed,
             status: task.status,
             created_at: task.created_at.to_string(),
-            pagination_token: pages.last().map(|c| c.id.clone()),
+            pagination_token: pages.last().map(|c| c.page),
             pages: Some(pages.into_iter().map(Chunk::from).collect()),
         }
     }

diff --git a/pdf2md/server/src/operators/clickhouse.rs b/pdf2md/server/src/operators/clickhouse.rs
@@ -170,18 +170,18 @@ pub async fn get_task(
 pub async fn get_task_pages(
     task: FileTaskClickhouse,
     limit: Option<u32>,
-    offset_id: Option<uuid::Uuid>,
+    offset_id: Option<u32>,
     clickhouse_client: &clickhouse::Client,
 ) -> Result<Vec<ChunkClickhouse>, ServiceError> {
     if FileTaskStatus::from(task.status.clone()) == FileTaskStatus::Completed || task.pages > 0 {
         let limit = limit.unwrap_or(20);
 
         let pages: Vec<ChunkClickhouse> = clickhouse_client
             .query(
-                "SELECT ?fields FROM file_chunks WHERE task_id = ? AND id > ? ORDER BY page LIMIT ?",
+                "SELECT ?fields FROM file_chunks WHERE task_id = ? AND page > ? ORDER BY page LIMIT ?",
             )
             .bind(task.id.clone())
-            .bind(offset_id.unwrap_or(uuid::Uuid::nil()))
+            .bind(offset_id.unwrap_or(0))
             .bind(limit)
             .fetch_all()
             .await

diff --git a/pdf2md/server/src/routes/jinja_templates.rs b/pdf2md/server/src/routes/jinja_templates.rs
@@ -3,7 +3,7 @@ use crate::{
     get_env, Templates,
 };
 use actix_web::{get, HttpResponse};
-use minijinja::{context, path_loader, Environment};
+use minijinja::context;
 
 #[utoipa::path(
   get,

diff --git a/server/src/bin/csv-jsonl-worker.rs b/server/src/bin/csv-jsonl-worker.rs
@@ -505,7 +505,7 @@ async fn process_csv_jsonl_file(
             rebalance_chunks: Some(false),
             split_delimiters: None,
             target_splits_per_chunk: None,
-            use_pdf2md_ocr: None,
+            pdf2md_options: None,
             base64_file: "".to_string(),
         },
         csv_jsonl_worker_message.dataset_id,