feat(assistants): add support for streaming

api.md

@@ -1,7 +1,7 @@
 # Shared Types
 
 ```python
-from openai.types import FunctionDefinition, FunctionParameters
+from openai.types import ErrorObject, FunctionDefinition, FunctionParameters
 ```
 
 # Completions
@@ -177,7 +177,19 @@ Methods:
 Types:
 
 ```python
-from openai.types.beta import Assistant, AssistantDeleted
+from openai.types.beta import (
+    Assistant,
+    AssistantDeleted,
+    AssistantStreamEvent,
+    AssistantTool,
+    CodeInterpreterTool,
+    FunctionTool,
+    MessageStreamEvent,
+    RetrievalTool,
+    RunStepStreamEvent,
+    RunStreamEvent,
+    ThreadStreamEvent,
+)
 ```
 
 Methods:
@@ -218,6 +230,7 @@ Methods:
 - <code title="post /threads/{thread_id}">client.beta.threads.<a href="./src/openai/resources/beta/threads/threads.py">update</a>(thread_id, \*\*<a href="src/openai/types/beta/thread_update_params.py">params</a>) -> <a href="./src/openai/types/beta/thread.py">Thread</a></code>
 - <code title="delete /threads/{thread_id}">client.beta.threads.<a href="./src/openai/resources/beta/threads/threads.py">delete</a>(thread_id) -> <a href="./src/openai/types/beta/thread_deleted.py">ThreadDeleted</a></code>
 - <code title="post /threads/runs">client.beta.threads.<a href="./src/openai/resources/beta/threads/threads.py">create_and_run</a>(\*\*<a href="src/openai/types/beta/thread_create_and_run_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/run.py">Run</a></code>
+- <code>client.beta.threads.<a href="./src/openai/resources/beta/threads/threads.py">create_and_run_stream</a>(\*args) -> AssistantStreamManager[AssistantEventHandler] | AssistantStreamManager[AssistantEventHandlerT]</code>
 
 ### Runs
 
@@ -235,18 +248,31 @@ Methods:
 - <code title="get /threads/{thread_id}/runs">client.beta.threads.runs.<a href="./src/openai/resources/beta/threads/runs/runs.py">list</a>(thread_id, \*\*<a href="src/openai/types/beta/threads/run_list_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/run.py">SyncCursorPage[Run]</a></code>
 - <code title="post /threads/{thread_id}/runs/{run_id}/cancel">client.beta.threads.runs.<a href="./src/openai/resources/beta/threads/runs/runs.py">cancel</a>(run_id, \*, thread_id) -> <a href="./src/openai/types/beta/threads/run.py">Run</a></code>
 - <code title="post /threads/{thread_id}/runs/{run_id}/submit_tool_outputs">client.beta.threads.runs.<a href="./src/openai/resources/beta/threads/runs/runs.py">submit_tool_outputs</a>(run_id, \*, thread_id, \*\*<a href="src/openai/types/beta/threads/run_submit_tool_outputs_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/run.py">Run</a></code>
+- <code>client.beta.threads.runs.<a href="./src/openai/resources/beta/threads/runs/runs.py">create_and_stream</a>(\*args) -> AssistantStreamManager[AssistantEventHandler] | AssistantStreamManager[AssistantEventHandlerT]</code>
+- <code>client.beta.threads.runs.<a href="./src/openai/resources/beta/threads/runs/runs.py">submit_tool_outputs_stream</a>(\*args) -> AssistantStreamManager[AssistantEventHandler] | AssistantStreamManager[AssistantEventHandlerT]</code>
 
 #### Steps
 
 Types:
 
 ```python
 from openai.types.beta.threads.runs import (
-    CodeToolCall,
+    CodeInterpreterLogs,
+    CodeInterpreterOutputImage,
+    CodeInterpreterToolCall,
+    CodeInterpreterToolCallDelta,
     FunctionToolCall,
+    FunctionToolCallDelta,
     MessageCreationStepDetails,
     RetrievalToolCall,
+    RetrievalToolCallDelta,
     RunStep,
+    RunStepDelta,
+    RunStepDeltaEvent,
+    RunStepDeltaMessageDelta,
+    ToolCall,
+    ToolCallDelta,
+    ToolCallDeltaObject,
     ToolCallsStepDetails,
 )
 ```
@@ -262,19 +288,35 @@ Types:
 
 ```python
 from openai.types.beta.threads import (
-    MessageContentImageFile,
-    MessageContentText,
-    ThreadMessage,
-    ThreadMessageDeleted,
+    Annotation,
+    AnnotationDelta,
+    FileCitationAnnotation,
+    FileCitationDeltaAnnotation,
+    FilePathAnnotation,
+    FilePathDeltaAnnotation,
+    ImageFile,
+    ImageFileContentBlock,
+    ImageFileDelta,
+    ImageFileDeltaBlock,
+    Message,
+    MessageContent,
+    MessageContentDelta,
+    MessageDeleted,
+    MessageDelta,
+    MessageDeltaEvent,
+    Text,
+    TextContentBlock,
+    TextDelta,
+    TextDeltaBlock,
 )
 ```
 
 Methods:
 
-- <code title="post /threads/{thread_id}/messages">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">create</a>(thread_id, \*\*<a href="src/openai/types/beta/threads/message_create_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/thread_message.py">ThreadMessage</a></code>
-- <code title="get /threads/{thread_id}/messages/{message_id}">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">retrieve</a>(message_id, \*, thread_id) -> <a href="./src/openai/types/beta/threads/thread_message.py">ThreadMessage</a></code>
-- <code title="post /threads/{thread_id}/messages/{message_id}">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">update</a>(message_id, \*, thread_id, \*\*<a href="src/openai/types/beta/threads/message_update_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/thread_message.py">ThreadMessage</a></code>
-- <code title="get /threads/{thread_id}/messages">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">list</a>(thread_id, \*\*<a href="src/openai/types/beta/threads/message_list_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/thread_message.py">SyncCursorPage[ThreadMessage]</a></code>
+- <code title="post /threads/{thread_id}/messages">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">create</a>(thread_id, \*\*<a href="src/openai/types/beta/threads/message_create_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/message.py">Message</a></code>
+- <code title="get /threads/{thread_id}/messages/{message_id}">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">retrieve</a>(message_id, \*, thread_id) -> <a href="./src/openai/types/beta/threads/message.py">Message</a></code>
+- <code title="post /threads/{thread_id}/messages/{message_id}">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">update</a>(message_id, \*, thread_id, \*\*<a href="src/openai/types/beta/threads/message_update_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/message.py">Message</a></code>
+- <code title="get /threads/{thread_id}/messages">client.beta.threads.messages.<a href="./src/openai/resources/beta/threads/messages/messages.py">list</a>(thread_id, \*\*<a href="src/openai/types/beta/threads/message_list_params.py">params</a>) -> <a href="./src/openai/types/beta/threads/message.py">SyncCursorPage[Message]</a></code>
 
 #### Files
 

examples/assistant_stream.py

@@ -0,0 +1,33 @@
+import openai
+
+# gets API Key from environment variable OPENAI_API_KEY
+client = openai.OpenAI()
+
+assistant = client.beta.assistants.create(
+    name="Math Tutor",
+    instructions="You are a personal math tutor. Write and run code to answer math questions.",
+    tools=[{"type": "code_interpreter"}],
+    model="gpt-4-1106-preview",
+)
+
+thread = client.beta.threads.create()
+
+message = client.beta.threads.messages.create(
+    thread_id=thread.id,
+    role="user",
+    content="I need to solve the equation `3x + 11 = 14`. Can you help me?",
+)
+
+print("starting run stream")
+
+stream = client.beta.threads.runs.create(
+    thread_id=thread.id,
+    assistant_id=assistant.id,
+    instructions="Please address the user as Jane Doe. The user has a premium account.",
+    stream=True,
+)
+
+for event in stream:
+    print(event.model_dump_json(indent=2, exclude_unset=True))
+
+client.beta.assistants.delete(assistant.id)

examples/assistant_stream_helpers.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from typing_extensions import override
+
+import openai
+from openai import AssistantEventHandler
+from openai.types.beta import AssistantStreamEvent
+from openai.types.beta.threads import Text, TextDelta
+from openai.types.beta.threads.runs import RunStep, RunStepDelta
+
+
+class EventHandler(AssistantEventHandler):
+    @override
+    def on_event(self, event: AssistantStreamEvent) -> None:
+        if event.event == "thread.run.step.created":
+            details = event.data.step_details
+            if details.type == "tool_calls":
+                print("Generating code to interpret:\n\n```py")
+        elif event.event == "thread.message.created":
+            print("\nResponse:\n")
+
+    @override
+    def on_text_delta(self, delta: TextDelta, snapshot: Text) -> None:
+        print(delta.value, end="", flush=True)
+
+    @override
+    def on_run_step_done(self, run_step: RunStep) -> None:
+        details = run_step.step_details
+        if details.type == "tool_calls":
+            for tool in details.tool_calls:
+                if tool.type == "code_interpreter":
+                    print("\n```\nExecuting code...")
+
+    @override
+    def on_run_step_delta(self, delta: RunStepDelta, snapshot: RunStep) -> None:
+        details = delta.step_details
+        if details is not None and details.type == "tool_calls":
+            for tool in details.tool_calls or []:
+                if tool.type == "code_interpreter" and tool.code_interpreter and tool.code_interpreter.input:
+                    print(tool.code_interpreter.input, end="", flush=True)
+
+
+def main() -> None:
+    client = openai.OpenAI()
+
+    assistant = client.beta.assistants.create(
+        name="Math Tutor",
+        instructions="You are a personal math tutor. Write and run code to answer math questions.",
+        tools=[{"type": "code_interpreter"}],
+        model="gpt-4-1106-preview",
+    )
+
+    try:
+        question = "I need to solve the equation `3x + 11 = 14`. Can you help me?"
+
+        thread = client.beta.threads.create(
+            messages=[
+                {
+                    "role": "user",
+                    "content": question,
+                },
+            ]
+        )
+        print(f"Question: {question}\n")
+
+        with client.beta.threads.runs.create_and_stream(
+            thread_id=thread.id,
+            assistant_id=assistant.id,
+            instructions="Please address the user as Jane Doe. The user has a premium account.",
+            event_handler=EventHandler(),
+        ) as stream:
+            stream.until_done()
+            print()
+    finally:
+        client.beta.assistants.delete(assistant.id)
+
+
+main()

pyproject.toml

@@ -60,6 +60,7 @@ dev-dependencies = [
     "nox",
     "dirty-equals>=0.6.0",
     "importlib-metadata>=6.7.0",
+    "inline-snapshot >=0.7.0",
     "azure-identity >=1.14.1",
     "types-tqdm > 4",
     "types-pyaudio > 0"

requirements-dev.lock

@@ -15,11 +15,15 @@ anyio==4.1.0
     # via openai
 argcomplete==3.1.2
     # via nox
+asttokens==2.4.1
+    # via inline-snapshot
 attrs==23.1.0
     # via pytest
 azure-core==1.30.1
     # via azure-identity
 azure-identity==1.15.0
+black==24.2.0
+    # via inline-snapshot
 certifi==2023.7.22
     # via httpcore
     # via httpx
@@ -28,6 +32,9 @@ cffi==1.16.0
     # via cryptography
 charset-normalizer==3.3.2
     # via requests
+click==8.1.7
+    # via black
+    # via inline-snapshot
 colorlog==6.7.0
     # via nox
 cryptography==42.0.5
@@ -41,6 +48,8 @@ distro==1.8.0
     # via openai
 exceptiongroup==1.1.3
     # via anyio
+executing==2.0.1
+    # via inline-snapshot
 filelock==3.12.4
     # via virtualenv
 h11==0.14.0
@@ -57,13 +66,15 @@ idna==3.4
 importlib-metadata==7.0.0
 iniconfig==2.0.0
     # via pytest
+inline-snapshot==0.7.0
 msal==1.27.0
     # via azure-identity
     # via msal-extensions
 msal-extensions==1.1.0
     # via azure-identity
 mypy==1.7.1
 mypy-extensions==1.0.0
+    # via black
     # via mypy
 nodeenv==1.8.0
     # via pyright
@@ -73,14 +84,18 @@ numpy==1.26.3
     # via pandas
     # via pandas-stubs
 packaging==23.2
+    # via black
     # via msal-extensions
     # via nox
     # via pytest
 pandas==2.1.4
     # via openai
 pandas-stubs==2.1.4.231227
     # via openai
+pathspec==0.12.1
+    # via black
 platformdirs==3.11.0
+    # via black
     # via virtualenv
 pluggy==1.3.0
     # via pytest
@@ -114,24 +129,31 @@ ruff==0.1.9
 setuptools==68.2.2
     # via nodeenv
 six==1.16.0
+    # via asttokens
     # via azure-core
     # via python-dateutil
 sniffio==1.3.0
     # via anyio
     # via httpx
     # via openai
 time-machine==2.9.0
+toml==0.10.2
+    # via inline-snapshot
 tomli==2.0.1
+    # via black
     # via mypy
     # via pytest
 tqdm==4.66.1
     # via openai
 types-pyaudio==0.2.16.20240106
 types-pytz==2024.1.0.20240203
     # via pandas-stubs
+types-toml==0.10.8.20240310
+    # via inline-snapshot
 types-tqdm==4.66.0.2
 typing-extensions==4.8.0
     # via azure-core
+    # via black
     # via mypy
     # via openai
     # via pydantic

src/openai/__init__.py

@@ -69,6 +69,10 @@
 from .version import VERSION as VERSION
 from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI
 from .lib._old_api import *
+from .lib.streaming import (
+    AssistantEventHandler as AssistantEventHandler,
+    AsyncAssistantEventHandler as AsyncAssistantEventHandler,
+)
 
 _setup_logging()