haiku.rag 0.11.3__py3-none-any.whl → 0.11.4__py3-none-any.whl

haiku/rag/client.py

@@ -46,6 +46,9 @@ class HaikuRAG:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):  # noqa: ARG002
         """Async context manager exit."""
+        # Wait for any pending vacuum to complete before closing
+        async with self.store._vacuum_lock:
+            pass
         self.close()
         return False
 
@@ -522,19 +525,22 @@ class HaikuRAG:
         merged.append(current)
         return merged
 
-    async def ask(self, question: str, cite: bool = False) -> str:
+    async def ask(
+        self, question: str, cite: bool = False, system_prompt: str | None = None
+    ) -> str:
         """Ask a question using the configured QA agent.
 
         Args:
             question: The question to ask.
             cite: Whether to include citations in the response.
+            system_prompt: Optional custom system prompt for the QA agent.
 
         Returns:
             The generated answer as a string.
         """
         from haiku.rag.qa import get_qa_agent
 
-        qa_agent = get_qa_agent(self, use_citations=cite)
+        qa_agent = get_qa_agent(self, use_citations=cite, system_prompt=system_prompt)
         return await qa_agent.answer(question)
 
     async def rebuild_database(self) -> AsyncGenerator[str, None]:
@@ -617,13 +623,13 @@ class HaikuRAG:
 
         # Final maintenance: centralized vacuum to curb disk usage
         try:
-            self.store.vacuum()
+            await self.store.vacuum()
         except Exception:
             pass
 
     async def vacuum(self) -> None:
         """Optimize and clean up old versions across all tables."""
-        self.store.vacuum()
+        await self.store.vacuum()
 
     def close(self):
         """Close the underlying store connection."""

haiku/rag/config.py

@@ -57,6 +57,11 @@ class AppConfig(BaseModel):
     # and error out when the database does not already exist.
     DISABLE_DB_AUTOCREATE: bool = False
 
+    # Vacuum retention threshold in seconds. Only versions older than this
+    # threshold will be removed during vacuum operations. Default is 60 seconds
+    # to allow concurrent connections to safely use recent versions.
+    VACUUM_RETENTION_SECONDS: int = 60
+
     @field_validator("MONITOR_DIRECTORIES", mode="before")
     @classmethod
     def parse_monitor_directories(cls, v):

haiku/rag/migration.py

@@ -27,7 +27,7 @@ class SQLiteToLanceDBMigrator:
         self.lancedb_path = lancedb_path
         self.console = Console()
 
