jcleee commited on
Commit
8472b68
·
verified ·
1 Parent(s): 352adda

Create agent.py

Browse files
Files changed (1) hide show
  1. agent.py +93 -0
agent.py ADDED
@@ -0,0 +1,93 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # agent.py
2
+ import os
3
+ import datetime
4
+ import requests
5
+ import pytz
6
+ import yaml
7
+ from typing import List
8
+
9
+ from langchain.text_splitter import CharacterTextSplitter
10
+ from tools.final_answer import FinalAnswerTool
11
+ from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
12
+
13
+ # === TOOLS ===
14
+
15
+ @tool
16
+ def web_search(query: str) -> str:
17
+ search_tool = DuckDuckGoSearchTool()
18
+ results = search_tool(query)
19
+ return "\n".join(results)
20
+
21
+ @tool
22
+ def get_current_time_in_timezone(timezone: str) -> str:
23
+ try:
24
+ tz = pytz.timezone(timezone)
25
+ local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
26
+ return f"The current local time in {timezone} is: {local_time}"
27
+ except Exception as e:
28
+ return f"Error fetching time for timezone '{timezone}': {str(e)}"
29
+
30
+ @tool
31
+ def visit_webpage(url: str) -> str:
32
+ try:
33
+ response = requests.get(url, timeout=5)
34
+ return response.text[:5000]
35
+ except Exception as e:
36
+ return f"[ERROR fetching {url}]: {str(e)}"
37
+
38
+ @tool
39
+ def text_splitter(text: str) -> List[str]:
40
+ splitter = CharacterTextSplitter(chunk_size=450, chunk_overlap=10)
41
+ return splitter.split_text(text)
42
+
43
+ # === FINAL ANSWER TOOL ===
44
+ final_answer = FinalAnswerTool()
45
+
46
+ # === LOAD PROMPT TEMPLATES ===
47
+ with open("prompts.yaml", "r") as stream:
48
+ prompt_templates = yaml.safe_load(stream)
49
+
50
+ # === LOAD agent.json CONFIG ===
51
+ with open("agent.json", "r") as f:
52
+ agent_config = yaml.safe_load(f)
53
+
54
+ model_config = agent_config["model"]["data"]
55
+
56
+ # === BUILD MODEL ===
57
+ model = LiteLLMModel(
58
+ model_id="gemini/gemini-2.0-flash-lite",
59
+ api_key=os.getenv("GEMINI_API_KEY"),
60
+ temperature=0.5,
61
+ max_tokens=1024,
62
+ )
63
+
64
+ # === IMPORT TOOL FROM HUB ===
65
+ image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
66
+
67
+ # === BUILD AGENT ===
68
+ agent = CodeAgent(
69
+ model=model,
70
+ tools=[
71
+ final_answer,
72
+ web_search,
73
+ get_current_time_in_timezone,
74
+ visit_webpage,
75
+ text_splitter,
76
+ image_generation_tool
77
+ ],
78
+ max_steps=6,
79
+ verbosity_level=1,
80
+ grammar=None,
81
+ planning_interval=None,
82
+ name=None,
83
+ description=None,
84
+ prompt_templates=prompt_templates
85
+ )
86
+
87
+ # === Export for import ===
88
+ def run_agent(question):
89
+ try:
90
+ result = agent(question)
91
+ return [str(result)]
92
+ except Exception as e:
93
+ return [f"Error: {e}"]