Make FileReader return a DoclingDocument instead of a md string

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})

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
+                # Fallback: try to read as text and convert to DoclingDocument
-                return path.read_text(encoding="utf-8")
+                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}")

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