diff --git a/src/haiku/rag/embeddings/base.py b/src/haiku/rag/embeddings/base.py index 781d090f..a62deffd 100644 --- a/src/haiku/rag/embeddings/base.py +++ b/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." diff --git a/src/haiku/rag/embeddings/ollama.py b/src/haiku/rag/embeddings/ollama.py index a9d1779a..6de861a9 100644 --- a/src/haiku/rag/embeddings/ollama.py +++ b/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) diff --git a/src/haiku/rag/store/models/document.py b/src/haiku/rag/store/models/document.py index 84ce81ea..00878212 100644 --- a/src/haiku/rag/store/models/document.py +++ b/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 diff --git a/src/haiku/rag/store/repositories/__init__.py b/src/haiku/rag/store/repositories/__init__.py new file mode 100644 index 00000000..109823b2 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/src/haiku/rag/store/repositories/base.py b/src/haiku/rag/store/repositories/base.py new file mode 100644 index 00000000..6facfb0e --- /dev/null +++ b/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 diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py new file mode 100644 index 00000000..5cb0d3e5 --- /dev/null +++ b/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 diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py new file mode 100644 index 00000000..027b76b7 --- /dev/null +++ b/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 diff --git a/tests/test_chunk.py b/tests/test_chunk.py new file mode 100644 index 00000000..a9263c7c --- /dev/null +++ b/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() diff --git a/tests/test_document.py b/tests/test_document.py index 529188b4..c6bc5485 100644 --- a/tests/test_document.py +++ b/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.""" - # Create an in-memory store + """Test creating a document with chunks from the qa_corpus using repository.""" + # 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): - """Test vector search functionality.""" - # Create an in-memory store +async def test_document_repository_crud(qa_corpus: Dataset): + """Test CRUD operations in DocumentRepository.""" + # 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 - search_query = "news" # Simple query - results = await Document.search_chunks(store, search_query, limit=3) + # Test get_by_id + assert created_document.id is not None + 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 - 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 == document.id for chunk in results) + # Test update (should regenerate chunks) + retrieved_document.content = "Updated content for testing" + updated_document = await doc_repo.update(retrieved_document) + assert updated_document.content == "Updated content for testing" + + # 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() \ No newline at end of file