haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

haiku/rag/store/repositories/document.py

@@ -1,16 +1,14 @@
-import asyncio
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING
 from uuid import uuid4
 
 from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.document import Document
 
-if TYPE_CHECKING:
-    from docling_core.types.doc.document import DoclingDocument
 
-    from haiku.rag.store.models.chunk import Chunk
+def _escape_sql_string(value: str) -> str:
+    """Escape single quotes in SQL string literals."""
+    return value.replace("'", "''")
 
 
 class DocumentRepository:
@@ -37,6 +35,8 @@ class DocumentRepository:
             uri=record.uri,
             title=record.title,
             metadata=json.loads(record.metadata),
+            docling_document_json=record.docling_document_json,
+            docling_version=record.docling_version,
             created_at=datetime.fromisoformat(record.created_at)
             if record.created_at
             else datetime.now(),
@@ -47,6 +47,7 @@ class DocumentRepository:
 
     async def create(self, entity: Document) -> Document:
         """Create a document in the database."""
+        self.store._assert_writable()
         # Generate new UUID
         doc_id = str(uuid4())
 
@@ -60,6 +61,8 @@ class DocumentRepository:
             uri=entity.uri,
             title=entity.title,
             metadata=json.dumps(entity.metadata),
+            docling_document_json=entity.docling_document_json,
+            docling_version=entity.docling_version,
             created_at=now,
             updated_at=now,
         )
@@ -88,8 +91,14 @@ class DocumentRepository:
 
     async def update(self, entity: Document) -> Document:
         """Update an existing document."""
+        self.store._assert_writable()
+        from haiku.rag.store.models.document import invalidate_docling_document_cache
+
         assert entity.id, "Document ID is required for update"
 
+        # Invalidate cache before update
+        invalidate_docling_document_cache(entity.id)
+
         # Update timestamp
         now = datetime.now().isoformat()
         entity.updated_at = datetime.fromisoformat(now)
@@ -102,6 +111,8 @@ class DocumentRepository:
                 "uri": entity.uri,
                 "title": entity.title,
                 "metadata": json.dumps(entity.metadata),
+                "docling_document_json": entity.docling_document_json,
+                "docling_version": entity.docling_version,
                 "updated_at": now,
             },
         )
@@ -110,11 +121,17 @@ class DocumentRepository:
 
     async def delete(self, entity_id: str) -> bool:
         """Delete a document by its ID."""
+        self.store._assert_writable()
+        from haiku.rag.store.models.document import invalidate_docling_document_cache
+
         # Check if document exists
         doc = await self.get_by_id(entity_id)
         if doc is None:
             return False
 
+        # Invalidate cache before delete
+        invalidate_docling_document_cache(entity_id)
+
         # Delete associated chunks first
         await self.chunk_repository.delete_by_document_id(entity_id)
 
@@ -152,9 +169,10 @@ class DocumentRepository:
 
     async def get_by_uri(self, uri: str) -> Document | None:
         """Get a document by its URI."""
+        escaped_uri = _escape_sql_string(uri)
         results = list(
             self.store.documents_table.search()
-            .where(f"uri = '{uri}'")
+            .where(f"uri = '{escaped_uri}'")
             .limit(1)
             .to_pydantic(DocumentRecord)
         )
@@ -166,6 +184,7 @@ class DocumentRepository:
 
     async def delete_all(self) -> None:
         """Delete all documents from the database."""
+        self.store._assert_writable()
         # Delete all chunks first
         await self.chunk_repository.delete_all()
 
@@ -181,81 +200,3 @@ class DocumentRepository:
             self.store.documents_table = self.store.db.create_table(
                 "documents", schema=DocumentRecord
             )
-
-    async def _create_and_chunk(
-        self,
-        entity: Document,
-        docling_document: "DoclingDocument | None",
-        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)
-
-        # Attempt to create chunks; on failure, prefer version rollback
-        try:
-            # Create chunks if not provided
-            if chunks is None:
-                assert docling_document is not None, (
-                    "docling_document is required when chunks are not provided"
-                )
-                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.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
-            self.store.restore_table_versions(versions)
-            raise
-
-    async def _update_and_rechunk(
-        self, entity: Document, docling_document: "DoclingDocument"
-    ) -> Document:
-        """Update a document and regenerate its 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)
-
-        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
-            )
-
-            # 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
-            self.store.restore_table_versions(versions)
-            raise

haiku/rag/store/repositories/settings.py

@@ -72,6 +72,7 @@ class SettingsRepository:
 
     def save_current_settings(self) -> None:
         """Save the current configuration to the database."""
+        self.store._assert_writable()
         current_config = self.store._config.model_dump(mode="json")
 
         # Check if settings exist
@@ -118,25 +119,21 @@ class SettingsRepository:
         current_config = self.store._config.model_dump(mode="json")
 
         # Check if embedding provider or model has changed
-        # Support both old flat structure and new nested structure for backward compatibility
+        # Both stored and current use nested structure: embeddings.model.{provider,name,vector_dim}
         stored_embeddings = stored_settings.get("embeddings", {})
         current_embeddings = current_config.get("embeddings", {})
 
-        # Try nested structure first, fall back to flat for old databases
-        stored_provider = stored_embeddings.get("provider") or stored_settings.get(
-            "EMBEDDINGS_PROVIDER"
-        )
-        current_provider = current_embeddings.get("provider")
+        stored_model_obj = stored_embeddings.get("model", {})
+        current_model_obj = current_embeddings.get("model", {})
 
-        stored_model = stored_embeddings.get("model") or stored_settings.get(
-            "EMBEDDINGS_MODEL"
-        )
-        current_model = current_embeddings.get("model")
+        stored_provider = stored_model_obj.get("provider")
+        current_provider = current_model_obj.get("provider")
 
-        stored_vector_dim = stored_embeddings.get("vector_dim") or stored_settings.get(
-            "EMBEDDINGS_VECTOR_DIM"
-        )
-        current_vector_dim = current_embeddings.get("vector_dim")
+        stored_model = stored_model_obj.get("name")
+        current_model = current_model_obj.get("name")
+
+        stored_vector_dim = stored_model_obj.get("vector_dim")
+        current_vector_dim = current_model_obj.get("vector_dim")
 
         # Check for incompatible changes
         incompatible_changes = []

haiku/rag/store/upgrades/__init__.py

@@ -53,10 +53,26 @@ def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> No
         logger.info("Completed upgrade %s", step.version)
 
 
-from .v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts  # noqa: E402
-from .v0_9_3 import upgrade_order as upgrade_0_9_3_order  # noqa: E402
-from .v0_10_1 import upgrade_add_title as upgrade_0_10_1_add_title  # noqa: E402
+# Import upgrade modules AFTER Upgrade class is defined to avoid circular imports
+# ruff: noqa: E402, I001
+from haiku.rag.store.upgrades.v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts
+from haiku.rag.store.upgrades.v0_9_3 import upgrade_order as upgrade_0_9_3_order
+from haiku.rag.store.upgrades.v0_10_1 import (
+    upgrade_add_title as upgrade_0_10_1_add_title,
+)
+from haiku.rag.store.upgrades.v0_19_6 import (
+    upgrade_embeddings_model_config as upgrade_0_19_6_embeddings,
+)
+from haiku.rag.store.upgrades.v0_20_0 import (
+    upgrade_add_docling_document as upgrade_0_20_0_docling,
+)
+from haiku.rag.store.upgrades.v0_23_1 import (
+    upgrade_contextualize_chunks as upgrade_0_23_1_contextualize,
+)
 
 upgrades.append(upgrade_0_9_3_order)
 upgrades.append(upgrade_0_9_3_fts)
 upgrades.append(upgrade_0_10_1_add_title)
+upgrades.append(upgrade_0_19_6_embeddings)
+upgrades.append(upgrade_0_20_0_docling)
+upgrades.append(upgrade_0_23_1_contextualize)

haiku/rag/store/upgrades/v0_10_1.py

@@ -7,7 +7,7 @@ from haiku.rag.store.engine import Store
 from haiku.rag.store.upgrades import Upgrade
 
 
-def _apply_add_document_title(store: Store) -> None:
+def _apply_add_document_title(store: Store) -> None:  # pragma: no cover
     """Add a nullable 'title' column to the documents table."""
 
     # Read existing rows using Arrow for schema-agnostic access

haiku/rag/store/upgrades/v0_19_6.py

@@ -0,0 +1,65 @@
+import json
+import logging
+
+from haiku.rag.store.engine import SettingsRecord, Store
+from haiku.rag.store.upgrades import Upgrade
+
+logger = logging.getLogger(__name__)
+
+
+def _apply_embeddings_model_config(store: Store) -> None:  # pragma: no cover
+    """Migrate embeddings config from flat to nested EmbeddingModelConfig structure."""
+    results = list(
+        store.settings_table.search()
+        .where("id = 'settings'")
+        .limit(1)
+        .to_pydantic(SettingsRecord)
+    )
+
+    if not results or not results[0].settings:
+        return
+
+    settings = json.loads(results[0].settings)
+    embeddings = settings.get("embeddings", {})
+
+    # Check if already migrated (model is a dict with nested structure)
+    if isinstance(embeddings.get("model"), dict):
+        return
+
+    # Migrate from flat structure to nested EmbeddingModelConfig
+    old_provider = embeddings.get("provider", "ollama")
+    old_model = embeddings.get("model", "qwen3-embedding:4b")
+    old_vector_dim = embeddings.get("vector_dim", 2560)
+
+    logger.info(
+        "Migrating embeddings config to new nested structure: "
+        "embeddings.{provider,model,vector_dim} -> embeddings.model.{provider,name,vector_dim}"
+    )
+
+    # Create new nested structure
+    settings["embeddings"] = {
+        "model": {
+            "provider": old_provider,
+            "name": old_model,
+            "vector_dim": old_vector_dim,
+        }
+    }
+
+    store.settings_table.update(
+        where="id = 'settings'",
+        values={"settings": json.dumps(settings)},
+    )
+
+    logger.info(
+        "Embeddings config migrated: provider=%s, name=%s, vector_dim=%d",
+        old_provider,
+        old_model,
+        old_vector_dim,
+    )
+
+
+upgrade_embeddings_model_config = Upgrade(
+    version="0.19.6",
+    apply=_apply_embeddings_model_config,
+    description="Migrate embeddings config to nested EmbeddingModelConfig structure",
+)

haiku/rag/store/upgrades/v0_20_0.py

@@ -0,0 +1,68 @@
+import json
+
+from lancedb.pydantic import LanceModel
+from pydantic import Field
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.upgrades import Upgrade
+
+
+def _apply_add_docling_document_columns(store: Store) -> None:  # pragma: no cover
+    """Add 'docling_document_json' and 'docling_version' columns to documents table."""
+
+    # Read existing rows using Arrow for schema-agnostic access
+    try:
+        docs_arrow = store.documents_table.search().to_arrow()
+        rows = docs_arrow.to_pylist()
+    except Exception:
+        rows = []
+
+    class DocumentRecordV3(LanceModel):
+        id: str
+        content: str
+        uri: str | None = None
+        title: str | None = None
+        metadata: str = Field(default="{}")
+        docling_document_json: str | None = None
+        docling_version: str | None = None
+        created_at: str = Field(default_factory=lambda: "")
+        updated_at: str = Field(default_factory=lambda: "")
+
+    # Drop and recreate documents table with the new schema
+    try:
+        store.db.drop_table("documents")
+    except Exception:
+        pass
+
+    store.documents_table = store.db.create_table("documents", schema=DocumentRecordV3)
+
+    # Reinsert previous rows with new columns as None
+    if rows:
+        backfilled = []
+        for row in rows:
+            backfilled.append(
+                DocumentRecordV3(
+                    id=row.get("id"),
+                    content=row.get("content", ""),
+                    uri=row.get("uri"),
+                    title=row.get("title"),
+                    metadata=(
+                        row.get("metadata")
+                        if isinstance(row.get("metadata"), str)
+                        else json.dumps(row.get("metadata") or {})
+                    ),
+                    docling_document_json=None,
+                    docling_version=None,
+                    created_at=row.get("created_at", ""),
+                    updated_at=row.get("updated_at", ""),
+                )
+            )
+
+        store.documents_table.add(backfilled)
+
+
+upgrade_add_docling_document = Upgrade(
+    version="0.20.0",
+    apply=_apply_add_docling_document_columns,
+    description="Add 'docling_document_json' and 'docling_version' columns to documents table",
+)

haiku/rag/store/upgrades/v0_23_1.py

@@ -0,0 +1,100 @@
+import json
+
+from lancedb.pydantic import LanceModel, Vector
+from pydantic import Field
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.upgrades import Upgrade
+
+
+def _apply_add_content_fts(store: Store) -> None:  # pragma: no cover
+    """Add content_fts column with contextualized content for better FTS."""
+    # Read existing chunks
+    try:
+        chunks_arrow = store.chunks_table.search().to_arrow()
+        rows = chunks_arrow.to_pylist()
+    except Exception:
+        return
+
+    if not rows:
+        return
+
+    # Infer vector dimensions from first row
+    vec = rows[0].get("vector")
+    if not isinstance(vec, list) or not vec:
+        return
+    vector_dim = len(vec)
+
+    class ChunkRecord(LanceModel):
+        id: str
+        document_id: str
+        content: str
+        content_fts: str = Field(default="")
+        metadata: str = Field(default="{}")
+        order: int = Field(default=0)
+        vector: Vector(vector_dim) = Field(  # type: ignore
+            default_factory=lambda: [0.0] * vector_dim
+        )
+
+    # Drop and recreate table with new schema
+    try:
+        store.db.drop_table("chunks")
+    except Exception:
+        pass
+
+    store.chunks_table = store.db.create_table("chunks", schema=ChunkRecord)
+
+    # Populate content_fts with contextualized content
+    new_records: list[ChunkRecord] = []
+    for row in rows:
+        metadata_raw = row.get("metadata") or "{}"
+        try:
+            metadata = (
+                json.loads(metadata_raw)
+                if isinstance(metadata_raw, str)
+                else metadata_raw
+            )
+        except Exception:
+            metadata = {}
+
+        headings = metadata.get("headings") if isinstance(metadata, dict) else None
+        content = row.get("content", "")
+
+        # Build contextualized content for FTS
+        if headings:
+            content_fts = "\n".join(headings) + "\n" + content
+        else:
+            content_fts = content
+
+        new_records.append(
+            ChunkRecord(
+                id=row.get("id"),
+                document_id=row.get("document_id"),
+                content=content,
+                content_fts=content_fts,
+                metadata=metadata_raw,
+                order=row.get("order", 0),
+                vector=row.get("vector") or [0.0] * vector_dim,
+            )
+        )
+
+    if new_records:
+        store.chunks_table.add(new_records)
+
+    # Drop old FTS index on content column if it exists
+    try:
+        store.chunks_table.drop_index("content_idx")
+    except Exception:
+        pass
+
+    # Create FTS index on content_fts
+    store.chunks_table.create_fts_index(
+        "content_fts", replace=True, with_position=True, remove_stop_words=False
+    )
+
+
+upgrade_contextualize_chunks = Upgrade(
+    version="0.23.1",
+    apply=_apply_add_content_fts,
+    description="Add content_fts column for contextualized FTS search",
+)

haiku/rag/store/upgrades/v0_9_3.py

@@ -7,7 +7,7 @@ from haiku.rag.store.engine import Store
 from haiku.rag.store.upgrades import Upgrade
 
 
-def _infer_vector_dim(store: Store) -> int:
+def _infer_vector_dim(store: Store) -> int:  # pragma: no cover
     """Infer vector dimension from existing data; fallback to embedder config."""
     try:
         arrow = store.chunks_table.search().limit(1).to_arrow()
@@ -22,7 +22,7 @@ def _infer_vector_dim(store: Store) -> int:
     return getattr(store.embedder, "_vector_dim", 1024)
 
 
-def _apply_chunk_order(store: Store) -> None:
+def _apply_chunk_order(store: Store) -> None:  # pragma: no cover
     """Add integer 'order' column to chunks and backfill from metadata."""
 
     vector_dim = _infer_vector_dim(store)
@@ -95,7 +95,7 @@ upgrade_order = Upgrade(
 )
 
 
-def _apply_fts_phrase_support(store: Store) -> None:
+def _apply_fts_phrase_support(store: Store) -> None:  # pragma: no cover
     """Recreate FTS index with phrase query support and no stop-word removal."""
     try:
         store.chunks_table.create_fts_index(