scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix(openrouter): track Responses API cost

Browse source
This commit is contained in:
Naineel Soyantar 2026-09-03 06:34:32 +00:00
parent 3cac5e5cd4
commit e720ffa61b
11 changed files with 284 additions and 130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -130,6 +130,7 @@ from litellm.types.utils import (
EmbeddingResponse,
FileTypes,
LiteLLMBatch,
LlmProviders,
TranscriptionResponse,
)
from litellm.types.vector_store_files import (
@ -242,6 +243,19 @@ def _responses_api_optional_request_param_names() -> frozenset[str]:
return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys())
def _ensure_openrouter_responses_usage(
data: dict, # mutable-ok: request payload accepts extra_body fields
responses_api_provider_config: BaseResponsesAPIConfig,
) -> None:
"""Require OpenRouter usage data after callers merge an extra request body."""
if responses_api_provider_config.custom_llm_provider != LlmProviders.OPENROUTER:
return
usage: Final = data.get("usage")
normalized_usage = {**usage, "include": True} if isinstance(usage, Mapping) else {"include": True}
data["usage"] = normalized_usage # rebind-ok: OpenRouter usage must override extra_body
def _custom_logger_callbacks(logging_obj: LiteLLMLoggingObj) -> list["CustomLogger"]:
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import (
@ -2661,6 +2675,7 @@ class BaseLLMHTTPHandler:
if extra_body:
data.update(extra_body)
_ensure_openrouter_responses_usage(data, responses_api_provider_config)
stream = bool(stream or data.get("stream"))
# Preserve the OpenAI-style request context (not sent to the provider) for streaming
@ -2839,6 +2854,7 @@ class BaseLLMHTTPHandler:
if extra_body:
data.update(extra_body)
_ensure_openrouter_responses_usage(data, responses_api_provider_config)
stream = bool(stream or data.get("stream"))
# Preserve the OpenAI-style request context (not sent to the provider) for streaming

62
litellm/llms/openrouter/responses/transformation.py
View file

@ -8,14 +8,26 @@ encrypted_content for multi-turn stateless workflows.
Docs: https://openrouter.ai/docs/api/reference/responses/overview
"""
from typing import Final
from collections.abc import Mapping
from json import JSONDecodeError
from typing import TYPE_CHECKING, Final
import httpx
import litellm
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIResponse
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
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = object
class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig):
"""
@ -75,3 +87,51 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig):
def supports_native_websocket(self) -> bool:
"""OpenRouter does not support native WebSocket for Responses API"""
return False
def transform_responses_api_request(
self,
model: str,
input: str | ResponseInputParam,
response_api_optional_request_params: dict, # mutable-ok: base method requires mutable payload
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: base method requires mutable payload
) -> dict: # mutable-ok: base method returns mutable JSON payload
"""Request usage details so OpenRouter returns the response cost."""
request: Final = super().transform_responses_api_request(
model=model,
input=input,
response_api_optional_request_params=response_api_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)
usage: Final = request.get("usage")
request["usage"] = {**usage, "include": True} if isinstance(usage, Mapping) else {"include": True}
return request
def transform_response_api_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> ResponsesAPIResponse:
"""Store the OpenRouter response cost for LiteLLM cost calculation."""
response: Final = super().transform_response_api_response(
model=model,
raw_response=raw_response,
logging_obj=logging_obj,
)
try:
response_json: Final = raw_response.json()
except JSONDecodeError:
return response
if not isinstance(response_json, Mapping):
return response
usage: Final = response_json.get("usage")
cost: Final = usage.get("cost") if isinstance(usage, Mapping) else None
try:
cost_value: Final = float(cost) if cost is not None else None
except (TypeError, ValueError):
return response
if cost_value is not None:
response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = cost_value
return response

2
litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py
View file

@ -471,7 +471,7 @@ def _open_claims(
def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int:
if upstream_expires_in is None:
return MAX_ENVELOPE_TTL_SECONDS
return upstream_expires_in
return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS)
def _refresh_ttl_seconds(upstream_refresh_expires_in: int | None) -> int:

13
litellm/proxy/search_endpoints/search_tool_management.py
View file

