From 4b950cd94f86dbe174d515a468cb9b5288bf0ab8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:01:14 -0700 Subject: [PATCH] feat(fireworks_ai): add native Responses API config --- basedpyright-code-budget.json | 4 +- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + litellm/llms/fireworks_ai/common_utils.py | 44 +-- .../fireworks_ai/responses/transformation.py | 68 +++++ litellm/utils.py | 2 + ...t_fireworks_ai_responses_transformation.py | 263 ++++++++++++++++++ .../test_responses_websocket_all_providers.py | 9 + type-discipline-budget.json | 4 +- 9 files changed, 382 insertions(+), 20 deletions(-) create mode 100644 litellm/llms/fireworks_ai/responses/transformation.py create mode 100644 tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 9b59480a0dc..e7f8aacf835 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38309 + "limit": 38307 }, "reportUnknownParameterType": { "limit": 19622 }, "reportUnknownVariableType": { - "limit": 29846 + "limit": 29840 }, "reportUnnecessaryCast": { "limit": 111 diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..fdc4435e5ff 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -2001,6 +2001,9 @@ if TYPE_CHECKING: from .llms.hosted_vllm.responses.transformation import ( HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig, ) + from .llms.fireworks_ai.responses.transformation import ( + FireworksAIResponsesAPIConfig as FireworksAIResponsesAPIConfig, + ) from .llms.github_copilot.chat.transformation import ( GithubCopilotConfig as GithubCopilotConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e9199e1ec80..dc323c8cc15 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -237,6 +237,7 @@ LLM_CONFIG_NAMES: Final = ( "XAIResponsesAPIConfig", "LiteLLMProxyResponsesAPIConfig", "HostedVLLMResponsesAPIConfig", + "FireworksAIResponsesAPIConfig", "VolcEngineResponsesAPIConfig", "PerplexityResponsesConfig", "DatabricksResponsesAPIConfig", @@ -957,6 +958,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.hosted_vllm.responses.transformation", "HostedVLLMResponsesAPIConfig", ), + "FireworksAIResponsesAPIConfig": ( + ".llms.fireworks_ai.responses.transformation", + "FireworksAIResponsesAPIConfig", + ), "VolcEngineResponsesAPIConfig": ( ".llms.volcengine.responses.transformation", "VolcEngineResponsesAPIConfig", diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index ac934ad0cb5..21a630a76d7 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from httpx import Headers @@ -13,7 +15,7 @@ class FireworksAIException(BaseLLMException): pass -def get_fireworks_session_id(litellm_params: dict) -> str | None: +def get_fireworks_session_id(litellm_params: Mapping[str, object]) -> str | None: """ Session id to send as `x-session-affinity`, or None when the caller gave none. @@ -23,19 +25,39 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: """ params: Final = litellm_params metadata: Final = params.get("metadata") - if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + if isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None for key in ("litellm_session_id", "session_id"): value = params.get(key) if value: return str(value) - if isinstance(metadata, dict): + if isinstance(metadata, Mapping): value = metadata.get("session_id") if value: return str(value) return None +def with_fireworks_session_affinity( + headers: Mapping[str, str], litellm_params: Mapping[str, object] +) -> Mapping[str, str]: + if any(key.lower() == "x-session-affinity" for key in headers): + return headers + session_id: Final = get_fireworks_session_id(litellm_params) + if not session_id: + return headers + return MappingProxyType({**headers, "x-session-affinity": session_id}) + + +def resolve_fireworks_api_key(api_key: str | None) -> str | None: + return api_key or ( + get_secret_str("FIREWORKS_API_KEY") + or get_secret_str("FIREWORKS_AI_API_KEY") + or get_secret_str("FIREWORKSAI_API_KEY") + or get_secret_str("FIREWORKS_AI_TOKEN") + ) + + AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-" @@ -63,13 +85,7 @@ class FireworksAIMixin: ) def _get_api_key(self, api_key: str | None) -> str | None: - dynamic_api_key: Final = api_key or ( - get_secret_str("FIREWORKS_API_KEY") - or get_secret_str("FIREWORKS_AI_API_KEY") - or get_secret_str("FIREWORKSAI_API_KEY") - or get_secret_str("FIREWORKS_AI_TOKEN") - ) - return dynamic_api_key + return resolve_fireworks_api_key(api_key) def validate_environment( self, @@ -92,9 +108,5 @@ class FireworksAIMixin: 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: Final = get_fireworks_session_id(litellm_params) - if not session_id: - return headers - return {**headers, "x-session-affinity": session_id} + pinned: Final = with_fireworks_session_affinity(headers, litellm_params) + return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py new file mode 100644 index 00000000000..f36030bb50a --- /dev/null +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -0,0 +1,68 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.llms.fireworks_ai.common_utils import ( + resolve_fireworks_api_key, + resolve_fireworks_resource_name, + with_fireworks_session_affinity, +) +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ResponseInputParam +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +FIREWORKS_AI_DEFAULT_API_BASE: Final = "https://api.fireworks.ai/inference/v1" + + +def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object]: + extras: Final[Mapping[str, object]] = litellm_params.model_extra or MappingProxyType({}) + return MappingProxyType( + {"litellm_session_id": extras.get("litellm_session_id"), "metadata": extras.get("litellm_metadata")} + ) + + +class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.FIREWORKS_AI + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict: # mutable-ok: overrides the base class signature + params: Final = litellm_params or GenericLiteLLMParams() + api_key: Final = resolve_fireworks_api_key(params.api_key) + if api_key is None: + raise ValueError("FIREWORKS_API_KEY is not set") + authorized: Final = MappingProxyType( + {"Content-Type": "application/json", **headers, "Authorization": f"Bearer {api_key}"} + ) + pinned: Final = with_fireworks_session_affinity(authorized, _session_params(params)) + return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place + + def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/") + return f"{base}/responses" + + def transform_responses_api_request( + self, + model: str, + input: str | ResponseInputParam, + response_api_optional_request_params: dict, # mutable-ok: overrides the base class signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: overrides the base class signature + ) -> dict: # mutable-ok: overrides the base class signature + return super().transform_responses_api_request( + model=resolve_fireworks_resource_name(model), + input=input, + response_api_optional_request_params=response_api_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + + def supports_native_websocket(self) -> bool: + return False diff --git a/litellm/utils.py b/litellm/utils.py index 8b1b32ea328..3043a502aaa 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8683,6 +8683,8 @@ class ProviderConfigManager: return litellm.OpenRouterResponsesAPIConfig() elif litellm.LlmProviders.HOSTED_VLLM == provider: return litellm.HostedVLLMResponsesAPIConfig() + elif litellm.LlmProviders.FIREWORKS_AI == provider: + return litellm.FireworksAIResponsesAPIConfig() elif litellm.LlmProviders.BEDROCK_MANTLE == provider: # Both decisions are data-driven from the model's price-map entry, with # no model-name logic. Capability (can it serve Responses?) comes from diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py new file mode 100644 index 00000000000..207fd518f89 --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -0,0 +1,263 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, TypedDict +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseOutputText, + ResponseReasoningItem, +) +from openai.types.responses.response_input_param import FunctionCallOutput +from openai.types.responses.tool_param import Mcp +from typing_extensions import ReadOnly + +import litellm +from litellm.llms.fireworks_ai.responses.transformation import FireworksAIResponsesAPIConfig +from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +FIREWORKS_RESPONSES_URL: Final = "https://api.fireworks.ai/inference/v1/responses" +HTTPX_CLIENT_FACTORY: Final = "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client" +NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) +NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _GeneratedSessionMetadata(TypedDict): + litellm_session_id_generated: ReadOnly[bool] + + +def _fireworks_response(model: str) -> Mapping[str, object]: + return ResponsesAPIResponse( + id="resp_0e946f2d46bf4b49bf8b29ff78083583", + object="response", + created_at=1788550000, + model=model, + status="completed", + output=( + ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"), + ResponseOutputMessage( + id="msg_1", + status="completed", + role="assistant", + type="message", + content=(ResponseOutputText(type="output_text", text="Paris is clear and 21C.", annotations=()),), + ), + ResponseFunctionToolCall( + id="fc_1", + call_id="call_abc123", + name="get_weather", + arguments='{"city": "Paris"}', + status="completed", + type="function_call", + ), + ), + usage=ResponseAPIUsage( + input_tokens=179, + output_tokens=100, + total_tokens=279, + input_tokens_details=InputTokensDetails(cached_tokens=0), + ), + ).model_dump(mode="json", exclude_none=True) + + +def _mock_http_client(response_body: Mapping[str, object]) -> MagicMock: + client: Final = MagicMock() + response: Final = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers((("content-type", "application/json"),)) + response.json.return_value = response_body + response.text = json.dumps(response_body) + client.post.return_value = response + return client + + +def _sent_request(client: MagicMock) -> tuple[str, Mapping[str, str], Mapping[str, object]]: + kwargs: Final = client.post.call_args.kwargs + body: Final = kwargs["json"] if "json" in kwargs else json.loads(kwargs["data"]) + return kwargs["url"], kwargs["headers"], body + + +@pytest.fixture(autouse=True) +def fireworks_env(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "FIREWORKS_API_KEY", + "FIREWORKS_AI_API_KEY", + "FIREWORKSAI_API_KEY", + "FIREWORKS_AI_TOKEN", + "FIREWORKS_API_BASE", + ): + monkeypatch.delenv(name, raising=False) + + +def test_fireworks_ai_provider_config_registration() -> None: + config: Final = ProviderConfigManager.get_provider_responses_api_config( + model="accounts/fireworks/models/kimi-k3", provider=LlmProviders.FIREWORKS_AI + ) + assert isinstance(config, FireworksAIResponsesAPIConfig) + assert config.custom_llm_provider == LlmProviders.FIREWORKS_AI + + +def test_responses_call_hits_native_endpoint_with_mcp_tool_untouched() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + mcp_tool: Final[Mcp] = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + response: Final = litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input="What is litellm?", + tools=[mcp_tool], # mutable-ok: the Responses API takes tools as a JSON list + api_key="fw-test-key", + ) + url, headers, body = _sent_request(client) + assert url == FIREWORKS_RESPONSES_URL + assert headers["Authorization"] == "Bearer fw-test-key" + assert body["model"] == "accounts/fireworks/models/kimi-k3" + assert tuple(body["tools"]) == (mcp_tool,) + assert "messages" not in body + assert isinstance(response, ResponsesAPIResponse) + function_calls: Final = tuple(item for item in response.output if getattr(item, "type", None) == "function_call") + assert getattr(function_calls[0], "call_id", None) == "call_abc123" + + +def test_responses_call_expands_bare_model_name_to_fireworks_resource() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/glm-5p3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/glm-5p3", input="hi", api_key="fw-test-key") + _, _, body = _sent_request(client) + assert body["model"] == "accounts/fireworks/models/glm-5p3" + + +def test_responses_call_forwards_previous_response_id_and_store() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + tool_output: Final[FunctionCallOutput] = { + "type": "function_call_output", + "call_id": "call_abc123", + "output": "{}", + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input=[tool_output], # mutable-ok: the Responses API takes input items as a JSON list + previous_response_id="resp_0e946f2d46bf4b49bf8b29ff78083583", + store=True, + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["previous_response_id"] == "resp_0e946f2d46bf4b49bf8b29ff78083583" + assert body["store"] is True + assert body["input"][0]["call_id"] == "call_abc123" + + +def test_responses_call_sends_session_affinity_for_caller_session_id() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key", litellm_session_id="sess-42") + _, headers, _ = _sent_request(client) + assert headers["x-session-affinity"] == "sess-42" + + +def test_responses_call_keeps_caller_supplied_session_affinity_header() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + pinned: Final[Mapping[str, str]] = MappingProxyType({"x-session-affinity": "explicit-node"}) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input="hi", + api_key="fw-test-key", + litellm_session_id="sess-42", + extra_headers=pinned, + ) + _, headers, _ = _sent_request(client) + assert headers["x-session-affinity"] == "explicit-node" + + +def test_responses_call_maps_provider_errors_to_fireworks_ai() -> None: + client: Final = MagicMock() + request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) + client.post.side_effect = httpx.HTTPStatusError( + "unauthorized", + request=request, + response=httpx.Response(401, text='{"error": {"message": "invalid api key"}}', request=request), + ) + with patch(HTTPX_CLIENT_FACTORY, return_value=client), pytest.raises(litellm.AuthenticationError) as raised: + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-bad-key") + assert raised.value.llm_provider == "fireworks_ai" + assert raised.value.status_code == 401 + assert "invalid api key" in str(raised.value) + + +def test_responses_call_skips_session_affinity_for_proxy_generated_session_id() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + generated: Final[_GeneratedSessionMetadata] = {"litellm_session_id_generated": True} + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input="hi", + api_key="fw-test-key", + litellm_session_id="generated-1", + litellm_metadata=generated, + ) + _, headers, _ = _sent_request(client) + assert "x-session-affinity" not in headers + + +@pytest.mark.parametrize( + "api_base, expected", + ( + (None, FIREWORKS_RESPONSES_URL), + ("https://api.fireworks.ai/inference/v1", FIREWORKS_RESPONSES_URL), + ("https://api.fireworks.ai/inference/v1/", FIREWORKS_RESPONSES_URL), + ("https://gateway.example.com/fireworks", "https://gateway.example.com/fireworks/responses"), + ), +) +def test_get_complete_url(api_base: str | None, expected: str) -> None: + assert FireworksAIResponsesAPIConfig().get_complete_url(api_base=api_base, litellm_params=NO_PARAMS) == expected + + +def test_responses_call_reads_fireworks_api_base_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FIREWORKS_API_BASE", "https://self-hosted.example.com/v1") + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key") + url, _, _ = _sent_request(client) + assert url == "https://self-hosted.example.com/v1/responses" + + +@pytest.mark.parametrize( + "env_name", ("FIREWORKS_API_KEY", "FIREWORKS_AI_API_KEY", "FIREWORKSAI_API_KEY", "FIREWORKS_AI_TOKEN") +) +def test_validate_environment_reads_every_fireworks_key_name(monkeypatch: pytest.MonkeyPatch, env_name: str) -> None: + monkeypatch.setenv(env_name, "env-key") + headers: Final = FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=GenericLiteLLMParams() + ) + assert headers["Authorization"] == "Bearer env-key" + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_prefers_explicit_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FIREWORKS_API_KEY", "env-key") + headers: Final = FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, + model="accounts/fireworks/models/kimi-k3", + litellm_params=GenericLiteLLMParams(api_key="explicit"), + ) + assert headers["Authorization"] == "Bearer explicit" + + +def test_validate_environment_without_any_key_raises() -> None: + with pytest.raises(ValueError, match="FIREWORKS_API_KEY"): + FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=None + ) diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index e8333214ea8..fe3c4a0640d 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -17,6 +17,9 @@ from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPICon from litellm.llms.databricks.responses.transformation import ( DatabricksResponsesAPIConfig, ) +from litellm.llms.fireworks_ai.responses.transformation import ( + FireworksAIResponsesAPIConfig, +) from litellm.llms.github_copilot.responses.transformation import ( GithubCopilotResponsesAPIConfig, ) @@ -102,6 +105,12 @@ class TestResponsesAPIWebSocketSupport: def test_openai_model_in_websocket_url_default(self): assert OpenAIResponsesAPIConfig().model_in_websocket_url() is True + def test_fireworks_ai_uses_managed_websocket(self): + """Fireworks AI should use managed websocket handler""" + assert ( + FireworksAIResponsesAPIConfig().supports_native_websocket() is False + ), "Fireworks AI should use managed websocket handler" + def test_xai_uses_managed_websocket(self): """XAI should use managed websocket handler""" config = XAIResponsesAPIConfig() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8589a9451cf..835b24b2b3c 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22328 + "limit": 22327 }, "LIT002": { - "limit": 26748 + "limit": 26747 }, "LIT003": { "limit": 261