Merge branch 'user-api-key'

src/plan/app_text2plan.py

@@ -97,9 +97,18 @@ class MarkdownBuilder:
 class SessionState:
     """
     In a multi-user environment (e.g. Hugging Face Spaces), this class hold each users state.
-    In a single-user environment, this class is used to hold the state of the current session.
+    In a single-user environment, this class is used to hold the state of that lonely user.
+
+    IDEA: Persist the user settings for longer. The settings survive a page refresh, but not a server restart.
+    The browser has a local storage. Can Gradio access that, so that the settings are remembered between sessions?
     """
     def __init__(self):
+        # Settings: the user's OpenRouter API key.
+        self.openrouter_api_key = "" # Initialize to empty string
+        # Settings: The model that the user has picked.
+        self.llm_model = available_model_names[0]
+        # Settings: The speedvsdetail that the user has picked.
+        self.speedvsdetail = SpeedVsDetailEnum.ALL_DETAILS_BUT_SLOW
         # Holds the subprocess.Popen object for the currently running pipeline process.
         self.active_proc = None
         # A threading.Event used to signal that the running process should stop.
@@ -117,14 +126,11 @@ class SessionState:
         """
         return self
 
-def run_planner(submit_or_retry_button, plan_prompt, llm_model, speedvsdetail, session_state: SessionState):
+def run_planner(submit_or_retry_button, plan_prompt, session_state: SessionState):
     """
     Generator function for launching the pipeline process and streaming updates.
     The session state is carried in a SessionState instance.
     """
-    # Initialize session_state if needed.
-    if session_state is None:
-        session_state = SessionState()
 
     # Clear any previous stop signal.
     session_state.stop_event.clear()
@@ -163,8 +169,15 @@ def run_planner(submit_or_retry_button, plan_prompt, llm_model, speedvsdetail, s
     # Set environment variables for the pipeline.
     env = os.environ.copy()
     env["RUN_ID"] = run_id
-    env["LLM_MODEL"] = llm_model
-    env["SPEED_VS_DETAIL"] = speedvsdetail
+    env["LLM_MODEL"] = session_state.llm_model
+    env["SPEED_VS_DETAIL"] = session_state.speedvsdetail
+
+    # If there is a non-empty OpenRouter API key, set it as an environment variable.
+    if session_state.openrouter_api_key and len(session_state.openrouter_api_key) > 0:
+        print("Setting OpenRouter API key as environment variable.")
+        env["OPENROUTER_API_KEY"] = session_state.openrouter_api_key
+    else:
+        print("No OpenRouter API key provided.")
 
     start_time = time.perf_counter()
     # Initialize the last zip creation time to be ZIP_INTERVAL_SECONDS in the past
@@ -278,8 +291,6 @@ def stop_planner(session_state: SessionState):
     """
     Sets a stop flag in the session_state and attempts to terminate the active process.
     """
-    if session_state is None:
-        session_state = SessionState()
 
     session_state.stop_event.set()
 
@@ -299,8 +310,6 @@ def open_output_dir(session_state: SessionState):
     """
     Opens the latest output directory in the native file explorer.
     """
-    if session_state is None:
-        session_state = SessionState()
 
     latest_run_dir = session_state.latest_run_dir
     if not latest_run_dir or not os.path.exists(latest_run_dir):
@@ -317,6 +326,29 @@ def open_output_dir(session_state: SessionState):
     except Exception as e:
         return f"Failed to open directory: {e}", session_state
 
+def update_openrouter_api_key(openrouter_api_key, session_state: SessionState):
+    """Updates the OpenRouter API key in the session state."""
+    session_state.openrouter_api_key = openrouter_api_key
+    return openrouter_api_key, session_state
+
+def update_model_radio(llm_model, session_state: SessionState):
+    """Updates the llm_model in the session state"""
+    session_state.llm_model = llm_model
+    return llm_model, session_state
+
+def update_speedvsdetail_radio(speedvsdetail, session_state: SessionState):
+    """Updates the speedvsdetail in the session state"""
+    session_state.speedvsdetail = speedvsdetail
+    return speedvsdetail, session_state
+
+def initialize_settings(session_state: SessionState):
+    """Initializes the settings from session_state, if available."""
+    return (session_state.openrouter_api_key,
+            session_state.llm_model,
+            session_state.speedvsdetail,
+            session_state)
+
+
 # Build the Gradio UI using Blocks.
 with gr.Blocks(title="PlanExe") as demo_text2plan:
     gr.Markdown("# PlanExe: crack open pandora’s box of ideas")
@@ -363,24 +395,31 @@ with gr.Blocks(title="PlanExe") as demo_text2plan:
             label="Speed vs Detail",
             interactive=True 
         )
+        openrouter_api_key_text = gr.Textbox(
+            label="OpenRouter API Key",
+            type="password",
+            placeholder="Enter your OpenRouter API key (optional)",
+            info="Sign up at [OpenRouter](https://openrouter.ai/) to get an API key. A small top-up (e.g. 5 USD) is needed to access paid models."
+        )
 
     with gr.Tab("Join the community"):
         gr.Markdown("""
 - [GitHub](https://github.com/neoneye/PlanExe) the source code.
 - [Discord](https://neoneye.github.io/PlanExe-web/discord) join the community. Suggestions, feedback, and questions are welcome.
 """)
-    # Create a hidden state to hold our SessionState instance.
+    
+    # Manage the state of the current user
     session_state = gr.State(SessionState())
 
     # Submit and Retry buttons call run_planner and update the state.
     submit_btn.click(
         fn=run_planner,
-        inputs=[submit_btn, prompt_input, model_radio, speedvsdetail_radio, session_state],
+        inputs=[submit_btn, prompt_input, session_state],
         outputs=[output_markdown, download_output, session_state]
     )
     retry_btn.click(
         fn=run_planner,
-        inputs=[retry_btn, prompt_input, model_radio, speedvsdetail_radio, session_state],
+        inputs=[retry_btn, prompt_input, session_state],
         outputs=[output_markdown, download_output, session_state]
     )
     # The Stop button uses the state to terminate the running process.
@@ -396,6 +435,29 @@ with gr.Blocks(title="PlanExe") as demo_text2plan:
         outputs=[status_markdown, session_state]
     )
 
+    openrouter_api_key_text.change(
+        fn=update_openrouter_api_key,
+        inputs=[openrouter_api_key_text, session_state],
+        outputs=[openrouter_api_key_text, session_state]
+    )
+    model_radio.change(
+        fn=update_model_radio,
+        inputs=[model_radio, session_state],
+        outputs=[model_radio, session_state]
+    )
+    speedvsdetail_radio.change(
+        fn=update_speedvsdetail_radio,
+        inputs=[speedvsdetail_radio, session_state],
+        outputs=[speedvsdetail_radio, session_state]
+    )
+
+
+    demo_text2plan.load(
+        fn=initialize_settings,
+        inputs=[session_state],
+        outputs=[openrouter_api_key_text, model_radio, speedvsdetail_radio, session_state]
+    )
+
 if __name__ == "__main__":
     # print("Environment variables Gradio:\n" + get_env_as_string() + "\n\n\n")