Merge pull request #1188 from stainless-sdks/pakrym/stream

Add support for streaming background Responses

Author: kwhinnery-openai

examples/responses/background.py

@@ -0,0 +1,46 @@
+from typing import List
+
+import rich
+from pydantic import BaseModel
+
+from openai import OpenAI
+
+
+class Step(BaseModel):
+    explanation: str
+    output: str
+
+
+class MathResponse(BaseModel):
+    steps: List[Step]
+    final_answer: str
+
+
+client = OpenAI()
+id = None
+
+with client.responses.create(
+    input="solve 8x + 31 = 2",
+    model="gpt-4o-2024-08-06",
+    background=True,
+    stream=True,
+) as stream:
+    for event in stream:
+        if event.type == "response.created":
+            id = event.response.id
+        if "output_text" in event.type:
+            rich.print(event)
+        if event.sequence_number == 10:
+            break
+
+print("Interrupted. Continuing...")
+
+assert id is not None
+with client.responses.retrieve(
+    response_id=id,
+    stream=True,
+    starting_after=10,
+) as stream:
+    for event in stream:
+        if "output_text" in event.type:
+            rich.print(event)

examples/responses/background_async.py

@@ -0,0 +1,52 @@
+import asyncio
+from typing import List
+
+import rich
+from pydantic import BaseModel
+
+from openai._client import AsyncOpenAI
+
+
+class Step(BaseModel):
+    explanation: str
+    output: str
+
+
+class MathResponse(BaseModel):
+    steps: List[Step]
+    final_answer: str
+
+
+async def main() -> None:
+    client = AsyncOpenAI()
+    id = None
+
+    async with await client.responses.create(
+        input="solve 8x + 31 = 2",
+        model="gpt-4o-2024-08-06",
+        background=True,
+        stream=True,
+    ) as stream:
+        async for event in stream:
+            if event.type == "response.created":
+                id = event.response.id
+            if "output_text" in event.type:
+                rich.print(event)
+            if event.sequence_number == 10:
+                break
+
+    print("Interrupted. Continuing...")
+
+    assert id is not None
+    async with await client.responses.retrieve(
+        response_id=id,
+        stream=True,
+        starting_after=10,
+    ) as stream:
+        async for event in stream:
+            if "output_text" in event.type:
+                rich.print(event)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/responses/background_streaming.py

@@ -0,0 +1,47 @@
+from typing import List
+
+import rich
+from pydantic import BaseModel
+
+from openai import OpenAI
+
+
+class Step(BaseModel):
+    explanation: str
+    output: str
+
+
+class MathResponse(BaseModel):
+    steps: List[Step]
+    final_answer: str
+
+
+client = OpenAI()
+id = None
+with client.responses.stream(
+    input="solve 8x + 31 = 2",
+    model="gpt-4o-2024-08-06",
+    text_format=MathResponse,
+    background=True,
+) as stream:
+    for event in stream:
+        if event.type == "response.created":
+            id = event.response.id
+        if "output_text" in event.type:
+            rich.print(event)
+        if event.sequence_number == 10:
+            break
+
+print("Interrupted. Continuing...")
+
+assert id is not None
+with client.responses.stream(
+    response_id=id,
+    starting_after=10,
+    text_format=MathResponse,
+) as stream:
+    for event in stream:
+        if "output_text" in event.type:
+            rich.print(event)
+
+    rich.print(stream.get_final_response())

examples/responses/background_streaming_async.py

@@ -0,0 +1,53 @@
+import asyncio
+from typing import List
+
+import rich
+from pydantic import BaseModel
+
+from openai import AsyncOpenAI
+
+
+class Step(BaseModel):
+    explanation: str
+    output: str
+
+
+class MathResponse(BaseModel):
+    steps: List[Step]
+    final_answer: str
+
+
+async def main() -> None:
+    client = AsyncOpenAI()
+    id = None
+    async with client.responses.stream(
+        input="solve 8x + 31 = 2",
+        model="gpt-4o-2024-08-06",
+        text_format=MathResponse,
+        background=True,
+    ) as stream:
+        async for event in stream:
+            if event.type == "response.created":
+                id = event.response.id
+            if "output_text" in event.type:
+                rich.print(event)
+            if event.sequence_number == 10:
+                break
+
+    print("Interrupted. Continuing...")
+
+    assert id is not None
+    async with client.responses.stream(
+        response_id=id,
+        starting_after=10,
+        text_format=MathResponse,
+    ) as stream:
+        async for event in stream:
+            if "output_text" in event.type:
+                rich.print(event)
+
+        rich.print(stream.get_final_response())
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/openai/lib/streaming/responses/_responses.py

