agents.py

"""Agent implementation with tool calling capabilities"""


class Agent:
    """
    Basic agent that uses an LLM and tools to answer questions.
    """
    
    def __init__(self, model, tools):
        self.model = model
        self.tools = {tool.name: tool for tool in tools}
        self.max_iterations = 10
    
    def __call__(self, question: str) -> str:
        """
        Run the agent on a question.
        Returns the final answer as a string.
        """
        print(f"\n{'='*60}")
        print(f"QUESTION: {question[:100]}...")
        print(f"{'='*60}")
        
        messages = [
            {
                "role": "system",
                "content": """You are a helpful AI assistant that answers questions accurately.

When you need information:
- Use web_search to find current information online
- Use wikipedia_search for encyclopedic knowledge
- Use calculate for mathematical operations

Always provide a direct, concise answer. If you don't know, search for the information.

For questions about files, images, or attachments you cannot access, respond with "File not found".

Format your final answer clearly and concisely."""
            },
            {
                "role": "user",
                "content": question
            }
        ]
        
        for iteration in range(self.max_iterations):
            # Get response from model
            response = self.model.generate(messages, tools=list(self.tools.values()))
            
            # Check if model wants to use a tool
            if response.get("tool_calls"):
                print(f"\nIteration {iteration + 1}: Tool calls requested")
                
                # Create assistant message with tool calls
                assistant_msg = {
                    "role": "assistant",
                    "content": response.get("content") or "Using tools...",
                }
                
                # Add tool_calls if the model supports it
                if "groq" in str(self.model.model_id).lower():
                    # For Groq, we need to include tool_calls in the message
                    assistant_msg["tool_calls"] = [
                        {
                            "id": tc.get("id", f"call_{tc['name']}"),
                            "type": "function",
                            "function": {
                                "name": tc["name"],
                                "arguments": str(tc["arguments"])
                            }
                        }
                        for tc in response["tool_calls"]
                    ]
                
                messages.append(assistant_msg)
                
                for tool_call in response["tool_calls"]:
                    tool_name = tool_call["name"]
                    tool_args = tool_call["arguments"]
                    tool_id = tool_call.get("id", f"call_{tool_name}")
                    
                    print(f"  → {tool_name}({tool_args})")
                    
                    if tool_name in self.tools:
                        try:
                            result = self.tools[tool_name](**tool_args)
                            print(f"  ← Result: {str(result)[:100]}...")
                            
                            # Add tool result to messages
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_id,
                                "name": tool_name,
                                "content": str(result)
                            })
                        except Exception as e:
                            print(f"  ← Error: {e}")
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_id,
                                "name": tool_name,
                                "content": f"Error: {str(e)}"
                            })
                    else:
                        print(f"  ← Tool not found!")
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_id,
                            "name": tool_name,
                            "content": "Tool not found"
                        })
            else:
                # Model provided final answer
                answer = response.get("content", "Unknown")
                print(f"\n✓ ANSWER: {answer}")
                print(f"{'='*60}\n")
                return answer.strip()
        
        print(f"\n⚠ Max iterations reached")
        return "Unknown"