"""Agentic RAG Agent for User Memory Evaluation This agent uses RAG-indexed conversation memories to answer questions about user interactions, following the 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 from tools import MemoryTools, get_tool_definitions from indexer import MemoryIndexer 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.""" m = str(model or "").lower().replace("/", "-") return 1 if ("kimi-k3" in m or "gpt-5" in m) else 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()) @dataclass class AgentTrajectory: """Tracks the agent's reasoning and tool usage""" test_id: str question: str iterations: List[Dict[str, Any]] = field(default_factory=list) final_answer: Optional[str] = None tool_calls: List[Dict[str, Any]] = field(default_factory=list) total_time: Optional[float] = None success: bool = False def to_dict(self) -> Dict[str, Any]: return { "test_id": self.test_id, "question": self.question, "iterations": self.iterations, "final_answer": self.final_answer, "tool_calls": self.tool_calls, "total_time": self.total_time, "success": self.success, "total_iterations": len(self.iterations), "total_tool_calls": len(self.tool_calls) } class UserMemoryRAGAgent: """Agent that uses RAG to answer questions about user conversation history""" def __init__(self, indexer: MemoryIndexer, config: Optional[Config] = None): """ Initialize the agent Args: indexer: The memory indexer with loaded conversations config: Configuration object """ self.config = config or Config.from_env() self.indexer = indexer self.memory_tools = MemoryTools(indexer) # Initialize LLM client self._init_llm_client() # Tool definitions self.tools = get_tool_definitions() # Conversation history self.conversation_history: List[Dict[str, Any]] = [] logger.info(f"Initialized UserMemoryRAGAgent 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, test_id: str) -> str: """Generate the system prompt""" return f"""You are an AI assistant with access to indexed conversation memories from user interactions. Your task is to answer questions about these conversations accurately based ONLY on the information you can find in the indexed memories. Current Test Case: {test_id} ## Important Guidelines: 1. **Memory Search Only**: You MUST only answer based on information found through the memory search tools. If the information is not available in the indexed conversations, clearly state that you cannot find it. 2. **Use Tools Effectively**: - Use `search_memory` to find relevant information across all conversations - Use `get_conversation_context` when you need more context around a search result - Use `get_full_conversation` to review entire conversation histories when needed 3. **Multiple Searches**: Don't hesitate to perform multiple searches with different queries to find all relevant information. Different phrasings may yield different results. 4. **Citations Required**: Always mention which conversation or chunk you found the information in when providing answers. 5. **Be Thorough**: For complex questions, gather information from multiple chunks and conversations before formulating your answer. 6. **Handle Ambiguity**: If you find conflicting information or multiple possible answers, report all of them with their sources. Remember: Your credibility depends on providing accurate, well-sourced information from the conversation memories only.""" def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any: """Execute a tool and return the result""" try: # Log tool call parameters to console logger.info("="*80) logger.info(f"TOOL CALL: {tool_name}") logger.info(f"PARAMETERS: {json.dumps(arguments, indent=2, ensure_ascii=False)}") logger.info("-"*80) if tool_name == "search_memory": query = arguments.get("query", "") filter_test_id = arguments.get("filter_test_id") result = self.memory_tools.search_memory( query, top_k=self.config.agent.max_search_results, filter_test_id=filter_test_id, ) # Log result to console result_dict = result.to_dict() logger.info("TOOL RESULT:") logger.info(json.dumps(result_dict, indent=2, ensure_ascii=False)) logger.info("="*80) return result_dict elif tool_name == "get_conversation_context": chunk_id = arguments.get("chunk_id", "") context_size = arguments.get("context_size", 2) result = self.memory_tools.get_conversation_context(chunk_id, context_size) # Log result to console result_dict = result.to_dict() logger.info("TOOL RESULT:") logger.info(json.dumps(result_dict, indent=2, ensure_ascii=False)) logger.info("="*80) return result_dict elif tool_name == "get_full_conversation": conversation_id = arguments.get("conversation_id", "") test_id = arguments.get("test_id", "") result = self.memory_tools.get_full_conversation(conversation_id, test_id) # Log result to console result_dict = result.to_dict() logger.info("TOOL RESULT:") logger.info(json.dumps(result_dict, indent=2, ensure_ascii=False)) logger.info("="*80) return result_dict else: return {"status": "error", "error": f"Unknown tool: {tool_name}"} except Exception as e: logger.error(f"Tool execution error: {e}") return {"status": "error", "error": str(e)} def answer_question(self, question: str, test_id: str, stream: bool = False) -> Dict[str, Any]: """ Answer a question about user conversation history using RAG Args: question: The question to answer test_id: The test case ID for context stream: Whether to stream the response Returns: Dictionary containing the answer and trajectory """ start_time = datetime.now() trajectory = AgentTrajectory(test_id=test_id, question=question) # Build initial messages messages = [ {"role": "system", "content": self._get_system_prompt(test_id)}, {"role": "user", "content": question} ] # Track iterations iterations = 0 max_iterations = self.config.evaluation.max_iterations # Process with ReAct loop while iterations < max_iterations: iterations += 1 iteration_data = {"iteration": iterations, "timestamp": datetime.now().isoformat()} if self.config.agent.enable_reasoning: logger.info(f"\n{'='*60}") logger.info(f"Iteration {iterations}/{max_iterations}") logger.info(f"{'='*60}") 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=self.config.llm.max_tokens, stream=False ) message = response.choices[0].message iteration_data["assistant_message"] = message.content or "" # Log the LLM response content if message.content: logger.info("-"*60) logger.info(f"LLM Response: {message.content}") logger.info("-"*60) # 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 ] iteration_data["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 = {} # Execute tool result = self._execute_tool(tool_name, arguments) # Track tool call tool_call_data = { "tool": tool_name, "arguments": arguments, "result": result, "timestamp": datetime.now().isoformat() } iteration_data["tool_calls"].append(tool_call_data) trajectory.tool_calls.append(tool_call_data) # Add tool result to messages tool_message = { "role": "tool", "tool_call_id": tool_call.id, "content": json.dumps(result, ensure_ascii=False) } messages.append(tool_message) # Continue loop for next iteration trajectory.iterations.append(iteration_data) continue else: # No tool calls, we have final answer trajectory.iterations.append(iteration_data) trajectory.final_answer = message.content or "" trajectory.success = True # Calculate total time end_time = datetime.now() trajectory.total_time = (end_time - start_time).total_seconds() # Return result result = { "answer": trajectory.final_answer, "success": True, "iterations": iterations, "tool_calls": len(trajectory.tool_calls), "trajectory": trajectory.to_dict() if self.config.evaluation.save_trajectories else None } if stream: return self._stream_response(result) else: return result except Exception as e: logger.error(f"Error in iteration {iterations}: {e}") trajectory.iterations.append({ "iteration": iterations, "error": str(e) }) # Continue to next iteration continue # Max iterations reached logger.warning(f"Max iterations ({max_iterations}) reached") # Calculate total time end_time = datetime.now() trajectory.total_time = (end_time - start_time).total_seconds() trajectory.success = False final_msg = "I was unable to find sufficient information to answer your question within the iteration limit. Please try rephrasing or breaking down your query." return { "answer": final_msg, "success": False, "iterations": iterations, "tool_calls": len(trajectory.tool_calls), "trajectory": trajectory.to_dict() if self.config.evaluation.save_trajectories else None } def _stream_response(self, result: Dict[str, Any]) -> Generator[str, None, None]: """Stream response content""" # Stream the answer character by character answer = result.get("answer", "") for char in answer: yield char def clear_history(self): """Clear conversation history""" self.conversation_history = [] logger.info("Conversation history cleared") def save_trajectory(self, trajectory: AgentTrajectory, filepath: str): """ Save agent trajectory to file Args: trajectory: The trajectory to save filepath: Path to save the trajectory """ with open(filepath, 'w', encoding='utf-8') as f: json.dump(trajectory.to_dict(), f, ensure_ascii=False, indent=2) logger.info(f"Trajectory saved to {filepath}")