Skip to content

Commit

Permalink
feat: add ability for user to specify their docling server
Browse files Browse the repository at this point in the history
  • Loading branch information
@shreyashankar
shreyashankar committed Jan 3, 2025
1 parent 242ae05 commit fcf7cb8
Show file tree
Hide file tree
Showing 3 changed files with 136 additions and 26 deletions.
60 changes: 59 additions & 1 deletion server/app/routes/convert.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand Down
44 changes: 23 additions & 21 deletions website/src/app/api/convertDocuments/route.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,37 +19,42 @@ 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 = {};
if (azureEndpoint && azureKey) {
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();
for (const file of files) {
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(() => ({}));
Expand All @@ -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 }
);
Expand Down
58 changes: 54 additions & 4 deletions website/src/components/FileExplorer.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ async function getAllFiles(entry: FileSystemEntry): Promise<FileWithPath[]> {
return files;
}

type ConversionMethod = "local" | "azure" | "docetl";
type ConversionMethod = "local" | "azure" | "docetl" | "custom-docling";

async function validateJsonDataset(file: Blob): Promise<void> {
const text = await file.text();
Expand Down Expand Up @@ -216,6 +216,7 @@ export const FileExplorer: React.FC<FileExplorerProps> = ({
useState<ConversionMethod>("local");
const [azureEndpoint, setAzureEndpoint] = useState("");
const [azureKey, setAzureKey] = useState("");
const [customDoclingUrl, setCustomDoclingUrl] = useState("");

const { uploadingFiles, uploadDataset } = useDatasetUpload({
namespace,
Expand Down Expand Up @@ -358,6 +359,8 @@ export const FileExplorer: React.FC<FileExplorerProps> = ({
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
Expand All @@ -368,7 +371,8 @@ export const FileExplorer: React.FC<FileExplorerProps> = ({
});

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();
Expand Down Expand Up @@ -430,7 +434,7 @@ export const FileExplorer: React.FC<FileExplorerProps> = ({
toast({
variant: "destructive",
title: "Error",
description: "Failed to process files. Please try again.",
description: error instanceof Error ? error.message : String(error),
});
} finally {
setIsConverting(false);
Expand Down Expand Up @@ -778,6 +782,33 @@ export const FileExplorer: React.FC<FileExplorerProps> = ({
Enterprise-grade cloud processing
</p>
</div>

<div className="flex flex-col space-y-1 p-2 rounded-md transition-colors hover:bg-gray-50 cursor-pointer border border-gray-100">
<div className="flex items-start space-x-2.5">
<RadioGroupItem
value="custom-docling"
id="custom-docling"
className="mt-0.5"
/>
<Label
htmlFor="custom-docling"
className="text-sm font-medium cursor-pointer"
>
Custom Docling Server{" "}
<a
href="https://github.com/DS4SD/docling-serve"
target="_blank"
rel="noopener noreferrer"
className="text-xs text-blue-600 hover:underline"
>
(GitHub ↗)
</a>
</Label>
</div>
<p className="text-xs text-muted-foreground pl-6">
Connect to your own Docling server instance
</p>
</div>
</RadioGroup>
</div>

Expand Down Expand Up @@ -826,6 +857,23 @@ export const FileExplorer: React.FC<FileExplorerProps> = ({
</div>
)}

{conversionMethod === "custom-docling" && (
<div className="grid gap-2 animate-in fade-in slide-in-from-top-1">
<div className="space-y-1">
<Label htmlFor="docling-url" className="text-sm">
Docling Server URL
</Label>
<Input
id="docling-url"
placeholder="http://hostname:port"
value={customDoclingUrl}
onChange={(e) => setCustomDoclingUrl(e.target.value)}
className="h-8"
/>
</div>
</div>
)}

<div
className={`
border-2 border-dashed rounded-lg transition-colors relative flex-shrink-0
Expand Down Expand Up @@ -978,7 +1026,9 @@ export const FileExplorer: React.FC<FileExplorerProps> = ({
disabled={
isConverting ||
(conversionMethod === "azure" &&
(!azureEndpoint || !azureKey))
(!azureEndpoint || !azureKey)) ||
(conversionMethod === "custom-docling" &&
!customDoclingUrl)
}
className="min-w-[100px]"
>
Expand Down

0 comments on commit fcf7cb8

Please sign in to comment.