Update agent.py

agent.py

@@ -181,13 +181,31 @@ def build_graph(provider: str = "groq"):
 
     def assistant(state: MessagesState):
         try:
-            result = llm_with_tools.invoke(state["messages"])
+            # Always prepend the system message
+            system_msg = SystemMessage(content=system_prompt.strip())
+            messages = [system_msg] + state["messages"]
+
+            result = llm_with_tools.invoke(messages)
             print("🤖 Raw LLM result:", repr(result))
-            final_response = result.content.strip()
-            return {"messages": [HumanMessage(content=final_response)]}
+
+            raw_output = result.content.strip()
+
+            # Extract FINAL ANSWER using regex
+            import re
+            match = re.search(r"FINAL ANSWER:\s*(.+)", raw_output, re.IGNORECASE)
+            if match:
+                final_answer = match.group(1).strip()
+                final_output = f"FINAL ANSWER: {final_answer}"
+            else:
+                print("⚠️ 'FINAL ANSWER:' prefix not found. Using fallback.")
+                final_output = "FINAL ANSWER: Unable to determine answer"
+
+            return {"messages": [HumanMessage(content=final_output)]}
+
         except Exception as e:
             print(f"🔥 Error in assistant node: {e}")
-            return {"messages": [HumanMessage(content=f"AGENT ERROR: {e}")] }
+            return {"messages": [HumanMessage(content=f"FINAL ANSWER: AGENT ERROR: {e}")]}
+