wehub-resource-sync
478 lines
15 KiB
Python
478 lines
15 KiB
Python
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 ""
|