Update app.py

app.py

@@ -11,6 +11,7 @@ from langgraph.graph import Graph, StateGraph
 from typing_extensions import TypedDict
 from openai import OpenAI
 from tools import simple_search
+import re
 
 # -------------------------
 # Utility helpers
@@ -105,15 +106,25 @@ class BasicAgent:
         state["current_step"] = "answer"
         return state
 
+    def _extract_boxed_answer(self, text: str) -> str:
+        """Extract answer from boxed format or return original text if no box found."""
+        # Look for text between [box] and [/box] tags
+        box_match = re.search(r'\[box\](.*?)\[/box\]', text, re.DOTALL)
+        if box_match:
+            return box_match.group(1).strip()
+        return text.strip()
+
     def _generate_answer(self, state: AgentState) -> AgentState:
         history_text = "\n".join(str(item) for item in state["history"])
         prompt = (
             f"Answer the user question as directly as possible. If sources were retrieved, incorporate them.\n"
             f"Question: {state['question']}\n\nContext:\n{history_text}\n\n"
-            "Give ONLY the final answer without extra formatting or explanation."
+            "Give ONLY the final answer without extra formatting or explanation.\n"
+            "Put your answer in a box using [box] and [/box] tags.\n"
+            "If you cannot find a definitive answer, say 'I cannot find a definitive answer to this question.'"
         )
         answer = self._call_llm(prompt, max_tokens=150)
-        state["final_answer"] = answer
+        state["final_answer"] = self._extract_boxed_answer(answer)
         state["history"].append({"step": "answer", "output": answer})
         state["logs"]["final_answer"] = {"prompt": prompt, "response": answer}
         state["current_step"] = "done"