@@ -34,11 +34,13 @@ def __init__(
         raw_stream: Stream[RawResponseStreamEvent],
         text_format: type[TextFormatT] | NotGiven,
         input_tools: Iterable[ToolParam] | NotGiven,
+        starting_after: int | None,
     ) -> None:
         self._raw_stream = raw_stream
         self._response = raw_stream.response
         self._iterator = self.__stream__()
         self._state = ResponseStreamState(text_format=text_format, input_tools=input_tools)
+        self._starting_after = starting_after
 
     def __next__(self) -> ResponseStreamEvent[TextFormatT]:
         return self._iterator.__next__()
@@ -54,7 +56,8 @@ def __stream__(self) -> Iterator[ResponseStreamEvent[TextFormatT]]:
         for sse_event in self._raw_stream:
             events_to_fire = self._state.handle_event(sse_event)
             for event in events_to_fire:
-                yield event
+                if self._starting_after is None or event.sequence_number > self._starting_after:
+                    yield event
 
     def __exit__(
         self,
@@ -96,11 +99,13 @@ def __init__(
         *,
         text_format: type[TextFormatT] | NotGiven,
         input_tools: Iterable[ToolParam] | NotGiven,
+        starting_after: int | None,
     ) -> None:
         self.__stream: ResponseStream[TextFormatT] | None = None
         self.__api_request = api_request
         self.__text_format = text_format
         self.__input_tools = input_tools
+        self.__starting_after = starting_after
 
     def __enter__(self) -> ResponseStream[TextFormatT]:
         raw_stream = self.__api_request()
@@ -109,6 +114,7 @@ def __enter__(self) -> ResponseStream[TextFormatT]:
             raw_stream=raw_stream,
             text_format=self.__text_format,
             input_tools=self.__input_tools,
+            starting_after=self.__starting_after,
         )
 
         return self.__stream
@@ -130,11 +136,13 @@ def __init__(
         raw_stream: AsyncStream[RawResponseStreamEvent],
         text_format: type[TextFormatT] | NotGiven,
         input_tools: Iterable[ToolParam] | NotGiven,
+        starting_after: int | None,
     ) -> None:
         self._raw_stream = raw_stream
         self._response = raw_stream.response
         self._iterator = self.__stream__()
         self._state = ResponseStreamState(text_format=text_format, input_tools=input_tools)
+        self._starting_after = starting_after
 
     async def __anext__(self) -> ResponseStreamEvent[TextFormatT]:
         return await self._iterator.__anext__()
@@ -147,7 +155,8 @@ async def __stream__(self) -> AsyncIterator[ResponseStreamEvent[TextFormatT]]:
         async for sse_event in self._raw_stream:
             events_to_fire = self._state.handle_event(sse_event)
             for event in events_to_fire:
-                yield event
+                if self._starting_after is None or event.sequence_number > self._starting_after:
+                    yield event
 
     async def __aenter__(self) -> Self:
         return self
@@ -192,11 +201,13 @@ def __init__(
         *,
         text_format: type[TextFormatT] | NotGiven,
         input_tools: Iterable[ToolParam] | NotGiven,
+        starting_after: int | None,
     ) -> None:
         self.__stream: AsyncResponseStream[TextFormatT] | None = None
         self.__api_request = api_request
         self.__text_format = text_format
         self.__input_tools = input_tools
+        self.__starting_after = starting_after
 
     async def __aenter__(self) -> AsyncResponseStream[TextFormatT]:
         raw_stream = await self.__api_request
@@ -205,6 +216,7 @@ async def __aenter__(self) -> AsyncResponseStream[TextFormatT]:
             raw_stream=raw_stream,
             text_format=self.__text_format,
             input_tools=self.__input_tools,
+            starting_after=self.__starting_after,
         )
 
         return self.__stream