GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/chatkit-python

Public

mirrored from https://github.com/openai/chatkit-pythonAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
800a66c9da1c80e59d9ac0e012b8fc5103cae4b7

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

chatkit-python/tests

tests/test_agents.py

1682lines · modecode

RawDownload
Latest commit unavailable.
1import asyncio
2import json
3from collections.abc import AsyncIterator
4from datetime import datetime
5from typing import cast
6from unittest.mock import AsyncMock, Mock
7
8import pytest
9from agents import (
10 Agent,
11 GuardrailFunctionOutput,
12 InputGuardrail,
13 InputGuardrailResult,
14 InputGuardrailTripwireTriggered,
15 OutputGuardrail,
16 OutputGuardrailResult,
17 OutputGuardrailTripwireTriggered,
18 RawResponsesStreamEvent,
19 RunContextWrapper,
20 RunItemStreamEvent,
21 Runner,
22 RunResultStreaming,
23 StreamEvent,
24 ToolCallItem,
25)
26from agents._run_impl import QueueCompleteSentinel
27from openai.types.responses import (
28 EasyInputMessageParam,
29 ResponseFileSearchToolCall,
30 ResponseInputContentParam,
31 ResponseInputTextParam,
32 ResponseOutputItemAddedEvent,
33 ResponseOutputItemDoneEvent,
34 ResponseOutputMessage,
35 ResponseReasoningItem,
36)
37from openai.types.responses.response_content_part_added_event import (
38 ResponseContentPartAddedEvent,
39)
40from openai.types.responses.response_file_search_tool_call import Result
41from openai.types.responses.response_input_item_param import Message
42from openai.types.responses.response_output_text import (
43 AnnotationContainerFileCitation as ResponsesAnnotationContainerFileCitation,
44)
45from openai.types.responses.response_output_text import (
46 AnnotationFileCitation as ResponsesAnnotationFileCitation,
47)
48from openai.types.responses.response_output_text import (
49 AnnotationFilePath as ResponsesAnnotationFilePath,
50)
51from openai.types.responses.response_output_text import (
52 AnnotationURLCitation as ResponsesAnnotationURLCitation,
53)
54from openai.types.responses.response_output_text import (
55 ResponseOutputText,
56)
57from openai.types.responses.response_text_delta_event import ResponseTextDeltaEvent
58from openai.types.responses.response_text_done_event import ResponseTextDoneEvent
59
60from chatkit.agents import (
61 AgentContext,
62 ResponseStreamConverter,
63 ThreadItemConverter,
64 accumulate_text,
65 simple_to_agent_input,
66 stream_agent_response,
67)
68from chatkit.types import (
69 Annotation,
70 AssistantMessageContent,
71 AssistantMessageContentPartAdded,
72 AssistantMessageContentPartAnnotationAdded,
73 AssistantMessageContentPartDone,
74 AssistantMessageContentPartTextDelta,
75 AssistantMessageItem,
76 Attachment,
77 ClientToolCallItem,
78 CustomSummary,
79 CustomTask,
80 DurationSummary,
81 FileSource,
82 GeneratedImage,
83 GeneratedImageItem,
84 GeneratedImageUpdated,
85 HiddenContextItem,
86 InferenceOptions,
87 Page,
88 TaskItem,
89 ThoughtTask,
90 Thread,
91 ThreadItemAddedEvent,
92 ThreadItemDoneEvent,
93 ThreadItemUpdatedEvent,
94 ThreadStreamEvent,
95 URLSource,
96 UserMessageItem,
97 UserMessageTagContent,
98 UserMessageTextContent,
99 WidgetItem,
100 Workflow,
101 WorkflowItem,
102 WorkflowTaskAdded,
103 WorkflowTaskUpdated,
104)
105from chatkit.widgets import Card, Text
106
107thread = Thread(id="123", title="Test", created_at=datetime.now(), items=Page())
108
109mock_store = Mock()
110mock_store.generate_item_id = lambda item_type, thread, context: f"{item_type}_id"
111mock_store.load_thread_items = AsyncMock(return_value=Page())
112mock_store.add_thread_item = AsyncMock()
113
114
115class RunResult(RunResultStreaming):
116 def add_event(self, event: StreamEvent):
117 self._event_queue.put_nowait(event)
118
119 def done(self):
120 self.is_complete = True
121 self._event_queue.put_nowait(QueueCompleteSentinel())
122
123 def throw_input_guardrails(self):
124 self._stored_exception = InputGuardrailTripwireTriggered(
125 InputGuardrailResult(
126 guardrail=Mock(spec=InputGuardrail),
127 output=GuardrailFunctionOutput(
128 output_info=None,
129 tripwire_triggered=True,
130 ),
131 )
132 )
133 self.is_complete = True
134 self._event_queue.put_nowait(QueueCompleteSentinel())
135
136 def throw_output_guardrails(self):
137 self._stored_exception = OutputGuardrailTripwireTriggered(
138 OutputGuardrailResult(
139 guardrail=Mock(spec=OutputGuardrail),
140 output=GuardrailFunctionOutput(
141 output_info=None,
142 tripwire_triggered=True,
143 ),
144 agent=Mock(spec=Agent),
145 agent_output=None,
146 )
147 )
148 self.is_complete = True
149 self._event_queue.put_nowait(QueueCompleteSentinel())
150
151
152def make_result() -> RunResult:
153 return RunResult(
154 context_wrapper=Mock(spec=RunContextWrapper),
155 input=[],
156 tool_input_guardrail_results=[],
157 tool_output_guardrail_results=[],
158 new_items=[],
159 raw_responses=[],
160 final_output=None,
161 current_agent=Agent(name="test"),
162 current_turn=0,
163 max_turns=10,
164 _current_agent_output_schema=None,
165 trace=None,
166 is_complete=False,
167 _event_queue=asyncio.Queue(),
168 _input_guardrail_queue=asyncio.Queue(),
169 _output_guardrails_task=None,
170 _run_impl_task=None,
171 _stored_exception=None,
172 output_guardrail_results=[],
173 input_guardrail_results=[],
174 )
175
176
177async def all_events(
178 events: AsyncIterator[ThreadStreamEvent],
179) -> list[ThreadStreamEvent]:
180 return [event async for event in events]
181
182
183async def test_returns_widget_item():
184 context = AgentContext(
185 previous_response_id=None, thread=thread, store=mock_store, request_context=None
186 )
187 result = make_result()
188 result.add_event(
189 RunItemStreamEvent(name="tool_called", item=Mock(spec=ToolCallItem))
190 )
191 await context.stream_widget(Card(children=[Text(value="Hello, world!")]))
192 result.done()
193
194 events = await all_events(
195 stream_agent_response(
196 context=context,
197 result=result,
198 )
199 )
200
201 assert len(events) == 1
202 assert isinstance(events[0], ThreadItemDoneEvent)
203 assert isinstance(events[0].item, WidgetItem)
204 assert events[0].item.widget == Card(children=[Text(value="Hello, world!")])
205
206
207async def test_returns_widget_item_generator():
208 context = AgentContext(
209 previous_response_id=None, thread=thread, store=mock_store, request_context=None
210 )
211 result = make_result()
212 result.add_event(
213 RunItemStreamEvent(name="tool_called", item=Mock(spec=ToolCallItem))
214 )
215
216 def render_widget(i: int) -> Card:
217 return Card(children=[Text(id="text", value="Hello, world"[:i])])
218
219 async def widget_generator():
220 yield render_widget(0)
221 yield render_widget(12)
222
223 await context.stream_widget(widget_generator())
224 result.done()
225
226 events = await all_events(
227 stream_agent_response(
228 context=context,
229 result=result,
230 )
231 )
232
233 assert len(events) == 3
234 assert isinstance(events[0], ThreadItemAddedEvent)
235 assert isinstance(events[0].item, WidgetItem)
236 assert events[0].item.widget == Card(children=[Text(id="text", value="")])
237
238 assert isinstance(events[1], ThreadItemUpdatedEvent)
239 assert events[1].update.type == "widget.streaming_text.value_delta"
240 assert events[1].update.component_id == "text"
241 assert events[1].update.delta == "Hello, world"
242
243 assert isinstance(events[2], ThreadItemDoneEvent)
244 assert isinstance(events[2].item, WidgetItem)
245 assert events[2].item.widget == Card(
246 children=[Text(id="text", value="Hello, world")]
247 )
248
249
250async def test_returns_widget_full_replace_generator():
251 context = AgentContext(
252 previous_response_id=None, thread=thread, store=mock_store, request_context=None
253 )
254 result = make_result()
255 result.add_event(
256 RunItemStreamEvent(name="tool_called", item=Mock(spec=ToolCallItem))
257 )
258
259 async def widget_generator():
260 yield Card(children=[Text(id="text", value="Hello!")])
261 yield Card(children=[Text(key="other text", value="World!", streaming=False)])
262
263 await context.stream_widget(widget_generator())
264 result.done()
265
266 events = await all_events(
267 stream_agent_response(
268 context=context,
269 result=result,
270 )
271 )
272
273 assert len(events) == 3
274 assert isinstance(events[0], ThreadItemAddedEvent)
275 assert isinstance(events[0].item, WidgetItem)
276 assert events[0].item.widget == Card(children=[Text(id="text", value="Hello!")])
277
278 assert isinstance(events[1], ThreadItemUpdatedEvent)
279 assert events[1].update.type == "widget.root.updated"
280 assert events[1].update.widget == Card(
281 children=[Text(key="other text", value="World!", streaming=False)]
282 )
283
284 assert isinstance(events[2], ThreadItemDoneEvent)
285 assert isinstance(events[2].item, WidgetItem)
286 assert events[2].item.widget == Card(
287 children=[Text(key="other text", value="World!", streaming=False)]
288 )
289
290
291async def test_accumulate_text():
292 def delta(text: str) -> RawResponsesStreamEvent:
293 return RawResponsesStreamEvent(
294 type="raw_response_event",
295 data=ResponseTextDeltaEvent(
296 type="response.output_text.delta",
297 delta=text,
298 content_index=0,
299 item_id="123",
300 logprobs=[],
301 output_index=0,
302 sequence_number=0,
303 ),
304 )
305
306 result = Runner.run_streamed(
307 Agent("Assistant", instructions="You are a helpful assistant."), "Say hello!"
308 )
309 result = make_result()
310 result.add_event(delta("Hello, "))
311 result.add_event(delta("world!"))
312
313 result.done()
314
315 events = [
316 event
317 async for event in accumulate_text(
318 result.stream_events(), Text(key="text", value="", streaming=True)
319 )
320 ]
321 assert events == [
322 Text(key="text", value="", streaming=True),
323 Text(key="text", value="Hello, ", streaming=True),
324 Text(key="text", value="Hello, world!", streaming=True),
325 Text(key="text", value="Hello, world!", streaming=False),
326 ]
327
328
329async def test_input_item_converter_quotes_last_user_message():
330 items = [
331 UserMessageItem(
332 id="123",
333 content=[UserMessageTextContent(text="Hello!")],
334 attachments=[],
335 inference_options=InferenceOptions(),
336 thread_id=thread.id,
337 quoted_text="Hi!",
338 created_at=datetime.now(),
339 ),
340 UserMessageItem(
341 id="123",
342 content=[UserMessageTextContent(text="I'm well, thank you!")],
343 attachments=[],
344 inference_options=InferenceOptions(),
345 thread_id=thread.id,
346 quoted_text="How are you doing?",
347 created_at=datetime.now(),
348 ),
349 ]
350
351 async def throw_exception(
352 _: Attachment,
353 ) -> ResponseInputContentParam:
354 raise Exception("Not implemented")
355
356 input_items = await simple_to_agent_input(items)
357 assert len(input_items) == 3
358 assert input_items[0] == {
359 "content": [
360 {
361 "text": "Hello!",
362 "type": "input_text",
363 },
364 ],
365 "role": "user",
366 "type": "message",
367 }
368 assert input_items[1] == {
369 "content": [
370 {
371 "text": "I'm well, thank you!",
372 "type": "input_text",
373 },
374 ],
375 "role": "user",
376 "type": "message",
377 }
378 assert input_items[2] == {
379 "content": [
380 {
381 "text": "The user is referring to this in particular: \nHow are you doing?",
382 "type": "input_text",
383 },
384 ],
385 "role": "user",
386 "type": "message",
387 }
388
389
390async def test_input_item_converter_to_input_items_mixed():
391 items = [
392 UserMessageItem(
393 id="123",
394 content=[UserMessageTextContent(text="Hello!")],
395 attachments=[],
396 inference_options=InferenceOptions(),
397 thread_id=thread.id,
398 quoted_text="Hi!",
399 created_at=datetime.now(),
400 ),
401 UserMessageItem(
402 id="123",
403 content=[UserMessageTextContent(text="I'm well, thank you!")],
404 attachments=[],
405 inference_options=InferenceOptions(),
406 thread_id=thread.id,
407 quoted_text="How are you doing?",
408 created_at=datetime.now(),
409 ),
410 AssistantMessageItem(
411 id="123",
412 content=[
413 AssistantMessageContent(text="How are you doing?"),
414 AssistantMessageContent(text="Can't do that"),
415 ],
416 thread_id=thread.id,
417 created_at=datetime.now(),
418 ),
419 WidgetItem(
420 id="wd_123",
421 widget=Card(children=[Text(value="Hello, world!")]),
422 thread_id=thread.id,
423 created_at=datetime.now(),
424 ),
425 ]
426
427 input_items = await simple_to_agent_input(items)
428 assert len(input_items) == 4
429 assert input_items[0] == {
430 "content": [
431 {
432 "text": "Hello!",
433 "type": "input_text",
434 },
435 ],
436 "role": "user",
437 "type": "message",
438 }
439 assert input_items[1] == {
440 "content": [
441 {
442 "text": "I'm well, thank you!",
443 "type": "input_text",
444 },
445 ],
446 "role": "user",
447 "type": "message",
448 }
449 assert input_items[2] == {
450 "content": [
451 {
452 "annotations": [],
453 "text": "How are you doing?",
454 "logprobs": None,
455 "type": "output_text",
456 },
457 {
458 "annotations": [],
459 "text": "Can't do that",
460 "logprobs": None,
461 "type": "output_text",
462 },
463 ],
464 "type": "message",
465 "role": "assistant",
466 }
467 assert "type" in input_items[3]
468 widget_item = cast(EasyInputMessageParam, input_items[3])
469 assert widget_item.get("type") == "message"
470 assert widget_item.get("role") == "user"
471 text = widget_item.get("content")[0]["text"] # type: ignore
472 assert (
473 "The following graphical UI widget (id: wd_123) was displayed to the user"
474 in text
475 )
476 assert "Hello, world!" in text
477 assert "created_at" not in text
478
479
480async def test_input_item_converter_user_input_with_tags():
481 class MyThreadItemConverter(ThreadItemConverter):
482 async def tag_to_message_content(self, tag):
483 return ResponseInputTextParam(
484 type="input_text", text=tag.text + " " + tag.data["key"]
485 )
486
487 items = [
488 UserMessageItem(
489 id="123",
490 content=[
491 UserMessageTagContent(
492 text="Hello!", type="input_tag", id="hello", data={"key": "value"}
493 )
494 ],
495 attachments=[],
496 inference_options=InferenceOptions(),
497 thread_id=thread.id,
498 created_at=datetime.now(),
499 )
500 ]
501 items = await MyThreadItemConverter().to_agent_input(items)
502
503 assert len(items) == 2
504 assert items[0] == {
505 "content": [
506 {
507 "text": "@Hello!",
508 "type": "input_text",
509 },
510 ],
511 "role": "user",
512 "type": "message",
513 }
514 assert items[1] == {
515 "content": [
516 {
517 "text": "# User-provided context for @-mentions\n- When referencing resolved entities, use their canonical names **without** '@'.\n"
518 + "- The '@' form appears only in user text and should not be echoed.",
519 "type": "input_text",
520 },
521 {
522 "text": "Hello! value",
523 "type": "input_text",
524 },
525 ],
526 "role": "user",
527 "type": "message",
528 }
529
530
531async def test_input_item_converter_user_input_with_tags_throws_by_default():
532 items = [
533 UserMessageItem(
534 id="123",
535 content=[
536 UserMessageTagContent(
537 text="Hello!", type="input_tag", id="hello", data={}
538 )
539 ],
540 attachments=[],
541 inference_options=InferenceOptions(),
542 thread_id=thread.id,
543 created_at=datetime.now(),
544 )
545 ]
546
547 with pytest.raises(NotImplementedError):
548 await simple_to_agent_input(items)
549
550
551async def test_input_item_converter_for_hidden_context_with_string_content():
552 items = [
553 HiddenContextItem(
554 id="123",
555 content="User pressed the red button",
556 thread_id=thread.id,
557 created_at=datetime.now(),
558 )
559 ]
560
561 # The default converter works for string content
562 items = await simple_to_agent_input(items)
563 assert len(items) == 1
564 assert items[0] == {
565 "content": [
566 {
567 "text": "Hidden context for the agent (not shown to the user):\n<HiddenContext>\nUser pressed the red button\n</HiddenContext>",
568 "type": "input_text",
569 },
570 ],
571 "role": "user",
572 "type": "message",
573 }
574
575
576async def test_input_item_converter_for_hidden_context_with_non_string_content():
577 items = [
578 HiddenContextItem(
579 id="123",
580 content={"harry": "potter", "hermione": "granger"},
581 thread_id=thread.id,
582 created_at=datetime.now(),
583 )
584 ]
585
586 # Default converter should throw
587 with pytest.raises(NotImplementedError):
588 await simple_to_agent_input(items)
589
590 class MyThreadItemConverter(ThreadItemConverter):
591 async def hidden_context_to_input(self, item: HiddenContextItem):
592 content = ",".join(item.content.keys())
593 return Message(
594 type="message",
595 content=[
596 ResponseInputTextParam(
597 type="input_text", text=f"<HIDDEN>{content}</HIDDEN>"
598 )
599 ],
600 role="user",
601 )
602
603 items = await MyThreadItemConverter().to_agent_input(items)
604 assert len(items) == 1
605 assert items[0] == {
606 "content": [
607 {
608 "text": "<HIDDEN>harry,hermione</HIDDEN>",
609 "type": "input_text",
610 },
611 ],
612 "role": "user",
613 "type": "message",
614 }
615
616
617async def test_input_item_converter_with_client_tool_call():
618 items = [
619 UserMessageItem(
620 id="123",
621 content=[UserMessageTextContent(text="Call a client tool call xyz")],
622 attachments=[],
623 inference_options=InferenceOptions(),
624 thread_id=thread.id,
625 quoted_text="Hi!",
626 created_at=datetime.now(),
627 ),
628 TaskItem(
629 id="tsk_123",
630 created_at=datetime.now(),
631 task=CustomTask(title="Called xyx"),
632 thread_id=thread.id,
633 ),
634 ClientToolCallItem(
635 id="ctc_123",
636 thread_id=thread.id,
637 created_at=datetime.now(),
638 name="xyz",
639 arguments={"foo": "bar"},
640 call_id="ctc_123",
641 ),
642 ClientToolCallItem(
643 id="ctc_123_done",
644 thread_id=thread.id,
645 created_at=datetime.now(),
646 name="xyz",
647 arguments={"foo": "bar"},
648 call_id="ctc_123",
649 status="completed",
650 output={"success": True},
651 ),
652 ]
653
654 input_items = await simple_to_agent_input(items)
655 assert len(input_items) == 4
656 assert input_items[0] == {
657 "content": [
658 {
659 "text": "Call a client tool call xyz",
660 "type": "input_text",
661 },
662 ],
663 "role": "user",
664 "type": "message",
665 }
666 assert input_items[1] == {
667 "content": [
668 {
669 "text": "A message was displayed to the user that the following task was performed:\n<Task>\nCalled xyx\n</Task>",
670 "type": "input_text",
671 },
672 ],
673 "type": "message",
674 "role": "user",
675 }
676 assert input_items[2] == {
677 "type": "function_call",
678 "name": "xyz",
679 "arguments": json.dumps({"foo": "bar"}),
680 "call_id": "ctc_123",
681 }
682 assert input_items[3] == {
683 "type": "function_call_output",
684 "call_id": "ctc_123",
685 "output": json.dumps({"success": True}),
686 }
687
688
689async def test_stream_agent_response_yields_context_events_without_streaming_events():
690 context = AgentContext(
691 previous_response_id=None, thread=thread, store=mock_store, request_context=None
692 )
693 result = make_result()
694
695 event = ThreadItemAddedEvent(
696 item=WidgetItem(
697 id="123",
698 created_at=datetime.now(),
699 thread_id=thread.id,
700 widget=Card(children=[Text(id="text", value="Hello, world!")]),
701 ),
702 )
703
704 await context.stream(event)
705
706 response_streamer = stream_agent_response(context, result)
707 event = await response_streamer.__anext__()
708
709 assert event.type == "thread.item.added"
710
711 future = asyncio.ensure_future(response_streamer.__anext__())
712 assert future.done() is False
713
714 result.done()
715
716 try:
717 await future
718 assert False, "expected StopAsyncIteration"
719 except StopAsyncIteration:
720 pass
721
722 assert future.done() is True
723
724
725async def test_stream_agent_response_maps_events():
726 context = AgentContext(
727 previous_response_id=None, thread=thread, store=mock_store, request_context=None
728 )
729 result = make_result()
730
731 event = ThreadItemAddedEvent(
732 item=WidgetItem(
733 id="123",
734 created_at=datetime.now(),
735 thread_id=thread.id,
736 widget=Card(children=[Text(id="text", value="Hello, world!")]),
737 ),
738 )
739
740 await context.stream(event)
741 result.add_event(
742 RawResponsesStreamEvent(
743 type="raw_response_event",
744 data=ResponseTextDeltaEvent(
745 type="response.output_text.delta",
746 delta="Hello, world!",
747 content_index=0,
748 item_id="123",
749 logprobs=[],
750 output_index=0,
751 sequence_number=0,
752 ),
753 )
754 )
755
756 response_streamer = stream_agent_response(context, result)
757 event1 = await response_streamer.__anext__()
758 event2 = await response_streamer.__anext__()
759
760 assert {event1.type, event2.type} == {
761 "thread.item.added",
762 "thread.item.updated",
763 }
764
765 future = asyncio.ensure_future(response_streamer.__anext__())
766 assert future.done() is False
767
768 result.done()
769
770 try:
771 await future
772 assert False, "expected StopAsyncIteration"
773 except StopAsyncIteration:
774 pass
775
776 assert future.done() is True
777
778
779@pytest.mark.parametrize(
780 "raw_event,expected_event",
781 [
782 (
783 RawResponsesStreamEvent(
784 type="raw_response_event",
785 data=ResponseTextDeltaEvent(
786 type="response.output_text.delta",
787 delta="Hello, world!",
788 content_index=0,
789 item_id="123",
790 logprobs=[],
791 output_index=0,
792 sequence_number=0,
793 ),
794 ),
795 ThreadItemUpdatedEvent(
796 item_id="123",
797 update=AssistantMessageContentPartTextDelta(
798 content_index=0,
799 delta="Hello, world!",
800 ),
801 ),
802 ),
803 (
804 RawResponsesStreamEvent(
805 type="raw_response_event",
806 data=ResponseContentPartAddedEvent(
807 type="response.content_part.added",
808 part=ResponseOutputText(
809 type="output_text",
810 text="New content",
811 annotations=[],
812 ),
813 content_index=1,
814 item_id="123",
815 output_index=0,
816 sequence_number=1,
817 ),
818 ),
819 ThreadItemUpdatedEvent(
820 item_id="123",
821 update=AssistantMessageContentPartAdded(
822 content_index=1,
823 content=AssistantMessageContent(text="New content", annotations=[]),
824 ),
825 ),
826 ),
827 (
828 RawResponsesStreamEvent(
829 type="raw_response_event",
830 data=ResponseTextDoneEvent(
831 type="response.output_text.done",
832 text="Final text",
833 content_index=0,
834 item_id="123",
835 logprobs=[],
836 output_index=0,
837 sequence_number=2,
838 ),
839 ),
840 ThreadItemUpdatedEvent(
841 item_id="123",
842 update=AssistantMessageContentPartDone(
843 content_index=0,
844 content=AssistantMessageContent(
845 text="Final text",
846 annotations=[],
847 ),
848 ),
849 ),
850 ),
851 (
852 RawResponsesStreamEvent(
853 type="raw_response_event",
854 data=Mock(
855 type="response.output_text.annotation.added",
856 annotation=ResponsesAnnotationFileCitation(
857 type="file_citation",
858 file_id="file_123",
859 filename="file.txt",
860 index=5,
861 ),
862 content_index=0,
863 item_id="123",
864 annotation_index=0,
865 output_index=0,
866 sequence_number=3,
867 ),
868 ),
869 ThreadItemUpdatedEvent(
870 item_id="123",
871 update=AssistantMessageContentPartAnnotationAdded(
872 content_index=0,
873 annotation_index=0,
874 annotation=Annotation(
875 source=FileSource(filename="file.txt", title="file.txt"),
876 index=5,
877 ),
878 ),
879 ),
880 ),
881 ],
882)
883async def test_event_mapping(raw_event, expected_event):
884 context = AgentContext(
885 previous_response_id=None, thread=thread, store=mock_store, request_context=None
886 )
887 result = make_result()
888
889 result.add_event(raw_event)
890 result.done()
891
892 events = await all_events(stream_agent_response(context, result))
893 if expected_event:
894 assert events == [expected_event]
895 else:
896 assert events == []
897
898
899async def test_stream_agent_response_emits_annotation_added_events():
900 context = AgentContext(
901 previous_response_id=None, thread=thread, store=mock_store, request_context=None
902 )
903 result = make_result()
904 item_id = "item_123"
905
906 def add_annotation_event(annotation, sequence_number):
907 result.add_event(
908 RawResponsesStreamEvent(
909 type="raw_response_event",
910 data=Mock(
911 type="response.output_text.annotation.added",
912 annotation=annotation,
913 content_index=0,
914 item_id=item_id,
915 annotation_index=sequence_number,
916 output_index=0,
917 sequence_number=sequence_number,
918 ),
919 )
920 )
921
922 add_annotation_event(
923 ResponsesAnnotationFileCitation(
924 type="file_citation",
925 file_id="file_invalid",
926 filename="",
927 index=0,
928 ),
929 sequence_number=0,
930 )
931 add_annotation_event(
932 ResponsesAnnotationContainerFileCitation(
933 type="container_file_citation",
934 container_id="container_1",
935 file_id="file_123",
936 filename="container.txt",
937 start_index=0,
938 end_index=3,
939 ),
940 sequence_number=1,
941 )
942 add_annotation_event(
943 ResponsesAnnotationURLCitation(
944 type="url_citation",
945 url="https://example.com",
946 title="Example",
947 start_index=1,
948 end_index=5,
949 ),
950 sequence_number=2,
951 )
952 result.done()
953
954 events = await all_events(stream_agent_response(context, result))
955 assert events == [
956 ThreadItemUpdatedEvent(
957 item_id=item_id,
958 update=AssistantMessageContentPartAnnotationAdded(
959 content_index=0,
960 annotation_index=0,
961 annotation=Annotation(
962 source=FileSource(filename="container.txt", title="container.txt"),
963 index=3,
964 ),
965 ),
966 ),
967 ThreadItemUpdatedEvent(
968 item_id=item_id,
969 update=AssistantMessageContentPartAnnotationAdded(
970 content_index=0,
971 annotation_index=1,
972 annotation=Annotation(
973 source=URLSource(
974 url="https://example.com",
975 title="Example",
976 ),
977 index=5,
978 ),
979 ),
980 ),
981 ]
982
983
984@pytest.mark.parametrize("throw_guardrail", ["input", "output"])
985async def test_stream_agent_response_yields_item_removed_event(throw_guardrail):
986 context = AgentContext(
987 previous_response_id=None, thread=thread, store=mock_store, request_context=None
988 )
989 result = make_result()
990 result.add_event(
991 RawResponsesStreamEvent(
992 type="raw_response_event",
993 data=ResponseOutputItemAddedEvent(
994 type="response.output_item.added",
995 item=ResponseOutputMessage(
996 id="1",
997 content=[
998 ResponseOutputText(
999 annotations=[], type="output_text", text="Hello, world!"
1000 )
1001 ],
1002 role="assistant",
1003 status="completed",
1004 type="message",
1005 ),
1006 output_index=0,
1007 sequence_number=0,
1008 ),
1009 )
1010 )
1011 await context.stream(
1012 ThreadItemAddedEvent(
1013 item=AssistantMessageItem(
1014 id="2",
1015 content=[AssistantMessageContent(text="Hello, world!")],
1016 thread_id=thread.id,
1017 created_at=datetime.now(),
1018 ),
1019 )
1020 )
1021
1022 await context.stream(
1023 ThreadItemDoneEvent(
1024 item=WidgetItem(
1025 id="3",
1026 created_at=datetime.now(),
1027 thread_id=thread.id,
1028 widget=Card(children=[Text(id="text", value="Hello, world!")]),
1029 ),
1030 )
1031 )
1032
1033 iterator = stream_agent_response(context, result)
1034
1035 n = 3
1036 events = []
1037 # Grab first 3 events to
1038 async for event in iterator:
1039 n -= 1
1040 events.append(event)
1041 if n == 0:
1042 break
1043
1044 if throw_guardrail == "input":
1045 result.throw_input_guardrails()
1046 else:
1047 result.throw_output_guardrails()
1048
1049 try:
1050 async for event in iterator:
1051 events.append(event)
1052 assert False, "Guardrail should have been thrown from stream_agent_response"
1053 except (InputGuardrailTripwireTriggered, OutputGuardrailTripwireTriggered):
1054 pass
1055 except Exception as e:
1056 assert False, f"Unexpected exception: {e}"
1057
1058 deleted_item_ids = {
1059 event.item_id for event in events if event.type == "thread.item.removed"
1060 }
1061 assert deleted_item_ids == {"1", "2", "3"}
1062
1063
1064async def test_stream_agent_response_assistant_message_content_types():
1065 AgentContext(
1066 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1067 )
1068 result = make_result()
1069
1070 result.add_event(
1071 RawResponsesStreamEvent(
1072 type="raw_response_event",
1073 data=ResponseOutputItemDoneEvent(
1074 type="response.output_item.done",
1075 item=ResponseFileSearchToolCall(
1076 id="fs_0",
1077 queries=["Hello, world!"],
1078 status="completed",
1079 type="file_search_call",
1080 results=[
1081 Result(
1082 file_id="f_123",
1083 filename="test.txt",
1084 text="Hello, world!",
1085 score=1.0,
1086 ),
1087 Result(
1088 file_id="f_123",
1089 filename="test.txt",
1090 text="Hello, friends!",
1091 score=0.5,
1092 ),
1093 ],
1094 ),
1095 output_index=0,
1096 sequence_number=0,
1097 ),
1098 )
1099 )
1100 result.add_event(
1101 RawResponsesStreamEvent(
1102 type="raw_response_event",
1103 data=ResponseOutputItemDoneEvent(
1104 type="response.output_item.done",
1105 item=ResponseOutputMessage(
1106 id="1",
1107 content=[
1108 ResponseOutputText(
1109 annotations=[
1110 ResponsesAnnotationFileCitation(
1111 type="file_citation",
1112 file_id="f_123",
1113 index=0,
1114 filename="test.txt",
1115 ),
1116 ResponsesAnnotationContainerFileCitation(
1117 type="container_file_citation",
1118 container_id="container_1",
1119 file_id="f_456",
1120 filename="container.txt",
1121 start_index=0,
1122 end_index=3,
1123 ),
1124 ResponsesAnnotationURLCitation(
1125 type="url_citation",
1126 url="https://www.google.com",
1127 title="Google",
1128 start_index=0,
1129 end_index=10,
1130 ),
1131 ResponsesAnnotationFilePath(
1132 type="file_path",
1133 file_id="123",
1134 index=0,
1135 ),
1136 ],
1137 text="Hello, world!",
1138 type="output_text",
1139 ),
1140 ResponseOutputText(
1141 annotations=[],
1142 text="Can't do that",
1143 type="output_text",
1144 ),
1145 ],
1146 role="assistant",
1147 status="completed",
1148 type="message",
1149 ),
1150 output_index=0,
1151 sequence_number=0,
1152 ),
1153 )
1154 )
1155
1156 result.done()
1157
1158 context = AgentContext(
1159 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1160 )
1161 events = await all_events(stream_agent_response(context, result))
1162 assert len(events) == 1
1163 assert isinstance(events[0], ThreadItemDoneEvent)
1164 message = events[0].item
1165 assert isinstance(message, AssistantMessageItem)
1166 assert message.content == [
1167 AssistantMessageContent(
1168 annotations=[
1169 Annotation(
1170 source=FileSource(
1171 filename="test.txt",
1172 title="test.txt",
1173 ),
1174 index=0,
1175 ),
1176 Annotation(
1177 source=FileSource(
1178 filename="container.txt",
1179 title="container.txt",
1180 ),
1181 index=3,
1182 ),
1183 Annotation(
1184 source=URLSource(
1185 url="https://www.google.com",
1186 title="Google",
1187 ),
1188 index=10,
1189 ),
1190 ],
1191 text="Hello, world!",
1192 ),
1193 AssistantMessageContent(text="Can't do that", annotations=[]),
1194 ]
1195 assert message.id == "1"
1196
1197
1198async def test_stream_agent_response_image_generation_events():
1199 context = AgentContext(
1200 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1201 )
1202 result = make_result()
1203
1204 result.add_event(
1205 RawResponsesStreamEvent(
1206 type="raw_response_event",
1207 data=Mock(
1208 type="response.output_item.added",
1209 item=Mock(type="image_generation_call", id="img_call_1"),
1210 output_index=0,
1211 sequence_number=0,
1212 ),
1213 )
1214 )
1215 result.add_event(
1216 RawResponsesStreamEvent(
1217 type="raw_response_event",
1218 data=Mock(
1219 type="response.output_item.done",
1220 item=Mock(
1221 type="image_generation_call", id="img_call_1", result="dGVzdA=="
1222 ),
1223 output_index=0,
1224 sequence_number=1,
1225 ),
1226 )
1227 )
1228 result.done()
1229
1230 stream = stream_agent_response(context, result)
1231 event1 = await stream.__anext__()
1232 assert isinstance(event1, ThreadItemAddedEvent)
1233 assert isinstance(event1.item, GeneratedImageItem)
1234 assert event1.item.type == "generated_image"
1235 assert event1.item.id == "message_id"
1236 assert event1.item.image is None
1237
1238 event2 = await stream.__anext__()
1239 assert isinstance(event2, ThreadItemDoneEvent)
1240 assert isinstance(event2.item, GeneratedImageItem)
1241 assert event2.item.id == event1.item.id
1242 assert event2.item.image == GeneratedImage(
1243 id="img_call_1", url="data:image/png;base64,dGVzdA=="
1244 )
1245
1246 with pytest.raises(StopAsyncIteration):
1247 await stream.__anext__()
1248
1249
1250async def test_stream_agent_response_image_generation_events_with_custom_converter():
1251 context = AgentContext(
1252 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1253 )
1254 result = make_result()
1255
1256 result.add_event(
1257 RawResponsesStreamEvent(
1258 type="raw_response_event",
1259 data=Mock(
1260 type="response.output_item.added",
1261 item=Mock(type="image_generation_call", id="img_call_1"),
1262 output_index=0,
1263 sequence_number=0,
1264 ),
1265 )
1266 )
1267 result.add_event(
1268 RawResponsesStreamEvent(
1269 type="raw_response_event",
1270 data=Mock(
1271 type="response.output_item.done",
1272 item=Mock(
1273 type="image_generation_call", id="img_call_1", result="dGVzdA=="
1274 ),
1275 output_index=0,
1276 sequence_number=1,
1277 ),
1278 )
1279 )
1280 result.done()
1281
1282 class CustomResponseStreamConverter(ResponseStreamConverter):
1283 def __init__(self):
1284 super().__init__()
1285 self.calls: list[tuple[str, str, int | None]] = []
1286
1287 async def base64_image_to_url(
1288 self,
1289 image_id: str,
1290 base64_image: str,
1291 partial_image_index: int | None = None,
1292 ) -> str:
1293 self.calls.append((image_id, base64_image, partial_image_index))
1294 return f"https://example.com/{image_id}"
1295
1296 converter = CustomResponseStreamConverter()
1297 stream = stream_agent_response(context, result, converter=converter)
1298 event1 = await stream.__anext__()
1299 assert isinstance(event1, ThreadItemAddedEvent)
1300 assert isinstance(event1.item, GeneratedImageItem)
1301 assert event1.item.image is None
1302
1303 event2 = await stream.__anext__()
1304 assert isinstance(event2, ThreadItemDoneEvent)
1305 assert isinstance(event2.item, GeneratedImageItem)
1306 assert converter.calls == [("img_call_1", "dGVzdA==", None)]
1307 assert event2.item.image == GeneratedImage(
1308 id="img_call_1", url="https://example.com/img_call_1"
1309 )
1310 with pytest.raises(StopAsyncIteration):
1311 await stream.__anext__()
1312
1313
1314async def test_stream_agent_response_image_generation_partial_progress():
1315 context = AgentContext(
1316 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1317 )
1318 result = make_result()
1319
1320 result.add_event(
1321 RawResponsesStreamEvent(
1322 type="raw_response_event",
1323 data=Mock(
1324 type="response.output_item.added",
1325 item=Mock(type="image_generation_call", id="img_call_1"),
1326 output_index=0,
1327 sequence_number=0,
1328 ),
1329 )
1330 )
1331 result.add_event(
1332 RawResponsesStreamEvent(
1333 type="raw_response_event",
1334 data=Mock(
1335 type="response.image_generation_call.partial_image",
1336 partial_image_b64="dGVzdA==",
1337 partial_image_index=1,
1338 item_id="img_call_1",
1339 output_index=0,
1340 sequence_number=1,
1341 ),
1342 )
1343 )
1344 result.add_event(
1345 RawResponsesStreamEvent(
1346 type="raw_response_event",
1347 data=Mock(
1348 type="response.output_item.done",
1349 item=Mock(
1350 type="image_generation_call", id="img_call_1", result="dGVzdA=="
1351 ),
1352 output_index=0,
1353 sequence_number=2,
1354 ),
1355 )
1356 )
1357 result.done()
1358
1359 converter = ResponseStreamConverter(partial_images=3)
1360 events = await all_events(
1361 stream_agent_response(context, result, converter=converter)
1362 )
1363
1364 assert len(events) == 3
1365 added_event, partial_event, done_event = events
1366
1367 assert isinstance(added_event, ThreadItemAddedEvent)
1368 assert isinstance(added_event.item, GeneratedImageItem)
1369
1370 assert isinstance(partial_event, ThreadItemUpdatedEvent)
1371 assert isinstance(partial_event.update, GeneratedImageUpdated)
1372 assert partial_event.update.progress == pytest.approx(1 / 3)
1373 assert partial_event.update.image == GeneratedImage(
1374 id="img_call_1", url="data:image/png;base64,dGVzdA=="
1375 )
1376
1377 assert isinstance(done_event, ThreadItemDoneEvent)
1378 assert isinstance(done_event.item, GeneratedImageItem)
1379 assert done_event.item.image == GeneratedImage(
1380 id="img_call_1", url="data:image/png;base64,dGVzdA=="
1381 )
1382
1383
1384async def test_workflow_streams_first_thought():
1385 context = AgentContext(
1386 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1387 )
1388 result = make_result()
1389
1390 # first thought
1391 result.add_event(
1392 RawResponsesStreamEvent(
1393 type="raw_response_event",
1394 data=ResponseOutputItemAddedEvent(
1395 type="response.output_item.added",
1396 item=ResponseReasoningItem(
1397 id="resp_1",
1398 summary=[],
1399 type="reasoning",
1400 ),
1401 output_index=0,
1402 sequence_number=0,
1403 ),
1404 )
1405 )
1406 result.add_event(
1407 RawResponsesStreamEvent(
1408 type="raw_response_event",
1409 data=Mock(
1410 type="response.reasoning_summary_text.delta",
1411 item_id="resp_1",
1412 summary_index=0,
1413 delta="Think",
1414 ),
1415 )
1416 )
1417 result.add_event(
1418 RawResponsesStreamEvent(
1419 type="raw_response_event",
1420 data=Mock(
1421 type="response.reasoning_summary_text.delta",
1422 item_id="resp_1",
1423 summary_index=0,
1424 delta="ing 1",
1425 ),
1426 )
1427 )
1428 result.add_event(
1429 RawResponsesStreamEvent(
1430 type="raw_response_event",
1431 data=Mock(
1432 type="response.reasoning_summary_text.done",
1433 item_id="resp_1",
1434 summary_index=0,
1435 text="Thinking 1",
1436 ),
1437 )
1438 )
1439
1440 # second thought
1441 result.add_event(
1442 RawResponsesStreamEvent(
1443 type="raw_response_event",
1444 data=Mock(
1445 type="response.reasoning_summary_text.delta",
1446 item_id="resp_1",
1447 summary_index=1,
1448 delta="Think",
1449 ),
1450 )
1451 )
1452 result.add_event(
1453 RawResponsesStreamEvent(
1454 type="raw_response_event",
1455 data=Mock(
1456 type="response.reasoning_summary_text.delta",
1457 item_id="resp_1",
1458 summary_index=1,
1459 delta="ing 2",
1460 ),
1461 )
1462 )
1463 result.add_event(
1464 RawResponsesStreamEvent(
1465 type="raw_response_event",
1466 data=Mock(
1467 type="response.reasoning_summary_text.done",
1468 item_id="resp_1",
1469 summary_index=1,
1470 text="Thinking 2",
1471 ),
1472 )
1473 )
1474
1475 result.done()
1476 stream = stream_agent_response(context, result)
1477
1478 # Workflow added
1479 event = await anext(stream)
1480 assert isinstance(event, ThreadItemAddedEvent)
1481 assert context.workflow_item is not None
1482 assert context.workflow_item.workflow.type == "reasoning"
1483 assert len(context.workflow_item.workflow.tasks) == 0
1484 assert event == ThreadItemAddedEvent(item=context.workflow_item)
1485
1486 # First thought added
1487 event = await anext(stream)
1488 assert context.workflow_item is not None
1489 assert len(context.workflow_item.workflow.tasks) == 1
1490 assert isinstance(event, ThreadItemUpdatedEvent)
1491 assert event == ThreadItemUpdatedEvent(
1492 item_id=context.workflow_item.id,
1493 update=WorkflowTaskAdded(
1494 task=ThoughtTask(content="Think"),
1495 task_index=0,
1496 ),
1497 )
1498
1499 # First thought delta
1500 event = await anext(stream)
1501 assert context.workflow_item is not None
1502 assert len(context.workflow_item.workflow.tasks) == 1
1503 assert isinstance(event, ThreadItemUpdatedEvent)
1504 assert event == ThreadItemUpdatedEvent(
1505 item_id=context.workflow_item.id,
1506 update=WorkflowTaskUpdated(
1507 task=ThoughtTask(content="Thinking 1"),
1508 task_index=0,
1509 ),
1510 )
1511
1512 # First thought done
1513 event = await anext(stream)
1514 assert context.workflow_item is not None
1515 assert len(context.workflow_item.workflow.tasks) == 1
1516 assert isinstance(event, ThreadItemUpdatedEvent)
1517 assert event == ThreadItemUpdatedEvent(
1518 item_id=context.workflow_item.id,
1519 update=WorkflowTaskUpdated(
1520 task=ThoughtTask(content="Thinking 1"),
1521 task_index=0,
1522 ),
1523 )
1524
1525 # Second thought added (not streamed)
1526 event = await anext(stream)
1527 assert context.workflow_item is not None
1528 assert len(context.workflow_item.workflow.tasks) == 2
1529 assert isinstance(event, ThreadItemUpdatedEvent)
1530 assert event == ThreadItemUpdatedEvent(
1531 item_id=context.workflow_item.id,
1532 update=WorkflowTaskAdded(
1533 task=ThoughtTask(content="Thinking 2"),
1534 task_index=1,
1535 ),
1536 )
1537
1538 try:
1539 while True:
1540 await anext(stream)
1541 except StopAsyncIteration:
1542 pass
1543
1544
1545async def test_workflow_ends_on_message():
1546 context = AgentContext(
1547 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1548 )
1549 result = make_result()
1550
1551 # first thought
1552 result.add_event(
1553 RawResponsesStreamEvent(
1554 type="raw_response_event",
1555 data=ResponseOutputItemAddedEvent(
1556 type="response.output_item.added",
1557 item=ResponseReasoningItem(
1558 id="resp_1",
1559 summary=[],
1560 type="reasoning",
1561 ),
1562 output_index=0,
1563 sequence_number=0,
1564 ),
1565 )
1566 )
1567 result.add_event(
1568 RawResponsesStreamEvent(
1569 type="raw_response_event",
1570 data=Mock(
1571 type="response.reasoning_summary_text.done",
1572 item_id="resp_1",
1573 summary_index=0,
1574 text="Thinking 1",
1575 ),
1576 )
1577 )
1578
1579 # not reasoning
1580 result.add_event(
1581 RawResponsesStreamEvent(
1582 type="raw_response_event",
1583 data=ResponseOutputItemAddedEvent(
1584 type="response.output_item.added",
1585 item=ResponseOutputMessage(
1586 id="m_1",
1587 content=[],
1588 role="assistant",
1589 status="in_progress",
1590 type="message",
1591 ),
1592 output_index=0,
1593 sequence_number=0,
1594 ),
1595 )
1596 )
1597
1598 result.done()
1599 stream = stream_agent_response(context, result)
1600
1601 # Workflow added
1602 event = await anext(stream)
1603 assert isinstance(event, ThreadItemAddedEvent)
1604 assert context.workflow_item is not None
1605 assert context.workflow_item.workflow.type == "reasoning"
1606 assert len(context.workflow_item.workflow.tasks) == 0
1607 assert event == ThreadItemAddedEvent(item=context.workflow_item)
1608
1609 # First thought done
1610 event = await anext(stream)
1611 assert context.workflow_item is not None
1612 assert len(context.workflow_item.workflow.tasks) == 1
1613 assert isinstance(event, ThreadItemUpdatedEvent)
1614 assert event == ThreadItemUpdatedEvent(
1615 item_id=context.workflow_item.id,
1616 update=WorkflowTaskAdded(
1617 task=ThoughtTask(content="Thinking 1"),
1618 task_index=0,
1619 ),
1620 )
1621
1622 # Workflow ended
1623 event = await anext(stream)
1624 assert isinstance(event, ThreadItemDoneEvent)
1625 assert event.item.type == "workflow"
1626 assert context.workflow_item is None
1627 # Summary and expanded are handled by the end_workflow method
1628 assert isinstance(event.item.workflow.summary, DurationSummary)
1629 assert event.item.workflow.expanded is False
1630
1631 try:
1632 while True:
1633 await anext(stream)
1634 except StopAsyncIteration:
1635 pass
1636
1637
1638async def test_existing_workflow_summary_not_overwritten_on_automatic_end():
1639 context = AgentContext(
1640 previous_response_id=None, thread=thread, store=mock_store, request_context=None
1641 )
1642 result = make_result()
1643 context.workflow_item = WorkflowItem(
1644 id="wf_1",
1645 created_at=datetime.now(),
1646 workflow=Workflow(type="custom", tasks=[], summary=CustomSummary(title="Test")),
1647 thread_id=thread.id,
1648 )
1649
1650 result.add_event(
1651 RawResponsesStreamEvent(
1652 type="raw_response_event",
1653 data=ResponseOutputItemAddedEvent(
1654 type="response.output_item.added",
1655 item=ResponseOutputMessage(
1656 id="m_1",
1657 content=[],
1658 role="assistant",
1659 status="in_progress",
1660 type="message",
1661 ),
1662 output_index=0,
1663 sequence_number=0,
1664 ),
1665 )
1666 )
1667
1668 result.done()
1669 stream = stream_agent_response(context, result)
1670
1671 event = await anext(stream)
1672
1673 assert isinstance(event, ThreadItemDoneEvent)
1674 assert context.workflow_item is None
1675 assert event.item.type == "workflow"
1676 assert event.item.workflow.summary == CustomSummary(title="Test")
1677
1678 try:
1679 while True:
1680 await anext(stream)
1681 except StopAsyncIteration:
1682 pass
1683