Files
2026-07-13 13:12:00 +08:00

114 lines
3.5 KiB
Python

"""Helpers for reading session-item wire shapes used by bench drivers."""
from __future__ import annotations
from collections.abc import Iterable, Mapping
from typing import Any
def item_data(item: Mapping[str, Any]) -> Mapping[str, Any]:
"""Return the nested item payload when one is present."""
data = item.get("data")
return data if isinstance(data, Mapping) else item
def item_type(item: Mapping[str, Any]) -> str | None:
"""Return an item's type from either supported envelope shape."""
data = item_data(item)
value = item.get("type") or data.get("type")
return value if isinstance(value, str) else None
def item_role(item: Mapping[str, Any]) -> str | None:
"""Return an item's role from either supported envelope shape."""
data = item_data(item)
value = data.get("role") or item.get("role")
return value if isinstance(value, str) else None
def assistant_text(
items: Mapping[str, Any] | Iterable[Mapping[str, Any]], *, separator: str = "\n"
) -> str:
"""Concatenate assistant text blocks from one item or an item collection."""
candidates = [items] if isinstance(items, Mapping) else items
output: list[str] = []
for item in candidates:
data = item_data(item)
if item_role(item) != "assistant":
continue
content = data.get("content")
if not isinstance(content, list):
continue
output.extend(
block["text"]
for block in content
if isinstance(block, Mapping)
and block.get("type") in (None, "output_text", "text")
and isinstance(block.get("text"), str)
)
return separator.join(text for text in output if text)
def function_calls(items: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]:
"""Return normalized function-call records from session items."""
calls: list[dict[str, Any]] = []
for item in items:
if item_type(item) != "function_call":
continue
data = item_data(item)
calls.append(
{
"call_id": data.get("call_id"),
"name": data.get("name"),
"arguments": data.get("arguments"),
}
)
return calls
def tool_output_states(
items: Iterable[Mapping[str, Any]], *, deny_marker: str
) -> tuple[bool, bool]:
"""Return whether tool outputs show an allowed or denied call."""
allowed = False
denied = False
for item in items:
if item_type(item) != "function_call_output":
continue
data = item_data(item)
output = str(data.get("output", ""))
if data.get("status") == "blocked" or deny_marker in output:
denied = True
else:
allowed = True
return allowed, denied
def contains_user_text(items: Iterable[Mapping[str, Any]], marker: str) -> bool:
"""Return whether a user message contains *marker*."""
for item in items:
data = item_data(item)
if item_type(item) != "message" or data.get("role") != "user":
continue
content = data.get("content")
if not isinstance(content, list):
continue
if any(
marker in block.get("text", "")
for block in content
if isinstance(block, Mapping) and isinstance(block.get("text"), str)
):
return True
return False
__all__ = [
"assistant_text",
"contains_user_text",
"function_calls",
"item_data",
"item_role",
"item_type",
"tool_output_states",
]