diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py index c3aa26ade35..0d0ab030132 100644 --- a/litellm/llms/cerebras/chat.py +++ b/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 diff --git a/tests/llm_translation/test_cerebras.py b/tests/llm_translation/test_cerebras.py new file mode 100644 index 00000000000..82f23dff25a --- /dev/null +++ b/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