ChrisNguyen05 commited on
Commit
534b80c
·
verified ·
1 Parent(s): 81917a3

Create agent_final.py

Browse files
Files changed (1) hide show
  1. agent_final.py +81 -0
agent_final.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ class GeminiAgent:
2
+ def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
3
+ # Suppress warnings
4
+ import warnings
5
+ warnings.filterwarnings("ignore", category=UserWarning)
6
+ warnings.filterwarnings("ignore", category=DeprecationWarning)
7
+ warnings.filterwarnings("ignore", message=".*will be deprecated.*")
8
+ warnings.filterwarnings("ignore", "LangChain.*")
9
+
10
+ self.api_key = api_key
11
+ self.model_name = model_name
12
+
13
+ # Configure Gemini
14
+ genai.configure(api_key=api_key)
15
+
16
+ # Initialize the LLM
17
+ self.llm = self._setup_llm()
18
+
19
+ # Setup tools
20
+ self.tools = [
21
+ SmolagentToolWrapper(WikipediaSearchTool()),
22
+ Tool(
23
+ name="analyze_video",
24
+ func=self._analyze_video,
25
+ description="Analyze YouTube video content directly"
26
+ ),
27
+ Tool(
28
+ name="analyze_image",
29
+ func=self._analyze_image,
30
+ description="Analyze image content"
31
+ ),
32
+ Tool(
33
+ name="analyze_table",
34
+ func=self._analyze_table,
35
+ description="Analyze table or matrix data"
36
+ ),
37
+ Tool(
38
+ name="analyze_list",
39
+ func=self._analyze_list,
40
+ description="Analyze and categorize list items"
41
+ ),
42
+ Tool(
43
+ name="web_search",
44
+ func=self._web_search,
45
+ description="Search the web for information"
46
+ )
47
+ ]
48
+
49
+ # Setup memory
50
+ self.memory = ConversationBufferMemory(
51
+ memory_key="chat_history",
52
+ return_messages=True
53
+ )
54
+
55
+ # Initialize agent
56
+ self.agent = self._setup_agent()
57
+
58
+
59
+ def run(self, query: str) -> str:
60
+ """Run the agent on a query with incremental retries."""
61
+ max_retries = 3
62
+ base_sleep = 1 # Start with 1 second sleep
63
+
64
+ for attempt in range(max_retries):
65
+ try:
66
+
67
+ # If no match found in answer bank, use the agent
68
+ response = self.agent.run(query)
69
+ return response
70
+
71
+ except Exception as e:
72
+ sleep_time = base_sleep * (attempt + 1) # Incremental sleep: 1s, 2s, 3s
73
+ if attempt < max_retries - 1:
74
+ print(f"Attempt {attempt + 1} failed. Retrying in {sleep_time} seconds...")
75
+ time.sleep(sleep_time)
76
+ continue
77
+ return f"Error processing query after {max_retries} attempts: {str(e)}"
78
+
79
+ print("Agent processed all queries!")
80
+
81
+ def _clean_response(self, response: str) -> str: