#!/usr/bin/env python3
"""
Session-context projection endpoint.
Returns the session context for a chat -- the app-defined projection of
*checkpointed graph state* via :meth:`BaseLangGraph.context_snapshot`
(ADR-0032) -- so frontends can restore the mode selector on
hydration without waiting for the next streamed query.
Design: the mode lives in checkpointed graph state, which is the single
source of truth (ADR-0030 / ADR-0032). This endpoint reads it from the
checkpoint directly (``graph.checkpointer`` / ``graph.graph``); nothing
is duplicated into the sessions database. An unsent query is not a chat
yet, so a thread with no checkpoint simply reports ``null`` context --
the chat entity exists, only its projection is unset, so ``200`` with
``null`` is returned rather than a 404.
File: klea_utils/api/context.py
Copyright 2026 Ankur Sinha
Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>
"""
import logging
from typing import Annotated
from fastapi import APIRouter, Request
from langgraph.types import RunnableConfig
from pydantic import Field
from klea_utils.api.chat_common import _graph_and_store, thread_id_for
from klea_utils.api.hitl import _snapshot_interrupts
from klea_utils.graph.base import (
BaseLangGraph,
_interrupt_event,
_normalise_state_snapshot,
)
logger = logging.getLogger(__name__)
[docs]
def create_context_router() -> APIRouter:
"""Create an APIRouter exposing the graph-level ``context`` projection.
``GET /chat/{user_id}/{chat_id}/context``
Return ``{"context": {...}, "pending_interrupt": {...} | null,
"usage": {...} | null}`` -- the checkpointed session context
projected by the app's ``context_snapshot`` hook, plus any pending
HITL interrupt (ADR-0046) so a reloaded frontend can re-present the
ask, and the checkpointed cumulative ``usage_metrics`` (ADR-0032
update) so the token totals survive a reload without storing the
per-node ``usage`` events. All are ``null`` when the thread has no
checkpoint yet (a chat that never ran a query) or the graph uses no
checkpointer.
"""
router = APIRouter(prefix="/chat", tags=["context"])
@router.get("/{user_id}/{chat_id}/context")
async def get_chat_context(
user_id: Annotated[str, Field(pattern=r"^[^:]*$")],
chat_id: Annotated[str, Field(pattern=r"^[^:]+$")],
request: Request,
):
# ``_graph_and_store`` also enforces service readiness (503).
graph: BaseLangGraph
graph, _ = _graph_and_store(request)
thread_id = thread_id_for(user_id, chat_id)
logger.debug("get_chat_context(%s, %s): thread=%s", user_id, chat_id, thread_id)
# No checkpointer (e.g. checkpoint="none") means there is no
# checkpointed context to project -- fall through to null.
if graph.checkpointer is None:
logger.debug(
"get_chat_context(%s, %s): graph has no checkpointer",
user_id,
chat_id,
)
return {"context": None, "pending_interrupt": None, "usage": None}
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
# ``aget_tuple`` returns None for a thread that never ran. That is
# the "no context yet" case (an unsent query is not a chat yet), so
# we surface null instead of projecting the state-schema defaults.
checkpoint = await graph.checkpointer.aget_tuple(config)
if checkpoint is None:
logger.debug(
"get_chat_context(%s, %s): no checkpoint yet", user_id, chat_id
)
return {"context": None, "pending_interrupt": None, "usage": None}
snapshot = await graph.graph.aget_state(config)
# ``_normalise_state_snapshot`` keeps this identical to the stream
# path (``run_graph_astream_events``), so the app hook always sees
# a plain dict whatever the state backend renders.
values = _normalise_state_snapshot(snapshot.values)
context = graph.context_snapshot(values)
# The checkpointed cumulative token totals ride the same snapshot, so
# the frontends can restore the session total without replaying the
# per-node ``usage`` events.
usage = values.get("usage_metrics")
if hasattr(usage, "model_dump"):
usage = usage.model_dump()
# The same snapshot carries any pending HITL interrupt (ADR-0046), so
# a reloaded frontend can re-present the ask without a second read.
interrupts = _snapshot_interrupts(snapshot)
interrupt = _interrupt_event(interrupts[0], "")["data"] if interrupts else None
logger.debug(
"get_chat_context(%s, %s): context=%s pending_interrupt=%s usage=%s",
user_id,
chat_id,
context,
bool(interrupts),
usage,
)
return {"context": context, "pending_interrupt": interrupt, "usage": usage}
return router