Update app.py

app.py

@@ -207,16 +207,32 @@ def submit_query(session_id, query):
             'apikey': api_key,
             '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
+        """
+        
         submit_query_body = {
             "endpointId": "predefined-openai-gpt4o",
-            "query": query,
+            "query": structured_query,
             "pluginIds": ["plugin-1712327325", "plugin-1713962163"],
             "responseMode": "sync"
         }
         
         logger.info(f"Submitting query for session {session_id}")
-        logger.info(f"Query content: {query}")
-        
         response = requests.post(submit_query_url, headers=submit_query_headers, json=submit_query_body)
         response.raise_for_status()
         
@@ -230,44 +246,39 @@ def submit_query(session_id, query):
             logger.error(f"Response content: {e.response.text}")
         raise
 
-def gradio_interface(patient_info, query_type):
+def format_json_response(json_str):
+    """Format the JSON response for better readability"""
+    try:
+        data = json.loads(json_str)
+        return json.dumps(data, indent=2)
+    except json.JSONDecodeError:
+        return json_str
+
+def gradio_interface(patient_info):
     try:
         # Create session
         session_id = create_chat_session()
         
-        # Construct query
-        query = f"Patient Info: {patient_info}\nQuery Type: {query_type}"
-        
         # Submit query and get response
-        llm_response = submit_query(session_id, query)
+        llm_response = submit_query(session_id, patient_info)
         
         # Enhanced response handling
-        if not llm_response:
-            logger.error("Empty response received from LLM")
-            return "Error: No response received from the LLM service"
-        
-        # Navigate the response structure with detailed logging
-        logger.info(f"Processing LLM response: {json.dumps(llm_response, indent=2)}")
-        
-        # Check for answer in the correct field
-        if 'data' not in llm_response or 'answer' not in llm_response['data']:
-            logger.error("Response missing required fields")
-            return f"Error: Unexpected response structure\nFull response: {json.dumps(llm_response, indent=2)}"
+        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']
-        if not answer:
-            logger.error("No answer found in response data")
-            return f"Error: No answer in response\nFull response: {json.dumps(llm_response, indent=2)}"
+        formatted_answer = format_json_response(answer)
         
-        # Format the response
-        response = f"Patient Info: {patient_info}\n\nQuery Type: {query_type}\n\nAnalysis:\n{answer}"
+        # Add processing metrics
+        metrics = llm_response['data'].get('metrics', {})
+        processing_time = metrics.get('totalTimeSec', 'N/A')
         
-        # Add completion metrics if available
-        if 'metrics' in llm_response['data']:
-            metrics = llm_response['data']['metrics']
-            response += f"\n\nProcessing Time: {metrics.get('totalTimeSec', 'N/A')} seconds"
+        # Combine the response
+        full_response = f"Analysis Results (Processing Time: {processing_time} seconds):\n\n{formatted_answer}"
         
-        return response
+        return full_response
     
     except Exception as e:
         logger.error(f"Error in gradio_interface: {str(e)}", exc_info=True)
@@ -279,18 +290,21 @@ iface = gr.Interface(
     inputs=[
         gr.Textbox(
             label="Patient Information",
-            placeholder="Enter patient details here...",
+            placeholder="Enter patient details including: symptoms, medical history, current medications, age, gender, and any relevant test results...",
             lines=5,
             max_lines=10
-        ),
-        gr.Textbox(
-            label="Query Type",
-            placeholder="Describe the type of diagnosis or information needed..."
-        ),
+        )
     ],
-    outputs=gr.Textbox(label="Response", placeholder="The response will appear here..."),
-    title="Medical Diagnosis with LLM",
-    description="Provide patient information and a query type for analysis by the LLM."
+    outputs=gr.Textbox(
+        label="Medical Analysis",
+        placeholder="The analysis will appear here in JSON format...",
+        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.
+    """
 )
 
 if __name__ == "__main__":