diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 8b2d3e25..81513c05 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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 ) diff --git a/haiku_rag_slim/haiku/rag/converters/base.py b/haiku_rag_slim/haiku/rag/converters/base.py index 10baa87e..713cbd94 100644 --- a/haiku_rag_slim/haiku/rag/converters/base.py +++ b/haiku_rag_slim/haiku/rag/converters/base.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/converters/docling_local.py b/haiku_rag_slim/haiku/rag/converters/docling_local.py index fc0dc6f4..a1ebfbc8 100644 --- a/haiku_rag_slim/haiku/rag/converters/docling_local.py +++ b/haiku_rag_slim/haiku/rag/converters/docling_local.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/converters/docling_serve.py b/haiku_rag_slim/haiku/rag/converters/docling_serve.py index d5b9af50..3b75b139 100644 --- a/haiku_rag_slim/haiku/rag/converters/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/converters/docling_serve.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/converters/text_utils.py b/haiku_rag_slim/haiku/rag/converters/text_utils.py index c8b90cd6..a149297c 100644 --- a/haiku_rag_slim/haiku/rag/converters/text_utils.py +++ b/haiku_rag_slim/haiku/rag/converters/text_utils.py @@ -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}") diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index dd251e6a..33623dd5 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -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: diff --git a/tests/test_chunk.py b/tests/test_chunk.py index ddcdf6af..08cf5984 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -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) diff --git a/tests/test_chunker.py b/tests/test_chunker.py index 205dd75c..7613560c 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -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) diff --git a/tests/test_converters.py b/tests/test_converters.py index d32902ec..27ed9632 100644 --- a/tests/test_converters.py +++ b/tests/test_converters.py @@ -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() diff --git a/tests/test_preprocessor.py b/tests/test_preprocessor.py index 42a36625..9f7d9b07 100644 --- a/tests/test_preprocessor.py +++ b/tests/test_preprocessor.py @@ -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) diff --git a/tests/test_reader.py b/tests/test_reader.py index cac67d70..34d0637c 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -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") diff --git a/tests/test_utils.py b/tests/test_utils.py index d40f5d17..f81b76b4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -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 diff --git a/tests/test_versioning.py b/tests/test_versioning.py index e702f895..e38986b6 100644 --- a/tests/test_versioning.py +++ b/tests/test_versioning.py @@ -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)