fix(oci): close remaining review gaps

Author: fede-kamel

src/cohere/oci_client.py

@@ -771,10 +771,34 @@ def transform_request_to_oci(
                 chat_request["temperature"] = cohere_body["temperature"]
             if "max_tokens" in cohere_body:
                 chat_request["maxTokens"] = cohere_body["max_tokens"]
+            if "k" in cohere_body:
+                chat_request["topK"] = cohere_body["k"]
+            if "p" in cohere_body:
+                chat_request["topP"] = cohere_body["p"]
+            if "seed" in cohere_body:
+                chat_request["seed"] = cohere_body["seed"]
+            if "stop_sequences" in cohere_body:
+                chat_request["stopSequences"] = cohere_body["stop_sequences"]
+            if "frequency_penalty" in cohere_body:
+                chat_request["frequencyPenalty"] = cohere_body["frequency_penalty"]
+            if "presence_penalty" in cohere_body:
+                chat_request["presencePenalty"] = cohere_body["presence_penalty"]
             if "preamble" in cohere_body:
                 chat_request["preambleOverride"] = cohere_body["preamble"]
             if "chat_history" in cohere_body:
                 chat_request["chatHistory"] = cohere_body["chat_history"]
+            if "documents" in cohere_body:
+                chat_request["documents"] = cohere_body["documents"]
+            if "tools" in cohere_body:
+                chat_request["tools"] = cohere_body["tools"]
+            if "tool_results" in cohere_body:
+                chat_request["toolResults"] = cohere_body["tool_results"]
+            if "response_format" in cohere_body:
+                chat_request["responseFormat"] = cohere_body["response_format"]
+            if "safety_mode" in cohere_body:
+                chat_request["safetyMode"] = cohere_body["safety_mode"]
+            if "priority" in cohere_body:
+                chat_request["priority"] = cohere_body["priority"]
 
         # Handle streaming for both versions
         if "stream" in endpoint or cohere_body.get("stream"):
@@ -988,6 +1012,8 @@ def transform_oci_stream_wrapper(
     generation_id = str(uuid.uuid4())
     emitted_start = False
     emitted_content_end = False
+    current_content_type: typing.Optional[str] = None
+    current_content_index = 0
     final_finish_reason = "COMPLETE"
     final_usage: typing.Optional[typing.Dict[str, typing.Any]] = None
     full_v1_text = ""
@@ -1000,18 +1026,23 @@ def _emit_v2_event(event: typing.Dict[str, typing.Any]) -> bytes:
     def _emit_v1_event(event: typing.Dict[str, typing.Any]) -> bytes:
         return json.dumps(event).encode("utf-8") + b"\n"
 
+    def _current_v2_content_type(oci_event: typing.Dict[str, typing.Any]) -> str:
+        message = oci_event.get("message")
+        if isinstance(message, dict):
+            content_list = message.get("content")
+            if content_list and isinstance(content_list, list) and len(content_list) > 0:
+                oci_type = content_list[0].get("type", "TEXT").upper()
+                if oci_type == "THINKING":
+                    return "thinking"
+        return "text"
+
     def _transform_v2_event(oci_event: typing.Dict[str, typing.Any]) -> typing.Iterator[bytes]:
-        nonlocal emitted_start, emitted_content_end, final_finish_reason, final_usage
+        nonlocal emitted_start, emitted_content_end, current_content_type, current_content_index
+        nonlocal final_finish_reason, final_usage
+
+        event_content_type = _current_v2_content_type(oci_event)
 
         if not emitted_start:
-            content_type = "text"
-            message = oci_event.get("message")
-            if isinstance(message, dict):
-                content_list = message.get("content")
-                if content_list and isinstance(content_list, list) and len(content_list) > 0:
-                    oci_type = content_list[0].get("type", "TEXT").upper()
-                    if oci_type == "THINKING":
-                        content_type = "thinking"
 
             yield _emit_v2_event(
                 {
@@ -1023,15 +1054,30 @@ def _transform_v2_event(oci_event: typing.Dict[str, typing.Any]) -> typing.Itera
             yield _emit_v2_event(
                 {
                     "type": "content-start",
-                    "index": 0,
-                    "delta": {"message": {"content": {"type": content_type}}},
+                    "index": current_content_index,
+                    "delta": {"message": {"content": {"type": event_content_type}}},
                 }
             )
             emitted_start = True
+            current_content_type = event_content_type
+        elif current_content_type != event_content_type:
+            yield _emit_v2_event({"type": "content-end", "index": current_content_index})
+            current_content_index += 1
+            yield _emit_v2_event(
+                {
+                    "type": "content-start",
+                    "index": current_content_index,
+                    "delta": {"message": {"content": {"type": event_content_type}}},
+                }
+            )
+            current_content_type = event_content_type
+            emitted_content_end = False
 
         for cohere_event in typing.cast(
             typing.List[typing.Dict[str, typing.Any]], transform_stream_event(endpoint, oci_event, is_v2=True)
         ):
+            if "index" in cohere_event:
+                cohere_event = {**cohere_event, "index": current_content_index}
             if cohere_event["type"] == "content-end":
                 emitted_content_end = True
                 final_finish_reason = oci_event.get("finishReason", final_finish_reason)
@@ -1069,6 +1115,7 @@ def _process_line(line: str) -> typing.Iterator[bytes]:
                 yield _emit_v1_event(
                     {
                         "event_type": "stream-end",
+                        "finish_reason": final_v1_finish_reason,
                         "response": {
                             "text": full_v1_text,
                             "generation_id": generation_id,

tests/test_oci_client.py

@@ -815,6 +815,46 @@ def test_transform_chat_request_optional_params(self):
         self.assertEqual(chat_req["toolChoice"], "REQUIRED")
         self.assertEqual(chat_req["priority"], 7)
 
+    def test_transform_v1_chat_request_optional_params(self):
+        """Test V1 chat request forwards the supported optional params."""
+        from cohere.oci_client import transform_request_to_oci
+
+        body = {
+            "model": "command-r-08-2024",
+            "message": "Hi",
+            "max_tokens": 100,
+            "temperature": 0.7,
+            "k": 10,
+            "p": 0.8,
+            "seed": 123,
+            "stop_sequences": ["END"],
+            "frequency_penalty": 0.5,
+            "presence_penalty": 0.2,
+            "documents": [{"title": "Doc", "text": "Body"}],
+            "tools": [{"name": "lookup"}],
+            "tool_results": [{"call": {"name": "lookup"}}],
+            "response_format": {"type": "json_object"},
+            "safety_mode": "NONE",
+            "priority": 4,
+        }
+        result = transform_request_to_oci("chat", body, "compartment-123")
+
+        chat_req = result["chatRequest"]
+        self.assertEqual(chat_req["maxTokens"], 100)
+        self.assertEqual(chat_req["temperature"], 0.7)
+        self.assertEqual(chat_req["topK"], 10)
+        self.assertEqual(chat_req["topP"], 0.8)
+        self.assertEqual(chat_req["seed"], 123)
+        self.assertEqual(chat_req["stopSequences"], ["END"])
+        self.assertEqual(chat_req["frequencyPenalty"], 0.5)
+        self.assertEqual(chat_req["presencePenalty"], 0.2)
+        self.assertEqual(chat_req["documents"], [{"title": "Doc", "text": "Body"}])
+        self.assertEqual(chat_req["tools"], [{"name": "lookup"}])
+        self.assertEqual(chat_req["toolResults"], [{"call": {"name": "lookup"}}])
+        self.assertEqual(chat_req["responseFormat"], {"type": "json_object"})
+        self.assertEqual(chat_req["safetyMode"], "NONE")
+        self.assertEqual(chat_req["priority"], 4)
+
     def test_transform_chat_request_tool_message_fields(self):
         """Test tool message fields are converted to OCI names."""
         from cohere.oci_client import transform_request_to_oci
@@ -1042,6 +1082,34 @@ def test_stream_wrapper_emits_full_event_lifecycle(self):
         self.assertEqual(events[1]["delta"]["message"]["content"]["type"], "text")
         self.assertEqual(events[5]["delta"]["finish_reason"], "COMPLETE")
 
+    def test_stream_wrapper_emits_new_content_block_on_thinking_transition(self):
+        """Test V2 streams emit a new content block when transitioning from thinking to text."""
+        import json
+        from cohere.oci_client import transform_oci_stream_wrapper
+
+        chunks = [
+            b'data: {"message": {"content": [{"type": "THINKING", "thinking": "Reasoning..."}]}}\n',
+            b'data: {"message": {"content": [{"type": "TEXT", "text": "Answer"}]}, "finishReason": "COMPLETE"}\n',
+            b"data: [DONE]\n",
+        ]
+
+        events = []
+        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:]))
+
+        self.assertEqual(events[1]["type"], "content-start")
+        self.assertEqual(events[1]["delta"]["message"]["content"]["type"], "thinking")
+        self.assertEqual(events[2]["type"], "content-delta")
+        self.assertEqual(events[2]["index"], 0)
+        self.assertEqual(events[3], {"type": "content-end", "index": 0})
+        self.assertEqual(events[4]["type"], "content-start")
+        self.assertEqual(events[4]["index"], 1)
+        self.assertEqual(events[4]["delta"]["message"]["content"]["type"], "text")
+        self.assertEqual(events[5]["type"], "content-delta")
+        self.assertEqual(events[5]["index"], 1)
+
     def test_stream_wrapper_skips_malformed_json_with_warning(self):
         """Test that malformed JSON in SSE stream is skipped (not silently swallowed)."""
         from cohere.oci_client import transform_oci_stream_wrapper
@@ -1080,6 +1148,7 @@ def test_v1_stream_wrapper_preserves_finish_reason_in_stream_end(self):
         ]
 
         self.assertEqual(events[2]["event_type"], "stream-end")
+        self.assertEqual(events[2]["finish_reason"], "MAX_TOKENS")
         self.assertEqual(events[2]["response"]["text"], "Hello world")
         self.assertEqual(events[2]["response"]["finish_reason"], "MAX_TOKENS")