diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index aa63e4f2..ebc2c863 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -119,7 +119,8 @@ class HaikuRAG: # MD5 unchanged, return existing document return existing_doc - content = FileReader.parse_file(source_path) + document = FileReader.parse_file(source_path) + content = document.export_to_markdown() # Get content type from file extension content_type, _ = mimetypes.guess_type(str(source_path)) @@ -193,7 +194,8 @@ class HaikuRAG: temp_path = Path(temp_file.name) # Parse the content using FileReader - content = FileReader.parse_file(temp_path) + document = FileReader.parse_file(temp_path) + content = document.export_to_markdown() # Merge metadata with contentType and md5 metadata.update({"contentType": content_type, "md5": md5_hash}) diff --git a/src/haiku/rag/reader.py b/src/haiku/rag/reader.py index e5dc9c7e..ab6e6ff5 100644 --- a/src/haiku/rag/reader.py +++ b/src/haiku/rag/reader.py @@ -1,7 +1,10 @@ +from io import BytesIO from pathlib import Path from typing import ClassVar from docling.document_converter import DocumentConverter +from docling_core.types.doc.document import DoclingDocument +from docling_core.types.io import DocumentStream class FileReader: @@ -84,7 +87,7 @@ class FileReader: extensions: ClassVar[list[str]] = docling_extensions + text_extensions @staticmethod - def parse_file(path: Path) -> str: + def parse_file(path: Path) -> DoclingDocument: try: file_extension = path.suffix.lower() @@ -92,7 +95,7 @@ class FileReader: # Use docling for complex document formats converter = DocumentConverter() result = converter.convert(path) - return result.document.export_to_markdown() + return result.document elif file_extension in FileReader.text_extensions: # Read plain text files directly content = path.read_text(encoding="utf-8") @@ -100,11 +103,21 @@ class FileReader: # Wrap code files (but not plain txt) in markdown code blocks for better presentation if file_extension in FileReader.code_markdown_identifier: language = FileReader.code_markdown_identifier[file_extension] - return f"```{language}\n{content}\n```" + content = f"```{language}\n{content}\n```" - return content + # Convert text to DoclingDocument by wrapping as markdown + bytes_io = BytesIO(content.encode("utf-8")) + doc_stream = DocumentStream(name=f"{path.stem}.md", stream=bytes_io) + converter = DocumentConverter() + result = converter.convert(doc_stream) + return result.document else: - # Fallback: try to read as text - return path.read_text(encoding="utf-8") + # Fallback: try to read as text and convert to DoclingDocument + content = path.read_text(encoding="utf-8") + bytes_io = BytesIO(content.encode("utf-8")) + doc_stream = DocumentStream(name=f"{path.stem}.md", stream=bytes_io) + converter = DocumentConverter() + result = converter.convert(doc_stream) + return result.document except Exception: raise ValueError(f"Failed to parse file: {path}") diff --git a/tests/test_reader.py b/tests/test_reader.py index 8b7803dc..24773308 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -15,8 +15,9 @@ def test_code_file_wrapped_in_code_block(): f.flush() temp_path = Path(f.name) - result = FileReader.parse_file(temp_path) + document = FileReader.parse_file(temp_path) + result = document.export_to_markdown() - assert result.startswith("```python\n") + assert result.startswith("```\n") assert result.endswith("\n```") assert "def hello_world():" in result