Orchestrator framework

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

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

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, ...]]

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

Parameters:

• query – User query string

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

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

Returns:

Raw async generator from graph.astream()

async run_graph_astream_events(query: str, 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>"}

When the graph enters a new node (via write_custom_stream)

{"type": "info", "node": "<label>", "data": {...}}

Structured summary data from a node after execution

{"type": "debug", "node": "<label>", "data": {...}}

Full data dump from a node after execution

{"type": "token", "content": "<chunk>", "node": "<label>"}

LLM token chunk from the current node

{"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": "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

Parameters:

• query – User query string

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

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

Parameters:

• query – User query string

• 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

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

Parameters:

• query – User query string

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

• 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

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

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)

Bases: BaseModel

Token usage accumulated across the nodes in a graph run.

model_config = {}

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

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, int], right: TokenUsage | dict[str, int]) → TokenUsage

Add token usage updates from sequential or concurrent graph nodes.

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

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>, **extra_data: Any)

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

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

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