Implement docling-serve converter
This commit is contained in:
parent
92def1c6a5
commit
fe7d0c5c24
7 changed files with 666 additions and 82 deletions
|
|
@ -23,9 +23,9 @@ def get_converter(config: AppConfig = Config) -> DocumentConverter:
|
||||||
|
|
||||||
return DoclingLocalConverter()
|
return DoclingLocalConverter()
|
||||||
|
|
||||||
# if config.processing.converter == "docling-serve":
|
if config.processing.converter == "docling-serve":
|
||||||
# from haiku.rag.converters.docling_serve import DoclingServeConverter
|
from haiku.rag.converters.docling_serve import DoclingServeConverter
|
||||||
#
|
|
||||||
# return DoclingServeConverter(config)
|
return DoclingServeConverter(config)
|
||||||
|
|
||||||
raise ValueError(f"Unsupported converter provider: {config.processing.converter}")
|
raise ValueError(f"Unsupported converter provider: {config.processing.converter}")
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,10 @@
|
||||||
"""Local docling converter implementation."""
|
"""Local docling converter implementation."""
|
||||||
|
|
||||||
from io import BytesIO
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, ClassVar
|
from typing import TYPE_CHECKING, ClassVar
|
||||||
|
|
||||||
from haiku.rag.converters.base import DocumentConverter
|
from haiku.rag.converters.base import DocumentConverter
|
||||||
|
from haiku.rag.converters.text_utils import TextFileHandler
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from docling_core.types.doc.document import DoclingDocument
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
@ -39,66 +39,10 @@ class DoclingLocalConverter(DocumentConverter):
|
||||||
".webp",
|
".webp",
|
||||||
]
|
]
|
||||||
|
|
||||||
# Plain text extensions that we'll read directly
|
|
||||||
text_extensions: ClassVar[list[str]] = [
|
|
||||||
".astro",
|
|
||||||
".c",
|
|
||||||
".cpp",
|
|
||||||
".css",
|
|
||||||
".go",
|
|
||||||
".h",
|
|
||||||
".hpp",
|
|
||||||
".java",
|
|
||||||
".js",
|
|
||||||
".json",
|
|
||||||
".kt",
|
|
||||||
".mdx",
|
|
||||||
".mjs",
|
|
||||||
".php",
|
|
||||||
".py",
|
|
||||||
".rb",
|
|
||||||
".rs",
|
|
||||||
".svelte",
|
|
||||||
".swift",
|
|
||||||
".ts",
|
|
||||||
".tsx",
|
|
||||||
".txt",
|
|
||||||
".vue",
|
|
||||||
".yaml",
|
|
||||||
".yml",
|
|
||||||
]
|
|
||||||
|
|
||||||
# Code file extensions with their markdown language identifiers for syntax highlighting
|
|
||||||
code_markdown_identifier: ClassVar[dict[str, str]] = {
|
|
||||||
".astro": "astro",
|
|
||||||
".c": "c",
|
|
||||||
".cpp": "cpp",
|
|
||||||
".css": "css",
|
|
||||||
".go": "go",
|
|
||||||
".h": "c",
|
|
||||||
".hpp": "cpp",
|
|
||||||
".java": "java",
|
|
||||||
".js": "javascript",
|
|
||||||
".json": "json",
|
|
||||||
".kt": "kotlin",
|
|
||||||
".mjs": "javascript",
|
|
||||||
".php": "php",
|
|
||||||
".py": "python",
|
|
||||||
".rb": "ruby",
|
|
||||||
".rs": "rust",
|
|
||||||
".svelte": "svelte",
|
|
||||||
".swift": "swift",
|
|
||||||
".ts": "typescript",
|
|
||||||
".tsx": "tsx",
|
|
||||||
".vue": "vue",
|
|
||||||
".yaml": "yaml",
|
|
||||||
".yml": "yaml",
|
|
||||||
}
|
|
||||||
|
|
||||||
@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_extensions + self.text_extensions
|
return self.docling_extensions + TextFileHandler.text_extensions
|
||||||
|
|
||||||
def convert_file(self, path: Path) -> "DoclingDocument":
|
def convert_file(self, path: Path) -> "DoclingDocument":
|
||||||
"""Convert a file to DoclingDocument using local docling.
|
"""Convert a file to DoclingDocument using local docling.
|
||||||
|
|
@ -122,17 +66,15 @@ class DoclingLocalConverter(DocumentConverter):
|
||||||
converter = DoclingDocConverter()
|
converter = DoclingDocConverter()
|
||||||
result = converter.convert(path)
|
result = converter.convert(path)
|
||||||
return result.document
|
return result.document
|
||||||
elif file_extension in self.text_extensions:
|
elif file_extension in TextFileHandler.text_extensions:
|
||||||
# Read plain text files directly
|
# Read plain text files directly
|
||||||
content = path.read_text(encoding="utf-8")
|
content = path.read_text(encoding="utf-8")
|
||||||
|
# Prepare content with code block wrapping if needed
|
||||||
# Wrap code files (but not plain txt) in markdown code blocks for better presentation
|
prepared_content = TextFileHandler.prepare_text_content(
|
||||||
if file_extension in self.code_markdown_identifier:
|
content, file_extension
|
||||||
language = self.code_markdown_identifier[file_extension]
|
)
|
||||||
content = f"```{language}\n{content}\n```"
|
|
||||||
|
|
||||||
# Convert text to DoclingDocument by wrapping as markdown
|
# Convert text to DoclingDocument by wrapping as markdown
|
||||||
return self.convert_text(content, name=f"{path.stem}.md")
|
return self.convert_text(prepared_content, name=f"{path.stem}.md")
|
||||||
else:
|
else:
|
||||||
# Fallback: try to read as text and convert to DoclingDocument
|
# Fallback: try to read as text and convert to DoclingDocument
|
||||||
content = path.read_text(encoding="utf-8")
|
content = path.read_text(encoding="utf-8")
|
||||||
|
|
@ -153,14 +95,4 @@ class DoclingLocalConverter(DocumentConverter):
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: If the text cannot be converted.
|
ValueError: If the text cannot be converted.
|
||||||
"""
|
"""
|
||||||
from docling.document_converter import DocumentConverter as DoclingDocConverter
|
return TextFileHandler.text_to_docling_document(text, name)
|
||||||
from docling_core.types.io import DocumentStream
|
|
||||||
|
|
||||||
try:
|
|
||||||
bytes_io = BytesIO(text.encode("utf-8"))
|
|
||||||
doc_stream = DocumentStream(name=name, stream=bytes_io)
|
|
||||||
converter = DoclingDocConverter()
|
|
||||||
result = converter.convert(doc_stream)
|
|
||||||
return result.document
|
|
||||||
except Exception as e:
|
|
||||||
raise ValueError(f"Failed to convert text to DoclingDocument: {e}")
|
|
||||||
|
|
|
||||||
181
haiku_rag_slim/haiku/rag/converters/docling_serve.py
Normal file
181
haiku_rag_slim/haiku/rag/converters/docling_serve.py
Normal file
|
|
@ -0,0 +1,181 @@
|
||||||
|
"""docling-serve remote converter implementation."""
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import TYPE_CHECKING, ClassVar
|
||||||
|
|
||||||
|
import requests
|
||||||
|
|
||||||
|
from haiku.rag.config import AppConfig
|
||||||
|
from haiku.rag.converters.base import DocumentConverter
|
||||||
|
from haiku.rag.converters.text_utils import TextFileHandler
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
|
|
||||||
|
class DoclingServeConverter(DocumentConverter):
|
||||||
|
"""Converter that uses docling-serve for document conversion.
|
||||||
|
|
||||||
|
This converter offloads document processing to a docling-serve instance,
|
||||||
|
which handles heavy operations like PDF parsing, OCR, and table extraction.
|
||||||
|
|
||||||
|
For plain text files, it reads them locally and converts to markdown format
|
||||||
|
before sending to docling-serve for DoclingDocument conversion.
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Extensions that docling-serve can handle
|
||||||
|
docling_serve_extensions: ClassVar[list[str]] = [
|
||||||
|
".adoc",
|
||||||
|
".asc",
|
||||||
|
".asciidoc",
|
||||||
|
".bmp",
|
||||||
|
".csv",
|
||||||
|
".docx",
|
||||||
|
".html",
|
||||||
|
".xhtml",
|
||||||
|
".jpeg",
|
||||||
|
".jpg",
|
||||||
|
".md",
|
||||||
|
".pdf",
|
||||||
|
".png",
|
||||||
|
".pptx",
|
||||||
|
".tiff",
|
||||||
|
".xlsx",
|
||||||
|
".xml",
|
||||||
|
".webp",
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, config: AppConfig):
|
||||||
|
"""Initialize the converter with configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Application configuration containing docling-serve settings.
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.base_url = config.providers.docling_serve.base_url.rstrip("/")
|
||||||
|
self.api_key = config.providers.docling_serve.api_key
|
||||||
|
self.timeout = config.providers.docling_serve.timeout
|
||||||
|
|
||||||
|
@property
|
||||||
|
def supported_extensions(self) -> list[str]:
|
||||||
|
"""Return list of file extensions supported by this converter."""
|
||||||
|
return self.docling_serve_extensions + TextFileHandler.text_extensions
|
||||||
|
|
||||||
|
def _make_request(self, files: dict, name: str) -> "DoclingDocument":
|
||||||
|
"""Make a request to docling-serve and return the DoclingDocument.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
files: Dictionary with files parameter for requests
|
||||||
|
name: Name of the document being converted (for error messages)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
DoclingDocument representation
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If conversion fails or service is unavailable
|
||||||
|
"""
|
||||||
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
|
try:
|
||||||
|
url = f"{self.base_url}/v1/convert/file"
|
||||||
|
data = {"to_formats": ["json"]}
|
||||||
|
headers = {}
|
||||||
|
|
||||||
|
if self.api_key:
|
||||||
|
headers["X-Api-Key"] = self.api_key
|
||||||
|
|
||||||
|
response = requests.post(
|
||||||
|
url,
|
||||||
|
files=files,
|
||||||
|
data=data,
|
||||||
|
headers=headers,
|
||||||
|
timeout=self.timeout,
|
||||||
|
)
|
||||||
|
|
||||||
|
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 requests.exceptions.ConnectionError 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 requests.exceptions.Timeout 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 requests.exceptions.HTTPError 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 convert via docling-serve: {e}")
|
||||||
|
|
||||||
|
def convert_file(self, path: Path) -> "DoclingDocument":
|
||||||
|
"""Convert a file to DoclingDocument using docling-serve.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
path: Path to the file to convert.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
DoclingDocument representation of the file.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If the file cannot be converted or service is unavailable.
|
||||||
|
"""
|
||||||
|
file_extension = path.suffix.lower()
|
||||||
|
|
||||||
|
# For plain text files, read locally and prepare content
|
||||||
|
if file_extension in TextFileHandler.text_extensions:
|
||||||
|
try:
|
||||||
|
content = path.read_text(encoding="utf-8")
|
||||||
|
prepared_content = TextFileHandler.prepare_text_content(
|
||||||
|
content, file_extension
|
||||||
|
)
|
||||||
|
return self.convert_text(prepared_content, name=f"{path.stem}.md")
|
||||||
|
except Exception as e:
|
||||||
|
raise ValueError(f"Failed to read text file {path}: {e}")
|
||||||
|
|
||||||
|
# For complex formats, send file to docling-serve
|
||||||
|
with open(path, "rb") as f:
|
||||||
|
files = {"files": f}
|
||||||
|
return self._make_request(files, path.name)
|
||||||
|
|
||||||
|
def convert_text(self, text: str, name: str = "content.md") -> "DoclingDocument":
|
||||||
|
"""Convert text content to DoclingDocument via docling-serve.
|
||||||
|
|
||||||
|
Sends the text as a markdown file to docling-serve for conversion.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: The text content to convert.
|
||||||
|
name: The name to use for the document (defaults to "content.md").
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
DoclingDocument representation of the text.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If the text cannot be converted.
|
||||||
|
"""
|
||||||
|
from io import BytesIO
|
||||||
|
|
||||||
|
text_bytes = text.encode("utf-8")
|
||||||
|
files = {"files": (name, BytesIO(text_bytes), "text/markdown")}
|
||||||
|
return self._make_request(files, name)
|
||||||
117
haiku_rag_slim/haiku/rag/converters/text_utils.py
Normal file
117
haiku_rag_slim/haiku/rag/converters/text_utils.py
Normal file
|
|
@ -0,0 +1,117 @@
|
||||||
|
"""Shared utilities for text file handling in converters."""
|
||||||
|
|
||||||
|
from io import BytesIO
|
||||||
|
from typing import TYPE_CHECKING, ClassVar
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
|
|
||||||
|
class TextFileHandler:
|
||||||
|
"""Handles conversion of text files to DoclingDocument format.
|
||||||
|
|
||||||
|
This class provides shared functionality for converting plain text and code files
|
||||||
|
to DoclingDocument format, with proper code block wrapping for syntax highlighting.
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Plain text extensions that we'll read directly
|
||||||
|
text_extensions: ClassVar[list[str]] = [
|
||||||
|
".astro",
|
||||||
|
".c",
|
||||||
|
".cpp",
|
||||||
|
".css",
|
||||||
|
".go",
|
||||||
|
".h",
|
||||||
|
".hpp",
|
||||||
|
".java",
|
||||||
|
".js",
|
||||||
|
".json",
|
||||||
|
".kt",
|
||||||
|
".mdx",
|
||||||
|
".mjs",
|
||||||
|
".php",
|
||||||
|
".py",
|
||||||
|
".rb",
|
||||||
|
".rs",
|
||||||
|
".svelte",
|
||||||
|
".swift",
|
||||||
|
".ts",
|
||||||
|
".tsx",
|
||||||
|
".txt",
|
||||||
|
".vue",
|
||||||
|
".yaml",
|
||||||
|
".yml",
|
||||||
|
]
|
||||||
|
|
||||||
|
# Code file extensions with their markdown language identifiers
|
||||||
|
code_markdown_identifier: ClassVar[dict[str, str]] = {
|
||||||
|
".astro": "astro",
|
||||||
|
".c": "c",
|
||||||
|
".cpp": "cpp",
|
||||||
|
".css": "css",
|
||||||
|
".go": "go",
|
||||||
|
".h": "c",
|
||||||
|
".hpp": "cpp",
|
||||||
|
".java": "java",
|
||||||
|
".js": "javascript",
|
||||||
|
".json": "json",
|
||||||
|
".kt": "kotlin",
|
||||||
|
".mjs": "javascript",
|
||||||
|
".php": "php",
|
||||||
|
".py": "python",
|
||||||
|
".rb": "ruby",
|
||||||
|
".rs": "rust",
|
||||||
|
".svelte": "svelte",
|
||||||
|
".swift": "swift",
|
||||||
|
".ts": "typescript",
|
||||||
|
".tsx": "tsx",
|
||||||
|
".vue": "vue",
|
||||||
|
".yaml": "yaml",
|
||||||
|
".yml": "yaml",
|
||||||
|
}
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def prepare_text_content(content: str, file_extension: str) -> str:
|
||||||
|
"""Prepare text content for conversion to DoclingDocument.
|
||||||
|
|
||||||
|
Wraps code files in markdown code blocks with appropriate language identifiers.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content: The text content.
|
||||||
|
file_extension: File extension (including dot, e.g., ".py").
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Prepared text content, possibly wrapped in code blocks.
|
||||||
|
"""
|
||||||
|
if file_extension in TextFileHandler.code_markdown_identifier:
|
||||||
|
language = TextFileHandler.code_markdown_identifier[file_extension]
|
||||||
|
return f"```{language}\n{content}\n```"
|
||||||
|
return content
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def text_to_docling_document(
|
||||||
|
text: str, name: str = "content.md"
|
||||||
|
) -> "DoclingDocument":
|
||||||
|
"""Convert text to DoclingDocument using docling's markdown parser.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: The text content to convert.
|
||||||
|
name: The name to use for the document.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
DoclingDocument representation of the text.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If the conversion fails.
|
||||||
|
"""
|
||||||
|
from docling.document_converter import DocumentConverter as DoclingDocConverter
|
||||||
|
from docling_core.types.io import DocumentStream
|
||||||
|
|
||||||
|
try:
|
||||||
|
bytes_io = BytesIO(text.encode("utf-8"))
|
||||||
|
doc_stream = DocumentStream(name=name, stream=bytes_io)
|
||||||
|
converter = DoclingDocConverter()
|
||||||
|
result = converter.convert(doc_stream)
|
||||||
|
return result.document
|
||||||
|
except Exception as e:
|
||||||
|
raise ValueError(f"Failed to convert text to DoclingDocument: {e}")
|
||||||
|
|
@ -28,10 +28,11 @@ class FileFilter(DefaultFilter):
|
||||||
if supported_extensions is None:
|
if supported_extensions is None:
|
||||||
# Default to docling-local extensions if not provided
|
# Default to docling-local extensions if not provided
|
||||||
from haiku.rag.converters.docling_local import DoclingLocalConverter
|
from haiku.rag.converters.docling_local import DoclingLocalConverter
|
||||||
|
from haiku.rag.converters.text_utils import TextFileHandler
|
||||||
|
|
||||||
supported_extensions = (
|
supported_extensions = (
|
||||||
DoclingLocalConverter.docling_extensions
|
DoclingLocalConverter.docling_extensions
|
||||||
+ DoclingLocalConverter.text_extensions
|
+ TextFileHandler.text_extensions
|
||||||
)
|
)
|
||||||
|
|
||||||
self.extensions = tuple(supported_extensions)
|
self.extensions = tuple(supported_extensions)
|
||||||
|
|
|
||||||
|
|
@ -102,6 +102,9 @@ asyncio_default_fixture_loop_scope = "session"
|
||||||
asyncio_mode = "auto"
|
asyncio_mode = "auto"
|
||||||
testpaths = ["tests"]
|
testpaths = ["tests"]
|
||||||
norecursedirs = ["examples", "docs", "evaluations", ".git", ".venv"]
|
norecursedirs = ["examples", "docs", "evaluations", ".git", ".venv"]
|
||||||
|
markers = [
|
||||||
|
"integration: marks tests as integration tests (require external services)",
|
||||||
|
]
|
||||||
|
|
||||||
# pyproject.toml
|
# pyproject.toml
|
||||||
filterwarnings = ["error", "ignore::UserWarning", "ignore::DeprecationWarning"]
|
filterwarnings = ["error", "ignore::UserWarning", "ignore::DeprecationWarning"]
|
||||||
|
|
|
||||||
350
tests/test_converters.py
Normal file
350
tests/test_converters.py
Normal file
|
|
@ -0,0 +1,350 @@
|
||||||
|
"""Tests for document converters."""
|
||||||
|
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import Mock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import requests
|
||||||
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
|
from haiku.rag.config import AppConfig
|
||||||
|
from haiku.rag.converters import get_converter
|
||||||
|
from haiku.rag.converters.docling_local import DoclingLocalConverter
|
||||||
|
from haiku.rag.converters.docling_serve import DoclingServeConverter
|
||||||
|
from haiku.rag.converters.text_utils import TextFileHandler
|
||||||
|
|
||||||
|
|
||||||
|
def is_docling_serve_available(base_url: str = "http://localhost:5001") -> bool:
|
||||||
|
"""Check if docling-serve is running and accessible."""
|
||||||
|
try:
|
||||||
|
response = requests.get(f"{base_url}/health", timeout=2)
|
||||||
|
return response.status_code == 200
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def create_mock_docling_document_json(name: str = "test") -> dict:
|
||||||
|
"""Create a minimal valid DoclingDocument JSON structure for mocking."""
|
||||||
|
return {
|
||||||
|
"schema_name": "DoclingDocument",
|
||||||
|
"version": "1.8.0",
|
||||||
|
"name": name,
|
||||||
|
"origin": {
|
||||||
|
"mimetype": "text/markdown",
|
||||||
|
"binary_hash": 12345,
|
||||||
|
"filename": f"{name}.md",
|
||||||
|
},
|
||||||
|
"furniture": {
|
||||||
|
"self_ref": "#/furniture",
|
||||||
|
"parent": None,
|
||||||
|
"children": [],
|
||||||
|
"content_layer": "furniture",
|
||||||
|
"name": "_root_",
|
||||||
|
"label": "unspecified",
|
||||||
|
},
|
||||||
|
"body": {
|
||||||
|
"self_ref": "#/body",
|
||||||
|
"parent": None,
|
||||||
|
"children": [],
|
||||||
|
"content_layer": "body",
|
||||||
|
"name": "_root_",
|
||||||
|
"label": "unspecified",
|
||||||
|
},
|
||||||
|
"groups": [],
|
||||||
|
"texts": [],
|
||||||
|
"pictures": [],
|
||||||
|
"tables": [],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class TestTextFileHandler:
|
||||||
|
"""Tests for TextFileHandler utility class."""
|
||||||
|
|
||||||
|
def test_text_extensions_defined(self):
|
||||||
|
"""Test that text extensions list is defined."""
|
||||||
|
assert len(TextFileHandler.text_extensions) > 0
|
||||||
|
assert ".py" in TextFileHandler.text_extensions
|
||||||
|
assert ".js" in TextFileHandler.text_extensions
|
||||||
|
assert ".txt" in TextFileHandler.text_extensions
|
||||||
|
|
||||||
|
def test_code_markdown_identifiers(self):
|
||||||
|
"""Test code language identifiers mapping."""
|
||||||
|
assert TextFileHandler.code_markdown_identifier[".py"] == "python"
|
||||||
|
assert TextFileHandler.code_markdown_identifier[".js"] == "javascript"
|
||||||
|
assert TextFileHandler.code_markdown_identifier[".ts"] == "typescript"
|
||||||
|
|
||||||
|
def test_prepare_text_content_with_code(self):
|
||||||
|
"""Test that code files are wrapped in markdown code blocks."""
|
||||||
|
code = "def hello():\n pass"
|
||||||
|
result = TextFileHandler.prepare_text_content(code, ".py")
|
||||||
|
assert result.startswith("```python\n")
|
||||||
|
assert result.endswith("\n```")
|
||||||
|
assert "def hello():" in result
|
||||||
|
|
||||||
|
def test_prepare_text_content_without_code(self):
|
||||||
|
"""Test that plain text files are not wrapped."""
|
||||||
|
text = "Hello world"
|
||||||
|
result = TextFileHandler.prepare_text_content(text, ".txt")
|
||||||
|
assert result == text
|
||||||
|
assert not result.startswith("```")
|
||||||
|
|
||||||
|
|
||||||
|
class TestConverterFactory:
|
||||||
|
"""Tests for converter factory function."""
|
||||||
|
|
||||||
|
def test_get_docling_local_converter(self):
|
||||||
|
"""Test getting docling-local converter."""
|
||||||
|
config = AppConfig()
|
||||||
|
config.processing.converter = "docling-local"
|
||||||
|
converter = get_converter(config)
|
||||||
|
assert isinstance(converter, DoclingLocalConverter)
|
||||||
|
|
||||||
|
def test_get_docling_serve_converter(self):
|
||||||
|
"""Test getting docling-serve converter."""
|
||||||
|
config = AppConfig()
|
||||||
|
config.processing.converter = "docling-serve"
|
||||||
|
converter = get_converter(config)
|
||||||
|
assert isinstance(converter, DoclingServeConverter)
|
||||||
|
|
||||||
|
def test_invalid_converter_raises_error(self):
|
||||||
|
"""Test that invalid converter name raises ValueError."""
|
||||||
|
config = AppConfig()
|
||||||
|
config.processing.converter = "invalid-converter"
|
||||||
|
with pytest.raises(ValueError, match="Unsupported converter provider"):
|
||||||
|
get_converter(config)
|
||||||
|
|
||||||
|
|
||||||
|
class TestDoclingLocalConverter:
|
||||||
|
"""Tests for DoclingLocalConverter."""
|
||||||
|
|
||||||
|
def test_supported_extensions(self):
|
||||||
|
"""Test that converter reports correct supported extensions."""
|
||||||
|
converter = DoclingLocalConverter()
|
||||||
|
extensions = converter.supported_extensions
|
||||||
|
assert ".pdf" in extensions
|
||||||
|
assert ".docx" in extensions
|
||||||
|
assert ".py" in extensions
|
||||||
|
assert ".txt" in extensions
|
||||||
|
|
||||||
|
def test_convert_text(self):
|
||||||
|
"""Test converting text to DoclingDocument."""
|
||||||
|
converter = DoclingLocalConverter()
|
||||||
|
doc = converter.convert_text("# Test\n\nContent here", name="test.md")
|
||||||
|
assert isinstance(doc, DoclingDocument)
|
||||||
|
assert doc.name == "test"
|
||||||
|
|
||||||
|
def test_convert_code_file(self):
|
||||||
|
"""Test that code files are wrapped in code blocks."""
|
||||||
|
python_code = "def hello():\n print('Hello')"
|
||||||
|
converter = DoclingLocalConverter()
|
||||||
|
|
||||||
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
|
||||||
|
f.write(python_code)
|
||||||
|
f.flush()
|
||||||
|
temp_path = Path(f.name)
|
||||||
|
doc = converter.convert_file(temp_path)
|
||||||
|
result = doc.export_to_markdown()
|
||||||
|
|
||||||
|
assert "```" in result
|
||||||
|
assert "def hello():" in result
|
||||||
|
|
||||||
|
|
||||||
|
class TestDoclingServeConverter:
|
||||||
|
"""Tests for DoclingServeConverter (mocked)."""
|
||||||
|
|
||||||
|
@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 = ""
|
||||||
|
config.providers.docling_serve.timeout = 300
|
||||||
|
return config
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def converter(self, config):
|
||||||
|
"""Create DoclingServeConverter instance."""
|
||||||
|
return DoclingServeConverter(config)
|
||||||
|
|
||||||
|
def test_initialization(self, converter):
|
||||||
|
"""Test converter initialization."""
|
||||||
|
assert converter.base_url == "http://localhost:5001"
|
||||||
|
assert converter.timeout == 300
|
||||||
|
|
||||||
|
def test_supported_extensions(self, converter):
|
||||||
|
"""Test that converter reports correct supported extensions."""
|
||||||
|
extensions = converter.supported_extensions
|
||||||
|
assert ".pdf" in extensions
|
||||||
|
assert ".docx" in extensions
|
||||||
|
assert ".py" in extensions
|
||||||
|
assert ".md" in extensions
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_text_success(self, mock_post, converter):
|
||||||
|
"""Test successful text conversion via docling-serve."""
|
||||||
|
mock_response = Mock()
|
||||||
|
mock_response.status_code = 200
|
||||||
|
mock_response.json.return_value = {
|
||||||
|
"status": "success",
|
||||||
|
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||||
|
}
|
||||||
|
mock_post.return_value = mock_response
|
||||||
|
|
||||||
|
doc = converter.convert_text("# Test", name="test.md")
|
||||||
|
assert isinstance(doc, DoclingDocument)
|
||||||
|
assert doc.version == "1.8.0"
|
||||||
|
mock_post.assert_called_once()
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_text_with_api_key(self, mock_post, config):
|
||||||
|
"""Test that API key is included in request headers."""
|
||||||
|
config.providers.docling_serve.api_key = "test-key"
|
||||||
|
converter = DoclingServeConverter(config)
|
||||||
|
|
||||||
|
mock_response = Mock()
|
||||||
|
mock_response.status_code = 200
|
||||||
|
mock_response.json.return_value = {
|
||||||
|
"status": "success",
|
||||||
|
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||||
|
}
|
||||||
|
mock_post.return_value = mock_response
|
||||||
|
|
||||||
|
converter.convert_text("# Test")
|
||||||
|
|
||||||
|
call_kwargs = mock_post.call_args.kwargs
|
||||||
|
assert "headers" in call_kwargs
|
||||||
|
assert call_kwargs["headers"]["X-Api-Key"] == "test-key"
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_text_connection_error(self, mock_post, converter):
|
||||||
|
"""Test handling of connection errors."""
|
||||||
|
import requests
|
||||||
|
|
||||||
|
mock_post.side_effect = requests.exceptions.ConnectionError("Connection failed")
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Could not connect to docling-serve"):
|
||||||
|
converter.convert_text("# Test")
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_text_timeout_error(self, mock_post, converter):
|
||||||
|
"""Test handling of timeout errors."""
|
||||||
|
import requests
|
||||||
|
|
||||||
|
mock_post.side_effect = requests.exceptions.Timeout("Timeout")
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="timed out"):
|
||||||
|
converter.convert_text("# Test")
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_text_auth_error(self, mock_post, converter):
|
||||||
|
"""Test handling of authentication errors."""
|
||||||
|
import requests
|
||||||
|
|
||||||
|
mock_response = Mock()
|
||||||
|
mock_response.status_code = 401
|
||||||
|
mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
|
||||||
|
response=mock_response
|
||||||
|
)
|
||||||
|
mock_post.return_value = mock_response
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Authentication failed"):
|
||||||
|
converter.convert_text("# Test")
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_text_no_json_content(self, mock_post, converter):
|
||||||
|
"""Test handling when docling-serve returns no JSON content."""
|
||||||
|
mock_response = Mock()
|
||||||
|
mock_response.status_code = 200
|
||||||
|
mock_response.json.return_value = {
|
||||||
|
"status": "success",
|
||||||
|
"document": {"json_content": None},
|
||||||
|
}
|
||||||
|
mock_post.return_value = mock_response
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="did not return JSON content"):
|
||||||
|
converter.convert_text("# Test")
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_file_pdf(self, mock_post, converter):
|
||||||
|
"""Test converting PDF file via docling-serve."""
|
||||||
|
mock_response = Mock()
|
||||||
|
mock_response.status_code = 200
|
||||||
|
mock_response.json.return_value = {
|
||||||
|
"status": "success",
|
||||||
|
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||||
|
}
|
||||||
|
mock_post.return_value = mock_response
|
||||||
|
|
||||||
|
with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
|
||||||
|
f.write(b"fake pdf content")
|
||||||
|
f.flush()
|
||||||
|
temp_path = Path(f.name)
|
||||||
|
doc = converter.convert_file(temp_path)
|
||||||
|
|
||||||
|
assert isinstance(doc, DoclingDocument)
|
||||||
|
mock_post.assert_called_once()
|
||||||
|
|
||||||
|
@patch("haiku.rag.converters.docling_serve.requests.post")
|
||||||
|
def test_convert_file_text(self, mock_post, converter):
|
||||||
|
"""Test converting text file (reads locally, sends to docling-serve)."""
|
||||||
|
mock_response = Mock()
|
||||||
|
mock_response.status_code = 200
|
||||||
|
mock_response.json.return_value = {
|
||||||
|
"status": "success",
|
||||||
|
"document": {"json_content": create_mock_docling_document_json("test")},
|
||||||
|
}
|
||||||
|
mock_post.return_value = mock_response
|
||||||
|
|
||||||
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
|
||||||
|
f.write("def hello():\n pass")
|
||||||
|
f.flush()
|
||||||
|
temp_path = Path(f.name)
|
||||||
|
doc = converter.convert_file(temp_path)
|
||||||
|
|
||||||
|
assert isinstance(doc, DoclingDocument)
|
||||||
|
# Should call docling-serve for conversion
|
||||||
|
mock_post.assert_called_once()
|
||||||
|
# Check that code was wrapped in code block
|
||||||
|
call_kwargs = mock_post.call_args.kwargs
|
||||||
|
assert "files" in call_kwargs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.integration
|
||||||
|
@pytest.mark.skipif(
|
||||||
|
not is_docling_serve_available(),
|
||||||
|
reason="docling-serve not available at http://localhost:5001",
|
||||||
|
)
|
||||||
|
class TestDoclingServeConverterIntegration:
|
||||||
|
"""Integration tests with real docling-serve (requires service running)."""
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def config(self):
|
||||||
|
"""Create configuration for integration tests."""
|
||||||
|
config = AppConfig()
|
||||||
|
config.providers.docling_serve.base_url = "http://localhost:5001"
|
||||||
|
return config
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def converter(self, config):
|
||||||
|
"""Create converter for integration tests."""
|
||||||
|
return DoclingServeConverter(config)
|
||||||
|
|
||||||
|
def test_convert_text_real_service(self, converter):
|
||||||
|
"""Test text conversion with real docling-serve (integration)."""
|
||||||
|
doc = converter.convert_text("# Test Document\n\nThis is a test.")
|
||||||
|
assert isinstance(doc, DoclingDocument)
|
||||||
|
assert doc.version == "1.8.0"
|
||||||
|
|
||||||
|
def test_convert_code_file_real_service(self, converter):
|
||||||
|
"""Test code file conversion with real docling-serve (integration)."""
|
||||||
|
code = "def test():\n return 42"
|
||||||
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
|
||||||
|
f.write(code)
|
||||||
|
f.flush()
|
||||||
|
temp_path = Path(f.name)
|
||||||
|
doc = converter.convert_file(temp_path)
|
||||||
|
|
||||||
|
assert isinstance(doc, DoclingDocument)
|
||||||
|
result = doc.export_to_markdown()
|
||||||
|
assert "def test():" in result
|
||||||
Loading…
Reference in a new issue