ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/retrieval-pipeline/retrieval_client.py

@@ -0,0 +1,201 @@
+"""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)}