ai-agent-book/chapter3/retrieval-pipeline/test_client.py

"""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())