Update gemini_agent.py

gemini_agent.py

@@ -5,6 +5,8 @@ from excel_parser import ExcelParser
 import re
 import time
 import asyncio
+import requests
+import json
 # Add LangChain tools for Wikipedia and DuckDuckGo
 from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
@@ -15,10 +17,14 @@ class GeminiAgent:
     def __init__(self):
         print("GeminiAgent initialized.")
         
-        # Get Google API key from environment variables
+        # Get API keys from environment variables
         api_key = os.getenv('GOOGLE_API_KEY')
         genai.configure(api_key=api_key)
         
+        # Google Custom Search API keys
+        self.google_search_api_key = os.getenv('GOOGLE_SEARCH_API_KEY')
+        self.google_search_cx = os.getenv('GOOGLE_SEARCH_CX')
+        
         self.model = genai.GenerativeModel('gemini-2.0-flash')
         self.last_request_time = 0
         self.min_request_interval = 8.0  # 7 seconds between requests (10 per minute limit, with margin)
@@ -45,9 +51,9 @@ class GeminiAgent:
             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)
+            # Check if question is about music discography or albums
+            if self._is_discography_question(question):
+                return await self._handle_discography_question(question)
             
             # Regular text-based question
             return await self._handle_text_question(question)
@@ -68,133 +74,66 @@ 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"""
+    def _is_discography_question(self, question: str) -> bool:
+        """Determine if a question is about music discography or albums"""
         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"
+            "album", "albums", "discography", "studio album", "published", "released",
+            "recorded", "track", "tracks", "song", "songs", "single", "singles"
         ]
+        artist_patterns = ["musician", "singer", "artist", "band", "composer"]
+        
         # Check for music-related terms
         has_music_term = any(pattern in q for pattern in music_patterns)
+        # Check for artist-related terms
+        has_artist_term = any(pattern in q for pattern in artist_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)
+        has_date_range = re.search(r'between\s+\d{4}\s+and\s+\d{4}', q) is not None or \
+                        re.search(r'from\s+\d{4}\s+to\s+\d{4}', q) is not None or \
+                        re.search(r'\d{4}\s*[-–]\s*\d{4}', q) is not None or \
+                        re.search(r'\d{4}\s+to\s+\d{4}', q) is not None
+        
+        # If it has a music term and either an artist term or a date range, it's likely a discography question
+        return has_music_term and (has_artist_term or has_date_range)
+        
+    async def _google_search(self, query: str, num_results: int = 5) -> str:
+        """Perform a Google search using the Custom Search API"""
+        if not self.google_search_api_key or not self.google_search_cx:
+            print("Google Search API key or CX not configured, falling back to DuckDuckGo")
+            return self.ddg_tool.run(query)
             
-            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.
+            url = "https://www.googleapis.com/customsearch/v1"
+            params = {
+                'key': self.google_search_api_key,
+                'cx': self.google_search_cx,
+                'q': query,
+                'num': num_results
+            }
             
-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)
+            response = requests.get(url, params=params)
+            if response.status_code != 200:
+                print(f"Google Search API error: {response.status_code}")
+                return self.ddg_tool.run(query)  # Fall back to DuckDuckGo
                 
-        return answer
+            results = response.json()
+            if 'items' not in results:
+                print("No search results found")
+                return self.ddg_tool.run(query)  # Fall back to DuckDuckGo
+                
+            # Extract and format search results
+            formatted_results = ""
+            for item in results['items']:
+                title = item.get('title', 'No title')
+                snippet = item.get('snippet', 'No description')
+                link = item.get('link', 'No link')
+                formatted_results += f"Title: {title}\nDescription: {snippet}\nURL: {link}\n\n"
+                
+            return formatted_results
+            
+        except Exception as e:
+            print(f"Google Search API error: {str(e)}")
+            return self.ddg_tool.run(query)  # Fall back to DuckDuckGo
         
     async def _handle_actor_show_question(self, question: str) -> str:
         """Handle questions about actors, TV shows, and movies with enhanced search"""