chore(client): refactor streaming slightly to better future proof it

Author: dtmeadows

src/openai/_streaming.py

@@ -59,14 +59,11 @@ def __stream__(self) -> Iterator[_T]:
             if sse.data.startswith("[DONE]"):
                 break
 
-            if sse.event is None or (
-                sse.event.startswith("response.") or 
-                sse.event.startswith("transcript.") or 
-                sse.event.startswith("image_edit.") or 
-                sse.event.startswith("image_generation.")
-            ):
+            # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
+            if sse.event and sse.event.startswith("thread."):
                 data = sse.json()
-                if is_mapping(data) and data.get("error"):
+
+                if sse.event == "error" and is_mapping(data) and data.get("error"):
                     message = None
                     error = data.get("error")
                     if is_mapping(error):
@@ -80,12 +77,10 @@ def __stream__(self) -> Iterator[_T]:
                         body=data["error"],
                     )
 
-                yield process_data(data=data, cast_to=cast_to, response=response)
-
+                yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response)
             else:
                 data = sse.json()
-
-                if sse.event == "error" and is_mapping(data) and data.get("error"):
+                if is_mapping(data) and data.get("error"):
                     message = None
                     error = data.get("error")
                     if is_mapping(error):
@@ -99,7 +94,7 @@ def __stream__(self) -> Iterator[_T]:
                         body=data["error"],
                     )
 
-                yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response)
+                yield process_data(data=data, cast_to=cast_to, response=response)
 
         # Ensure the entire stream is consumed
         for _sse in iterator:
@@ -166,9 +161,11 @@ async def __stream__(self) -> AsyncIterator[_T]:
             if sse.data.startswith("[DONE]"):
                 break
 
-            if sse.event is None or sse.event.startswith("response.") or sse.event.startswith("transcript."):
+            # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
+            if sse.event and sse.event.startswith("thread."):
                 data = sse.json()
-                if is_mapping(data) and data.get("error"):
+
+                if sse.event == "error" and is_mapping(data) and data.get("error"):
                     message = None
                     error = data.get("error")
                     if is_mapping(error):
@@ -182,12 +179,10 @@ async def __stream__(self) -> AsyncIterator[_T]:
                         body=data["error"],
                     )
 
-                yield process_data(data=data, cast_to=cast_to, response=response)
-
+                yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response)
             else:
                 data = sse.json()
-
-                if sse.event == "error" and is_mapping(data) and data.get("error"):
+                if is_mapping(data) and data.get("error"):
                     message = None
                     error = data.get("error")
                     if is_mapping(error):
@@ -201,7 +196,7 @@ async def __stream__(self) -> AsyncIterator[_T]:
                         body=data["error"],
                     )
 
-                yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response)
+                yield process_data(data=data, cast_to=cast_to, response=response)
 
         # Ensure the entire stream is consumed
         async for _sse in iterator: