|
class GeminiAgent: |
|
def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"): |
|
|
|
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 |
|
|
|
|
|
genai.configure(api_key=api_key) |
|
|
|
|
|
self.llm = self._setup_llm() |
|
|
|
|
|
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" |
|
) |
|
] |
|
|
|
|
|
self.memory = ConversationBufferMemory( |
|
memory_key="chat_history", |
|
return_messages=True |
|
) |
|
|
|
|
|
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 |
|
|
|
for attempt in range(max_retries): |
|
try: |
|
|
|
|
|
response = self.agent.run(query) |
|
return response |
|
|
|
except Exception as e: |
|
sleep_time = base_sleep * (attempt + 1) |
|
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: |