When rebuilding if a document does not exist anymore, re-embed the chunks instead of failing
This commit is contained in:
parent
43b9cd50ba
commit
05c382fcb5
3 changed files with 85 additions and 23 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue