Update app.py

app.py

@@ -13,28 +13,28 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
-    """A langgraph agent."""
     def __init__(self):
         print("BasicAgent initialized.")
-        self.graph = build_graph()
 
+    def post_process_answer(self, answer: str):
+        if "FINAL ANSWER:" not in answer:
+            raise ValueError("The answer does not contain 'FINAL ANSWER:' keyword")
+        key_answer_start_idx = answer.find("FINAL ANSWER:") + len("FINAL ANSWER:")
+        key_answer = answer[key_answer_start_idx:].strip()
+        return key_answer
+    
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        messages = [HumanMessage(content=question)]
-        result = self.graph.invoke({"messages": messages})
-        answer = result['messages'][-1].content
-        if isinstance(answer, str):
-            # If the tool already returned a clean value (like "RSA" or "EGY")
-            # strip whitespace and formatting
-            cleaned = answer.strip()
-            # If model included "Answer:" or markdown, extract last token
-            if "Answer:" in cleaned:
-                cleaned = cleaned.split("Answer:")[-1].strip()
-            if cleaned.startswith("**") and cleaned.endswith("**"):
-                cleaned = cleaned.strip("*").strip()
-            return cleaned
-
-        return str(answer)
+        try:
+            result = agent.invoke({"messages": [HumanMessage(content=question)]})
+            answer = result["messages"][-1].content
+            key_answer = self.post_process_answer(answer)
+            print("Question:", question)
+            print("Answer:", key_answer)
+        except Exception as e:
+            print(e)
+            key_answer = str(e)
+        return key_answer
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """