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

chapter3/retrieval-pipeline/demo.py

@@ -0,0 +1,281 @@
+"""Demo script showcasing dense vs sparse embedding strengths.
+
+Service Configuration:
+- Dense Embedding: http://localhost:4240
+- Sparse Embedding: http://localhost:4241  
+- Retrieval Pipeline: http://localhost:4242
+"""
+
+import asyncio
+import httpx
+from typing import Dict, List
+import json
+
+class RetrievalDemo:
+    """Demo for the retrieval pipeline."""
+    
+    def __init__(self, pipeline_url: str = "http://localhost:4242"):
+        self.pipeline_url = pipeline_url
+        
+    async def index_document(self, text: str, doc_id: str, metadata: Dict = None):
+        """Index a document."""
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                f"{self.pipeline_url}/index",
+                json={"text": text, "doc_id": doc_id, "metadata": metadata or {}}
+            )
+            return response.json()
+    
+    async def search(self, query: str, mode: str = "hybrid"):
+        """Search for documents."""
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                f"{self.pipeline_url}/search",
+                json={"query": query, "mode": mode, "top_k": 10, "rerank_top_k": 5}
+            )
+            return response.json()
+    
+    async def clear(self):
+        """Clear all documents."""
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.delete(f"{self.pipeline_url}/clear")
+            return response.json()
+
+async def main():
+    """Run the demonstration."""
+    demo = RetrievalDemo()
+    
+    print("="*80)
+    print("RETRIEVAL PIPELINE DEMONSTRATION")
+    print("Showcasing Dense vs Sparse Embedding Strengths")
+    print("="*80)
+    
+    # Clear existing documents
+    print("\nClearing existing documents...")
+    await demo.clear()
+    
+    # Create diverse test documents
+    documents = [
+        # Category 1: Programming Languages (for semantic similarity)
+        {
+            "doc_id": "prog_python",
+            "text": "Python is renowned for its clean syntax and readability, making it ideal for beginners and experts alike.",
+            "metadata": {"category": "programming", "subcategory": "languages"}
+        },
+        {
+            "doc_id": "prog_javascript",
+            "text": "JavaScript powers interactive web applications and runs in browsers worldwide.",
+            "metadata": {"category": "programming", "subcategory": "languages"}
+        },
+        {
+            "doc_id": "prog_rust",
+            "text": "Rust provides memory safety without garbage collection through its ownership system.",
+            "metadata": {"category": "programming", "subcategory": "languages"}
+        },
+        
+        # Category 2: Machine Learning (for concept matching)
+        {
+            "doc_id": "ml_intro",
+            "text": "Artificial intelligence enables computers to learn from data and make decisions.",
+            "metadata": {"category": "AI", "subcategory": "intro"}
+        },
+        {
+            "doc_id": "ml_deep",
+            "text": "Deep neural networks consist of multiple layers that progressively extract features.",
+            "metadata": {"category": "AI", "subcategory": "deep_learning"}
+        },
+        {
+            "doc_id": "ml_nlp",
+            "text": "Natural language processing helps machines understand and generate human text.",
+            "metadata": {"category": "AI", "subcategory": "NLP"}
+        },
+        
+        # Category 3: Specific Technical Terms (for exact matching)
+        {
+            "doc_id": "error_404",
+            "text": "HTTP status code 404 indicates that the requested resource was not found on the server.",
+            "metadata": {"category": "errors", "code": "404"}
+        },
+        {
+            "doc_id": "error_500",
+            "text": "HTTP status code 500 represents an internal server error that prevented request fulfillment.",
+            "metadata": {"category": "errors", "code": "500"}
+        },
+        {
+            "doc_id": "api_key",
+            "text": "The API key XK9-2B4-7Q1 provides access to premium features of the service.",
+            "metadata": {"category": "authentication", "type": "api_key"}
+        },
+        
+        # Category 4: Multilingual Content
+        {
+            "doc_id": "ml_chinese",
+            "text": "机器学习是人工智能的核心技术，通过数据训练模型来解决问题。",
+            "metadata": {"category": "AI", "language": "chinese"}
+        },
+        {
+            "doc_id": "ml_spanish",
+            "text": "El aprendizaje automático permite a las computadoras aprender sin programación explícita.",
+            "metadata": {"category": "AI", "language": "spanish"}
+        },
+        
+        # Category 5: People and Names
+        {
+            "doc_id": "person_turing",
+            "text": "Alan Turing pioneered computer science and artificial intelligence in the 20th century.",
+            "metadata": {"category": "people", "field": "computer_science"}
+        },
+        {
+            "doc_id": "person_lecun",
+            "text": "Yann LeCun developed convolutional neural networks that revolutionized computer vision.",
+            "metadata": {"category": "people", "field": "deep_learning"}
+        }
+    ]
+    
+    # Index all documents
+    print(f"\nIndexing {len(documents)} documents...")
+    for doc in documents:
+        result = await demo.index_document(
+            text=doc["text"],
+            doc_id=doc["doc_id"],
+            metadata=doc["metadata"]
+        )
+        success = result.get("success", False)
+        status = "✓" if success else "✗"
+        print(f"  {status} {doc['doc_id']}: {doc['text'][:60]}...")
+    
+    print("\n" + "="*80)
+    print("DEMONSTRATION QUERIES")
+    print("="*80)
+    
+    # Test queries demonstrating different strengths
+    test_queries = [
+        {
+            "query": "code readability and simplicity",
+            "description": "Semantic similarity - Dense should excel",
+            "expected_strong": "dense",
+            "explanation": "Dense embeddings understand 'readability' relates to Python even without exact match"
+        },
+        {
+            "query": "XK9-2B4-7Q1",
+            "description": "Exact code match - Sparse should excel",
+            "expected_strong": "sparse",
+            "explanation": "Sparse search finds exact API key string"
+        },
+        {
+            "query": "AI learning from examples",
+            "description": "Conceptual understanding - Dense should excel",
+            "expected_strong": "dense",
+            "explanation": "Dense understands AI/ML concepts without exact terminology"
+        },
+        {
+            "query": "404",
+            "description": "Specific error code - Sparse should excel",
+            "expected_strong": "sparse",
+            "explanation": "Sparse matches exact error code"
+        },
+        {
+            "query": "人工智能",
+            "description": "Cross-lingual search (Chinese for AI) - Dense should excel",
+            "expected_strong": "dense",
+            "explanation": "Dense embeddings (BGE-M3) handle multiple languages"
+        },
+        {
+            "query": "Yann LeCun",
+            "description": "Exact name search - Sparse should excel",
+            "expected_strong": "sparse",
+            "explanation": "Sparse finds exact person name"
+        },
+        {
+            "query": "web browser programming",
+            "description": "Semantic context - Dense should excel",
+            "expected_strong": "dense",
+            "explanation": "Dense connects 'web browser' with JavaScript"
+        }
+    ]
+    
+    # Run each test query
+    for test in test_queries:
+        print(f"\n{'='*60}")
+        print(f"Query: '{test['query']}'")
+        print(f"Type: {test['description']}")
+        print(f"Expected winner: {test['expected_strong']}")
+        print(f"Reason: {test['explanation']}")
+        print("-"*60)
+        
+        # Run search in all three modes
+        results = {}
+        for mode in ["dense", "sparse", "hybrid"]:
+            result = await demo.search(test["query"], mode=mode)
+            
+            # Extract top results
+            if mode == "dense":
+                top_docs = result.get("dense_results", [])[:3]
+            elif mode == "sparse":
+                top_docs = result.get("sparse_results", [])[:3]
+            else:  # hybrid
+                top_docs = result.get("reranked_results", [])[:3]
+            
+            results[mode] = top_docs
+            
+            # Print results for this mode
+            print(f"\n{mode.upper()} Results:")
+            for i, doc in enumerate(top_docs, 1):
+                doc_id = doc.get("doc_id", "unknown")
+                score = doc.get("score") or doc.get("rerank_score", 0)
+                print(f"  {i}. {doc_id} (score: {score:.4f})")
+        
+        # Analyze which mode performed best
+        print(f"\nAnalysis:")
+        if test["expected_strong"] == "dense":
+            if results["dense"] and results["sparse"]:
+                dense_top = results["dense"][0]["doc_id"] if results["dense"] else None
+                sparse_top = results["sparse"][0]["doc_id"] if results["sparse"] else None
+                if dense_top != sparse_top:
+                    print(f"  ✓ Dense found different (likely better) result: {dense_top}")
+                    print(f"  ✓ Sparse found: {sparse_top}")
+        elif test["expected_strong"] == "sparse":
+            if results["sparse"]:
+                print(f"  ✓ Sparse found exact match: {results['sparse'][0]['doc_id']}")
+        
+        # Show hybrid performance
+        if results["hybrid"]:
+            print(f"  ✓ Hybrid (reranked) top result: {results['hybrid'][0]['doc_id']}")
+    
+    print("\n" + "="*80)
+    print("DEMONSTRATION COMPLETE")
+    print("="*80)
+    print("\nKey Takeaways:")
+    print("1. Dense embeddings excel at semantic similarity and concepts")
+    print("2. Sparse search excels at exact matches and specific terms")
+    print("3. Hybrid search with reranking combines the best of both")
+    print("4. BGE-M3 dense embeddings support multilingual search")
+    print("5. BM25 sparse search is unbeatable for exact string matching")
+
+if __name__ == "__main__":
+    import sys
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Retrieval pipeline demonstration")
+    parser.add_argument("--url", default="http://localhost:4242",
+                       help="Pipeline service URL")
+    
+    args = parser.parse_args()
+    
+    # Check if services are running
+    print("Checking if services are available...")
+    print(f"Pipeline URL: {args.url}")
+    
+    try:
+        asyncio.run(main())
+    except httpx.ConnectError:
+        print("\nError: Could not connect to the retrieval pipeline service.")
+        print("Please ensure all services are running:")
+        print("  1. Dense embedding service (port 4240)")
+        print("  2. Sparse embedding service (port 4241)")
+        print("  3. Retrieval pipeline (port 4242)")
+        print("\nRun: ./restart_services.sh")
+        sys.exit(1)
+    except KeyboardInterrupt:
+        print("\nDemo interrupted by user")
+        sys.exit(0)