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

async versions for convert_file() convert_text()

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-26 17:46:25 +02:00
parent 6e24d5772b
commit 14bdb7cb27
No known key found for this signature in database
13 changed files with 215 additions and 176 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
haiku_rag_slim/haiku/rag/client.py
View file

@ -114,7 +114,7 @@ class HaikuRAG:
if chunks is None:
# Use converter to convert text
converter = get_converter(self._config)
docling_document = converter.convert_text(content)
docling_document = await converter.convert_text(content)
else:
# Chunks already provided, no conversion needed
docling_document = None
@ -241,7 +241,7 @@ class HaikuRAG:
# Parse file only when content changed or new document
converter = get_converter(self._config)
docling_document = converter.convert_file(source_path)
docling_document = await converter.convert_file(source_path)
if existing_doc:
# Update existing document
@ -333,7 +333,7 @@ class HaikuRAG:
temp_path = Path(temp_file.name)
# Parse the content using converter
docling_document = converter.convert_file(temp_path)
docling_document = await converter.convert_file(temp_path)
# Merge metadata with contentType and md5
metadata.update({"contentType": content_type, "md5": md5_hash})
@ -411,7 +411,7 @@ class HaikuRAG:
"""Update an existing document."""
# Convert content to DoclingDocument
converter = get_converter(self._config)
docling_document = converter.convert_text(document.content)
docling_document = await converter.convert_text(document.content)
return await self.document_repository._update_and_rechunk(
document, docling_document
@ -472,7 +472,7 @@ class HaikuRAG:
else:
# Auto-generate chunks from content
converter = get_converter(self._config)
docling_document = converter.convert_text(existing_doc.content)
docling_document = await converter.convert_text(existing_doc.content)
return await self.document_repository._update_and_rechunk(
existing_doc, docling_document
)
@ -762,14 +762,14 @@ class HaikuRAG:
logger.warning(
"Source missing for %s, re-embedding from content", doc.uri
)
docling_document = converter.convert_text(doc.content)
docling_document = await converter.convert_text(doc.content)
await self.chunk_repository.create_chunks_for_document(
doc.id, docling_document
)
yield doc.id
else:
# Document without URI - re-create chunks from existing content
docling_document = converter.convert_text(doc.content)
docling_document = await converter.convert_text(doc.content)
await self.chunk_repository.create_chunks_for_document(
doc.id, docling_document
)

6
haiku_rag_slim/haiku/rag/converters/base.py
View file

@ -26,7 +26,7 @@ class DocumentConverter(ABC):
pass
@abstractmethod
def convert_file(self, path: Path) -> "DoclingDocument":
async def convert_file(self, path: Path) -> "DoclingDocument":
"""Convert a file to DoclingDocument format.
Args:
@ -41,7 +41,9 @@ class DocumentConverter(ABC):
pass
@abstractmethod
def convert_text(self, text: str, name: str = "content.md") -> "DoclingDocument":
async def convert_text(
self, text: str, name: str = "content.md"
) -> "DoclingDocument":
"""Convert text content to DoclingDocument format.
Args:

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

@ -1,5 +1,6 @@
"""Local docling converter implementation."""
import asyncio
from pathlib import Path
from typing import TYPE_CHECKING, ClassVar, cast
@ -53,18 +54,8 @@ class DoclingLocalConverter(DocumentConverter):
"""Return list of file extensions supported by this converter."""
return self.docling_extensions + TextFileHandler.text_extensions
def convert_file(self, path: Path) -> "DoclingDocument":
"""Convert a file to DoclingDocument using local docling.
Args:
path: Path to the file to convert.
Returns:
DoclingDocument representation of the file.
Raises:
ValueError: If the file cannot be converted.
"""
def _sync_convert_docling_file(self, path: Path) -> "DoclingDocument":
"""Synchronous conversion of docling-supported files."""
from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import (
@ -81,64 +72,72 @@ class DoclingLocalConverter(DocumentConverter):
PdfFormatOption,
)
opts = self.config.processing.conversion_options
pipeline_options = PdfPipelineOptions(
do_ocr=opts.do_ocr,
do_table_structure=opts.do_table_structure,
images_scale=opts.images_scale,
table_structure_options=TableStructureOptions(
do_cell_matching=opts.table_cell_matching,
mode=(
TableFormerMode.FAST
if opts.table_mode == "fast"
else TableFormerMode.ACCURATE
),
),
ocr_options=OcrOptions(
force_full_page_ocr=opts.force_ocr,
lang=opts.ocr_lang if opts.ocr_lang else [],
),
)
format_options = cast(
dict[InputFormat, FormatOption],
{
InputFormat.PDF: PdfFormatOption(
pipeline_options=pipeline_options,
backend=DoclingParseDocumentBackend,
)
},
)
converter = DoclingDocConverter(format_options=format_options)
result = converter.convert(path)
return result.document
async def convert_file(self, path: Path) -> "DoclingDocument":
"""Convert a file to DoclingDocument using local docling.
Args:
path: Path to the file to convert.
Returns:
DoclingDocument representation of the file.
Raises:
ValueError: If the file cannot be converted.
"""
try:
file_extension = path.suffix.lower()
if file_extension in self.docling_extensions:
# Get conversion options from config
opts = self.config.processing.conversion_options
# Build pipeline options for PDF conversion
pipeline_options = PdfPipelineOptions(
do_ocr=opts.do_ocr,
do_table_structure=opts.do_table_structure,
images_scale=opts.images_scale,
table_structure_options=TableStructureOptions(
do_cell_matching=opts.table_cell_matching,
mode=(
TableFormerMode.FAST
if opts.table_mode == "fast"
else TableFormerMode.ACCURATE
),
),
ocr_options=OcrOptions(
force_full_page_ocr=opts.force_ocr,
lang=opts.ocr_lang if opts.ocr_lang else [],
),
)
# Create format options for PDF
format_options = cast(
dict[InputFormat, FormatOption],
{
InputFormat.PDF: PdfFormatOption(
pipeline_options=pipeline_options,
backend=DoclingParseDocumentBackend,
)
},
)
# Use docling for complex document formats
converter = DoclingDocConverter(format_options=format_options)
result = converter.convert(path)
return result.document
return await asyncio.to_thread(self._sync_convert_docling_file, path)
elif file_extension in TextFileHandler.text_extensions:
# Read plain text files directly
content = path.read_text(encoding="utf-8")
# Prepare content with code block wrapping if needed
content = await asyncio.to_thread(path.read_text, encoding="utf-8")
prepared_content = TextFileHandler.prepare_text_content(
content, file_extension
)
# Convert text to DoclingDocument by wrapping as markdown
return self.convert_text(prepared_content, name=f"{path.stem}.md")
return await 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")
return self.convert_text(content, name=f"{path.stem}.md")
content = await asyncio.to_thread(path.read_text, encoding="utf-8")
return await self.convert_text(content, name=f"{path.stem}.md")
except Exception:
raise ValueError(f"Failed to parse file: {path}")
def convert_text(self, text: str, name: str = "content.md") -> "DoclingDocument":
async def convert_text(
self, text: str, name: str = "content.md"
) -> "DoclingDocument":
"""Convert text content to DoclingDocument using local docling.
Args:
@ -151,4 +150,4 @@ class DoclingLocalConverter(DocumentConverter):
Raises:
ValueError: If the text cannot be converted.
"""
return TextFileHandler.text_to_docling_document(text, name)
return await TextFileHandler.text_to_docling_document(text, name)

93
haiku_rag_slim/haiku/rag/converters/docling_serve.py
View file

@ -1,5 +1,6 @@
"""docling-serve remote converter implementation."""
import asyncio
from pathlib import Path
from typing import TYPE_CHECKING, ClassVar
@ -61,7 +62,40 @@ class DoclingServeConverter(DocumentConverter):
"""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":
def _sync_make_request(
self, files: dict, name: str, data: dict, headers: dict
) -> "DoclingDocument":
"""Synchronous HTTP request to docling-serve."""
from docling_core.types.doc.document import DoclingDocument
url = f"{self.base_url}/v1/convert/file"
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)
async def _make_request(self, files: dict, name: str) -> "DoclingDocument":
"""Make a request to docling-serve and return the DoclingDocument.
Args:
@ -74,27 +108,19 @@ class DoclingServeConverter(DocumentConverter):
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"
opts = self.config.processing.conversion_options
# Build data dict with conversion options
data = {
"to_formats": ["json"],
# OCR options
"do_ocr": opts.do_ocr,
"force_ocr": opts.force_ocr,
# Table options
"do_table_structure": opts.do_table_structure,
"table_mode": opts.table_mode,
"table_cell_matching": opts.table_cell_matching,
# Image options
"images_scale": opts.images_scale,
}
# Add OCR language if specified
if opts.ocr_lang:
data["ocr_lang"] = opts.ocr_lang
@ -102,32 +128,10 @@ class DoclingServeConverter(DocumentConverter):
if self.api_key:
headers["X-Api-Key"] = self.api_key
response = requests.post(
url,
files=files,
data=data,
headers=headers,
timeout=self.timeout,
return await asyncio.to_thread(
self._sync_make_request, files, name, data, 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 requests.exceptions.ConnectionError as e:
raise ValueError(
f"Could not connect to docling-serve at {self.base_url}. "
@ -147,7 +151,7 @@ class DoclingServeConverter(DocumentConverter):
except Exception as e:
raise ValueError(f"Failed to convert via docling-serve: {e}")
def convert_file(self, path: Path) -> "DoclingDocument":
async def convert_file(self, path: Path) -> "DoclingDocument":
"""Convert a file to DoclingDocument using docling-serve.
Args:
@ -161,23 +165,26 @@ class DoclingServeConverter(DocumentConverter):
"""
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")
content = await asyncio.to_thread(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")
return await 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 read_and_prepare_files():
with open(path, "rb") as f:
return {"files": (path.name, f.read(), "application/octet-stream")}
def convert_text(self, text: str, name: str = "content.md") -> "DoclingDocument":
files = await asyncio.to_thread(read_and_prepare_files)
return await self._make_request(files, path.name)
async 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.
@ -196,4 +203,4 @@ class DoclingServeConverter(DocumentConverter):
text_bytes = text.encode("utf-8")
files = {"files": (name, BytesIO(text_bytes), "text/markdown")}
return self._make_request(files, name)
return await self._make_request(files, name)

28
haiku_rag_slim/haiku/rag/converters/text_utils.py
View file

@ -1,5 +1,6 @@
"""Shared utilities for text file handling in converters."""
import asyncio
from io import BytesIO
from typing import TYPE_CHECKING, ClassVar
@ -89,7 +90,21 @@ class TextFileHandler:
return content
@staticmethod
def text_to_docling_document(
def _sync_text_to_docling_document(
text: str, name: str = "content.md"
) -> "DoclingDocument":
"""Synchronous implementation of text to DoclingDocument conversion."""
from docling.document_converter import DocumentConverter as DoclingDocConverter
from docling_core.types.io import DocumentStream
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
@staticmethod
async def text_to_docling_document(
text: str, name: str = "content.md"
) -> "DoclingDocument":
"""Convert text to DoclingDocument using docling's markdown parser.
@ -104,14 +119,9 @@ class TextFileHandler:
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
return await asyncio.to_thread(
TextFileHandler._sync_text_to_docling_document, text, name
)
except Exception as e:
raise ValueError(f"Failed to convert text to DoclingDocument: {e}")

2
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -216,7 +216,7 @@ class ChunkRepository:
if not isinstance(processed_markdown, str):
raise ValueError("Preprocessor must return a markdown string")
converter = get_converter(self.store._config)
processed_document = converter.convert_text(
processed_document = await converter.convert_text(
processed_markdown, name="content.md"
)
except Exception as e:

2
tests/test_chunk.py
View file

@ -73,7 +73,7 @@ async def test_create_chunks_for_document(qa_corpus: Dataset, temp_db_path):
# Convert text to DoclingDocument
converter = get_converter(Config)
docling_document = converter.convert_text(document_text, name="test.md")
docling_document = await converter.convert_text(document_text, name="test.md")
# Test creating chunks for the document
chunks = await chunk_repo.create_chunks_for_document(document_id, docling_document)

20
tests/test_chunker.py
View file

@ -19,7 +19,7 @@ async def test_local_chunker(qa_corpus: Dataset):
# Convert text to DoclingDocument
converter = get_converter(Config)
doc = converter.convert_text(doc_text, name="test.md")
doc = await converter.convert_text(doc_text, name="test.md")
chunks = await chunker.chunk(doc)
@ -87,7 +87,7 @@ async def test_local_chunker_hierarchical(qa_corpus: Dataset):
doc_text = qa_corpus[0]["document_extracted"]
converter = get_converter(Config)
doc = converter.convert_text(doc_text, name="test.md")
doc = await converter.convert_text(doc_text, name="test.md")
chunks = await chunker.chunk(doc)
@ -118,7 +118,7 @@ async def test_local_chunker_markdown_tables():
"""
converter = get_converter(Config)
doc = converter.convert_text(markdown_with_table, name="test.md")
doc = await converter.convert_text(markdown_with_table, name="test.md")
# Test with markdown tables enabled
config_md = AppConfig()
@ -185,7 +185,7 @@ class TestDoclingServeChunker:
# Create a simple document
converter = get_converter(Config)
doc = converter.convert_text("# Test\n\nContent", name="test.md")
doc = await converter.convert_text("# Test\n\nContent", name="test.md")
chunks = await chunker.chunk(doc)
assert len(chunks) == 2
@ -208,7 +208,7 @@ class TestDoclingServeChunker:
mock_post.return_value = mock_response
converter = get_converter(Config)
doc = converter.convert_text("# Test", name="test.md")
doc = await converter.convert_text("# Test", name="test.md")
await chunker.chunk(doc)
call_kwargs = mock_post.call_args.kwargs
@ -230,7 +230,7 @@ class TestDoclingServeChunker:
mock_post.return_value = mock_response
converter = get_converter(Config)
doc = converter.convert_text("# Test", name="test.md")
doc = await converter.convert_text("# Test", name="test.md")
await chunker.chunk(doc)
call_args = mock_post.call_args
@ -253,7 +253,7 @@ class TestDoclingServeChunker:
mock_post.return_value = mock_response
converter = get_converter(Config)
doc = converter.convert_text("# Test", name="test.md")
doc = await converter.convert_text("# Test", name="test.md")
await chunker.chunk(doc)
call_kwargs = mock_post.call_args.kwargs
@ -271,7 +271,7 @@ class TestDoclingServeChunker:
mock_post.side_effect = requests.exceptions.ConnectionError("Connection failed")
converter = get_converter(Config)
doc = converter.convert_text("# Test", name="test.md")
doc = await converter.convert_text("# Test", name="test.md")
with pytest.raises(ValueError, match="Could not connect to docling-serve"):
await chunker.chunk(doc)
@ -285,7 +285,7 @@ class TestDoclingServeChunker:
mock_post.side_effect = requests.exceptions.Timeout("Timeout")
converter = get_converter(Config)
doc = converter.convert_text("# Test", name="test.md")
doc = await converter.convert_text("# Test", name="test.md")
with pytest.raises(ValueError, match="timed out"):
await chunker.chunk(doc)
@ -304,7 +304,7 @@ class TestDoclingServeChunker:
mock_post.return_value = mock_response
converter = get_converter(Config)
doc = converter.convert_text("# Test", name="test.md")
doc = await converter.convert_text("# Test", name="test.md")
with pytest.raises(ValueError, match="Authentication failed"):
await chunker.chunk(doc)

65
tests/test_converters.py
View file

@ -136,13 +136,15 @@ class TestDoclingLocalConverter:
assert ".py" in extensions
assert ".txt" in extensions
def test_convert_text(self, converter):
@pytest.mark.asyncio
async def test_convert_text(self, converter):
"""Test converting text to DoclingDocument."""
doc = converter.convert_text("# Test\n\nContent here", name="test.md")
doc = await converter.convert_text("# Test\n\nContent here", name="test.md")
assert isinstance(doc, DoclingDocument)
assert doc.name == "test"
def test_convert_code_file(self, converter):
@pytest.mark.asyncio
async def test_convert_code_file(self, converter):
"""Test that code files are wrapped in code blocks."""
python_code = "def hello():\n print('Hello')"
@ -150,7 +152,7 @@ class TestDoclingLocalConverter:
f.write(python_code)
f.flush()
temp_path = Path(f.name)
doc = converter.convert_file(temp_path)
doc = await converter.convert_file(temp_path)
result = doc.export_to_markdown()
assert "```" in result
@ -198,8 +200,9 @@ class TestDoclingServeConverter:
assert ".py" in extensions
assert ".md" in extensions
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_text_success(self, mock_post, converter):
async def test_convert_text_success(self, mock_post, converter):
"""Test successful text conversion via docling-serve."""
mock_response = Mock()
mock_response.status_code = 200
@ -209,13 +212,14 @@ class TestDoclingServeConverter:
}
mock_post.return_value = mock_response
doc = converter.convert_text("# Test", name="test.md")
doc = await converter.convert_text("# Test", name="test.md")
assert isinstance(doc, DoclingDocument)
assert doc.version == "1.8.0"
mock_post.assert_called_once()
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_text_with_api_key(self, mock_post, config):
async 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)
@ -228,14 +232,15 @@ class TestDoclingServeConverter:
}
mock_post.return_value = mock_response
converter.convert_text("# Test")
await 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"
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_conversion_options_passed_to_api(self, mock_post, config):
async def test_conversion_options_passed_to_api(self, mock_post, config):
"""Test that conversion options are passed to docling-serve API."""
config.processing.conversion_options.do_ocr = False
config.processing.conversion_options.force_ocr = True
@ -254,7 +259,7 @@ class TestDoclingServeConverter:
}
mock_post.return_value = mock_response
converter.convert_text("# Test")
await converter.convert_text("# Test")
call_kwargs = mock_post.call_args.kwargs
assert "data" in call_kwargs
@ -268,28 +273,31 @@ class TestDoclingServeConverter:
assert data["do_table_structure"] is False
assert data["images_scale"] == 3.0
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_text_connection_error(self, mock_post, converter):
async 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")
await converter.convert_text("# Test")
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_text_timeout_error(self, mock_post, converter):
async 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")
await converter.convert_text("# Test")
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_text_auth_error(self, mock_post, converter):
async def test_convert_text_auth_error(self, mock_post, converter):
"""Test handling of authentication errors."""
import requests
@ -301,10 +309,11 @@ class TestDoclingServeConverter:
mock_post.return_value = mock_response
with pytest.raises(ValueError, match="Authentication failed"):
converter.convert_text("# Test")
await converter.convert_text("# Test")
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_text_no_json_content(self, mock_post, converter):
async 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
@ -315,10 +324,11 @@ class TestDoclingServeConverter:
mock_post.return_value = mock_response
with pytest.raises(ValueError, match="did not return JSON content"):
converter.convert_text("# Test")
await converter.convert_text("# Test")
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_file_pdf(self, mock_post, converter):
async def test_convert_file_pdf(self, mock_post, converter):
"""Test converting PDF file via docling-serve."""
mock_response = Mock()
mock_response.status_code = 200
@ -332,13 +342,14 @@ class TestDoclingServeConverter:
f.write(b"fake pdf content")
f.flush()
temp_path = Path(f.name)
doc = converter.convert_file(temp_path)
doc = await converter.convert_file(temp_path)
assert isinstance(doc, DoclingDocument)
mock_post.assert_called_once()
@pytest.mark.asyncio
@patch("haiku.rag.converters.docling_serve.requests.post")
def test_convert_file_text(self, mock_post, converter):
async 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
@ -352,7 +363,7 @@ class TestDoclingServeConverter:
f.write("def hello():\n pass")
f.flush()
temp_path = Path(f.name)
doc = converter.convert_file(temp_path)
doc = await converter.convert_file(temp_path)
assert isinstance(doc, DoclingDocument)
# Should call docling-serve for conversion
@ -382,20 +393,22 @@ class TestDoclingServeConverterIntegration:
"""Create converter for integration tests."""
return DoclingServeConverter(config)
def test_convert_text_real_service(self, converter):
@pytest.mark.asyncio
async 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.")
doc = await 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):
@pytest.mark.asyncio
async 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)
doc = await converter.convert_file(temp_path)
assert isinstance(doc, DoclingDocument)
result = doc.export_to_markdown()

