Fix: Add logger and adjustment to update Column and stripes

callbackmanager.py

@@ -4,6 +4,12 @@ import json
 import os
 import tempfile
 from datetime import datetime
+import traceback
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Import PDF utilities
 from pdfutils import PDFGenerator, generate_discharge_summary
@@ -31,12 +37,81 @@ class CallbackManager:
         return "Authentication failed. Please check the code."
 
     def get_patient_data(self) -> str:
-        if not self.access_token:
-            return "Not authenticated. Please provide a valid authorization code first."
-        patients = self.api.get_patients()
-        if patients is not None:
-            return json.dumps(patients, indent=2) if patients else "No patient data returned."
-        return "Failed to retrieve patient data."
+        """Fetch patient data from MeldRx"""
+        try:
+            if not self.access_token:
+                logger.warning("Not authenticated when getting patient data")
+                return "Not authenticated. Please provide a valid authorization code first."
+            
+            # For demo purposes, if there's no actual API connected, return mock data
+            # Remove this in production and use the real API call
+            if not hasattr(self.api, 'get_patients') or self.api.get_patients is None:
+                logger.info("Using mock patient data (no API connection)")
+                # Return mock FHIR bundle with patient data
+                mock_data = {
+                    "resourceType": "Bundle",
+                    "type": "searchset",
+                    "total": 2,
+                    "link": [],
+                    "entry": [
+                        {
+                            "resource": {
+                                "resourceType": "Patient",
+                                "id": "patient1",
+                                "name": [
+                                    {
+                                        "use": "official",
+                                        "family": "Smith",
+                                        "given": ["John"]
+                                    }
+                                ],
+                                "gender": "male",
+                                "birthDate": "1970-01-01",
+                                "address": [
+                                    {
+                                        "city": "Boston",
+                                        "state": "MA",
+                                        "postalCode": "02108"
+                                    }
+                                ]
+                            }
+                        },
+                        {
+                            "resource": {
+                                "resourceType": "Patient",
+                                "id": "patient2",
+                                "name": [
+                                    {
+                                        "use": "official",
+                                        "family": "Johnson",
+                                        "given": ["Jane"]
+                                    }
+                                ],
+                                "gender": "female",
+                                "birthDate": "1985-05-15",
+                                "address": [
+                                    {
+                                        "city": "Cambridge",
+                                        "state": "MA",
+                                        "postalCode": "02139"
+                                    }
+                                ]
+                            }
+                        }
+                    ]
+                }
+                return json.dumps(mock_data, indent=2)
+                
+            # Real implementation with API call
+            logger.info("Calling Meldrx API to get patients")
+            patients = self.api.get_patients()
+            if patients is not None:
+                return json.dumps(patients, indent=2) if patients else "No patient data returned."
+            return "Failed to retrieve patient data."
+        except Exception as e:
+            error_msg = f"Error in get_patient_data: {str(e)}"
+            logger.error(error_msg)
+            return f"Error retrieving patient data: {str(e)}"
 
     def get_patient_documents(self, patient_id: str = None):
         """Fetch patient documents from MeldRx"""
@@ -156,13 +231,13 @@ def create_patient_stripe(patient):
     def update_components(selected_status):
         visible = selected_status in ["Reviewed", "Approved"]
         return (
-            gr.Button(visible=visible),
-            gr.Label(visible=visible, value={"Low": 0.3, "Medium": 0.6, "High": 0.9}.get(selected_status, 0.5))
+            {"visible": visible},  # For generate_btn
+            {"visible": visible, "value": {"Low": 0.3, "Medium": 0.6, "High": 0.9}.get(selected_status, 0.5)}  # For probability
         )
 
     def show_documents(patient_id):
         docs = CALLBACK_MANAGER.get_patient_documents(patient_id)
