test search

tests/test_chunk.py

@@ -8,35 +8,6 @@ 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."""

tests/test_search.py

@@ -0,0 +1,50 @@
+import pytest
+from datasets import Dataset
+
+from haiku.rag.store.engine import Store
+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_qa_corpus(qa_corpus: Dataset):
+    """Test that documents can be found by searching with their associated questions."""
+    # Create an in-memory store and repositories
+    store = Store(":memory:")
+    doc_repo = DocumentRepository(store)
+    chunk_repo = ChunkRepository(store)
+    num_documents = 10
+    # Load first 10 documents with embeddings (reduced for faster testing)
+    documents = []
+    for i in range(num_documents):
+        doc_data = qa_corpus[i]
+        document_text = doc_data["document_extracted"]
+
+        # Create a Document instance
+        document = Document(
+            content=document_text,
+            metadata={
+                "source": "qa_corpus",
+                "topic": doc_data.get("document_topic", ""),
+                "document_id": doc_data.get("document_id", ""),
+                "question": doc_data["question"],
+            },
+        )
+
+        # Create the document with chunks and embeddings
+        created_document = await doc_repo.create(document)
+        documents.append((created_document, doc_data))
+
+    for i in range(num_documents):  # Test with first few documents
+        target_document, doc_data = documents[i]
+        question = doc_data["question"]
+
+        # Search for chunks using the question
+        search_results = await chunk_repo.search_chunks(question, limit=5)
+
+        # Check if target document is in results
+        target_document_ids = {chunk.document_id for chunk in search_results}
+        assert target_document.id in target_document_ids
+
+    store.close()