diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py index 9929e2ab9a2..e3a89cc95bb 100644 --- a/litellm/llms/cerebras/chat.py +++ b/litellm/llms/cerebras/chat.py @@ -68,6 +68,8 @@ class CerebrasConfig(OpenAIGPTConfig): "tool_choice", "tools", "user", + "max_retries", + "extra_headers", ] # Only add reasoning_effort for models that support it diff --git a/tests/test_litellm/llms/cerebras/__init__.py b/tests/test_litellm/llms/cerebras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat.py b/tests/test_litellm/llms/cerebras/test_cerebras_chat.py new file mode 100644 index 00000000000..e72664ede18 --- /dev/null +++ b/tests/test_litellm/llms/cerebras/test_cerebras_chat.py @@ -0,0 +1,36 @@ +import pytest + +from litellm.llms.cerebras.chat import CerebrasConfig + + +class TestCerebrasConfig: + def test_get_supported_openai_params_includes_max_retries(self): + config = CerebrasConfig() + params = config.get_supported_openai_params(model="cerebras/llama-4-scout-17b-16e") + assert "max_retries" in params + + def test_get_supported_openai_params_includes_extra_headers(self): + config = CerebrasConfig() + params = config.get_supported_openai_params(model="cerebras/llama-4-scout-17b-16e") + assert "extra_headers" in params + + def test_map_openai_params_passes_through_max_retries(self): + config = CerebrasConfig() + result = config.map_openai_params( + non_default_params={"max_retries": 0}, + optional_params={}, + model="cerebras/llama-4-scout-17b-16e", + drop_params=False, + ) + assert result["max_retries"] == 0 + + def test_map_openai_params_passes_through_extra_headers(self): + config = CerebrasConfig() + headers = {"X-Custom": "value"} + result = config.map_openai_params( + non_default_params={"extra_headers": headers}, + optional_params={}, + model="cerebras/llama-4-scout-17b-16e", + drop_params=False, + ) + assert result["extra_headers"] == headers