diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 397f4117..24f2bb68 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/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 @@ -249,64 +253,56 @@ class ChunkRepository: """ if not query.strip(): return [] - - chunk_where_clause = None + filtered_doc_ids = None if filter: # We perform filtering as a two-step process, first filtering documents, then # filtering chunks based on those document IDs. # This is because LanceDB does not support joins directly in search queries. - matching_doc_ids = self._get_filtered_document_ids(filter) - - if not matching_doc_ids: + docs_df = ( + self.store.documents_table.search() + .select(["id"]) + .where(filter) + .to_pandas() + ) + # Early exit if no documents match the filter + if docs_df.empty: return [] + # Keep as pandas Series for efficient vectorized operations + filtered_doc_ids = docs_df["id"] - # Build WHERE clause for chunks table - # Use IN clause with document IDs - id_list = "', '".join(matching_doc_ids) - chunk_where_clause = f"document_id IN ('{id_list}')" - + # Prepare search query based on search type if search_type == "vector": query_embedding = await self.embedder.embed(query) - results = self.store.chunks_table.search( query_embedding, query_type="vector", vector_column_name="vector" ) - if chunk_where_clause: - results = results.where(chunk_where_clause) - - results = results.limit(limit) - - return await self._process_search_results(results) - elif search_type == "fts": results = self.store.chunks_table.search(query, query_type="fts") - if chunk_where_clause: - results = results.where(chunk_where_clause) - - results = results.limit(limit) - return await self._process_search_results(results) - else: # hybrid (default) query_embedding = await self.embedder.embed(query) - # Create RRF reranker reranker = RRFReranker() - # Perform native hybrid search with RRF reranking - results = self.store.chunks_table.search(query_type="hybrid") - - if chunk_where_clause: - results = results.where(chunk_where_clause) - results = ( - results.vector(query_embedding) + self.store.chunks_table.search(query_type="hybrid") + .vector(query_embedding) .text(query) .rerank(reranker) - .limit(limit) ) - return await self._process_search_results(results) + + # Apply filtering if needed (common for all search types) + if filtered_doc_ids is not None: + chunks_df = results.to_pandas() + filtered_chunks_df = chunks_df.loc[ + chunks_df["document_id"].isin(filtered_doc_ids) + ].head(limit) + return await self._process_search_results(filtered_chunks_df) + + # No filtering needed, apply limit and return + results = results.limit(limit) + 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.""" @@ -364,39 +360,51 @@ class ChunkRepository: return adjacent_chunks - def _get_filtered_document_ids(self, filter: str) -> list[str]: - """Query documents table with filter and return matching document IDs.""" - filtered_docs = ( - self.store.documents_table.search() - .where(filter) - .to_pydantic(DocumentRecord) - ) - return [doc.id for doc in filtered_docs] + async def _process_search_results( + self, query_result: "pd.DataFrame | LanceQueryBuilder" + ) -> list[tuple[Chunk, float]]: + """Process search results into chunks with document info and scores. - 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 = [] + Args: + query_result: Either a pandas DataFrame or a LanceDB query result + """ + import pandas as pd - # 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)) + 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 ((df["_distance"] + 1).rdiv(1)).clip(lower=0.0).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") - # 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() + # Convert everything to DataFrame for uniform processing + if isinstance(query_result, pd.DataFrame): + df = query_result else: - raise ValueError("Unknown search result format, cannot extract scores") + # Convert LanceDB query result to DataFrame + df = query_result.to_pandas() + + # Extract scores + scores = extract_scores(df) + + # Convert DataFrame rows to ChunkRecords + pydantic_results = [ + 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, + ) + for _, row in df.iterrows() + ] # Collect all unique document IDs for batch lookup document_ids = list(set(chunk.document_id for chunk in pydantic_results)) @@ -404,8 +412,9 @@ class ChunkRepository: # Batch fetch all documents at once documents_map = {} if document_ids: - # Create a WHERE clause for all document IDs - where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids) + # Use IN clause for efficient batch lookup + id_list = "', '".join(document_ids) + where_clause = f"id IN ('{id_list}')" doc_results = list( self.store.documents_table.search() .where(where_clause) @@ -413,29 +422,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_title=doc_title, - document_meta=json.loads(doc_meta), + document_uri=doc.uri if doc else None, + document_title=doc.title if doc else None, + 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