test_agent.py

#!/usr/bin/env python3

"""
Simple test script to verify that our agents work correctly with final_answer()
"""

import os
import pathlib

import dotenv
from smolagents import CodeAgent, OpenAIServerModel

# Load environment
if pathlib.Path(".env").exists():
    dotenv.load_dotenv(".env")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", None)

if not OPENAI_API_KEY:
    print("Please set OPENAI_API_KEY in your .env file")
    exit(1)

# Create model
model = OpenAIServerModel(
    model_id="gpt-4o-mini",  # Using a more cost-effective model for testing
    api_key=OPENAI_API_KEY,
)

def test_simple_calculation():
    """Test that calc agent can solve a simple math problem and call final_answer()"""
    from prompt import calc_agent_prompt
    
    calc_agent = CodeAgent(
        model=model,
        prompt_templates=calc_agent_prompt,
        additional_authorized_imports=["pandas", "numpy"],
        tools=[]  # No special tools needed for basic calculations
    )
    
    question = "What is the result of 5 + 3 + 1294.678?"
    print(f"\n=== Testing Calculation Agent ===")
    print(f"Question: {question}")
    print("Running agent...")
    
    try:
        result = calc_agent.run(question)
        print(f"Result: {result}")
        print("✅ Calculation agent test passed!")
        return True
    except Exception as e:
        print(f"❌ Calculation agent test failed: {e}")
        return False

def test_simple_search():
    """Test that web search agent can answer a simple question and call final_answer()"""
    from smolagents import VisitWebpageTool, WebSearchTool, WikipediaSearchTool

    from prompt import web_search_agent_prompt
    
    web_agent = CodeAgent(
        model=model,
        prompt_templates=web_search_agent_prompt,
        additional_authorized_imports=["requests", "beautifulsoup4"],
        tools=[
            WebSearchTool(),
            VisitWebpageTool(),
            WikipediaSearchTool(),
        ],
    )
    
    question = "Who is the current president of France?"
    print(f"\n=== Testing Web Search Agent ===")
    print(f"Question: {question}")
    print("Running agent...")
    
    try:
        result = web_agent.run(question)
        print(f"Result: {result}")
        print("✅ Web search agent test passed!")
        return True
    except Exception as e:
        print(f"❌ Web search agent test failed: {e}")
        return False

if __name__ == "__main__":
    print("Testing our updated agents with final_answer() support...")
    
    calc_success = test_simple_calculation()
    search_success = test_simple_search()
    
    if calc_success and search_success:
        print("\n🎉 All tests passed! Agents are working correctly with final_answer()")
    else:
        print("\n⚠️  Some tests failed. Check the error messages above.")