Orchestrator framework

class klea_utils.graph.base.BaseLangGraph(logging_level: int = 20, checkpoint: str = 'inmemory', log_file: bool = True)[source]

Bases: 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.

config_class: type[BaseModel]

Pydantic BaseModel class for configuration loading. Subclasses must set this to their AppConfig class.

config_file_default: 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.

context_snapshot(state: dict[str, Any]) → dict[str, Any] | None[source]

Return the session context to surface as a graph-level event.

Called by 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.

Parameters:

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.

env_class

Pydantic BaseSettings class for env loading. Subclasses need not set this: it is generated at load time from self.llm_models (see _build_env_class()).

alias of BaseModel

env_file_default: str = 'config.env'

Default config file name if the environment variable is not set.

env_prefix: str = ''

Prefix prepended to generated env var names (e.g. "KLEA_AGENT_"). Each model role r becomes the env var <env_prefix>R_MODEL.

env_var: str = 'ENV_FILE'

Name of the environment variable that controls the env file path.

get_allowed_msgpack_modules() → list[type | tuple[str, ...]][source]

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.

graph_name: str = 'BaseLangGraph'

Logger name for this orchestrator, also used as the app name for platformdirs data/cache directories.

async graph_stream(query: str | None, thread_id: str = 'default_thread', *, extra_state: dict[str, Any] | None = None, context: KleaRunContext | dict[str, Any] | None = None) → Any[source]

Run the graph and return the raw astream result.

Parameters:

  • query – User query string, or None to resume the thread from its last checkpoint (see run_graph_invoke()).

  • thread_id – Session/thread identifier for checkpointing

  • extra_state – Optional initial state fields merged into the invocation alongside query (e.g. an app-specific mode request). Ignored on resume.

  • context – Per-run runtime context (ADR-0033), forwarded verbatim to the graph run (see run_graph_invoke()).

Returns:

Raw async generator from graph.astream()

mcp_session_pool: Any | 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.

new_mcp_client(client_context: ClientContext | None = None) → Client | None[source]

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 _create_mcp_client().

Parameters:

client_context – Optional client-supplied environment for the bundled tools server.

Returns:

A new MCP client, or None when nothing is configured.

async run_graph_astream_events(query: str | Command | None, thread_id: str = 'default_thread', *, extra_state: dict[str, Any] | None = None, context: KleaRunContext | dict[str, Any] | None = None)[source]

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 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

Parameters:

  • query – User query string, None to resume the thread from its last checkpoint (see the fault-tolerance note), or a 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.

  • thread_id – Session/thread identifier for checkpointing

  • 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.

  • context – Per-run runtime context (ADR-0033), forwarded verbatim to the graph run (see run_graph_invoke()).

Yields:

Structured event dicts

async run_graph_invoke(query: str | Command | None, thread_id: str = 'default_thread', *, extra_state: dict[str, Any] | None = None, context: KleaRunContext | dict[str, Any] | None = None) → str[source]

Run the graph with a simple string query.

Parameters:

  • 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 Command (e.g. Command(resume=...)) to continue a paused thread. On a None resume extra_state is ignored; a Command is passed through verbatim.

  • thread_id – Session/thread identifier for checkpointing

  • 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.

  • 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 run_graph_astream_events() (context stream events) or the hydration endpoint GET /chat/{user_id}/{chat_id}/context (ADR-0032).

async run_graph_invoke_state(state: dict, thread_id: str = 'default_thread') → dict[source]

Run the graph, accepting and returning full state dicts.

Parameters:

  • state – Initial graph state (must contain query key)

  • thread_id – Session/thread identifier for checkpointing

Returns:

Final graph state

async run_graph_stream(query: str | None, thread_id: str = 'default_thread', *, extra_state: dict[str, Any] | None = None, context: KleaRunContext | dict[str, Any] | None = None)[source]

Run the graph and yield intermediate message_for_user values.

Parameters:

  • query – User query string, or None to resume the thread from its last checkpoint (see run_graph_invoke()).

  • thread_id – Session/thread identifier for checkpointing

  • extra_state – Optional initial state fields merged into the invocation alongside query (e.g. an app-specific mode request). Ignored on resume.

  • context – Per-run runtime context (ADR-0033), forwarded verbatim to the graph run (see run_graph_invoke()).

Yields:

message_for_user strings from each node

final async setup() → None[source]

Set up the orchestrator.

Calls hooks and template methods in this order:

  1. _pre_setup()

  2. _setup_checkpointer()

  3. _setup_models(): build self.llm_models (roles and required flags) that the env schema is generated from.

  4. _load_env(): parse the env into self.app_env using the schema generated from llm_models, then load the JSON config.

  5. _configure_resources()

  6. _check_required_models()

  7. _create_mcp_client()

  8. _pre_graph()

  9. _create_graph()

  10. _post_setup()

