Update app.py

app.py

@@ -1,80 +1,76 @@
 """
-app.py — Gradio front‑end + smolagents CodeAgent
-================================================
-
-This version avoids private/gated models and works on any free Hugging Face
-Space **without extra secrets**.  It relies on:
-
-* `mcp_server.py` sitting next to this file
-* A public chat‑completion capable model exposed via the HF Inference API
-  (defaults to **microsoft/Phi‑3‑mini‑4k‑instruct**, ~3 B params, free‑tier‑OK)
-* `smolagents[mcp]` for the agent loop
-* **Optional**: set `HF_MODEL_ID` or `HF_API_TOKEN` in **Settings → Secrets**
-  if you want a different (or gated) model.
-
-If you hit the free‑tier rate‑limit you can still point to OpenAI by setting the
-env var `OPENAI_API_KEY` — the code will auto‑switch to OpenAI chat.
+Gradio front-end + smolagents CodeAgent
+---------------------------------------
+• If you set an OPENAI_API_KEY the agent will call OpenAI (GPT-4o by default).
+• Otherwise it falls back to a free Hugging Face chat-completion model
+  (defaults to microsoft/Phi-3-mini-4k-instruct which the public Inference
+  API exposes for the Chat Completion task).
+• You can override the fallback by defining HF_MODEL_ID and, if needed,
+  HF_API_TOKEN in the Space → Settings → Secrets.
 """
-import os
-import pathlib
-import gradio as gr
+
+import os, pathlib, gradio as gr
 from mcp import StdioServerParameters
 from smolagents import MCPClient, CodeAgent, InferenceClientModel
 
-# ---------- Tool server ------------------------------------------------------
+# Path to your MCP tool server
 SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
 
-# ---------- Model selection --------------------------------------------------
-# 1) Use OpenAI automatically if OPENAI_API_KEY is set.
-# 2) Otherwise fall back to a public HF Inference model that supports chat‑completion.
-OPENAI_KEY = os.getenv("OPENAI_API_KEY")
-HF_MODEL_ID = os.getenv("HF_MODEL_ID", "microsoft/Phi-3-mini-4k-instruct")
-
-if OPENAI_KEY:
-    from smolagents.models import OpenAIChatModel  # lazy import only if needed
-    BASE_MODEL = OpenAIChatModel()  # defaults gpt‑4o‑preview
-else:
-    BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID)
+# Decide which base model to use
+OPENAI_KEY   = os.getenv("OPENAI_API_KEY")
+HF_MODEL_ID  = os.getenv("HF_MODEL_ID", "microsoft/Phi-3-mini-4k-instruct")
 
+if OPENAI_KEY:                          # --- OpenAI branch -------------
+    from smolagents.models import OpenAIChat           # <- works in all versions
+    BASE_MODEL = OpenAIChat(                           # gpt-4o by default
+        model=os.getenv("OPENAI_MODEL", "gpt-4o-preview"),
+        temperature=0.3,
+    )
+else:                                     # --- Hugging Face branch ----
+    BASE_MODEL = InferenceClientModel(               # uses HF Inference API
+        model_id=HF_MODEL_ID,
+        hf_api_token=os.getenv("HF_API_TOKEN"),      # optional for gated repos
+        timeout=90,
+    )
 
-# ---------- Gradio callback ---------------------------------------------------
+# ----------------- callback ---------------------------------------------------
 def respond(message: str, history: list):
-    """Run the user prompt through a CodeAgent that can call MCP SQL tools."""
+    """Send user prompt → CodeAgent → SQL tools → natural-language answer."""
     params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])
-
     with MCPClient(params) as tools:
-        agent = CodeAgent(tools=tools, model=BASE_MODEL)
+        agent  = CodeAgent(tools=tools, model=BASE_MODEL)
         answer = agent.run(message)
 
-    # Append to chat history (OpenAI messages format)
-    history.append({"role": "user", "content": message})
-    history.append({"role": "assistant", "content": answer})
+    history += [
+        {"role": "user",      "content": message},
+        {"role": "assistant", "content": answer},
+    ]
     return history, history
 
-
-# ---------- UI ---------------------------------------------------------------
+# ----------------- UI ---------------------------------------------------------
 with gr.Blocks(title="Enterprise SQL Agent") as demo:
-    chat_state = gr.State([])
-    gr.Markdown("## Enterprise SQL Agent — ask natural‑language questions about your data 🏢➡️📊")
+    state   = gr.State([])
+    gr.Markdown("## Enterprise SQL Agent — ask questions about your data 🏢➡️📊")
 
-    chatbot = gr.Chatbot(type="messages", label="Chat")
-    textbox = gr.Textbox(placeholder="e.g. Who are my inactive Northeast customers?", show_label=False)
-    textbox.submit(respond, [textbox, chat_state], [chatbot, chat_state])
+    chat    = gr.Chatbot(type="messages", label="Chat")
+    box     = gr.Textbox(
+        placeholder="e.g. Who are my inactive Northeast customers?",
+        show_label=False,
+    )
+    box.submit(respond, [box, state], [chat, state])
 
     with gr.Accordion("Example prompts"):
         gr.Markdown(
-            """
-            * Who are my **Northeast** customers with no orders in 6 months?
-            * List customers sorted by **LastOrderDate**.
-            * Draft re‑engagement emails for inactive accounts.
-            """
+            "* Who are my **Northeast** customers with no orders in 6 months?\n"
+            "* List customers sorted by **LastOrderDate**.\n"
+            "* Draft re-engagement emails for inactive accounts."
         )
 
-    gr.Markdown(
-        "_Powered by MCP + smolagents + Gradio • Model: {}_".format(
-            "OpenAI (gpt‑4o)" if OPENAI_KEY else HF_MODEL_ID
-        )
+    footer = (
+        f"_Powered by MCP + smolagents + Gradio • Model: "
+        f"{'OpenAI' if OPENAI_KEY else HF_MODEL_ID}_"
     )
+    gr.Markdown(footer)
 
 if __name__ == "__main__":
     demo.launch()