# Standard libraries import json import os from dotenv import load_dotenv from typing import Dict, List, Any, Optional, Annotated from typing_extensions import TypedDict # Langchain and langgraph from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage, AnyMessage from langgraph.graph import StateGraph, START, END, add_messages from langgraph.prebuilt import ToolNode, tools_condition # Custom modules from prompts import MAIN_SYSTEM_PROMPT, QUESTION_DECOMPOSITION_PROMPT, TOOL_USE_INSTRUCTION, EXECUTION_INSTRUCTION from utils import check_api_keys, setup_llm from tools import ( calculator_tool, extract_text_from_image, transcribe_audio, execute_python_code, read_file, web_search, wikipedia_search, arxiv_search, chess_board_image_analysis, find_phrase_in_text, download_youtube_audio, web_content_extract, analyse_tabular_data ) # AGENT STATE class AgentState(TypedDict): task_id: Optional[str] file_name: Optional[str] file_type: Optional[str] file_path: Optional[str] question_decomposition: Optional[str] messages: Annotated[list[AnyMessage], add_messages] tool_results: Dict error_message: Optional[str] # WORKFLOW CREATION def create_workflow_for_final_agent(): """ Creates and compiles the LangGraph workflow. """ llm_agent_management, llm_question_decomposition, _, _, _ = setup_llm() tools = [ web_search, web_content_extract, wikipedia_search, calculator_tool, extract_text_from_image, transcribe_audio, execute_python_code, read_file, arxiv_search, chess_board_image_analysis, find_phrase_in_text, download_youtube_audio, analyse_tabular_data ] llm_agent_management_with_tools = llm_agent_management.bind_tools(tools) # Define nodes def question_decomposition_node(state: AgentState): new_state = state.copy() messages = new_state.get("messages", []) # Use .get for safety, ensure it's a list question = None for msg in messages: if isinstance(msg, HumanMessage): question = msg.content break if not question: new_state["error_message"] = "No question found for decomposition." # Ensure messages list exists even if we return early if "messages" not in new_state or not isinstance(new_state["messages"], list): new_state["messages"] = [] return new_state question_decomposition_prompt_messages = [ SystemMessage(content=QUESTION_DECOMPOSITION_PROMPT), HumanMessage(content=f"Decompose this question: {question}") ] question_decomposition_object = llm_question_decomposition.invoke(question_decomposition_prompt_messages) question_decomposition_response = question_decomposition_object.content new_state["question_decomposition"] = question_decomposition_response # Ensure messages list exists if "messages" not in new_state or not isinstance(new_state["messages"], list): new_state["messages"] = [] return new_state def call_model_node(state: AgentState): new_state = state.copy() messages = new_state.get("messages", []) # Use .get for safety question_decomposition = new_state.get("question_decomposition", "") llm_messages = list(messages) # Ensure it's a mutable list add_decomposition = question_decomposition and (not llm_messages or not isinstance(llm_messages[-1], ToolMessage)) if add_decomposition: decomposition_message = SystemMessage(content=f"Question decomposition: {question_decomposition}\\nUse this analysis to guide your actions.") llm_messages.append(decomposition_message) response = llm_agent_management_with_tools.invoke(llm_messages) # Ensure new_state["messages"] exists and is a list before extending current_messages = new_state.get("messages", []) if not isinstance(current_messages, list): current_messages = [] new_state["messages"] = current_messages + [response] return new_state workflow = StateGraph(AgentState) workflow.add_node("decomposition", question_decomposition_node) workflow.add_node("agent", call_model_node) workflow.add_node("tools", ToolNode(tools)) workflow.add_edge(START, "decomposition") workflow.add_edge("decomposition", "agent") workflow.add_conditional_edges("agent", tools_condition) workflow.add_edge("tools", "agent") return workflow.compile() class FinalAgent: def __init__(self): print("FinalAgent initializing...") load_dotenv() if not os.path.exists('.config'): print("Warning: .config file not found. Using default values or expecting environment variables.") self.config = {} # Default to empty config else: with open('.config', 'r') as f: self.config = json.load(f) self.base_url = self.config.get('BASE_URL', os.getenv('BASE_URL')) self.debug_mode = self.config.get('DEBUG_MODE', str(os.getenv('DEBUG_MODE', 'False')).lower() == 'true') if not check_api_keys(): # check_api_keys itself prints messages raise ValueError("API keys are missing or invalid. Please set the required environment variables.") self.workflow = create_workflow_for_final_agent() print("FinalAgent initialized successfully.") def __call__(self, question: str, task_id: Optional[str] = None) -> str: print(f"FinalAgent received question for task_id '{task_id}': {question[:100]}...") initial_messages = [ SystemMessage(content=MAIN_SYSTEM_PROMPT + "\\n\\n" + TOOL_USE_INSTRUCTION + "\\n\\n" + EXECUTION_INSTRUCTION), HumanMessage(content=question) ] initial_state: AgentState = { "messages": initial_messages, "task_id": task_id, "file_name": None, "file_path": None, "file_type": None, "question_decomposition": None, "tool_results": {}, "error_message": None } try: result_state = self.workflow.invoke(initial_state) except Exception as e: print(f"Error invoking workflow for task {task_id}: {e}") import traceback traceback.print_exc() return f"AGENT ERROR: Failed to process question due to an internal error: {e}" messages = result_state.get("messages", []) final_answer = "" if not messages: print(f"No messages found in the result state for task {task_id}.") return "AGENT ERROR: No messages returned by the agent." for msg in reversed(messages): if hasattr(msg, "content") and msg.content: content = msg.content if isinstance(content, str): if "FINAL ANSWER:" in content: final_answer = content.split("FINAL ANSWER:", 1)[1].strip() break elif isinstance(msg, AIMessage): # If it's an AIMessage and no "FINAL ANSWER:" has been found yet, # tentatively set it. This will be overridden if a "FINAL ANSWER:" is found later. if not final_answer: final_answer = content # If after checking all messages, final_answer is still from a non-"FINAL ANSWER:" AIMessage, that's our best guess. # If final_answer is empty, it means no AIMessage with content or "FINAL ANSWER:" was found. if not final_answer: # This means no "FINAL ANSWER:" and no AIMessage content was suitable final_answer = "AGENT ERROR: Could not extract a final answer from the agent's messages." print(f"Could not extract final answer for task {task_id}. Messages: {messages}") print(f"FinalAgent returning answer for task_id '{task_id}': {final_answer[:100]}...") return final_answer if __name__ == '__main__': print("Running a simple test for FinalAgent...") if not os.path.exists('.config'): print("Creating a dummy .config file for testing.") with open('.config', 'w') as f: json.dump({"DEBUG_MODE": True, "BASE_URL": "http://localhost:8000"}, f) # Check for .env and API keys if not load_dotenv(): # Attempts to load .env and returns True if successful print("Warning: .env file not found or failed to load. API keys might be missing.") if not (os.getenv("OPENAI_API_KEY") or os.getenv("DEEPSEEK_API_KEY") or os.getenv("TAVILY_API_KEY")): print("\\nWARNING: Required API key no found in environment variables (OPENAI_API_KEY, DEEPSEEK_API_KEY, TAVILY_API_KEY).") print("The agent will likely fail to initialize or run properly without at least one.") print("Please set them in your .env file or environment for testing.\\n") try: agent = FinalAgent() test_question = "What is the capital of France? And what is the weather like there today?" print(f"Test Question 1: {test_question}") answer = agent(test_question, task_id="test_001") print(f"Test Answer 1: {answer}") test_question_calc = "What is 123 * 4 / 2 + 6?" print(f"\\nTest Question 2 (Calc): {test_question_calc}") answer_calc = agent(test_question_calc, task_id="test_002") print(f"Test Answer 2 (Calc): {answer_calc}") except ValueError as ve: print(f"Initialization Error: {ve}") except Exception as e: print(f"An error occurred during the test: {e}") import traceback traceback.print_exc()