Update callbackmanager.py

callbackmanager.py

@@ -7,6 +7,7 @@ from datetime import datetime
 import traceback
 import logging
 from huggingface_hub import InferenceClient # Import InferenceClient
+from urllib.parse import urlparse, parse_qs # Import URL parsing utilities
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -737,6 +738,17 @@ def generate_pdf_from_meldrx_with_ai_content(patient_data, llm_content):
     except Exception as e:
         return None, f"Error generating PDF with AI content: {str(e)}"
 
+def extract_auth_code_from_url(redirected_url):
+    """Extracts the authorization code from the redirected URL."""
+    try:
+        parsed_url = urlparse(redirected_url)
+        query_params = parse_qs(parsed_url.query)
+        if "code" in query_params:
+            return query_params["code"][0], None # Return code and no error
+        else:
+            return None, "Authorization code not found in URL." # Return None and error message
+    except Exception as e:
+        return None, f"Error parsing URL: {e}" # Return None and error message
 
 # Create a simplified interface to avoid complex component interactions
 CALLBACK_MANAGER = CallbackManager(
@@ -751,9 +763,13 @@ with gr.Blocks() as demo:
     with gr.Tab("Authenticate with MeldRx"):
         gr.Markdown("## SMART on FHIR Authentication")
         auth_url_output = gr.Textbox(label="Authorization URL", value=CALLBACK_MANAGER.get_auth_url(), interactive=False)
-        gr.Markdown("Copy the URL above, open it in a browser, log in, and paste the 'code' from the redirect URL below.")
-        auth_code_input = gr.Textbox(label="Authorization Code")
-        auth_submit = gr.Button("Submit Code")
+        gr.Markdown("Copy the URL above, open it in a browser, log in, and paste the **entire redirected URL** from your browser's address bar below.") # Updated instructions
+        redirected_url_input = gr.Textbox(label="Redirected URL") # New textbox for redirected URL
+        extract_code_button = gr.Button("Extract Authorization Code") # Button to extract code
+        extracted_code_output = gr.Textbox(label="Extracted Authorization Code", interactive=False) # Textbox to show extracted code
+
+        auth_code_input = gr.Textbox(label="Authorization Code (from above, or paste manually if extraction fails)", interactive=True) # Updated label to be clearer
+        auth_submit = gr.Button("Submit Code for Authentication")
         auth_result = gr.Textbox(label="Authentication Result")
 
         patient_data_button = gr.Button("Fetch Patient Data")
@@ -764,7 +780,28 @@ with gr.Blocks() as demo:
         meldrx_pdf_status = gr.Textbox(label="PDF Generation Status (No AI)") # Renamed status
         meldrx_pdf_download = gr.File(label="Download Generated PDF (No AI)") # Renamed download
 
-        auth_submit.click(fn=CALLBACK_MANAGER.set_auth_code, inputs=auth_code_input, outputs=auth_result)
+        def process_redirected_url(redirected_url):
+            """Processes the redirected URL to extract and display the authorization code."""
+            auth_code, error_message = extract_auth_code_from_url(redirected_url)
+            if auth_code:
+                return auth_code, "Authorization code extracted!"
+            else:
+                return "", error_message or "Could not extract authorization code."
+
+
+        extract_code_button.click(
+            fn=process_redirected_url,
+            inputs=redirected_url_input,
+            outputs=[extracted_code_output, auth_result] # Reusing auth_result for extraction status
+        )
+
+
+        auth_submit.click(
+            fn=CALLBACK_MANAGER.set_auth_code,
+            inputs=extracted_code_output, # Using extracted code as input for authentication
+            outputs=auth_result
+        )
+
 
     with gr.Tab("Patient Dashboard"):
         gr.Markdown("## Patient Data")