Skip to content

Python: fix: streamed chat completions with n>1 produce separate messages (Fixes #9779) #9805

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Dhruv Shaswat (Dhruv-Sh-git) wants to merge 1 commit into from
Closed

Python: fix: streamed chat completions with n>1 produce separate messages (Fixes #9779) #9805

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 68 additions & 18 deletions python/packages/openai/agent_framework_openai/_chat_completion_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
84 changes: 84 additions & 0 deletions repro_9677.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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))
Loading