File size: 2,797 Bytes
534b80c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
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: |