Update agent.py

agent.py

@@ -1,231 +1,353 @@
 import requests
 import json
-from transformers import Tool
-from huggingface_hub import login
-import os
 import re
-from typing import Dict, Any
+import os
+import math
+from typing import Dict, Any, List, Union
+from datetime import datetime, timedelta
+import urllib.parse
+from transformers import Tool
 
-class WikipediaSearchTool(Tool):
-    name = "wikipedia_search"
-    description = "Search Wikipedia for information about a specific topic"
+class AdvancedSearchTool(Tool):
+    name = "advanced_search"
+    description = "Advanced search tool for Wikipedia and web content"
     inputs = {
-        "query": {
-            "type": "text", 
-            "description": "The search query for Wikipedia"
-        }
+        "query": {"type": "text", "description": "The search query"},
+        "search_type": {"type": "text", "description": "Type of search: 'wikipedia', 'general'"}
     }
     output_type = "text"
 
-    def forward(self, query: str) -> str:
-        """Search Wikipedia using the API"""
+    def forward(self, query: str, search_type: str = "wikipedia") -> str:
+        try:
+            if search_type == "wikipedia":
+                return self._search_wikipedia(query)
+            else:
+                return self._search_wikipedia(query)  # Fallback to Wikipedia
+        except Exception as e:
+            return f"Search error: {str(e)}"
+
+    def _search_wikipedia(self, query: str) -> str:
         try:
-            # Use Wikipedia API to search
-            search_url = "https://en.wikipedia.org/api/rest_v1/page/summary/"
-            # Clean the query
+            # Try direct page lookup first
             clean_query = query.replace(" ", "_")
+            summary_url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{clean_query}"
             
-            response = requests.get(f"{search_url}{clean_query}")
+            response = requests.get(summary_url, timeout=10)
             if response.status_code == 200:
                 data = response.json()
-                return f"Title: {data.get('title', '')}\nSummary: {data.get('extract', '')}"
-            else:
-                # Try search API if direct lookup fails
-                search_api_url = "https://en.wikipedia.org/w/api.php"
-                params = {
-                    'action': 'query',
-                    'format': 'json',
-                    'list': 'search',
-                    'srsearch': query,
-                    'srlimit': 3
-                }
+                extract = data.get('extract', '')
+                if extract and len(extract) > 50:
+                    return f"Title: {data.get('title', '')}\nContent: {extract}"
+            
+            # Search API if direct lookup fails
+            search_url = "https://en.wikipedia.org/w/api.php"
+            search_params = {
+                'action': 'query',
+                'format': 'json',
+                'list': 'search',
+                'srsearch': query,
+                'srlimit': 5
+            }
+            
+            search_response = requests.get(search_url, params=search_params, timeout=10)
+            if search_response.status_code == 200:
+                search_data = search_response.json()
+                results = []
                 
-                search_response = requests.get(search_api_url, params=params)
-                if search_response.status_code == 200:
-                    search_data = search_response.json()
-                    results = []
-                    for result in search_data['query']['search'][:2]:
-                        title = result['title']
-                        snippet = result['snippet']
-                        results.append(f"Title: {title}\nSnippet: {snippet}")
-                    return "\n\n".join(results)
+                for result in search_data['query']['search'][:3]:
+                    title = result['title']
+                    # Get page content
+                    page_params = {
+                        'action': 'query',
+                        'format': 'json',
+                        'titles': title,
+                        'prop': 'extracts',
+                        'exintro': True,
+                        'explaintext': True,
+                        'exsectionformat': 'plain'
+                    }
+                    
+                    page_response = requests.get(search_url, params=page_params, timeout=10)
+                    if page_response.status_code == 200:
+                        page_data = page_response.json()
+                        pages = page_data.get('query', {}).get('pages', {})
+                        for page_id, page_info in pages.items():
+                            extract = page_info.get('extract', '')
+                            if extract:
+                                results.append(f"Title: {title}\nContent: {extract[:1000]}")
+                                break
                 
-                return f"No information found for query: {query}"
+                return "\n\n".join(results) if results else f"No detailed results found for: {query}"
+            
+            return f"No Wikipedia results found for: {query}"
+            
         except Exception as e:
-            return f"Error searching Wikipedia: {str(e)}"
+            return f"Wikipedia search error: {str(e)}"
 
-class WebSearchTool(Tool):
-    name = "web_search"
-    description = "Search the web for current information"
-    inputs = {
-        "query": {
-            "type": "text",
-            "description": "The search query"
-        }
-    }
-    output_type = "text"
-
-    def forward(self, query: str) -> str:
-        """Search the web using a search API or fallback method"""
+class MathCalculator:
+    @staticmethod
+    def evaluate_expression(expression: str) -> Union[float, int, str]:
+        """Safely evaluate mathematical expressions"""
         try:
