feat(streaming): add Executor.execute_stream() returning AsyncIterator[str] for token-by-token delivery #45

Merged
hurui200320 merged 1 commits from feature/streaming-execute-stream into master 2026-06-12 12:29:49 +00:00
8 changed files with 9314 additions and 12 deletions
+1
View File
@@ -19,6 +19,7 @@ and this project adheres to [Clever Semantic Versioning](https://www.w3.org/subm
### Added
- **Registry Error Hierarchy** (`cleveractors.registry.exceptions`): Typed exception hierarchy per Package Registry Standard §13.2. `RegistryError(CleverAgentsException)` base carries `message`, optional `details: dict`, and optional `original_reference: str`; `__str__` includes the reference when present. Nine leaf exceptions: `PackageNotFoundError` (404), `InvalidPackageIdError` (400), `InvalidPackageReferenceError` (400), `VersionNotFoundError` (404), `ValidationError` (400), `AuthenticationRequiredError` (401), `AccessDeniedError` (403), `ConflictError` (409), and `RegistryNetworkError` (5xx / connection / timeout) which additionally carries `status_code` and `url`. `exception_for_status()` maps HTTP codes to typed exceptions; `_ERROR_TYPE_MAP` enables error-type parsing from structured JSON error bodies. 23 Behave BDD scenarios and 12 Robot Framework integration tests; exceptions.py achieves 100% coverage.
- **`Executor.execute_stream()` — token-by-token streaming delivery** (`cleveractors.runtime.Executor`, `cleveractors.agents.llm.LLMAgent`, `cleveractors.langgraph.nodes.Node`, `cleveractors.langgraph.pure_graph.PureLangGraph`, `cleveractors.runtime_dispatch`): `LLMAgent` gains `stream_message(message, context)` using `self.chat_model.astream(messages)`, yielding token chunks via the LangChain streaming API with `_safe_int()` / fallback chain for token counts from the final chunk's `usage_metadata`. `Node` gains `stream_agent(state)` which delegates to `stream_message()` for `LLMAgent` instances or falls back to `process_message()` for non-LLM agents. `PureLangGraph.execute_stream()` mirrors `execute()` but uses `_stream_from_node()`, which buffers tokens for intermediate AGENT nodes (only yielding from the terminal node) while running non-AGENT nodes with `ainvoke()`. `Executor` gains `last_result: ActorResult | None = None` (populated after stream exhaustion for billing) and `execute_stream(message)` dispatching to `_execute_llm_stream()` (LLM actors) or `_execute_graph_stream()` (graph actors). All existing execution limits (`timeout_ms`, `max_model_calls`, `max_tool_calls`) are enforced in the streaming path. `execute_stream()` raises `ConfigurationError` for unsupported actor types (`tool`, `multi_actor`). No new top-level package export required. (issue #16)
- **Structured `ExecutionError` fields + 5-limit enforcement in `PureLangGraph`** (`cleveractors.core.exceptions.ExecutionError`, `cleveractors.langgraph.pure_graph.PureLangGraph`): `ExecutionError` gains `kind` (categorical: `depth`, `model_calls`, `tool_calls`, `timeout`, `cost`) and `reason` (sub-code: `budget_exhausted` or `missing_pricing_entry`) fields, both defaulting to `""` for backward compatibility with existing `raise ExecutionError(msg)` call sites. `PureLangGraph` now accepts `limits` and `pricing` constructor arguments. When `limits["max_depth"]` is supplied, a depth breach raises `ExecutionError(kind="depth")` instead of silently returning the current message. `max_model_calls` is checked before each AGENT-type node; `max_tool_calls` before each TOOL-type node. `execute()` wraps `_execute_from_node()` in `asyncio.wait_for()` when `limits["timeout_ms"]` is set, mapping `asyncio.TimeoutError` to `ExecutionError(kind="timeout")`. After each LLM node, cost is computed from the supplied `pricing` table (rates are USD per million tokens per ADR-2029); breach raises `ExecutionError(kind="cost", reason="budget_exhausted")`; a missing provider or model entry raises `ExecutionError(kind="cost", reason="missing_pricing_entry")`. `runtime_dispatch._execute_graph()` now passes `executor.limits` and `executor.pricing` to `PureLangGraph`. `ExecutionError` is exported from `cleveractors.__init__` and `__all__`. (ADR-2029, issue #15)
- **`create_executor()` router-facing API** (`cleveractors.create_executor`): New module-level factory function that constructs an `Executor` wrapping `PureLangGraph` and `AgentFactory`. Accepts `config_dict` (validated actor configuration), `credentials` (per-provider credential dict for per-request injection), `limits` (execution budget), and `pricing` (per-model cost table). `executor.execute(message)` runs the actor graph and returns an `ActorResult(response, prompt_tokens, completion_tokens, nodes)`. All four execution paths are supported — `llm`, `graph`, `tool`, and `multi_actor`. Credentials are passed to `AgentFactory` and never injected into the stored `config_dict` (ADR-2026 AC8). Exported from `cleveractors.__init__` and `__all__` (ADR-2024, ADR-2026, ADR-2029).
- **`ActorResult` and `NodeUsage` types** (`cleveractors.result`): Canonical dataclasses for the router-facing result API, now defined in `cleveractors.result` (ADR-2027). `ActorResult` carries the response string, aggregated `prompt_tokens`/`completion_tokens`, a non-empty `nodes: list[NodeUsage]` breakdown for per-model billing, and an optional opaque `state` blob for stateless graph resumption (ADR-2026). Re-exported from `cleveractors.runtime` for backward compatibility.
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+300
View File
@@ -26,6 +26,7 @@ from __future__ import annotations
import contextvars
import logging
import threading
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any, ClassVar, Literal
if TYPE_CHECKING:
@@ -667,6 +668,305 @@ class LLMAgent(AgentWithMemory):
)
self._chat_model.temperature = saved_temperature
async def stream_message(
self,
message: str,
context: dict[str, Any] | None = None,
) -> AsyncGenerator[str, None]:
"""Stream tokens from the LLM using astream().
This is the streaming counterpart to :meth:`process_message`. It
builds the same LangChain message list (system + history + user) and
calls ``self.chat_model.astream(messages)`` instead of ``ainvoke()``,
yielding each token as a ``str`` as it arrives.
Token counts are captured from the **final chunk's** ``usage_metadata``
(primary) or ``response_metadata["token_usage"]`` (fallback) using the
same three-tier ``_safe_int()`` fallback chain as
:meth:`process_message`. After the generator is exhausted,
``self._last_token_usage`` and ``last_token_usage_var`` are set to the
captured counts.
If ``_temperature_override`` is present in ``context``, the configured
temperature is replaced for the duration of this call and restored
afterwards (spec §4.4.5).
After a successful stream, memory is updated if ``memory_enabled`` is
set in the agent config (spec §4.4.4).
Args:
message: The user's input message (plain string).
context: Same context dict accepted by :meth:`process_message`
(``conversation_history``, ``_temperature_override``, graph
state, etc.).
Yields:
Each token chunk's content as a ``str``.
Note:
``stream_message()`` resets ``_last_token_usage`` and
``last_token_usage_var`` to ``(0, 0)`` at entry, then sets
them to the captured token counts after the last chunk is yielded.
If the stream is abandoned before exhaustion (caller breaks out of
the ``async for``), the counts remain at ``(0, 0)`` callers must
exhaust the iterator to get accurate billing data.
"""
# Sentinel variables for billing-integrity (mirrors process_message):
# set to non-None only after the astream loop completes so that the
# except handler can distinguish pre-stream failures (no tokens
# consumed → reset to (0,0)) from post-stream failures (tokens
# consumed → preserve captured counts).
_captured_prompt: int | None = None
_captured_completion: int | None = None
# Temperature override — applied before template rendering so that
# any template that inspects the model temperature sees the override.
# Mirrors the process_message() block (spec §4.4.5).
saved_temperature: float | None = None
try:
# Reset at start so an abandoned stream never leaks stale counts from a
# previous successful call. Placed inside the try block (matching
# process_message()) so any pre-stream code added later cannot race
# against the (0, 0) state (m2 fix from review).
self._last_token_usage = (0, 0)
last_token_usage_var.set((0, 0))
# Apply temperature override if present in context (spec §4.4.5).
# NOTE: not thread-safe — see process_message() docstring.
if context and "_temperature_override" in context:
temperature_override = context["_temperature_override"]
if not isinstance(temperature_override, (int, float)):
raise ConfigurationError(
f"_temperature_override must be a number, "
f"got {type(temperature_override).__name__}"
)
current_temp = self.chat_model.temperature
if temperature_override != current_temp:
logger.debug(
"Agent %s: Applying temperature override %.2f (was %.2f)"
" for streaming",
self.name,
temperature_override,
current_temp,
)
saved_temperature = current_temp
self.chat_model.temperature = temperature_override
# Process template if specified (mirrors process_message logic)
if "template" in self.config:
template_name = self.config["template"]
template_vars: dict[str, Any] = {
"message": message,
"context": context or {},
**self.config.get("template_vars", {}),
}
processed_message = self.template_renderer.render(
template_name, template_vars
)
else:
processed_message = message
# Build the LangChain message list (same logic as process_message).
# Typed as list[Any] because the LangChain message types are loaded
# lazily; they are not available at module level for static annotation.
lc_messages: list[Any] = []
# Add system message
if self.system_message:
try:
template_context: dict[str, Any] = {
"context": context or {},
"message": message,
}
rendered_system_message = self.template_renderer.render_string(
self.system_message,
template_context,
source_description="system prompt",
)
except Exception as _sp_err: # pylint: disable=broad-exception-caught
# m1 fix: log the render failure at WARNING level so operators
# see the same signal from stream_message() as from
# process_message() for the same condition.
logger.warning(
"Agent %s: Failed to render system prompt: %s",
self.name,
type(_sp_err).__name__,
)
rendered_system_message = self.system_message
lc_messages.append(SystemMessage(content=rendered_system_message))
# Add conversation history from context (mirrors process_message)
history: list[dict[str, str]] | None = None
if context and "conversation_history" in context:
history = context["conversation_history"]
elif self.config.get("memory_enabled", False):
history = await self.get_memory("conversation_history", [])
if history:
for msg in history:
if msg.get("role") == "user":
lc_messages.append(HumanMessage(content=msg.get("content", "")))
elif msg.get("role") == "assistant":
lc_messages.append(AIMessage(content=msg.get("content", "")))
# Add current user message
lc_messages.append(HumanMessage(content=processed_message))
# Stream tokens — accumulate full response for memory update only
# when memory_enabled is True. Accumulating unconditionally would
# grow agent_response_parts to 100K+ entries for long responses
# with memory disabled, wasting ~12 MB of memory that is never
# consumed (m1 fix).
# Use a list accumulator and join at the end to avoid O(n²) string
# allocations from repeated += on immutable Python strings.
_memory_enabled: bool = self.config.get("memory_enabled", False)
last_chunk: Any = None
agent_response_parts: list[str] = []
async for chunk in self.chat_model.astream(lc_messages):
last_chunk = chunk
# Guard against metadata-only chunks where content is None:
# str(None) would yield the literal string "None" to the caller.
token = str(chunk.content) if chunk.content is not None else ""
if _memory_enabled:
agent_response_parts.append(token)
yield token
# Extract token counts from the final chunk.
# Three-tier fallback chain (AC3, mirrors process_message()):
# 1. usage_metadata (LangChain standard field)
# 2. response_metadata["token_usage"] (provider-specific)
# 3. 0 with warning
_prompt_tokens: int = 0
_completion_tokens: int = 0
if last_chunk is not None:
_usage_raw: object = getattr(last_chunk, "usage_metadata", None)
_usage: dict[str, Any] | None = (
_usage_raw if isinstance(_usage_raw, dict) else None
)
if _usage is not None and _usage:
_prompt_tokens = self._safe_int(
_usage.get("input_tokens"), "input_tokens"
)
_completion_tokens = self._safe_int(
_usage.get("output_tokens"), "output_tokens"
)
elif _usage is not None:
# usage_metadata present but empty
self._log_no_usage_metadata(_CAUSE_USAGE_METADATA_EMPTY)
elif (
hasattr(last_chunk, "response_metadata")
and last_chunk.response_metadata is not None
):
# Fallback to response_metadata["token_usage"] (tier 2).
# Mirrors process_message() lines for the same fallback.
_rm: object = last_chunk.response_metadata
if not isinstance(_rm, dict):
self._log_no_usage_metadata(_CAUSE_RESPONSE_METADATA_NOT_DICT)
else:
_token_usage: dict[str, Any] = _rm.get("token_usage", {})
if _token_usage:
_prompt_tokens = self._safe_int(
_token_usage.get("prompt_tokens"), "prompt_tokens"
)
_completion_tokens = self._safe_int(
_token_usage.get("completion_tokens"),
"completion_tokens",
)
else:
self._log_no_usage_metadata(
_CAUSE_RESPONSE_METADATA_NO_TOKEN_USAGE
)
else:
# Neither usage_metadata nor response_metadata available.
self._log_no_usage_metadata(_CAUSE_RESPONSE_METADATA_MISSING)
# Mark astream() as successfully completed so the except handler
# preserves these counts on any post-stream failure.
_captured_prompt = _prompt_tokens
_captured_completion = _completion_tokens
self._last_token_usage = (_captured_prompt, _captured_completion)
last_token_usage_var.set((_captured_prompt, _captured_completion))
# Update memory if enabled (spec §4.4.4).
# Mirrors process_message() memory block.
if self.config.get("memory_enabled", False):
agent_response: str = "".join(agent_response_parts)
await self.update_memory("last_message", message)
await self.update_memory("last_response", agent_response)
mem_history: list[dict[str, str]] = await self.get_memory(
"conversation_history", []
)
mem_history.append({"role": "user", "content": processed_message})
mem_history.append({"role": "assistant", "content": agent_response})
max_history: int = self.config.get("max_history", DEFAULT_MAX_HISTORY)
if len(mem_history) > max_history:
mem_history = mem_history[-max_history:]
await self.update_memory("conversation_history", mem_history)
except ConfigurationError:
# Reset token usage so a failed call never leaks counts from a
# previous successful call. ConfigurationError is re-raised
# without wrapping (mirrors process_message()).
self._last_token_usage = (0, 0)
last_token_usage_var.set((0, 0))
raise
except LangChainException as e:
# Mirrors the process_message() LangChainException handler.
# LangChainException is a subclass of Exception, so it must be
# caught before the broad `except Exception` arm to produce a
# distinct log message. Log analytics that filter on
# "LangChain streaming error" will correctly identify LangChain-
# specific failures in the streaming path (symmetric with the
# "LangChain error" message in process_message()).
if _captured_prompt is not None and _captured_completion is not None:
self._last_token_usage = (_captured_prompt, _captured_completion)
last_token_usage_var.set((_captured_prompt, _captured_completion))
else:
self._last_token_usage = (0, 0)
last_token_usage_var.set((0, 0))
logger.error(
"LLM agent %s LangChain streaming error: %s",
self.name,
type(e).__name__,
)
logger.debug(
"Raw LangChain streaming exception (sanitized): type=%s",
type(e).__name__,
)
raise ExecutionError("LLM streaming failed") from None
except Exception as e: # pylint: disable=broad-exception-caught
# Billing integrity: if astream() already completed
# (_captured_prompt is not None), the LLM provider has already
# billed for those tokens. Preserve the captured counts so the
# router receives accurate billing data even when a post-stream
# step (e.g. update_memory()) raises. Only reset to (0, 0) when
# the exception occurred before the capture point (astream() itself
# failed), in which case no tokens were consumed.
if _captured_prompt is not None and _captured_completion is not None:
self._last_token_usage = (_captured_prompt, _captured_completion)
last_token_usage_var.set((_captured_prompt, _captured_completion))
else:
self._last_token_usage = (0, 0)
last_token_usage_var.set((0, 0))
logger.error(
"LLM agent %s streaming failed: %s", self.name, type(e).__name__
)
raise ExecutionError("LLM streaming failed") from None
finally:
# Restore original temperature if it was overridden.
# Mirrors process_message() finally block (spec §4.4.5).
if saved_temperature is not None and self._chat_model is not None:
logger.debug(
"Agent %s: Restoring temperature to %.2f after streaming",
self.name,
saved_temperature,
)
self._chat_model.temperature = saved_temperature
def _log_no_usage_metadata(
self,
cause: Literal[
+148
View File
@@ -6,6 +6,7 @@ from __future__ import annotations
import asyncio
import logging
from collections.abc import AsyncGenerator
from copy import deepcopy
from dataclasses import dataclass, field
from enum import Enum
@@ -136,6 +137,12 @@ class Node: # pylint: disable=too-many-instance-attributes
self.last_execution_time: Optional[float] = None
self.last_error: Optional[Exception] = None
# Per-stream token usage snapshot populated by _stream_agent() after
# the async generator is exhausted. Callers (_stream_from_node() in
# pure_graph.py) read this to append to _node_usages. None means
# _stream_agent() has not been called yet or it raised an exception.
self._last_stream_usage: dict[str, Any] | None = None
def _prepare_conversation_history(
self, messages: List[dict[str, Any]]
) -> tuple[List[dict[str, Any]], bool]:
@@ -450,6 +457,147 @@ class Node: # pylint: disable=too-many-instance-attributes
return state_updates
async def _stream_agent(self, state: GraphState) -> AsyncGenerator[str, None]:
"""Stream tokens from an agent node using astream().
This is the streaming counterpart to :meth:`_execute_agent`. It
builds the same context and calls ``agent.stream_message()`` for
:class:`~cleveractors.agents.llm.LLMAgent` instances, yielding each
token chunk. For non-LLM agents (``ToolAgent``, etc.) that do not
support streaming, it falls back to ``process_message()`` and yields
the complete response as a single token.
After the generator is exhausted, per-node token usage is stored in
``self._last_stream_usage`` (a dict) or ``None`` on error, for the
caller to read and append to ``_node_usages``.
.. note::
Unlike :meth:`_execute_agent`, this method does **not** propagate
context changes made by the agent back to ``state.metadata``.
This is an intentional asymmetry with the non-streaming path;
propagation is deferred to a follow-up PR.
Args:
state: Current :class:`~cleveractors.langgraph.state.GraphState`.
Yields:
Each token as a ``str``.
Raises:
ValueError: If no agent is configured on this node or the agent
is not found in ``self.agents``.
"""
if not self.config.agent:
raise ValueError(f"Agent node {self.name} has no agent specified")
agent = self.agents.get(self.config.agent)
if not agent:
raise ValueError(f"Agent {self.config.agent} not found")
# Build agent_input (mirrors _execute_agent logic)
if state.messages:
if isinstance(agent, ToolAgent):
current_msg = state.metadata.get("current_message", "")
if current_msg:
agent_input = str(current_msg)
else:
agent_input = state.messages[-1].get("content", "")
else:
current_msg = state.metadata.get("current_message")
if current_msg is not None:
agent_input = str(current_msg)
else:
last_user_message = None
for msg in reversed(state.messages):
if msg.get("role") == "user":
last_user_message = msg.get("content", "")
break
agent_input = last_user_message or state.messages[-1].get(
"content", ""
)
else:
agent_input = ""
trimmed_history, history_truncated = self._prepare_conversation_history(
state.messages
)
graph_state_dict = state.to_dict()
graph_state_dict["messages"] = trimmed_history
context: dict[str, Any] = {
"graph_state": graph_state_dict,
"conversation_history": trimmed_history,
"full_context": True,
}
if history_truncated:
context["_history_truncated"] = True
context["_history_original_length"] = len(state.messages)
if state.metadata:
context.update(state.metadata)
nested_context = context.get("context")
if isinstance(nested_context, dict):
for key, value in nested_context.items():
context.setdefault(key, value)
# Reset per-stream usage snapshot for this invocation.
# Note: stream_message() resets last_token_usage_var internally at the
# start of its try block, so a pre-call reset here would be redundant
# (unlike _execute_agent, where process_message() does NOT reset the
# ContextVar internally). The double-reset is therefore removed
# (n3 fix from review).
self._last_stream_usage = None
try:
if isinstance(agent, LLMAgent):
# LLMAgent: use stream_message() for real token-by-token streaming
async for token in agent.stream_message(agent_input, context):
yield token
else:
# Non-LLM agent: fall back to process_message and yield as one token
response_str = await agent.process_message(agent_input, context)
agent_response = str(response_str)
yield agent_response
# Capture per-node token usage (mirrors _execute_agent success path).
# For LLMAgent: use ContextVar (authoritative, race-free for parallel).
# For non-LLM: use instance attribute.
_tok_from_var: tuple[int, int] = last_token_usage_var.get((0, 0))
_last_tok_inst: object = getattr(agent, "_last_token_usage", None)
if _tok_from_var != (0, 0):
_last_tok: object = _tok_from_var
elif not isinstance(agent, LLMAgent):
_last_tok = _last_tok_inst
else:
_last_tok = (0, 0)
if isinstance(_last_tok, tuple) and len(_last_tok) == 2:
self._last_stream_usage = {
"node_id": self.name,
"provider": str(getattr(agent, "provider", "unknown")),
"model": str(getattr(agent, "model", "unknown")),
"prompt_tokens": _safe_node_token_int(_last_tok[0]),
"completion_tokens": _safe_node_token_int(_last_tok[1]),
}
except Exception as e: # pylint: disable=broad-exception-caught
self.logger.error(
"Agent %s streaming failed: %s", agent.name, type(e).__name__
)
# _last_stream_usage is already None (set before the try block).
# Re-assigning it here would be redundant and inconsistent with
# _execute_agent() which does not re-assign _node_token_usage in
# its except block (Minor #7 fix).
# M1 fix: re-raise so the caller (_stream_from_node) can map the
# exception to the correct HTTP status code. Yielding the error
# as a token is a contract violation — the user would see the
# exception message as the assistant's answer and billing data
# would be populated against a node that produced no real tokens.
raise
async def _execute_function(self, state: GraphState) -> dict[str, Any]:
"""Execute a function node."""
if not self.config.function:
File diff suppressed because it is too large Load Diff
+116 -7
View File
@@ -13,13 +13,16 @@ into a clean, stable interface.
from __future__ import annotations
from collections.abc import AsyncGenerator
from typing import Any
from cleveractors.core.exceptions import ConfigurationError
from cleveractors.result import ActorResult, NodeUsage
from cleveractors.runtime_dispatch import (
_execute_graph,
_execute_graph_stream,
_execute_llm,
_execute_llm_stream,
_execute_multi_actor,
_execute_tool,
)
@@ -65,6 +68,22 @@ class Executor:
self.credentials = credentials
self.limits = limits
self.pricing = pricing
# Populated by execute_stream() after the async iterator is exhausted.
# Remains None until the stream completes (AC4, issue #16).
self.last_result: ActorResult | None = None
def _detect_actor_type(self) -> str:
"""Determine the actor type from the config dict.
Returns:
One of ``"graph"``, ``"llm"``, ``"tool"``, or ``"multi_actor"``.
"""
# CleverAgents v2.0: "routes" key indicates a graph even when "type" is missing
if "routes" in self.config:
return "graph"
return self.config.get(
"type", "multi_actor" if "actors" in self.config else "llm"
)
async def execute(
self,
@@ -87,13 +106,7 @@ class Executor:
An :class:`ActorResult` with the response, token usage, and
updated state (for graph actors).
"""
# CleverAgents v2.0: "routes" key indicates a graph even when "type" is missing
if "routes" in self.config:
actor_type = "graph"
else:
actor_type = self.config.get(
"type", "multi_actor" if "actors" in self.config else "llm"
)
actor_type = self._detect_actor_type()
# Delegate to module-level dispatch functions in runtime_dispatch.
# These functions implement AC2 (credential injection via AgentFactory),
@@ -111,6 +124,102 @@ class Executor:
else:
raise ConfigurationError(f"Cannot execute actor of type {actor_type!r}")
async def execute_stream(
self,
message: str,
messages: list[dict[str, Any]] | None = None,
state: dict[str, Any] | None = None,
) -> AsyncGenerator[str, None]:
"""Stream the actor response token-by-token.
Calls ``astream()`` on the underlying LLM or on the terminal agent
node of a graph actor, yielding each token as it arrives.
After the iterator is **fully exhausted**, :attr:`last_result` is
populated with an :class:`~cleveractors.result.ActorResult` containing
the concatenated response, aggregated token counts, and per-node usage
for billing. While the stream is in progress, :attr:`last_result`
remains ``None``.
On exception paths (e.g. ``ExecutionError`` for a limit breach or
``ConfigurationError`` for an invalid config), :attr:`last_result` is
also populated with a partial ``ActorResult`` (billing-integrity
guarantee) before the exception propagates to the caller. This
guarantee applies only when the dispatch function is reached (i.e. for
``"llm"`` and ``"graph"`` actor types). If the actor type is
unsupported (``"tool"``, ``"multi_actor"``), :attr:`last_result`
remains ``None`` because no LLM call was attempted.
.. note::
If the caller abandons the iterator before exhaustion (e.g. by
breaking out of the ``async for`` loop), :attr:`last_result`
remains ``None``. Callers must exhaust the iterator to obtain
complete billing data.
.. note::
When ``timeout_ms`` is set, the stream is collected under
``asyncio.wait_for``. If the timeout fires, any tokens already
generated by the LLM but not yet yielded are discarded and
``executor.last_result.response`` will be ``""`` (or contain only
the tokens yielded before the timeout in the graph path). This is
a known limitation: partial token recovery on timeout is not
supported.
Supported actor types:
- ``"llm"`` delegates to ``_execute_llm_stream()``
- ``"graph"`` delegates to ``_execute_graph_stream()``
Unsupported (raises :class:`~cleveractors.core.exceptions.ConfigurationError`):
- ``"tool"`` and ``"multi_actor"``
Args:
message: The user's input message (plain string).
messages: Full conversation history for multi-turn context.
state: Opaque graph state blob from a previous call (graph actors
only; ignored by LLM actors).
Yields:
Token strings as they arrive from the LLM.
Raises:
:class:`~cleveractors.core.exceptions.ConfigurationError`: For
unsupported actor types or invalid configuration.
:class:`~cleveractors.core.exceptions.ExecutionError`: On
execution failures or limit violations.
"""
# Reset last_result so it's None while the stream is in progress.
self.last_result = None
actor_type = self._detect_actor_type()
if actor_type == "llm":
# Explicitly close the inner generator in a finally block so that
# agent.cleanup() (which closes httpx.AsyncClient instances) is
# called promptly even when the caller abandons the iterator early
# (e.g. by breaking out of the async for loop). Without this,
# abandoned async generators are closed non-deterministically by
# the GC, which may not happen before the event loop closes.
_llm_gen = _execute_llm_stream(self, message, messages=messages)
try:
async for token in _llm_gen:
yield token
finally:
await _llm_gen.aclose()
elif actor_type == "graph":
_graph_gen = _execute_graph_stream(
self, message, state=state, messages=messages
)
try:
async for token in _graph_gen:
yield token
finally:
await _graph_gen.aclose()
else:
raise ConfigurationError(
f"Streaming is not supported for actor type {actor_type!r}. "
"Only 'llm' and 'graph' actor types support execute_stream()."
)
# ---------------------------------------------------------------------------
# Factory function
+789
View File
@@ -16,9 +16,11 @@ as its first argument so it can access ``self.config``, ``self.credentials``,
from __future__ import annotations
import asyncio
import copy
import logging
import re
from collections.abc import AsyncGenerator
from dataclasses import replace
from typing import TYPE_CHECKING, Any
@@ -627,3 +629,790 @@ async def _execute_multi_actor(
prompt_tokens=sum(n.prompt_tokens for n in prefixed_nodes),
completion_tokens=sum(n.completion_tokens for n in prefixed_nodes),
)
# -- streaming LLM ------------------------------------------------------------
async def _collect_llm_stream_tokens(
agent: Any,
message: str,
llm_context: dict[str, Any] | None,
) -> list[str]:
"""Collect all tokens from ``agent.stream_message()`` into a list.
Used by :func:`_execute_llm_stream` when ``timeout_ms`` is set: the entire
stream is run under ``asyncio.wait_for`` and the collected tokens are
yielded afterwards. This mirrors the graph path's
:meth:`~cleveractors.langgraph.pure_graph.PureLangGraph._collect_stream_tokens`
helper.
"""
tokens: list[str] = []
async for token in agent.stream_message(message, llm_context):
tokens.append(token)
return tokens
async def _execute_llm_stream(
executor: Executor,
message: str,
messages: list[dict[str, Any]] | None = None,
) -> AsyncGenerator[str, None]:
"""Stream tokens from a single LLM actor using AgentFactory.
Mirrors :func:`_execute_llm` but uses
:meth:`~cleveractors.agents.llm.LLMAgent.stream_message` (which calls
``astream()``) instead of ``ainvoke()``.
Execution limits applied (AC5 of issue #16):
- ``timeout_ms``: wraps the entire stream in ``asyncio.wait_for``,
converting ``asyncio.TimeoutError`` to ``ExecutionError(kind="timeout")``.
- ``max_cost_usd``: after the stream completes, computes node cost from
``executor.pricing[provider][model]`` and raises
``ExecutionError(kind="cost", reason="budget_exhausted")`` if exceeded.
After the generator is exhausted:
- ``executor.last_result`` is set to an :class:`~cleveractors.result.ActorResult`
with the concatenated response and token counts from the final chunk.
Cleanup (``agent.cleanup()``) is called in a ``finally`` block so it runs
whether the stream is exhausted normally or an exception is raised.
"""
config_block: dict[str, Any] = executor.config.get("config", {})
top_provider: str | None = executor.config.get("provider")
provider: str = (
top_provider
if top_provider is not None
else config_block.get("provider", "openai")
)
top_model: str | None = executor.config.get("model")
model: str = (
top_model if top_model is not None else config_block.get("model", DEFAULT_MODEL)
)
top_sp: str | None = executor.config.get("system_prompt")
system_prompt: str = (
top_sp
if top_sp is not None
else config_block.get("system_prompt", DEFAULT_SYSTEM_MESSAGE)
)
temperature_raw: Any = executor.config.get("temperature")
if temperature_raw is None:
temperature_raw = config_block.get("temperature", DEFAULT_TEMPERATURE)
agent_name: str = executor.config.get("name", "llm")
# -- Early config-validation wrapper (billing-integrity) ------------------
# The temperature, max_tokens, and timeout_ms validations below may raise
# ConfigurationError or ExecutionError before the agent is created and
# before the main try/except block that populates executor.last_result.
# To satisfy the billing-integrity guarantee (executor.last_result is always
# set on exception paths), we wrap these early validations in a try/except
# that sets a <no_llm> placeholder before re-raising. This mirrors the
# factory.create_agent() handler further below.
try:
try:
temperature: float = float(temperature_raw)
except (TypeError, ValueError, OverflowError) as err:
raise ConfigurationError(
f"Invalid temperature value: {temperature_raw!r}"
) from err
max_tokens_raw: Any = executor.config.get("max_tokens")
if max_tokens_raw is None:
max_tokens_raw = config_block.get("max_tokens", DEFAULT_MAX_TOKENS)
try:
max_tokens: int = int(max_tokens_raw)
except (TypeError, ValueError, OverflowError) as err:
raise ConfigurationError(
f"Invalid max_tokens value: {max_tokens_raw!r}"
) from err
# -- Validate timeout_ms limit (AC5) ----------------------------------
# Mirrors the validation in PureLangGraph.execute_stream() so that
# invalid limit values are caught early (before the agent is created)
# and produce the same ExecutionError(kind="timeout") as the graph path.
timeout_ms: Any = executor.limits.get("timeout_ms")
_timeout_float: float | None = None
if timeout_ms is not None:
if isinstance(timeout_ms, bool):
raise ExecutionError(
f"Invalid timeout_ms value {timeout_ms!r}: bool is not a valid timeout",
kind="timeout",
)
try:
_timeout_float = float(timeout_ms) / 1000.0
except (TypeError, ValueError) as _to_err:
raise ExecutionError(
f"Invalid timeout_ms value {timeout_ms!r}: {_to_err}",
kind="timeout",
) from _to_err
if _timeout_float <= 0:
raise ExecutionError(
f"Invalid timeout_ms value {timeout_ms!r}: "
"must be a positive number of milliseconds",
kind="timeout",
)
except (ConfigurationError, ExecutionError):
# Billing-integrity: set a <no_llm> placeholder so executor.last_result
# is always populated on exception paths, even for early config errors
# that fire before the agent is created.
executor.last_result = ActorResult(
response="",
prompt_tokens=0,
completion_tokens=0,
nodes=[
NodeUsage(
node_id=agent_name,
provider=provider,
model="<no_llm>",
prompt_tokens=0,
completion_tokens=0,
)
],
)
raise
agent_inner_config: dict[str, Any] = copy.deepcopy(
executor.config.get("config", {})
)
agent_inner_config["provider"] = provider
agent_inner_config["model"] = model
agent_inner_config["system_prompt"] = system_prompt
agent_inner_config["temperature"] = temperature
agent_inner_config["max_tokens"] = max_tokens
factory_cfg: dict[str, Any] = {
"agents": {
agent_name: {
"type": "llm",
"provider": provider,
"config": agent_inner_config,
}
}
}
renderer = TemplateRenderer()
factory = AgentFactory(
config=factory_cfg,
credentials=executor.credentials,
template_renderer=renderer,
)
llm_context: dict[str, Any] | None = None
if messages:
conversation_history = [
{
"role": m.get("role", "user"),
"content": m.get("content", ""),
}
for m in messages
]
llm_context = {"conversation_history": conversation_history}
last_usage: tuple[int, int] = (0, 0)
response_parts: list[str] = []
try:
agent = factory.create_agent(agent_name)
except (ConfigurationError, AgentCreationError):
# m1 fix: populate executor.last_result with a <no_llm> placeholder
# before re-raising so that executor.last_result is always set on the
# exception path, mirroring the graph path's N5 fix.
executor.last_result = ActorResult(
response="",
prompt_tokens=0,
completion_tokens=0,
nodes=[
NodeUsage(
node_id=agent_name,
provider=provider,
model="<no_llm>",
prompt_tokens=0,
completion_tokens=0,
)
],
)
raise
except Exception as exc:
executor.last_result = ActorResult(
response="",
prompt_tokens=0,
completion_tokens=0,
nodes=[
NodeUsage(
node_id=agent_name,
provider=provider,
model="<no_llm>",
prompt_tokens=0,
completion_tokens=0,
)
],
)
logger.exception(
"Failed to create LLM agent for streaming: %s", type(exc).__name__
)
raise ExecutionError("LLM execution failed") from None
try:
# -- timeout_ms enforcement (AC5) -------------------------------------
# Collect all tokens under asyncio.wait_for when timeout_ms is set,
# then yield them. This mirrors the graph path's approach in
# PureLangGraph.execute_stream() which buffers via _collect_stream_tokens
# under asyncio.wait_for. For the single-LLM path, we collect into a
# list and yield after the wait_for completes.
if _timeout_float is not None:
try:
buffered_tokens: list[str] = await asyncio.wait_for(
_collect_llm_stream_tokens(agent, message, llm_context),
timeout=_timeout_float,
)
except asyncio.TimeoutError:
raise ExecutionError(
f"Execution timed out after {timeout_ms} ms",
kind="timeout",
) from None
for token in buffered_tokens:
response_parts.append(token)
yield token
else:
async for token in agent.stream_message(message, llm_context):
response_parts.append(token)
yield token
# Capture token usage after stream exhaustion (same fallback as _execute_llm).
_tok_from_var: tuple[int, int] = last_token_usage_var.get((0, 0))
_tok_from_inst: object = getattr(agent, "_last_token_usage", (0, 0))
if _tok_from_var != (0, 0):
last_usage = _tok_from_var
elif not isinstance(agent, LLMAgent):
_tok_inst_val = _tok_from_inst
last_usage = (
_tok_inst_val
if (
isinstance(_tok_inst_val, tuple)
and len(_tok_inst_val) == 2
and isinstance(_tok_inst_val[0], int)
and isinstance(_tok_inst_val[1], int)
)
else (0, 0)
)
else:
last_usage = (0, 0)
# -- max_cost_usd enforcement (AC5) -----------------------------------
# IMPORTANT: This block is intentionally inside the try/except so that
# when it raises ExecutionError (e.g. budget_exhausted, missing pricing
# entry, invalid rate), the except (ConfigurationError, ExecutionError,
# AgentCreationError) handler below populates executor.last_result
# before re-raising — satisfying the billing-integrity guarantee
# documented in runtime.py. Placing this block outside the try/except
# would bypass that handler and leave executor.last_result as None on
# cost-check failures.
#
# Mirrors the cost block in PureLangGraph._stream_from_node() so that
# single-LLM actor streams are subject to the same budget guardrail as
# graph actor streams. Cost is computed after the stream completes
# (token counts are only available then).
prompt_tokens, completion_tokens = last_usage
node_usage = NodeUsage(
node_id=agent_name,
provider=provider,
model=model,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
)
if executor.pricing:
_llm_provider_pricing = executor.pricing.get(provider)
if _llm_provider_pricing is None or not isinstance(
_llm_provider_pricing, dict
):
raise ExecutionError(
f"Missing pricing entry for provider '{provider}'",
kind="cost",
reason="missing_pricing_entry",
)
_llm_model_pricing = _llm_provider_pricing.get(model)
if _llm_model_pricing is None or not isinstance(_llm_model_pricing, dict):
raise ExecutionError(
f"Missing pricing entry for model '{model}' "
f"under provider '{provider}'",
kind="cost",
reason="missing_pricing_entry",
)
_llm_prompt_rate_raw = _llm_model_pricing.get("prompt")
_llm_completion_rate_raw = _llm_model_pricing.get("completion")
if _llm_prompt_rate_raw is None or _llm_completion_rate_raw is None:
raise ExecutionError(
f"Incomplete pricing entry for model '{model}' "
f"under provider '{provider}': missing 'prompt' "
f"or 'completion' rate key",
kind="cost",
reason="missing_pricing_entry",
)
try:
_llm_prompt_rate = float(_llm_prompt_rate_raw)
_llm_completion_rate = float(_llm_completion_rate_raw)
except (TypeError, ValueError) as _llm_rate_err:
raise ExecutionError(
f"Invalid pricing rate for model '{model}' "
f"under provider '{provider}': {_llm_rate_err}",
kind="cost",
reason="missing_pricing_entry",
) from _llm_rate_err
_llm_node_cost = (
prompt_tokens / 1_000_000.0 * _llm_prompt_rate
+ completion_tokens / 1_000_000.0 * _llm_completion_rate
)
_llm_max_cost = executor.limits.get("max_cost_usd")
if _llm_max_cost is not None:
if isinstance(_llm_max_cost, bool):
raise ExecutionError(
f"Invalid max_cost_usd value {_llm_max_cost!r}: "
"bool is not a valid cost limit",
kind="cost",
)
try:
_llm_max_cost_float = float(_llm_max_cost)
except (TypeError, ValueError) as _llm_cost_err:
raise ExecutionError(
f"Invalid max_cost_usd value {_llm_max_cost!r}: {_llm_cost_err}",
kind="cost",
) from _llm_cost_err
if _llm_node_cost > _llm_max_cost_float:
raise ExecutionError(
f"Cost limit exceeded: "
f"{_llm_node_cost:.6f} USD "
f"> {_llm_max_cost} USD",
kind="cost",
reason="budget_exhausted",
)
executor.last_result = ActorResult(
response="".join(response_parts),
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
nodes=[node_usage],
)
except (ConfigurationError, ExecutionError, AgentCreationError):
# Billing integrity — read whatever token counts stream_message() had
# captured before re-raising. Note: when chat_model.astream() raises
# mid-stream, token counts are (0, 0) because LangChain's astream does
# not surface partial token counts before the final chunk. The captured
# counts are non-zero only when the exception occurred in a post-stream
# step (e.g. update_memory()). Mirrors the fix in _execute_graph_stream.
# Also catches ExecutionError raised by the cost check block above,
# ensuring executor.last_result is always set before re-raising.
_exc_tok_var: tuple[int, int] = last_token_usage_var.get((0, 0))
_exc_tok_inst: object = getattr(agent, "_last_token_usage", (0, 0))
_exc_usage: tuple[int, int]
if _exc_tok_var != (0, 0):
_exc_usage = _exc_tok_var
elif not isinstance(agent, LLMAgent) and (
isinstance(_exc_tok_inst, tuple)
and len(_exc_tok_inst) == 2
and isinstance(_exc_tok_inst[0], int)
and isinstance(_exc_tok_inst[1], int)
):
# Mirror the isinstance(agent, LLMAgent) guard from the success
# path — only fall back to the instance attribute for non-LLMAgent
# types, consistent with the established pattern.
_exc_usage = (_exc_tok_inst[0], _exc_tok_inst[1])
else:
_exc_usage = (0, 0)
_exc_pt, _exc_ct = _exc_usage
executor.last_result = ActorResult(
response="".join(response_parts),
prompt_tokens=_exc_pt,
completion_tokens=_exc_ct,
nodes=[
NodeUsage(
node_id=agent_name,
provider=provider,
model=model,
prompt_tokens=_exc_pt,
completion_tokens=_exc_ct,
)
],
)
raise
except Exception as exc:
# Billing-integrity treatment for unexpected exceptions.
# Use (0, 0) since the stream failed unexpectedly.
executor.last_result = ActorResult(
response="".join(response_parts),
prompt_tokens=0,
completion_tokens=0,
nodes=[
NodeUsage(
node_id=agent_name,
provider=provider,
model=model,
prompt_tokens=0,
completion_tokens=0,
)
],
)
logger.exception("LLM agent streaming failed: %s", type(exc).__name__)
raise ExecutionError("LLM execution failed") from None
finally:
if hasattr(agent, "cleanup"):
try:
await agent.cleanup()
except Exception as e:
logger.warning(
"cleanup failed for agent %s: %s",
getattr(agent, "name", "?"),
e,
)
# -- streaming graph ----------------------------------------------------------
async def _execute_graph_stream(
executor: Executor,
message: str,
state: dict[str, Any] | None = None,
messages: list[dict[str, Any]] | None = None,
) -> AsyncGenerator[str, None]:
"""Stream tokens from a graph actor using PureLangGraph.execute_stream().
Mirrors :func:`_execute_graph` but uses
:meth:`~cleveractors.langgraph.pure_graph.PureLangGraph.execute_stream`
instead of :meth:`~cleveractors.langgraph.pure_graph.PureLangGraph.execute`.
After the generator is exhausted:
- ``executor.last_result`` is set from the graph's post-stream state and
per-node token usage (stored on the graph as ``_last_stream_state`` and
``_last_stream_node_usages``).
Limit enforcement (depth, model_calls, tool_calls, timeout) is handled
entirely within :meth:`PureLangGraph.execute_stream` and
:meth:`PureLangGraph._stream_from_node`.
"""
# Graph config normalisation (mirrors _execute_graph)
route = executor.config.get("route", {})
nodes_cfg = list(route.get("nodes", [])) if route else []
edges_cfg = list(route.get("edges", [])) if route else []
entry_point = route.get("entry_node", "start") if route else "start"
if not route:
routes = executor.config.get("routes", {})
main = routes.get("main", {})
raw_nodes = main.get("nodes", {})
if isinstance(raw_nodes, dict):
nodes_cfg = []
for node_id, node_def in raw_nodes.items():
node_def = dict(node_def) if node_def else {}
node_def["id"] = node_id
nodes_cfg.append(node_def)
elif isinstance(raw_nodes, list):
nodes_cfg = list(raw_nodes)
edges_cfg = list(main.get("edges", []))
entry_point = main.get("entry_point", "start")
# -- Early config-normalisation wrapper (billing-integrity) ---------------
# The node/edge validation loop below may raise ConfigurationError before
# the main try/except block that populates executor.last_result. To satisfy
# the billing-integrity guarantee (executor.last_result is always set on
# exception paths), we wrap these early validations in a try/except that
# sets a <no_llm> placeholder before re-raising. This mirrors the
# equivalent wrapper in _execute_llm_stream.
try:
pg_nodes: dict[str, NodeConfig] = {}
all_actors: dict[str, Any] = executor.config.get("actors", {})
for node_def in nodes_cfg:
if not isinstance(node_def, dict) or "id" not in node_def:
raise ConfigurationError(
f"Invalid node definition in graph route: {node_def!r}"
)
node_id: str = node_def["id"]
if node_id in pg_nodes:
raise ConfigurationError(f"Duplicate node ID in graph: {node_id!r}")
agent_name = node_def.get("agent")
if not agent_name and node_id in all_actors:
agent_name = node_id
node_type_str: str = node_def.get("type", "function")
if agent_name:
node_type_str = "agent"
try:
node_type = NodeType(node_type_str)
except (ValueError, TypeError):
node_type = NodeType.FUNCTION
pg_nodes[node_id] = NodeConfig(
name=node_def["id"],
type=node_type,
agent=agent_name,
function=node_def.get("function"),
tools=node_def.get("tools", []),
retry_policy=node_def.get("retry_policy"),
timeout=node_def.get("timeout"),
parallel=node_def.get("parallel", False),
condition=node_def.get("condition"),
subgraph=node_def.get("subgraph"),
metadata=node_def.get("metadata", {}),
)
pg_edges: list[Edge] = []
for edge_def in edges_cfg:
if (
not isinstance(edge_def, dict)
or "source" not in edge_def
or "target" not in edge_def
):
raise ConfigurationError(
f"Invalid edge definition in graph route: {edge_def!r}"
)
pg_edges.append(
Edge(
source=edge_def["source"],
target=edge_def["target"],
condition=edge_def.get("condition"),
metadata=edge_def.get("metadata", {}),
)
)
if route:
parallel_execution: bool = route.get("parallel_execution", True)
else:
parallel_execution = (
executor.config.get("routes", {})
.get("main", {})
.get("parallel_execution", True)
)
pg_config = PureGraphConfig(
name=executor.config.get("name", "graph"),
nodes=pg_nodes,
edges=pg_edges,
entry_point=entry_point,
parallel_execution=parallel_execution,
)
except ConfigurationError:
# Billing-integrity: set a <no_llm> placeholder so executor.last_result
# is always populated on exception paths, even for early config errors
# that fire before the graph is built.
_early_graph_name = executor.config.get("name", "graph")
executor.last_result = ActorResult(
response="",
prompt_tokens=0,
completion_tokens=0,
nodes=[
NodeUsage(
node_id=f"<{_early_graph_name}:no_llm>",
provider="graph",
model="<no_llm>",
prompt_tokens=0,
completion_tokens=0,
)
],
)
raise
factory_config = copy.deepcopy(executor.config)
factory_config.setdefault("agents", {})
factory_config["agents"].update(factory_config.get("actors", {}))
renderer = TemplateRenderer()
factory = AgentFactory(
config=factory_config,
credentials=executor.credentials,
template_renderer=renderer,
)
agents: dict[str, Any] = {}
response_parts: list[str] = []
captured_state: dict[str, Any] = {}
raw_node_usages: list[Any] = []
# Declared before the try block so the except handlers can access it for
# billing-integrity state capture (M2 fix from review).
graph: PureLangGraph | None = None
try:
for node_def in nodes_cfg:
agent_name = node_def.get("agent")
node_id = node_def["id"]
if not agent_name and node_id in all_actors:
agent_name = node_id
if agent_name and agent_name not in agents:
try:
agents[agent_name] = factory.create_agent(agent_name)
except ConfigurationError:
raise
except AgentCreationError:
raise
except Exception as exc:
# Sanitised: do not embed exc in the message to avoid
# leaking sensitive provider details (API keys, URLs).
# The cause chain is preserved via `from exc`.
raise ConfigurationError(
f"Failed to create agent '{agent_name}'"
) from exc
conversation_history: list[dict[str, Any]] | None = None
if messages:
conversation_history = [
{
"role": m.get("role", "user"),
"content": m.get("content", ""),
}
for m in messages
]
actor_context: dict[str, Any] = executor.config.get("context", {})
global_context: dict[str, Any] = {}
if "global" in actor_context:
global_context.update(actor_context["global"])
elif actor_context:
global_context.update(actor_context)
if conversation_history:
global_context["conversation_history"] = conversation_history
graph = PureLangGraph(
config=pg_config,
agents=agents,
limits=executor.limits,
pricing=executor.pricing,
)
async for token in graph.execute_stream(
input_message=message,
global_context=global_context if global_context else None,
conversation_history=conversation_history,
initial_state=state,
):
response_parts.append(token)
yield token
# After stream exhaustion: read state/usages stored by execute_stream()
captured_state = dict(graph._last_stream_state)
raw_node_usages = list(graph._last_stream_node_usages)
except (ConfigurationError, AgentCreationError, ExecutionError):
# M2 fix: billing integrity — read whatever state/usages the graph's
# finally block already captured before re-raising. The execute_stream()
# finally block populates _last_stream_state/_last_stream_node_usages even
# on exception, so partial token counts are preserved for the router.
#
# N5 fix: when graph is None (agent creation failed before the graph was
# built), populate a synthetic <no_llm> placeholder so executor.last_result
# is always set on the exception path, mirroring _execute_graph behaviour.
if graph is not None:
try:
captured_state = dict(graph._last_stream_state)
raw_node_usages = list(graph._last_stream_node_usages)
except Exception as _state_err: # pylint: disable=broad-exception-caught
# State capture failed; billing data will be (0, 0) for this
# call. Log at debug level so the original exception is not
# masked (the outer except re-raises it).
logger.debug(
"Failed to capture post-stream state on exception path: %s",
_state_err,
)
_exc_nodes: list[NodeUsage] = []
for _exc_usage in raw_node_usages:
try:
_exc_nodes.append(NodeUsage(*_exc_usage))
except (TypeError, ValueError):
pass
if not _exc_nodes:
_exc_graph_name = executor.config.get("name", "graph")
_exc_nodes = [
NodeUsage(
node_id=f"<{_exc_graph_name}:no_llm>",
provider="graph",
model="<no_llm>",
prompt_tokens=0,
completion_tokens=0,
)
]
executor.last_result = ActorResult(
response="".join(response_parts),
prompt_tokens=sum(n.prompt_tokens for n in _exc_nodes),
completion_tokens=sum(n.completion_tokens for n in _exc_nodes),
nodes=_exc_nodes,
state=captured_state if captured_state else None,
)
raise
except Exception as exc:
# Billing integrity — mirror the (ConfigurationError,
# AgentCreationError, ExecutionError) block above. Read whatever
# state/usages the graph's finally block already captured before
# re-raising so executor.last_result is always set on the exception
# path, even for unexpected (non-ExecutionError) exceptions such as
# RuntimeError from a function node or KeyError from a malformed config.
logger.exception("Graph streaming failed: %s", type(exc).__name__)
if graph is not None:
try:
captured_state = dict(graph._last_stream_state)
raw_node_usages = list(graph._last_stream_node_usages)
except Exception as _state_err: # pylint: disable=broad-exception-caught
logger.debug(
"Failed to capture post-stream state on unexpected exception path: %s",
_state_err,
)
_unexp_nodes: list[NodeUsage] = []
for _unexp_usage in raw_node_usages:
try:
_unexp_nodes.append(NodeUsage(*_unexp_usage))
except (TypeError, ValueError):
pass
if not _unexp_nodes:
_unexp_graph_name = executor.config.get("name", "graph")
_unexp_nodes = [
NodeUsage(
node_id=f"<{_unexp_graph_name}:no_llm>",
provider="graph",
model="<no_llm>",
prompt_tokens=0,
completion_tokens=0,
)
]
executor.last_result = ActorResult(
response="".join(response_parts),
prompt_tokens=sum(n.prompt_tokens for n in _unexp_nodes),
completion_tokens=sum(n.completion_tokens for n in _unexp_nodes),
nodes=_unexp_nodes,
state=captured_state if captured_state else None,
)
raise ExecutionError("Graph execution failed") from exc
finally:
for ag in agents.values():
if hasattr(ag, "cleanup"):
try:
await ag.cleanup()
except Exception as e:
logger.warning(
"cleanup failed for agent %s: %s",
getattr(ag, "name", "?"),
e,
)
# Build NodeUsage objects (mirrors _execute_graph)
nodes: list[NodeUsage] = []
for usage_tuple in raw_node_usages:
try:
nodes.append(NodeUsage(*usage_tuple))
except (TypeError, ValueError) as _node_usage_err:
logger.warning(
"Skipping malformed node usage tuple %r: %s",
usage_tuple,
_node_usage_err,
)
if not nodes:
graph_name = executor.config.get("name", "graph")
nodes = [
NodeUsage(
node_id=f"<{graph_name}:no_llm>",
provider="graph",
model="<no_llm>",
prompt_tokens=0,
completion_tokens=0,
)
]
executor.last_result = ActorResult(
response="".join(response_parts),
prompt_tokens=sum(n.prompt_tokens for n in nodes),
completion_tokens=sum(n.completion_tokens for n in nodes),
nodes=nodes,
state=captured_state if captured_state else None,
)