From fcf7cb8c5cb0d4ce8d9eb3c244b5eff2d67a44c6 Mon Sep 17 00:00:00 2001 From: Shreya Shankar Date: Fri, 3 Jan 2025 14:55:13 -0800 Subject: [PATCH] feat: add ability for user to specify their docling server --- server/app/routes/convert.py | 60 ++++++++++++++++++- website/src/app/api/convertDocuments/route.ts | 44 +++++++------- website/src/components/FileExplorer.tsx | 58 ++++++++++++++++-- 3 files changed, 136 insertions(+), 26 deletions(-) diff --git a/server/app/routes/convert.py b/server/app/routes/convert.py index 8472dc29..9ad18aac 100644 --- a/server/app/routes/convert.py +++ b/server/app/routes/convert.py @@ -10,6 +10,7 @@ import asyncio from concurrent.futures import ThreadPoolExecutor from dotenv import load_dotenv +import base64 from docling.datamodel.base_models import InputFormat from docling.document_converter import DocumentConverter, PdfFormatOption @@ -48,8 +49,63 @@ def process_document_with_azure(file_path: str, endpoint: str, key: str) -> str: return f"Error processing document: {str(e)}" @router.post("/api/convert-documents") -async def convert_documents(files: List[UploadFile] = File(...), use_docetl_server: str = "false"): +async def convert_documents( + files: List[UploadFile] = File(...), + use_docetl_server: str = "false", + custom_docling_url: Optional[str] = Header(None) +): use_docetl_server = use_docetl_server.lower() == "true" # TODO: make this a boolean + + # If custom Docling URL is provided, forward the request there + if custom_docling_url: + try: + async with aiohttp.ClientSession() as session: + results = [] + for file in files: + # Read file content and encode as base64 + content = await file.read() + base64_content = base64.b64encode(content).decode('utf-8') + + # Prepare request payload according to Docling server spec + payload = { + "file_source": { + "base64_string": base64_content, + "filename": file.filename + }, + "options": { + "output_docling_document": False, + "output_markdown": True, + "output_html": False, + "do_ocr": True, + "do_table_structure": True, + "include_images": True + } + } + + async with session.post( + f"{custom_docling_url}/convert", + json=payload, + timeout=120 + ) as response: + if response.status == 200: + result = await response.json() + if result["status"] == "success": + results.append({ + "filename": file.filename, + "markdown": result["document"]["markdown"] + }) + else: + return {"error": f"Docling server failed to convert {file.filename}: {result.get('errors', [])}"} + else: + error_msg = await response.text() + return {"error": f"Custom Docling server returned error for {file.filename}: {error_msg}"} + + return {"documents": results} + + except Exception as e: + print(f"Custom Docling server failed: {str(e)}. Falling back to local processing...") + return {"error": f"Failed to connect to custom Docling server: {str(e)}"} + # Only try Modal endpoint if use_docetl_server is true and there are no txt files all_txt_files = all(file.filename.lower().endswith('.txt') or file.filename.lower().endswith('.md') for file in files) if use_docetl_server and not all_txt_files: @@ -58,6 +114,8 @@ async def convert_documents(files: List[UploadFile] = File(...), use_docetl_serv # Prepare files for multipart upload data = aiohttp.FormData() for file in files: + # Reset file position since we might have read it in the custom Docling attempt + await file.seek(0) data.add_field('files', await file.read(), filename=file.filename, diff --git a/website/src/app/api/convertDocuments/route.ts b/website/src/app/api/convertDocuments/route.ts index bce02769..954c0d01 100644 --- a/website/src/app/api/convertDocuments/route.ts +++ b/website/src/app/api/convertDocuments/route.ts @@ -19,12 +19,7 @@ export async function POST(request: NextRequest) { // 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"; + const customDoclingUrl = request.headers.get("custom-docling-url"); // Prepare headers for the backend request const headers: HeadersInit = {}; @@ -32,6 +27,9 @@ export async function POST(request: NextRequest) { headers["azure-endpoint"] = azureEndpoint; headers["azure-key"] = azureKey; } + if (customDoclingUrl) { + headers["custom-docling-url"] = customDoclingUrl; + } // Create FormData since FastAPI expects multipart/form-data const backendFormData = new FormData(); @@ -39,17 +37,24 @@ export async function POST(request: NextRequest) { backendFormData.append("files", file); } - // Forward the request to the Python backend - const response = await fetch( - `${FASTAPI_URL}${endpoint}?use_docetl_server=${ - conversionMethod === "docetl" ? "true" : "false" - }`, - { - method: "POST", - body: backendFormData, - headers, - } - ); + // Determine which endpoint to use and construct the URL + let targetUrl: string; + if (azureEndpoint && azureKey) { + targetUrl = `${FASTAPI_URL}/api/azure-convert-documents`; + } else if (customDoclingUrl) { + targetUrl = `${customDoclingUrl}/convert`; + } else { + targetUrl = `${FASTAPI_URL}/api/convert-documents${ + conversionMethod === "docetl" ? "?use_docetl_server=true" : "" + }`; + } + + // Forward the request to the appropriate backend + const response = await fetch(targetUrl, { + method: "POST", + body: backendFormData, + headers, + }); if (!response.ok) { const errorData = await response.json().catch(() => ({})); @@ -66,10 +71,7 @@ export async function POST(request: NextRequest) { console.error("Error converting documents:", error); return NextResponse.json( { - error: - error instanceof Error - ? error.message - : "Failed to convert documents", + error: error instanceof Error ? error.message : String(error), }, { status: 500 } ); diff --git a/website/src/components/FileExplorer.tsx b/website/src/components/FileExplorer.tsx index 6026a910..8b91dfbe 100644 --- a/website/src/components/FileExplorer.tsx +++ b/website/src/components/FileExplorer.tsx @@ -147,7 +147,7 @@ async function getAllFiles(entry: FileSystemEntry): Promise { return files; } -type ConversionMethod = "local" | "azure" | "docetl"; +type ConversionMethod = "local" | "azure" | "docetl" | "custom-docling"; async function validateJsonDataset(file: Blob): Promise { const text = await file.text(); @@ -216,6 +216,7 @@ export const FileExplorer: React.FC = ({ useState("local"); const [azureEndpoint, setAzureEndpoint] = useState(""); const [azureKey, setAzureKey] = useState(""); + const [customDoclingUrl, setCustomDoclingUrl] = useState(""); const { uploadingFiles, uploadDataset } = useDatasetUpload({ namespace, @@ -358,6 +359,8 @@ export const FileExplorer: React.FC = ({ if (conversionMethod === "azure") { headers["azure-endpoint"] = azureEndpoint; headers["azure-key"] = azureKey; + } else if (conversionMethod === "custom-docling") { + headers["custom-docling-url"] = customDoclingUrl; } // Then proceed with conversion @@ -368,7 +371,8 @@ export const FileExplorer: React.FC = ({ }); if (!response.ok) { - throw new Error("Failed to convert documents"); + const errorData = await response.json(); + throw new Error(errorData.error || "Internal Server Error"); } const result = await response.json(); @@ -430,7 +434,7 @@ export const FileExplorer: React.FC = ({ toast({ variant: "destructive", title: "Error", - description: "Failed to process files. Please try again.", + description: error instanceof Error ? error.message : String(error), }); } finally { setIsConverting(false); @@ -778,6 +782,33 @@ export const FileExplorer: React.FC = ({ Enterprise-grade cloud processing

+ +
+
+ + +
+

+ Connect to your own Docling server instance +

+
@@ -826,6 +857,23 @@ export const FileExplorer: React.FC = ({ )} + {conversionMethod === "custom-docling" && ( +
+
+ + setCustomDoclingUrl(e.target.value)} + className="h-8" + /> +
+
+ )} +
= ({ disabled={ isConverting || (conversionMethod === "azure" && - (!azureEndpoint || !azureKey)) + (!azureEndpoint || !azureKey)) || + (conversionMethod === "custom-docling" && + !customDoclingUrl) } className="min-w-[100px]" >