-        return gr.JSON.update(value=docs, visible=True)
+        return {"value": docs, "visible": True}  # For doc_output
 
     status.change(
         update_components,
@@ -190,22 +265,48 @@ def update_patient_display(patient_data_json):
     """
     try:
         # Handle the case where patient_data_json is an error message
-        if patient_data_json.startswith("Not authenticated") or patient_data_json.startswith("Failed"):
+        if isinstance(patient_data_json, str) and (patient_data_json.startswith("Not authenticated") or patient_data_json.startswith("Failed")):
+            logger.warning(f"Error in patient data: {patient_data_json}")
+            return (
+                {"visible": False},  # For pagination_row
+                [],  # For patient_stripes
+                {"value": f"Error: {patient_data_json}"},  # For metadata_md
+                None  # For patient_data_state
+            )
+        
+        # Parse JSON if it's a string
+        if isinstance(patient_data_json, str):
+            try:
+                fhir_bundle = json.loads(patient_data_json)
+            except json.JSONDecodeError as e:
+                logger.error(f"JSON decode error: {str(e)} - Data: {patient_data_json[:100]}...")
+                return (
+                    {"visible": False},
+                    [],
+                    {"value": f"Invalid JSON data: {str(e)}"},
+                    None
+                )
+        else:
+            fhir_bundle = patient_data_json
+        
+        # Safety check
+        if not isinstance(fhir_bundle, dict):
+            logger.error(f"Unexpected data type: {type(fhir_bundle)}")
             return (
-                gr.Row.update(visible=False),
-                gr.Column.update([]),
-                gr.Markdown.update(value=f"Error: {patient_data_json}"),
+                {"visible": False},
+                [],
+                {"value": f"Expected dictionary but got {type(fhir_bundle)}"},
                 None
             )
             
-        fhir_bundle = json.loads(patient_data_json) if isinstance(patient_data_json, str) else patient_data_json
         patients = process_fhir_bundle(fhir_bundle)
         
         if not patients:
+            logger.warning("No patients found in data")
             return (
-                gr.Row.update(visible=False),
-                gr.Column.update([]),
-                gr.Markdown.update(value="No patients found in the data"),
+                {"visible": False},
+                [],
+                {"value": "No patients found in the data"},
                 None
             )
         
@@ -213,26 +314,36 @@ def update_patient_display(patient_data_json):
         has_prev = any(link.get("relation") == "previous" for link in fhir_bundle.get("link", []))
         has_next = any(link.get("relation") == "next" for link in fhir_bundle.get("link", []))
         
-        # Create stripes
-        stripes = [create_patient_stripe(patient) for patient in patients]
+        # Create stripes - handle list vs single patient
+        stripes = []
+        for patient in patients:
+            try:
+                stripe = create_patient_stripe(patient)
+                stripes.append(stripe)
+            except Exception as e:
+                logger.error(f"Error creating patient stripe: {str(e)}")
+                # Continue with other patients if one fails
+        
+        logger.info(f"Created {len(stripes)} patient stripes successfully")
         
         return (
-            gr.Row.update(visible=has_prev or has_next),
-            gr.Column.update(value=stripes),
-            gr.Markdown.update(value=f"""
+            {"visible": has_prev or has_next},
+            stripes,
+            {"value": f"""
             **Bundle Metadata**  
             Type: {fhir_bundle.get('type', 'Unknown')}  
             Total Patients: {fhir_bundle.get('total', len(patients))}  
             Last Updated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
-            """),
+            """},
             {"patients": patients, "bundle": fhir_bundle}
         )
     except Exception as e:
-        import traceback
+        error_msg = f"Error loading patient data: {str(e)}\n{traceback.format_exc()}"
+        logger.error(error_msg)
         return (
-            gr.Row.update(visible=False),
-            gr.Column.update([]),
-            gr.Markdown.update(value=f"Error loading patient data: {str(e)}\n{traceback.format_exc()}"),
+            {"visible": False},
+            [],
+            {"value": error_msg},
             None
         )
 
@@ -479,20 +590,58 @@ with gr.Blocks() as demo:
             outputs=pdf_output
         )
     
-    # Define a function to handle fetching and displaying patients
+    # Define a function to handle fetching and displaying patients with better error handling
     def fetch_and_display_patients():
-        data = CALLBACK_MANAGER.get_patient_data()
-        
-        # Process the results for UI display
-        pagination_update, stripes_update, metadata_update, state_update = update_patient_display(data)
-        
-        return [
-            data,  # For the raw JSON output
-            pagination_update,
-            stripes_update,
-            metadata_update,
-            state_update
-        ]
+        try:
+            logger.info("Fetching patient data...")
+            # Get data from API
+            data = CALLBACK_MANAGER.get_patient_data()
+            logger.info(f"Received patient data: {data[:100]}..." if isinstance(data, str) else "Received non-string data")
+            
+            # Handle case when the data is just an error message
+            if data.startswith("Not authenticated") or data.startswith("Failed"):
+                logger.warning(f"Authentication issue: {data}")
+                return [
+                    data,  # For the raw JSON output
+                    {"visible": False},  # For pagination_row
+                    [],  # For patient_stripes
+                    {"value": f"Error: {data}"},  # For metadata_md
+                    None  # For patient_data_state
+                ]
+            
+            # Process the results for UI display
+            try:
+                pagination_update, stripes_update, metadata_update, state_update = update_patient_display(data)
+                logger.info("Successfully processed patient display data")
+                
+                return [
+                    data,  # For the raw JSON output
+                    pagination_update,
+                    stripes_update,
+                    metadata_update,
+                    state_update
+                ]
+            except Exception as e:
+                logger.error(f"Error in update_patient_display: {str(e)}\n{traceback.format_exc()}")
+                # Return at least the raw data, even if display processing fails
+                return [
+                    data,
+                    {"visible": False},
+                    [],
+                    {"value": f"Error processing patient data: {str(e)}"},
+                    None
+                ]
+                
+        except Exception as e:
+            error_msg = f"Unexpected error fetching patient data: {str(e)}\n{traceback.format_exc()}"
+            logger.error(error_msg)
+            return [
+                f"Error: {str(e)}",
+                {"visible": False},
+                [],
+                {"value": error_msg},
+                None
+            ]
 
     # Connect the patient data button to both display the raw JSON and update the patient display
     patient_data_button.click(