storage_root

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 storage_root_env.

storage_root_env: str | None = None

Environment variable that relocates this graph’s storage root (data and runtime artifacts). None disables relocation; the agent sets this to KLEA_AGENT_HOME.

token_stream_nodes: set[str]

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.

user_agent

Klea identity sent to OpenAI-compatible endpoints (User-Agent).

Shared state schemas

Schemas shared by LangGraph orchestrators.

File: klea_utils/graph/schemas.py

Copyright 2026 Ankur Sinha Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>

class klea_utils.graph.schemas.TokenUsage(*, input_tokens: int = 0, output_tokens: int = 0, total_tokens: int = 0, reasoning_tokens: int = 0, cached_tokens: int = 0, role: str = '')[source]

Bases: BaseModel

Token usage for one node, or accumulated across a graph run.

role, reasoning_tokens and cached_tokens are populated on the per-node usage stream event so a run can be benchmarked per node and per role. When this schema is accumulated into the graph state’s usage_metrics (see add_token_usage()) only the numeric fields are summed; role is identity, not a quantity, so it is dropped from the aggregate.

cached_tokens: int

Prompt tokens served from the provider’s prompt cache (cache reads only; cache-creation/write tokens remain part of input_tokens).

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

reasoning_tokens: int

Hidden reasoning/thinking tokens reported by the provider (subset of output_tokens when the provider exposes them).

role: str

Model role that produced this usage (e.g. chat, plan, guard).

Shared base state

All application graph states subclass klea_utils.graph.state.BaseGraphSchema (utils_pkg/klea_utils/graph/state.py), which holds the fields the shared nodes and the streaming layer rely on: query, messages, guard_decision, context_summary, summarised_till, message_for_user, tool_calls, tool_results and the reduced usage_metrics. Applications add their own fields (the agent’s mode/plan/…; the RAG’s query_domains and reference material).

Reducers

Reducers shared by LangGraph orchestrators.

File: klea_utils/graph/reducers.py

Copyright 2026 Ankur Sinha Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>

klea_utils.graph.reducers.add_token_usage(left: TokenUsage | dict[str, Any], right: TokenUsage | dict[str, Any]) → TokenUsage[source]

Add token usage updates from sequential or concurrent graph nodes.

Only the numeric counters are summed. role identifies the node that produced a single update, so it is not meaningful for the aggregate and is dropped (left empty) rather than concatenated.

Per-run runtime context

Per-run runtime context for Klea graphs.

Carried by LangGraph’s Runtime context mechanism (ADR-0033): the run methods accept a context= value forwarded to ainvoke/astream/ astream_events, and nodes read it inside execution via ambient get_runtime(). This is the framework-native replacement for the hand-rolled model_overrides_ctx contextvar removed in ADR-0033; the LLM’s per-invocation RunnableConfig merge is unchanged (ADR-0014).

See devdocs/system/streams.md (context event) and devdocs/system/api-sse-sequence.md.

File: klea_utils/graph/context.py

Copyright 2026 Ankur Sinha Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>

class klea_utils.graph.context.KleaRunContext(*, model_overrides: dict[str, dict[str, ~typing.Any]]=<factory>, client_context: ClientContext | None = None, **extra_data: Any)[source]

Bases: BaseModel

Per-run runtime context for one graph invocation.

model_overrides is the conventional key the shared nodes in klea_utils.nodes consume (chat_core populates it from the sessions database). client_context carries the runtime environment a client (currently the TUI) wants its per-session tool server to use; it is accepted from loopback peers only and never persisted. extra="allow" keeps the schema generic (ADR-0031): an app that wants its own keys validated at the run boundary subclasses this model and registers the subclass as its context_schema (LangGraph coerces the context dict via context_schema(**context)); an app that prefers to validate custom keys itself reads them from model_extra.

client_context: ClientContext | None

Client-supplied environment for the per-session bundled MCP server (cwd, interpreter, env). Runtime-only and never checkpointed.

model_config = {'extra': 'allow'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_overrides: dict[str, dict[str, Any]]

Per-role model/API overrides; populated by chat_core from the sessions database (ADR-0014 decisions on sourcing, merge and key masking remain governing; ADR-0033 only changed the transport).

klea_utils.graph.context.model_overrides_from_context(context: KleaRunContext | dict[str, Any] | None) → dict[str, dict[str, Any]][source]

Return the model_overrides slice from a run context, or {}.

Tolerates a KleaRunContext (attribute access), a plain dict (context["model_overrides"]), or None – Runtime.context is None when no context= was passed (ADR-0033, probe-verified), so the shared nodes must never see None.

Parameters:

context – The run context, e.g. get_runtime().context.

Returns:

The overrides dict (never None).