AlessandroMasala commited on
Commit
043aa11
·
verified ·
1 Parent(s): 806af71

Delete agent.py

Browse files
Files changed (1) hide show
  1. agent.py +0 -700
agent.py DELETED
@@ -1,700 +0,0 @@
1
- import os
2
- import re
3
- from pathlib import Path
4
- from typing import Optional, Union, Dict, List, Any
5
- from enum import Enum
6
- import requests
7
- import tempfile
8
- import ast
9
-
10
- from dotenv import load_dotenv
11
- from langgraph.graph import StateGraph, END
12
- from langchain.tools import Tool as LangTool
13
- from langchain_core.runnables import RunnableLambda
14
-
15
- from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
16
- from langchain_community.llms import HuggingFacePipeline
17
- import torch
18
-
19
- from langchain.tools import StructuredTool
20
-
21
- from tools import (
22
- EnhancedSearchTool,
23
- EnhancedWikipediaTool,
24
- excel_to_markdown,
25
- image_file_info,
26
- audio_file_info,
27
- code_file_read,
28
- extract_youtube_info)
29
-
30
- def get_zephyr_llm():
31
- model_id = "HuggingFaceH4/zephyr-7b-alpha"
32
- tokenizer = AutoTokenizer.from_pretrained(model_id)
33
- model = AutoModelForCausalLM.from_pretrained(
34
- model_id, torch_dtype=torch.float16, device_map="auto"
35
- )
36
- gen = pipeline(
37
- "text-generation", model=model, tokenizer=tokenizer,
38
- max_new_tokens=256, temperature=0.7, top_p=0.9
39
- )
40
- return HuggingFacePipeline(pipeline=gen)
41
-
42
- # LLM Instance
43
- llm = get_zephyr_llm()
44
-
45
- # Load environment variables
46
- load_dotenv()
47
-
48
- # --- Constants ---
49
- DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
50
- QUESTIONS_URL = f"{DEFAULT_API_URL}/questions"
51
- SUBMIT_URL = f"{DEFAULT_API_URL}/submit"
52
- FILE_PATH = f"{DEFAULT_API_URL}/files/"
53
-
54
-
55
- # ----------- Enhanced State Management -----------
56
- from typing import TypedDict
57
-
58
- class AgentState(TypedDict):
59
- """Enhanced state tracking for the agent - using TypedDict for LangGraph compatibility"""
60
- question: str
61
- original_question: str
62
- conversation_history: List[Dict[str, str]]
63
- selected_tools: List[str]
64
- tool_results: Dict[str, Any]
65
- final_answer: str
66
- current_step: str
67
- error_count: int
68
- max_errors: int
69
-
70
- class AgentStep(Enum):
71
- ANALYZE_QUESTION = "analyze_question"
72
- SELECT_TOOLS = "select_tools"
73
- EXECUTE_TOOLS = "execute_tools"
74
- SYNTHESIZE_ANSWER = "synthesize_answer"
75
- ERROR_RECOVERY = "error_recovery"
76
- COMPLETE = "complete"
77
-
78
- # ----------- Helper Functions -----------
79
- def initialize_state(question: str) -> AgentState:
80
- """Initialize agent state with default values"""
81
- return {
82
- "question": question,
83
- "original_question": question,
84
- "conversation_history": [],
85
- "selected_tools": [],
86
- "tool_results": {},
87
- "final_answer": "",
88
- "current_step": "start",
89
- "error_count": 0,
90
- "max_errors": 3
91
- }
92
-
93
- # Initialize vanilla tools
94
- from langchain.tools import DuckDuckGoSearchResults, WikipediaQueryRun
95
- from langchain.utilities import WikipediaAPIWrapper
96
-
97
- duckduckgo_tool = DuckDuckGoSearchResults()
98
- wiki_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
99
-
100
-
101
- # Initialize enhanced tools
102
- enhanced_search_tool = LangTool.from_function(
103
- name="enhanced_web_search",
104
- func=EnhancedSearchTool().run,
105
- description="Enhanced web search with intelligent query processing, multiple search strategies, and result filtering. Provides comprehensive and relevant search results."
106
- )
107
-
108
- enhanced_wiki_tool = LangTool.from_function(
109
- name="enhanced_wikipedia",
110
- func=EnhancedWikipediaTool().run,
111
- description="Enhanced Wikipedia search with entity extraction, multi-term search, and relevant content filtering. Provides detailed encyclopedic information."
112
- )
113
-
114
- excel_tool = StructuredTool.from_function(
115
- name="excel_to_text",
116
- func=excel_to_markdown,
117
- description="Enhanced Excel analysis with metadata, statistics, and structured data preview. Inputs: 'excel_path' (str), 'sheet_name' (str, optional).",
118
- )
119
-
120
- image_tool = StructuredTool.from_function(
121
- name="image_file_info",
122
- func=image_file_info,
123
- description="Enhanced image file analysis with detailed metadata and properties."
124
- )
125
-
126
- audio_tool = LangTool.from_function(
127
- name="audio_file_info",
128
- func=audio_file_info,
129
- description="Enhanced audio processing with transcription, language detection, and timestamped segments."
130
- )
131
-
132
- code_tool = LangTool.from_function(
133
- name="code_file_read",
134
- func=code_file_read,
135
- description="Enhanced code file analysis with language-specific insights and structure analysis."
136
- )
137
-
138
- youtube_tool = LangTool.from_function(
139
- name="extract_youtube_info",
140
- func=extract_youtube_info,
141
- description="Extracts transcription from the youtube link"
142
- )
143
-
144
- # Enhanced tool registry
145
- AVAILABLE_TOOLS = {
146
- "excel": excel_tool,
147
- "search": wiki_tool,
148
- "wikipedia": duckduckgo_tool,
149
- "image": image_tool,
150
- "audio": audio_tool,
151
- "code": code_tool,
152
- "youtube": youtube_tool
153
- }
154
-
155
- # ----------- Intelligent Tool Selection -----------
156
- def analyze_question(state: AgentState) -> AgentState:
157
- """Enhanced question analysis with better tool recommendation"""
158
- analysis_prompt = f"""
159
- Analyze this question and determine the best tools and approach:
160
- Question: {state["question"]}
161
-
162
- Available enhanced tools:
163
- 1. excel - Enhanced Excel/CSV analysis with statistics and metadata
164
- 2. search - Enhanced web search with intelligent query processing and result filtering
165
- 3. wikipedia - Enhanced Wikipedia search with entity extraction and content filtering
166
- 4. image - Enhanced image analysis with what the image contains
167
- 5. audio - Enhanced audio processing with transcription
168
- 6. code - Enhanced code analysis with language-specific insights
169
- 7. youtube - Extracts transcription from the youtube link
170
-
171
- Consider:
172
- - Question type (factual, analytical, current events, technical)
173
- - Required information sources (files, web, encyclopedic)
174
- - Time sensitivity (current vs historical information)
175
- - Complexity level
176
-
177
- Respond with:
178
- 1. Question type: <type>
179
- 2. Primary tools needed: <tools>
180
- 3. Search strategy: <strategy>
181
- 4. Expected answer format: <format>
182
-
183
- Format: TYPE: <type> | TOOLS: <tools> | STRATEGY: <strategy> | FORMAT: <format>
184
- """
185
-
186
- try:
187
- response = llm.invoke(analysis_prompt)
188
- state["conversation_history"].append({"role": "analysis", "content": response})
189
- state["current_step"] = AgentStep.SELECT_TOOLS.value
190
- except Exception as e:
191
- state["error_count"] += 1
192
- state["conversation_history"].append({"role": "error", "content": f"Analysis failed: {e}"})
193
- state["current_step"] = AgentStep.ERROR_RECOVERY.value
194
-
195
- return state
196
-
197
- def select_tools(state: AgentState) -> AgentState:
198
- """Enhanced tool selection with smarter logic"""
199
- question = state["question"].lower()
200
- selected_tools = []
201
-
202
- # File-based tool selection
203
- if any(keyword in question for keyword in ["excel", "csv", "spreadsheet", ".xlsx", ".xls"]):
204
- selected_tools.append("excel")
205
- if any(keyword in question for keyword in [".png", ".jpg", ".jpeg", ".bmp", ".gif", "image"]):
206
- selected_tools.append("image")
207
- if any(keyword in question for keyword in [".mp3", ".wav", ".ogg", "audio", "transcribe"]):
208
- selected_tools.append("audio")
209
- if any(keyword in question for keyword in [".py", ".ipynb", "code", "script", "function"]):
210
- selected_tools.append("code")
211
- if any(keyword in question for keyword in ["youtube"]):
212
- selected_tools.append("youtube")
213
-
214
- print(f"File-based tools selected: {selected_tools}")
215
-
216
- tools_prompt = f"""
217
- You are a smart assistant that selects relevant tools based on the user's natural language question.
218
-
219
- Available tools:
220
- - "search" → Use for real-time, recent, or broad web information.
221
- - "wikipedia" → Use for factual or encyclopedic knowledge.
222
- - "excel" → Use for spreadsheet-related questions (.xlsx, .csv).
223
- - "image" → Use for image files (.png, .jpg, etc.) or image-based tasks.
224
- - "audio" → Use for sound files (.mp3, .wav, etc.) or transcription.
225
- - "code" → Use for programming-related questions or when files like .py are mentioned.
226
- - "youtube" → Use for questions involving YouTube videos.
227
-
228
- Return the result as a **Python list of strings**, no explanation. Use only the relevant tools.
229
- If not relevant tool is found, return an empty list such as [].
230
-
231
- ### Examples:
232
-
233
- Q: "Show me recent news about elections in 2025"
234
- A: ["search"]
235
-
236
- Q: "Summarize this Wikipedia article about Einstein"
237
- A: ["wikipedia"]
238
-
239
- Q: "Analyze this .csv file"
240
- A: ["excel"]
241
-
242
- Q: "Transcribe this .wav audio file"
243
- A: ["audio"]
244
-
245
- Q: "Generate Python code from this prompt"
246
- A: ["code"]
247
-
248
- Q: "Who was the president of USA in 1945?"
249
- A: ["wikipedia"]
250
-
251
- Q: "Give me current weather updates"
252
- A: ["search"]
253
-
254
- Q: "Look up the history of space exploration"
255
- A: ["search", "wikipedia"]
256
-
257
- Q: "What is 2 + 2?"
258
- A: []
259
-
260
- ### Now answer:
261
-
262
- Q: {state["question"]}
263
- A:
264
- """
265
-
266
- llm_tools = ast.literal_eval(llm.invoke(tools_prompt).strip())
267
- if not isinstance(llm_tools, list):
268
- llm_tools = []
269
- print(f"LLM suggested tools: {llm_tools}")
270
- selected_tools.extend(llm_tools)
271
- selected_tools = list(set(selected_tools)) # Remove duplicates
272
-
273
- print(f"Final selected tools after LLM suggestion: {selected_tools}")
274
-
275
-
276
- # # Information-based tool selection
277
- # current_indicators = ["recent", "current", "news", "today", "2025", "now"]
278
- # encyclopedia_indicators = ["wiki", "wikipedia"]
279
-
280
- # if any(indicator in question for indicator in current_indicators):
281
- # selected_tools.append("search")
282
- # elif any(indicator in question for indicator in encyclopedia_indicators):
283
- # selected_tools.append("wikipedia")
284
- # elif any(keyword in question for keyword in ["search", "find", "look up", "information about"]):
285
- # # Use both for comprehensive coverage
286
- # selected_tools.extend(["search", "wikipedia"])
287
-
288
- # # Default fallback
289
- # if not selected_tools:
290
- # if any(word in question for word in ["who", "what", "when", "where"]):
291
- # selected_tools.append("wikipedia")
292
- # selected_tools.append("search")
293
-
294
- # # Remove duplicates while preserving order
295
- # selected_tools = list(dict.fromkeys(selected_tools))
296
-
297
- state["selected_tools"] = selected_tools
298
- state["current_step"] = AgentStep.EXECUTE_TOOLS.value
299
- return state
300
-
301
- def execute_tools(state: AgentState) -> AgentState:
302
- """Enhanced tool execution with better error handling"""
303
- results = {}
304
-
305
- # Enhanced file detection
306
- file_path = None
307
- downloaded_file_marker = "A file was downloaded for this task and saved locally at:"
308
- if downloaded_file_marker in state["question"]:
309
- lines = state["question"].splitlines()
310
- for i, line in enumerate(lines):
311
- if downloaded_file_marker in line:
312
- if i + 1 < len(lines):
313
- file_path_candidate = lines[i + 1].strip()
314
- if Path(file_path_candidate).exists():
315
- file_path = file_path_candidate
316
- print(f"Detected file path: {file_path}")
317
- break
318
-
319
- for tool_name in state["selected_tools"]:
320
- try:
321
- print(f"Executing tool: {tool_name}")
322
-
323
- # File-based tools
324
- if tool_name in ["excel", "image", "audio", "code"] and file_path:
325
- if tool_name == "excel":
326
- result = AVAILABLE_TOOLS["excel"].run({"excel_path": file_path, "sheet_name": None})
327
- elif tool_name == "image":
328
- result = AVAILABLE_TOOLS["image"].run({"image_path": file_path, "question": state["question"]})
329
- elif tool_name == "youtube":
330
- print(f"Running YouTube tool with file path: {file_path}")
331
- result = AVAILABLE_TOOLS["youtube"].run(state["question"])
332
- else:
333
- result = AVAILABLE_TOOLS[tool_name].run(file_path)
334
- # Information-based tools
335
- else:
336
- # Extract clean query for search tools
337
- clean_query = state["question"]
338
- if downloaded_file_marker in clean_query:
339
- clean_query = clean_query.split(downloaded_file_marker)[0].strip()
340
-
341
- result = AVAILABLE_TOOLS[tool_name].run(clean_query)
342
-
343
- results[tool_name] = result
344
-
345
- print(f"Tool {tool_name} completed successfully.")
346
- print(f"Output for {tool_name}: {result}")
347
-
348
- except Exception as e:
349
- error_msg = f"Error using {tool_name}: {str(e)}"
350
- results[tool_name] = error_msg
351
- state["error_count"] += 1
352
- print(error_msg)
353
-
354
- state["tool_results"] = results
355
- state["current_step"] = AgentStep.SYNTHESIZE_ANSWER.value
356
- return state
357
-
358
- def synthesize_answer(state: AgentState) -> AgentState:
359
- """Enhanced answer synthesis with better formatting"""
360
-
361
- tool_results_str = "\n".join([f"=== {tool.upper()} RESULTS ===\n{result}\n" for tool, result in state["tool_results"].items()])
362
-
363
- cot_prompt = f"""You are a precise assistant tasked with analyzing the user's question{" using the available tool outputs" if state["tool_results"] else ""}.
364
-
365
- Question:
366
- {state["question"]}
367
-
368
- {f"Available tool outputs: {tool_results_str}" if state["tool_results"] else ""}
369
-
370
- Instructions:
371
- - Think step-by-step to determine the best strategy to answer the question.
372
- - Use only the given information; do not hallucinate or infer from external knowledge.
373
- - If decoding, logical deduction, counting, or interpretation is required, show each step clearly.
374
- - If any part of the tool output is unclear or incomplete, mention it and its impact.
375
- - Do not guess. If the information is insufficient, say so clearly.
376
- - Finish with a clearly marked line: `---END OF ANALYSIS---`
377
-
378
- Your step-by-step analysis:"""
379
-
380
- cot_response = llm.invoke(cot_prompt)
381
-
382
- final_answer_prompt = f"""You are a precise assistant tasked with deriving the **final answer** from the step-by-step analysis below.
383
-
384
- Question:
385
- {state["question"]}
386
-
387
- Step-by-step analysis:
388
- {cot_response}
389
-
390
- Instructions:
391
- - Read the analysis thoroughly before responding.
392
- - Output ONLY the final answer. Do NOT include any reasoning or explanation.
393
- - Remove any punctuation at the corners of the answer unless it is explicitly mentioned in the question.
394
- - The answer must be concise and factual.
395
- - If the analysis concluded that a definitive answer cannot be determined, respond with: `NA` (exactly).
396
-
397
- Final answer:"""
398
-
399
-
400
- try:
401
- response = llm.invoke(final_answer_prompt)
402
- state["final_answer"] = response
403
- state["current_step"] = AgentStep.COMPLETE.value
404
- except Exception as e:
405
- state["error_count"] += 1
406
- state["final_answer"] = f"Error synthesizing answer: {e}"
407
- state["current_step"] = AgentStep.ERROR_RECOVERY.value
408
-
409
- return state
410
-
411
- def error_recovery(state: AgentState) -> AgentState:
412
- """Enhanced error recovery with multiple fallback strategies"""
413
- if state["error_count"] >= state["max_errors"]:
414
- state["final_answer"] = "I encountered multiple errors and cannot complete this task reliably."
415
- state["current_step"] = AgentStep.COMPLETE.value
416
- else:
417
- # Enhanced fallback: try with simplified approach
418
- try:
419
- fallback_prompt = f"""
420
- Answer this question directly using your knowledge:
421
- {state["original_question"]}
422
-
423
- Provide a helpful response even if you cannot access external tools.
424
- Be clear about any limitations in your answer.
425
- """
426
- response = llm.invoke(fallback_prompt)
427
- state["final_answer"] = f"Using available knowledge (some tools unavailable): {response}"
428
- state["current_step"] = AgentStep.COMPLETE.value
429
- except Exception as e:
430
- state["final_answer"] = f"All approaches failed. Error: {e}"
431
- state["current_step"] = AgentStep.COMPLETE.value
432
-
433
- return state
434
-
435
- # ----------- Enhanced LangGraph Workflow -----------
436
- def route_next_step(state: AgentState) -> str:
437
- """Route to next step based on current state"""
438
- step_routing = {
439
- "start": AgentStep.ANALYZE_QUESTION.value,
440
- AgentStep.ANALYZE_QUESTION.value: AgentStep.SELECT_TOOLS.value,
441
- AgentStep.SELECT_TOOLS.value: AgentStep.EXECUTE_TOOLS.value,
442
- AgentStep.EXECUTE_TOOLS.value: AgentStep.SYNTHESIZE_ANSWER.value,
443
- AgentStep.SYNTHESIZE_ANSWER.value: AgentStep.COMPLETE.value,
444
- AgentStep.ERROR_RECOVERY.value: AgentStep.COMPLETE.value,
445
- AgentStep.COMPLETE.value: END,
446
- }
447
-
448
- return step_routing.get(state["current_step"], END)
449
-
450
- # Create enhanced workflow
451
- workflow = StateGraph(AgentState)
452
-
453
- # Add nodes
454
- workflow.add_node("analyze_question", RunnableLambda(analyze_question))
455
- workflow.add_node("select_tools", RunnableLambda(select_tools))
456
- workflow.add_node("execute_tools", RunnableLambda(execute_tools))
457
- workflow.add_node("synthesize_answer", RunnableLambda(synthesize_answer))
458
- workflow.add_node("error_recovery", RunnableLambda(error_recovery))
459
-
460
- # Set entry point
461
- workflow.set_entry_point("analyze_question")
462
-
463
- # Add conditional edges
464
- workflow.add_conditional_edges(
465
- "analyze_question",
466
- lambda state: "select_tools" if state["current_step"] == AgentStep.SELECT_TOOLS.value else "error_recovery"
467
- )
468
- workflow.add_edge("select_tools", "execute_tools")
469
- workflow.add_conditional_edges(
470
- "execute_tools",
471
- lambda state: "synthesize_answer" if state["current_step"] == AgentStep.SYNTHESIZE_ANSWER.value else "error_recovery"
472
- )
473
- workflow.add_conditional_edges(
474
- "synthesize_answer",
475
- lambda state: END if state["current_step"] == AgentStep.COMPLETE.value else "error_recovery"
476
- )
477
- workflow.add_edge("error_recovery", END)
478
-
479
- # Compile the enhanced graph
480
- graph = workflow.compile()
481
-
482
- # ----------- Agent Class -----------
483
- class GaiaAgent:
484
- """GAIA Agent with tools and intelligent processing"""
485
-
486
- def __init__(self):
487
- self.graph = graph
488
- self.tool_usage_stats = {}
489
- print("Enhanced GAIA Agent initialized with:")
490
- print("✓ Intelligent multi-query web search")
491
- print("✓ Entity-aware Wikipedia search")
492
- print("✓ Enhanced file processing tools")
493
- print("✓ Advanced error recovery")
494
- print("✓ Comprehensive result synthesis")
495
-
496
- def get_tool_stats(self) -> Dict[str, int]:
497
- """Get usage statistics for tools"""
498
- return self.tool_usage_stats.copy()
499
-
500
- def __call__(self, task_id: str, question: str) -> str:
501
- print(f"\n{'='*60}")
502
- print(f"[{task_id}] ENHANCED PROCESSING: {question}")
503
-
504
- # Initialize state
505
- processed_question = process_file(task_id, question)
506
- initial_state = initialize_state(processed_question)
507
-
508
- try:
509
- # Execute the enhanced workflow
510
- result = self.graph.invoke(initial_state)
511
-
512
- # Extract results
513
- answer = result.get("final_answer", "No answer generated")
514
- selected_tools = result.get("selected_tools", [])
515
- conversation_history = result.get("conversation_history", [])
516
- tool_results = result.get("tool_results", {})
517
- error_count = result.get("error_count", 0)
518
-
519
- # Update tool usage statistics
520
- for tool in selected_tools:
521
- self.tool_usage_stats[tool] = self.tool_usage_stats.get(tool, 0) + 1
522
-
523
- # Enhanced logging
524
- print(f"[{task_id}] Selected tools: {selected_tools}")
525
- print(f"[{task_id}] Tools executed: {list(tool_results.keys())}")
526
- print(f"[{task_id}] Processing steps: {len(conversation_history)}")
527
- print(f"[{task_id}] Errors encountered: {error_count}")
528
-
529
- # Log tool result sizes for debugging
530
- for tool, result in tool_results.items():
531
- result_size = len(str(result)) if result else 0
532
- print(f"[{task_id}] {tool} result size: {result_size} chars")
533
-
534
- print(f"[{task_id}] FINAL ANSWER: {answer}")
535
- print(f"{'='*60}")
536
-
537
- return answer
538
-
539
- except Exception as e:
540
- error_msg = f"Critical error in enhanced agent execution: {str(e)}"
541
- print(f"[{task_id}] {error_msg}")
542
-
543
- # Try fallback direct LLM response
544
- try:
545
- fallback_response = llm.invoke(f"Please answer this question: {question}")
546
- return f"Fallback response: {fallback_response}"
547
- except:
548
- return error_msg
549
-
550
- # ----------- Enhanced File Processing -----------
551
- def detect_file_type(file_path: str) -> Optional[str]:
552
- """Enhanced file type detection with more formats"""
553
- ext = Path(file_path).suffix.lower()
554
-
555
- file_type_mapping = {
556
- # Spreadsheets
557
- '.xlsx': 'excel', '.xls': 'excel', '.csv': 'excel',
558
- # Images
559
- '.png': 'image', '.jpg': 'image', '.jpeg': 'image',
560
- '.bmp': 'image', '.gif': 'image', '.tiff': 'image', '.webp': 'image',
561
- # Audio
562
- '.mp3': 'audio', '.wav': 'audio', '.ogg': 'audio',
563
- '.flac': 'audio', '.m4a': 'audio', '.aac': 'audio',
564
- # Code
565
- '.py': 'code', '.ipynb': 'code', '.js': 'code', '.html': 'code',
566
- '.css': 'code', '.java': 'code', '.cpp': 'code', '.c': 'code',
567
- '.sql': 'code', '.r': 'code', '.json': 'code', '.xml': 'code',
568
- # Documents
569
- '.txt': 'text', '.md': 'text', '.pdf': 'document',
570
- '.doc': 'document', '.docx': 'document'
571
- }
572
-
573
- return file_type_mapping.get(ext)
574
-
575
- def process_file(task_id: str, question_text: str) -> str:
576
- """Enhanced file processing with better error handling and metadata"""
577
- file_url = f"{FILE_PATH}{task_id}"
578
-
579
- try:
580
- print(f"[{task_id}] Attempting to download file from: {file_url}")
581
- response = requests.get(file_url, timeout=30)
582
- response.raise_for_status()
583
- print(f"[{task_id}] File download successful. Status: {response.status_code}")
584
-
585
- except requests.exceptions.RequestException as exc:
586
- print(f"[{task_id}] File download failed: {str(exc)}")
587
- return question_text # Return original question if no file
588
-
589
- # Enhanced filename extraction
590
- content_disposition = response.headers.get("content-disposition", "")
591
- filename = task_id # Default fallback
592
-
593
- # Try to extract filename from Content-Disposition header
594
- filename_match = re.search(r'filename[*]?=(?:"([^"]+)"|([^;]+))', content_disposition)
595
- if filename_match:
596
- filename = filename_match.group(1) or filename_match.group(2)
597
- filename = filename.strip()
598
-
599
- # Create enhanced temp directory structure
600
- temp_storage_dir = Path(tempfile.gettempdir()) / "gaia_enhanced_files" / task_id
601
- temp_storage_dir.mkdir(parents=True, exist_ok=True)
602
-
603
- file_path = temp_storage_dir / filename
604
- file_path.write_bytes(response.content)
605
-
606
- # Get file metadata
607
- file_size = len(response.content)
608
- file_type = detect_file_type(filename)
609
-
610
- print(f"[{task_id}] File saved: {filename} ({file_size:,} bytes, type: {file_type})")
611
-
612
- # Enhanced question augmentation
613
- enhanced_question = f"{question_text}\n\n"
614
- enhanced_question += f"{'='*50}\n"
615
- enhanced_question += f"FILE INFORMATION:\n"
616
- enhanced_question += f"A file was downloaded for this task and saved locally at:\n"
617
- enhanced_question += f"{str(file_path)}\n"
618
- enhanced_question += f"File details:\n"
619
- enhanced_question += f"- Name: {filename}\n"
620
- enhanced_question += f"- Size: {file_size:,} bytes ({file_size/1024:.1f} KB)\n"
621
- enhanced_question += f"- Type: {file_type or 'unknown'}\n"
622
- enhanced_question += f"{'='*50}\n\n"
623
-
624
- return enhanced_question
625
-
626
- # ----------- Usage Examples and Testing -----------
627
- def run_enhanced_tests():
628
- """Run comprehensive tests of the enhanced agent"""
629
- agent = GaiaAgent()
630
-
631
- test_cases = [
632
- {
633
- "id": "test_search_1",
634
- "question": "What are the latest developments in artificial intelligence in 2024?",
635
- "expected_tools": ["search"]
636
- },
637
- {
638
- "id": "test_wiki_1",
639
- "question": "Tell me about Albert Einstein's contributions to physics",
640
- "expected_tools": ["wikipedia"]
641
- },
642
- {
643
- "id": "test_combined_1",
644
- "question": "What is machine learning and what are recent breakthroughs?",
645
- "expected_tools": ["wikipedia", "search"]
646
- },
647
- {
648
- "id": "test_excel_1",
649
- "question": "Analyze the data in the Excel file sales_data.xlsx",
650
- "expected_tools": ["excel"]
651
- }
652
- ]
653
-
654
- print("\n" + "="*80)
655
- print("RUNNING ENHANCED AGENT TESTS")
656
- print("="*80)
657
-
658
- for test_case in test_cases:
659
- print(f"\nTest Case: {test_case['id']}")
660
- print(f"Question: {test_case['question']}")
661
- print(f"Expected tools: {test_case['expected_tools']}")
662
-
663
- try:
664
- result = agent(test_case['id'], test_case['question'])
665
- print(f"Result length: {len(result)} characters")
666
- print(f"Result preview: {result[:200]}...")
667
- except Exception as e:
668
- print(f"Test failed: {e}")
669
-
670
- print("-" * 60)
671
-
672
- # Print tool usage statistics
673
- print(f"\nTool Usage Statistics:")
674
- for tool, count in agent.get_tool_stats().items():
675
- print(f" {tool}: {count} times")
676
-
677
- # Usage example
678
- if __name__ == "__main__":
679
- # Create enhanced agent
680
- agent = GaiaAgent()
681
-
682
- # Example usage
683
- sample_questions = [
684
- "What is the current population of Tokyo and how has it changed recently?",
685
- "Explain quantum computing and its recent developments",
686
- "Tell me about the history of machine learning and current AI trends",
687
- ]
688
-
689
- print("\n" + "="*80)
690
- print("ENHANCED GAIA AGENT DEMONSTRATION")
691
- print("="*80)
692
-
693
- for i, question in enumerate(sample_questions):
694
- print(f"\nExample {i+1}: {question}")
695
- result = agent(f"demo_{i}", question)
696
- print(f"Answer: {result[:300]}...")
697
- print("-" * 60)
698
-
699
- # Uncomment to run comprehensive tests
700
- # run_enhanced_tests()