wehub-resource-sync
2044 lines
64 KiB
Python
2044 lines
64 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Any, cast
|
|
|
|
import httpx
|
|
import pytest
|
|
from openai import APIConnectionError, BadRequestError
|
|
from openai.types.responses import (
|
|
ResponseCompletedEvent,
|
|
ResponseErrorEvent,
|
|
ResponseFailedEvent,
|
|
ResponseFunctionToolCall,
|
|
ResponseIncompleteEvent,
|
|
)
|
|
from openai.types.responses.response_reasoning_item import ResponseReasoningItem, Summary
|
|
from typing_extensions import TypedDict
|
|
|
|
from agents import (
|
|
Agent,
|
|
GuardrailFunctionOutput,
|
|
Handoff,
|
|
HandoffInputData,
|
|
InputGuardrail,
|
|
InputGuardrailTripwireTriggered,
|
|
MaxTurnsExceeded,
|
|
ModelBehaviorError,
|
|
ModelRetrySettings,
|
|
ModelSettings,
|
|
OpenAIResponsesWSModel,
|
|
OutputGuardrail,
|
|
OutputGuardrailTripwireTriggered,
|
|
RunContextWrapper,
|
|
Runner,
|
|
UserError,
|
|
function_tool,
|
|
handoff,
|
|
retry_policies,
|
|
)
|
|
from agents.items import RunItem, ToolApprovalItem, TResponseInputItem
|
|
from agents.memory.openai_conversations_session import OpenAIConversationsSession
|
|
from agents.run import RunConfig
|
|
from agents.run_internal import run_loop
|
|
from agents.run_internal.run_loop import QueueCompleteSentinel
|
|
from agents.stream_events import AgentUpdatedStreamEvent, RawResponsesStreamEvent, StreamEvent
|
|
from agents.usage import Usage
|
|
|
|
from .fake_model import FakeModel, get_response_obj
|
|
from .test_responses import (
|
|
get_final_output_message,
|
|
get_function_tool,
|
|
get_function_tool_call,
|
|
get_handoff_tool_call,
|
|
get_text_input_item,
|
|
get_text_message,
|
|
)
|
|
from .utils.hitl import (
|
|
consume_stream,
|
|
make_model_and_agent,
|
|
queue_function_call_and_text,
|
|
resume_streamed_after_first_approval,
|
|
)
|
|
from .utils.simple_session import SimpleListSession
|
|
|
|
|
|
def _conversation_locked_error() -> BadRequestError:
|
|
request = httpx.Request("POST", "https://example.com")
|
|
response = httpx.Response(
|
|
400,
|
|
request=request,
|
|
json={"error": {"code": "conversation_locked", "message": "locked"}},
|
|
)
|
|
error = BadRequestError(
|
|
"locked",
|
|
response=response,
|
|
body={"error": {"code": "conversation_locked"}},
|
|
)
|
|
error.code = "conversation_locked"
|
|
return error
|
|
|
|
|
|
def _find_reasoning_input_item(
|
|
items: str | list[TResponseInputItem] | Any,
|
|
) -> dict[str, Any] | None:
|
|
if not isinstance(items, list):
|
|
return None
|
|
for item in items:
|
|
if isinstance(item, dict) and item.get("type") == "reasoning":
|
|
return cast(dict[str, Any], item)
|
|
return None
|
|
|
|
|
|
def _ws_terminal_response_frame(event_type: str, response_id: str, sequence_number: int) -> str:
|
|
response = get_response_obj([get_text_message("partial final")], response_id=response_id)
|
|
return json.dumps(
|
|
{
|
|
"type": event_type,
|
|
"response": response.model_dump(),
|
|
"sequence_number": sequence_number,
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_simple_first_run():
|
|
model = FakeModel()
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
model.set_next_output([get_text_message("first")])
|
|
|
|
result = Runner.run_streamed(agent, input="test")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.input == "test"
|
|
assert len(result.new_items) == 1, "exactly one item should be generated"
|
|
assert result.final_output == "first"
|
|
assert len(result.raw_responses) == 1, "exactly one model response should be generated"
|
|
assert result.raw_responses[0].output == [get_text_message("first")]
|
|
assert result.last_agent == agent
|
|
|
|
assert len(result.to_input_list()) == 2, "should have original input and generated item"
|
|
|
|
model.set_next_output([get_text_message("second")])
|
|
|
|
result = Runner.run_streamed(
|
|
agent, input=[get_text_input_item("message"), get_text_input_item("another_message")]
|
|
)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert len(result.new_items) == 1, "exactly one item should be generated"
|
|
assert result.final_output == "second"
|
|
assert len(result.raw_responses) == 1, "exactly one model response should be generated"
|
|
assert len(result.to_input_list()) == 3, "should have original input and generated item"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_tool_not_found_behavior_returns_error_to_model() -> None:
|
|
model = FakeModel()
|
|
agent = Agent(name="test", model=model)
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("missing_tool", "{}", call_id="call_missing")],
|
|
[get_text_message("recovered")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(
|
|
agent,
|
|
input="start",
|
|
run_config=RunConfig(tool_not_found_behavior="return_error_to_model"),
|
|
)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == "recovered"
|
|
second_turn_input = model.last_turn_args["input"]
|
|
assert isinstance(second_turn_input, list)
|
|
assert {
|
|
item.get("call_id"): item.get("output")
|
|
for item in second_turn_input
|
|
if isinstance(item, dict) and item.get("type") == "function_call_output"
|
|
} == {"call_missing": "Tool 'missing_tool' not found."}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("terminal_event_type", "terminal_event_cls"),
|
|
[
|
|
("response.incomplete", ResponseIncompleteEvent),
|
|
("response.failed", ResponseFailedEvent),
|
|
],
|
|
)
|
|
async def test_streamed_run_rejects_failed_terminal_response_payload_events(
|
|
terminal_event_type: str, terminal_event_cls: type[Any]
|
|
) -> None:
|
|
class TerminalPayloadFakeModel(FakeModel):
|
|
async def stream_response(
|
|
self,
|
|
system_instructions,
|
|
input,
|
|
model_settings,
|
|
tools,
|
|
output_schema,
|
|
handoffs,
|
|
tracing,
|
|
*,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
prompt=None,
|
|
):
|
|
self.last_turn_args = {
|
|
"system_instructions": system_instructions,
|
|
"input": input,
|
|
"model_settings": model_settings,
|
|
"tools": tools,
|
|
"output_schema": output_schema,
|
|
"previous_response_id": previous_response_id,
|
|
"conversation_id": conversation_id,
|
|
}
|
|
if self.first_turn_args is None:
|
|
self.first_turn_args = self.last_turn_args.copy()
|
|
|
|
response = get_response_obj(
|
|
[get_text_message("partial final")], response_id="resp-partial"
|
|
)
|
|
yield terminal_event_cls(
|
|
type=terminal_event_type,
|
|
response=response,
|
|
sequence_number=0,
|
|
)
|
|
|
|
model = TerminalPayloadFakeModel()
|
|
agent = Agent(name="test", model=model)
|
|
|
|
result = Runner.run_streamed(agent, input="test")
|
|
stream_events: list[StreamEvent] = []
|
|
with pytest.raises(ModelBehaviorError, match=terminal_event_type):
|
|
async for event in result.stream_events():
|
|
stream_events.append(event)
|
|
|
|
assert len(stream_events) == 2
|
|
assert isinstance(stream_events[0], AgentUpdatedStreamEvent)
|
|
assert isinstance(stream_events[1], RawResponsesStreamEvent)
|
|
assert stream_events[1].data.type == terminal_event_type
|
|
assert result.final_output is None
|
|
assert result.raw_responses == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_run_rejects_response_error_terminal_event() -> None:
|
|
class TerminalErrorFakeModel(FakeModel):
|
|
async def stream_response(
|
|
self,
|
|
system_instructions,
|
|
input,
|
|
model_settings,
|
|
tools,
|
|
output_schema,
|
|
handoffs,
|
|
tracing,
|
|
*,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
prompt=None,
|
|
):
|
|
self.last_turn_args = {
|
|
"system_instructions": system_instructions,
|
|
"input": input,
|
|
"model_settings": model_settings,
|
|
"tools": tools,
|
|
"output_schema": output_schema,
|
|
"previous_response_id": previous_response_id,
|
|
"conversation_id": conversation_id,
|
|
}
|
|
if self.first_turn_args is None:
|
|
self.first_turn_args = self.last_turn_args.copy()
|
|
|
|
yield ResponseErrorEvent(
|
|
type="error",
|
|
code="invalid_request_error",
|
|
message="bad request",
|
|
param=None,
|
|
sequence_number=0,
|
|
)
|
|
|
|
model = TerminalErrorFakeModel()
|
|
agent = Agent(name="test", model=model)
|
|
|
|
result = Runner.run_streamed(agent, input="test")
|
|
stream_events: list[StreamEvent] = []
|
|
with pytest.raises(ModelBehaviorError, match="error"):
|
|
async for event in result.stream_events():
|
|
stream_events.append(event)
|
|
|
|
assert len(stream_events) == 2
|
|
assert isinstance(stream_events[0], AgentUpdatedStreamEvent)
|
|
assert isinstance(stream_events[1], RawResponsesStreamEvent)
|
|
assert stream_events[1].data.type == "error"
|
|
assert stream_events[1].data.code == "invalid_request_error"
|
|
assert stream_events[1].data.message == "bad request"
|
|
assert result.final_output is None
|
|
assert result.raw_responses == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_run_exposes_request_id_on_raw_responses() -> None:
|
|
class RequestIdTerminalFakeModel(FakeModel):
|
|
async def stream_response(
|
|
self,
|
|
system_instructions,
|
|
input,
|
|
model_settings,
|
|
tools,
|
|
output_schema,
|
|
handoffs,
|
|
tracing,
|
|
*,
|
|
previous_response_id=None,
|
|
conversation_id=None,
|
|
prompt=None,
|
|
):
|
|
response = get_response_obj(
|
|
[get_text_message("partial final")], response_id="resp-partial"
|
|
)
|
|
response._request_id = "req_streamed_result_123"
|
|
yield ResponseCompletedEvent(
|
|
type="response.completed",
|
|
response=response,
|
|
sequence_number=0,
|
|
)
|
|
|
|
model = RequestIdTerminalFakeModel()
|
|
agent = Agent(name="test", model=model)
|
|
|
|
result = Runner.run_streamed(agent, input="test")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert len(result.raw_responses) == 1
|
|
assert result.raw_responses[0].request_id == "req_streamed_result_123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_run_preserves_request_usage_entries_after_retry() -> None:
|
|
model = FakeModel()
|
|
model.set_hardcoded_usage(
|
|
Usage(
|
|
requests=1,
|
|
input_tokens=10,
|
|
output_tokens=5,
|
|
total_tokens=15,
|
|
)
|
|
)
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
APIConnectionError(
|
|
message="connection error",
|
|
request=httpx.Request("POST", "https://example.com"),
|
|
),
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
model_settings=ModelSettings(
|
|
retry=ModelRetrySettings(
|
|
max_retries=1,
|
|
policy=retry_policies.network_error(),
|
|
)
|
|
),
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="test")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
usage = result.context_wrapper.usage
|
|
assert usage.requests == 2
|
|
assert len(usage.request_usage_entries) == 2
|
|
assert usage.request_usage_entries[0].total_tokens == 0
|
|
assert usage.request_usage_entries[1].input_tokens == 10
|
|
assert usage.request_usage_entries[1].output_tokens == 5
|
|
assert usage.request_usage_entries[1].total_tokens == 15
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_run_preserves_request_usage_entries_after_conversation_locked_retry() -> (
|
|
None
|
|
):
|
|
model = FakeModel()
|
|
model.set_hardcoded_usage(
|
|
Usage(
|
|
requests=1,
|
|
input_tokens=10,
|
|
output_tokens=5,
|
|
total_tokens=15,
|
|
)
|
|
)
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
_conversation_locked_error(),
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
model_settings=ModelSettings(
|
|
retry=ModelRetrySettings(
|
|
max_retries=1,
|
|
policy=retry_policies.network_error(),
|
|
)
|
|
),
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="test")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
usage = result.context_wrapper.usage
|
|
assert usage.requests == 2
|
|
assert len(usage.request_usage_entries) == 2
|
|
assert usage.request_usage_entries[0].total_tokens == 0
|
|
assert usage.request_usage_entries[1].input_tokens == 10
|
|
assert usage.request_usage_entries[1].output_tokens == 5
|
|
assert usage.request_usage_entries[1].total_tokens == 15
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("terminal_event_type", ["response.incomplete", "response.failed"])
|
|
async def test_streamed_run_rejects_failed_terminal_response_payload_events_from_ws_model(
|
|
monkeypatch, terminal_event_type: str
|
|
) -> None:
|
|
class DummyWSConnection:
|
|
def __init__(self, frames: list[str]):
|
|
self._frames = frames
|
|
self.close_code: int | None = None
|
|
|
|
async def send(self, payload: str) -> None:
|
|
return None
|
|
|
|
async def recv(self) -> str:
|
|
if not self._frames:
|
|
raise RuntimeError("No more websocket frames configured")
|
|
return self._frames.pop(0)
|
|
|
|
async def close(self) -> None:
|
|
if self.close_code is None:
|
|
self.close_code = 1000
|
|
|
|
class DummyWSClient:
|
|
def __init__(self) -> None:
|
|
self.base_url = httpx.URL("https://api.openai.com/v1/")
|
|
self.websocket_base_url = None
|
|
self.default_query: dict[str, Any] = {}
|
|
self.default_headers = {
|
|
"Authorization": "Bearer test-key",
|
|
"User-Agent": "AsyncOpenAI/Python test",
|
|
}
|
|
self.timeout: Any = None
|
|
|
|
async def _refresh_api_key(self) -> None:
|
|
return None
|
|
|
|
ws = DummyWSConnection([_ws_terminal_response_frame(terminal_event_type, "resp-ws", 1)])
|
|
model = OpenAIResponsesWSModel(model="gpt-4", openai_client=DummyWSClient()) # type: ignore[arg-type]
|
|
|
|
async def fake_open(
|
|
_ws_url: str,
|
|
_headers: dict[str, str],
|
|
*,
|
|
connect_timeout: float | None = None,
|
|
) -> DummyWSConnection:
|
|
return ws
|
|
|
|
monkeypatch.setattr(model, "_open_websocket_connection", fake_open)
|
|
|
|
agent = Agent(name="test", model=model)
|
|
result = Runner.run_streamed(agent, input="test")
|
|
stream_events: list[StreamEvent] = []
|
|
with pytest.raises(ModelBehaviorError, match=terminal_event_type):
|
|
async for event in result.stream_events():
|
|
stream_events.append(event)
|
|
|
|
assert len(stream_events) == 2
|
|
assert isinstance(stream_events[0], AgentUpdatedStreamEvent)
|
|
assert isinstance(stream_events[1], RawResponsesStreamEvent)
|
|
assert stream_events[1].data.type == terminal_event_type
|
|
assert result.final_output is None
|
|
assert result.raw_responses == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_subsequent_runs():
|
|
model = FakeModel()
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
model.set_next_output([get_text_message("third")])
|
|
|
|
result = Runner.run_streamed(agent, input="test")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.input == "test"
|
|
assert len(result.new_items) == 1, "exactly one item should be generated"
|
|
assert len(result.to_input_list()) == 2, "should have original input and generated item"
|
|
|
|
model.set_next_output([get_text_message("fourth")])
|
|
|
|
result = Runner.run_streamed(agent, input=result.to_input_list())
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert len(result.input) == 2, f"should have previous input but got {result.input}"
|
|
assert len(result.new_items) == 1, "exactly one item should be generated"
|
|
assert result.final_output == "fourth"
|
|
assert len(result.raw_responses) == 1, "exactly one model response should be generated"
|
|
assert result.raw_responses[0].output == [get_text_message("fourth")]
|
|
assert result.last_agent == agent
|
|
assert len(result.to_input_list()) == 3, "should have original input and generated items"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_call_runs():
|
|
model = FakeModel()
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[get_function_tool("foo", "tool_result")],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
# First turn: a message and tool call
|
|
[get_text_message("a_message"), get_function_tool_call("foo", json.dumps({"a": "b"}))],
|
|
# Second turn: text message
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == "done"
|
|
assert len(result.raw_responses) == 2, (
|
|
"should have two responses: the first which produces a tool call, and the second which"
|
|
"handles the tool result"
|
|
)
|
|
|
|
assert len(result.to_input_list()) == 5, (
|
|
"should have five inputs: the original input, the message, the tool call, the tool result "
|
|
"and the done message"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_parallel_tool_call_with_cancelled_sibling_reaches_final_output() -> None:
|
|
async def _ok_tool() -> str:
|
|
return "ok"
|
|
|
|
async def _cancel_tool() -> str:
|
|
raise asyncio.CancelledError("tool-cancelled")
|
|
|
|
model = FakeModel()
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[
|
|
function_tool(_ok_tool, name_override="ok_tool"),
|
|
function_tool(_cancel_tool, name_override="cancel_tool"),
|
|
],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[
|
|
get_function_tool_call("ok_tool", "{}", call_id="call_ok"),
|
|
get_function_tool_call("cancel_tool", "{}", call_id="call_cancel"),
|
|
],
|
|
[get_text_message("final answer")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
await consume_stream(result)
|
|
|
|
assert result.final_output == "final answer"
|
|
assert len(result.raw_responses) == 2
|
|
|
|
second_turn_input = cast(list[dict[str, Any]], model.last_turn_args["input"])
|
|
tool_outputs = [
|
|
item for item in second_turn_input if item.get("type") == "function_call_output"
|
|
]
|
|
assert tool_outputs == [
|
|
{"call_id": "call_ok", "output": "ok", "type": "function_call_output"},
|
|
{
|
|
"call_id": "call_cancel",
|
|
"output": (
|
|
"An error occurred while running the tool. Please try again. Error: tool-cancelled"
|
|
),
|
|
"type": "function_call_output",
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_single_tool_call_with_cancelled_tool_reaches_final_output() -> None:
|
|
async def _cancel_tool() -> str:
|
|
raise asyncio.CancelledError("tool-cancelled")
|
|
|
|
model = FakeModel()
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[function_tool(_cancel_tool, name_override="cancel_tool")],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("cancel_tool", "{}", call_id="call_cancel")],
|
|
[get_text_message("final answer")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
await consume_stream(result)
|
|
|
|
assert result.final_output == "final answer"
|
|
assert len(result.raw_responses) == 2
|
|
|
|
second_turn_input = cast(list[dict[str, Any]], model.last_turn_args["input"])
|
|
tool_outputs = [
|
|
item for item in second_turn_input if item.get("type") == "function_call_output"
|
|
]
|
|
assert tool_outputs == [
|
|
{
|
|
"call_id": "call_cancel",
|
|
"output": (
|
|
"An error occurred while running the tool. Please try again. Error: tool-cancelled"
|
|
),
|
|
"type": "function_call_output",
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_reasoning_item_id_policy_omits_follow_up_reasoning_ids() -> None:
|
|
model = FakeModel()
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[get_function_tool("foo", "tool_result")],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[
|
|
ResponseReasoningItem(
|
|
id="rs_stream",
|
|
type="reasoning",
|
|
summary=[Summary(text="Thinking...", type="summary_text")],
|
|
),
|
|
get_function_tool_call("foo", json.dumps({"a": "b"}), call_id="call_stream"),
|
|
],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(
|
|
agent,
|
|
input="hello",
|
|
run_config=RunConfig(reasoning_item_id_policy="omit"),
|
|
)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == "done"
|
|
second_request_reasoning = _find_reasoning_input_item(model.last_turn_args.get("input"))
|
|
assert second_request_reasoning is not None
|
|
assert "id" not in second_request_reasoning
|
|
|
|
history_reasoning = _find_reasoning_input_item(result.to_input_list())
|
|
assert history_reasoning is not None
|
|
assert "id" not in history_reasoning
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_run_again_persists_tool_items_to_session():
|
|
model = FakeModel()
|
|
call_id = "call-session-run-again"
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[get_function_tool("foo", "tool_result")],
|
|
)
|
|
session = SimpleListSession()
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("foo", json.dumps({"a": "b"}), call_id=call_id)],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="user_message", session=session)
|
|
await consume_stream(result)
|
|
|
|
saved_items = await session.get_items()
|
|
assert any(
|
|
isinstance(item, dict)
|
|
and item.get("type") == "function_call"
|
|
and item.get("call_id") == call_id
|
|
for item in saved_items
|
|
)
|
|
assert any(
|
|
isinstance(item, dict)
|
|
and item.get("type") == "function_call_output"
|
|
and item.get("call_id") == call_id
|
|
for item in saved_items
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handoffs():
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
agent_2 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
agent_3 = Agent(
|
|
name="test",
|
|
model=model,
|
|
handoffs=[agent_1, agent_2],
|
|
tools=[get_function_tool("some_function", "result")],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
# First turn: a tool call
|
|
[get_function_tool_call("some_function", json.dumps({"a": "b"}))],
|
|
# Second turn: a message and a handoff
|
|
[get_text_message("a_message"), get_handoff_tool_call(agent_1)],
|
|
# Third turn: text message
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent_3, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == "done"
|
|
assert len(result.raw_responses) == 3, "should have three model responses"
|
|
assert len(result.to_input_list()) == 7, (
|
|
"should have 7 inputs: summary message, tool call, tool result, message, handoff, "
|
|
"handoff result, and done message"
|
|
)
|
|
assert result.last_agent == agent_1, "should have handed off to agent_1"
|
|
|
|
|
|
class Foo(TypedDict):
|
|
bar: str
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_structured_output():
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[get_function_tool("bar", "bar_result")],
|
|
output_type=Foo,
|
|
)
|
|
|
|
agent_2 = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[get_function_tool("foo", "foo_result")],
|
|
handoffs=[agent_1],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
# First turn: a tool call
|
|
[get_function_tool_call("foo", json.dumps({"bar": "baz"}))],
|
|
# Second turn: a message and a handoff
|
|
[get_text_message("a_message"), get_handoff_tool_call(agent_1)],
|
|
# Third turn: tool call with preamble message
|
|
[
|
|
get_text_message(json.dumps(Foo(bar="preamble"))),
|
|
get_function_tool_call("bar", json.dumps({"bar": "baz"})),
|
|
],
|
|
# Fourth turn: structured output
|
|
[get_final_output_message(json.dumps(Foo(bar="baz")))],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(
|
|
agent_2,
|
|
input=[
|
|
get_text_input_item("user_message"),
|
|
get_text_input_item("another_message"),
|
|
],
|
|
run_config=RunConfig(nest_handoff_history=True),
|
|
)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == Foo(bar="baz")
|
|
assert len(result.raw_responses) == 4, "should have four model responses"
|
|
assert len(result.to_input_list()) == 10, (
|
|
"should have input: conversation summary, function call, function call result, message, "
|
|
"handoff, handoff output, preamble message, tool call, tool call result, final output"
|
|
)
|
|
assert len(result.to_input_list(mode="normalized")) == 6, (
|
|
"should have normalized replay input: conversation summary, carried-forward message, "
|
|
"preamble message, tool call, tool call result, final output"
|
|
)
|
|
|
|
assert result.last_agent == agent_1, "should have handed off to agent_1"
|
|
assert result.final_output == Foo(bar="baz"), "should have structured output"
|
|
|
|
|
|
def remove_new_items(handoff_input_data: HandoffInputData) -> HandoffInputData:
|
|
return HandoffInputData(
|
|
input_history=handoff_input_data.input_history,
|
|
pre_handoff_items=(),
|
|
new_items=(),
|
|
run_context=handoff_input_data.run_context,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handoff_filters():
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
agent_2 = Agent(
|
|
name="test",
|
|
model=model,
|
|
handoffs=[
|
|
handoff(
|
|
agent=agent_1,
|
|
input_filter=remove_new_items,
|
|
)
|
|
],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("1"), get_text_message("2"), get_handoff_tool_call(agent_1)],
|
|
[get_text_message("last")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent_2, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == "last"
|
|
assert len(result.raw_responses) == 2, "should have two model responses"
|
|
assert len(result.to_input_list()) == 2, (
|
|
"should only have 2 inputs: orig input and last message"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamed_nested_handoff_filters_reasoning_items_from_model_input():
|
|
model = FakeModel()
|
|
delegate = Agent(
|
|
name="delegate",
|
|
model=model,
|
|
)
|
|
triage = Agent(
|
|
name="triage",
|
|
model=model,
|
|
handoffs=[delegate],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[
|
|
ResponseReasoningItem(
|
|
id="reasoning_1",
|
|
type="reasoning",
|
|
summary=[Summary(text="Thinking about a handoff.", type="summary_text")],
|
|
),
|
|
get_handoff_tool_call(delegate),
|
|
],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
|
|
captured_inputs: list[list[dict[str, Any]]] = []
|
|
|
|
def capture_model_input(data):
|
|
if isinstance(data.model_data.input, list):
|
|
captured_inputs.append(
|
|
[item for item in data.model_data.input if isinstance(item, dict)]
|
|
)
|
|
return data.model_data
|
|
|
|
result = Runner.run_streamed(
|
|
triage,
|
|
input="user_message",
|
|
run_config=RunConfig(
|
|
nest_handoff_history=True,
|
|
call_model_input_filter=capture_model_input,
|
|
),
|
|
)
|
|
await consume_stream(result)
|
|
|
|
assert result.final_output == "done"
|
|
assert len(captured_inputs) >= 2
|
|
handoff_input = captured_inputs[1]
|
|
handoff_input_types = [
|
|
item["type"] for item in handoff_input if isinstance(item.get("type"), str)
|
|
]
|
|
assert "reasoning" not in handoff_input_types
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_input_filter_supported():
|
|
# DO NOT rename this without updating pyproject.toml
|
|
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
|
|
async def on_invoke_handoff(_ctx: RunContextWrapper[Any], _input: str) -> Agent[Any]:
|
|
return agent_1
|
|
|
|
async def async_input_filter(data: HandoffInputData) -> HandoffInputData:
|
|
return data # pragma: no cover
|
|
|
|
agent_2 = Agent[None](
|
|
name="test",
|
|
model=model,
|
|
handoffs=[
|
|
Handoff(
|
|
tool_name=Handoff.default_tool_name(agent_1),
|
|
tool_description=Handoff.default_tool_description(agent_1),
|
|
input_json_schema={},
|
|
on_invoke_handoff=on_invoke_handoff,
|
|
agent_name=agent_1.name,
|
|
input_filter=async_input_filter,
|
|
)
|
|
],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("1"), get_text_message("2"), get_handoff_tool_call(agent_1)],
|
|
[get_text_message("last")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent_2, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalid_input_filter_fails():
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
|
|
async def on_invoke_handoff(_ctx: RunContextWrapper[Any], _input: str) -> Agent[Any]:
|
|
return agent_1
|
|
|
|
def invalid_input_filter(data: HandoffInputData) -> HandoffInputData:
|
|
# Purposely returning a string to simulate invalid output
|
|
return "foo" # type: ignore
|
|
|
|
agent_2 = Agent[None](
|
|
name="test",
|
|
model=model,
|
|
handoffs=[
|
|
Handoff(
|
|
tool_name=Handoff.default_tool_name(agent_1),
|
|
tool_description=Handoff.default_tool_description(agent_1),
|
|
input_json_schema={},
|
|
on_invoke_handoff=on_invoke_handoff,
|
|
agent_name=agent_1.name,
|
|
input_filter=invalid_input_filter,
|
|
)
|
|
],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("1"), get_text_message("2"), get_handoff_tool_call(agent_1)],
|
|
[get_text_message("last")],
|
|
]
|
|
)
|
|
|
|
with pytest.raises(UserError):
|
|
result = Runner.run_streamed(agent_2, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_callable_input_filter_causes_error():
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
|
|
async def on_invoke_handoff(_ctx: RunContextWrapper[Any], _input: str) -> Agent[Any]:
|
|
return agent_1
|
|
|
|
agent_2 = Agent[None](
|
|
name="test",
|
|
model=model,
|
|
handoffs=[
|
|
Handoff(
|
|
tool_name=Handoff.default_tool_name(agent_1),
|
|
tool_description=Handoff.default_tool_description(agent_1),
|
|
input_json_schema={},
|
|
on_invoke_handoff=on_invoke_handoff,
|
|
agent_name=agent_1.name,
|
|
# Purposely ignoring the type error here to simulate invalid input
|
|
input_filter="foo", # type: ignore
|
|
)
|
|
],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("1"), get_text_message("2"), get_handoff_tool_call(agent_1)],
|
|
[get_text_message("last")],
|
|
]
|
|
)
|
|
|
|
with pytest.raises(UserError):
|
|
result = Runner.run_streamed(agent_2, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handoff_on_input():
|
|
call_output: str | None = None
|
|
|
|
def on_input(_ctx: RunContextWrapper[Any], data: Foo) -> None:
|
|
nonlocal call_output
|
|
call_output = data["bar"]
|
|
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
|
|
agent_2 = Agent(
|
|
name="test",
|
|
model=model,
|
|
handoffs=[
|
|
handoff(
|
|
agent=agent_1,
|
|
on_handoff=on_input,
|
|
input_type=Foo,
|
|
)
|
|
],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[
|
|
get_text_message("1"),
|
|
get_text_message("2"),
|
|
get_handoff_tool_call(agent_1, args=json.dumps(Foo(bar="test_input"))),
|
|
],
|
|
[get_text_message("last")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent_2, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == "last"
|
|
|
|
assert call_output == "test_input", "should have called the handoff with the correct input"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_handoff_on_input():
|
|
call_output: str | None = None
|
|
|
|
async def on_input(_ctx: RunContextWrapper[Any], data: Foo) -> None:
|
|
nonlocal call_output
|
|
call_output = data["bar"]
|
|
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
|
|
agent_2 = Agent(
|
|
name="test",
|
|
model=model,
|
|
handoffs=[
|
|
handoff(
|
|
agent=agent_1,
|
|
on_handoff=on_input,
|
|
input_type=Foo,
|
|
)
|
|
],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[
|
|
get_text_message("1"),
|
|
get_text_message("2"),
|
|
get_handoff_tool_call(agent_1, args=json.dumps(Foo(bar="test_input"))),
|
|
],
|
|
[get_text_message("last")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent_2, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert result.final_output == "last"
|
|
|
|
assert call_output == "test_input", "should have called the handoff with the correct input"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_guardrail_tripwire_triggered_causes_exception_streamed():
|
|
def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], input: Any
|
|
) -> GuardrailFunctionOutput:
|
|
return GuardrailFunctionOutput(
|
|
output_info=None,
|
|
tripwire_triggered=True,
|
|
)
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
input_guardrails=[InputGuardrail(guardrail_function=guardrail_function)],
|
|
model=FakeModel(),
|
|
)
|
|
|
|
with pytest.raises(InputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_guardrail_streamed_does_not_save_assistant_message_to_session():
|
|
async def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], input: Any
|
|
) -> GuardrailFunctionOutput:
|
|
await asyncio.sleep(0.01)
|
|
return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True)
|
|
|
|
session = SimpleListSession()
|
|
|
|
model = FakeModel()
|
|
model.set_next_output([get_text_message("should_not_be_saved")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
input_guardrails=[InputGuardrail(guardrail_function=guardrail_function)],
|
|
)
|
|
|
|
with pytest.raises(InputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(agent, input="user_message", session=session)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
items = await session.get_items()
|
|
|
|
assert len(items) == 1
|
|
first_item = cast(dict[str, Any], items[0])
|
|
assert "role" in first_item
|
|
assert first_item["role"] == "user"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_guardrail_streamed_persists_user_input_for_sequential_guardrail():
|
|
def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], input: Any
|
|
) -> GuardrailFunctionOutput:
|
|
return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True)
|
|
|
|
session = SimpleListSession()
|
|
|
|
model = FakeModel()
|
|
model.set_next_output([get_text_message("should_not_be_saved")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
input_guardrails=[
|
|
InputGuardrail(guardrail_function=guardrail_function, run_in_parallel=False)
|
|
],
|
|
)
|
|
|
|
with pytest.raises(InputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(agent, input="user_message", session=session)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
items = await session.get_items()
|
|
|
|
assert len(items) == 1
|
|
first_item = cast(dict[str, Any], items[0])
|
|
assert "role" in first_item
|
|
assert first_item["role"] == "user"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_guardrail_streamed_persists_user_input_for_async_sequential_guardrail():
|
|
async def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], input: Any
|
|
) -> GuardrailFunctionOutput:
|
|
await asyncio.sleep(0)
|
|
return GuardrailFunctionOutput(output_info=None, tripwire_triggered=True)
|
|
|
|
session = SimpleListSession()
|
|
|
|
model = FakeModel()
|
|
model.set_next_output([get_text_message("should_not_be_saved")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
input_guardrails=[
|
|
InputGuardrail(guardrail_function=guardrail_function, run_in_parallel=False)
|
|
],
|
|
)
|
|
|
|
with pytest.raises(InputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(agent, input="user_message", session=session)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
items = await session.get_items()
|
|
|
|
assert len(items) == 1
|
|
first_item = cast(dict[str, Any], items[0])
|
|
assert "role" in first_item
|
|
assert first_item["role"] == "user"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_input_persistence_strips_ids_for_openai_conversation_session():
|
|
class DummyOpenAIConversationsSession(OpenAIConversationsSession):
|
|
def __init__(self) -> None:
|
|
self.saved: list[list[TResponseInputItem]] = []
|
|
|
|
async def _get_session_id(self) -> str:
|
|
return "conv_test"
|
|
|
|
async def add_items(self, items: list[TResponseInputItem]) -> None:
|
|
for item in items:
|
|
if isinstance(item, dict):
|
|
assert "id" not in item, "IDs should be stripped before saving"
|
|
assert "provider_data" not in item, (
|
|
"provider_data should be stripped before saving"
|
|
)
|
|
self.saved.append(items)
|
|
|
|
async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
|
|
return []
|
|
|
|
async def pop_item(self) -> TResponseInputItem | None:
|
|
return None
|
|
|
|
async def clear_session(self) -> None:
|
|
return None
|
|
|
|
session = DummyOpenAIConversationsSession()
|
|
|
|
model = FakeModel()
|
|
model.set_next_output([get_text_message("ok")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
|
|
run_config = RunConfig(session_input_callback=lambda existing, new: existing + new)
|
|
|
|
input_items = [
|
|
cast(
|
|
TResponseInputItem,
|
|
{
|
|
"id": "message-1",
|
|
"type": "message",
|
|
"role": "user",
|
|
"content": "hello",
|
|
"provider_data": {"model": "litellm/test"},
|
|
},
|
|
)
|
|
]
|
|
|
|
result = Runner.run_streamed(agent, input=input_items, session=session, run_config=run_config)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert session.saved, "input items should be persisted via save_result_to_session"
|
|
assert len(session.saved[0]) == 1
|
|
saved_item = session.saved[0][0]
|
|
assert isinstance(saved_item, dict)
|
|
assert "id" not in saved_item, "saved input items should not include IDs"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_input_persistence_saves_only_new_turn_input(monkeypatch: pytest.MonkeyPatch):
|
|
session = SimpleListSession()
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first")],
|
|
[get_text_message("second")],
|
|
]
|
|
)
|
|
agent = Agent(name="test", model=model)
|
|
|
|
from agents.run_internal import session_persistence as sp
|
|
|
|
real_save_result = sp.save_result_to_session
|
|
input_saves: list[list[TResponseInputItem]] = []
|
|
|
|
async def save_wrapper(
|
|
sess: Any,
|
|
original_input: Any,
|
|
new_items: list[RunItem],
|
|
run_state: Any = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
if isinstance(original_input, list) and original_input:
|
|
input_saves.append(list(original_input))
|
|
await real_save_result(sess, original_input, new_items, run_state, **kwargs)
|
|
|
|
monkeypatch.setattr(
|
|
"agents.run_internal.session_persistence.save_result_to_session", save_wrapper
|
|
)
|
|
monkeypatch.setattr("agents.run_internal.run_loop.save_result_to_session", save_wrapper)
|
|
|
|
run_config = RunConfig(session_input_callback=lambda existing, new: existing + new)
|
|
|
|
first = Runner.run_streamed(
|
|
agent, input=[get_text_input_item("hello")], session=session, run_config=run_config
|
|
)
|
|
async for _ in first.stream_events():
|
|
pass
|
|
|
|
second = Runner.run_streamed(
|
|
agent, input=[get_text_input_item("next")], session=session, run_config=run_config
|
|
)
|
|
async for _ in second.stream_events():
|
|
pass
|
|
|
|
assert len(input_saves) == 2, "each turn should persist only the turn input once"
|
|
assert all(len(saved) == 1 for saved in input_saves), (
|
|
"each persisted input should contain only the new turn items"
|
|
)
|
|
first_saved = input_saves[0][0]
|
|
second_saved = input_saves[1][0]
|
|
assert isinstance(first_saved, dict) and first_saved.get("content") == "hello"
|
|
assert isinstance(second_saved, dict) and second_saved.get("content") == "next"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_slow_input_guardrail_still_raises_exception_streamed():
|
|
async def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], input: Any
|
|
) -> GuardrailFunctionOutput:
|
|
# Simulate a slow guardrail that completes after model streaming ends.
|
|
await asyncio.sleep(0.05)
|
|
return GuardrailFunctionOutput(
|
|
output_info=None,
|
|
tripwire_triggered=True,
|
|
)
|
|
|
|
model = FakeModel()
|
|
# Ensure the model finishes streaming quickly.
|
|
model.set_next_output([get_text_message("ok")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
input_guardrails=[InputGuardrail(guardrail_function=guardrail_function)],
|
|
model=model,
|
|
)
|
|
|
|
# Even though the guardrail is slower than the model stream, the exception should still raise.
|
|
with pytest.raises(InputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_guardrail_tripwire_triggered_causes_exception_streamed():
|
|
def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any
|
|
) -> GuardrailFunctionOutput:
|
|
return GuardrailFunctionOutput(
|
|
output_info=None,
|
|
tripwire_triggered=True,
|
|
)
|
|
|
|
model = FakeModel(initial_output=[get_text_message("first_test")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
output_guardrails=[OutputGuardrail(guardrail_function=guardrail_function)],
|
|
model=model,
|
|
)
|
|
|
|
with pytest.raises(OutputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_guardrail_tripwire_raises_from_run_loop_task_before_stream_consumption():
|
|
def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any
|
|
) -> GuardrailFunctionOutput:
|
|
return GuardrailFunctionOutput(
|
|
output_info=None,
|
|
tripwire_triggered=True,
|
|
)
|
|
|
|
model = FakeModel(initial_output=[get_text_message("first_test")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
output_guardrails=[OutputGuardrail(guardrail_function=guardrail_function)],
|
|
model=model,
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
|
|
assert result.run_loop_task is not None
|
|
with pytest.raises(OutputGuardrailTripwireTriggered):
|
|
await result.run_loop_task
|
|
|
|
assert result.final_output is None
|
|
assert result.is_complete is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_guardrail_exception_raises_from_run_loop_task_before_stream_consumption():
|
|
def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any
|
|
) -> GuardrailFunctionOutput:
|
|
raise RuntimeError("guardrail failed")
|
|
|
|
model = FakeModel(initial_output=[get_text_message("first_test")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
output_guardrails=[OutputGuardrail(guardrail_function=guardrail_function)],
|
|
model=model,
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="user_message")
|
|
|
|
assert result.run_loop_task is not None
|
|
with pytest.raises(RuntimeError, match="guardrail failed"):
|
|
await result.run_loop_task
|
|
|
|
assert result.final_output is None
|
|
assert result.is_complete is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_input_guardrail_tripwire_triggered_causes_exception_streamed():
|
|
def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], input: Any
|
|
) -> GuardrailFunctionOutput:
|
|
return GuardrailFunctionOutput(
|
|
output_info=None,
|
|
tripwire_triggered=True,
|
|
)
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
model=FakeModel(),
|
|
)
|
|
|
|
with pytest.raises(InputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(
|
|
agent,
|
|
input="user_message",
|
|
run_config=RunConfig(
|
|
input_guardrails=[InputGuardrail(guardrail_function=guardrail_function)]
|
|
),
|
|
)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_output_guardrail_tripwire_triggered_causes_exception_streamed():
|
|
def guardrail_function(
|
|
context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any
|
|
) -> GuardrailFunctionOutput:
|
|
return GuardrailFunctionOutput(
|
|
output_info=None,
|
|
tripwire_triggered=True,
|
|
)
|
|
|
|
model = FakeModel(initial_output=[get_text_message("first_test")])
|
|
|
|
agent = Agent(
|
|
name="test",
|
|
model=model,
|
|
)
|
|
|
|
with pytest.raises(OutputGuardrailTripwireTriggered):
|
|
result = Runner.run_streamed(
|
|
agent,
|
|
input="user_message",
|
|
run_config=RunConfig(
|
|
output_guardrails=[OutputGuardrail(guardrail_function=guardrail_function)]
|
|
),
|
|
)
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_events():
|
|
model = FakeModel()
|
|
agent_1 = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[get_function_tool("bar", "bar_result")],
|
|
output_type=Foo,
|
|
)
|
|
|
|
agent_2 = Agent(
|
|
name="test",
|
|
model=model,
|
|
tools=[get_function_tool("foo", "foo_result")],
|
|
handoffs=[agent_1],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
# First turn: a tool call
|
|
[get_function_tool_call("foo", json.dumps({"bar": "baz"}))],
|
|
# Second turn: a message and a handoff
|
|
[get_text_message("a_message"), get_handoff_tool_call(agent_1)],
|
|
# Third turn: tool call
|
|
[get_function_tool_call("bar", json.dumps({"bar": "baz"}))],
|
|
# Fourth turn: structured output
|
|
[get_final_output_message(json.dumps(Foo(bar="baz")))],
|
|
]
|
|
)
|
|
|
|
# event_type: (count, event)
|
|
event_counts: dict[str, int] = {}
|
|
item_data: list[RunItem] = []
|
|
agent_data: list[AgentUpdatedStreamEvent] = []
|
|
|
|
result = Runner.run_streamed(
|
|
agent_2,
|
|
input=[
|
|
get_text_input_item("user_message"),
|
|
get_text_input_item("another_message"),
|
|
],
|
|
run_config=RunConfig(nest_handoff_history=True),
|
|
)
|
|
async for event in result.stream_events():
|
|
event_counts[event.type] = event_counts.get(event.type, 0) + 1
|
|
if event.type == "run_item_stream_event":
|
|
item_data.append(event.item)
|
|
elif event.type == "agent_updated_stream_event":
|
|
agent_data.append(event)
|
|
|
|
assert result.final_output == Foo(bar="baz")
|
|
assert len(result.raw_responses) == 4, "should have four model responses"
|
|
assert len(result.to_input_list()) == 9, (
|
|
"should have input: conversation summary, function call, function call result, message, "
|
|
"handoff, handoff output, tool call, tool call result, final output"
|
|
)
|
|
assert len(result.to_input_list(mode="normalized")) == 5, (
|
|
"should have normalized replay input: conversation summary, carried-forward message, "
|
|
"tool call, tool call result, final output"
|
|
)
|
|
|
|
assert result.last_agent == agent_1, "should have handed off to agent_1"
|
|
assert result.final_output == Foo(bar="baz"), "should have structured output"
|
|
|
|
# Now lets check the events
|
|
|
|
expected_item_type_map = {
|
|
# 3 tool_call_item events:
|
|
# 1. get_function_tool_call("foo", ...)
|
|
# 2. get_handoff_tool_call(agent_1) because handoffs are implemented via tool calls too
|
|
# 3. get_function_tool_call("bar", ...)
|
|
"tool_call": 3,
|
|
# Only 2 outputs, handoff tool call doesn't have corresponding tool_call_output event
|
|
"tool_call_output": 2,
|
|
"message": 2, # get_text_message("a_message") + get_final_output_message(...)
|
|
"handoff": 1, # get_handoff_tool_call(agent_1)
|
|
"handoff_output": 1, # handoff_output_item
|
|
}
|
|
|
|
total_expected_item_count = sum(expected_item_type_map.values())
|
|
|
|
assert event_counts["run_item_stream_event"] == total_expected_item_count, (
|
|
f"Expected {total_expected_item_count} events, got {event_counts['run_item_stream_event']}"
|
|
f"Expected events were: {expected_item_type_map}, got {event_counts}"
|
|
)
|
|
|
|
assert len(item_data) == total_expected_item_count, (
|
|
f"should have {total_expected_item_count} run items"
|
|
)
|
|
assert len(agent_data) == 2, "should have 2 agent updated events"
|
|
assert agent_data[0].new_agent == agent_2, "should have started with agent_2"
|
|
assert agent_data[1].new_agent == agent_1, "should have handed off to agent_1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dynamic_tool_addition_run_streamed() -> None:
|
|
model = FakeModel()
|
|
|
|
executed: dict[str, bool] = {"called": False}
|
|
|
|
agent = Agent(name="test", model=model, tool_use_behavior="run_llm_again")
|
|
|
|
@function_tool(name_override="tool2")
|
|
def tool2() -> str:
|
|
executed["called"] = True
|
|
return "result2"
|
|
|
|
@function_tool(name_override="add_tool")
|
|
async def add_tool() -> str:
|
|
agent.tools.append(tool2)
|
|
return "added"
|
|
|
|
agent.tools.append(add_tool)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("add_tool", json.dumps({}))],
|
|
[get_function_tool_call("tool2", json.dumps({}))],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="start")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
assert executed["called"] is True
|
|
assert result.final_output == "done"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_step_items_to_queue_handles_tool_approval_item():
|
|
"""Test that stream_step_items_to_queue handles ToolApprovalItem."""
|
|
_, agent = make_model_and_agent(name="test")
|
|
tool_call = get_function_tool_call("test_tool", "{}")
|
|
assert isinstance(tool_call, ResponseFunctionToolCall)
|
|
approval_item = ToolApprovalItem(agent=agent, raw_item=tool_call)
|
|
|
|
queue: asyncio.Queue[StreamEvent | QueueCompleteSentinel] = asyncio.Queue()
|
|
|
|
# ToolApprovalItem should not be streamed
|
|
run_loop.stream_step_items_to_queue([approval_item], queue)
|
|
|
|
# Queue should be empty since ToolApprovalItem is not streamed
|
|
assert queue.empty()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_hitl_resume_with_approved_tools():
|
|
"""Test resuming streaming run from RunState with approved tools executes them."""
|
|
tool_called = False
|
|
|
|
async def test_tool() -> str:
|
|
nonlocal tool_called
|
|
tool_called = True
|
|
return "tool_result"
|
|
|
|
# Create a tool that requires approval
|
|
tool = function_tool(test_tool, name_override="test_tool", needs_approval=True)
|
|
model, agent = make_model_and_agent(name="test", tools=[tool])
|
|
|
|
# First run - tool call that requires approval
|
|
queue_function_call_and_text(
|
|
model,
|
|
get_function_tool_call("test_tool", json.dumps({})),
|
|
followup=[get_text_message("done")],
|
|
)
|
|
|
|
first = Runner.run_streamed(agent, input="Use test_tool")
|
|
await consume_stream(first)
|
|
|
|
# Resume from state - should execute approved tool
|
|
result2 = await resume_streamed_after_first_approval(agent, first)
|
|
|
|
# Tool should have been called
|
|
assert tool_called is True
|
|
assert result2.final_output == "done"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_resume_with_session_does_not_duplicate_items():
|
|
"""Ensure session persistence does not duplicate tool items after streaming resume."""
|
|
|
|
async def test_tool() -> str:
|
|
return "tool_result"
|
|
|
|
tool = function_tool(test_tool, name_override="test_tool", needs_approval=True)
|
|
model, agent = make_model_and_agent(name="test", tools=[tool])
|
|
session = SimpleListSession()
|
|
|
|
queue_function_call_and_text(
|
|
model,
|
|
get_function_tool_call("test_tool", json.dumps({}), call_id="call-resume"),
|
|
followup=[get_text_message("done")],
|
|
)
|
|
|
|
first = Runner.run_streamed(agent, input="Use test_tool", session=session)
|
|
await consume_stream(first)
|
|
assert first.interruptions
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = Runner.run_streamed(agent, state, session=session)
|
|
await consume_stream(resumed)
|
|
assert resumed.final_output == "done"
|
|
|
|
saved_items = await session.get_items()
|
|
call_count = sum(
|
|
1
|
|
for item in saved_items
|
|
if isinstance(item, dict)
|
|
and item.get("type") == "function_call"
|
|
and item.get("call_id") == "call-resume"
|
|
)
|
|
output_count = sum(
|
|
1
|
|
for item in saved_items
|
|
if isinstance(item, dict)
|
|
and item.get("type") == "function_call_output"
|
|
and item.get("call_id") == "call-resume"
|
|
)
|
|
|
|
assert call_count == 1
|
|
assert output_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_resume_preserves_filtered_model_input_after_handoff():
|
|
model = FakeModel()
|
|
|
|
@function_tool(name_override="approval_tool", needs_approval=True)
|
|
def approval_tool() -> str:
|
|
return "ok"
|
|
|
|
delegate = Agent(
|
|
name="delegate",
|
|
model=model,
|
|
tools=[approval_tool],
|
|
)
|
|
triage = Agent(
|
|
name="triage",
|
|
model=model,
|
|
handoffs=[delegate],
|
|
tools=[get_function_tool("some_function", "result")],
|
|
)
|
|
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[
|
|
get_function_tool_call(
|
|
"some_function", json.dumps({"a": "b"}), call_id="triage-call"
|
|
)
|
|
],
|
|
[get_text_message("a_message"), get_handoff_tool_call(delegate)],
|
|
[get_function_tool_call("approval_tool", json.dumps({}), call_id="delegate-call")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
|
|
model_input_call_ids: list[set[str]] = []
|
|
model_input_output_call_ids: list[set[str]] = []
|
|
|
|
def capture_model_input(data):
|
|
call_ids: set[str] = set()
|
|
output_call_ids: set[str] = set()
|
|
for item in data.model_data.input:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
item_type = item.get("type")
|
|
call_id = item.get("call_id")
|
|
if not isinstance(call_id, str):
|
|
continue
|
|
if item_type == "function_call":
|
|
call_ids.add(call_id)
|
|
elif item_type == "function_call_output":
|
|
output_call_ids.add(call_id)
|
|
model_input_call_ids.append(call_ids)
|
|
model_input_output_call_ids.append(output_call_ids)
|
|
return data.model_data
|
|
|
|
run_config = RunConfig(
|
|
nest_handoff_history=True,
|
|
call_model_input_filter=capture_model_input,
|
|
)
|
|
|
|
first = Runner.run_streamed(triage, input="user_message", run_config=run_config)
|
|
await consume_stream(first)
|
|
assert first.interruptions
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = Runner.run_streamed(triage, state, run_config=run_config)
|
|
await consume_stream(resumed)
|
|
|
|
last_call_ids = model_input_call_ids[-1]
|
|
last_output_call_ids = model_input_output_call_ids[-1]
|
|
assert "triage-call" not in last_call_ids
|
|
assert "triage-call" not in last_output_call_ids
|
|
assert "delegate-call" in last_call_ids
|
|
assert "delegate-call" in last_output_call_ids
|
|
assert resumed.final_output == "done"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_resume_persists_tool_outputs_on_run_again():
|
|
"""Approved tool outputs should be persisted before streaming resumes the next turn."""
|
|
|
|
async def test_tool() -> str:
|
|
return "tool_result"
|
|
|
|
tool = function_tool(test_tool, name_override="test_tool", needs_approval=True)
|
|
model, agent = make_model_and_agent(name="test", tools=[tool])
|
|
session = SimpleListSession()
|
|
|
|
queue_function_call_and_text(
|
|
model,
|
|
get_function_tool_call("test_tool", json.dumps({}), call_id="call-resume"),
|
|
followup=[get_text_message("done")],
|
|
)
|
|
|
|
first = Runner.run_streamed(agent, input="Use test_tool", session=session)
|
|
await consume_stream(first)
|
|
|
|
assert first.interruptions
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = Runner.run_streamed(agent, state, session=session)
|
|
await consume_stream(resumed)
|
|
|
|
saved_items = await session.get_items()
|
|
assert any(
|
|
isinstance(item, dict)
|
|
and item.get("type") == "function_call_output"
|
|
and item.get("call_id") == "call-resume"
|
|
for item in saved_items
|
|
), "approved tool outputs should be persisted on resume"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_resume_carries_persisted_count(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Ensure resumed streaming preserves the persisted count for session saves."""
|
|
|
|
async def test_tool() -> str:
|
|
return "tool_result"
|
|
|
|
tool = function_tool(test_tool, name_override="test_tool", needs_approval=True)
|
|
model, agent = make_model_and_agent(name="test", tools=[tool])
|
|
session = SimpleListSession()
|
|
|
|
queue_function_call_and_text(
|
|
model,
|
|
get_function_tool_call("test_tool", json.dumps({}), call_id="call-resume"),
|
|
followup=[get_text_message("done")],
|
|
)
|
|
|
|
first = Runner.run_streamed(agent, input="Use test_tool", session=session)
|
|
await consume_stream(first)
|
|
assert first.interruptions
|
|
|
|
persisted_count = first._current_turn_persisted_item_count
|
|
assert persisted_count > 0
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
observed_counts: list[int] = []
|
|
run_loop_any = cast(Any, run_loop)
|
|
real_save_resumed = run_loop_any.save_resumed_turn_items
|
|
|
|
async def save_wrapper(
|
|
*,
|
|
session: Any,
|
|
items: list[RunItem],
|
|
persisted_count: int,
|
|
response_id: str | None,
|
|
reasoning_item_id_policy: str | None = None,
|
|
store: bool | None = None,
|
|
) -> int:
|
|
observed_counts.append(persisted_count)
|
|
result = await real_save_resumed(
|
|
session=session,
|
|
items=items,
|
|
persisted_count=persisted_count,
|
|
response_id=response_id,
|
|
reasoning_item_id_policy=reasoning_item_id_policy,
|
|
store=store,
|
|
)
|
|
return int(result)
|
|
|
|
monkeypatch.setattr(run_loop_any, "save_resumed_turn_items", save_wrapper)
|
|
|
|
resumed = Runner.run_streamed(agent, state, session=session)
|
|
await consume_stream(resumed)
|
|
|
|
assert observed_counts, "expected resumed save to capture persisted count"
|
|
assert all(count == persisted_count for count in observed_counts)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_hitl_resume_enforces_max_turns():
|
|
"""Test that streamed resumes advance turn counts for max_turns enforcement."""
|
|
|
|
async def test_tool() -> str:
|
|
return "tool_result"
|
|
|
|
tool = function_tool(test_tool, name_override="test_tool", needs_approval=True)
|
|
model, agent = make_model_and_agent(name="test", tools=[tool])
|
|
|
|
queue_function_call_and_text(
|
|
model,
|
|
get_function_tool_call("test_tool", json.dumps({})),
|
|
followup=[get_text_message("done")],
|
|
)
|
|
|
|
first = Runner.run_streamed(agent, input="Use test_tool", max_turns=1)
|
|
await consume_stream(first)
|
|
|
|
assert first.interruptions
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = Runner.run_streamed(agent, state)
|
|
with pytest.raises(MaxTurnsExceeded):
|
|
async for _ in resumed.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_max_turns_emits_pending_tool_output_events() -> None:
|
|
async def test_tool() -> str:
|
|
return "tool_result"
|
|
|
|
tool = function_tool(test_tool, name_override="test_tool")
|
|
model, agent = make_model_and_agent(name="test", tools=[tool])
|
|
|
|
queue_function_call_and_text(
|
|
model,
|
|
get_function_tool_call("test_tool", json.dumps({})),
|
|
followup=[get_text_message("done")],
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="Use test_tool", max_turns=1)
|
|
streamed_item_types: list[str] = []
|
|
|
|
with pytest.raises(MaxTurnsExceeded):
|
|
async for event in result.stream_events():
|
|
if event.type == "run_item_stream_event":
|
|
streamed_item_types.append(event.item.type)
|
|
|
|
assert "tool_call_item" in streamed_item_types
|
|
assert "tool_call_output_item" in streamed_item_types
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_non_max_turns_exception_does_not_emit_queued_events() -> None:
|
|
model, agent = make_model_and_agent(name="test")
|
|
model.set_next_output([get_text_message("done")])
|
|
|
|
result = Runner.run_streamed(agent, input="hello")
|
|
result.cancel()
|
|
await asyncio.sleep(0)
|
|
|
|
while not result._event_queue.empty():
|
|
result._event_queue.get_nowait()
|
|
result._event_queue.task_done()
|
|
|
|
result._stored_exception = RuntimeError("guardrail-triggered")
|
|
result._event_queue.put_nowait(AgentUpdatedStreamEvent(new_agent=agent))
|
|
|
|
streamed_events: list[StreamEvent] = []
|
|
with pytest.raises(RuntimeError, match="guardrail-triggered"):
|
|
async for event in result.stream_events():
|
|
streamed_events.append(event)
|
|
|
|
assert streamed_events == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_hitl_server_conversation_tracker_priming():
|
|
"""Test that resuming streaming run from RunState primes server conversation tracker."""
|
|
model, agent = make_model_and_agent(name="test")
|
|
|
|
# First run with conversation_id
|
|
model.set_next_output([get_text_message("First response")])
|
|
result1 = Runner.run_streamed(
|
|
agent, input="test", conversation_id="conv123", previous_response_id="resp123"
|
|
)
|
|
await consume_stream(result1)
|
|
|
|
# Create state from result
|
|
state = result1.to_state()
|
|
|
|
# Resume with same conversation_id - should not duplicate messages
|
|
model.set_next_output([get_text_message("Second response")])
|
|
result2 = Runner.run_streamed(
|
|
agent, state, conversation_id="conv123", previous_response_id="resp123"
|
|
)
|
|
await consume_stream(result2)
|
|
|
|
# Should complete successfully without message duplication
|
|
assert result2.final_output == "Second response"
|
|
assert len(result2.new_items) >= 1
|