diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 319f03fea89..eeae8c76888 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -133,6 +133,32 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): def get_config(cls): return super().get_config() + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + api_key = self._get_api_key(api_key) + if api_key is None: + raise ValueError("FIREWORKS_API_KEY is not set") + + validated_headers = OpenAIGPTConfig.validate_environment( + self, + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + return self._add_session_affinity_header(validated_headers, litellm_params) + def get_supported_openai_params(self, model: str): # Base parameters supported by all models supported_params = [ diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 4e22445bcc0..51ed8afbbd2 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -64,9 +64,16 @@ class FireworksAIMixin: if api_key is None: raise ValueError("FIREWORKS_API_KEY is not set") - validated_headers = {"Authorization": "Bearer {}".format(api_key), **headers} - if not any(key.lower() == "x-session-affinity" for key in validated_headers): - session_id = get_fireworks_session_id(litellm_params) - if session_id: - validated_headers["x-session-affinity"] = session_id - return validated_headers + auth_headers = {"Authorization": "Bearer {}".format(api_key), **headers} + content_type_header = ( + {} if any(key.lower() == "content-type" for key in auth_headers) else {"Content-Type": "application/json"} + ) + return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params) + + def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict: + if any(key.lower() == "x-session-affinity" for key in headers): + return headers + session_id = get_fireworks_session_id(litellm_params) + if not session_id: + return headers + return {**headers, "x-session-affinity": session_id} diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 6809799d34f..94945ed4bfb 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -123,6 +123,90 @@ def test_validate_environment_preserves_explicit_session_affinity_header(): assert headers["x-session-affinity"] == "explicit-session" +def test_validate_environment_sets_json_content_type(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_preserves_explicit_content_type(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={"content-type": "multipart/form-data"}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert headers["content-type"] == "multipart/form-data" + assert "Content-Type" not in headers + + +def test_validate_environment_sets_json_content_type_with_session_affinity(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_session_id": "session-123"}, + api_key="test-key", + ) + + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + assert headers["x-session-affinity"] == "session-123" + + +def test_validate_environment_resolves_api_key_from_env_and_sets_content_type(monkeypatch): + monkeypatch.setenv("FIREWORKS_API_KEY", "fw-env-key") + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert headers["Authorization"] == "Bearer fw-env-key" + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_raises_without_api_key(monkeypatch): + for env_var in ( + "FIREWORKS_API_KEY", + "FIREWORKS_AI_API_KEY", + "FIREWORKSAI_API_KEY", + "FIREWORKS_AI_TOKEN", + ): + monkeypatch.delenv(env_var, raising=False) + config = FireworksAIConfig() + + with pytest.raises(ValueError, match="FIREWORKS_API_KEY is not set"): + config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + ) + + def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): assert ( get_fireworks_session_id(