scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

fix(oci): support cohere tool-result turns and citation document ids

Browse source 
OCI rejects a cohere chat request that sets `message` while the last
chatHistory entry carries tool results with a 400 ("cannot specify
message if the last entry in chat history contains tool results"). Lift
trailing tool-result messages into top-level `toolResults` with an empty
`message`, keep the user and chatbot turns in `chatHistory`, and convert
every message 1:1 instead of reordering or dropping turns.

Also fix a crash parsing cohere citations: OCI's SDK exposes the field
as `documentIds` while the model read `document_ids`, so any citation
broke attribute access. Bind the alias explicitly.

Regression tests cover both directions: the request payload shape for a
tool-result turn, and the citation document ids alias.
This commit is contained in:
aashutosh-yadav 2026-10-13 14:22:30 +05:30
parent a5d05355b6
commit 1e5807b53a
5 changed files with 248 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
litellm/llms/oci/chat/cohere.py
View file

@ -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)))

8
litellm/llms/oci/chat/transformation.py
View file

@ -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(

15
litellm/types/llms/oci.py
View file

@ -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):

187
tests/unit/llms/oci/chat/test_oci_cohere_tool_calls.py
View file

@ -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

9
tests/unit/llms/oci/test_oci_coverage_boost.py
View file

@ -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"