-            # For this implementation, we'll focus on Wikipedia since the question specifically mentions it
-            wiki_tool = WikipediaSearchTool()
-            return wiki_tool.forward(query)
-        except Exception as e:
-            return f"Error in web search: {str(e)}"
+            # Clean the expression
+            expression = re.sub(r'[^\d\+\-\*\/\.\(\)\s]', '', expression)
+            if not expression.strip():
+                return "Invalid expression"
+            
+            # Use eval cautiously with limited scope
+            result = eval(expression, {"__builtins__": {}}, {
+                "abs": abs, "round": round, "min": min, "max": max,
+                "sum": sum, "len": len, "pow": pow, "sqrt": math.sqrt,
+                "sin": math.sin, "cos": math.cos, "tan": math.tan,
+                "log": math.log, "exp": math.exp, "pi": math.pi, "e": math.e
+            })
+            
+            # Return integer if it's a whole number
+            if isinstance(result, float) and result.is_integer():
+                return int(result)
+            return result
+        except:
+            return "Calculation error"
 
-class MusicAgent:
+class ComprehensiveAgent:
     def __init__(self):
-        self.tools = {
-            "wikipedia_search": WikipediaSearchTool(),
-            "web_search": WebSearchTool()
-        }
-    
-    def extract_years_from_text(self, text: str, start_year: int, end_year: int) -> list:
-        """Extract years within the specified range from text"""
-        year_pattern = r'\b(19|20)\d{2}\b'
-        years = re.findall(year_pattern, text)
-        valid_years = []
-        for match in re.finditer(year_pattern, text):
+        self.search_tool = AdvancedSearchTool()
+        self.calculator = MathCalculator()
+        
+    def extract_numbers(self, text: str) -> List[Union[int, float]]:
+        """Extract numbers from text"""
+        numbers = []
+        # Find integers and floats
+        for match in re.finditer(r'\b\d+(?:\.\d+)?\b', text):
+            try:
+                num_str = match.group()
+                if '.' in num_str:
+                    numbers.append(float(num_str))
+                else:
+                    numbers.append(int(num_str))
+            except:
+                continue
+        return numbers
+
+    def extract_years(self, text: str, start_year: int = 1900, end_year: int = 2025) -> List[int]:
+        """Extract years within a reasonable range"""
+        years = []
+        for match in re.finditer(r'\b(19|20)\d{2}\b', text):
             year = int(match.group())
             if start_year <= year <= end_year:
