wlchee's picture
Upload agent.py
96f906d verified
raw
history blame
3.3 kB
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}")