wehub-resource-sync
241 lines
8.1 KiB
Python
241 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, cast
|
|
|
|
from openai.types.responses import ResponseFunctionToolCall
|
|
|
|
from agents import Agent, ModelSettings, function_tool, tool_namespace
|
|
from agents.items import ToolCallItem, ToolCallOutputItem, ToolSearchCallItem, ToolSearchOutputItem
|
|
from agents.run_internal.run_loop import maybe_reset_tool_choice
|
|
from agents.run_internal.run_steps import ProcessedResponse, ToolRunFunction
|
|
from agents.run_internal.tool_use_tracker import (
|
|
AgentToolUseTracker,
|
|
hydrate_tool_use_tracker,
|
|
serialize_tool_use_tracker,
|
|
)
|
|
|
|
from .test_responses import get_function_tool_call
|
|
|
|
|
|
def test_tool_use_tracker_as_serializable_uses_agent_map_or_runtime_snapshot() -> None:
|
|
tracker = AgentToolUseTracker()
|
|
tracker.agent_map = {"agent-a": {"tool-b", "tool-a"}}
|
|
assert tracker.as_serializable() == {"agent-a": ["tool-a", "tool-b"]}
|
|
|
|
runtime_tracker = AgentToolUseTracker()
|
|
agent = Agent(name="runtime-agent")
|
|
runtime_tracker.add_tool_use(agent, ["beta", "alpha"])
|
|
assert runtime_tracker.as_serializable() == {"runtime-agent": ["alpha", "beta"]}
|
|
|
|
|
|
def test_tool_use_tracker_from_and_serialize_snapshots() -> None:
|
|
hydrated = AgentToolUseTracker.from_serializable({"agent": ["tool-2", "tool-1"]})
|
|
assert hydrated.agent_map == {"agent": {"tool-1", "tool-2"}}
|
|
|
|
runtime_tracker = AgentToolUseTracker()
|
|
agent = Agent(name="serialize-agent")
|
|
runtime_tracker.add_tool_use(agent, ["one"])
|
|
runtime_tracker.add_tool_use(agent, ["two"])
|
|
assert serialize_tool_use_tracker(runtime_tracker) == {"serialize-agent": ["one", "two"]}
|
|
|
|
|
|
def test_serialize_and_hydrate_tool_use_tracker_preserves_duplicate_agent_identity() -> None:
|
|
second = Agent(name="duplicate")
|
|
first = Agent(name="duplicate", handoffs=[second])
|
|
second.handoffs = [first]
|
|
|
|
tracker = AgentToolUseTracker()
|
|
tracker.add_tool_use(second, ["approval_tool"])
|
|
|
|
snapshot = serialize_tool_use_tracker(tracker, starting_agent=first)
|
|
assert snapshot == {"duplicate#2": ["approval_tool"]}
|
|
|
|
class _RunState:
|
|
def get_tool_use_tracker_snapshot(self) -> dict[str, list[str]]:
|
|
return snapshot
|
|
|
|
hydrated = AgentToolUseTracker()
|
|
hydrate_tool_use_tracker(
|
|
tool_use_tracker=hydrated,
|
|
run_state=_RunState(),
|
|
starting_agent=first,
|
|
)
|
|
|
|
assert hydrated.agent_to_tools == [(second, ["approval_tool"])]
|
|
|
|
|
|
def test_tool_use_tracker_handles_literal_suffix_names_without_collision() -> None:
|
|
literal_suffix = Agent(name="sandbox#2")
|
|
first = Agent(name="sandbox", handoffs=[literal_suffix])
|
|
second = Agent(name="sandbox")
|
|
literal_suffix.handoffs = [first, second]
|
|
first.handoffs = [literal_suffix, second]
|
|
second.handoffs = [first, literal_suffix]
|
|
|
|
tracker = AgentToolUseTracker()
|
|
tracker.add_tool_use(second, ["approval_tool"])
|
|
|
|
snapshot = serialize_tool_use_tracker(tracker, starting_agent=first)
|
|
assert snapshot == {"sandbox#3": ["approval_tool"]}
|
|
|
|
class _RunState:
|
|
def get_tool_use_tracker_snapshot(self) -> dict[str, list[str]]:
|
|
return snapshot
|
|
|
|
hydrated = AgentToolUseTracker()
|
|
hydrate_tool_use_tracker(
|
|
tool_use_tracker=hydrated,
|
|
run_state=_RunState(),
|
|
starting_agent=first,
|
|
)
|
|
|
|
assert hydrated.agent_to_tools == [(second, ["approval_tool"])]
|
|
|
|
|
|
def test_record_used_tools_uses_trace_names_for_namespaced_and_deferred_functions() -> None:
|
|
agent = Agent(name="tracked-agent")
|
|
tracker = AgentToolUseTracker()
|
|
|
|
billing_tool = tool_namespace(
|
|
name="billing",
|
|
description="Billing tools",
|
|
tools=[function_tool(lambda customer_id: customer_id, name_override="lookup_account")],
|
|
)[0]
|
|
deferred_tool = function_tool(
|
|
lambda city: city,
|
|
name_override="get_weather",
|
|
defer_loading=True,
|
|
)
|
|
|
|
tracker.record_used_tools(
|
|
agent,
|
|
[
|
|
ToolRunFunction(
|
|
function_tool=billing_tool,
|
|
tool_call=cast(
|
|
ResponseFunctionToolCall,
|
|
get_function_tool_call("lookup_account", namespace="billing"),
|
|
),
|
|
),
|
|
ToolRunFunction(
|
|
function_tool=deferred_tool,
|
|
tool_call=cast(
|
|
ResponseFunctionToolCall,
|
|
get_function_tool_call("get_weather", namespace="get_weather"),
|
|
),
|
|
),
|
|
],
|
|
)
|
|
|
|
assert tracker.as_serializable() == {"tracked-agent": ["billing.lookup_account", "get_weather"]}
|
|
|
|
|
|
def test_record_processed_response_ignores_hosted_tool_search_for_resets():
|
|
agent = Agent(name="tracked-agent")
|
|
tracker = AgentToolUseTracker()
|
|
processed_response = ProcessedResponse(
|
|
new_items=[
|
|
ToolSearchCallItem(agent=agent, raw_item={"type": "tool_search_call"}),
|
|
ToolSearchOutputItem(agent=agent, raw_item={"type": "tool_search_output"}),
|
|
],
|
|
handoffs=[],
|
|
functions=[],
|
|
computer_actions=[],
|
|
local_shell_calls=[],
|
|
shell_calls=[],
|
|
apply_patch_calls=[],
|
|
tools_used=["tool_search", "tool_search"],
|
|
mcp_approval_requests=[],
|
|
interruptions=[],
|
|
)
|
|
|
|
tracker.record_processed_response(agent, processed_response)
|
|
|
|
assert tracker.has_used_tools(agent) is False
|
|
assert tracker.as_serializable() == {}
|
|
assert maybe_reset_tool_choice(
|
|
agent, tracker, ModelSettings(tool_choice="required")
|
|
).tool_choice == ("required")
|
|
|
|
|
|
def test_record_processed_response_keeps_function_named_tool_search():
|
|
agent = Agent(name="tracked-agent")
|
|
tracker = AgentToolUseTracker()
|
|
processed_response = ProcessedResponse(
|
|
new_items=[
|
|
ToolSearchCallItem(agent=agent, raw_item={"type": "tool_search_call"}),
|
|
ToolSearchOutputItem(agent=agent, raw_item={"type": "tool_search_output"}),
|
|
ToolCallItem(
|
|
raw_item=cast(ResponseFunctionToolCall, get_function_tool_call("tool_search")),
|
|
agent=agent,
|
|
),
|
|
],
|
|
handoffs=[],
|
|
functions=[],
|
|
computer_actions=[],
|
|
local_shell_calls=[],
|
|
shell_calls=[],
|
|
apply_patch_calls=[],
|
|
tools_used=["tool_search", "tool_search", "tool_search"],
|
|
mcp_approval_requests=[],
|
|
interruptions=[],
|
|
)
|
|
|
|
tracker.record_processed_response(agent, processed_response)
|
|
|
|
assert tracker.as_serializable() == {"tracked-agent": ["tool_search"]}
|
|
|
|
|
|
def test_record_processed_response_counts_output_only_tools_without_shifting_names() -> None:
|
|
agent = Agent(name="tracked-agent")
|
|
tracker = AgentToolUseTracker()
|
|
processed_response = ProcessedResponse(
|
|
new_items=[
|
|
ToolCallOutputItem(
|
|
agent=agent,
|
|
raw_item=cast(
|
|
Any,
|
|
{"type": "shell_call_output", "call_id": "shell-1", "output": []},
|
|
),
|
|
output=[],
|
|
),
|
|
ToolCallItem(
|
|
raw_item=cast(ResponseFunctionToolCall, get_function_tool_call("lookup_account")),
|
|
agent=agent,
|
|
),
|
|
],
|
|
handoffs=[],
|
|
functions=[],
|
|
computer_actions=[],
|
|
local_shell_calls=[],
|
|
shell_calls=[],
|
|
apply_patch_calls=[],
|
|
tools_used=["shell", "lookup_account"],
|
|
mcp_approval_requests=[],
|
|
interruptions=[],
|
|
)
|
|
|
|
tracker.record_processed_response(agent, processed_response)
|
|
|
|
assert tracker.has_used_tools(agent)
|
|
assert tracker.as_serializable() == {"tracked-agent": ["lookup_account", "shell"]}
|
|
|
|
|
|
def test_hydrate_tool_use_tracker_skips_unknown_agents() -> None:
|
|
class _RunState:
|
|
def get_tool_use_tracker_snapshot(self) -> dict[str, list[str]]:
|
|
return {"known-agent": ["known_tool"], "missing-agent": ["missing_tool"]}
|
|
|
|
starting_agent = Agent(name="known-agent")
|
|
tracker = AgentToolUseTracker()
|
|
|
|
hydrate_tool_use_tracker(
|
|
tool_use_tracker=tracker,
|
|
run_state=_RunState(),
|
|
starting_agent=starting_agent,
|
|
)
|
|
|
|
assert tracker.has_used_tools(starting_agent)
|
|
assert tracker.as_serializable() == {"known-agent": ["known_tool"]}
|
|
assert "missing-agent" not in tracker.as_serializable()
|