Fix scores depending on search type

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 = (
+            results = self.store.chunks_table.search(
-                self.store.chunks_table.search(query_embedding, query_type="vector")
+                query_embedding, query_type="vector"
-                .limit(limit)
+            ).limit(limit)
-                .to_pydantic(self.store.ChunkRecord)
-            )
 
             return await self._process_search_results(results)
 
@@ -197,10 +195,8 @@ class ChunkRepository:
             # Ensure FTS index exists
             self._ensure_fts_index()
 
-            results = (
+            results = self.store.chunks_table.search(query, query_type="fts").limit(
-                self.store.chunks_table.search(query, query_type="fts")
+                limit
-                .limit(limit)
-                .to_pydantic(self.store.ChunkRecord)
             )
             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)
+            # Get score from arrow result
-            distance = getattr(chunk_record, "_distance", 1.0)
+            score = scores[i] if i < len(scores) else 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))
 
             chunks_with_scores.append((chunk, score))
 

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)
+    # Load unique documents (limited to 10)
-    num_documents = 20
+    seen_documents = set()
     documents = []
-    for i in range(num_documents):
+
-        doc_data = qa_corpus[i]
+    for doc_data in qa_corpus:
-        document_text = doc_data["document_extracted"]
+        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(
+        document = Document(content=document_text)
-            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
         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
+    # Test with first few unique documents
-        target_document, doc_data = documents[i]
+
+    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()