From 61b99a0475f3992eb202acb88626bd886625a026 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 21 May 2026 06:40:53 +0000 Subject: [PATCH] fix(oci): harden Cohere stream/finish-reason and dedupe maxTokens param mapping - Cohere stream: track per-stream tool-call emission and only suppress the terminal consolidation chunk's tool calls once they've been seen earlier. Prevents silent drop if tool calls are delivered exclusively on the terminal chunk. - Cohere stream: emit content=None (not "") on non-terminal text-free chunks (e.g. tool-call-only / keep-alive) so downstream consumers that distinguish missing vs explicitly-empty deltas behave correctly. - Generic handlers: accept singular TOOL_CALL finish reason in addition to TOOL_CALLS, matching the Cohere handlers. - _get_optional_params: when both max_tokens and max_completion_tokens are provided, explicitly prefer max_completion_tokens instead of relying on dict iteration order. Co-authored-by: Yassin Kortam --- litellm/llms/oci/chat/cohere.py | 40 ++++++++++++++++++------- litellm/llms/oci/chat/generic.py | 4 +-- litellm/llms/oci/chat/transformation.py | 29 +++++++++++++++++- 3 files changed, 59 insertions(+), 14 deletions(-) diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 7e380cf0e7f..c95a44236a3 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -279,8 +279,18 @@ def handle_cohere_response( return model_response -def handle_cohere_stream_chunk(dict_chunk: dict) -> ModelResponseStream: - """Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream.""" +def handle_cohere_stream_chunk( + dict_chunk: dict, prior_tool_calls_emitted: bool = False +) -> ModelResponseStream: + """Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream. + + ``prior_tool_calls_emitted`` lets the caller signal whether tool calls + were already emitted in earlier chunks of the same stream. When set, the + terminal consolidation chunk's tool calls are suppressed (they would + duplicate prior deltas); otherwise they are passed through so a stream + that delivers tool calls only on the terminal chunk doesn't silently + drop them. + """ try: typed_chunk = CohereStreamChunk(**dict_chunk) except (TypeError, ValidationError) as e: @@ -301,17 +311,25 @@ def handle_cohere_stream_chunk(dict_chunk: dict) -> ModelResponseStream: is_terminal_consolidation = ( typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None ) - text: Optional[str] = ( - None if is_terminal_consolidation else (typed_chunk.text or "") - ) + # On non-terminal text-free chunks (e.g. tool-call-only or keep-alive + # chunks) emit ``content=None`` rather than ``content=""`` so downstream + # stream-mergers that distinguish "no text in this delta" from "an + # explicitly empty text delta" behave correctly. + text: Optional[str] = None if is_terminal_consolidation else typed_chunk.text # Tool calls on the terminal consolidation chunk (whether from - # `typed_chunk.toolCalls` or from `chatHistory`) restate what was already - # streamed in intermediate chunks. Re-emitting them here would mint fresh - # `uuid4` IDs and cause downstream consumers to execute each tool call - # twice. Suppress them on the terminal chunk for the same reason `text` - # is suppressed above. - cohere_tool_calls = None if is_terminal_consolidation else typed_chunk.toolCalls + # `typed_chunk.toolCalls` or from `chatHistory`) typically restate what + # was already streamed in intermediate chunks. Re-emitting them would + # mint fresh `uuid4` IDs and cause downstream consumers to execute each + # tool call twice. We only suppress when the caller has confirmed that + # tool calls were already emitted earlier — otherwise (e.g. a short + # response that delivers tool calls exclusively on the terminal chunk), + # passing them through is the only chance to surface them. + cohere_tool_calls = ( + None + if (is_terminal_consolidation and prior_tool_calls_emitted) + else typed_chunk.toolCalls + ) tool_calls: Optional[List[Dict[str, Any]]] = None if cohere_tool_calls: diff --git a/litellm/llms/oci/chat/generic.py b/litellm/llms/oci/chat/generic.py index 92b050004dc..bc9e30c1cd5 100644 --- a/litellm/llms/oci/chat/generic.py +++ b/litellm/llms/oci/chat/generic.py @@ -347,7 +347,7 @@ def handle_generic_response( model_response.choices[0].finish_reason = "stop" # type: ignore[union-attr] elif oci_finish_reason == "MAX_TOKENS": model_response.choices[0].finish_reason = "length" # type: ignore[union-attr] - elif oci_finish_reason == "TOOL_CALLS": + elif oci_finish_reason in ("TOOL_CALL", "TOOL_CALLS"): model_response.choices[0].finish_reason = "tool_calls" # type: ignore[union-attr] elif oci_finish_reason is not None: # OCI GENERIC can emit non-OpenAI finish reasons (e.g. ``ERROR``, @@ -418,7 +418,7 @@ def handle_generic_stream_chunk(dict_chunk: dict) -> ModelResponseStream: finish_reason: Optional[str] = "stop" elif oci_finish_reason == "MAX_TOKENS": finish_reason = "length" - elif oci_finish_reason == "TOOL_CALLS": + elif oci_finish_reason in ("TOOL_CALL", "TOOL_CALLS"): finish_reason = "tool_calls" elif oci_finish_reason is not None: # OCI can emit error/cancel finish reasons (e.g. ``ERROR``, diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 2b928220fb6..28ff708cc10 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -355,8 +355,21 @@ class OCIChatConfig(BaseConfig): else "maxTokens" ) + # Both ``max_tokens`` and ``max_completion_tokens`` map to OCI's + # ``maxTokens`` (or ``maxCompletionTokens`` for reasoning models), so + # if both are provided explicitly prefer ``max_completion_tokens`` + # rather than relying on dict iteration order to pick a winner. + prefer_max_completion = ( + "max_tokens" in optional_params + and "max_completion_tokens" in optional_params + and param_map.get("max_tokens") == "maxTokens" + and param_map.get("max_completion_tokens") == "maxTokens" + ) + for openai_key, oci_key in param_map.items(): if oci_key and openai_key in optional_params: + if prefer_max_completion and openai_key == "max_tokens": + continue target = max_tokens_key if oci_key == "maxTokens" else oci_key selected_params[target] = optional_params[openai_key] # type: ignore[index] @@ -654,6 +667,11 @@ class OCIStreamWrapper(CustomStreamWrapper): def __init__(self, **kwargs: Any): super().__init__(**kwargs) + # Tracks whether any prior Cohere chunk in this stream has emitted + # tool calls. The Cohere handler uses this to decide whether the + # terminal consolidation chunk's tool calls are duplicates (suppress) + # or the only copy of the tool calls (pass through). + self._cohere_tool_calls_emitted = False def chunk_creator(self, chunk: Any) -> ModelResponseStream: if not isinstance(chunk, str): @@ -663,7 +681,16 @@ class OCIStreamWrapper(CustomStreamWrapper): dict_chunk = json.loads(chunk[5:]) if dict_chunk.get("apiFormat") == "COHERE": - return handle_cohere_stream_chunk(dict_chunk) + result = handle_cohere_stream_chunk( + dict_chunk, + prior_tool_calls_emitted=self._cohere_tool_calls_emitted, + ) + if not self._cohere_tool_calls_emitted: + for choice in result.choices: + if getattr(choice.delta, "tool_calls", None): + self._cohere_tool_calls_emitted = True + break + return result return handle_generic_stream_chunk(dict_chunk)