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

chapter3/agentic-rag-for-user-memory/tools.py

@@ -0,0 +1,320 @@
+"""Tool definitions for the User Memory RAG Agent
+
+This module provides tool definitions and implementations for searching
+and retrieving information from indexed conversation memories.
+"""
+
+import json
+import logging
+from typing import Dict, Any, List, Optional
+from dataclasses import dataclass
+
+from indexer import MemoryIndexer, SearchResult
+from config import IndexConfig
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ToolResult:
+    """Result from a tool execution"""
+    success: bool
+    data: Any
+    error: Optional[str] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        if self.success:
+            return {"status": "success", "data": self.data}
+        else:
+            return {"status": "error", "error": self.error}
+
+
+class MemoryTools:
+    """Tools for searching and retrieving user memory information"""
+    
+    def __init__(self, indexer: MemoryIndexer):
+        """
+        Initialize memory tools
+        
+        Args:
+            indexer: The memory indexer instance
+        """
+        self.indexer = indexer
+        logger.info("Initialized memory tools")
+    
+    def search_memory(self, 
+                     query: str, 
+                     top_k: int = 3,
+                     filter_test_id: Optional[str] = None) -> ToolResult:
+        """
+        Search user memory for relevant information
+        
+        Args:
+            query: Natural language search query
+            top_k: Number of results to return
+            filter_test_id: Optional test ID to filter results
+            
+        Returns:
+            ToolResult with search results
+        """
+        try:
+            # Perform search
+            results = self.indexer.search(query, top_k=top_k)
+            
+            # Filter by test ID if specified
+            if filter_test_id:
+                results = [r for r in results if r.chunk.test_id == filter_test_id]
+            
+            # Format results
+            formatted_results = []
+            for result in results:
+                # Extract key information from the chunk
+                chunk_info = {
+                    "chunk_id": result.chunk_id,
+                    "score": round(result.score, 4),
+                    "test_id": result.chunk.test_id,
+                    "conversation_id": result.chunk.conversation_id,
+                    "rounds": f"{result.chunk.start_round}-{result.chunk.end_round}",
+                    "metadata": result.chunk.metadata,
+                    "content": result.chunk.to_text(),  # FULL content, not truncated
+                    "match_type": result.match_type
+                }
+                formatted_results.append(chunk_info)
+            
+            logger.info(f"Search query: '{query}' returned {len(formatted_results)} results")
+            
+            return ToolResult(
+                success=True,
+                data={
+                    "query": query,
+                    "total_results": len(formatted_results),
+                    "results": formatted_results
+                }
+            )
+            
+        except Exception as e:
+            logger.error(f"Error in search_memory: {e}")
+            return ToolResult(
+                success=False,
+                data=None,
+                error=str(e)
+            )
+    
+    def get_conversation_context(self,
+                                chunk_id: str,
+                                context_size: int = 2) -> ToolResult:
+        """
+        Get surrounding context for a specific chunk
+        
+        Args:
+            chunk_id: The chunk ID to get context for
+            context_size: Number of chunks before/after to include
+            
+        Returns:
+            ToolResult with conversation context
+        """
+        try:
+            # Get the target chunk
+            if chunk_id not in self.indexer.chunks:
+                return ToolResult(
+                    success=False,
+                    data=None,
+                    error=f"Chunk {chunk_id} not found"
+                )
+            
+            target_chunk = self.indexer.chunks[chunk_id]
+            
+            # Find related chunks from same conversation
+            related_chunks = []
+            for cid, chunk in self.indexer.chunks.items():
+                if (chunk.conversation_id == target_chunk.conversation_id and 
+                    chunk.test_id == target_chunk.test_id):
+                    related_chunks.append(chunk)
+            
+            # Sort by chunk index
+            related_chunks.sort(key=lambda x: x.chunk_index)
+            
+            # Find target index
+            target_idx = next(
+                (i for i, c in enumerate(related_chunks) if c.chunk_id == chunk_id),
+                None
+            )
+            
+            if target_idx is None:
+                return ToolResult(
+                    success=False,
+                    data=None,
+                    error="Could not locate chunk in conversation"
+                )
+            
+            # Get context chunks
+            start_idx = max(0, target_idx - context_size)
+            end_idx = min(len(related_chunks), target_idx + context_size + 1)
+            context_chunks = related_chunks[start_idx:end_idx]
+            
+            # Format result
+            context_data = {
+                "target_chunk": {
+                    "chunk_id": target_chunk.chunk_id,
+                    "rounds": f"{target_chunk.start_round}-{target_chunk.end_round}",
+                    "content": target_chunk.to_text()
+                },
+                "context_chunks": []
+            }
+            
+            for chunk in context_chunks:
+                if chunk.chunk_id != chunk_id:
+                    context_data["context_chunks"].append({
+                        "chunk_id": chunk.chunk_id,
+                        "rounds": f"{chunk.start_round}-{chunk.end_round}",
+                        "position": "before" if chunk.chunk_index < target_chunk.chunk_index else "after",
+                        "content": chunk.to_text()
+                    })
+            
+            return ToolResult(
+                success=True,
+                data=context_data
+            )
+            
+        except Exception as e:
+            logger.error(f"Error in get_conversation_context: {e}")
+            return ToolResult(
+                success=False,
+                data=None,
+                error=str(e)
+            )
+    
+    def get_full_conversation(self,
+                            conversation_id: str,
+                            test_id: str) -> ToolResult:
+        """
+        Retrieve all chunks from a specific conversation
+        
+        Args:
+            conversation_id: Conversation identifier
+            test_id: Test case identifier
+            
+        Returns:
+            ToolResult with full conversation
+        """
+        try:
+            # Find all chunks for this conversation
+            conversation_chunks = []
+            for chunk_id, chunk in self.indexer.chunks.items():
+                if (chunk.conversation_id == conversation_id and 
+                    chunk.test_id == test_id):
+                    conversation_chunks.append(chunk)
+            
+            if not conversation_chunks:
+                return ToolResult(
+                    success=False,
+                    data=None,
+                    error=f"No chunks found for conversation {conversation_id}"
+                )
+            
+            # Sort by chunk index
+            conversation_chunks.sort(key=lambda x: x.chunk_index)
+            
+            # Format result
+            conversation_data = {
+                "conversation_id": conversation_id,
+                "test_id": test_id,
+                "total_chunks": len(conversation_chunks),
+                "total_rounds": max(c.end_round for c in conversation_chunks),
+                "metadata": conversation_chunks[0].metadata if conversation_chunks else {},
+                "chunks": []
+            }
+            
+            for chunk in conversation_chunks:
+                conversation_data["chunks"].append({
+                    "chunk_id": chunk.chunk_id,
+                    "chunk_index": chunk.chunk_index,
+                    "rounds": f"{chunk.start_round}-{chunk.end_round}",
+                    "content": chunk.to_text()
+                })
+            
+            return ToolResult(
+                success=True,
+                data=conversation_data
+            )
+            
+        except Exception as e:
+            logger.error(f"Error in get_full_conversation: {e}")
+            return ToolResult(
+                success=False,
+                data=None,
+                error=str(e)
+            )
+    
+
+
+def get_tool_definitions() -> List[Dict[str, Any]]:
+    """
+    Get OpenAI function calling tool definitions
+    
+    Returns:
+        List of tool definitions for OpenAI API
+    """
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": "search_memory",
+                "description": "Search user conversation memory for relevant information. Use this to find specific details from past conversations.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "Natural language search query describing what information to find"
+                        },
+                    },
+                    "required": ["query"]
+                }
+            }
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "get_conversation_context",
+                "description": "Get surrounding context for a specific conversation chunk. Use this when you need more context around a search result.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "chunk_id": {
+                            "type": "string",
+                            "description": "The chunk ID to get context for"
+                        },
+                        "context_size": {
+                            "type": "integer",
+                            "description": "Number of chunks before/after to include (default: 2)",
+                            "default": 2
+                        }
+                    },
+                    "required": ["chunk_id"]
+                }
+            }
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "get_full_conversation",
+                "description": "Retrieve all chunks from a specific conversation. Use this when you need to review an entire conversation history.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "conversation_id": {
+                            "type": "string",
+                            "description": "The conversation identifier"
+                        },
+                        "test_id": {
+                            "type": "string",
+                            "description": "The test case identifier"
+                        }
+                    },
+                    "required": ["conversation_id", "test_id"]
+                }
+            }
+        }
+    ]