#!/usr/bin/env python3
"""
Base class for LangGraph-based orchestrators
See ``devdocs/system/streams.md`` (event emission via
``run_graph_astream_events`` and ``context_snapshot``),
``devdocs/system/graph-resume.md`` (fault tolerance and resume), and the
generated graph topology embedded in ``devdocs/system/c4-component-*.md``;
ADRs 0016, 0032, 0033 and 0043.
File: klea_utils/graph/base.py
Copyright 2026 Ankur Sinha
Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>
"""
from __future__ import annotations
import json
import logging
import os
import time
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Literal, cast, final
from fastmcp import Client
from fastmcp.mcp_config import MCPConfig
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.graph.state import CompiledStateGraph
from langgraph.stream import StreamTransformer
from langgraph.types import Command, RunnableConfig
from mcp.types import Tool
from platformdirs import PlatformDirs
from pydantic import BaseModel, ConfigDict, Field, create_model
from klea_utils.client_context import ClientContext
from klea_utils.graph.context import KleaRunContext
from klea_utils.llm import LLMModel, missing_required_roles, resolve_user_agent
from klea_utils.mcp import spawn
from klea_utils.mcp.access import (
AccessLevel,
ToolAccessOverride,
resolve_access_level,
resolve_capability,
)
from klea_utils.mcp.privilege import warn_if_root
from klea_utils.mcp.schemas import ToolCallSchema, ToolInfo
from klea_utils.paths import (
get_config_dir,
init_dir,
resolve_app_config_path,
resolve_storage_root,
)
from klea_utils.stores.config import RetrieverConfig
from klea_utils.stores.retrieval.bm25 import BM25RetrieverManager
from klea_utils.stores.retrieval.vs import VSRetriever
from klea_utils.tools import build_tool_description, clean_tool_meta
# Per-request model overrides are carried by the LangGraph Runtime context
# (KleaRunContext, ADR-0033), not a contextvar: the API layer builds the
# context from the sessions DB per invocation and the nodes read it via
# ambient ``get_runtime()``. The old ``model_overrides_ctx`` contextvar
# is gone -- see ``klea_utils.graph.context``.
def _normalise_state_snapshot(state: Any) -> dict[str, Any]:
"""Normalize a checkpoint ``values`` snapshot into a plain dict.
LangGraph state snapshots arrive either as a dict (TypedDict state)
or as the pydantic state instance. ``context_snapshot`` and the
hydration endpoint (``klea_utils.api.context``) both need a plain
dict, so the normalization lives here: one code path shared by the
stream loop and the direct checkpoint read, guaranteeing the hook
sees identical input wherever it is called.
:param state: A checkpoint/superstep ``values`` snapshot.
:returns: A plain dict view of *state*.
"""
if not isinstance(state, dict) and hasattr(state, "model_dump"):
return state.model_dump()
return state
def _interrupt_event(intr: Any, node: str) -> dict[str, Any]:
"""Build an ``interrupt`` stream event from a LangGraph ``Interrupt``.
The interrupting node authors the interrupt *value*. When it is a
mapping (the convention: ``kind`` and ``question``) its keys are carried
into ``data``; any other value is surfaced as ``question``. The
framework's ``id`` is always added as ``interrupt_id`` so a client can
echo it when resuming (stale-answer safety), and a
``hitl_response_schema`` (when the node supplied one) is forwarded so
clients can render a typed form.
:param intr: The ``langgraph.types.Interrupt`` from the paused run.
:param node: Label of the node that interrupted.
:returns: A ``{"type": "interrupt", "node": ..., "data": {...}}`` event.
"""
value = getattr(intr, "value", None)
data: dict[str, Any] = (
dict(value) if isinstance(value, dict) else {"question": value}
)
data["interrupt_id"] = getattr(intr, "id", None)
# ``Interrupt.response_schema`` is LangGraph's API attribute (set by
# ``interrupt(value, response_schema=...)``). Forward it under our own
# ``hitl_response_schema`` stream key so it is not confused with an
# LLM/API response elsewhere.
schema = getattr(intr, "response_schema", None)
if schema is not None:
data["hitl_response_schema"] = schema
return {"type": "interrupt", "node": node, "data": data}
#: SQLite PRAGMAs applied to every Klea SQLite database so the checkpoint
#: and session stores share one durability and concurrency policy: WAL
#: journalling (readers never block the single writer), ``NORMAL``
#: synchronous mode (durable enough under WAL), a five second busy timeout
#: (a concurrent writer waits instead of failing with "database is locked"),
#: and enforced foreign keys.
SQLITE_PRAGMAS: tuple[str, ...] = (
"PRAGMA journal_mode=WAL",
"PRAGMA synchronous=NORMAL",
"PRAGMA busy_timeout=5000",
"PRAGMA foreign_keys=ON",
)
def _apply_sqlite_pragmas(conn) -> None:
"""Apply :data:`SQLITE_PRAGMAS` to a synchronous ``sqlite3`` connection.
The session store (:class:`~klea_utils.api.sessions_db.SessionStore`)
uses the standard-library ``sqlite3`` module, whose ``execute`` is
synchronous. The checkpoint store uses ``aiosqlite`` and applies the
same policy through :func:`_apply_sqlite_pragmas_async`, so both
databases share one connection policy without duplicating the SQL.
:param conn: An open :class:`sqlite3.Connection`; PRAGMAs are applied
and committed immediately after connecting.
"""
for pragma in SQLITE_PRAGMAS:
conn.execute(pragma)
conn.commit()
async def _apply_sqlite_pragmas_async(conn) -> None:
"""Apply :data:`SQLITE_PRAGMAS` to an ``aiosqlite`` connection.
:param conn: An open :class:`aiosqlite.Connection`; PRAGMAs are applied
and committed immediately after connecting.
"""
for pragma in SQLITE_PRAGMAS:
await conn.execute(pragma)
await conn.commit()
class _CustomChannelEnabler(StreamTransformer):
"""Enables the ``custom`` channel in LangGraph v3 event streams.
Non-LLM nodes use ``get_stream_writer()`` to emit progress events on the
custom channel. LangGraph requires a ``StreamTransformer`` declaring
``required_stream_modes = ("custom",)`` for that channel to be enabled.
This no-op transformer satisfies that requirement.
"""
required_stream_modes = ("custom",)
def __init__(self, scope=()):
super().__init__(scope)
def init(self):
return {}
def process(self, event):
return True
[docs]
class BaseLangGraph(ABC):
"""Abstract base class for LangGraph-based orchestrators.
Provides common infrastructure for:
- Configuration loading from env files
- MCP client creation from JSON config
- LLM model setup (delegated to subclasses)
- LangGraph compilation and execution
- Session checkpointing
- Dual-stream logging
Subclasses must implement:
- ``_setup_models``: create ``self.llm_models``; the env schema is
generated from its roles.
- ``_create_graph``: build and compile the LangGraph.
- ``config_class``: the Pydantic class for the JSON configuration.
"""
#: Pydantic BaseSettings class for env loading. Subclasses need not
#: set this: it is generated at load time from ``self.llm_models`` (see
#: :meth:`_build_env_class`).
env_class: type[BaseModel] = BaseModel
#: Pydantic BaseModel class for configuration loading.
#: Subclasses must set this to their AppConfig class.
config_class: type[BaseModel]
#: Name of the environment variable that controls the env file path.
env_var: str = "ENV_FILE"
#: Default config file name if the environment variable is not set.
env_file_default: str = "config.env"
#: Prefix prepended to generated env var names (e.g. ``"KLEA_AGENT_"``).
#: Each model role ``r`` becomes the env var ``<env_prefix>R_MODEL``.
env_prefix: str = ""
#: Default JSON config file name, used as the ``app_config_file`` env
#: field default when none is set in the env file or process env.
config_file_default: str = ""
#: Logger name for this orchestrator, also used as the app name
#: for ``platformdirs`` data/cache directories.
graph_name: str = "BaseLangGraph"
#: Environment variable that relocates this graph's storage root (data
#: and runtime artifacts). ``None`` disables relocation; the agent sets
#: this to ``KLEA_AGENT_HOME``.
storage_root_env: str | None = None
def __init__(
self,
logging_level: int = logging.INFO,
checkpoint: str = "inmemory",
log_file: bool = True,
):
"""Initialise the base orchestrator.
:param logging_level: Logging level for the orchestrator. The
:data:`KLEA_LOG_LEVEL` environment variable (and the ``--debug``
flag that sets it) takes precedence over this constructor
argument: when it resolves to ``DEBUG``, full debug logging is
enabled regardless of *logging_level*; otherwise *logging_level*
is used (default ``INFO``).
:param checkpoint: Checkpointer mode --- ``"inmemory"`` (volatile, default),
``"sqlite"`` (persistent via ``self.storage_root``),
or ``"none"`` (no checkpointing). When set to ``"none"``, nodes
that need conversation history receive ``memory=False``.
:param log_file: When ``True`` (default), configure the process-wide
root logger with a rotating file handler writing to
``{self.storage_root}/{self.graph_name}.log``, alongside
the checkpoints and session database. Set to ``False`` in
short-lived processes (e.g. tests) to avoid writing log files.
"""
self.env_file = os.getenv(self.env_var, self.env_file_default)
self.app_env: BaseModel
# Graph-level default models per role. Per-request model
# overrides arrive via the LangGraph Runtime context (KleaRunContext,
# ADR-0033) and do NOT change this dict.
self.llm_models: dict[str, LLMModel] = {}
self.tools_info: dict[str, dict[str, ToolInfo]] = {}
self.domain_mcp_configs: dict[str, MCPConfig] = {}
self.checkpointer_mode = checkpoint
self.memory = checkpoint != "none"
self.checkpointer: None | AsyncSqliteSaver | InMemorySaver = None
self.paths = PlatformDirs(self.graph_name.lower())
#: Storage root for the checkpoint/session databases, lockfile, and
#: logs. Defaults to this graph's per-app data directory; a graph
#: may opt into an environment-variable override via
#: :attr:`storage_root_env`.
self.storage_root = resolve_storage_root(
default_root=self.paths.user_data_dir,
env_name=type(self).storage_root_env,
)
#: Klea identity sent to OpenAI-compatible endpoints (User-Agent).
self.user_agent = resolve_user_agent(self.graph_name)
self.config_dict: dict[str, Any]
self.graph: CompiledStateGraph[Any, Any, Any, Any] | None = None
#: Labels of nodes whose LLM token deltas should be forwarded as
#: ``token`` stream events. A node opts in with
#: ``stream_tokens = True``; the app adds its label here at graph
#: build time. Empty means no token events are emitted, so
#: structured-output nodes never flood the stream with JSON fragments.
self.token_stream_nodes: set[str] = set()
self.mcp_config: MCPConfig | None = None
self.mcp_client: Client | None = None
#: Optional per-session MCP client pool (set by the API app before
#: ``setup()`` so the tool-caller node can use it). ``None`` keeps
#: the shared ``mcp_client`` for non-server callers and tests.
self.mcp_session_pool: Any | None = None
self.mcp_tools: list[Tool] | None = None
self.retriever_config: RetrieverConfig | None = None
self.stores: VSRetriever | None = None
self.bm25_stores: BM25RetrieverManager | None = None
# Graph-wide fallback retrieval settings. Individual vector stores
# may override these in the config with their own default_k / k_max /
# k_inc values.
self.default_k: int = 5
self.k_max: int = 10
self.k_inc: int = 1
self.QueryDomainSchema: type[BaseModel] | None = None
from klea_utils.plogging import resolve_log_level, setup_root_logger
# The KLEA_LOG_LEVEL env var (set by --debug) overrides the
# constructor's logging_level: when it resolves to DEBUG we show
# full debug output regardless of the passed level.
if resolve_log_level() == logging.DEBUG:
stderr_level = logging.DEBUG
else:
stderr_level = logging_level
setup_root_logger(
self.graph_name,
stderr_level=stderr_level,
log_dir=self.storage_root if log_file else None,
)
self.logger = logging.getLogger(self.graph_name)
def _build_env_class(self) -> type[BaseModel]:
"""Build the pydantic-settings class for this graph's environment.
The schema is derived from ``self.llm_models``: each model role ``r``
becomes a ``{r}_model`` string field (env var ``<env_prefix>R_MODEL``),
plus ``app_config_file``. Generating the class from the model
declaration keeps the env schema and the model roles as a single
source of truth, so they cannot drift apart.
:returns: A ``BaseSettings`` subclass configured with ``env_prefix``.
"""
# Lazy: pydantic-settings is only needed when loading the env.
from pydantic_settings import BaseSettings, SettingsConfigDict
fields: dict[str, Any] = {
f"{role}_model": (str, "") for role in self.llm_models
}
fields["app_config_file"] = (str, self.config_file_default)
return create_model(
f"{self.graph_name}Env",
__base__=BaseSettings,
__config__=cast(ConfigDict, SettingsConfigDict(env_prefix=self.env_prefix)),
**fields,
)
def _load_env(self) -> None:
"""Load env file, and configuration
Builds the env settings class from ``self.llm_models`` (see
:meth:`_build_env_class`) and parses the env file (optional --
when missing, process environment variables and class defaults are
used), then resolves the application config file (from
``self.app_env.app_config_file``) via
:func:`klea_utils.paths.resolve_app_config_path` -- the working
directory first, then the per-app config directory.
The JSON config file is optional when no default is declared (see
``config_file_default``): graphs that can run without one (e.g. the
agent, whose models/access settings come from env vars and the web
UI) start from ``config_class`` defaults. A graph with a non-empty
default (e.g. RAG, which needs its store wiring) raises
``FileNotFoundError`` when the file does not exist.
"""
self.env_class = self._build_env_class()
env_file_path = Path(self.env_file)
if not env_file_path.exists():
# Env files are optional. With no file, pydantic-settings
# falls back to process environment variables and class
# defaults, so a clean machine can run with only a JSON config
# (models are then set via shell env vars or the web UI).
self.logger.warning(
f"Env file not found: {self.env_file}. "
f"Using process environment variables and defaults only. "
f"Set {self.env_var} to specify an env file."
)
self.app_env = self.env_class()
else:
self.app_env = self.env_class(_env_file=self.env_file)
assert self.app_env
self.logger.debug(f"env file: {self.env_file}")
self.logger.debug(f"env: {self.app_env}")
# An explicit empty value means "not set" -- fall back to the field
# default so a line like ``KLEA_AGENT_APP_CONFIG_FILE=`` does not
# resolve to the current directory itself.
app_config_file = self.app_env.app_config_file
if not app_config_file:
app_config_file = self.env_class.model_fields["app_config_file"].default
self.logger.debug(
f"empty app_config_file -- using default {app_config_file}"
)
if not app_config_file:
# No config file selected and no default declared: start from
# the config class defaults (env vars and the web UI still
# configure the app).
self.logger.info(
"No app config file selected; using %s defaults.",
self.config_class.__name__,
)
self.app_config = self.config_class()
self.logger.debug(f"{self.app_config = }")
return
config_file = resolve_app_config_path(
app_config_file, get_config_dir(self.paths)
)
self.logger.debug(f"config file: {config_file}")
with open(config_file, "r") as f:
config_dict = json.load(f)
self.logger.debug(f"{config_dict = }")
self.app_config = self.config_class(**config_dict)
self.logger.debug(f"{self.app_config = }")
def _create_mcp_client(self) -> None:
"""Create MCP client from the JSON config file.
Reads the MCP server configurations from ``self.app_env.mcp_config_file``
and creates a ``fastmcp.Client`` instance.
"""
self.mcp_client = self.new_mcp_client()
if self.mcp_client is None:
self.logger.warning("No MCP server configured.")
[docs]
def new_mcp_client(
self, client_context: ClientContext | None = None
) -> Client | None:
"""Return a new MCP client over the configured servers.
The returned client owns its own connections and server subprocesses;
it is used to give each chat session an isolated MCP session. When
*client_context* is given, only the ``bundled`` server entry is
rebuilt with that cwd/env/interpreter; external MCP servers keep
their configured invocation. Returns ``None`` when no MCP server is
configured, matching :meth:`_create_mcp_client`.
:param client_context: Optional client-supplied environment for the
bundled tools server.
:returns: A new MCP client, or ``None`` when nothing is configured.
"""
if not self.mcp_config or not self.mcp_config.mcpServers:
return None
if client_context is None or (
client_context.cwd is None
and client_context.interpreter is None
and not client_context.env
):
self.logger.debug(f"{self.mcp_config = }")
return Client(self.mcp_config)
servers: dict[str, Any] = dict(self.mcp_config.mcpServers)
if "bundled" in servers:
bundled = self._bundled_server_config(client_context=client_context)
if bundled is not None:
servers["bundled"] = bundled
self.logger.debug(
"Creating MCP client with client context: cwd=%s interpreter=%s "
"env_keys=%d",
client_context.cwd,
client_context.interpreter,
len(client_context.env),
)
return Client(MCPConfig(mcpServers=servers))
async def _get_mcp_tools(self) -> None:
"""Get MCP tools."""
if self.mcp_client:
async with self.mcp_client:
self.mcp_tools = await self.mcp_client.list_tools()
self.logger.debug(f"{self.mcp_tools =}")
self._build_tools_info()
def _build_tools_info(self) -> None:
"""Build per-domain tool metadata from fetched MCP tools."""
self.tools_info = {}
if not self.mcp_tools or not self.domain_mcp_configs:
return
# map server names to domains
domain_servers: dict[str, list[str]] = {}
num_servers = 0
for domain, config in self.domain_mcp_configs.items():
if config.mcpServers:
domain_servers[domain] = list(config.mcpServers.keys())
num_servers += len(list(config.mcpServers.keys()))
# Operator capability overrides (ADR-0037), applied over annotations.
access_overrides = self._tool_access_overrides()
for domain, server_names in domain_servers.items():
domain_tools_info: dict[str, ToolInfo] = {}
for t in self.mcp_tools:
# tools will be prefixed with server names
if num_servers > 1 and not any(
t.name.startswith(s + "_") for s in server_names
):
continue
# otherwise, there's only one server
# Klea expects MCP tools to follow the docstring-first
# convention (summary + Use when / Do not use for bullets +
# one example; params via Args:), see build_tool_description
# and docs/concepts/mcp.rst.
full_description, short_description = build_tool_description(t)
# Carry the standard MCP annotations so the tool access level
# (ADR-0037) can classify the tool client-side; an absent
# annotation stays ``None`` (fail-closed in read_only). A
# configured ``general.tool_access`` override takes precedence.
annotations = getattr(t, "annotations", None)
read_only, destructive = resolve_capability(
getattr(annotations, "readOnlyHint", None),
getattr(annotations, "destructiveHint", None),
access_overrides.get(t.name),
)
domain_tools_info[t.name] = ToolInfo(
title=t.title,
description=full_description,
short_description=short_description,
input_schema=t.inputSchema,
meta=clean_tool_meta(t.meta),
read_only=read_only,
destructive=destructive,
)
self.tools_info[domain] = domain_tools_info
self.logger.debug(f"{self.tools_info = }")
def _tool_access_overrides(self) -> dict[str, ToolAccessOverride]:
"""Return the configured per-tool capability overrides (ADR-0037).
Reads ``app_config.general.tool_access``; returns an empty mapping
when the app config does not declare it.
:returns: ``{tool_name: ToolAccessOverride}``.
"""
general = getattr(getattr(self, "app_config", None), "general", None)
overrides = getattr(general, "tool_access", None)
return dict(overrides) if overrides else {}
def _configured_access_level(self) -> AccessLevel | None:
"""Return the config-declared access level, or ``None``.
Reads ``app_config.general.access_level``. Apps that do not declare
the field (e.g. RAG, whose level is fixed) return ``None`` so the
state default applies.
:returns: The configured level, or ``None`` when unset.
"""
general = getattr(getattr(self, "app_config", None), "general", None)
raw = getattr(general, "access_level", None)
if raw is None:
return None
return resolve_access_level(raw)
def _apply_access_level_default(
self, input_state: dict[str, Any], extra_state: dict[str, Any] | None
) -> None:
"""Seed the config access level unless the request overrides it.
Precedence is request (``extra_state``) > ``general.access_level`` >
the state field default. Mutates *input_state* in place.
:param input_state: The initial-state dict being built for the run.
:param extra_state: The caller's per-run state extras, or ``None``.
"""
if extra_state and extra_state.get("access_level") is not None:
return
default = self._configured_access_level()
if default is not None:
input_state["access_level"] = default
def _tools_by_name(self) -> dict[str, ToolInfo]:
"""Return ``{tool_name: ToolInfo}`` flattened from ``tools_info``.
The per-domain map is flattened for the dispatch gate, so the
client-side path and access-level checks (ADR-0037) read one canonical
``ToolInfo`` per tool. Tool names are already server-prefixed when
several servers are configured, so a name repeated across domains is
the same tool with the same metadata.
:returns: Capability/metadata map keyed by tool name.
"""
tools = {
name: info
for domain_tools in self.tools_info.values()
for name, info in domain_tools.items()
}
self.logger.debug(f"{sorted(tools) = }")
return tools
async def _get_vector_stores(self) -> None:
"""Get vector stores"""
emb = self.llm_models.get("embedding")
if self.retriever_config and emb and emb.model_name:
self.stores = VSRetriever(
config=self.retriever_config,
logger=self.logger,
embedding_model=emb.model_name,
default_k=self.default_k,
k_max=self.k_max,
k_inc=self.k_inc,
)
self.stores.setup()
self.logger.info(f"Vector stores loaded: {self.stores.domains}")
elif self.retriever_config:
self.logger.warning(
"No vector stores configured (missing embedding model)."
)
# BM25 keyword stores need no embedding model, so build them whenever
# any domain configures bm25_stores.
if self.retriever_config and any(
domain.bm25_stores for domain in self.retriever_config.domains.values()
):
self.bm25_stores = BM25RetrieverManager(
config=self.retriever_config,
logger=self.logger,
default_k=self.default_k,
k_max=self.k_max,
k_inc=self.k_inc,
)
self.logger.info(f"BM25 stores loaded: {self.bm25_stores.domains}")
# Build QueryDomainSchema from all configured domains (vector or BM25)
# so BM25-only deployments still get structured classification.
if self.retriever_config and self.retriever_config.domains:
all_domains = list(self.retriever_config.domains.keys())
if "undefined" not in all_domains:
all_domains.append("undefined")
self.QueryDomainSchema = create_model(
"QueryDomainSchema",
query_domains=(
list[Literal[tuple(all_domains)]],
Field(default=["undefined"], validate_default=True),
),
)
elif self.QueryDomainSchema is None:
self.logger.warning("No domains configured for QueryDomainSchema")
def _export_graph_png(self, filename: str) -> None:
"""Export the LangGraph as a Mermaid PNG diagram and its Mermaid source.
A development aid (it keeps the C4 topology diagrams in sync, see
ADR-0016/0019), so it runs only under debug logging (``--debug`` /
``KLEA_LOG_LEVEL=debug``). Skipped inside Docker
(``RUNNING_IN_DOCKER`` env var set). Normal runs avoid the Mermaid
render (a network call) and the stray PNG/MMD files in the working
directory.
:param filename: Output file path for the PNG
"""
if os.environ.get("RUNNING_IN_DOCKER"):
return
# Lazy: plogging is otherwise imported inside ``__init__``.
from klea_utils.plogging import resolve_log_level
if resolve_log_level() != logging.DEBUG:
self.logger.debug(
"Skipping graph export; pass --debug (or set "
"KLEA_LOG_LEVEL=debug) to write the graph diagram."
)
return
try:
assert self.graph
self.graph.get_graph().draw_mermaid_png(output_file_path=filename)
# cancellation-contract: exempt (dev-only graph diagram export)
except BaseException as e: # noqa: BLE001
self.logger.error("Something went wrong generating lang graph png")
self.logger.error(e)
# Also write the Mermaid source alongside the PNG so the explicit
# C4 Level-3 diagrams (``devdocs/system/c4-component-*.md``) can
# embed the auto-generated node/edge topology as their core and
# avoid drift between code and docs (see ADR-0016/0019).
try:
assert self.graph
mermaid = self.graph.get_graph().draw_mermaid()
mermaid_path = Path(filename)
if mermaid_path.suffix == ".png":
mermaid_path = mermaid_path.with_suffix(".mmd")
else:
mermaid_path = Path(str(mermaid_path) + ".mmd")
mermaid_path.write_text(mermaid)
self.logger.debug(f"Wrote Mermaid source to {mermaid_path}")
# cancellation-contract: exempt (dev-only graph diagram export)
except BaseException as e: # noqa: BLE001
self.logger.error("Something went wrong generating lang graph mermaid")
self.logger.error(e)
# ------------------------------------------------------------------
# Abstract methods -- subclasses must implement these
# ------------------------------------------------------------------
@abstractmethod
def _configure_resources(self) -> None:
"""Configure vector stores and MCP servers
Subclasses should implement this to populate ``self.retriever_config``,
``self.mcp_config``, and ``self.domain_mcp_configs``, which will be used
to create the vector store class, mcp client, and per-domain tool descriptions.
"""
...
@abstractmethod
def _setup_models(self) -> None:
"""Set up LLM model instances.
Subclasses must populate ``self.llm_models`` with ``LLMModel``
entries keyed by role (e.g. ``"chat"``, ``"plan"``, ``"guard"``),
setting ``required`` on each. This runs *before* ``_load_env()``,
so the roles declared here define the env schema generated by
:meth:`_build_env_class`. ``model_name`` is populated from the env
by :meth:`_apply_model_names` after the env is loaded. These are
graph-wide defaults; per-request overrides arrive via the LangGraph
Runtime context (KleaRunContext, ADR-0033).
"""
...
def _apply_model_names(self) -> None:
"""Populate ``model_name`` on each LLMModel from the loaded env.
Called by :meth:`setup` after ``_load_env()``. Each role ``r``
reads its default model from the generated ``app_env.<r>_model``
field (env var ``<env_prefix>R_MODEL``).
"""
for role, entry in self.llm_models.items():
# The env field is generated from the role declaration, so the
# attribute is not statically known on the settings instance.
entry.model_name = getattr(self.app_env, f"{role}_model")
def _apply_provider_defaults(self) -> None:
"""Populate per-provider defaults on each LLMModel from the config.
Called by :meth:`setup` after ``_load_env()`` (which loads
``app_config``), since ``_setup_models`` runs before the config is
available.
"""
for role, entry in self.llm_models.items():
entry.provider_defaults = self._provider_defaults_for_role(role)
def _apply_user_agent(self) -> None:
"""Set the resolved Klea User-Agent on each ``LLMModel``.
Gives every role the app identity sent to OpenAI-compatible endpoints
(``klea-agent/<version>`` / ``klea-rag/<version>``).
"""
for entry in self.llm_models.values():
entry.user_agent = self.user_agent
def _check_required_models(self) -> None:
"""Warn at startup when required model roles have no default model.
Missing models are not fatal: the server still starts so the
frontends can be used to set models per chat. The warning lists
all model env vars (derived as ``<env_prefix><ROLE>_MODEL``) with
their current state so users can see exactly what to set.
"""
missing = missing_required_roles(self.llm_models)
if not missing:
return
lines = [
(
f"The following models have not been set: {', '.join(missing)}. "
"Please set them before running queries (or from the web UI, "
"Choose models)."
),
"Current configured models:",
]
for role, entry in self.llm_models.items():
env_var = f"{self.env_prefix}{role.upper()}_MODEL"
lines.append(f" {env_var}={entry.model_name or '<not set>'}")
self.logger.warning("\n".join(lines))
def _provider_defaults_for_role(self, role: str) -> dict[str, dict[str, Any]]:
"""Return per-provider default params for *role* from the config.
Reads the optional ``providers`` config section, e.g.::
{"huggingface": {"chat": {"max_output_tokens": 2048}}}
and returns just the entries relevant to *role*, so each
``LLMModel`` carries the defaults for its own role. Shared by all
graphs so ``_setup_models`` implementations do not repeat it.
:param role: Model role (``chat``, ``plan``, ``guard``, ...).
:returns: ``{provider: {param: value}}`` for *role*.
"""
result: dict[str, dict[str, Any]] = {}
# ``app_config`` is the concrete ``config_class`` instance; access
# via getattr so this works for any subclass without extra typing.
providers = getattr(self.app_config, "providers", {})
for provider, role_configs in providers.items():
if role in role_configs:
result[provider] = dict(role_configs[role])
return result
def _bundled_server_config(
self, client_context: ClientContext | None = None
) -> dict[str, Any] | None:
"""Build the bundled tools server config entry from the app config.
Reads ``app_config.general.bundled_tools`` and returns ``None`` when
the bundled tools are disabled. The entry is a *dict*, not a
``StdioMCPServer`` instance, so ``MCPConfig``'s union validator picks
``TransformingStdioMCPServer`` whenever tag filters are present and
applies them; without filters it stays a plain ``StdioMCPServer``.
Apps merge the result into their ``mcpServers`` mapping; when merged
into a domain's config the bundled tools become available to that
domain's tool picker.
The spawn mapping itself is built by
:func:`klea_utils.mcp.spawn.bundled_server_config`, so the bundled
tools server is described in exactly one place. Without
*client_context* the ``env``, ``cwd`` and ``interpreter`` are left
unset, which keeps the previous behaviour: the inherited environment
is used and the server runs with the interpreter running this
process. With a context, only the fields it sets are forwarded.
:param client_context: Optional client-supplied environment for the
bundled tools server.
:returns: Dict-form stdio server config, or ``None`` when disabled.
"""
general = getattr(self.app_config, "general", None)
bundled = getattr(general, "bundled_tools", None)
if bundled is None or not bundled.enabled:
self.logger.debug("Bundled tools server disabled or not configured")
return None
# Tag filters are forwarded only when set, so a disabled/empty filter
# leaves the ``include_tags``/``exclude_tags`` keys out exactly as
# before.
include_tags = list(bundled.include_tags) if bundled.include_tags else None
exclude_tags = list(bundled.exclude_tags) if bundled.exclude_tags else None
# An explicitly empty env is treated as "unset" so a cwd/interpreter
# override alone keeps the inherited environment unchanged.
stores_config = getattr(general, "stores_config", None)
env = {
**(client_context.env if client_context and client_context.env else {}),
**({"KLEA_STORES_CONFIG_FILE": stores_config} if stores_config else {}),
} or None
config: dict[str, Any] = spawn.bundled_server_config(
include_tags=include_tags,
exclude_tags=exclude_tags,
env=env,
cwd=client_context.cwd if client_context else None,
interpreter=client_context.interpreter if client_context else None,
)
self.logger.debug(f"Bundled tools server config: {config}")
return config
@abstractmethod
async def _create_graph(self) -> None:
"""Build and compile the LangGraph, storing it in ``self.graph``.
This is where subclasses define their nodes, edges, and conditional
routing logic.
"""
...
# ------------------------------------------------------------------
# Hook methods -- override for pre/post setup work
# ------------------------------------------------------------------
[docs]
def get_allowed_msgpack_modules(self) -> list[type | tuple[str, ...]]:
"""Return types allowed for checkpoint msgpack deserialization.
Subclasses should override to add their state schemas
(e.g. ``EvaluateAnswerSchema``, ``RetrievalQueryOutput``). The base
list covers shared utils types checkpointed by all graphs.
"""
from fastmcp.client.client import CallToolResult as FastMCPCallToolResult
from mcp.types import EmbeddedResource, ImageContent, TextContent
from klea_utils.graph.schemas import TokenUsage
modules: list[type | tuple[str, ...]] = [
TokenUsage,
ToolCallSchema,
FastMCPCallToolResult,
TextContent,
ImageContent,
EmbeddedResource,
]
try:
from mcp.types import AudioContent
modules.append(AudioContent)
except ImportError:
pass
try:
from mcp.types import CallToolResult as McpCallToolResult
modules.append(McpCallToolResult)
except ImportError:
pass
return modules
[docs]
def context_snapshot(self, state: dict[str, Any]) -> dict[str, Any] | None:
"""Return the session context to surface as a graph-level event.
Called by :meth:`run_graph_astream_events` on every per-superstep
``values`` state snapshot. Returning a dict publishes a
``context`` stream event (change-deduped); returning ``None``
(the default) publishes nothing.
Session context is a *projection of state*, not a node-authored
message: the app defines what the snapshot is by overriding this
hook (e.g. the agent's operating mode), while the emission itself
is structural -- nodes cannot write a ``context`` event through
the ``custom`` channel.
:param state: The current graph state snapshot (a dict).
:returns: A JSON-serializable dict for the ``context`` event, or
``None`` if there is no session context to surface.
"""
return None
def _pre_setup(self) -> None:
"""Hook called before the standard setup sequence.
Override to perform subclass-specific initialisation before
config loading and model setup.
"""
pass
async def _setup_checkpointer(self) -> None:
"""Set up the checkpointer.
``setup()`` calls this hook before ``_load_env()``.
The checkpointer is ``None`` when ``checkpoint="none"``.
"""
if self.checkpointer_mode == "sqlite":
import aiosqlite
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
# ``checkpoints.db`` uses WAL journalling, which keeps the
# ``checkpoints.db-wal`` (write-ahead log) and
# ``checkpoints.db-shm`` (shared-memory index) sidecar files next
# to the database. They are live state for the running backend
# and must not be removed while it is running.
db_path = init_dir(self.storage_root) / "checkpoints.db"
self.logger.debug("Opening sqlite checkpointer at %s", db_path)
conn = await aiosqlite.connect(str(db_path))
# One-off migration: a database created before WAL was enabled
# keeps its old journal mode until it is forced, so verify it here
# and switch it over once, logging the change.
mode_cursor = await conn.execute("PRAGMA journal_mode")
mode_row = await mode_cursor.fetchone()
current_mode = str(mode_row[0]).lower() if mode_row else ""
if current_mode != "wal":
await conn.execute("PRAGMA journal_mode=WAL")
self.logger.info(
"Migrated checkpoints database %s to WAL journal mode (was %s)",
db_path,
current_mode or "unknown",
)
await _apply_sqlite_pragmas_async(conn)
serde = JsonPlusSerializer(
allowed_msgpack_modules=self.get_allowed_msgpack_modules()
)
self.checkpointer = AsyncSqliteSaver(conn, serde=serde)
# Keep raw connection for lifespan cleanup (AsyncSqliteSaver holds it as .conn)
self._checkpointer_conn = conn # type: ignore[attr-defined]
self.logger.debug("Sqlite checkpointer ready")
elif self.checkpointer_mode == "inmemory":
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
serde = JsonPlusSerializer(
allowed_msgpack_modules=self.get_allowed_msgpack_modules()
)
self.checkpointer = InMemorySaver(serde=serde)
self.logger.debug("In-memory checkpointer ready")
def _post_setup(self) -> None:
"""Hook called after the standard setup sequence.
Override to perform subclass-specific finalisation after
the LangGraph has been compiled.
"""
pass
# ------------------------------------------------------------------
# Template method
# ------------------------------------------------------------------
async def _pre_graph(self) -> None:
"""Hook called after MCP client setup but before graph compilation.
Override to perform subclass-specific initialisation that depends
on config and MCP client but must happen before the LangGraph is built.
"""
pass
[docs]
@final
async def setup(self) -> None:
"""Set up the orchestrator.
Calls hooks and template methods in this order:
#. ``_pre_setup()``
#. ``_setup_checkpointer()``
#. ``_setup_models()``: build ``self.llm_models`` (roles and required
flags) that the env schema is generated from.
#. ``_load_env()``: parse the env into ``self.app_env`` using the
schema generated from ``llm_models``, then load the JSON config.
#. ``_configure_resources()``
#. ``_check_required_models()``
#. ``_create_mcp_client()``
#. ``_pre_graph()``
#. ``_create_graph()``
#. ``_post_setup()``
"""
self._pre_setup()
warn_if_root(self.logger)
await self._setup_checkpointer()
self._setup_models()
self._load_env()
self._apply_model_names()
self._apply_provider_defaults()
self._apply_user_agent()
# ``llm_models`` is now fully populated (roles, model names, required
# flags, provider defaults) -- log the resolved config once.
self.logger.debug(
"Resolved model configuration:\n"
+ "\n".join(
f" {role}: "
f"model={entry.model_name or '<not set>'}, "
f"required={entry.required}, "
f"modifiable={entry.modifiable}"
for role, entry in self.llm_models.items()
)
)
self._configure_resources()
self._check_required_models()
self._create_mcp_client()
await self._get_mcp_tools()
await self._get_vector_stores()
await self._pre_graph()
await self._create_graph()
self._post_setup()
# ------------------------------------------------------------------
# Execution methods -- identical across all implementations
# ------------------------------------------------------------------
[docs]
async def run_graph_invoke_state(
self, state: dict, thread_id: str = "default_thread"
) -> dict:
"""Run the graph, accepting and returning full state dicts.
:param state: Initial graph state (must contain ``query`` key)
:param thread_id: Session/thread identifier for checkpointing
:returns: Final graph state
"""
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
if "query" not in state:
raise ValueError("state must contain 'query' key")
if self.graph is None:
raise RuntimeError("Graph not compiled. Call setup() first.")
final_state = await self.graph.ainvoke(state, config=config)
self.logger.debug(final_state)
return final_state
[docs]
async def run_graph_invoke(
self,
query: str | Command | None,
thread_id: str = "default_thread",
*,
extra_state: dict[str, Any] | None = None,
context: KleaRunContext | dict[str, Any] | None = None,
) -> str:
"""Run the graph with a simple string query.
:param query: User query string, ``None`` to **resume** the thread
from its last checkpoint (the failed node is re-run; the entry
node is not), or a :class:`~langgraph.types.Command` (e.g.
``Command(resume=...)``) to continue a paused thread. On a
``None`` resume ``extra_state`` is ignored; a ``Command`` is
passed through verbatim.
:param thread_id: Session/thread identifier for checkpointing
:param extra_state: Optional initial state fields merged into the
invocation (e.g. an app-specific ``mode`` request). These are
passed to ``graph.ainvoke`` alongside ``query`` and validated
against the graph's state schema.
:param context: Per-run runtime context (ADR-0033), forwarded
verbatim to the graph run. Accepts a ``KleaRunContext`` (or
app subclass) instance or a plain dict -- dicts are
coerced/validated against the graph's ``context_schema`` at
the run boundary. The base contract is generic:
``model_overrides`` is a framework-provided convention (apps
may diverge at the node layer) and apps add their own keys.
:returns: The ``message_for_user`` field from the final state
:note: This is a bare ``ainvoke`` with no ``values``-event loop, so
``context_snapshot`` is never invoked here and the session
context projection is not produced on this path. Callers that
need it should use :meth:`run_graph_astream_events` (``context``
stream events) or the hydration endpoint
``GET /chat/{user_id}/{chat_id}/context`` (ADR-0032).
"""
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
if self.graph is None:
raise RuntimeError("Graph not compiled. Call setup() first.")
graph_input: dict[str, Any] | Command | None
if query is None:
# Resume: pass ``None`` so LangGraph continues from the checkpoint
# at the failed node (it does not re-enter the entry node).
self.logger.info("Resuming thread %s from checkpoint", thread_id)
graph_input = None
elif isinstance(query, Command):
# A Command (e.g. ``Command(resume=...)``) is passed through
# verbatim; it carries its own update/resume payload.
self.logger.info("Continuing thread %s with a Command", thread_id)
graph_input = query
else:
graph_input = {"query": query}
if extra_state:
graph_input.update(extra_state)
self._apply_access_level_default(graph_input, extra_state)
final_state = await self.graph.ainvoke(
graph_input, config=config, context=context
)
self.logger.debug(f"{final_state =}")
if message := final_state.get("message_for_user", None):
return message
else:
return "I was unable to answer"
[docs]
async def run_graph_stream(
self,
query: str | None,
thread_id: str = "default_thread",
*,
extra_state: dict[str, Any] | None = None,
context: KleaRunContext | dict[str, Any] | None = None,
):
"""Run the graph and yield intermediate ``message_for_user`` values.
:param query: User query string, or ``None`` to resume the thread
from its last checkpoint (see :meth:`run_graph_invoke`).
:param thread_id: Session/thread identifier for checkpointing
:param extra_state: Optional initial state fields merged into the
invocation alongside ``query`` (e.g. an app-specific ``mode``
request). Ignored on resume.
:param context: Per-run runtime context (ADR-0033), forwarded
verbatim to the graph run (see :meth:`run_graph_invoke`).
:yields: ``message_for_user`` strings from each node
"""
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
if self.graph is None:
raise RuntimeError("Graph not compiled. Call setup() first.")
graph_input: dict[str, Any] | None
if query is None:
self.logger.info("Resuming thread %s from checkpoint", thread_id)
graph_input = None
else:
graph_input = {"query": query}
if extra_state:
graph_input.update(extra_state)
self._apply_access_level_default(graph_input, extra_state)
async for chunk in self.graph.astream(
graph_input, config=config, context=context
):
for node, state in chunk.items():
self.logger.debug(f"{node}: {state!r}")
if message := state.get("message_for_user", None):
self.logger.info(f"User message: {message}")
yield message
else:
self.logger.debug(f"Working in node: {node}")
[docs]
async def graph_stream(
self,
query: str | None,
thread_id: str = "default_thread",
*,
extra_state: dict[str, Any] | None = None,
context: KleaRunContext | dict[str, Any] | None = None,
) -> Any:
"""Run the graph and return the raw astream result.
:param query: User query string, or ``None`` to resume the thread
from its last checkpoint (see :meth:`run_graph_invoke`).
:param thread_id: Session/thread identifier for checkpointing
:param extra_state: Optional initial state fields merged into the
invocation alongside ``query`` (e.g. an app-specific ``mode``
request). Ignored on resume.
:param context: Per-run runtime context (ADR-0033), forwarded
verbatim to the graph run (see :meth:`run_graph_invoke`).
:returns: Raw async generator from ``graph.astream()``
"""
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
if self.graph is None:
raise RuntimeError("Graph not compiled. Call setup() first.")
graph_input: dict[str, Any] | None
if query is None:
self.logger.info("Resuming thread %s from checkpoint", thread_id)
graph_input = None
else:
graph_input = {"query": query}
if extra_state:
graph_input.update(extra_state)
self._apply_access_level_default(graph_input, extra_state)
res = self.graph.astream(graph_input, config=config, context=context)
return res
[docs]
async def run_graph_astream_events(
self,
query: str | Command | None,
thread_id: str = "default_thread",
*,
extra_state: dict[str, Any] | None = None,
context: KleaRunContext | dict[str, Any] | None = None,
):
"""Run the graph and yield structured streaming events.
Yields dicts with:
``{"type": "progress", "node": "<label>", "data": {...}}``
When the graph enters a new node, or when the node's progress
heading changes (e.g. an LLM invoke retry). ``data["heading"]``
is the line shown while the node runs.
``{"type": "inspect", "node": "<label>", "data": {...}}``
Inspection data from a node after execution (summary + details)
``{"type": "token", "content": "<chunk>", "node": "<label>"}``
LLM token chunk from the current node. Emitted only for nodes in
:attr:`token_stream_nodes` (free-text answer nodes), not for
structured-output nodes.
``{"type": "usage", "node": "<label>", "data": {...}}``
Per-node token usage (input / output / total tokens)
``{"type": "context", "data": {...}}``
Session-level context (app-defined, e.g. an operating mode and
its assurance), emitted by a node via a ``context`` custom event
``{"type": "interrupt", "node": "<label>", "data": {...}}``
The run paused at a node's ``interrupt()``. ``data`` carries the
node-authored payload (``kind``, ``question``) plus the
framework's ``interrupt_id``; when a pause is emitted no
``complete`` event follows.
``{"type": "complete", "message_for_user": "<answer>"}``
Final answer from the completed graph
Uses LangGraph's ``astream_events`` v3 protocol. Progress events
from all nodes (LLM and non-LLM) arrive via the ``custom`` channel.
LLM token output is read from the ``messages`` channel.
A ``StreamTransformer`` enables the custom channel so those events
flow through.
Reference: https://docs.langchain.com/oss/python/langgraph/event-streaming
Interrupts: https://docs.langchain.com/oss/python/langgraph/interrupts
:param query: User query string, ``None`` to **resume** the thread
from its last checkpoint (see the fault-tolerance note), or a
:class:`~langgraph.types.Command` (e.g. ``Command(resume=...)``) to
continue a paused thread. On a ``None`` resume the graph re-enters
the failed node rather than the entry node, so per-turn state is
preserved; a ``Command`` is passed through verbatim.
:param thread_id: Session/thread identifier for checkpointing
:param extra_state: Optional initial state fields merged into the
invocation alongside ``query`` (e.g. an app-specific ``mode``
request). Ignored when *query* is ``None`` or a ``Command``.
:param context: Per-run runtime context (ADR-0033), forwarded
verbatim to the graph run (see :meth:`run_graph_invoke`).
:yields: Structured event dicts
"""
config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
if self.graph is None:
raise RuntimeError("Graph not compiled. Call setup() first.")
graph_input: dict[str, Any] | Command | None
if query is None:
# Resume: pass ``None`` so LangGraph continues from the checkpoint
# at the failed node (it does not re-enter the entry node).
self.logger.info("Resuming thread %s from checkpoint", thread_id)
graph_input = None
elif isinstance(query, Command):
# A Command (e.g. ``Command(resume=...)``) is passed through
# verbatim; it carries its own update/resume payload.
self.logger.info("Continuing thread %s with a Command", thread_id)
graph_input = query
else:
graph_input = {"query": query}
if extra_state:
graph_input.update(extra_state)
self._apply_access_level_default(graph_input, extra_state)
stream = await self.graph.astream_events(
graph_input,
config=config,
version="v3",
transformers=[_CustomChannelEnabler],
context=context,
)
current_node = ""
current_heading = ""
node_start = time.monotonic()
total_start = time.monotonic()
last_values: dict = {}
# Last emitted session-context snapshot, for change-dedup.
last_context: dict[str, Any] | None = None
async for event in stream:
method = event["method"]
if method == "custom":
data = event["params"]["data"]
if not isinstance(data, dict):
continue
event_type = data.get("type")
if event_type == "progress":
node = data.get("node")
if not node:
continue
payload = data.get("data") or {}
heading = payload.get("heading", "")
if node != current_node:
now = time.monotonic()
if current_node:
self.logger.debug(
"Node [%s] took %.2fs",
current_node,
now - node_start,
)
node_start = now
current_node = node
current_heading = heading
self.logger.debug(f"Progress: {current_node}")
yield {
"type": "progress",
"node": current_node,
"data": payload,
}
elif heading and heading != current_heading:
# Same node, changed heading (e.g. an LLM invoke retry):
# forward it without resetting the node timer.
current_heading = heading
yield {
"type": "progress",
"node": current_node,
"data": payload,
}
elif event_type in ("inspect", "state", "usage", "tool", "notice"):
node = data.get("node")
if not node:
continue
data_out = data.get("data", {}).copy()
data_out["timing_seconds"] = round(time.monotonic() - node_start, 2)
yield {
"type": event_type,
"node": node,
"data": data_out,
}
elif method == "messages":
# Token deltas are only useful for user-facing free-text
# nodes; structured-output nodes emit JSON fragments and are
# not in ``token_stream_nodes`` (see its docstring).
if current_node not in self.token_stream_nodes:
continue
data = event["params"]["data"]
for item in data:
if not isinstance(item, dict):
continue
if item.get("event") == "content-block-delta":
delta = item.get("delta", {})
if "text" in delta:
yield {
"type": "token",
"content": delta["text"],
"node": current_node,
}
elif method == "values":
last_values = event["params"]["data"]
# ``values`` snapshots arrive as a dict (TypedDict state) or
# as the pydantic state instance; ``_normalise_state_snapshot``
# gives ``context_snapshot`` a plain dict in the stream path --
# the same normalization the hydration endpoint applies on a
# direct checkpoint read (see ``klea_utils.api.context``).
snapshot = _normalise_state_snapshot(last_values)
# Graph-level session context is a *projection of state*: it is
# derived here from the per-superstep ``values`` snapshot via
# ``context_snapshot`` (default None), never written by a node.
# This keeps the event structural -- nodes cannot author a
# ``context`` event -- while the app defines what context means.
# (Local named context_data: the method's ``context`` parameter
# is the ADR-0033 per-run runtime context.)
context_data = self.context_snapshot(snapshot)
if context_data is not None and context_data != last_context:
last_context = context_data
yield {"type": "context", "data": context_data}
total_elapsed = time.monotonic() - total_start
if current_node:
self.logger.debug(
"Node [%s] took %.2fs",
current_node,
time.monotonic() - node_start,
)
# Documented v3 interrupt projection (ADR-0046): after the stream
# drains, ``stream.interrupted`` says whether the run paused and
# ``stream.interrupts`` carries the payloads. On the async lane these
# are async methods; on the sync lane they are properties -- handle
# both. A pause must not emit a terminal ``complete``.
interrupted = getattr(stream, "interrupted", None)
if callable(interrupted):
interrupted = await interrupted()
if interrupted:
interrupts = getattr(stream, "interrupts", None)
if callable(interrupts):
interrupts = await interrupts()
for intr in interrupts or ():
self.logger.info(
"Graph paused at interrupt %s (node %s)",
getattr(intr, "id", None),
current_node,
)
yield _interrupt_event(intr, current_node)
self.logger.info("Graph paused at an interrupt after %.2fs", total_elapsed)
return
self.logger.info("Graph completed in %.2fs", total_elapsed)
message = ""
if last_values:
if isinstance(last_values, dict):
message = last_values.get("message_for_user", "")
elif hasattr(last_values, "message_for_user"):
message = last_values.message_for_user
yield {"type": "complete", "message_for_user": message}