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

chapter3/retrieval-pipeline/test_client.py

@@ -0,0 +1,497 @@
+"""Test client with educational test cases for dense vs sparse retrieval."""
+
+import asyncio
+import httpx
+import json
+from typing import List, Dict, Any
+import logging
+from datetime import datetime
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class TestClient:
+    """Test client for the retrieval pipeline."""
+    
+    def __init__(self, base_url: str = "http://localhost:4242"):
+        self.base_url = base_url.rstrip('/')
+        self.test_results = []
+    
+    async def index_document(self, text: str, doc_id: str = None, metadata: Dict = None) -> Dict:
+        """Index a document."""
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                f"{self.base_url}/index",
+                json={"text": text, "doc_id": doc_id, "metadata": metadata or {}}
+            )
+            return response.json()
+    
+    async def search(self, query: str, mode: str = "hybrid", top_k: int = 20, rerank_top_k: int = 10) -> Dict:
+        """Search for documents."""
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                f"{self.base_url}/search",
+                json={
+                    "query": query,
+                    "mode": mode,
+                    "top_k": top_k,
+                    "rerank_top_k": rerank_top_k
+                }
+            )
+            return response.json()
+    
+    async def clear_documents(self) -> Dict:
+        """Clear all documents."""
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.delete(f"{self.base_url}/clear")
+            return response.json()
+    
+    def print_results(self, results: Dict, title: str = "Search Results"):
+        """Pretty print search results."""
+        print(f"\n{'='*80}")
+        print(f"{title}")
+        print(f"{'='*80}")
+        print(f"Query: {results.get('query', 'N/A')}")
+        print(f"Mode: {results.get('mode', 'N/A')}")
+        print(f"Times: Retrieval={results.get('retrieval_time_ms', 0):.1f}ms, "
+              f"Rerank={results.get('rerank_time_ms', 0):.1f}ms, "
+              f"Total={results.get('total_time_ms', 0):.1f}ms")
+        
+        # Show top dense results
+        if results.get('dense_results'):
+            print(f"\nTop Dense Results:")
+            for r in results['dense_results'][:5]:
+                print(f"  #{r['rank']}: {r['doc_id']} (score: {r['score']:.4f})")
+        
+        # Show top sparse results
+        if results.get('sparse_results'):
+            print(f"\nTop Sparse Results:")
+            for r in results['sparse_results'][:5]:
+                matched = r.get('matched_terms', [])
+                print(f"  #{r['rank']}: {r['doc_id']} (score: {r['score']:.4f}, matched: {matched})")
+        
+        # Show reranked results
+        if results.get('reranked_results'):
+            print(f"\nReranked Results:")
+            for r in results['reranked_results'][:5]:
+                changes = r.get('rank_changes', [])
+                print(f"  #{r['rank']}: {r['doc_id']} (score: {r['rerank_score']:.4f})")
+                if changes:
+                    print(f"    Rank changes: {', '.join(changes)}")
+        
+        # Show statistics
+        if results.get('statistics'):
+            stats = results['statistics']
+            print(f"\nStatistics:")
+            print(f"  Dense retrieved: {stats.get('dense_retrieved', 0)}")
+            print(f"  Sparse retrieved: {stats.get('sparse_retrieved', 0)}")
+            print(f"  Overlap: {stats.get('overlap_count', 0)} ({stats.get('overlap_percentage', 0):.1f}%)")
+    
+    async def run_test_case(self, name: str, documents: List[Dict], queries: List[Dict]) -> Dict:
+        """Run a complete test case."""
+        print(f"\n{'='*80}")
+        print(f"TEST CASE: {name}")
+        print(f"{'='*80}")
+        
+        test_result = {
+            "name": name,
+            "timestamp": datetime.now().isoformat(),
+            "documents": len(documents),
+            "queries": len(queries),
+            "results": []
+        }
+        
+        # Index documents
+        print(f"\nIndexing {len(documents)} documents...")
+        for doc in documents:
+            result = await self.index_document(
+                text=doc["text"],
+                doc_id=doc.get("doc_id"),
+                metadata=doc.get("metadata", {})
+            )
+            print(f"  Indexed: {doc.get('doc_id', 'auto')} - {doc['text'][:50]}...")
+        
+        # Run queries
+        print(f"\nRunning {len(queries)} queries...")
+        for query_spec in queries:
+            query = query_spec["query"]
+            expected = query_spec.get("expected", [])
+            explanation = query_spec.get("explanation", "")
+            
+            print(f"\nQuery: '{query}'")
+            if explanation:
+                print(f"Explanation: {explanation}")
+            if expected:
+                print(f"Expected top results: {expected}")
+            
+            # Test all modes
+            for mode in ["dense", "sparse", "hybrid"]:
+                print(f"\n--- Mode: {mode} ---")
+                result = await self.search(query, mode=mode, top_k=10, rerank_top_k=5)
+                
+                # Extract top results
+                top_results = []
+                if mode == "hybrid" and result.get("reranked_results"):
+                    top_results = [r["doc_id"] for r in result["reranked_results"][:3]]
+                elif mode == "dense" and result.get("dense_results"):
+                    top_results = [r["doc_id"] for r in result["dense_results"][:3]]
+                elif mode == "sparse" and result.get("sparse_results"):
+                    top_results = [r["doc_id"] for r in result["sparse_results"][:3]]
+                
+                print(f"Top 3: {top_results}")
+                
+                # Check if expected results are in top positions
+                if expected:
+                    matches = [doc_id in top_results for doc_id in expected]
+                    accuracy = sum(matches) / len(expected) * 100
+                    print(f"Accuracy: {accuracy:.0f}% ({sum(matches)}/{len(expected)} expected found)")
+                
+                test_result["results"].append({
+                    "query": query,
+                    "mode": mode,
+                    "top_results": top_results,
+                    "expected": expected,
+                    "time_ms": result.get("total_time_ms", 0)
+                })
+        
+        self.test_results.append(test_result)
+        return test_result
+
+# Test cases demonstrating dense vs sparse strengths
+async def run_educational_tests():
+    """Run educational test cases."""
+    client = TestClient()
+    
+    # Clear existing documents
+    await client.clear_documents()
+    
+    # Test Case 1: Semantic Similarity (Dense is better)
+    semantic_docs = [
+        {
+            "doc_id": "cat_1",
+            "text": "The feline jumped onto the couch and purred contentedly.",
+            "metadata": {"category": "animals", "type": "behavior"}
+        },
+        {
+            "doc_id": "cat_2", 
+            "text": "A tabby cat sleeps on the windowsill in the afternoon sun.",
+            "metadata": {"category": "animals", "type": "description"}
+        },
+        {
+            "doc_id": "dog_1",
+            "text": "The puppy barked excitedly and wagged its tail.",
+            "metadata": {"category": "animals", "type": "behavior"}
+        },
+        {
+            "doc_id": "car_1",
+            "text": "The vehicle accelerated down the highway.",
+            "metadata": {"category": "transportation", "type": "action"}
+        }
+    ]
+    
+    semantic_queries = [
+        {
+            "query": "kitty behavior",  # Uses different words but same concept
+            "expected": ["cat_1", "cat_2"],
+            "explanation": "Dense should find cat documents despite using 'kitty' instead of 'cat/feline'"
+        },
+        {
+            "query": "automobile speed",  # Semantic similarity to car/vehicle
+            "expected": ["car_1"],
+            "explanation": "Dense should match 'automobile' to 'vehicle' and 'speed' to 'accelerated'"
+        }
+    ]
+    
+    await client.run_test_case(
+        "Semantic Similarity (Dense Advantage)",
+        semantic_docs,
+        semantic_queries
+    )
+    
+    # Test Case 2: Exact Terms and Names (Sparse is better)
+    exact_docs = [
+        {
+            "doc_id": "person_1",
+            "text": "Dr. Alexander Humphrey published groundbreaking research on quantum computing.",
+            "metadata": {"type": "person", "field": "science"}
+        },
+        {
+            "doc_id": "person_2",
+            "text": "Professor Smith teaches computer science at the university.",
+            "metadata": {"type": "person", "field": "education"}
+        },
+        {
+            "doc_id": "company_1",
+            "text": "XR-7000 is a new model released by TechCorp Industries.",
+            "metadata": {"type": "product", "company": "TechCorp"}
+        },
+        {
+            "doc_id": "company_2",
+            "text": "The latest smartphone features advanced technology.",
+            "metadata": {"type": "product", "category": "electronics"}
+        }
+    ]
+    
+    exact_queries = [
+        {
+            "query": "Alexander Humphrey",  # Exact name match
+            "expected": ["person_1"],
+            "explanation": "Sparse should excel at finding exact name 'Alexander Humphrey'"
+        },
+        {
+            "query": "XR-7000",  # Specific product code
+            "expected": ["company_1"],
+            "explanation": "Sparse should find exact product code 'XR-7000'"
+        }
+    ]
+    
+    await client.run_test_case(
+        "Exact Terms and Names (Sparse Advantage)",
+        exact_docs,
+        exact_queries
+    )
+    
+    # Test Case 3: Multilingual (Dense is better)
+    multilingual_docs = [
+        {
+            "doc_id": "ml_en_1",
+            "text": "Machine learning is a subset of artificial intelligence.",
+            "metadata": {"language": "english", "topic": "AI"}
+        },
+        {
+            "doc_id": "ml_zh_1",
+            "text": "机器学习是人工智能的一个子集。",  # Same content in Chinese
+            "metadata": {"language": "chinese", "topic": "AI"}
+        },
+        {
+            "doc_id": "ml_es_1",
+            "text": "El aprendizaje automático es un subconjunto de la inteligencia artificial.",  # Spanish
+            "metadata": {"language": "spanish", "topic": "AI"}
+        },
+        {
+            "doc_id": "other_1",
+            "text": "Database systems store and retrieve information efficiently.",
+            "metadata": {"language": "english", "topic": "database"}
+        }
+    ]
+    
+    multilingual_queries = [
+        {
+            "query": "AI learning",  # English query
+            "expected": ["ml_en_1", "ml_zh_1", "ml_es_1"],
+            "explanation": "Dense embeddings (BGE-M3) should find similar content across languages"
+        },
+        {
+            "query": "人工智能",  # Chinese query for "artificial intelligence"
+            "expected": ["ml_zh_1", "ml_en_1"],
+            "explanation": "Dense should match Chinese query to related documents in any language"
+        }
+    ]
+    
+    await client.run_test_case(
+        "Multilingual Matching (Dense Advantage)",
+        multilingual_docs,
+        multilingual_queries
+    )
+    
+    # Test Case 4: Technical Terms and Codes (Sparse is better)
+    technical_docs = [
+        {
+            "doc_id": "error_1",
+            "text": "Error code HTTP-403 indicates forbidden access to the resource.",
+            "metadata": {"type": "error", "category": "http"}
+        },
+        {
+            "doc_id": "error_2",
+            "text": "The system returned status 500 for internal server problems.",
+            "metadata": {"type": "error", "category": "http"}
+        },
+        {
+            "doc_id": "config_1",
+            "text": "Set parameter MAX_BUFFER_SIZE=8192 in the configuration file.",
+            "metadata": {"type": "configuration"}
+        },
+        {
+            "doc_id": "generic_1",
+            "text": "The application encountered an issue during startup.",
+            "metadata": {"type": "error", "category": "general"}
+        }
+    ]
+    
+    technical_queries = [
+        {
+            "query": "HTTP-403",  # Exact error code
+            "expected": ["error_1"],
+            "explanation": "Sparse should match exact error code 'HTTP-403'"
+        },
+        {
+            "query": "MAX_BUFFER_SIZE",  # Exact parameter name
+            "expected": ["config_1"],
+            "explanation": "Sparse should find exact configuration parameter"
+        }
+    ]
+    
+    await client.run_test_case(
+        "Technical Terms and Codes (Sparse Advantage)",
+        technical_docs,
+        technical_queries
+    )
+    
+    # Test Case 5: Conceptual Understanding (Dense is better)
+    conceptual_docs = [
+        {
+            "doc_id": "happy_1",
+            "text": "She was filled with joy and couldn't stop smiling.",
+            "metadata": {"emotion": "positive"}
+        },
+        {
+            "doc_id": "happy_2",
+            "text": "His elation was evident as he celebrated the victory.",
+            "metadata": {"emotion": "positive"}
+        },
+        {
+            "doc_id": "sad_1",
+            "text": "Tears rolled down her face as she felt overwhelmed with sorrow.",
+            "metadata": {"emotion": "negative"}
+        },
+        {
+            "doc_id": "neutral_1",
+            "text": "The meeting proceeded according to the scheduled agenda.",
+            "metadata": {"emotion": "neutral"}
+        }
+    ]
+    
+    conceptual_queries = [
+        {
+            "query": "happiness and excitement",  # Concept not exact words
+            "expected": ["happy_1", "happy_2"],
+            "explanation": "Dense should understand happiness concept despite different words (joy, elation)"
+        },
+        {
+            "query": "melancholy mood",  # Related to sadness
+            "expected": ["sad_1"],
+            "explanation": "Dense should connect 'melancholy' with 'sorrow' conceptually"
+        }
+    ]
+    
+    await client.run_test_case(
+        "Conceptual Understanding (Dense Advantage)",
+        conceptual_docs,
+        conceptual_queries
+    )
+    
+    # Print summary
+    print(f"\n{'='*80}")
+    print("TEST SUMMARY")
+    print(f"{'='*80}")
+    print(f"Total test cases: {len(client.test_results)}")
+    
+    for test in client.test_results:
+        print(f"\n{test['name']}:")
+        print(f"  Documents: {test['documents']}")
+        print(f"  Queries: {test['queries']}")
+        print(f"  Total searches: {len(test['results'])}")
+
+async def run_interactive_demo():
+    """Run an interactive demonstration."""
+    client = TestClient()
+    
+    print("\n" + "="*80)
+    print("INTERACTIVE RETRIEVAL PIPELINE DEMO")
+    print("="*80)
+    print("\nThis demo shows how dense and sparse retrieval work differently.")
+    print("Dense is better for: semantic similarity, concepts, multilingual")
+    print("Sparse is better for: exact names, codes, technical terms")
+    
+    # Sample documents for interactive demo
+    sample_docs = [
+        {
+            "doc_id": "python_intro",
+            "text": "Python is a high-level programming language known for its simplicity and readability.",
+            "metadata": {"category": "programming", "language": "english"}
+        },
+        {
+            "doc_id": "python_syntax",
+            "text": "def hello_world(): print('Hello, World!') is a simple Python function.",
+            "metadata": {"category": "code", "language": "english"}
+        },
+        {
+            "doc_id": "ml_basics",
+            "text": "Machine learning algorithms learn patterns from data without explicit programming.",
+            "metadata": {"category": "AI", "language": "english"}
+        },
+        {
+            "doc_id": "深度学习",
+            "text": "深度学习是机器学习的一个分支，使用神经网络处理复杂数据。",
+            "metadata": {"category": "AI", "language": "chinese"}
+        },
+        {
+            "doc_id": "api_error",
+            "text": "API returned error code E-2001: Invalid authentication token provided.",
+            "metadata": {"category": "error", "type": "api"}
+        }
+    ]
+    
+    print("\nIndexing sample documents...")
+    await client.clear_documents()
+    
+    for doc in sample_docs:
+        await client.index_document(
+            text=doc["text"],
+            doc_id=doc["doc_id"],
+            metadata=doc.get("metadata", {})
+        )
+        print(f"  ✓ {doc['doc_id']}: {doc['text'][:60]}...")
+    
+    # Interactive queries
+    queries = [
+        ("coding simplicity", "Should find Python docs via semantic similarity"),
+        ("E-2001", "Should find exact error code via sparse search"),
+        ("neural networks", "Should find ML/DL docs including Chinese via dense"),
+        ("hello_world", "Should find exact function name via sparse")
+    ]
+    
+    print("\n" + "="*80)
+    print("RUNNING COMPARISON QUERIES")
+    print("="*80)
+    
+    for query, explanation in queries:
+        print(f"\nQuery: '{query}'")
+        print(f"Expected: {explanation}")
+        print("-" * 40)
+        
+        # Compare all three modes
+        modes_results = {}
+        for mode in ["dense", "sparse", "hybrid"]:
+            result = await client.search(query, mode=mode, top_k=5, rerank_top_k=3)
+            
+            # Get top results based on mode
+            if mode == "hybrid" and result.get("reranked_results"):
+                top = [r["doc_id"] for r in result["reranked_results"][:3]]
+            elif mode == "dense" and result.get("dense_results"):
+                top = [r["doc_id"] for r in result["dense_results"][:3]]
+            elif mode == "sparse" and result.get("sparse_results"):
+                top = [r["doc_id"] for r in result["sparse_results"][:3]]
+            else:
+                top = []
+            
+            modes_results[mode] = top
+            print(f"{mode:8}: {top}")
+        
+        # Show which mode performed best
+        print("-" * 40)
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Test client for retrieval pipeline")
+    parser.add_argument("--url", default="http://localhost:4242", help="Pipeline service URL")
+    parser.add_argument("--mode", choices=["test", "demo"], default="test",
+                       help="Run mode: test (all test cases) or demo (interactive)")
+    
+    args = parser.parse_args()
+    
+    if args.mode == "test":
+        asyncio.run(run_educational_tests())
+    else:
+        asyncio.run(run_interactive_demo())