b60439152b
- harness/es.py: append-only event store + проекции (образец — кооп-портал) - harness/levels.py: уровни 1с/1м/1ч/1д, контекст ∝ масштабу - harness/scheduler.py: уровневые потоки, идемпотентные тики, absorbed-стабы - harness/isolated_runner.py + _llm_worker.py: LLM-тик в сабпроцессе (глобальный asyncio-рантайм Strands не переживает 2+ create_harness) - cli.py: init/task/tasks/status/events + requeue (написан самим минутным агентом) - 9router: отдельный ключ sk-simbiot, модель auto (glm-5.3 первым)
573 lines
22 KiB
Python
573 lines
22 KiB
Python
"""Utilities for handling streaming responses from language models."""
|
|
|
|
import copy
|
|
import json
|
|
import logging
|
|
import threading
|
|
import time
|
|
import warnings
|
|
from collections.abc import AsyncGenerator, AsyncIterable
|
|
from typing import Any
|
|
|
|
from ..agent.agent_metadata import AgentMetadata
|
|
from ..models.model import Model
|
|
from ..tools import InvalidToolUseNameException
|
|
from ..tools.tools import validate_tool_use_name
|
|
from ..types._events import (
|
|
CitationStreamEvent,
|
|
ModelStopReason,
|
|
ModelStreamChunkEvent,
|
|
ModelStreamEvent,
|
|
ReasoningRedactedContentStreamEvent,
|
|
ReasoningSignatureStreamEvent,
|
|
ReasoningTextStreamEvent,
|
|
TextStreamEvent,
|
|
ToolUseStreamEvent,
|
|
TypedEvent,
|
|
)
|
|
from ..types.citations import CitationsContentBlock
|
|
from ..types.content import ContentBlock, Message, Messages, SystemContentBlock
|
|
from ..types.streaming import (
|
|
ContentBlockDeltaEvent,
|
|
ContentBlockStart,
|
|
ContentBlockStartEvent,
|
|
MessageStartEvent,
|
|
MessageStopEvent,
|
|
MetadataEvent,
|
|
Metrics,
|
|
RedactContentEvent,
|
|
StopReason,
|
|
StreamEvent,
|
|
Usage,
|
|
)
|
|
from ..types.tools import ToolSpec, ToolUse
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _normalize_messages(messages: Messages) -> Messages:
|
|
"""Remove or replace blank text in message content.
|
|
|
|
Args:
|
|
messages: Conversation messages to update.
|
|
|
|
Returns:
|
|
Updated messages.
|
|
"""
|
|
removed_blank_message_content_text = False
|
|
replaced_blank_message_content_text = False
|
|
replaced_tool_names = False
|
|
|
|
# Deep copy up front so downstream normalization can mutate freely without
|
|
# affecting the caller's message history.
|
|
messages = copy.deepcopy(messages)
|
|
|
|
for message in messages:
|
|
# only modify assistant messages
|
|
if "role" in message and message["role"] != "assistant":
|
|
continue
|
|
if "content" in message:
|
|
content = message["content"]
|
|
if len(content) == 0:
|
|
content.append({"text": "[blank text]"})
|
|
continue
|
|
|
|
has_tool_use = False
|
|
|
|
# Ensure the tool-uses always have valid names before sending
|
|
# https://github.com/strands-agents/harness-sdk/issues/1069
|
|
for item in content:
|
|
if "toolUse" in item:
|
|
has_tool_use = True
|
|
tool_use: ToolUse = item["toolUse"]
|
|
|
|
try:
|
|
validate_tool_use_name(tool_use)
|
|
except InvalidToolUseNameException:
|
|
tool_use["name"] = "INVALID_TOOL_NAME"
|
|
replaced_tool_names = True
|
|
|
|
if has_tool_use:
|
|
# Remove blank or None 'text' items for assistant messages
|
|
before_len = len(content)
|
|
content[:] = [
|
|
item
|
|
for item in content
|
|
if "text" not in item or (item["text"] is not None and item["text"].strip())
|
|
]
|
|
if not removed_blank_message_content_text and before_len != len(content):
|
|
removed_blank_message_content_text = True
|
|
else:
|
|
# Replace blank or None 'text' with '[blank text]' for assistant messages
|
|
for item in content:
|
|
if "text" in item and (item["text"] is None or not item["text"].strip()):
|
|
replaced_blank_message_content_text = True
|
|
item["text"] = "[blank text]"
|
|
|
|
if removed_blank_message_content_text:
|
|
logger.debug("removed blank message context text")
|
|
if replaced_blank_message_content_text:
|
|
logger.debug("replaced blank message context text")
|
|
if replaced_tool_names:
|
|
logger.debug("replaced invalid tool name")
|
|
|
|
return messages
|
|
|
|
|
|
def remove_blank_messages_content_text(messages: Messages) -> Messages:
|
|
"""Remove or replace blank text in message content.
|
|
|
|
!!deprecated!!
|
|
This function is deprecated and will be removed in a future version.
|
|
|
|
Args:
|
|
messages: Conversation messages to update.
|
|
|
|
Returns:
|
|
Updated messages.
|
|
"""
|
|
warnings.warn(
|
|
"remove_blank_messages_content_text is deprecated and will be removed in a future version.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
removed_blank_message_content_text = False
|
|
replaced_blank_message_content_text = False
|
|
|
|
for message in messages:
|
|
# only modify assistant messages
|
|
if "role" in message and message["role"] != "assistant":
|
|
continue
|
|
if "content" in message:
|
|
content = message["content"]
|
|
has_tool_use = any("toolUse" in item for item in content)
|
|
if len(content) == 0:
|
|
content.append({"text": "[blank text]"})
|
|
continue
|
|
|
|
if has_tool_use:
|
|
# Remove blank or None 'text' items for assistant messages
|
|
before_len = len(content)
|
|
content[:] = [
|
|
item
|
|
for item in content
|
|
if "text" not in item or (item["text"] is not None and item["text"].strip())
|
|
]
|
|
if not removed_blank_message_content_text and before_len != len(content):
|
|
removed_blank_message_content_text = True
|
|
else:
|
|
# Replace blank or None 'text' with '[blank text]' for assistant messages
|
|
for item in content:
|
|
if "text" in item and (item["text"] is None or not item["text"].strip()):
|
|
replaced_blank_message_content_text = True
|
|
item["text"] = "[blank text]"
|
|
|
|
if removed_blank_message_content_text:
|
|
logger.debug("removed blank message context text")
|
|
if replaced_blank_message_content_text:
|
|
logger.debug("replaced blank message context text")
|
|
|
|
return messages
|
|
|
|
|
|
def handle_message_start(event: MessageStartEvent, message: Message) -> Message:
|
|
"""Handles the start of a message by setting the role in the message dictionary.
|
|
|
|
Args:
|
|
event: A message start event.
|
|
message: The message dictionary being constructed.
|
|
|
|
Returns:
|
|
Updated message dictionary with the role set.
|
|
"""
|
|
message["role"] = event["role"]
|
|
return message
|
|
|
|
|
|
def handle_content_block_start(event: ContentBlockStartEvent) -> dict[str, Any]:
|
|
"""Handles the start of a content block by extracting tool usage information if any.
|
|
|
|
Args:
|
|
event: Start event.
|
|
|
|
Returns:
|
|
Dictionary with tool use id and name if tool use request, empty dictionary otherwise.
|
|
"""
|
|
start: ContentBlockStart = event["start"]
|
|
current_tool_use = {}
|
|
|
|
if "toolUse" in start and start["toolUse"]:
|
|
tool_use_data = start["toolUse"]
|
|
current_tool_use["toolUseId"] = tool_use_data["toolUseId"]
|
|
current_tool_use["name"] = tool_use_data["name"]
|
|
current_tool_use["input"] = ""
|
|
if "reasoningSignature" in tool_use_data:
|
|
current_tool_use["reasoningSignature"] = tool_use_data["reasoningSignature"]
|
|
|
|
return current_tool_use
|
|
|
|
|
|
def handle_content_block_delta(
|
|
event: ContentBlockDeltaEvent, state: dict[str, Any]
|
|
) -> tuple[dict[str, Any], ModelStreamEvent]:
|
|
"""Handles content block delta updates by appending text, tool input, or reasoning content to the state.
|
|
|
|
Args:
|
|
event: Delta event.
|
|
state: The current state of message processing.
|
|
|
|
Returns:
|
|
Updated state with appended text or tool input.
|
|
"""
|
|
delta_content = event["delta"]
|
|
|
|
typed_event: ModelStreamEvent = ModelStreamEvent({})
|
|
|
|
if "toolUse" in delta_content:
|
|
tool_use_delta = delta_content["toolUse"]
|
|
if "input" not in state["current_tool_use"]:
|
|
state["current_tool_use"]["input"] = ""
|
|
|
|
state["current_tool_use"]["input"] += tool_use_delta.get("input", "")
|
|
|
|
# Some models emit toolUseId/name in the delta instead of contentBlockStart; keep values already set.
|
|
for field in ("toolUseId", "name"):
|
|
if field not in state["current_tool_use"] and field in tool_use_delta:
|
|
state["current_tool_use"][field] = tool_use_delta[field]
|
|
|
|
typed_event = ToolUseStreamEvent(delta_content, state["current_tool_use"])
|
|
|
|
elif "text" in delta_content:
|
|
state["text"] += delta_content["text"]
|
|
typed_event = TextStreamEvent(text=delta_content["text"], delta=delta_content)
|
|
|
|
elif "citation" in delta_content:
|
|
if "citationsContent" not in state:
|
|
state["citationsContent"] = []
|
|
|
|
state["citationsContent"].append(delta_content["citation"])
|
|
typed_event = CitationStreamEvent(delta=delta_content, citation=delta_content["citation"])
|
|
|
|
elif "reasoningContent" in delta_content:
|
|
if "text" in delta_content["reasoningContent"]:
|
|
if "reasoningText" not in state:
|
|
state["reasoningText"] = ""
|
|
|
|
state["reasoningText"] += delta_content["reasoningContent"]["text"]
|
|
typed_event = ReasoningTextStreamEvent(
|
|
reasoning_text=delta_content["reasoningContent"]["text"],
|
|
delta=delta_content,
|
|
)
|
|
|
|
elif "signature" in delta_content["reasoningContent"]:
|
|
if "signature" not in state:
|
|
state["signature"] = ""
|
|
|
|
state["signature"] += delta_content["reasoningContent"]["signature"]
|
|
typed_event = ReasoningSignatureStreamEvent(
|
|
reasoning_signature=delta_content["reasoningContent"]["signature"],
|
|
delta=delta_content,
|
|
)
|
|
|
|
elif redacted_content := delta_content["reasoningContent"].get("redactedContent"):
|
|
state["redactedContent"] = state.get("redactedContent", b"") + redacted_content
|
|
typed_event = ReasoningRedactedContentStreamEvent(redacted_content=redacted_content, delta=delta_content)
|
|
|
|
return state, typed_event
|
|
|
|
|
|
def handle_content_block_stop(state: dict[str, Any]) -> dict[str, Any]:
|
|
"""Handles the end of a content block by finalizing tool usage, text content, or reasoning content.
|
|
|
|
Args:
|
|
state: The current state of message processing.
|
|
|
|
Returns:
|
|
Updated state with finalized content block.
|
|
"""
|
|
content: list[ContentBlock] = state["content"]
|
|
|
|
current_tool_use = state["current_tool_use"]
|
|
text = state["text"]
|
|
reasoning_text = state["reasoningText"]
|
|
citations_content = state["citationsContent"]
|
|
redacted_content = state.get("redactedContent")
|
|
|
|
if current_tool_use:
|
|
if "input" not in current_tool_use:
|
|
current_tool_use["input"] = ""
|
|
|
|
# Handle empty or whitespace-only input (common for zero-argument tools)
|
|
raw_input = current_tool_use["input"]
|
|
if isinstance(raw_input, str) and not raw_input.strip():
|
|
current_tool_use["input"] = {}
|
|
else:
|
|
try:
|
|
current_tool_use["input"] = json.loads(raw_input)
|
|
except ValueError:
|
|
logger.warning(
|
|
"tool_name=<%s>, raw_input=<%s> | failed to parse tool input json, defaulting to empty dict",
|
|
current_tool_use.get("name", "unknown"),
|
|
raw_input[:200] if isinstance(raw_input, str) else "",
|
|
)
|
|
current_tool_use["input"] = {}
|
|
|
|
tool_use_id = current_tool_use.get("toolUseId", "")
|
|
tool_use_name = current_tool_use.get("name", "")
|
|
|
|
if not tool_use_id or not tool_use_name:
|
|
# Skip, don't raise: an empty tool_uses list still appends a valid toolResult, so the loop continues.
|
|
logger.warning(
|
|
"tool_use_id=<%s>, tool_name=<%s> | incomplete tool use block, skipping content block "
|
|
"(model may be using a non-standard streaming format)",
|
|
tool_use_id,
|
|
tool_use_name,
|
|
)
|
|
state["current_tool_use"] = {}
|
|
return state
|
|
|
|
tool_use = ToolUse(
|
|
toolUseId=tool_use_id,
|
|
name=tool_use_name,
|
|
input=current_tool_use["input"],
|
|
)
|
|
if "reasoningSignature" in current_tool_use:
|
|
tool_use["reasoningSignature"] = current_tool_use["reasoningSignature"]
|
|
content.append({"toolUse": tool_use})
|
|
state["current_tool_use"] = {}
|
|
|
|
elif text:
|
|
if citations_content:
|
|
citations_block: CitationsContentBlock = {"citations": citations_content, "content": [{"text": text}]}
|
|
content.append({"citationsContent": citations_block})
|
|
state["citationsContent"] = []
|
|
else:
|
|
content.append({"text": text})
|
|
state["text"] = ""
|
|
|
|
elif reasoning_text or "signature" in state:
|
|
content_block: ContentBlock = {
|
|
"reasoningContent": {
|
|
"reasoningText": {
|
|
"text": state["reasoningText"],
|
|
}
|
|
}
|
|
}
|
|
|
|
# Consume the signature so it belongs to exactly this block and does not leak into the next one.
|
|
if (signature := state.pop("signature", None)) is not None:
|
|
content_block["reasoningContent"]["reasoningText"]["signature"] = signature
|
|
|
|
content.append(content_block)
|
|
state["reasoningText"] = ""
|
|
elif redacted_content:
|
|
content.append({"reasoningContent": {"redactedContent": redacted_content}})
|
|
state["redactedContent"] = b""
|
|
|
|
return state
|
|
|
|
|
|
def handle_message_stop(event: MessageStopEvent, content: list[dict[str, Any]]) -> StopReason:
|
|
"""Handles the end of a message by returning the stop reason.
|
|
|
|
Some models return "end_turn" even when tool calls are present, which prevents the event loop from processing
|
|
those tool calls. This function overrides to "tool_use" so tool execution proceeds correctly.
|
|
|
|
Args:
|
|
event: Stop event.
|
|
content: The message content blocks accumulated during streaming.
|
|
|
|
Returns:
|
|
The reason for stopping the stream.
|
|
"""
|
|
stop_reason = event["stopReason"]
|
|
|
|
if stop_reason == "end_turn" and any("toolUse" in item for item in content):
|
|
logger.warning(
|
|
"original_stop_reason=<%s>, new_stop_reason=<%s> | "
|
|
"overriding stop reason due to toolUse blocks in response",
|
|
"end_turn",
|
|
"tool_use",
|
|
)
|
|
stop_reason = "tool_use"
|
|
|
|
return stop_reason
|
|
|
|
|
|
def handle_redact_content(event: RedactContentEvent, state: dict[str, Any]) -> None:
|
|
"""Handles redacting content from the input or output.
|
|
|
|
Args:
|
|
event: Redact Content Event.
|
|
state: The current state of message processing.
|
|
"""
|
|
if event.get("redactAssistantContentMessage") is not None:
|
|
state["message"]["content"] = [{"text": event["redactAssistantContentMessage"]}]
|
|
|
|
|
|
def extract_usage_metrics(event: MetadataEvent, time_to_first_byte_ms: int | None = None) -> tuple[Usage, Metrics]:
|
|
"""Extracts usage metrics from the metadata chunk.
|
|
|
|
Args:
|
|
event: metadata.
|
|
time_to_first_byte_ms: time to get the first byte from the model in milliseconds
|
|
|
|
Returns:
|
|
The extracted usage metrics and latency.
|
|
"""
|
|
# MetadataEvent has total=False, making all fields optional, but Usage and Metrics types
|
|
# have Required fields. Provide defaults to handle cases where custom models don't
|
|
# provide usage/metrics (e.g., when latency info is unavailable).
|
|
usage = Usage(**{"inputTokens": 0, "outputTokens": 0, "totalTokens": 0, **event.get("usage", {})})
|
|
metrics = Metrics(**{"latencyMs": 0, **event.get("metrics", {})})
|
|
if time_to_first_byte_ms:
|
|
metrics["timeToFirstByteMs"] = time_to_first_byte_ms
|
|
|
|
return usage, metrics
|
|
|
|
|
|
async def process_stream(
|
|
chunks: AsyncIterable[StreamEvent],
|
|
start_time: float | None = None,
|
|
cancel_signal: threading.Event | None = None,
|
|
) -> AsyncGenerator[TypedEvent, None]:
|
|
"""Processes the response stream from the API, constructing the final message and extracting usage metrics.
|
|
|
|
Args:
|
|
chunks: The chunks of the response stream from the model.
|
|
start_time: Time when the model request is initiated
|
|
cancel_signal: Optional threading.Event to check for cancellation during streaming.
|
|
|
|
Yields:
|
|
The reason for stopping, the constructed message, and the usage metrics.
|
|
"""
|
|
stop_reason: StopReason = "end_turn"
|
|
saw_message_stop = False
|
|
first_byte_time = None
|
|
|
|
state: dict[str, Any] = {
|
|
"message": {"role": "assistant", "content": []},
|
|
"text": "",
|
|
"current_tool_use": {},
|
|
"reasoningText": "",
|
|
"citationsContent": [],
|
|
}
|
|
state["content"] = state["message"]["content"]
|
|
|
|
usage: Usage = Usage(inputTokens=0, outputTokens=0, totalTokens=0)
|
|
metrics: Metrics = Metrics(latencyMs=0, timeToFirstByteMs=0)
|
|
|
|
async for chunk in chunks:
|
|
# Check for cancellation during stream processing
|
|
if cancel_signal and cancel_signal.is_set():
|
|
logger.debug("cancellation detected during stream processing")
|
|
# Return cancelled stop reason with cancellation message
|
|
# The incomplete message in state["message"] is discarded and never added to agent.messages
|
|
yield ModelStopReason(
|
|
stop_reason="cancelled",
|
|
message={"role": "assistant", "content": [{"text": "Cancelled by user"}]},
|
|
usage=usage,
|
|
metrics=metrics,
|
|
)
|
|
return
|
|
|
|
# Track first byte time when we get first content
|
|
if first_byte_time is None and ("contentBlockDelta" in chunk or "contentBlockStart" in chunk):
|
|
first_byte_time = time.time()
|
|
yield ModelStreamChunkEvent(chunk=chunk)
|
|
|
|
if "messageStart" in chunk:
|
|
state["message"] = handle_message_start(chunk["messageStart"], state["message"])
|
|
elif "contentBlockStart" in chunk:
|
|
state["current_tool_use"] = handle_content_block_start(chunk["contentBlockStart"])
|
|
elif "contentBlockDelta" in chunk:
|
|
state, typed_event = handle_content_block_delta(chunk["contentBlockDelta"], state)
|
|
yield typed_event
|
|
elif "contentBlockStop" in chunk:
|
|
state = handle_content_block_stop(state)
|
|
elif "messageStop" in chunk:
|
|
saw_message_stop = True
|
|
stop_reason = handle_message_stop(chunk["messageStop"], state["message"].get("content", []))
|
|
elif "metadata" in chunk:
|
|
time_to_first_byte_ms = (
|
|
int(1000 * (first_byte_time - start_time)) if (start_time and first_byte_time) else None
|
|
)
|
|
usage, metrics = extract_usage_metrics(chunk["metadata"], time_to_first_byte_ms)
|
|
elif "redactContent" in chunk:
|
|
handle_redact_content(chunk["redactContent"], state)
|
|
|
|
# A provider that aborts in flight ends the stream without messageStop; the in-loop check
|
|
# above cannot see it, so report the cancellation here rather than a truncated end_turn.
|
|
if not saw_message_stop and cancel_signal and cancel_signal.is_set():
|
|
yield ModelStopReason(
|
|
stop_reason="cancelled",
|
|
message={"role": "assistant", "content": [{"text": "Cancelled by user"}]},
|
|
usage=usage,
|
|
metrics=metrics,
|
|
)
|
|
return
|
|
|
|
yield ModelStopReason(stop_reason=stop_reason, message=state["message"], usage=usage, metrics=metrics)
|
|
|
|
|
|
async def stream_messages(
|
|
model: Model,
|
|
system_prompt: str | None,
|
|
messages: Messages,
|
|
tool_specs: list[ToolSpec],
|
|
*,
|
|
tool_choice: Any | None = None,
|
|
system_prompt_content: list[SystemContentBlock] | None = None,
|
|
invocation_state: dict[str, Any] | None = None,
|
|
model_state: dict[str, Any] | None = None,
|
|
dynamic_trailing_blocks: int = 0,
|
|
cancel_signal: threading.Event | None = None,
|
|
agent_metadata: AgentMetadata | None = None,
|
|
**kwargs: Any,
|
|
) -> AsyncGenerator[TypedEvent, None]:
|
|
"""Streams messages to the model and processes the response.
|
|
|
|
Args:
|
|
model: Model provider.
|
|
system_prompt: The system prompt string, used for backwards compatibility with models that expect it.
|
|
messages: List of messages to send.
|
|
tool_specs: The list of tool specs.
|
|
tool_choice: Optional tool choice constraint for forcing specific tool usage.
|
|
system_prompt_content: The authoritative system prompt content blocks that always contains the
|
|
system prompt data.
|
|
invocation_state: Caller-provided state/context that was passed to the agent when it was invoked.
|
|
model_state: Runtime state for model providers (e.g., server-side response ids).
|
|
dynamic_trailing_blocks: How many trailing blocks of the last user message are rebuilt on every
|
|
call, so a provider placing cache points keeps its own ahead of them.
|
|
cancel_signal: Optional threading.Event to check for cancellation during streaming. Also
|
|
forwarded to the model so a provider can abort an in-flight request.
|
|
agent_metadata: Metadata of the invoking agent, forwarded to the model.
|
|
**kwargs: Additional keyword arguments for future extensibility.
|
|
|
|
Yields:
|
|
The reason for stopping, the final message, and the usage metrics
|
|
"""
|
|
logger.debug("model=<%s> | streaming messages", model)
|
|
|
|
messages = _normalize_messages(messages)
|
|
# Whitelist only role and content before sending to the model provider.
|
|
# This ensures metadata (and any future non-model fields) never leak to providers.
|
|
messages = [Message(role=msg["role"], content=msg["content"]) for msg in messages]
|
|
start_time = time.time()
|
|
|
|
chunks = model.stream(
|
|
messages,
|
|
tool_specs if tool_specs else None,
|
|
system_prompt,
|
|
tool_choice=tool_choice,
|
|
system_prompt_content=system_prompt_content,
|
|
invocation_state=invocation_state,
|
|
model_state=model_state,
|
|
cancel_signal=cancel_signal,
|
|
agent_metadata=agent_metadata,
|
|
# Omitted when zero, so an ordinary call's arguments are unchanged.
|
|
**({"dynamic_trailing_blocks": dynamic_trailing_blocks} if dynamic_trailing_blocks else {}),
|
|
)
|
|
|
|
async for event in process_stream(chunks, start_time, cancel_signal):
|
|
yield event
|