Chat UI Protocols

data: {"type":"start","messageId":"msg_01"}

data: {"type":"start-step"}

data: {"type":"text-start","id":"t1"}

data: {"type":"text-delta","id":"t1","delta":"Use the retriever "}

data: {"type":"text-delta","id":"t1","delta":"tool [1]."}

data: {"type":"text-end","id":"t1"}

data: {"type":"source-url","sourceId":"c1","url":"https://example.org/docs/rag","title":"RAG overview"}

data: {"type":"finish-step"}

data: {"type":"finish"}

data: [DONE]

import json
import uuid
from collections.abc import AsyncIterator

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse

from docsassistant.core import Citation, DocsAssistantCore, Token  # see Channel-Agnostic Architecture

app = FastAPI()


def sse(part: dict) -> str:
    return f"data: {json.dumps(part, separators=(',', ':'))}\n\n"


def last_user_text(messages: list[dict]) -> str:
    # UI messages carry typed parts; the question is the text parts of the last user message.
    last = next(m for m in reversed(messages) if m["role"] == "user")
    return "".join(p["text"] for p in last["parts"] if p["type"] == "text")


async def ui_message_stream(core: DocsAssistantCore, conversation_id: str, user_id: str,
                            text: str) -> AsyncIterator[str]:
    text_id = f"t-{uuid.uuid4().hex[:8]}"
    yield sse({"type": "start", "messageId": f"msg-{uuid.uuid4().hex[:8]}"})
    yield sse({"type": "start-step"})
    yield sse({"type": "text-start", "id": text_id})
    try:
        async for event in core.ask(conversation_id, user_id, text):
            match event:
                case Token(delta):
                    yield sse({"type": "text-delta", "id": text_id, "delta": delta})
                case Citation(_marker, chunk_id, title, url):
                    yield sse({"type": "source-url", "sourceId": chunk_id, "url": url, "title": title})
    except Exception:
        yield sse({"type": "error", "errorText": "The assistant could not answer. Try again."})
    yield sse({"type": "text-end", "id": text_id})
    yield sse({"type": "finish-step"})
    yield sse({"type": "finish"})
    yield "data: [DONE]\n\n"


@app.post("/chat")
async def chat(request: Request) -> StreamingResponse:
    body = await request.json()
    user_id = request.state.user_id  # set by your auth dependency or middleware
    return StreamingResponse(
        ui_message_stream(request.app.state.core, body["id"], user_id, last_user_text(body["messages"])),
        media_type="text/event-stream",
        headers={
            "x-vercel-ai-ui-message-stream": "v1",
            "cache-control": "no-cache",
            "x-accel-buffering": "no",
        },
    )

@RestController
class UiMessageStreamController {

    private final AssistantCore core; // wraps ChatClient + RAG advisor + JDBC chat memory + ACL filter

    UiMessageStreamController(AssistantCore core) {
        this.core = core;
    }

