diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 4ed8023b..87cd86a5 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -6,7 +6,6 @@ from uuid import uuid4 import lancedb from lancedb.pydantic import LanceModel, Vector from pydantic import Field -from rich.console import Console from haiku.rag.config import Config from haiku.rag.embeddings import get_embedder @@ -105,12 +104,9 @@ class Store: self.settings_table.search().limit(1).to_pydantic(SettingsRecord) ) if existing_settings: - console = Console() - db_version = self.get_haiku_version() - # Future: Add upgrade logic here similar to SQLite version - console.print( - f"[green]LanceDB store initialized (version: {db_version})[/green]" - ) + db_version = self.get_haiku_version() # noqa: F841 + # XXX Add upgrade logic here similar to SQLite version + except Exception: pass diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 6c8bfaca..b77dca55 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -8,6 +8,7 @@ from haiku.rag.chunker import chunker from haiku.rag.embeddings import get_embedder from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.models.chunk import Chunk +from haiku.rag.utils import debounce class ChunkRepository: @@ -24,6 +25,14 @@ class ChunkRepository: except Exception: pass + @debounce(1.0) + async def _optimize(self) -> None: + """Optimize the chunks table to refresh indexes.""" + try: + self.store.chunks_table.optimize() + except RuntimeError: + pass + async def create(self, entity: Chunk) -> Chunk: """Create a chunk in the database.""" assert entity.document_id, "Chunk must have a document_id to be created" @@ -49,7 +58,7 @@ class ChunkRepository: entity.id = chunk_id # Optimize table after insert to update indexes - self.store.chunks_table.optimize() + await self._optimize() return entity async def get_by_id(self, entity_id: str) -> Chunk | None: @@ -88,7 +97,7 @@ class ChunkRepository: }, ) # Optimize table after update to refresh indexes - self.store.chunks_table.optimize() + await self._optimize() return entity diff --git a/src/haiku/rag/utils.py b/src/haiku/rag/utils.py index edf381a5..05f9056e 100644 --- a/src/haiku/rag/utils.py +++ b/src/haiku/rag/utils.py @@ -1,4 +1,7 @@ +import asyncio import sys +from collections.abc import Callable +from functools import wraps from importlib import metadata from io import BytesIO from pathlib import Path @@ -10,6 +13,42 @@ from docling_core.types.io import DocumentStream from packaging.version import Version, parse +def debounce(wait: float) -> Callable: + """ + A decorator to debounce a function, ensuring it is called only after a specified delay + and always executes after the last call. + + Args: + wait (float): The debounce delay in seconds. + + Returns: + Callable: The decorated function. + """ + + def decorator(func: Callable) -> Callable: + last_call = None + task = None + + @wraps(func) + async def debounced(*args, **kwargs): + nonlocal last_call, task + last_call = asyncio.get_event_loop().time() + + if task: + task.cancel() + + async def call_func(): + await asyncio.sleep(wait) + if asyncio.get_event_loop().time() - last_call >= wait: # type: ignore + await func(*args, **kwargs) + + task = asyncio.create_task(call_func()) + + return debounced + + return decorator + + def get_default_data_dir() -> Path: """Get the user data directory for the current system platform.