Use a lock for optimize(), when bulk creating chunks optimize once.
This commit is contained in:
parent
fe850e778d
commit
767d36b494
1 changed files with 35 additions and 12 deletions
|
|
@ -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:
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue