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 hashlib
|
||||||
|
import logging
|
||||||
import mimetypes
|
import mimetypes
|
||||||
import tempfile
|
import tempfile
|
||||||
from collections.abc import AsyncGenerator
|
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.store.repositories.settings import SettingsRepository
|
||||||
from haiku.rag.utils import text_to_docling_document
|
from haiku.rag.utils import text_to_docling_document
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
class HaikuRAG:
|
class HaikuRAG:
|
||||||
"""High-level haiku-rag client."""
|
"""High-level haiku-rag client."""
|
||||||
|
|
@ -538,8 +541,8 @@ class HaikuRAG:
|
||||||
"""Rebuild the database by deleting all chunks and re-indexing all documents.
|
"""Rebuild the database by deleting all chunks and re-indexing all documents.
|
||||||
|
|
||||||
For documents with URIs:
|
For documents with URIs:
|
||||||
- Deletes the document and re-adds it from source if source exists
|
- Re-adds from source if source exists
|
||||||
- Skips documents where source no longer exists
|
- Re-embeds from existing content if source is missing
|
||||||
|
|
||||||
For documents without URIs:
|
For documents without URIs:
|
||||||
- Re-creates chunks from existing content
|
- Re-creates chunks from existing content
|
||||||
|
|
@ -559,29 +562,51 @@ class HaikuRAG:
|
||||||
for doc in documents:
|
for doc in documents:
|
||||||
assert doc.id is not None, "Document ID should not be None"
|
assert doc.id is not None, "Document ID should not be None"
|
||||||
if doc.uri:
|
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:
|
try:
|
||||||
# Delete the old document first
|
if parsed_url.scheme == "file":
|
||||||
await self.delete_document(doc.id)
|
# 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)
|
if source_accessible:
|
||||||
new_doc = await self.create_document_from_source(
|
# Source exists - delete and recreate from source
|
||||||
source=doc.uri, metadata=doc.metadata or {}
|
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
|
||||||
)
|
)
|
||||||
|
docling_document = text_to_docling_document(doc.content)
|
||||||
assert new_doc.id is not None, "New document ID should not be None"
|
await self.chunk_repository.create_chunks_for_document(
|
||||||
yield new_doc.id
|
doc.id, docling_document
|
||||||
|
|
||||||
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
|
yield doc.id
|
||||||
else:
|
else:
|
||||||
# Document without URI - re-create chunks from existing content
|
# Document without URI - re-create chunks from existing content
|
||||||
docling_document = text_to_docling_document(doc.content)
|
docling_document = text_to_docling_document(doc.content)
|
||||||
|
|
|
||||||
|
|
@ -1,13 +1,13 @@
|
||||||
|
import logging
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from watchfiles import Change, DefaultFilter, awatch
|
from watchfiles import Change, DefaultFilter, awatch
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.logging import get_logger
|
|
||||||
from haiku.rag.reader import FileReader
|
from haiku.rag.reader import FileReader
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
||||||
logger = get_logger()
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
class FileFilter(DefaultFilter):
|
class FileFilter(DefaultFilter):
|
||||||
|
|
|
||||||
|
|
@ -47,3 +47,40 @@ async def test_rebuild_database(qa_corpus: Dataset, temp_db_path):
|
||||||
chunks_after.extend(doc_chunks)
|
chunks_after.extend(doc_chunks)
|
||||||
|
|
||||||
assert len(chunks_after) > 0
|
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