Update app.py

app.py

@@ -1,112 +1,108 @@
-import gradio as gr
 import os
-from dotenv import load_dotenv
-
-# Import necessary functions/variables from run.py
-from run import create_agent
-
-load_dotenv()
-
-CONFIG_FILE = ".user_config.env"
-
-
-def save_env_vars_to_file(env_vars):
-    """Saves environment variables to a file."""
-    print("[DEBUG] Saving user config to file")
-    with open(CONFIG_FILE, "w") as f:
-        for key, value in env_vars.items():
-            f.write(f"{key}={value}\n")
-
-
-def launch_interface():
-    """Launches the Gradio interface for configuring and running the agent."""
-
-    def setup_agent(question, model_id, hf_token, serpapi_key, use_custom_endpoint,
-                    custom_api_endpoint, custom_api_key, search_provider, search_api_key, custom_search_url):
-        """Sets up the agent with the given configuration and runs it with the provided question."""
-        print("[DEBUG] Setting up agent with input question:", question)
-
-        if question.strip() == "":
-            return "Please enter a question."
-
-        endpoint = custom_api_endpoint if use_custom_endpoint else None
-        api_key = custom_api_key if use_custom_endpoint else None
-
-        save_env_vars_to_file({
-            "HF_TOKEN": hf_token,
-            "SERPAPI_API_KEY": serpapi_key,
-        })
-
-        print("[DEBUG] Instantiating agent with UI configuration")
-        agent = create_agent(
-            model_id=model_id,
-            hf_token=hf_token,
-            serpapi_key=serpapi_key,
-            custom_api_endpoint=endpoint,
-            custom_api_key=api_key,
-            search_provider=search_provider,
-            search_api_key=search_api_key,
-            custom_search_url=custom_search_url
+import gradio as gr
+from run import create_agent, log_formatter, MODEL_CONFIGS
+import io
+import logging
+import contextlib
+
+# Simple dark theme styling.
+CSS = """
+body {
+    background-color: #2c2c2c;
+    color: #ffffff;
+}
+.gradio-container {
+    background-color: #3a3a3a;
+    border-radius: 10px;
+    padding: 20px;
+}
+h1, h2, h3 {
+    color: #79c0ff;
+}
+"""
+
+def set_keys(openai_api_key, serper_api_key, hf_token, gemini_api_key, groq_api_key):
+    os.environ["OPENAI_API_KEY"] = openai_api_key
+    os.environ["SERPER_API_KEY"] = serper_api_key
+    os.environ["HF_TOKEN"] = hf_token
+    os.environ["GEMINI_API_KEY"] = gemini_api_key
+    os.environ["GROQ_API_KEY"] = groq_api_key
+    return "API keys have been updated successfully! Please restart the agent for changes to take effect."
+
+def get_answer(question, model_name):
+    log_buffer = io.StringIO()
+    stdout_buffer = io.StringIO()
+    stream_handler = logging.StreamHandler(log_buffer)
+    stream_handler.setFormatter(log_formatter)
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.DEBUG)
+    root_logger.addHandler(stream_handler)
+
+    conversation = []
+    try:
+        agent = create_agent(model_name=model_name)
+        answer = agent.run(question)
+        if isinstance(answer, str):
+            conversation = [
+                {"role": "user", "content": question},
+                {"role": "assistant", "content": answer}
+            ]
+        else:
+            result = ""
+            for chunk in answer:
+                result += chunk
+            conversation = [
+                {"role": "user", "content": question},
+                {"role": "assistant", "content": result}
+            ]
+    except Exception as e:
+        conversation = [
+            {"role": "user", "content": question},
+            {"role": "assistant", "content": f"An error occurred: {e}"}
+        ]
+    finally:
+        root_logger.removeHandler(stream_handler)
+        stream_handler.close()
+        logs = log_buffer.getvalue()
+    return conversation, f"```\n{stdout_buffer.getvalue()}\n{logs}\n```"
+
+def build_app():
+    with gr.Blocks(css=CSS) as demo:
+        # Title and header.
+        gr.HTML("<h1>SmolAgents Open Deep Search 🥳</h1>")
+        gr.Markdown("## Enhanced Agent UI")
+
+        # Configuration Accordion.
+        with gr.Accordion("Configuration (Click to Expand)", open=False):
+            openai_field = gr.Textbox(label="OPENAI_API_KEY", type="password", placeholder="Enter your OpenAI API key")
+            serper_field = gr.Textbox(label="SERPER_API_KEY", type="password", placeholder="Enter your Serper API key")
+            hf_field = gr.Textbox(label="HF_TOKEN", type="password", placeholder="Enter your Hugging Face Token")
+            gemini_field = gr.Textbox(label="GEMINI_API_KEY", type="password", placeholder="Enter your Gemini API key")
+            groq_field = gr.Textbox(label="GROQ_API_KEY", type="password", placeholder="Enter your Groq API key")
+            update_btn = gr.Button("Update Keys")
+            status_box = gr.Markdown("*(No keys set yet)*")
+            update_btn.click(fn=set_keys, inputs=[openai_field, serper_field, hf_field, gemini_field, groq_field], outputs=status_box)
+
+        # Placeholder for agent logs.
+        log_markdown = gr.Markdown(label="Agent Logs")
+
+        gr.Markdown("### Select a model and ask your question below:")
+        model_names = list(MODEL_CONFIGS.keys())
+        model_dropdown = gr.Dropdown(choices=model_names, label="Select Model", value="o1", allow_custom_value=True)
+        question_input = gr.Textbox(label="Your Question", placeholder="Enter your question here...")
+        submit_btn = gr.Button("Get Answer")
+        # Use Chatbot component with type "messages" for OpenAI-style dict format.
+        chatbot = gr.Chatbot(label="Answer", type="messages")
+
+        # Connect the answer function.
+        submit_btn.click(
+            fn=get_answer,
+            inputs=[question_input, model_dropdown],
+            outputs=[chatbot, log_markdown]
         )
 
-        return agent.run(question)
-
-    with gr.Blocks(theme=gr.themes.Base(), css=".gr-box { border-radius: 15px; padding: 20px; }") as demo:
-        gr.Markdown("# 🤖 SmolAgent Configurable Interface")
-
-        with gr.Row():
-            with gr.Column(scale=2):
-                question = gr.Textbox(label="🧠 Question", placeholder="Ask me anything...")
-                model_id = gr.Textbox(value="gpt-4o-mini", label="🧬 Model ID")
-                hf_token = gr.Textbox(value=os.getenv("HF_TOKEN", ""), label="🔑 HuggingFace Token", type="password")
-                serpapi_key = gr.Textbox(value=os.getenv("SERPAPI_API_KEY", ""), label="🔍 Serper API Key",
-                                         type="password", visible=True)
-                use_custom_endpoint = gr.Checkbox(label="🌐 Use Custom API Endpoint")
-                custom_api_endpoint = gr.Textbox(label="🔌 Custom API Endpoint URL",
-                                                 placeholder="https://your-api-endpoint.com", visible=False)
-                custom_api_key = gr.Textbox(label="🔐 Custom API Endpoint Key", type="password", visible=False)
-                search_provider = gr.Dropdown(label="🔎 Search Provider", choices=["serper", "searxng"], value="serper")
-                search_api_key = gr.Textbox(label="🔑 Search Provider API Key (optional)", type="password",
-                                            visible=False)
-                custom_search_url = gr.Textbox(label="🌐 Custom SearxNG Instance URL",
-                                               placeholder="https://your-searxng-instance/search", visible=False)
-                submit_btn = gr.Button("🚀 Run Agent")
-
-            with gr.Column(scale=1):
-                output = gr.Textbox(label="📤 Answer", lines=15)
-
-
-        def update_search_visibility(provider):
-            """Updates the visibility of search-related textboxes based on the selected search provider."""
-            return {
-                serpapi_key: gr.update(visible=(provider == "serper")),
-                custom_search_url: gr.update(visible=(provider == "searxng")),
-                search_api_key: gr.update(visible=(provider == "searxng")),
-            }
-
-
-        def update_custom_endpoint_visibility(checked):
-            """Updates the visibility of custom API endpoint textboxes based on the checkbox state."""
-            return {
-                custom_api_endpoint: gr.update(visible=checked),
-                custom_api_key: gr.update(visible=checked),
-            }
-
-
-        search_provider.change(fn=update_search_visibility, inputs=search_provider,
-                               outputs=[serpapi_key, custom_search_url, search_api_key])
-        use_custom_endpoint.change(fn=update_custom_endpoint_visibility, inputs=use_custom_endpoint,
-                                   outputs=[custom_api_endpoint, custom_api_key])
-
-        submit_btn.click(fn=setup_agent,
-                         inputs=[question, model_id, hf_token, serpapi_key,
-                                 use_custom_endpoint, custom_api_endpoint, custom_api_key,
-                                 search_provider, search_api_key, custom_search_url],
-                         outputs=output)
-
-    print("[DEBUG] Launching Gradio interface")
-    demo.launch()
-
+    return demo
 
 if __name__ == "__main__":
-    launch_interface()
+    demo = build_app()
+    demo.launch(server_name="0.0.0.0")