AI_SQL / app.py
mgbam's picture
Update app.py
a2913a9 verified
raw
history blame
3.3 kB
"""
app.py — Gradio front‑end + smolagents CodeAgent
================================================
This version avoids private/gated models and works on any free Hugging Face
Space **without extra secrets**. It relies on:
* `mcp_server.py` sitting next to this file
* A public chat‑completion capable model exposed via the HF Inference API
(defaults to **microsoft/Phi‑3‑mini‑4k‑instruct**, ~3 B params, free‑tier‑OK)
* `smolagents[mcp]` for the agent loop
* **Optional**: set `HF_MODEL_ID` or `HF_API_TOKEN` in **Settings → Secrets**
if you want a different (or gated) model.
If you hit the free‑tier rate‑limit you can still point to OpenAI by setting the
env var `OPENAI_API_KEY` — the code will auto‑switch to OpenAI chat.
"""
import os
import pathlib
import gradio as gr
from mcp import StdioServerParameters
from smolagents import MCPClient, CodeAgent, InferenceClientModel
# ---------- Tool server ------------------------------------------------------
SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
# ---------- Model selection --------------------------------------------------
# 1) Use OpenAI automatically if OPENAI_API_KEY is set.
# 2) Otherwise fall back to a public HF Inference model that supports chat‑completion.
OPENAI_KEY = os.getenv("OPENAI_API_KEY")
HF_MODEL_ID = os.getenv("HF_MODEL_ID", "microsoft/Phi-3-mini-4k-instruct")
if OPENAI_KEY:
from smolagents.models import OpenAIChatModel # lazy import only if needed
BASE_MODEL = OpenAIChatModel() # defaults gpt‑4o‑preview
else:
BASE_MODEL = InferenceClientModel(model_id=HF_MODEL_ID)
# ---------- Gradio callback ---------------------------------------------------
def respond(message: str, history: list):
"""Run the user prompt through a CodeAgent that can call MCP SQL tools."""
params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])
with MCPClient(params) as tools:
agent = CodeAgent(tools=tools, model=BASE_MODEL)
answer = agent.run(message)
# Append to chat history (OpenAI messages format)
history.append({"role": "user", "content": message})
history.append({"role": "assistant", "content": answer})
return history, history
# ---------- UI ---------------------------------------------------------------
with gr.Blocks(title="Enterprise SQL Agent") as demo:
chat_state = gr.State([])
gr.Markdown("## Enterprise SQL Agent — ask natural‑language questions about your data 🏢➡️📊")
chatbot = gr.Chatbot(type="messages", label="Chat")
textbox = gr.Textbox(placeholder="e.g. Who are my inactive Northeast customers?", show_label=False)
textbox.submit(respond, [textbox, chat_state], [chatbot, chat_state])
with gr.Accordion("Example prompts"):
gr.Markdown(
"""
* Who are my **Northeast** customers with no orders in 6 months?
* List customers sorted by **LastOrderDate**.
* Draft re‑engagement emails for inactive accounts.
"""
)
gr.Markdown(
"_Powered by MCP + smolagents + Gradio • Model: {}_".format(
"OpenAI (gpt‑4o)" if OPENAI_KEY else HF_MODEL_ID
)
)
if __name__ == "__main__":
demo.launch()