AI_SQL / app.py
mgbam's picture
Update app.py
8c14ec2 verified
"""
app.py – Enterprise SQL Agent (Gradio + smolagents + MCP)
HubSpot Integration Only
"""
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 ──────────────────────────
OPENAI_KEY = os.getenv("OPENAI_API_KEY")
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")
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, timeout=90)
ACTIVE = f"Hugging Face Β· {HF_MODEL_ID}"
# ───────────────────────── 2. MCP server path ──────────────────────────
SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")
# ───────────────────────── 3. Chat callback ────────────────────────────
def respond(message: str, history: list):
"""Prompt β†’ CodeAgent β†’ MCP tools β†’ string reply."""
params = StdioServerParameters(command="python", args=[str(SERVER_PATH)])
try:
with MCPClient(params) as tools:
answer = CodeAgent(tools=tools, model=BASE_MODEL).run(message)
except Exception as e:
answer = f"Error while querying tools: {e}"
# ensure plain-text output
if not isinstance(answer, str):
try:
answer = json.dumps(answer, indent=2, ensure_ascii=False)
except Exception:
answer = pprint.pformat(answer, width=100)
history += [
{"role": "user", "content": message},
{"role": "assistant", "content": answer},
]
return history, history
# ───────────────────────── 4. Gradio UI ────────────────────────────────
with gr.Blocks(title="Enterprise SQL Agent") as demo:
state = gr.State([])
gr.Markdown("## 🏒 Enterprise SQL Agent β€” query your data with natural language")
chat = gr.Chatbot(type="messages", label="Conversation")
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):
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 Β· smolagents Β· Gradio β€’ Active model β†’ **{ACTIVE}**_")
if __name__ == "__main__":
demo.launch()