Update app.py

app.py

@@ -80,7 +80,7 @@ def submit_query(session_id, query):
         raise
 
 
-def extract_json_from_answer(answer):
+def extract_json_from_answer(answer, image_analysis):
     """Extract and clean JSON from the LLM response"""
     try:
         # First try to parse the answer directly
@@ -91,10 +91,10 @@ def extract_json_from_answer(answer):
         try:
             # If that fails, try to find JSON content and parse it
             start_idx = answer.find('{')
-            end_idx = answer.rfind('}') + 1
+            end_idx = answer.rfind('}')
             if start_idx != -1 and end_idx != 0:
                 json_str = answer[start_idx:end_idx]
-
+                json_str += f"image_analysis: ['prediction': {image_analysis['prediction']}, 'confidence': {image_analysis['confidence']}%.]
                 return json.loads(json_str)
         except (json.JSONDecodeError, ValueError):
             logger.error("Failed to parse JSON from response")
@@ -179,7 +179,7 @@ def gradio_interface(patient_info, image):
                 "status": "error"
             }, indent=2)
 
-        classes = ["Alzheimer's", "Stroke", "Tumor", "Normal"]
+        classes = ["Alzheimer's", "Normal", "Stroke", "Tumor"]
         # Process image if provided
         image_analysis = None
         if image is not None:
@@ -208,7 +208,7 @@ def gradio_interface(patient_info, image):
             raise ValueError("Invalid response structure from LLM")
         
         # Extract and clean JSON from the response
-        json_data = extract_json_from_answer(llm_response['data']['answer'])
+        json_data = extract_json_from_answer(llm_response['data']['answer'], image_analysis)
         logger.info(f"llm_response: {llm_response}")
         logger.info(f"llm_response[data]: {llm_response['data']}")
         logger.info(f"llm_response[data][answer]: {llm_response['data']['answer']}")