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

chapter3/contextual-retrieval/contextual_agent.py

@@ -0,0 +1,356 @@
+"""Contextual Agentic RAG System
+
+This module extends the base AgenticRAG to use contextual retrieval,
+demonstrating the improved answer quality from better retrieval.
+"""
+
+import json
+import logging
+from typing import List, Dict, Any, Optional, Generator
+from datetime import datetime
+
+from agent import AgenticRAG, Message
+from config import Config
+from contextual_tools import ContextualKnowledgeBaseTools, ContextualSearchResult
+
+logger = logging.getLogger(__name__)
+
+
+class ContextualAgenticRAG(AgenticRAG):
+    """
+    Enhanced Agentic RAG with contextual retrieval support.
+    
+    Educational Features:
+    - Shows how better retrieval leads to better answers
+    - Logs retrieval quality metrics
+    - Compares contextual vs non-contextual results
+    """
+    
+    def __init__(self, 
+                 config: Optional[Config] = None,
+                 kb_tools: Optional[ContextualKnowledgeBaseTools] = None,
+                 use_contextual: bool = True):
+        """
+        Initialize contextual agent.
+        
+        Args:
+            config: Configuration object
+            kb_tools: Contextual knowledge base tools
+            use_contextual: Whether to use contextual retrieval
+        """
+        self.use_contextual = use_contextual
+        
+        # Initialize base class but skip KB tools initialization
+        self.config = config or Config.from_env()
+        self._init_llm_client()
+        
+        # Use provided contextual KB tools or create new ones
+        if kb_tools:
+            self.kb_tools = kb_tools
+        else:
+            self.kb_tools = ContextualKnowledgeBaseTools(
+                self.config.knowledge_base,
+                use_contextual=use_contextual,
+                enable_comparison=False
+            )
+        
+        # Initialize other components
+        self.conversation_history = []
+        self.tools = self._get_contextual_tool_definitions()
+        
+        # Track retrieval metrics
+        self.retrieval_metrics = {
+            "queries_executed": 0,
+            "total_chunks_retrieved": 0,
+            "avg_retrieval_score": 0.0,
+            "contextual_chunks_used": 0,
+            "non_contextual_chunks_used": 0
+        }
+        
+        mode = "CONTEXTUAL" if use_contextual else "NON-CONTEXTUAL"
+        logger.info(f"Initialized ContextualAgenticRAG in {mode} mode")
+    
+    def _get_system_prompt(self) -> str:
+        """Enhanced system prompt that leverages contextual information"""
+        base_prompt = super()._get_system_prompt()
+        
+        if self.use_contextual:
+            contextual_addition = """
+
+## Using Contextual Information:
+
+When you receive search results, pay attention to the contextual information provided. Each chunk may include:
+- Context that situates the chunk within the larger document
+- Information about what section or topic the chunk belongs to
+- Related entities and concepts mentioned
+
+Use this contextual information to:
+1. Better understand the relevance of each chunk
+2. Identify connections between different pieces of information
+3. Provide more accurate and complete answers
+
+Remember that contextual chunks have been enhanced with additional information to improve retrieval accuracy."""
+            
+            return base_prompt + contextual_addition
+        
+        return base_prompt
+    
+    def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Execute tool with enhanced logging for educational purposes.
+        
+        This override adds detailed logging to show how contextual
+        retrieval improves the search results.
+        """
+        logger.debug(f"Executing tool: {tool_name} with args: {arguments}")
+        
+        try:
+            if tool_name == "contextual_knowledge_search":
+                query = arguments.get("query", "")
+                method = arguments.get("method", "hybrid")
+                top_k = arguments.get("top_k", 20)
+                
+                logger.info(f"\n{'='*60}")
+                logger.info(f"CONTEXTUAL SEARCH EXECUTION")
+                logger.info(f"{'='*60}")
+                logger.info(f"Query: {query}")
+                logger.info(f"Method: {method}")
+                logger.info(f"Top-K: {top_k}")
+                logger.info(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
+                
+                # Perform search
+                results = self.kb_tools.contextual_search(query, method, top_k)
+                
+                if not results:
+                    logger.info("No results found")
+                    return {"status": "no_results", "message": "No relevant documents found"}
+                
+                # Log retrieval quality
+                avg_score = sum(r.score for r in results) / len(results)
+                logger.info(f"\nRetrieval Statistics:")
+                logger.info(f"  Results found: {len(results)}")
+                logger.info(f"  Average score: {avg_score:.4f}")
+                logger.info(f"  Top score: {results[0].score:.4f}")
+                
+                # Show distribution of retrieval methods if hybrid
+                if method == "hybrid":
+                    bm25_only = sum(1 for r in results if r.bm25_score > 0 and r.embedding_score == 0)
+                    embedding_only = sum(1 for r in results if r.embedding_score > 0 and r.bm25_score == 0)
+                    both = sum(1 for r in results if r.bm25_score > 0 and r.embedding_score > 0)
+                    
+                    logger.info(f"\nRetrieval Method Distribution:")
+                    logger.info(f"  BM25 only: {bm25_only}")
+                    logger.info(f"  Embedding only: {embedding_only}")
+                    logger.info(f"  Both methods: {both}")
+                
+                # Update metrics
+                self.retrieval_metrics["queries_executed"] += 1
+                self.retrieval_metrics["total_chunks_retrieved"] += len(results)
+                
+                # Calculate running average score
+                prev_avg = self.retrieval_metrics["avg_retrieval_score"]
+                n = self.retrieval_metrics["queries_executed"]
+                self.retrieval_metrics["avg_retrieval_score"] = (
+                    (prev_avg * (n - 1) + avg_score) / n
+                )
+                
+                if self.use_contextual:
+                    self.retrieval_metrics["contextual_chunks_used"] += len(results)
+                else:
+                    self.retrieval_metrics["non_contextual_chunks_used"] += len(results)
+                
+                # Format results for agent
+                formatted_results = []
+                for i, r in enumerate(results, 1):
+                    result_dict = {
+                        "rank": i,
+                        "doc_id": r.doc_id,
+                        "chunk_id": r.chunk_id,
+                        "text": r.text,
+                        "score": r.score
+                    }
+                    
+                    formatted_results.append(result_dict)
+                    
+                    # Log top results
+                    if i <= 3:
+                        logger.info(f"\nTop Result {i}:")
+                        logger.info(f"  Score: {r.score:.4f}")
+                        logger.info(f"  Text: {r.text}")
+                
+                logger.info(f"{'='*60}\n")
+                
+                return {
+                    "status": "success",
+                    "results": formatted_results,
+                    "total_found": len(results),
+                    "avg_score": avg_score,
+                    "search_method": method,
+                    "is_contextual": self.use_contextual
+                }
+            
+            elif tool_name == "get_document":
+                # Use parent implementation
+                return super()._execute_tool(tool_name, arguments)
+            
+            else:
+                return {"status": "error", "message": f"Unknown tool: {tool_name}"}
+                
+        except Exception as e:
+            logger.error(f"Tool execution error: {e}")
+            return {"status": "error", "message": str(e)}
+    
+    def _get_contextual_tool_definitions(self) -> List[Dict[str, Any]]:
+        """Get tool definitions with contextual search capabilities"""
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "contextual_knowledge_search",
+                    "description": "Search the knowledge base using contextual or non-contextual retrieval. Returns chunks with optional context information that helps understand their relevance.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string",
+                                "description": "Natural language search query"
+                            },
+                            "method": {
+                                "type": "string",
+                                "enum": ["bm25", "embedding", "hybrid"],
+                                "description": "Search method to use (default: hybrid)",
+                                "default": "hybrid"
+                            },
+                            "top_k": {
+                                "type": "integer",
+                                "description": "Number of results to retrieve (default: 20)",
+                                "default": 20
+                            }
+                        },
+                        "required": ["query"]
+                    }
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_document",
+                    "description": "Retrieve the complete content of a specific document",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "doc_id": {
+                                "type": "string",
+                                "description": "Document ID to retrieve"
+                            }
+                        },
+                        "required": ["doc_id"]
+                    }
+                }
+            }
+        ]
+    
+    def query(self, user_query: str, stream: bool = None) -> Any:
+        """
+        Process query with contextual retrieval and detailed logging.
+        
+        This override adds educational logging to show the retrieval process.
+        """
+        logger.info(f"\n{'='*80}")
+        logger.info(f"PROCESSING QUERY: {user_query}")
+        logger.info(f"Mode: {'CONTEXTUAL' if self.use_contextual else 'NON-CONTEXTUAL'}")
+        logger.info(f"{'='*80}\n")
+        
+        # Call parent implementation
+        result = super().query(user_query, stream)
+        
+        # Log final metrics
+        logger.info(f"\n{'='*60}")
+        logger.info("QUERY PROCESSING COMPLETE")
+        logger.info(f"{'='*60}")
+        logger.info("Retrieval Metrics:")
+        logger.info(f"  Queries executed: {self.retrieval_metrics['queries_executed']}")
+        logger.info(f"  Total chunks retrieved: {self.retrieval_metrics['total_chunks_retrieved']}")
+        logger.info(f"  Average retrieval score: {self.retrieval_metrics['avg_retrieval_score']:.4f}")
+        
+        if self.use_contextual:
+            logger.info(f"  Contextual chunks used: {self.retrieval_metrics['contextual_chunks_used']}")
+        else:
+            logger.info(f"  Non-contextual chunks used: {self.retrieval_metrics['non_contextual_chunks_used']}")
+        
+        logger.info(f"{'='*60}\n")
+        
+        return result
+    
+    def get_retrieval_metrics(self) -> Dict[str, Any]:
+        """Get detailed retrieval metrics for analysis"""
+        metrics = self.retrieval_metrics.copy()
+        
+        # Add KB statistics
+        kb_stats = self.kb_tools.get_statistics()
+        metrics["kb_stats"] = kb_stats
+        
+        # Add mode information
+        metrics["mode"] = "contextual" if self.use_contextual else "non_contextual"
+        metrics["timestamp"] = datetime.now().isoformat()
+        
+        return metrics
+    
+    def compare_modes(self, user_query: str) -> Dict[str, Any]:
+        """
+        Compare contextual vs non-contextual modes for the same query.
+        
+        Educational method to demonstrate the difference in retrieval quality.
+        """
+        comparison = {
+            "query": user_query,
+            "timestamp": datetime.now().isoformat(),
+            "modes": {}
+        }
+        
+        # Test contextual mode
+        logger.info("Testing CONTEXTUAL mode...")
+        self.use_contextual = True
+        self.kb_tools.use_contextual = True
+        contextual_response = self.query(user_query, stream=False)
+        comparison["modes"]["contextual"] = {
+            "response": contextual_response,
+            "metrics": self.get_retrieval_metrics()
+        }
+        
+        # Reset metrics
+        self.retrieval_metrics = {
+            "queries_executed": 0,
+            "total_chunks_retrieved": 0,
+            "avg_retrieval_score": 0.0,
+            "contextual_chunks_used": 0,
+            "non_contextual_chunks_used": 0
+        }
+        
+        # Test non-contextual mode
+        logger.info("Testing NON-CONTEXTUAL mode...")
+        self.use_contextual = False
+        self.kb_tools.use_contextual = False
+        non_contextual_response = self.query(user_query, stream=False)
+        comparison["modes"]["non_contextual"] = {
+            "response": non_contextual_response,
+            "metrics": self.get_retrieval_metrics()
+        }
+        
+        # Analyze differences
+        contextual_score = comparison["modes"]["contextual"]["metrics"]["avg_retrieval_score"]
+        non_contextual_score = comparison["modes"]["non_contextual"]["metrics"]["avg_retrieval_score"]
+        
+        improvement = ((contextual_score - non_contextual_score) / non_contextual_score * 100) if non_contextual_score > 0 else 0
+        
+        comparison["analysis"] = {
+            "score_improvement_pct": improvement,
+            "contextual_avg_score": contextual_score,
+            "non_contextual_avg_score": non_contextual_score,
+            "recommendation": (
+                "Use contextual retrieval" if improvement > 10
+                else "Consider cost-benefit of contextual retrieval"
+            )
+        }
+        
+        return comparison