2
tests/test_preprocessor.py
View file

@ -64,7 +64,7 @@ def add_marker(text: str) -> str:
chunk_repo.embedder.embed = fake_embed # type: ignore[assignment]
converter = get_converter(Config)
docling = converter.convert_text(document.content, name="test.md")
docling = await converter.convert_text(document.content, name="test.md")
chunks = await chunk_repo.create_chunks_for_document(created_doc.id, docling)
assert any(marker in c.content for c in chunks)

7
tests/test_reader.py
View file

@ -1,11 +1,14 @@
import tempfile
from pathlib import Path
import pytest
from haiku.rag.config import Config
from haiku.rag.converters import get_converter
def test_code_file_wrapped_in_code_block():
@pytest.mark.asyncio
async def test_code_file_wrapped_in_code_block():
"""Test that code files are wrapped in markdown code blocks."""
python_code = '''def hello_world():
print("Hello, World!")
@ -17,7 +20,7 @@ def test_code_file_wrapped_in_code_block():
temp_path = Path(f.name)
converter = get_converter(Config)
document = converter.convert_file(temp_path)
document = await converter.convert_file(temp_path)
result = document.export_to_markdown()
assert result.startswith("```\n")

25
tests/test_utils.py
View file

@ -15,12 +15,13 @@ HAS_GROQ = importlib.util.find_spec("groq") is not None
HAS_BEDROCK = importlib.util.find_spec("botocore") is not None
def test_text_to_docling_document():
@pytest.mark.asyncio
async def test_text_to_docling_document():
"""Test text to DoclingDocument conversion."""
# Test basic text conversion
simple_text = "This is a simple text document."
converter = get_converter(Config)
doc = converter.convert_text(simple_text)
doc = await converter.convert_text(simple_text)
# Verify it returns a DoclingDocument
from docling_core.types.doc.document import DoclingDocument
@ -32,7 +33,8 @@ def test_text_to_docling_document():
assert "This is a simple text document." in markdown
def test_text_to_docling_document_with_custom_name():
@pytest.mark.asyncio
async def test_text_to_docling_document_with_custom_name():
"""Test text to DoclingDocument conversion with custom name parameter."""
code_text = """# Python Code
@ -43,7 +45,7 @@ def hello():
```"""
converter = get_converter(Config)
doc = converter.convert_text(code_text, name="hello.md")
doc = await converter.convert_text(code_text, name="hello.md")
# Verify it's a valid DoclingDocument
from docling_core.types.doc.document import DoclingDocument
@ -56,7 +58,8 @@ def hello():
assert "Hello, World!" in markdown
def test_text_to_docling_document_markdown_content():
@pytest.mark.asyncio
async def test_text_to_docling_document_markdown_content():
"""Test text to DoclingDocument conversion with markdown content."""
markdown_text = """# Test Document
@ -75,7 +78,7 @@ def test():
**Bold text** and *italic text*."""
converter = get_converter(Config)
doc = converter.convert_text(markdown_text, name="test.md")
doc = await converter.convert_text(markdown_text, name="test.md")
# Verify it's a DoclingDocument
from docling_core.types.doc.document import DoclingDocument
@ -89,10 +92,11 @@ def test():
assert "def test():" in result_markdown
def test_text_to_docling_document_empty_content():
@pytest.mark.asyncio
async def test_text_to_docling_document_empty_content():
"""Test text to DoclingDocument conversion with empty content."""
converter = get_converter(Config)
doc = converter.convert_text("")
doc = await converter.convert_text("")
# Should still create a valid DoclingDocument
from docling_core.types.doc.document import DoclingDocument
@ -104,7 +108,8 @@ def test_text_to_docling_document_empty_content():
assert isinstance(markdown, str)
def test_text_to_docling_document_unicode_content():
@pytest.mark.asyncio
async def test_text_to_docling_document_unicode_content():
"""Test text to DoclingDocument conversion with unicode content."""
unicode_text = """# 测试文档
@ -120,7 +125,7 @@ function saludar() {
Emoji test: 🚀 ✅ 📝"""
converter = get_converter(Config)
doc = converter.convert_text(unicode_text, name="unicode.md")
doc = await converter.convert_text(unicode_text, name="unicode.md")
# Verify it's a DoclingDocument
from docling_core.types.doc.document import DoclingDocument

10
tests/test_versioning.py
View file

@ -36,7 +36,7 @@ async def test_version_rollback_on_create_failure(temp_db_path):
content = "Hello, rollback!"
doc = Document(content=content)
converter = get_converter(Config)
dl_doc = converter.convert_text(content, name="test.md")
dl_doc = await converter.convert_text(content, name="test.md")
with pytest.raises(RuntimeError):
await repo._create_and_chunk(doc, dl_doc)
@ -68,7 +68,7 @@ async def test_version_rollback_on_update_failure(temp_db_path):
base_content = "Base content"
base_doc = Document(content=base_content)
converter = get_converter(Config)
base_dl = converter.convert_text(base_content, name="base.md")
base_dl = await converter.convert_text(base_content, name="base.md")
created = await repo._create_and_chunk(base_doc, base_dl)
# Force new chunk creation to fail during update after writing
@ -83,7 +83,7 @@ async def test_version_rollback_on_update_failure(temp_db_path):
# Attempt update
updated_content = "Updated content"
created.content = updated_content
updated_dl = converter.convert_text(updated_content, name="updated.md")
updated_dl = await converter.convert_text(updated_content, name="updated.md")
with pytest.raises(RuntimeError):
await repo._update_and_rechunk(created, updated_dl)
@ -145,12 +145,12 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
# Create first document
converter = get_converter(Config)
doc1 = Document(content="First document")
dl_doc1 = converter.convert_text("First document", name="doc1.md")
dl_doc1 = await converter.convert_text("First document", name="doc1.md")
await repo._create_and_chunk(doc1, dl_doc1)
# Create second document
doc2 = Document(content="Second document")
dl_doc2 = converter.convert_text("Second document", name="doc2.md")
dl_doc2 = await converter.convert_text("Second document", name="doc2.md")
await repo._create_and_chunk(doc2, dl_doc2)
# Get initial version counts (should have multiple versions from creates)