app.py

import gradio as gr
from typing import List, Dict, Set, Callable, Any
import json
import uuid
from pathlib import Path

# Azure AI Agents SDK (API key auth)
from azure.core.credentials import AzureKeyCredential
from azure.ai.agents import AgentsClient
from azure.ai.agents.models import (
    FunctionTool,
    ToolSet,
    ListSortOrder,
    MessageRole,
)

# ---------- Inline custom function(s) so no separate user_functions.py is needed ----------

def submit_support_ticket(email_address: str, description: str) -> str:
    """
    Generates a ticket number and saves a support ticket as a text file
    in the current app directory. Returns a small JSON message string.
    """
    script_dir = Path(__file__).parent
    ticket_number = str(uuid.uuid4()).replace("-", "")[:6]
    file_name = f"ticket-{ticket_number}.txt"
    file_path = script_dir / file_name

    text = (
        f"Support ticket: {ticket_number}\n"
        f"Submitted by: {email_address}\n"
        f"Description:\n{description}\n"
    )
    file_path.write_text(text, encoding="utf-8")

    message_json = json.dumps({
        "message": f"Support ticket {ticket_number} submitted. The ticket file is saved as {file_name}"
    })
    return message_json

# Define the callable tool set
user_functions: Set[Callable[..., Any]] = { submit_support_ticket }

# ---------- Core Agent Helpers ----------

def init_agent(endpoint: str, api_key: str, model_deployment: str) -> dict:
    """
    Initialize an Azure AI Agent with a custom FunctionTool.
    Returns a session dict containing client, agent_id, thread_id, etc.
    """
    if not endpoint or not api_key or not model_deployment:
        raise ValueError("Please provide endpoint, key, and model deployment name.")

    client = AgentsClient(
        endpoint=endpoint.strip(),
        credential=AzureKeyCredential(api_key.strip()),
    )

    # Build toolset with your functions and enable auto function calls
    functions_tool = FunctionTool(user_functions)
    toolset = ToolSet()
    toolset.add(functions_tool)
    client.enable_auto_function_calls(toolset)

    # Create the agent
    agent = client.create_agent(
        model=model_deployment.strip(),
        name="support-agent",
        instructions=(
            "You are a technical support agent. "
            "When a user has a technical issue, ask for their email address and a description "
            "of the issue. Then submit a support ticket using the available function. "
            "If a file is saved, tell the user the file name."
        ),
        toolset=toolset,
    )

    # Create a thread for the conversation
    thread = client.threads.create()

    # Session we keep in Gradio state
    return {
        "endpoint": endpoint.strip(),
        "api_key": api_key.strip(),
        "model": model_deployment.strip(),
        "client": client,
        "agent_id": agent.id,
        "thread_id": thread.id,
    }


def send_to_agent(user_msg: str, session: dict):
    """
    Send message to the agent thread and return:
    - agent_reply (str)
    - history_str (str) chronological log
    """
    if not session or "client" not in session:
        raise ValueError("Agent is not initialized. Click 'Connect & Prepare' first.")

    client: AgentsClient = session["client"]
    agent_id = session["agent_id"]
    thread_id = session["thread_id"]

    # Add user message
    client.messages.create(
        thread_id=thread_id,
        role="user",
        content=user_msg,
    )

    # Run and wait (auto function-calls enabled)
    run = client.runs.create_and_process(thread_id=thread_id, agent_id=agent_id)
    if getattr(run, "status", None) == "failed":
        last_error = getattr(run, "last_error", "Unknown error")
        return f"Run failed: {last_error}", ""

    # Last agent message
    last_msg = client.messages.get_last_message_text_by_role(
        thread_id=thread_id,
        role=MessageRole.AGENT,
    )
    agent_reply = last_msg.text.value if last_msg else "(No reply text found.)"

    # Build readable history (chronological)
    history_lines = []
    messages = client.messages.list(thread_id=thread_id, order=ListSortOrder.ASCENDING)
    for m in messages:
        if m.text_messages:
            last_text = m.text_messages[-1].text.value
            history_lines.append(f"{m.role}: {last_text}")
    history_str = "\n\n".join(history_lines)

    return agent_reply, history_str


def teardown(session: dict) -> str:
    """
    Delete the agent to reduce costs. (Threads are retained by the service.)
    """
    if not session:
        return "Nothing to clean up."

    notes = []
    try:
        client: AgentsClient = session.get("client")
        agent_id = session.get("agent_id")
        if client and agent_id:
            client.delete_agent(agent_id)
            notes.append("Deleted agent.")
    except Exception as e:
        notes.append(f"Cleanup warning: {e}")

    return " ".join(notes) if notes else "Cleanup complete."


# ---------- Gradio UI ----------

with gr.Blocks(title="Azure AI Support Agent (Functions) — Gradio") as demo:
    gr.Markdown(
        "## Azure AI Support Agent (Custom Function Tool)\n"
        "Enter your **Project Endpoint** and **Key**, set your **Model Deployment** (e.g., `gpt-4o`), then chat.\n"
        "The agent can call a custom function to **submit a support ticket** and save it as a file."
    )

    with gr.Row():
        endpoint = gr.Textbox(label="Project Endpoint", placeholder="https://<your-project-endpoint>")
        api_key = gr.Textbox(label="Project Key", placeholder="paste your key", type="password")

    with gr.Row():
        model = gr.Textbox(label="Model Deployment Name", value="gpt-4o")

    session_state = gr.State(value=None)

    connect_btn = gr.Button("🔌 Connect & Prepare Agent", variant="primary")
    connect_status = gr.Markdown("")

    with gr.Row():
        chatbot = gr.Chatbot(
            label="Conversation",
            height=420,
            type="messages",  # openai-style dicts: {"role": "...", "content": "..."}
        )

    user_input = gr.Textbox(label="Your message", placeholder="e.g., I have a technical problem")
    with gr.Row():
        send_btn = gr.Button("Send ▶")
        cleanup_btn = gr.Button("Delete Agent & Cleanup 🧹")

    history = gr.Textbox(label="Conversation Log (chronological)", lines=12)

    # ----- Callbacks -----

    def on_connect(ep, key, mdl):
        try:
            sess = init_agent(ep, key, mdl)
            return sess, "✅ Connected. Support agent and thread are ready."
        except Exception as e:
            return None, f"❌ Connection error: {e}"

    connect_btn.click(
        fn=on_connect,
        inputs=[endpoint, api_key, model],
        outputs=[session_state, connect_status],
    )

    def on_send(msg: str, session: dict, chat_msgs: List[Dict[str, str]]):
        if not msg:
            return gr.update(), gr.update(), gr.update(value="Please enter a message.")
        try:
            agent_reply, log = send_to_agent(msg, session)
            chat_msgs = (chat_msgs or []) + [
                {"role": "user", "content": msg},
                {"role": "assistant", "content": agent_reply},
            ]
            return chat_msgs, "", gr.update(value=log)  # clear input, update log
        except Exception as e:
            return chat_msgs, msg, gr.update(value=f"❌ Error: {e}")

    send_btn.click(
        fn=on_send,
        inputs=[user_input, session_state, chatbot],
        outputs=[chatbot, user_input, history],
    )

    def on_cleanup(session):
        try:
            msg = teardown(session)
            return None, f"🧹 {msg}"
        except Exception as e:
            return session, f"⚠️ Cleanup error: {e}"

    cleanup_btn.click(
        fn=on_cleanup,
        inputs=[session_state],
        outputs=[session_state, connect_status],
    )

if __name__ == "__main__":
    demo.launch()