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

chapter3/retrieval-pipeline/main.py

@@ -0,0 +1,289 @@
+"""FastAPI server for the retrieval pipeline."""
+
+import logging
+import sys
+from typing import Dict, Any, Optional, List
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+import uvicorn
+import asyncio
+
+from config import PipelineConfig, SearchMode
+from retrieval_pipeline import RetrievalPipeline
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout)
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# Initialize pipeline
+config = PipelineConfig()
+pipeline: Optional[RetrievalPipeline] = None
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan context manager for startup and shutdown events."""
+    global pipeline
+    # Startup
+    try:
+        logger.info("Starting retrieval pipeline...")
+        pipeline = RetrievalPipeline(config)
+        logger.info("Pipeline initialized successfully")
+    except Exception as e:
+        logger.error(f"Failed to initialize pipeline: {e}")
+        raise
+    
+    yield
+    
+    # Shutdown (cleanup if needed)
+    logger.info("Shutting down retrieval pipeline...")
+
+# Create FastAPI app with lifespan
+app = FastAPI(
+    title="Hybrid Retrieval Pipeline",
+    description="Educational retrieval pipeline combining dense embeddings, sparse search, and neural reranking",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Pydantic models
+class IndexRequest(BaseModel):
+    """Request model for document indexing."""
+    text: str = Field(..., description="Document text to index")
+    doc_id: Optional[str] = Field(None, description="Optional document ID")
+    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Optional metadata")
+
+class SearchRequest(BaseModel):
+    """Request model for search."""
+    query: str = Field(..., description="Search query")
+    mode: SearchMode = Field(SearchMode.HYBRID, description="Search mode: dense, sparse, or hybrid")
+    top_k: int = Field(20, ge=1, le=100, description="Number of candidates to retrieve")
+    rerank_top_k: int = Field(10, ge=1, le=50, description="Number of results after reranking")
+    skip_reranking: bool = Field(False, description="Skip reranking for comparison")
+
+class DeleteRequest(BaseModel):
+    """Request model for document deletion."""
+    doc_id: str = Field(..., description="Document ID to delete")
+
+@app.get("/")
+async def root():
+    """Root endpoint with service information."""
+    return {
+        "service": "Hybrid Retrieval Pipeline",
+        "status": "running" if pipeline else "not initialized",
+        "endpoints": {
+            "index": "/index",
+            "search": "/search",
+            "delete": "/delete",
+            "stats": "/stats",
+            "health": "/health"
+        },
+        "modes": ["dense", "sparse", "hybrid"],
+        "features": [
+            "Parallel dense and sparse indexing",
+            "Hybrid search with both embedding types",
+            "Neural reranking with BGE-Reranker-v2",
+            "Educational score visualization",
+            "Rank change analysis"
+        ]
+    }
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    return {"status": "healthy"}
+
+@app.post("/index")
+async def index_document(request: IndexRequest):
+    """Index a document in both dense and sparse services.
+    
+    This endpoint:
+    1. Stores the document locally
+    2. Indexes in dense embedding service (semantic search)
+    3. Indexes in sparse service (BM25 keyword search)
+    4. Returns status from both services
+    """
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    
+    try:
+        logger.info(f"Indexing document: {request.doc_id or 'auto-generated'}")
+        result = await pipeline.index_document(
+            text=request.text,
+            doc_id=request.doc_id,
+            metadata=request.metadata
+        )
+        return result
+    except Exception as e:
+        logger.error(f"Indexing failed: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/search")
+async def search_documents(request: SearchRequest):
+    """Search for documents using specified mode with optional reranking.
+    
+    Educational features:
+    - Shows original rankings from dense and sparse search
+    - Displays similarity scores from each method
+    - Shows final reranked results with score changes
+    - Provides statistics on rank changes and overlap
+    
+    Modes:
+    - dense: Semantic search using dense embeddings
+    - sparse: Keyword search using BM25
+    - hybrid: Both methods combined with reranking
+    """
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    
+    try:
+        logger.info(f"Search request: mode={request.mode}, query='{request.query[:50]}...'")
+        result = await pipeline.search(
+            query=request.query,
+            mode=request.mode,
+            top_k=request.top_k,
+            rerank_top_k=request.rerank_top_k,
+            skip_reranking=request.skip_reranking
+        )
+        return result
+    except Exception as e:
+        logger.error(f"Search failed: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.delete("/delete")
+async def delete_document(request: DeleteRequest):
+    """Delete a document from all services."""
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    
+    try:
+        logger.info(f"Deleting document: {request.doc_id}")
+        result = await pipeline.delete_document(request.doc_id)
+        return result
+    except Exception as e:
+        logger.error(f"Deletion failed: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/stats")
+async def get_statistics():
+    """Get pipeline statistics and configuration."""
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    
+    try:
+        stats = pipeline.get_statistics()
+        return stats
+    except Exception as e:
+        logger.error(f"Failed to get statistics: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/documents")
+async def list_documents(limit: int = 100, offset: int = 0):
+    """List indexed documents."""
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    
+    try:
+        docs = pipeline.document_store.list_documents(limit=limit, offset=offset)
+        return {
+            "documents": docs,
+            "total": pipeline.document_store.size(),
+            "limit": limit,
+            "offset": offset
+        }
+    except Exception as e:
+        logger.error(f"Failed to list documents: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/documents/{doc_id}")
+async def get_document(doc_id: str):
+    """Get a specific document by ID."""
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    
+    try:
+        doc = pipeline.document_store.get_document(doc_id)
+        if doc is None:
+            raise HTTPException(status_code=404, detail=f"Document {doc_id} not found")
+        
+        # Return in the format expected by agentic RAG
+        return {
+            "doc_id": doc_id,
+            "content": doc.get("text", ""),
+            "metadata": doc.get("metadata", {})
+        }
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Failed to get document {doc_id}: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.delete("/clear")
+async def clear_all():
+    """Clear all documents from the pipeline (for testing)."""
+    if not pipeline:
+        raise HTTPException(status_code=503, detail="Pipeline not initialized")
+    
+    try:
+        pipeline.document_store.clear()
+        # Note: This only clears local store, not remote services
+        return {
+            "success": True,
+            "message": "Local document store cleared. Note: Remote services not cleared."
+        }
+    except Exception as e:
+        logger.error(f"Failed to clear documents: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+def main():
+    """Run the server."""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Retrieval Pipeline Server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind")
+    parser.add_argument("--port", type=int, default=4242, help="Port to bind (default: 4242)")
+    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
+    parser.add_argument("--dense-url", default="http://localhost:4240", help="Dense service URL")
+    parser.add_argument("--sparse-url", default="http://localhost:4241", help="Sparse service URL")
+    
+    args = parser.parse_args()
+    
+    # Update config
+    config.services.dense_service_url = args.dense_url
+    config.services.sparse_service_url = args.sparse_url
+    config.debug = args.debug
+    
+    if args.debug:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    logger.info(f"Starting server on {args.host}:{args.port}")
+    logger.info(f"Dense service: {config.services.dense_service_url}")
+    logger.info(f"Sparse service: {config.services.sparse_service_url}")
+    
+    uvicorn.run(
+        app,
+        host=args.host,
+        port=args.port,
+        log_level="debug" if args.debug else "info"
+    )
+
+if __name__ == "__main__":
+    main()