#!/usr/bin/env python3
"""
Base node classes for LangGraph processing nodes
See ``devdocs/system/structured-output-fallback.md`` and
``devdocs/system/prompt-conventions.md``; ADRs 0017 and 0044.
File: klea_utils/nodes/base.py
Copyright 2026 Ankur Sinha
Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>
"""
from __future__ import annotations
import asyncio
import inspect
import json
import logging
from collections.abc import Awaitable, Callable
from functools import cached_property
from pathlib import Path
from textwrap import dedent
from typing import Any, ClassVar, Literal, cast
from langchain_core.exceptions import OutputParserException
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.prompt_values import PromptValue
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.utils.function_calling import convert_to_json_schema
from langgraph.runtime import get_runtime
from pydantic import BaseModel, ValidationError
from klea_utils.graph.context import model_overrides_from_context
from klea_utils.plogging import mask_sensitive
from ..errors import LLMInvocationErrorCategory, PromptTemplateError
from ..llm import (
DEFAULT_MAX_OUTPUT_TOKENS,
add_memory_to_prompt,
classify_llm_invocation_error,
escape_prompt_braces,
estimate_input_tokens,
get_provider_allowed_fields,
get_recent_messages,
get_token_limit_param,
is_empty_structured_parse_error,
is_output_empty,
is_output_truncated,
is_structured_capability_rejection,
is_structured_output_failure,
load_prompt,
mark_structured_output_unsupported,
parse_output_with_thought,
resolve_langchain_endpoint,
resolve_output_token_limit,
structured_output_known_unsupported,
)
from ..models_catalog import probe_endpoint_model_limits
from .abstract import AbstractLLMNode, NodeStreamData, NodeStreamEvent
from .context import LLMNodeContext
#: Product name reported to the model by user-facing nodes.
KLEA_AGENT_NAME = "Klea"
#: Constant identity line prepended to user-facing system prompts. Keeping it
#: identical across nodes makes it part of the cacheable prompt prefix.
_AGENT_IDENTITY = f"Agent: {KLEA_AGENT_NAME}"
#: Max times to retry an invoke that overflowed the context window, each
#: time shrinking the reserved output window to free headroom.
MAX_CONTEXT_OVERFLOW_RETRIES = 3
#: Max times to retry an invoke that returned an empty response. Some
#: providers (notably HuggingFace) intermittently return a successful but
#: blank message; the retry is a plain re-invoke with no window change.
MAX_EMPTY_OUTPUT_RETRIES = 2
#: Character budget for the verbatim recent-conversation window injected into
#: memory-enabled node prompts. Older history is represented by
#: ``context_summary`` instead.
DEFAULT_RECENT_MEMORY_CHARS = 10_000
#: Max times to retry a transiently failed invoke (a timeout, e.g. a stalled
#: stream that never produced a chunk). Rate limits, auth and model-not-found
#: failures are deliberately *not* retried here: the provider SDK already
#: retries ``429``/``5xx`` and honours ``Retry-After``, so retrying again
#: would multiply the attempts and can amplify a rate limit.
MAX_TRANSIENT_RETRIES = 2
#: Base backoff (seconds) between transient retries; doubles each attempt.
TRANSIENT_RETRY_BACKOFF_S = 1.0
def _current_session_id() -> str | None:
"""Return the current LangGraph ``thread_id``, or ``None`` outside a run.
The thread id is Klea's stable per-conversation session id. It is sent as
the opencode ``x-opencode-session`` request header so the backend can
optimise routing and prompt caching (see
:func:`klea_utils.llm.apply_provider_overrides`).
"""
try:
from langgraph.config import get_config
config = get_config()
except Exception: # noqa: BLE001 - no run context (e.g. unit tests)
return None
thread_id = (config or {}).get("configurable", {}).get("thread_id")
return thread_id or None
#: Max times to retry an invoke whose output was truncated (``finish_reason
#: == "length"``), each time growing the reserved output window. The
#: budget covers the full climb from the smallest node window to the
#: largest advertised context (262144) so heavy-thinking models are not
#: cut short, while still bounding pathological repeated truncations.
MAX_TRUNCATION_RETRIES = 15
#: Linear-phase step for truncation retries: the window grows by
#: ``TRUNCATION_LINEAR_STEP`` while below ``TRUNCATION_LINEAR_CAP``.
#: Covers the common case -- a model cut off a few hundred to a couple of
#: thousand tokens short of finishing -- with small, predictable
#: reservations.
TRUNCATION_LINEAR_STEP = 2048
#: Ceiling for the linear phase; above it the window grows by
#: :data:`TRUNCATION_PHASE2_STEP`.
TRUNCATION_LINEAR_CAP = 16384
#: Large-step phase for truncation retries: once the linear phase is
#: exhausted the window grows by this fixed step up to the output ceiling.
#: Fixed-size steps (not exponential) keep reservations predictable for
#: inference engines that allocate resources against ``max_tokens`` (e.g.
#: vLLM KV-cache), avoiding the spikes of a doubling ladder.
TRUNCATION_PHASE2_STEP = 32768
#: Output-window ceiling for truncation retries. Reasoning-capable models
#: can spend a large, unpredictable number of tokens thinking before
#: producing a small final answer, so the grow ladder clamps at a generous
#: ceiling (mirrors opencode's ``OUTPUT_TOKEN_MAX``). This is the
#: fallback when the model's context is unknown; for a custom endpoint
#: that advertises a larger ``max_model_len``,
#: :meth:`BaseLLMNode._jump_output_target` raises the ceiling to the
#: endpoint's remaining context. It is still clamped to the model's
#: catalog output limit / total budget by ``resolve_output_token_limit``.
MAX_OUTPUT_TOKENS_CEILING = 65536
#: Floor for the reserved output window when shrinking it on overflow.
MIN_OUTPUT_TOKENS = 64
def _resolve_ref(ref: str, root: dict[str, Any]) -> dict[str, Any] | None:
"""Resolve a JSON-schema ``$ref`` (``#/defs/...``) against *root*.
:param ref: The ``$ref`` string (only local ``#/...`` pointers).
:param root: The root schema carrying ``$defs``.
:returns: The referenced schema fragment, or ``None`` when unresolvable.
"""
if not ref.startswith("#/"):
return None
node: Any = root
for part in ref[2:].split("/"):
if not isinstance(node, dict) or part not in node:
return None
node = node[part]
return node if isinstance(node, dict) else None
def _schema_to_example(
schema: dict[str, Any], root: dict[str, Any] | None = None
) -> Any:
"""Generate a placeholder example value from a JSON schema fragment.
Walks a JSON Schema fragment (as produced by
``convert_to_json_schema``) and returns a placeholder value for each
type, so the prompt can show the model a concrete instance to imitate
instead of the abstract schema definition (which invites the model to
echo the schema back verbatim instead of producing an instance).
A concrete ``examples`` entry is preferred when present (e.g. from
``json_schema_extra``); ``$ref`` is resolved against *root*; and
``anyOf``/``oneOf`` pick the first non-null branch.
:param schema: JSON Schema fragment (a ``{"type": ...}`` dict)
:param root: The root schema, for resolving ``$ref`` pointers.
:returns: A placeholder value matching the schema's type
"""
if not isinstance(schema, dict):
return None
examples = schema.get("examples")
if isinstance(examples, list) and examples:
return examples[0]
ref = schema.get("$ref")
if isinstance(ref, str) and root is not None:
resolved = _resolve_ref(ref, root)
if resolved is not None:
return _schema_to_example(resolved, root)
for keyword in ("anyOf", "oneOf"):
branches = schema.get(keyword)
if isinstance(branches, list) and branches:
non_null = [
branch
for branch in branches
if not (isinstance(branch, dict) and branch.get("type") == "null")
]
chosen = non_null[0] if non_null else branches[0]
return _schema_to_example(chosen, root)
if schema.get("enum"):
return schema["enum"][0]
match schema.get("type"):
case "string":
return "text"
case "integer" | "number":
return 0
case "boolean":
return True
case "null":
return None
case "array":
return [_schema_to_example(schema.get("items", {}), root)]
case "object":
return {
key: _schema_to_example(value, root)
for key, value in schema.get("properties", {}).items()
}
case _:
return None
#: User-facing fallback when an LLM produced no usable answer and the node
#: has no more specific fallback of its own.
EMPTY_RESULT_FALLBACK = "Something went wrong. Please retry."
def _is_empty_result(result: Any, schema: type[BaseModel] | None = None) -> bool:
"""Return True if *result* carries no usable content.
A structured output that parsed to an all-default instance (the model
echoed the schema back instead of producing an instance) compares equal
to a freshly-constructed default; a non-structured response is empty
when its content is blank. Used to flag silently-degraded LLM output.
:param result: Processed output from :meth:`BaseLLMNode._process_output`
:param schema: The node's output schema, or ``None`` for non-structured
:returns: True when nothing usable was produced
"""
if schema is not None:
return isinstance(result, schema) and result == schema()
if isinstance(result, AIMessage):
return is_output_empty(result)
return False
[docs]
class BaseLLMNode[TState: BaseModel, TOutput: BaseModel](
AbstractLLMNode[TState, TOutput]
):
"""Base class for LangGraph nodes that load prompts from files.
Extends AbstractLLMNode with:
- File-based prompt loading via load_prompt()
- Optional memory support (appends memory content to system prompt)
- Auto-derived prompt registry location from subclass file path
Prompt files are expected to be named ``{prefix}_system.md`` and
``{prefix}_user.md``.
Subclasses can override ``prompt_prefix`` or ``prompt_registry_location``
via the setter if the defaults (lowercase class name / sibling ``prompts/``)
are not appropriate.
"""
#: Report the agent and effective model identity in the system prompt.
#: Opt in on nodes that compose a user-facing reply, so the model can
#: answer "which agent/model are you running" with the values that
#: actually produced the reply. The constant agent line is prepended
#: (cacheable prefix); the per-run model line is appended after the
#: stable prompt sections. Off by default: most nodes never need it and
#: it costs prompt tokens.
include_model_identity: ClassVar[bool] = False
def __init__(
self,
logger: logging.Logger,
label: str,
llm_models: dict[str, Any],
output_schema: type[TOutput] | None,
memory: bool = False,
):
"""Initialize with file-based prompt loading and memory support.
:param logger: Logger instance
:param label: Human-readable label for UI progress display
:param llm_models: ``{role: LLMModel}`` dict (from ``BaseLangGraph.llm_models``)
:param output_schema: Pydantic schema for structured output
:param memory: Whether to append memory content to the system prompt
"""
super().__init__(logger, label, llm_models, output_schema=output_schema)
self._prompt_prefix: str | None = None
self._prompt_registry_location: Path | None = None
self.memory = memory
@property
def prompt_prefix(self) -> str:
"""Return the prompt file prefix.
Falls back to the lowercase class name if not explicitly set.
"""
if self._prompt_prefix is not None:
return self._prompt_prefix
return self.__class__.__name__
@prompt_prefix.setter
def prompt_prefix(self, value: str) -> None:
"""Set the prompt file prefix."""
self._prompt_prefix = value
@property
def prompt_registry_location(self) -> Path:
"""Return path to the prompts directory.
Falls back to a sibling ``prompts/`` directory relative to the
subclass file if not explicitly set.
"""
if self._prompt_registry_location is not None:
return self._prompt_registry_location
subclass_file = inspect.getfile(self.__class__)
loc = Path(subclass_file).parent / "prompts"
self.logger.debug(f"No prompt registry location set. Falling back to {loc}")
return loc
@prompt_registry_location.setter
def prompt_registry_location(self, value: Path) -> None:
"""Set the prompts directory path."""
self._prompt_registry_location = value
@property
def output_schema(self) -> type[TOutput] | None:
"""Return Pydantic schema for structured output if required"""
return self._output_schema
@output_schema.setter
def output_schema(self, value: type[TOutput] | None) -> None:
"""Set Pydantic schema for structured output"""
self._output_schema = value
[docs]
@cached_property
def output_schema_json(self) -> dict[str, Any]:
"""Return JSON schema string for use in prompts."""
return convert_to_json_schema(self.output_schema) if self.output_schema else {}
def _configure_llm(
self, ctx: LLMNodeContext[TOutput]
) -> tuple[Runnable, RunnableConfig]:
"""Configure LLM and build per-invoke config.
Returns the raw ``instance`` (a ``_ConfigurableModel``) without
wrapping it --- structured output is applied inside ``_invoke_llm``
so that providers that reject ``response_format`` can fall back to
prompt-based structured output.
:param ctx: Per-run node context.
:returns: (llm_instance, config_dict) where config_dict is a
``RunnableConfig`` with ``configurable`` populated.
"""
inst = self._llm_entry.instance
config = self._build_invoke_config(ctx)
self.logger.debug(
f"{self.model_role = }\n{mask_sensitive(config.get('configurable', {})) = }"
)
return inst, config
def _build_invoke_config(self, ctx: LLMNodeContext[TOutput]) -> RunnableConfig:
"""Build the per-invoke RunnableConfig.
Delegates the full merge (role defaults -> context overrides ->
node defaults -> provider defaults, including model-string parsing)
to ``LLMModel.build_config()``, then resolves the bounded max-output
token param (translated + clamped to the catalog's output/context
limits) before applying provider field filtering to strip fields
invalid for the resolved provider.
"""
# Per-run model overrides come from the LangGraph Runtime context
# (get_runtime().context, ADR-0033), not an ad-hoc contextvar.
# Ambient access: no node signature takes a ``runtime`` parameter;
# the value is static for the whole run. Outside a graph run
# (unit tests) get_runtime() raises, so tests drive the merge via
# a Runtime-context harness instead of calling this directly.
ctx_val = model_overrides_from_context(get_runtime().context)
role_overrides = ctx_val.get(self.model_role, {})
self.logger.debug(
f"{mask_sensitive(ctx_val) = }\n"
f"{self.model_role = }\n"
f"{mask_sensitive(role_overrides) = }\n"
f"{self.model_defaults = }"
)
# Delegate merge + model parsing to LLMModel.
config = self._llm_entry.build_config(
context_overrides=role_overrides,
node_defaults=self.model_defaults,
session_id=_current_session_id(),
)
# Get the merged configurable dict for provider field filtering.
overrides: dict[str, Any] = config["configurable"]
# A missing model (no env default and no per-chat override) cannot
# invoke the LLM. Raise a clear, actionable error instead of the
# confusing provider-level "Missing credentials" error that an empty
# model would otherwise produce.
if not overrides.get("model"):
raise RuntimeError(
f"No model configured for role '{self.model_role}'. "
f"Set the {self.model_role.upper()}_MODEL environment "
"variable (e.g. KLEA_AGENT_CHAT_MODEL) or set a model for "
"this chat from the web UI (Choose models)."
)
# --- Bounded output tokens ---
# Guarantee a finite max-output token param for the resolved
# provider, clamped to the model's catalog output limit and total
# budget (input + output <= context). Must run before provider
# field filtering so the translated provider token param survives.
# This is the normal (non-retry) path: use_endpoint stays False so
# we use the fast, offline models.dev value (the purpose here is a
# finite budget, mainly for the HuggingFace whole-window
# reservation). The retry path in _update_output_window opts into
# the live endpoint for an accurate cap.
input_chars = len(ctx.prompt.to_string()) if ctx.prompt else None
resolve_output_token_limit(
overrides,
provider=overrides.get("model_provider") or "openai",
role=self.model_role,
input_chars=input_chars,
)
# --- Provider field filtering ---
active_provider = overrides.get("model_provider") or "openai"
provider_allowed = get_provider_allowed_fields(active_provider)
allowed = provider_allowed | {"model", "model_provider"}
overrides = {k: v for k, v in overrides.items() if k in allowed}
self.logger.debug(
f"After provider field filtering ({active_provider = }):\n{mask_sensitive(overrides) = }"
)
return cast(RunnableConfig, {"configurable": overrides})
def _add_cache_control(
self, prompt: PromptValue, config: RunnableConfig
) -> PromptValue:
"""Add Anthropic cache_control to the system message if applicable.
Only Anthropic supports ``cache_control: ephemeral`` on system blocks,
and it requires the ``model_provider`` to be known (resolved at
invoke time, not at prompt creation). Other providers ignore it, so
we only set it for ``anthropic``.
Anthropic prompt caching needs ``cache_control`` *inside* the
structured system content block. langchain-anthropic serialises a
plain-string ``SystemMessage`` as the bare ``system`` field and drops
``additional_kwargs``, so tagging there is a silent no-op --- the
string content must be wrapped in a text block carrying the flag.
"""
provider = config.get("configurable", {}).get("model_provider")
if provider != "anthropic":
return prompt
# PromptValue -> messages -> move cache_control into the first
# SystemMessage's structured content block
try:
messages = prompt.to_messages()
if messages and messages[0].type == "system":
system = messages[0]
content = system.content
if isinstance(content, str):
system.content = [
{
"type": "text",
"text": content,
"cache_control": {"type": "ephemeral"},
}
]
elif isinstance(content, list):
for i in range(len(content) - 1, -1, -1):
block = content[i]
if isinstance(block, str):
content[i] = {
"type": "text",
"text": block,
"cache_control": {"type": "ephemeral"},
}
break
if isinstance(block, dict) and block.get("type") == "text":
content[i] = {
**block,
"cache_control": {"type": "ephemeral"},
}
break
else:
return prompt
system.content = content
else:
return prompt
self.logger.debug("Added cache_control for Anthropic call")
# Rebuild PromptValue from modified messages
from langchain_core.prompt_values import ChatPromptValue
return ChatPromptValue(messages=messages)
except Exception as exc: # noqa: BLE001
self.logger.debug(f"Failed to add cache_control: {exc}")
return prompt
async def _invoke_llm(
self, ctx: LLMNodeContext[TOutput]
) -> AIMessage | dict[str, Any]:
"""Async invoke LLM with optional structured output + fallback.
When an output schema exists, the model is first invoked through
``with_structured_output``, unless this ``(provider, model, base_url)``
was already seen refusing the parameter (process-local cache), in which
case the plain invoke is used directly. Any structured-only failure
--- a schema validation/JSON parse error, or the provider rejecting
the ``response_format`` parameter --- latches the call to a plain
invoke for the rest of this invocation; the prompt already contains the
JSON schema as text instructions and :meth:`_process_output` parses it
tolerantly. Only a stable capability rejection is remembered in the
cache; parse failures are per-response and are never cached.
Both strategies share one :meth:`_invoke_with_retries` loop so
context-overflow / truncation / empty retries and infra-error
classification apply uniformly. Blank responses are the one exception:
they stay on the structured path so its existing empty-retry budget is
used before degrading.
"""
assert ctx.prompt is not None
assert ctx.config is not None
prompt = self._add_cache_control(ctx.prompt, ctx.config)
config = ctx.config
output_schema = ctx.output_schema
self.logger.debug(f"{prompt = }")
inst = self._llm_entry.instance
configurable = config.get("configurable") or {}
provider = configurable.get("model_provider") or "openai"
model = configurable.get("model") or self._llm_entry.model_name
base_url = configurable.get("base_url")
if not output_schema:
output = await self._invoke_with_retries(inst.ainvoke, prompt, config)
self.logger.debug(f"{output = }")
return output
if structured_output_known_unsupported(provider, model, base_url):
self.logger.debug(
f"Structured output known unsupported for {provider}:{model}; "
"using the plain invoke"
)
output = await self._invoke_with_retries(inst.ainvoke, prompt, config)
self.logger.debug(f"{output = }")
return output
structured = inst.with_structured_output(
output_schema, method="json_schema", include_raw=True
)
use_structured = True
async def invoke(prompt: PromptValue, config: RunnableConfig):
nonlocal use_structured
if use_structured:
try:
return await structured.ainvoke(prompt, config=config)
except Exception as exc:
# Blank responses keep their structured retry budget: let
# the retry loop see the error and apply the empty-output
# backoff instead of switching strategy.
if is_empty_structured_parse_error(exc):
raise
if not is_structured_output_failure(exc):
raise
if is_structured_capability_rejection(exc):
mark_structured_output_unsupported(provider, model, base_url)
self.logger.warning(
f"Structured output unusable ({type(exc).__name__}: "
f"{exc}); falling back to the plain invoke for this call"
)
use_structured = False
return await inst.ainvoke(prompt, config=config)
output = await self._invoke_with_retries(invoke, prompt, config)
self.logger.debug(f"{output = }")
return output
def _update_output_window(
self,
config: RunnableConfig,
direction: Literal["shrink", "grow"],
prompt: PromptValue | None = None,
) -> bool:
"""Resize the reserved output window.
Updates ``config["configurable"]`` in place. ``"shrink"`` halves
the window (context-overflow retry); ``"grow"`` extends it
(truncation retry) in two bounded phases: a linear step-up
(:data:`TRUNCATION_LINEAR_STEP` while below
:data:`TRUNCATION_LINEAR_CAP`) for models cut off just short of
finishing, then fixed large steps (:data:`TRUNCATION_PHASE2_STEP`)
climbing toward the output ceiling (:meth:`_jump_output_target`).
No exponential doubling: inference engines allocate resources
against ``max_tokens``, so doubling spikes KV-cache usage,
overshoots the real need, and overpays hallucination loops. Both
grow paths re-clamp to the model's catalog output limit and total
budget via ``resolve_output_token_limit`` (endpoint-first context
on these retries).
:param config: The per-invoke RunnableConfig to update in place.
:param direction: ``"shrink"`` or ``"grow"``.
:returns: True if the window actually changed, False if it was
already at a bound (no point retrying).
"""
overrides = config["configurable"]
# Native providers (``mistral:``, ``anthropic:``, ``deepseek:``,
# ...) carry no ``base_url`` in the configurable dict -- the
# provider resolves its own default endpoint internally. Resolve
# it by materialising the concrete model (already built on every
# invoke, no network) and reading its resolved endpoint attribute,
# so the endpoint-aware retry probes target the right server.
# Storing it back in *overrides* is safe: every provider that
# exposes one of the endpoint attributes aliases ``base_url``, so
# the retried invoke accepts it.
if not overrides.get("base_url"):
resolved_base_url = resolve_langchain_endpoint(
self._llm_entry.instance, config
)
if resolved_base_url:
overrides["base_url"] = resolved_base_url
provider = overrides.get("model_provider") or "openai"
token_param = get_token_limit_param(provider)
current = int(overrides.get(token_param, DEFAULT_MAX_OUTPUT_TOKENS))
if direction == "shrink":
target = max(MIN_OUTPUT_TOKENS, current // 2)
phase = "-"
else:
# Two-phase grow (see docstring). Clamp to the output ceiling
# so the ladder never exceeds what the endpoint (or the fixed
# fallback) allows; the resolver below re-applies the catalog
# output / total-budget clamps.
if current < TRUNCATION_LINEAR_CAP:
step = TRUNCATION_LINEAR_STEP
phase = "linear"
else:
step = TRUNCATION_PHASE2_STEP
phase = "large-step"
target = min(current + step, self._jump_output_target(overrides, prompt))
# Set the provider token param directly (rather than the generic
# key) so the resolver's "explicit value wins" precedence does not
# pick up a stale explicit value; resolve then re-applies the
# catalog output / total-budget clamps to *target*. Shrink and
# grow are retries, so use the live endpoint for an accurate
# context cap (falling back to models.dev), unlike the normal path.
overrides[token_param] = target
input_chars = len(prompt.to_string()) if prompt else None
resolve_output_token_limit(
overrides,
provider=provider,
role=self.model_role,
input_chars=input_chars,
use_endpoint=True,
)
new_value = int(overrides[token_param])
self.logger.warning(
"Output window %s (%s): %d -> %d (%s)",
direction,
phase,
current,
new_value,
provider,
)
return new_value != current
def _jump_output_target(
self, overrides: dict[str, Any], prompt: PromptValue | None = None
) -> int:
"""Return the output-window ceiling for truncation retries.
The ceiling is :data:`MAX_OUTPUT_TOKENS_CEILING` by default, but
for an OpenAI-compatible custom endpoint that advertises a larger
``max_model_len`` we raise it to the endpoint's remaining context
(``max_model_len - estimated input``) so a long reasoning trace is
not capped by the fixed fallback. The grow ladder in
:meth:`_update_output_window` clamps to this ceiling; the live
endpoint is trusted as the source of truth for the retry ceiling,
so no models.dev ``output`` clamp is applied below it -- the
resolver re-applies that clamp to the *target*, not the ceiling.
:param overrides: The merged ``configurable`` dict.
:returns: The ceiling for the grow ladder.
"""
target = MAX_OUTPUT_TOKENS_CEILING
provider = overrides.get("model_provider") or "openai"
limits = probe_endpoint_model_limits(
provider,
overrides.get("model", ""),
overrides.get("base_url"),
overrides.get("api_key"),
)
if limits and limits.context:
input_chars = len(prompt.to_string()) if prompt else None
if input_chars is not None:
headroom = limits.context - estimate_input_tokens(input_chars)
target = max(target, headroom)
self.logger.debug(
f"Output ceiling {target = } "
f"(endpoint_context = {limits.context if limits else None}, "
f"{provider = })"
)
return target
async def _invoke_with_retries(
self,
invoke: Callable[..., Awaitable[Any]],
prompt: PromptValue,
config: RunnableConfig,
) -> AIMessage | dict[str, Any]:
"""Invoke an LLM with adaptive retries on length-related failures.
Four retry behaviours, all bounded:
* ``context_overflow`` errors (request rejected because input plus
the reserved output exceeds the window) retry up to
:data:`MAX_CONTEXT_OVERFLOW_RETRIES` times, shrinking the
output window each time.
* Truncation retries by growing the output window in two bounded
phases -- a linear step-up, then fixed large steps up to the
output ceiling (see :meth:`_update_output_window`) -- capped at
:data:`MAX_TRUNCATION_RETRIES` attempts. Truncation is detected
two ways: a successful call that was cut off (``finish_reason ==
"length"``), or a raised ``length_truncation`` exception (some
SDKs, e.g. the OpenAI streaming / structured-output path, raise
instead of returning a truncated message).
* Empty responses (some providers, notably HuggingFace,
intermittently return a blank successful message) retry up to
:data:`MAX_EMPTY_OUTPUT_RETRIES` times with a plain re-invoke.
* Transient timeouts (e.g. a stream that stalls after the response
starts) retry up to :data:`MAX_TRANSIENT_RETRIES` times with an
exponential backoff. The provider SDK's own retries do not cover
a post-response stall, so this is the only layer that can.
All other failures (rate limits, auth, model-not-found, ...) are
re-raised immediately. Retrying stops early if resizing the
window makes no progress (already at a bound).
:param invoke: Async callable ``(prompt, config) -> output``.
:param prompt: The prompt to invoke.
:param config: Per-invoke RunnableConfig (mutated between attempts).
:returns: The (non-truncated) LLM output.
"""
overflow_retries = 0
truncation_retries = 0
empty_retries = 0
transient_retries = 0
# Attempts that return a usable-but-useless response (truncated or
# blank) are retried, but the provider still bills them. Their usage
# is accumulated here and folded into the surviving response so the
# single ``_extract_usage`` call downstream reports all of it.
discarded_usage: dict[str, int] = {}
def _fold_discarded(output: AIMessage | dict[str, Any]) -> None:
"""Record the usage of an attempt that is about to be retried."""
target = output if isinstance(output, AIMessage) else output.get("raw")
if not isinstance(target, AIMessage):
return
usage = self._extract_usage(LLMNodeContext(output=target))
if usage is None:
return
for field in (
"input_tokens",
"output_tokens",
"total_tokens",
"reasoning_tokens",
"cached_tokens",
):
discarded_usage[field] = discarded_usage.get(field, 0) + (
getattr(usage, field, 0) or 0
)
def _merge_discarded(
output: AIMessage | dict[str, Any],
) -> AIMessage | dict[str, Any]:
"""Fold retried-attempt usage into the surviving response."""
if not discarded_usage:
return output
target = output if isinstance(output, AIMessage) else output.get("raw")
if not isinstance(target, AIMessage):
return output
usage = self._extract_usage(LLMNodeContext(output=target))
own_input = usage.input_tokens if usage else 0
own_output = usage.output_tokens if usage else 0
own_total = usage.total_tokens if usage else 0
own_cached = usage.cached_tokens if usage else 0
own_reasoning = usage.reasoning_tokens if usage else 0
usage_metadata: dict[str, Any] = {
"input_tokens": own_input + discarded_usage.get("input_tokens", 0),
"output_tokens": own_output + discarded_usage.get("output_tokens", 0),
"total_tokens": own_total + discarded_usage.get("total_tokens", 0),
}
cached_tokens = own_cached + discarded_usage.get("cached_tokens", 0)
if cached_tokens:
usage_metadata["input_token_details"] = {"cache_read": cached_tokens}
reasoning_tokens = own_reasoning + discarded_usage.get(
"reasoning_tokens", 0
)
if reasoning_tokens:
usage_metadata["output_token_details"] = {"reasoning": reasoning_tokens}
merged = target.model_copy(update={"usage_metadata": usage_metadata})
if isinstance(output, dict):
return {**output, "raw": merged}
return merged
while True:
try:
output = await invoke(prompt, config=config)
except Exception as exc:
self.logger.error(f"{exc = }")
category = classify_llm_invocation_error(exc)
if (
category is LLMInvocationErrorCategory.CONTEXT_OVERFLOW
and overflow_retries < MAX_CONTEXT_OVERFLOW_RETRIES
):
overflow_retries += 1
if not self._update_output_window(config, "shrink", prompt):
self.logger.warning(
"Context overflow but output window cannot shrink further"
)
raise
self.logger.warning(
"Context overflow, retrying with smaller output window (%d/%d)",
overflow_retries,
MAX_CONTEXT_OVERFLOW_RETRIES,
)
self._emit_retry(
reason="context window exceeded",
attempt=overflow_retries,
max_attempts=MAX_CONTEXT_OVERFLOW_RETRIES,
action="shrink",
max_tokens=self._current_output_window(config),
)
continue
if category is LLMInvocationErrorCategory.LENGTH_TRUNCATION:
grow, message = self._next_truncation_retry(
config, truncation_retries, prompt
)
if grow:
truncation_retries += 1
self.logger.warning(message)
self._emit_retry(
reason="output truncated",
attempt=truncation_retries,
max_attempts=MAX_TRUNCATION_RETRIES,
action="grow",
max_tokens=self._current_output_window(config),
)
continue
raise
# A blank structured response makes the parser raise instead
# of returning an empty message, so route it into the same
# empty-retry budget a plain blank response gets.
if (
is_empty_structured_parse_error(exc)
and empty_retries < MAX_EMPTY_OUTPUT_RETRIES
):
empty_retries += 1
self.logger.warning(
"Empty structured LLM output, retrying (%d/%d)",
empty_retries,
MAX_EMPTY_OUTPUT_RETRIES,
)
self._emit_retry(
reason="empty structured response",
attempt=empty_retries,
max_attempts=MAX_EMPTY_OUTPUT_RETRIES,
)
continue
if (
category is LLMInvocationErrorCategory.TIMEOUT
and transient_retries < MAX_TRANSIENT_RETRIES
):
transient_retries += 1
delay = TRANSIENT_RETRY_BACKOFF_S * (2 ** (transient_retries - 1))
self.logger.warning(
"Transient LLM timeout, retrying in %.1fs (%d/%d): %s",
delay,
transient_retries,
MAX_TRANSIENT_RETRIES,
exc,
)
self._emit_retry(
reason="timed out",
attempt=transient_retries,
max_attempts=MAX_TRANSIENT_RETRIES,
)
await asyncio.sleep(delay)
continue
raise
if is_output_truncated(output):
grow, message = self._next_truncation_retry(
config, truncation_retries, prompt
)
if grow:
truncation_retries += 1
self.logger.warning(message)
self._emit_retry(
reason="output truncated",
attempt=truncation_retries,
max_attempts=MAX_TRUNCATION_RETRIES,
action="grow",
max_tokens=self._current_output_window(config),
)
_fold_discarded(output)
continue
return _merge_discarded(output)
if is_output_empty(output):
if empty_retries < MAX_EMPTY_OUTPUT_RETRIES:
empty_retries += 1
self.logger.warning(
"Empty LLM output, retrying (%d/%d)",
empty_retries,
MAX_EMPTY_OUTPUT_RETRIES,
)
self._emit_retry(
reason="empty response",
attempt=empty_retries,
max_attempts=MAX_EMPTY_OUTPUT_RETRIES,
)
_fold_discarded(output)
continue
self.logger.warning(
"Empty LLM output after %d retries; returning as-is",
MAX_EMPTY_OUTPUT_RETRIES,
)
return _merge_discarded(output)
return _merge_discarded(output)
def _next_truncation_retry(
self,
config: RunnableConfig,
truncation_retries: int,
prompt: PromptValue | None = None,
) -> tuple[bool, str]:
"""Return whether to retry a truncated invoke, and the log message.
Grows the reserved output window on each truncation -- linear
steps while below :data:`TRUNCATION_LINEAR_CAP`, then
:data:`TRUNCATION_PHASE2_STEP` steps up to the output ceiling
(:meth:`_jump_output_target`) -- capped at
:data:`MAX_TRUNCATION_RETRIES` attempts. Stops retrying once the
window cannot grow further (already at the ceiling) or the retry
budget is spent.
:param config: The per-invoke RunnableConfig to update in place.
:param truncation_retries: Number of truncation retries so far.
:returns: ``(grow, message)``; ``grow`` is False when the window
cannot grow further or the budget is exhausted.
"""
if truncation_retries >= MAX_TRUNCATION_RETRIES:
self.logger.warning(
"Output truncated but truncation retry budget exhausted"
)
return False, ""
if self._update_output_window(config, "grow", prompt):
return (
True,
(
"Output truncated, retrying with larger output window "
f"({truncation_retries + 1}/{MAX_TRUNCATION_RETRIES})"
),
)
self.logger.warning("Output truncated but output window cannot grow further")
return False, ""
def _current_output_window(self, config: RunnableConfig) -> int | None:
"""Return the output-token window currently set in *config*, or None."""
overrides = config["configurable"]
provider = overrides.get("model_provider") or "openai"
value = overrides.get(get_token_limit_param(provider))
return int(value) if isinstance(value, int) else None
def _emit_retry(
self,
*,
reason: str,
attempt: int,
max_attempts: int,
action: str = "",
max_tokens: int | None = None,
) -> None:
"""Report an LLM invoke retry to the live UI and the inspect tab.
Emits a ``progress`` event (the live line, via the shared
``_emit_progress`` helper) and an ``inspect`` event (a permanent
entry in the inspect tab) carrying the structured retry details.
Guarded because ``_invoke_with_retries`` is called directly by unit
tests outside a LangGraph run, where no stream writer exists; the
caller has already logged the retry.
:param reason: Short human-readable cause (e.g. ``"timed out"``).
:param attempt: 1-based retry number.
:param max_attempts: Total retry budget for this kind.
:param action: Output-window action (``"shrink"``/``"grow"``) or "".
:param max_tokens: The new output window, when it changed.
"""
summary = f"{self.label} (retry {attempt}/{max_attempts}: {reason})"
details: dict[str, Any] = {
"attempt": attempt,
"max": max_attempts,
"reason": reason,
}
if action:
details["action"] = action
if max_tokens is not None:
details["max_tokens"] = max_tokens
try:
self._emit_progress(heading=summary)
event = NodeStreamEvent(
type="inspect",
node=self.label,
data=NodeStreamData(heading="Retry", summary=summary, details=details),
)
self.write_custom_stream(event.model_dump())
except RuntimeError:
# No active run (direct unit-test call): the retry is still logged.
pass
def _process_output(self, ctx: LLMNodeContext[TOutput]) -> Any:
"""Common output processing with error handling.
NOTE: structured output is best-effort. A model can return a valid
JSON object that is not an instance of the schema (e.g. it echoes the
schema definition back, or returns only defaults); the parser then
yields an all-default instance without any ``parsing_error``. An
empty result here is therefore a possible failure mode, not a normal
"the model had nothing to say" response. The prompt (example instance
+ directive) reduces the odds; if empty results recur for a model,
revisit the prompt/model rather than expecting a loud invocation
error.
"""
output = ctx.output
result: TOutput | AIMessage | None = None
schema = ctx.output_schema
if schema:
# but answer is returned as message instead of json/dict
if isinstance(output, AIMessage):
result = self._parse_structured_message(output, ctx)
else:
assert isinstance(output, dict)
if output["parsing_error"]:
self.logger.warning(
f"LLM parsing error, using fallback: {output['parsing_error']}"
)
result = self._parse_structured_message(output["raw"], ctx)
else:
result = output["parsed"]
if isinstance(result, dict):
result = schema(**result)
else:
if not isinstance(result, schema):
self.logger.critical(
f"Unexpected output type: {type(result)}"
)
result = self._get_default_error_result(ctx)
self.logger.debug(f"Processed output: {result}")
else:
assert isinstance(output, AIMessage)
result = output
self.logger.debug(
f"No output schema. Returning unprocessed output: {result}"
)
if _is_empty_result(result, ctx.output_schema):
self.logger.warning(
f"Empty LLM output from {self.label}: nothing usable was "
f"produced (all-default structured result or blank message)"
)
# Structured nodes define a typed fail-closed default (e.g. the
# router's ``task`` route, the evaluator's ``need_replan``), so
# surface it instead of the empty instance. Non-structured
# nodes keep the raw output: their ``_get_default_error_result``
# is not necessarily an ``AIMessage`` (the guard returns a bare
# ``str``), and user-facing ones supply their own fallback text.
if ctx.output_schema is not None:
result = self._get_default_error_result(ctx)
return result
def _parse_structured_message(
self, message: AIMessage, ctx: LLMNodeContext[TOutput]
) -> Any:
"""Parse a structured-output message, degrading to the typed default.
``parse_output_with_thought`` can raise when the model returned blank
or unrecoverable text (for example a blank message that ``repair_json``
cannot turn into JSON). Structured nodes define a typed fail-closed
default, so use it rather than letting the parser error escape the
node and abort the run.
:param message: The raw model message to parse.
:param ctx: Per-run node context (supplies the resolved schema).
:returns: A schema instance, or ``_get_default_error_result()`` on
parse/validation failure.
"""
assert ctx.output_schema is not None
try:
result, _ = parse_output_with_thought(message, ctx.output_schema)
except (OutputParserException, ValidationError, ValueError, TypeError) as exc:
self.logger.warning(
f"Structured output could not be parsed; using default result: {exc}"
)
return self._get_default_error_result(ctx)
if isinstance(result, dict):
try:
return ctx.output_schema(**result)
except (ValidationError, ValueError, TypeError) as exc:
self.logger.warning(
"Parsed output failed schema validation; using default "
f"result: {exc}"
)
return self._get_default_error_result(ctx)
return result
def _invoke_prompt(
self,
prompt_template: ChatPromptTemplate,
variables: Any | dict[str, Any],
ctx: LLMNodeContext[TOutput],
) -> PromptValue:
"""Format prompt with state-specific parameters."""
prompt = prompt_template.invoke(variables)
self.logger.debug(f"{prompt =}")
return prompt
def _format_output_schema_prompt(self, schema: type[BaseModel]) -> str:
"""Return the ``Output schema (strict)`` prompt block.
The raw JSON Schema (``title``/``type``/``properties``) invites
models to echo the schema definition back instead of producing an
instance (the observed failure mode), so the prompt shows a
sanitized schema (top-level ``title``/``description`` dropped), an
explicit directive, and a generated example instance.
:param schema: The run's resolved output schema (static or dynamic).
:returns: Prompt text describing the required JSON output
"""
schema_json = convert_to_json_schema(schema)
sanitized = {
key: value
for key, value in schema_json.items()
if key not in ("title", "description")
}
example = _schema_to_example(schema_json, schema_json)
return dedent(
f"""
## Output schema (strict)
Respond in JSON following this schema:
{escape_prompt_braces(json.dumps(sanitized))}
The response must be a raw, valid JSON object like this example
(replace the placeholder values with real content):
{escape_prompt_braces(json.dumps(example))}
Do not output the schema definition itself, or notes/comments, or
the 'title'/'type'/'properties' keys.
"""
)
def _effective_model_name(self) -> str:
"""Return the model string this node will invoke.
Mirrors the model-selection layers of :meth:`_build_invoke_config`
(frozen node default > per-run role override > graph/env default) so a
node can report the model that will actually produce its output without
building the full invoke config. Reads the LangGraph runtime context
when available and tolerates its absence (unit tests). Never reads
credentials.
:returns: The effective model string (provider-prefixed when known), or
``""`` when no model is configured.
"""
model = str(self.model_defaults.get("model", "") or "")
if not model:
try:
overrides = model_overrides_from_context(get_runtime().context)
except Exception: # noqa: BLE001 - no run context (e.g. unit tests)
overrides = {}
model = str(overrides.get(self.model_role, {}).get("model", "") or "")
return model or str(getattr(self._llm_entry, "model_name", "") or "")
def _get_system_prompt(
self, state: TState, ctx: LLMNodeContext[TOutput]
) -> str | list:
"""Load system prompt from file, optionally adding memory and schema.
When memory is enabled, returns a list of ``("system", text)`` plus
the recent conversation as real message objects (in interleaved
order), so ``_create_prompt_template`` can place them between the
system and human prompts. Otherwise returns the system prompt text
as a plain string.
:param state: Graph state.
:param ctx: Per-run node context.
:returns: System prompt text, or a list of system text + history
message objects when memory is enabled.
"""
system_prompt = self._load_prompt_file(f"{self.prompt_prefix}_system")
if self.include_model_identity:
# Constant prefix: identical across nodes, so it stays in the
# cacheable part of the system prompt.
system_prompt = f"{_AGENT_IDENTITY}\n\n{system_prompt}"
if ctx.output_schema:
system_prompt += self._format_output_schema_prompt(ctx.output_schema)
if self.include_model_identity:
model = self._effective_model_name()
if model:
system_prompt += f"\n\nModel in use: {model}"
if self.memory:
memory_addition = self._get_memory_addition(state)
system_prompt += memory_addition
system_messages: list[Any] = [("system", system_prompt)]
system_messages += self._get_recent_memory_messages(state)
self.logger.debug(f"{system_messages =}")
return system_messages
self.logger.debug(f"{system_prompt =}")
return system_prompt
def _get_recent_memory_messages(self, state: TState) -> list[BaseMessage]:
"""Return the recent verbatim conversation for the prompt.
The window is bounded by :data:`DEFAULT_RECENT_MEMORY_CHARS`, so a
long run cannot inject its entire history. ``context_summary``
carries the older history. Only user-visible ``messages`` are injected;
``messages_internal`` is summarised separately and never placed in a
prompt verbatim.
:param state: Graph state.
:returns: Ordered recent human/ai messages within the budget.
"""
messages = list(getattr(state, "messages", None) or [])
return get_recent_messages(
messages,
DEFAULT_RECENT_MEMORY_CHARS,
keep_at_least=2,
)
def _get_human_prompt(self, state: TState, ctx: LLMNodeContext[TOutput]) -> str:
"""Load human prompt from file."""
human_prompt = self._load_prompt_file(f"{self.prompt_prefix}_user")
self.logger.debug(f"{human_prompt =}")
return human_prompt
def _load_prompt_file(self, prompt_name: str) -> str:
"""Load a prompt file from the registry.
:param prompt_name: Prompt file name (without extension)
:returns: Prompt text content
"""
return load_prompt(
prompt_name=prompt_name,
prompt_registry_location=self.prompt_registry_location,
)
@staticmethod
def _optional_section(title: str, body: str, level: int = 2) -> str:
"""Render an optional prompt section, or ``""`` when *body* is empty.
``ChatPromptTemplate`` has no conditionals, so every placeholder always
renders; a bare ``Heading:\\n{value}`` line would leave an empty label
on the calls where the value does not apply. Nodes therefore compose a
conditional section whole -- heading included -- into one variable and
leave it out entirely when there is nothing to say (prompt
conventions, ``devdocs/system/prompt-conventions.md``). The user
prompt then carries a bare ``{..._block}`` placeholder.
:param title: Section heading text, without the ``#`` markers.
:param body: Section content; whitespace-only means "omit".
:param level: Markdown heading level (default ``2``).
:returns: The rendered section, or ``""`` when *body* is empty.
"""
text = str(body).strip()
if not text:
return ""
return f"{'#' * level} {title}\n\n{text}"
def _create_prompt_template(
self,
system_prompt: str | list[Any],
human_prompt: str,
ctx: LLMNodeContext[TOutput],
) -> ChatPromptTemplate:
"""Create ChatPromptTemplate with system and human messages.
*system_prompt* may be a plain string (memory disabled) or a list of
``("system", text)`` plus recent history message objects (memory
enabled); the human prompt is appended after it.
:param system_prompt: System prompt text or a system-side list
including recent history messages.
:param human_prompt: Human prompt text.
:param ctx: Per-run node context.
"""
system_messages = (
[("system", system_prompt)]
if isinstance(system_prompt, str)
else system_prompt
)
if len(system_messages) and len(human_prompt):
prompt_template = ChatPromptTemplate(
[*system_messages, ("human", human_prompt)]
)
elif len(system_messages) and not len(human_prompt):
prompt_template = ChatPromptTemplate(system_messages)
elif len(human_prompt) and not len(system_messages):
prompt_template = ChatPromptTemplate([("human", human_prompt)])
else:
raise PromptTemplateError(
"No prompts provided. Cannot create prompt template!"
)
self.logger.debug(f"{prompt_template =}")
return prompt_template
def _get_memory_addition(self, state: TState) -> str:
"""Hook for subclasses to append memory content into the system prompt.
Override this method to provide memory-specific content.
The default implementation returns an empty string.
"""
return add_memory_to_prompt(
context_summary=state.context_summary, # type: ignore
)