Extract shared DoclingServeClient for async workflow & use it for VLM support, chunking, converting
This commit is contained in:
parent
3b7fd440cc
commit
f861664656
11 changed files with 30996 additions and 14952 deletions
|
|
@ -2,10 +2,9 @@ import re
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
import httpx
|
|
||||||
|
|
||||||
from haiku.rag.chunkers.base import DocumentChunker
|
from haiku.rag.chunkers.base import DocumentChunker
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, Config
|
||||||
|
from haiku.rag.providers.docling_serve import DoclingServeClient
|
||||||
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
|
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
|
@ -56,11 +55,25 @@ class DoclingServeChunker(DocumentChunker):
|
||||||
|
|
||||||
def __init__(self, config: AppConfig = Config):
|
def __init__(self, config: AppConfig = Config):
|
||||||
self.config = config
|
self.config = config
|
||||||
self.base_url = config.providers.docling_serve.base_url.rstrip("/")
|
self.client = DoclingServeClient(
|
||||||
self.api_key = config.providers.docling_serve.api_key
|
base_url=config.providers.docling_serve.base_url,
|
||||||
self.timeout = config.providers.docling_serve.timeout
|
api_key=config.providers.docling_serve.api_key,
|
||||||
|
)
|
||||||
self.chunker_type = config.processing.chunker_type
|
self.chunker_type = config.processing.chunker_type
|
||||||
|
|
||||||
|
def _build_chunking_data(self) -> dict[str, str]:
|
||||||
|
"""Build form data for chunking request."""
|
||||||
|
return {
|
||||||
|
"chunking_max_tokens": str(self.config.processing.chunk_size),
|
||||||
|
"chunking_tokenizer": self.config.processing.chunking_tokenizer,
|
||||||
|
"chunking_merge_peers": str(
|
||||||
|
self.config.processing.chunking_merge_peers
|
||||||
|
).lower(),
|
||||||
|
"chunking_use_markdown_tables": str(
|
||||||
|
self.config.processing.chunking_use_markdown_tables
|
||||||
|
).lower(),
|
||||||
|
}
|
||||||
|
|
||||||
async def _call_chunk_api(self, document: "DoclingDocument") -> list[dict]:
|
async def _call_chunk_api(self, document: "DoclingDocument") -> list[dict]:
|
||||||
"""Call docling-serve chunking API and return raw chunk data.
|
"""Call docling-serve chunking API and return raw chunk data.
|
||||||
|
|
||||||
|
|
@ -75,9 +88,9 @@ class DoclingServeChunker(DocumentChunker):
|
||||||
"""
|
"""
|
||||||
# Determine endpoint based on chunker_type
|
# Determine endpoint based on chunker_type
|
||||||
if self.chunker_type == "hierarchical":
|
if self.chunker_type == "hierarchical":
|
||||||
url = f"{self.base_url}/v1/chunk/hierarchical/file"
|
endpoint = "/v1/chunk/hierarchical/file/async"
|
||||||
else:
|
else:
|
||||||
url = f"{self.base_url}/v1/chunk/hybrid/file"
|
endpoint = "/v1/chunk/hybrid/file/async"
|
||||||
|
|
||||||
# Export document to JSON
|
# Export document to JSON
|
||||||
doc_json = document.model_dump_json()
|
doc_json = document.model_dump_json()
|
||||||
|
|
@ -85,53 +98,16 @@ class DoclingServeChunker(DocumentChunker):
|
||||||
|
|
||||||
# Prepare multipart request with DoclingDocument JSON
|
# Prepare multipart request with DoclingDocument JSON
|
||||||
files = {"files": ("document.json", BytesIO(doc_bytes), "application/json")}
|
files = {"files": ("document.json", BytesIO(doc_bytes), "application/json")}
|
||||||
|
data = self._build_chunking_data()
|
||||||
|
|
||||||
# Build form data with chunking parameters
|
result = await self.client.submit_and_poll(
|
||||||
data = {
|
endpoint=endpoint,
|
||||||
"chunking_max_tokens": str(self.config.processing.chunk_size),
|
files=files,
|
||||||
"chunking_tokenizer": self.config.processing.chunking_tokenizer,
|
data=data,
|
||||||
"chunking_merge_peers": str(
|
name="document",
|
||||||
self.config.processing.chunking_merge_peers
|
)
|
||||||
).lower(),
|
|
||||||
"chunking_use_markdown_tables": str(
|
|
||||||
self.config.processing.chunking_use_markdown_tables
|
|
||||||
).lower(),
|
|
||||||
}
|
|
||||||
|
|
||||||
headers = {}
|
return result.get("chunks", [])
|
||||||
if self.api_key:
|
|
||||||
headers["X-Api-Key"] = self.api_key
|
|
||||||
|
|
||||||
try:
|
|
||||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
||||||
response = await client.post(
|
|
||||||
url,
|
|
||||||
files=files,
|
|
||||||
data=data,
|
|
||||||
headers=headers,
|
|
||||||
)
|
|
||||||
response.raise_for_status()
|
|
||||||
result = response.json()
|
|
||||||
return result.get("chunks", [])
|
|
||||||
|
|
||||||
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. "
|
|
||||||
f"Consider increasing the timeout in configuration. 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 Exception as e:
|
|
||||||
raise ValueError(f"Failed to chunk via docling-serve: {e}")
|
|
||||||
|
|
||||||
async def chunk(self, document: "DoclingDocument") -> list[Chunk]:
|
async def chunk(self, document: "DoclingDocument") -> list[Chunk]:
|
||||||
"""Split the document into chunks with metadata via docling-serve.
|
"""Split the document into chunks with metadata via docling-serve.
|
||||||
|
|
|
||||||
|
|
@ -164,7 +164,6 @@ class OllamaConfig(BaseModel):
|
||||||
class DoclingServeConfig(BaseModel):
|
class DoclingServeConfig(BaseModel):
|
||||||
base_url: str = "http://localhost:5001"
|
base_url: str = "http://localhost:5001"
|
||||||
api_key: str = ""
|
api_key: str = ""
|
||||||
timeout: int = 300
|
|
||||||
|
|
||||||
|
|
||||||
class ProvidersConfig(BaseModel):
|
class ProvidersConfig(BaseModel):
|
||||||
|
|
|
||||||
|
|
@ -1,18 +1,20 @@
|
||||||
"""docling-serve remote converter implementation."""
|
"""docling-serve remote converter implementation."""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import json
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, ClassVar
|
from typing import TYPE_CHECKING, ClassVar
|
||||||
|
|
||||||
import httpx
|
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig
|
from haiku.rag.config import AppConfig
|
||||||
from haiku.rag.converters.base import DocumentConverter
|
from haiku.rag.converters.base import DocumentConverter
|
||||||
from haiku.rag.converters.text_utils import TextFileHandler
|
from haiku.rag.converters.text_utils import TextFileHandler
|
||||||
|
from haiku.rag.providers.docling_serve import DoclingServeClient
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from docling_core.types.doc.document import DoclingDocument
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
|
||||||
class DoclingServeConverter(DocumentConverter):
|
class DoclingServeConverter(DocumentConverter):
|
||||||
"""Converter that uses docling-serve for document conversion.
|
"""Converter that uses docling-serve for document conversion.
|
||||||
|
|
@ -53,17 +55,70 @@ class DoclingServeConverter(DocumentConverter):
|
||||||
config: Application configuration containing docling-serve settings.
|
config: Application configuration containing docling-serve settings.
|
||||||
"""
|
"""
|
||||||
self.config = config
|
self.config = config
|
||||||
self.base_url = config.providers.docling_serve.base_url.rstrip("/")
|
self.client = DoclingServeClient(
|
||||||
self.api_key = config.providers.docling_serve.api_key
|
base_url=config.providers.docling_serve.base_url,
|
||||||
self.timeout = config.providers.docling_serve.timeout
|
api_key=config.providers.docling_serve.api_key,
|
||||||
|
)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def supported_extensions(self) -> list[str]:
|
def supported_extensions(self) -> list[str]:
|
||||||
"""Return list of file extensions supported by this converter."""
|
"""Return list of file extensions supported by this converter."""
|
||||||
return self.docling_serve_extensions + TextFileHandler.text_extensions
|
return self.docling_serve_extensions + TextFileHandler.text_extensions
|
||||||
|
|
||||||
|
def _get_vlm_api_url(self, model: "ModelConfig") -> str:
|
||||||
|
"""Construct VLM API URL from model config."""
|
||||||
|
if model.base_url:
|
||||||
|
base = model.base_url.rstrip("/")
|
||||||
|
return f"{base}/v1/chat/completions"
|
||||||
|
|
||||||
|
if model.provider == "ollama":
|
||||||
|
base = self.config.providers.ollama.base_url.rstrip("/")
|
||||||
|
return f"{base}/v1/chat/completions"
|
||||||
|
|
||||||
|
if model.provider == "openai":
|
||||||
|
return "https://api.openai.com/v1/chat/completions"
|
||||||
|
|
||||||
|
raise ValueError(f"Unsupported VLM provider: {model.provider}")
|
||||||
|
|
||||||
|
def _build_conversion_data(self) -> dict[str, str | list[str]]:
|
||||||
|
"""Build form data for conversion request."""
|
||||||
|
opts = self.config.processing.conversion_options
|
||||||
|
pic_desc = opts.picture_description
|
||||||
|
|
||||||
|
data: dict[str, str | list[str]] = {
|
||||||
|
"to_formats": "json",
|
||||||
|
"do_ocr": str(opts.do_ocr).lower(),
|
||||||
|
"force_ocr": str(opts.force_ocr).lower(),
|
||||||
|
"do_table_structure": str(opts.do_table_structure).lower(),
|
||||||
|
"table_mode": opts.table_mode,
|
||||||
|
"table_cell_matching": str(opts.table_cell_matching).lower(),
|
||||||
|
"images_scale": str(opts.images_scale),
|
||||||
|
"generate_picture_images": str(
|
||||||
|
opts.generate_picture_images or pic_desc.enabled
|
||||||
|
).lower(),
|
||||||
|
"do_picture_description": str(pic_desc.enabled).lower(),
|
||||||
|
}
|
||||||
|
|
||||||
|
if opts.ocr_lang:
|
||||||
|
data["ocr_lang"] = opts.ocr_lang
|
||||||
|
|
||||||
|
if pic_desc.enabled:
|
||||||
|
prompt = self.config.prompts.picture_description
|
||||||
|
picture_description_api = {
|
||||||
|
"url": self._get_vlm_api_url(pic_desc.model),
|
||||||
|
"params": {
|
||||||
|
"model": pic_desc.model.name,
|
||||||
|
"max_completion_tokens": pic_desc.max_tokens,
|
||||||
|
},
|
||||||
|
"prompt": prompt,
|
||||||
|
"timeout": pic_desc.timeout,
|
||||||
|
}
|
||||||
|
data["picture_description_api"] = json.dumps(picture_description_api)
|
||||||
|
|
||||||
|
return data
|
||||||
|
|
||||||
async def _make_request(self, files: dict, name: str) -> "DoclingDocument":
|
async def _make_request(self, files: dict, name: str) -> "DoclingDocument":
|
||||||
"""Make a request to docling-serve and return the DoclingDocument.
|
"""Make an async request to docling-serve and poll for results.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
files: Dictionary with files parameter for httpx
|
files: Dictionary with files parameter for httpx
|
||||||
|
|
@ -77,73 +132,27 @@ class DoclingServeConverter(DocumentConverter):
|
||||||
"""
|
"""
|
||||||
from docling_core.types.doc.document import DoclingDocument
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
try:
|
data = self._build_conversion_data()
|
||||||
opts = self.config.processing.conversion_options
|
result = await self.client.submit_and_poll(
|
||||||
|
endpoint="/v1/convert/file/async",
|
||||||
|
files=files,
|
||||||
|
data=data,
|
||||||
|
name=name,
|
||||||
|
)
|
||||||
|
|
||||||
data: dict[str, str | list[str]] = {
|
if result.get("status") not in ("success", "partial_success", None):
|
||||||
"to_formats": "json",
|
errors = result.get("errors", [])
|
||||||
"do_ocr": str(opts.do_ocr).lower(),
|
raise ValueError(f"Conversion failed: {errors}")
|
||||||
"force_ocr": str(opts.force_ocr).lower(),
|
|
||||||
"do_table_structure": str(opts.do_table_structure).lower(),
|
|
||||||
"table_mode": opts.table_mode,
|
|
||||||
"table_cell_matching": str(opts.table_cell_matching).lower(),
|
|
||||||
"images_scale": str(opts.images_scale),
|
|
||||||
"generate_picture_images": str(opts.generate_picture_images).lower(),
|
|
||||||
}
|
|
||||||
|
|
||||||
if opts.ocr_lang:
|
json_content = result.get("document", {}).get("json_content")
|
||||||
data["ocr_lang"] = opts.ocr_lang
|
|
||||||
|
|
||||||
headers = {}
|
if json_content is None:
|
||||||
if self.api_key:
|
|
||||||
headers["X-Api-Key"] = self.api_key
|
|
||||||
|
|
||||||
url = f"{self.base_url}/v1/convert/file"
|
|
||||||
|
|
||||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
||||||
response = await client.post(
|
|
||||||
url,
|
|
||||||
files=files,
|
|
||||||
data=data,
|
|
||||||
headers=headers,
|
|
||||||
)
|
|
||||||
response.raise_for_status()
|
|
||||||
result = response.json()
|
|
||||||
|
|
||||||
if result["status"] not in ("success", "partial_success"):
|
|
||||||
errors = result.get("errors", [])
|
|
||||||
raise ValueError(f"Conversion failed: {errors}")
|
|
||||||
|
|
||||||
json_content = result["document"]["json_content"]
|
|
||||||
|
|
||||||
if json_content is None:
|
|
||||||
raise ValueError(
|
|
||||||
f"docling-serve did not return JSON content for {name}. "
|
|
||||||
"This may indicate an unsupported file format."
|
|
||||||
)
|
|
||||||
|
|
||||||
return DoclingDocument.model_validate(json_content)
|
|
||||||
|
|
||||||
except httpx.ConnectError as e:
|
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Could not connect to docling-serve at {self.base_url}. "
|
f"docling-serve did not return JSON content for {name}. "
|
||||||
f"Ensure the service is running and accessible. Error: {e}"
|
"This may indicate an unsupported file format."
|
||||||
)
|
)
|
||||||
except httpx.TimeoutException as e:
|
|
||||||
raise ValueError(
|
return DoclingDocument.model_validate(json_content)
|
||||||
f"Request to docling-serve timed out after {self.timeout}s. "
|
|
||||||
f"Consider increasing the timeout in configuration. 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 convert via docling-serve: {e}")
|
|
||||||
|
|
||||||
async def convert_file(self, path: Path) -> "DoclingDocument":
|
async def convert_file(self, path: Path) -> "DoclingDocument":
|
||||||
"""Convert a file to DoclingDocument using docling-serve.
|
"""Convert a file to DoclingDocument using docling-serve.
|
||||||
|
|
|
||||||
5
haiku_rag_slim/haiku/rag/providers/__init__.py
Normal file
5
haiku_rag_slim/haiku/rag/providers/__init__.py
Normal file
|
|
@ -0,0 +1,5 @@
|
||||||
|
"""Provider clients for external services."""
|
||||||
|
|
||||||
|
from haiku.rag.providers.docling_serve import DoclingServeClient
|
||||||
|
|
||||||
|
__all__ = ["DoclingServeClient"]
|
||||||
108
haiku_rag_slim/haiku/rag/providers/docling_serve.py
Normal file
108
haiku_rag_slim/haiku/rag/providers/docling_serve.py
Normal file
|
|
@ -0,0 +1,108 @@
|
||||||
|
"""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}")
|
||||||
File diff suppressed because one or more lines are too long
|
|
@ -1,13 +1,14 @@
|
||||||
interactions:
|
interactions:
|
||||||
- request:
|
- request:
|
||||||
body: "--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition: form-data; name=\"to_formats\"\r\n\r\njson\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition:
|
body: "--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition: form-data; name=\"to_formats\"\r\n\r\njson\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition:
|
||||||
form-data; name=\"do_ocr\"\r\n\r\ntrue\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition: form-data; name=\"force_ocr\"\r\n\r\nfalse\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition:
|
form-data; name=\"do_ocr\"\r\n\r\ntrue\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition: form-data; name=\"force_ocr\"\r\n\r\nfalse\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition:
|
||||||
form-data; name=\"do_table_structure\"\r\n\r\ntrue\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition: form-data;
|
form-data; name=\"do_table_structure\"\r\n\r\ntrue\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition: form-data;
|
||||||
name=\"table_mode\"\r\n\r\naccurate\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition: form-data; name=\"table_cell_matching\"\r\n\r\ntrue\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition:
|
name=\"table_mode\"\r\n\r\naccurate\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition: form-data; name=\"table_cell_matching\"\r\n\r\ntrue\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition:
|
||||||
form-data; name=\"images_scale\"\r\n\r\n2.0\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition: form-data;
|
form-data; name=\"images_scale\"\r\n\r\n2.0\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition: form-data;
|
||||||
name=\"generate_picture_images\"\r\n\r\nfalse\r\n--313fb28cff30a333a81f610f9c6bdaf3\r\nContent-Disposition: form-data;
|
name=\"generate_picture_images\"\r\n\r\nfalse\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition: form-data;
|
||||||
name=\"files\"; filename=\"tmpgc9za8mg.md\"\r\nContent-Type: text/markdown\r\n\r\n```python\ndef test():\n return
|
name=\"do_picture_description\"\r\n\r\nfalse\r\n--26e45846cd2f5c82db7b770428d74912\r\nContent-Disposition: form-data;
|
||||||
42\n```\r\n--313fb28cff30a333a81f610f9c6bdaf3--\r\n"
|
name=\"files\"; filename=\"tmp0aa4633w.md\"\r\nContent-Type: text/markdown\r\n\r\n```python\ndef test():\n return
|
||||||
|
42\n```\r\n--26e45846cd2f5c82db7b770428d74912--\r\n"
|
||||||
headers:
|
headers:
|
||||||
accept:
|
accept:
|
||||||
- '*/*'
|
- '*/*'
|
||||||
|
|
@ -16,23 +17,107 @@ interactions:
|
||||||
connection:
|
connection:
|
||||||
- keep-alive
|
- keep-alive
|
||||||
content-length:
|
content-length:
|
||||||
- '1011'
|
- '1119'
|
||||||
content-type:
|
content-type:
|
||||||
- multipart/form-data; boundary=313fb28cff30a333a81f610f9c6bdaf3
|
- multipart/form-data; boundary=26e45846cd2f5c82db7b770428d74912
|
||||||
host:
|
host:
|
||||||
- localhost:5001
|
- localhost:5001
|
||||||
method: POST
|
method: POST
|
||||||
uri: http://localhost:5001/v1/convert/file
|
uri: http://localhost:5001/v1/convert/file/async
|
||||||
response:
|
response:
|
||||||
headers:
|
headers:
|
||||||
content-length:
|
content-length:
|
||||||
- '1113'
|
- '131'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
task_id: 7436648c-490d-4938-b42b-9755823f8d78
|
||||||
|
task_meta: null
|
||||||
|
task_position: 1
|
||||||
|
task_status: pending
|
||||||
|
task_type: convert
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
- request:
|
||||||
|
body: ''
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- '*/*'
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
host:
|
||||||
|
- localhost:5001
|
||||||
|
method: GET
|
||||||
|
uri: http://localhost:5001/v1/status/poll/7436648c-490d-4938-b42b-9755823f8d78
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '134'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
task_id: 7436648c-490d-4938-b42b-9755823f8d78
|
||||||
|
task_meta: null
|
||||||
|
task_position: null
|
||||||
|
task_status: started
|
||||||
|
task_type: convert
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
- request:
|
||||||
|
body: ''
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- '*/*'
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
host:
|
||||||
|
- localhost:5001
|
||||||
|
method: GET
|
||||||
|
uri: http://localhost:5001/v1/status/poll/7436648c-490d-4938-b42b-9755823f8d78
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '134'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
task_id: 7436648c-490d-4938-b42b-9755823f8d78
|
||||||
|
task_meta: null
|
||||||
|
task_position: null
|
||||||
|
task_status: success
|
||||||
|
task_type: convert
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
- request:
|
||||||
|
body: ''
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- '*/*'
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
host:
|
||||||
|
- localhost:5001
|
||||||
|
method: GET
|
||||||
|
uri: http://localhost:5001/v1/result/7436648c-490d-4938-b42b-9755823f8d78
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '1114'
|
||||||
content-type:
|
content-type:
|
||||||
- application/json
|
- application/json
|
||||||
parsed_body:
|
parsed_body:
|
||||||
document:
|
document:
|
||||||
doctags_content: null
|
doctags_content: null
|
||||||
filename: tmpgc9za8mg.md
|
filename: tmp0aa4633w.md
|
||||||
html_content: null
|
html_content: null
|
||||||
json_content:
|
json_content:
|
||||||
body:
|
body:
|
||||||
|
|
@ -55,10 +140,10 @@ interactions:
|
||||||
self_ref: '#/furniture'
|
self_ref: '#/furniture'
|
||||||
groups: []
|
groups: []
|
||||||
key_value_items: []
|
key_value_items: []
|
||||||
name: tmpgc9za8mg
|
name: tmp0aa4633w
|
||||||
origin:
|
origin:
|
||||||
binary_hash: 9008975733065065710
|
binary_hash: 9008975733065065710
|
||||||
filename: tmpgc9za8mg.md
|
filename: tmp0aa4633w.md
|
||||||
mimetype: text/markdown
|
mimetype: text/markdown
|
||||||
uri: null
|
uri: null
|
||||||
pages: {}
|
pages: {}
|
||||||
|
|
@ -91,7 +176,7 @@ interactions:
|
||||||
md_content: null
|
md_content: null
|
||||||
text_content: null
|
text_content: null
|
||||||
errors: []
|
errors: []
|
||||||
processing_time: 0.002258999999980915
|
processing_time: 0.0027105000044684857
|
||||||
status: success
|
status: success
|
||||||
timings: {}
|
timings: {}
|
||||||
status:
|
status:
|
||||||
|
|
|
||||||
|
|
@ -1,12 +1,13 @@
|
||||||
interactions:
|
interactions:
|
||||||
- request:
|
- request:
|
||||||
body: "--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition: form-data; name=\"to_formats\"\r\n\r\njson\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition:
|
body: "--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition: form-data; name=\"to_formats\"\r\n\r\njson\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition:
|
||||||
form-data; name=\"do_ocr\"\r\n\r\ntrue\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition: form-data; name=\"force_ocr\"\r\n\r\nfalse\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition:
|
form-data; name=\"do_ocr\"\r\n\r\ntrue\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition: form-data; name=\"force_ocr\"\r\n\r\nfalse\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition:
|
||||||
form-data; name=\"do_table_structure\"\r\n\r\ntrue\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition: form-data;
|
form-data; name=\"do_table_structure\"\r\n\r\ntrue\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition: form-data;
|
||||||
name=\"table_mode\"\r\n\r\naccurate\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition: form-data; name=\"table_cell_matching\"\r\n\r\ntrue\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition:
|
name=\"table_mode\"\r\n\r\naccurate\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition: form-data; name=\"table_cell_matching\"\r\n\r\ntrue\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition:
|
||||||
form-data; name=\"images_scale\"\r\n\r\n2.0\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition: form-data;
|
form-data; name=\"images_scale\"\r\n\r\n2.0\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition: form-data;
|
||||||
name=\"generate_picture_images\"\r\n\r\nfalse\r\n--a906caa5bbe5d30e17b44fb63c456240\r\nContent-Disposition: form-data;
|
name=\"generate_picture_images\"\r\n\r\nfalse\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition: form-data;
|
||||||
name=\"files\"; filename=\"content.md\"\r\nContent-Type: text/markdown\r\n\r\n# Test Document\n\nThis is a test.\r\n--a906caa5bbe5d30e17b44fb63c456240--\r\n"
|
name=\"do_picture_description\"\r\n\r\nfalse\r\n--9df242c136eddc0594e7eb126b1f39cd\r\nContent-Disposition: form-data;
|
||||||
|
name=\"files\"; filename=\"content.md\"\r\nContent-Type: text/markdown\r\n\r\n# Test Document\n\nThis is a test.\r\n--9df242c136eddc0594e7eb126b1f39cd--\r\n"
|
||||||
headers:
|
headers:
|
||||||
accept:
|
accept:
|
||||||
- '*/*'
|
- '*/*'
|
||||||
|
|
@ -15,13 +16,97 @@ interactions:
|
||||||
connection:
|
connection:
|
||||||
- keep-alive
|
- keep-alive
|
||||||
content-length:
|
content-length:
|
||||||
- '1000'
|
- '1108'
|
||||||
content-type:
|
content-type:
|
||||||
- multipart/form-data; boundary=a906caa5bbe5d30e17b44fb63c456240
|
- multipart/form-data; boundary=9df242c136eddc0594e7eb126b1f39cd
|
||||||
host:
|
host:
|
||||||
- localhost:5001
|
- localhost:5001
|
||||||
method: POST
|
method: POST
|
||||||
uri: http://localhost:5001/v1/convert/file
|
uri: http://localhost:5001/v1/convert/file/async
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '131'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
task_id: 9f7f1455-1ab6-4830-b5a5-3e63419379ae
|
||||||
|
task_meta: null
|
||||||
|
task_position: 1
|
||||||
|
task_status: pending
|
||||||
|
task_type: convert
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
- request:
|
||||||
|
body: ''
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- '*/*'
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
host:
|
||||||
|
- localhost:5001
|
||||||
|
method: GET
|
||||||
|
uri: http://localhost:5001/v1/status/poll/9f7f1455-1ab6-4830-b5a5-3e63419379ae
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '134'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
task_id: 9f7f1455-1ab6-4830-b5a5-3e63419379ae
|
||||||
|
task_meta: null
|
||||||
|
task_position: null
|
||||||
|
task_status: started
|
||||||
|
task_type: convert
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
- request:
|
||||||
|
body: ''
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- '*/*'
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
host:
|
||||||
|
- localhost:5001
|
||||||
|
method: GET
|
||||||
|
uri: http://localhost:5001/v1/status/poll/9f7f1455-1ab6-4830-b5a5-3e63419379ae
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
content-length:
|
||||||
|
- '134'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
parsed_body:
|
||||||
|
task_id: 9f7f1455-1ab6-4830-b5a5-3e63419379ae
|
||||||
|
task_meta: null
|
||||||
|
task_position: null
|
||||||
|
task_status: success
|
||||||
|
task_type: convert
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
- request:
|
||||||
|
body: ''
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- '*/*'
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
host:
|
||||||
|
- localhost:5001
|
||||||
|
method: GET
|
||||||
|
uri: http://localhost:5001/v1/result/9f7f1455-1ab6-4830-b5a5-3e63419379ae
|
||||||
response:
|
response:
|
||||||
headers:
|
headers:
|
||||||
content-length:
|
content-length:
|
||||||
|
|
@ -94,7 +179,7 @@ interactions:
|
||||||
md_content: null
|
md_content: null
|
||||||
text_content: null
|
text_content: null
|
||||||
errors: []
|
errors: []
|
||||||
processing_time: 0.004331874999934371
|
processing_time: 0.004173958994215354
|
||||||
status: success
|
status: success
|
||||||
timings: {}
|
timings: {}
|
||||||
status:
|
status:
|
||||||
|
|
|
||||||
File diff suppressed because it is too large
Load diff
|
|
@ -232,6 +232,28 @@ def test_get_chunker_docling_serve():
|
||||||
assert isinstance(chunker, DoclingServeChunker)
|
assert isinstance(chunker, DoclingServeChunker)
|
||||||
|
|
||||||
|
|
||||||
|
def create_async_workflow_mocks(
|
||||||
|
result_data: dict, task_id: str = "test-task-123"
|
||||||
|
) -> tuple[Mock, Mock, Mock]:
|
||||||
|
"""Create mock responses for docling-serve async workflow."""
|
||||||
|
submit_response = Mock()
|
||||||
|
submit_response.status_code = 200
|
||||||
|
submit_response.json.return_value = {"task_id": task_id, "task_status": "pending"}
|
||||||
|
submit_response.raise_for_status = Mock()
|
||||||
|
|
||||||
|
poll_response = Mock()
|
||||||
|
poll_response.status_code = 200
|
||||||
|
poll_response.json.return_value = {"task_id": task_id, "task_status": "success"}
|
||||||
|
poll_response.raise_for_status = Mock()
|
||||||
|
|
||||||
|
result_response = Mock()
|
||||||
|
result_response.status_code = 200
|
||||||
|
result_response.json.return_value = result_data
|
||||||
|
result_response.raise_for_status = Mock()
|
||||||
|
|
||||||
|
return submit_response, poll_response, result_response
|
||||||
|
|
||||||
|
|
||||||
class TestDoclingServeChunker:
|
class TestDoclingServeChunker:
|
||||||
"""Tests for DoclingServeChunker (mocked)."""
|
"""Tests for DoclingServeChunker (mocked)."""
|
||||||
|
|
||||||
|
|
@ -241,7 +263,6 @@ class TestDoclingServeChunker:
|
||||||
config = AppConfig()
|
config = AppConfig()
|
||||||
config.providers.docling_serve.base_url = "http://localhost:5001"
|
config.providers.docling_serve.base_url = "http://localhost:5001"
|
||||||
config.providers.docling_serve.api_key = ""
|
config.providers.docling_serve.api_key = ""
|
||||||
config.providers.docling_serve.timeout = 300
|
|
||||||
config.processing.chunk_size = 256
|
config.processing.chunk_size = 256
|
||||||
config.processing.chunking_tokenizer = "Qwen/Qwen3-Embedding-0.6B"
|
config.processing.chunking_tokenizer = "Qwen/Qwen3-Embedding-0.6B"
|
||||||
return config
|
return config
|
||||||
|
|
@ -252,21 +273,20 @@ class TestDoclingServeChunker:
|
||||||
return DoclingServeChunker(config)
|
return DoclingServeChunker(config)
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_success(self, mock_client_class, chunker):
|
async def test_chunk_success(self, mock_client_class, chunker):
|
||||||
"""Test successful chunking via docling-serve."""
|
"""Test successful chunking via docling-serve async workflow."""
|
||||||
mock_response = Mock()
|
result_data = {
|
||||||
mock_response.status_code = 200
|
|
||||||
mock_response.json.return_value = {
|
|
||||||
"chunks": [
|
"chunks": [
|
||||||
{"text": "Chunk 1", "chunk_index": 0},
|
{"text": "Chunk 1", "chunk_index": 0},
|
||||||
{"text": "Chunk 2", "chunk_index": 1},
|
{"text": "Chunk 2", "chunk_index": 1},
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
mock_response.raise_for_status = Mock()
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(result_data)
|
||||||
|
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post.return_value = mock_response
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
# Create a simple document
|
# Create a simple document
|
||||||
|
|
@ -280,21 +300,18 @@ class TestDoclingServeChunker:
|
||||||
mock_client.post.assert_called_once()
|
mock_client.post.assert_called_once()
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_with_api_key(self, mock_client_class, config):
|
async def test_chunk_with_api_key(self, mock_client_class, config):
|
||||||
"""Test that API key is included in request headers."""
|
"""Test that API key is included in request headers."""
|
||||||
config.providers.docling_serve.api_key = "test-key"
|
config.providers.docling_serve.api_key = "test-key"
|
||||||
chunker = DoclingServeChunker(config)
|
chunker = DoclingServeChunker(config)
|
||||||
|
|
||||||
mock_response = Mock()
|
result_data = {"chunks": [{"text": "Chunk 1", "chunk_index": 0}]}
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(result_data)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"chunks": [{"text": "Chunk 1", "chunk_index": 0}]
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post.return_value = mock_response
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(Config)
|
||||||
|
|
@ -306,21 +323,18 @@ class TestDoclingServeChunker:
|
||||||
assert call_kwargs["headers"]["X-Api-Key"] == "test-key"
|
assert call_kwargs["headers"]["X-Api-Key"] == "test-key"
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_hierarchical_endpoint(self, mock_client_class, config):
|
async def test_chunk_hierarchical_endpoint(self, mock_client_class, config):
|
||||||
"""Test that hierarchical chunker uses correct endpoint."""
|
"""Test that hierarchical chunker uses correct endpoint."""
|
||||||
config.processing.chunker_type = "hierarchical"
|
config.processing.chunker_type = "hierarchical"
|
||||||
chunker = DoclingServeChunker(config)
|
chunker = DoclingServeChunker(config)
|
||||||
|
|
||||||
mock_response = Mock()
|
result_data = {"chunks": [{"text": "Chunk 1", "chunk_index": 0}]}
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(result_data)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"chunks": [{"text": "Chunk 1", "chunk_index": 0}]
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post.return_value = mock_response
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(Config)
|
||||||
|
|
@ -328,10 +342,10 @@ class TestDoclingServeChunker:
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
call_args = mock_client.post.call_args
|
call_args = mock_client.post.call_args
|
||||||
assert "/v1/chunk/hierarchical/file" in call_args[0][0]
|
assert "/v1/chunk/hierarchical/file/async" in call_args[0][0]
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_passes_config_parameters(self, mock_client_class, config):
|
async def test_chunk_passes_config_parameters(self, mock_client_class, config):
|
||||||
"""Test that all config parameters are passed to API."""
|
"""Test that all config parameters are passed to API."""
|
||||||
config.processing.chunk_size = 512
|
config.processing.chunk_size = 512
|
||||||
|
|
@ -339,15 +353,12 @@ class TestDoclingServeChunker:
|
||||||
config.processing.chunking_use_markdown_tables = True
|
config.processing.chunking_use_markdown_tables = True
|
||||||
chunker = DoclingServeChunker(config)
|
chunker = DoclingServeChunker(config)
|
||||||
|
|
||||||
mock_response = Mock()
|
result_data = {"chunks": [{"text": "Chunk 1", "chunk_index": 0}]}
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(result_data)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"chunks": [{"text": "Chunk 1", "chunk_index": 0}]
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post.return_value = mock_response
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
converter = get_converter(Config)
|
converter = get_converter(Config)
|
||||||
|
|
@ -361,7 +372,7 @@ class TestDoclingServeChunker:
|
||||||
assert data["chunking_use_markdown_tables"] == "true"
|
assert data["chunking_use_markdown_tables"] == "true"
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_connection_error(self, mock_client_class, chunker):
|
async def test_chunk_connection_error(self, mock_client_class, chunker):
|
||||||
"""Test handling of connection errors."""
|
"""Test handling of connection errors."""
|
||||||
import httpx
|
import httpx
|
||||||
|
|
@ -377,7 +388,7 @@ class TestDoclingServeChunker:
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_timeout_error(self, mock_client_class, chunker):
|
async def test_chunk_timeout_error(self, mock_client_class, chunker):
|
||||||
"""Test handling of timeout errors."""
|
"""Test handling of timeout errors."""
|
||||||
import httpx
|
import httpx
|
||||||
|
|
@ -393,7 +404,7 @@ class TestDoclingServeChunker:
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_auth_error(self, mock_client_class, chunker):
|
async def test_chunk_auth_error(self, mock_client_class, chunker):
|
||||||
"""Test handling of authentication errors."""
|
"""Test handling of authentication errors."""
|
||||||
import httpx
|
import httpx
|
||||||
|
|
@ -416,18 +427,14 @@ class TestDoclingServeChunker:
|
||||||
await chunker.chunk(doc)
|
await chunker.chunk(doc)
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
|
@patch("haiku.rag.providers.docling_serve.httpx.AsyncClient")
|
||||||
async def test_chunk_metadata_extraction(self, mock_client_class, chunker):
|
async def test_chunk_metadata_extraction(self, mock_client_class, chunker):
|
||||||
"""Test that metadata is correctly extracted from API response.
|
"""Test that metadata is correctly extracted from API response.
|
||||||
|
|
||||||
Labels are resolved from the DoclingDocument using the refs, so we need
|
Labels are resolved from the DoclingDocument using the refs, so we need
|
||||||
to create a document with matching structure for the mocked API response.
|
to create a document with matching structure for the mocked API response.
|
||||||
"""
|
"""
|
||||||
mock_response = Mock()
|
result_data = {
|
||||||
mock_response.status_code = 200
|
|
||||||
# docling-serve returns doc_items as list of ref strings
|
|
||||||
# We'll reference texts[0], texts[1], and tables[0]
|
|
||||||
mock_response.json.return_value = {
|
|
||||||
"chunks": [
|
"chunks": [
|
||||||
{
|
{
|
||||||
"text": "Chapter 1\nThis is content.",
|
"text": "Chapter 1\nThis is content.",
|
||||||
|
|
@ -443,10 +450,11 @@ class TestDoclingServeChunker:
|
||||||
},
|
},
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
mock_response.raise_for_status = Mock()
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(result_data)
|
||||||
|
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post.return_value = mock_response
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client_class.return_value.__aenter__.return_value = mock_client
|
mock_client_class.return_value.__aenter__.return_value = mock_client
|
||||||
|
|
||||||
# Create a document with texts and tables that match the mocked refs
|
# Create a document with texts and tables that match the mocked refs
|
||||||
|
|
|
||||||
|
|
@ -54,6 +54,31 @@ def create_mock_docling_document_json(name: str = "test") -> dict:
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def create_async_workflow_mocks(
|
||||||
|
doc_json: dict, task_id: str = "test-task-123"
|
||||||
|
) -> tuple[Mock, Mock, Mock]:
|
||||||
|
"""Create mock responses for docling-serve async workflow.
|
||||||
|
|
||||||
|
Returns tuple of (submit_response, poll_response, result_response).
|
||||||
|
"""
|
||||||
|
submit_response = Mock()
|
||||||
|
submit_response.status_code = 200
|
||||||
|
submit_response.json.return_value = {"task_id": task_id, "task_status": "pending"}
|
||||||
|
submit_response.raise_for_status = Mock()
|
||||||
|
|
||||||
|
poll_response = Mock()
|
||||||
|
poll_response.status_code = 200
|
||||||
|
poll_response.json.return_value = {"task_id": task_id, "task_status": "success"}
|
||||||
|
poll_response.raise_for_status = Mock()
|
||||||
|
|
||||||
|
result_response = Mock()
|
||||||
|
result_response.status_code = 200
|
||||||
|
result_response.json.return_value = {"document": {"json_content": doc_json}}
|
||||||
|
result_response.raise_for_status = Mock()
|
||||||
|
|
||||||
|
return submit_response, poll_response, result_response
|
||||||
|
|
||||||
|
|
||||||
class TestTextFileHandler:
|
class TestTextFileHandler:
|
||||||
"""Tests for TextFileHandler utility class."""
|
"""Tests for TextFileHandler utility class."""
|
||||||
|
|
||||||
|
|
@ -429,7 +454,6 @@ class TestDoclingServeConverter:
|
||||||
config = AppConfig()
|
config = AppConfig()
|
||||||
config.providers.docling_serve.base_url = "http://localhost:5001"
|
config.providers.docling_serve.base_url = "http://localhost:5001"
|
||||||
config.providers.docling_serve.api_key = ""
|
config.providers.docling_serve.api_key = ""
|
||||||
config.providers.docling_serve.timeout = 300
|
|
||||||
return config
|
return config
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
@ -439,8 +463,7 @@ class TestDoclingServeConverter:
|
||||||
|
|
||||||
def test_initialization(self, converter):
|
def test_initialization(self, converter):
|
||||||
"""Test converter initialization."""
|
"""Test converter initialization."""
|
||||||
assert converter.base_url == "http://localhost:5001"
|
assert converter.client.base_url == "http://localhost:5001"
|
||||||
assert converter.timeout == 300
|
|
||||||
|
|
||||||
def test_supported_extensions(self, converter):
|
def test_supported_extensions(self, converter):
|
||||||
"""Test that converter reports correct supported extensions."""
|
"""Test that converter reports correct supported extensions."""
|
||||||
|
|
@ -452,18 +475,14 @@ class TestDoclingServeConverter:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_convert_text_success(self, converter):
|
async def test_convert_text_success(self, converter):
|
||||||
"""Test successful text conversion via docling-serve."""
|
"""Test successful text conversion via docling-serve async workflow."""
|
||||||
mock_response = Mock()
|
doc_json = create_mock_docling_document_json("test")
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"status": "success",
|
|
||||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient") as mock_client_class:
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post = AsyncMock(return_value=mock_response)
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
mock_client_class.return_value = mock_client
|
mock_client_class.return_value = mock_client
|
||||||
|
|
@ -479,17 +498,13 @@ class TestDoclingServeConverter:
|
||||||
config.providers.docling_serve.api_key = "test-key"
|
config.providers.docling_serve.api_key = "test-key"
|
||||||
converter = DoclingServeConverter(config)
|
converter = DoclingServeConverter(config)
|
||||||
|
|
||||||
mock_response = Mock()
|
doc_json = create_mock_docling_document_json("test")
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"status": "success",
|
|
||||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient") as mock_client_class:
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post = AsyncMock(return_value=mock_response)
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
mock_client_class.return_value = mock_client
|
mock_client_class.return_value = mock_client
|
||||||
|
|
@ -512,17 +527,13 @@ class TestDoclingServeConverter:
|
||||||
config.processing.conversion_options.images_scale = 3.0
|
config.processing.conversion_options.images_scale = 3.0
|
||||||
converter = DoclingServeConverter(config)
|
converter = DoclingServeConverter(config)
|
||||||
|
|
||||||
mock_response = Mock()
|
doc_json = create_mock_docling_document_json("test")
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"status": "success",
|
|
||||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient") as mock_client_class:
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post = AsyncMock(return_value=mock_response)
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
mock_client_class.return_value = mock_client
|
mock_client_class.return_value = mock_client
|
||||||
|
|
@ -591,17 +602,16 @@ class TestDoclingServeConverter:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_convert_text_no_json_content(self, converter):
|
async def test_convert_text_no_json_content(self, converter):
|
||||||
"""Test handling when docling-serve returns no JSON content."""
|
"""Test handling when docling-serve returns no JSON content."""
|
||||||
mock_response = Mock()
|
submit_resp, poll_resp, _ = create_async_workflow_mocks({})
|
||||||
mock_response.status_code = 200
|
result_resp = Mock()
|
||||||
mock_response.json.return_value = {
|
result_resp.status_code = 200
|
||||||
"status": "success",
|
result_resp.json.return_value = {"document": {"json_content": None}}
|
||||||
"document": {"json_content": None},
|
result_resp.raise_for_status = Mock()
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient") as mock_client_class:
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post = AsyncMock(return_value=mock_response)
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
mock_client_class.return_value = mock_client
|
mock_client_class.return_value = mock_client
|
||||||
|
|
@ -611,18 +621,14 @@ class TestDoclingServeConverter:
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_convert_file_pdf(self, converter):
|
async def test_convert_file_pdf(self, converter):
|
||||||
"""Test converting PDF file via docling-serve."""
|
"""Test converting PDF file via docling-serve async workflow."""
|
||||||
mock_response = Mock()
|
doc_json = create_mock_docling_document_json("test")
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"status": "success",
|
|
||||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient") as mock_client_class:
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post = AsyncMock(return_value=mock_response)
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
mock_client_class.return_value = mock_client
|
mock_client_class.return_value = mock_client
|
||||||
|
|
@ -639,17 +645,13 @@ class TestDoclingServeConverter:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_convert_file_text(self, converter):
|
async def test_convert_file_text(self, converter):
|
||||||
"""Test converting text file (reads locally, sends to docling-serve)."""
|
"""Test converting text file (reads locally, sends to docling-serve)."""
|
||||||
mock_response = Mock()
|
doc_json = create_mock_docling_document_json("test")
|
||||||
mock_response.status_code = 200
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
|
||||||
mock_response.json.return_value = {
|
|
||||||
"status": "success",
|
|
||||||
"document": {"json_content": create_mock_docling_document_json("test")},
|
|
||||||
}
|
|
||||||
mock_response.raise_for_status = Mock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient") as mock_client_class:
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
mock_client = AsyncMock()
|
mock_client = AsyncMock()
|
||||||
mock_client.post = AsyncMock(return_value=mock_response)
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
mock_client_class.return_value = mock_client
|
mock_client_class.return_value = mock_client
|
||||||
|
|
@ -666,6 +668,124 @@ class TestDoclingServeConverter:
|
||||||
assert "files" in call_kwargs
|
assert "files" in call_kwargs
|
||||||
|
|
||||||
|
|
||||||
|
class TestDoclingServeConverterPictureDescription:
|
||||||
|
"""Tests for DoclingServeConverter picture description support."""
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def config(self):
|
||||||
|
"""Create test configuration."""
|
||||||
|
config = AppConfig()
|
||||||
|
config.providers.docling_serve.base_url = "http://localhost:5001"
|
||||||
|
config.providers.docling_serve.api_key = ""
|
||||||
|
return config
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_with_ollama(self, config):
|
||||||
|
"""Test VLM API URL construction for Ollama provider."""
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(provider="ollama", name="ministral-3")
|
||||||
|
url = converter._get_vlm_api_url(model)
|
||||||
|
assert url == "http://localhost:11434/v1/chat/completions"
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_with_custom_base_url(self, config):
|
||||||
|
"""Test VLM API URL construction with custom base_url."""
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(
|
||||||
|
provider="openai", name="gpt-4-vision", base_url="http://my-vllm:8000"
|
||||||
|
)
|
||||||
|
url = converter._get_vlm_api_url(model)
|
||||||
|
assert url == "http://my-vllm:8000/v1/chat/completions"
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_with_openai(self, config):
|
||||||
|
"""Test VLM API URL construction for OpenAI provider."""
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(provider="openai", name="gpt-4-vision")
|
||||||
|
url = converter._get_vlm_api_url(model)
|
||||||
|
assert url == "https://api.openai.com/v1/chat/completions"
|
||||||
|
|
||||||
|
def test_get_vlm_api_url_unsupported_provider(self, config):
|
||||||
|
"""Test VLM API URL construction raises error for unsupported provider."""
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
from haiku.rag.config.models import ModelConfig
|
||||||
|
|
||||||
|
model = ModelConfig(provider="unsupported", name="test")
|
||||||
|
with pytest.raises(ValueError, match="Unsupported VLM provider"):
|
||||||
|
converter._get_vlm_api_url(model)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_picture_description_options_passed_to_api(self, config):
|
||||||
|
"""Test that picture description options are passed to docling-serve API."""
|
||||||
|
import json
|
||||||
|
|
||||||
|
config.processing.conversion_options.picture_description.enabled = True
|
||||||
|
config.processing.conversion_options.picture_description.model.provider = (
|
||||||
|
"ollama"
|
||||||
|
)
|
||||||
|
config.processing.conversion_options.picture_description.model.name = (
|
||||||
|
"ministral-3"
|
||||||
|
)
|
||||||
|
config.processing.conversion_options.picture_description.timeout = 120
|
||||||
|
config.processing.conversion_options.picture_description.max_tokens = 300
|
||||||
|
config.prompts.picture_description = "Test prompt for picture description"
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
|
||||||
|
doc_json = create_mock_docling_document_json("test")
|
||||||
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
|
||||||
|
|
||||||
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
|
mock_client = AsyncMock()
|
||||||
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
|
mock_client_class.return_value = mock_client
|
||||||
|
|
||||||
|
await converter.convert_text("# Test")
|
||||||
|
|
||||||
|
call_kwargs = mock_client.post.call_args.kwargs
|
||||||
|
assert "data" in call_kwargs
|
||||||
|
data = call_kwargs["data"]
|
||||||
|
|
||||||
|
assert data["do_picture_description"] == "true"
|
||||||
|
assert data["generate_picture_images"] == "true"
|
||||||
|
assert "picture_description_api" in data
|
||||||
|
|
||||||
|
api_config = json.loads(data["picture_description_api"])
|
||||||
|
assert api_config["url"] == "http://localhost:11434/v1/chat/completions"
|
||||||
|
assert api_config["params"]["model"] == "ministral-3"
|
||||||
|
assert api_config["params"]["max_completion_tokens"] == 300
|
||||||
|
assert api_config["prompt"] == "Test prompt for picture description"
|
||||||
|
assert api_config["timeout"] == 120
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_picture_description_disabled_by_default(self, config):
|
||||||
|
"""Test that picture description is disabled by default."""
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
|
||||||
|
doc_json = create_mock_docling_document_json("test")
|
||||||
|
submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
|
||||||
|
|
||||||
|
with patch("httpx.AsyncClient") as mock_client_class:
|
||||||
|
mock_client = AsyncMock()
|
||||||
|
mock_client.post = AsyncMock(return_value=submit_resp)
|
||||||
|
mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
|
||||||
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||||
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||||
|
mock_client_class.return_value = mock_client
|
||||||
|
|
||||||
|
await converter.convert_text("# Test")
|
||||||
|
|
||||||
|
call_kwargs = mock_client.post.call_args.kwargs
|
||||||
|
data = call_kwargs["data"]
|
||||||
|
assert data["do_picture_description"] == "false"
|
||||||
|
assert "picture_description_api" not in data
|
||||||
|
|
||||||
|
|
||||||
class TestDoclingServeConverterIntegration:
|
class TestDoclingServeConverterIntegration:
|
||||||
"""Integration tests with real docling-serve recorded via VCR."""
|
"""Integration tests with real docling-serve recorded via VCR."""
|
||||||
|
|
||||||
|
|
@ -702,3 +822,48 @@ class TestDoclingServeConverterIntegration:
|
||||||
assert isinstance(doc, DoclingDocument)
|
assert isinstance(doc, DoclingDocument)
|
||||||
result = doc.export_to_markdown()
|
result = doc.export_to_markdown()
|
||||||
assert "def test():" in result
|
assert "def test():" in result
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.integration
|
||||||
|
async def test_picture_description_end_to_end(self, config):
|
||||||
|
"""End-to-end test: convert PDF with VLM picture descriptions via docling-serve.
|
||||||
|
|
||||||
|
Note: Not using VCR because this test involves polling with changing task IDs.
|
||||||
|
"""
|
||||||
|
pdf_path = Path("tests/data/doclaynet.pdf")
|
||||||
|
if not pdf_path.exists():
|
||||||
|
pytest.skip("doclaynet.pdf not found")
|
||||||
|
|
||||||
|
config.processing.conversion_options.picture_description.enabled = True
|
||||||
|
config.processing.conversion_options.picture_description.model.provider = (
|
||||||
|
"ollama"
|
||||||
|
)
|
||||||
|
config.processing.conversion_options.picture_description.model.name = (
|
||||||
|
"ministral-3"
|
||||||
|
)
|
||||||
|
# Use host.docker.internal so docling-serve in Docker can reach host's Ollama
|
||||||
|
config.processing.conversion_options.picture_description.model.base_url = (
|
||||||
|
"http://host.docker.internal:11434"
|
||||||
|
)
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
|
||||||
|
doc = await converter.convert_file(pdf_path)
|
||||||
|
|
||||||
|
assert doc.pictures, "Document should have pictures"
|
||||||
|
|
||||||
|
from docling_core.types.doc.document import PictureDescriptionData
|
||||||
|
|
||||||
|
pictures_with_descriptions = []
|
||||||
|
markdown = doc.export_to_markdown()
|
||||||
|
for pic in doc.pictures:
|
||||||
|
for ann in pic.annotations:
|
||||||
|
if isinstance(ann, PictureDescriptionData):
|
||||||
|
pictures_with_descriptions.append(pic)
|
||||||
|
assert ann.text in markdown, (
|
||||||
|
f"Picture description '{ann.text[:50]}...' should be in markdown"
|
||||||
|
)
|
||||||
|
break
|
||||||
|
|
||||||
|
assert pictures_with_descriptions, (
|
||||||
|
"At least one picture should have a VLM description"
|
||||||
|
)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue