scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Revert "fix(oci): make Cohere {{trace}} judges work (tool param types + agentic tool-calling continuation) (#30646)"

Browse source 
This reverts commit 50f34e0b15.
This commit is contained in:
Sameer Kankute 2026-06-18 19:17:24 +05:30
parent bba8bc6f1b
commit 15a7c68728
No known key found for this signature in database
8 changed files with 176 additions and 323 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -32,12 +32,10 @@ from litellm.types.llms.oci import (
CohereStreamChunk,
CohereTool,
CohereToolCall,
CohereToolMessage,
CohereToolResult,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantToolCall,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import (
Choices,
Delta,
@ -63,76 +61,42 @@ def _extract_text_content(content: Any) -> str:
return str(content)
def _message_text(msg: AllMessageValues) -> str:
"""Plain-text content of a message (OpenAI's content field is loosely typed)."""
return _extract_text_content(msg.get("content")) # any-ok: OpenAI content union
def _assistant_tool_calls(
msg: AllMessageValues,
) -> list[ChatCompletionAssistantToolCall]:
"""Tool calls on an assistant message, empty for any other role."""
if msg.get("role") != "assistant":
return []
return msg.get("tool_calls") or [] # any-ok: union .get() is loosely typed
def _tool_call_to_cohere(tc: ChatCompletionAssistantToolCall) -> CohereToolCall:
"""Convert one OpenAI tool call into a Cohere ``CohereToolCall``."""
fn = tc["function"]
name = fn.get("name") or ""
raw = fn.get("arguments") or "{}"
try:
args = json.loads(raw) # any-ok: tool-call args are arbitrary JSON
except json.JSONDecodeError:
args = {} # any-ok: arbitrary JSON
return CohereToolCall(name=name, parameters=args) # any-ok: arbitrary JSON
def adapt_messages_to_cohere_standard(
messages: List[AllMessageValues],
) -> List[CohereMessage]:
"""Build a Cohere ``chatHistory`` (USER and CHATBOT turns) from an OpenAI array.
"""Build a Cohere ``chatHistory`` list from an OpenAI-format message array.
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 a 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``).
- 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.
"""
chat_history: list[CohereMessage] = []
for msg in messages[:-1]:
role = msg.get("role")
if role == "user":
chat_history.append(CohereMessage(role="USER", message=_message_text(msg)))
elif role == "assistant":
calls = [_tool_call_to_cohere(tc) for tc in _assistant_tool_calls(msg)]
chat_history.append(
CohereMessage(
role="CHATBOT",
message=_message_text(msg),
toolCalls=calls or None,
# First pass: build tool_call_id → CohereToolCall so tool-result messages can
# reference the originating call by name and parameters.
tool_call_lookup: Dict[str, CohereToolCall] = {}
for msg in messages:
if msg.get("role") == "assistant":
tool_calls_raw: Any = msg.get("tool_calls") or []
for tc in tool_calls_raw:
tc_id = tc.get("id", "")
raw_args: Any = tc.get("function", {}).get("arguments", "{}")
try:
params: Dict[str, Any] = (
json.loads(raw_args) if isinstance(raw_args, str) else raw_args
)
except json.JSONDecodeError:
params = {}
tool_call_lookup[tc_id] = CohereToolCall(
name=str(tc.get("function", {}).get("name", "")),
parameters=params,
)
)
return chat_history
def extract_cohere_tool_results(
messages: list[AllMessageValues],
) -> Optional[list[CohereToolResult]]:
"""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.
"""
lookup: dict[str, CohereToolCall] = {
tc.get("id") or "": _tool_call_to_cohere(tc)
for msg in messages
for tc in _assistant_tool_calls(msg)
}
last_user_index = next(
(
@ -142,17 +106,62 @@ def extract_cohere_tool_results(
),
None,
)
current_turn = messages if last_user_index is None else messages[last_user_index:]
history_source = (
messages
if last_user_index is None
else [m for i, m in enumerate(messages) if i != last_user_index]
)
unknown_call = CohereToolCall(name="", parameters={}) # any-ok: arbitrary JSON
results: list[CohereToolResult] = []
for msg in current_turn:
if msg.get("role") != "tool":
continue
call = lookup.get(str(msg.get("tool_call_id") or ""), unknown_call)
outputs = [{"output": _message_text(msg)}] # any-ok: outputs are arbitrary
results.append(CohereToolResult(call=call, outputs=outputs)) # any-ok: outputs
return results or None
chat_history: List[CohereMessage] = []
for msg in history_source:
role = msg.get("role")
content = _extract_text_content(msg.get("content"))
tool_calls: Optional[List[CohereToolCall]] = None
if role == "assistant" and msg.get("tool_calls"): # type: ignore[union-attr,typeddict-item]
tool_calls = []
for tc in msg["tool_calls"]: # type: ignore[union-attr,typeddict-item]
raw_arguments: Any = tc.get("function", {}).get("arguments", {})
if isinstance(raw_arguments, str):
try:
arguments: Dict[str, Any] = json.loads(raw_arguments)
except json.JSONDecodeError:
arguments = {}
else:
arguments = raw_arguments
tool_calls.append(
CohereToolCall(
name=str(tc.get("function", {}).get("name", "")),
parameters=arguments,
)
)
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
def adapt_tool_definitions_to_cohere_standard(

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

@ -37,10 +37,8 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.llms.oci.chat.cohere import (
_extract_text_content,
_message_text,
adapt_messages_to_cohere_standard,
adapt_tool_definitions_to_cohere_standard,
extract_cohere_tool_results,
handle_cohere_response,
handle_cohere_stream_chunk,
)
@ -596,19 +594,12 @@ class OCIChatConfig(BaseConfig):
if preamble:
preamble_override = preamble
non_system_messages = [m for m in messages if m.get("role") != "system"]
tool_results = extract_cohere_tool_results(non_system_messages)
# 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 = "" if tool_results else _message_text(non_system_messages[-1])
chat_request = CohereChatRequest(
apiFormat="COHERE",
message=message,
chatHistory=adapt_messages_to_cohere_standard(non_system_messages),
toolResults=tool_results,
message=_extract_text_content(user_messages[-1]["content"]),
chatHistory=adapt_messages_to_cohere_standard(
[m for m in messages if m.get("role") != "system"]
),
preambleOverride=preamble_override,
**self._get_optional_params(OCIVendors.COHERE, optional_params, model),
)

8
litellm/llms/oci/common_utils.py
View file

@ -432,16 +432,14 @@ 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 collection
# types must be the lowercase builtins: OCI's Cohere backend returns HTTP 500
# on a bare ``List`` (it expects ``list`` / ``List[...]``).
# the OCI Cohere API's CohereParameterDefinition.type field.
OCI_JSON_TO_PYTHON_TYPES: Dict[str, str] = {
"string": "str",
"number": "float",
"boolean": "bool",
"integer": "int",
"array": "list",
"object": "dict",
"array": "List",
"object": "Dict",
"any": "any",
}

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

@ -313,10 +313,6 @@ class CohereChatRequest(BaseModel):
stopSequences: Optional[List[str]] = None
seed: Optional[int] = None
tools: Optional[List[CohereTool]] = 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: Optional[List[CohereToolResult]] = 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
@ -358,18 +354,14 @@ class CohereCitation(BaseModel):
start: int
end: int
text: 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: Optional[List[str]] = None
document_ids: List[str]
class CohereSearchQuery(BaseModel):
"""Search query generated by Cohere."""
text: str
generation_id: Optional[str] = None
generation_id: str
class CohereChatResponse(BaseModel):

89
tests/integration/test_oci_integration.py
View file

@ -387,95 +387,6 @@ 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)
# ---------------------------------------------------------------------------

53
tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py
View file

@ -1089,58 +1089,6 @@ 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
@ -1530,7 +1478,6 @@ 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

106
tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py
View file

@ -296,14 +296,10 @@ class TestOCICohereToolCalls:
assert transformed_request["chatRequest"]["responseFormat"] == {"type": "TEXT"}
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."""
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."""
config = OCIChatConfig()
messages = [
{"role": "user", "content": "What's the weather?"},
{
@ -320,43 +316,81 @@ class TestOCICohereToolCalls:
}
],
},
{"role": "tool", "content": "Sunny, 25C", "tool_call_id": "call_1"},
{
"role": "tool",
"content": "Sunny, 25C",
"tool_call_id": "call_1",
},
]
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)
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"]
# 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
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"
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"
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."""
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."""
config = OCIChatConfig()
messages = [{"role": "user", "content": "What's the weather like in Tokyo?"}]
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"]
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",
},
]
assert chat_request["message"] == "What's the weather like in Tokyo?"
assert "toolResults" not in chat_request
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."
)
def test_cohere_streaming_chunk_handling(self):
"""Test Cohere streaming chunk handling"""

49
tests/test_litellm/llms/oci/test_oci_coverage_boost.py
View file

@ -485,11 +485,8 @@ def test_adapt_cohere_assistant_with_tool_calls_in_history():
assert history[0].toolCalls[0].name == "calc"
def test_adapt_cohere_tool_result_goes_to_tool_results_not_history():
from litellm.llms.oci.chat.cohere import extract_cohere_tool_results
def test_adapt_cohere_tool_result_in_history():
messages = [
{"role": "user", "content": "compute"},
{
"role": "assistant",
"content": None,
@ -501,17 +498,17 @@ def test_adapt_cohere_tool_result_goes_to_tool_results_not_history():
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "result: 42"},
{
"role": "tool",
"tool_call_id": "call_1",
"content": "result: 42",
},
{"role": "user", "content": "ok"},
]
# Tool results never appear in chatHistory; they ride the top-level field.
history = adapt_messages_to_cohere_standard(messages)
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"
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"
# ===========================================================================
@ -548,32 +545,6 @@ def test_handle_cohere_response_complete():
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,