Agent chat UI

A useful way to interact with an agent is through a chat interface. Because Union.ai can host apps behind a URL, you can serve a chat UI for your agent with no separate infrastructure. There are two approaches:

  1. AgentChatAppEnvironment — the fastest path. Any agent that implements the AgentProtocol (including the built-in Agent and CodeModeAgent) gets a hosted chat shell, tool sidebar, and streaming for free.
  2. A custom FastAPI app — full control over the UI. Wrap the agent in a FastAPIAppEnvironment and serve your own HTML/CSS/JS.

Both reuse the same agent object, so you can start with the built-in shell and graduate to a custom UI later.

Option 1: the built-in chat UI

flyte.ai.chat.AgentChatAppEnvironment wraps an agent in a hosted chat app. Since Agent implements the AgentProtocol, it plugs straight in:

agent_chat_ui.py 
import flyte
from flyte.ai.agents import Agent
from flyte.ai.chat import AgentChatAppEnvironment, CustomTheme

task_env = flyte.TaskEnvironment(
    name="chat-agent-tools",
    image=flyte.Image.from_debian_base().with_pip_packages("litellm", "httpx"),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    secrets=[flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY")],
)


@task_env.task
async def search_docs(query: str, max_results: int = 3) -> list[dict[str, str]]:
    """Search internal documentation (stub) and return matching snippets."""
    corpus = [
        {"title": "Tasks", "body": "Define a task by decorating an async function with @env.task."},
        {"title": "Triggers", "body": "Schedule a task by attaching a flyte.Trigger with a flyte.Cron automation."},
        {"title": "Secrets", "body": "Mount cluster-managed secrets into a task with flyte.Secret(...)."},
    ]
    needle = query.lower()
    matches = [d for d in corpus if needle in d["body"].lower() or needle in d["title"].lower()]
    return matches[:max_results]


agent = Agent(
    name="docs-helper",
    instructions=(
        "You are a friendly internal docs assistant. Use search_docs to find "
        "relevant snippets. Always cite the doc title in your final answer."
    ),
    model="claude-haiku-4-5",
    tools=[search_docs],
    max_turns=8,
)


@task_env.task(report=True)
async def chat_entrypoint(message: str, history: list[dict]) -> dict:
    """Parent task that owns the agent loop and the nested tool tasks."""
    result = await agent.run.aio(message, history=history)
    return {
        "summary": result.summary,
        "error": result.error,
        "attempts": result.attempts,
        "charts": [],
        "code": "",
    }


env = AgentChatAppEnvironment(
    name="docs-agent-chat-ui",
    agent=agent,
    task_entrypoint=chat_entrypoint,
    title="Internal docs assistant",
    subtitle="Backed by a flyte.ai.agents.Agent + durable Flyte task tools.",
    theme=CustomTheme(accent_color="#6F2AEF", accent_hover_color="#8B52F2"),
    prompt_nudges=[
        {"label": "Basics", "prompt": "Can you show me a hello world example?"},
        {"label": "Triggers", "prompt": "How do I schedule a task?"},
    ],
    depends_on=[task_env],
    image=flyte.Image.from_debian_base().with_pip_packages("litellm", "fastapi", "uvicorn"),
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    secrets=flyte.Secret("internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
)

The task_entrypoint is a parent task that owns the agent loop, so the nested durable tool tasks run correctly under it. The chat shell streams progress by subscribing to the agent’s agent_progress_cb events.

Option 2: a custom FastAPI chat app

When you want to control the look and feel, wrap any AgentProtocol-compatible agent in a FastAPIAppEnvironment and serve your own UI. This is the pattern used by the CodeModeAgent chat app: a single LLM call generates Python code, runs it in a sandbox, and returns charts + a summary, all behind a conversational web interface.

The architecture is small:

Browser (Chat UI)
  ├── GET  /            -> embedded HTML/CSS/JS chat interface
  ├── GET  /api/tools   -> JSON list of available tool descriptions
  └── POST /api/chat    -> { message, history } -> { code, charts, summary, error }
           └── CodeModeAgent.run(message, history)

The app itself is just a FastAPI server. The endpoints call the agent’s run.aio and tool_descriptions methods:

import pathlib

from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel

import flyte
from flyte.ai.agents import CodeModeAgent
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI(title="Chat Data Analytics Agent")

env = FastAPIAppEnvironment(
    name="chat-analytics-agent",
    app=app,
    image=flyte.Image.from_debian_base().with_pip_packages(
        "fastapi", "uvicorn", "httpx", "pydantic-monty", "litellm",
    ),
    secrets=flyte.Secret(key="internal-anthropic-api-key", as_env_var="ANTHROPIC_API_KEY"),
    scaling=flyte.app.Scaling(replicas=1),
)

agent = CodeModeAgent(tools=ALL_TOOLS, max_retries=2)


class ChatRequest(BaseModel):
    message: str
    history: list[dict] = []


class ChatResponse(BaseModel):
    code: str = ""
    charts: list[str] = []
    summary: str = ""
    error: str = ""


@app.get("/api/tools")
async def get_tools() -> list[dict]:
    """Return JSON descriptions of available tool functions (for the sidebar)."""
    return agent.tool_descriptions()


@app.post("/api/chat")
async def chat(req: ChatRequest) -> ChatResponse:
    """Generate code, run it in the sandbox, and return results."""
    result = await agent.run.aio(req.message, req.history)
    return ChatResponse(code=result.code, charts=result.charts,
                        summary=result.summary, error=result.error)


@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    """Serve the embedded chat UI."""
    return HTMLResponse(content=CHAT_HTML)


if __name__ == "__main__":
    flyte.init_from_config(root_dir=pathlib.Path(__file__).parent)
    app_handle = flyte.serve(env)
    print(f"Deployed Chat Analytics Agent: {app_handle.url}")

CHAT_HTML is the embedded front-end (a single HTML string with the chat markup, styles, and a small fetch-based client that POSTs to /api/chat and renders the returned charts and summary). ALL_TOOLS is the agent’s tool registry. Keeping both in their own modules means adding a tool is the only change required — the agent auto-generates its system prompt from each tool’s signature and docstring.

Run it locally during development, then deploy with one command:

# Local development
python chat_app.py

# Deploy to Union.ai
flyte deploy chat_app.py env

Union.ai assigns a URL, handles TLS, and auto-scales the app.

Swap CodeModeAgent for the Agent harness (or any object implementing the AgentProtocol) to serve a tool-use agent behind the same UI. The endpoints only depend on run.aio and tool_descriptions.

Next steps