from transformers import Tool, HfAgent
from huggingface_hub import list_models
import requests
from typing import Optional, List
import random

# First, let's define some custom tools the agent can use

class WebSearchTool(Tool):
    name = "web_search"
    description = ("A tool that performs a web search using a search engine API. "
                   "Input should be a search query. Output will be search results.")
    
    inputs = ["text"]
    outputs = ["text"]
    
    def __call__(self, query: str):
        # In a real implementation, you would call a search API here
        # For demonstration, we'll return mock results
        return f"Search results for '{query}': 1. Relevant result 1, 2. Relevant result 2"

class CalculatorTool(Tool):
    name = "calculator"
    description = ("A tool for performing mathematical calculations. "
                   "Input should be a mathematical expression. Output will be the result.")
    
    inputs = ["text"]
    outputs = ["text"]
    
    def __call__(self, expression: str):
        try:
            result = eval(expression)  # Note: In production, use a safer eval method
            return str(result)
        except:
            return "Error: Could not evaluate the expression"

class CurrentTimeTool(Tool):
    name = "get_current_time"
    description = ("A tool that returns the current time in UTC. "
                   "No input needed. Output will be the current time.")
    
    inputs = []
    outputs = ["text"]
    
    def __call__(self):
        from datetime import datetime
        return datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")

class WikipediaTool(Tool):
    name = "wikipedia_search"
    description = ("A tool that searches Wikipedia. "
                   "Input should be a search term. Output will be a summary from Wikipedia.")
    
    inputs = ["text"]
    outputs = ["text"]
    
    def __call__(self, term: str):
        # Mock implementation - in real use, you'd call the Wikipedia API
        return f"Wikipedia summary for '{term}': This is a summary about {term}."

# Now let's create the agent with these tools and some pre-trained tools

def create_agent():
    # Load pre-trained tools from the Hub
    agent = HfAgent(
        "https://api-inference.huggingface.co/models/bigcode/starcoder",
        additional_tools=[
            WebSearchTool(),
            CalculatorTool(),
            CurrentTimeTool(),
            WikipediaTool()
        ],
        # These parameters help with performance
        max_new_tokens=200,
        temperature=0.7,
        top_p=0.9,
    )
    return agent

# Example usage of the agent
if __name__ == "__main__":
    agent = create_agent()
    
    # Test the agent with some sample queries
    queries = [
        "What's the capital of France?",
        "Calculate 123 * 45",
        "What time is it now?",
        "Tell me about Albert Einstein",
        "Search for the latest news about AI"
    ]
    
    for query in queries:
        print(f"Query: {query}")
        result = agent.run(query)
        print(f"Result: {result}\n")
    
    # To evaluate on the benchmark, you would use:
    # from transformers.benchmarks import evaluate_agent
    # benchmark_score = evaluate_agent(agent)
    # print(f"Benchmark score: {benchmark_score}")