Merge c8b4898963 into 1df25e26cf

litellm/llms/cerebras/chat.py

@@ -6,6 +6,10 @@ this is OpenAI compatible - no translation needed / occurs
 
 from typing import Final
 
+from openai.types.chat.chat_completion_prediction_content_param import (
+    ChatCompletionPredictionContentParam,
+)
+
 from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
 from litellm.utils import supports_reasoning
 
@@ -17,29 +21,49 @@ class CerebrasConfig(OpenAIGPTConfig):
     Below are the parameters:
     """
 
+    max_completion_tokens: int | None = None
     max_tokens: int | None = None
     response_format: dict | None = None
     seed: int | None = None
     stream: bool | None = None
-    top_p: int | None = None
-    tool_choice: str | None = None
+    top_p: float | None = None  # pyright: ignore[reportIncompatibleVariableOverride]  # Fractions supported.
+    tool_choice: str | dict | None = None  # mutable-ok: API payload.
     tools: list | None = None
     user: str | None = None
     reasoning_effort: str | None = None
+    parallel_tool_calls: bool | None = None
+    logprobs: bool | None = None
+    top_logprobs: int | None = None
+    frequency_penalty: float | None = None  # pyright: ignore[reportIncompatibleVariableOverride]  # Fractions supported.
+    presence_penalty: float | None = None  # pyright: ignore[reportIncompatibleVariableOverride]  # Fractions supported.
+    logit_bias: dict[str, float] | None = None  # mutable-ok: API payload.
+    service_tier: str | None = None
+    prompt_cache_key: str | None = None
+    prediction: ChatCompletionPredictionContentParam | None = None
 
     def __init__(
         self,
+        max_completion_tokens: int | None = None,
         max_tokens: int | None = None,
         response_format: dict | None = None,
         seed: int | None = None,
-        stop: str | None = None,
+        stop: str | list[str] | None = None,  # mutable-ok: API payload.
         stream: bool | None = None,
         temperature: float | None = None,
-        top_p: int | None = None,
-        tool_choice: str | None = None,
+        top_p: float | None = None,
+        tool_choice: str | dict | None = None,  # mutable-ok: API payload.
         tools: list | None = None,
         user: str | None = None,
         reasoning_effort: str | None = None,
+        parallel_tool_calls: bool | None = None,
+        logprobs: bool | None = None,
+        top_logprobs: int | None = None,
+        frequency_penalty: float | None = None,
+        presence_penalty: float | None = None,
+        logit_bias: dict[str, float] | None = None,  # mutable-ok: API payload.
+        service_tier: str | None = None,
+        prompt_cache_key: str | None = None,
+        prediction: ChatCompletionPredictionContentParam | None = None,
     ) -> None:
         locals_: Final = locals().copy()
         for key, value in locals_.items():
@@ -70,6 +94,15 @@ class CerebrasConfig(OpenAIGPTConfig):
             "user",
             "max_retries",
             "extra_headers",
+            "parallel_tool_calls",
+            "logprobs",
+            "top_logprobs",
+            "frequency_penalty",
+            "presence_penalty",
+            "logit_bias",
+            "service_tier",
+            "prompt_cache_key",
+            "prediction",
         ]
 
         # Only add reasoning_effort for models that support it
@@ -86,9 +119,7 @@ class CerebrasConfig(OpenAIGPTConfig):
         drop_params: bool,
     ) -> dict:
         supported_openai_params: Final = self.get_supported_openai_params(model=model)
-        for param, value in non_default_params.items():
-            if param == "max_completion_tokens":
-                optional_params["max_tokens"] = value
-            elif param in supported_openai_params:
-                optional_params[param] = value
+        for param in supported_openai_params:
+            if param in non_default_params:
+                optional_params[param] = non_default_params[param]
         return optional_params

tests/llm_translation/test_cerebras.py

@@ -0,0 +1,46 @@
+import pytest
+
+from litellm.llms.cerebras.chat import CerebrasConfig
+
+
+@pytest.mark.parametrize(
+    ("parameter", "value"),
+    [
+        ("max_completion_tokens", 64),
+        ("max_tokens", 32),
+        ("parallel_tool_calls", False),
+        ("logprobs", True),
+        ("top_logprobs", 3),
+        ("frequency_penalty", 0.2),
+        ("presence_penalty", 0.3),
+        ("logit_bias", {"42": -1}),
+        ("service_tier", "default"),
+        ("prompt_cache_key", "conversation-1"),
+        ("prediction", {"type": "content", "content": "expected"}),
+    ],
+)
+def test_cerebras_preserves_supported_parameters(parameter: str, value: object) -> None:
+    config = CerebrasConfig()
+
+    mapped = config.map_openai_params(
+        non_default_params={parameter: value},
+        optional_params={},
+        model="gpt-oss-120b",
+        drop_params=False,
+    )
+
+    assert mapped == {parameter: value}
+
+
+def test_cerebras_does_not_alias_max_completion_tokens() -> None:
+    config = CerebrasConfig()
+
+    mapped = config.map_openai_params(
+        non_default_params={"max_completion_tokens": 64},
+        optional_params={},
+        model="gpt-oss-120b",
+        drop_params=False,
+    )
+
+    assert mapped["max_completion_tokens"] == 64
+    assert "max_tokens" not in mapped