File size: 3,842 Bytes
a2913a9 05f1269 a2913a9 05f1269 3b69291 c3c6e00 05f1269 5997785 05f1269 3b69291 05f1269 3b69291 05f1269 131111c 865daca 05f1269 865daca 05f1269 865daca 05f1269 865daca 05f1269 99137ea 05f1269 c3c6e00 3b69291 05f1269 3b69291 05f1269 c3c6e00 05f1269 99137ea 05f1269 3b69291 5997785 05f1269 a2913a9 05f1269 c3c6e00 99137ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
"""
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
from smolagents.models import InferenceClientModel, LiteLLMModel
# --------------------------------------------------------------------- #
# 1. Resolve the base LLM model, preferring OpenAI → Gemini → HF
# --------------------------------------------------------------------- #
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")
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}"
elif GEMINI_KEY:
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}"
# --------------------------------------------------------------------- #
# 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
# --------------------------------------------------------------------- #
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)
reply = agent.run(msg)
chat_history += [
{"role": "user", "content": msg},
{"role": "assistant", "content": reply},
]
return chat_history, chat_history
# --------------------------------------------------------------------- #
# 4. Build the 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?",
)
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."
)
gr.Markdown(f"_Powered by MCP tools + smolagents · Active model → **{ACTIVE}**_")
if __name__ == "__main__":
demo.launch()
|