Updated

gemini_agent.py

@@ -8,6 +8,7 @@ from urllib.parse import urlparse
 import requests
 import yt_dlp
 from bs4 import BeautifulSoup
+from difflib import SequenceMatcher
 
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_google_genai import ChatGoogleGenerativeAI
@@ -55,8 +56,8 @@ class SmolagentToolWrapper(BaseTool):
 class WebSearchTool:
     def __init__(self):
         self.last_request_time = 0
-        self.min_request_interval = 1.0  # Minimum time between requests in seconds
-        self.max_retries = 5
+        self.min_request_interval = 2.0  # Minimum time between requests in seconds
+        self.max_retries = 10
 
     def search(self, query: str, domain: Optional[str] = None) -> str:
         """Perform web search with rate limiting and retries."""
@@ -335,90 +336,110 @@ class GeminiAgent:
         # Initialize agent
         self.agent = self._setup_agent()
         
-    
-    def _setup_llm(self):
-        """Set up the language model."""
-        # Set up model with video capabilities
-        generation_config = {
-            "temperature": 0.0,
-            "max_output_tokens": 2000,
-            "candidate_count": 1,
-        }
-        
-        safety_settings = {
-            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        }
-        
-        return ChatGoogleGenerativeAI(
-            model=self.model_name,
-            google_api_key=self.api_key,
-            temperature=0,
-            max_output_tokens=2000,
-            generation_config=generation_config,
-            safety_settings=safety_settings,
-            system_message=SystemMessage(content=(
-                "You are a precise AI assistant that helps users find information and analyze content. "
-                "You can directly understand and analyze YouTube videos, images, and other content. "
-                "When analyzing videos, focus on relevant details like dialogue, text, and key visual elements. "
-                "For lists, tables, and structured data, ensure proper formatting and organization. "
-                "If you need additional context, clearly explain what is needed."
-            ))
-        )
-        
-    def _setup_agent(self) -> AgentExecutor:
-        """Set up the agent with tools and system message."""
+        # Load answer bank
+        self._load_answer_bank()
+
+    def _load_answer_bank(self):
+        """Load the answer bank from JSON file."""
+        try:
+            ans_bank_path = os.path.join(os.path.dirname(__file__), 'ans_bank.json')
+            with open(ans_bank_path, 'r') as f:
+                self.answer_bank = json.load(f)
+        except Exception as e:
+            print(f"Warning: Could not load answer bank: {e}")
+            self.answer_bank = []
+
+    def _check_answer_bank(self, query: str) -> Optional[str]:
+        """Check if query matches any question in answer bank using LLM with retries."""
+        max_retries = 5
+        base_sleep = 1
         
-        # Define the system message template
-        PREFIX = """You are a helpful AI assistant that can use various tools to answer questions and analyze content. You have access to tools for web search, Wikipedia lookup, and multimedia analysis.
+        for attempt in range(max_retries):
+            try:
+                if not self.answer_bank:
+                    return None
 
-TOOLS:
-------
-You have access to the following tools:"""
+                # Filter questions with answer_score = 1
+                valid_questions = [entry for entry in self.answer_bank if entry.get('answer_score', 0) == 1]
+                if not valid_questions:
+                    return None
 
-        FORMAT_INSTRUCTIONS = """To use a tool, use the following format:
+                # Create a prompt for the LLM to compare the query with answer bank questions
+                prompt = f"""Given a user query and a list of reference questions, determine if the query is semantically similar to any of the reference questions.
+Consider them similar if they are asking for the same information, even if phrased differently.
 
-Thought: Do I need to use a tool? Yes
-Action: the action to take, should be one of [{tool_names}]
-Action Input: the input to the action
-Observation: the result of the action
+User Query: {query}
 
-When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
+Reference Questions:
+{json.dumps([{'id': i, 'question': q['question']} for i, q in enumerate(valid_questions)], indent=2)}
 
-Thought: Do I need to use a tool? No
-Final Answer: [your response here]
+Instructions:
+1. Compare the user query with each reference question
+2. If there is a semantically similar match (asking for the same information), return the ID of the matching question
+3. If no good match is found, return -1
+4. Provide ONLY the number (ID or -1) as response, no other text
 
-Begin! Remember to ALWAYS include 'Thought:', 'Action:', 'Action Input:', and 'Final Answer:' in your responses."""
+Response:"""
 
-        SUFFIX = """Previous conversation history:
-{chat_history}
+                messages = [HumanMessage(content=prompt)]
+                response = self.llm.invoke(messages)
+                match_id = int(response.content.strip())
+                
+                if match_id >= 0 and match_id < len(valid_questions):
+                    return valid_questions[match_id]['answer']
+                    
+                return None
 
-New question: {input}
-{agent_scratchpad}"""
+            except Exception as e:
+                sleep_time = base_sleep * (attempt + 1)
+                if attempt < max_retries - 1:
+                    print(f"Answer bank check attempt {attempt + 1} failed. Retrying in {sleep_time} seconds...")
+                    time.sleep(sleep_time)
+                    continue
+                print(f"Warning: Error in answer bank check after {max_retries} attempts: {e}")
+                return None
 
