mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #39826 from BerriAI/litellm_lit_6348_fireworks_responses_api
feat(fireworks_ai): add native Responses API config
This commit is contained in:
commit
bb8fe4a32f
7 changed files with 604 additions and 16 deletions
|
|
@ -2002,6 +2002,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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
102
litellm/llms/fireworks_ai/responses/transformation.py
Normal file
102
litellm/llms/fireworks_ai/responses/transformation.py
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
import httpx
|
||||
from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam
|
||||
|
||||
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.responses.main import DeleteResponseResult
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
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")}
|
||||
)
|
||||
|
||||
|
||||
def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam:
|
||||
if "role" not in item or item["role"] != "developer":
|
||||
return item
|
||||
return EasyInputMessageParam(role="system", content=item["content"], type="message")
|
||||
|
||||
|
||||
def _developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam:
|
||||
if isinstance(input, str):
|
||||
return input
|
||||
return [_developer_item_as_system(item) for item in input]
|
||||
|
||||
|
||||
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 _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
|
||||
return _developer_items_as_system(super()._validate_input_param(input))
|
||||
|
||||
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 transform_delete_response_api_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
) -> DeleteResponseResult:
|
||||
deleted_id: Final = unquote(raw_response.request.url.path.rsplit("/", 1)[-1])
|
||||
return DeleteResponseResult(id=deleted_id, object="response", deleted=True)
|
||||
|
||||
def supports_native_websocket(self) -> bool:
|
||||
return False
|
||||
|
||||
def supports_native_file_search(self) -> bool:
|
||||
return False
|
||||
|
|
@ -8697,6 +8697,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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,455 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypedDict, cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai.types.responses import (
|
||||
EasyInputMessage,
|
||||
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.responses.file_search.emulated_handler import should_use_emulated_file_search
|
||||
from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponseInputParam, 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_developer_items_as_system_messages() -> 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/accounts/fireworks/models/kimi-k3",
|
||||
input=[ # mutable-ok: the Responses API takes input as a JSON list
|
||||
{"role": "user", "content": "Hi there"},
|
||||
{"role": "developer", "content": "Answer with exactly one word."},
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]},
|
||||
],
|
||||
api_key="fw-test-key",
|
||||
)
|
||||
_, _, body = _sent_request(client)
|
||||
assert tuple(body["input"]) == (
|
||||
{"role": "user", "content": "Hi there"},
|
||||
{"role": "system", "content": "Answer with exactly one word.", "type": "message"},
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]},
|
||||
)
|
||||
|
||||
|
||||
def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None:
|
||||
client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3"))
|
||||
pydantic_input: Final = cast(
|
||||
ResponseInputParam,
|
||||
[ # mutable-ok: the Responses API takes input as a JSON list
|
||||
EasyInputMessage(role="developer", content="Answer with exactly one word.", type="message"),
|
||||
ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"),
|
||||
ResponseFunctionToolCall(
|
||||
id="fc_1",
|
||||
call_id="call_abc123",
|
||||
name="get_weather",
|
||||
arguments='{"city": "Paris"}',
|
||||
status="completed",
|
||||
type="function_call",
|
||||
),
|
||||
FunctionCallOutput(type="function_call_output", call_id="call_abc123", output="21C"),
|
||||
],
|
||||
)
|
||||
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
|
||||
litellm.responses(
|
||||
model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key"
|
||||
)
|
||||
_, _, body = _sent_request(client)
|
||||
assert tuple(body["input"]) == (
|
||||
{"role": "system", "content": "Answer with exactly one word.", "type": "message"},
|
||||
{"id": "rs_1", "summary": [], "type": "reasoning"},
|
||||
{
|
||||
"id": "fc_1",
|
||||
"call_id": "call_abc123",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
"status": "completed",
|
||||
"type": "function_call",
|
||||
},
|
||||
{"type": "function_call_output", "call_id": "call_abc123", "output": "21C"},
|
||||
)
|
||||
|
||||
|
||||
def test_file_search_tools_take_litellm_emulated_search_not_fireworks() -> None:
|
||||
config: Final = FireworksAIResponsesAPIConfig()
|
||||
file_search: Final = ({"type": "file_search", "vector_store_ids": ("vs_kb",)},)
|
||||
function_tool: Final = ({"type": "function", "name": "get_weather", "parameters": {"type": "object"}},)
|
||||
assert should_use_emulated_file_search(tools=file_search, provider_config=config)
|
||||
assert not should_use_emulated_file_search(tools=function_tool, provider_config=config)
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
def test_delete_responses_maps_fireworks_message_only_body_to_deleted_result() -> None:
|
||||
response_id: Final = (
|
||||
"resp_xFaIJR9Nc_OXmqKRqL78UuAGj2Te5GY5BT_knpZiMrYoNOVmu5oc2mQW1HI7hCtEYB4mcx2lEYS0DYP1U5yEQskHunuB4=="
|
||||
)
|
||||
request: Final = httpx.Request("DELETE", f"{FIREWORKS_RESPONSES_URL}/{quote(response_id, safe='')}")
|
||||
client: Final = MagicMock()
|
||||
client.delete.return_value = httpx.Response(200, json={"message": "Response deleted successfully"}, request=request)
|
||||
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
|
||||
result: Final = litellm.delete_responses(
|
||||
response_id=response_id, custom_llm_provider="fireworks_ai", api_key="fw-test-key"
|
||||
)
|
||||
assert client.delete.call_args.kwargs["url"] == str(request.url)
|
||||
assert (result.id, result.object, result.deleted) == (response_id, "response", True)
|
||||
|
||||
|
||||
def _fireworks_stream_response(status: str, output: tuple[Mapping[str, object], ...]) -> Mapping[str, object]:
|
||||
return {
|
||||
"id": "resp_htnkJ8piNKeOHkn9LfAusC38O2OgcDQs4S8trSOJ6anLeqjUDGqu2PkWmg5N",
|
||||
"object": "response",
|
||||
"created_at": 1788567245,
|
||||
"model": "accounts/fireworks/models/kimi-k3",
|
||||
"status": status,
|
||||
"output": output,
|
||||
"usage": None
|
||||
if status == "in_progress"
|
||||
else {
|
||||
"input_tokens": 95,
|
||||
"output_tokens": 89,
|
||||
"total_tokens": 184,
|
||||
"input_tokens_details": {"cached_tokens": 94},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
FIREWORKS_SSE_EVENTS: Final[tuple[Mapping[str, object], ...]] = (
|
||||
{"type": "response.created", "sequence_number": 0, "response": _fireworks_stream_response("in_progress", ())},
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"sequence_number": 1,
|
||||
"output_index": 0,
|
||||
"item": {"id": "rs_1", "type": "reasoning", "summary": []},
|
||||
},
|
||||
{
|
||||
"type": "response.reasoning_summary_text.delta",
|
||||
"sequence_number": 2,
|
||||
"item_id": "rs_1",
|
||||
"output_index": 0,
|
||||
"summary_index": 0,
|
||||
"delta": "pong",
|
||||
},
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"sequence_number": 3,
|
||||
"output_index": 1,
|
||||
"item": {"id": "msg_1", "type": "message", "role": "assistant", "status": "in_progress", "content": []},
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 4,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 1,
|
||||
"content_index": 0,
|
||||
"delta": "po",
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 5,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 1,
|
||||
"content_index": 0,
|
||||
"delta": "ng",
|
||||
},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"sequence_number": 6,
|
||||
"response": _fireworks_stream_response(
|
||||
"completed",
|
||||
(
|
||||
{"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "pong"}]},
|
||||
{
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "pong", "annotations": []}],
|
||||
},
|
||||
),
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _sse_body(events: tuple[Mapping[str, object], ...]) -> bytes:
|
||||
return b"".join(f"data: {json.dumps(dict(event))}\n\n".encode() for event in events) + b"data: [DONE]\n\n"
|
||||
|
||||
|
||||
def test_streaming_responses_call_hits_native_endpoint_and_yields_every_fireworks_event() -> None:
|
||||
request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL)
|
||||
client: Final = MagicMock()
|
||||
client.post.return_value = httpx.Response(
|
||||
200, content=_sse_body(FIREWORKS_SSE_EVENTS), headers={"content-type": "text/event-stream"}, request=request
|
||||
)
|
||||
with patch(HTTPX_CLIENT_FACTORY, return_value=client):
|
||||
received: Final = tuple(
|
||||
litellm.responses(
|
||||
model="fireworks_ai/kimi-k3",
|
||||
input="Reply with the single word pong.",
|
||||
stream=True,
|
||||
api_key="fw-test-key",
|
||||
)
|
||||
)
|
||||
url, _, body = _sent_request(client)
|
||||
assert (url, body["model"], body["stream"], client.post.call_args.kwargs["stream"]) == (
|
||||
FIREWORKS_RESPONSES_URL,
|
||||
"accounts/fireworks/models/kimi-k3",
|
||||
True,
|
||||
True,
|
||||
)
|
||||
assert tuple(event.type for event in received) == tuple(event["type"] for event in FIREWORKS_SSE_EVENTS)
|
||||
assert "".join(event.delta for event in received if event.type == "response.output_text.delta") == "pong"
|
||||
assert received[-1].response.usage.output_tokens == 89
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue