"""Client for communicating with dense and sparse embedding services.""" import httpx import asyncio from typing import Dict, Any, List, Optional, Tuple import logging from dataclasses import dataclass logger = logging.getLogger(__name__) @dataclass class SearchResult: """Unified search result from embedding services.""" doc_id: str score: float text: Optional[str] = None metadata: Optional[Dict[str, Any]] = None source: str = "" # "dense" or "sparse" rank: Optional[int] = None debug_info: Optional[Dict[str, Any]] = None class RetrievalClient: """Client for parallel retrieval from dense and sparse services.""" def __init__(self, dense_url: str, sparse_url: str, timeout: float = 30.0): self.dense_url = dense_url.rstrip('/') self.sparse_url = sparse_url.rstrip('/') self.timeout = timeout async def index_document_dense(self, text: str, doc_id: str, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Index a document in the dense embedding service.""" async with httpx.AsyncClient(timeout=self.timeout) as client: try: payload = { "text": text, "doc_id": doc_id, "metadata": metadata or {} } response = await client.post(f"{self.dense_url}/index", json=payload) response.raise_for_status() result = response.json() logger.debug(f"Dense indexing successful for doc {doc_id}") return result except Exception as e: logger.error(f"Dense indexing failed for doc {doc_id}: {e}") return {"success": False, "error": str(e)} async def index_document_sparse(self, text: str, doc_id: str, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Index a document in the sparse embedding service.""" async with httpx.AsyncClient(timeout=self.timeout) as client: try: # Sparse service now accepts doc_id directly payload = { "text": text, "doc_id": doc_id, # Pass doc_id directly "metadata": metadata or {} } response = await client.post(f"{self.sparse_url}/index", json=payload) response.raise_for_status() result = response.json() logger.debug(f"Sparse indexing successful for doc {doc_id}") return result except Exception as e: logger.error(f"Sparse indexing failed for doc {doc_id}: {e}") return {"success": False, "error": str(e)} async def index_document(self, text: str, doc_id: str, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Index a document in both services in parallel.""" logger.info(f"Indexing document {doc_id} in parallel...") # Run both indexing operations in parallel dense_task = self.index_document_dense(text, doc_id, metadata) sparse_task = self.index_document_sparse(text, doc_id, metadata) dense_result, sparse_result = await asyncio.gather(dense_task, sparse_task) return { "doc_id": doc_id, "dense": dense_result, "sparse": sparse_result, "success": dense_result.get("success", False) and sparse_result.get("success", False) } async def search_dense(self, query: str, top_k: int = 20) -> List[SearchResult]: """Search using dense embeddings.""" async with httpx.AsyncClient(timeout=self.timeout) as client: try: payload = { "query": query, "top_k": top_k, "return_documents": True } response = await client.post(f"{self.dense_url}/search", json=payload) response.raise_for_status() data = response.json() results = [] for item in data.get("results", []): results.append(SearchResult( doc_id=item["doc_id"], score=item["score"], text=item.get("text"), metadata=item.get("metadata"), source="dense", rank=item.get("rank"), debug_info={ "original_score": item["score"], "original_rank": item.get("rank", 0) } )) logger.debug(f"Dense search returned {len(results)} results") return results except Exception as e: logger.error(f"Dense search failed: {e}") return [] async def search_sparse(self, query: str, top_k: int = 20) -> List[SearchResult]: """Search using sparse embeddings (BM25).""" async with httpx.AsyncClient(timeout=self.timeout) as client: try: payload = { "query": query, "top_k": top_k } response = await client.post(f"{self.sparse_url}/search", json=payload) response.raise_for_status() data = response.json() results = [] for idx, item in enumerate(data): # Now doc_id is returned directly from the sparse service doc_id = item.get("doc_id", f"doc_{idx}") results.append(SearchResult( doc_id=doc_id, score=item["score"], text=item.get("text"), metadata=item.get("metadata"), source="sparse", rank=idx + 1, debug_info={ "bm25_score": item["score"], "matched_terms": item.get("debug", {}).get("matched_terms", []) if item.get("debug") else [], "doc_length": item.get("debug", {}).get("doc_length", 0) if item.get("debug") else 0, "original_rank": idx + 1 } )) logger.debug(f"Sparse search returned {len(results)} results") return results except Exception as e: logger.error(f"Sparse search failed: {e}") return [] async def search(self, query: str, top_k: int = 20, mode: str = "hybrid") -> Tuple[List[SearchResult], List[SearchResult]]: """Search using specified mode (dense, sparse, or hybrid).""" logger.info(f"Searching with mode: {mode}, query: '{query[:50]}...'") dense_results = [] sparse_results = [] if mode == "dense": dense_results = await self.search_dense(query, top_k) elif mode == "sparse": sparse_results = await self.search_sparse(query, top_k) elif mode == "hybrid": # Run both searches in parallel dense_task = self.search_dense(query, top_k) sparse_task = self.search_sparse(query, top_k) dense_results, sparse_results = await asyncio.gather(dense_task, sparse_task) else: raise ValueError(f"Invalid search mode: {mode}") return dense_results, sparse_results async def delete_document(self, doc_id: str) -> Dict[str, Any]: """Delete a document from both services.""" logger.info(f"Deleting document {doc_id} from both services...") async with httpx.AsyncClient(timeout=self.timeout) as client: # Delete from dense service dense_task = client.delete(f"{self.dense_url}/index", json={"doc_id": doc_id}) # For sparse service, we need to check if it supports deletion # If not, we'll need to rebuild the index sparse_task = client.delete(f"{self.sparse_url}/index") # Clear all for now try: dense_response, sparse_response = await asyncio.gather(dense_task, sparse_task) return { "doc_id": doc_id, "dense": dense_response.json() if dense_response.status_code == 200 else {"success": False}, "sparse": sparse_response.json() if sparse_response.status_code == 200 else {"success": False} } except Exception as e: logger.error(f"Failed to delete document {doc_id}: {e}") return {"success": False, "error": str(e)}