Update gemini_agent.py

gemini_agent.py

@@ -44,6 +44,10 @@ class GeminiAgent:
             # 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)
+                
+            # Check if question is about music, albums, or discography
+            if self._is_music_question(question):
+                return await self._handle_music_question(question)
             
             # Regular text-based question
             return await self._handle_text_question(question)
@@ -64,6 +68,134 @@ class GeminiAgent:
         ]
         return any(pattern in q for pattern in actor_show_patterns)
         
+    def _is_music_question(self, question: str) -> bool:
+        """Determine if a question is about music, albums, or discography"""
+        q = question.lower()
+        music_patterns = [
+            "album", "albums", "discography", "song", "songs", "track", "tracks",
+            "musician", "singer", "artist", "band", "composer", "recorded", "released",
+            "studio album", "live album", "compilation", "single", "ep"
+        ]
+        # Check for music-related terms
+        has_music_term = any(pattern in q for pattern in music_patterns)
+        # Check for date ranges which are common in discography questions
+        has_date_range = re.search(r'between\s+\d{4}\s+and\s+\d{4}', q) is not None
+        return has_music_term or has_date_range
+        
+    async def _handle_music_question(self, question: str) -> str:
+        """Handle questions about music, albums, and discographies with enhanced search"""
+        print(f"Processing music/discography question: {question[:50]}...")
+        
+        # Always try both Wikipedia and DuckDuckGo for these questions
+        wiki_context = ""
+        ddg_context = ""
+        
+        # Extract date range if present
+        date_range = re.search(r'between\s+(\d{4})\s+and\s+(\d{4})', question.lower())
+        start_year = None
+        end_year = None
+        if date_range:
+            start_year = int(date_range.group(1))
+            end_year = int(date_range.group(2))
+        else:
+            # Check for other date range formats
+            date_range = re.search(r'from\s+(\d{4})\s+to\s+(\d{4})', question.lower())
+            if date_range:
+                start_year = int(date_range.group(1))
+                end_year = int(date_range.group(2))
+            else:
+                # Check for single year with inclusion indicator
+                included_year = re.search(r'(\d{4})\s*\(included\)', question.lower())
+                if included_year:
+                    end_year = int(included_year.group(1))
+        
+        # Specifically mention Wikipedia version if specified
+        wiki_version = None
+        if "version of" in question.lower() and "wikipedia" in question.lower():
+            version_match = re.search(r'(\d{4})\s+version\s+of\s+\w+\s+wikipedia', question.lower())
+            if version_match:
+                wiki_version = version_match.group(1)
+                print(f"Using Wikipedia version from {wiki_version}")
+        
+        # Construct a more specific search query for discography questions
+        search_query = question
+        if "how many" in question.lower() and "album" in question.lower():
+            # Extract artist name - look for patterns like "by [Artist Name]" or just the name before "albums"
+            artist_match = re.search(r'by\s+([\w\s]+)\s+between', question)
+            if not artist_match:
+                artist_match = re.search(r'([\w\s]+)\s+albums', question)
+            
+            if artist_match:
+                artist_name = artist_match.group(1).strip()
+                search_query = f"{artist_name} complete discography studio albums {start_year if start_year else ''} to {end_year if end_year else ''}"
+                print(f"Enhanced search query: {search_query}")
+        
+        try:
+            wiki_context = self.wiki_tool.run(search_query)
+            print("Wikipedia search completed")
+        except Exception as e:
+            print(f"Wikipedia tool failed: {e}")
+            
+        try:
+            ddg_context = self.ddg_tool.run(search_query)
+            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 music/discography questions
+        prompt = f"""Based on the following context, answer this question about music or discography:
+
+Question: {question}
+
+{combined_context}
+
+"""
+        
+        # Add specific instructions for counting albums in a date range
+        if "how many" in question.lower() and "album" in question.lower() and start_year and end_year:
+            prompt += f"""Count ONLY the studio albums released between {start_year} and {end_year}, inclusive of both years.
+            
+Provide ONLY the numeric count as your answer, with no additional text.
+            
+Make sure to count each album only once, and only count studio albums (not compilations, live albums, or EPs) unless specifically asked for those."""
+        else:
+            prompt += "Provide ONLY the specific information requested. No explanations or additional context."
+        
+        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 number or information
+        # Remove common prefixes
+        prefixes = ['The answer is', 'Based on', 'According to', 'The number is', 'There were']
+        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 is asking for a count, extract just the number
+        if "how many" in question.lower():
+            # Try to extract just the number
+            number_match = re.search(r'\b(\d+)\b', answer)
+            if number_match:
+                answer = number_match.group(1)
+                
+        return answer
+        
     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]}...")