fix(oci): preserve v1 stream end finish reason

Author: fede-kamel

src/cohere/aws_client.py

@@ -4,14 +4,15 @@
 import typing
 
 import httpx
-from httpx import URL, SyncByteStream, ByteStream
+from httpx import URL, ByteStream
 
 from . import GenerateStreamedResponse, Generation, \
     NonStreamedChatResponse, EmbedResponse, StreamedChatResponse, RerankResponse, ApiMeta, ApiMetaTokens, \
     ApiMetaBilledUnits
 from .client import Client, ClientEnvironment
 from .core import construct_type
 from .manually_maintained.lazy_aws_deps import lazy_boto3, lazy_botocore
+from .manually_maintained.streaming import Streamer
 from .client_v2 import ClientV2
 
 class AwsClient(Client):
@@ -112,16 +113,6 @@ def get_event_hooks(
                               })
 
 
-class Streamer(SyncByteStream):
-    lines: typing.Iterator[bytes]
-
-    def __init__(self, lines: typing.Iterator[bytes]):
-        self.lines = lines
-
-    def __iter__(self) -> typing.Iterator[bytes]:
-        return self.lines
-
-
 response_mapping: typing.Dict[str, typing.Any] = {
     "chat": NonStreamedChatResponse,
     "embed": EmbedResponse,
@@ -291,4 +282,4 @@ def get_api_version(*, version: str):
         "v2": 2,
     }
 
-    return int_version.get(version, 1)
+    return int_version.get(version, 1)

src/cohere/manually_maintained/streaming.py

@@ -0,0 +1,15 @@
+import typing
+
+from httpx import SyncByteStream
+
+
+class Streamer(SyncByteStream):
+    """Wrap an iterator of bytes for httpx streaming responses."""
+
+    lines: typing.Iterator[bytes]
+
+    def __init__(self, lines: typing.Iterator[bytes]):
+        self.lines = lines
+
+    def __iter__(self) -> typing.Iterator[bytes]:
+        return self.lines

src/cohere/oci_client.py

@@ -12,7 +12,8 @@
 from .client import Client, ClientEnvironment
 from .client_v2 import ClientV2
 from .manually_maintained.lazy_oci_deps import lazy_oci
-from httpx import URL, ByteStream, SyncByteStream
+from .manually_maintained.streaming import Streamer
+from httpx import URL, ByteStream
 
 
 class OciClient(Client):
@@ -232,20 +233,6 @@ def __init__(
 EventHook = typing.Callable[..., typing.Any]
 
 
-
-
-class Streamer(SyncByteStream):
-    """Wraps an iterator of bytes for streaming responses."""
-
-    lines: typing.Iterator[bytes]
-
-    def __init__(self, lines: typing.Iterator[bytes]):
-        self.lines = lines
-
-    def __iter__(self) -> typing.Iterator[bytes]:
-        return self.lines
-
-
 def _load_oci_config(
     auth_type: str,
     config_path: typing.Optional[str],
@@ -1006,6 +993,7 @@ def transform_oci_stream_wrapper(
     final_finish_reason = "COMPLETE"
     final_usage: typing.Optional[typing.Dict[str, typing.Any]] = None
     full_v1_text = ""
+    final_v1_finish_reason = "COMPLETE"
     buffer = b""
 
     def _emit_v2_event(event: typing.Dict[str, typing.Any]) -> bytes:
@@ -1053,10 +1041,12 @@ def _transform_v2_event(oci_event: typing.Dict[str, typing.Any]) -> typing.Itera
             yield _emit_v2_event(cohere_event)
 
     def _transform_v1_event(oci_event: typing.Dict[str, typing.Any]) -> bytes:
-        nonlocal full_v1_text
+        nonlocal full_v1_text, final_v1_finish_reason
         event = typing.cast(typing.Dict[str, typing.Any], transform_stream_event(endpoint, oci_event, is_v2=False))
         if event.get("event_type") == "text-generation" and event.get("text"):
             full_v1_text += typing.cast(str, event["text"])
+        if "finishReason" in oci_event:
+            final_v1_finish_reason = oci_event.get("finishReason", final_v1_finish_reason)
         return _emit_v1_event(event)
 
     def _process_line(line: str) -> typing.Iterator[bytes]:
@@ -1083,7 +1073,7 @@ def _process_line(line: str) -> typing.Iterator[bytes]:
                         "response": {
                             "text": full_v1_text,
                             "generation_id": generation_id,
-                            "finish_reason": final_finish_reason,
+                            "finish_reason": final_v1_finish_reason,
                         },
                     }
                 )

tests/test_oci_client.py

@@ -1039,6 +1039,26 @@ def test_stream_wrapper_skips_malformed_json_with_warning(self):
         # Should get message-start + content-start + content-delta + content-end + message-end.
         self.assertEqual(len(events), 5)
 
+    def test_v1_stream_wrapper_preserves_finish_reason_in_stream_end(self):
+        """Test that V1 stream-end uses the OCI finish reason from the final event."""
+        import json
+        from cohere.oci_client import transform_oci_stream_wrapper
+
+        chunks = [
+            b'data: {"text": "Hello", "isFinished": false}\n',
+            b'data: {"text": " world", "isFinished": true, "finishReason": "MAX_TOKENS"}\n',
+            b"data: [DONE]\n",
+        ]
+
+        events = [
+            json.loads(raw.decode("utf-8"))
+            for raw in transform_oci_stream_wrapper(iter(chunks), "chat_stream", is_v2=False)
+        ]
+
+        self.assertEqual(events[2]["event_type"], "stream-end")
+        self.assertEqual(events[2]["response"]["text"], "Hello world")
+        self.assertEqual(events[2]["response"]["finish_reason"], "MAX_TOKENS")
+
     def test_stream_wrapper_raises_on_transform_error(self):
         """Test that transform errors in stream produce OCI-specific error, not opaque httpx error."""
         from cohere.oci_client import transform_oci_stream_wrapper