Per-agent event streaming

LangGraph exposes astream, which yields an event every time a node finishes. We can turn those events into Server-Sent Events the browser understands. Each SSE message carries the node name and whatever fields that node wrote, so the UI can show a "supervisor thinking..." indicator then switch to "policy agent answering..." when the route is decided.

service.py
python
import json
from typing import Any, AsyncGenerator, Dict, Optional

from agent_graph import get_graph


def _sse(payload: Dict[str, Any]) -> str:
    """Format a payload as a Server-Sent Events message."""
    return f'data: {json.dumps(payload, default=str)}\n\n'


def _thread_config(thread_id: str) -> Dict[str, Any]:
    return {'configurable': {'thread_id': thread_id}}

SSE format is deceptively simple: 'data: <json>\n\n'. The double newline is the message terminator. default=str in json.dumps handles datetime and decimal fields from the database without extra encoding logic.

service.py
python
async def stream_chat(
    message: str,
    thread_id: Optional[str] = None,
    customer_id: Optional[str] = None,
) -> AsyncGenerator[str, None]:
    """Yield SSE-formatted events as the graph progresses."""
    thread_id = thread_id or _new_thread_id()
    graph = get_graph()

    snapshot = graph.get_state(_thread_config(thread_id))
    history = (snapshot.values or {}).get('messages', []) if snapshot else []
    history = history + [{'role': 'user', 'content': message}]

    initial: Dict[str, Any] = {
        'messages': history,
        'user_input': message,
        'customer_id': customer_id,
    }

    yield _sse({'event': 'start', 'thread_id': thread_id})

    final_state: Dict[str, Any] = {}
    try:
        async for event in graph.astream(initial, config=_thread_config(thread_id)):
            for node_name, node_state in event.items():
                final_state.update(node_state or {})
                yield _sse({
                    'event': 'node',
                    'node': node_name,
                    'route': final_state.get('route'),
                    'partial_answer': (node_state or {}).get('answer'),
                })
    except Exception as exc:
        yield _sse({'event': 'error', 'error': str(exc)})
        return

    history.append({'role': 'assistant', 'content': final_state.get('answer', '')})
    graph.update_state(_thread_config(thread_id), {'messages': history})

    yield _sse({
        'event': 'done',
        'thread_id': thread_id,
        'route': final_state.get('route'),
        'answer': final_state.get('answer', ''),
        'citations': final_state.get('citations', []) or [],
    })

Three event types. start signals the thread has begun. node fires every time a graph node finishes, carrying the node name so the UI can show which agent spoke. done wraps up with the final answer and citations. Errors emit their own event so the client can render a friendly failure instead of silently stalling.

So the client gets the thread_id immediately, before any LLM work happens. The UI needs that id to associate the response with the current session and to pass it back on the next turn for memory. If the client had to wait for the first node event, it would have no session handle during the supervisor's classification call.

Quiz: Quiz

Loading practice…