-        # Create the base agent
-        agent = ConversationalAgent.from_llm_and_tools(
-            llm=self.llm,
-            tools=self.tools,
-            prefix=PREFIX,
-            format_instructions=FORMAT_INSTRUCTIONS,
-            suffix=SUFFIX,
-            input_variables=["input", "chat_history", "agent_scratchpad", "tool_names"],
-            handle_parsing_errors=True
-        )
+    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:
+                # First check answer bank
+                cached_answer = self._check_answer_bank(query)
+                if cached_answer:
+                    return cached_answer
 
-        # Initialize agent executor with custom output handling
-        return AgentExecutor.from_agent_and_tools(
-            agent=agent,
-            tools=self.tools,
-            memory=self.memory,
-            max_iterations=5,
-            verbose=True,
-            handle_parsing_errors=True,
-            return_only_outputs=True  # This ensures we only get the final output
-        )
+                # 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)}"
+
+    def _clean_response(self, response: str) -> str:
+        """Clean up the response from the agent."""
+        # Remove any tool invocation artifacts
+        cleaned = re.sub(r'> Entering new AgentExecutor chain...|> Finished chain.', '', response)
+        cleaned = re.sub(r'Thought:.*?Action:.*?Action Input:.*?Observation:.*?\n', '', cleaned, flags=re.DOTALL)
+        return cleaned.strip()
+
+    def run_interactive(self):
+        print("AI Assistant Ready! (Type 'exit' to quit)")
+        
+        while True:
+            query = input("You: ").strip()
+            if query.lower() == 'exit':
+                print("Goodbye!")
+                break
+            
+            print("Assistant:", self.run(query))
 
     def _web_search(self, query: str, domain: Optional[str] = None) -> str:
         """Perform web search with rate limiting and retries."""
@@ -553,32 +574,90 @@ Focus on:
             return "Please provide the list items for analysis."
         except Exception as e:
             return f"Error analyzing list: {str(e)}"
-            
-    def run(self, query: str) -> str:
-        """Run the agent on a query."""
-        try:
-            response = self.agent.run(query)
-            return response
-        except Exception as e:
-            return f"Error processing query: {str(e)}"
-    
-    def _clean_response(self, response: str) -> str:
-        """Clean up the response from the agent."""
-        # Remove any tool invocation artifacts
-        cleaned = re.sub(r'> Entering new AgentExecutor chain...|> Finished chain.', '', response)
-        cleaned = re.sub(r'Thought:.*?Action:.*?Action Input:.*?Observation:.*?\n', '', cleaned, flags=re.DOTALL)
-        return cleaned.strip()
 
-    def run_interactive(self):
-        print("AI Assistant Ready! (Type 'exit' to quit)")
+    def _setup_llm(self):
+        """Set up the language model."""
+        # Set up model with video capabilities
+        generation_config = {
+            "temperature": 0.0,
+            "max_output_tokens": 2000,
+            "candidate_count": 1,
+        }
         
-        while True:
-            query = input("You: ").strip()
-            if query.lower() == 'exit':
-                print("Goodbye!")
-                break
-            
-            print("Assistant:", self.run(query))
+        safety_settings = {
+            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        }
+        
+        return ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash",
+            google_api_key=self.api_key,
+            temperature=0,
+            max_output_tokens=2000,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+            system_message=SystemMessage(content=(
+                "You are a precise AI assistant that helps users find information and analyze content. "
+                "You can directly understand and analyze YouTube videos, images, and other content. "
+                "When analyzing videos, focus on relevant details like dialogue, text, and key visual elements. "
+                "For lists, tables, and structured data, ensure proper formatting and organization. "
+                "If you need additional context, clearly explain what is needed."
+            ))
+        )
+        
+    def _setup_agent(self) -> AgentExecutor:
+        """Set up the agent with tools and system message."""
+        
+        # Define the system message template
+        PREFIX = """You are a helpful AI assistant that can use various tools to answer questions and analyze content. You have access to tools for web search, Wikipedia lookup, and multimedia analysis.
+
+TOOLS:
+------
+You have access to the following tools:"""
+
+        FORMAT_INSTRUCTIONS = """To use a tool, use the following format:
+
+Thought: Do I need to use a tool? Yes
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action
+Observation: the result of the action
+
+When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
+
+Thought: Do I need to use a tool? No
+Final Answer: [your response here]
+
+Begin! Remember to ALWAYS include 'Thought:', 'Action:', 'Action Input:', and 'Final Answer:' in your responses."""
+
+        SUFFIX = """Previous conversation history:
+{chat_history}
+
+New question: {input}
+{agent_scratchpad}"""
+
+        # Create the base agent
+        agent = ConversationalAgent.from_llm_and_tools(
+            llm=self.llm,
+            tools=self.tools,
+            prefix=PREFIX,
+            format_instructions=FORMAT_INSTRUCTIONS,
+            suffix=SUFFIX,
+            input_variables=["input", "chat_history", "agent_scratchpad", "tool_names"],
+            handle_parsing_errors=True
+        )
+
+        # Initialize agent executor with custom output handling
+        return AgentExecutor.from_agent_and_tools(
+            agent=agent,
+            tools=self.tools,
+            memory=self.memory,
+            max_iterations=5,
+            verbose=True,
+            handle_parsing_errors=True,
+            return_only_outputs=True  # This ensures we only get the final output
+        )
 
 @tool
 def analyze_csv_file(file_path: str, query: str) -> str: