From a5e087aa536e1c3ba1885e50c02ca30fce5e0580 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 1 Sep 2025 12:40:46 +0300 Subject: [PATCH] Fix scores depending on search type --- src/haiku/rag/store/repositories/chunk.py | 47 ++++++---- tests/test_search.py | 108 ++++++++++++++++++---- 2 files changed, 120 insertions(+), 35 deletions(-) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 8de40b6a..acc7384e 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -185,11 +185,9 @@ class ChunkRepository: 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) - ) + results = self.store.chunks_table.search( + query_embedding, query_type="vector" + ).limit(limit) return await self._process_search_results(results) @@ -197,10 +195,8 @@ class ChunkRepository: # Ensure FTS index exists self._ensure_fts_index() - results = ( - self.store.chunks_table.search(query, query_type="fts") - .limit(limit) - .to_pydantic(self.store.ChunkRecord) + results = self.store.chunks_table.search(query, query_type="fts").limit( + limit ) return await self._process_search_results(results) @@ -220,7 +216,6 @@ class ChunkRepository: .text(query) .rerank(reranker) .limit(limit) - .to_pydantic(self.store.ChunkRecord) ) return await self._process_search_results(results) @@ -279,11 +274,31 @@ class ChunkRepository: return adjacent_chunks - async def _process_search_results(self, results) -> list[tuple[Chunk, float]]: + async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]: """Process search results into chunks with document info and scores.""" chunks_with_scores = [] - for chunk_record in results: + # Get both arrow and pydantic results to access scores + arrow_result = query_result.to_arrow() + pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord)) + # Extract scores from arrow result based on search type + scores = [] + column_names = arrow_result.column_names + + if "_distance" in column_names: + # Vector search - distance (lower is better, convert to similarity) + distances = arrow_result.column("_distance").to_pylist() + scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances] + elif "_relevance_score" in column_names: + # Hybrid search - relevance score (higher is better) + scores = arrow_result.column("_relevance_score").to_pylist() + elif "_score" in column_names: + # FTS search - score (higher is better) + scores = arrow_result.column("_score").to_pylist() + else: + raise ValueError("Unknown search result format, cannot extract scores") + + for i, chunk_record in enumerate(pydantic_results): # Get document info doc_results = list( self.store.documents_table.search() @@ -306,12 +321,8 @@ class ChunkRepository: document_meta=json.loads(doc_meta) if doc_meta else {}, ) - # Get distance score - LanceDB returns _distance (lower is better) - distance = getattr(chunk_record, "_distance", 1.0) - - # Convert distance to similarity score (higher is better) - # Using exponential decay to convert distance to similarity - score = max(0.0, 1.0 / (1.0 + distance)) + # Get score from arrow result + score = scores[i] if i < len(scores) else 1.0 chunks_with_scores.append((chunk, score)) diff --git a/tests/test_search.py b/tests/test_search.py index 94b93a99..f97086e9 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -15,23 +15,22 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path): doc_repo = DocumentRepository(store) chunk_repo = ChunkRepository(store) - # Load first 20 documents with embeddings (reduced for faster testing) - num_documents = 20 + # Load unique documents (limited to 10) + seen_documents = set() documents = [] - for i in range(num_documents): - doc_data = qa_corpus[i] - document_text = doc_data["document_extracted"] + + for doc_data in qa_corpus: + if len(seen_documents) >= 10: + break + document_text = doc_data["document_extracted"] # type: ignore + document_id = doc_data.get("document_id", "") # type: ignore + + if document_id in seen_documents: + continue + seen_documents.add(document_id) # 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"], - }, - ) + document = Document(content=document_text) # Create the document with chunks and embeddings from haiku.rag.utils import text_to_docling_document @@ -42,8 +41,9 @@ async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path): ) documents.append((created_document, doc_data)) - for i in range(3): # Test with first few documents - target_document, doc_data = documents[i] + # Test with first few unique documents + + for target_document, doc_data in documents: question = doc_data["question"] # Test vector search @@ -92,7 +92,11 @@ async def test_chunks_include_document_info(temp_db_path): results = await chunk_repo.search("test document", limit=1, search_type="hybrid") assert len(results) > 0 - chunk, _ = results[0] + chunk, score = results[0] + + # Test that score is valid + assert isinstance(score, int | float), f"Score should be numeric, got {type(score)}" + assert score >= 0, f"Score should be non-negative, got {score}" # Verify the chunk includes document information assert chunk.document_uri == "https://example.com/test.html" @@ -100,3 +104,73 @@ async def test_chunks_include_document_info(temp_db_path): assert chunk.document_id == created_document.id store.close() + + +@pytest.mark.asyncio +async def test_search_score_types(temp_db_path): + """Test that different search types return appropriate score ranges.""" + store = Store(temp_db_path) + doc_repo = DocumentRepository(store) + chunk_repo = ChunkRepository(store) + + # Create multiple documents with different content + documents_content = [ + "Machine learning algorithms are powerful tools for data analysis and pattern recognition.", + "Deep learning neural networks can process complex datasets and identify hidden patterns.", + "Natural language processing enables computers to understand and generate human text.", + "Computer vision systems can interpret and analyze visual information from images.", + ] + + for content in documents_content: + document = Document(content=content) + from haiku.rag.utils import text_to_docling_document + + docling_document = text_to_docling_document(content, name="test.md") + await doc_repo._create_with_docling(document, docling_document) + + query = "machine learning" + + # Test vector search scores (should be converted from distances) + vector_results = await chunk_repo.search(query, limit=3, search_type="vector") + assert len(vector_results) > 0 + vector_scores = [score for _, score in vector_results] + + # Test FTS search scores (should be native LanceDB FTS scores) + fts_results = await chunk_repo.search(query, limit=3, search_type="fts") + assert len(fts_results) > 0 + fts_scores = [score for _, score in fts_results] + + # Test hybrid search scores (should be native LanceDB relevance scores) + hybrid_results = await chunk_repo.search(query, limit=3, search_type="hybrid") + assert len(hybrid_results) > 0 + hybrid_scores = [score for _, score in hybrid_results] + + # All scores should be numeric and non-negative + for scores, search_type in [ + (vector_scores, "vector"), + (fts_scores, "fts"), + (hybrid_scores, "hybrid"), + ]: + for score in scores: + assert isinstance(score, int | float), ( + f"{search_type} score should be numeric" + ) + assert score >= 0, f"{search_type} score should be non-negative" + + # Vector scores should typically be small (0-1 range due to distance conversion) + assert all(0 <= score <= 1 for score in vector_scores), ( + "Vector scores should be in 0-1 range" + ) + + # Scores should be sorted in descending order (most relevant first) + for scores, search_type in [ + (vector_scores, "vector"), + (fts_scores, "fts"), + (hybrid_scores, "hybrid"), + ]: + for i in range(len(scores) - 1): + assert scores[i] >= scores[i + 1], ( + f"{search_type} results should be sorted by score descending" + ) + + store.close()