Final_Assignment_Template / agent_final.py
ChrisNguyen05's picture
Create agent_final.py
534b80c verified
class GeminiAgent:
def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
# Suppress warnings
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", message=".*will be deprecated.*")
warnings.filterwarnings("ignore", "LangChain.*")
self.api_key = api_key
self.model_name = model_name
# Configure Gemini
genai.configure(api_key=api_key)
# Initialize the LLM
self.llm = self._setup_llm()
# Setup tools
self.tools = [
SmolagentToolWrapper(WikipediaSearchTool()),
Tool(
name="analyze_video",
func=self._analyze_video,
description="Analyze YouTube video content directly"
),
Tool(
name="analyze_image",
func=self._analyze_image,
description="Analyze image content"
),
Tool(
name="analyze_table",
func=self._analyze_table,
description="Analyze table or matrix data"
),
Tool(
name="analyze_list",
func=self._analyze_list,
description="Analyze and categorize list items"
),
Tool(
name="web_search",
func=self._web_search,
description="Search the web for information"
)
]
# Setup memory
self.memory = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True
)
# Initialize agent
self.agent = self._setup_agent()
def run(self, query: str) -> str:
"""Run the agent on a query with incremental retries."""
max_retries = 3
base_sleep = 1 # Start with 1 second sleep
for attempt in range(max_retries):
try:
# If no match found in answer bank, use the agent
response = self.agent.run(query)
return response
except Exception as e:
sleep_time = base_sleep * (attempt + 1) # Incremental sleep: 1s, 2s, 3s
if attempt < max_retries - 1:
print(f"Attempt {attempt + 1} failed. Retrying in {sleep_time} seconds...")
time.sleep(sleep_time)
continue
return f"Error processing query after {max_retries} attempts: {str(e)}"
print("Agent processed all queries!")
def _clean_response(self, response: str) -> str: