Sauten commited on
Commit
331db0c
·
verified ·
1 Parent(s): 27a2031

Create agent.py

Browse files
Files changed (1) hide show
  1. agent.py +56 -0
agent.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # from google import genai
2
+ # from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
3
+ import os
4
+ from dotenv import load_dotenv
5
+ from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool, VisitWebpageTool, FinalAnswerTool
6
+ from tools import visitWikipedia, visitWikipediaTable
7
+ import time
8
+
9
+ load_dotenv()
10
+ google = os.getenv("GOOGLE")
11
+
12
+ # client = genai.Client(api_key=google)
13
+ # google_search_tool = Tool(
14
+ # google_search = GoogleSearch()
15
+ # )
16
+ # response = client.models.generate_content(
17
+ # model="gemini-2.0-flash",
18
+ # contents = prompt,
19
+ # config=GenerateContentConfig(
20
+ # tools=[google_search_tool],
21
+ # response_modalities=["TEXT"],
22
+ # )
23
+ # )
24
+ # print(response.text)
25
+
26
+ model = LiteLLMModel(
27
+ model_id="gemini/gemini-2.0-flash",
28
+ api_key=google,
29
+ max_tokens=1000,
30
+ temperature=0.0,
31
+ )
32
+ tools = [FinalAnswerTool(), DuckDuckGoSearchTool(), VisitWebpageTool(), visitWikipedia, visitWikipediaTable]
33
+ agent = CodeAgent(
34
+ tools=tools,
35
+ model=model,
36
+ additional_authorized_imports=["math", "datetime", "re", "json", "pandas", "numpy"],
37
+ max_steps=15,
38
+ )
39
+
40
+ prompt = """
41
+ Do not assume your code is correct. Always check your code and the output of your code.
42
+ You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the FinalAnswerTool(). YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
43
+ Question:
44
+ """
45
+
46
+ def run(question):
47
+ try:
48
+ return agent.run(task=prompt+"\n"+question)
49
+ except Exception as e:
50
+ time.sleep(60)
51
+ return agent.run(task=prompt+"\n"+question)
52
+
53
+ question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
54
+
55
+ # answer = run(question)
56
+ # print(answer)