Update app.py

app.py

@@ -51,6 +51,23 @@ def submit_query(session_id, query):
     response = requests.post(submit_query_url, headers=submit_query_headers, json=submit_query_body)
     return response.json()
 
+def extract_json_from_answer(answer):
+    """Extract and clean JSON from the LLM response"""
+    try:
+        # First try to parse the answer directly
+        return json.loads(answer)
+    except json.JSONDecodeError:
+        try:
+            # If that fails, try to find JSON content and parse it
+            start_idx = answer.find('{')
+            end_idx = answer.rfind('}') + 1
+            if start_idx != -1 and end_idx != 0:
+                json_str = answer[start_idx:end_idx]
+                return json.loads(json_str)
+        except (json.JSONDecodeError, ValueError):
+            logger.error("Failed to parse JSON from response")
+            raise
+
 def load_model():
     try:
         model_path = 'model_epoch_01.h5.keras'