danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Batch operations & small optimization

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-03 10:50:08 +03:00
parent 75e878308e
commit ed604f6a1d
No known key found for this signature in database
3 changed files with 143 additions and 126 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

116
src/haiku/rag/store/engine.py
View file

@ -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."""

72
src/haiku/rag/store/repositories/chunk.py
View file

@ -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,

81
src/haiku/rag/store/repositories/document.py
View file

@ -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(