fix: ensure streams are always closed

Author: stainless-app[bot]

src/gradient/_streaming.py

@@ -55,29 +55,30 @@ def __stream__(self) -> Iterator[_T]:
         process_data = self._client._process_response_data
         iterator = self._iter_events()
 
-        for sse in iterator:
-            if sse.data.startswith("[DONE]"):
-                break
-
-            data = sse.json()
-            if is_mapping(data) and data.get("error"):
-                message = None
-                error = data.get("error")
-                if is_mapping(error):
-                    message = error.get("message")
-                if not message or not isinstance(message, str):
-                    message = "An error occurred during streaming"
-
-                raise APIError(
-                    message=message,
-                    request=self.response.request,
-                    body=data["error"],
-                )
-
-            yield process_data(data=data, cast_to=cast_to, response=response)
-
-        # As we might not fully consume the response stream, we need to close it explicitly
-        response.close()
+        try:
+            for sse in iterator:
+                if sse.data.startswith("[DONE]"):
+                    break
+
+                data = sse.json()
+                if is_mapping(data) and data.get("error"):
+                    message = None
+                    error = data.get("error")
+                    if is_mapping(error):
+                        message = error.get("message")
+                    if not message or not isinstance(message, str):
+                        message = "An error occurred during streaming"
+
+                    raise APIError(
+                        message=message,
+                        request=self.response.request,
+                        body=data["error"],
+                    )
+
+                yield process_data(data=data, cast_to=cast_to, response=response)
+        finally:
+            # Ensure the response is closed even if the consumer doesn't read all data
+            response.close()
 
     def __enter__(self) -> Self:
         return self
@@ -136,29 +137,30 @@ async def __stream__(self) -> AsyncIterator[_T]:
         process_data = self._client._process_response_data
         iterator = self._iter_events()
 
-        async for sse in iterator:
-            if sse.data.startswith("[DONE]"):
-                break
-
-            data = sse.json()
-            if is_mapping(data) and data.get("error"):
-                message = None
-                error = data.get("error")
-                if is_mapping(error):
-                    message = error.get("message")
-                if not message or not isinstance(message, str):
-                    message = "An error occurred during streaming"
-
-                raise APIError(
-                    message=message,
-                    request=self.response.request,
-                    body=data["error"],
-                )
-
-            yield process_data(data=data, cast_to=cast_to, response=response)
-
-        # As we might not fully consume the response stream, we need to close it explicitly
-        await response.aclose()
+        try:
+            async for sse in iterator:
+                if sse.data.startswith("[DONE]"):
+                    break
+
+                data = sse.json()
+                if is_mapping(data) and data.get("error"):
+                    message = None
+                    error = data.get("error")
+                    if is_mapping(error):
+                        message = error.get("message")
+                    if not message or not isinstance(message, str):
+                        message = "An error occurred during streaming"
+
+                    raise APIError(
+                        message=message,
+                        request=self.response.request,
+                        body=data["error"],
+                    )
+
+                yield process_data(data=data, cast_to=cast_to, response=response)
+        finally:
+            # Ensure the response is closed even if the consumer doesn't read all data
+            await response.aclose()
 
     async def __aenter__(self) -> Self:
         return self