wehub-resource-sync
107 lines
3.4 KiB
Python
107 lines
3.4 KiB
Python
from unittest.mock import Mock
|
|
|
|
import aisuite as ai
|
|
from aisuite.framework.message import ChatCompletionMessageToolCall, Function, Message
|
|
from tests.agents.helpers import chat_response
|
|
|
|
|
|
def tool_call(name, arguments, call_id):
|
|
return ChatCompletionMessageToolCall(
|
|
id=call_id,
|
|
type="function",
|
|
function=Function(name=name, arguments=arguments),
|
|
)
|
|
|
|
|
|
def test_parent_agent_subagent_flow_is_reconstructable_from_jsonl_store(tmp_path):
|
|
client = ai.Client()
|
|
provider = Mock()
|
|
first_response = chat_response(None)
|
|
first_response.choices[0].message = Message(
|
|
role="assistant",
|
|
tool_calls=[
|
|
tool_call(
|
|
"research_topic",
|
|
'{"input": "agent observability"}',
|
|
"call_research",
|
|
)
|
|
],
|
|
)
|
|
provider.chat_completions_create.side_effect = [
|
|
first_response,
|
|
chat_response("research notes"),
|
|
chat_response("final recommendation"),
|
|
]
|
|
client.providers["openai"] = provider
|
|
trace_file = tmp_path / "events.jsonl"
|
|
sink = ai.tracing.LocalTraceSink(trace_file)
|
|
researcher = ai.Agent(name="researcher", model="openai:gpt-4o-mini")
|
|
writer = ai.Agent(
|
|
name="writer",
|
|
model="openai:gpt-4o",
|
|
tools=[ai.agent_tool(researcher, name="research_topic")],
|
|
)
|
|
|
|
parent = ai.Runner.run_sync(
|
|
writer,
|
|
"Write a recommendation",
|
|
client=client,
|
|
run_name="recommendation",
|
|
group_id="workflow_1",
|
|
tags=["integration"],
|
|
metadata={"request_id": "req_1"},
|
|
trace_sinks=[sink],
|
|
)
|
|
|
|
runs = ai.tracing.JsonlTraceStore(trace_file).list_runs()
|
|
parent_run = next(run for run in runs if run["trace_id"] == parent.trace_id)
|
|
child_run = next(run for run in runs if run["agent_name"] == "researcher")
|
|
|
|
assert parent_run["final_output"] == "final recommendation"
|
|
assert child_run["final_output"] == "research notes"
|
|
assert child_run["parent_run_id"] == parent.trace_id
|
|
assert {run["group_id"] for run in runs} == {"workflow_1"}
|
|
assert [
|
|
call.args[0] for call in provider.chat_completions_create.call_args_list
|
|
] == [
|
|
"gpt-4o",
|
|
"gpt-4o-mini",
|
|
"gpt-4o",
|
|
]
|
|
|
|
|
|
def test_denied_tool_policy_flow_emits_denial_event():
|
|
client = ai.Client()
|
|
provider = Mock()
|
|
first_response = chat_response(None)
|
|
first_response.choices[0].message = Message(
|
|
role="assistant",
|
|
tool_calls=[tool_call("lookup_secret", '{"key": "token"}', "call_secret")],
|
|
)
|
|
provider.chat_completions_create.side_effect = [
|
|
first_response,
|
|
chat_response("I cannot access that."),
|
|
]
|
|
client.providers["openai"] = provider
|
|
sink = ai.tracing.InMemoryTraceSink()
|
|
|
|
def lookup_secret(key: str) -> str:
|
|
return f"secret: {key}"
|
|
|
|
def deny_policy(context):
|
|
return ai.ToolPolicyDecision(allowed=False, reason="blocked")
|
|
|
|
result = ai.Runner.run_sync(
|
|
ai.Agent(name="assistant", model="openai:gpt-4o", tools=[lookup_secret]),
|
|
"Find the token",
|
|
client=client,
|
|
tool_policy=deny_policy,
|
|
trace_sinks=[sink],
|
|
)
|
|
|
|
assert result.final_output == "I cannot access that."
|
|
assert "tool.denied" in [event.event_type for event in sink.events]
|
|
denied = next(event for event in sink.events if event.event_type == "tool.denied")
|
|
assert denied.data["tool_name"] == "lookup_secret"
|
|
assert denied.data["reason"] == "blocked"
|