diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index 27da4fea33..f4dc3ac70f 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -669,25 +669,75 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: async for chunk in completion_stream: if len(chunk.choices) == 0 and chunk.usage is None: continue - update = self._parse_response_update_from_openai(chunk) - for content in update.contents: - if content.type != "function_call": + # Emit usage as a separate update if present. + if chunk.usage: + usage_update = ChatResponseUpdate( + created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + contents=[ + Content.from_usage( + usage_details=self._parse_usage_from_openai(chunk.usage), + merge_mode="replace", + raw_representation=chunk, + ) + ], + role="assistant", + model=chunk.model, + additional_properties=self._get_metadata_from_streaming_chat_response(chunk), + finish_reason=None, + raw_representation=chunk, + response_id=chunk.id, + message_id=f"{chunk.id}-usage", + ) + yield usage_update + # Emit one update per choice to keep messages separate. + for choice in chunk.choices: + if choice.delta is None: continue - choice_index = content.additional_properties.get("tool_call_choice_index") - tool_index = content.additional_properties.get("tool_call_index") - if not isinstance(choice_index, int) or not isinstance(tool_index, int): - continue - index_key = (choice_index, tool_index) - identity = tool_call_identities.get(index_key) - if identity is None: - identity = (f"af-call-{uuid4().hex}", content.call_id or "") - occurrence_id, provider_call_id = identity - if content.call_id: - provider_call_id = content.call_id - tool_call_identities[index_key] = (occurrence_id, provider_call_id) - content.id = occurrence_id - content.call_id = provider_call_id - yield update + choice_contents: list[Content] = [] + choice_contents.extend(self._parse_tool_calls_from_openai(choice)) + if text_content := self._parse_text_from_openai(choice): + choice_contents.append(text_content) + if refusal_content := self._parse_refusal_from_openai(choice): + choice_contents.append(refusal_content) + if reasoning_details := getattr(choice.delta, "reasoning_details", None): + choice_contents.append( + Content.from_text_reasoning( + protected_data=json.dumps(reasoning_details), + merge_mode="append", + ) + ) + if self.response_parser is not None: + choice_contents = list(self.response_parser(choice.delta, choice_contents)) + choice_update = ChatResponseUpdate( + created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + contents=choice_contents, + role="assistant", + model=chunk.model, + additional_properties=self._get_metadata_from_streaming_chat_response(chunk), + finish_reason=choice.finish_reason if choice.finish_reason else None, + raw_representation=chunk, + response_id=chunk.id, + message_id=f"{chunk.id}-{choice.index}", + ) + # Preserve tool call identity mapping as before. + for content in choice_update.contents: + if content.type != "function_call": + continue + choice_index = content.additional_properties.get("tool_call_choice_index") + tool_index = content.additional_properties.get("tool_call_index") + if not isinstance(choice_index, int) or not isinstance(tool_index, int): + continue + index_key = (choice_index, tool_index) + identity = tool_call_identities.get(index_key) + if identity is None: + identity = (f"af-call-{uuid4().hex}", content.call_id or "") + occurrence_id, provider_call_id = identity + if content.call_id: + provider_call_id = content.call_id + tool_call_identities[index_key] = (occurrence_id, provider_call_id) + content.id = occurrence_id + content.call_id = provider_call_id + yield choice_update except BadRequestError as ex: if ex.code == "content_filter": raise OpenAIContentFilterException( diff --git a/repro_9677.py b/repro_9677.py new file mode 100644 index 0000000000..a5d7da351f --- /dev/null +++ b/repro_9677.py @@ -0,0 +1,84 @@ +import asyncio +from collections.abc import AsyncIterable +from agent_framework import ( + Agent, + AgentSession, + BaseChatClient, + ChatMiddlewareLayer, + ChatResponse, + ChatResponseUpdate, + Content, + FunctionInvocationLayer, + Message, + ResponseStream, + ToolApprovalMiddleware, + tool, +) + +executed = [] + +@tool +def lookup(query: str) -> str: + executed.append(query) + return f"result for {query}" + +class ScriptedClient(FunctionInvocationLayer, ChatMiddlewareLayer, BaseChatClient): + def __init__(self, turns, **kwargs): + super().__init__(**kwargs) + self.turns = list(turns) + self.tool_choices = [] + + def _next(self, options): + self.tool_choices.append(options.get("tool_choice")) + turn = self.turns.pop(0) + if isinstance(turn, Exception): + raise turn + return turn + + def _inner_get_response(self, *, messages, stream, options, **kwargs): + if stream: + async def updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[self._next(options)]) + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + async def get() -> ChatResponse: + return ChatResponse(messages=[Message(role="assistant", contents=[self._next(options)])]) + return get() + +async def main(stream: bool) -> None: + executed.clear() + client = ScriptedClient( + [ + Content.from_function_call(call_id="c1", name="lookup", arguments='{"query": "first"}'), + RuntimeError("transient provider error"), + Content.from_function_call(call_id="c2", name="lookup", arguments='{"query": "retry"}'), + Content.from_text("final answer"), + ], + function_invocation_configuration={"max_iterations": 2}, + ) + agent = Agent(client=client, tools=[lookup], middleware=[ToolApprovalMiddleware()]) + session = AgentSession() + + async def run(text: str): + if stream: + response_stream = agent.run(text, session=session, stream=True) + async for _ in response_stream: + pass + return await response_stream.get_final_response() + return await agent.run(text, session=session) + + print(f"stream={stream}") + try: + await run("look it up") + except RuntimeError as exc: + print(" run 1 raised:", exc) + budget = session.state.get("_function_invocation_budget_state") or {} + print(" budget left in session:", {k: budget[k] for k in ("attempt_count", "total_function_calls") if k in budget}) + calls_before = len(client.tool_choices) + response = await run("look it up again") + print(" run 2 tool_choice per provider call:", client.tool_choices[calls_before:]) + print(" tools executed:", executed) + print(" run 2 text:", repr(response.text)) + +asyncio.run(main(stream=False)) +asyncio.run(main(stream=True))