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 from langgraph.graph import StateGraph, MessagesState, START, END from langgraph.prebuilt import ToolNode from langgraph.checkpoint.memory import MemorySaver import json # 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)}" # Define the agent state class AgentState(TypedDict): messages: Annotated[list, "The messages in the conversation"] 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 tools only if we have Tavily key self.tools = [] if self.has_search: self.search_tool = TavilySearchResults( max_results=5, tavily_api_key=tavily_key ) self.tools = [self.search_tool] # Create LLM with tools (only if we have tools) if self.tools: self.llm_with_tools = self.llm.bind_tools(self.tools) else: self.llm_with_tools = self.llm # Build the graph self.graph = self._build_graph() self.system_prompt = read_system_prompt() def _build_graph(self): """Build the LangGraph workflow""" def agent_node(state: AgentState): """Main agent reasoning node""" messages = state["messages"] # Add system message if not present at the beginning if not any(isinstance(msg, SystemMessage) for msg in messages): system_msg = SystemMessage(content=self.system_prompt) messages = [system_msg] + messages # Get the original question (the first HumanMessage) original_question = None for msg in messages: if isinstance(msg, HumanMessage): original_question = msg.content break # Check if this is a fresh question (not after tool calls) last_msg = messages[-1] is_fresh_question = isinstance(last_msg, HumanMessage) # Only do special processing for fresh questions if is_fresh_question and original_question: # Check if this is a math problem if self._is_math_problem(original_question): try: math_result = math_calculator(original_question) enhanced_msg = f"Question: {original_question}\n\nMath calculation result: {math_result}\n\nBased on this calculation, provide your final answer using the format: FINAL ANSWER: [your answer]" messages[-1] = HumanMessage(content=enhanced_msg) except Exception as e: print(f"Math calculation error: {e}") # Check if this is a date/time problem elif self._is_datetime_problem(original_question): try: datetime_result = date_time_processor(original_question) enhanced_msg = f"Question: {original_question}\n\nDate/time processing result: {datetime_result}\n\nBased on this information, provide your final answer using the format: FINAL ANSWER: [your answer]" messages[-1] = HumanMessage(content=enhanced_msg) except Exception as e: print(f"DateTime processing error: {e}") try: response = self.llm_with_tools.invoke(messages) return {"messages": messages + [response]} except Exception as e: print(f"LLM invocation error: {e}") # Return a simple response on error error_response = HumanMessage(content=f"FINAL ANSWER: Error processing question: {str(e)}") return {"messages": messages + [error_response]} def tool_node(state: AgentState): """Tool execution node""" try: tool_node_instance = ToolNode(self.tools) result = tool_node_instance.invoke(state) return result except Exception as e: print(f"Tool execution error: {e}") # Add an error message and continue messages = state["messages"] error_msg = HumanMessage(content=f"Tool execution failed: {str(e)}. Please provide your best answer without tools.") return {"messages": messages + [error_msg]} def should_continue(state: AgentState): """Decide whether to continue or end""" try: last_message = state["messages"][-1] # If we don't have tools, just end if not self.tools: return "end" # If the last message has tool calls, continue to tools if hasattr(last_message, 'tool_calls') and last_message.tool_calls: return "tools" # If we have a final answer, end if (hasattr(last_message, 'content') and last_message.content and "FINAL ANSWER:" in str(last_message.content)): return "end" # Check if we've had too many iterations (prevent infinite loops) if len(state["messages"]) > 10: return "end" # Otherwise end (be conservative) return "end" except Exception as e: print(f"Should continue error: {e}") return "end" # Build the graph workflow = StateGraph(AgentState) # Add nodes workflow.add_node("agent", agent_node) workflow.add_node("tools", tool_node) # Add edges workflow.add_edge(START, "agent") workflow.add_conditional_edges("agent", should_continue, { "tools": "tools", "end": END }) workflow.add_edge("tools", "agent") # Compile without checkpointer to avoid state issues return workflow.compile() 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" # Create initial state initial_state = { "messages": [HumanMessage(content=question)] } # Run the graph final_state = self.graph.invoke(initial_state) # Extract the final answer last_message = final_state["messages"][-1] response_content = last_message.content if hasattr(last_message, 'content') else str(last_message) # Extract just the final answer part 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 _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()