diff --git a/chatkit/agents.py b/chatkit/agents.py index 7f435e3..38e674b 100644 --- a/chatkit/agents.py +++ b/chatkit/agents.py @@ -438,8 +438,11 @@ async def _convert_content( async def _convert_annotation( raw_annotation: object, converter: ResponseStreamConverter ) -> Annotation | None: - # There is a bug in the OpenAPI client that sometimes parses the annotation delta event into the wrong class - # resulting into annotation being a dict or untyped object instead instead of a ResponsesAnnotation + # OpenAI clients can use event-specific models for annotation delta events. + # Normalize those models before validating them as response annotations. + if isinstance(raw_annotation, BaseModel): + raw_annotation = raw_annotation.model_dump() + annotation = TypeAdapter[ResponsesAnnotation](ResponsesAnnotation).validate_python( raw_annotation ) diff --git a/tests/test_agents.py b/tests/test_agents.py index 1611162..afdf3e6 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -3,7 +3,7 @@ from collections.abc import AsyncIterator from datetime import datetime from importlib import import_module -from typing import cast +from typing import Literal, cast from unittest.mock import AsyncMock, Mock import pytest @@ -57,6 +57,7 @@ ) from openai.types.responses.response_text_delta_event import ResponseTextDeltaEvent from openai.types.responses.response_text_done_event import ResponseTextDoneEvent +from pydantic import BaseModel from chatkit.agents import ( AgentContext, @@ -1128,6 +1129,14 @@ def add_annotation_event(annotation, sequence_number): async def test_stream_agent_response_annotation_added_normalizes_annotations(): + # OpenAI 3.x defines this event model separately from the response-content model. + class EventAnnotationURLCitation(BaseModel): + end_index: int + start_index: int + title: str + type: Literal["url_citation"] + url: str + context = AgentContext( previous_response_id=None, thread=thread, store=mock_store, request_context=None ) @@ -1174,13 +1183,13 @@ def add_annotation_event(annotation, sequence_number): sequence_number=1, ) add_annotation_event( - { - "type": "url_citation", - "url": "https://example.com", - "title": "Example", - "start_index": 1, - "end_index": 5, - }, + EventAnnotationURLCitation( + type="url_citation", + url="https://example.com", + title="Example", + start_index=1, + end_index=5, + ), sequence_number=2, )