From 628bdc8151f6a2cec6fbd9f32b427578d2a351c4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 1 Sep 2025 11:21:42 +0300 Subject: [PATCH] Cleanup --- docs/installation.md | 1 - mkdocs.yml | 2 +- src/haiku/rag/client.py | 24 +---- src/haiku/rag/store/engine.py | 4 +- src/haiku/rag/store/repositories/chunk.py | 101 +++++++++------------- tests/generate_benchmark_db.py | 2 +- tests/test_chunk.py | 2 +- tests/test_search.py | 12 ++- 8 files changed, 58 insertions(+), 90 deletions(-) diff --git a/docs/installation.md b/docs/installation.md index fd27e3a2..727fadc4 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -31,5 +31,4 @@ uv pip install haiku.rag[mxbai] ## Requirements - Python 3.10+ -- SQLite 3.38+ - Ollama (for default embeddings) diff --git a/mkdocs.yml b/mkdocs.yml index 07bcb523..a1ed910e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,5 +1,5 @@ site_name: haiku.rag -site_description: Retrieval-Augmented Generation (RAG) library on SQLite. +site_description: Retrieval-Augmented Generation (RAG) library on LanceDB. site_url: https://ggozad.github.io/haiku.rag/ theme: name: material diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index e721af09..e46ce00e 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -334,27 +334,13 @@ class HaikuRAG: if reranker is None: # No reranking - return direct search results - if search_type == "vector": - return await self.chunk_repository.search_chunks(query, limit) - elif search_type == "fts": - return await self.chunk_repository.search_chunks_fts(query, limit) - else: # hybrid (default) - return await self.chunk_repository.search_chunks_hybrid(query, limit) + return await self.chunk_repository.search(query, limit, search_type) # Get more initial results (3X) for reranking search_limit = limit * 3 - if search_type == "vector": - search_results = await self.chunk_repository.search_chunks( - query, search_limit - ) - elif search_type == "fts": - search_results = await self.chunk_repository.search_chunks_fts( - query, search_limit - ) - else: # hybrid (default) - search_results = await self.chunk_repository.search_chunks_hybrid( - query, search_limit - ) + search_results = await self.chunk_repository.search( + query, search_limit, search_type + ) # Apply reranking chunks = [chunk for chunk, _ in search_results] @@ -564,8 +550,6 @@ class HaikuRAG: ) yield doc.id - # LanceDB doesn't need explicit commits - def close(self): """Close the underlying store connection.""" self.store.close() diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 541b2f79..78f3d7d0 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -47,9 +47,7 @@ class Store: # Create the ChunkRecord model with the correct vector dimension self.ChunkRecord = create_chunk_model(self.embedder._vector_dim) - # For file paths, create a LanceDB directory structure - lance_path = str(db_path).replace(".sqlite", ".lancedb") - self.db = lancedb.connect(lance_path) + self.db = lancedb.connect(db_path) self.create_or_update_db() diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 29216bdc..8de40b6a 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -28,16 +28,14 @@ class ChunkRepository: """Create a chunk in the database.""" assert entity.document_id, "Chunk must have a document_id to be created" + chunk_id = str(uuid4()) + # Generate embedding if not provided if entity.embedding is not None: embedding = entity.embedding else: embedding = await self.embedder.embed(entity.content) - # Generate new UUID - chunk_id = str(uuid4()) - - # Create chunk record chunk_record = self.store.ChunkRecord( id=chunk_id, document_id=entity.document_id, @@ -46,7 +44,6 @@ class ChunkRepository: vector=embedding, ) - # Add to table self.store.chunks_table.add([chunk_record]) entity.id = chunk_id @@ -76,10 +73,8 @@ class ChunkRepository: """Update an existing chunk.""" assert entity.id, "Chunk ID is required for update" - # Generate new embedding embedding = await self.embedder.embed(entity.content) - # Update the record self.store.chunks_table.update( where=f"id = '{entity.id}'", values={ @@ -94,12 +89,10 @@ class ChunkRepository: async def delete(self, entity_id: str) -> bool: """Delete a chunk by its ID.""" - # Check if chunk exists chunk = await self.get_by_id(entity_id) if chunk is None: return False - # Delete the chunk self.store.chunks_table.delete(f"id = '{entity_id}'") return True @@ -145,7 +138,6 @@ class ChunkRepository: async def delete_all(self) -> bool: """Delete all chunks from the database.""" try: - # Get count before deletion count = len( list( self.store.chunks_table.search() @@ -171,74 +163,66 @@ class ChunkRepository: if not chunks: return False - # Delete chunks by document_id self.store.chunks_table.delete(f"document_id = '{document_id}'") return True - async def search_chunks( - self, query: str, limit: int = 5 + async def search( + self, query: str, limit: int = 5, search_type: str = "hybrid" ) -> list[tuple[Chunk, float]]: - """Search for relevant chunks using vector similarity.""" - # Generate embedding for the query - query_embedding = await self.embedder.embed(query) + """Search for relevant chunks using the specified search method. - # Perform vector search with proper query type - results = ( - self.store.chunks_table.search(query_embedding, query_type="vector") - .limit(limit) - .to_pydantic(self.store.ChunkRecord) - ) + Args: + query: The search query string. + limit: Maximum number of results to return. + search_type: Type of search - "vector", "fts", or "hybrid" (default). - return await self._process_search_results(results) - - async def search_chunks_fts( - self, query: str, limit: int = 5 - ) -> list[tuple[Chunk, float]]: - """Search for chunks using full-text search.""" + Returns: + List of (chunk, score) tuples ordered by relevance. + """ if not query.strip(): return [] - # Ensure FTS index exists - self._ensure_fts_index() + if search_type == "vector": + query_embedding = await self.embedder.embed(query) + + results = ( + self.store.chunks_table.search(query_embedding, query_type="vector") + .limit(limit) + .to_pydantic(self.store.ChunkRecord) + ) + + return await self._process_search_results(results) + + elif search_type == "fts": + # Ensure FTS index exists + self._ensure_fts_index() - # Use LanceDB's native full-text search - try: results = ( self.store.chunks_table.search(query, query_type="fts") .limit(limit) .to_pydantic(self.store.ChunkRecord) ) return await self._process_search_results(results) - except Exception: - # Fallback to vector search if FTS is not available or fails - return await self.search_chunks(query, limit) - async def search_chunks_hybrid( - self, query: str, limit: int = 5 - ) -> list[tuple[Chunk, float]]: - """Hybrid search combining vector and full-text search with native LanceDB RRFReranker.""" - if not query.strip(): - return [] + else: # hybrid (default) + # Ensure FTS index exists for hybrid search + self._ensure_fts_index() - # Ensure FTS index exists for hybrid search - self._ensure_fts_index() + query_embedding = await self.embedder.embed(query) - # Generate embedding for the query since LanceDB doesn't have embedding function configured - query_embedding = await self.embedder.embed(query) + # Create RRF reranker + reranker = RRFReranker() - # Create RRF reranker (k parameter is handled internally) - reranker = RRFReranker() - - # Perform native hybrid search with RRF reranking - results = ( - self.store.chunks_table.search(query_type="hybrid") - .vector(query_embedding) - .text(query) - .rerank(reranker) - .limit(limit) - .to_pydantic(self.store.ChunkRecord) - ) - return await self._process_search_results(results) + # Perform native hybrid search with RRF reranking + results = ( + self.store.chunks_table.search(query_type="hybrid") + .vector(query_embedding) + .text(query) + .rerank(reranker) + .limit(limit) + .to_pydantic(self.store.ChunkRecord) + ) + return await self._process_search_results(results) async def get_by_document_id(self, document_id: str) -> list[Chunk]: """Get all chunks for a specific document.""" @@ -272,7 +256,6 @@ class ChunkRepository: for chunk in results ] - # Sort by order if available chunks.sort(key=lambda c: c.metadata.get("order", 0)) return chunks diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index bb43a2fc..5404f6b9 100644 --- a/tests/generate_benchmark_db.py +++ b/tests/generate_benchmark_db.py @@ -11,7 +11,7 @@ from haiku.rag.qa import get_qa_agent console = Console() -db_path = Path(__file__).parent / "data" / "benchmark.sqlite" +db_path = Path(__file__).parent / "data" / "benchmark.lancedb" async def populate_db(): diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 68880b2f..6f4e7bbc 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -35,7 +35,7 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path): assert all(chunk.document_id == created_document.id for chunk in chunks) # Test chunk search - results = await chunk_repo.search_chunks("election", limit=2) + results = await chunk_repo.search("election", limit=2, search_type="vector") assert len(results) <= 2 assert all(hasattr(chunk, "content") for chunk, _ in results) diff --git a/tests/test_search.py b/tests/test_search.py index f3e4c41f..94b93a99 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -47,17 +47,21 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path): question = doc_data["question"] # Test vector search - vector_results = await chunk_repo.search_chunks(question, limit=5) + vector_results = await chunk_repo.search( + question, limit=5, search_type="vector" + ) target_document_ids = {chunk.document_id for chunk, _ in vector_results} assert target_document.id in target_document_ids # Test FTS search - fts_results = await chunk_repo.search_chunks_fts(question, limit=5) + fts_results = await chunk_repo.search(question, limit=5, search_type="fts") target_document_ids = {chunk.document_id for chunk, _ in fts_results} assert target_document.id in target_document_ids # Test hybrid search - hybrid_results = await chunk_repo.search_chunks_hybrid(question, limit=5) + hybrid_results = await chunk_repo.search( + question, limit=5, search_type="hybrid" + ) target_document_ids = {chunk.document_id for chunk, _ in hybrid_results} assert target_document.id in target_document_ids @@ -85,7 +89,7 @@ async def test_chunks_include_document_info(temp_db_path): created_document = await doc_repo._create_with_docling(document, docling_document) # Search for chunks - results = await chunk_repo.search_chunks_hybrid("test document", limit=1) + results = await chunk_repo.search("test document", limit=1, search_type="hybrid") assert len(results) > 0 chunk, _ = results[0]