Use db repositories

src/haiku/rag/embeddings/base.py

@@ -1,11 +1,7 @@
-import functools
-
-
 class EmbedderBase:
     _model: str = ""
     _vector_dim: int = 0
 
-    @functools.lru_cache(maxsize=128)
     async def embed(self, text: str) -> list[float]:
         raise NotImplementedError(
             "Embedder is an abstract class. Please implement the embed method in a subclass."

src/haiku/rag/embeddings/ollama.py

@@ -1,5 +1,3 @@
-import functools
-
 from ollama import AsyncClient
 
 from haiku.rag.config import Config
@@ -10,7 +8,6 @@ class Embedder(EmbedderBase):
     _model: str = Config.EMBEDDING_MODEL
     _vector_dim: int = 1024
 
-    @functools.lru_cache(maxsize=128)
     async def embed(self, text: str) -> list[float]:
         client = AsyncClient(host=Config.OLLAMA_BASE_URL)
         res = await client.embeddings(model=self._model, prompt=text)

src/haiku/rag/store/models/document.py

@@ -1,16 +1,7 @@
-import json
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, Field
 
-from haiku.rag.chunker import chunker
-from haiku.rag.embeddings.ollama import Embedder
-from haiku.rag.store.models.chunk import Chunk
-
-if TYPE_CHECKING:
-    from haiku.rag.store.engine import Store
-
 
 class Document(BaseModel):
     """
@@ -22,124 +13,3 @@ class Document(BaseModel):
     metadata: dict = {}
     created_at: datetime = Field(default_factory=datetime.now)
     updated_at: datetime = Field(default_factory=datetime.now)
-
-    async def create_with_chunks(self, store: "Store") -> "Document":
-        """
-        Create a document in the database along with its chunks and embeddings.
-
-        Args:
-            store: The Store instance to use for database operations
-
-        Returns:
-            Document: The created document with updated id
-        """
-        if store._connection is None:
-            raise ValueError("Store connection is not available")
-
-        cursor = store._connection.cursor()
-        embedder = Embedder()
-
-        # Insert the document
-        cursor.execute(
-            """
-            INSERT INTO documents (content, metadata, created_at, updated_at)
-            VALUES (?, ?, ?, ?)
-            """,
-            (
-                self.content,
-                json.dumps(self.metadata),
-                self.created_at,
-                self.updated_at,
-            ),
-        )
-
-        document_id = cursor.lastrowid
-        assert document_id is not None, "Failed to create document in database"
-        self.id = document_id
-
-        # Chunk the document content
-        chunk_texts = await chunker.chunk(self.content)
-
-        # Create chunks with embeddings
-        for order, chunk_text in enumerate(chunk_texts):
-            # Create chunk with order in metadata
-            chunk = Chunk(
-                document_id=document_id, content=chunk_text, metadata={"order": order}
-            )
-
-            cursor.execute(
-                """
-                INSERT INTO chunks (document_id, content, metadata)
-                VALUES (?, ?, ?)
-                """,
-                (chunk.document_id, chunk.content, json.dumps(chunk.metadata)),
-            )
-            chunk_id = cursor.lastrowid
-
-            # Generate and store embedding
-            embedding = await embedder.embed(chunk_text)
-            serialized_embedding = store.serialize_embedding(embedding)
-            cursor.execute(
-                """
-                INSERT INTO chunk_embeddings (chunk_id, embedding)
-                VALUES (?, ?)
-                """,
-                (chunk_id, serialized_embedding),
-            )
-
-        store._connection.commit()
-        return self
-
-    @classmethod
-    async def search_chunks(
-        cls, store: "Store", query: str, limit: int = 5
-    ) -> list[Chunk]:
-        """
-        Search for relevant chunks using vector similarity with sqlite-vec.
-
-        Args:
-            store: The Store instance to use for database operations
-            query: The text query to search for
-            limit: Maximum number of chunks to return
-
-        Returns:
-            List of relevant Chunk objects ordered by similarity
-        """
-        if store._connection is None:
-            raise ValueError("Store connection is not available")
-
-        embedder = Embedder()
-        cursor = store._connection.cursor()
-
-        # Generate embedding for the query
-        query_embedding = await embedder.embed(query)
-        serialized_query_embedding = store.serialize_embedding(query_embedding)
-
-        # Search for similar chunks using sqlite-vec
-        cursor.execute(
-            """
-            SELECT c.id, c.document_id, c.content, c.metadata, distance
-            FROM chunk_embeddings
-            JOIN chunks c ON c.id = chunk_embeddings.chunk_id
-            WHERE embedding MATCH ? AND k = ?
-            ORDER BY distance
-            """,
-            (serialized_query_embedding, limit),
-        )
-
-        results = cursor.fetchall()
-        chunks = []
-
-        for row in results:
-            chunk_id, document_id, content, metadata_json, distance = row
-            metadata = json.loads(metadata_json) if metadata_json else {}
-            chunks.append(
-                Chunk(
-                    id=chunk_id,
-                    document_id=document_id,
-                    content=content,
-                    metadata=metadata,
-                )
-            )
-
-        return chunks

src/haiku/rag/store/repositories/__init__.py

@@ -0,0 +1,5 @@
+from haiku.rag.store.repositories.base import BaseRepository
+from haiku.rag.store.repositories.chunk import ChunkRepository
+from haiku.rag.store.repositories.document import DocumentRepository
+
+__all__ = ["BaseRepository", "DocumentRepository", "ChunkRepository"]

src/haiku/rag/store/repositories/base.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from haiku.rag.store.engine import Store
+
+T = TypeVar("T")
+
+
+class BaseRepository(ABC, Generic[T]):
+    """Base repository interface for database operations."""
+
+    def __init__(self, store: Store):
+        self.store = store
+
+    @abstractmethod
+    async def create(self, entity: T) -> T:
+        """Create a new entity in the database."""
+        pass
+
+    @abstractmethod
+    async def get_by_id(self, entity_id: int) -> T | None:
+        """Get an entity by its ID."""
+        pass
+
+    @abstractmethod
+    async def update(self, entity: T) -> T:
+        """Update an existing entity."""
+        pass
+
+    @abstractmethod
+    async def delete(self, entity_id: int) -> bool:
+        """Delete an entity by its ID."""
+        pass
+
+    @abstractmethod
+    async def list_all(
+        self, limit: int | None = None, offset: int | None = None
+    ) -> list[T]:
+        """List all entities with optional pagination."""
+        pass

src/haiku/rag/store/repositories/chunk.py

@@ -0,0 +1,270 @@
+import json
+
+from haiku.rag.chunker import chunker
+from haiku.rag.embeddings.ollama import Embedder
+from haiku.rag.store.models.chunk import Chunk
+from haiku.rag.store.repositories.base import BaseRepository
+
+
+class ChunkRepository(BaseRepository[Chunk]):
+    """Repository for Chunk database operations."""
+
+    def __init__(self, store, embedder: Embedder | None = None):
+        super().__init__(store)
+        self.embedder = embedder or Embedder()
+
+    async def create(self, entity: Chunk, commit: bool = True) -> Chunk:
+        """Create a chunk in the database."""
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self.store._connection.cursor()
+        cursor.execute(
+            """
+            INSERT INTO chunks (document_id, content, metadata)
+            VALUES (?, ?, ?)
+            """,
+            (entity.document_id, entity.content, json.dumps(entity.metadata)),
+        )
+
+        entity.id = cursor.lastrowid
+
+        # Generate and store embedding
+        embedding = await self.embedder.embed(entity.content)
+        serialized_embedding = self.store.serialize_embedding(embedding)
+        cursor.execute(
+            """
+            INSERT INTO chunk_embeddings (chunk_id, embedding)
+            VALUES (?, ?)
+            """,
+            (entity.id, serialized_embedding),
+        )
+
+        if commit:
+            self.store._connection.commit()
+        return entity
+
+    async def get_by_id(self, entity_id: int) -> Chunk | None:
+        """Get a chunk 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, document_id, content, metadata
+            FROM chunks WHERE id = ?
+            """,
+            (entity_id,),
+        )
+
+        row = cursor.fetchone()
+        if row is None:
+            return None
+
+        chunk_id, document_id, content, metadata_json = row
+        metadata = json.loads(metadata_json) if metadata_json else {}
+
+        return Chunk(
+            id=chunk_id, document_id=document_id, content=content, metadata=metadata
+        )
+
+    async def update(self, entity: Chunk) -> Chunk:
+        """Update an existing chunk."""
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+        if entity.id is None:
+            raise ValueError("Chunk ID is required for update")
+
+        cursor = self.store._connection.cursor()
+        cursor.execute(
+            """
+            UPDATE chunks
+            SET document_id = ?, content = ?, metadata = ?
+            WHERE id = ?
+            """,
+            (
+                entity.document_id,
+                entity.content,
+                json.dumps(entity.metadata),
+                entity.id,
+            ),
+        )
+
+        # Regenerate and update embedding
+        embedding = await self.embedder.embed(entity.content)
+        serialized_embedding = self.store.serialize_embedding(embedding)
+        cursor.execute(
+            """
+            UPDATE chunk_embeddings
+            SET embedding = ?
+            WHERE chunk_id = ?
+            """,
+            (serialized_embedding, entity.id),
+        )
+
+        self.store._connection.commit()
+        return entity
+
+    async def delete(self, entity_id: int, commit: bool = True) -> bool:
+        """Delete a chunk by its ID."""
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self.store._connection.cursor()
+
+        # Delete the embedding first
+        cursor.execute("DELETE FROM chunk_embeddings WHERE chunk_id = ?", (entity_id,))
+
+        # Delete the chunk
+        cursor.execute("DELETE FROM chunks WHERE id = ?", (entity_id,))
+
+        deleted = cursor.rowcount > 0
+        if commit:
+            self.store._connection.commit()
+        return deleted
+
+    async def list_all(
+        self, limit: int | None = None, offset: int | None = None
+    ) -> list[Chunk]:
+        """List all chunks with optional pagination."""
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self.store._connection.cursor()
+        query = "SELECT id, document_id, content, metadata FROM chunks ORDER BY document_id, id"
+        params = []
+
+        if limit is not None:
+            query += " LIMIT ?"
+            params.append(limit)
+
+        if offset is not None:
+            query += " OFFSET ?"
+            params.append(offset)
+
+        cursor.execute(query, params)
+        rows = cursor.fetchall()
+
+        chunks = []
+        for row in rows:
+            chunk_id, document_id, content, metadata_json = row
+            metadata = json.loads(metadata_json) if metadata_json else {}
+            chunks.append(
+                Chunk(
+                    id=chunk_id,
+                    document_id=document_id,
+                    content=content,
+                    metadata=metadata,
+                )
+            )
+
+        return chunks
+
+    async def create_chunks_for_document(
+        self, document_id: int, content: str, commit: bool = True
+    ) -> list[Chunk]:
+        """Create chunks and embeddings for a document."""
+        # Chunk the document content
+        chunk_texts = await chunker.chunk(content)
+        created_chunks = []
+
+        # Create chunks with embeddings using the create method
+        for order, chunk_text in enumerate(chunk_texts):
+            # Create chunk with order in metadata
+            chunk = Chunk(
+                document_id=document_id, content=chunk_text, metadata={"order": order}
+            )
+
+            created_chunk = await self.create(chunk, commit=commit)
+            created_chunks.append(created_chunk)
+
+        return created_chunks
+
+    async def delete_by_document_id(
+        self, document_id: int, commit: bool = True
+    ) -> bool:
+        """Delete all chunks for a document."""
+        chunks = await self.get_by_document_id(document_id)
+
+        deleted_any = False
+        for chunk in chunks:
+            if chunk.id is not None:
+                deleted = await self.delete(chunk.id, commit=False)
+                deleted_any = deleted_any or deleted
+
+        if commit and deleted_any and self.store._connection:
+            self.store._connection.commit()
+        return deleted_any
+
+    async def search_chunks(self, query: str, limit: int = 5) -> list[Chunk]:
+        """Search for relevant chunks using vector similarity."""
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self.store._connection.cursor()
+
+        # Generate embedding for the query
+        query_embedding = await self.embedder.embed(query)
+        serialized_query_embedding = self.store.serialize_embedding(query_embedding)
+
+        # Search for similar chunks using sqlite-vec
+        cursor.execute(
+            """
+            SELECT c.id, c.document_id, c.content, c.metadata, distance
+            FROM chunk_embeddings
+            JOIN chunks c ON c.id = chunk_embeddings.chunk_id
+            WHERE embedding MATCH ? AND k = ?
+            ORDER BY distance
+            """,
+            (serialized_query_embedding, limit),
+        )
+
+        results = cursor.fetchall()
+        chunks = []
+
+        for row in results:
+            chunk_id, document_id, content, metadata_json, _ = row
+            metadata = json.loads(metadata_json) if metadata_json else {}
+            chunks.append(
+                Chunk(
+                    id=chunk_id,
+                    document_id=document_id,
+                    content=content,
+                    metadata=metadata,
+                )
+            )
+
+        return chunks
+
+    async def get_by_document_id(self, document_id: int) -> list[Chunk]:
+        """Get all chunks for a specific document."""
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self.store._connection.cursor()
+        cursor.execute(
+            """
+            SELECT id, document_id, content, metadata
+            FROM chunks WHERE document_id = ?
+            ORDER BY JSON_EXTRACT(metadata, '$.order')
+            """,
+            (document_id,),
+        )
+
+        rows = cursor.fetchall()
+        chunks = []
+
+        for row in rows:
+            chunk_id, document_id, content, metadata_json = row
+            metadata = json.loads(metadata_json) if metadata_json else {}
+            chunks.append(
+                Chunk(
+                    id=chunk_id,
+                    document_id=document_id,
+                    content=content,
+                    metadata=metadata,
+                )
+            )
+
+        return chunks

src/haiku/rag/store/repositories/document.py

@@ -0,0 +1,181 @@
+import json
+
+from haiku.rag.store.models.document import Document
+from haiku.rag.store.repositories.base import BaseRepository
+
+
+class DocumentRepository(BaseRepository[Document]):
+    """Repository for Document database 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:
+            from haiku.rag.store.repositories.chunk import ChunkRepository
+
+            chunk_repository = ChunkRepository(store)
+        self.chunk_repository = chunk_repository
+
+    async def create(self, entity: Document) -> 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, metadata, created_at, updated_at)
+                VALUES (?, ?, ?, ?)
+                """,
+                (
+                    entity.content,
+                    json.dumps(entity.metadata),
+                    entity.created_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 and embeddings using ChunkRepository
+            await self.chunk_repository.create_chunks_for_document(
+                document_id, entity.content, commit=False
+            )
+
+            cursor.execute("COMMIT")
+            return entity
+
+        except Exception:
+            cursor.execute("ROLLBACK")
+            raise
+
+    async def get_by_id(self, entity_id: int) -> 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, metadata, created_at, updated_at
+            FROM documents WHERE id = ?
+            """,
+            (entity_id,),
+        )
+
+        row = cursor.fetchone()
+        if row is None:
+            return None
+
+        document_id, content, metadata_json, created_at, updated_at = row
+        metadata = json.loads(metadata_json) if metadata_json else {}
+
+        return Document(
+            id=document_id,
+            content=content,
+            metadata=metadata,
+            created_at=created_at,
+            updated_at=updated_at,
+        )
+
+    async def update(self, entity: Document) -> 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 = ?, metadata = ?, updated_at = ?
+                WHERE id = ?
+                """,
+                (
+                    entity.content,
+                    json.dumps(entity.metadata),
+                    entity.updated_at,
+                    entity.id,
+                ),
+            )
+
+            # Delete existing chunks and regenerate using ChunkRepository
+            await self.chunk_repository.delete_by_document_id(entity.id, commit=False)
+            await self.chunk_repository.create_chunks_for_document(
+                entity.id, entity.content, commit=False
+            )
+
+            cursor.execute("COMMIT")
+            return entity
+
+        except Exception:
+            cursor.execute("ROLLBACK")
+            raise
+
+    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)
+
+        if self.store._connection is None:
+            raise ValueError("Store connection is not available")
+
+        cursor = self.store._connection.cursor()
+        cursor.execute("DELETE FROM documents WHERE id = ?", (entity_id,))
+
+        deleted = cursor.rowcount > 0
+        self.store._connection.commit()
+        return deleted
+
+    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")
+
+        cursor = self.store._connection.cursor()
+        query = "SELECT id, content, metadata, created_at, updated_at FROM documents ORDER BY created_at DESC"
+        params = []
+
+        if limit is not None:
+            query += " LIMIT ?"
+            params.append(limit)
+
+        if offset is not None:
+            query += " OFFSET ?"
+            params.append(offset)
+
+        cursor.execute(query, params)
+        rows = cursor.fetchall()
+
+        documents = []
+        for row in rows:
+            document_id, content, metadata_json, created_at, updated_at = row
+            metadata = json.loads(metadata_json) if metadata_json else {}
+            documents.append(
+                Document(
+                    id=document_id,
+                    content=content,
+                    metadata=metadata,
+                    created_at=created_at,
+                    updated_at=updated_at,
+                )
+            )
+
+        return documents

tests/test_chunk.py

@@ -0,0 +1,186 @@
+import pytest
+from datasets import Dataset
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.models.chunk import Chunk
+from haiku.rag.store.models.document import Document
+from haiku.rag.store.repositories.chunk import ChunkRepository
+from haiku.rag.store.repositories.document import DocumentRepository
+
+
+@pytest.mark.asyncio
+async def test_search_chunks(qa_corpus: Dataset):
+    """Test vector search functionality using ChunkRepository."""
+    # Create an in-memory store and repositories
+    store = Store(":memory:")
+    doc_repo = DocumentRepository(store)
+    chunk_repo = ChunkRepository(store)
+
+    # Get the first document from the corpus
+    first_doc = qa_corpus[0]
+    document_text = first_doc["document_extracted"]
+
+    # Create and store a document
+    document = Document(content=document_text, metadata={"source": "qa_corpus"})
+    created_document = await doc_repo.create(document)
+
+    # Perform a search using ChunkRepository
+    search_query = "news"  # Simple query
+    results = await chunk_repo.search_chunks(search_query, limit=3)
+
+    # Verify search results
+    assert len(results) <= 3
+    assert all(hasattr(chunk, "content") for chunk in results)
+    assert all(hasattr(chunk, "document_id") for chunk in results)
+    assert all(chunk.document_id == created_document.id for chunk in results)
+
+    store.close()
+
+
+@pytest.mark.asyncio
+async def test_chunk_repository_operations(qa_corpus: Dataset):
+    """Test ChunkRepository operations."""
+    # Create an in-memory store and repositories
+    store = Store(":memory:")
+    doc_repo = DocumentRepository(store)
+    chunk_repo = ChunkRepository(store)
+
+    # Get the first document from the corpus
+    first_doc = qa_corpus[0]
+    document_text = first_doc["document_extracted"]
+
+    # Create a document first
+    document = Document(content=document_text, metadata={"source": "test"})
+    created_document = await doc_repo.create(document)
+    assert created_document.id is not None
+
+    # Test getting chunks by document ID
+    chunks = await chunk_repo.get_by_document_id(created_document.id)
+    assert len(chunks) > 0
+    assert all(chunk.document_id == created_document.id for chunk in chunks)
+
+    # Test chunk search
+    results = await chunk_repo.search_chunks("election", limit=2)
+    assert len(results) <= 2
+    assert all(hasattr(chunk, "content") for chunk in results)
+
+    # Test deleting chunks by document ID
+    deleted = await chunk_repo.delete_by_document_id(created_document.id)
+    assert deleted is True
+
+    # Verify chunks are gone
+    chunks_after_delete = await chunk_repo.get_by_document_id(created_document.id)
+    assert len(chunks_after_delete) == 0
+
+    store.close()
+
+
+@pytest.mark.asyncio
+async def test_create_chunks_for_document(qa_corpus: Dataset):
+    """Test creating chunks for a document."""
+    # Create an in-memory store and repositories
+    store = Store(":memory:")
+    chunk_repo = ChunkRepository(store)
+
+    # Get the first document from the corpus
+    first_doc = qa_corpus[0]
+    document_text = first_doc["document_extracted"]
+
+    # Create a document first (without chunks)
+    document = Document(content=document_text, metadata={"source": "test"})
+
+    # Insert document manually to test chunk creation independently
+    document_id = None
+    if store._connection is not None:
+        cursor = store._connection.cursor()
+        cursor.execute(
+            """
+            INSERT INTO documents (content, metadata, created_at, updated_at)
+            VALUES (?, ?, ?, ?)
+            """,
+            (document.content, "{}", document.created_at, document.updated_at),
+        )
+        document_id = cursor.lastrowid
+        document.id = document_id
+        store._connection.commit()
+
+    assert document_id is not None, "Document ID should not be None"
+
+    # Test creating chunks for the document
+    chunks = await chunk_repo.create_chunks_for_document(document_id, document_text)
+
+    # Verify chunks were created
+    assert len(chunks) > 0
+    assert all(chunk.document_id == document_id for chunk in chunks)
+    assert all(chunk.id is not None for chunk in chunks)
+
+    # Verify chunk order metadata
+    for i, chunk in enumerate(chunks):
+        assert chunk.metadata.get("order") == i
+
+    # Verify chunks exist in database
+    db_chunks = await chunk_repo.get_by_document_id(document_id)
+    assert len(db_chunks) == len(chunks)
+
+    store.close()
+
+
+@pytest.mark.asyncio
+async def test_chunk_repository_crud():
+    """Test basic CRUD operations in ChunkRepository."""
+    # Create an in-memory store
+    store = Store(":memory:")
+    chunk_repo = ChunkRepository(store)
+
+    # First create a document to reference
+    document_id = None
+    if store._connection is not None:
+        cursor = store._connection.cursor()
+        cursor.execute(
+            """
+            INSERT INTO documents (content, metadata, created_at, updated_at)
+            VALUES (?, ?, datetime('now'), datetime('now'))
+            """,
+            ("Test document content", "{}"),
+        )
+        document_id = cursor.lastrowid
+        store._connection.commit()
+
+    assert document_id is not None, "Document ID should not be None"
+
+    # Test create chunk manually
+    chunk = Chunk(
+        document_id=document_id,
+        content="Test chunk content",
+        metadata={"test": "value"},
+    )
+
+    created_chunk = await chunk_repo.create(chunk)
+    assert created_chunk.id is not None
+    assert created_chunk.content == "Test chunk content"
+
+    # Test get by ID
+    retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
+    assert retrieved_chunk is not None
+    assert retrieved_chunk.content == "Test chunk content"
+    assert retrieved_chunk.metadata["test"] == "value"
+
+    # Test update
+    retrieved_chunk.content = "Updated chunk content"
+    updated_chunk = await chunk_repo.update(retrieved_chunk)
+    assert updated_chunk.content == "Updated chunk content"
+
+    # Test list all
+    all_chunks = await chunk_repo.list_all()
+    assert len(all_chunks) >= 1
+    assert any(chunk.id == created_chunk.id for chunk in all_chunks)
+
+    # Test delete
+    deleted = await chunk_repo.delete(created_chunk.id)
+    assert deleted is True
+
+    # Verify chunk is gone
+    retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
+    assert retrieved_chunk is None
+
+    store.close()

tests/test_document.py

@@ -3,13 +3,15 @@ from datasets import Dataset
 
 from haiku.rag.store.engine import Store
 from haiku.rag.store.models.document import Document
+from haiku.rag.store.repositories.document import DocumentRepository
 
 
 @pytest.mark.asyncio
 async def test_create_document_with_chunks(qa_corpus: Dataset):
-    """Test creating a document with chunks from the qa_corpus."""
+    """Test creating a document with chunks from the qa_corpus using repository."""
-    # Create an in-memory store
+    # Create an in-memory store and repository
     store = Store(":memory:")
+    doc_repo = DocumentRepository(store)
     
     # Get the first document from the corpus
     first_doc = qa_corpus[0]
@@ -22,7 +24,7 @@ async def test_create_document_with_chunks(qa_corpus: Dataset):
     )
     
     # Create the document with chunks in the database
-    created_document = await document.create_with_chunks(store)
+    created_document = await doc_repo.create(document)
     
     # Verify the document was created
     assert created_document.id is not None
@@ -60,30 +62,45 @@ async def test_create_document_with_chunks(qa_corpus: Dataset):
 
 
 @pytest.mark.asyncio
-async def test_search_chunks(qa_corpus: Dataset):
+async def test_document_repository_crud(qa_corpus: Dataset):
-    """Test vector search functionality."""
+    """Test CRUD operations in DocumentRepository."""
-    # Create an in-memory store
+    # Create an in-memory store and repository
     store = Store(":memory:")
+    doc_repo = DocumentRepository(store)
     
     # Get the first document from the corpus
     first_doc = qa_corpus[0]
     document_text = first_doc["document_extracted"]
     
-    # Create and store a document
+    # Create a document
     document = Document(
         content=document_text,
-        metadata={"source": "qa_corpus"}
+        metadata={"source": "test"}
     )
-    await document.create_with_chunks(store)
+    created_document = await doc_repo.create(document)
     
-    # Perform a search
+    # Test get_by_id
-    search_query = "news"  # Simple query
+    assert created_document.id is not None
-    results = await Document.search_chunks(store, search_query, limit=3)
+    retrieved_document = await doc_repo.get_by_id(created_document.id)
+    assert retrieved_document is not None
+    assert retrieved_document.content == document_text
     
-    # Verify search results
+    # Test update (should regenerate chunks)
-    assert len(results) <= 3
+    retrieved_document.content = "Updated content for testing"
-    assert all(hasattr(chunk, "content") for chunk in results)
+    updated_document = await doc_repo.update(retrieved_document)
-    assert all(hasattr(chunk, "document_id") for chunk in results)
+    assert updated_document.content == "Updated content for testing"
-    assert all(chunk.document_id == document.id for chunk in results)
+    
+    # Test list_all
+    all_documents = await doc_repo.list_all()
+    assert len(all_documents) == 1
+    assert all_documents[0].id == created_document.id
+    
+    # Test delete
+    deleted = await doc_repo.delete(created_document.id)
+    assert deleted is True
+    
+    # Verify document is gone
+    retrieved_document = await doc_repo.get_by_id(created_document.id)
+    assert retrieved_document is None
     
     store.close()