-    def migrate(self) -> bool:
+    async def migrate(self) -> bool:
         """Perform the migration."""
         try:
             self.console.print(
@@ -94,7 +94,7 @@ class SQLiteToLanceDBMigrator:
             # Optimize and cleanup using centralized vacuum
             self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
             try:
-                lance_store.vacuum()
+                await lance_store.vacuum()
                 self.console.print("[green]✅ Optimization completed[/green]")
             except Exception as e:
                 self.console.print(
@@ -313,4 +313,4 @@ async def migrate_sqlite_to_lancedb(
         lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
 
     migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
-    return migrator.migrate()
+    return await migrator.migrate()

haiku/rag/qa/__init__.py

@@ -3,7 +3,11 @@ from haiku.rag.config import Config
 from haiku.rag.qa.agent import QuestionAnswerAgent
 
 
-def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent:
+def get_qa_agent(
+    client: HaikuRAG,
+    use_citations: bool = False,
+    system_prompt: str | None = None,
+) -> QuestionAnswerAgent:
     provider = Config.QA_PROVIDER
     model_name = Config.QA_MODEL
 
@@ -12,4 +16,5 @@ def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswe
         provider=provider,
         model=model_name,
         use_citations=use_citations,
+        system_prompt=system_prompt,
     )

haiku/rag/qa/agent.py

@@ -30,18 +30,21 @@ class QuestionAnswerAgent:
         model: str,
         use_citations: bool = False,
         q: float = 0.0,
+        system_prompt: str | None = None,
     ):
         self._client = client
 
-        system_prompt = (
-            QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
-        )
+        if system_prompt is None:
+            system_prompt = (
+                QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
+            )
         model_obj = self._get_model(provider, model)
 
         self._agent = Agent(
             model=model_obj,
             deps_type=Dependencies,
             system_prompt=system_prompt,
+            retries=3,
         )
 
         @self._agent.tool

haiku/rag/store/engine.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import logging
 from datetime import timedelta
@@ -51,6 +52,7 @@ class Store:
     def __init__(self, db_path: Path, skip_validation: bool = False):
         self.db_path: Path = db_path
         self.embedder = get_embedder()
+        self._vacuum_lock = asyncio.Lock()
 
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
@@ -78,14 +80,40 @@ class Store:
         if not skip_validation:
             self._validate_configuration()
 
-    def vacuum(self) -> None:
-        """Optimize and clean up old versions across all tables to reduce disk usage."""
+    async def vacuum(self, retention_seconds: int | None = None) -> None:
+        """Optimize and clean up old versions across all tables to reduce disk usage.
+
+        Args:
+            retention_seconds: Retention threshold in seconds. Only versions older
+                              than this will be removed. If None, uses Config.VACUUM_RETENTION_SECONDS.
+
+        Note:
+            If vacuum is already running, this method returns immediately without blocking.
+            Use asyncio.create_task(store.vacuum()) for non-blocking background execution.
+        """
         if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
             return
 
-        # Perform maintenance per table using optimize() with cleanup_older_than 0
-        for table in [self.documents_table, self.chunks_table, self.settings_table]:
-            table.optimize(cleanup_older_than=timedelta(0))
+        # Skip if already running (non-blocking)
+        if self._vacuum_lock.locked():
+            return
+
+        async with self._vacuum_lock:
+            try:
+                # Evaluate config at runtime to allow dynamic changes
+                if retention_seconds is None:
+                    retention_seconds = Config.VACUUM_RETENTION_SECONDS
+                # Perform maintenance per table using optimize() with configurable retention
+                retention = timedelta(seconds=retention_seconds)
+                for table in [
+                    self.documents_table,
+                    self.chunks_table,
+                    self.settings_table,
+                ]:
+                    table.optimize(cleanup_older_than=retention)
+            except (RuntimeError, OSError) as e:
+                # Handle resource errors gracefully
+                logger.debug(f"Vacuum skipped due to resource constraints: {e}")
 
     def _connect_to_lancedb(self, db_path: Path):
         """Establish connection to LanceDB (local, cloud, or object storage)."""

haiku/rag/store/repositories/chunk.py

@@ -1,4 +1,3 @@
-import asyncio
 import inspect
 import json
 import logging
@@ -23,7 +22,6 @@ class ChunkRepository:
     def __init__(self, store: Store) -> None:
         self.store = store
         self.embedder = get_embedder()
-        self._optimize_lock = asyncio.Lock()
 
     def _ensure_fts_index(self) -> None:
         """Ensure FTS index exists on the content column."""
@@ -35,21 +33,6 @@ class ChunkRepository:
             # 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."""
-        # Skip optimization for LanceDB Cloud as it handles this automatically
-        if Config.LANCEDB_URI and Config.LANCEDB_URI.startswith("db://"):
-            return
-
-        async with self._optimize_lock:
-            try:
-                self.store.chunks_table.optimize()
-            except (RuntimeError, OSError) as e:
-                # Handle "too many open files" and other resource errors gracefully
-                logger.debug(
-                    f"Table optimization skipped due to resource constraints: {e}"
-                )
-
     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"
@@ -77,11 +60,6 @@ class ChunkRepository:
         self.store.chunks_table.add([chunk_record])
 
         entity.id = chunk_id
-
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
         return entity
 
     async def get_by_id(self, entity_id: str) -> Chunk | None:
@@ -125,10 +103,6 @@ class ChunkRepository:
                 "vector": embedding,
             },
         )
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
         return entity
 
     async def delete(self, entity_id: str) -> bool:
@@ -227,8 +201,6 @@ class ChunkRepository:
         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
 
     async def delete_all(self) -> None:

haiku/rag/store/repositories/document.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from datetime import datetime
 from typing import TYPE_CHECKING
@@ -200,6 +201,9 @@ class DocumentRepository:
                     chunk.order = order
                     await self.chunk_repository.create(chunk)
 
