fix: address Greptile review — oci/ prefix fallback and empty final message guard

- _get_max_tokens_for_model now retries with oci/ prefix when the bare
  model name is not found in the cost map. This handles the case where
  litellm strips the provider prefix before calling transform_request.
- Filter empty user messages when extracting the final Cohere message
  field, not just in chat history.
- Added test for bare model name lookup (without oci/ prefix).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Rami Salman 2026-03-30 18:50:09 +03:00
parent 4f7ac6e72b
commit 5efccbd06b
2 changed files with 36 additions and 10 deletions

View file

@ -671,15 +671,25 @@ class OCIChatConfig(BaseConfig):
Uses litellm's model cost map for a dynamic lookup, falling back to
DEFAULT_MAX_TOKENS (configurable via the DEFAULT_MAX_TOKENS env var)
when the model is not found.
The model cost map stores OCI entries under ``oci/`` prefixed keys
(e.g. ``oci/cohere.command-a-03-2025``), but litellm strips the
provider prefix before calling ``transform_request``. We therefore
try the bare name first and, if that misses, retry with the
``oci/`` prefix.
"""
if model is None:
return DEFAULT_MAX_TOKENS
try:
max_tokens = get_max_tokens(model)
if max_tokens is not None:
return max_tokens
except Exception:
pass
# Try the bare model name first, then with the oci/ provider prefix
# since the cost map keys are prefixed but transform_request receives
# the name without the prefix.
for candidate in (model, f"oci/{model}"):
try:
max_tokens = get_max_tokens(candidate)
if max_tokens is not None:
return max_tokens
except Exception:
continue
return DEFAULT_MAX_TOKENS
def _get_optional_params(
@ -907,8 +917,15 @@ class OCIChatConfig(BaseConfig):
# Build request based on vendor type
if vendor == OCIVendors.COHERE:
# For Cohere, we need to use the specific Cohere format
# Extract the last user message as the main message
user_messages = [msg for msg in messages if msg.get("role") == "user"]
# Extract the last user message as the main message.
# Filter out user messages with empty content — the Cohere API
# requires the main message field to be non-empty.
user_messages = [
msg
for msg in messages
if msg.get("role") == "user"
and self._extract_text_content(msg.get("content", "")).strip()
]
if not user_messages:
raise Exception("No user message found for Cohere model")

View file

@ -74,11 +74,20 @@ class TestOCICohereMaxTokens:
assert params["maxTokens"] == 1024
def test_get_max_tokens_for_model_known_model(self):
"""_get_max_tokens_for_model should return the correct value for a known model."""
def test_get_max_tokens_for_model_with_oci_prefix(self):
"""_get_max_tokens_for_model should work with the oci/ prefix."""
result = OCIChatConfig._get_max_tokens_for_model("oci/cohere.command-a-03-2025")
assert result == 4000
def test_get_max_tokens_for_model_without_oci_prefix(self):
"""_get_max_tokens_for_model should also work without the oci/ prefix.
litellm strips the provider prefix before calling transform_request,
so the method must retry with the oci/ prefix to find the cost map entry.
"""
result = OCIChatConfig._get_max_tokens_for_model("cohere.command-a-03-2025")
assert result == 4000
def test_get_max_tokens_for_model_unknown_model(self):
"""_get_max_tokens_for_model should return DEFAULT_MAX_TOKENS for an unknown model."""
result = OCIChatConfig._get_max_tokens_for_model("oci/unknown-model-xyz")