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

chapter3/agentic-rag/agent.py

@@ -0,0 +1,459 @@
+"""Agentic RAG System with ReAct Pattern"""
+
+import json
+import logging
+from typing import List, Dict, Any, Optional, Generator
+from dataclasses import dataclass, field
+from datetime import datetime
+from openai import OpenAI
+
+from config import Config, LLMConfig, AgentConfig
+from tools import KnowledgeBaseTools, get_tool_definitions
+
+
+def _is_reasoning_model(model) -> bool:
+    """Whether the model is a reasoning model (Kimi K3, GPT-5, ...)."""
+    m = str(model or "").lower().replace("/", "-")
+    return "kimi-k3" in m or "gpt-5" in m
+
+
+def _reasoning_safe_temperature(model, requested=1.0):
+    """Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
+    Return 1 for those; otherwise the requested value so non-reasoning
+    providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
+    return 1 if _is_reasoning_model(model) else requested
+
+
+def _reasoning_safe_max_tokens(model, requested=1024, floor=4096):
+    """Reasoning models spend part of their budget on hidden reasoning tokens,
+    so a small ``max_tokens`` (e.g. 1024) silently truncates the visible answer.
+    Ensure reasoning models get at least ``floor`` tokens; leave other providers
+    at the requested value."""
+    if _is_reasoning_model(model):
+        return max(requested, floor)
+    return requested
+
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Message:
+    """Represents a message in the conversation"""
+    role: str  # "user", "assistant", "tool"
+    content: str
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    tool_call_id: Optional[str] = None
+    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
+
+
+class AgenticRAG:
+    """Agentic RAG system with ReAct pattern and multiple LLM provider support"""
+    
+    def __init__(self, config: Optional[Config] = None):
+        """Initialize the agent"""
+        self.config = config or Config.from_env()
+        
+        # Initialize LLM client
+        self._init_llm_client()
+        
+        # Initialize knowledge base tools
+        self.kb_tools = KnowledgeBaseTools(self.config.knowledge_base)
+        
+        # Conversation history
+        self.conversation_history: List[Dict[str, Any]] = []
+        
+        # Tool definitions
+        self.tools = get_tool_definitions()
+        
+        logger.info(f"Initialized AgenticRAG with provider: {self.config.llm.provider}")
+    
+    def _init_llm_client(self):
+        """Initialize the LLM client based on provider"""
+        client_config, model = self.config.llm.get_client_config()
+        
+        # Extract base_url if present
+        base_url = client_config.pop("base_url", None)
+        
+        # Create OpenAI client
+        if base_url:
+            self.client = OpenAI(base_url=base_url, **client_config)
+        else:
+            self.client = OpenAI(**client_config)
+        
+        self.model = model
+        logger.info(f"Using model: {self.model}")
+    
+    def _get_system_prompt(self) -> str:
+        """Generate the system prompt"""
+        return """You are an intelligent assistant with access to a knowledge base. Your primary role is to answer questions accurately based on the information available in the knowledge base.
+
+## Important Guidelines:
+
+1. **Knowledge Base Only**: You MUST only answer questions based on information found in the knowledge base. If the information is not available, clearly state that you cannot answer based on the available knowledge.
+
+2. **Use Tools Effectively**:
+   - Use `knowledge_base_search` to search for relevant information
+   - Use `get_document` to retrieve complete documents when you need more context
+   - You may need multiple searches with different queries to fully answer a question
+
+3. **Citations Required**: Always include citations in your answers. Format citations as [Doc: document_id] or [Chunk: chunk_id] inline with your response.
+
+4. **Reasoning Process**: Think step-by-step:
+   - First, understand what information is needed
+   - Search for relevant information
+   - If needed, retrieve full documents for context
+   - Synthesize the information to answer the question
+   - Include proper citations
+
+5. **Handle Follow-ups**: For follow-up questions, consider the conversation context but always verify information from the knowledge base.
+
+6. **Be Accurate**: Never make up information. If something is unclear or not found, say so explicitly.
+
+Remember: Your credibility depends on providing accurate, well-cited information from the knowledge base only."""
+    
+    def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """Execute a tool and return the result"""
+        try:
+            if tool_name == "knowledge_base_search":
+                query = arguments.get("query", "")
+                results = self.kb_tools.knowledge_base_search(query)
+                
+                # Log full trajectory when verbose
+                if self.config.agent.verbose:
+                    logger.info("=" * 80)
+                    logger.info(f"TOOL EXECUTION: {tool_name}")
+                    logger.info("-" * 80)
+                    logger.info(f"Query: {query}")
+                    logger.info("-" * 80)
+                
+                if not results:
+                    if self.config.agent.verbose:
+                        logger.info("Results: No relevant documents found")
+                        logger.info("=" * 80)
+                    return {"status": "no_results", "message": "No relevant documents found"}
+                
+                # Format results for agent - KEEP ALL RESULTS
+                formatted_results = []
+                for i, r in enumerate(results, 1):
+                    formatted_results.append({
+                        "doc_id": r["doc_id"],
+                        "chunk_id": r["chunk_id"],
+                        "text": r["text"],
+                        "score": r["score"]
+                    })
+                    
+                    # Log each result in full detail
+                    if self.config.agent.verbose:
+                        logger.info(f"Result {i}/{len(results)}:")
+                        logger.info(f"  Document ID: {r['doc_id']}")
+                        logger.info(f"  Chunk ID: {r['chunk_id']}")
+                        logger.info(f"  Score: {r['score']:.4f}")
+                        logger.info(f"  Text (full):\n{'-' * 40}")
+                        logger.info(r['text'])
+                        logger.info("-" * 40)
+                
+                if self.config.agent.verbose:
+                    logger.info(f"Total results found: {len(results)}")
+                    logger.info("=" * 80)
+                
+                return {
+                    "status": "success",
+                    "results": formatted_results[:3],  # Limit to top 3 for LLM context
+                    "total_found": len(results),
+                    "all_results": formatted_results  # Keep all for logging
+                }
+                
+            elif tool_name == "get_document":
+                doc_id = arguments.get("doc_id", "")
+                
+                # Log full trajectory when verbose
+                if self.config.agent.verbose:
+                    logger.info("=" * 80)
+                    logger.info(f"TOOL EXECUTION: {tool_name}")
+                    logger.info("-" * 80)
+                    logger.info(f"Document ID: {doc_id}")
+                    logger.info("-" * 80)
+                
+                document = self.kb_tools.get_document(doc_id)
+                
+                if "error" in document:
+                    if self.config.agent.verbose:
+                        logger.info(f"Error: {document['error']}")
+                        logger.info("=" * 80)
+                    return {"status": "error", "message": document["error"]}
+                
+                # Log full document content
+                if self.config.agent.verbose:
+                    logger.info("Document Retrieved:")
+                    logger.info(f"  Doc ID: {document.get('doc_id', doc_id)}")
+                    if document.get('metadata'):
+                        logger.info(f"  Metadata: {json.dumps(document['metadata'], indent=2, ensure_ascii=False)}")
+                    logger.info("  Content (full):\n" + "=" * 40)
+                    logger.info(document.get('content', ''))
+                    logger.info("=" * 80)
+                
+                return {
+                    "status": "success",
+                    "document": {
+                        "doc_id": document.get("doc_id", doc_id),
+                        "content": document.get("content", ""),
+                        "metadata": document.get("metadata", {})
+                    }
+                }
+            
+            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 _build_messages(self, user_query: str) -> List[Dict[str, Any]]:
+        """Build messages for the LLM including conversation history"""
+        messages = [{"role": "system", "content": self._get_system_prompt()}]
+        
+        # Add conversation history (limited)
+        history_limit = self.config.agent.conversation_history_limit
+        # limit<=0 → no history; list[-0:] would include all turns.
+        if history_limit > 0:
+            if len(self.conversation_history) > history_limit:
+                messages.extend(self.conversation_history[-history_limit:])
+            else:
+                messages.extend(self.conversation_history)
+        
+        # Add current user query
+        messages.append({"role": "user", "content": user_query})
+        
+        return messages
+    
+    def query(self, user_query: str, stream: bool = None) -> Any:
+        """
+        Process a user query using the ReAct pattern.
+        
+        Args:
+            user_query: The user's question
+            stream: Whether to stream the response
+            
+        Returns:
+            The agent's response (string or generator for streaming)
+        """
+        if stream is None:
+            stream = self.config.llm.stream
+        
+        # Build messages
+        messages = self._build_messages(user_query)
+        
+        # Track iterations
+        iterations = 0
+        max_iterations = self.config.agent.max_iterations
+        
+        # Process with ReAct loop
+        while iterations < max_iterations:
+            iterations += 1
+            
+            if self.config.agent.verbose:
+                logger.info("\n" + "=" * 100)
+                logger.info(f"ITERATION {iterations}/{max_iterations}")
+                logger.info("=" * 100)
+            
+            try:
+                # Call LLM with tools
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    tools=self.tools,
+                    tool_choice="auto",
+                    temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
+                    max_tokens=_reasoning_safe_max_tokens(self.model, self.config.llm.max_tokens),
+                    stream=False  # We handle streaming separately
+                )
+                
+                message = response.choices[0].message
+                
+                # Add assistant message to history
+                assistant_msg = {"role": "assistant", "content": message.content or ""}
+                if message.tool_calls:
+                    assistant_msg["tool_calls"] = [
+                        {
+                            "id": tc.id,
+                            "type": tc.type,
+                            "function": {
+                                "name": tc.function.name,
+                                "arguments": tc.function.arguments
+                            }
+                        } for tc in message.tool_calls
+                    ]
+                messages.append(assistant_msg)
+                
+                # Process tool calls if present
+                if message.tool_calls:
+                    for tool_call in message.tool_calls:
+                        tool_name = tool_call.function.name
+                        try:
+                            arguments = json.loads(tool_call.function.arguments)
+                        except json.JSONDecodeError:
+                            arguments = {}
+                        
+                        if self.config.agent.verbose:
+                            logger.info("\n" + "#" * 80)
+                            logger.info(f"TOOL CALL: {tool_name}")
+                            logger.info(f"Arguments: {json.dumps(arguments, indent=2, ensure_ascii=False)}")
+                            logger.info("#" * 80)
+                        
+                        # Execute tool
+                        result = self._execute_tool(tool_name, arguments)
+                        
+                        # Log full tool result when verbose
+                        if self.config.agent.verbose:
+                            logger.info("\n" + "*" * 80)
+                            logger.info("TOOL RESULT:")
+                            logger.info("*" * 80)
+                            # Show full result including all_results if present
+                            if 'all_results' in result:
+                                logger.info("All Search Results (Complete):")
+                                for idx, res in enumerate(result['all_results'], 1):
+                                    logger.info(f"\nResult {idx}:")
+                                    logger.info(json.dumps(res, indent=2, ensure_ascii=False))
+                            else:
+                                logger.info(json.dumps(result, indent=2, ensure_ascii=False))
+                            logger.info("*" * 80 + "\n")
+                        
+                        # For messages, don't include all_results to avoid overloading LLM
+                        result_for_llm = {k: v for k, v in result.items() if k != 'all_results'}
+                        
+                        # Add tool result to messages
+                        tool_message = {
+                            "role": "tool",
+                            "tool_call_id": tool_call.id,
+                            "content": json.dumps(result_for_llm, ensure_ascii=False)
+                        }
+                        messages.append(tool_message)
+                    
+                    # Continue loop for next iteration
+                    continue
+                else:
+                    # No tool calls, we have final answer
+                    # Update conversation history
+                    self.conversation_history.append({"role": "user", "content": user_query})
+                    self.conversation_history.append(assistant_msg)
+                    
+                    # Return response
+                    if stream:
+                        return self._stream_response(message.content or "")
+                    else:
+                        return message.content or ""
+                        
+            except Exception as e:
+                logger.error(f"Error in query processing: {e}")
+                error_msg = f"Error processing query: {str(e)}"
+                if stream:
+                    return self._stream_response(error_msg)
+                else:
+                    return error_msg
+        
+        # Max iterations reached
+        logger.warning(f"Max iterations ({max_iterations}) reached")
+        final_msg = "I need more iterations to fully answer your question. Please try rephrasing or breaking down your query."
+        
+        if stream:
+            return self._stream_response(final_msg)
+        else:
+            return final_msg
+    
+    def _stream_response(self, content: str) -> Generator[str, None, None]:
+        """Stream response content"""
+        # Simple character streaming for demonstration
+        for char in content:
+            yield char
+    
+    def query_non_agentic(self, user_query: str, stream: bool = None) -> Any:
+        """
+        Non-agentic RAG mode: Simple retrieval + LLM response.
+        
+        Args:
+            user_query: The user's question
+            stream: Whether to stream the response
+            
+        Returns:
+            The response (string or generator for streaming)
+        """
+        if stream is None:
+            stream = self.config.llm.stream
+        
+        try:
+            # Simple retrieval
+            search_results = self.kb_tools.knowledge_base_search(user_query)
+            
+            # Build context from search results
+            context_parts = []
+            for i, result in enumerate(search_results[:3], 1):  # Top 3 results
+                context_parts.append(
+                    f"[Document {i}] (ID: {result['doc_id']}, Chunk: {result['chunk_id']})\n{result['text']}\n"
+                )
+            
+            if not context_parts:
+                context = "No relevant information found in the knowledge base."
+            else:
+                context = "\n".join(context_parts)
+            
+            # Build prompt
+            system_prompt = """You are an assistant that answers questions based on provided context from a knowledge base.
+
+IMPORTANT RULES:
+1. Only answer based on the provided context
+2. Include citations in format [Doc: document_id] 
+3. If the context doesn't contain the answer, say so clearly
+4. Be accurate and don't make up information"""
+
+            user_prompt = f"""Context from knowledge base:
+{context}
+
+User Question: {user_query}
+
+Please answer the question based only on the provided context. Include citations."""
+
+            # Call LLM
+            messages = [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_prompt}
+            ]
+            
+            if stream:
+                response_stream = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
+                    max_tokens=_reasoning_safe_max_tokens(self.model, self.config.llm.max_tokens),
+                    stream=True
+                )
+                
+                def response_generator():
+                    for chunk in response_stream:
+                        if chunk.choices[0].delta.content:
+                            yield chunk.choices[0].delta.content
+                
+                return response_generator()
+            else:
+                response = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=messages,
+                    temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
+                    max_tokens=_reasoning_safe_max_tokens(self.model, self.config.llm.max_tokens),
+                    stream=False
+                )
+                return response.choices[0].message.content
+                
+        except Exception as e:
+            logger.error(f"Error in non-agentic query: {e}")
+            error_msg = f"Error processing query: {str(e)}"
+            if stream:
+                return self._stream_response(error_msg)
+            else:
+                return error_msg
+    
+    def clear_history(self):
+        """Clear conversation history"""
+        self.conversation_history = []
+        logger.info("Conversation history cleared")