Update app.py

app.py

@@ -1,58 +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)
+
+ENV VARS SUPPORTED
+------------------
+OPENAI_API_KEY   → use any OpenAI chat model (default  = gpt-4o)
+OPENAI_MODEL     → override the OpenAI model ID   (e.g. gpt-4-turbo)
+GOOGLE_API_KEY   → use Gemini-Pro via Google AI
+GOOGLE_MODEL     → override Gemini model ID      (e.g. gemini-1.5-pro)
+
+HF_MODEL_ID      → fallback model repo (must expose Chat-Completion)
+HF_API_TOKEN     → token if that repo is gated
+
+If no provider keys are set, the code falls back to a free HF model
+`microsoft/Phi-3-mini-4k-instruct`.
+
+MCP SERVER
+----------
+`mcp_server.py` must live in the same folder and expose your SQL tools.
 """
+
 import os, pathlib, gradio as gr
 from mcp import StdioServerParameters
-from smolagents import MCPClient, CodeAgent, InferenceClientModel, LiteLLMModel
-
-# Path to mcp_server.py (must be beside this file)
-SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
+from smolagents import MCPClient, CodeAgent
+from smolagents.models import InferenceClientModel, LiteLLMModel
 
-# ---------- Model resolution ----------
+# --------------------------------------------------------------------- #
+# 1.  Resolve the base LLM model, preferring OpenAI → Gemini → HF
+# --------------------------------------------------------------------- #
 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") or None            # only needed 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 your MCP server (must be alongside this file)
+# --------------------------------------------------------------------- #
+SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
 
-# ---------- Callback ----------
-def respond(message: str, history: list):
+# --------------------------------------------------------------------- #
+# 3.  Gradio callback – run prompt → CodeAgent → SQL tools
+# --------------------------------------------------------------------- #
+def respond(msg: str, chat_history: list):
     params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])
+
     with MCPClient(params) as tools:
         agent  = CodeAgent(tools=tools, model=BASE_MODEL)
-        answer = agent.run(message)
-    history += [
-        {"role": "user", "content": message},
-        {"role": "assistant", "content": answer},
+        reply  = agent.run(msg)
+
+    chat_history += [
+        {"role": "user",      "content": msg},
+        {"role": "assistant", "content": reply},
     ]
-    return history, history
+    return chat_history, chat_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?")
+    state  = gr.State([])
+    gr.Markdown("## 🏢 Enterprise SQL Agent – ask natural-language questions about your data")
+
+    chat   = gr.Chatbot(type="messages", label="Chat")
+    box    = gr.Textbox(
+        show_label=False,
+        placeholder="e.g. Who are my Northeast customers with no orders in 6 months?",
+    )
     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 tools + smolagents · Active model → **{ACTIVE}**_")
 
 if __name__ == "__main__":
     demo.launch()