98 lines
3.3 KiB
Python
98 lines
3.3 KiB
Python
"""Visibility rules for transport-only tools used inside the agent graph.
|
|
|
|
The graph/checkpoint must retain these calls because they are part of the
|
|
model's causal history. Public serializers use the projection helpers below
|
|
without mutating that history.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable, Mapping
|
|
from typing import Any
|
|
|
|
INTERNAL_TOOL_NAMES = frozenset({"begin_write_file"})
|
|
|
|
|
|
def is_internal_tool_name(name: object) -> bool:
|
|
"""Return whether *name* is an exact match for a transport-only tool."""
|
|
|
|
return isinstance(name, str) and name in INTERNAL_TOOL_NAMES
|
|
|
|
|
|
def _serialized_tool_name(tool_call: object) -> object:
|
|
if not isinstance(tool_call, Mapping):
|
|
return None
|
|
name = tool_call.get("name")
|
|
if name is not None:
|
|
return name
|
|
function = tool_call.get("function")
|
|
return function.get("name") if isinstance(function, Mapping) else None
|
|
|
|
|
|
def internal_tool_call_ids(messages: Iterable[object]) -> set[str]:
|
|
"""Collect ids belonging to exact internal calls in serialized messages."""
|
|
|
|
ids: set[str] = set()
|
|
for message in messages:
|
|
if not isinstance(message, Mapping):
|
|
continue
|
|
tool_calls = message.get("tool_calls")
|
|
if not isinstance(tool_calls, (list, tuple)):
|
|
continue
|
|
for tool_call in tool_calls:
|
|
if not is_internal_tool_name(_serialized_tool_name(tool_call)):
|
|
continue
|
|
call_id = tool_call.get("id") if isinstance(tool_call, Mapping) else None
|
|
if call_id:
|
|
ids.add(str(call_id))
|
|
return ids
|
|
|
|
|
|
def project_outward_messages(
|
|
messages: Iterable[Any],
|
|
*,
|
|
known_internal_tool_call_ids: Iterable[str] = (),
|
|
) -> list[Any]:
|
|
"""Remove transport-only call/result pairs from serialized messages."""
|
|
|
|
materialized = list(messages)
|
|
internal_ids = {str(call_id) for call_id in known_internal_tool_call_ids if call_id}
|
|
projected: list[Any] = []
|
|
|
|
for message in materialized:
|
|
if not isinstance(message, Mapping):
|
|
projected.append(message)
|
|
continue
|
|
|
|
kind = str(message.get("type") or message.get("role") or "").lower()
|
|
if kind in {"ai", "assistant", "aimessage", "aimessagechunk"}:
|
|
tool_calls = message.get("tool_calls")
|
|
if isinstance(tool_calls, (list, tuple)) and tool_calls:
|
|
internal_ids.update(internal_tool_call_ids([message]))
|
|
visible_calls = [
|
|
tool_call
|
|
for tool_call in tool_calls
|
|
if not is_internal_tool_name(_serialized_tool_name(tool_call))
|
|
]
|
|
if not visible_calls:
|
|
continue
|
|
if len(visible_calls) != len(tool_calls):
|
|
message = {**message, "tool_calls": visible_calls}
|
|
elif kind in {"tool", "toolmessage"}:
|
|
call_id = str(message.get("tool_call_id") or "")
|
|
if is_internal_tool_name(message.get("name")) or call_id in internal_ids:
|
|
internal_ids.discard(call_id)
|
|
continue
|
|
|
|
projected.append(message)
|
|
|
|
return projected
|
|
|
|
|
|
__all__ = [
|
|
"INTERNAL_TOOL_NAMES",
|
|
"internal_tool_call_ids",
|
|
"is_internal_tool_name",
|
|
"project_outward_messages",
|
|
]
|