Refactor answer_message function to prepend a SystemMessage for context, handle tool calls more robustly, and ensure final responses are generated without tool invocation. This enhances the overall response accuracy and flow.

agent.py

@@ -3,7 +3,7 @@ from typing import TypedDict, List, Dict, Any, Optional, Union
 from langchain_core import tools
 from langgraph.graph import StateGraph, START, END
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFacePipeline
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage, BaseMessage
 from langchain_core.tools import tool
 from ddgs import DDGS
 from dotenv import load_dotenv
@@ -50,7 +50,7 @@ def web_search(keywords: str, max_results:int = 5) -> str:
 
 
 class AgentState(TypedDict):
-    messages: List[Union[HumanMessage, AIMessage]]
+    messages: List[BaseMessage]
 
 
 def read_message(state: AgentState) -> AgentState:
@@ -82,39 +82,55 @@ model_with_tools = model.bind_tools(tools)
 
 def answer_message(state: AgentState) -> AgentState:
     messages = state["messages"]
-    prompt = f"""
-    You are a GAIA question answering expert. 
-    Your task is to provide an answer to a question. 
-    Think carefully before answering the question. 
-    Do not include any thought process before answering the question, and only response exactly what was being asked of you.
-    If you are not able to provide an answer, please state the limitation that you're facing instead.
-
-    Example question: How many hours are there in a day?
-    Response: 24
-
-    Here is the question:
-    {messages}
-    
-    """
-    ai_msg = model_with_tools.invoke(prompt)
+
+    # Prepend system message if not already there
+    if not messages or not isinstance(messages[0], SystemMessage):
+        system_message = SystemMessage(
+            content=(
+                "You are a GAIA question answering expert. "
+                "Your task is to provide an answer to a question. "
+                "Think carefully before answering the question. "
+                "Do not include any thought process before answering the question, "
+                "and only respond exactly what was being asked of you. "
+                "If you are not able to provide an answer, use the available tools and "
+                "state the limitation that you're facing instead."
+            )
+        )
+        messages = [system_message] + messages
+
+    # Step 1: First call (tools enabled). Often returns tool_calls with empty content.
+    ai_msg = model_with_tools.invoke(messages)
     messages.append(ai_msg)
 
-    # Step 2: Execute tools and collect results
-    for tool_call in ai_msg.tool_calls:
-        # Execute the tool with the generated arguments
-        name = tool_call['name']
-        args = tool_call['args']
-        tool = tools_by_name[name]
-        tool_result = tool.invoke(args)
-        # prompt.append(tool_result)
-
-    # Step 3: Pass results back to model for final response
-    print(f"Messages: {messages}")
-    final_response = model_with_tools.invoke(prompt + tool_result)
-    print(f"Final response: {final_response}")
-
-    # Append the model's answer to the messages list
-    return {"messages": [final_response]}
+    tool_calls = getattr(ai_msg, "tool_calls", None) or []
+    if not tool_calls:
+        # Model answered directly
+        return {"messages": messages}
+
+    # Step 2: Execute tools and add ToolMessages
+    for tool_call in tool_calls:
+        name = tool_call["name"]
+        args = tool_call["args"]
+        tool_obj = tools_by_name[name]
+        tool_result = tool_obj.invoke(args)
+        messages.append(
+            ToolMessage(
+                content=str(tool_result),
+                tool_call_id=tool_call["id"],
+            )
+        )
+
+    # Step 3: Force final answer with tools DISABLED.
+    # Some HF chat models keep emitting tool_calls forever; calling the base model avoids blank .content.
+    final_instruction = HumanMessage(
+        content=(
+            "Using the tool results above, provide the FINAL answer now. "
+            "Do not call any tools. Respond with only the answer."
+        )
+    )
+    final_response = model.invoke(messages + [final_instruction])
+    messages.append(final_response)
+    return {"messages": messages}