app.py

"""
app.py – Enterprise SQL Agent (Gradio + smolagents + MCP)

Provider priority
──────────────────
1. OpenAI   → set OPENAI_API_KEY   (override model with OPENAI_MODEL, default = gpt-4o)
2. Gemini   → set GOOGLE_API_KEY  (override model with GOOGLE_MODEL, default = gemini-pro)
3. Hugging Face Inference fallback
     • HF_MODEL_ID   (default = microsoft/Phi-3-mini-4k-instruct)
     • HF_API_TOKEN  (only if the repo is gated)

File layout
────────────
app.py
mcp_server.py
connectors/
    └─ salesforce_connector.py
requirements.txt
"""

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")          # stable id

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")  # optional

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)])
    with MCPClient(params) as tools:
        answer = CodeAgent(tools=tools, model=BASE_MODEL).run(message)

    # 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()