|
|
|
@@ -7,8 +7,11 @@ import base64
|
|
|
|
|
import json
|
|
|
|
|
import logging
|
|
|
|
|
import os
|
|
|
|
|
import tempfile
|
|
|
|
|
import threading
|
|
|
|
|
from collections.abc import AsyncIterable, AsyncIterator, Generator, Mapping, Sequence
|
|
|
|
|
from typing import cast
|
|
|
|
|
from contextlib import suppress
|
|
|
|
|
from typing import Protocol, cast
|
|
|
|
|
|
|
|
|
|
from agent_framework import (
|
|
|
|
|
ChatOptions,
|
|
|
|
@@ -109,11 +112,105 @@ from typing_extensions import Any
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ApprovalStorage(Protocol):
|
|
|
|
|
"""Storage for saving function approval requests."""
|
|
|
|
|
|
|
|
|
|
async def save_approval_request(self, approval_request_id: str, request: Content) -> None:
|
|
|
|
|
"""Save a function approval request under the given ID."""
|
|
|
|
|
...
|
|
|
|
|
|
|
|
|
|
async def load_approval_request(self, approval_request_id: str) -> Content:
|
|
|
|
|
"""Load a function approval request by its ID."""
|
|
|
|
|
...
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class InMemoryFunctionApprovalStorage:
|
|
|
|
|
"""An in-memory storage for function approval requests."""
|
|
|
|
|
|
|
|
|
|
def __init__(self) -> None:
|
|
|
|
|
self._store: dict[str, Content] = {}
|
|
|
|
|
|
|
|
|
|
async def save_approval_request(self, approval_request_id: str, request: Content) -> None:
|
|
|
|
|
if approval_request_id in self._store:
|
|
|
|
|
raise ValueError(f"Approval request with ID '{approval_request_id}' already exists.")
|
|
|
|
|
self._store[approval_request_id] = request
|
|
|
|
|
|
|
|
|
|
async def load_approval_request(self, approval_request_id: str) -> Content:
|
|
|
|
|
if approval_request_id not in self._store:
|
|
|
|
|
raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.")
|
|
|
|
|
return self._store[approval_request_id]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class FileBasedFunctionApprovalStorage:
|
|
|
|
|
"""A simple file-based storage for function approval requests.
|
|
|
|
|
|
|
|
|
|
Concurrent writes from multiple threads in the same process are
|
|
|
|
|
serialized by a ``threading.Lock``, and the on-disk JSON file is
|
|
|
|
|
updated atomically (write to a temp file, then ``os.replace``) so a
|
|
|
|
|
crash mid-write cannot leave a partially written file behind.
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
def __init__(self, storage_path: str) -> None:
|
|
|
|
|
self._storage_path = storage_path
|
|
|
|
|
self._lock = threading.Lock()
|
|
|
|
|
|
|
|
|
|
def _create_storage_file_if_not_exists_sync(self) -> None:
|
|
|
|
|
"""Lazy-create the storage file (and its parent directory) if it does not already exist.
|
|
|
|
|
|
|
|
|
|
Uses exclusive-create mode (``"x"``) so a concurrent creator cannot
|
|
|
|
|
be truncated by an ``open(..., "w")`` after a stale existence check.
|
|
|
|
|
"""
|
|
|
|
|
os.makedirs(os.path.dirname(self._storage_path) or ".", exist_ok=True)
|
|
|
|
|
with suppress(FileExistsError), open(self._storage_path, "x") as f:
|
|
|
|
|
json.dump({}, f)
|
|
|
|
|
|
|
|
|
|
def _atomic_write(self, data: dict[str, Any]) -> None:
|
|
|
|
|
"""Atomically replace the storage file with the serialized ``data``."""
|
|
|
|
|
directory = os.path.dirname(self._storage_path) or "."
|
|
|
|
|
# Serialize first so any error doesn't leave a partial file behind.
|
|
|
|
|
serialized = json.dumps(data)
|
|
|
|
|
fd, tmp_path = tempfile.mkstemp(prefix=".approvals-", suffix=".tmp", dir=directory)
|
|
|
|
|
try:
|
|
|
|
|
with os.fdopen(fd, "w") as tmp:
|
|
|
|
|
tmp.write(serialized)
|
|
|
|
|
os.replace(tmp_path, self._storage_path)
|
|
|
|
|
except BaseException:
|
|
|
|
|
with suppress(OSError):
|
|
|
|
|
os.unlink(tmp_path)
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
def _save_sync(self, approval_request_id: str, request: Content) -> None:
|
|
|
|
|
with self._lock:
|
|
|
|
|
self._create_storage_file_if_not_exists_sync()
|
|
|
|
|
with open(self._storage_path) as f:
|
|
|
|
|
data = json.load(f)
|
|
|
|
|
if approval_request_id in data:
|
|
|
|
|
raise ValueError(f"Approval request with ID '{approval_request_id}' already exists.")
|
|
|
|
|
data[approval_request_id] = request.to_dict()
|
|
|
|
|
self._atomic_write(data)
|
|
|
|
|
|
|
|
|
|
def _load_sync(self, approval_request_id: str) -> Content:
|
|
|
|
|
with self._lock:
|
|
|
|
|
self._create_storage_file_if_not_exists_sync()
|
|
|
|
|
with open(self._storage_path) as f:
|
|
|
|
|
data = json.load(f)
|
|
|
|
|
if approval_request_id not in data:
|
|
|
|
|
raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.")
|
|
|
|
|
return Content.from_dict(data[approval_request_id])
|
|
|
|
|
|
|
|
|
|
async def save_approval_request(self, approval_request_id: str, request: Content) -> None:
|
|
|
|
|
await asyncio.to_thread(self._save_sync, approval_request_id, request)
|
|
|
|
|
|
|
|
|
|
async def load_approval_request(self, approval_request_id: str) -> Content:
|
|
|
|
|
return await asyncio.to_thread(self._load_sync, approval_request_id)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ResponsesHostServer(ResponsesAgentServerHost):
|
|
|
|
|
"""A responses server host for an agent."""
|
|
|
|
|
|
|
|
|
|
# TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally
|
|
|
|
|
CHECKPOINT_STORAGE_PATH = "/.checkpoints"
|
|
|
|
|
FUNCTION_APPROVAL_STORAGE_PATH = "/.function_approvals/approval_requests.json"
|
|
|
|
|
|
|
|
|
|
def __init__(
|
|
|
|
|
self,
|
|
|
|
@@ -171,6 +268,11 @@ class ResponsesHostServer(ResponsesAgentServerHost):
|
|
|
|
|
self._is_workflow_agent = True
|
|
|
|
|
|
|
|
|
|
self._agent = agent
|
|
|
|
|
self._approval_storage = (
|
|
|
|
|
FileBasedFunctionApprovalStorage(self.FUNCTION_APPROVAL_STORAGE_PATH)
|
|
|
|
|
if self.config.is_hosted
|
|
|
|
|
else InMemoryFunctionApprovalStorage()
|
|
|
|
|
)
|
|
|
|
|
self.response_handler(self._handle_response) # pyright: ignore[reportUnknownMemberType]
|
|
|
|
|
|
|
|
|
|
async def _handle_response(
|
|
|
|
@@ -192,10 +294,15 @@ class ResponsesHostServer(ResponsesAgentServerHost):
|
|
|
|
|
) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]:
|
|
|
|
|
"""Handle the creation of a response for a regular (non-workflow) agent."""
|
|
|
|
|
input_items = await context.get_input_items()
|
|
|
|
|
input_messages = _items_to_messages(input_items)
|
|
|
|
|
input_messages = await _items_to_messages(input_items, approval_storage=self._approval_storage)
|
|
|
|
|
|
|
|
|
|
history = await context.get_history()
|
|
|
|
|
run_kwargs: dict[str, Any] = {"messages": [*_output_items_to_messages(history), *input_messages]}
|
|
|
|
|
run_kwargs: dict[str, Any] = {
|
|
|
|
|
"messages": [
|
|
|
|
|
*(await _output_items_to_messages(history, approval_storage=self._approval_storage)),
|
|
|
|
|
*input_messages,
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
is_streaming_request = request.stream is not None and request.stream is True
|
|
|
|
|
|
|
|
|
|
chat_options, are_options_set = _to_chat_options(request)
|
|
|
|
@@ -216,7 +323,11 @@ class ResponsesHostServer(ResponsesAgentServerHost):
|
|
|
|
|
|
|
|
|
|
for message in response.messages:
|
|
|
|
|
for content in message.contents:
|
|
|
|
|
async for item in _to_outputs(response_event_stream, content):
|
|
|
|
|
async for item in _to_outputs(
|
|
|
|
|
response_event_stream,
|
|
|
|
|
content,
|
|
|
|
|
approval_storage=self._approval_storage,
|
|
|
|
|
):
|
|
|
|
|
yield item
|
|
|
|
|
|
|
|
|
|
yield response_event_stream.emit_completed()
|
|
|
|
@@ -232,7 +343,11 @@ class ResponsesHostServer(ResponsesAgentServerHost):
|
|
|
|
|
for event in tracker.handle(content):
|
|
|
|
|
yield event
|
|
|
|
|
if tracker.needs_async:
|
|
|
|
|
async for item in _to_outputs(response_event_stream, content):
|
|
|
|
|
async for item in _to_outputs(
|
|
|
|
|
response_event_stream,
|
|
|
|
|
content,
|
|
|
|
|
approval_storage=self._approval_storage,
|
|
|
|
|
):
|
|
|
|
|
yield item
|
|
|
|
|
tracker.needs_async = False
|
|
|
|
|
|
|
|
|
@@ -254,7 +369,7 @@ class ResponsesHostServer(ResponsesAgentServerHost):
|
|
|
|
|
by the hosting infrastructure or files will be preserved upon deactivation.
|
|
|
|
|
"""
|
|
|
|
|
input_items = await context.get_input_items()
|
|
|
|
|
input_messages = _items_to_messages(input_items)
|
|
|
|
|
input_messages = await _items_to_messages(input_items)
|
|
|
|
|
is_streaming_request = request.stream is not None and request.stream is True
|
|
|
|
|
|
|
|
|
|
_, are_options_set = _to_chat_options(request)
|
|
|
|
@@ -581,26 +696,32 @@ def _to_chat_options(request: CreateResponse) -> tuple[ChatOptions, bool]:
|
|
|
|
|
# region Input Message Conversion
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _items_to_messages(input_items: Sequence[Item]) -> list[Message]:
|
|
|
|
|
async def _items_to_messages(
|
|
|
|
|
input_items: Sequence[Item], *, approval_storage: ApprovalStorage | None = None
|
|
|
|
|
) -> list[Message]:
|
|
|
|
|
"""Converts a sequence of input items to a list of Messages, one per item.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
input_items: The input items to convert.
|
|
|
|
|
approval_storage: An optional ApprovalStorage instance used to look up
|
|
|
|
|
approval requests when converting MCP approval response items.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
A list of Messages, one per supported input item.
|
|
|
|
|
"""
|
|
|
|
|
messages: list[Message] = []
|
|
|
|
|
for item in input_items:
|
|
|
|
|
messages.append(_item_to_message(item))
|
|
|
|
|
messages.append(await _item_to_message(item, approval_storage=approval_storage))
|
|
|
|
|
return messages
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _item_to_message(item: Item) -> Message:
|
|
|
|
|
async def _item_to_message(item: Item, *, approval_storage: ApprovalStorage | None = None) -> Message:
|
|
|
|
|
"""Converts an Item to a Message.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
item: The Item to convert.
|
|
|
|
|
approval_storage: An optional ApprovalStorage instance used to look up
|
|
|
|
|
approval requests when converting MCP approval response items.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
The converted Message.
|
|
|
|
@@ -659,27 +780,26 @@ def _item_to_message(item: Item) -> Message:
|
|
|
|
|
|
|
|
|
|
if item.type == "mcp_approval_request":
|
|
|
|
|
mcp_req = cast(ItemMcpApprovalRequest, item)
|
|
|
|
|
mcp_call_content = Content.from_mcp_server_tool_call(
|
|
|
|
|
mcp_req.id,
|
|
|
|
|
mcp_req.name,
|
|
|
|
|
server_name=mcp_req.server_label,
|
|
|
|
|
arguments=mcp_req.arguments,
|
|
|
|
|
)
|
|
|
|
|
if approval_storage is not None:
|
|
|
|
|
function_approval_request_content = await approval_storage.load_approval_request(mcp_req.id)
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("ApprovalStorage is required to load approval request.")
|
|
|
|
|
return Message(
|
|
|
|
|
role="assistant",
|
|
|
|
|
contents=[Content.from_function_approval_request(mcp_req.id, mcp_call_content)],
|
|
|
|
|
contents=[function_approval_request_content],
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
if item.type == "mcp_approval_response":
|
|
|
|
|
mcp_resp = cast(MCPApprovalResponse, item)
|
|
|
|
|
placeholder_content = Content.from_function_call(mcp_resp.approval_request_id, "mcp_approval")
|
|
|
|
|
if approval_storage is not None:
|
|
|
|
|
function_approval_request_content = await approval_storage.load_approval_request(
|
|
|
|
|
mcp_resp.approval_request_id
|
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("ApprovalStorage is required to load approval request.")
|
|
|
|
|
return Message(
|
|
|
|
|
role="user",
|
|
|
|
|
contents=[
|
|
|
|
|
Content.from_function_approval_response(
|
|
|
|
|
mcp_resp.approve, mcp_resp.approval_request_id, placeholder_content
|
|
|
|
|
)
|
|
|
|
|
],
|
|
|
|
|
contents=[function_approval_request_content.to_function_approval_response(mcp_resp.approve)],
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
if item.type == "code_interpreter_call":
|
|
|
|
@@ -846,26 +966,34 @@ def _item_to_message(item: Item) -> Message:
|
|
|
|
|
raise ValueError(f"Unsupported Item type: {item.type}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _output_items_to_messages(history: Sequence[OutputItem]) -> list[Message]:
|
|
|
|
|
async def _output_items_to_messages(
|
|
|
|
|
history: Sequence[OutputItem],
|
|
|
|
|
*,
|
|
|
|
|
approval_storage: ApprovalStorage | None = None,
|
|
|
|
|
) -> list[Message]:
|
|
|
|
|
"""Converts a sequence of OutputItem objects to a list of Message objects.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
history (Sequence[OutputItem]): The sequence of OutputItem objects to convert.
|
|
|
|
|
approval_storage (ApprovalStorage | None, optional): The approval storage to use for
|
|
|
|
|
resolving MCP approval requests. Defaults to None.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
list[Message]: The list of Message objects.
|
|
|
|
|
"""
|
|
|
|
|
messages: list[Message] = []
|
|
|
|
|
for item in history:
|
|
|
|
|
messages.append(_output_item_to_message(item))
|
|
|
|
|
messages.append(await _output_item_to_message(item, approval_storage=approval_storage))
|
|
|
|
|
return messages
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _output_item_to_message(item: OutputItem) -> Message:
|
|
|
|
|
async def _output_item_to_message(item: OutputItem, *, approval_storage: ApprovalStorage | None = None) -> Message:
|
|
|
|
|
"""Converts an OutputItem to a Message.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
item (OutputItem): The OutputItem to convert.
|
|
|
|
|
approval_storage (ApprovalStorage | None, optional): The approval storage to use for
|
|
|
|
|
resolving MCP approval requests. Defaults to None.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
Message: The converted Message.
|
|
|
|
@@ -922,24 +1050,27 @@ def _output_item_to_message(item: OutputItem) -> Message:
|
|
|
|
|
|
|
|
|
|
if item.type == "mcp_approval_request":
|
|
|
|
|
mcp_req = cast(OutputItemMcpApprovalRequest, item)
|
|
|
|
|
mcp_call_content = Content.from_mcp_server_tool_call(
|
|
|
|
|
mcp_req.id,
|
|
|
|
|
mcp_req.name,
|
|
|
|
|
server_name=mcp_req.server_label,
|
|
|
|
|
arguments=mcp_req.arguments,
|
|
|
|
|
)
|
|
|
|
|
if approval_storage is not None:
|
|
|
|
|
function_approval_request_content = await approval_storage.load_approval_request(mcp_req.id)
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("ApprovalStorage is required to load approval request.")
|
|
|
|
|
return Message(
|
|
|
|
|
role="assistant",
|
|
|
|
|
contents=[Content.from_function_approval_request(mcp_req.id, mcp_call_content)],
|
|
|
|
|
contents=[function_approval_request_content],
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
if item.type == "mcp_approval_response":
|
|
|
|
|
mcp_resp = cast(OutputItemMcpApprovalResponseResource, item)
|
|
|
|
|
# Build a placeholder function_call Content since the original call details are not available
|
|
|
|
|
placeholder_content = Content.from_function_call(mcp_resp.approval_request_id, "mcp_approval")
|
|
|
|
|
if approval_storage is not None:
|
|
|
|
|
function_approval_request_content = await approval_storage.load_approval_request(
|
|
|
|
|
mcp_resp.approval_request_id
|
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("ApprovalStorage is required to load approval request.")
|
|
|
|
|
|
|
|
|
|
return Message(
|
|
|
|
|
role="user",
|
|
|
|
|
contents=[Content.from_function_approval_response(mcp_resp.approve, mcp_resp.id, placeholder_content)],
|
|
|
|
|
contents=[function_approval_request_content.to_function_approval_response(mcp_resp.approve)],
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
if item.type == "code_interpreter_call":
|
|
|
|
@@ -1237,12 +1368,18 @@ def _arguments_to_str(arguments: str | Mapping[str, Any] | None) -> str:
|
|
|
|
|
return json.dumps(arguments)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def _to_outputs(stream: ResponseEventStream, content: Content) -> AsyncIterator[ResponseStreamEvent]:
|
|
|
|
|
async def _to_outputs(
|
|
|
|
|
stream: ResponseEventStream,
|
|
|
|
|
content: Content,
|
|
|
|
|
*,
|
|
|
|
|
approval_storage: ApprovalStorage | None = None,
|
|
|
|
|
) -> AsyncIterator[ResponseStreamEvent]:
|
|
|
|
|
"""Converts a Content object to an async sequence of ResponseStreamEvent objects.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
stream: The ResponseEventStream to use for building events.
|
|
|
|
|
content: The Content to convert.
|
|
|
|
|
approval_storage: An optional ApprovalStorage instance to use for saving and loading function approval requests.
|
|
|
|
|
|
|
|
|
|
Yields:
|
|
|
|
|
ResponseStreamEvent: The converted event objects.
|
|
|
|
@@ -1320,6 +1457,31 @@ async def _to_outputs(stream: ResponseEventStream, content: Content) -> AsyncIte
|
|
|
|
|
max_output_length=content.max_output_length,
|
|
|
|
|
):
|
|
|
|
|
yield event
|
|
|
|
|
elif content.type == "function_approval_request":
|
|
|
|
|
function_call: Content = content.function_call # type: ignore
|
|
|
|
|
server_label = function_call.additional_properties.get("server_label", "agent_framework")
|
|
|
|
|
request_saved = False
|
|
|
|
|
async for event in stream.aoutput_item_mcp_approval_request(
|
|
|
|
|
server_label,
|
|
|
|
|
function_call.name, # type: ignore
|
|
|
|
|
_arguments_to_str(function_call.arguments),
|
|
|
|
|
):
|
|
|
|
|
if approval_storage is not None and not request_saved:
|
|
|
|
|
# Extract the approval request ID generated by the infrastructure
|
|
|
|
|
# when the approval request item is added to the stream. Save the
|
|
|
|
|
# approval request to the approval storage so it can be retrieved later
|
|
|
|
|
# for round trips where the original approval request needs to be looked up.
|
|
|
|
|
item = getattr(event, "item", None)
|
|
|
|
|
if item is not None and getattr(item, "id", None) is not None:
|
|
|
|
|
approval_request_id = cast(str, item.id) # type: ignore
|
|
|
|
|
await approval_storage.save_approval_request(approval_request_id, content)
|
|
|
|
|
request_saved = True
|
|
|
|
|
yield event
|
|
|
|
|
if approval_storage is not None and not request_saved:
|
|
|
|
|
logger.warning(
|
|
|
|
|
"Approval request was not saved to approval storage because the approval request ID "
|
|
|
|
|
"could not be extracted from the stream event."
|
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
# Log a warning for unsupported content types instead of raising an error to avoid breaking the response stream.
|
|
|
|
|
logger.warning(f"Content type '{content.type}' is not supported yet. This is usually safe to ignore.")
|
|
|
|
|