Update app.py

app.py

@@ -184,20 +184,12 @@ def create_chat_session():
             "externalUserId": external_user_id
         }
         
-        logger.info("Creating chat session...")
         response = requests.post(create_session_url, headers=create_session_headers, json=create_session_body)
         response.raise_for_status()
-        
-        response_data = response.json()
-        logger.info(f"Session created successfully: {json.dumps(response_data, indent=2)}")
-        
-        session_id = response_data['data']['id']
-        return session_id
+        return response.json()['data']['id']
     
     except requests.exceptions.RequestException as e:
         logger.error(f"Error creating chat session: {str(e)}")
-        if hasattr(e.response, 'text'):
-            logger.error(f"Response content: {e.response.text}")
         raise
 
 def submit_query(session_id, query):
@@ -208,21 +200,20 @@ def submit_query(session_id, query):
             'Content-Type': 'application/json'
         }
         
-        # Structured prompt to get JSON response
         structured_query = f"""
         Based on the following patient information, provide a detailed medical analysis in JSON format:
 
         {query}
 
-        Please structure your response in valid JSON format with the following fields:
-        - diagnosis_details: Overall diagnosis analysis
-        - probable_diagnoses: Array of most probable diagnoses, ordered by likelihood
-        - treatment_plans: Array of recommended treatment plans
-        - lifestyle_modifications: Array of recommended lifestyle changes
-        - medications: Array of recommended medications with dosages
-        - additional_tests: Array of recommended additional tests or examinations
-        - precautions: Array of important precautions or warnings
-        - follow_up: Recommended follow-up timeline and actions
+        Return only valid JSON with these fields:
+        - diagnosis_details
+        - probable_diagnoses (array)
+        - treatment_plans (array)
+        - lifestyle_modifications (array)
+        - medications (array of objects with name and dosage)
+        - additional_tests (array)
+        - precautions (array)
+        - follow_up (string)
         """
         
         submit_query_body = {
@@ -232,57 +223,48 @@ def submit_query(session_id, query):
             "responseMode": "sync"
         }
         
-        logger.info(f"Submitting query for session {session_id}")
         response = requests.post(submit_query_url, headers=submit_query_headers, json=submit_query_body)
         response.raise_for_status()
-        
-        response_data = response.json()
-        logger.info(f"Query response received: {json.dumps(response_data, indent=2)}")
-        return response_data
+        return response.json()
     
     except requests.exceptions.RequestException as e:
         logger.error(f"Error submitting query: {str(e)}")
-        if hasattr(e.response, 'text'):
-            logger.error(f"Response content: {e.response.text}")
         raise
 
-def format_json_response(json_str):
-    """Format the JSON response for better readability"""
+def extract_json_from_answer(answer):
+    """Extract and clean JSON from the LLM response"""
     try:
-        data = json.loads(json_str)
-        return json.dumps(data, indent=2)
+        # First try to parse the answer directly
+        return json.loads(answer)
     except json.JSONDecodeError:
-        return json_str
+        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 gradio_interface(patient_info):
     try:
-        # Create session
         session_id = create_chat_session()
-        
-        # Submit query and get response
         llm_response = submit_query(session_id, patient_info)
         
-        # Enhanced response handling
         if not llm_response or 'data' not in llm_response or 'answer' not in llm_response['data']:
-            logger.error("Invalid response structure")
-            return "Error: Invalid response from the LLM service"
-        
-        # Get the answer and format it
-        answer = llm_response['data']['answer']
-        formatted_answer = format_json_response(answer)
-        
-        # Add processing metrics
-        metrics = llm_response['data'].get('metrics', {})
-        processing_time = metrics.get('totalTimeSec', 'N/A')
+            raise ValueError("Invalid response structure")
         
-        # Combine the response
-        full_response = f"Analysis Results (Processing Time: {processing_time} seconds):\n\n{formatted_answer}"
+        # Extract and clean JSON from the response
+        json_data = extract_json_from_answer(llm_response['data']['answer'])
         
-        return full_response
+        # Return clean JSON string without extra formatting
+        return json.dumps(json_data)
     
     except Exception as e:
-        logger.error(f"Error in gradio_interface: {str(e)}", exc_info=True)
-        return f"Error processing request: {str(e)}"
+        logger.error(f"Error in gradio_interface: {str(e)}")
+        return json.dumps({"error": str(e)})
 
 # Gradio interface
 iface = gr.Interface(
@@ -297,14 +279,11 @@ iface = gr.Interface(
     ],
     outputs=gr.Textbox(
         label="Medical Analysis",
-        placeholder="The analysis will appear here in JSON format...",
+        placeholder="JSON analysis will appear here...",
         lines=15
     ),
     title="Medical Diagnosis Assistant",
-    description="""
-    Enter detailed patient information to receive a structured medical analysis.
-    The system will provide diagnosis possibilities, treatment recommendations, and other relevant medical advice in JSON format.
-    """
+    description="Enter detailed patient information to receive a structured medical analysis in JSON format."
 )
 
 if __name__ == "__main__":