From 8a37c6f66f04cdd5d9cb49188b0d2649df818768 Mon Sep 17 00:00:00 2001 From: jacobbabula Date: Sun, 16 Aug 2026 17:09:22 -0500 Subject: [PATCH] fix: generate IDs for Chat Completions messages Fixes #36 --- chatkit/agents.py | 33 ++++++++----- tests/test_agents.py | 115 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 12 deletions(-) diff --git a/chatkit/agents.py b/chatkit/agents.py index 7f435e3..daad683 100644 --- a/chatkit/agents.py +++ b/chatkit/agents.py @@ -91,6 +91,8 @@ ) from .widgets import Markdown, Text, WidgetRoot +_AGENTS_SDK_FAKE_RESPONSES_ID = "__fake_id__" + class ClientToolCall(BaseModel): """ @@ -501,6 +503,7 @@ async def stream_agent_response( thread = context.thread queue_iterator = _AsyncQueueIterator(context._events) produced_items = set() + synthetic_message_ids: dict[str, str] = {} streaming_thought: None | StreamingThoughtTracker = None # item_id -> content_index -> annotation count item_annotation_count: defaultdict[str, defaultdict[int, int]] = defaultdict( @@ -592,7 +595,7 @@ def end_workflow(item: WorkflowItem): continue content = await _convert_content(event.part, converter) yield ThreadItemUpdatedEvent( - item_id=event.item_id, + item_id=synthetic_message_ids.get(event.item_id, event.item_id), update=AssistantMessageContentPartAdded( content_index=event.content_index, content=content, @@ -600,7 +603,7 @@ def end_workflow(item: WorkflowItem): ) elif event.type == "response.output_text.delta": yield ThreadItemUpdatedEvent( - item_id=event.item_id, + item_id=synthetic_message_ids.get(event.item_id, event.item_id), update=AssistantMessageContentPartTextDelta( content_index=event.content_index, delta=event.delta, @@ -608,7 +611,7 @@ def end_workflow(item: WorkflowItem): ) elif event.type == "response.output_text.done": yield ThreadItemUpdatedEvent( - item_id=event.item_id, + item_id=synthetic_message_ids.get(event.item_id, event.item_id), update=AssistantMessageContentPartDone( content_index=event.content_index, content=AssistantMessageContent( @@ -620,16 +623,17 @@ def end_workflow(item: WorkflowItem): elif event.type == "response.output_text.annotation.added": annotation = await _convert_annotation(event.annotation, converter) if annotation: + item_id = synthetic_message_ids.get(event.item_id, event.item_id) # Manually track annotation indices per content part in case we drop an annotation that # we can't convert to our internal representation (e.g. missing filename). - annotation_index = item_annotation_count[event.item_id][ + annotation_index = item_annotation_count[item_id][ event.content_index ] - item_annotation_count[event.item_id][event.content_index] = ( + item_annotation_count[item_id][event.content_index] = ( annotation_index + 1 ) yield ThreadItemUpdatedEvent( - item_id=event.item_id, + item_id=item_id, update=AssistantMessageContentPartAnnotationAdded( content_index=event.content_index, annotation_index=annotation_index, @@ -651,11 +655,14 @@ def end_workflow(item: WorkflowItem): if item.type == "message": if ctx.workflow_item: yield end_workflow(ctx.workflow_item) - produced_items.add(item.id) + item_id = item.id + if item_id == _AGENTS_SDK_FAKE_RESPONSES_ID: + item_id = ctx.generate_id("message") + synthetic_message_ids[item.id] = item_id + produced_items.add(item_id) yield ThreadItemAddedEvent( item=AssistantMessageItem( - # Reusing the Responses message ID - id=item.id, + id=item_id, thread_id=thread.id, content=[ await _convert_content(c, converter) @@ -763,11 +770,13 @@ def end_workflow(item: WorkflowItem): elif event.type == "response.output_item.done": item = event.item if item.type == "message": - produced_items.add(item.id) + item_id = synthetic_message_ids.pop(item.id, item.id) + if item_id == _AGENTS_SDK_FAKE_RESPONSES_ID: + item_id = ctx.generate_id("message") + produced_items.add(item_id) yield ThreadItemDoneEvent( item=AssistantMessageItem( - # Reusing the Responses message ID - id=item.id, + id=item_id, thread_id=thread.id, content=[ await _convert_content(c, converter) diff --git a/tests/test_agents.py b/tests/test_agents.py index 1611162..5b626c8 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -922,6 +922,121 @@ async def test_stream_agent_response_maps_events(): assert future.done() is True +async def test_stream_agent_response_replaces_synthetic_chat_completion_ids(): + generated_ids = iter(["message_1", "message_2"]) + store = Mock() + store.generate_item_id = lambda item_type, thread, context: next(generated_ids) + store.load_thread_items = AsyncMock(return_value=Page()) + store.add_thread_item = AsyncMock() + context = AgentContext( + previous_response_id=None, thread=thread, store=store, request_context=None + ) + result = make_result() + + for sequence_number, text in enumerate(["First", "Second"]): + result.add_event( + RawResponsesStreamEvent( + type="raw_response_event", + data=ResponseOutputItemAddedEvent( + type="response.output_item.added", + item=ResponseOutputMessage( + id="__fake_id__", + content=[], + role="assistant", + status="in_progress", + type="message", + ), + output_index=0, + sequence_number=sequence_number * 5, + ), + ) + ) + result.add_event( + RawResponsesStreamEvent( + type="raw_response_event", + data=ResponseContentPartAddedEvent( + type="response.content_part.added", + part=ResponseOutputText( + type="output_text", + text="", + annotations=[], + ), + content_index=0, + item_id="__fake_id__", + output_index=0, + sequence_number=sequence_number * 5 + 1, + ), + ) + ) + result.add_event( + RawResponsesStreamEvent( + type="raw_response_event", + data=ResponseTextDeltaEvent( + type="response.output_text.delta", + delta=text, + content_index=0, + item_id="__fake_id__", + logprobs=[], + output_index=0, + sequence_number=sequence_number * 5 + 2, + ), + ) + ) + result.add_event( + RawResponsesStreamEvent( + type="raw_response_event", + data=ResponseTextDoneEvent( + type="response.output_text.done", + text=text, + content_index=0, + item_id="__fake_id__", + logprobs=[], + output_index=0, + sequence_number=sequence_number * 5 + 3, + ), + ) + ) + result.add_event( + RawResponsesStreamEvent( + type="raw_response_event", + data=ResponseOutputItemDoneEvent( + type="response.output_item.done", + item=ResponseOutputMessage( + id="__fake_id__", + content=[ + ResponseOutputText( + annotations=[], type="output_text", text=text + ) + ], + role="assistant", + status="completed", + type="message", + ), + output_index=0, + sequence_number=sequence_number * 5 + 4, + ), + ) + ) + + result.done() + events = await all_events(stream_agent_response(context, result)) + + added = [event for event in events if isinstance(event, ThreadItemAddedEvent)] + updated = [event for event in events if isinstance(event, ThreadItemUpdatedEvent)] + done = [event for event in events if isinstance(event, ThreadItemDoneEvent)] + + assert [event.item.id for event in added] == ["message_1", "message_2"] + assert [event.item_id for event in updated] == [ + "message_1", + "message_1", + "message_1", + "message_2", + "message_2", + "message_2", + ] + assert [event.item.id for event in done] == ["message_1", "message_2"] + + @pytest.mark.parametrize( "raw_event,expected_event", [