diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 9163ddfa5de..3d1c263cc6b 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -9,6 +9,7 @@ response parsing, and streaming chunk parsing for models served with import datetime import json from collections.abc import Iterable, Mapping, Sequence +from itertools import chain from typing import Final import httpx @@ -106,16 +107,15 @@ 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. + """Convert an OpenAI-format message array into a Cohere ``chatHistory``. + + Every message is converted in order, including the last user turn and any + trailing tool result. The caller decides which message becomes the request's + top-level ``message`` and lifts a trailing tool result into the top-level + ``toolResults`` field, see :func:`build_cohere_request_messages`. - - 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``. + ``preambleOverride`` separately) and 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. @@ -129,16 +129,8 @@ def adapt_messages_to_cohere_standard( for tool_call in msg["tool_calls"] or [] } - 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] - ) - chat_history: Final[list[CohereMessage]] = [] - for msg in history_source: + for msg in messages: role = msg.get("role") content = extract_text_content(msg.get("content")) @@ -172,6 +164,48 @@ def adapt_messages_to_cohere_standard( return chat_history +def _trailing_tool_split_index(chat_history: Sequence[CohereMessage]) -> int: + """Index where the trailing run of TOOL entries starts (``len`` if there is none).""" + trailing_tool_count: Final = next( + (i for i, entry in enumerate(reversed(chat_history)) if not isinstance(entry, CohereToolMessage)), + len(chat_history), + ) + return len(chat_history) - trailing_tool_count + + +def _flatten_tool_results(tool_messages: Sequence[CohereMessage]) -> Sequence[CohereToolResult]: + return list( + chain.from_iterable(entry.toolResults for entry in tool_messages if isinstance(entry, CohereToolMessage)) + ) + + +def build_cohere_request_messages( + messages: Sequence[AllMessageValues], +) -> tuple[str, Sequence[CohereMessage], Sequence[CohereToolResult]]: + """Split OpenAI-format messages into Cohere's ``message``, ``chatHistory`` and ``toolResults``. + + A tool-result turn, the standard agentic continuation where the last message + is a tool result, must send those results in the request's top-level + ``toolResults`` and leave ``message`` empty. OCI rejects a request that sets + ``message`` while the last ``chatHistory`` entry holds tool results, and it + rejects an empty ``message`` unless ``toolResults`` is present. The user turn + that triggered the call stays in ``chatHistory``. + + Any other turn puts the last user message's text in ``message`` and every + preceding message in ``chatHistory``. + """ + non_system_messages: Final = [m for m in messages if m.get("role") != "system"] + converted: Final = adapt_messages_to_cohere_standard(non_system_messages) + + split_index: Final = _trailing_tool_split_index(converted) + if split_index < len(converted): + return "", converted[:split_index], _flatten_tool_results(converted[split_index:]) + + user_messages: Final = [m for m in non_system_messages if m.get("role") == "user"] + message: Final = extract_text_content(user_messages[-1]["content"]) if user_messages else "" + return message, converted[:-1], [] + + def _resolved_oci_parameter_schema(raw_parameters: dict[str, JsonValue]) -> JsonValue: return sanitize_oci_schema(resolve_oci_schema_anyof(resolve_oci_schema_refs(raw_parameters))) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 14ef6225666..5f3410bb8c8 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -28,8 +28,8 @@ from litellm.llms.custom_httpx.http_handler import ( version, ) from litellm.llms.oci.chat.cohere import ( - adapt_messages_to_cohere_standard, adapt_tool_definitions_to_cohere_standard, + build_cohere_request_messages, extract_text_content, handle_cohere_response, handle_cohere_stream_chunk, @@ -567,11 +567,13 @@ class OCIChatConfig(BaseConfig): if preamble: preamble_override = preamble + message, chat_history, tool_results = build_cohere_request_messages(messages) 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=list(chat_history), preambleOverride=preamble_override, + toolResults=list(tool_results) or None, **self._get_optional_params(OCIVendors.COHERE, optional_params, model), ) data = OCICompletionPayload( diff --git a/litellm/types/llms/oci.py b/litellm/types/llms/oci.py index bd3a0af1dca..85121beb4df 100644 --- a/litellm/types/llms/oci.py +++ b/litellm/types/llms/oci.py @@ -1,9 +1,10 @@ from __future__ import annotations +from collections.abc import Sequence from enum import Enum from typing import Literal -from pydantic import SerializeAsAny +from pydantic import Field, SerializeAsAny from litellm.types.llms.base import LiteLLMBaseModel @@ -313,6 +314,10 @@ class CohereChatRequest(LiteLLMBaseModel): stopSequences: list[str] | None = None seed: int | None = None tools: list[CohereTool] | None = None + # Top-level results of the tool calls in the last assistant turn. OCI + # requires these here (with ``message`` empty) for a tool-result turn; + # trailing tool results must not stay in ``chatHistory``. + toolResults: Sequence[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 @@ -349,12 +354,16 @@ class CohereUsage(LiteLLMBaseModel): class CohereCitation(LiteLLMBaseModel): - """Citation in Cohere response.""" + """Citation in Cohere response. + + OCI writes the referenced document ids as ``documentIds`` (camelCase), so + the field accepts that wire name as an alias. + """ start: int end: int text: str - document_ids: list[str] + document_ids: Sequence[str] = Field(alias="documentIds") class CohereSearchQuery(LiteLLMBaseModel): 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..037ed684fab 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 @@ -325,18 +325,18 @@ class TestOCICohereToolCalls: chat_history = adapt_messages_to_cohere_standard(messages) - # 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 + # The converter keeps every message; splitting it into the request's + # top-level `message`/`toolResults` is the caller's job. + assert [msg.role for msg in chat_history] == ["USER", "CHATBOT", "TOOL"] - assistant_msg = chat_history[0] + assistant_msg = chat_history[1] 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_msg = chat_history[1] + tool_msg = chat_history[2] assert tool_msg.role == "TOOL" assert tool_msg.toolResults[0].call.name == "get_weather" assert tool_msg.toolResults[0].outputs[0]["output"] == "Sunny, 25C" @@ -370,12 +370,11 @@ class TestOCICohereToolCalls: 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 + # The converter keeps every message in order, including the last user + # turn and the trailing tool result. + assert [msg.role for msg in chat_history] == ["USER", "CHATBOT", "TOOL"] - assistant_msg = chat_history[0] + assistant_msg = chat_history[1] assert assistant_msg.role == "CHATBOT" assert assistant_msg.message == "I will look up the weather in Tokyo." assert assistant_msg.toolCalls is not None @@ -383,7 +382,7 @@ class TestOCICohereToolCalls: assert assistant_msg.toolCalls[0].name == "get_weather" assert assistant_msg.toolCalls[0].parameters == {"location": "Tokyo"} - tool_msg = chat_history[1] + tool_msg = chat_history[2] assert tool_msg.role == "TOOL" assert tool_msg.toolResults[0].call.name == "get_weather" assert tool_msg.toolResults[0].call.parameters == {"location": "Tokyo"} @@ -392,6 +391,172 @@ class TestOCICohereToolCalls: == "The weather in Tokyo is 22°C with partly cloudy skies." ) + def test_cohere_tool_result_turn_uses_top_level_tool_results(self): + """The call after a tool call must move the tool result into the top-level + ``toolResults`` field and clear ``message``, keeping the user turn in + ``chatHistory``. OCI rejects ``message`` set next to trailing tool results + (issue #46157). + """ + config = OCIChatConfig() + messages = [ + {"role": "user", "content": "How is support ticket priority decided?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lookup_policy", + "arguments": json.dumps({"topic": "ticket priority"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "P1 for an outage, P2 for degraded service, P3 for questions.", + }, + ] + optional_params = {"oci_compartment_id": TEST_COMPARTMENT_ID} + + chat_request = config.transform_request( + model="cohere.command-a-03-2025", + messages=messages, # type: ignore + optional_params=optional_params, + litellm_params={}, + headers={}, + )["chatRequest"] + + assert chat_request["message"] == "" + assert chat_request["toolResults"] == [ + { + "call": {"name": "lookup_policy", "parameters": {"topic": "ticket priority"}}, + "outputs": [{"output": "P1 for an outage, P2 for degraded service, P3 for questions."}], + } + ] + # The user question and the assistant tool call stay in history so the + # model still has the turn the tool result answers. + assert [msg["role"] for msg in chat_request["chatHistory"]] == ["USER", "CHATBOT"] + assert chat_request["chatHistory"][0]["message"] == "How is support ticket priority decided?" + assert chat_request["chatHistory"][1]["toolCalls"][0]["name"] == "lookup_policy" + + def test_cohere_parallel_tool_result_turn_uses_top_level_tool_results(self): + """Parallel tool results trailing the turn are all lifted into the + top-level ``toolResults`` in call order.""" + config = OCIChatConfig() + messages = [ + {"role": "user", "content": "Compare the weather in Paris and London"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"location": "Paris"})}, + }, + { + "id": "call_2", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"location": "London"})}, + }, + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Paris: 18C"}, + {"role": "tool", "tool_call_id": "call_2", "content": "London: 15C"}, + ] + optional_params = {"oci_compartment_id": TEST_COMPARTMENT_ID} + + chat_request = config.transform_request( + model="cohere.command-a-03-2025", + messages=messages, # type: ignore + optional_params=optional_params, + litellm_params={}, + headers={}, + )["chatRequest"] + + assert chat_request["message"] == "" + assert [result["call"]["parameters"] for result in chat_request["toolResults"]] == [ + {"location": "Paris"}, + {"location": "London"}, + ] + assert [msg["role"] for msg in chat_request["chatHistory"]] == ["USER", "CHATBOT"] + + def test_cohere_citation_parses_document_ids_wire_alias(self): + """OCI writes citation ids as ``documentIds``; the response model must + accept that wire name (issue #46157).""" + from litellm.types.llms.oci import CohereChatResult + + payload = { + "modelId": "cohere.command-a-03-2025", + "modelVersion": "1.0", + "chatResponse": { + "apiFormat": "COHERE", + "text": "P1 is for an outage.", + "finishReason": "COMPLETE", + "citations": [ + { + "start": 44, + "end": 70, + "text": "P1 is for an outage.", + "documentIds": ["lookup_policy:0:1:0"], + } + ], + "usage": {"promptTokens": 20, "completionTokens": 8, "totalTokens": 28}, + }, + } + + result = CohereChatResult.model_validate(payload) + + assert result.chatResponse.citations is not None + assert result.chatResponse.citations[0].document_ids == ["lookup_policy:0:1:0"] + + def test_cohere_response_with_citations_parses(self): + """A tool-answer response whose citations use ``documentIds`` must not + raise in transform_response (issue #46157).""" + config = OCIChatConfig() + mock_response = { + "modelId": "cohere.command-a-03-2025", + "modelVersion": "1.0", + "chatResponse": { + "apiFormat": "COHERE", + "text": "P1 is for an outage.", + "finishReason": "COMPLETE", + "citations": [ + { + "start": 0, + "end": 20, + "text": "P1 is for an outage.", + "documentIds": ["lookup_policy:0:1:0"], + } + ], + "usage": {"promptTokens": 20, "completionTokens": 8, "totalTokens": 28}, + }, + } + response = httpx.Response( + status_code=200, + json=mock_response, + headers={"Content-Type": "application/json"}, + ) + + result = config.transform_response( + model="cohere.command-a-03-2025", + raw_response=response, + model_response=ModelResponse(), + logging_obj={}, # type: ignore + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding={}, + ) + + assert result.choices[0].message.content == "P1 is for an outage." + assert result.choices[0].finish_reason == "stop" + def test_cohere_streaming_chunk_handling(self): """Test Cohere streaming chunk handling""" # Mock the required parameters diff --git a/tests/unit/llms/oci/test_oci_coverage_boost.py b/tests/unit/llms/oci/test_oci_coverage_boost.py index 251df5e4cad..29e68f0b1f4 100644 --- a/tests/unit/llms/oci/test_oci_coverage_boost.py +++ b/tests/unit/llms/oci/test_oci_coverage_boost.py @@ -449,9 +449,9 @@ def test_adapt_cohere_user_in_history(): {"role": "user", "content": "current question"}, ] history = adapt_messages_to_cohere_standard(messages) - assert len(history) == 1 - assert history[0].role == "USER" + assert [msg.role for msg in history] == ["USER", "USER"] assert history[0].message == "first question" + assert history[1].message == "current question" def test_adapt_cohere_assistant_in_history(): @@ -461,7 +461,7 @@ def test_adapt_cohere_assistant_in_history(): {"role": "user", "content": "follow-up"}, ] history = adapt_messages_to_cohere_standard(messages) - assert len(history) == 2 + assert [msg.role for msg in history] == ["USER", "CHATBOT", "USER"] chatbot_msg = history[1] assert chatbot_msg.role == "CHATBOT" assert chatbot_msg.message == "answer" @@ -483,8 +483,7 @@ def test_adapt_cohere_assistant_with_tool_calls_in_history(): {"role": "user", "content": "thanks"}, ] history = adapt_messages_to_cohere_standard(messages) - assert len(history) == 1 - assert history[0].role == "CHATBOT" + assert [msg.role for msg in history] == ["CHATBOT", "USER"] assert history[0].toolCalls is not None assert history[0].toolCalls[0].name == "calc"