danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Implement docling-serve converter

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-07 18:02:13 +02:00
parent 92def1c6a5
commit fe7d0c5c24
No known key found for this signature in database
7 changed files with 666 additions and 82 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
haiku_rag_slim/haiku/rag/converters/__init__.py
View file

@ -23,9 +23,9 @@ def get_converter(config: AppConfig = Config) -> DocumentConverter:
return DoclingLocalConverter()
# if config.processing.converter == "docling-serve":
# from haiku.rag.converters.docling_serve import DoclingServeConverter
#
# return DoclingServeConverter(config)
if config.processing.converter == "docling-serve":
from haiku.rag.converters.docling_serve import DoclingServeConverter
return DoclingServeConverter(config)
raise ValueError(f"Unsupported converter provider: {config.processing.converter}")

86
haiku_rag_slim/haiku/rag/converters/docling_local.py
View file

@ -1,10 +1,10 @@
"""Local docling converter implementation."""
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, ClassVar
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
@ -39,66 +39,10 @@ class DoclingLocalConverter(DocumentConverter):
".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
def supported_extensions(self) -> list[str]:
"""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":
"""Convert a file to DoclingDocument using local docling.
@ -122,17 +66,15 @@ class DoclingLocalConverter(DocumentConverter):
converter = DoclingDocConverter()
result = converter.convert(path)
return result.document
elif file_extension in self.text_extensions:
elif file_extension in TextFileHandler.text_extensions:
# Read plain text files directly
content = path.read_text(encoding="utf-8")
# Wrap code files (but not plain txt) in markdown code blocks for better presentation
if file_extension in self.code_markdown_identifier:
language = self.code_markdown_identifier[file_extension]
content = f"```{language}\n{content}\n```"
# Prepare content with code block wrapping if needed
prepared_content = TextFileHandler.prepare_text_content(
content, file_extension
)
# 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:
# Fallback: try to read as text and convert to DoclingDocument
content = path.read_text(encoding="utf-8")
@ -153,14 +95,4 @@ class DoclingLocalConverter(DocumentConverter):
Raises:
ValueError: If the text cannot be converted.
"""
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}")
return TextFileHandler.text_to_docling_document(text, name)

181
haiku_rag_slim/haiku/rag/converters/docling_serve.py Normal file
View 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
View 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}")

3
haiku_rag_slim/haiku/rag/monitor.py
View file

@ -28,10 +28,11 @@ class FileFilter(DefaultFilter):
if supported_extensions is None:
# Default to docling-local extensions if not provided
from haiku.rag.converters.docling_local import DoclingLocalConverter
from haiku.rag.converters.text_utils import TextFileHandler
supported_extensions = (
DoclingLocalConverter.docling_extensions
+ DoclingLocalConverter.text_extensions
+ TextFileHandler.text_extensions
)
self.extensions = tuple(supported_extensions)

3
pyproject.toml
View file

@ -102,6 +102,9 @@ asyncio_default_fixture_loop_scope = "session"
asyncio_mode = "auto"
testpaths = ["tests"]
norecursedirs = ["examples", "docs", "evaluations", ".git", ".venv"]
markers = [
"integration: marks tests as integration tests (require external services)",
]
# pyproject.toml
filterwarnings = ["error", "ignore::UserWarning", "ignore::DeprecationWarning"]

350
tests/test_converters.py Normal file
View 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