try to submit only final answer

app.py

@@ -18,19 +18,49 @@ class WalidzAgent:
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        # Wrap the question in a HumanMessage from langchain_core
-        messages = [HumanMessage(content=question)]
-        messages = self.graph.invoke({"messages": messages})
-        answer = messages['messages'][-1].content
-        # Find the last occurrence of "FINAL ANSWER:" in the response
-        if "FINAL ANSWER:" in answer:
-            # Split by "FINAL ANSWER:" and take the last part
-            final_answer_parts = answer.split("FINAL ANSWER:")
+        
+        # Handle both string and dict input formats
+        if isinstance(question, dict) and "messages" in question and "question_text" in question:
+            # We got a properly formatted input
+            input_state = question
+        else:
+            # Create a proper input state from the question string
+            human_msg = HumanMessage(content=question)
+            input_state = {
+                "messages": [human_msg],
+                "question_text": question
+            }
+            
+        # Invoke the graph
+        messages = self.graph.invoke(input_state)
+        
+        # Extract the answer based on message structure
+        if "assistant" in messages and "messages" in messages["assistant"]:
+            content = messages["assistant"]["messages"].content
+        elif "messages" in messages and len(messages["messages"]) > 0:
+            content = messages["messages"][-1].content
+        else:
+            print("Warning: Unexpected message structure")
+            return "Error: Could not extract answer due to unexpected message structure"
+            
+        # First try to extract from XML tags
+        import re
+        xml_pattern = re.compile(r'<FINAL_ANSWER>(.*?)</FINAL_ANSWER>', re.DOTALL)
+        match = xml_pattern.search(content)
+        
+        if match:
+            # Found XML tags
+            answer = match.group(1).strip()
+        elif "FINAL ANSWER:" in content:
+            # Fall back to "FINAL ANSWER:" format
+            final_answer_parts = content.split("FINAL ANSWER:")
             answer = final_answer_parts[-1].strip()
         else:
-            # If "FINAL ANSWER:" is not found, return the original answer
-            print("Warning: 'FINAL ANSWER:' not found in response")
-        print(answer)
+            # If no explicit final answer format is found, use the original content
+            print("Warning: No final answer format found in response")
+            answer = content
+            
+        print(f"Extracted final answer: {answer}")
         return answer
 
 # (Keep Constants as is)
@@ -99,7 +129,14 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
-            submitted_answer = agent(question_text)
+            # Prepare input state
+            human_msg = HumanMessage(content=question_text)
+            input_state = {
+                "messages": [human_msg],
+                "question_text": question_text
+            }
+            
+            submitted_answer = agent(input_state)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e:
@@ -222,9 +259,8 @@ def run_evaluation_only(profile: gr.OAuthProfile | None):
                 "messages": [human_msg],
                 "question_text": question_text
             }
-
+            
             submitted_answer = agent(input_state)
-            #submitted_answer = agent(question_text)
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e:
             print(f"Error running agent on task {task_id}: {e}")
@@ -235,7 +271,7 @@ def run_evaluation_only(profile: gr.OAuthProfile | None):
 
 # --- Build Gradio Interface using Blocks ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Walid's Basic Agent Evaluation Runner (v4)")
+    gr.Markdown("# Walid's Basic Agent Evaluation Runner (v5)")
     gr.Markdown(
         """
         **Instructions:**