@ -17,6 +17,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
from litellm.proxy.search_endpoints.search_tool_registry import SearchToolRegistry
from litellm.types.search import (
ListSearchToolsResponse,
@ -31,6 +32,12 @@ router: Final = APIRouter()
SEARCH_TOOL_REGISTRY: Final = SearchToolRegistry()
async def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)) -> None:
"""Allow search tool changes only for proxy administrators."""
if not is_proxy_admin(user_api_key_dict):
raise HTTPException(status_code=403, detail="Only proxy administrators can manage search tools")
def _convert_datetime_to_str(value: datetime | str | None) -> str | None:
"""
Convert datetime object to ISO format string.
@ -267,7 +274,7 @@ class CreateSearchToolRequest(BaseModel):
@router.post(
"/search_tools",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
dependencies=[Depends(_require_proxy_admin)],
)
async def create_search_tool(request: CreateSearchToolRequest):
"""
@ -339,7 +346,7 @@ class UpdateSearchToolRequest(BaseModel):
@router.put(
"/search_tools/{search_tool_id}",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
dependencies=[Depends(_require_proxy_admin)],
)
async def update_search_tool(search_tool_id: str, request: UpdateSearchToolRequest):
"""
@ -422,7 +429,7 @@ async def update_search_tool(search_tool_id: str, request: UpdateSearchToolReque
@router.delete(
"/search_tools/{search_tool_id}",
tags=["Search Tools"],
dependencies=[Depends(user_api_key_auth)],
dependencies=[Depends(_require_proxy_admin)],
)
async def delete_search_tool(search_tool_id: str):
"""

6
litellm/responses/main.py
View file

@ -632,7 +632,11 @@ def _resolve_prompt_swapped_provider(
prompt_id: str | None,
) -> str:
swapped_provider: Final = litellm.get_llm_provider(model=swapped_model)[1]
if kwargs.get("api_key") is None and kwargs.get("api_base") is None:
extra_headers: Final = kwargs.get("extra_headers")
has_authorization_header: Final = isinstance(extra_headers, Mapping) and any(
str(header_name).lower() == "authorization" for header_name in extra_headers
)
if kwargs.get("api_key") is None and kwargs.get("api_base") is None and not has_authorization_header:
return swapped_provider
try:
original_provider: Final = custom_llm_provider or litellm.get_llm_provider(model=original_model)[1]

1
tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
View file

