haiku.rag 0.7.7__py3-none-any.whl → 0.8.1__py3-none-any.whl

haiku/rag/app.py

@@ -102,6 +102,15 @@ class HaikuRAGApp:
             except Exception as e:
                 self.console.print(f"[red]Error rebuilding database: {e}[/red]")
 
+    async def vacuum(self):
+        """Run database maintenance: optimize and cleanup table history."""
+        try:
+            async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
+                await client.vacuum()
+            self.console.print("[b]Vacuum completed successfully.[/b]")
+        except Exception as e:
+            self.console.print(f"[red]Error during vacuum: {e}[/red]")
+
     def show_settings(self):
         """Display current configuration settings."""
         self.console.print("[bold]haiku.rag configuration[/bold]")

haiku/rag/cli.py

@@ -256,6 +256,18 @@ def rebuild(
     asyncio.run(app.rebuild())
 
 
+@cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage")
+def vacuum(
+    db: Path = typer.Option(
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        "--db",
+        help="Path to the LanceDB database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    asyncio.run(app.vacuum())
+
+
 @cli.command(
     "serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)"
 )

haiku/rag/client.py

@@ -550,6 +550,16 @@ class HaikuRAG:
                 )
                 yield doc.id
 
+        # Final maintenance: centralized vacuum to curb disk usage
+        try:
+            self.store.vacuum()
+        except Exception:
+            pass
+
+    async def vacuum(self) -> None:
+        """Optimize and clean up old versions across all tables."""
+        self.store.vacuum()
+
     def close(self):
         """Close the underlying store connection."""
         self.store.close()

haiku/rag/config.py

@@ -32,6 +32,10 @@ class AppConfig(BaseModel):
     CHUNK_SIZE: int = 256
     CONTEXT_CHUNK_RADIUS: int = 0
 
+    # Optional dotted path or file path to a callable that preprocesses
+    # markdown content before chunking. Examples:
+    MARKDOWN_PREPROCESSOR: str = ""
+
     OLLAMA_BASE_URL: str = "http://localhost:11434"
     VLLM_EMBEDDINGS_BASE_URL: str = ""
     VLLM_RERANK_BASE_URL: str = ""

haiku/rag/logging.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -50,4 +51,6 @@ def configure_cli_logging(level: int = logging.INFO) -> logging.Logger:
     logger = get_logger()
     logger.setLevel(level)
     logger.propagate = False
+
+    warnings.filterwarnings("ignore")
     return logger

haiku/rag/migration.py

@@ -47,7 +47,7 @@ class SQLiteToLanceDBMigrator:
 
             # Load the sqlite-vec extension
             try:
-                import sqlite_vec
+                import sqlite_vec  # type: ignore
 
                 sqlite_conn.enable_load_extension(True)
                 sqlite_vec.load(sqlite_conn)
@@ -91,10 +91,10 @@ class SQLiteToLanceDBMigrator:
 
             sqlite_conn.close()
 
-            # Optimize the chunks table after migration
+            # Optimize and cleanup using centralized vacuum
             self.console.print("[blue]Optimizing LanceDB...[/blue]")
             try:
-                lance_store.chunks_table.optimize()
+                lance_store.vacuum()
                 self.console.print("[green]✅ Optimization completed[/green]")
             except Exception as e:
                 self.console.print(

haiku/rag/store/engine.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from datetime import timedelta
 from importlib import metadata
 from pathlib import Path
 from uuid import uuid4
@@ -62,6 +63,15 @@ 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."""
+        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))
+
     def _connect_to_lancedb(self, db_path: Path):
         """Establish connection to LanceDB (local, cloud, or object storage)."""
         # Check if we have cloud configuration
@@ -159,16 +169,18 @@ class Store:
             self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
         )
         if settings_records:
-            settings = (
+            # Only write if version actually changes to avoid creating new table versions
+            current = (
                 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)}
-            )
+            if current.get("version") != version:
+                current["version"] = version
+                self.settings_table.update(
+                    where="id = 'settings'",
+                    values={"settings": json.dumps(current)},
+                )
         else:
             # Create new settings record
             settings_data = Config.model_dump(mode="json")
@@ -197,6 +209,21 @@ class Store:
         # LanceDB connections are automatically managed
         pass
 
+    def current_table_versions(self) -> dict[str, int]:
+        """Capture current versions of key tables for rollback using LanceDB's API."""
+        return {
+            "documents": int(self.documents_table.version),
+            "chunks": int(self.chunks_table.version),
+            "settings": int(self.settings_table.version),
+        }
+
+    def restore_table_versions(self, versions: dict[str, int]) -> bool:
+        """Restore tables to the provided versions using LanceDB's API."""
+        self.documents_table.restore(int(versions["documents"]))
+        self.chunks_table.restore(int(versions["chunks"]))
+        self.settings_table.restore(int(versions["settings"]))
+        return True
+
     @property
     def _connection(self):
         """Compatibility property for repositories expecting _connection."""

haiku/rag/store/repositories/chunk.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import json
 import logging
 from uuid import uuid4
@@ -11,6 +12,7 @@ from haiku.rag.config import Config
 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 load_callable, text_to_docling_document
 
 logger = logging.getLogger(__name__)
 
@@ -152,7 +154,28 @@ class ChunkRepository:
         self, document_id: str, document: DoclingDocument
     ) -> list[Chunk]:
         """Create chunks and embeddings for a document from DoclingDocument."""
-        chunk_texts = await chunker.chunk(document)
+        # Optionally preprocess markdown before chunking
+        processed_document = document
+        preprocessor_path = Config.MARKDOWN_PREPROCESSOR
+        if preprocessor_path:
+            try:
+                pre_fn = load_callable(preprocessor_path)
+                markdown = document.export_to_markdown()
+                result = pre_fn(markdown)
+                if inspect.isawaitable(result):
+                    result = await result  # type: ignore[assignment]
+                processed_markdown = result
+                if not isinstance(processed_markdown, str):
+                    raise ValueError("Preprocessor must return a markdown string")
+                processed_document = text_to_docling_document(
+                    processed_markdown, name="content.md"
+                )
+            except Exception as e:
+                logger.warning(
+                    f"Failed to apply MARKDOWN_PREPROCESSOR '{preprocessor_path}': {e}. Proceeding without preprocessing."
+                )
+
+        chunk_texts = await chunker.chunk(processed_document)
 
         embeddings = await self.embedder.embed(chunk_texts)
 

haiku/rag/store/repositories/document.py

@@ -171,44 +171,64 @@ class DocumentRepository:
         chunks: list["Chunk"] | None = None,
     ) -> Document:
         """Create a document with its chunks and embeddings."""
+        # Snapshot table versions for versioned rollback (if supported)
+        versions = self.store.current_table_versions()
+
         # Create the document
         created_doc = await self.create(entity)
 
-        # Create chunks if not provided
-        if chunks is None:
-            assert created_doc.id is not None, (
-                "Document ID should not be None after creation"
-            )
-            await self.chunk_repository.create_chunks_for_document(
-                created_doc.id, docling_document
-            )
-        else:
-            # Use provided chunks, set order from list position
-            assert created_doc.id is not None, (
-                "Document ID should not be None after creation"
-            )
-            for order, chunk in enumerate(chunks):
-                chunk.document_id = created_doc.id
-                chunk.metadata["order"] = order
-                await self.chunk_repository.create(chunk)
-
-        return created_doc
+        # Attempt to create chunks; on failure, prefer version rollback
+        try:
+            # Create chunks if not provided
+            if chunks is None:
+                assert created_doc.id is not None, (
+                    "Document ID should not be None after creation"
+                )
+                await self.chunk_repository.create_chunks_for_document(
+                    created_doc.id, docling_document
+                )
+            else:
+                # Use provided chunks, set order from list position
+                assert created_doc.id is not None, (
+                    "Document ID should not be None after creation"
+                )
+                for order, chunk in enumerate(chunks):
+                    chunk.document_id = created_doc.id
+                    chunk.metadata["order"] = order
+                    await self.chunk_repository.create(chunk)
+
+            return created_doc
+        except Exception:
+            # Roll back to the captured versions and re-raise
+            self.store.restore_table_versions(versions)
+            raise
 
     async def _update_with_docling(
         self, entity: Document, docling_document: DoclingDocument
     ) -> Document:
         """Update a document and regenerate its chunks."""
-        # Delete existing chunks
         assert entity.id is not None, "Document ID is required for update"
+
+        # Snapshot table versions for versioned rollback
+        versions = self.store.current_table_versions()
+
+        # Delete existing chunks before writing new ones
         await self.chunk_repository.delete_by_document_id(entity.id)
 
-        # Update the document
-        updated_doc = await self.update(entity)
+        try:
+            # Update the document
+            updated_doc = await self.update(entity)
 
-        # Create new chunks
-        assert updated_doc.id is not None, "Document ID should not be None after update"
-        await self.chunk_repository.create_chunks_for_document(
-            updated_doc.id, docling_document
-        )
+            # Create new chunks
+            assert updated_doc.id is not None, (
+                "Document ID should not be None after update"
+            )
+            await self.chunk_repository.create_chunks_for_document(
+                updated_doc.id, docling_document
+            )
 
-        return updated_doc
+            return updated_doc
+        except Exception:
+            # Roll back to the captured versions and re-raise
+            self.store.restore_table_versions(versions)
+            raise

haiku/rag/store/repositories/settings.py

@@ -84,10 +84,15 @@ class SettingsRepository:
         )
 
         if existing:
