Update app.py

app.py

@@ -3,7 +3,8 @@ import datetime
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
-import gradio as gr
+
+from Gradio_UI import GradioUI
 
 #  Tool to create a personalized greeting repeated a specific number of times.
 @tool
@@ -42,8 +43,11 @@ image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_co
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
-# ----- Initialize agent with API key -----
+# ----- Initialize agent function -----
 def initialize_agent(api_key: str):
+    """
+    Initialize the CodeAgent with the user-provided OpenAI API key.
+    """
     model = OpenAIServerModel(
         max_tokens=2096,
         temperature=0.5,
@@ -65,39 +69,8 @@ def initialize_agent(api_key: str):
     )
     return agent
 
-# ----- Gradio functions -----
-# Store agent globally after API key input
-AGENT = None
-
-def set_api_key(api_key: str):
-    global AGENT
-    AGENT = initialize_agent(api_key)
-    return "Agent initialized! You can now enter your queries below."
-
-def ask_agent(user_query: str):
-    global AGENT
-    if AGENT is None:
-        return "Please enter your OpenAI API key first."
-    response = AGENT.run(user_query)
-    return response
-
-# ----- Gradio interface -----
-with gr.Blocks() as demo:
-    gr.Markdown("## Enter your OpenAI API key and start querying the agent")
-    
-    with gr.Row():
-        api_key_input = gr.Textbox(label="OpenAI API Key", placeholder="sk-...", type="password")
-        api_key_button = gr.Button("Initialize Agent")
-    
-    init_output = gr.Textbox(label="Status", interactive=False)
-
-    with gr.Row():
-        user_query_input = gr.Textbox(label="Enter your query here", placeholder="Ask something...")
-        submit_query_button = gr.Button("Submit Query")
-    
-    query_output = gr.Textbox(label="Agent Response", interactive=False)
-
-    api_key_button.click(set_api_key, inputs=api_key_input, outputs=init_output)
-    submit_query_button.click(ask_agent, inputs=user_query_input, outputs=query_output)
+# Launch Gradio interface with API key input
+def run_agent_with_api_key(api_key: str):
+    agent = initialize_agent(api_key)
+    GradioUI(agent).launch()
 
-demo.launch()