fix(oci): tighten mixed v1 v2 stream coverage

Author: fede-kamel

src/cohere/oci_client.py

@@ -566,7 +566,12 @@ def get_oci_url(
         "rerank": "rerankText",  # OCI uses rerankText, not rerank
     }
 
-    action = action_map.get(endpoint, endpoint)
+    action = action_map.get(endpoint)
+    if action is None:
+        raise ValueError(
+            f"Endpoint '{endpoint}' is not supported by OCI Generative AI. "
+            f"Supported endpoints: {list(action_map.keys())}"
+        )
     return f"{base}/{api_version}/actions/{action}"
 
 
@@ -1089,11 +1094,14 @@ def _process_line(line: str) -> typing.Iterator[bytes]:
         except json.JSONDecodeError:
             return
 
-        if is_v2:
-            for event_bytes in _transform_v2_event(oci_event):
-                yield event_bytes
-        else:
-            yield _transform_v1_event(oci_event)
+        try:
+            if is_v2:
+                for event_bytes in _transform_v2_event(oci_event):
+                    yield event_bytes
+            else:
+                yield _transform_v1_event(oci_event)
+        except Exception as exc:
+            raise RuntimeError(f"OCI stream event transformation failed for endpoint '{endpoint}': {exc}") from exc
 
     for chunk in stream:
         buffer += chunk
@@ -1129,6 +1137,9 @@ def transform_stream_event(
             content_value = ""
             message = oci_event.get("message")
 
+            if "message" in oci_event and not isinstance(message, dict):
+                raise TypeError("OCI V2 stream event message must be an object")
+
             if isinstance(message, dict) and "content" in message:
                 content_list = message["content"]
                 if content_list and isinstance(content_list, list) and len(content_list) > 0:

tests/test_oci_client.py

@@ -956,6 +956,7 @@ def test_embed_response_lowercases_embedding_keys(self):
         result = transform_oci_response_to_cohere(
             "embed",
             {"id": "embed-id", "embeddings": {"FLOAT": [[0.1, 0.2]], "INT8": [[1, 2]]}},
+            is_v2=True,
         )
 
         self.assertIn("float", result["embeddings"])
@@ -1003,7 +1004,7 @@ def test_stream_wrapper_emits_full_event_lifecycle(self):
         ]
 
         events = []
-        for raw in transform_oci_stream_wrapper(iter(chunks), "chat"):
+        for raw in transform_oci_stream_wrapper(iter(chunks), "chat", is_v2=True):
             line = raw.decode("utf-8").strip()
             if line.startswith("data: "):
                 events.append(json.loads(line[6:]))
@@ -1034,7 +1035,7 @@ def test_stream_wrapper_skips_malformed_json_with_warning(self):
             b'data: {"message": {"content": [{"type": "TEXT", "text": "hello"}]}}\n',
             b'data: [DONE]\n',
         ]
-        events = list(transform_oci_stream_wrapper(iter(chunks), "chat"))
+        events = list(transform_oci_stream_wrapper(iter(chunks), "chat", is_v2=True))
         # Should get message-start + content-start + content-delta + content-end + message-end.
         self.assertEqual(len(events), 5)
 
@@ -1048,7 +1049,7 @@ def test_stream_wrapper_raises_on_transform_error(self):
             b'data: {"message": null}\n',
         ]
         with self.assertRaises(RuntimeError) as ctx:
-            list(transform_oci_stream_wrapper(iter(chunks), "chat"))
+            list(transform_oci_stream_wrapper(iter(chunks), "chat", is_v2=True))
         self.assertIn("OCI stream event transformation failed", str(ctx.exception))
 
     def test_stream_event_finish_reason_keeps_final_text(self):
@@ -1061,6 +1062,7 @@ def test_stream_event_finish_reason_keeps_final_text(self):
                 "message": {"content": [{"type": "TEXT", "text": " world"}]},
                 "finishReason": "COMPLETE",
             },
+            is_v2=True,
         )
 
         self.assertEqual(events[0]["type"], "content-delta")