feat: add user api input

app.py

@@ -540,6 +540,7 @@ def start_simulation(
     personality: str,
     recall: str,
     confusion: str,
+    user_api_key: str = "",
     request: gr.Request = None,
 ):
     if not hadm_id:
@@ -548,14 +549,20 @@ def start_simulation(
     if model not in BACKEND_MODELS:
         return _setup_error("Invalid model selection.")
 
-    client_key = get_client_key(request)
-    allowed, limit_msg = _rate_limiter.check_simulation_start(client_key)
-    if not allowed:
-        return _setup_error(limit_msg)
+    using_own_key = bool(user_api_key.strip())
+
+    # Only apply rate limiting when using the shared demo key
+    if not using_own_key:
+        client_key = get_client_key(request)
+        allowed, limit_msg = _rate_limiter.check_simulation_start(client_key)
+        if not allowed:
+            return _setup_error(limit_msg)
 
     is_openai = "gpt" in model.lower()
 
-    if is_openai:
+    if using_own_key:
+        api_key = user_api_key.strip()
+    elif is_openai:
         api_key = os.environ.get("OPENAI_API_KEY", "")
     else:
         api_key = os.environ.get("GENAI_API_KEY", "") or os.environ.get("GOOGLE_API_KEY", "")
@@ -597,6 +604,7 @@ def start_simulation(
         "patient": patient,
         "model": model,
         "recap_html": recap,
+        "user_api_key": user_api_key.strip(),  # empty string = using shared demo key
     }
 
     return (
@@ -660,7 +668,7 @@ _INJECTION_PATTERNS = re.compile(
 )
 
 
-def chat(message: str, history: list, agent, request: gr.Request = None):
+def chat(message: str, history: list, agent, sim_config: dict, request: gr.Request = None):
     if agent is None:
         raise gr.Error("No simulation running. Please start a simulation first.")
     if not message.strip():
@@ -676,10 +684,12 @@ def chat(message: str, history: list, agent, request: gr.Request = None):
         _logger.warning("Prompt injection attempt detected from key=%s", get_client_key(request))
         raise gr.Error("Invalid input detected. Please enter a valid clinical question.")
 
-    client_key = get_client_key(request)
-    allowed, limit_msg = _rate_limiter.check_chat_message(client_key)
-    if not allowed:
-        raise gr.Error(limit_msg)
+    using_own_key = bool(sim_config and sim_config.get("user_api_key"))
+    if not using_own_key:
+        client_key = get_client_key(request)
+        allowed, limit_msg = _rate_limiter.check_chat_message(client_key)
+        if not allowed:
+            raise gr.Error(limit_msg)
 
     response = agent(user_prompt=message, using_multi_turn=True, verbose=False)
     history = history + [
@@ -736,11 +746,13 @@ def start_auto(agent, sim_config: dict, request: gr.Request = None):
         yield _auto_fallback_outputs()
         return
 
-    allowed, limit_msg = _rate_limiter.check_auto_run(client_key)
-    if not allowed:
-        gr.Warning(limit_msg)
-        yield _auto_fallback_outputs()
-        return
+    using_own_key = bool(sim_config.get("user_api_key"))
+    if not using_own_key:
+        allowed, limit_msg = _rate_limiter.check_auto_run(client_key)
+        if not allowed:
+            gr.Warning(limit_msg)
+            yield _auto_fallback_outputs()
+            return
 
     try:
         agent.reset_history(verbose=False)
@@ -758,7 +770,9 @@ def start_auto(agent, sim_config: dict, request: gr.Request = None):
 
         model = sim_config["model"]
         is_openai = "gpt" in model.lower()
-        if is_openai:
+        if using_own_key:
+            api_key = sim_config["user_api_key"]
+        elif is_openai:
             api_key = os.environ.get("OPENAI_API_KEY", "")
         else:
             api_key = os.environ.get("GENAI_API_KEY", "") or os.environ.get("GOOGLE_API_KEY", "")
@@ -914,12 +928,26 @@ with gr.Blocks(title="PatientSim", theme=gr.themes.Soft(), css=CUSTOM_CSS) as de
 
         # ── Connection card ──────────────────────────────────────────────────
         with gr.Group(elem_classes=["form-card"]):
-            # gr.HTML("<span class='card-title'>Model Selection</span>")
-            model_dd = gr.Dropdown(
-                choices=BACKEND_MODELS,
-                value=BACKEND_MODELS[0],
-                label="Model Selection",
+            gr.Markdown(
+                "**🔑 Model & API Key**\n\n"
+                "This demo runs on a shared API key with a limited number of free calls. "
+                "If the free quota has been exhausted, please enter your own API key below "
+                "(OpenAI or Google Gemini) to continue without restrictions. "
+                "Your key is used only for this session and is never stored on our servers."
             )
+            with gr.Row(equal_height=True):
+                model_dd = gr.Dropdown(
+                    choices=BACKEND_MODELS,
+                    value=BACKEND_MODELS[0],
+                    label="Model",
+                    scale=1,
+                )
+                api_key_input = gr.Textbox(
+                    label="API Key (optional)",
+                    placeholder="Leave blank to use the shared demo key  ·  sk-... or paste your Gemini key",
+                    type="password",
+                    scale=2,
+                )
 
         # ── Patient Case card (one gr.Button per patient) ────────────────────
         with gr.Group(elem_classes=["form-card"]):
@@ -1165,7 +1193,7 @@ with gr.Blocks(title="PatientSim", theme=gr.themes.Soft(), css=CUSTOM_CSS) as de
     # Start simulation → mode selection
     start_btn.click(
         fn=start_simulation,
-        inputs=[selected_patient_state, model_dd, cefr_radio, personality_dd, recall_radio, confusion_radio],
+        inputs=[selected_patient_state, model_dd, cefr_radio, personality_dd, recall_radio, confusion_radio, api_key_input],
         outputs=[patient_agent_state, sim_config_state, setup_section, mode_section, recap_display, setup_error_display],
     )
 
@@ -1203,12 +1231,12 @@ with gr.Blocks(title="PatientSim", theme=gr.themes.Soft(), css=CUSTOM_CSS) as de
     )
     chat_event_send = send_btn.click(
         fn=chat,
-        inputs=[msg_box, chatbot, patient_agent_state],
+        inputs=[msg_box, chatbot, patient_agent_state, sim_config_state],
         outputs=[chatbot, msg_box],
     )
     chat_event_submit = msg_box.submit(
         fn=chat,
-        inputs=[msg_box, chatbot, patient_agent_state],
+        inputs=[msg_box, chatbot, patient_agent_state, sim_config_state],
         outputs=[chatbot, msg_box],
     )
     back_from_chat_to_mode_btn.click(