108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
"""Shared client for docling-serve async API."""
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
|
|
class DoclingServeClient:
|
|
"""Client for docling-serve async workflow.
|
|
|
|
Handles the submit → poll → fetch pattern used by both conversion and chunking.
|
|
"""
|
|
|
|
def __init__(self, base_url: str, api_key: str | None = None, timeout: float = 300):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.api_key = api_key
|
|
self.timeout = timeout
|
|
|
|
def _get_headers(self) -> dict[str, str]:
|
|
"""Get headers for API requests."""
|
|
headers: dict[str, str] = {}
|
|
if self.api_key:
|
|
headers["X-Api-Key"] = self.api_key
|
|
return headers
|
|
|
|
async def submit_and_poll(
|
|
self,
|
|
endpoint: str,
|
|
files: dict[str, Any],
|
|
data: dict[str, Any],
|
|
name: str = "document",
|
|
) -> dict[str, Any]:
|
|
"""Submit a task and poll until completion.
|
|
|
|
Args:
|
|
endpoint: The async endpoint path (e.g., "/v1/convert/file/async")
|
|
files: Files to upload
|
|
data: Form data parameters
|
|
name: Name for error messages
|
|
|
|
Returns:
|
|
The result dictionary from the completed task
|
|
|
|
Raises:
|
|
ValueError: If the task fails or service is unavailable
|
|
"""
|
|
headers = self._get_headers()
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
# Submit async task
|
|
submit_url = f"{self.base_url}{endpoint}"
|
|
response = await client.post(
|
|
submit_url,
|
|
files=files,
|
|
data=data,
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
submit_result = response.json()
|
|
task_id = submit_result.get("task_id")
|
|
|
|
if not task_id:
|
|
raise ValueError("docling-serve did not return a task_id")
|
|
|
|
# Poll for completion
|
|
poll_url = f"{self.base_url}/v1/status/poll/{task_id}"
|
|
while True:
|
|
poll_response = await client.get(poll_url, headers=headers)
|
|
poll_response.raise_for_status()
|
|
poll_result = poll_response.json()
|
|
status = poll_result.get("task_status")
|
|
|
|
if status == "success":
|
|
break
|
|
elif status in ("failure", "error"):
|
|
raise ValueError(
|
|
f"docling-serve task failed for {name}: {poll_result}"
|
|
)
|
|
|
|
await asyncio.sleep(1)
|
|
|
|
# Fetch result
|
|
result_url = f"{self.base_url}/v1/result/{task_id}"
|
|
result_response = await client.get(result_url, headers=headers)
|
|
result_response.raise_for_status()
|
|
return result_response.json()
|
|
|
|
except httpx.ConnectError as e:
|
|
raise ValueError(
|
|
f"Could not connect to docling-serve at {self.base_url}. "
|
|
f"Ensure the service is running and accessible. Error: {e}"
|
|
)
|
|
except httpx.TimeoutException as e:
|
|
raise ValueError(
|
|
f"Request to docling-serve timed out after {self.timeout}s. Error: {e}"
|
|
)
|
|
except httpx.HTTPStatusError as e:
|
|
if e.response.status_code == 401:
|
|
raise ValueError(
|
|
"Authentication failed. Check your API key configuration."
|
|
)
|
|
raise ValueError(f"HTTP error from docling-serve: {e}")
|
|
except ValueError:
|
|
raise
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to process via docling-serve: {e}")
|