Update handling of SSE parsing in AsyncRawV2Client

Author: aditya-arolkar-swe

src/cohere/v2/raw_client.py

@@ -245,11 +245,9 @@ def _iter():
                                         ),
                                     )
                                 except json.JSONDecodeError as e:
-                                    # Log the problematic data for debugging
                                     print(f"JSON decode error: {e}, data: {repr(_sse.data)}")
                                     continue
                                 except Exception as e:
-                                    # Log other parsing errors
                                     print(f"Parsing error: {e}, event: {_sse.event}, data: {repr(_sse.data)}")
                                     continue
 
@@ -1335,20 +1333,33 @@ async def _stream() -> AsyncHttpResponse[typing.AsyncIterator[V2ChatStreamRespon
                     if 200 <= _response.status_code < 300:
 
                         async def _iter():
-                            _event_source = httpx_sse.EventSource(_response)
+                            _event_source = EventSource(_response)
                             async for _sse in _event_source.aiter_sse():
-                                if _sse.data == None:
-                                    return
                                 try:
+                                    # Skip empty events
+                                    if not _sse.data or _sse.data.strip() == "":
+                                        continue
+                                        
+                                    # Handle [DONE] token from OpenAI-style APIs
+                                    if _sse.data.strip() == '[DONE]':
+                                        continue
+                                        
+                                    parsed_data = json.loads(_sse.data)
+
                                     yield typing.cast(
                                         V2ChatStreamResponse,
                                         construct_type(
                                             type_=V2ChatStreamResponse,  # type: ignore
-                                            object_=json.loads(_sse.data),
+                                            object_=parsed_data,
                                         ),
                                     )
-                                except Exception:
-                                    pass
+                                except json.JSONDecodeError as e:
+                                    print(f"JSON decode error: {e}, data: {repr(_sse.data)}")
+                                    continue
+                                except Exception as e:
+                                    print(f"Parsing error: {e}, event: {_sse.event}, data: {repr(_sse.data)}")
+                                    continue
+                                
                             return
 
                         return AsyncHttpResponse(response=_response, data=_iter())