ai-agent-book/chapter3/contextual-retrieval/contextual_agent.py

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