    @PostMapping(path = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    ResponseEntity<Flux<ServerSentEvent<Map<String, Object>>>> chat(
            @RequestBody ChatRequest request, Principal principal) {

        String textId = "t-" + UUID.randomUUID().toString().substring(0, 8);
        Flux<Map<String, Object>> parts = core
                .ask(request.id(), principal.getName(), request.lastUserText())
                .map(event -> switch (event) {
                    case CoreEvent.Token t -> Map.<String, Object>of("type", "text-delta", "id", textId,
                            "delta", t.text());
                    case CoreEvent.Citation c -> Map.<String, Object>of("type", "source-url", "sourceId", c.chunkId(),
                            "url", c.url(), "title", c.title());
                    case CoreEvent.Done d -> Map.<String, Object>of("type", "text-end", "id", textId);
                });

        Flux<Map<String, Object>> stream = Flux.concat(
                Flux.just(Map.<String, Object>of("type", "start", "messageId", "msg-" + textId),
                          Map.<String, Object>of("type", "text-start", "id", textId)),
                parts,
                Flux.just(Map.<String, Object>of("type", "finish")))
                .onErrorResume(e -> Flux.just(Map.<String, Object>of(
                        "type", "error", "errorText", "The assistant could not answer.")));

        Flux<ServerSentEvent<Map<String, Object>>> sse = stream
                .map(p -> ServerSentEvent.builder(p).build());
        // The literal [DONE] marker is a raw string, not JSON; see the text below for how to append it.

        return ResponseEntity.ok()
                .header("x-vercel-ai-ui-message-stream", "v1")
                .header(HttpHeaders.CACHE_CONTROL, "no-cache")
                .body(sse);
    }

    record ChatRequest(String id, List<Map<String, Object>> messages) {
        String lastUserText() { /* concatenate the text parts of the last user message */ return ""; }
    }
}

import { useChat } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import { useState } from 'react';

export function DocsChat() {
  const [input, setInput] = useState('');
  const { messages, sendMessage, status, error } = useChat({
    transport: new DefaultChatTransport({
      api: '/chat',
      // Add a bearer token from your session layer; never hard-code one.
      headers: () => ({ Authorization: `Bearer ${getAccessToken()}` }),
    }),
  });

  return (
    <div>
      {messages.map((m) => (
        <article key={m.id} data-role={m.role}>
          {m.parts.map((part, i) => {
            if (part.type === 'text') return <p key={i}>{part.text}</p>;
            if (part.type === 'source-url')
              return <a key={i} href={part.url}>{part.title ?? part.url}</a>;
            return null;
          })}
        </article>
      ))}
      {error && <p role="alert">Something went wrong.</p>}
      <form
        onSubmit={(e) => {
          e.preventDefault();
          sendMessage({ text: input });
          setInput('');
        }}
      >
        <input value={input} onChange={(e) => setInput(e.target.value)} disabled={status !== 'ready'} />
      </form>
    </div>
  );
}

RUN_STARTED          { threadId: "th_1", runId: "run_1" }
STEP_STARTED         { stepName: "retrieve" }
TOOL_CALL_START      { toolCallId: "tc_1", toolCallName: "search_docs" }
TOOL_CALL_ARGS       { toolCallId: "tc_1", delta: "{\"query\":\"rag api\"}" }
TOOL_CALL_END        { toolCallId: "tc_1" }
TOOL_CALL_RESULT     { messageId: "m_t", toolCallId: "tc_1", content: "3 chunks" }
STATE_DELTA          { delta: [ { op: "add", path: "/citations/-", value: { id: "c1", url: "https://example.org/docs/rag" } } ] }
STEP_FINISHED        { stepName: "retrieve" }
TEXT_MESSAGE_START   { messageId: "m_1", role: "assistant" }
TEXT_MESSAGE_CONTENT { messageId: "m_1", delta: "Use the retriever tool [1]." }
TEXT_MESSAGE_END     { messageId: "m_1" }
RUN_FINISHED         { threadId: "th_1", runId: "run_1" }

import os

from ag_ui_langgraph import LangGraphAgent, add_langgraph_fastapi_endpoint
from fastapi import FastAPI
from langchain.agents import create_agent
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver

from docsassistant.tools import search_docs  # retriever tool

# The checkpointer is the same Postgres one used by the other channels.
# Opened at application start-up, for example in a FastAPI lifespan handler.
checkpointer: AsyncPostgresSaver = ...

graph = create_agent(
    model=os.environ["DOCSASSISTANT_MODEL"],   # for example "provider:model-name"
    tools=[search_docs],
    system_prompt="Answer only from the retrieved documentation and cite sources.",
    checkpointer=checkpointer,
)

agui_agent = LangGraphAgent(
    name="docs_assistant",
    description="Answers questions about the documentation with citations.",
    graph=graph,
)

app = FastAPI()
add_langgraph_fastapi_endpoint(app, agui_agent, "/agent")