From ed604f6a1db79c45a232fc884d5ccb46b7ffd910 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 3 Sep 2025 10:50:08 +0300 Subject: [PATCH] Batch operations & small optimization --- src/haiku/rag/store/engine.py | 116 +++++++++++-------- src/haiku/rag/store/repositories/chunk.py | 72 ++++++++---- src/haiku/rag/store/repositories/document.py | 81 +++++-------- 3 files changed, 143 insertions(+), 126 deletions(-) diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index e23fbb49..db51f22e 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -1,4 +1,5 @@ import json +import logging from importlib import metadata from pathlib import Path from uuid import uuid4 @@ -10,6 +11,8 @@ from pydantic import Field from haiku.rag.config import Config from haiku.rag.embeddings import get_embedder +logger = logging.getLogger(__name__) + class DocumentRecord(LanceModel): id: str = Field(default_factory=lambda: str(uuid4())) @@ -21,14 +24,17 @@ class DocumentRecord(LanceModel): def create_chunk_model(vector_dim: int): - """Create a ChunkRecord model with the specified vector dimension.""" + """Create a ChunkRecord model with the specified vector dimension. + + This creates a model with proper vector typing for LanceDB. + """ class ChunkRecord(LanceModel): id: str = Field(default_factory=lambda: str(uuid4())) document_id: str content: str metadata: str = Field(default="{}") - vector: Vector(vector_dim) = Field(default_factory=list) # type: ignore + vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore return ChunkRecord @@ -46,27 +52,41 @@ class Store: # Create the ChunkRecord model with the correct vector dimension self.ChunkRecord = create_chunk_model(self.embedder._vector_dim) - # Connect to LanceDB (local, cloud, or object storage) - if Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION: - self.db = lancedb.connect( + # Connect to LanceDB + self.db = self._connect_to_lancedb(db_path) + + # Initialize tables + self.create_or_update_db() + + # Validate config compatibility after connection is established + if not skip_validation: + self._validate_configuration() + + def _connect_to_lancedb(self, db_path: Path): + """Establish connection to LanceDB (local, cloud, or object storage).""" + # Check if we have cloud configuration + if self._has_cloud_config(): + return lancedb.connect( uri=Config.LANCEDB_URI, api_key=Config.LANCEDB_API_KEY, region=Config.LANCEDB_REGION, ) else: # Local file system connection - self.db = lancedb.connect(db_path) + return lancedb.connect(db_path) - self.create_or_update_db() + def _has_cloud_config(self) -> bool: + """Check if cloud configuration is complete.""" + return bool( + Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION + ) - # Validate config compatibility after connection is established - if not skip_validation: - from haiku.rag.store.repositories.settings import ( - SettingsRepository, - ) + def _validate_configuration(self) -> None: + """Validate that the configuration is compatible with the database.""" + from haiku.rag.store.repositories.settings import SettingsRepository - settings_repo = SettingsRepository(self) - settings_repo.validate_config_compatibility() + settings_repo = SettingsRepository(self) + settings_repo.validate_config_compatibility() def create_or_update_db(self): """Create the database tables.""" @@ -114,54 +134,48 @@ class Store: ) if existing_settings: db_version = self.get_haiku_version() # noqa: F841 - # XXX Add upgrade logic here similar to SQLite version - + # TODO: Add upgrade logic here similar to SQLite version when needed except Exception: + # Settings table might not exist yet in fresh databases pass def get_haiku_version(self) -> str: """Returns the user version stored in settings.""" - try: - settings_records = list( - self.settings_table.search().limit(1).to_pydantic(SettingsRecord) + settings_records = list( + self.settings_table.search().limit(1).to_pydantic(SettingsRecord) + ) + if settings_records: + settings = ( + json.loads(settings_records[0].settings) + if settings_records[0].settings + else {} ) - if settings_records: - settings = ( - json.loads(settings_records[0].settings) - if settings_records[0].settings - else {} - ) - return settings.get("version", "0.0.0") - except Exception: - pass + return settings.get("version", "0.0.0") return "0.0.0" def set_haiku_version(self, version: str) -> None: """Updates the user version in settings.""" - try: - settings_records = list( - self.settings_table.search().limit(1).to_pydantic(SettingsRecord) + settings_records = list( + self.settings_table.search().limit(1).to_pydantic(SettingsRecord) + ) + if settings_records: + settings = ( + json.loads(settings_records[0].settings) + if settings_records[0].settings + else {} + ) + settings["version"] = version + # Update the record + self.settings_table.update( + where="id = 'settings'", values={"settings": json.dumps(settings)} + ) + else: + # Create new settings record + settings_data = Config.model_dump(mode="json") + settings_data["version"] = version + self.settings_table.add( + [SettingsRecord(id="settings", settings=json.dumps(settings_data))] ) - if settings_records: - settings = ( - json.loads(settings_records[0].settings) - if settings_records[0].settings - else {} - ) - settings["version"] = version - # Update the record - self.settings_table.update( - where="id = 'settings'", values={"settings": json.dumps(settings)} - ) - else: - # Create new settings record - settings_data = Config.model_dump(mode="json") - settings_data["version"] = version - self.settings_table.add( - [SettingsRecord(id="settings", settings=json.dumps(settings_data))] - ) - except Exception: - pass def recreate_embeddings_table(self) -> None: """Recreate the chunks table with current vector dimensions.""" diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index c757935e..870c697d 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -1,5 +1,6 @@ import asyncio import json +import logging from uuid import uuid4 from docling_core.types.doc.document import DoclingDocument @@ -11,6 +12,8 @@ from haiku.rag.embeddings import get_embedder from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.models.chunk import Chunk +logger = logging.getLogger(__name__) + class ChunkRepository: """Repository for Chunk operations.""" @@ -24,8 +27,9 @@ class ChunkRepository: """Ensure FTS index exists on the content column.""" try: self.store.chunks_table.create_fts_index("content", replace=True) - except Exception: - pass + except Exception as e: + # Log the error but don't fail - FTS might already exist + logger.debug(f"FTS index creation skipped: {e}") async def _optimize(self) -> None: """Optimize the chunks table to refresh indexes.""" @@ -36,9 +40,11 @@ class ChunkRepository: async with self._optimize_lock: try: self.store.chunks_table.optimize() - except (RuntimeError, OSError): + except (RuntimeError, OSError) as e: # Handle "too many open files" and other resource errors gracefully - pass + logger.debug( + f"Table optimization skipped due to resource constraints: {e}" + ) async def create(self, entity: Chunk) -> Chunk: """Create a chunk in the database.""" @@ -147,18 +153,21 @@ class ChunkRepository: ) -> list[Chunk]: """Create chunks and embeddings for a document from DoclingDocument.""" chunk_texts = await chunker.chunk(document) + + # Generate embeddings in parallel for all chunks + embeddings_tasks = [] + for chunk_text in chunk_texts: + embeddings_tasks.append(self.embedder.embed(chunk_text)) + + # Wait for all embeddings to complete + embeddings = await asyncio.gather(*embeddings_tasks) + + # Prepare all chunk records for batch insertion + chunk_records = [] created_chunks = [] - for order, chunk_text in enumerate(chunk_texts): - chunk = Chunk( - document_id=document_id, content=chunk_text, metadata={"order": order} - ) - # Use create but don't trigger individual optimizations + for order, (chunk_text, embedding) in enumerate(zip(chunk_texts, embeddings)): chunk_id = str(uuid4()) - if chunk.embedding is not None: - embedding = chunk.embedding - else: - embedding = await self.embedder.embed(chunk.content) chunk_record = self.store.ChunkRecord( id=chunk_id, @@ -167,11 +176,20 @@ class ChunkRepository: metadata=json.dumps({"order": order}), vector=embedding, ) - self.store.chunks_table.add([chunk_record]) + chunk_records.append(chunk_record) - chunk.id = chunk_id + chunk = Chunk( + id=chunk_id, + document_id=document_id, + content=chunk_text, + metadata={"order": order}, + ) created_chunks.append(chunk) + # Batch insert all chunks at once + if chunk_records: + self.store.chunks_table.add(chunk_records) + # Force optimization once at the end for bulk operations await self._optimize() return created_chunks @@ -216,7 +234,7 @@ class ChunkRepository: query_embedding = await self.embedder.embed(query) results = self.store.chunks_table.search( - query_embedding, query_type="vector" + query_embedding, query_type="vector", vector_column_name="vector" ).limit(limit) return await self._process_search_results(results) @@ -305,6 +323,7 @@ class ChunkRepository: # Get both arrow and pydantic results to access scores arrow_result = query_result.to_arrow() pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord)) + # Extract scores from arrow result based on search type scores = [] column_names = arrow_result.column_names @@ -322,17 +341,26 @@ class ChunkRepository: else: raise ValueError("Unknown search result format, cannot extract scores") - for i, chunk_record in enumerate(pydantic_results): - # Get document info + # Collect all unique document IDs for batch lookup + document_ids = list(set(chunk.document_id for chunk in pydantic_results)) + + # Batch fetch all documents at once + documents_map = {} + if document_ids: + # Create a WHERE clause for all document IDs + where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids) doc_results = list( self.store.documents_table.search() - .where(f"id = '{chunk_record.document_id}'") - .limit(1) + .where(where_clause) .to_pydantic(DocumentRecord) ) + documents_map = {doc.id: doc for doc in doc_results} - doc_uri = doc_results[0].uri if doc_results else None - doc_meta = doc_results[0].metadata if doc_results else "{}" + for i, chunk_record in enumerate(pydantic_results): + # Get document info from pre-fetched map + doc = documents_map.get(chunk_record.document_id) + doc_uri = doc.uri if doc else None + doc_meta = doc.metadata if doc else "{}" chunk = Chunk( id=chunk_record.id, diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index d9c20490..1a218bb9 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/src/haiku/rag/store/repositories/document.py @@ -17,11 +17,31 @@ class DocumentRepository: def __init__(self, store: Store) -> None: self.store = store + self._chunk_repository = None - from haiku.rag.store.repositories.chunk import ChunkRepository + @property + def chunk_repository(self): + """Lazy-load ChunkRepository when needed.""" + if self._chunk_repository is None: + from haiku.rag.store.repositories.chunk import ChunkRepository - chunk_repository = ChunkRepository(store) - self.chunk_repository = chunk_repository + self._chunk_repository = ChunkRepository(self.store) + return self._chunk_repository + + def _record_to_document(self, record: DocumentRecord) -> Document: + """Convert a DocumentRecord to a Document model.""" + return Document( + id=record.id, + content=record.content, + uri=record.uri, + metadata=json.loads(record.metadata) if record.metadata else {}, + created_at=datetime.fromisoformat(record.created_at) + if record.created_at + else datetime.now(), + updated_at=datetime.fromisoformat(record.updated_at) + if record.updated_at + else datetime.now(), + ) async def create(self, entity: Document) -> Document: """Create a document in the database.""" @@ -61,19 +81,7 @@ class DocumentRepository: if not results: return None - doc_record = results[0] - return Document( - id=doc_record.id, - content=doc_record.content, - uri=doc_record.uri, - metadata=json.loads(doc_record.metadata) if doc_record.metadata else {}, - created_at=datetime.fromisoformat(doc_record.created_at) - if doc_record.created_at - else datetime.now(), - updated_at=datetime.fromisoformat(doc_record.updated_at) - if doc_record.updated_at - else datetime.now(), - ) + return self._record_to_document(results[0]) async def update(self, entity: Document) -> Document: """Update an existing document.""" @@ -104,10 +112,7 @@ class DocumentRepository: return False # Delete associated chunks first - from haiku.rag.store.repositories.chunk import ChunkRepository - - chunk_repo = ChunkRepository(self.store) - await chunk_repo.delete_by_document_id(entity_id) + await self.chunk_repository.delete_by_document_id(entity_id) # Delete the document self.store.documents_table.delete(f"id = '{entity_id}'") @@ -125,22 +130,7 @@ class DocumentRepository: query = query.limit(limit) results = list(query.to_pydantic(DocumentRecord)) - - return [ - Document( - id=doc.id, - content=doc.content, - uri=doc.uri, - metadata=json.loads(doc.metadata) if doc.metadata else {}, - created_at=datetime.fromisoformat(doc.created_at) - if doc.created_at - else datetime.now(), - updated_at=datetime.fromisoformat(doc.updated_at) - if doc.updated_at - else datetime.now(), - ) - for doc in results - ] + return [self._record_to_document(doc) for doc in results] async def get_by_uri(self, uri: str) -> Document | None: """Get a document by its URI.""" @@ -154,27 +144,12 @@ class DocumentRepository: if not results: return None - doc_record = results[0] - return Document( - id=doc_record.id, - content=doc_record.content, - uri=doc_record.uri, - metadata=json.loads(doc_record.metadata) if doc_record.metadata else {}, - created_at=datetime.fromisoformat(doc_record.created_at) - if doc_record.created_at - else datetime.now(), - updated_at=datetime.fromisoformat(doc_record.updated_at) - if doc_record.updated_at - else datetime.now(), - ) + return self._record_to_document(results[0]) async def delete_all(self) -> None: """Delete all documents from the database.""" # Delete all chunks first - from haiku.rag.store.repositories.chunk import ChunkRepository - - chunk_repo = ChunkRepository(self.store) - await chunk_repo.delete_all() + await self.chunk_repository.delete_all() # Get count before deletion count = len(