-            # Update existing settings
-            self.store.settings_table.update(
-                where="id = 'settings'", values={"settings": json.dumps(current_config)}
+            # Only update when configuration actually changed to avoid needless new versions
+            existing_payload = (
+                json.loads(existing[0].settings) if existing[0].settings else {}
             )
+            if existing_payload != current_config:
+                self.store.settings_table.update(
+                    where="id = 'settings'",
+                    values={"settings": json.dumps(current_config)},
+                )
         else:
             # Create new settings
             settings_record = SettingsRecord(

haiku/rag/utils.py

@@ -1,10 +1,13 @@
 import asyncio
+import importlib
+import importlib.util
 import sys
 from collections.abc import Callable
 from functools import wraps
 from importlib import metadata
 from io import BytesIO
 from pathlib import Path
+from types import ModuleType
 
 import httpx
 from docling.document_converter import DocumentConverter
@@ -106,3 +109,54 @@ def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocu
     converter = DocumentConverter()
     result = converter.convert(doc_stream)
     return result.document
+
+
+def load_callable(path: str):
+    """Load a callable from a dotted path or file path.
+
+    Supported formats:
+    - "package.module:func" or "package.module.func"
+    - "path/to/file.py:func"
+
+    Returns the loaded callable. Raises ValueError on failure.
+    """
+    if not path:
+        raise ValueError("Empty callable path provided")
+
+    module_part = None
+    func_name = None
+
+    if ":" in path:
+        module_part, func_name = path.split(":", 1)
+    else:
+        # split by last dot for module.attr
+        if "." in path:
+            module_part, func_name = path.rsplit(".", 1)
+        else:
+            raise ValueError(
+                "Invalid callable path format. Use 'module:func' or 'module.func' or 'file.py:func'."
+            )
+
+    # Try file path first
+    mod: ModuleType | None = None
+    module_path = Path(module_part)
+    if module_path.suffix == ".py" and module_path.exists():
+        spec = importlib.util.spec_from_file_location(module_path.stem, module_path)
+        if spec and spec.loader:
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+    else:
+        # Import as a module path
+        try:
+            mod = importlib.import_module(module_part)
+        except Exception as e:
+            raise ValueError(f"Failed to import module '{module_part}': {e}")
+
+    if not hasattr(mod, func_name):
+        raise ValueError(f"Callable '{func_name}' not found in module '{module_part}'")
+    func = getattr(mod, func_name)
+    if not callable(func):
+        raise ValueError(
+            f"Attribute '{func_name}' in module '{module_part}' is not callable"
+        )
+    return func

{haiku_rag-0.7.7.dist-info → haiku_rag-0.8.1.dist-info}/METADATA

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

{haiku_rag-0.7.7.dist-info → haiku_rag-0.8.1.dist-info}/RECORD

@@ -1,15 +1,15 @@
 haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-haiku/rag/app.py,sha256=n9Y-4wHnxGeatCE-scqbv5JBnPo-5WOJ6URdfZveHC8,7837
+haiku/rag/app.py,sha256=XlL6PNPSqeBKF6bemvdSfXEnQghywudwZv-C116NuZU,8254
 haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
-haiku/rag/cli.py,sha256=HqFHU9x2tR1yTR74V3NPndqE4R2Yn-ohASyHp334pAg,8597
-haiku/rag/client.py,sha256=N4zkWjE9Rsw9YgPvNo83xptHUQR2ognfOnjkoV_w6hc,20999
-haiku/rag/config.py,sha256=3H41da9BU1R1y2JJHD0cOSErX_VSM1UXA7M2JSOxFXE,1795
-haiku/rag/logging.py,sha256=a0ELyeMqb85ebeOTN8OQCTL1PiMWiiV9R_OOH-VZoA8,1665
+haiku/rag/cli.py,sha256=houkHTeVc89BA3zPksCjUooEnScSg1Ez_BIHBH6cmJQ,8920
+haiku/rag/client.py,sha256=NJVGXzVzpoVy1sttz_xEU7mXWtObKT8pGpvo5pZyzwc,21288
+haiku/rag/config.py,sha256=k5SSh7nYIFKX5LcWYu4bP-4GV5Y-Wq1UzrLUtRAM5Pw,1954
+haiku/rag/logging.py,sha256=dm65AwADpcQsH5OAPtRA-4hsw0w5DK-sGOvzYkj6jzw,1720
 haiku/rag/mcp.py,sha256=bR9Y-Nz-hvjiql20Y0KE0hwNGwyjmPGX8K9d-qmXptY,4683
-haiku/rag/migration.py,sha256=n5G6SDhTo8wTf0uCYbWGegq1LqIgILDLNjWcGvSj-SQ,11053
+haiku/rag/migration.py,sha256=M--KnSF3lxgKjxmokb4vuzGH-pV8eg0C_8e7jvPqW8Y,11058
 haiku/rag/monitor.py,sha256=r386nkhdlsU8UECwIuVwnrSlgMk3vNIuUZGNIzkZuec,2770
 haiku/rag/reader.py,sha256=qkPTMJuQ_o4sK-8zpDl9WFYe_MJ7aL_gUw6rczIpW-g,3274
-haiku/rag/utils.py,sha256=c8F0ECsFSqvQxzxINAOAnvShoOnJPLsOaNE3JEY2JSc,3230
+haiku/rag/utils.py,sha256=aiuPu_rrfpyIvJJq0o5boUIIvCdNzdpKwAIPYYn3iG8,4965
 haiku/rag/embeddings/__init__.py,sha256=44IfDITGIFTflGT6UEmiYOwpWFVbYv5smLY59D0YeCs,1419
 haiku/rag/embeddings/base.py,sha256=BnSviKrlzjv3L0sZJs_T-pxfawd-bcTak-rsX-D2f3A,497
 haiku/rag/embeddings/ollama.py,sha256=LuLlHH6RGoO9_gFCIlbmesuXOj017gTw6z-p8Ez0CfE,595
@@ -25,17 +25,17 @@ haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c
 haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,867
 haiku/rag/reranking/vllm.py,sha256=xVGH9ss-ISWdJ5SKUUHUbTqBo7PIEmA_SQv0ScdJ6XA,1479
 haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
-haiku/rag/store/engine.py,sha256=XHGo5Xl-dCFdQHrOdMo64xVK5n0k8-LoUl5V-tlA0HI,7131
+haiku/rag/store/engine.py,sha256=fNrykqMX7PRSCt4LSRfuJ66OLrb8BVYq2bpbfI2iaWU,8455
 haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
 haiku/rag/store/models/chunk.py,sha256=ZNyTfO6lh3rXWLVYO3TZcitbL4LSUGr42fR6jQQ5iQc,364
 haiku/rag/store/models/document.py,sha256=zSSpt6pyrMJAIXGQvIcqojcqUzwZnhp3WxVokaWxNRc,396
 haiku/rag/store/repositories/__init__.py,sha256=Olv5dLfBQINRV3HrsfUpjzkZ7Qm7goEYyMNykgo_DaY,291
-haiku/rag/store/repositories/chunk.py,sha256=v4y4eh4yIf6zJaWfHxljvnmb12dmvwdinzmxQt8Lvhs,13343
-haiku/rag/store/repositories/document.py,sha256=lP8Lo82KTP-qwXFRpYZ46WjeAdAsHwZ5pJcrXdz4g0U,6988
-haiku/rag/store/repositories/settings.py,sha256=dqnAvm-98nQrWpLBbf9QghJw673QD80-iqQhRMP5t0c,5025
+haiku/rag/store/repositories/chunk.py,sha256=n4VMVFPhKj7K2V7llehrpH0wGa-3XYvl6gPYc5H09Vw,14445
+haiku/rag/store/repositories/document.py,sha256=XoLCrMrZqs0iCZoHlDOfRDaVUux77Vdu5iZczduF1rY,7812
+haiku/rag/store/repositories/settings.py,sha256=wx3fuP_5CpPflZHRrIkeoer6ml-iD0qXERh5k6MQRzI,5291
 haiku/rag/store/upgrades/__init__.py,sha256=wUiEoSiHTahvuagx93E4FB07v123AhdbOjwUkPusiIg,14
-haiku_rag-0.7.7.dist-info/METADATA,sha256=JDihlKWhbK0GCvpesZOwqW1PnOychzHwMfsO-UIIfUY,4610
-haiku_rag-0.7.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.7.7.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.7.7.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.7.7.dist-info/RECORD,,
+haiku_rag-0.8.1.dist-info/METADATA,sha256=jPmTXHAXvT99zs0mFw_UXt80j8APLNkJuP7KspJaVro,4610
+haiku_rag-0.8.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.8.1.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.8.1.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.8.1.dist-info/RECORD,,