-                valid_years.append(year)
-        return valid_years
-    
-    def extract_albums_from_text(self, text: str) -> list:
-        """Extract album information from text"""
-        albums = []
-        # Look for common album indicators
-        album_patterns = [
-            r'album[s]?\s*[":]\s*([^,\n\.]+)',
-            r'released\s+([^,\n\.]+?)\s+in\s+(\d{4})',
-            r'(\d{4})[:\s]+([^,\n\.]+)',
-            r'"([^"]+)"\s*\((\d{4})\)',
+                years.append(year)
+        return list(set(years))  # Remove duplicates
+
+    def answer_counting_question(self, question: str, context: str) -> str:
+        """Handle questions that ask 'how many'"""
+        question_lower = question.lower()
+        
+        # Extract what we're counting
+        counting_patterns = [
+            r'how many (.*?) (?:were|are|did|have|has)',
+            r'how many (.*?)(?:\?|$)',
+            r'number of (.*?)(?:\?|$)'
         ]
         
-        for pattern in album_patterns:
-            matches = re.findall(pattern, text, re.IGNORECASE)
-            albums.extend(matches)
+        counting_target = ""
+        for pattern in counting_patterns:
+            match = re.search(pattern, question_lower)
+            if match:
+                counting_target = match.group(1).strip()
+                break
         
-        return albums
-    
-    def count_studio_albums(self, artist_name: str, start_year: int, end_year: int) -> int:
-        """Count studio albums for an artist within a year range"""
-        try:
-            # Search for the artist's discography
-            discography_queries = [
-                f"{artist_name} discography",
-                f"{artist_name} studio albums",
-                f"{artist_name} albums {start_year}-{end_year}",
-                f"{artist_name} complete discography"
-            ]
-            
-            all_text = ""
-            for query in discography_queries:
-                try:
-                    result = self.tools["wikipedia_search"].forward(query)
-                    all_text += result + "\n"
-                except:
-                    continue
-            
-            if not all_text.strip():
-                return 0
+        if not counting_target:
+            return "Could not identify counting target"
+        
+        # Look for numbers in context that might be the answer
+        numbers = self.extract_numbers(context)
+        
+        # Special handling for common counting scenarios
+        if "albums" in counting_target:
+            return self._count_albums(question, context)
+        elif "years" in counting_target or "year" in counting_target:
+            years = self.extract_years(context)
+            if years:
+                return str(len(years))
+        elif "countries" in counting_target or "states" in counting_target:
+            # Look for country/state names or numbers
+            if numbers:
+                return str(numbers[0])
+        
+        # Default: return first reasonable number found
+        reasonable_numbers = [n for n in numbers if 0 <= n <= 1000]
+        if reasonable_numbers:
+            return str(reasonable_numbers[0])
+        
+        return "Unable to determine count"
+
+    def _count_albums(self, question: str, context: str) -> str:
+        """Specifically handle album counting questions"""
+        # Extract years from question
+        years_in_question = self.extract_years(question)
+        
+        if len(years_in_question) >= 2:
+            start_year = min(years_in_question)
+            end_year = max(years_in_question)
             
-            # Count albums within the year range
-            # Look for year patterns and album mentions
-            year_pattern = r'\b(19|20)\d{2}\b'
-            years_in_text = re.findall(year_pattern, all_text)
+            # Count years in the context that fall within range
+            context_years = self.extract_years(context, start_year, end_year)
             
-            # Simple heuristic: count unique years in range that likely represent album releases
-            valid_years = set()
-            for year_match in re.finditer(year_pattern, all_text):
-                year = int(year_match.group())
-                if start_year <= year <= end_year:
-                    # Check if this year is associated with album context
-                    context_start = max(0, year_match.start() - 100)
-                    context_end = min(len(all_text), year_match.end() + 100)
-                    context = all_text[context_start:context_end].lower()
+            # Look for album-related keywords near years
+            album_count = 0
+            for year in context_years:
+                year_str = str(year)
+                year_pos = context.lower().find(year_str)
+                if year_pos != -1:
+                    # Check surrounding context for album keywords
+                    start_context = max(0, year_pos - 200)
+                    end_context = min(len(context), year_pos + 200)
+                    surrounding = context[start_context:end_context].lower()
                     
-                    album_keywords = ['album', 'studio', 'released', 'record', 'disc']
-                    if any(keyword in context for keyword in album_keywords):
-                        valid_years.add(year)
-            
-            return len(valid_years)
+                    if any(word in surrounding for word in ['album', 'studio', 'released', 'record']):
+                        album_count += 1
             
-        except Exception as e:
-            print(f"Error counting albums: {str(e)}")
-            return 0
-    
+            return str(album_count) if album_count > 0 else "1"
+        
+        # Fallback: look for explicit numbers
+        numbers = self.extract_numbers(context)
+        small_numbers = [n for n in numbers if 0 <= n <= 50]
+        return str(small_numbers[0]) if small_numbers else "0"
+
+    def answer_calculation_question(self, question: str) -> str:
+        """Handle mathematical calculation questions"""
+        # Extract mathematical expressions
+        math_patterns = [
+            r'(\d+(?:\.\d+)?)\s*[\+\-\*\/]\s*(\d+(?:\.\d+)?)',
+            r'what is\s+(.+?)(?:\?|$)',
+            r'calculate\s+(.+?)(?:\?|$)'
+        ]
+        
+        for pattern in math_patterns:
+            match = re.search(pattern, question.lower())
+            if match:
+                expression = match.group(1) if len(match.groups()) == 1 else f"{match.group(1)} {match.group(2)}"
+                result = self.calculator.evaluate_expression(expression)
+                if result != "Calculation error":
+                    return str(result)
+        
+        return "Could not parse mathematical expression"
+
+    def answer_factual_question(self, question: str) -> str:
+        """Handle general factual questions"""
+        # Search for information
+        search_result = self.search_tool.forward(question, "wikipedia")
+        
+        if "error" in search_result.lower():
+            return "Information not available"
+        
+        # Extract potential answers based on question type
+        question_lower = question.lower()
+        
+        if question_lower.startswith("when"):
+            # Look for years or dates
+            years = self.extract_years(search_result)
+            if years:
+                return str(years[0])
+        
+        elif question_lower.startswith("where"):
+            # Look for place names (simplified)
+            sentences = search_result.split('.')
+            for sentence in sentences[:3]:
+                if any(word in sentence.lower() for word in ['located', 'in', 'at', 'city', 'country']):
+                    return sentence.strip()[:100]
+        
+        elif question_lower.startswith("who"):
+            # Return first meaningful sentence
+            sentences = search_result.split('.')
+            for sentence in sentences[:2]:
+                if len(sentence.strip()) > 20:
+                    return sentence.strip()[:100]
+        
+        elif question_lower.startswith("what"):
+            # Return definition or explanation
+            sentences = search_result.split('.')
+            for sentence in sentences[:2]:
+                if len(sentence.strip()) > 30:
+                    return sentence.strip()[:150]
+        
+        # Default: return first substantial sentence
+        sentences = search_result.split('.')
+        for sentence in sentences[:3]:
+            if len(sentence.strip()) > 20:
+                return sentence.strip()[:100]
+        
+        return "Answer not found"
+
     def answer_question(self, question: str) -> str:
-        """Answer a question using available tools"""
+        """Main method to answer various types of questions"""
         try:
+            question = question.strip()
             question_lower = question.lower()
             
-            # Check if this is about Mercedes Sosa albums
-            if "mercedes sosa" in question_lower and "studio albums" in question_lower:
-                # Extract year range from question
-                year_matches = re.findall(r'\b(19|20)\d{2}\b', question)
-                if len(year_matches) >= 2:
-                    start_year = int(year_matches[0])
-                    end_year = int(year_matches[1])
-                else:
-                    start_year = 2000
-                    end_year = 2009
-                
-                count = self.count_studio_albums("Mercedes Sosa", start_year, end_year)
-                
-                # If we got 0, try alternative searches
-                if count == 0:
-                    # Try more specific searches
-                    specific_queries = [
-                        "Mercedes Sosa 2000 2001 2002 2003 2004 2005 2006 2007 2008 2009 albums",
-                        "Mercedes Sosa studio albums 2000s"
-                    ]
-                    
-                    for query in specific_queries:
-                        try:
-                            result = self.tools["wikipedia_search"].forward(query)
-                            # Manual check for known albums in that period
-                            if any(year in result for year in ["2000", "2001", "2002", "2003", "2004", "2005", "2006", "2007", "2008", "2009"]):
-                                # Mercedes Sosa had limited studio album releases in 2000-2009
-                                # Based on typical discography patterns, estimate
-                                return "2"
-                        except:
-                            continue
-                
-                return str(max(count, 1))  # Ensure at least 1 if we found some evidence
+            # Handle different question types
+            if question_lower.startswith("how many"):
+                # Get relevant context first
+                search_context = self.search_tool.forward(question, "wikipedia")
+                return self.answer_counting_question(question, search_context)
             
-            # For other questions, use general search
-            search_result = self.tools["wikipedia_search"].forward(question)
+            elif any(op in question for op in ['+', '-', '*', '/', 'calculate', 'what is']):
+                return self.answer_calculation_question(question)
             
-            # Try to extract a simple answer
-            if "how many" in question_lower:
-                numbers = re.findall(r'\b\d+\b', search_result)
-                if numbers:
-                    return numbers[0]
+            elif question_lower.startswith(("when", "where", "who", "what", "which")):
+                return self.answer_factual_question(question)
             
-            # Return first meaningful sentence
-            sentences = search_result.split('.')
-            for sentence in sentences[:3]:
-                if len(sentence.strip()) > 10:
-                    return sentence.strip()
-            
-            return "Unable to determine answer from available information"
+            elif "year" in question_lower and "born" in question_lower:
+                search_result = self.search_tool.forward(question, "wikipedia")
+                years = self.extract_years(search_result)
+                return str(years[0]) if years else "Year not found"
             
+            else:
+                # General question handling
+                return self.answer_factual_question(question)
+                
         except Exception as e:
-            print(f"Error answering question: {str(e)}")
-            return "Error processing question"
+            print(f"Error processing question: {str(e)}")
+            return "Processing error"
 
-# Initialize the agent
-agent = MusicAgent()
+# Initialize the comprehensive agent
+agent = ComprehensiveAgent()
 
 def answer_question(question: str) -> str:
     """Main function to answer questions"""
-    return agent.answer_question(question)
+    try:
+        result = agent.answer_question(question)
+        # Ensure result is clean and concise
+        if isinstance(result, str):
+            result = result.strip()
+            # Remove common prefixes that might interfere with exact matching
+            prefixes_to_remove = [
+                "the answer is ", "answer: ", "result: ", "final answer: ",
+                "title: ", "content: "
+            ]
+            result_lower = result.lower()
+            for prefix in prefixes_to_remove:
+                if result_lower.startswith(prefix):
+                    result = result[len(prefix):].strip()
+                    break
+        
+        return result
+    except Exception as e:
+        print(f"Error in answer_question: {str(e)}")
+        return "Error processing question"
 
-# Test the specific question
+# Test with various question types
 if __name__ == "__main__":
-    test_question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
-    result = answer_question(test_question)
-    print(f"Question: {test_question}")
-    print(f"Answer: {result}")
+    test_questions = [
+        "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia.",
+        "What is 15 + 27?",
+        "When was Albert Einstein born?",
+        "Where is the Eiffel Tower located?",
+        "How many continents are there?"
+    ]
+    
+    for question in test_questions:
+        result = answer_question(question)
+        print(f"Q: {question}")
+        print(f"A: {result}")
+        print("-" * 50)