Update app.py

app.py

@@ -1,82 +1,75 @@
 """
-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
+app.py – Enterprise SQL Agent (Gradio + smolagents + MCP)
+
+Secrets / ENV
+─────────────
+OPENAI_API_KEY   → calls OpenAI  (default model gpt-4o, override via OPENAI_MODEL)
+GOOGLE_API_KEY   → calls Gemini (default model gemini-pro, override via GOOGLE_MODEL)
+HF_MODEL_ID      → Hugging Face chat-completion model (fallback if no keys)
+HF_API_TOKEN     → token for gated HF repo (optional)
 """
 
-import os, pathlib, gradio as gr
+import os, pathlib, json, pprint, gradio as gr
 from mcp import StdioServerParameters
 from smolagents import MCPClient, CodeAgent
 from smolagents.models import LiteLLMModel, InferenceClientModel
 
-# ─────────── 1. Choose base LLM ──────────────────────────────────────────
+# ─── 1. Pick the base LLM ───────────────────────────────────────────────
 OPENAI_KEY   = os.getenv("OPENAI_API_KEY")
-OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")          # change if not whitelisted
+OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")
 
 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
+HF_TOKEN     = os.getenv("HF_API_TOKEN")  # only if the repo is gated
 
 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=f"openai/{OPENAI_MODEL}",
+                              api_key=OPENAI_KEY)
+    ACTIVE = f"OpenAI · {OPENAI_MODEL}"
 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=f"google/{GEM_MODEL}",
+                              api_key=GEMINI_KEY)
+    ACTIVE = f"Gemini · {GEM_MODEL}"
 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,
+                                      hf_api_token=HF_TOKEN,
+                                      timeout=90)
+    ACTIVE = f"Hugging Face · {HF_MODEL_ID}"
 
-# ─────────── 2. Path to MCP tool server ──────────────────────────────────
+# ─── 2. Path to MCP server ──────────────────────────────────────────────
 SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
 
-# ─────────── 3. Gradio callback ──────────────────────────────────────────
-def respond(msg: str, history: list):
-    """Run prompt → CodeAgent → MCP tools → safe string reply."""
+# ─── 3. Chat callback ───────────────────────────────────────────────────
+def respond(message: str, history: list):
+    """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)
+        answer = CodeAgent(tools=tools, model=BASE_MODEL).run(message)
 
-    # Ensure reply is always string for Chatbot
-    if not isinstance(raw, str):
-        import json, pprint
+    # Always stringify for Gradio
+    if not isinstance(answer, str):
         try:
-            raw = json.dumps(raw, indent=2, ensure_ascii=False)
-        except (TypeError, ValueError):
-            raw = pprint.pformat(raw)
-    reply = raw
+            answer = json.dumps(answer, indent=2, ensure_ascii=False)
+        except Exception:
+            answer = pprint.pformat(answer, width=100)
 
     history += [
-        {"role": "user",      "content": msg},
-        {"role": "assistant", "content": reply},
+        {"role": "user",      "content": message},
+        {"role": "assistant", "content": answer},
     ]
     return history, history
 
-# ───��─────── 4. Build the UI ─────────────────────────────────────────────
+# ─── 4. Build 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")
+    gr.Markdown("## 🏢 Enterprise SQL Agent — ask 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  = 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", open=False):