Update app.py

app.py

@@ -1,63 +1,50 @@
 """
-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.
+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).
 """
-
 import os, pathlib, gradio as gr
 from mcp import StdioServerParameters
-from smolagents import MCPClient, CodeAgent
-from smolagents.models import InferenceClientModel, LiteLLMModel
+from smolagents import MCPClient, CodeAgent, InferenceClientModel, LiteLLMModel
 
-# --------------------------------------------------------------------- #
-# 1.  Resolve the base LLM model, preferring OpenAI → Gemini → HF
-# --------------------------------------------------------------------- #
+# Path to mcp_server.py (must be beside this file)
+SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
+
+# ---------- Model resolution ----------
 OPENAI_KEY   = os.getenv("OPENAI_API_KEY")
-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")
+GEMINI_KEY   = os.getenv("GOOGLE_API_KEY")  # for Gemini via LiteLLM
 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=f"openai/{OPENAI_MODEL}", api_key=OPENAI_KEY)
-    ACTIVE     = f"OpenAI · {OPENAI_MODEL}"
+    BASE_MODEL = LiteLLMModel(model_id="openai/gpt-4o-preview", api_key=OPENAI_KEY)
+    MODEL_NAME = "openai/gpt-4o-preview"
 elif GEMINI_KEY:
-    BASE_MODEL = LiteLLMModel(model_id=f"google/{GEM_MODEL}", api_key=GEMINI_KEY)
-    ACTIVE     = f"Gemini · {GEM_MODEL}"
+    BASE_MODEL = LiteLLMModel(model_id="google/gemini-pro", api_key=GEMINI_KEY)
+    MODEL_NAME = "google/gemini-pro"
 else:
-    BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID, hf_api_token=HF_TOKEN)
-    ACTIVE     = f"Hugging Face · {HF_MODEL_ID}"
+    BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID, timeout=90)
+    MODEL_NAME = HF_MODEL_ID
 
-# --------------------------------------------------------------------- #
-# 2.  Path to your MCP server (must be alongside this file)
-# --------------------------------------------------------------------- #
-SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
-
-# --------------------------------------------------------------------- #
-# 3.  Gradio callback – run prompt → CodeAgent → SQL tools
-# --------------------------------------------------------------------- #
+# ---------- Callback ----------
 def respond(msg: str, chat_history: list):
+    """Run prompt → CodeAgent → string reply safe for Gradio Chatbot."""
     params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])
 
     with MCPClient(params) as tools:
         agent  = CodeAgent(tools=tools, model=BASE_MODEL)
-        reply  = agent.run(msg)
+        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 += [
         {"role": "user",      "content": msg},
@@ -65,28 +52,18 @@ def respond(msg: str, chat_history: list):
     ]
     return chat_history, chat_history
 
-# --------------------------------------------------------------------- #
-# 4.  Build the UI
-# --------------------------------------------------------------------- #
+# ---------- UI ----------
 with gr.Blocks(title="Enterprise SQL Agent") as demo:
-    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?",
-    )
+    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?")
     box.submit(respond, [box, state], [chat, state])
 
-    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."
-        )
+    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.""")
 
-    gr.Markdown(f"_Powered by MCP tools + smolagents · Active model → **{ACTIVE}**_")
+    gr.Markdown(f"_Powered by MCP + smolagents • Model: **{MODEL_NAME}**_")
 
 if __name__ == "__main__":
     demo.launch()