From c74f8bddf33bf5a8ee7599b4d276994203d5eab4 Mon Sep 17 00:00:00 2001
From: Ryan Loney <ryan.loney@cerebras.net>
Date: Thu, 20 Aug 2026 13:50:07 -0700
Subject: [PATCH] fix(cerebras): strengthen request parameter types

---
 litellm/llms/cerebras/chat.py          | 12 ++++++++----
 tests/llm_translation/test_cerebras.py |  4 ++--
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py
index 062955211aa..d3dd426173a 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
 
@@ -32,10 +36,10 @@ class CerebrasConfig(OpenAIGPTConfig):
     top_logprobs: int | None = None
     frequency_penalty: float | None = None
     presence_penalty: float | None = None
-    logit_bias: dict | None = None
+    logit_bias: dict[str, float] | None = None
     service_tier: str | None = None
     prompt_cache_key: str | None = None
-    prediction: dict | None = None
+    prediction: ChatCompletionPredictionContentParam | None = None
 
     def __init__(
         self,
@@ -56,10 +60,10 @@ class CerebrasConfig(OpenAIGPTConfig):
         top_logprobs: int | None = None,
         frequency_penalty: float | None = None,
         presence_penalty: float | None = None,
-        logit_bias: dict | None = None,
+        logit_bias: dict[str, float] | None = None,
         service_tier: str | None = None,
         prompt_cache_key: str | None = None,
-        prediction: dict | None = None,
+        prediction: ChatCompletionPredictionContentParam | None = None,
     ) -> None:
         locals_: Final = locals().copy()
         for key, value in locals_.items():
diff --git a/tests/llm_translation/test_cerebras.py b/tests/llm_translation/test_cerebras.py
index 15a500fb1e6..82f23dff25a 100644
--- a/tests/llm_translation/test_cerebras.py
+++ b/tests/llm_translation/test_cerebras.py
@@ -19,7 +19,7 @@ from litellm.llms.cerebras.chat import CerebrasConfig
         ("prediction", {"type": "content", "content": "expected"}),
     ],
 )
-def test_cerebras_preserves_supported_parameters(parameter, value):
+def test_cerebras_preserves_supported_parameters(parameter: str, value: object) -> None:
     config = CerebrasConfig()
 
     mapped = config.map_openai_params(
@@ -32,7 +32,7 @@ def test_cerebras_preserves_supported_parameters(parameter, value):
     assert mapped == {parameter: value}
 
 
-def test_cerebras_does_not_alias_max_completion_tokens():
+def test_cerebras_does_not_alias_max_completion_tokens() -> None:
     config = CerebrasConfig()
 
     mapped = config.map_openai_params(