Update app.py

app.py

@@ -1,76 +1,58 @@
 """
-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.
+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, InferenceClientModel
+from smolagents import MCPClient, CodeAgent, InferenceClientModel, LiteLLMModel
 
-# Path to your MCP tool server
+# Path to mcp_server.py (must be beside this file)
 SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
 
-# Decide which base model to use
+# ---------- Model resolution ----------
 OPENAI_KEY   = os.getenv("OPENAI_API_KEY")
+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")
 
-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,
-    )
-
-# ----------------- callback ---------------------------------------------------
+if OPENAI_KEY:
+    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="google/gemini-pro", api_key=GEMINI_KEY)
+    MODEL_NAME = "google/gemini-pro"
+else:
+    BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID, timeout=90)
+    MODEL_NAME = HF_MODEL_ID
+
+# ---------- Callback ----------
 def respond(message: str, history: list):
-    """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)
         answer = agent.run(message)
-
     history += [
-        {"role": "user",      "content": message},
+        {"role": "user", "content": message},
         {"role": "assistant", "content": answer},
     ]
     return history, history
 
-# ----------------- UI ---------------------------------------------------------
+# ---------- UI ----------
 with gr.Blocks(title="Enterprise SQL Agent") as demo:
-    state   = gr.State([])
-    gr.Markdown("## Enterprise SQL Agent — ask questions about your data 🏢➡️📊")
-
-    chat    = gr.Chatbot(type="messages", label="Chat")
-    box     = gr.Textbox(
-        placeholder="e.g. Who are my inactive Northeast customers?",
-        show_label=False,
-    )
+    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"):
-        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("""* 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.""")
 
-    footer = (
-        f"_Powered by MCP + smolagents + Gradio • Model: "
-        f"{'OpenAI' if OPENAI_KEY else HF_MODEL_ID}_"
-    )
-    gr.Markdown(footer)
+    gr.Markdown(f"_Powered by MCP + smolagents • Model: **{MODEL_NAME}**_")
 
 if __name__ == "__main__":
     demo.launch()