feat(fireworks_ai): add native Responses API config

This commit is contained in:
mateo-berri 2026-09-04 17:01:14 -07:00
parent d23bec84c4
commit 4b950cd94f
9 changed files with 382 additions and 20 deletions

View file

@ -105,13 +105,13 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38309
"limit": 38307
},
"reportUnknownParameterType": {
"limit": 19622
},
"reportUnknownVariableType": {
"limit": 29846
"limit": 29840
},
"reportUnnecessaryCast": {
"limit": 111

View file

@ -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,
)

View file

@ -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",

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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
)

View file

@ -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()

View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22328
"limit": 22327
},
"LIT002": {
"limit": 26748
"limit": 26747
},
"LIT003": {
"limit": 261