@ -376,7 +376,6 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
timeout=None,
client=None,
_is_async=False,
embedding_executor=None,
):
litellm_params_dict = (
litellm_params.model_dump(exclude_none=False)

72
tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py
View file

@ -9,6 +9,9 @@ reasoning.encrypted_content for multi-turn stateless workflows.
Related issue: https://github.com/BerriAI/litellm/issues/22189
"""
from unittest.mock import MagicMock
import httpx
import pytest
import litellm
@ -19,6 +22,21 @@ from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
def _successful_response_body(cost: float | None = None) -> dict:
usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
if cost is not None:
usage["cost"] = cost
return {
"id": "resp_123",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "openai/gpt-4o-mini",
"output": [],
"usage": usage,
}
class TestOpenRouterResponsesAPIConfig:
"""Test OpenRouter Responses API configuration."""
@ -57,9 +75,7 @@ class TestOpenRouterResponsesAPIConfig:
from litellm.types.router import GenericLiteLLMParams
params = GenericLiteLLMParams(api_key="sk-or-test-key")
headers = config.validate_environment(
headers={}, model="openai/o4-mini", litellm_params=params
)
headers = config.validate_environment(headers={}, model="openai/o4-mini", litellm_params=params)
assert headers["Authorization"] == "Bearer sk-or-test-key"
def test_validate_environment_raises_without_key(self, monkeypatch):
@ -81,6 +97,53 @@ class TestOpenRouterResponsesAPIConfig:
e = exc_info.value
assert "OpenRouter API key is required" in str(e)
def test_transform_request_includes_usage(self):
"""OpenRouter must return usage details for response cost accounting."""
from litellm.types.router import GenericLiteLLMParams
request = OpenRouterResponsesAPIConfig().transform_responses_api_request(
model="openai/gpt-4o-mini",
input="hello",
response_api_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert request["usage"] == {"include": True}
def test_transform_response_stores_usage_cost(self):
"""OpenRouter usage.cost must reach the LiteLLM cost calculator."""
raw_response = httpx.Response(200, json=_successful_response_body(cost=0.0125))
response = OpenRouterResponsesAPIConfig().transform_response_api_response(
model="openai/gpt-4o-mini",
raw_response=raw_response,
logging_obj=MagicMock(),
)
assert response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.0125
def test_transform_response_allows_missing_usage_cost(self):
"""A response without cost data must remain usable."""
raw_response = httpx.Response(200, json=_successful_response_body())
response = OpenRouterResponsesAPIConfig().transform_response_api_response(
model="openai/gpt-4o-mini",
raw_response=raw_response,
logging_obj=MagicMock(),
)
assert "llm_provider-x-litellm-response-cost" not in response._hidden_params["additional_headers"]
def test_extra_body_usage_cannot_disable_openrouter_cost_tracking(self):
"""The final request merge must force usage.include to true."""
from litellm.llms.custom_httpx.llm_http_handler import _ensure_openrouter_responses_usage
request = {"usage": {"include": False, "detail": "all"}}
_ensure_openrouter_responses_usage(request, OpenRouterResponsesAPIConfig())
assert request["usage"] == {"include": True, "detail": "all"}
class TestOpenRouterResponsesAPIRegistration:
"""Test that OpenRouter is properly registered as a native Responses API provider."""
@ -97,8 +160,7 @@ class TestOpenRouterResponsesAPIRegistration:
provider=LlmProviders.OPENROUTER,
)
assert config is not None, (
"OpenRouter must be registered as a native Responses API provider "
"to preserve reasoning.encrypted_content"
"OpenRouter must be registered as a native Responses API provider to preserve reasoning.encrypted_content"
)
assert isinstance(config, OpenRouterResponsesAPIConfig)

16
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py
View file

@ -30,7 +30,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import
DecryptFailed,
EnvelopeIdentity,
EnvelopeKeys,
EnvelopeLifetimeUnrepresentable,
EnvelopeMintError,
EnvelopeTooLarge,
Expired,
@ -161,7 +160,7 @@ def test_claim_layout_and_no_plaintext_token_in_envelope():
assert _REFRESH_TOKEN not in json.dumps(claims)
def test_unrepresentable_access_lifetime_is_a_typed_mint_error():
def test_large_access_lifetime_is_capped():
grant = UpstreamTokenGrant(
access_token=SecretStr(_ACCESS_TOKEN),
token_type="Bearer",
@ -170,9 +169,8 @@ def test_unrepresentable_access_lifetime_is_a_typed_mint_error():
result = mint_envelope(_IDENTITY, grant, _KEYS, _NOW)
assert isinstance(result, EnvelopeLifetimeUnrepresentable)
assert result.tag == "envelope_lifetime_unrepresentable"
assert result.expires_in == 10**30
assert isinstance(result, SealedEnvelope)
assert result.expires_at == _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS)
def _refresh_credential() -> RefreshCredential:
@ -259,7 +257,7 @@ def test_refresh_envelope_never_leaks_the_refresh_token_in_plaintext():
"expires_in, expected_ttl",
[
(600, 600),
(MAX_ENVELOPE_TTL_SECONDS + 82800, MAX_ENVELOPE_TTL_SECONDS + 82800),
(MAX_ENVELOPE_TTL_SECONDS + 82800, MAX_ENVELOPE_TTL_SECONDS),
(None, MAX_ENVELOPE_TTL_SECONDS),
],
)
@ -277,11 +275,11 @@ def test_expiry_honored_against_injected_clock():
assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=601)), Expired)
def test_upstream_token_lifetime_is_enforced_on_open():
def test_envelope_lifetime_is_capped_on_open():
grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=86400)
token = _sealed_token(grant)
just_before_expiry = _NOW + timedelta(seconds=86399)
at_expiry = _NOW + timedelta(seconds=86400)
just_before_expiry = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS - 1)
at_expiry = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS)
assert isinstance(open_envelope(token, _KEYS, just_before_expiry), OpenedEnvelope)
assert isinstance(open_envelope(token, _KEYS, at_expiry), Expired)

11
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -5690,7 +5690,7 @@ async def test_bridge_envelope_too_large_upstream_token_is_502():
@pytest.mark.asyncio
async def test_bridge_envelope_unrepresentable_upstream_lifetime_is_502():
async def test_bridge_envelope_large_upstream_lifetime_is_capped():
from litellm.types.mcp import MCPAuth
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
@ -5706,11 +5706,10 @@ async def test_bridge_envelope_unrepresentable_upstream_lifetime_is_502():
key_hash="hashed-litellm-key-77",
)
assert response.status_code == 502
assert json.loads(response.body) == {
"error": "server_error",
"error_description": "the upstream token response reports an unrepresentable lifetime",
}
assert response.status_code == 200
body = json.loads(response.body)
assert body["access_token"].startswith("llm_env_")
assert body["expires_in"] <= 3600
@pytest.mark.asyncio

171
tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py
View file

@ -114,9 +114,7 @@ async def test_list_search_tools_config_only(monkeypatch):
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
# Mock proxy_config
mock_proxy_config = MagicMock()
mock_proxy_config.get_config = AsyncMock(
return_value={"search_tools": config_tools}
)
mock_proxy_config.get_config = AsyncMock(return_value={"search_tools": config_tools})
mock_proxy_config.parse_search_tools = MagicMock(return_value=config_tools)
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
# Mock auth
@ -190,9 +188,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch):
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
# Mock proxy_config
mock_proxy_config = MagicMock()
mock_proxy_config.get_config = AsyncMock(
return_value={"search_tools": config_tools}
)
mock_proxy_config.get_config = AsyncMock(return_value={"search_tools": config_tools})
mock_proxy_config.parse_search_tools = MagicMock(return_value=config_tools)
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
# Mock auth
@ -213,11 +209,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch):
# Verify DB tool is present
db_tool = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "existing-tool"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "existing-tool"),
None,
)
assert db_tool is not None
@ -230,11 +222,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch):
# Verify unique config tool is present
config_tool = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "unique-config-tool"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "unique-config-tool"),
None,
)
assert config_tool is not None
@ -245,8 +233,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch):
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "existing-tool"
and t["is_from_config"] is True
if t["search_tool_name"] == "existing-tool" and t["is_from_config"] is True
),
None,
)
@ -321,11 +308,7 @@ async def test_list_search_tools_datetime_conversion(monkeypatch):
# Test datetime conversion for tool 1
tool1 = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "datetime-test-tool"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "datetime-test-tool"),
None,
)
assert tool1 is not None
@ -339,11 +322,7 @@ async def test_list_search_tools_datetime_conversion(monkeypatch):
# Test None handling for tool 2
tool2 = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "null-datetime-tool"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "null-datetime-tool"),
None,
)
assert tool2 is not None
@ -355,11 +334,7 @@ async def test_list_search_tools_datetime_conversion(monkeypatch):
# Test string passthrough for tool 3
tool3 = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "string-datetime-tool"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "string-datetime-tool"),
None,
)
assert tool3 is not None
@ -399,9 +374,7 @@ async def test_list_search_tools_config_error_handling(monkeypatch):
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
# Mock proxy_config to raise an error
mock_proxy_config = MagicMock()
mock_proxy_config.get_config = AsyncMock(
side_effect=Exception("Config error")
)
mock_proxy_config.get_config = AsyncMock(side_effect=Exception("Config error"))
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
# Mock auth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -420,13 +393,8 @@ async def test_list_search_tools_config_error_handling(monkeypatch):
assert len(data["search_tools"]) == 1
assert data["search_tools"][0]["search_tool_name"] == "db-tool-1"
# Verify masking of sensitive values
assert (
data["search_tools"][0]["litellm_params"]["api_key"]
!= "sk-test"
)
assert (
"****" in data["search_tools"][0]["litellm_params"]["api_key"]
)
assert data["search_tools"][0]["litellm_params"]["api_key"] != "sk-test"
assert "****" in data["search_tools"][0]["litellm_params"]["api_key"]
finally:
app.dependency_overrides.pop(user_api_key_auth, None)
@ -541,31 +509,18 @@ async def test_list_search_tools_db_masking_sensitive_values(monkeypatch):
# Test tool 1: api_key should be masked
tool1 = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "perplexity-tool"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "perplexity-tool"),
None,
)
assert tool1 is not None
assert (
tool1["litellm_params"]["api_key"] != "pplx-sk-1234567890abcdef"
)
assert tool1["litellm_params"]["api_key"] != "pplx-sk-1234567890abcdef"
assert "****" in tool1["litellm_params"]["api_key"]
assert tool1["litellm_params"]["search_provider"] == "perplexity"
assert (
tool1["litellm_params"]["api_base"]
== "https://api.perplexity.ai"
)
assert tool1["litellm_params"]["api_base"] == "https://api.perplexity.ai"
# Test tool 2: api_key should be masked
tool2 = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "tavily-tool"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "tavily-tool"),
None,
)
assert tool2 is not None
@ -575,29 +530,18 @@ async def test_list_search_tools_db_masking_sensitive_values(monkeypatch):
# Test tool 3: access_token and secret_key should be masked
tool3 = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "tool-with-token"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "tool-with-token"),
None,
)
assert tool3 is not None
assert (
tool3["litellm_params"]["access_token"]
!= "token-abcdefghijklmnop"
)
assert tool3["litellm_params"]["access_token"] != "token-abcdefghijklmnop"
assert "****" in tool3["litellm_params"]["access_token"]
assert tool3["litellm_params"]["secret_key"] != "secret-xyz123"
assert "****" in tool3["litellm_params"]["secret_key"]
# Test tool 4: non-sensitive fields should remain unmasked
tool4 = next(
(
t
for t in data["search_tools"]
if t["search_tool_name"] == "tool-with-non-sensitive"
),
(t for t in data["search_tools"] if t["search_tool_name"] == "tool-with-non-sensitive"),
None,
)
assert tool4 is not None
@ -626,25 +570,18 @@ async def test_get_all_search_tools_from_db_retries_on_transport_error():
return []
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_searchtoolstable.find_many = AsyncMock(
side_effect=_flaky_find_many
)
mock_prisma_client.db.litellm_searchtoolstable.find_many = AsyncMock(side_effect=_flaky_find_many)
mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0
mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1
result = await SearchToolRegistry.get_all_search_tools_from_db(
prisma_client=mock_prisma_client
)
result = await SearchToolRegistry.get_all_search_tools_from_db(prisma_client=mock_prisma_client)
assert result == []
assert len(invocations) == 2
mock_prisma_client.attempt_db_reconnect.assert_awaited_once()
reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs
assert (
reconnect_kwargs["reason"]
== "get_all_search_tools_from_db_lookup_failure"
)
assert reconnect_kwargs["reason"] == "get_all_search_tools_from_db_lookup_failure"
@contextlib.contextmanager
@ -715,6 +652,46 @@ def _override_auth(user):
app.dependency_overrides.pop(user_api_key_auth, None)
@pytest.mark.parametrize(
("method", "path", "body", "registry_method"),
[
(
"post",
"/search_tools",
{"search_tool": {"search_tool_name": "blocked", "litellm_params": {"search_provider": "tavily"}}},
"add_search_tool_to_db",
),
(
"put",
"/search_tools/blocked-id",
{"search_tool": {"search_tool_name": "blocked", "litellm_params": {"search_provider": "tavily"}}},
"update_search_tool_in_db",
),
("delete", "/search_tools/blocked-id", None, "delete_search_tool_from_db"),
],
)
def test_search_tool_mutations_require_proxy_admin(method, path, body, registry_method):
"""Internal users cannot configure proxy-originated search requests."""
registry = _fake_registry([])
internal_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user")
with (
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY", registry
),
_override_auth(internal_user),
):
client = TestClient(app)
request_kwargs = {} if body is None else {"json": body}
response = getattr(client, method)(path, **request_kwargs)
assert response.status_code == 403
getattr(registry, registry_method).assert_not_awaited()
@pytest.mark.asyncio
async def test_list_search_tools_scoped_to_key_object_permission():
"""
@ -747,9 +724,7 @@ async def test_list_search_tools_scoped_to_key_object_permission():
@pytest.mark.asyncio
async def test_list_search_tools_unrestricted_internal_user_sees_all():
"""An internal user with no search_tools allowlist is unrestricted and sees every tool."""
unrestricted_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user"
)
unrestricted_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user")
with (
_mock_search_tool_backend(_scoping_db_tools()),
@ -789,9 +764,7 @@ async def test_list_search_tools_scoped_to_team_object_permission():
response = TestClient(app).get("/search_tools/list")
assert response.status_code == 200
assert [t["search_tool_name"] for t in response.json()["search_tools"]] == [
"db-tool-2"
]
assert [t["search_tool_name"] for t in response.json()["search_tools"]] == ["db-tool-2"]
@pytest.mark.asyncio
@ -1059,9 +1032,21 @@ def _live_router_and_db(db_rows: list):
fake_router.search_tools = list(db_rows)
with contextlib.ExitStack() as stack:
stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", MagicMock())) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", proxy_config)) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", fake_router)) # test-quality-ok: proxy globals are the only seam; see the module note above
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.proxy_server.prisma_client", MagicMock()
)
)
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.proxy_server.proxy_config", proxy_config
)
)
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.proxy_server.llm_router", fake_router
)
)
stack.enter_context(
patch( # test-quality-ok: proxy globals are the only seam; see the module note above
"litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY",

44
tests/test_litellm/responses/test_responses_prompt_management.py
View file

@ -67,18 +67,15 @@ def _patch_responses_dispatch():
side_effect=_provider_by_model,
),
patch(
"litellm.responses.mcp.litellm_proxy_mcp_handler."
"LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway",
"litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway",
return_value=False,
),
patch(
"litellm.responses.main.ProviderConfigManager"
".get_provider_responses_api_config",
"litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config",
return_value=None,
),
patch(
"litellm.responses.main.litellm_completion_transformation_handler"
".response_api_handler",
"litellm.responses.main.litellm_completion_transformation_handler.response_api_handler",
return_value=MagicMock(),
),
]
@ -140,7 +137,6 @@ def _make_cache_control_case() -> tuple[
class TestResponsesAPIPromptManagement:
def test_str_input_coerced_and_merged(self):
"""[A] str input is wrapped into a message list before being passed to the hook."""
template_messages: List[AllMessageValues] = [
@ -171,9 +167,7 @@ class TestResponsesAPIPromptManagement:
logging_obj.get_chat_completion_prompt.assert_called_once()
call_kwargs = logging_obj.get_chat_completion_prompt.call_args.kwargs
# str was coerced to a single user message before being passed to the hook
assert call_kwargs["messages"] == [
{"role": "user", "content": "Tell me about AI."}
]
assert call_kwargs["messages"] == [{"role": "user", "content": "Tell me about AI."}]
assert call_kwargs["prompt_id"] == "summariser-prompt"
def test_list_input_merged_with_template(self):
@ -579,6 +573,36 @@ def test_resolve_prompt_swapped_provider_allows_swap_without_credentials():
)
@pytest.mark.parametrize("authorization_header", ["Authorization", "authorization"])
def test_resolve_prompt_swapped_provider_rejects_authorization_extra_headers(authorization_header: str):
import litellm
from litellm.responses.main import _resolve_prompt_swapped_provider
with pytest.raises(litellm.BadRequestError, match="Refusing to send"):
_resolve_prompt_swapped_provider(
original_model="anthropic/claude-haiku-4-5",
swapped_model="gpt-4o-mini",
custom_llm_provider="anthropic",
kwargs={"extra_headers": {authorization_header: "Bearer secret"}},
prompt_id="p1",
)
def test_resolve_prompt_swapped_provider_allows_harmless_extra_headers():
from litellm.responses.main import _resolve_prompt_swapped_provider
assert (
_resolve_prompt_swapped_provider(
original_model="anthropic/claude-haiku-4-5",
swapped_model="gpt-4o-mini",
custom_llm_provider="anthropic",
kwargs={"extra_headers": {"X-Request-ID": "request-1"}},
prompt_id="p1",
)
== "openai"
)
def test_resolve_prompt_swapped_provider_allows_same_provider_swap_with_credentials():
from litellm.responses.main import _resolve_prompt_swapped_provider