diff --git a/src/haiku/rag/utils.py b/src/haiku/rag/utils.py index f5f2b5eb..3cde88b3 100644 --- a/src/haiku/rag/utils.py +++ b/src/haiku/rag/utils.py @@ -1,8 +1,12 @@ import sys from importlib import metadata +from io import BytesIO from pathlib import Path import httpx +from docling.document_converter import DocumentConverter +from docling_core.types.doc.document import DoclingDocument +from docling_core.types.io import DocumentStream from packaging.version import Version, parse @@ -77,3 +81,20 @@ async def is_up_to_date() -> tuple[bool, Version, Version]: # If no network connection, do not raise alarms. pypi_version = running_version return running_version >= pypi_version, running_version, pypi_version + + +def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocument: + """Convert text content to a DoclingDocument. + + Args: + text: The text content to convert. + name: The name to use for the document stream (defaults to "content.md"). + + Returns: + A DoclingDocument created from the text content. + """ + bytes_io = BytesIO(text.encode("utf-8")) + doc_stream = DocumentStream(name=name, stream=bytes_io) + converter = DocumentConverter() + result = converter.convert(doc_stream) + return result.document diff --git a/tests/test_utils.py b/tests/test_utils.py index cd5b742c..51201b89 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,4 +1,8 @@ -from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int +from haiku.rag.utils import ( + int_to_semantic_version, + semantic_version_to_int, + text_to_docling_document, +) def test_sqlite_user_version(): @@ -13,3 +17,117 @@ def test_sqlite_user_version(): version = "255.255.255" assert semantic_version_to_int(version) == 16777215 assert int_to_semantic_version(16777215) == version + + +def test_text_to_docling_document(): + """Test the text_to_docling_document utility function.""" + # Test basic text conversion + simple_text = "This is a simple text document." + doc = text_to_docling_document(simple_text) + + # Verify it returns a DoclingDocument + from docling_core.types.doc.document import DoclingDocument + + assert isinstance(doc, DoclingDocument) + + # Verify the content can be exported back to markdown + markdown = doc.export_to_markdown() + assert "This is a simple text document." in markdown + + +def test_text_to_docling_document_with_custom_name(): + """Test text_to_docling_document with custom name parameter.""" + code_text = """# Python Code + +```python +def hello(): + print("Hello, World!") + return True +```""" + + doc = text_to_docling_document(code_text, name="hello.md") + + # Verify it's a valid DoclingDocument + from docling_core.types.doc.document import DoclingDocument + + assert isinstance(doc, DoclingDocument) + + # Verify the content is preserved + markdown = doc.export_to_markdown() + assert "def hello():" in markdown + assert "Hello, World!" in markdown + + +def test_text_to_docling_document_markdown_content(): + """Test text_to_docling_document with markdown content.""" + markdown_text = """# Test Document + +This is a test document with: + +- List item 1 +- List item 2 + +## Code Example + +```python +def test(): + return "Hello" +``` + +**Bold text** and *italic text*.""" + + doc = text_to_docling_document(markdown_text, name="test.md") + + # Verify it's a DoclingDocument + from docling_core.types.doc.document import DoclingDocument + + assert isinstance(doc, DoclingDocument) + + # Verify the markdown structure is preserved + result_markdown = doc.export_to_markdown() + assert "# Test Document" in result_markdown + assert "List item 1" in result_markdown + assert "def test():" in result_markdown + + +def test_text_to_docling_document_empty_content(): + """Test text_to_docling_document with empty content.""" + doc = text_to_docling_document("") + + # Should still create a valid DoclingDocument + from docling_core.types.doc.document import DoclingDocument + + assert isinstance(doc, DoclingDocument) + + # Export should work even with empty content + markdown = doc.export_to_markdown() + assert isinstance(markdown, str) + + +def test_text_to_docling_document_unicode_content(): + """Test text_to_docling_document with unicode content.""" + unicode_text = """# 测试文档 + +这是一个包含中文的测试文档。 + +## Código en Español +```javascript +function saludar() { + return "¡Hola mundo!"; +} +``` + +Emoji test: 🚀 ✅ 📝""" + + doc = text_to_docling_document(unicode_text, name="unicode.md") + + # Verify it's a DoclingDocument + from docling_core.types.doc.document import DoclingDocument + + assert isinstance(doc, DoclingDocument) + + # Verify unicode content is preserved + result_markdown = doc.export_to_markdown() + assert "测试文档" in result_markdown + assert "¡Hola mundo!" in result_markdown + assert "🚀" in result_markdown