mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Add long-running agents and background responses support (#3808)
* Python: Add long-running agents and background responses support - Add ContinuationToken TypedDict to core types - Add continuation_token field to ChatResponse, ChatResponseUpdate, AgentResponse, and AgentResponseUpdate - Add background and continuation_token options to OpenAIResponsesOptions - Implement polling via responses.retrieve() and streaming resumption in RawOpenAIResponsesClient - Propagate continuation tokens through agent run() and map_chat_to_agent_update - Fix streaming telemetry 'Failed to detach context' error in both ChatTelemetryLayer and AgentTelemetryLayer by avoiding trace.use_span() context attachment for async-managed spans - Add 14 unit tests for continuation token types and background flows - Add background_responses sample showing polling and stream resumption Fixes #2478 * Python: Add A2A long-running task support via ContinuationToken - Make ContinuationToken provider-agnostic (total=False, optional task_id/context_id fields) - Add background param to A2AAgent.run() controlling token emission - Add poll_task() for single-request task state retrieval - Add resubscribe support via continuation_token param on run() - Extract _updates_from_task() and _map_a2a_stream() for cleaner code - Streamline run()/streaming by removing intermediate _stream_updates wrapper - Update A2A sample to show background=False (default) with link to background_responses sample - Remove stale BareAgent from __all__ - Add 12 new A2A continuation token tests * fix logic for overriding continuation token when done * refactored ContinuationToken setup
This commit is contained in:
committed by
GitHub
Unverified
parent
32ba81e990
commit
35097d8c75
@@ -165,7 +165,7 @@ class _RunContext(TypedDict):
|
||||
finalize_kwargs: dict[str, Any]
|
||||
|
||||
|
||||
__all__ = ["BareAgent", "BaseAgent", "ChatAgent", "RawChatAgent", "SupportsAgentRun"]
|
||||
__all__ = ["BaseAgent", "ChatAgent", "RawChatAgent", "SupportsAgentRun"]
|
||||
|
||||
|
||||
# region Agent Protocol
|
||||
@@ -523,10 +523,6 @@ class BaseAgent(SerializationMixin):
|
||||
return agent_tool
|
||||
|
||||
|
||||
# Backward compatibility alias
|
||||
BareAgent = BaseAgent
|
||||
|
||||
|
||||
# region ChatAgent
|
||||
|
||||
|
||||
@@ -908,6 +904,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
|
||||
usage_details=response.usage_details,
|
||||
value=response.value,
|
||||
response_format=response_format,
|
||||
continuation_token=response.continuation_token,
|
||||
raw_representation=response,
|
||||
additional_properties=response.additional_properties,
|
||||
)
|
||||
|
||||
@@ -36,6 +36,7 @@ __all__ = [
|
||||
"ChatResponse",
|
||||
"ChatResponseUpdate",
|
||||
"Content",
|
||||
"ContinuationToken",
|
||||
"FinalT",
|
||||
"FinishReason",
|
||||
"FinishReasonLiteral",
|
||||
@@ -1760,6 +1761,7 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse
|
||||
response.finish_reason = update.finish_reason
|
||||
if update.model_id is not None:
|
||||
response.model_id = update.model_id
|
||||
response.continuation_token = update.continuation_token
|
||||
|
||||
|
||||
def _coalesce_text_content(contents: list[Content], type_str: Literal["text", "text_reasoning"]) -> None:
|
||||
@@ -1796,6 +1798,39 @@ def _finalize_response(response: ChatResponse | AgentResponse) -> None:
|
||||
_coalesce_text_content(msg.contents, "text_reasoning")
|
||||
|
||||
|
||||
# region ContinuationToken
|
||||
|
||||
|
||||
class ContinuationToken(TypedDict):
|
||||
"""Opaque token for resuming long-running agent operations.
|
||||
|
||||
A JSON-serializable dict used to poll for completion or resume a
|
||||
streaming response. Presence on a response indicates the operation
|
||||
is still in progress; ``None`` means the operation is complete.
|
||||
|
||||
Each provider subclasses this with its own fields; consumers should
|
||||
treat the token as opaque and simply pass it back to the same agent.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
import json
|
||||
|
||||
# Persist token across restarts
|
||||
token_json = json.dumps(response.continuation_token)
|
||||
|
||||
# Restore and resume
|
||||
token = json.loads(token_json)
|
||||
response = await agent.run(
|
||||
thread=thread,
|
||||
options={"continuation_token": token},
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
"""Represents the response to a chat request.
|
||||
|
||||
@@ -1861,6 +1896,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
usage_details: UsageDetails | None = None,
|
||||
value: ResponseModelT | None = None,
|
||||
response_format: type[BaseModel] | None = None,
|
||||
continuation_token: ContinuationToken | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
raw_representation: Any | None = None,
|
||||
) -> None:
|
||||
@@ -1876,6 +1912,8 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
usage_details: Optional usage details for the chat response.
|
||||
value: Optional value of the structured output.
|
||||
response_format: Optional response format for the chat response.
|
||||
continuation_token: Optional token for resuming a long-running background operation.
|
||||
When present, indicates the operation is still in progress.
|
||||
additional_properties: Optional additional properties associated with the chat response.
|
||||
raw_representation: Optional raw representation of the chat response from an underlying implementation.
|
||||
"""
|
||||
@@ -1907,6 +1945,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
self._response_format: type[BaseModel] | None = response_format
|
||||
self._value_parsed: bool = value is not None
|
||||
self.additional_properties = additional_properties or {}
|
||||
self.continuation_token = continuation_token
|
||||
self.raw_representation: Any | list[Any] | None = raw_representation
|
||||
|
||||
@overload
|
||||
@@ -2109,6 +2148,7 @@ class ChatResponseUpdate(SerializationMixin):
|
||||
model_id: str | None = None,
|
||||
created_at: CreatedAtT | None = None,
|
||||
finish_reason: FinishReasonLiteral | FinishReason | None = None,
|
||||
continuation_token: ContinuationToken | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
raw_representation: Any | None = None,
|
||||
) -> None:
|
||||
@@ -2124,6 +2164,8 @@ class ChatResponseUpdate(SerializationMixin):
|
||||
model_id: Optional model ID associated with this response update.
|
||||
created_at: Optional timestamp for the chat response update.
|
||||
finish_reason: Optional finish reason for the operation.
|
||||
continuation_token: Optional token for resuming a long-running background operation.
|
||||
When present, indicates the operation is still in progress.
|
||||
additional_properties: Optional additional properties associated with the chat response update.
|
||||
raw_representation: Optional raw representation of the chat response update
|
||||
from an underlying implementation.
|
||||
@@ -2151,6 +2193,7 @@ class ChatResponseUpdate(SerializationMixin):
|
||||
self.model_id = model_id
|
||||
self.created_at = created_at
|
||||
self.finish_reason = finish_reason
|
||||
self.continuation_token = continuation_token
|
||||
self.additional_properties = additional_properties
|
||||
self.raw_representation = raw_representation
|
||||
|
||||
@@ -2222,6 +2265,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
usage_details: UsageDetails | None = None,
|
||||
value: ResponseModelT | None = None,
|
||||
response_format: type[BaseModel] | None = None,
|
||||
continuation_token: ContinuationToken | None = None,
|
||||
raw_representation: Any | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
@@ -2236,6 +2280,8 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
usage_details: The usage details for the chat response.
|
||||
value: The structured output of the agent run response, if applicable.
|
||||
response_format: Optional response format for the agent response.
|
||||
continuation_token: Optional token for resuming a long-running background operation.
|
||||
When present, indicates the operation is still in progress.
|
||||
additional_properties: Any additional properties associated with the chat response.
|
||||
raw_representation: The raw representation of the chat response from an underlying implementation.
|
||||
"""
|
||||
@@ -2262,6 +2308,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
|
||||
self._response_format: type[BaseModel] | None = response_format
|
||||
self._value_parsed: bool = value is not None
|
||||
self.additional_properties = additional_properties or {}
|
||||
self.continuation_token = continuation_token
|
||||
self.raw_representation = raw_representation
|
||||
|
||||
@property
|
||||
@@ -2444,6 +2491,7 @@ class AgentResponseUpdate(SerializationMixin):
|
||||
response_id: str | None = None,
|
||||
message_id: str | None = None,
|
||||
created_at: CreatedAtT | None = None,
|
||||
continuation_token: ContinuationToken | None = None,
|
||||
additional_properties: dict[str, Any] | None = None,
|
||||
raw_representation: Any | None = None,
|
||||
) -> None:
|
||||
@@ -2458,6 +2506,8 @@ class AgentResponseUpdate(SerializationMixin):
|
||||
response_id: Optional ID of the response of which this update is a part.
|
||||
message_id: Optional ID of the message of which this update is a part.
|
||||
created_at: Optional timestamp for the chat response update.
|
||||
continuation_token: Optional token for resuming a long-running background operation.
|
||||
When present, indicates the operation is still in progress.
|
||||
additional_properties: Optional additional properties associated with the chat response update.
|
||||
raw_representation: Optional raw representation of the chat response update.
|
||||
|
||||
@@ -2486,6 +2536,7 @@ class AgentResponseUpdate(SerializationMixin):
|
||||
self.response_id = response_id
|
||||
self.message_id = message_id
|
||||
self.created_at = created_at
|
||||
self.continuation_token = continuation_token
|
||||
self.additional_properties = additional_properties
|
||||
self.raw_representation: Any | list[Any] | None = raw_representation
|
||||
|
||||
@@ -2514,6 +2565,7 @@ def map_chat_to_agent_update(update: ChatResponseUpdate, agent_name: str | None)
|
||||
response_id=update.response_id,
|
||||
message_id=update.message_id,
|
||||
created_at=update.created_at,
|
||||
continuation_token=update.continuation_token,
|
||||
additional_properties=update.additional_properties,
|
||||
raw_representation=update,
|
||||
)
|
||||
|
||||
@@ -1139,8 +1139,14 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
|
||||
else:
|
||||
raise RuntimeError("Streaming telemetry requires a ResponseStream result.")
|
||||
|
||||
span_cm = _get_span(attributes=attributes, span_name_attribute=SpanAttributes.LLM_REQUEST_MODEL)
|
||||
span = span_cm.__enter__()
|
||||
# Create span directly without trace.use_span() context attachment.
|
||||
# Streaming spans are closed asynchronously in cleanup hooks, which run
|
||||
# in a different async context than creation — using use_span() would
|
||||
# cause "Failed to detach context" errors from OpenTelemetry.
|
||||
operation = attributes.get(OtelAttr.OPERATION, "operation")
|
||||
span_name = attributes.get(SpanAttributes.LLM_REQUEST_MODEL, "unknown")
|
||||
span = get_tracer().start_span(f"{operation} {span_name}")
|
||||
span.set_attributes(attributes)
|
||||
if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages:
|
||||
_capture_messages(
|
||||
span=span,
|
||||
@@ -1157,7 +1163,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
|
||||
if span_state["closed"]:
|
||||
return
|
||||
span_state["closed"] = True
|
||||
span_cm.__exit__(None, None, None)
|
||||
span.end()
|
||||
|
||||
def _record_duration() -> None:
|
||||
duration_state["duration"] = perf_counter() - start_time
|
||||
@@ -1326,8 +1332,14 @@ class AgentTelemetryLayer:
|
||||
else:
|
||||
raise RuntimeError("Streaming telemetry requires a ResponseStream result.")
|
||||
|
||||
span_cm = _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME)
|
||||
span = span_cm.__enter__()
|
||||
# Create span directly without trace.use_span() context attachment.
|
||||
# Streaming spans are closed asynchronously in cleanup hooks, which run
|
||||
# in a different async context than creation — using use_span() would
|
||||
# cause "Failed to detach context" errors from OpenTelemetry.
|
||||
operation = attributes.get(OtelAttr.OPERATION, "operation")
|
||||
span_name = attributes.get(OtelAttr.AGENT_NAME, "unknown")
|
||||
span = get_tracer().start_span(f"{operation} {span_name}")
|
||||
span.set_attributes(attributes)
|
||||
if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages:
|
||||
_capture_messages(
|
||||
span=span,
|
||||
@@ -1344,7 +1356,7 @@ class AgentTelemetryLayer:
|
||||
if span_state["closed"]:
|
||||
return
|
||||
span_state["closed"] = True
|
||||
span_cm.__exit__(None, None, None)
|
||||
span.end()
|
||||
|
||||
def _record_duration() -> None:
|
||||
duration_state["duration"] = perf_counter() - start_time
|
||||
|
||||
@@ -56,6 +56,7 @@ from .._types import (
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
ContinuationToken,
|
||||
ResponseStream,
|
||||
Role,
|
||||
TextSpanRegion,
|
||||
@@ -98,7 +99,14 @@ if TYPE_CHECKING:
|
||||
logger = get_logger("agent_framework.openai")
|
||||
|
||||
|
||||
__all__ = ["OpenAIResponsesClient", "OpenAIResponsesOptions", "RawOpenAIResponsesClient"]
|
||||
__all__ = ["OpenAIContinuationToken", "OpenAIResponsesClient", "OpenAIResponsesOptions", "RawOpenAIResponsesClient"]
|
||||
|
||||
|
||||
class OpenAIContinuationToken(ContinuationToken):
|
||||
"""Continuation token for OpenAI Responses API background operations."""
|
||||
|
||||
response_id: str
|
||||
"""OpenAI Responses API response ID."""
|
||||
|
||||
|
||||
# region OpenAI Responses Options TypedDict
|
||||
@@ -190,6 +198,17 @@ class OpenAIResponsesOptions(ChatOptions[ResponseFormatT], Generic[ResponseForma
|
||||
- 'auto': Truncate from beginning if exceeds context
|
||||
- 'disabled': Fail with 400 error if exceeds context"""
|
||||
|
||||
background: bool
|
||||
"""Whether to run the model response in the background.
|
||||
When True, the response returns immediately with a continuation token
|
||||
that can be used to poll for the result.
|
||||
See: https://platform.openai.com/docs/guides/background"""
|
||||
|
||||
continuation_token: OpenAIContinuationToken
|
||||
"""Token for resuming or polling a long-running background operation.
|
||||
Pass the ``continuation_token`` from a previous response to poll for
|
||||
completion or resume a streaming response."""
|
||||
|
||||
|
||||
OpenAIResponsesOptionsT = TypeVar(
|
||||
"OpenAIResponsesOptionsT",
|
||||
@@ -266,33 +285,60 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
stream: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
|
||||
continuation_token: OpenAIContinuationToken | None = options.get("continuation_token") # type: ignore[assignment]
|
||||
|
||||
if stream:
|
||||
function_call_ids: dict[int, tuple[str, str]] = {}
|
||||
validated_options: dict[str, Any] | None = None
|
||||
|
||||
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
|
||||
nonlocal validated_options
|
||||
client, run_options, validated_options = await self._prepare_request(messages, options, **kwargs)
|
||||
try:
|
||||
if "text_format" in run_options:
|
||||
async with client.responses.stream(**run_options) as response:
|
||||
async for chunk in response:
|
||||
yield self._parse_chunk_from_openai(
|
||||
chunk, options=validated_options, function_call_ids=function_call_ids
|
||||
)
|
||||
else:
|
||||
async for chunk in await client.responses.create(stream=True, **run_options):
|
||||
if continuation_token is not None:
|
||||
# Resume a background streaming response by retrieving with stream=True
|
||||
client = await self._ensure_client()
|
||||
validated_options = await self._validate_options(options)
|
||||
try:
|
||||
stream_response = await client.responses.retrieve(
|
||||
continuation_token["response_id"],
|
||||
stream=True,
|
||||
)
|
||||
async for chunk in stream_response:
|
||||
yield self._parse_chunk_from_openai(
|
||||
chunk, options=validated_options, function_call_ids=function_call_ids
|
||||
)
|
||||
except Exception as ex:
|
||||
self._handle_request_error(ex)
|
||||
except Exception as ex:
|
||||
self._handle_request_error(ex)
|
||||
else:
|
||||
client, run_options, validated_options = await self._prepare_request(messages, options, **kwargs)
|
||||
try:
|
||||
if "text_format" in run_options:
|
||||
async with client.responses.stream(**run_options) as response:
|
||||
async for chunk in response:
|
||||
yield self._parse_chunk_from_openai(
|
||||
chunk, options=validated_options, function_call_ids=function_call_ids
|
||||
)
|
||||
else:
|
||||
async for chunk in await client.responses.create(stream=True, **run_options):
|
||||
yield self._parse_chunk_from_openai(
|
||||
chunk, options=validated_options, function_call_ids=function_call_ids
|
||||
)
|
||||
except Exception as ex:
|
||||
self._handle_request_error(ex)
|
||||
|
||||
response_format = validated_options.get("response_format") if validated_options else None
|
||||
return self._build_response_stream(_stream(), response_format=response_format)
|
||||
|
||||
# Non-streaming
|
||||
async def _get_response() -> ChatResponse:
|
||||
if continuation_token is not None:
|
||||
# Poll a background response by retrieving without stream
|
||||
client = await self._ensure_client()
|
||||
validated_options = await self._validate_options(options)
|
||||
try:
|
||||
response = await client.responses.retrieve(continuation_token["response_id"])
|
||||
except Exception as ex:
|
||||
self._handle_request_error(ex)
|
||||
return self._parse_response_from_openai(response, options=validated_options)
|
||||
client, run_options, validated_options = await self._prepare_request(messages, options, **kwargs)
|
||||
try:
|
||||
if "text_format" in run_options:
|
||||
@@ -538,6 +584,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
"response_format", # handled separately
|
||||
"conversation_id", # handled separately
|
||||
"tool_choice", # handled separately
|
||||
"continuation_token", # handled separately in _inner_get_response
|
||||
}
|
||||
run_options: dict[str, Any] = {k: v for k, v in options.items() if k not in exclude_keys and v is not None}
|
||||
|
||||
@@ -1070,6 +1117,9 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
# Only pass response_format to ChatResponse if it's a Pydantic model type,
|
||||
# not a runtime JSON schema dict
|
||||
args["response_format"] = response_format
|
||||
# Set continuation_token when background operation is still in progress
|
||||
if response.status and response.status in ("in_progress", "queued"):
|
||||
args["continuation_token"] = OpenAIContinuationToken(response_id=response.id)
|
||||
return ChatResponse(**args)
|
||||
|
||||
def _parse_chunk_from_openai(
|
||||
@@ -1083,6 +1133,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
contents: list[Content] = []
|
||||
conversation_id: str | None = None
|
||||
response_id: str | None = None
|
||||
continuation_token: OpenAIContinuationToken | None = None
|
||||
model = self.model_id
|
||||
match event.type:
|
||||
# types:
|
||||
@@ -1211,9 +1262,12 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
case "response.created":
|
||||
response_id = event.response.id
|
||||
conversation_id = self._get_conversation_id(event.response, options.get("store"))
|
||||
if event.response.status and event.response.status in ("in_progress", "queued"):
|
||||
continuation_token = OpenAIContinuationToken(response_id=event.response.id)
|
||||
case "response.in_progress":
|
||||
response_id = event.response.id
|
||||
conversation_id = self._get_conversation_id(event.response, options.get("store"))
|
||||
continuation_token = OpenAIContinuationToken(response_id=event.response.id)
|
||||
case "response.completed":
|
||||
response_id = event.response.id
|
||||
conversation_id = self._get_conversation_id(event.response, options.get("store"))
|
||||
@@ -1454,6 +1508,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
|
||||
response_id=response_id,
|
||||
role="assistant",
|
||||
model_id=model,
|
||||
continuation_token=continuation_token,
|
||||
additional_properties=metadata,
|
||||
raw_representation=event,
|
||||
)
|
||||
|
||||
@@ -2434,3 +2434,263 @@ async def test_integration_streaming_file_search() -> None:
|
||||
|
||||
assert "sunny" in full_message.lower()
|
||||
assert "75" in full_message
|
||||
|
||||
|
||||
# region Background Response / ContinuationToken Tests
|
||||
|
||||
|
||||
def test_continuation_token_json_serializable() -> None:
|
||||
"""Test that OpenAIContinuationToken is a plain dict and JSON-serializable."""
|
||||
from agent_framework.openai import OpenAIContinuationToken
|
||||
|
||||
token = OpenAIContinuationToken(response_id="resp_abc123")
|
||||
assert token["response_id"] == "resp_abc123"
|
||||
|
||||
# JSON round-trip
|
||||
serialized = json.dumps(token)
|
||||
restored = json.loads(serialized)
|
||||
assert restored["response_id"] == "resp_abc123"
|
||||
|
||||
|
||||
def test_chat_response_with_continuation_token() -> None:
|
||||
"""Test that ChatResponse accepts and stores continuation_token."""
|
||||
from agent_framework.openai import OpenAIContinuationToken
|
||||
|
||||
token = OpenAIContinuationToken(response_id="resp_123")
|
||||
response = ChatResponse(
|
||||
messages=ChatMessage(role="assistant", contents=[Content.from_text(text="Hello")]),
|
||||
response_id="resp_123",
|
||||
continuation_token=token,
|
||||
)
|
||||
assert response.continuation_token is not None
|
||||
assert response.continuation_token["response_id"] == "resp_123"
|
||||
|
||||
|
||||
def test_chat_response_without_continuation_token() -> None:
|
||||
"""Test that ChatResponse defaults continuation_token to None."""
|
||||
response = ChatResponse(
|
||||
messages=ChatMessage(role="assistant", contents=[Content.from_text(text="Hello")]),
|
||||
)
|
||||
assert response.continuation_token is None
|
||||
|
||||
|
||||
def test_chat_response_update_with_continuation_token() -> None:
|
||||
"""Test that ChatResponseUpdate accepts and stores continuation_token."""
|
||||
from agent_framework.openai import OpenAIContinuationToken
|
||||
|
||||
token = OpenAIContinuationToken(response_id="resp_456")
|
||||
update = ChatResponseUpdate(
|
||||
contents=[Content.from_text(text="chunk")],
|
||||
role="assistant",
|
||||
continuation_token=token,
|
||||
)
|
||||
assert update.continuation_token is not None
|
||||
assert update.continuation_token["response_id"] == "resp_456"
|
||||
|
||||
|
||||
def test_agent_response_with_continuation_token() -> None:
|
||||
"""Test that AgentResponse accepts and stores continuation_token."""
|
||||
from agent_framework import AgentResponse
|
||||
from agent_framework.openai import OpenAIContinuationToken
|
||||
|
||||
token = OpenAIContinuationToken(response_id="resp_789")
|
||||
response = AgentResponse(
|
||||
messages=ChatMessage(role="assistant", contents=[Content.from_text(text="done")]),
|
||||
continuation_token=token,
|
||||
)
|
||||
assert response.continuation_token is not None
|
||||
assert response.continuation_token["response_id"] == "resp_789"
|
||||
|
||||
|
||||
def test_agent_response_update_with_continuation_token() -> None:
|
||||
"""Test that AgentResponseUpdate accepts and stores continuation_token."""
|
||||
from agent_framework import AgentResponseUpdate
|
||||
from agent_framework.openai import OpenAIContinuationToken
|
||||
|
||||
token = OpenAIContinuationToken(response_id="resp_012")
|
||||
update = AgentResponseUpdate(
|
||||
contents=[Content.from_text(text="streaming")],
|
||||
role="assistant",
|
||||
continuation_token=token,
|
||||
)
|
||||
assert update.continuation_token is not None
|
||||
assert update.continuation_token["response_id"] == "resp_012"
|
||||
|
||||
|
||||
def test_parse_response_from_openai_with_background_in_progress() -> None:
|
||||
"""Test that _parse_response_from_openai sets continuation_token when status is in_progress."""
|
||||
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.output_parsed = None
|
||||
mock_response.metadata = {}
|
||||
mock_response.usage = None
|
||||
mock_response.id = "resp_bg_123"
|
||||
mock_response.model = "test-model"
|
||||
mock_response.created_at = 1000000000
|
||||
mock_response.status = "in_progress"
|
||||
|
||||
mock_message = MagicMock()
|
||||
mock_message.type = "message"
|
||||
mock_message.content = []
|
||||
mock_response.output = [mock_message]
|
||||
|
||||
options: dict[str, Any] = {"store": False}
|
||||
result = client._parse_response_from_openai(mock_response, options=options)
|
||||
|
||||
assert result.continuation_token is not None
|
||||
assert result.continuation_token["response_id"] == "resp_bg_123"
|
||||
|
||||
|
||||
def test_parse_response_from_openai_with_background_queued() -> None:
|
||||
"""Test that _parse_response_from_openai sets continuation_token when status is queued."""
|
||||
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.output_parsed = None
|
||||
mock_response.metadata = {}
|
||||
mock_response.usage = None
|
||||
mock_response.id = "resp_bg_456"
|
||||
mock_response.model = "test-model"
|
||||
mock_response.created_at = 1000000000
|
||||
mock_response.status = "queued"
|
||||
|
||||
mock_message = MagicMock()
|
||||
mock_message.type = "message"
|
||||
mock_message.content = []
|
||||
mock_response.output = [mock_message]
|
||||
|
||||
options: dict[str, Any] = {"store": False}
|
||||
result = client._parse_response_from_openai(mock_response, options=options)
|
||||
|
||||
assert result.continuation_token is not None
|
||||
assert result.continuation_token["response_id"] == "resp_bg_456"
|
||||
|
||||
|
||||
def test_parse_response_from_openai_with_background_completed() -> None:
|
||||
"""Test that _parse_response_from_openai does NOT set continuation_token when status is completed."""
|
||||
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.output_parsed = None
|
||||
mock_response.metadata = {}
|
||||
mock_response.usage = None
|
||||
mock_response.id = "resp_bg_789"
|
||||
mock_response.model = "test-model"
|
||||
mock_response.created_at = 1000000000
|
||||
mock_response.status = "completed"
|
||||
|
||||
mock_text_content = MagicMock()
|
||||
mock_text_content.type = "output_text"
|
||||
mock_text_content.text = "Final answer"
|
||||
mock_text_content.annotations = []
|
||||
mock_text_content.logprobs = None
|
||||
|
||||
mock_message = MagicMock()
|
||||
mock_message.type = "message"
|
||||
mock_message.content = [mock_text_content]
|
||||
mock_response.output = [mock_message]
|
||||
|
||||
options: dict[str, Any] = {"store": False}
|
||||
result = client._parse_response_from_openai(mock_response, options=options)
|
||||
|
||||
assert result.continuation_token is None
|
||||
|
||||
|
||||
def test_streaming_response_in_progress_sets_continuation_token() -> None:
|
||||
"""Test that _parse_chunk_from_openai sets continuation_token for in_progress events."""
|
||||
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
|
||||
chat_options: dict[str, Any] = {}
|
||||
function_call_ids: dict[int, tuple[str, str]] = {}
|
||||
|
||||
mock_event = MagicMock()
|
||||
mock_event.type = "response.in_progress"
|
||||
mock_event.response = MagicMock()
|
||||
mock_event.response.id = "resp_stream_123"
|
||||
mock_event.response.conversation = MagicMock()
|
||||
mock_event.response.conversation.id = "conv_456"
|
||||
mock_event.response.status = "in_progress"
|
||||
|
||||
update = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
|
||||
|
||||
assert update.continuation_token is not None
|
||||
assert update.continuation_token["response_id"] == "resp_stream_123"
|
||||
|
||||
|
||||
def test_streaming_response_created_with_in_progress_status_sets_continuation_token() -> None:
|
||||
"""Test that response.created with in_progress status sets continuation_token."""
|
||||
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
|
||||
chat_options: dict[str, Any] = {}
|
||||
function_call_ids: dict[int, tuple[str, str]] = {}
|
||||
|
||||
mock_event = MagicMock()
|
||||
mock_event.type = "response.created"
|
||||
mock_event.response = MagicMock()
|
||||
mock_event.response.id = "resp_created_123"
|
||||
mock_event.response.conversation = MagicMock()
|
||||
mock_event.response.conversation.id = "conv_789"
|
||||
mock_event.response.status = "in_progress"
|
||||
|
||||
update = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
|
||||
|
||||
assert update.continuation_token is not None
|
||||
assert update.continuation_token["response_id"] == "resp_created_123"
|
||||
|
||||
|
||||
def test_streaming_response_completed_no_continuation_token() -> None:
|
||||
"""Test that response.completed does NOT set continuation_token."""
|
||||
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
|
||||
chat_options: dict[str, Any] = {}
|
||||
function_call_ids: dict[int, tuple[str, str]] = {}
|
||||
|
||||
mock_event = MagicMock()
|
||||
mock_event.type = "response.completed"
|
||||
mock_event.response = MagicMock()
|
||||
mock_event.response.id = "resp_done_123"
|
||||
mock_event.response.conversation = MagicMock()
|
||||
mock_event.response.conversation.id = "conv_done"
|
||||
mock_event.response.model = "test-model"
|
||||
mock_event.response.usage = None
|
||||
|
||||
update = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
|
||||
|
||||
assert update.continuation_token is None
|
||||
|
||||
|
||||
def test_map_chat_to_agent_update_preserves_continuation_token() -> None:
|
||||
"""Test that map_chat_to_agent_update propagates continuation_token."""
|
||||
from agent_framework._types import map_chat_to_agent_update
|
||||
|
||||
token = {"response_id": "resp_map_123"}
|
||||
chat_update = ChatResponseUpdate(
|
||||
contents=[Content.from_text(text="chunk")],
|
||||
role="assistant",
|
||||
response_id="resp_map_123",
|
||||
continuation_token=token,
|
||||
)
|
||||
|
||||
agent_update = map_chat_to_agent_update(chat_update, agent_name="test-agent")
|
||||
|
||||
assert agent_update.continuation_token is not None
|
||||
assert agent_update.continuation_token["response_id"] == "resp_map_123"
|
||||
|
||||
|
||||
async def test_prepare_options_excludes_continuation_token() -> None:
|
||||
"""Test that _prepare_options does not pass continuation_token to OpenAI API."""
|
||||
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
|
||||
|
||||
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
|
||||
options: dict[str, Any] = {
|
||||
"model_id": "test-model",
|
||||
"continuation_token": {"response_id": "resp_123"},
|
||||
"background": True,
|
||||
}
|
||||
|
||||
run_options = await client._prepare_options(messages, options)
|
||||
|
||||
assert "continuation_token" not in run_options
|
||||
assert "background" in run_options
|
||||
assert run_options["background"] is True
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
Reference in New Issue
Block a user