Update gemini_agent.py

gemini_agent.py

@@ -41,12 +41,92 @@ class GeminiAgent:
             if '.xlsx' in question or '.xls' in question or 'excel' in question.lower():
                 return await self._handle_excel_question(question)
             
+            # Check if question is about actors, TV shows, or movies
+            if self._is_actor_or_show_question(question):
+                return await self._handle_actor_show_question(question)
+            
             # Regular text-based question
             return await self._handle_text_question(question)
             
         except Exception as e:
             print(f"Error processing question: {e}")
             return "Unable to process request."
+            
+    def _is_actor_or_show_question(self, question: str) -> bool:
+        """Determine if a question is about actors, TV shows, or movies"""
+        q = question.lower()
+        actor_show_patterns = [
+            "who played", "who did", "who was the actor", "who was the actress",
+            "what role", "what character", "what part", 
+            "which actor", "which actress",
+            "in the movie", "in the show", "in the series", "in the film",
+            "version of", "language version", "dubbed version"
+        ]
+        return any(pattern in q for pattern in actor_show_patterns)
+        
+    async def _handle_actor_show_question(self, question: str) -> str:
+        """Handle questions about actors, TV shows, and movies with enhanced search"""
+        print(f"Processing actor/show question: {question[:50]}...")
+        
+        # Always try both Wikipedia and DuckDuckGo for these questions
+        wiki_context = ""
+        ddg_context = ""
+        
+        try:
+            wiki_context = self.wiki_tool.run(question)
+            print("Wikipedia search completed")
+        except Exception as e:
+            print(f"Wikipedia tool failed: {e}")
+            
+        try:
+            ddg_context = self.ddg_tool.run(question)
+            print("DuckDuckGo search completed")
+        except Exception as e:
+            print(f"DuckDuckGo tool failed: {e}")
+            
+        # Combine contexts if available
+        combined_context = ""
+        if wiki_context and not any(x in wiki_context.lower() for x in ["not found", "no results", "does not contain"]):
+            combined_context += f"Wikipedia context: {wiki_context}\n\n"
+        if ddg_context and not any(x in ddg_context.lower() for x in ["not found", "no results", "does not contain"]):
+            combined_context += f"Web search context: {ddg_context}\n\n"
+            
+        # Create a specialized prompt for actor/show questions
+        prompt = f"""Based on the following context, answer this question about an actor or TV show:
+
+{combined_context}
+
+Question: {question}
+
+Provide ONLY the specific name or information requested. No explanations or additional context.
+If the answer is a person's name, provide ONLY their first name as requested."""
+        
+        await self._rate_limit()
+        response = self.model.generate_content(
+            prompt,
+            generation_config=genai.types.GenerationConfig(
+                max_output_tokens=50,
+                temperature=0.0
+            )
+        )
+        answer = response.text.strip()
+        
+        # Clean up the answer to extract just the name or information
+        # Remove common prefixes
+        prefixes = ['The answer is', 'Based on', 'According to', 'The actor is', 'The actress is']
+        for prefix in prefixes:
+            if answer.lower().startswith(prefix.lower()):
+                answer = answer[len(prefix):].strip()
+                if answer.startswith(','):
+                    answer = answer[1:].strip()
+                    
+        # If the question asks for just a first name, extract it
+        if "give only the first name" in question.lower() or "only the first name" in question.lower():
+            name_parts = answer.split()
+            if name_parts:
+                answer = name_parts[0].rstrip(',.')
+                
+        return answer
     
     async def _handle_video_question(self, question: str) -> str:
         """Handle questions that require video analysis"""
@@ -163,7 +243,7 @@ Provide only the direct answer. If it's a quote, give just the quoted text. If i
     async def _handle_text_question(self, question: str) -> str:
         """Handle regular text-based questions"""
         prompt = ""
-        # Only use retrieval for explicit web/Wikipedia questions
+        # Check for different types of questions that need retrieval
         def is_explicit_retrieval_question(question):
             q = question.lower()
             return (
@@ -173,17 +253,36 @@ Provide only the direct answer. If it's a quote, give just the quoted text. If i
                 "duckduckgo" in q or
                 "web search" in q
             )
+            
+        def is_factual_question(question):
+            q = question.lower()
+            # Check for factual question patterns about people, shows, movies, etc.
+            factual_patterns = [
+                "who played", "who did", "who was", "who is", 
+                "what role", "what character", "what part", 
+                "which actor", "which actress",
+                "in the movie", "in the show", "in the series", "in the film",
+                "version of"
+            ]
+            return any(pattern in q for pattern in factual_patterns)
+            
         wiki_context = ""
         ddg_context = ""
-        if is_explicit_retrieval_question(question):
-            if "wikipedia" in question.lower():
+        
+        # Use retrieval for explicit web/Wikipedia questions OR factual questions about actors/shows
+        if is_explicit_retrieval_question(question) or is_factual_question(question):
+            # For factual questions, always try both Wikipedia and DuckDuckGo
+            if is_factual_question(question) or "wikipedia" in question.lower():
                 try:
                     wiki_context = self.wiki_tool.run(question)
+                    print(f"Wikipedia search completed for: {question[:50]}...")
                 except Exception as e:
                     print(f"Wikipedia tool failed: {e}")
-            if "duckduckgo" in question.lower() or "web search" in question.lower():
+                    
+            if is_factual_question(question) or "duckduckgo" in question.lower() or "web search" in question.lower():
                 try:
                     ddg_context = self.ddg_tool.run(question)
+                    print(f"DuckDuckGo search completed for: {question[:50]}...")
                 except Exception as e:
                     print(f"DuckDuckGo tool failed: {e}")
         # Handle attached file questions with enhanced prompts
@@ -226,10 +325,23 @@ Provide only the direct answer. If it's a quote, give just the quoted text. If i
         # Prepend context to the prompt if available and likely relevant
         def is_good_context(context):
             return context and not any(x in context.lower() for x in ["not found", "no results", "does not contain information"])
-        if wiki_context and is_good_context(wiki_context):
-            prompt = f"Use the following Wikipedia context to answer the question:\n{wiki_context}\n\n{prompt}"
-        elif ddg_context and is_good_context(ddg_context):
-            prompt = f"Use the following web search context to answer the question:\n{ddg_context}\n\n{prompt}"
+            
+        # For factual questions about actors/shows, try to use both Wikipedia and web search results
+        if is_factual_question(question):
+            combined_context = ""
+            if wiki_context and is_good_context(wiki_context):
+                combined_context += f"Wikipedia context: {wiki_context}\n\n"
+            if ddg_context and is_good_context(ddg_context):
+                combined_context += f"Web search context: {ddg_context}\n\n"
+                
+            if combined_context:
+                prompt = f"Use the following context to answer the question accurately. Focus on finding the exact name or information requested:\n{combined_context}\n{prompt}"
+        else:
+            # For non-factual questions, use the first good context available
+            if wiki_context and is_good_context(wiki_context):
+                prompt = f"Use the following Wikipedia context to answer the question:\n{wiki_context}\n\n{prompt}"
+            elif ddg_context and is_good_context(ddg_context):
+                prompt = f"Use the following web search context to answer the question:\n{ddg_context}\n\n{prompt}"
         
         # Use the constructed prompt for all cases
         await self._rate_limit()