haiku.rag 0.6.0__py3-none-any.whl → 0.7.1__py3-none-any.whl

haiku/rag/store/repositories/document.py

@@ -1,248 +1,214 @@
 import json
+from datetime import datetime
 from typing import TYPE_CHECKING
+from uuid import uuid4
 
 from docling_core.types.doc.document import DoclingDocument
 
+from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.document import Document
-from haiku.rag.store.repositories.base import BaseRepository
-from haiku.rag.utils import text_to_docling_document
 
 if TYPE_CHECKING:
     from haiku.rag.store.models.chunk import Chunk
 
 
-class DocumentRepository(BaseRepository[Document]):
-    """Repository for Document database operations."""
+class DocumentRepository:
+    """Repository for Document operations."""
 
-    def __init__(self, store, chunk_repository=None):
-        super().__init__(store)
-        # Avoid circular import by using late import if not provided
-        if chunk_repository is None:
+    def __init__(self, store: Store) -> None:
+        self.store = store
+        self._chunk_repository = None
+
+    @property
+    def chunk_repository(self):
+        """Lazy-load ChunkRepository when needed."""
+        if self._chunk_repository is None:
             from haiku.rag.store.repositories.chunk import ChunkRepository
 
-            chunk_repository = ChunkRepository(store)
-        self.chunk_repository = chunk_repository
+            self._chunk_repository = ChunkRepository(self.store)
+        return self._chunk_repository
 
-    async def _create_with_docling(
-        self,
-        entity: Document,
-        docling_document: DoclingDocument,
-        chunks: list["Chunk"] | None = None,
-    ) -> Document:
-        """Create a document with its chunks and embeddings."""
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
-
-        cursor = self.store._connection.cursor()
-
-        # Start transaction
-        cursor.execute("BEGIN TRANSACTION")
-
-        try:
-            # Insert the document
-            cursor.execute(
-                """
-                INSERT INTO documents (content, uri, metadata, created_at, updated_at)
-                VALUES (:content, :uri, :metadata, :created_at, :updated_at)
-                """,
-                {
-                    "content": entity.content,
-                    "uri": entity.uri,
-                    "metadata": json.dumps(entity.metadata),
-                    "created_at": entity.created_at,
-                    "updated_at": entity.updated_at,
-                },
-            )
-
-            document_id = cursor.lastrowid
-            assert document_id is not None, "Failed to create document in database"
-            entity.id = document_id
-
-            # Create chunks - either use provided chunks or generate from content
-            if chunks is not None:
-                # Use provided chunks, but update their document_id and set order from list position
-                for order, chunk in enumerate(chunks):
-                    chunk.document_id = document_id
-                    # Ensure order is set from list position
-                    chunk.metadata = chunk.metadata.copy() if chunk.metadata else {}
-                    chunk.metadata["order"] = order
-                    await self.chunk_repository.create(chunk, commit=False)
-            else:
-                # Create chunks and embeddings using DoclingDocument
-                await self.chunk_repository.create_chunks_for_document(
-                    document_id, docling_document, commit=False
-                )
-
-            cursor.execute("COMMIT")
-            return entity
-
-        except Exception:
-            cursor.execute("ROLLBACK")
-            raise
+    def _record_to_document(self, record: DocumentRecord) -> Document:
+        """Convert a DocumentRecord to a Document model."""
+        return Document(
+            id=record.id,
+            content=record.content,
+            uri=record.uri,
+            metadata=json.loads(record.metadata) if record.metadata else {},
+            created_at=datetime.fromisoformat(record.created_at)
+            if record.created_at
+            else datetime.now(),
+            updated_at=datetime.fromisoformat(record.updated_at)
+            if record.updated_at
+            else datetime.now(),
+        )
 
     async def create(self, entity: Document) -> Document:
-        """Create a document with its chunks and embeddings."""
-        # Convert content to DoclingDocument
-        docling_document = text_to_docling_document(entity.content)
+        """Create a document in the database."""
+        # Generate new UUID
+        doc_id = str(uuid4())
+
+        # Create timestamp
+        now = datetime.now().isoformat()
+
+        # Create document record
+        doc_record = DocumentRecord(
+            id=doc_id,
+            content=entity.content,
+            uri=entity.uri,
+            metadata=json.dumps(entity.metadata),
+            created_at=now,
+            updated_at=now,
+        )
+
+        # Add to table
+        self.store.documents_table.add([doc_record])
 
-        return await self._create_with_docling(entity, docling_document)
+        entity.id = doc_id
+        entity.created_at = datetime.fromisoformat(now)
+        entity.updated_at = datetime.fromisoformat(now)
+        return entity
 
-    async def get_by_id(self, entity_id: int) -> Document | None:
+    async def get_by_id(self, entity_id: str) -> Document | None:
         """Get a document by its ID."""
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
-
-        cursor = self.store._connection.cursor()
-        cursor.execute(
-            """
-            SELECT id, content, uri, metadata, created_at, updated_at
-            FROM documents WHERE id = :id
-            """,
-            {"id": entity_id},
+        results = list(
+            self.store.documents_table.search()
+            .where(f"id = '{entity_id}'")
+            .limit(1)
+            .to_pydantic(DocumentRecord)
         )
 
-        row = cursor.fetchone()
-        if row is None:
+        if not results:
             return None
 
-        document_id, content, uri, metadata_json, created_at, updated_at = row
-        metadata = json.loads(metadata_json) if metadata_json else {}
-
-        return Document(
-            id=document_id,
-            content=content,
-            uri=uri,
-            metadata=metadata,
-            created_at=created_at,
-            updated_at=updated_at,
-        )
+        return self._record_to_document(results[0])
 
-    async def get_by_uri(self, uri: str) -> Document | None:
-        """Get a document by its URI."""
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
-
-        cursor = self.store._connection.cursor()
-        cursor.execute(
-            """
-            SELECT id, content, uri, metadata, created_at, updated_at
-            FROM documents WHERE uri = :uri
-            """,
-            {"uri": uri},
+    async def update(self, entity: Document) -> Document:
+        """Update an existing document."""
+        assert entity.id, "Document ID is required for update"
+
+        # Update timestamp
+        now = datetime.now().isoformat()
+        entity.updated_at = datetime.fromisoformat(now)
+
+        # Update the record
+        self.store.documents_table.update(
+            where=f"id = '{entity.id}'",
+            values={
+                "content": entity.content,
+                "uri": entity.uri,
+                "metadata": json.dumps(entity.metadata),
+                "updated_at": now,
+            },
         )
 
-        row = cursor.fetchone()
-        if row is None:
-            return None
+        return entity
 
-        document_id, content, uri, metadata_json, created_at, updated_at = row
-        metadata = json.loads(metadata_json) if metadata_json else {}
+    async def delete(self, entity_id: str) -> bool:
+        """Delete a document by its ID."""
+        # Check if document exists
+        doc = await self.get_by_id(entity_id)
+        if doc is None:
+            return False
 
-        return Document(
-            id=document_id,
-            content=content,
-            uri=uri,
-            metadata=metadata,
-            created_at=created_at,
-            updated_at=updated_at,
-        )
+        # Delete associated chunks first
+        await self.chunk_repository.delete_by_document_id(entity_id)
 
-    async def _update_with_docling(
-        self, entity: Document, docling_document: DoclingDocument
-    ) -> Document:
-        """Update an existing document and regenerate its chunks and embeddings."""
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
-        if entity.id is None:
-            raise ValueError("Document ID is required for update")
-
-        cursor = self.store._connection.cursor()
-
-        # Start transaction
-        cursor.execute("BEGIN TRANSACTION")
-
-        try:
-            # Update the document
-            cursor.execute(
-                """
-                UPDATE documents
-                SET content = :content, uri = :uri, metadata = :metadata, updated_at = :updated_at
-                WHERE id = :id
-                """,
-                {
-                    "content": entity.content,
-                    "uri": entity.uri,
-                    "metadata": json.dumps(entity.metadata),
-                    "updated_at": entity.updated_at,
-                    "id": entity.id,
-                },
-            )
+        # Delete the document
+        self.store.documents_table.delete(f"id = '{entity_id}'")
+        return True
 
-            # Delete existing chunks and regenerate using DoclingDocument
-            await self.chunk_repository.delete_by_document_id(entity.id, commit=False)
-            await self.chunk_repository.create_chunks_for_document(
-                entity.id, docling_document, commit=False
-            )
+    async def list_all(
+        self, limit: int | None = None, offset: int | None = None
+    ) -> list[Document]:
+        """List all documents with optional pagination."""
+        query = self.store.documents_table.search()
 
-            cursor.execute("COMMIT")
-            return entity
+        if offset is not None:
+            query = query.offset(offset)
+        if limit is not None:
+            query = query.limit(limit)
 
-        except Exception:
-            cursor.execute("ROLLBACK")
-            raise
+        results = list(query.to_pydantic(DocumentRecord))
+        return [self._record_to_document(doc) for doc in results]
 
-    async def update(self, entity: Document) -> Document:
-        """Update an existing document and regenerate its chunks and embeddings."""
-        # Convert content to DoclingDocument
-        docling_document = text_to_docling_document(entity.content)
+    async def get_by_uri(self, uri: str) -> Document | None:
+        """Get a document by its URI."""
+        results = list(
+            self.store.documents_table.search()
+            .where(f"uri = '{uri}'")
+            .limit(1)
+            .to_pydantic(DocumentRecord)
+        )
 
-        return await self._update_with_docling(entity, docling_document)
+        if not results:
+            return None
 
-    async def delete(self, entity_id: int) -> bool:
-        """Delete a document and all its associated chunks and embeddings."""
-        # Delete chunks and embeddings first
-        await self.chunk_repository.delete_by_document_id(entity_id)
+        return self._record_to_document(results[0])
 
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
+    async def delete_all(self) -> None:
+        """Delete all documents from the database."""
+        # Delete all chunks first
+        await self.chunk_repository.delete_all()
 
-        cursor = self.store._connection.cursor()
-        cursor.execute("DELETE FROM documents WHERE id = :id", {"id": entity_id})
+        # Get count before deletion
+        count = len(
+            list(
+                self.store.documents_table.search().limit(1).to_pydantic(DocumentRecord)
+            )
+        )
+        if count > 0:
+            # Drop and recreate table to clear all data
+            self.store.db.drop_table("documents")
+            self.store.documents_table = self.store.db.create_table(
+                "documents", schema=DocumentRecord
+            )
 
-        deleted = cursor.rowcount > 0
-        self.store._connection.commit()
-        return deleted
+    async def _create_with_docling(
+        self,
+        entity: Document,
+        docling_document: DoclingDocument,
+        chunks: list["Chunk"] | None = None,
+    ) -> Document:
+        """Create a document with its chunks and embeddings."""
+        # Create the document
+        created_doc = await self.create(entity)
 
-    async def list_all(
-        self, limit: int | None = None, offset: int | None = None
-    ) -> list[Document]:
-        """List all documents with optional pagination."""
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
+        # 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)
 
-        cursor = self.store._connection.cursor()
-        query = "SELECT id, content, uri, metadata, created_at, updated_at FROM documents ORDER BY created_at DESC"
-        params = {}
+        return created_doc
 
-        if limit is not None:
-            query += " LIMIT :limit"
-            params["limit"] = limit
+    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"
+        await self.chunk_repository.delete_by_document_id(entity.id)
+
+        # 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
+        )
 
-        if offset is not None:
-            query += " OFFSET :offset"
-            params["offset"] = offset
-
-        cursor.execute(query, params)
-        rows = cursor.fetchall()
-
-        return [
-            Document(
-                id=document_id,
-                content=content,
-                uri=uri,
-                metadata=json.loads(metadata_json) if metadata_json else {},
-                created_at=created_at,
-                updated_at=updated_at,
-            )
-            for document_id, content, uri, metadata_json, created_at, updated_at in rows
-        ]
+        return updated_doc

haiku/rag/store/repositories/settings.py

@@ -1,77 +1,143 @@
 import json
-from typing import Any
 
-from haiku.rag.store.engine import Store
+from haiku.rag.config import Config
+from haiku.rag.store.engine import SettingsRecord, Store
 
 
 class ConfigMismatchError(Exception):
-    """Raised when current config doesn't match stored settings."""
+    """Raised when stored config doesn't match current config."""
 
     pass
 
 
 class SettingsRepository:
-    def __init__(self, store: Store):
+    """Repository for Settings operations."""
+
+    def __init__(self, store: Store) -> None:
         self.store = store
 
-    def get(self) -> dict[str, Any]:
-        """Get all settings from the database."""
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
+    async def create(self, entity: dict) -> dict:
+        """Create settings in the database."""
+        settings_record = SettingsRecord(id="settings", settings=json.dumps(entity))
+        self.store.settings_table.add([settings_record])
+        return entity
+
+    async def get_by_id(self, entity_id: str) -> dict | None:
+        """Get settings by ID."""
+        results = list(
+            self.store.settings_table.search()
+            .where(f"id = '{entity_id}'")
+            .limit(1)
+            .to_pydantic(SettingsRecord)
+        )
 
-        cursor = self.store._connection.execute("SELECT settings FROM settings LIMIT 1")
-        row = cursor.fetchone()
-        if row:
-            return json.loads(row[0])
-        return {}
+        if not results:
+            return None
 
-    def save(self) -> None:
-        """Sync settings from the current AppConfig to database."""
-        if self.store._connection is None:
-            raise ValueError("Store connection is not available")
+        return json.loads(results[0].settings) if results[0].settings else {}
 
-        from haiku.rag.config import Config
+    async def update(self, entity: dict) -> dict:
+        """Update existing settings."""
+        self.store.settings_table.update(
+            where="id = 'settings'", values={"settings": json.dumps(entity)}
+        )
+        return entity
+
+    async def delete(self, entity_id: str) -> bool:
+        """Delete settings by ID."""
+        self.store.settings_table.delete(f"id = '{entity_id}'")
+        return True
+
+    async def list_all(
+        self, limit: int | None = None, offset: int | None = None
+    ) -> list[dict]:
+        """List all settings."""
+        results = list(self.store.settings_table.search().to_pydantic(SettingsRecord))
+        return [
+            json.loads(record.settings) if record.settings else {} for record in results
+        ]
 
-        settings_json = Config.model_dump_json()
+    def get_current_settings(self) -> dict:
+        """Get the current settings."""
+        results = list(
+            self.store.settings_table.search()
+            .where("id = 'settings'")
+            .limit(1)
+            .to_pydantic(SettingsRecord)
+        )
+
+        if not results:
+            return {}
+
+        return json.loads(results[0].settings) if results[0].settings else {}
+
+    def save_current_settings(self) -> None:
+        """Save the current configuration to the database."""
+        current_config = Config.model_dump(mode="json")
 
-        self.store._connection.execute(
-            "INSERT INTO settings (id, settings) VALUES (1, ?) ON CONFLICT(id) DO UPDATE SET settings = excluded.settings",
-            (settings_json,),
+        # Check if settings exist
+        existing = list(
+            self.store.settings_table.search()
+            .where("id = 'settings'")
+            .limit(1)
+            .to_pydantic(SettingsRecord)
         )
 
-        self.store._connection.commit()
+        if existing:
+            # Update existing settings
+            self.store.settings_table.update(
+                where="id = 'settings'", values={"settings": json.dumps(current_config)}
+            )
+        else:
+            # Create new settings
+            settings_record = SettingsRecord(
+                id="settings", settings=json.dumps(current_config)
+            )
+            self.store.settings_table.add([settings_record])
 
     def validate_config_compatibility(self) -> None:
-        """Check if current config is compatible with stored settings.
-
-        Raises ConfigMismatchError if there are incompatible differences.
-        If no settings exist, saves current config.
-        """
-        db_settings = self.get()
-        if not db_settings:
-            # No settings in DB, save current config
-            self.save()
-            return
+        """Validate that the current configuration is compatible with stored settings."""
+        stored_settings = self.get_current_settings()
 
-        from haiku.rag.config import Config
+        # If no stored settings, this is a new database - save current config and return
+        if not stored_settings:
+            self.save_current_settings()
+            return
 
         current_config = Config.model_dump(mode="json")
 
-        # Critical settings that must match
-        critical_settings = [
-            "EMBEDDINGS_PROVIDER",
-            "EMBEDDINGS_MODEL",
-            "EMBEDDINGS_VECTOR_DIM",
-            "CHUNK_SIZE",
-        ]
+        # Check if embedding provider or model has changed
+        stored_provider = stored_settings.get("EMBEDDINGS_PROVIDER")
+        current_provider = current_config.get("EMBEDDINGS_PROVIDER")
+
+        stored_model = stored_settings.get("EMBEDDINGS_MODEL")
+        current_model = current_config.get("EMBEDDINGS_MODEL")
+
+        stored_vector_dim = stored_settings.get("EMBEDDINGS_VECTOR_DIM")
+        current_vector_dim = current_config.get("EMBEDDINGS_VECTOR_DIM")
+
+        # Check for incompatible changes
+        incompatible_changes = []
+
+        if stored_provider and stored_provider != current_provider:
+            incompatible_changes.append(
+                f"Embedding provider changed from '{stored_provider}' to '{current_provider}'"
+            )
+
+        if stored_model and stored_model != current_model:
+            incompatible_changes.append(
+                f"Embedding model changed from '{stored_model}' to '{current_model}'"
+            )
 
-        errors = []
-        for setting in critical_settings:
-            if db_settings.get(setting) != current_config.get(setting):
-                errors.append(
-                    f"{setting}: current={current_config.get(setting)}, stored={db_settings.get(setting)}"
-                )
+        if stored_vector_dim and stored_vector_dim != current_vector_dim:
+            incompatible_changes.append(
+                f"Vector dimension changed from {stored_vector_dim} to {current_vector_dim}"
+            )
 
-        if errors:
-            error_msg = f"Config mismatch detected: {'; '.join(errors)}. Consider rebuilding the database with the current configuration."
+        if incompatible_changes:
+            error_msg = (
+                "Database configuration is incompatible with current settings:\n"
+                + "\n".join(f"  - {change}" for change in incompatible_changes)
+            )
+            error_msg += "\n\nPlease rebuild the database using: haiku-rag rebuild"
             raise ConfigMismatchError(error_msg)

haiku/rag/store/upgrades/__init__.py

@@ -1,3 +1 @@
-from haiku.rag.store.upgrades.v0_3_4 import upgrades as v0_3_4_upgrades
-
-upgrades = v0_3_4_upgrades
+upgrades = []