Update app.py

app.py

@@ -388,15 +388,17 @@ QUESTION:
 MATERIALS:
 {search_block}
 
-Write ANSWER: <answer> on its own line.
+Return ONLY a JSON object in this format:
+{{"ANSWER": "your answer here"}}
 """
         try:
             raw = self._call_llm(prompt, 300)
-            answer = raw.split("ANSWER:")[-1].strip()
-            
-            # If first attempt fails or is empty, try a more direct prompt
-            if not answer or "ANSWER:" not in raw or any(k in answer.lower() for k in ["cannot", "sorry", "don't know"]):
-                print("\nFirst attempt failed, trying direct prompt...")
+            try:
+                result = json.loads(raw)
+                answer = result.get("ANSWER", "")
+            except json.JSONDecodeError:
+                print("\nJSON parse error, trying direct prompt...")
+                # If first attempt fails, try a more direct prompt
                 direct_prompt = f"""
 Answer this question directly and concisely. Use the materials provided.
 
@@ -406,28 +408,33 @@ QUESTION:
 MATERIALS:
 {search_block}
 
-If you cannot find an exact answer, provide the most relevant information from the materials.
-Write ANSWER: <answer> on its own line.
+Return ONLY a JSON object in this format:
+{{"ANSWER": "your answer here"}}
 """
                 raw = self._call_llm(direct_prompt, 300)
-                answer = raw.split("ANSWER:")[-1].strip()
-            
-            # Final validation and fallback
-            if not answer or "ANSWER:" not in raw:
-                print("\nBoth attempts failed, using fallback answer...")
-                if materials:
-                    # If we have materials but no answer, summarize what we know
-                    summary_prompt = f"""
-Summarize the key information from these materials in one sentence:
+                try:
+                    result = json.loads(raw)
+                    answer = result.get("ANSWER", "")
+                except json.JSONDecodeError:
+                    print("\nBoth attempts failed, using fallback answer...")
+                    if materials:
+                        # If we have materials but no answer, summarize what we know
+                        summary_prompt = f"""
+Summarize the key information from these materials in one sentence.
+Return ONLY a JSON object in this format:
+{{"ANSWER": "your summary here"}}
 
+MATERIALS:
 {search_block}
-
-Write ANSWER: <answer> on its own line.
 """
-                    raw = self._call_llm(summary_prompt, 150)
-                    answer = raw.split("ANSWER:")[-1].strip()
-                else:
-                    answer = "I cannot provide a definitive answer at this time."
+                        raw = self._call_llm(summary_prompt, 150)
+                        try:
+                            result = json.loads(raw)
+                            answer = result.get("ANSWER", "")
+                        except json.JSONDecodeError:
+                            answer = "I cannot provide a definitive answer at this time."
+                    else:
+                        answer = "I cannot provide a definitive answer at this time."
             
             state["final_answer"] = answer
             state["current_step"] = "done"
@@ -508,7 +515,7 @@ Write ANSWER: <answer> on its own line.
             }
         
         final_state = self.workflow.invoke(state)
-        return final_state["final_answer"]
+        return final_state["final_answer"]  # Return the answer string directly, not JSON encoded
 
     def _download_file(self, url: str) -> bytes:
         """Download a file from a URL."""