Updated

gemini_agent.py

@@ -1,10 +1,5 @@
-from typing import List
-from langchain_community.tools.ddg_search import DuckDuckGoSearchRun
-from langchain.agents import AgentType, initialize_agent
-from langchain.tools.base import BaseTool
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.messages import SystemMessage
-from langchain.memory import ConversationBufferMemory
 
 class GeminiAgent:
     def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
@@ -21,39 +16,19 @@ class GeminiAgent:
 
     def _setup_agent(self):
         # Initialize model with system message
-        model = ChatGoogleGenerativeAI(
+        return ChatGoogleGenerativeAI(
             model=self.model_name,
             google_api_key=self.api_key,
-            temperature=0,  # Lower temperature for faster, more focused responses
-            max_output_tokens=200,  # Limit response length
-            convert_system_message_to_human=True,  # Faster processing of system message
-            stream=True,  # Enable streaming for faster initial response
-            system_message=SystemMessage(content="You are a concise AI assistant. Provide a short and accurate answer. Preferable answer should be in one word or line. Unless if query asked expects an elaborate answer.")
-        )
-
-        # Setup tools
-        tools: List[BaseTool] = [DuckDuckGoSearchRun()]
-
-        # Setup memory
-        memory = ConversationBufferMemory(
-            memory_key="chat_history",
-            return_messages=True
-        )
-
-        # Create and return agent
-        return initialize_agent(
-            tools,
-            model,
-            agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
-            memory=memory,
-            verbose=False,
-            handle_parsing_errors=True
+            temperature=0,  # Lower temperature for focused responses
+            max_output_tokens=2000,  # Increased for more detailed responses
+            convert_system_message_to_human=True,
+            system_message=SystemMessage(content="You are a helpful AI assistant. For the Wikipedia question, use the latest 2022 English Wikipedia version as your knowledge source. For the YouTube video question, analyze the video content carefully and count the maximum number of different bird species visible simultaneously in any frame.")
         )
 
     def run(self, query: str) -> str:
         try:
-            result = self.agent.invoke({"input": query})
-            return result["output"]
+            response = self.agent.invoke(query)
+            return response.content
         except Exception as e:
             return f"Error: {e}"