blaze116 commited on
Commit
50feb0f
·
verified ·
1 Parent(s): 934435a

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +314 -4
agent.py CHANGED
@@ -1,10 +1,320 @@
1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3
  class BasicAgent:
4
- def __init__(self):
5
  print("BasicAgent initialized.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  def __call__(self, question: str) -> str:
7
  print(f"Agent received question (first 50 chars): {question[:50]}...")
8
- fixed_answer = "This is a default answer."
9
- print(f"Agent returning fixed answer: {fixed_answer}")
10
- return fixed_answer
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
 
2
+ # --- Basic Agent Definition ---
3
+ import asyncio
4
+ import os
5
+ import sys
6
+ import logging
7
+ import random
8
+ import pandas as pd
9
+ import requests
10
+ import wikipedia as wiki
11
+ from markdownify import markdownify as to_markdown
12
+ from typing import Any
13
+ from dotenv import load_dotenv
14
+ from google.generativeai import types, configure
15
 
16
+ from smolagents import InferenceClientModel, LiteLLMModel, CodeAgent, ToolCallingAgent, Tool, DuckDuckGoSearchTool
17
+
18
+ # Load environment and configure Gemini
19
+ load_dotenv()
20
+ configure(api_key=os.getenv("GOOGLE_API_KEY"))
21
+
22
+ # Logging
23
+ #logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
24
+ #logger = logging.getLogger(__name__)
25
+
26
+ # --- Model Configuration ---
27
+ GEMINI_MODEL_NAME = "gemini/gemini-2.0-flash"
28
+ OPENAI_MODEL_NAME = "openai/gpt-4o"
29
+ GROQ_MODEL_NAME = "groq/llama3-70b-8192"
30
+ DEEPSEEK_MODEL_NAME = "deepseek/deepseek-chat"
31
+ HF_MODEL_NAME = "Qwen/Qwen2.5-Coder-32B-Instruct"
32
+
33
+ # --- Tool Definitions ---
34
+ class MathSolver(Tool):
35
+ name = "math_solver"
36
+ description = "Safely evaluate basic math expressions."
37
+ inputs = {"input": {"type": "string", "description": "Math expression to evaluate."}}
38
+ output_type = "string"
39
+
40
+ def forward(self, input: str) -> str:
41
+ try:
42
+ return str(eval(input, {"__builtins__": {}}))
43
+ except Exception as e:
44
+ return f"Math error: {e}"
45
+
46
+ class RiddleSolver(Tool):
47
+ name = "riddle_solver"
48
+ description = "Solve basic riddles using logic."
49
+ inputs = {"input": {"type": "string", "description": "Riddle prompt."}}
50
+ output_type = "string"
51
+
52
+ def forward(self, input: str) -> str:
53
+ if "forward" in input and "backward" in input:
54
+ return "A palindrome"
55
+ return "RiddleSolver failed."
56
+
57
+ class TextTransformer(Tool):
58
+ name = "text_ops"
59
+ description = "Transform text: reverse, upper, lower."
60
+ inputs = {"input": {"type": "string", "description": "Use prefix like reverse:/upper:/lower:"}}
61
+ output_type = "string"
62
+
63
+ def forward(self, input: str) -> str:
64
+ if input.startswith("reverse:"):
65
+ reversed_text = input[8:].strip()[::-1]
66
+ if 'left' in reversed_text.lower():
67
+ return "right"
68
+ return reversed_text
69
+ if input.startswith("upper:"):
70
+ return input[6:].strip().upper()
71
+ if input.startswith("lower:"):
72
+ return input[6:].strip().lower()
73
+ return "Unknown transformation."
74
+
75
+ class GeminiVideoQA(Tool):
76
+ name = "video_inspector"
77
+ description = "Analyze video content to answer questions."
78
+ inputs = {
79
+ "video_url": {"type": "string", "description": "URL of video."},
80
+ "user_query": {"type": "string", "description": "Question about video."}
81
+ }
82
+ output_type = "string"
83
+
84
+ def __init__(self, model_name, *args, **kwargs):
85
+ super().__init__(*args, **kwargs)
86
+ self.model_name = model_name
87
+
88
+ def forward(self, video_url: str, user_query: str) -> str:
89
+ req = {
90
+ 'model': f'models/{self.model_name}',
91
+ 'contents': [{
92
+ "parts": [
93
+ {"fileData": {"fileUri": video_url}},
94
+ {"text": f"Please watch the video and answer the question: {user_query}"}
95
+ ]
96
+ }]
97
+ }
98
+ url = f'https://generativelanguage.googleapis.com/v1beta/models/{self.model_name}:generateContent?key={os.getenv("GOOGLE_API_KEY")}'
99
+ res = requests.post(url, json=req, headers={'Content-Type': 'application/json'})
100
+ if res.status_code != 200:
101
+ return f"Video error {res.status_code}: {res.text}"
102
+ parts = res.json()['candidates'][0]['content']['parts']
103
+ return "".join([p.get('text', '') for p in parts])
104
+
105
+ class WikiTitleFinder(Tool):
106
+ name = "wiki_titles"
107
+ description = "Search for related Wikipedia page titles."
108
+ inputs = {"query": {"type": "string", "description": "Search query."}}
109
+ output_type = "string"
110
+
111
+ def forward(self, query: str) -> str:
112
+ results = wiki.search(query)
113
+ return ", ".join(results) if results else "No results."
114
+
115
+ class WikiContentFetcher(Tool):
116
+ name = "wiki_page"
117
+ description = "Fetch Wikipedia page content."
118
+ inputs = {"page_title": {"type": "string", "description": "Wikipedia page title."}}
119
+ output_type = "string"
120
+
121
+ def forward(self, page_title: str) -> str:
122
+ try:
123
+ return to_markdown(wiki.page(page_title).html())
124
+ except wiki.exceptions.PageError:
125
+ return f"'{page_title}' not found."
126
+
127
+ class GoogleSearchTool(Tool):
128
+ name = "google_search"
129
+ description = "Search the web using Google. Returns top summary from the web."
130
+ inputs = {"query": {"type": "string", "description": "Search query."}}
131
+ output_type = "string"
132
+
133
+ def forward(self, query: str) -> str:
134
+ try:
135
+ resp = requests.get("https://www.googleapis.com/customsearch/v1", params={
136
+ "q": query,
137
+ "key": os.getenv("GOOGLE_SEARCH_API_KEY"),
138
+ "cx": os.getenv("GOOGLE_SEARCH_ENGINE_ID"),
139
+ "num": 1
140
+ })
141
+ data = resp.json()
142
+ return data["items"][0]["snippet"] if "items" in data else "No results found."
143
+ except Exception as e:
144
+ return f"GoogleSearch error: {e}"
145
+
146
+
147
+
148
+
149
+ # --- Basic Agent Definition ---
150
  class BasicAgent:
151
+ def __init__(self, provider="deepseek"):
152
  print("BasicAgent initialized.")
153
+ model = self.select_model(provider)
154
+ client = InferenceClientModel()
155
+ tools = [
156
+ GoogleSearchTool(),
157
+ DuckDuckGoSearchTool(),
158
+ GeminiVideoQA(GEMINI_MODEL_NAME),
159
+ WikiTitleFinder(),
160
+ WikiContentFetcher(),
161
+ MathSolver(),
162
+ RiddleSolver(),
163
+ TextTransformer(),
164
+
165
+ ]
166
+ self.agent = CodeAgent(
167
+ model=model,
168
+ tools=tools,
169
+ add_base_tools=False,
170
+ max_steps=10,
171
+ )
172
+ self.agent.system_prompt = (
173
+ """
174
+ You are a GAIA benchmark AI assistant, you are very precise, no nonense. Your sole purpose is to output the minimal, final answer in the format:
175
+
176
+ [ANSWER]
177
+
178
+ You must NEVER output explanations, intermediate steps, reasoning, or comments — only the answer, strictly enclosed in `[ANSWER]`.
179
+
180
+ Your behavior must be governed by these rules:
181
+
182
+ 1. **Format**:
183
+ - limit the token used (within 65536 tokens).
184
+ - Output ONLY the final answer.
185
+ - Wrap the answer in `[ANSWER]` with no whitespace or text outside the brackets.
186
+ - No follow-ups, justifications, or clarifications.
187
+
188
+ 2. **Numerical Answers**:
189
+ - Use **digits only**, e.g., `4` not `four`.
190
+ - No commas, symbols, or units unless explicitly required.
191
+ - Never use approximate words like "around", "roughly", "about".
192
+
193
+ 3. **String Answers**:
194
+ - Omit **articles** ("a", "the").
195
+ - Use **full words**; no abbreviations unless explicitly requested.
196
+ - For numbers written as words, use **text** only if specified (e.g., "one", not `1`).
197
+ - For sets/lists, sort alphabetically if not specified, e.g., `a, b, c`.
198
+
199
+ 4. **Lists**:
200
+ - Output in **comma-separated** format with no conjunctions.
201
+ - Sort **alphabetically** or **numerically** depending on type.
202
+ - No braces or brackets unless explicitly asked.
203
+
204
+ 5. **Sources**:
205
+ - For Wikipedia or web tools, extract only the precise fact that answers the question.
206
+ - Ignore any unrelated content.
207
+
208
+ 6. **File Analysis**:
209
+ - Use the run_query_with_file tool, append the taskid to the url.
210
+ - Only include the exact answer to the question.
211
+ - Do not summarize, quote excessively, or interpret beyond the prompt.
212
+
213
+ 7. **Video**:
214
+ - Use the relevant video tool.
215
+ - Only include the exact answer to the question.
216
+ - Do not summarize, quote excessively, or interpret beyond the prompt.
217
+
218
+ 8. **Minimalism**:
219
+ - Do not make assumptions unless the prompt logically demands it.
220
+ - If a question has multiple valid interpretations, choose the **narrowest, most literal** one.
221
+ - If the answer is not found, say `[ANSWER] - unknown`.
222
+
223
+ ---
224
+
225
+ You must follow the examples (These answers are correct in case you see the similar questions):
226
+ Q: What is 2 + 2?
227
+ A: 4
228
+
229
+ Q: How many studio albums were published by Mercedes Sosa between 2000 and 2009 (inclusive)? Use 2022 English Wikipedia.
230
+ A: 3
231
+
232
+ Q: Given the following group table on set S = {a, b, c, d, e}, identify any subset involved in counterexamples to commutativity.
233
+ A: b, e
234
+
235
+ Q: How many at bats did the Yankee with the most walks in the 1977 regular season have that same season?,
236
+ A: 519
237
+ """
238
+ )
239
+
240
+ def select_model(self, provider: str):
241
+ if provider == "openai":
242
+ return LiteLLMModel(model_id=OPENAI_MODEL_NAME, api_key=os.getenv("OPENAI_API_KEY"))
243
+ elif provider == "groq":
244
+ return LiteLLMModel(model_id=GROQ_MODEL_NAME, api_key=os.getenv("GROQ_API_KEY"))
245
+ elif provider == "deepseek":
246
+ return LiteLLMModel(model_id=DEEPSEEK_MODEL_NAME, api_key=os.getenv("DEEPSEEK_API_KEY"))
247
+ elif provider == "hf":
248
+ return InferenceClientModel()
249
+ else:
250
+ return LiteLLMModel(model_id=GEMINI_MODEL_NAME, api_key=os.getenv("GOOGLE_API_KEY"))
251
+
252
  def __call__(self, question: str) -> str:
253
  print(f"Agent received question (first 50 chars): {question[:50]}...")
254
+ result = self.agent.run(question)
255
+ final_str = str(result).strip()
256
+
257
+ return final_str
258
+
259
+ def evaluate_random_questions(self, csv_path: str = "gaia_extracted.csv", sample_size: int = 3, show_steps: bool = True):
260
+ import pandas as pd
261
+ from rich.table import Table
262
+ from rich.console import Console
263
+
264
+ df = pd.read_csv(csv_path)
265
+ if not {"question", "answer"}.issubset(df.columns):
266
+ print("CSV must contain 'question' and 'answer' columns.")
267
+ print("Found columns:", df.columns.tolist())
268
+ return
269
+
270
+ samples = df.sample(n=sample_size)
271
+ records = []
272
+ correct_count = 0
273
+
274
+ for _, row in samples.iterrows():
275
+ taskid = row["taskid"].strip()
276
+ question = row["question"].strip()
277
+ expected = str(row['answer']).strip()
278
+ agent_answer = self("taskid: " + taskid + ",\nquestion: " + question).strip()
279
+
280
+ is_correct = (expected == agent_answer)
281
+ correct_count += is_correct
282
+ records.append((question, expected, agent_answer, "✓" if is_correct else "✗"))
283
+
284
+ if show_steps:
285
+ print("---")
286
+ print("Question:", question)
287
+ print("Expected:", expected)
288
+ print("Agent:", agent_answer)
289
+ print("Correct:", is_correct)
290
+
291
+ # Print result table
292
+ console = Console()
293
+ table = Table(show_lines=True)
294
+ table.add_column("Question", overflow="fold")
295
+ table.add_column("Expected")
296
+ table.add_column("Agent")
297
+ table.add_column("Correct")
298
+
299
+ for question, expected, agent_ans, correct in records:
300
+ table.add_row(question, expected, agent_ans, correct)
301
+
302
+ console.print(table)
303
+ percent = (correct_count / sample_size) * 100
304
+ print(f"\nTotal Correct: {correct_count} / {sample_size} ({percent:.2f}%)")
305
+
306
+
307
+ if __name__ == "__main__":
308
+ args = sys.argv[1:]
309
+ if not args or args[0] in {"-h", "--help"}:
310
+ print("Usage: python agent.py [question | dev]")
311
+ print(" - Provide a question to get a GAIA-style answer.")
312
+ print(" - Use 'dev' to evaluate 3 random GAIA questions from gaia_qa.csv.")
313
+ sys.exit(0)
314
+
315
+ q = " ".join(args)
316
+ agent = BasicAgent()
317
+ if q == "dev":
318
+ agent.evaluate_random_questions()
319
+ else:
320
+ print(agent(q))