Update agent.py

agent.py

@@ -3,7 +3,8 @@ import os
 import re
 
 from dotenv import load_dotenv
-from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
+from langchain_core.messages import (AIMessage, HumanMessage, SystemMessage,
+                                     ToolMessage)
 from langchain_huggingface import (ChatHuggingFace, HuggingFaceEmbeddings,
                                    HuggingFaceEndpoint)
 from langgraph.graph import START, MessagesState, StateGraph
@@ -71,18 +72,25 @@ def build_graph():
     llm_with_tools = llm.bind_tools(tools)
 
     # --- Nodes ---
+    def extract_answer(llm_output):
+        # Try to parse as JSON if possible
+        try:
+            # If the LLM output is a JSON list, extract the answer
+            parsed = json.loads(llm_output.strip().split('\n')[0])
+            if isinstance(parsed, list) and isinstance(parsed[0], dict) and "submitted_answer" in parsed[0]:
+                return parsed[0]["submitted_answer"]
+        except Exception:
+            pass
+        # Otherwise, just return the first line (before any explanation)
+        return llm_output.strip().split('\n')[0]
+
     def assistant(state: MessagesState):
-        """Assistant node"""
         messages_with_system_prompt = [sys_msg] + state["messages"]
         llm_response = llm_with_tools.invoke(messages_with_system_prompt)
-        # Extract the answer text (strip any "FINAL ANSWER:" if present)
-        answer_text = llm_response.content
-        if answer_text.strip().lower().startswith("final answer:"):
-            answer_text = answer_text.split(":", 1)[1].strip()
-        # Get task_id from state or set a placeholder
-        task_id = state.get("task_id", "1")  # Replace with actual logic if needed
+        answer_text = extract_answer(llm_response.content)
+        task_id = str(state.get("task_id", "1"))  # Ensure task_id is a string
         formatted = [{"task_id": task_id, "submitted_answer": answer_text}]
-        return {"messages": [formatted]}
+        return {"messages": [AIMessage(content=json.dumps(formatted, ensure_ascii=False))]}
 
     # --- Graph Definition ---
     builder = StateGraph(MessagesState)
@@ -154,9 +162,12 @@ if __name__ == "__main__":
             print(message.content)
             print("-----------------------")
         else:
-            output = str(message)
-            print("Agent Output:", output)
-            if is_valid_agent_output(output):
-                print("✅ Output is in the correct format!")
-            else:
-                print("❌ Output is NOT in the correct format!")
+            output = message.content  # This is a string
+            try:
+                parsed = json.loads(output)
+                if isinstance(parsed, list) and "task_id" in parsed[0] and "submitted_answer" in parsed[0]:
+                    print("✅ Output is in the correct format!")
+                else:
+                    print("❌ Output is NOT in the correct format!")
+            except Exception as e:
+                print("❌ Output is NOT in the correct format!", e)