+            # Vacuum old versions in background (non-blocking)
+            asyncio.create_task(self.store.vacuum())
+
             return created_doc
         except Exception:
             # Roll back to the captured versions and re-raise
@@ -230,6 +234,9 @@ class DocumentRepository:
                 updated_doc.id, docling_document
             )
 
+            # Vacuum old versions in background (non-blocking)
+            asyncio.create_task(self.store.vacuum())
+
             return updated_doc
         except Exception:
             # Roll back to the captured versions and re-raise

{haiku_rag-0.11.3.dist-info → haiku_rag-0.11.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.11.3
+Version: 0.11.4
 Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT

{haiku_rag-0.11.3.dist-info → haiku_rag-0.11.4.dist-info}/RECORD

@@ -2,11 +2,11 @@ haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haiku/rag/app.py,sha256=B5BZaUgC9u3hz17tRKu_xKefPu5xqdZQPCxfa5K8_EI,19180
 haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
 haiku/rag/cli.py,sha256=6GmWfs30uUytcnyBMv-OE4tCkkuKnfYeJFU5DIrZ_vU,13212
-haiku/rag/client.py,sha256=fz5bZP1KNWbc9cvpEC8puMBHEvt-vVtFjRMItt_WD0M,23920
-haiku/rag/config.py,sha256=c2WoaieI3-HAWb6lCmVnJHY22NXl2SGLsndRbiqCzeA,2305
+haiku/rag/client.py,sha256=tk3BWa9u2mtUEMEOEcgF2ebUqdZRv4jSnu4S6Zho-Xg,24210
+haiku/rag/config.py,sha256=rJ-xi66gc94F_QA6BvgRgNq8CNiFTN0CCIaOv1TMMFw,2569
 haiku/rag/logging.py,sha256=dm65AwADpcQsH5OAPtRA-4hsw0w5DK-sGOvzYkj6jzw,1720
 haiku/rag/mcp.py,sha256=H7XibtSNUviFeaJVsXzHiRqUm0nJCpA7A1QHuBv6SKQ,5057
-haiku/rag/migration.py,sha256=zm0-60PiS1hIQnZz65B7qfsgM7GwZVXFqMFowjpVBs8,11058
+haiku/rag/migration.py,sha256=XldX0CTHPXNGrkdQ-gocr4kQGBsz-316WcE0ZDRfb48,11076
 haiku/rag/monitor.py,sha256=VP3bqY0mEodOP60eN4RMldgrL1ti5gMjuDuQ-_vBvFc,2759
 haiku/rag/reader.py,sha256=aW8LG0X31kVWS7kU2tKVpe8RqP3Ne_oIidd_X3UDLH0,3307
 haiku/rag/utils.py,sha256=dBzhKaOHI9KRiJqHErcXUnqtnXY2AgOK8PCLA3rhO0A,6115
@@ -26,8 +26,8 @@ haiku/rag/graph/nodes/analysis.py,sha256=g-Aw3nPuCHWo0CXM96Ixa4vQI4TpI6tg6ooHT_J
 haiku/rag/graph/nodes/plan.py,sha256=Bb6Fva9vwArCU-5xBr24N4pM3wfLP-Vwufgss8HfXMQ,2622
 haiku/rag/graph/nodes/search.py,sha256=DdHhEY7fmWUqis6Nk0bj-di56-ML262B51N9zytzKYk,3699
 haiku/rag/graph/nodes/synthesize.py,sha256=WF0D44SwLP1OK8C6ViOAhFOtGQ0mj3aO54z5bemJb4E,1828
-haiku/rag/qa/__init__.py,sha256=Sl7Kzrg9CuBOcMF01wc1NtQhUNWjJI0MhIHfCWrb8V4,434
-haiku/rag/qa/agent.py,sha256=rtUkEmnD8lMHIxpPPVY6TdmF4aSlZnLjad5eDefrlBw,3145
+haiku/rag/qa/__init__.py,sha256=eFRV5GFwe1UsqniEqOLdzAMT2J6QhSiHq5_Li7c6Fs4,520
+haiku/rag/qa/agent.py,sha256=A4FrzoYP4pRzJOOJQGlNFp48yRWMSICH4d8JfxFabqk,3256
 haiku/rag/qa/prompts.py,sha256=Lqwn3m4zCsu_CJiC4s9cLsuPNbb9nq6j2PqEF3lw1eA,3380
 haiku/rag/qa/deep/__init__.py,sha256=SnCpWxWip-TaFzVKlFyrOgYeXEqT_gpIlaSItEEJ6r0,50
 haiku/rag/qa/deep/dependencies.py,sha256=AKFqcC1D3N1VPudnFmLH29K5eJWEC5wtwUGkO4FM4jc,998
@@ -50,19 +50,19 @@ haiku/rag/research/prompts.py,sha256=opz4MXjoDHH1wjG6bPyiqT0LVzk3pBA6y_a9zpBW8yM
 haiku/rag/research/state.py,sha256=P8RXJMi3wA3l1j6yo8dsAyso6S27FgqS7fvZUUY447A,917
 haiku/rag/research/stream.py,sha256=amyGDimkNp_FHYUXCqtpbeDOx7sC1jQ-7DwoxuNOL1g,5576
 haiku/rag/store/__init__.py,sha256=R2IRcxtkFDxqa2sgMirqLq3l2-FPdWr6ydYStaqm5OQ,104
-haiku/rag/store/engine.py,sha256=i2t15j1hhEY0SMNXk2uUSeCio-mHnTpnN5S83afVx6o,10257
+haiku/rag/store/engine.py,sha256=n2IxztyN2UpLLSUVXurjL-e_ANthKUpWyB1gdHfgBMM,11468
 haiku/rag/store/models/__init__.py,sha256=kc7Ctf53Jr483tk4QTIrcgqBbXDz4ZoeYSkFXfPnpks,89
 haiku/rag/store/models/chunk.py,sha256=3EuZav4QekJIeHBCub48EM8SjNX8HEJ6wVDXGot4PEQ,421
 haiku/rag/store/models/document.py,sha256=cZXy_jEti-hnhq7FKhuhCfd99ccY9fIHMLovB_Thbb8,425
 haiku/rag/store/repositories/__init__.py,sha256=Olv5dLfBQINRV3HrsfUpjzkZ7Qm7goEYyMNykgo_DaY,291
-haiku/rag/store/repositories/chunk.py,sha256=UfajEWf5VmMuSozGRDlWBjJNR0ngvOVFDrp6_augzBg,15217
-haiku/rag/store/repositories/document.py,sha256=C9GbIl8sa2-Djaml4hlaPTtjV2HwHaz_Wzs35sdbdhg,7876
+haiku/rag/store/repositories/chunk.py,sha256=B0CowrBNy0fd8GLnVJVfqDaLoWxEPPJK3SODya0I0OI,14093
+haiku/rag/store/repositories/document.py,sha256=EtgD5pDjghXf6dloBOOEVJp8DI9O_celc_FbYzOywAE,8125
 haiku/rag/store/repositories/settings.py,sha256=ObrDrzxHn-yA1WcbgIoJoVmAbVvQHAFvEdRyJFt5Opc,5685
 haiku/rag/store/upgrades/__init__.py,sha256=RQ8A6rEXBASLb5PD9vdDnEas_m_GgRzzdVu4B88Snqc,1975
 haiku/rag/store/upgrades/v0_10_1.py,sha256=qNGnxj6hoHaHJ1rKTiALfw0c9NQOi0KAK-VZCD_073A,1959
 haiku/rag/store/upgrades/v0_9_3.py,sha256=NrjNilQSgDtFWRbL3ZUtzQzJ8tf9u0dDRJtnDFwwbdw,3322
-haiku_rag-0.11.3.dist-info/METADATA,sha256=kohzziLYTajCvyfMxnJkaY7XeEfqPnmVuC_20NpkeLA,6748
-haiku_rag-0.11.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.11.3.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.11.3.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.11.3.dist-info/RECORD,,
+haiku_rag-0.11.4.dist-info/METADATA,sha256=YA7Fr6OnWYeOH139aZkLRP9Yj0S0KXEayzfKsgFaz08,6748
+haiku_rag-0.11.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.11.4.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.11.4.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.11.4.dist-info/RECORD,,