documents resposeDocument compatible - Adithya S K

Author: adithya-s-k

omniparse/documents/__init__.py

@@ -3,8 +3,9 @@
 import subprocess
 from omniparse.documents.parse import parse_single_pdf
 from omniparse.utils import encode_images
+from omniparse.models import responseDocument
 # Function to handle PDF parsing
-def parse_pdf(input_data , model_state) -> dict:
+def parse_pdf(input_data , model_state) -> responseDocument:
     try:
         if isinstance(input_data, bytes):
             with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf_file:
@@ -22,18 +23,23 @@ def parse_pdf(input_data , model_state) -> dict:
             raise ValueError("Invalid input data format. Expected bytes or PDF file path.")
 
         full_text, images, out_meta = parse_single_pdf(input_path, model_state.model_list)
-        images = encode_images(images)
+        
+        parse_pdf_result = responseDocument(
+            text=full_text,
+            metadata=out_meta
+        )
+        encode_images(images,parse_pdf_result)
 
         if cleanup_tempfile:
             os.remove(input_path)
 
-        return {"message": "PDF parsed successfully", "markdown": full_text, "metadata": out_meta, "images": images}
+        return parse_pdf_result
 
     except Exception as e:
         raise RuntimeError(f"Error parsing PPT: {str(e)}")
 
 # Function to handle PPT and DOC parsing
-def parse_ppt(input_data ,model_state) -> dict:
+def parse_ppt(input_data ,model_state) -> responseDocument:
     try:
         if isinstance(input_data, bytes):
             print("Recieved ppt file")
@@ -58,15 +64,21 @@ def parse_ppt(input_data ,model_state) -> dict:
         full_text, images, out_meta = parse_single_pdf(input_path, model_state.model_list)
         images = encode_images(images)
 
+        parse_ppt_result = responseDocument(
+            text=full_text,
+            metadata=out_meta
+        )
+        encode_images(images,parse_ppt_result)
+        
         if input_data != input_path:
             os.remove(input_path)
 
-        return {"message": "Document parsed successfully", "markdown": full_text, "metadata": out_meta, "images": images}
+        return parse_ppt_result
 
     except Exception as e:
         raise RuntimeError(f"Error parsing PPT: {str(e)}")
 
-def parse_doc(input_data ,model_state) -> dict:
+def parse_doc(input_data ,model_state) -> responseDocument:
     try:
         if isinstance(input_data, bytes):
             with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
@@ -90,10 +102,16 @@ def parse_doc(input_data ,model_state) -> dict:
         full_text, images, out_meta = parse_single_pdf(input_path, model_state.model_list)
         images = encode_images(images)
 
+        parse_doc_result = responseDocument(
+            text=full_text,
+            metadata=out_meta
+        )
+        encode_images(images,parse_doc_result)
+        
         if input_data != input_path:
             os.remove(input_path)
 
-        return {"message": "Document parsed successfully", "markdown": full_text, "metadata": out_meta, "images": images}
+        return parse_doc_result
 
     except Exception as e:
         raise RuntimeError(f"Error parsing PPT: {str(e)}")

omniparse/documents/router.py

@@ -8,6 +8,7 @@
 # from omniparse.documents import parse_pdf , parse_ppt , parse_doc
 from omniparse.documents import parse_pdf
 from omniparse.utils import encode_images
+from omniparse.models import responseDocument
 
 document_router = APIRouter()
 model_state = get_shared_state()
@@ -17,9 +18,9 @@
 async def parse_pdf_endpoint(file: UploadFile = File(...)):
     try:
         file_bytes = await file.read()
-        result = parse_pdf(file_bytes , model_state)
+        result : responseDocument = parse_pdf(file_bytes , model_state)
 
-        return JSONResponse(content=result)
+        return JSONResponse(content=result.model_dump())
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -43,13 +44,18 @@ async def parse_ppt_endpoint(file: UploadFile = File(...)):
         pdf_bytes = pdf_file.read()
 
     full_text, images, out_meta = parse_single_pdf(pdf_bytes, model_state.model_list)
-    images = encode_images(images)
 
     os.remove(input_path)
     os.remove(output_pdf_path)
     os.rmdir(output_dir)
 
-    return JSONResponse(content={"message": "PPT parsed successfully", "markdown": full_text, "metadata": out_meta, "images": images})
+    result = responseDocument(
+        text=full_text,
+        metadata=out_meta
+    )
+    encode_images(images,result)
+    
+    return JSONResponse(content=result.model_dump())
 
 @document_router.post("/docs")
 async def parse_doc_endpoint(file: UploadFile = File(...)):
@@ -68,13 +74,14 @@ async def parse_doc_endpoint(file: UploadFile = File(...)):
         pdf_bytes = pdf_file.read()
 
     full_text, images, out_meta = parse_single_pdf(pdf_bytes, model_state.model_list)
-    images = encode_images(images)
 
-    os.remove(input_path)
-    os.remove(output_pdf_path)
-    os.rmdir(output_dir)
+    result = responseDocument(
+        text=full_text,
+        metadata=out_meta
+    )
+    encode_images(images,result)
 
-    return JSONResponse(content={"message": "PPT parsed successfully", "markdown": full_text, "metadata": out_meta, "images": images})
+    return JSONResponse(content=result.model_dump())
 
 @document_router.post("")
 async def parse_any_endpoint(file: UploadFile = File(...)):
@@ -98,11 +105,16 @@ async def parse_any_endpoint(file: UploadFile = File(...)):
 
     # Common parsing logic
     full_text, images, out_meta = parse_single_pdf(input_path, model_state.model_list)
-    images = encode_images(images)
 
     os.remove(input_path)
 
-    return JSONResponse(content={"message": "Document parsed successfully", "markdown": full_text, "metadata": out_meta, "images": images})
+    result = responseDocument(
+        text=full_text,
+        metadata=out_meta
+    )
+    encode_images(images,result)
+    
+    return JSONResponse(content=result.model_dump())
 
 
 # @document_router.post("/docs")