From 89d82ad7cb0bf5ec6ee85c839f68ac326c2d9874 Mon Sep 17 00:00:00 2001 From: Federico Kamelhar Date: Thu, 18 Jun 2026 15:43:02 -0400 Subject: [PATCH] fix(oci): make Cohere {{trace}} judges work (tool param types + agentic tool-calling continuation) Re-lands #30646, which was rolled back in a batch stabilization revert (commit 15a7c687), with the greptile-review empty-message guard folded in. Three OCI Cohere transform bugs broke MLflow {{trace}} judges: 1. Tool parameter types: a JSON-schema `array` param mapped to the bare string "List" in OCI_JSON_TO_PYTHON_TYPES, and OCI's Cohere backend returns HTTP 500 on a bare "List". Mapped the collection types to the lowercase builtins list/dict. 2. Agentic tool-calling continuation (request side): after a tool runs, litellm pulled the last user message into the top-level `message` and emitted the tool result as a TOOL entry in chatHistory. OCI rejects that ("cannot specify message if the last entry in chat history contains tool results"), and an empty message alone is also rejected ("message must be at least 1 token long or tool results must be specified"). OCI carries the current turn's results in a dedicated top-level `toolResults` field. The transform now sends an empty message, keeps the user turn in chatHistory, and puts results in `toolResults`, matching the langchain-oracle reference. Tool results are no longer chatHistory entries. 3. Response parsing: tool-grounded answers carry citations with `documentIds` (camelCase) and no `document_ids`, which made the required CohereCitation.document_ids field sink the whole response parse. That field (and CohereSearchQuery.generation_id) is now optional. Verified live against us-chicago-1 (cohere.command-a-03-2025 and command-latest), single and multi-round tool loops. Adds unit regressions on the transformed request shape and citation parsing, plus gated integration tests for the continuation. --- litellm/llms/oci/chat/cohere.py | 106 ++++++------- litellm/llms/oci/chat/transformation.py | 20 ++- litellm/llms/oci/common_utils.py | 8 +- litellm/types/llms/oci.py | 31 ++-- tests/integration/test_oci_integration.py | 103 +++++++++++-- .../oci/chat/test_oci_chat_transformation.py | 131 ++++++++-------- .../oci/chat/test_oci_cohere_tool_calls.py | 141 ++++++------------ .../unit/llms/oci/test_oci_coverage_boost.py | 139 ++++++++--------- 8 files changed, 355 insertions(+), 324 deletions(-) diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 9163ddfa5de..b0048ac3583 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -33,7 +33,6 @@ from litellm.types.llms.oci import ( CohereStreamChunk, CohereTool, CohereToolCall, - CohereToolMessage, CohereToolResult, ) from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantToolCall @@ -106,70 +105,73 @@ def _to_cohere_tool_call(tool_call: ChatCompletionAssistantToolCall) -> CohereTo def adapt_messages_to_cohere_standard( messages: list[AllMessageValues], ) -> list[CohereMessage]: - """Build a Cohere ``chatHistory`` list from an OpenAI-format message array. + """Build a Cohere ``chatHistory`` (USER and CHATBOT turns) from an OpenAI array. - - All messages except the *last user message* are included. The caller pulls - the last user message into the request's top-level ``message`` field, so - trailing tool results (the standard agentic continuation pattern) still - appear in ``chatHistory`` and reach the model. - - If no user message exists, every message is included (no slice). - - System messages must be filtered out by the caller (they are routed into - ``preambleOverride`` separately) — they are not represented in - ``chatHistory``. - - Tool results are expressed as OCI ``CohereToolMessage.toolResults`` entries, - with the originating call's name and parameters resolved from the preceding - assistant message via a ``tool_call_id`` lookup. + The final message is omitted: the caller sends it as the top-level ``message`` + (a normal user turn) or via ``toolResults`` (a tool-result continuation). + Tool-result messages are never represented in ``chatHistory``: OCI carries the + current turn's results in the separate top-level ``toolResults`` field and + rejects a request whose last history entry is a tool result. System messages + must be filtered out by the caller (they are routed into ``preambleOverride``). + """ + return [entry for msg in messages[:-1] if (entry := _cohere_history_entry(msg)) is not None] + + +def cohere_message_text(msg: AllMessageValues) -> str: + """Plain-text content of a message.""" + return extract_text_content(msg.get("content")) + + +def _assistant_tool_calls(msg: AllMessageValues) -> Sequence[ChatCompletionAssistantToolCall]: + """Tool calls carried by an assistant message; empty for other roles or when absent.""" + if msg.get("role") != "assistant" or "tool_calls" not in msg: + return [] + return list(msg["tool_calls"] or []) + + +def _cohere_history_entry(msg: AllMessageValues) -> CohereMessage | None: + """One ``chatHistory`` entry for a message, ``None`` for non-history roles.""" + role: Final = msg.get("role") + if role == "user": + return CohereMessage(role="USER", message=cohere_message_text(msg)) + if role != "assistant": + return None + tool_calls: Final = [_to_cohere_tool_call(tool_call) for tool_call in _assistant_tool_calls(msg)] + return CohereMessage(role="CHATBOT", message=cohere_message_text(msg), toolCalls=tool_calls or None) + + +def extract_cohere_tool_results( + messages: Sequence[AllMessageValues], +) -> list[CohereToolResult] | None: + """Return the current turn's tool results for OCI's top-level ``toolResults``. + + The current turn spans from the last user message to the end. Each tool + message is matched to its originating call by ``tool_call_id`` so OCI sees the + call name and parameters alongside the output. Returns ``None`` when there are + no tool results so the field is omitted. """ - # First pass: build tool_call_id → CohereToolCall so tool-result messages can - # reference the originating call by name and parameters. tool_call_lookup: Final = { tool_call.get("id", ""): _to_cohere_tool_call(tool_call) for msg in messages - if msg.get("role") == "assistant" and "tool_calls" in msg - for tool_call in msg["tool_calls"] or [] + for tool_call in _assistant_tool_calls(msg) } last_user_index: Final = next( (i for i in range(len(messages) - 1, -1, -1) if messages[i].get("role") == "user"), None, ) - history_source: Final = ( - messages if last_user_index is None else [m for i, m in enumerate(messages) if i != last_user_index] - ) + current_turn: Final = messages if last_user_index is None else messages[last_user_index:] - chat_history: Final[list[CohereMessage]] = [] - for msg in history_source: - role = msg.get("role") - content = extract_text_content(msg.get("content")) - - tool_calls = ( - [_to_cohere_tool_call(tool_call) for tool_call in msg["tool_calls"]] - if role == "assistant" and "tool_calls" in msg and msg["tool_calls"] - else None + unknown_call: Final = CohereToolCall(name="", parameters={}) + tool_results: Final = [ + CohereToolResult( + call=tool_call_lookup.get(str(msg.get("tool_call_id", "") or ""), unknown_call), + outputs=[{"output": cohere_message_text(msg)}], ) - - if role == "user": - chat_history.append(CohereMessage(role="USER", message=content)) - elif role == "assistant": - chat_history.append(CohereMessage(role="CHATBOT", message=content, toolCalls=tool_calls)) - elif role == "tool": - tool_call_id = str(msg.get("tool_call_id", "") or "") - cohere_call = tool_call_lookup.get(tool_call_id, CohereToolCall(name="", parameters={})) - tool_result = CohereToolResult( - call=cohere_call, - outputs=[{"output": content}], - ) - # OpenAI emits one tool-role message per parallel tool call, but - # the OCI Cohere API expects all results from a single assistant - # turn to share one TOOL history entry with multiple toolResults. - # Merge consecutive tool messages so the model sees the parallel - # call/result pairing correctly during agentic loops. - if chat_history and isinstance(chat_history[-1], CohereToolMessage): - chat_history[-1].toolResults.append(tool_result) - else: - chat_history.append(CohereToolMessage(toolResults=[tool_result])) - - return chat_history + for msg in current_turn + if msg.get("role") == "tool" + ] + return tool_results or None def _resolved_oci_parameter_schema(raw_parameters: dict[str, JsonValue]) -> JsonValue: diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index f97f8f4ca59..b2053845566 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -40,6 +40,8 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.llms.oci.chat.cohere import ( adapt_messages_to_cohere_standard, adapt_tool_definitions_to_cohere_standard, + cohere_message_text, + extract_cohere_tool_results, extract_text_content, handle_cohere_response, handle_cohere_stream_chunk, @@ -684,10 +686,24 @@ class OCIChatConfig(BaseConfig): if preamble: preamble_override = preamble + non_system_messages: Final = [m for m in messages if m.get("role") != "system"] + tool_results: Final = extract_cohere_tool_results(non_system_messages) + if not tool_results and not non_system_messages: + raise OCIError( + status_code=400, + message="No non-system messages found — Cohere models require at least one non-system message", + ) + # OCI carries the current turn's tool results in the top-level + # ``toolResults`` field and rejects a populated ``message`` alongside + # them. After tool execution (the agentic continuation MLflow judges + # drive) there is no new user turn, so the message is empty. + message: Final = "" if tool_results else cohere_message_text(non_system_messages[-1]) + chat_request: Final = CohereChatRequest( apiFormat="COHERE", - message=extract_text_content(user_messages[-1]["content"]), - chatHistory=adapt_messages_to_cohere_standard([m for m in messages if m.get("role") != "system"]), + message=message, + chatHistory=adapt_messages_to_cohere_standard(non_system_messages), + toolResults=tool_results, preambleOverride=preamble_override, **self._get_optional_params(OCIVendors.COHERE, optional_params, model), ) diff --git a/litellm/llms/oci/common_utils.py b/litellm/llms/oci/common_utils.py index fe8d0e1c1e7..36415dcd102 100644 --- a/litellm/llms/oci/common_utils.py +++ b/litellm/llms/oci/common_utils.py @@ -546,14 +546,16 @@ def validate_oci_environment( # --------------------------------------------------------------------------- # Mapping from JSON Schema type names to Python type names, as expected by -# the OCI Cohere API's CohereParameterDefinition.type field. +# the OCI Cohere API's CohereParameterDefinition.type field. The collection +# types must be the lowercase builtins: OCI's Cohere backend returns HTTP 500 +# on a bare ``List`` (it expects ``list`` / ``List[...]``). OCI_JSON_TO_PYTHON_TYPES: Final[dict[str, str]] = { "string": "str", "number": "float", "boolean": "bool", "integer": "int", - "array": "List", - "object": "Dict", + "array": "list", + "object": "dict", "any": "any", } diff --git a/litellm/types/llms/oci.py b/litellm/types/llms/oci.py index bd3a0af1dca..7305b14bb2b 100644 --- a/litellm/types/llms/oci.py +++ b/litellm/types/llms/oci.py @@ -3,8 +3,6 @@ from __future__ import annotations from enum import Enum from typing import Literal -from pydantic import SerializeAsAny - from litellm.types.llms.base import LiteLLMBaseModel OCIRoles = Literal["SYSTEM", "USER", "ASSISTANT", "TOOL"] @@ -246,17 +244,6 @@ class CohereSystemMessage(CohereMessage): role: Literal["SYSTEM"] = "SYSTEM" -class CohereToolMessage(CohereMessage): - """Tool message in Cohere chat. - - The OCI Cohere API represents tool results via a ``toolResults`` list on the - TOOL-role history entry — not via a ``toolCallId`` string. - """ - - role: Literal["TOOL"] = "TOOL" - toolResults: list[CohereToolResult] - - class CohereParameterDefinition(LiteLLMBaseModel): """Parameter definition for Cohere tools.""" @@ -299,11 +286,7 @@ class CohereChatRequest(LiteLLMBaseModel): apiFormat: Literal["COHERE"] = "COHERE" # Optional fields - # ``SerializeAsAny`` preserves subclass-specific fields (e.g. ``toolResults`` - # on ``CohereToolMessage``) when this request is serialized via ``model_dump``. - # Without it, Pydantic v2 would serialize each element using the declared - # ``CohereMessage`` schema and silently drop subclass fields. - chatHistory: list[SerializeAsAny[CohereMessage]] | None = None + chatHistory: list[CohereMessage] | None = None maxTokens: int | None = None temperature: float | None = None topP: float | None = None @@ -313,6 +296,10 @@ class CohereChatRequest(LiteLLMBaseModel): stopSequences: list[str] | None = None seed: int | None = None tools: list[CohereTool] | None = None + # Current-turn tool results for an agentic continuation. OCI carries these in + # this top-level field (not as a chatHistory entry) and requires an empty + # ``message`` when it is set. + toolResults: list[CohereToolResult] | None = None # NOTE: OCI's Cohere chat endpoint does not accept ``toolChoice`` — see # ``OCIChatConfig.openai_to_oci_cohere_param_map`` which marks # ``tool_choice`` as unsupported. The field is intentionally absent here @@ -354,14 +341,18 @@ class CohereCitation(LiteLLMBaseModel): start: int end: int text: str - document_ids: list[str] + # OCI returns ``documentIds`` (camelCase) and omits it entirely for + # tool-grounded citations (which reference tool calls, not documents). These + # citations are never surfaced by the response handler, so keep the field + # optional rather than failing the whole response parse. + document_ids: list[str] | None = None class CohereSearchQuery(LiteLLMBaseModel): """Search query generated by Cohere.""" text: str - generation_id: str + generation_id: str | None = None class CohereChatResponse(LiteLLMBaseModel): diff --git a/tests/integration/test_oci_integration.py b/tests/integration/test_oci_integration.py index 686facc6e30..928c8c6f903 100644 --- a/tests/integration/test_oci_integration.py +++ b/tests/integration/test_oci_integration.py @@ -20,7 +20,7 @@ Run only these tests: import math import os -from typing import NamedTuple, Optional +from typing import NamedTuple import pytest @@ -59,9 +59,7 @@ def oci_signer(): if "security_token_file" in config: with open(os.path.expanduser(config["security_token_file"])) as f: token = f.read().strip() - private_key = oci.signer.load_private_key_from_file( - config["key_file"], config.get("pass_phrase") - ) + private_key = oci.signer.load_private_key_from_file(config["key_file"], config.get("pass_phrase")) return oci.auth.signers.SecurityTokenSigner(token, private_key) return oci.Signer( @@ -103,7 +101,7 @@ class _M(NamedTuple): # reasoning budget is exhausted before the answer token budget starts. reasoning: bool = False # tool_choice value to send; None means omit the parameter entirely. - tool_choice: Optional[str] = "auto" + tool_choice: str | None = "auto" # Whether to include the model in tool-use parametrize list. supports_tool_use: bool = True @@ -385,6 +383,95 @@ async def test_async_tool_use(m: _M, oci_params): _assert_tool_call(resp) +# --------------------------------------------------------------------------- +# Cohere tool-schema regression +# +# MLflow {{trace}} judge tools have params typed as JSON-schema arrays. The +# array type used to map to a bare ``List``, which OCI's Cohere backend rejected +# with HTTP 500. +# --------------------------------------------------------------------------- + +# A tool whose parameter is a JSON-schema array — the shape that 500'd on Cohere. +_TRACE_TOOL = { + "type": "function", + "function": { + "name": "get_span", + "description": "Fetch a span and selected attributes.", + "parameters": { + "type": "object", + "properties": { + "attributes_to_fetch": { + "type": "array", + "items": {"type": "string"}, + "description": "Attribute names to fetch.", + }, + }, + "required": [], + }, + }, +} + + +@pytest.mark.parametrize("model", ["cohere.command-latest", "cohere.command-a-03-2025"]) +def test_cohere_tool_with_array_param(model, oci_params): + """A Cohere tool whose parameter is a JSON-schema array (e.g. MLflow's + {{trace}} judge tools). The array type used to map to a bare ``List``, which + OCI's Cohere backend rejected with HTTP 500.""" + import litellm + + resp = litellm.completion( + model=f"oci/{model}", + messages=[ + { + "role": "user", + "content": "Use get_span to fetch the 'inputs' attribute.", + } + ], + tools=[_TRACE_TOOL], + max_tokens=150, + **oci_params, + ) + # The model may call the tool or answer directly; either proves OCI accepted + # the request (no 400/500). + assert resp.choices[0].finish_reason is not None + + +@pytest.mark.parametrize("model", ["cohere.command-latest", "cohere.command-a-03-2025"]) +def test_cohere_tool_result_continuation(model, oci_params): + """The agentic continuation MLflow {{trace}} judges drive: the conversation + ends with a tool result. OCI rejects a populated `message` alongside tool + results and requires them in the top-level `toolResults` field, so this used + to 400/500. The model should now produce a final answer.""" + import litellm + + resp = litellm.completion( + model=f"oci/{model}", + messages=[ + {"role": "user", "content": "How many spans are in the trace?"}, + { + "role": "assistant", + "content": "I will look at the root span.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_span", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": '{"span_count": 47}', + }, + ], + tools=[_TRACE_TOOL], + max_tokens=150, + **oci_params, + ) + assert resp.choices[0].finish_reason is not None + + # --------------------------------------------------------------------------- # Reasoning-effort tests (reasoning models only) # --------------------------------------------------------------------------- @@ -448,7 +535,6 @@ def test_reasoning_tokens_in_usage(oci_params): class TestOCIEmbeddings: - def test_english_v3_basic(self, oci_params): import litellm @@ -530,9 +616,7 @@ class TestOCIEmbeddings: stock = resp.data[2]["embedding"] sim_cats = cosine(cat1, cat2) sim_diff = cosine(cat1, stock) - assert ( - sim_cats > sim_diff - ), f"Expected similar sentences to score higher ({sim_cats:.3f} vs {sim_diff:.3f})" + assert sim_cats > sim_diff, f"Expected similar sentences to score higher ({sim_cats:.3f} vs {sim_diff:.3f})" def test_embed_v4(self, oci_params): import litellm @@ -565,7 +649,6 @@ class TestOCIEmbeddings: class TestOCIAsyncEmbeddings: - @pytest.mark.asyncio async def test_async_embedding_basic(self, oci_params): import litellm diff --git a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py index 5093d01f153..bfd235620d8 100644 --- a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py @@ -178,10 +178,7 @@ class TestOCIChatConfig: headers={}, ) assert "tools" in transformed_request["chatRequest"] - assert ( - transformed_request["chatRequest"]["tools"][0]["name"] - == "get_current_weather" - ) + assert transformed_request["chatRequest"]["tools"][0]["name"] == "get_current_weather" assert transformed_request["chatRequest"]["tools"][0]["type"] == "FUNCTION" assert ( transformed_request["chatRequest"]["tools"][0]["description"] @@ -320,18 +317,11 @@ class TestOCIChatConfig: # Verify Cohere-specific request structure assert "message" in transformed_request["chatRequest"] # Cohere uses "message" - assert ( - "chatHistory" in transformed_request["chatRequest"] - ) # Cohere uses "chatHistory" - assert ( - "messages" not in transformed_request["chatRequest"] - ) # Generic uses "messages" + assert "chatHistory" in transformed_request["chatRequest"] # Cohere uses "chatHistory" + assert "messages" not in transformed_request["chatRequest"] # Generic uses "messages" # Verify the message content - assert ( - transformed_request["chatRequest"]["message"] - == "What is quantum computing?" - ) + assert transformed_request["chatRequest"]["message"] == "What is quantum computing?" def test_transform_request_response_format_json_object(self): """ @@ -527,11 +517,7 @@ class TestOCIChatConfig: are handled correctly (fields are optional). """ config = OCIChatConfig() - created_time = ( - datetime.datetime.now(datetime.timezone.utc) - .isoformat() - .replace("+00:00", "Z") - ) + created_time = datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z") mock_oci_response = { "modelId": TEST_MODEL_NAME, "modelVersion": "1.0", @@ -583,11 +569,7 @@ class TestOCIChatConfig: Tests if a simple text response is transformed correctly. """ config = OCIChatConfig() - created_time = ( - datetime.datetime.now(datetime.timezone.utc) - .isoformat() - .replace("+00:00", "Z") - ) + created_time = datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z") mock_oci_response = { "modelId": TEST_MODEL_NAME, "modelVersion": "1.0", @@ -598,9 +580,7 @@ class TestOCIChatConfig: "index": 0, "message": { "role": "ASSISTANT", - "content": [ - {"type": "TEXT", "text": "I am doing well, thank you!"} - ], + "content": [{"type": "TEXT", "text": "I am doing well, thank you!"}], }, "finishReason": "COMPLETE", } @@ -660,11 +640,7 @@ class TestOCIChatConfig: Tests if a response with tool calls is transformed correctly. """ config = OCIChatConfig() - created_time = ( - datetime.datetime.now(datetime.timezone.utc) - .isoformat() - .replace("+00:00", "Z") - ) + created_time = datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z") mock_oci_response = { "modelId": TEST_MODEL_NAME, "modelVersion": "1.0", @@ -704,9 +680,7 @@ class TestOCIChatConfig: }, } response = httpx.Response(status_code=200, json=mock_oci_response) - model_response = ModelResponse( - choices=[litellm.Choices(index=0, message=litellm.Message())] - ) + model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())]) result = config.transform_response( model=TEST_MODEL_NAME, @@ -998,9 +972,7 @@ class TestOCICohereParamMapping: "topP", "frequencyPenalty", ): - assert ( - injected not in result - ), f"'{injected}' should not be injected when user did not provide it" + assert injected not in result, f"'{injected}' should not be injected when user did not provide it" def test_cohere_explicit_params_still_passed(self): """User-provided Cohere params must still be forwarded correctly.""" @@ -1015,6 +987,58 @@ class TestOCICohereParamMapping: assert result.get("temperature") == 0.5 +class TestOCICohereToolParameterTypes: + """OCI's Cohere backend returns HTTP 500 on a bare ``List`` parameter type + (e.g. MLflow judge tools whose params are JSON-schema arrays). Collection + types must map to the lowercase builtins. Verified live against us-chicago-1. + """ + + def _parameter_definitions(self) -> dict: + config = OCIChatConfig() + body = config.transform_request( + model="cohere.command-latest", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + **BASE_OCI_PARAMS, + "tools": [ + { + "type": "function", + "function": { + "name": "get_span", + "description": "Fetch a span.", + "parameters": { + "type": "object", + "properties": { + "attributes": { + "type": "array", + "items": {"type": "string"}, + "description": "attribute names", + }, + "filters": { + "type": "object", + "description": "filter mapping", + }, + }, + "required": [], + }, + }, + } + ], + }, + litellm_params={}, + headers={}, + ) + return body["chatRequest"]["tools"][0]["parameterDefinitions"] + + def test_array_param_maps_to_lowercase_list(self): + params = self._parameter_definitions() + assert params["attributes"]["type"] == "list" + + def test_object_param_maps_to_lowercase_dict(self): + params = self._parameter_definitions() + assert params["filters"]["type"] == "dict" + + class TestOCIDefaultMaxTokens: """Regression for OCI's tiny server-side token cap (~20 tokens), which silently truncated responses mid-string whenever the caller omitted @@ -1033,16 +1057,12 @@ class TestOCIDefaultMaxTokens: ) return body["chatRequest"] - @pytest.mark.parametrize( - "model", ["cohere.command-latest", "meta.llama-3.3-70b-instruct"] - ) + @pytest.mark.parametrize("model", ["cohere.command-latest", "meta.llama-3.3-70b-instruct"]) def test_default_injected_when_max_tokens_omitted(self, model): chat_request = self._chat_request(model, {}) assert chat_request["maxTokens"] == DEFAULT_OCI_CHAT_MAX_TOKENS - @pytest.mark.parametrize( - "model", ["cohere.command-latest", "meta.llama-3.3-70b-instruct"] - ) + @pytest.mark.parametrize("model", ["cohere.command-latest", "meta.llama-3.3-70b-instruct"]) def test_explicit_max_tokens_not_overridden(self, model): chat_request = self._chat_request(model, {"max_tokens": 256}) assert chat_request["maxTokens"] == 256 @@ -1146,11 +1166,7 @@ class TestOCIReasoningEffort: Usage.completion_tokens_details.reasoning_tokens.""" from litellm.llms.oci.chat.generic import handle_generic_response - created_time = ( - datetime.datetime.now(datetime.timezone.utc) - .isoformat() - .replace("+00:00", "Z") - ) + created_time = datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z") oci_response = { "modelId": "xai.grok-4-fast-reasoning", "modelVersion": "1.0", @@ -1190,11 +1206,7 @@ class TestOCIReasoningEffort: """When OCI omits completionTokensDetails, Usage has no reasoning_tokens.""" from litellm.llms.oci.chat.generic import handle_generic_response - created_time = ( - datetime.datetime.now(datetime.timezone.utc) - .isoformat() - .replace("+00:00", "Z") - ) + created_time = datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z") oci_response = { "modelId": "xai.grok-4", "modelVersion": "1.0", @@ -1274,9 +1286,7 @@ class TestOCIStreamingSignedBody: signed_json_body=signed_bytes, ) - assert ( - posted_data["data"] == signed_bytes - ), "Streaming must use signed_json_body, not re-serialize data" + assert posted_data["data"] == signed_bytes, "Streaming must use signed_json_body, not re-serialize data" def test_get_custom_stream_wrapper_fallback_without_signed_body(self, monkeypatch): """When signed_json_body is None, fall back to json.dumps(data).""" @@ -1314,9 +1324,9 @@ class TestOCIStreamingSignedBody: signed_json_body=None, ) - assert posted_data["data"] == json.dumps( - payload - ), "Without signed_json_body, must fall back to json.dumps(data)" + assert posted_data["data"] == json.dumps(payload), ( + "Without signed_json_body, must fall back to json.dumps(data)" + ) # --------------------------------------------------------------------------- @@ -1416,6 +1426,7 @@ class TestOCIChatConfigErrorPaths: drop_params=False, ) assert "max_retries" not in result + def test_map_openai_params_cohere_n_default_dropped(self): """Cohere has no numGenerations field, but n=1 (and None) is the OpenAI default single-generation request. It must be dropped silently rather diff --git a/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py index 30b574e040e..cc98e20453f 100644 --- a/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -186,9 +186,7 @@ class TestOCICohereToolCalls: "apiFormat": "COHERE", "text": "I will look up the weather in Tokyo.", "finishReason": "COMPLETE", - "toolCalls": [ - {"name": "get_weather", "parameters": {"location": "Tokyo"}} - ], + "toolCalls": [{"name": "get_weather", "parameters": {"location": "Tokyo"}}], "usage": { "promptTokens": 26, "completionTokens": 22, @@ -218,9 +216,7 @@ class TestOCICohereToolCalls: # Verify response structure assert isinstance(result, ModelResponse) assert result.model == "cohere.command-latest" - assert ( - result.choices[0].message.content == "I will look up the weather in Tokyo." - ) + assert result.choices[0].message.content == "I will look up the weather in Tokyo." # Verify tool calls are present assert result.choices[0].message.tool_calls is not None @@ -296,10 +292,14 @@ class TestOCICohereToolCalls: assert transformed_request["chatRequest"]["responseFormat"] == {"type": "TEXT"} - def test_cohere_tool_call_only_message_no_text(self): - """Test chat history with an assistant message that has tool calls but no text content.""" + def test_cohere_tool_continuation_uses_top_level_tool_results(self): + """A tool-result continuation must send an empty top-level `message`, keep + the user turn in chatHistory, carry the tool result in the top-level + `toolResults` field, and never emit a TOOL chatHistory entry. OCI 500s + ("message must be ... or tool results must be specified") and 400s + ("cannot specify message if the last entry in chat history contains tool + results") otherwise. Verified live against us-chicago-1.""" config = OCIChatConfig() - messages = [ {"role": "user", "content": "What's the weather?"}, { @@ -316,81 +316,43 @@ class TestOCICohereToolCalls: } ], }, - { - "role": "tool", - "content": "Sunny, 25C", - "tool_call_id": "call_1", - }, + {"role": "tool", "content": "Sunny, 25C", "tool_call_id": "call_1"}, ] - chat_history = adapt_messages_to_cohere_standard(messages) + chat_request = config.transform_request( + model="cohere.command-latest", + messages=messages, # type: ignore[arg-type] + optional_params={"oci_compartment_id": TEST_COMPARTMENT_ID}, + litellm_params={}, + headers={}, + )["chatRequest"] - # The last user message is consumed by the request's top-level `message` - # field, so chatHistory carries the assistant tool call and tool result. - assert len(chat_history) == 2 + assert chat_request["message"] == "" + assert all(m["role"] != "TOOL" for m in chat_request["chatHistory"]) + assert [m["role"] for m in chat_request["chatHistory"]] == ["USER", "CHATBOT"] - assistant_msg = chat_history[0] - assert assistant_msg.role == "CHATBOT" - assert assistant_msg.message is None or assistant_msg.message == "" - assert assistant_msg.toolCalls is not None - assert len(assistant_msg.toolCalls) == 1 - assert assistant_msg.toolCalls[0].name == "get_weather" + tool_results = chat_request["toolResults"] + assert len(tool_results) == 1 + assert tool_results[0]["call"]["name"] == "get_weather" + assert tool_results[0]["call"]["parameters"] == {"location": "Paris"} + assert tool_results[0]["outputs"][0]["output"] == "Sunny, 25C" - tool_msg = chat_history[1] - assert tool_msg.role == "TOOL" - assert tool_msg.toolResults[0].call.name == "get_weather" - assert tool_msg.toolResults[0].outputs[0]["output"] == "Sunny, 25C" - - def test_cohere_chat_history_with_tool_calls(self): - """Tool results trailing the last user turn must be preserved in chatHistory.""" + def test_cohere_first_turn_uses_message_not_tool_results(self): + """Without a tool-result continuation the last user message is the top-level + `message` and no `toolResults` field is emitted.""" config = OCIChatConfig() + messages = [{"role": "user", "content": "What's the weather like in Tokyo?"}] - messages = [ - {"role": "user", "content": "What's the weather like in Tokyo?"}, - { - "role": "assistant", - "content": "I will look up the weather in Tokyo.", - "tool_calls": [ - { - "id": "call_0", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"location": "Tokyo"}', - }, - } - ], - }, - { - "role": "tool", - "content": "The weather in Tokyo is 22°C with partly cloudy skies.", - "tool_call_id": "call_0", - }, - ] + chat_request = config.transform_request( + model="cohere.command-latest", + messages=messages, # type: ignore[arg-type] + optional_params={"oci_compartment_id": TEST_COMPARTMENT_ID}, + litellm_params={}, + headers={}, + )["chatRequest"] - chat_history = adapt_messages_to_cohere_standard(messages) - - # The last user message becomes the request's top-level `message`. - # Everything else — including the trailing tool result — must remain in - # chatHistory so the model can see the tool output. - assert len(chat_history) == 2 - - assistant_msg = chat_history[0] - assert assistant_msg.role == "CHATBOT" - assert assistant_msg.message == "I will look up the weather in Tokyo." - assert assistant_msg.toolCalls is not None - assert len(assistant_msg.toolCalls) == 1 - assert assistant_msg.toolCalls[0].name == "get_weather" - assert assistant_msg.toolCalls[0].parameters == {"location": "Tokyo"} - - tool_msg = chat_history[1] - assert tool_msg.role == "TOOL" - assert tool_msg.toolResults[0].call.name == "get_weather" - assert tool_msg.toolResults[0].call.parameters == {"location": "Tokyo"} - assert ( - tool_msg.toolResults[0].outputs[0]["output"] - == "The weather in Tokyo is 22°C with partly cloudy skies." - ) + assert chat_request["message"] == "What's the weather like in Tokyo?" + assert "toolResults" not in chat_request def test_cohere_streaming_chunk_handling(self): """Test Cohere streaming chunk handling""" @@ -399,9 +361,7 @@ class TestOCICohereToolCalls: mock_model = "cohere.command-latest" mock_logging = MagicMock() - stream_wrapper = OCIStreamWrapper( - completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging - ) + stream_wrapper = OCIStreamWrapper(completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging) # Mock Cohere streaming chunk cohere_chunk = { @@ -425,9 +385,7 @@ class TestOCICohereToolCalls: mock_model = "cohere.command-latest" mock_logging = MagicMock() - stream_wrapper = OCIStreamWrapper( - completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging - ) + stream_wrapper = OCIStreamWrapper(completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging) # Mock Cohere finish chunk cohere_finish_chunk = { @@ -588,9 +546,7 @@ class TestOCICohereToolCalls: "apiFormat": "COHERE", "text": "", "finishReason": "TOOL_CALL", - "toolCalls": [ - {"name": "get_weather", "parameters": {"location": "London"}} - ], + "toolCalls": [{"name": "get_weather", "parameters": {"location": "London"}}], "usage": { "promptTokens": 20, "completionTokens": 10, @@ -672,9 +628,7 @@ class TestOCICohereToolCalls: assert get_vendor_from_model("cohere.command-latest") == OCIVendors.COHERE assert get_vendor_from_model("cohere.command-a-03-2025") == OCIVendors.COHERE assert get_vendor_from_model("cohere.command-plus-latest") == OCIVendors.COHERE - assert ( - get_vendor_from_model("cohere.command-r-plus-08-2024") == OCIVendors.COHERE - ) + assert get_vendor_from_model("cohere.command-r-plus-08-2024") == OCIVendors.COHERE assert get_vendor_from_model("cohere.command-r-08-2024") == OCIVendors.COHERE def test_cohere_error_handling_invalid_tool_format(self): @@ -787,10 +741,7 @@ class TestOCICoherePreambleOverride: ) chat_request = result["chatRequest"] - assert ( - chat_request["preambleOverride"] - == "You are a helpful assistant.\nAlways respond in JSON." - ) + assert chat_request["preambleOverride"] == "You are a helpful assistant.\nAlways respond in JSON." def test_system_message_with_content_array(self): """Test system message with list-style content (text blocks)""" @@ -919,8 +870,6 @@ class TestCohereMessageAdaptationEdgeCases: """Coverage for adapt_messages_to_cohere_standard error paths.""" def test_json_decode_error_in_tool_args_defaults_to_empty(self): - from litellm.llms.oci.chat.cohere import adapt_messages_to_cohere_standard - messages = [ { "role": "assistant", @@ -962,9 +911,7 @@ class TestOCICohereStreaming: mock_model = "cohere.command-latest" mock_logging = MagicMock() - return OCIStreamWrapper( - completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging - ) + return OCIStreamWrapper(completion_stream=mock_stream, model=mock_model, logging_obj=mock_logging) def test_cohere_streaming_chunk_parsing(self): diff --git a/tests/unit/llms/oci/test_oci_coverage_boost.py b/tests/unit/llms/oci/test_oci_coverage_boost.py index 251df5e4cad..5f356dbee19 100644 --- a/tests/unit/llms/oci/test_oci_coverage_boost.py +++ b/tests/unit/llms/oci/test_oci_coverage_boost.py @@ -70,17 +70,13 @@ _GENERIC_MODEL = "meta.llama-3-70b-instruct" @patch("litellm.llms.oci.common_utils.load_private_key_from_str") @patch("litellm.llms.oci.common_utils.padding") @patch("litellm.llms.oci.common_utils.hashes") -def test_sign_with_manual_credentials_inline_key( - mock_hashes, mock_padding, mock_load_key -): +def test_sign_with_manual_credentials_inline_key(mock_hashes, mock_padding, mock_load_key): """sign_with_manual_credentials succeeds with an inline oci_key string.""" mock_key = MagicMock() mock_key.sign.return_value = b"fake_signature" mock_load_key.return_value = mock_key - result_headers, body = sign_with_manual_credentials( - {}, _MANUAL_CREDS, {"key": "val"}, _API_BASE - ) + result_headers, body = sign_with_manual_credentials({}, _MANUAL_CREDS, {"key": "val"}, _API_BASE) assert "authorization" in result_headers assert result_headers["authorization"].startswith('Signature version="1"') @@ -93,9 +89,7 @@ def test_sign_with_manual_credentials_inline_key( @patch("litellm.llms.oci.common_utils.load_private_key_from_file") @patch("litellm.llms.oci.common_utils.padding") @patch("litellm.llms.oci.common_utils.hashes") -def test_sign_with_manual_credentials_key_file( - mock_hashes, mock_padding, mock_load_file -): +def test_sign_with_manual_credentials_key_file(mock_hashes, mock_padding, mock_load_file): """sign_with_manual_credentials falls back to oci_key_file when oci_key absent.""" mock_key = MagicMock() mock_key.sign.return_value = b"sig_from_file" @@ -104,9 +98,7 @@ def test_sign_with_manual_credentials_key_file( creds = {**_MANUAL_CREDS, "oci_key_file": "/tmp/key.pem"} creds_no_inline = {k: v for k, v in creds.items() if k != "oci_key"} - result_headers, body = sign_with_manual_credentials( - {}, creds_no_inline, {}, _API_BASE - ) + result_headers, body = sign_with_manual_credentials({}, creds_no_inline, {}, _API_BASE) assert "authorization" in result_headers mock_load_file.assert_called_once_with("/tmp/key.pem") @@ -116,9 +108,7 @@ def test_sign_with_manual_credentials_key_file( @patch("litellm.llms.oci.common_utils.load_private_key_from_str") @patch("litellm.llms.oci.common_utils.padding") @patch("litellm.llms.oci.common_utils.hashes") -def test_sign_with_manual_credentials_authorization_contains_key_id( - mock_hashes, mock_padding, mock_load_key -): +def test_sign_with_manual_credentials_authorization_contains_key_id(mock_hashes, mock_padding, mock_load_key): """Authorization header encodes tenancy/user/fingerprint as key ID.""" mock_key = MagicMock() mock_key.sign.return_value = b"sig" @@ -157,9 +147,7 @@ def test_sign_oci_request_routes_to_signer_when_present(): @patch("litellm.llms.oci.common_utils.load_private_key_from_str") @patch("litellm.llms.oci.common_utils.padding") @patch("litellm.llms.oci.common_utils.hashes") -def test_sign_oci_request_routes_to_manual_when_no_signer( - mock_hashes, mock_padding, mock_load_key -): +def test_sign_oci_request_routes_to_manual_when_no_signer(mock_hashes, mock_padding, mock_load_key): """sign_oci_request delegates to sign_with_manual_credentials when oci_signer absent.""" mock_key = MagicMock() mock_key.sign.return_value = b"sig" @@ -274,9 +262,7 @@ def test_adapt_generic_multipart_content(): def test_adapt_generic_tool_response_direct(): - result = adapt_messages_to_generic_oci_standard_tool_response( - "tool", "call_999", "The answer is 42" - ) + result = adapt_messages_to_generic_oci_standard_tool_response("tool", "call_999", "The answer is 42") assert result.role == "TOOL" assert result.toolCallId == "call_999" assert result.content[0].text == "The answer is 42" @@ -489,8 +475,11 @@ def test_adapt_cohere_assistant_with_tool_calls_in_history(): assert history[0].toolCalls[0].name == "calc" -def test_adapt_cohere_tool_result_in_history(): +def test_adapt_cohere_tool_result_goes_to_tool_results_not_history(): + from litellm.llms.oci.chat.cohere import extract_cohere_tool_results + messages = [ + {"role": "user", "content": "compute"}, { "role": "assistant", "content": None, @@ -502,17 +491,17 @@ def test_adapt_cohere_tool_result_in_history(): } ], }, - { - "role": "tool", - "tool_call_id": "call_1", - "content": "result: 42", - }, - {"role": "user", "content": "ok"}, + {"role": "tool", "tool_call_id": "call_1", "content": "result: 42"}, ] + # Tool results never appear in chatHistory; they ride the top-level field. history = adapt_messages_to_cohere_standard(messages) - tool_msg = next(m for m in history if m.role == "TOOL") - assert tool_msg.toolResults[0].call.name == "calc" - assert tool_msg.toolResults[0].outputs[0]["output"] == "result: 42" + assert all(m.role != "TOOL" for m in history) + + results = extract_cohere_tool_results(messages) + assert results is not None + assert results[0].call.name == "calc" + assert results[0].call.parameters == {"x": 1} + assert results[0].outputs[0]["output"] == "result: 42" # =========================================================================== @@ -541,14 +530,36 @@ _COHERE_RAW_RESPONSE = httpx.Response(200, request=httpx.Request("POST", "https: def test_handle_cohere_response_complete(): model_response = ModelResponse() - result = handle_cohere_response( - _COHERE_RESPONSE_JSON, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE - ) + result = handle_cohere_response(_COHERE_RESPONSE_JSON, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE) assert result.choices[0].finish_reason == "stop" assert result.choices[0].message["content"] == "Hello from Cohere!" assert result.usage.prompt_tokens == 10 +def test_handle_cohere_response_tool_grounded_citations(): + """Tool-grounded answers carry citations with ``documentIds`` (camelCase) and + no ``document_ids``. Response parsing must not choke on them (they are never + surfaced). Regression for the MLflow {{trace}} agentic loop.""" + resp = { + **_COHERE_RESPONSE_JSON, + "chatResponse": { + **_COHERE_RESPONSE_JSON["chatResponse"], + "text": "There are 47 spans.", + "citations": [ + { + "start": 10, + "end": 18, + "text": "47 spans", + "documentIds": ["get_root_span:0:2:0"], + } + ], + }, + } + model_response = ModelResponse() + result = handle_cohere_response(resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE) + assert result.choices[0].message["content"] == "There are 47 spans." + + def test_handle_cohere_response_max_tokens(): resp = { **_COHERE_RESPONSE_JSON, @@ -558,9 +569,7 @@ def test_handle_cohere_response_max_tokens(): }, } model_response = ModelResponse() - result = handle_cohere_response( - resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE - ) + result = handle_cohere_response(resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE) assert result.choices[0].finish_reason == "length" @@ -574,9 +583,7 @@ def test_handle_cohere_response_tool_call(): }, } model_response = ModelResponse() - result = handle_cohere_response( - resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE - ) + result = handle_cohere_response(resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE) assert result.choices[0].finish_reason == "tool_calls" tool_calls = result.choices[0].message["tool_calls"] assert tool_calls is not None @@ -586,16 +593,10 @@ def test_handle_cohere_response_tool_call(): def test_handle_cohere_response_missing_usage(): resp = { **_COHERE_RESPONSE_JSON, - "chatResponse": { - k: v - for k, v in _COHERE_RESPONSE_JSON["chatResponse"].items() - if k != "usage" - }, + "chatResponse": {k: v for k, v in _COHERE_RESPONSE_JSON["chatResponse"].items() if k != "usage"}, } model_response = ModelResponse() - result = handle_cohere_response( - resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE - ) + result = handle_cohere_response(resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE) assert result.usage.prompt_tokens == 0 assert result.usage.completion_tokens == 0 assert result.usage.total_tokens == 0 @@ -763,10 +764,7 @@ class TestOCIChatConfigGetCompleteUrl: optional_params={"oci_region": "eu-frankfurt-1"}, litellm_params={}, ) - assert url == ( - "https://inference.generativeai.eu-frankfurt-1.oci.oraclecloud.com" - "/20231130/actions/chat" - ) + assert url == ("https://inference.generativeai.eu-frankfurt-1.oci.oraclecloud.com/20231130/actions/chat") def test_respects_explicit_api_base(self): config = OCIChatConfig() @@ -781,10 +779,7 @@ class TestOCIChatConfigGetCompleteUrl: def test_full_chat_url_is_not_doubled(self): config = OCIChatConfig() - full_url = ( - "https://inference.generativeai.us-chicago-1.oci.oraclecloud.com" - "/20231130/actions/chat" - ) + full_url = "https://inference.generativeai.us-chicago-1.oci.oraclecloud.com/20231130/actions/chat" url = config.get_complete_url( api_base=full_url, api_key=None, @@ -868,23 +863,17 @@ class TestOCIChatConfigGetOptionalParams: def test_tool_choice_string_auto_converted_to_dict(self): config = self._config() - result = config._get_optional_params( - OCIVendors.GENERIC, {"tool_choice": "auto"} - ) + result = config._get_optional_params(OCIVendors.GENERIC, {"tool_choice": "auto"}) assert result["toolChoice"] == {"type": "AUTO"} def test_tool_choice_string_none_converted_to_dict(self): config = self._config() - result = config._get_optional_params( - OCIVendors.GENERIC, {"tool_choice": "none"} - ) + result = config._get_optional_params(OCIVendors.GENERIC, {"tool_choice": "none"}) assert result["toolChoice"] == {"type": "NONE"} def test_tool_choice_string_required_converted_to_dict(self): config = self._config() - result = config._get_optional_params( - OCIVendors.GENERIC, {"tool_choice": "required"} - ) + result = config._get_optional_params(OCIVendors.GENERIC, {"tool_choice": "required"}) assert result["toolChoice"] == {"type": "REQUIRED"} def test_tool_choice_openai_function_dict_converted_to_oci_form(self): @@ -910,16 +899,12 @@ class TestOCIChatConfigGetOptionalParams: def test_tool_choice_dict_auto_uppercased(self): config = self._config() - result = config._get_optional_params( - OCIVendors.GENERIC, {"tool_choice": {"type": "auto"}} - ) + result = config._get_optional_params(OCIVendors.GENERIC, {"tool_choice": {"type": "auto"}}) assert result["toolChoice"] == {"type": "AUTO"} def test_response_format_json_generic(self): config = self._config() - result = config._get_optional_params( - OCIVendors.GENERIC, {"response_format": {"type": "json_object"}} - ) + result = config._get_optional_params(OCIVendors.GENERIC, {"response_format": {"type": "json_object"}}) assert result["responseFormat"]["type"] == "JSON_OBJECT" def test_tools_adapted_for_cohere(self): @@ -1039,9 +1024,7 @@ class TestOCIStreamWrapperChunkCreator: def test_cohere_chunk_dispatched_correctly(self): wrapper = self._make_wrapper(_COHERE_MODEL) - payload = json.dumps( - {"apiFormat": "COHERE", "text": "hi", "finishReason": None} - ) + payload = json.dumps({"apiFormat": "COHERE", "text": "hi", "finishReason": None}) result = wrapper.chunk_creator(f"data:{payload}") assert result.choices[0].delta.content == "hi" @@ -1072,9 +1055,7 @@ class TestOCIStreamWrapperChunkCreator: # consolidation chunk would have its real text suppressed as a # "duplicate" and the response would be lost. wrapper = self._make_wrapper(_COHERE_MODEL) - empty_payload = json.dumps( - {"apiFormat": "COHERE", "text": "", "finishReason": None} - ) + empty_payload = json.dumps({"apiFormat": "COHERE", "text": "", "finishReason": None}) wrapper.chunk_creator(f"data:{empty_payload}") assert wrapper._cohere_text_emitted is False @@ -1102,9 +1083,7 @@ def test_get_sync_custom_stream_wrapper_returns_wrapper(): mock_response = MagicMock() mock_response.status_code = 200 - mock_response.iter_text.return_value = iter( - ['data:{"finishReason":"COMPLETE","index":0}'] - ) + mock_response.iter_text.return_value = iter(['data:{"finishReason":"COMPLETE","index":0}']) mock_client = MagicMock() mock_client.post.return_value = mock_response