Update final_agent.py

final_agent.py

@@ -0,0 +1,231 @@
+# 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()