chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,477 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
from openai.types.responses import ResponseFunctionToolCall
|
||||
|
||||
from agents import (
|
||||
Agent,
|
||||
Model,
|
||||
ModelResponse,
|
||||
ModelSettings,
|
||||
OpenAIConversationsSession,
|
||||
Runner,
|
||||
Usage,
|
||||
function_tool,
|
||||
)
|
||||
from agents.items import TResponseInputItem, TResponseStreamEvent
|
||||
from tests.test_responses import get_text_message
|
||||
from tests.utils.hitl import HITL_REJECTION_MSG
|
||||
from tests.utils.simple_session import SimpleListSession
|
||||
|
||||
TOOL_ECHO = "approved_echo"
|
||||
TOOL_NOTE = "approved_note"
|
||||
USER_MESSAGES = [
|
||||
"Fetch profile for customer 104.",
|
||||
"Update note for customer 104.",
|
||||
"Delete note for customer 104.",
|
||||
]
|
||||
|
||||
execute_counts: dict[str, int] = {}
|
||||
|
||||
|
||||
@function_tool(
|
||||
name_override=TOOL_ECHO,
|
||||
description_override="Echoes back the provided query after approval.",
|
||||
needs_approval=True,
|
||||
)
|
||||
def approval_echo(query: str) -> str:
|
||||
execute_counts[TOOL_ECHO] = execute_counts.get(TOOL_ECHO, 0) + 1
|
||||
return f"approved:{query}"
|
||||
|
||||
|
||||
@function_tool(
|
||||
name_override=TOOL_NOTE,
|
||||
description_override="Records the provided query after approval.",
|
||||
needs_approval=True,
|
||||
)
|
||||
def approval_note(query: str) -> str:
|
||||
execute_counts[TOOL_NOTE] = execute_counts.get(TOOL_NOTE, 0) + 1
|
||||
return f"approved_note:{query}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ScenarioStep:
|
||||
label: str
|
||||
message: str
|
||||
tool_name: str
|
||||
approval: str
|
||||
expected_output: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ScenarioResult:
|
||||
approval_item: Any
|
||||
items: list[TResponseInputItem]
|
||||
|
||||
|
||||
class ScenarioModel(Model):
|
||||
def __init__(self) -> None:
|
||||
self._counter = 0
|
||||
|
||||
async def get_response(
|
||||
self,
|
||||
system_instructions: str | None,
|
||||
input: str | list[TResponseInputItem],
|
||||
model_settings: ModelSettings,
|
||||
tools: list[Any],
|
||||
output_schema: Any,
|
||||
handoffs: list[Any],
|
||||
tracing: Any,
|
||||
*,
|
||||
previous_response_id: str | None,
|
||||
conversation_id: str | None,
|
||||
prompt: Any | None,
|
||||
) -> ModelResponse:
|
||||
if input_has_rejection(input):
|
||||
return ModelResponse(
|
||||
output=[get_text_message(HITL_REJECTION_MSG)],
|
||||
usage=Usage(),
|
||||
response_id="resp-test",
|
||||
)
|
||||
tool_choice = model_settings.tool_choice
|
||||
tool_name = tool_choice if isinstance(tool_choice, str) else TOOL_ECHO
|
||||
self._counter += 1
|
||||
call_id = f"call_{self._counter}"
|
||||
query = extract_user_message(input)
|
||||
tool_call = ResponseFunctionToolCall(
|
||||
type="function_call",
|
||||
name=tool_name,
|
||||
call_id=call_id,
|
||||
arguments=json.dumps({"query": query}),
|
||||
)
|
||||
return ModelResponse(output=[tool_call], usage=Usage(), response_id="resp-test")
|
||||
|
||||
async def stream_response(
|
||||
self,
|
||||
system_instructions: str | None,
|
||||
input: str | list[TResponseInputItem],
|
||||
model_settings: ModelSettings,
|
||||
tools: list[Any],
|
||||
output_schema: Any,
|
||||
handoffs: list[Any],
|
||||
tracing: Any,
|
||||
*,
|
||||
previous_response_id: str | None,
|
||||
conversation_id: str | None,
|
||||
prompt: Any | None,
|
||||
) -> AsyncIterator[TResponseStreamEvent]:
|
||||
if False:
|
||||
yield cast(TResponseStreamEvent, {})
|
||||
raise RuntimeError("Streaming is not supported in this scenario.")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_session_hitl_scenario() -> None:
|
||||
execute_counts.clear()
|
||||
session = SimpleListSession(session_id="memory")
|
||||
model = ScenarioModel()
|
||||
|
||||
steps = [
|
||||
ScenarioStep(
|
||||
label="turn 1",
|
||||
message=USER_MESSAGES[0],
|
||||
tool_name=TOOL_ECHO,
|
||||
approval="approve",
|
||||
expected_output=f"approved:{USER_MESSAGES[0]}",
|
||||
),
|
||||
ScenarioStep(
|
||||
label="turn 2 (rehydrated)",
|
||||
message=USER_MESSAGES[1],
|
||||
tool_name=TOOL_NOTE,
|
||||
approval="approve",
|
||||
expected_output=f"approved_note:{USER_MESSAGES[1]}",
|
||||
),
|
||||
ScenarioStep(
|
||||
label="turn 3 (rejected)",
|
||||
message=USER_MESSAGES[2],
|
||||
tool_name=TOOL_ECHO,
|
||||
approval="reject",
|
||||
expected_output=HITL_REJECTION_MSG,
|
||||
),
|
||||
]
|
||||
|
||||
rehydrated: SimpleListSession | None = None
|
||||
|
||||
try:
|
||||
first = await run_scenario_step(session, model, steps[0])
|
||||
assert_counts(first.items, 1)
|
||||
assert_step_output(first.items, first.approval_item, steps[0])
|
||||
|
||||
rehydrated = SimpleListSession(
|
||||
session_id=session.session_id,
|
||||
history=first.items,
|
||||
)
|
||||
second = await run_scenario_step(rehydrated, model, steps[1])
|
||||
assert_counts(second.items, 2)
|
||||
assert_step_output(second.items, second.approval_item, steps[1])
|
||||
|
||||
third = await run_scenario_step(rehydrated, model, steps[2])
|
||||
assert_counts(third.items, 3)
|
||||
assert_step_output(third.items, third.approval_item, steps[2])
|
||||
|
||||
assert execute_counts.get(TOOL_ECHO) == 1
|
||||
assert execute_counts.get(TOOL_NOTE) == 1
|
||||
finally:
|
||||
await (rehydrated or session).clear_session()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_conversations_session_hitl_scenario() -> None:
|
||||
execute_counts.clear()
|
||||
stored_items: list[dict[str, Any]] = []
|
||||
|
||||
async def create_items(*, conversation_id: str, items: list[Any]) -> None:
|
||||
stored_items.extend(items)
|
||||
|
||||
def list_items(*, conversation_id: str, order: str, limit: int | None = None):
|
||||
class StoredItem:
|
||||
def __init__(self, payload: dict[str, Any]) -> None:
|
||||
self._payload = payload
|
||||
|
||||
def model_dump(self, exclude_unset: bool = True) -> dict[str, Any]:
|
||||
return self._payload
|
||||
|
||||
async def iterator():
|
||||
if order == "desc":
|
||||
items_iter = list(reversed(stored_items))
|
||||
else:
|
||||
items_iter = list(stored_items)
|
||||
if limit is not None:
|
||||
items_iter = items_iter[:limit]
|
||||
for item in items_iter:
|
||||
yield StoredItem(item)
|
||||
|
||||
return iterator()
|
||||
|
||||
class ConversationsItems:
|
||||
create = staticmethod(create_items)
|
||||
list = staticmethod(list_items)
|
||||
|
||||
async def delete(self, *args: Any, **kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
class Conversations:
|
||||
items = ConversationsItems()
|
||||
|
||||
async def create(self, *args: Any, **kwargs: Any) -> Any:
|
||||
return type("Response", (), {"id": "conv_test"})()
|
||||
|
||||
async def delete(self, *args: Any, **kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
class Client:
|
||||
conversations = Conversations()
|
||||
|
||||
client = Client()
|
||||
typed_client = cast(Any, client)
|
||||
session = OpenAIConversationsSession(conversation_id="conv_test", openai_client=typed_client)
|
||||
rehydrated_session = OpenAIConversationsSession(
|
||||
conversation_id="conv_test", openai_client=typed_client
|
||||
)
|
||||
model = ScenarioModel()
|
||||
|
||||
steps = [
|
||||
ScenarioStep(
|
||||
label="turn 1",
|
||||
message=USER_MESSAGES[0],
|
||||
tool_name=TOOL_ECHO,
|
||||
approval="approve",
|
||||
expected_output=f"approved:{USER_MESSAGES[0]}",
|
||||
),
|
||||
ScenarioStep(
|
||||
label="turn 2 (rehydrated)",
|
||||
message=USER_MESSAGES[1],
|
||||
tool_name=TOOL_NOTE,
|
||||
approval="approve",
|
||||
expected_output=f"approved_note:{USER_MESSAGES[1]}",
|
||||
),
|
||||
ScenarioStep(
|
||||
label="turn 3 (rejected)",
|
||||
message=USER_MESSAGES[2],
|
||||
tool_name=TOOL_ECHO,
|
||||
approval="reject",
|
||||
expected_output=HITL_REJECTION_MSG,
|
||||
),
|
||||
]
|
||||
|
||||
offset = 0
|
||||
first = await run_scenario_step(session, model, steps[0])
|
||||
first_items = stored_items[offset:]
|
||||
offset = len(stored_items)
|
||||
assert_step_items(first_items, steps[0], first.approval_item)
|
||||
|
||||
second = await run_scenario_step(rehydrated_session, model, steps[1])
|
||||
second_items = stored_items[offset:]
|
||||
offset = len(stored_items)
|
||||
assert_step_items(second_items, steps[1], second.approval_item)
|
||||
|
||||
third = await run_scenario_step(rehydrated_session, model, steps[2])
|
||||
third_items = stored_items[offset:]
|
||||
assert_step_items(third_items, steps[2], third.approval_item)
|
||||
|
||||
assert execute_counts.get(TOOL_ECHO) == 1
|
||||
assert execute_counts.get(TOOL_NOTE) == 1
|
||||
|
||||
|
||||
async def run_scenario_step(
|
||||
session: Any,
|
||||
model: ScenarioModel,
|
||||
step: ScenarioStep,
|
||||
) -> ScenarioResult:
|
||||
agent = Agent(
|
||||
name=f"Scenario {step.label}",
|
||||
instructions=f"Always call {step.tool_name} before responding.",
|
||||
model=model,
|
||||
tools=[approval_echo, approval_note],
|
||||
model_settings=ModelSettings(tool_choice=step.tool_name),
|
||||
tool_use_behavior="stop_on_first_tool",
|
||||
)
|
||||
|
||||
first_run = await Runner.run(agent, step.message, session=session)
|
||||
assert len(first_run.interruptions) == 1
|
||||
|
||||
approval = first_run.interruptions[0]
|
||||
state = first_run.to_state()
|
||||
if step.approval == "reject":
|
||||
state.reject(approval)
|
||||
else:
|
||||
state.approve(approval)
|
||||
|
||||
resumed = await Runner.run(agent, state, session=session)
|
||||
assert resumed.interruptions == []
|
||||
assert resumed.final_output == step.expected_output
|
||||
|
||||
return ScenarioResult(approval_item=approval, items=await session.get_items())
|
||||
|
||||
|
||||
def assert_counts(items: list[TResponseInputItem], turn: int) -> None:
|
||||
assert count_user_messages(items) == turn
|
||||
assert count_function_calls(items) == turn
|
||||
assert count_function_outputs(items) == turn
|
||||
|
||||
|
||||
def assert_step_output(
|
||||
items: list[TResponseInputItem],
|
||||
approval_item: Any,
|
||||
step: ScenarioStep,
|
||||
) -> None:
|
||||
last_user = get_last_user_text(items)
|
||||
assert last_user == step.message
|
||||
|
||||
last_call = find_last_function_call(items)
|
||||
last_result = find_last_function_output(items)
|
||||
|
||||
approval_call_id = extract_call_id(approval_item.raw_item)
|
||||
assert last_call is not None
|
||||
assert last_call.get("name") == step.tool_name
|
||||
assert last_call.get("call_id") == approval_call_id
|
||||
|
||||
assert last_result is not None
|
||||
assert last_result.get("call_id") == approval_call_id
|
||||
assert extract_output_text(last_result) == step.expected_output
|
||||
|
||||
|
||||
def assert_step_items(
|
||||
items: list[dict[str, Any]],
|
||||
step: ScenarioStep,
|
||||
approval_item: Any,
|
||||
) -> None:
|
||||
user_items = [item for item in items if item.get("role") == "user"]
|
||||
function_calls = [item for item in items if item.get("type") == "function_call"]
|
||||
function_outputs = [item for item in items if item.get("type") == "function_call_output"]
|
||||
|
||||
assert len(user_items) == 1
|
||||
assert len(function_calls) == 1
|
||||
assert len(function_outputs) == 1
|
||||
|
||||
assert extract_user_text(user_items[0]) == step.message
|
||||
assert function_calls[0].get("name") == step.tool_name
|
||||
|
||||
approval_call_id = extract_call_id(approval_item.raw_item)
|
||||
assert function_calls[0].get("call_id") == approval_call_id
|
||||
assert function_outputs[0].get("call_id") == approval_call_id
|
||||
assert extract_output_text(function_outputs[0]) == step.expected_output
|
||||
|
||||
|
||||
def extract_user_message(input: str | list[TResponseInputItem]) -> str:
|
||||
if isinstance(input, str):
|
||||
return input
|
||||
|
||||
for item in reversed(input):
|
||||
if isinstance(item, dict) and item.get("role") == "user":
|
||||
content = item.get("content")
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
text = "".join(
|
||||
part.get("text", "")
|
||||
for part in content
|
||||
if isinstance(part, dict) and part.get("type") == "input_text"
|
||||
)
|
||||
if text:
|
||||
return text
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def input_has_rejection(input: str | list[TResponseInputItem]) -> bool:
|
||||
if not isinstance(input, list):
|
||||
return False
|
||||
for item in input:
|
||||
if not isinstance(item, dict) or item.get("type") != "function_call_output":
|
||||
continue
|
||||
output = item.get("output")
|
||||
if output == HITL_REJECTION_MSG:
|
||||
return True
|
||||
if isinstance(output, dict) and output.get("type") == "input_text":
|
||||
if output.get("text") == HITL_REJECTION_MSG:
|
||||
return True
|
||||
if isinstance(output, list):
|
||||
for entry in output:
|
||||
if isinstance(entry, dict) and entry.get("type") == "input_text":
|
||||
if entry.get("text") == HITL_REJECTION_MSG:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def count_user_messages(items: list[TResponseInputItem]) -> int:
|
||||
return sum(1 for item in items if isinstance(item, dict) and item.get("role") == "user")
|
||||
|
||||
|
||||
def count_function_calls(items: list[TResponseInputItem]) -> int:
|
||||
return sum(
|
||||
1 for item in items if isinstance(item, dict) and item.get("type") == "function_call"
|
||||
)
|
||||
|
||||
|
||||
def count_function_outputs(items: list[TResponseInputItem]) -> int:
|
||||
return sum(
|
||||
1 for item in items if isinstance(item, dict) and item.get("type") == "function_call_output"
|
||||
)
|
||||
|
||||
|
||||
def find_last_function_call(
|
||||
items: list[TResponseInputItem],
|
||||
) -> dict[str, Any] | None:
|
||||
for item in reversed(items):
|
||||
if isinstance(item, dict) and item.get("type") == "function_call":
|
||||
return cast(dict[str, Any], item)
|
||||
return None
|
||||
|
||||
|
||||
def find_last_function_output(
|
||||
items: list[TResponseInputItem],
|
||||
) -> dict[str, Any] | None:
|
||||
for item in reversed(items):
|
||||
if isinstance(item, dict) and item.get("type") == "function_call_output":
|
||||
return cast(dict[str, Any], item)
|
||||
return None
|
||||
|
||||
|
||||
def get_last_user_text(items: list[TResponseInputItem]) -> str | None:
|
||||
for item in reversed(items):
|
||||
if isinstance(item, dict) and item.get("role") == "user":
|
||||
return extract_user_text(cast(dict[str, Any], item))
|
||||
return None
|
||||
|
||||
|
||||
def extract_user_text(item: dict[str, Any]) -> str:
|
||||
content = item.get("content")
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
return "".join(
|
||||
part.get("text", "")
|
||||
for part in content
|
||||
if isinstance(part, dict) and part.get("type") == "input_text"
|
||||
)
|
||||
return ""
|
||||
|
||||
|
||||
def extract_call_id(item: Any) -> str | None:
|
||||
if isinstance(item, dict):
|
||||
return item.get("call_id") or item.get("id")
|
||||
return getattr(item, "call_id", None) or getattr(item, "id", None)
|
||||
|
||||
|
||||
def extract_output_text(item: dict[str, Any] | None) -> str:
|
||||
if not item:
|
||||
return ""
|
||||
|
||||
output = item.get("output")
|
||||
if isinstance(output, str):
|
||||
return output
|
||||
if isinstance(output, list):
|
||||
for entry in output:
|
||||
if isinstance(entry, dict) and entry.get("type") == "input_text":
|
||||
text = entry.get("text")
|
||||
return text if isinstance(text, str) else ""
|
||||
if isinstance(output, dict) and output.get("type") == "input_text":
|
||||
text = output.get("text")
|
||||
return text if isinstance(text, str) else ""
|
||||
return ""
|
||||
Reference in New Issue
Block a user