diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 9cbb6e8d..1c4377ed 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -1,4 +1,5 @@ import hashlib +import logging import mimetypes import tempfile from collections.abc import AsyncGenerator @@ -18,6 +19,8 @@ from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.settings import SettingsRepository from haiku.rag.utils import text_to_docling_document +logger = logging.getLogger(__name__) + class HaikuRAG: """High-level haiku-rag client.""" @@ -538,8 +541,8 @@ class HaikuRAG: """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 + - Re-adds from source if source exists + - Re-embeds from existing content if source is missing For documents without URIs: - Re-creates chunks from existing content @@ -559,29 +562,51 @@ class HaikuRAG: for doc in documents: 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 + # Document has a URI - check if source is accessible + source_accessible = False + parsed_url = urlparse(doc.uri) + try: - # Delete the old document first - await self.delete_document(doc.id) + if parsed_url.scheme == "file": + # Check if file exists + source_path = Path(parsed_url.path) + source_accessible = source_path.exists() + elif parsed_url.scheme in ("http", "https"): + # For URLs, we'll try to create and catch errors + source_accessible = True + else: + source_accessible = False + except Exception: + source_accessible = False - # Try to re-create from source (this creates the document with chunks) - new_doc = await self.create_document_from_source( - source=doc.uri, metadata=doc.metadata or {} + if source_accessible: + # Source exists - delete and recreate from source + try: + await self.delete_document(doc.id) + new_doc = await self.create_document_from_source( + source=doc.uri, metadata=doc.metadata or {} + ) + assert new_doc.id is not None, ( + "New document ID should not be None" + ) + yield new_doc.id + except Exception as e: + logger.error( + "Error recreating document from source %s: %s", + doc.uri, + e, + ) + continue + else: + # Source missing - re-embed from existing content + logger.warning( + "Source missing for %s, re-embedding from content", doc.uri ) - - 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}" + docling_document = text_to_docling_document(doc.content) + await self.chunk_repository.create_chunks_for_document( + doc.id, docling_document ) - continue + yield doc.id else: # Document without URI - re-create chunks from existing content docling_document = text_to_docling_document(doc.content) diff --git a/src/haiku/rag/monitor.py b/src/haiku/rag/monitor.py index 618bb32a..713c5834 100644 --- a/src/haiku/rag/monitor.py +++ b/src/haiku/rag/monitor.py @@ -1,13 +1,13 @@ +import logging from pathlib import Path from watchfiles import Change, DefaultFilter, awatch from haiku.rag.client import HaikuRAG -from haiku.rag.logging import get_logger from haiku.rag.reader import FileReader from haiku.rag.store.models.document import Document -logger = get_logger() +logger = logging.getLogger(__name__) class FileFilter(DefaultFilter): diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 888bb8bc..0ec81e75 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -47,3 +47,40 @@ async def test_rebuild_database(qa_corpus: Dataset, temp_db_path): chunks_after.extend(doc_chunks) assert len(chunks_after) > 0 + + +@pytest.mark.asyncio +async def test_rebuild_with_missing_source(qa_corpus: Dataset, temp_db_path): + """Test rebuild functionality when document source is missing.""" + async with HaikuRAG(temp_db_path) as client: + # Create document with content + content = qa_corpus["document_extracted"][0] + doc = await client.create_document(content=content) + + # Manually set a URI that doesn't exist + assert doc.id is not None + doc_with_uri = await client.document_repository.get_by_id(doc.id) + assert doc_with_uri is not None + doc_with_uri.uri = "file:///nonexistent/path.txt" + await client.document_repository.update(doc_with_uri) + + # Verify chunks exist before rebuild + chunks_before = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks_before) > 0 + + # Perform rebuild + processed_doc_ids = [] + async for doc_id in client.rebuild_database(): + processed_doc_ids.append(doc_id) + + # Document should still be processed (not skipped) + assert doc.id in processed_doc_ids + + # Verify document still exists + doc_after = await client.document_repository.get_by_id(doc.id) + assert doc_after is not None + assert doc_after.content == content + + # Verify chunks were recreated from content + chunks_after = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks_after) > 0