Seb1101's picture
Update agent.py
4cc5535 verified
raw
history blame
13.5 kB
import os
import re
from datetime import datetime, timedelta
from typing import TypedDict, Annotated
import sympy as sp
import math
from langchain_openai import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import HumanMessage, SystemMessage
# Load environment variables
from dotenv import load_dotenv
load_dotenv()
def read_system_prompt():
"""Read the system prompt from file"""
try:
with open('system_prompt.txt', 'r') as f:
return f.read().strip()
except FileNotFoundError:
return """You are a helpful assistant tasked with answering questions using a set of tools.
Now, I will ask you a question. Report your thoughts, and finish your answer with the following template:
FINAL ANSWER: [YOUR FINAL ANSWER].
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Your answer should only start with "FINAL ANSWER: ", then follows with the answer."""
def math_calculator(expression: str) -> str:
"""
Advanced mathematical calculator that can handle complex expressions,
equations, symbolic math, calculus, and more using SymPy.
"""
try:
# Clean the expression
expression = expression.strip()
# Handle common mathematical operations and functions
expression = expression.replace('^', '**') # Convert ^ to **
expression = expression.replace('ln', 'log') # Natural log
# Try to evaluate as a symbolic expression first
try:
result = sp.sympify(expression)
# If it's a symbolic expression that can be simplified
simplified = sp.simplify(result)
# Try to get numerical value
try:
numerical = float(simplified.evalf())
return str(numerical)
except:
return str(simplified)
except:
# Fall back to basic evaluation
# Replace common math functions
safe_expression = expression
for func in ['sin', 'cos', 'tan', 'sqrt', 'log', 'exp', 'abs']:
safe_expression = safe_expression.replace(func, f'math.{func}')
# Evaluate safely
result = eval(safe_expression, {"__builtins__": {}}, {
"math": math,
"pi": math.pi,
"e": math.e
})
return str(result)
except Exception as e:
return f"Error calculating '{expression}': {str(e)}"
def date_time_processor(query: str) -> str:
"""
Process date and time related queries, calculations, and conversions.
"""
try:
current_time = datetime.now()
query_lower = query.lower()
# Current date/time queries
if 'current' in query_lower or 'today' in query_lower or 'now' in query_lower:
if 'date' in query_lower:
return current_time.strftime('%Y-%m-%d')
elif 'time' in query_lower:
return current_time.strftime('%H:%M:%S')
else:
return current_time.strftime('%Y-%m-%d %H:%M:%S')
# Day of week queries
if 'day of week' in query_lower or 'what day' in query_lower:
return current_time.strftime('%A')
# Year queries
if 'year' in query_lower and 'current' in query_lower:
return str(current_time.year)
# Month queries
if 'month' in query_lower and 'current' in query_lower:
return current_time.strftime('%B')
# Date arithmetic (simple cases)
if 'days ago' in query_lower:
days_match = re.search(r'(\d+)\s+days?\s+ago', query_lower)
if days_match:
days = int(days_match.group(1))
past_date = current_time - timedelta(days=days)
return past_date.strftime('%Y-%m-%d')
if 'days from now' in query_lower or 'days later' in query_lower:
days_match = re.search(r'(\d+)\s+days?\s+(?:from now|later)', query_lower)
if days_match:
days = int(days_match.group(1))
future_date = current_time + timedelta(days=days)
return future_date.strftime('%Y-%m-%d')
# If no specific pattern matched, return current datetime
return f"Current date and time: {current_time.strftime('%Y-%m-%d %H:%M:%S')}"
except Exception as e:
return f"Error processing date/time query: {str(e)}"
# Removed LangGraph dependencies - using simpler approach
class GAIAAgent:
def __init__(self):
# Check for required API keys
openai_key = os.getenv("OPENAI_API_KEY")
tavily_key = os.getenv("TAVILY_API_KEY")
if not openai_key:
raise ValueError("OPENAI_API_KEY environment variable is required")
if not tavily_key:
print("⚠️ TAVILY_API_KEY not found - web search will be disabled")
self.has_search = False
else:
self.has_search = True
print("✅ Initializing GAIA agent...")
# Initialize LLM (using OpenAI GPT-4)
self.llm = ChatOpenAI(
model="gpt-4o-mini",
temperature=0,
openai_api_key=openai_key
)
# Initialize search tool if available
if self.has_search:
self.search_tool = TavilySearchResults(
max_results=5,
tavily_api_key=tavily_key
)
else:
self.search_tool = None
self.system_prompt = read_system_prompt()
def _search_web(self, query: str) -> str:
"""Perform web search if available"""
if not self.search_tool:
return "Web search not available (no Tavily API key)"
try:
results = self.search_tool.invoke({"query": query})
if results and len(results) > 0:
# Format the results nicely
formatted_results = []
for i, result in enumerate(results[:3], 1): # Top 3 results
if isinstance(result, dict):
title = result.get('title', 'No title')
content = result.get('content', 'No content')
url = result.get('url', 'No URL')
formatted_results.append(f"{i}. {title}\n {content}\n Source: {url}")
else:
formatted_results.append(f"{i}. {str(result)}")
return "\n\n".join(formatted_results)
else:
return "No search results found"
except Exception as e:
return f"Search error: {str(e)}"
def _is_math_problem(self, text: str) -> bool:
"""Check if the text contains mathematical expressions"""
math_indicators = [
'+', '-', '*', '/', '^', '=', 'calculate', 'compute',
'solve', 'equation', 'integral', 'derivative', 'sum',
'sqrt', 'log', 'sin', 'cos', 'tan', 'exp'
]
text_lower = text.lower()
return any(indicator in text_lower for indicator in math_indicators) or \
re.search(r'\d+[\+\-\*/\^]\d+', text) is not None
def _is_datetime_problem(self, text: str) -> bool:
"""Check if the text contains date/time related queries"""
datetime_indicators = [
'date', 'time', 'day', 'month', 'year', 'today', 'yesterday',
'tomorrow', 'current', 'now', 'ago', 'later', 'when'
]
text_lower = text.lower()
return any(indicator in text_lower for indicator in datetime_indicators)
def __call__(self, question: str) -> str:
"""Process a question and return the answer"""
try:
print(f"Processing question: {question[:100]}...")
# Check for file/media requirements that we can't handle
if any(indicator in question.lower() for indicator in [
'attached', 'audio', 'video', 'image', 'file', 'mp3', 'pdf',
'excel', 'spreadsheet', 'listen to', 'watch', 'download'
]):
return "Unable to process files or media attachments"
# Build the prompt based on question type
enhanced_question = self._enhance_question(question)
# Create messages
messages = [
SystemMessage(content=self.system_prompt),
HumanMessage(content=enhanced_question)
]
# Get response from LLM
response = self.llm.invoke(messages)
response_content = response.content if hasattr(response, 'content') else str(response)
# Extract the final answer
final_answer = self._extract_final_answer(response_content)
print(f"Final answer: {final_answer}")
return final_answer
except Exception as e:
print(f"Error processing question: {e}")
# Try to provide a meaningful fallback
if "api" in str(e).lower() or "key" in str(e).lower():
return "Error: API key configuration issue"
elif "tool" in str(e).lower():
return "Error: Tool execution issue"
else:
return f"Unable to process question due to technical error"
def _enhance_question(self, question: str) -> str:
"""Enhance the question with relevant context and tools"""
try:
# Check if this is a math problem
if self._is_math_problem(question):
try:
math_result = math_calculator(question)
return f"Question: {question}\n\nMath calculation result: {math_result}\n\nBased on this calculation, provide your final answer using the format: FINAL ANSWER: [your answer]"
except Exception as e:
print(f"Math calculation error: {e}")
# Check if this is a date/time problem
elif self._is_datetime_problem(question):
try:
datetime_result = date_time_processor(question)
return f"Question: {question}\n\nDate/time processing result: {datetime_result}\n\nBased on this information, provide your final answer using the format: FINAL ANSWER: [your answer]"
except Exception as e:
print(f"DateTime processing error: {e}")
# Check if this needs web search
elif self._needs_web_search(question):
try:
search_result = self._search_web(question)
return f"Question: {question}\n\nWeb search results:\n{search_result}\n\nBased on this information, provide your final answer using the format: FINAL ANSWER: [your answer]"
except Exception as e:
print(f"Web search error: {e}")
# For other questions, just add the format instruction
return f"Question: {question}\n\nProvide your final answer using the format: FINAL ANSWER: [your answer]"
except Exception as e:
print(f"Question enhancement error: {e}")
return f"Question: {question}\n\nProvide your final answer using the format: FINAL ANSWER: [your answer]"
def _needs_web_search(self, text: str) -> bool:
"""Check if the question likely needs web search"""
search_indicators = [
'who', 'what', 'when', 'where', 'which', 'published', 'article',
'wikipedia', 'latest', 'recent', 'current', 'news', 'website',
'url', 'http', 'www', 'competition', 'olympics', 'award',
'winner', 'recipient', 'author', 'published in', 'paper',
'study', 'research', 'species', 'city', 'country'
]
text_lower = text.lower()
return any(indicator in text_lower for indicator in search_indicators)
def _extract_final_answer(self, response: str) -> str:
"""Extract the final answer from the response"""
if "FINAL ANSWER:" in response:
# Find the final answer part
parts = response.split("FINAL ANSWER:")
if len(parts) > 1:
answer = parts[-1].strip()
# Remove any trailing punctuation or explanations
answer = answer.split('\n')[0].strip()
return answer
# If no FINAL ANSWER format found, return the whole response
return response.strip()
# Create a function to get the agent (for use in app.py)
def create_agent():
"""Factory function to create the GAIA agent"""
return GAIAAgent()