feat: Add azure doc intelligence to PDF upload in the UI

pyproject.toml

@@ -38,10 +38,11 @@ fastapi = { version = "^0.115.0", optional = true }
 uvicorn = { version = "^0.31.0", optional = true }
 websockets = "^13.1"
 docling = { version  = "^2.5.2", optional = true }
+azure-ai-formrecognizer = { version = "^3.3.3", optional = true }
 
 [tool.poetry.extras]
 parsing = ["python-docx", "openpyxl", "pydub", "python-pptx", "azure-ai-documentintelligence", "paddlepaddle", "pymupdf"]
-server = ["fastapi", "uvicorn", "docling"]
+server = ["fastapi", "uvicorn", "docling", "azure-ai-formrecognizer", "azure-ai-documentintelligence"]
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.2"

server/app/routes/convert.py

@@ -1,14 +1,46 @@
-from fastapi import APIRouter, UploadFile, File
-from typing import List
+from fastapi import APIRouter, UploadFile, File, Header
+from typing import List, Optional
 import tempfile
 import os
 import aiohttp
 from pathlib import Path
+from azure.ai.documentintelligence.models import AnalyzeDocumentRequest, ContentFormat, AnalyzeResult
+from azure.ai.documentintelligence import DocumentIntelligenceClient
+from azure.core.credentials import AzureKeyCredential
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Add Azure credentials
+AZURE_ENDPOINT = os.getenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
+AZURE_KEY = os.getenv("AZURE_DOCUMENT_INTELLIGENCE_KEY")
 
 router = APIRouter()
 
 MODAL_ENDPOINT = "https://ucbepic--docling-converter-convert-documents.modal.run"
 
+def process_document_with_azure(file_path: str, endpoint: str, key: str) -> str:
+    """Process a single document with Azure Document Intelligence"""
+    try:
+        document_analysis_client = DocumentIntelligenceClient(
+            endpoint=endpoint, 
+            credential=AzureKeyCredential(key)
+        )
+
+        with open(file_path, "rb") as f:
+            poller = document_analysis_client.begin_analyze_document(
+                "prebuilt-layout", AnalyzeDocumentRequest(bytes_source=f.read()), output_content_format=ContentFormat.MARKDOWN,
+            )
+        result = poller.result()
+
+        return result.content
+    except Exception as e:
+        print(f"Error processing document: {str(e)}")
+        return f"Error processing document: {str(e)}"
+
 @router.post("/api/convert-documents")
 async def convert_documents(files: List[UploadFile] = File(...)):
     # First try Modal endpoint
@@ -77,4 +109,59 @@ async def convert_documents(files: List[UploadFile] = File(...)):
                     "markdown": conv_result.document.export_to_markdown()
                 })
 
-        return {"documents": results}
+        return {"documents": results}
+
+@router.post("/api/azure-convert-documents")
+async def azure_convert_documents(
+    files: List[UploadFile] = File(...),
+    azure_endpoint: Optional[str] = Header(None),
+    azure_key: Optional[str] = Header(None)
+):
+    if not azure_endpoint or not azure_key:
+        return {"error": "Azure credentials are required"}
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Save uploaded files and prepare for processing
+        file_paths = []
+        original_filenames = []
+
+        for file in files:
+            file_path = os.path.join(temp_dir, file.filename)
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+            with open(file_path, "wb") as buffer:
+                content = await file.read()
+                buffer.write(content)
+
+            file_paths.append(file_path)
+            original_filenames.append(file.filename)
+
+        # Process documents concurrently using ThreadPoolExecutor
+        with ThreadPoolExecutor() as executor:
+            futures = []
+            for file_path in file_paths:
+                future = executor.submit(
+                    process_document_with_azure,
+                    file_path,
+                    azure_endpoint,
+                    azure_key
+                )
+                futures.append(future)
+
+            # Collect results as they complete
+            results = []
+            for future in futures:
+                results.append(future.result())
+
+        # Format results to match the existing endpoint's schema
+        formatted_results = [
+            {
+                "filename": filename,
+                "markdown": content
+            }
+            for filename, content in zip(original_filenames, results)
+        ]
+
+        return {"documents": formatted_results}
+
+

website/src/app/api/convertDocuments/route.ts

@@ -17,17 +17,36 @@ export async function POST(request: NextRequest) {
       backendFormData.append("files", file);
     });
 
+    // Get Azure credentials from headers if they exist
+    const azureEndpoint = request.headers.get("azure-endpoint");
+    const azureKey = request.headers.get("azure-key");
+
+    // Determine which endpoint to use
+    const endpoint =
+      azureEndpoint && azureKey
+        ? "/api/azure-convert-documents"
+        : "/api/convert-documents";
+
+    // Prepare headers for the backend request
+    const headers: HeadersInit = {};
+    if (azureEndpoint && azureKey) {
+      headers["azure-endpoint"] = azureEndpoint;
+      headers["azure-key"] = azureKey;
+    }
+
     // Forward the request to the Python backend
     const response = await fetch(
-      `http://${process.env.NEXT_PUBLIC_BACKEND_HOST}:${process.env.NEXT_PUBLIC_BACKEND_PORT}/api/convert-documents`,
+      `http://${process.env.NEXT_PUBLIC_BACKEND_HOST}:${process.env.NEXT_PUBLIC_BACKEND_PORT}${endpoint}`,
       {
         method: "POST",
         body: backendFormData,
+        headers,
       }
     );
 
     if (!response.ok) {
-      throw new Error(`Backend returned ${response.status}`);
+      const errorData = await response.json().catch(() => ({}));
+      throw new Error(errorData.error || `Backend returned ${response.status}`);
     }
 
     const data = await response.json();
@@ -39,7 +58,12 @@ export async function POST(request: NextRequest) {
   } catch (error) {
     console.error("Error converting documents:", error);
     return NextResponse.json(
-      { error: "Failed to convert documents" },
+      {
+        error:
+          error instanceof Error
+            ? error.message
+            : "Failed to convert documents",
+      },
       { status: 500 }
     );
   }