fix(oci): align types and transformation with official OCI SDK

- Remove OCIVendors.GEMINI — apiFormat="GEMINI" is invalid; all non-Cohere
  models use apiFormat="GENERIC"
- Add toolChoice, logitBias, logProbs to OCIChatRequestPayload so params
  present in the mapping are no longer silently dropped by Pydantic
- Exclude n→numGenerations from Cohere param map (not a Cohere API field)
- Fix CohereToolResult: change callId/result to call/outputs matching
  the OCI SDK's CohereToolResult structure
- Fix CohereToolMessage: replace non-existent toolCallId with toolResults
  list; update adapt_messages_to_cohere_standard to build proper tool-result
  history entries by resolving tool call name+params from preceding assistant
  messages
- Map generic-model stream finish reasons to OpenAI convention
  (COMPLETE→stop, MAX_TOKENS→length, TOOL_CALLS→tool_calls), consistent
  with the existing Cohere streaming path
- Add optional id field to OCIEmbedResponse so valid API responses
  carrying an id are not rejected by the Pydantic model
This commit is contained in:
Federico Kamelhar 2026-04-05 12:18:38 -04:00
parent 836d7e6cf3
commit 2b055ed5f8
2 changed files with 75 additions and 19 deletions

View file

@ -41,6 +41,8 @@ from litellm.types.llms.oci import (
CohereStreamChunk,
CohereTool,
CohereToolCall,
CohereToolMessage,
CohereToolResult,
OCIChatRequestPayload,
OCICompletionPayload,
OCICompletionResponse,
@ -149,14 +151,15 @@ class OCIChatConfig(BaseConfig):
"response_format": "responseFormat",
}
# Cohere uses the same parameter keys as GENERIC with two differences:
# Cohere uses the same parameter keys as GENERIC with three differences:
# - tool_choice is unsupported
# - stop sequences are named "stopSequences" not "stop"
# - n (numGenerations) is a GenericChatRequest-only field; CohereChatRequest has no equivalent
# Build a *separate* frozen reference map so callers never mutate the canonical dict.
self._openai_to_oci_cohere_param_map = {
k: ("stopSequences" if k == "stop" else v)
for k, v in self.openai_to_oci_generic_param_map.items()
if k not in ("tool_choice", "max_retries")
if k not in ("tool_choice", "max_retries", "n")
}
def get_supported_openai_params(self, model: str) -> List[str]:
@ -355,8 +358,32 @@ class OCIChatConfig(BaseConfig):
def adapt_messages_to_cohere_standard(
self, messages: List[AllMessageValues]
) -> List[CohereMessage]:
"""Build chat history for Cohere models."""
chat_history = []
"""Build chat history for Cohere models.
Tool results are represented as OCI Cohere ``toolResults`` entries, where each
entry carries the originating tool call (name + parameters resolved from the
preceding assistant message) and the output text.
"""
# First pass: build tool_call_id -> CohereToolCall lookup 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":
for tc in msg.get("tool_calls") or []: # type: ignore[union-attr]
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,
)
chat_history: List[CohereMessage] = []
for msg in messages[:-1]: # All messages except the last one
role = msg.get("role")
content = msg.get("content")
@ -407,15 +434,22 @@ class OCIChatConfig(BaseConfig):
CohereMessage(role="CHATBOT", message=content, toolCalls=tool_calls)
)
elif role == "tool":
# Tool result messages: include the tool_call_id so Cohere can correlate
# the result back to the right tool call in the conversation history.
tool_call_id = msg.get("tool_call_id") # type: ignore[union-attr]
# Construct a proper OCI Cohere tool-result message.
# The API expects toolResults with the originating call (name + params)
# and a list of output objects — not a flat toolCallId string.
tool_call_id = msg.get("tool_call_id", "") # type: ignore[union-attr]
cohere_call = tool_call_lookup.get(
tool_call_id,
CohereToolCall(name="", parameters={}),
)
chat_history.append(
CohereMessage(
role="TOOL",
message=content,
toolCalls=None,
toolCallId=tool_call_id,
CohereToolMessage(
toolResults=[
CohereToolResult(
call=cohere_call,
outputs=[{"result": content}],
)
]
)
)
@ -1134,6 +1168,17 @@ class OCIStreamWrapper(CustomStreamWrapper):
if typed_chunk.message and typed_chunk.message.toolCalls:
tool_calls = adapt_tools_to_openai_standard(typed_chunk.message.toolCalls)
# Map OCI finish reasons to OpenAI convention (same as Cohere path)
oci_finish_reason = typed_chunk.finishReason
if oci_finish_reason == "COMPLETE":
finish_reason: Optional[str] = "stop"
elif oci_finish_reason == "MAX_TOKENS":
finish_reason = "length"
elif oci_finish_reason == "TOOL_CALLS":
finish_reason = "tool_calls"
else:
finish_reason = oci_finish_reason # None while streaming; unknown passed through
return ModelResponseStream(
choices=[
StreamingChoices(
@ -1149,7 +1194,7 @@ class OCIStreamWrapper(CustomStreamWrapper):
thinking_blocks=None, # OCI does not have thinking blocks in the response
reasoning_content=None, # OCI does not have reasoning content in the response
),
finish_reason=typed_chunk.finishReason,
finish_reason=finish_reason,
)
]
)

View file

@ -15,7 +15,6 @@ class OCIVendors(Enum):
"""
COHERE = "COHERE"
GEMINI = "GEMINI"
GENERIC = "GENERIC"
@ -103,6 +102,9 @@ class OCIChatRequestPayload(BaseModel):
frequencyPenalty: Optional[float] = None
presencePenalty: Optional[float] = None
responseFormat: Optional[Dict[str, Any]] = None
toolChoice: Optional[Union[str, Dict[str, Any]]] = None
logitBias: Optional[Dict[str, Any]] = None
logProbs: Optional[int] = None
class OCIServingMode(BaseModel):
@ -238,10 +240,14 @@ class CohereSystemMessage(CohereMessage):
class CohereToolMessage(CohereMessage):
"""Tool message in Cohere chat."""
"""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"
toolCallId: str
toolResults: List[CohereToolResult]
class CohereParameterDefinition(BaseModel):
@ -268,10 +274,14 @@ class CohereToolCall(BaseModel):
class CohereToolResult(BaseModel):
"""Result of a tool call."""
"""Result of a tool call.
callId: str
result: str
Matches the OCI SDK's CohereToolResult: each result carries the originating
tool call (name + parameters) and a list of output objects.
"""
call: CohereToolCall
outputs: List[Dict[str, Any]]
class CohereResponseFormat(BaseModel):
@ -424,6 +434,7 @@ class OCIEmbedUsage(BaseModel):
class OCIEmbedResponse(BaseModel):
"""Response body from POST /20231130/actions/embedText."""
id: Optional[str] = None # present in the official SDK response
embeddings: List[List[float]]
modelId: str
modelVersion: str