app.py

"""
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, json, pprint, gradio as gr
from mcp import StdioServerParameters
from smolagents import MCPClient, CodeAgent
from smolagents.models import LiteLLMModel, InferenceClientModel

# ─── 1. Pick the 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")  # 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}"
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. Path to MCP server ──────────────────────────────────────────────
SERVER_PATH = pathlib.Path(__file__).with_name("mcp_server.py")

# ─── 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:
        answer = CodeAgent(tools=tools, model=BASE_MODEL).run(message)

    # Always stringify for Gradio
    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. Build 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="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()