Use a lock for optimize(), when bulk creating chunks optimize once.

This commit is contained in:
Yiorgis Gozadinos 2025-09-02 12:10:54 +03:00
parent fe850e778d
commit 767d36b494
No known key found for this signature in database

View file

@ -1,3 +1,4 @@
import asyncio
import json import json
from uuid import uuid4 from uuid import uuid4
@ -8,7 +9,6 @@ from haiku.rag.chunker import chunker
from haiku.rag.embeddings import get_embedder from haiku.rag.embeddings import get_embedder
from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.engine import DocumentRecord, Store
from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.chunk import Chunk
from haiku.rag.utils import debounce
class ChunkRepository: class ChunkRepository:
@ -17,6 +17,7 @@ class ChunkRepository:
def __init__(self, store: Store) -> None: def __init__(self, store: Store) -> None:
self.store = store self.store = store
self.embedder = get_embedder() self.embedder = get_embedder()
self._optimize_lock = asyncio.Lock()
def _ensure_fts_index(self) -> None: def _ensure_fts_index(self) -> None:
"""Ensure FTS index exists on the content column.""" """Ensure FTS index exists on the content column."""
@ -25,13 +26,14 @@ class ChunkRepository:
except Exception: except Exception:
pass pass
@debounce(1.0)
async def _optimize(self) -> None: async def _optimize(self) -> None:
"""Optimize the chunks table to refresh indexes.""" """Optimize the chunks table to refresh indexes."""
try: async with self._optimize_lock:
self.store.chunks_table.optimize() try:
except RuntimeError: self.store.chunks_table.optimize()
pass except (RuntimeError, OSError):
# Handle "too many open files" and other resource errors gracefully
pass
async def create(self, entity: Chunk) -> Chunk: async def create(self, entity: Chunk) -> Chunk:
"""Create a chunk in the database.""" """Create a chunk in the database."""
@ -57,8 +59,10 @@ class ChunkRepository:
entity.id = chunk_id entity.id = chunk_id
# Optimize table after insert to update indexes # Try to optimize if not currently locked (non-blocking)
await self._optimize() if not self._optimize_lock.locked():
asyncio.create_task(self._optimize())
return entity return entity
async def get_by_id(self, entity_id: str) -> Chunk | None: async def get_by_id(self, entity_id: str) -> Chunk | None:
@ -96,8 +100,9 @@ class ChunkRepository:
"vector": embedding, "vector": embedding,
}, },
) )
# Optimize table after update to refresh indexes # Try to optimize if not currently locked (non-blocking)
await self._optimize() if not self._optimize_lock.locked():
asyncio.create_task(self._optimize())
return entity return entity
@ -144,9 +149,27 @@ class ChunkRepository:
chunk = Chunk( chunk = Chunk(
document_id=document_id, content=chunk_text, metadata={"order": order} document_id=document_id, content=chunk_text, metadata={"order": order}
) )
created_chunk = await self.create(chunk) # Use create but don't trigger individual optimizations
created_chunks.append(created_chunk) 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,
document_id=document_id,
content=chunk_text,
metadata=json.dumps({"order": order}),
vector=embedding,
)
self.store.chunks_table.add([chunk_record])
chunk.id = chunk_id
created_chunks.append(chunk)
# Force optimization once at the end for bulk operations
await self._optimize()
return created_chunks return created_chunks
async def delete_all(self) -> None: async def delete_all(self) -> None: