Update document_generator_v2.py

document_generator_v2.py

@@ -172,6 +172,7 @@ import psycopg2
 from datetime import datetime
 import base64
 from fastapi import Form 
+from llama_parse import LlamaParse
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -448,6 +449,45 @@ class MarkdownConverter:
         markdown += "</div>"
         return markdown
 
+async def load_documents(documents: List[UploadFile]) -> List[str]:
+    """
+    Load and parse documents using LlamaParse.
+
+    Args:
+    documents (List[UploadFile]): List of uploaded document files.
+
+    Returns:
+    List[str]: List of parsed document contents.
+    """
+    parser = LlamaParse(
+        api_key=os.getenv("LLAMA_PARSE_API_KEY"),
+        result_type="markdown",
+        num_workers=4,
+        verbose=True,
+        language="en",
+    )
+
+    # Save uploaded files temporarily
+    temp_files = []
+    for doc in documents:
+        temp_file_path = f"/tmp/{doc.filename}"
+        with open(temp_file_path, "wb") as buffer:
+            content = await doc.read()
+            buffer.write(content)
+        temp_files.append(temp_file_path)
+
+    try:
+        # Use LlamaParse to extract content
+        parsed_documents = await parser.aload_data(temp_files)
+        documents_list = [doc.text for doc in parsed_documents]
+        return documents_list
+    finally:
+        # Clean up temporary files
+        for temp_file in temp_files:
+            os.remove(temp_file)
+
+
+    
 router = APIRouter()
 
 class JsonDocumentResponse(BaseModel):