fix: address OCI streaming review feedback

Author: fede-kamel

src/cohere/oci_client.py

@@ -384,6 +384,7 @@ def map_request_to_oci(
                 "OCI config profile is missing 'key_file'. "
                 "Session-based auth requires a key_file entry in your OCI config profile."
             )
+        key_file = os.path.expanduser(key_file)
         private_key = oci.signer.load_private_key_from_file(key_file)
 
         signer = oci.auth.signers.SecurityTokenSigner(
@@ -411,7 +412,6 @@ def _event_hook(request: httpx.Request) -> None:
         url = get_oci_url(
             region=oci_region,
             endpoint=endpoint,
-            stream="stream" in endpoint or body.get("stream", False),
         )
 
         # Transform request body to OCI format
@@ -517,16 +517,13 @@ def _hook(response: httpx.Response) -> None:
 def get_oci_url(
     region: str,
     endpoint: str,
-    stream: bool = False,
 ) -> str:
     """
     Map Cohere endpoints to OCI Generative AI endpoints.
 
     Args:
         region: OCI region (e.g., "us-chicago-1")
         endpoint: Cohere endpoint name
-        stream: Whether this is a streaming request
-
     Returns:
         Full OCI Generative AI endpoint URL
     """
@@ -972,6 +969,21 @@ def transform_oci_stream_wrapper(
 
     generation_id = str(uuid.uuid4())
     emitted_start = False
+    v1_text_parts: typing.List[str] = []
+    v1_finish_reason = "COMPLETE"
+    v1_response: typing.Dict[str, typing.Any] = {
+        "text": "",
+        "generation_id": generation_id,
+        "response_id": None,
+        "citations": [],
+        "documents": [],
+        "is_search_required": None,
+        "search_queries": [],
+        "search_results": [],
+        "chat_history": [],
+        "meta": {"api_version": {"version": "1"}},
+    }
+    v2_message_end_delta: typing.Dict[str, typing.Any] = {}
     buffer = b""
     for chunk in stream:
         buffer += chunk
@@ -982,10 +994,25 @@ def transform_oci_stream_wrapper(
             if line.startswith("data: "):
                 data_str = line[6:]  # Remove "data: " prefix
                 if data_str.strip() == "[DONE]":
-                    # Emit message-end event for V2 before stopping
                     if is_v2:
-                        message_end_event = {"type": "message-end"}
+                        message_end_event: typing.Dict[str, typing.Any] = {
+                            "type": "message-end",
+                            "id": generation_id,
+                        }
+                        if v2_message_end_delta:
+                            message_end_event["delta"] = v2_message_end_delta
                         yield b"data: " + json.dumps(message_end_event).encode("utf-8") + b"\n\n"
+                    elif endpoint in ["chat_stream", "chat"]:
+                        stream_end_event = {
+                            "event_type": "stream-end",
+                            "finish_reason": v1_finish_reason,
+                            "response": {
+                                **v1_response,
+                                "text": "".join(v1_text_parts),
+                                "finish_reason": v1_finish_reason,
+                            },
+                        }
+                        yield json.dumps(stream_end_event).encode("utf-8") + b"\n"
                     # Return to stop the generator completely
                     return
 
@@ -1023,10 +1050,61 @@ def transform_oci_stream_wrapper(
                         yield b"data: " + json.dumps(content_start).encode("utf-8") + b"\n\n"
                         emitted_start = True
 
-                    cohere_event = transform_stream_event(endpoint, oci_event, is_v2)
                     if is_v2:
-                        yield b"data: " + json.dumps(cohere_event).encode("utf-8") + b"\n\n"
+                        if endpoint in ["chat_stream", "chat"] and "finishReason" in oci_event:
+                            content_event = transform_stream_event(
+                                endpoint,
+                                {k: v for k, v in oci_event.items() if k != "finishReason"},
+                                is_v2,
+                            )
+                            content_payload = (
+                                content_event.get("delta", {})
+                                .get("message", {})
+                                .get("content", {})
+                            )
+                            if content_payload.get("text") or content_payload.get("thinking"):
+                                yield b"data: " + json.dumps(content_event).encode("utf-8") + b"\n\n"
+
+                            usage_data = oci_event.get("usage", {})
+                            usage: typing.Dict[str, typing.Any] = {
+                                "tokens": {
+                                    "input_tokens": usage_data.get("inputTokens", 0),
+                                    "output_tokens": usage_data.get("completionTokens", 0),
+                                }
+                            }
+                            if usage_data.get("inputTokens") or usage_data.get("completionTokens"):
+                                usage["billed_units"] = {
+                                    "input_tokens": usage_data.get("inputTokens", 0),
+                                    "output_tokens": usage_data.get("completionTokens", 0),
+                                }
+                            v2_message_end_delta = {
+                                "finish_reason": oci_event.get("finishReason"),
+                                "usage": usage,
+                            }
+
+                            content_end_event = {"type": "content-end", "index": 0}
+                            yield b"data: " + json.dumps(content_end_event).encode("utf-8") + b"\n\n"
+                        else:
+                            cohere_event = transform_stream_event(endpoint, oci_event, is_v2)
+                            yield b"data: " + json.dumps(cohere_event).encode("utf-8") + b"\n\n"
                     else:
+                        if endpoint in ["chat_stream", "chat"]:
+                            text = oci_event.get("text", "")
+                            if text:
+                                v1_text_parts.append(text)
+                            if "finishReason" in oci_event:
+                                v1_finish_reason = oci_event.get("finishReason", v1_finish_reason)
+                            if "chatHistory" in oci_event:
+                                v1_response["chat_history"] = oci_event.get("chatHistory", [])
+                            if "citations" in oci_event:
+                                v1_response["citations"] = oci_event.get("citations", [])
+                            if "documents" in oci_event:
+                                v1_response["documents"] = oci_event.get("documents", [])
+                            if "searchQueries" in oci_event:
+                                v1_response["search_queries"] = oci_event.get("searchQueries", [])
+                            if "searchResults" in oci_event:
+                                v1_response["search_results"] = oci_event.get("searchResults", [])
+                        cohere_event = transform_stream_event(endpoint, oci_event, is_v2)
                         yield json.dumps(cohere_event).encode("utf-8") + b"\n"
                 except Exception as e:
                     raise RuntimeError(

tests/test_oci_client.py

@@ -15,6 +15,7 @@
 import unittest
 
 import cohere
+from cohere.errors import NotFoundError
 
 
 @unittest.skipIf(os.getenv("TEST_OCI") is None, "TEST_OCI not set")
@@ -467,11 +468,14 @@ def test_embed_english_v3(self):
 
     def test_embed_light_v3(self):
         """Test embed-english-light-v3.0 model."""
-        response = self.client.embed(
-            model="embed-english-light-v3.0",
-            texts=["Test"],
-            input_type="search_document",
-        )
+        try:
+            response = self.client.embed(
+                model="embed-english-light-v3.0",
+                texts=["Test"],
+                input_type="search_document",
+            )
+        except NotFoundError:
+            self.skipTest("embed-english-light-v3.0 is not available in this OCI region/profile")
         self.assertIsNotNone(response.embeddings)
         self.assertEqual(len(response.embeddings[0]), 384)
 
@@ -709,7 +713,7 @@ def test_get_oci_url_known_endpoints(self):
         url = get_oci_url("us-chicago-1", "chat")
         self.assertIn("/actions/chat", url)
 
-        url = get_oci_url("us-chicago-1", "chat_stream", stream=True)
+        url = get_oci_url("us-chicago-1", "chat_stream")
         self.assertIn("/actions/chat", url)
 
     def test_get_oci_url_unknown_endpoint_raises(self):
@@ -744,7 +748,7 @@ def test_stream_wrapper_emits_full_event_lifecycle(self):
 
         chunks = [
             b'data: {"message": {"content": [{"type": "TEXT", "text": "Hello"}]}}\n',
-            b'data: {"message": {"content": [{"type": "TEXT", "text": " world"}]}, "finishReason": "COMPLETE"}\n',
+            b'data: {"message": {"content": [{"type": "TEXT", "text": " world"}]}, "finishReason": "COMPLETE", "usage": {"inputTokens": 3, "completionTokens": 2}}\n',
             b"data: [DONE]\n",
         ]
 
@@ -758,13 +762,19 @@ def test_stream_wrapper_emits_full_event_lifecycle(self):
         self.assertEqual(event_types[0], "message-start")
         self.assertEqual(event_types[1], "content-start")
         self.assertEqual(event_types[2], "content-delta")
-        self.assertEqual(event_types[3], "content-end")
-        self.assertEqual(event_types[4], "message-end")
+        self.assertEqual(event_types[3], "content-delta")
+        self.assertEqual(event_types[4], "content-end")
+        self.assertEqual(event_types[5], "message-end")
 
         self.assertIn("id", events[0])
         self.assertEqual(events[0]["delta"]["message"]["role"], "assistant")
         self.assertEqual(events[1]["index"], 0)
         self.assertEqual(events[1]["delta"]["message"]["content"]["type"], "text")
+        self.assertEqual(events[2]["delta"]["message"]["content"]["text"], "Hello")
+        self.assertEqual(events[3]["delta"]["message"]["content"]["text"], " world")
+        self.assertEqual(events[5]["delta"]["finish_reason"], "COMPLETE")
+        self.assertEqual(events[5]["delta"]["usage"]["tokens"]["input_tokens"], 3)
+        self.assertEqual(events[5]["delta"]["usage"]["tokens"]["output_tokens"], 2)
 
     def test_stream_wrapper_skips_malformed_json_with_warning(self):
         """Test that malformed JSON in SSE streams is skipped with a warning."""
@@ -779,6 +789,60 @@ def test_stream_wrapper_skips_malformed_json_with_warning(self):
         events = list(transform_oci_stream_wrapper(iter(chunks), "chat", is_v2=True))
         self.assertEqual(len(events), 4)
 
+    def test_v1_stream_wrapper_emits_stream_end(self):
+        """Test that V1 chat streams end with a stream-end event containing the response payload."""
+        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": "COMPLETE"}\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[0]["event_type"], "text-generation")
+        self.assertEqual(events[1]["event_type"], "text-generation")
+        self.assertEqual(events[2]["event_type"], "stream-end")
+        self.assertEqual(events[2]["finish_reason"], "COMPLETE")
+        self.assertEqual(events[2]["response"]["text"], "Hello world")
+
+    def test_session_auth_expands_key_file_path(self):
+        """Test that session-based auth expands key_file paths before loading them."""
+        from unittest.mock import MagicMock, patch
+        from cohere.oci_client import map_request_to_oci
+
+        mock_private_key = object()
+        mock_security_token_signer = MagicMock()
+        mock_oci = MagicMock()
+        mock_oci.signer.load_private_key_from_file.return_value = mock_private_key
+        mock_oci.auth.signers.SecurityTokenSigner.return_value = mock_security_token_signer
+
+        oci_config = {
+            "security_token_file": "~/.oci/sessions/TEST/token",
+            "key_file": "~/.oci/sessions/TEST/oci_api_key.pem",
+        }
+
+        with patch("cohere.oci_client.lazy_oci", return_value=mock_oci), patch(
+            "builtins.open", create=True
+        ) as mock_open:
+            mock_open.return_value.__enter__.return_value.read.return_value = "token"
+            with patch("os.path.expanduser", side_effect=lambda p: p.replace("~", "/Users/test")):
+                with patch("cohere.oci_client.transform_request_to_oci", return_value={"compartmentId": "c"}):
+                    hook = map_request_to_oci(oci_config, "us-chicago-1", "compartment-123", is_v2_client=False)
+                    request = MagicMock()
+                    request.url.path = "/v1/chat"
+                    request.read.return_value = b'{"message":"hi"}'
+                    request.method = "POST"
+                    request.extensions = {}
+                    hook(request)
+
+        mock_oci.signer.load_private_key_from_file.assert_called_with("/Users/test/.oci/sessions/TEST/oci_api_key.pem")
+
     def test_stream_wrapper_raises_on_transform_error(self):
         """Test that stream transformation errors are re-raised with OCI-specific context."""
         from cohere.oci_client import transform_oci_stream_wrapper