mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
Merge remote-tracking branch 'origin/main' into litellm_failure_path_redaction
# Conflicts: # litellm/proxy/utils.py
This commit is contained in:
yucheng
commit
0d6d09a00a
16 changed files with 478 additions and 90 deletions
|
|
@ -9,6 +9,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -218,6 +219,15 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None =
|
|||
)
|
||||
|
||||
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuardrailParams]]:
|
||||
if not isinstance(metadata, dict):
|
||||
return []
|
||||
return metadata.get("guardrails") or []
|
||||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
|
@ -791,28 +801,22 @@ class CustomGuardrail(CustomLogger):
|
|||
return True
|
||||
raise
|
||||
|
||||
def get_guardrail_from_metadata(self, data: dict) -> list[str] | list[dict[str, DynamicGuardrailParams]]:
|
||||
def get_guardrail_from_metadata(self, data: dict) -> list[str | dict[str, DynamicGuardrailParams]]:
|
||||
"""
|
||||
Returns the guardrail(s) to be run from the metadata or root
|
||||
Returns the guardrail(s) to be run from the root, else from both metadata fields together
|
||||
"""
|
||||
|
||||
if "guardrails" in data:
|
||||
return data["guardrails"]
|
||||
# Check both metadata locations. For regular endpoints move_guardrails_to_metadata
|
||||
# writes to "metadata"; for thread/assistant endpoints it writes to
|
||||
# "litellm_metadata". We check the one that actually contains the "guardrails"
|
||||
# key so that a non-empty litellm_metadata without guardrails does not shadow
|
||||
# the merged list stored in metadata (which would cause team guardrails to be
|
||||
# silently skipped while default_on=True policy guardrails still fire).
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key) or {}
|
||||
if isinstance(meta, dict) and "guardrails" in meta:
|
||||
return meta.get("guardrails") or []
|
||||
return []
|
||||
request_fields: Final = _REQUEST_FIELDS.validate_python(data)
|
||||
return [
|
||||
*_guardrails_listed_in(request_fields.get("metadata")),
|
||||
*_guardrails_listed_in(request_fields.get("litellm_metadata")),
|
||||
]
|
||||
|
||||
def _guardrail_is_in_requested_guardrails(
|
||||
self,
|
||||
requested_guardrails: list[str] | list[dict[str, DynamicGuardrailParams]],
|
||||
requested_guardrails: Sequence[str | dict[str, DynamicGuardrailParams]],
|
||||
) -> bool:
|
||||
for _guardrail in requested_guardrails:
|
||||
if isinstance(_guardrail, dict):
|
||||
|
|
|
|||
|
|
@ -446,7 +446,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request,
|
|||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
_trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser
|
||||
get_chain_id_from_headers,
|
||||
metadata_variable_name_for_route,
|
||||
metadata_variable_name_for_request,
|
||||
)
|
||||
|
||||
if get_chain_id_from_headers(dict(request.headers)):
|
||||
|
|
@ -460,7 +460,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request,
|
|||
traceparent: Final[str | None] = request.headers.get("traceparent")
|
||||
if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent):
|
||||
return True
|
||||
metadata_variable_name: Final = metadata_variable_name_for_route(get_request_route(request))
|
||||
metadata_variable_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body)
|
||||
selected_metadata: Final[object] = request_body.get(metadata_variable_name)
|
||||
if isinstance(selected_metadata, dict):
|
||||
selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata)
|
||||
|
|
@ -1367,11 +1367,15 @@ async def common_checks(
|
|||
if not skip_all_budget_checks:
|
||||
# Key metadata.tags are injected into request_body here so the tag budget
|
||||
# check can read them; this mutation must run before the gathered checks.
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
|
||||
|
||||
if valid_token is not None:
|
||||
for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_body):
|
||||
LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route)
|
||||
LiteLLMProxyRequestSetup.merge_caller_body_tags(
|
||||
request_data=tag_target,
|
||||
metadata_variable_name=metadata_variable_name_for_request(route, tag_target),
|
||||
)
|
||||
LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(request_data=tag_target, user_api_key_dict=valid_token)
|
||||
LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body)
|
||||
|
||||
|
|
|
|||
|
|
@ -2401,9 +2401,10 @@ class ProxyBaseLLMRequestProcessing:
|
|||
version=version,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
request_fields: Final = _REQUEST_FIELDS.validate_python(self.data)
|
||||
metadata_variable_name: Final = get_metadata_variable_name(request, request_fields)
|
||||
if logging_auth is not None:
|
||||
metadata_variable_name: Final = get_metadata_variable_name(request)
|
||||
request_metadata: Final = self.data.get(metadata_variable_name)
|
||||
request_metadata: Final = request_fields.get(metadata_variable_name)
|
||||
if isinstance(request_metadata, dict):
|
||||
self.data[metadata_variable_name] = {
|
||||
**TypeAdapter(dict[str, object]).validate_python(request_metadata),
|
||||
|
|
@ -2431,12 +2432,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
# Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved
|
||||
if queue_time_seconds is not None:
|
||||
_metadata_variable_name: Final = get_metadata_variable_name(request)
|
||||
if _metadata_variable_name not in self.data:
|
||||
self.data[_metadata_variable_name] = {}
|
||||
if not isinstance(self.data[_metadata_variable_name], dict):
|
||||
self.data[_metadata_variable_name] = {}
|
||||
self.data[_metadata_variable_name]["queue_time_seconds"] = queue_time_seconds
|
||||
if metadata_variable_name not in self.data:
|
||||
self.data[metadata_variable_name] = {}
|
||||
if not isinstance(self.data[metadata_variable_name], dict):
|
||||
self.data[metadata_variable_name] = {}
|
||||
self.data[metadata_variable_name]["queue_time_seconds"] = queue_time_seconds
|
||||
|
||||
if isinstance(model, str):
|
||||
reject_url_valued_destination("model", model)
|
||||
|
|
|
|||
|
|
@ -3312,7 +3312,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if key_agent_id:
|
||||
return key_agent_id
|
||||
metadata: Final = data.get("metadata") or {}
|
||||
return metadata.get("agent_id")
|
||||
litellm_metadata: Final = data.get("litellm_metadata") or {}
|
||||
return metadata.get("agent_id") or litellm_metadata.get("agent_id")
|
||||
|
||||
def _get_session_id_from_data(self, data: dict) -> str | None:
|
||||
"""Extract session_id from request metadata or litellm_session_id."""
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@ from litellm.proxy.db.db_spend_update_writer import (
|
|||
get_llm_router,
|
||||
)
|
||||
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope
|
||||
from litellm.proxy.spend_tracking.spend_event import (
|
||||
|
|
@ -140,19 +140,34 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
|||
_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping)
|
||||
|
||||
|
||||
def _proxy_stamped_used_client_oauth_token(
|
||||
request_data: Mapping[str, object], request_route: str | None
|
||||
) -> bool | None:
|
||||
def _proxy_stamped_metadata(request_data: Mapping[str, object], request_route: str | None) -> Mapping[str, object]:
|
||||
proxy_bucket: Final = (
|
||||
get_metadata_variable_name_from_kwargs(request_data)
|
||||
if request_route is None
|
||||
else metadata_variable_name_for_route(request_route)
|
||||
else metadata_variable_name_for_request(request_route, request_data)
|
||||
)
|
||||
proxy_metadata: Final = request_data.get(proxy_bucket)
|
||||
stamped: Final = proxy_metadata.get("used_client_oauth_token") if isinstance(proxy_metadata, dict) else None
|
||||
return _str_keyed_fields(request_data.get(proxy_bucket))
|
||||
|
||||
|
||||
def _proxy_stamped_used_client_oauth_token(
|
||||
request_data: Mapping[str, object], request_route: str | None
|
||||
) -> bool | None:
|
||||
stamped: Final = _proxy_stamped_metadata(request_data, request_route).get("used_client_oauth_token")
|
||||
return stamped if isinstance(stamped, bool) else None
|
||||
|
||||
|
||||
_PROXY_STAMPED_REQUEST_FIELDS: Final = ("requester_ip_address", "user_agent")
|
||||
|
||||
|
||||
def _proxy_stamped_request_fields(
|
||||
request_data: Mapping[str, object], request_route: str | None
|
||||
) -> Mapping[str, object]:
|
||||
if request_route is None:
|
||||
return {}
|
||||
proxy_metadata: Final = _proxy_stamped_metadata(request_data, request_route)
|
||||
return {field: proxy_metadata[field] for field in _PROXY_STAMPED_REQUEST_FIELDS if field in proxy_metadata}
|
||||
|
||||
|
||||
def _str_keyed_fields(value: object) -> Mapping[str, object]:
|
||||
return value if is_str_keyed_mapping(value) else {}
|
||||
|
||||
|
|
@ -461,7 +476,13 @@ class _ProxyDBLogger(CustomLogger):
|
|||
_metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route)
|
||||
|
||||
existing_metadata: Final[dict] = request_data.get("metadata", None) or {}
|
||||
existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata})
|
||||
existing_metadata.update(
|
||||
{
|
||||
**_routed_metadata_fallback(request_fields),
|
||||
**_proxy_stamped_request_fields(request_fields, request_route),
|
||||
**_metadata,
|
||||
}
|
||||
)
|
||||
|
||||
litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
|
||||
litellm_metadata: Final[ObjectMapping] = (
|
||||
|
|
|
|||
|
|
@ -379,6 +379,8 @@ LITELLM_METADATA_ROUTES: Final = (
|
|||
"interactions",
|
||||
)
|
||||
|
||||
CURSOR_CHAT_COMPLETIONS_ROUTE: Final = "/cursor/chat/completions"
|
||||
|
||||
LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
|
||||
{
|
||||
"mask_input",
|
||||
|
|
@ -636,7 +638,7 @@ def _has_tag_budget(tag: "LiteLLM_TagTable | None") -> bool:
|
|||
return tag.litellm_budget_table.max_budget is not None
|
||||
|
||||
|
||||
def _normalized_metadata_slot(
|
||||
def normalized_metadata_slot(
|
||||
request_data: MutableMapping[str, object], metadata_variable_name: str
|
||||
) -> dict[str, object]:
|
||||
"""Return the request's metadata slot as a dict, normalising it in place first.
|
||||
|
|
@ -653,6 +655,12 @@ def _normalized_metadata_slot(
|
|||
return normalized
|
||||
|
||||
|
||||
def _holds_metadata_object(value: object) -> bool:
|
||||
if value is None or isinstance(value, dict):
|
||||
return True
|
||||
return isinstance(value, str) and isinstance(_DECODED_JSON.validate_python(safe_json_loads(value)), dict)
|
||||
|
||||
|
||||
def _metadata_object_or_error(field: str, value: object) -> Mapping[str, object] | ProxyException:
|
||||
if value is None:
|
||||
return {}
|
||||
|
|
@ -863,7 +871,7 @@ def _strip_router_reserved_metadata(
|
|||
)
|
||||
|
||||
|
||||
def get_metadata_variable_name(request: Request) -> str:
|
||||
def get_metadata_variable_name(request: Request, data: Mapping[str, object] | None = None) -> str:
|
||||
"""
|
||||
Helper to return what the "metadata" field should be called in the request data
|
||||
|
||||
|
|
@ -874,7 +882,10 @@ def get_metadata_variable_name(request: Request) -> str:
|
|||
# Inline imports — auth_utils/route_checks participate in a proxy import cycle.
|
||||
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415
|
||||
|
||||
return metadata_variable_name_for_route(get_request_route(request))
|
||||
route: Final = get_request_route(request)
|
||||
if data is None:
|
||||
return metadata_variable_name_for_route(route)
|
||||
return metadata_variable_name_for_request(route, data)
|
||||
|
||||
|
||||
_get_metadata_variable_name: Final = get_metadata_variable_name
|
||||
|
|
@ -890,6 +901,25 @@ def metadata_variable_name_for_route(route: str) -> Literal["metadata", "litellm
|
|||
return "metadata"
|
||||
|
||||
|
||||
def is_chat_completions_body(data: Mapping[str, object]) -> bool:
|
||||
messages: Final = data.get("messages")
|
||||
if isinstance(messages, list) and messages:
|
||||
return True
|
||||
return "messages" in data and "input" not in data
|
||||
|
||||
|
||||
def is_cursor_responses_request(route: str, body: Mapping[str, object]) -> bool:
|
||||
return CURSOR_CHAT_COMPLETIONS_ROUTE in route and not is_chat_completions_body(body)
|
||||
|
||||
|
||||
def metadata_variable_name_for_request(
|
||||
route: str, body: Mapping[str, object]
|
||||
) -> Literal["metadata", "litellm_metadata"]:
|
||||
if is_cursor_responses_request(route, body):
|
||||
return "litellm_metadata"
|
||||
return metadata_variable_name_for_route(route)
|
||||
|
||||
|
||||
def _promoted_trace_control_fields(
|
||||
requester_metadata: Mapping[str, object],
|
||||
litellm_metadata: Mapping[str, object],
|
||||
|
|
@ -2309,7 +2339,9 @@ class LiteLLMProxyRequestSetup:
|
|||
``_tag_max_budget_check`` resolve to ``litellm_metadata``, leaving
|
||||
header tags invisible to per-tag budget enforcement.
|
||||
"""
|
||||
if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES):
|
||||
if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES) or is_cursor_responses_request(
|
||||
route, request_data
|
||||
):
|
||||
request_data.setdefault("litellm_metadata", {})
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -2380,7 +2412,7 @@ class LiteLLMProxyRequestSetup:
|
|||
return
|
||||
tags: Final = _STRING_TAGS_ADAPTER.validate_python(raw_tags)
|
||||
slot: Final = metadata_key if metadata_key is not None else get_metadata_variable_name_from_kwargs(request_data)
|
||||
metadata: Final = _normalized_metadata_slot(request_data, slot)
|
||||
metadata: Final = normalized_metadata_slot(request_data, slot)
|
||||
existing: Final = metadata.get("tags")
|
||||
if existing and not isinstance(existing, list):
|
||||
return
|
||||
|
|
@ -2393,6 +2425,34 @@ class LiteLLMProxyRequestSetup:
|
|||
tags=tuple(tag for tag in tags if tag not in sent_tags)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def merge_caller_body_tags(
|
||||
request_data: MutableMapping[str, object], # mutable-ok: the request body the tag merges write into
|
||||
metadata_variable_name: str | None = None,
|
||||
) -> None:
|
||||
slot: Final = (
|
||||
metadata_variable_name
|
||||
if metadata_variable_name is not None
|
||||
else get_metadata_variable_name_from_kwargs(request_data)
|
||||
)
|
||||
if slot == "metadata" or not _holds_metadata_object(request_data.get(slot)):
|
||||
return
|
||||
raw_caller_metadata: Final = request_data.get("metadata")
|
||||
caller_metadata: Final = (
|
||||
_DECODED_JSON.validate_python(safe_json_loads(raw_caller_metadata))
|
||||
if isinstance(raw_caller_metadata, str)
|
||||
else raw_caller_metadata
|
||||
)
|
||||
caller_tags: Final = caller_metadata.get("tags") if isinstance(caller_metadata, dict) else None
|
||||
if not isinstance(caller_tags, list) or not caller_tags:
|
||||
return
|
||||
metadata: Final = normalized_metadata_slot(request_data, slot)
|
||||
existing_tags: Final = metadata.get("tags")
|
||||
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
|
||||
request_tags=existing_tags if isinstance(existing_tags, list) else None,
|
||||
tags_to_add=caller_tags,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def apply_key_tags_pre_auth(
|
||||
request_data: dict,
|
||||
|
|
@ -2408,7 +2468,7 @@ class LiteLLMProxyRequestSetup:
|
|||
return
|
||||
|
||||
_metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data)
|
||||
metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
|
||||
existing_tags: Final = metadata.get("tags")
|
||||
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
|
||||
|
|
@ -2460,7 +2520,7 @@ class LiteLLMProxyRequestSetup:
|
|||
# from (litellm_metadata vs metadata) so the merged tags are visible
|
||||
# to _tag_max_budget_check.
|
||||
_metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data)
|
||||
metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
|
||||
existing_tags: Final = metadata.get("tags")
|
||||
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
|
||||
|
|
@ -2637,7 +2697,7 @@ async def add_litellm_data_to_request(
|
|||
}
|
||||
|
||||
safe_add_api_version_from_query_params(data, request)
|
||||
_metadata_variable_name: Final = get_metadata_variable_name(request)
|
||||
_metadata_variable_name: Final = get_metadata_variable_name(request, _REQUEST_FIELDS.validate_python(data))
|
||||
if data.get(_metadata_variable_name, None) is None:
|
||||
data[_metadata_variable_name] = {}
|
||||
|
||||
|
|
@ -2974,17 +3034,7 @@ async def add_litellm_data_to_request(
|
|||
tags_to_add=tags,
|
||||
)
|
||||
|
||||
_caller_body_metadata: Final = data.get("metadata") if _metadata_variable_name != "metadata" else None
|
||||
_caller_body_tags: Final = (
|
||||
_caller_body_metadata.get("tags")
|
||||
if isinstance(_caller_body_metadata, dict) and isinstance(_caller_body_metadata.get("tags"), list)
|
||||
else None
|
||||
)
|
||||
if _caller_body_tags:
|
||||
data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup.merge_tags( # rebind-ok: matches file idiom
|
||||
request_tags=data[_metadata_variable_name].get("tags"),
|
||||
tags_to_add=_caller_body_tags,
|
||||
)
|
||||
LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data, metadata_variable_name=_metadata_variable_name)
|
||||
|
||||
# caller_tags: the tags the request's headers and body carry (x-litellm-tags, body "tags", and body
|
||||
# "metadata.tags" / "litellm_metadata.tags" read before the mapped tags are added below),
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry
|
|||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
_CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
|
||||
_UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
|
||||
is_chat_completions_body,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
|
||||
request_tags_from_request_data,
|
||||
|
|
@ -213,13 +214,6 @@ def _normalize_tool_dialect(
|
|||
return {**data, **{key: value for key, value in replaceable if key in data}}
|
||||
|
||||
|
||||
def _is_chat_completions_body(data: Mapping[str, object]) -> bool:
|
||||
messages: Final = data.get("messages")
|
||||
if isinstance(messages, list) and messages:
|
||||
return True
|
||||
return "messages" in data and "input" not in data
|
||||
|
||||
|
||||
_CURSOR_THINKING_SEPARATOR: Final = "-thinking-"
|
||||
_CURSOR_FAST_SUFFIX: Final = "-fast"
|
||||
_CURSOR_THINKING_LEVELS: Final[frozenset[str]] = frozenset(get_args(REASONING_EFFORT))
|
||||
|
|
@ -261,7 +255,7 @@ def _resolve_cursor_model_variant(
|
|||
resolved: Final = {**data, "model": variant.base_model}
|
||||
if variant.reasoning_effort is None:
|
||||
return resolved
|
||||
if _is_chat_completions_body(data):
|
||||
if is_chat_completions_body(data):
|
||||
if "reasoning_effort" in data:
|
||||
return resolved
|
||||
return {**resolved, "reasoning_effort": variant.reasoning_effort}
|
||||
|
|
@ -646,7 +640,7 @@ async def cursor_chat_completions(
|
|||
|
||||
raw_body: Final = await read_request_body(request=request)
|
||||
|
||||
if _is_chat_completions_body(raw_body):
|
||||
if is_chat_completions_body(raw_body):
|
||||
# Genuine chat completions body (Cursor sends these for models whose BYOK it
|
||||
# already fixed); delegate so behavior matches /chat/completions exactly.
|
||||
# Keyed on messages CONTENT, not key presence: Cursor can send a null or
|
||||
|
|
|
|||
|
|
@ -255,6 +255,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
LiteLLMProxyRequestSetup,
|
||||
add_guardrails_from_auth_metadata,
|
||||
key_or_team_allows_client_message_redaction_opt_out,
|
||||
normalized_metadata_slot,
|
||||
)
|
||||
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
||||
from litellm.proxy.policy_engine.pipeline_executor import (
|
||||
|
|
@ -3986,9 +3987,7 @@ class ProxyLogging:
|
|||
)
|
||||
request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller
|
||||
request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP})
|
||||
if "metadata" not in request_data:
|
||||
request_data["metadata"] = {}
|
||||
request_data["metadata"].update(user_api_key_logged_metadata)
|
||||
normalized_metadata_slot(request_data, "metadata").update(user_api_key_logged_metadata)
|
||||
|
||||
if litellm_logging_obj is not None:
|
||||
## UPDATE LOGGING INPUT
|
||||
|
|
|
|||
|
|
@ -3,11 +3,14 @@ from __future__ import annotations
|
|||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from hashlib import sha256
|
||||
from importlib.resources import files
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.client import Gateway, Scenario, eventually, list_value, object_value
|
||||
from integration._support.responses_vendor import same_response
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration._support.worker_readiness import wait_for_worker_readiness
|
||||
|
|
@ -38,6 +41,7 @@ _CHAT_SUFFIXES: Final[tuple[tuple[str, str], ...]] = (("", ""),) + tuple(
|
|||
(f"-thinking-{level}", level) for level in _CHAT_THINKING_LEVELS
|
||||
)
|
||||
_PROMPT: Final = "Inspect the requested change."
|
||||
_REQUEST_TAG: Final = "cursor-integration"
|
||||
_INPUT: Final = [{"role": "user", "content": [{"type": "input_text", "text": _PROMPT}]}]
|
||||
_CUSTOM_TOOL: Final = {
|
||||
"type": "custom",
|
||||
|
|
@ -107,16 +111,30 @@ def _stream(identity: str) -> tuple[bytes, ...]:
|
|||
) + (b"data: [DONE]\n\n",)
|
||||
|
||||
|
||||
def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]:
|
||||
def _chat_completion(stream: bool, identity: str) -> bytes | tuple[bytes, ...]:
|
||||
if stream:
|
||||
return (
|
||||
b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{"content":"Patch inspected."},"finish_reason":null}]}\n\n',
|
||||
b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
return tuple(
|
||||
f"data: {json.dumps(chunk)}\n\n".encode()
|
||||
for chunk in (
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _MODEL,
|
||||
"choices": [{"index": 0, "delta": {"content": "Patch inspected."}, "finish_reason": None}],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _MODEL,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
},
|
||||
)
|
||||
) + (b"data: [DONE]\n\n",)
|
||||
return json.dumps(
|
||||
{
|
||||
"id": "chat_cursor_1",
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": _MODEL,
|
||||
|
|
@ -133,11 +151,12 @@ def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]:
|
|||
).encode()
|
||||
|
||||
|
||||
def _deployment(scenario: Scenario, wire: Wire) -> str:
|
||||
def _deployment(scenario: Scenario, wire: Wire, **pricing: JsonValue) -> str:
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_MODEL}",
|
||||
api_key=_API_KEY,
|
||||
api_base=f"{wire.url}/v1",
|
||||
**pricing,
|
||||
)
|
||||
wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS)
|
||||
return model
|
||||
|
|
@ -167,7 +186,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
|
|||
assert request.method == "POST"
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}", dict(request.headers)
|
||||
if request.target == "/v1/chat/completions":
|
||||
completion: Final = _chat_completion(stream)
|
||||
completion: Final = _chat_completion(stream, raw_id)
|
||||
return Reply(
|
||||
content_type="text/event-stream" if stream else "application/json",
|
||||
chunks=completion if isinstance(completion, tuple) else None,
|
||||
|
|
@ -175,7 +194,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
|
|||
)
|
||||
assert request.target == "/v1/responses", request.target
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
assert {key: value for key, value in body.items() if key != "metadata"} == {
|
||||
assert body == {
|
||||
"model": _MODEL,
|
||||
"input": _INPUT,
|
||||
**reasoning,
|
||||
|
|
@ -259,41 +278,69 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
|
|||
assert [request.target for request in drain_contract_requests(wire)] == ["/v1/responses"]
|
||||
|
||||
|
||||
def test_cursor_responses_body_carries_no_proxy_metadata(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: /cursor/chat/completions forwards proxy-internal metadata (key hash, team, user, client ip, user agent) "
|
||||
"to the provider in the Responses body"
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"client_metadata", ({}, {"cursor_conversation": "integration"}), ids=("no_client_metadata", "client_metadata")
|
||||
)
|
||||
@pytest.mark.parametrize("stream", (False, True))
|
||||
def test_cursor_responses_body_sends_only_client_metadata_and_bills_the_key(
|
||||
gateway: Gateway,
|
||||
stream: bool,
|
||||
client_metadata: Mapping[str, JsonValue],
|
||||
) -> None:
|
||||
raw_id: Final = f"resp_cursor_metadata_{uuid.uuid4().hex}"
|
||||
alias: Final = f"integration-cursor-{uuid.uuid4().hex}"
|
||||
sent_metadata: Final = {"metadata": client_metadata} if client_metadata else {}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
assert request.method == "POST" and request.target == "/v1/responses", request.target
|
||||
if stream:
|
||||
return Reply(content_type="text/event-stream", chunks=_stream(raw_id))
|
||||
return Reply(body=json.dumps(_response(raw_id)).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire)
|
||||
key: Final = scenario.key(key_alias=alias)
|
||||
hashed_key: Final = sha256(key.encode()).hexdigest()
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/cursor/chat/completions",
|
||||
{
|
||||
"model": f"{model}-thinking-high",
|
||||
"model": model,
|
||||
"input": _INPUT,
|
||||
"messages": [],
|
||||
"tools": [_CUSTOM_TOOL],
|
||||
"stream": stream,
|
||||
**sent_metadata,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
requests: Final = drain_contract_requests(wire)
|
||||
assert len(requests) == 1, requests
|
||||
body: Final = _JSON_OBJECT.validate_json(requests[0].body)
|
||||
assert body == {
|
||||
assert _JSON_OBJECT.validate_json(requests[0].body) == {
|
||||
"model": _MODEL,
|
||||
"input": _INPUT,
|
||||
"reasoning": {"effort": "high"},
|
||||
"tools": [_CUSTOM_TOOL],
|
||||
"stream": stream,
|
||||
**sent_metadata,
|
||||
}, requests[0].body
|
||||
(row,) = eventually(
|
||||
lambda: gateway.rows(
|
||||
'SELECT call_type, spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (hashed_key,)
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert row["call_type"] == "aresponses", row
|
||||
assert object_value(row["metadata"])["user_api_key_alias"] == alias, row
|
||||
assert isinstance(row["spend"], float) and row["spend"] > 0, row
|
||||
eventually(
|
||||
lambda: gateway.rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed_key,)),
|
||||
lambda tokens: len(tokens) == 1 and tokens[0]["spend"] == pytest.approx(row["spend"]),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def test_cursor_stream_reports_usage_when_the_client_asks_for_it(gateway: Gateway) -> None:
|
||||
|
|
@ -341,10 +388,12 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
suffix: str,
|
||||
reasoning_effort: str,
|
||||
) -> None:
|
||||
chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
completion: Final = _chat_completion(stream)
|
||||
completion: Final = _chat_completion(stream, chat_id)
|
||||
return Reply(
|
||||
content_type="text/event-stream" if stream else "application/json",
|
||||
chunks=completion if isinstance(completion, tuple) else None,
|
||||
|
|
@ -367,6 +416,7 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
"stream": stream,
|
||||
},
|
||||
key=key,
|
||||
headers={"x-litellm-tags": _REQUEST_TAG},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
if stream:
|
||||
|
|
@ -393,7 +443,7 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
payload: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
identity: Final = payload.get("id")
|
||||
created: Final = payload.get("created")
|
||||
assert isinstance(identity, str) and same_response(identity, "chat_cursor_1"), response.text
|
||||
assert isinstance(identity, str) and same_response(identity, chat_id), response.text
|
||||
assert isinstance(created, int) and created > 0, response.text
|
||||
assert payload == {
|
||||
"id": identity,
|
||||
|
|
@ -456,3 +506,72 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
**({"stream_options": {"include_usage": True}} if stream else {}),
|
||||
**({"reasoning_effort": reasoning_effort} if reasoning_effort else {}),
|
||||
}, requests[0].body
|
||||
(row,) = eventually(
|
||||
lambda: gateway.rows(
|
||||
'SELECT call_type, team_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert (row["call_type"], row["team_id"]) == ("acompletion", team), row
|
||||
assert _REQUEST_TAG in list_value(row["request_tags"]), row
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body_shape", "metadata_as_json_string"), (("messages", False), ("input", False), ("input", True))
|
||||
)
|
||||
def test_cursor_body_keeps_client_and_key_tags_and_enforces_the_client_tag_budget(
|
||||
gateway: Gateway, body_shape: str, metadata_as_json_string: bool
|
||||
) -> None:
|
||||
chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}"
|
||||
response_id: Final = f"resp_cursor_tags_{uuid.uuid4().hex}"
|
||||
budget_tag: Final = f"cursor-tag-budget-{uuid.uuid4().hex}"
|
||||
key_tag: Final = f"cursor-key-tag-{uuid.uuid4().hex}"
|
||||
client_metadata: Final = {"tags": [budget_tag]}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
if request.target == "/v1/responses":
|
||||
return Reply(body=json.dumps(_response(response_id)).encode())
|
||||
completion: Final = _chat_completion(False, chat_id)
|
||||
return Reply(body=completion if isinstance(completion, bytes) else b"")
|
||||
|
||||
def delete_tag() -> None:
|
||||
gateway.post("/tag/delete", {"name": budget_tag})
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, input_cost_per_token=0.01, output_cost_per_token=0.01)
|
||||
key: Final = scenario.key(models=[model], metadata={"tags": [key_tag]})
|
||||
gateway.post("/tag/new", {"name": budget_tag, "max_budget": 0.0001})
|
||||
scenario.cleanups.callback(delete_tag)
|
||||
|
||||
def tagged_chat(text: str) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/cursor/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
body_shape: [{"role": "user", "content": text}],
|
||||
"metadata": json.dumps(client_metadata) if metadata_as_json_string else client_metadata,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
first: Final = tagged_chat("first tagged chat")
|
||||
assert first.status_code == 200, first.text
|
||||
(row,) = eventually(
|
||||
lambda: gateway.rows(
|
||||
'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert {budget_tag, key_tag} <= set(list_value(row["request_tags"])), row
|
||||
blocked: Final = eventually(
|
||||
lambda: tagged_chat("tag budget probe"), lambda response: response.status_code != 200, seconds=70
|
||||
)
|
||||
assert blocked.status_code == 422, blocked.text
|
||||
assert "budget" in blocked.text.lower(), blocked.text
|
||||
|
|
|
|||
|
|
@ -3699,6 +3699,20 @@ def test_get_guardrail_from_metadata() -> None:
|
|||
) == [{"test-guardrail": {"extra_body": {"key": "value"}}}]
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"caller_metadata",
|
||||
[{"guardrails": []}, {"guardrails": ["not-a-guardrail"]}, {"guardrails": None}],
|
||||
)
|
||||
def test_caller_metadata_guardrails_never_hide_the_proxy_guardrails_in_litellm_metadata(caller_metadata: dict) -> None:
|
||||
guardrail: Final = CustomGuardrail(
|
||||
guardrail_name="key-guardrail", event_hook=GuardrailEventHooks.pre_call, default_on=False
|
||||
)
|
||||
request_data: Final = {"metadata": caller_metadata, "litellm_metadata": {"guardrails": ["key-guardrail"]}}
|
||||
|
||||
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is True
|
||||
|
||||
|
||||
def test_guardrail_is_in_requested_guardrails() -> None:
|
||||
guardrail: Final = CustomGuardrail(guardrail_name="test-guardrail")
|
||||
|
||||
|
|
|
|||
|
|
@ -2591,7 +2591,8 @@ async def test_agent_session_rate_limit_skipped_without_session_id():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_rate_limit_from_metadata_agent_id():
|
||||
@pytest.mark.parametrize("metadata_field", ("metadata", "litellm_metadata"))
|
||||
async def test_agent_rate_limit_from_metadata_agent_id(metadata_field):
|
||||
"""
|
||||
Test that agent rate limits work when agent_id comes from
|
||||
request metadata (header) rather than from the API key.
|
||||
|
|
@ -2639,7 +2640,7 @@ async def test_agent_rate_limit_from_metadata_agent_id():
|
|||
cache=local_cache,
|
||||
data={
|
||||
"model": "gpt-4",
|
||||
"metadata": {"agent_id": _agent_id},
|
||||
metadata_field: {"agent_id": _agent_id},
|
||||
},
|
||||
call_type="",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -623,6 +623,56 @@ async def test_async_post_call_failure_hook_reads_used_client_oauth_token_from_t
|
|||
assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"request_route, body",
|
||||
[
|
||||
(
|
||||
"/cursor/chat/completions",
|
||||
{
|
||||
"input": "Hello",
|
||||
"metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"},
|
||||
"litellm_metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"},
|
||||
},
|
||||
),
|
||||
(
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"},
|
||||
"litellm_metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"},
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_async_post_call_failure_hook_logs_the_proxy_stamped_ip_and_user_agent(request_route: str, body: dict):
|
||||
request_data = {
|
||||
"model": "gpt-4o",
|
||||
"proxy_server_request": {"request_id": "test_request_id"},
|
||||
**body,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_database:
|
||||
await ProxyDBLogger().async_post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("invalid api key"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key", request_route=request_route),
|
||||
)
|
||||
|
||||
payload = get_logging_payload(
|
||||
kwargs=mock_update_database.call_args[1]["kwargs"],
|
||||
response_obj={},
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
assert payload["requester_ip_address"] == "127.0.0.1"
|
||||
assert json.loads(payload["metadata"])["user_agent"] == "Cursor/1.0"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_failure_hook_bills_guardrail_cost_on_blocked_request():
|
||||
"""LIT-5651: a request blocked by a guardrail never reaches the LLM, but the
|
||||
|
|
|
|||
|
|
@ -3938,7 +3938,7 @@ class TestCursorInputArmFlattening:
|
|||
|
||||
class TestChatCompletionsBodyDetection:
|
||||
def test_routing_matrix(self):
|
||||
from litellm.proxy.response_api_endpoints.endpoints import _is_chat_completions_body
|
||||
from litellm.proxy.litellm_pre_call_utils import is_chat_completions_body as _is_chat_completions_body
|
||||
|
||||
assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}]}) is True
|
||||
assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}], "input": []}) is True
|
||||
|
|
|
|||
|
|
@ -9711,6 +9711,40 @@ async def test_saved_weights_override_caller_input_and_preserve_key_precedence(
|
|||
assert logging.pre_call_hook.call_args.kwargs["data"].get("_router_weights") == expected
|
||||
|
||||
|
||||
def _hook_data(**kwargs: object) -> object:
|
||||
return kwargs["data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("prompt, proxy_field, client_field", [
|
||||
({"input": "hi"}, "litellm_metadata", "metadata"),
|
||||
({"messages": [{"role": "user", "content": "hi"}]}, "metadata", "litellm_metadata"),
|
||||
])
|
||||
async def test_cursor_route_writes_proxy_fields_only_to_the_field_its_body_shape_picks(
|
||||
prompt: dict[str, str | list[dict[str, str]]],
|
||||
proxy_field: str,
|
||||
client_field: str,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
logging: Final = MagicMock(spec=ProxyLogging)
|
||||
logging.pre_call_hook = AsyncMock(side_effect=_hook_data)
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-x", "metadata": {"client": "cursor"}, **prompt})
|
||||
data, _ = await processor.common_processing_pre_call_logic(
|
||||
request=Request({"type": "http", "method": "POST", "path": "/cursor/chat/completions", "headers": []}),
|
||||
general_settings={},
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash"),
|
||||
proxy_logging_obj=logging,
|
||||
proxy_config=proxy_server.ProxyConfig(),
|
||||
route_type="acompletion",
|
||||
logging_auth=ProxyUserAPIKeyAuth(api_key="logging-hash", key_alias="logging-alias"),
|
||||
)
|
||||
assert data[proxy_field]["user_api_key_alias"] == "logging-alias"
|
||||
assert "queue_time_seconds" in data[proxy_field]
|
||||
assert "user_api_key_alias" not in data.get(client_field, {})
|
||||
assert "queue_time_seconds" not in data.get(client_field, {})
|
||||
assert data["metadata"]["client"] == "cursor"
|
||||
|
||||
class TestPerRequestModelGroupAlias:
|
||||
"""``router_settings.model_group_alias`` on a key or team has to be resolved
|
||||
by the proxy: the Router resolves aliases from its own shared instance
|
||||
|
|
|
|||
|
|
@ -121,6 +121,23 @@ class TestGetMetadataVariableName:
|
|||
request = self._make_request("/chat/completions")
|
||||
assert get_metadata_variable_name(request) == "metadata"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body, expected",
|
||||
[
|
||||
({"messages": [{"role": "user", "content": "hi"}]}, "metadata"),
|
||||
({"messages": []}, "metadata"),
|
||||
({"input": [{"role": "user", "content": "hi"}]}, "litellm_metadata"),
|
||||
({"input": "hi", "messages": []}, "litellm_metadata"),
|
||||
],
|
||||
)
|
||||
def test_cursor_route_picks_the_metadata_field_by_body_shape(self, body, expected):
|
||||
request = self._make_request("/cursor/chat/completions")
|
||||
assert get_metadata_variable_name(request, body) == expected
|
||||
|
||||
def test_cursor_route_without_a_body_keeps_the_route_answer(self):
|
||||
request = self._make_request("/cursor/chat/completions")
|
||||
assert get_metadata_variable_name(request) == "metadata"
|
||||
|
||||
def test_returns_metadata_for_completions(self):
|
||||
request = self._make_request("/v1/completions")
|
||||
assert get_metadata_variable_name(request) == "metadata"
|
||||
|
|
@ -681,6 +698,45 @@ async def test_add_litellm_data_to_request_rejects_non_object_metadata(field, va
|
|||
assert field not in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/responses", "/cursor/chat/completions"])
|
||||
@pytest.mark.parametrize("value, received_type", [("junk", "a string"), (5, "an integer"), ([1], "an array")])
|
||||
async def test_invalid_litellm_metadata_answers_400_when_metadata_carries_tags(route, value, received_type):
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
data = {"model": "gpt-4o", "input": "hi", "metadata": {"tags": ["client-tag"]}, "litellm_metadata": value}
|
||||
request = _batches_request_mock()
|
||||
await common_checks(
|
||||
request_body=data,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route=route,
|
||||
llm_router=None,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
|
||||
valid_token=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
request=request,
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
assert exc_info.value.code == "400"
|
||||
assert exc_info.value.message == (
|
||||
f"Invalid type for 'litellm_metadata': expected an object, but got {received_type} instead."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_removes_every_invalid_metadata_field_before_raising():
|
||||
"""When both fields are invalid, the raise for the first must not leave the
|
||||
|
|
@ -11098,3 +11154,15 @@ def test_mapped_baggage_tags_keep_non_string_client_tags_as_sent():
|
|||
data: Final[dict[str, object]] = {"metadata": {"tags": [1, {"team": "x"}, "client-tag", "client-tag"]}}
|
||||
LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data)
|
||||
assert data["metadata"]["tags"] == [1, {"team": "x"}, "client-tag", "client-tag", "tenant:a"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("caller_metadata", ({"tags": ["client-tag"]}, json.dumps({"tags": ["client-tag"]})))
|
||||
def test_caller_tags_reach_litellm_metadata_whether_metadata_is_a_dict_or_a_json_string(caller_metadata):
|
||||
data: Final[dict[str, object]] = {
|
||||
"input": "hi",
|
||||
"metadata": caller_metadata,
|
||||
"litellm_metadata": {"tags": ["key-tag"]},
|
||||
}
|
||||
LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data)
|
||||
assert set(data["litellm_metadata"]["tags"]) == {"key-tag", "client-tag"}
|
||||
assert data["metadata"] == caller_metadata
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``."""
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
|
|
@ -1192,3 +1193,31 @@ async def test_handle_logging_proxy_only_path_flags_requests_that_died_before_re
|
|||
blocked_by_guardrail.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP,
|
||||
blocked_by_guardrail["litellm_call_id"],
|
||||
) == (True, "client-supplied", False, "call-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_failure_on_a_responses_body_with_a_json_string_metadata_is_logged_against_the_key(
|
||||
proxy_logging, monkeypatch
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
recorded: list[dict] = []
|
||||
|
||||
class _RecordingLogger(CustomLogger):
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
recorded.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm, "callbacks", [_RecordingLogger()])
|
||||
proxy_logging.alert_types = []
|
||||
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data={"model": "internal-model", "input": "hi", "metadata": json.dumps({"tags": ["checkout"]})},
|
||||
original_exception=litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-hashed", key_alias="checkout-key", request_route="/v1/responses"),
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
route="/v1/responses",
|
||||
)
|
||||
|
||||
assert len(recorded) == 1
|
||||
assert recorded[0]["standard_logging_object"]["metadata"]["user_api_key_alias"] == "checkout-key"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue