Update app.py

app.py

@@ -1,69 +1,92 @@
 """
-app.py — robust multimodel agent
---------------------------------
-* Supports **OpenAI** (if `OPENAI_API_KEY`) or **Gemini** (if `GOOGLE_API_KEY`) via
-  `LiteLLMModel`.
-* Otherwise falls back to a free HF Inference chat model
-  (`microsoft/Phi-3-mini-4k-instruct`).
-* No version‑specific imports (avoids `OpenAIChat` errors).
+app.py  –  Enterprise SQL Agent (Gradio + smolagents + MCP)
+
+SECRETS / ENV VARS
+------------------
+OPENAI_API_KEY   ← use OpenAI (default model gpt-4o, override with OPENAI_MODEL)
+GOOGLE_API_KEY   ← use Gemini-Pro (override model with GOOGLE_MODEL)
+HF_MODEL_ID      ← repo that exposes Chat-Completion (fallback if no keys)
+HF_API_TOKEN     ← token if that repo is gated
+
+FILE LAYOUT
+-----------
+app.py
+mcp_server.py      # your FastMCP SQL tool server
+requirements.txt   # see bottom of this file
 """
+
 import os, pathlib, gradio as gr
 from mcp import StdioServerParameters
-from smolagents import MCPClient, CodeAgent, InferenceClientModel, LiteLLMModel
+from smolagents import MCPClient, CodeAgent
+from smolagents.models import LiteLLMModel, InferenceClientModel
 
-# Path to mcp_server.py (must be beside this file)
-SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
-
-# ---------- Model resolution ----------
+# ─────────── 1. Choose base LLM ──────────────────────────────────────────
 OPENAI_KEY   = os.getenv("OPENAI_API_KEY")
-GEMINI_KEY   = os.getenv("GOOGLE_API_KEY")  # for Gemini via LiteLLM
+OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")          # change if not whitelisted
+
+GEMINI_KEY   = os.getenv("GOOGLE_API_KEY")
+GEM_MODEL    = os.getenv("GOOGLE_MODEL", "gemini-pro")
+
 HF_MODEL_ID  = os.getenv("HF_MODEL_ID", "microsoft/Phi-3-mini-4k-instruct")
+HF_TOKEN     = os.getenv("HF_API_TOKEN")  # only for gated repos
 
 if OPENAI_KEY:
-    BASE_MODEL = LiteLLMModel(model_id="openai/gpt-4o-preview", api_key=OPENAI_KEY)
-    MODEL_NAME = "openai/gpt-4o-preview"
+    BASE_MODEL = LiteLLMModel(model_id=f"openai/{OPENAI_MODEL}", api_key=OPENAI_KEY)
+    ACTIVE     = f"OpenAI · {OPENAI_MODEL}"
 elif GEMINI_KEY:
-    BASE_MODEL = LiteLLMModel(model_id="google/gemini-pro", api_key=GEMINI_KEY)
-    MODEL_NAME = "google/gemini-pro"
+    BASE_MODEL = LiteLLMModel(model_id=f"google/{GEM_MODEL}", api_key=GEMINI_KEY)
+    ACTIVE     = f"Gemini · {GEM_MODEL}"
 else:
-    BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID, timeout=90)
-    MODEL_NAME = HF_MODEL_ID
+    BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID, hf_api_token=HF_TOKEN)
+    ACTIVE     = f"Hugging Face · {HF_MODEL_ID}"
+
+# ─────────── 2. Path to MCP tool server ──────────────────────────────────
+SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
 
-# ---------- Callback ----------
-def respond(msg: str, chat_history: list):
-    """Run prompt → CodeAgent → string reply safe for Gradio Chatbot."""
+# ─────────── 3. Gradio callback ──────────────────────────────────────────
+def respond(msg: str, history: list):
+    """Run prompt → CodeAgent → MCP tools → safe string reply."""
     params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])
 
     with MCPClient(params) as tools:
         agent  = CodeAgent(tools=tools, model=BASE_MODEL)
         raw    = agent.run(msg)
-        # Ensure the reply is always a string to avoid Gradio validation errors
-        if not isinstance(raw, str):
-            import json, pprint
-            try:
-                raw = json.dumps(raw, indent=2)
-            except TypeError:
-                raw = pprint.pformat(raw)
-        reply = raw
-
-    chat_history += [
+
+    # Ensure reply is always string for Chatbot
+    if not isinstance(raw, str):
+        import json, pprint
+        try:
+            raw = json.dumps(raw, indent=2, ensure_ascii=False)
+        except (TypeError, ValueError):
+            raw = pprint.pformat(raw)
+    reply = raw
+
+    history += [
         {"role": "user",      "content": msg},
         {"role": "assistant", "content": reply},
     ]
-    return chat_history, chat_history
+    return history, history
 
-# ---------- UI ----------
+# ─────────── 4. Build the UI ─────────────────────────────────────────────
 with gr.Blocks(title="Enterprise SQL Agent") as demo:
     state = gr.State([])
-    gr.Markdown("## Enterprise SQL Agent — natural‑language to SQL via MCP")
-    chat  = gr.Chatbot(type="messages", label="Chat")
-    box   = gr.Textbox(show_label=False, placeholder="Ask: Who are my inactive Northeast customers?")
+    gr.Markdown("## 🏢 Enterprise SQL Agent — ask natural-language questions about your data")
+
+    chat = gr.Chatbot(type="messages", label="Conversation")
+    box  = gr.Textbox(
+        placeholder="e.g. Who are my Northeast customers with no orders in 6 months?",
+        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?\n* List customers sorted by **LastOrderDate**.\n* Draft re‑engagement emails for inactive accounts.""")
+    with gr.Accordion("Example prompts", open=False):
+        gr.Markdown(
+            "* 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(f"_Powered by MCP + smolagents • Model: **{MODEL_NAME}**_")
+    gr.Markdown(f"_Powered by MCP + smolagents + Gradio • Active model → **{ACTIVE}**_")
 
 if __name__ == "__main__":
     demo.launch()