wehub-resource-sync
139 lines
4.2 KiB
Python
139 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any, Literal, TypeVar, cast
|
|
|
|
from openai.types.responses import (
|
|
ResponseFunctionToolCall,
|
|
ResponseOutputMessage,
|
|
ResponseOutputText,
|
|
)
|
|
|
|
from agents import Agent
|
|
from agents._tool_identity import FunctionToolLookupKey, get_function_tool_lookup_key
|
|
from agents.items import ToolApprovalItem
|
|
from agents.run_context import RunContextWrapper
|
|
from agents.run_state import RunState
|
|
from agents.sandbox.session.sandbox_session_state import SandboxSessionState
|
|
|
|
TContext = TypeVar("TContext")
|
|
_AUTO_LOOKUP_KEY = object()
|
|
|
|
|
|
class TestSessionState(SandboxSessionState):
|
|
"""Concrete ``SandboxSessionState`` subclass for tests that don't need a real backend."""
|
|
|
|
__test__ = False
|
|
type: Literal["test"] = "test"
|
|
|
|
|
|
def make_tool_call(
|
|
call_id: str = "call_1",
|
|
*,
|
|
name: str = "test_tool",
|
|
namespace: str | None = None,
|
|
status: Literal["in_progress", "completed", "incomplete"] | None = "completed",
|
|
arguments: str = "{}",
|
|
call_type: Literal["function_call"] = "function_call",
|
|
) -> ResponseFunctionToolCall:
|
|
"""Build a ResponseFunctionToolCall with common defaults."""
|
|
|
|
kwargs: dict[str, Any] = {
|
|
"type": call_type,
|
|
"name": name,
|
|
"call_id": call_id,
|
|
"status": status,
|
|
"arguments": arguments,
|
|
}
|
|
if namespace is not None:
|
|
kwargs["namespace"] = namespace
|
|
return ResponseFunctionToolCall(**kwargs)
|
|
|
|
|
|
def make_tool_approval_item(
|
|
agent: Agent[Any],
|
|
*,
|
|
call_id: str = "call_1",
|
|
name: str = "test_tool",
|
|
namespace: str | None = None,
|
|
allow_bare_name_alias: bool = False,
|
|
status: Literal["in_progress", "completed", "incomplete"] | None = "completed",
|
|
arguments: str = "{}",
|
|
tool_lookup_key: FunctionToolLookupKey | None | object = _AUTO_LOOKUP_KEY,
|
|
) -> ToolApprovalItem:
|
|
"""Create a ToolApprovalItem backed by a function call."""
|
|
|
|
resolved_tool_lookup_key: FunctionToolLookupKey | None
|
|
if tool_lookup_key is _AUTO_LOOKUP_KEY:
|
|
resolved_tool_lookup_key = get_function_tool_lookup_key(name, namespace)
|
|
else:
|
|
resolved_tool_lookup_key = cast(FunctionToolLookupKey | None, tool_lookup_key)
|
|
|
|
return ToolApprovalItem(
|
|
agent=agent,
|
|
raw_item=make_tool_call(
|
|
call_id=call_id,
|
|
name=name,
|
|
namespace=namespace,
|
|
status=status,
|
|
arguments=arguments,
|
|
),
|
|
tool_namespace=namespace,
|
|
tool_lookup_key=resolved_tool_lookup_key,
|
|
_allow_bare_name_alias=allow_bare_name_alias,
|
|
)
|
|
|
|
|
|
def make_message_output(
|
|
*,
|
|
message_id: str = "msg_1",
|
|
text: str = "Hello",
|
|
role: Literal["assistant"] = "assistant",
|
|
status: Literal["in_progress", "completed", "incomplete"] = "completed",
|
|
) -> ResponseOutputMessage:
|
|
"""Create a minimal ResponseOutputMessage."""
|
|
|
|
return ResponseOutputMessage(
|
|
id=message_id,
|
|
type="message",
|
|
role=role,
|
|
status=status,
|
|
content=[ResponseOutputText(type="output_text", text=text, annotations=[], logprobs=[])],
|
|
)
|
|
|
|
|
|
def make_run_state(
|
|
agent: Agent[Any],
|
|
*,
|
|
context: RunContextWrapper[TContext] | dict[str, Any] | None = None,
|
|
original_input: Any = "input",
|
|
max_turns: int | None = 3,
|
|
) -> RunState[TContext, Agent[Any]]:
|
|
"""Create a RunState with sensible defaults for tests."""
|
|
|
|
wrapper: RunContextWrapper[TContext]
|
|
if isinstance(context, RunContextWrapper):
|
|
wrapper = context
|
|
else:
|
|
wrapper = RunContextWrapper(context=context or {}) # type: ignore[arg-type]
|
|
|
|
return RunState(
|
|
context=wrapper,
|
|
original_input=original_input,
|
|
starting_agent=agent,
|
|
max_turns=max_turns,
|
|
)
|
|
|
|
|
|
async def roundtrip_state(
|
|
agent: Agent[Any],
|
|
state: RunState[TContext, Agent[Any]],
|
|
mutate_json: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
|
|
) -> RunState[TContext, Agent[Any]]:
|
|
"""Serialize and restore a RunState, optionally mutating the JSON in between."""
|
|
|
|
json_data = state.to_json()
|
|
if mutate_json is not None:
|
|
json_data = mutate_json(json_data)
|
|
return await RunState.from_json(agent, json_data)
|