From b92f25d6fbf224f5b30082b5f1355fbce35ca765 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 8 Aug 2025 14:45:24 +0200 Subject: [PATCH 1/5] Make FileReader return a DoclingDocument instead of a md string --- src/haiku/rag/client.py | 6 ++++-- src/haiku/rag/reader.py | 25 +++++++++++++++++++------ tests/test_reader.py | 5 +++-- 3 files changed, 26 insertions(+), 10 deletions(-) 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 From 6a3c85e6aeb7e2e567cc3b4261db9842a755df21 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 8 Aug 2025 15:38:57 +0200 Subject: [PATCH 2/5] text_to_docling_document() util --- src/haiku/rag/utils.py | 21 ++++++++ tests/test_utils.py | 120 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 140 insertions(+), 1 deletion(-) 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 From ad4b655a313d60d9a01e71135c41326ffa7c7259 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 8 Aug 2025 17:24:34 +0200 Subject: [PATCH 3/5] Operate on DoclingDocument instead of Document --- src/haiku/rag/chunker.py | 21 ++----- src/haiku/rag/client.py | 66 +++++++++++++++----- src/haiku/rag/reader.py | 16 ++--- src/haiku/rag/store/repositories/chunk.py | 8 ++- src/haiku/rag/store/repositories/document.py | 36 ++++++++--- tests/test_chunk.py | 6 +- tests/test_chunker.py | 11 +++- 7 files changed, 107 insertions(+), 57 deletions(-) diff --git a/src/haiku/rag/chunker.py b/src/haiku/rag/chunker.py index ab29f749..31cb2717 100644 --- a/src/haiku/rag/chunker.py +++ b/src/haiku/rag/chunker.py @@ -1,11 +1,9 @@ -from io import BytesIO from typing import ClassVar import tiktoken from docling.chunking import HybridChunker # type: ignore -from docling.document_converter import DocumentConverter from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer -from docling_core.types.io import DocumentStream +from docling_core.types.doc.document import DoclingDocument from haiku.rag.config import Config @@ -33,27 +31,20 @@ class Chunker: self.chunker = HybridChunker(tokenizer=tokenizer) # type: ignore - async def chunk(self, text: str) -> list[str]: - """Split the text into chunks using docling's structure-aware chunking. + async def chunk(self, document: DoclingDocument) -> list[str]: + """Split the document into chunks using docling's structure-aware chunking. Args: - text: The text to be split into chunks. + document: The DoclingDocument to be split into chunks. Returns: A list of text chunks with semantic boundaries. """ - if not text: + if document is None: return [] - # Convert to docling document - bytes_io = BytesIO(text.encode("utf-8")) - doc_stream = DocumentStream(name="text.md", stream=bytes_io) - converter = DocumentConverter() - result = converter.convert(doc_stream) - doc = result.document - # Chunk using docling's hybrid chunker - chunks = list(self.chunker.chunk(doc)) + chunks = list(self.chunker.chunk(document)) return [self.chunker.contextualize(chunk) for chunk in chunks] diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index ebc2c863..58e815df 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -16,6 +16,7 @@ from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.utils import text_to_docling_document class HaikuRAG: @@ -49,6 +50,24 @@ class HaikuRAG: self.close() return False + async def _create_document_with_docling( + self, + docling_document, + uri: str | None = None, + metadata: dict | None = None, + chunks: list[Chunk] | None = None, + ) -> Document: + """Create a new document from DoclingDocument.""" + content = docling_document.export_to_markdown() + document = Document( + content=content, + uri=uri, + metadata=metadata or {}, + ) + return await self.document_repository._create_with_docling( + document, docling_document, chunks + ) + async def create_document( self, content: str, @@ -67,12 +86,17 @@ class HaikuRAG: Returns: The created Document instance. """ + # Convert content to DoclingDocument for processing + docling_document = text_to_docling_document(content) + document = Document( content=content, uri=uri, metadata=metadata or {}, ) - return await self.document_repository.create(document, chunks) + return await self.document_repository._create_with_docling( + document, docling_document, chunks + ) async def create_document_from_source( self, source: str | Path, metadata: dict = {} @@ -119,8 +143,7 @@ class HaikuRAG: # MD5 unchanged, return existing document return existing_doc - document = FileReader.parse_file(source_path) - content = document.export_to_markdown() + docling_document = FileReader.parse_file(source_path) # Get content type from file extension content_type, _ = mimetypes.guess_type(str(source_path)) @@ -132,13 +155,15 @@ class HaikuRAG: if existing_doc: # Update existing document - existing_doc.content = content + existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata - return await self.update_document(existing_doc) + return await self.document_repository._update_with_docling( + existing_doc, docling_document + ) else: - # Create new document - return await self.create_document( - content=content, uri=uri, metadata=metadata + # Create new document using DoclingDocument + return await self._create_document_with_docling( + docling_document=docling_document, uri=uri, metadata=metadata ) async def _create_or_update_document_from_url( @@ -194,19 +219,20 @@ class HaikuRAG: temp_path = Path(temp_file.name) # Parse the content using FileReader - document = FileReader.parse_file(temp_path) - content = document.export_to_markdown() + docling_document = FileReader.parse_file(temp_path) # Merge metadata with contentType and md5 metadata.update({"contentType": content_type, "md5": md5_hash}) if existing_doc: - existing_doc.content = content + existing_doc.content = docling_document.export_to_markdown() existing_doc.metadata = metadata - return await self.update_document(existing_doc) + return await self.document_repository._update_with_docling( + existing_doc, docling_document + ) else: - return await self.create_document( - content=content, uri=url, metadata=metadata + return await self._create_document_with_docling( + docling_document=docling_document, uri=url, metadata=metadata ) def _get_extension_from_content_type_or_url( @@ -264,7 +290,12 @@ class HaikuRAG: async def update_document(self, document: Document) -> Document: """Update an existing document.""" - return await self.document_repository.update(document) + # Convert content to DoclingDocument + docling_document = text_to_docling_document(document.content) + + return await self.document_repository._update_with_docling( + document, docling_document + ) async def delete_document(self, document_id: int) -> bool: """Delete a document by its ID.""" @@ -346,8 +377,11 @@ class HaikuRAG: for doc in documents: if doc.id is not None: + # Convert content to DoclingDocument for rebuild + docling_document = text_to_docling_document(doc.content) + await self.chunk_repository.create_chunks_for_document( - doc.id, doc.content, commit=False + doc.id, docling_document, commit=False ) yield doc.id diff --git a/src/haiku/rag/reader.py b/src/haiku/rag/reader.py index ab6e6ff5..cd2c6128 100644 --- a/src/haiku/rag/reader.py +++ b/src/haiku/rag/reader.py @@ -1,10 +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 + +from haiku.rag.utils import text_to_docling_document class FileReader: @@ -106,18 +106,10 @@ class FileReader: content = f"```{language}\n{content}\n```" # 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 + return text_to_docling_document(content, name=f"{path.stem}.md") else: # 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 + return text_to_docling_document(content, name=f"{path.stem}.md") except Exception: raise ValueError(f"Failed to parse file: {path}") diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 3a5c1f78..bb9e1c2b 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -1,6 +1,8 @@ import json import re +from docling_core.types.doc.document import DoclingDocument + from haiku.rag.chunker import chunker from haiku.rag.embeddings import get_embedder from haiku.rag.store.models.chunk import Chunk @@ -197,11 +199,11 @@ class ChunkRepository(BaseRepository[Chunk]): ] async def create_chunks_for_document( - self, document_id: int, content: str, commit: bool = True + self, document_id: int, document: DoclingDocument, commit: bool = True ) -> list[Chunk]: - """Create chunks and embeddings for a document.""" + """Create chunks and embeddings for a document from DoclingDocument.""" # Chunk the document content - chunk_texts = await chunker.chunk(content) + chunk_texts = await chunker.chunk(document) created_chunks = [] # Create chunks with embeddings using the create method diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index dd5c17c8..750b8d3c 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/src/haiku/rag/store/repositories/document.py @@ -1,8 +1,11 @@ import json from typing import TYPE_CHECKING +from docling_core.types.doc.document import DoclingDocument + from haiku.rag.store.models.document import Document from haiku.rag.store.repositories.base import BaseRepository +from haiku.rag.utils import text_to_docling_document if TYPE_CHECKING: from haiku.rag.store.models.chunk import Chunk @@ -20,8 +23,11 @@ class DocumentRepository(BaseRepository[Document]): chunk_repository = ChunkRepository(store) self.chunk_repository = chunk_repository - async def create( - self, entity: Document, chunks: list["Chunk"] | None = None + async def _create_with_docling( + self, + entity: Document, + docling_document: DoclingDocument, + chunks: list["Chunk"] | None = None, ) -> Document: """Create a document with its chunks and embeddings.""" if self.store._connection is None: @@ -62,9 +68,9 @@ class DocumentRepository(BaseRepository[Document]): chunk.metadata["order"] = order await self.chunk_repository.create(chunk, commit=False) else: - # Create chunks and embeddings using ChunkRepository + # Create chunks and embeddings using DoclingDocument await self.chunk_repository.create_chunks_for_document( - document_id, entity.content, commit=False + document_id, docling_document, commit=False ) cursor.execute("COMMIT") @@ -74,6 +80,13 @@ class DocumentRepository(BaseRepository[Document]): cursor.execute("ROLLBACK") raise + async def create(self, entity: Document) -> Document: + """Create a document with its chunks and embeddings.""" + # Convert content to DoclingDocument + docling_document = text_to_docling_document(entity.content) + + return await self._create_with_docling(entity, docling_document) + async def get_by_id(self, entity_id: int) -> Document | None: """Get a document by its ID.""" if self.store._connection is None: @@ -134,7 +147,9 @@ class DocumentRepository(BaseRepository[Document]): updated_at=updated_at, ) - async def update(self, entity: Document) -> Document: + async def _update_with_docling( + self, entity: Document, docling_document: DoclingDocument + ) -> Document: """Update an existing document and regenerate its chunks and embeddings.""" if self.store._connection is None: raise ValueError("Store connection is not available") @@ -163,10 +178,10 @@ class DocumentRepository(BaseRepository[Document]): }, ) - # Delete existing chunks and regenerate using ChunkRepository + # Delete existing chunks and regenerate using DoclingDocument await self.chunk_repository.delete_by_document_id(entity.id, commit=False) await self.chunk_repository.create_chunks_for_document( - entity.id, entity.content, commit=False + entity.id, docling_document, commit=False ) cursor.execute("COMMIT") @@ -176,6 +191,13 @@ class DocumentRepository(BaseRepository[Document]): cursor.execute("ROLLBACK") raise + async def update(self, entity: Document) -> Document: + """Update an existing document and regenerate its chunks and embeddings.""" + # Convert content to DoclingDocument + docling_document = text_to_docling_document(entity.content) + + return await self._update_with_docling(entity, docling_document) + async def delete(self, entity_id: int) -> bool: """Delete a document and all its associated chunks and embeddings.""" # Delete chunks and embeddings first diff --git a/tests/test_chunk.py b/tests/test_chunk.py index da0216a4..902dbdf3 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -6,6 +6,7 @@ from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.document import DocumentRepository +from haiku.rag.utils import text_to_docling_document @pytest.mark.asyncio @@ -77,8 +78,11 @@ async def test_create_chunks_for_document(qa_corpus: Dataset): assert document_id is not None, "Document ID should not be None" + # Convert text to DoclingDocument + docling_document = text_to_docling_document(document_text, name="test.md") + # Test creating chunks for the document - chunks = await chunk_repo.create_chunks_for_document(document_id, document_text) + chunks = await chunk_repo.create_chunks_for_document(document_id, docling_document) # Verify chunks were created assert len(chunks) > 0 diff --git a/tests/test_chunker.py b/tests/test_chunker.py index 8afe44de..759d07b1 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -2,13 +2,18 @@ import pytest from datasets import Dataset from haiku.rag.chunker import Chunker +from haiku.rag.utils import text_to_docling_document @pytest.mark.asyncio async def test_chunker(qa_corpus: Dataset): chunker = Chunker() - doc = qa_corpus[0]["document_extracted"] - chunks = await Chunker().chunk(doc) + doc_text = qa_corpus[0]["document_extracted"] + + # Convert text to DoclingDocument + doc = text_to_docling_document(doc_text, name="test.md") + + chunks = await chunker.chunk(doc) # Ensure that the text is split into multiple chunks assert len(chunks) > 1 @@ -27,7 +32,7 @@ async def test_chunker(qa_corpus: Dataset): assert token_count > 5 # Ensure chunks aren't too small # Ensure that all chunks together contain roughly the same content as original - original_tokens = len(Chunker.encoder.encode(doc, disallowed_special=())) + original_tokens = len(Chunker.encoder.encode(doc_text, disallowed_special=())) # Due to structure-aware chunking, we might have some variation in token count # but it should be reasonable From 22ed9ca51e2191889c22e554ee933dc638c04a23 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 8 Aug 2025 17:31:42 +0200 Subject: [PATCH 4/5] Disable python warning unless in development --- src/haiku/rag/cli.py | 4 ++++ src/haiku/rag/config.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 087e0acf..ae377b1c 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -1,4 +1,5 @@ import asyncio +import warnings from importlib.metadata import version from pathlib import Path @@ -9,6 +10,9 @@ from haiku.rag.app import HaikuRAGApp from haiku.rag.config import Config from haiku.rag.utils import is_up_to_date +if not Config.ENV == "development": + warnings.filterwarnings("ignore") + cli = typer.Typer( context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True ) diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index babd3d31..3671ed19 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -10,7 +10,7 @@ load_dotenv() class AppConfig(BaseModel): - ENV: str = "development" + ENV: str = "production" DEFAULT_DATA_DIR: Path = get_default_data_dir() MONITOR_DIRECTORIES: list[Path] = [] From 0b3477c46c403afe062c0f410f784461facbeeac Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 8 Aug 2025 18:30:53 +0200 Subject: [PATCH 5/5] Rebuild now re-adds files. --- src/haiku/rag/client.py | 46 +++++++++++++++++++++++++++++++++++------ 1 file changed, 40 insertions(+), 6 deletions(-) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 58e815df..1e7ebaae 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -125,9 +125,12 @@ class HaikuRAG: parsed_url = urlparse(source_str) if parsed_url.scheme in ("http", "https"): return await self._create_or_update_document_from_url(source_str, metadata) - - # Handle as file path - source_path = Path(source) if isinstance(source, str) else source + elif parsed_url.scheme == "file": + # Handle file:// URI by converting to path + source_path = Path(parsed_url.path) + else: + # Handle as regular file path + source_path = Path(source) if isinstance(source, str) else source if source_path.suffix.lower() not in FileReader.extensions: raise ValueError(f"Unsupported file extension: {source_path.suffix}") @@ -361,6 +364,13 @@ class HaikuRAG: async def rebuild_database(self) -> AsyncGenerator[int, None]: """Rebuild the database by deleting all chunks and re-indexing all documents. + For documents with URIs: + - Deletes the document and re-adds it from source if source exists + - Skips documents where source no longer exists + + For documents without URIs: + - Re-creates chunks from existing content + Yields: int: The ID of the document currently being processed """ @@ -376,10 +386,34 @@ class HaikuRAG: documents = await self.list_documents() for doc in documents: - if doc.id is not None: - # Convert content to DoclingDocument for rebuild - docling_document = text_to_docling_document(doc.content) + assert doc.id is not None, "Document ID should not be None" + if doc.uri: + # Document has a URI - delete and try to re-add from source + try: + # Delete the old document first + await self.delete_document(doc.id) + # Try to re-create from source (this creates the document with chunks) + new_doc = await self.create_document_from_source( + doc.uri, doc.metadata or {} + ) + + assert new_doc.id is not None, "New document ID should not be None" + yield new_doc.id + + except (FileNotFoundError, ValueError, OSError) as e: + # Source doesn't exist or can't be accessed - document already deleted, skip + print(f"Skipping document with URI {doc.uri}: {e}") + continue + except Exception as e: + # Unexpected error - log it and skip + print( + f"Unexpected error processing document with URI {doc.uri}: {e}" + ) + continue + else: + # Document without URI - re-create chunks from existing content + docling_document = text_to_docling_document(doc.content) await self.chunk_repository.create_chunks_for_document( doc.id, docling_document, commit=False )