Simplify _process_search_results

src/haiku/rag/store/repositories/chunk.py

@@ -4,6 +4,10 @@ import logging
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
+if TYPE_CHECKING:
+    import pandas as pd
+    from lancedb.query import LanceQueryBuilder
+
 from lancedb.rerankers import RRFReranker
 
 from haiku.rag.store.engine import DocumentRecord, Store
@@ -351,62 +355,51 @@ class ChunkRepository:
 
         return adjacent_chunks
 
-    async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]:
+    async def _process_search_results(
-        """Process search results into chunks with document info and scores."""
+        self, query_result: "pd.DataFrame | LanceQueryBuilder"
-        chunks_with_scores = []
+    ) -> list[tuple[Chunk, float]]:
+        """Process search results into chunks with document info and scores.
 
-        # Handle pandas DataFrame (from filtered results)
+        Args:
+            query_result: Either a pandas DataFrame or a LanceDB query result
+        """
         import pandas as pd
 
+        def extract_scores(df: pd.DataFrame) -> list[float]:
+            """Extract scores from DataFrame columns based on search type."""
+            if "_distance" in df.columns:
+                # Vector search - convert distance to similarity
+                return [max(0.0, 1.0 / (1.0 + d)) for d in df["_distance"].tolist()]
+            elif "_relevance_score" in df.columns:
+                # Hybrid search - relevance score (higher is better)
+                return df["_relevance_score"].tolist()
+            elif "_score" in df.columns:
+                # FTS search - score (higher is better)
+                return df["_score"].tolist()
+            else:
+                raise ValueError("Unknown search result format, cannot extract scores")
+
+        # Convert everything to DataFrame for uniform processing
         if isinstance(query_result, pd.DataFrame):
-            # DataFrame already contains the data we need
+            df = query_result
-            pydantic_results = []
-            for _, row in query_result.iterrows():
-                chunk_record = self.store.ChunkRecord(
-                    id=str(row["id"]),
-                    document_id=str(row["document_id"]),
-                    content=str(row["content"]),
-                    metadata=str(row["metadata"]),
-                    order=int(row["order"]) if "order" in row else 0,
-                )
-                pydantic_results.append(chunk_record)
-
-            # Extract scores from DataFrame columns
-            scores = []
-            if "_distance" in query_result.columns:
-                # Vector search - distance (lower is better, convert to similarity)
-                distances = query_result["_distance"].tolist()
-                scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances]
-            elif "_relevance_score" in query_result.columns:
-                # Hybrid search - relevance score (higher is better)
-                scores = query_result["_relevance_score"].tolist()
-            elif "_score" in query_result.columns:
-                # FTS search - score (higher is better)
-                scores = query_result["_score"].tolist()
-            else:
-                raise ValueError("Unknown search result format, cannot extract scores")
         else:
-            # Handle LanceDB query result (original behavior)
+            # Convert LanceDB query result to DataFrame
-            # Get both arrow and pydantic results to access scores
+            df = query_result.to_pandas()
-            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
+        # Extract scores
-            scores = []
+        scores = extract_scores(df)
-            column_names = arrow_result.column_names
 
-            if "_distance" in column_names:
+        # Convert DataFrame rows to ChunkRecords
-                # Vector search - distance (lower is better, convert to similarity)
+        pydantic_results = [
-                distances = arrow_result.column("_distance").to_pylist()
+            self.store.ChunkRecord(
-                scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances]
+                id=str(row["id"]),
-            elif "_relevance_score" in column_names:
+                document_id=str(row["document_id"]),
-                # Hybrid search - relevance score (higher is better)
+                content=str(row["content"]),
-                scores = arrow_result.column("_relevance_score").to_pylist()
+                metadata=str(row["metadata"]),
-            elif "_score" in column_names:
+                order=int(row["order"]) if "order" in row else 0,
-                # FTS search - score (higher is better)
+            )
-                scores = arrow_result.column("_score").to_pylist()
+            for _, row in df.iterrows()
-            else:
+        ]
-                raise ValueError("Unknown search result format, cannot extract scores")
 
         # Collect all unique document IDs for batch lookup
         document_ids = list(set(chunk.document_id for chunk in pydantic_results))
@@ -423,29 +416,21 @@ class ChunkRepository:
             )
             documents_map = {doc.id: doc for doc in doc_results}
 
+        # Build final results with document info
+        chunks_with_scores = []
         for i, chunk_record in enumerate(pydantic_results):
-            # Get document info from pre-fetched map
             doc = documents_map.get(chunk_record.document_id)
-            doc_uri = doc.uri if doc else None
-            doc_title = doc.title if doc else None
-            doc_meta = doc.metadata if doc else "{}"
-
-            md = json.loads(chunk_record.metadata)
-
             chunk = Chunk(
                 id=chunk_record.id,
                 document_id=chunk_record.document_id,
                 content=chunk_record.content,
-                metadata=md,
+                metadata=json.loads(chunk_record.metadata),
                 order=chunk_record.order,
-                document_uri=doc_uri,
+                document_uri=doc.uri if doc else None,
-                document_title=doc_title,
+                document_title=doc.title if doc else None,
-                document_meta=json.loads(doc_meta),
+                document_meta=json.loads(doc.metadata if doc else "{}"),
             )
-
-            # Get score from arrow result
             score = scores[i] if i < len(scores) else 1.0
-
             chunks_with_scores.append((chunk, score))
 
         return chunks_with_scores