diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index f441ce71ab9..1fddc527ec8 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -11,17 +11,35 @@ import sys sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +import functools import tempfile -from typing import Optional +from contextvars import ContextVar +from typing import TYPE_CHECKING, ClassVar, Literal, Optional from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import walk_user_text +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj GUARDRAIL_NAME = "hide_secrets" +GUARDRAIL_PROVIDER = "hide-secrets" + +# Per-invocation tally of redacted secrets by detect-secrets plugin type; None +# means the guardrail did not run, so _process_response records nothing. +_masked_entity_count: ContextVar[Optional[dict]] = ContextVar( + "hide_secrets_masked_entity_count", default=None +) + _custom_plugins_path = "file://" + os.path.join( os.path.dirname(os.path.abspath(__file__)), "secrets_plugins" ) @@ -422,6 +440,10 @@ _default_detect_secrets_config = { class _ENTERPRISE_SecretDetection(CustomGuardrail): + # Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail + # path skips should_run_check and never sees data["prompt"]). + use_native_lifecycle_hooks: ClassVar[bool] = True + def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs): self.user_defined_detect_secrets_config = detect_secrets_config super().__init__(**kwargs) @@ -455,6 +477,26 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): return detected_secrets + def redact_text(self, text: str, source: str = "message") -> str: + """Replace every detected secret in ``text`` with ``[REDACTED]`` and + tally the detected types into the per-invocation masked-entity count.""" + detected_secrets = self.scan_message_for_secrets(text) + if not detected_secrets: + return text + counts = _masked_entity_count.get() + if counts is not None: + for secret in detected_secrets: + counts[secret["type"]] = counts.get(secret["type"], 0) + 1 + secret_types = [secret["type"] for secret in detected_secrets] + verbose_proxy_logger.warning( + f"Detected and redacted secrets in {source}: {secret_types}" + ) + return functools.reduce( + lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"), + detected_secrets, + text, + ) + async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool: if user_api_key_dict.permissions is not None: if GUARDRAIL_NAME in user_api_key_dict.permissions: @@ -463,7 +505,45 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): return True + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Unified-interface entrypoint, used by /guardrails/apply_guardrail + (the UI test playground). Proxied traffic keeps using + ``async_pre_call_hook``, see ``use_native_lifecycle_hooks``.""" + texts = inputs.get("texts") + if not texts or not any(texts): + return inputs + _masked_entity_count.set({}) + return {**inputs, "texts": [self.redact_text(text) for text in texts]} + + def _redact_prompt(self, data: dict) -> int: + """Redact ``data["prompt"]`` (the text-completion shape, which + ``walk_user_text`` does not cover) and return how many non-empty + strings were inspected.""" + prompt = data.get("prompt") + if isinstance(prompt, str): + if not prompt: + return 0 + data["prompt"] = self.redact_text(prompt, source="prompt") + return 1 + if isinstance(prompt, list): + data["prompt"] = [ # mutable-ok: data["prompt"] is a list on the wire + self.redact_text(item, source="prompt") + if isinstance(item, str) and item + else item + for item in prompt + ] + return sum(1 for item in prompt if isinstance(item, str) and item) + return 0 + #### CALL HOOKS - proxy only #### + @log_guardrail_information async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -471,53 +551,84 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): data: dict, call_type: str, # "completion", "embeddings", "image_generation", "moderation" ): + _masked_entity_count.set(None) if await self.should_run_check(user_api_key_dict) is False: return + _masked_entity_count.set({}) + # Covers multimodal list content + Responses-API input. - def _redact_message_text(text: str) -> str: - detected_secrets = self.scan_message_for_secrets(text) - for secret in detected_secrets: - text = text.replace(secret["value"], "[REDACTED]") - if detected_secrets: - secret_types = [secret["type"] for secret in detected_secrets] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in message: {secret_types}" - ) - return text + inspected = walk_user_text(data, self.redact_text) + self._redact_prompt(data) - walk_user_text(data, _redact_message_text) + if inspected == 0: + # Image-only, empty-text, and unsupported payloads inspected + # nothing, so recording "allow" would count a run that never + # looked at any content. + _masked_entity_count.set(None) - if "prompt" in data: - if isinstance(data["prompt"], str): - detected_secrets = self.scan_message_for_secrets(data["prompt"]) - for secret in detected_secrets: - data["prompt"] = data["prompt"].replace( - secret["value"], "[REDACTED]" - ) - if len(detected_secrets) > 0: - secret_types = [secret["type"] for secret in detected_secrets] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in prompt: {secret_types}" - ) - elif isinstance(data["prompt"], list): - # Index back into the list — assigning to ``item`` would only - # rebind the loop variable and leave ``data["prompt"]`` - # carrying the unredacted secret. - for idx, item in enumerate(data["prompt"]): - if isinstance(item, str): - detected_secrets = self.scan_message_for_secrets(item) - for secret in detected_secrets: - item = item.replace(secret["value"], "[REDACTED]") - data["prompt"][idx] = item - if len(detected_secrets) > 0: - secret_types = [ - secret["type"] for secret in detected_secrets - ] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in prompt: {secret_types}" - ) - - # ``data["input"]`` (Responses API and embeddings/moderation) is - # already covered by ``walk_user_text`` above. return + + def _process_response( + self, + response: Optional[dict], + request_data: dict, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, + original_inputs: Optional[dict] = None, + ): + """Record allow/mask plus the masked-entity tally for a completed run. + + Records nothing when the guardrail inspected nothing (opted-out key, + empty inputs) or when the instance has no guardrail_name (legacy + ``litellm_settings.callbacks`` deployments, which predate guardrail + telemetry and stay without it). + """ + counts = _masked_entity_count.get() + _masked_entity_count.set(None) + if counts is None or self.guardrail_name is None: + return response + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response="mask" if counts else "allow", + request_data=request_data, + guardrail_status="success", + duration=duration, + start_time=start_time, + end_time=end_time, + event_type=event_type, + guardrail_provider=GUARDRAIL_PROVIDER, + masked_entity_count=counts, + ) + return response + + def _process_error( + self, + e: Exception, + request_data: dict, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, + ): + """Label the failed run with this guardrail's provider so error rows + group with the successful ones in the monitor. Nameless legacy + instances record nothing, matching ``_process_response``.""" + _masked_entity_count.set(None) + if self.guardrail_name is None: + raise e + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=e, + request_data=request_data, + guardrail_status=( + "guardrail_intervened" + if self._is_guardrail_intervention(e) + else "guardrail_failed_to_respond" + ), + duration=duration, + start_time=start_time, + end_time=end_time, + event_type=event_type, + guardrail_provider=GUARDRAIL_PROVIDER, + ) + raise e diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b83e9b395a8..cb7da32f857 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -4,6 +4,7 @@ import logging import time from collections.abc import Mapping, Sequence from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response @@ -1164,6 +1165,12 @@ def _store_cost_breakdown_in_logging_obj( # Don't fail the main cost calculation if breakdown storage fails +def _without_provider_stated_cost(usage: Usage | None) -> Usage | None: + if usage is None or getattr(usage, "cost", None) is None: + return usage + return usage.model_copy(update=MappingProxyType({"cost": None})) + + def completion_cost( completion_response: object | None = None, model: str | None = None, @@ -1243,7 +1250,10 @@ def completion_cost( cache_creation_input_tokens: int | None = None cache_read_input_tokens: int | None = None audio_transcription_file_duration: float = 0.0 - cost_per_token_usage_object: Final[Usage | None] = _get_usage_object(completion_response=completion_response) + provider_usage_object: Final = _get_usage_object(completion_response=completion_response) + cost_per_token_usage_object: Final[Usage | None] = ( + _without_provider_stated_cost(provider_usage_object) if custom_pricing else provider_usage_object + ) rerank_billed_units: RerankBilledUnits | None = None # Extract service_tier from optional_params if not provided directly diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 372c9bf6b91..c8e5610b4a9 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1485,7 +1485,7 @@ def log_guardrail_information(func): if func.__name__ == "apply_guardrail" and "inputs" in kwargs: original_inputs = kwargs.get("inputs") - logging_obj: Final = kwargs.get("logging_obj") + logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj") self_recorded_token: Final = _guardrail_self_recorded.set(False) try: response: Final = await func(*args, **kwargs) @@ -1527,7 +1527,7 @@ def log_guardrail_information(func): if func.__name__ == "apply_guardrail" and "inputs" in kwargs: original_inputs = kwargs.get("inputs") - logging_obj: Final = kwargs.get("logging_obj") + logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj") self_recorded_token: Final = _guardrail_self_recorded.set(False) try: response: Final = func(*args, **kwargs) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 480b1921c18..355faa41e68 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -54,6 +54,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call" _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) +_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) def _next_sync_or_exhausted(it: Any) -> object: @@ -1886,8 +1887,8 @@ class CustomStreamWrapper: @staticmethod def _resolve_provider_reported_cost(usage_cost: object) -> float | None: """ - Providers report usage.cost either as a number or, for Perplexity, as a - breakdown object whose total lives under ``total_cost``. + Providers report usage.cost either as a number or as a breakdown object + whose total lives under ``total_cost``. """ if isinstance(usage_cost, bool): return None @@ -1900,12 +1901,10 @@ class CustomStreamWrapper: @staticmethod def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", + custom_llm_provider: str | None, ) -> None: - """ - If the assembled response carries a provider-reported cost on - usage.cost, copy it into _hidden_params so litellm's cost - calculator uses it instead of a token-based estimate. - """ + if custom_llm_provider not in _USAGE_COST_HEADER_PROVIDERS: + return _usage: Final[Usage | None] = getattr(response, "usage", None) _cost: Final = CustomStreamWrapper._resolve_provider_reported_cost(getattr(_usage, "cost", None)) if _cost is not None: @@ -2020,7 +2019,7 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) setattr( response, @@ -2270,7 +2269,7 @@ class CustomStreamWrapper: response: Final = self.model_response_creator() if complete_streaming_response is not None: - self._propagate_usage_cost_to_hidden_params(complete_streaming_response) + self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) setattr( response, diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index db1a0fc89a3..7638229bc32 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -24,7 +24,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -121,11 +120,10 @@ class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, @@ -145,11 +143,10 @@ class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 9624a721870..c8d2b7fe522 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -99,12 +99,9 @@ class RouterVectorStoreEmbeddingExecutor: ) return bool(resolved) or model in deployment_models - def _embeds_through_sdk(self, model: str, configuration: Mapping[str, object]) -> bool: - return bool(configuration) and not self._router_serves(model) - def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if self._embeds_through_sdk(model, configuration): + if not self._router_serves(model): return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, embedding_kwargs) return self.router.embedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, @@ -114,7 +111,7 @@ class RouterVectorStoreEmbeddingExecutor: async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if self._embeds_through_sdk(model, configuration): + if not self._router_serves(model): return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, embedding_kwargs) return await self.router.aembedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, @@ -153,7 +150,6 @@ class BaseVectorStoreConfig: litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: Router | None = None, ) -> tuple[str, dict]: pass @@ -166,7 +162,6 @@ class BaseVectorStoreConfig: litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: Router | None = None, ) -> tuple[str, dict]: """ Optional async version of transform_search_vector_store_request. @@ -182,7 +177,6 @@ class BaseVectorStoreConfig: litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, extra_body=extra_body, - router=router, ) @abstractmethod @@ -271,7 +265,6 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: pass @@ -285,7 +278,6 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: return self.transform_search_vector_store_request( @@ -296,7 +288,6 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, extra_body=extra_body, - router=router, embedding_executor=embedding_executor, ) @@ -338,11 +329,10 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): query_text: str, litellm_params: Mapping[str, object], embedding_executor: VectorStoreEmbeddingExecutor | None, - router: Router | None = None, ) -> Sequence[float]: model: Final = self.query_embedding_model(litellm_params) configuration: Final = self.query_embedding_configuration(litellm_params) - executor: Final = self.query_embedding_executor(embedding_executor, router) + executor: Final = self.query_embedding_executor(embedding_executor, None) try: response: Final = executor.embed(model, query_text, configuration) except Exception as e: @@ -354,11 +344,10 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): query_text: str, litellm_params: Mapping[str, object], embedding_executor: VectorStoreEmbeddingExecutor | None, - router: Router | None = None, ) -> Sequence[float]: model: Final = self.query_embedding_model(litellm_params) configuration: Final = self.query_embedding_configuration(litellm_params) - executor: Final = self.query_embedding_executor(embedding_executor, router) + executor: Final = self.query_embedding_executor(embedding_executor, None) try: response: Final = await executor.aembed(model, query_text, configuration) except Exception as e: @@ -408,7 +397,6 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, ) -> NoReturn: raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP request shape") diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index bad17a2181d..2d72db0cdba 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -27,7 +27,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -197,7 +196,6 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: if isinstance(query, list): query = " ".join(query) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 0f6966b0ae2..71a598a6fe7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -184,7 +184,6 @@ if TYPE_CHECKING: AnthropicMessagesStreamingResponse, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig - from litellm.router import Router from litellm.types.llms.openai_evals import ( CancelEvalResponse, CancelRunResponse, @@ -9709,7 +9708,6 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - router: "Router | None" = None, ) -> VectorStoreSearchResponse: if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig): self._pre_call_direct_vector_store_search( @@ -9760,7 +9758,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, embedding_executor=embedding_executor, ) else: @@ -9775,7 +9772,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) @@ -9828,7 +9824,6 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, - router: "Router | None" = None, ) -> VectorStoreSearchResponse | Coroutine[object, object, VectorStoreSearchResponse]: if _is_async: return self.async_vector_store_search_handler( @@ -9844,7 +9839,6 @@ class BaseLLMHTTPHandler: extra_body=extra_body, timeout=timeout, client=client, - router=router, ) if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig): @@ -9893,7 +9887,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, embedding_executor=embedding_executor, ) else: @@ -9908,7 +9901,6 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, - router=router, ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 82586b1f638..f6525a449b6 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -33,7 +33,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -169,7 +168,6 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Mapping[str, object] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: """ Transform search request to Gemini's generateContent format. diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 4f3c366d8c1..70d3649debd 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -24,7 +24,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -129,11 +128,10 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, @@ -153,11 +151,10 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: Mapping[str, object], extra_body: Mapping[str, object] | None = None, - router: Router | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, ) -> tuple[str, dict[str, object]]: query_text: Final = self.query_text(query) - query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor, router) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) return self._search_request( vector_store_id, query_text, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 4e925494039..f6c093f2e2a 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -21,7 +21,6 @@ from litellm.utils import add_openai_metadata if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -100,7 +99,6 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = f"{api_base}/{encoded_vector_store_id}/search" diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 9de1f589ae4..e4b06c36bf4 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -8,7 +8,6 @@ from litellm.types.vector_stores import VectorStoreSearchOptionalRequestParams if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -81,7 +80,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url: Final = f"{api_base}/{encoded_vector_store_id}/search" diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ffa6c9e1076..282cb7a92a7 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -17,7 +17,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -93,7 +92,6 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict]: """RAGFlow vector stores are management-only, search is not supported.""" raise NotImplementedError("RAGFlow vector stores support dataset management only, not search/retrieval") diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 733358381fe..a9902a0d27c 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -1,9 +1,12 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx -from litellm.caching._embedding_router import resolve_embedding_router -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -18,16 +21,18 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.router import Router else: LiteLLMLoggingObj = Any +_DEFAULT_QUERY_EMBEDDING_MODEL: Final = "text-embedding-3-small" +_DEFAULT_TOP_K: Final = 5 -class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): + +class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS S3 Vectors.""" def __init__(self) -> None: - BaseVectorStoreConfig.__init__(self) + BaseQueryEmbeddingVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: @@ -59,141 +64,94 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): return headers def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str: - # Resolve region the same way the ingestion path does: - # dynamic param -> AWS_REGION_NAME -> AWS_REGION -> default (us-west-2) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(litellm_params.get("aws_region_name")) return f"https://s3vectors.{aws_region_name}.api.aws" - def _resolve_query_embedding_router(self, embedding_model: str, router: "Router | None") -> "Router | None": - """Return the router iff it serves ``embedding_model`` as a deployment.""" - if router is None: - return None - model_list: Final = [ - dict(m) for m in (router.get_model_list() or ()) - ] # mutable-ok: resolve_embedding_router requires list[dict] - return resolve_embedding_router(embedding_model=embedding_model, llm_router=router, llm_model_list=model_list) + @staticmethod + def query_embedding_model(litellm_params: Mapping[str, object]) -> str: + configured: Final = litellm_params.get("litellm_embedding_model") or litellm_params.get("embedding_model") + return configured if isinstance(configured, str) and configured else _DEFAULT_QUERY_EMBEDDING_MODEL + + @staticmethod + def _query_target(vector_store_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]: + if ":" in vector_store_id: + bucket_name, index_name = vector_store_id.split(":", 1) + return bucket_name, index_name + bucket_name_from_params: Final = litellm_params.get("vector_bucket_name") + if not isinstance(bucket_name_from_params, str) or not bucket_name_from_params: + raise ValueError( + "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " + "or vector_bucket_name must be provided in litellm_params" + ) + return bucket_name_from_params, vector_store_id + + @staticmethod + def _query_request( + bucket_name: str, + index_name: str, + query_text: str, + query_vector: Sequence[float], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + ) -> tuple[str, dict[str, object]]: + litellm_logging_obj.model_call_details["query"] = query_text + return f"{api_base}/QueryVectors", { + "vectorBucketName": bucket_name, + "indexName": index_name, + "queryVector": {"float32": list(query_vector)}, + "topK": vector_store_search_optional_params.get("max_num_results", _DEFAULT_TOP_K), + "returnDistance": True, + "returnMetadata": True, + } def transform_search_vector_store_request( self, vector_store_id: str, - query: str | list[str], + query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, - ) -> tuple[str, dict]: - """Sync version - generates embedding synchronously.""" - # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name - # If not in that format, try to construct it from litellm_params - bucket_name: str - index_name: str - - if ":" in vector_store_id: - bucket_name, index_name = vector_store_id.split(":", 1) - else: - # Try to get bucket_name from litellm_params - bucket_name_from_params: Final = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): - raise ValueError( - "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " - "or vector_bucket_name must be provided in litellm_params" - ) - bucket_name = bucket_name_from_params - index_name = vector_store_id - - if isinstance(query, list): - query = " ".join(query) - - # Generate embedding for the query - embedding_model: Final = litellm_params.get("embedding_model", "text-embedding-3-small") - embedding_router: Final = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router) - - import litellm as litellm_module - - embedding_input: Final = [query] # mutable-ok: the embedding API takes list input - embedding_response: Final = ( - embedding_router.embedding(model=embedding_model, input=embedding_input) - if embedding_router is not None - else litellm_module.embedding(model=embedding_model, input=embedding_input) + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + bucket_name, index_name = self._query_target(vector_store_id, litellm_params) + query_text: Final = self.query_text(query) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) + return self._query_request( + bucket_name, + index_name, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, ) - query_embedding: Final = embedding_response.data[0]["embedding"] - - url: Final = f"{api_base}/QueryVectors" - - request_body: Final[dict[str, Any]] = { - "vectorBucketName": bucket_name, - "indexName": index_name, - "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 - "returnDistance": True, - "returnMetadata": True, - } - - litellm_logging_obj.model_call_details["query"] = query - return url, request_body async def atransform_search_vector_store_request( self, vector_store_id: str, - query: str | list[str], + query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - router: "Router | None" = None, - ) -> tuple[str, dict]: - """Async version - generates embedding asynchronously.""" - # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name - # If not in that format, try to construct it from litellm_params - bucket_name: str - index_name: str - - if ":" in vector_store_id: - bucket_name, index_name = vector_store_id.split(":", 1) - else: - # Try to get bucket_name from litellm_params - bucket_name_from_params: Final = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): - raise ValueError( - "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " - "or vector_bucket_name must be provided in litellm_params" - ) - bucket_name = bucket_name_from_params - index_name = vector_store_id - - if isinstance(query, list): - query = " ".join(query) - - # Generate embedding for the query asynchronously - embedding_model: Final = litellm_params.get("embedding_model", "text-embedding-3-small") - embedding_router: Final = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router) - - import litellm as litellm_module - - embedding_input: Final = [query] # mutable-ok: the embedding API takes list input - embedding_response: Final = ( - await embedding_router.aembedding(model=embedding_model, input=embedding_input) - if embedding_router is not None - else await litellm_module.aembedding(model=embedding_model, input=embedding_input) + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + bucket_name, index_name = self._query_target(vector_store_id, litellm_params) + query_text: Final = self.query_text(query) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) + return self._query_request( + bucket_name, + index_name, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, ) - query_embedding: Final = embedding_response.data[0]["embedding"] - - url: Final = f"{api_base}/QueryVectors" - - request_body: Final[dict[str, Any]] = { - "vectorBucketName": bucket_name, - "indexName": index_name, - "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 - "returnDistance": True, - "returnMetadata": True, - } - - litellm_logging_obj.model_call_details["query"] = query - return url, request_body def sign_request( self, @@ -226,21 +184,13 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if not source_text: continue - # Extract file information from metadata chunk_index = metadata.get("chunk_index", "0") file_id = f"s3-vectors-chunk-{chunk_index}" filename = metadata.get("filename", f"document-{chunk_index}") - # S3 Vectors returns distance, convert to similarity score (0-1) - # Lower distance = higher similarity - # We'll normalize using 1 / (1 + distance) to get a 0-1 score distance = item.get("distance") score = None if distance is not None: - # Convert distance to similarity score between 0 and 1 - # For cosine distance: similarity = 1 - distance - # For euclidean: use 1 / (1 + distance) - # Assuming cosine distance here score = max(0.0, min(1.0, 1.0 - float(distance))) results.append( @@ -265,7 +215,6 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): headers=response.headers, ) - # Vector store creation is not yet implemented def transform_create_vector_store_request( self, vector_store_create_optional_params, diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 36b57e7c995..5c250fc1a7e 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -21,7 +21,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -162,7 +161,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Mapping[str, object] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict[str, object]]: """ Transform search request for Vertex AI RAG API diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index f0812e3ed9f..0bcf16ee06f 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -25,7 +25,6 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -246,7 +245,6 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Mapping[str, object] | None = None, - router: "Router | None" = None, ) -> tuple[str, dict[str, object]]: """ Transform a search request for the Vertex AI Search (Discovery Engine) API. diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index ae5849812bf..09af8544b28 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -11,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, strip_name_from_messages, ) -from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.llms.xai.cost_calculator import ( apply_server_side_tool_usage_details_to_usage, ) @@ -30,6 +31,13 @@ from ...openai.chat.gpt_transformation import ( ) +def _usage_restated_from_xai_ticks(usage: Usage | None) -> Usage | None: + reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) + if usage is None or reported_cost is None: + return None + return usage.model_copy(update=MappingProxyType({"cost": reported_cost})) + + class XAIChatConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> str | None: @@ -283,6 +291,9 @@ class XAIChatConfig(OpenAIGPTConfig): self._fold_reasoning_tokens_into_completion(response) self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) + restated_usage: Final = _usage_restated_from_xai_ticks(getattr(response, "usage", None)) + if restated_usage is not None: + response.usage = restated_usage return response @staticmethod @@ -411,4 +422,8 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) - return super().chunk_parser(chunk) + parsed_chunk: Final = super().chunk_parser(chunk) + restated_usage: Final = _usage_restated_from_xai_ticks(getattr(parsed_chunk, "usage", None)) + if restated_usage is not None: + parsed_chunk.usage = restated_usage + return parsed_chunk diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index f122098332e..cf76a851a86 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -8,6 +8,17 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ProviderSpecificModelInfo +USD_TICKS_PER_DOLLAR: Final = 10_000_000_000 + + +def xai_reported_cost_in_usd(cost_in_usd_ticks: object) -> float | None: + """xAI bills in ticks of a dollar: https://docs.x.ai/developers/cost-tracking""" + if not isinstance(cost_in_usd_ticks, int) or isinstance(cost_in_usd_ticks, bool): + return None + if cost_in_usd_ticks < 0: + return None + return cost_in_usd_ticks / USD_TICKS_PER_DOLLAR + class XAIModelInfo(BaseLLMModelInfo): def get_provider_info( diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index dd77b8d5d09..164568451d4 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -1,9 +1,11 @@ """ Helper util for handling XAI-specific cost calculation +- Prefers the cost xAI reports on the response over recomputing it locally - Uses the generic cost calculator which already handles tiered pricing correctly - Handles XAI-specific reasoning token billing (billed as part of completion tokens) """ +import math from collections.abc import Mapping from typing import TYPE_CHECKING, Final @@ -36,6 +38,17 @@ def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping usage.prompt_tokens_details = prompt_tokens_details # rebind-ok: write details onto caller usage +def _cost_reported_by_xai(usage: "Usage") -> float | None: + reported_cost: Final[object] = getattr(usage, "cost", None) + if not isinstance(reported_cost, (int, float)) or isinstance(reported_cost, bool): + return None + if not math.isfinite(reported_cost): + return None + if reported_cost < 0: + return None + return float(reported_cost) + + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. @@ -48,6 +61,10 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ + reported_cost: Final = _cost_reported_by_xai(usage) + if reported_cost is not None: + return 0.0, reported_cost + # XAI-specific completion cost: completion is billed as visible + reasoning # tokens. Detect when the transformation layer already folded them so we # don't double-count; fall back to raw xAI shape for callers that bypass @@ -112,6 +129,9 @@ def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> floa Per-call rate comes from model_info.search_context_cost_per_query when set, otherwise the default xAI tools rate ($5 / 1k calls). """ + if _cost_reported_by_xai(usage) is not None: + return 0.0 + details: Final = getattr(usage, "server_side_tool_usage_details", None) if not isinstance(details, Mapping): return 0.0 diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index d79e7d4c146..acf03d88911 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,17 +1,44 @@ -from typing import Any, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final + +import httpx import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig -from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool 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 = Any + + +def _usage_restated_from_xai_ticks(usage: ResponseAPIUsage | None) -> ResponseAPIUsage | None: + reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) + if usage is None or reported_cost is None: + return None + return usage.model_copy(update=MappingProxyType({"cost": reported_cost})) + class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): """ @@ -250,6 +277,41 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return f"{api_base}/responses" + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + response: Final = super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + + restated_usage: Final = _usage_restated_from_xai_ticks(response.usage) + if restated_usage is not None: + response.usage = restated_usage + return response + + def transform_streaming_response( + self, + model: str, + parsed_chunk: dict, # mutable-ok: overrides the base class signature + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIStreamingResponse: + event: Final = super().transform_streaming_response( + model=model, + parsed_chunk=parsed_chunk, + logging_obj=logging_obj, + ) + if not isinstance(event, (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent)): + return event + restated_usage: Final = _usage_restated_from_xai_ticks(event.response.usage) + if restated_usage is not None: + event.response.usage = restated_usage + return event + def supports_native_websocket(self) -> bool: """XAI does not support native WebSocket for Responses API""" return False diff --git a/litellm/main.py b/litellm/main.py index 0128e4defe5..0bca4a7350e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8595,9 +8595,19 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N return TextCompletionResponse(**response) +_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) + + +def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool: + if logging_obj is None: + return False + provider: Final[object] = logging_obj.model_call_details.get("custom_llm_provider") + return provider in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS + + def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None: usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None) - if isinstance(usage_cost, (int, float)): + if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj): return float(usage_cost) if logging_obj is not None: return None diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 3d2ed641a30..744b2959c73 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -8,7 +8,7 @@ import json import os from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone -from types import UnionType +from types import MappingProxyType, UnionType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin from urllib.parse import urlparse @@ -51,6 +51,9 @@ from litellm.types.guardrails import ( SupportedGuardrailIntegrations, ToolPermissionGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.hide_secrets import ( + HideSecretsGuardrailConfigModel, +) if TYPE_CHECKING: from types import CodeType @@ -1401,7 +1404,11 @@ async def get_guardrail_ui_settings(): provider: [hook.value for hook in hooks] for provider, guardrail_class in guardrail_class_registry.items() if (hooks := guardrail_class.get_supported_event_hooks()) is not None - } + } | MappingProxyType( + # hide-secrets lives in the enterprise package, not in the registry + # above; it only runs on pre_call. + {SupportedGuardrailIntegrations.HIDE_SECRETS.value: [GuardrailEventHooks.pre_call.value]} + ) return GuardrailUIAddGuardrailSettings( supported_entities=[entity.value for entity in PiiEntityType], @@ -1953,12 +1960,18 @@ async def get_provider_specific_params(): tool_permission_fields["ui_friendly_name"] = ToolPermissionGuardrailConfigModel.ui_friendly_name() + # hide-secrets lives in the enterprise package, not in the registry loop below. + hide_secrets_fields: Final = _get_fields_from_model(HideSecretsGuardrailConfigModel) + + hide_secrets_fields["ui_friendly_name"] = HideSecretsGuardrailConfigModel.ui_friendly_name() + # Return the provider-specific parameters provider_params: Final = { SupportedGuardrailIntegrations.BEDROCK.value: bedrock_fields, SupportedGuardrailIntegrations.PRESIDIO.value: presidio_fields, SupportedGuardrailIntegrations.LAKERA_V2.value: lakera_v2_fields, SupportedGuardrailIntegrations.TOOL_PERMISSION.value: tool_permission_fields, + SupportedGuardrailIntegrations.HIDE_SECRETS.value: hide_secrets_fields, } ### get the config model for the guardrail - go through the registry and get the config model for the guardrail diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 466b100ea1f..12e1f1a03f3 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -443,7 +443,9 @@ class SAMLAuthHandler: last_name: Final = SAMLAuthHandler._attribute_value( attributes, "SAML_ATTRIBUTE_LAST_NAME", _LAST_NAME_ATTRIBUTE_CANDIDATES ) - role_value = SAMLAuthHandler._attribute_value(attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES) + role_values: Final = SAMLAuthHandler._attribute_values( + attributes, "SAML_ATTRIBUTE_ROLE", _ROLE_ATTRIBUTE_CANDIDATES + ) team_ids: Final = SAMLAuthHandler._attribute_values( attributes, "SAML_ATTRIBUTE_TEAM_IDS", _TEAM_IDS_ATTRIBUTE_CANDIDATES ) @@ -464,7 +466,7 @@ class SAMLAuthHandler: picture=None, provider="saml", team_ids=team_ids, - user_role=get_litellm_user_role(role_value) if role_value else None, + user_role=get_litellm_user_role(role_values), ) except ValidationError as e: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 070df97d09a..4414eed97b2 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -4,12 +4,44 @@ Types for the management endpoints Might include fastapi/proxy requirements.txt related imports """ +from collections.abc import Iterable, Sequence from typing import Any, Final, cast from fastapi_sso.sso.base import OpenID from litellm.proxy._types import LitellmUserRoles +# Ordered highest to lowest privilege +LITELLM_USER_ROLE_HIERARCHY: Final = ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, +) + + +def highest_privilege_role(roles: Iterable[LitellmUserRoles]) -> LitellmUserRoles | None: + """ + Pick the highest privilege role out of the roles an IdP asserted for one user. + + IdPs do not guarantee ordering within a multi-valued role claim, so a user holding + several roles resolves to the most privileged one rather than whichever came first. + Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay + deterministic. + + Args: + roles: The roles resolved from the claim + + Returns: + The highest privilege role, or None if `roles` is empty + """ + resolved: Final = frozenset(roles) + if not resolved: + return None + + ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) + return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + def is_valid_litellm_user_role(role_str: str) -> bool: """ @@ -28,12 +60,22 @@ def is_valid_litellm_user_role(role_str: str) -> bool: return False -def get_litellm_user_role(role_str) -> LitellmUserRoles | None: +def _role_from_claim_value(role_str: object) -> LitellmUserRoles | None: + if not isinstance(role_str, str): + return None + # Use _value2member_map_ for O(1) lookup, case-insensitive + result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) + return cast(LitellmUserRoles | None, result) + + +def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: """ Convert a string (or list of strings) to a LitellmUserRoles enum if valid (case-insensitive). Handles list inputs since some SSO providers (e.g., Keycloak) return roles - as arrays like ["proxy_admin"] instead of plain strings. + as arrays like ["proxy_admin"] instead of plain strings. A claim carrying several + roles resolves to the highest privilege one, so a user does not lose access just + because the IdP listed a weaker role first. Args: role_str: String or list to convert (e.g., "proxy_admin", ["proxy_admin"]) @@ -41,16 +83,12 @@ def get_litellm_user_role(role_str) -> LitellmUserRoles | None: Returns: LitellmUserRoles enum if valid, None otherwise """ - try: - if isinstance(role_str, list): - if len(role_str) == 0: - return None - role_str = role_str[0] - # Use _value2member_map_ for O(1) lookup, case-insensitive - result: Final = LitellmUserRoles._value2member_map_.get(role_str.lower()) - return cast(LitellmUserRoles | None, result) - except Exception: - return None + if isinstance(role_str, (list, tuple)): + entries: Final = cast(Sequence[object], role_str) # cast-ok: isinstance narrows the claim, not its elements + return highest_privilege_role( + role for role in (_role_from_claim_value(entry) for entry in entries) if role is not None + ) + return _role_from_claim_value(role_str) class CustomOpenID(OpenID): diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1feefa5725d..84150ef7935 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.sso_helper_utils import ( ) from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add from litellm.proxy.management_endpoints.types import ( + LITELLM_USER_ROLE_HIERARCHY, CustomOpenID, get_litellm_user_role, is_valid_litellm_user_role, @@ -809,15 +810,6 @@ def normalize_email(email: str | None) -> str | None: return email.lower() if isinstance(email, str) else email -# Ordered highest to lowest privilege -LITELLM_USER_ROLE_HIERARCHY: Final = ( - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, -) - - def determine_role_from_groups( user_groups: list[str], role_mappings: "RoleMappings", @@ -4312,14 +4304,7 @@ class MicrosoftSSOHandler: listed first. Roles the hierarchy does not rank (org_admin, team, customer) resolve by name to stay deterministic """ - resolved: Final = frozenset( - role for role in (get_litellm_user_role(role_str) for role_str in app_roles or ()) if role is not None - ) - if not resolved: - return None - - ranked: Final = next((role for role in LITELLM_USER_ROLE_HIERARCHY if role in resolved), None) - return ranked if ranked is not None else min(resolved, key=lambda role: role.value) + return get_litellm_user_role(tuple(app_roles or ())) @staticmethod def get_app_roles_from_id_token(id_token: str | None) -> list[str]: diff --git a/litellm/proxy/public_endpoints/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json index c2b13b81542..6a6642d4211 100644 --- a/litellm/proxy/public_endpoints/autorouter_presets.json +++ b/litellm/proxy/public_endpoints/autorouter_presets.json @@ -1,4 +1,26 @@ { + "1m_context": { + "label": "1M Context", + "description": "Routes across models with 1M-token context windows: Luna for simple queries, Terra for medium, Opus 5 for complex, Opus 5 at high thinking for reasoning.", + "complexity_router_config": { + "tiers": { + "SIMPLE": ["gpt-5.6-luna"], + "MEDIUM": ["gpt-5.6-terra"], + "COMPLEX": ["claude-opus-5"], + "REASONING": ["claude-opus-5"] + }, + "tier_model_configs": { + "REASONING": [{ "model_name": "claude-opus-5", "litellm_params": { "reasoning_effort": "high" } }] + }, + "classifier_type": "heuristic_v2", + "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", + "session_affinity": false, + "modality_routing": false, + "modality_pin_override": false, + "deployment_affinity": true + } + }, "anthropic_family": { "label": "Anthropic Family", "description": "Routes across the Claude model family: Haiku for simple queries, Sonnet for medium, Opus for complex, Opus at high thinking for reasoning.", diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/hide_secrets.py b/litellm/types/proxy/guardrails/guardrail_hooks/hide_secrets.py new file mode 100644 index 00000000000..3100968e8b1 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/hide_secrets.py @@ -0,0 +1,20 @@ +"""Types for the Hide Secrets guardrail.""" + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class HideSecretsGuardrailConfigModel(GuardrailConfigModel): + """Configuration for the Hide Secrets guardrail. Detection runs in-process + on the detect-secrets library; ``detect_secrets_config`` overrides the + bundled plugin set.""" + + detect_secrets_config: dict | None = Field( # mutable-ok: UI type derivation maps dict to "object" + default=None, + description="Optional detect-secrets configuration (plugins_used, filters_used) overriding the bundled plugin set", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Hide Secrets" diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 636bdd4b52e..2fe1965a192 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -482,7 +482,6 @@ def search( timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), - router=router, ) return response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 4fcf650a8bc..be2b30fc189 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2985 + "limit": 2984 }, "ANN002": { "limit": 71 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2917 + "limit": 2916 }, "C401": { "limit": 8 diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 6b162ea90f1..63bad6d0e22 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -376,7 +376,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( timeout=None, client=None, _is_async=False, - router: "litellm.Router | None" = None, + embedding_executor=None, ): litellm_params_dict = ( litellm_params.model_dump(exclude_none=False) diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 2cc9914c9b3..e10ba0f0962 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -187,7 +187,7 @@ class TestRouterEmbeddingIntegration: assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-large", ["async query"]) @pytest.mark.asyncio - async def test_router_executor_rejects_unserved_models_without_explicit_config( + async def test_router_executor_embeds_unserved_models_through_the_sdk( self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -198,12 +198,13 @@ class TestRouterEmbeddingIntegration: metadata={"user_api_key_team_id": "team-a"}, ) - with pytest.raises(litellm.BadRequestError): - executor.embed("openai/text-embedding-3-large", "sync query", {}) - with pytest.raises(litellm.BadRequestError): - await executor.aembed("openai/text-embedding-3-large", "async query", {}) + sync_response = executor.embed("text-embedding-3-large", "sync query", {}) + async_response = await executor.aembed("text-embedding-3-large", "async query", {}) - assert openai_route.call_count == 0 + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer env-key", "text-embedding-3-large", ["sync query"]) + assert _sent(openai_route, 1) == ("Bearer env-key", "text-embedding-3-large", ["async query"]) def test_router_executor_routes_deployment_model_names_through_the_router( self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py new file mode 100644 index 00000000000..dc1cbb9983e --- /dev/null +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -0,0 +1,274 @@ +"""Tests for the hide-secrets guardrail (LIT-3548). + +Covers the three defects from the ticket: +- ``apply_guardrail`` (the UI test playground path) must redact, not echo. +- Guardrail runs must record ``standard_logging_guardrail_information`` so + Spend Logs / the guardrails monitor show activity, with hits ("mask" + + masked_entity_count) distinguishable from clean requests ("allow"). +- Defining ``apply_guardrail`` must NOT reroute proxied traffic off the + native ``async_pre_call_hook`` (per-key opt-out and ``data["prompt"]`` + handling live only on the native path). +""" + +import pytest + +from litellm_enterprise.enterprise_callbacks.secret_detection import ( + _ENTERPRISE_SecretDetection, +) +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth + +AWS_KEY = "AKIAIOSFODNN7EXAMPLE" + + +def _guardrail() -> _ENTERPRISE_SecretDetection: + return _ENTERPRISE_SecretDetection( + guardrail_name="hide-secrets", event_hook="pre_call", default_on=True + ) + + +def _recorded(request_data: dict) -> dict: + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(entries) == 1 + return entries[0] + + +@pytest.mark.asyncio +async def test_apply_guardrail_redacts_secrets(): + """Playground path: the returned texts must carry [REDACTED], not the secret.""" + guardrail = _guardrail() + request_data: dict = {"metadata": {}} + + result = await guardrail.apply_guardrail( + inputs={"texts": [f"my key is {AWS_KEY}, keep it safe"]}, + request_data=request_data, + input_type="request", + ) + + assert result["texts"] == ["my key is [REDACTED], keep it safe"] + + recorded = _recorded(request_data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "mask" + assert recorded["guardrail_provider"] == "hide-secrets" + assert recorded["masked_entity_count"] == {"AWS Access Key": 1} + + +@pytest.mark.asyncio +async def test_apply_guardrail_clean_text_records_allow(): + guardrail = _guardrail() + request_data: dict = {"metadata": {}} + + result = await guardrail.apply_guardrail( + inputs={"texts": ["nothing sensitive here"]}, + request_data=request_data, + input_type="request", + ) + + assert result["texts"] == ["nothing sensitive here"] + + recorded = _recorded(request_data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "allow" + assert recorded["masked_entity_count"] == {} + + +@pytest.mark.asyncio +async def test_pre_call_hook_records_mask_with_entity_count(): + """Live-traffic path: a redaction must be visible in spend-log telemetry.""" + guardrail = _guardrail() + data = { + "messages": [{"role": "user", "content": f"use {AWS_KEY} for auth"}], + "metadata": {}, + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["messages"][0]["content"] == "use [REDACTED] for auth" + + recorded = _recorded(data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "mask" + assert recorded["guardrail_provider"] == "hide-secrets" + assert recorded["masked_entity_count"] == {"AWS Access Key": 1} + + +@pytest.mark.asyncio +async def test_pre_call_hook_clean_request_records_allow(): + """A request with no secrets must be distinguishable from a redacted one.""" + guardrail = _guardrail() + data = { + "messages": [{"role": "user", "content": "what's the weather"}], + "metadata": {}, + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + recorded = _recorded(data) + assert recorded["guardrail_status"] == "success" + assert recorded["guardrail_response"] == "allow" + assert recorded["masked_entity_count"] == {} + + +@pytest.mark.asyncio +async def test_pre_call_hook_opt_out_records_nothing(): + """A key with permissions={"hide_secrets": False} skips redaction, so no + telemetry is recorded: every reader of a recorded entry (guardrail usage + tracking, compliance checks, the spend-log viewer) counts it as a run.""" + guardrail = _guardrail() + content = f"my key is {AWS_KEY}" + data = {"messages": [{"role": "user", "content": content}], "metadata": {}} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(permissions={"hide_secrets": False}), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["messages"][0]["content"] == content # untouched + assert "standard_logging_guardrail_information" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_still_redacts_text_completion_prompt(): + """data["prompt"] (str and list) is a native-hook-only surface; it must + keep redacting now that the class also implements apply_guardrail.""" + guardrail = _guardrail() + data = {"prompt": f"key {AWS_KEY} end", "metadata": {}} + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + assert data["prompt"] == "key [REDACTED] end" + + guardrail = _guardrail() + data = {"prompt": [f"key {AWS_KEY}", "clean"], "metadata": {}} + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + assert data["prompt"] == ["key [REDACTED]", "clean"] + + +def test_proxied_traffic_stays_on_native_hooks(): + """Implementing apply_guardrail must not reroute proxied requests onto the + unified path: that path skips ``should_run_check`` (per-key opt-out) and + never sees ``data["prompt"]``.""" + guardrail = _guardrail() + assert guardrail.uses_apply_guardrail_interface() is True + assert guardrail._deployment_pre_call_target() is guardrail + + +@pytest.mark.asyncio +async def test_apply_guardrail_without_texts_records_nothing(): + """No inputs means nothing was inspected, so no "allow" row is recorded. + Empty strings count as no input: there is no content to inspect.""" + guardrail = _guardrail() + + empty_variants: list[list[str]] = [[], ["", ""]] + for texts in empty_variants: + request_data: dict = {"metadata": {}} + result = await guardrail.apply_guardrail( + inputs={"texts": texts}, request_data=request_data, input_type="request" + ) + assert result == {"texts": texts} + assert "standard_logging_guardrail_information" not in request_data["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data", + [ + pytest.param( + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://x/y.png"}} + ], + } + ], + "metadata": {}, + }, + id="image_only", + ), + pytest.param( + {"messages": [{"role": "user", "content": ""}], "metadata": {}}, + id="empty_message", + ), + pytest.param({"prompt": "", "metadata": {}}, id="empty_prompt"), + pytest.param({"prompt": ["", ""], "metadata": {}}, id="empty_prompt_list"), + ], +) +async def test_pre_call_hook_without_inspectable_text_records_nothing(data: dict): + """A payload the guardrail could not inspect (image-only content, empty + strings) must not record an "allow" run: monitoring would count a check + that never looked at any text.""" + guardrail = _guardrail() + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert "standard_logging_guardrail_information" not in data["metadata"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_mixed_prompt_list_still_redacts_and_records(): + """A prompt list mixing empty and real strings is inspected, so the run is + recorded and the non-empty entry is still redacted.""" + guardrail = _guardrail() + data = {"prompt": ["", f"key {AWS_KEY}"], "metadata": {}} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["prompt"] == ["", "key [REDACTED]"] + recorded = _recorded(data) + assert recorded["guardrail_response"] == "mask" + assert recorded["masked_entity_count"] == {"AWS Access Key": 1} + + +@pytest.mark.asyncio +async def test_legacy_nameless_instance_records_nothing(): + """``litellm_settings.callbacks: ["hide_secrets"]`` builds an arg-less + instance with no guardrail_name. It still redacts, but recording a nameless + entry would flip every spend row's guardrail status with nothing to join on.""" + guardrail = _ENTERPRISE_SecretDetection() + data = { + "messages": [{"role": "user", "content": f"use {AWS_KEY} for auth"}], + "metadata": {}, + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert data["messages"][0]["content"] == "use [REDACTED] for auth" + assert "standard_logging_guardrail_information" not in data["metadata"] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 4c99bce2b0f..4b2a3f9247c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.streaming_handler import ( from litellm.types.utils import ( CompletionTokensDetailsWrapper, Delta, + ModelResponse, ModelResponseStream, PromptTokensDetailsWrapper, StandardLoggingPayload, @@ -1750,7 +1751,7 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert complete_response.usage.cost == 0.00025 # Use the real propagation method from CustomStreamWrapper - CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openrouter") assert "additional_headers" in complete_response._hidden_params assert ( @@ -1769,14 +1770,12 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert provider_cost == 0.00025 -def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): - """ - Regression: Perplexity reports usage.cost as a breakdown object, which used to - blow up the end of the stream with - `float() argument must be a string or a real number, not 'dict'`. - """ +def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator(): import litellm - from litellm.cost_calculator import get_response_cost_from_hidden_params + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, + ) chunks = [ ModelResponseStream( @@ -1828,13 +1827,81 @@ def test_perplexity_streaming_dict_cost_propagates_to_hidden_params(): assert complete_response is not None - CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "perplexity") - assert ( - get_response_cost_from_hidden_params(complete_response._hidden_params) - == 0.00503 + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + assert response_cost_calculator( + response_object=complete_response, + model="perplexity/sonar", + custom_llm_provider="perplexity", + call_type="completion", + optional_params={}, + ) == pytest.approx(0.00503) + + +def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map(): + import litellm + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, ) + model = "openai/streams-cost-in-nanodollars" + litellm.register_model( + { + model: { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + complete_response = ModelResponse( + id="chatcmpl-openai-compatible", + model=model, + choices=[], + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=3_144_000), + ) + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openai") + + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + assert response_cost_calculator( + response_object=complete_response, + model=model, + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) == pytest.approx(2e-5) + + +def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch): + import litellm + from litellm.cost_calculator import ( + get_response_cost_from_hidden_params, + response_cost_calculator, + ) + + complete_response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[], + usage=Usage(completion_tokens=353, prompt_tokens=198, total_tokens=551, cost=0.0009956), + ) + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "xai") + + assert get_response_cost_from_hidden_params(complete_response._hidden_params) is None + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) + assert response_cost_calculator( + response_object=complete_response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + call_type="completion", + optional_params={}, + ) == pytest.approx(0.0009956 * 1.5) + def test_provider_reported_cost_ignores_unusable_shapes(): assert CustomStreamWrapper._resolve_provider_reported_cost(None) is None diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 4b58d220623..e313b749d06 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -1,40 +1,70 @@ +from collections.abc import Mapping from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, +) from litellm.llms.s3_vectors.vector_stores.transformation import ( S3VectorsVectorStoreConfig, ) +from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import VectorStoreSearchResponse +QUERY_VECTOR = [0.1, 0.2, 0.3] -def _mock_router(model_names, sync=False): - """Router mock serving the given embedding model names.""" - router = MagicMock() - router.get_model_list.return_value = [{"model_name": name} for name in model_names] - embedding_response = Mock(data=[{"embedding": [0.1, 0.2, 0.3]}]) - if sync: - router.embedding = MagicMock(return_value=embedding_response) - else: - router.aembedding = AsyncMock(return_value=embedding_response) - return router + +def _embedding_response(vector): + return EmbeddingResponse(data=[{"embedding": vector, "index": 0, "object": "embedding"}]) + + +class _RecordingExecutor: + def __init__(self, vector=QUERY_VECTOR): + self.vector = vector + self.calls = [] + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + self.calls.append((model, query, dict(configuration))) + return _embedding_response(self.vector) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + self.calls.append((model, query, dict(configuration))) + return _embedding_response(self.vector) + + +def _logging_obj(): + logging_obj = Mock() + logging_obj.model_call_details = {} + return logging_obj + + +def _search_kwargs(**overrides): + kwargs = { + "vector_store_id": "test-bucket:test-index", + "query": "test query", + "vector_store_search_optional_params": {}, + "api_base": "https://s3vectors.us-west-2.api.aws", + "litellm_logging_obj": _logging_obj(), + "litellm_params": {}, + "extra_body": None, + } + kwargs.update(overrides) + return kwargs class TestS3VectorsVectorStoreConfig: def test_init(self): - """Test that S3VectorsVectorStoreConfig initializes correctly""" config = S3VectorsVectorStoreConfig() assert config is not None def test_get_supported_openai_params(self): - """Test that supported OpenAI params are returned""" config = S3VectorsVectorStoreConfig() params = config.get_supported_openai_params("test-model") assert "max_num_results" in params def test_get_complete_url(self): - """Test URL generation for S3 Vectors""" config = S3VectorsVectorStoreConfig() litellm_params = {"aws_region_name": "us-west-2"} url = config.get_complete_url(None, litellm_params) @@ -57,180 +87,170 @@ class TestS3VectorsVectorStoreConfig: assert url == "https://s3vectors.eu-west-1.api.aws" def test_get_complete_url_invalid_region_format(self): - """Invalid region format raises""" config = S3VectorsVectorStoreConfig() with pytest.raises(ValueError, match="Invalid AWS region format"): config.get_complete_url(None, {"aws_region_name": "Bad_Region!"}) def test_transform_search_request(self): - """Full request-body transformation with a router-injected embedding""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["text-embedding-3-small"], sync=True) + logging_obj = _logging_obj() + executor = _RecordingExecutor() url, request_body = config.transform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={"max_num_results": 7}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, - router=router, + **_search_kwargs( + vector_store_search_optional_params={"max_num_results": 7}, + litellm_logging_obj=logging_obj, + embedding_executor=executor, + ) ) assert url == "https://s3vectors.us-west-2.api.aws/QueryVectors" assert request_body == { "vectorBucketName": "test-bucket", "indexName": "test-index", - "queryVector": {"float32": [0.1, 0.2, 0.3]}, + "queryVector": {"float32": QUERY_VECTOR}, "topK": 7, "returnDistance": True, "returnMetadata": True, } - assert mock_logging_obj.model_call_details["query"] == "test query" + assert executor.calls == [("text-embedding-3-small", "test query", {})] + assert logging_obj.model_call_details["query"] == "test query" + + @pytest.mark.parametrize( + ("litellm_params", "expected_model"), + [ + ({}, "text-embedding-3-small"), + ({"embedding_model": ""}, "text-embedding-3-small"), + ({"embedding_model": "my-embedding-model"}, "my-embedding-model"), + ({"litellm_embedding_model": "shared-key-model"}, "shared-key-model"), + ( + {"litellm_embedding_model": "shared-key-model", "embedding_model": "legacy-alias"}, + "shared-key-model", + ), + ], + ) + def test_query_embedding_model_accepts_embedding_model_alias(self, litellm_params, expected_model): + assert S3VectorsVectorStoreConfig.query_embedding_model(litellm_params) == expected_model @pytest.mark.asyncio - async def test_atransform_search_uses_router_for_virtual_model(self): - """Regression: router-served embedding models must resolve via the router, - not a bare litellm.aembedding call (which has no deployment credentials).""" + async def test_atransform_search_embeds_alias_and_store_config_through_executor(self): config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["my-embedding-model"]) + executor = _RecordingExecutor(vector=[0.4, 0.5]) - with patch("litellm.aembedding", new=AsyncMock()) as mock_bare_aembedding: # test-quality-ok: guards that the bare-embedding path is not taken; dispatch seam is the behavior under test - url, request_body = await config.atransform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={"embedding_model": "my-embedding-model"}, - extra_body=None, - router=router, + _, request_body = await config.atransform_search_vector_store_request( + **_search_kwargs( + query=["test", "query"], + litellm_params={ + "embedding_model": "my-embedding-model", + "litellm_embedding_config": {"api_key": "store-key"}, + }, + embedding_executor=executor, ) + ) - router.aembedding.assert_awaited_once_with(model="my-embedding-model", input=["test query"]) - mock_bare_aembedding.assert_not_awaited() - assert request_body["queryVector"]["float32"] == [0.1, 0.2, 0.3] - assert request_body["topK"] == 5 # default - - @pytest.mark.asyncio - async def test_atransform_search_falls_back_when_router_does_not_serve_model(self): - """Router present but embedding_model is not a router deployment -> - bare litellm.aembedding keeps working (provider-prefixed + env creds stores).""" - config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["some-other-model"]) - - mock_bare = AsyncMock(return_value=Mock(data=[{"embedding": [0.4, 0.5]}])) - with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on - _, request_body = await config.atransform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={"embedding_model": "azure/text-embedding-3-small"}, - extra_body=None, - router=router, - ) - - mock_bare.assert_awaited_once_with(model="azure/text-embedding-3-small", input=["test query"]) - router.aembedding.assert_not_awaited() + assert executor.calls == [("my-embedding-model", "test query", {"api_key": "store-key"})] assert request_body["queryVector"]["float32"] == [0.4, 0.5] + assert request_body["topK"] == 5 @pytest.mark.asyncio - async def test_atransform_search_without_router_uses_bare_embedding(self): - """Backward compat: no router -> bare litellm.aembedding as before""" + async def test_atransform_search_router_executor_carries_request_metadata(self): + """Regression (LIT-6750): a bare Router alias resolves through the Router with the request's + team metadata on the embedding call, so the embedding is attributed to the calling key and team.""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} + router = MagicMock() + router.aembedding = AsyncMock(return_value=_embedding_response(QUERY_VECTOR)) + request_metadata = {"user_api_key_team_id": "team-a", "user_api_key": "hashed-key"} - mock_bare = AsyncMock(return_value=Mock(data=[{"embedding": [0.6, 0.7]}])) - with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on - _, request_body = await config.atransform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, + _, request_body = await config.atransform_search_vector_store_request( + **_search_kwargs( + litellm_params={"embedding_model": "team-embeddings"}, + embedding_executor=RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata), ) + ) + + router.aembedding.assert_awaited_once_with( + model="team-embeddings", input=["test query"], metadata=request_metadata + ) + assert request_body["queryVector"]["float32"] == QUERY_VECTOR + + @pytest.mark.asyncio + async def test_atransform_search_default_model_falls_back_to_the_sdk(self): + """Regression (LIT-6750): a store that never named an embedding model keeps working on a proxy + whose model list has no text-embedding-3-small, embedding through the SDK instead of erroring.""" + config = S3VectorsVectorStoreConfig() + router = MagicMock() + router.get_model_list.return_value = [ + {"model_name": "team-embeddings", "litellm_params": {"model": "openai/text-embedding-3-small"}} + ] + router.resolved_litellm_models.return_value = [] + router.aembedding = AsyncMock(side_effect=AssertionError("unserved model must not reach the Router")) + request_metadata = {"user_api_key_team_id": "team-a"} + + mock_bare = AsyncMock(return_value=_embedding_response(QUERY_VECTOR)) + with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose call the test asserts on + _, request_body = await config.atransform_search_vector_store_request( + **_search_kwargs( + embedding_executor=RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata) + ) + ) + + mock_bare.assert_awaited_once_with( + model="text-embedding-3-small", input=["test query"], metadata=request_metadata + ) + assert request_body["queryVector"]["float32"] == QUERY_VECTOR + + @pytest.mark.asyncio + async def test_atransform_search_without_executor_uses_bare_embedding(self): + config = S3VectorsVectorStoreConfig() + + mock_bare = AsyncMock(return_value=_embedding_response([0.6, 0.7])) + with patch("litellm.aembedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on + _, request_body = await config.atransform_search_vector_store_request(**_search_kwargs()) mock_bare.assert_awaited_once_with(model="text-embedding-3-small", input=["test query"]) assert request_body["queryVector"]["float32"] == [0.6, 0.7] - def test_transform_search_uses_router_for_virtual_model_sync(self): - """Sync twin: router-served embedding model resolves via router.embedding""" + def test_transform_search_without_executor_uses_bare_embedding_sync(self): config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - router = _mock_router(["my-embedding-model"], sync=True) - with patch("litellm.embedding", new=MagicMock()) as mock_bare_embedding: # test-quality-ok: guards that the bare-embedding path is not taken; dispatch seam is the behavior under test - _, request_body = config.transform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={"embedding_model": "my-embedding-model"}, - extra_body=None, - router=router, - ) - - router.embedding.assert_called_once_with(model="my-embedding-model", input=["test query"]) - mock_bare_embedding.assert_not_called() - assert request_body["queryVector"]["float32"] == [0.1, 0.2, 0.3] - - def test_transform_search_without_router_uses_bare_embedding_sync(self): - """Sync twin: no router -> bare litellm.embedding as before""" - config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} - - mock_bare = MagicMock(return_value=Mock(data=[{"embedding": [0.8, 0.9]}])) + mock_bare = MagicMock(return_value=_embedding_response([0.8, 0.9])) with patch("litellm.embedding", new=mock_bare): # test-quality-ok: stubs the bare-embedding fallback whose request body the test asserts on _, request_body = config.transform_search_vector_store_request( - vector_store_id="test-bucket:test-index", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, + **_search_kwargs(litellm_params={"embedding_model": "my-embedding-model"}) ) - mock_bare.assert_called_once_with(model="text-embedding-3-small", input=["test query"]) + mock_bare.assert_called_once_with(model="my-embedding-model", input=["test query"]) assert request_body["queryVector"]["float32"] == [0.8, 0.9] def test_transform_search_request_invalid_vector_store_id(self): - """Test that invalid vector_store_id format raises error""" config = S3VectorsVectorStoreConfig() - mock_logging_obj = Mock() - mock_logging_obj.model_call_details = {} + executor = _RecordingExecutor() with pytest.raises( ValueError, match="vector_store_id must be in format 'bucket_name:index_name'", ): config.transform_search_vector_store_request( - vector_store_id="invalid-format", - query="test query", - vector_store_search_optional_params={}, - api_base="https://s3vectors.us-west-2.api.aws", - litellm_logging_obj=mock_logging_obj, - litellm_params={}, - extra_body=None, + **_search_kwargs(vector_store_id="invalid-format", embedding_executor=executor) ) + assert executor.calls == [] + + def test_transform_search_request_bucket_from_litellm_params(self): + config = S3VectorsVectorStoreConfig() + + _, request_body = config.transform_search_vector_store_request( + **_search_kwargs( + vector_store_id="only-index", + litellm_params={"vector_bucket_name": "params-bucket"}, + embedding_executor=_RecordingExecutor(), + ) + ) + + assert request_body["vectorBucketName"] == "params-bucket" + assert request_body["indexName"] == "only-index" + def test_transform_search_response(self): - """Test search response transformation""" config = S3VectorsVectorStoreConfig() mock_logging_obj = Mock() mock_logging_obj.model_call_details = {"query": "test query"} @@ -239,7 +259,7 @@ class TestS3VectorsVectorStoreConfig: mock_response.json.return_value = { "vectors": [ { - "distance": 0.05, # S3 Vectors returns distance, not score + "distance": 0.05, "metadata": { "source_text": "This is test content", "chunk_index": "0", @@ -258,23 +278,18 @@ class TestS3VectorsVectorStoreConfig: mock_response.status_code = 200 mock_response.headers = {} - result = config.transform_search_vector_store_response( - mock_response, mock_logging_obj - ) + result = config.transform_search_vector_store_response(mock_response, mock_logging_obj) - # VectorStoreSearchResponse is a TypedDict, so check structure instead of isinstance assert result["object"] == "vector_store.search_results.page" assert result["search_query"] == "test query" assert len(result["data"]) == 2 - # Score should be 1 - distance (cosine similarity) - assert result["data"][0]["score"] == 0.95 # 1 - 0.05 + assert result["data"][0]["score"] == 0.95 assert result["data"][0]["content"][0]["text"] == "This is test content" assert result["data"][0]["filename"] == "test.pdf" - assert result["data"][1]["score"] == 0.85 # 1 - 0.15 + assert result["data"][1]["score"] == 0.85 assert result["data"][1]["content"][0]["text"] == "More test content" def test_map_openai_params(self): - """Test OpenAI parameter mapping""" config = S3VectorsVectorStoreConfig() non_default_params = {"max_num_results": 5} optional_params = {} diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index befd4c5ffbd..4cff5c76b9e 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -7,12 +7,13 @@ transformations for the Responses API. Source: litellm/llms/xai/responses/transformation.py """ -from unittest.mock import MagicMock - +from unittest.mock import MagicMock, Mock +import httpx import pytest import litellm +from litellm.llms.xai.cost_calculator import cost_per_token from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.llms.openai import ( @@ -400,3 +401,94 @@ class TestXAIResponsesWebSearchBilling: bridged = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage) assert getattr(bridged, "server_side_tool_usage_details") == self._TOOL_DETAILS + + +class TestXAIResponsesReportedCost: + """xAI reports what it charged; the transformation moves it to where litellm bills from. + + ``ResponseAPILoggingUtils`` copies ``usage.cost`` onto the chat Usage that cost + tracking prices, so restating ``cost_in_usd_ticks`` there is what makes /v1/responses + bill the reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756. + """ + + @staticmethod + def _response_body(usage: dict) -> dict: + return { + "id": "resp_xai", + "object": "response", + "created_at": 0, + "model": "grok-4-latest", + "status": "completed", + "output": [], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": usage, + } + + def _transformed_usage(self, usage: dict) -> ResponseAPIUsage | None: + raw_response = httpx.Response(status_code=200, json=self._response_body(usage)) + + response = XAIResponsesAPIConfig().transform_response_api_response( + model="grok-4-latest", + raw_response=raw_response, + logging_obj=Mock(), + ) + return response.usage + + def test_reported_cost_reaches_the_cost_calculator(self): + usage = self._transformed_usage( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ) + + assert usage.cost == 0.0037756 + + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756) + + def test_streamed_reported_cost_reaches_the_cost_calculator(self): + event = XAIResponsesAPIConfig().transform_streaming_response( + model="grok-4-latest", + parsed_chunk={ + "type": "response.completed", + "sequence_number": 7, + "response": self._response_body( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ), + }, + logging_obj=Mock(), + ) + + assert isinstance(event, ResponseCompletedEvent) + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage) + assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756) + + def test_usage_without_a_reported_cost_is_left_alone(self): + usage = self._transformed_usage( + {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300} + ) + + assert usage.cost is None + + def test_negative_reported_cost_is_not_carried(self): + """A caller who can set api_base must not be able to report negative spend.""" + usage = self._transformed_usage( + { + "input_tokens": 100, + "output_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": -37756000, + } + ) + + assert usage.cost is None diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index e5e853ec82f..c67dca11a56 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -1,9 +1,14 @@ +from unittest.mock import Mock - +import httpx import pytest import litellm -from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.chat.transformation import ( + XAIChatCompletionStreamingHandler, + XAIChatConfig, +) +from litellm.llms.xai.cost_calculator import cost_per_token from litellm.types.utils import ( CompletionTokensDetailsWrapper, ModelResponse, @@ -195,3 +200,113 @@ class TestXAIChatWebSearchBilling: ) assert with_search - without_search == pytest.approx(3 * 5.0 / 1000.0) + + +class TestXAIReportedCost: + """xAI reports what it charged; the transformation moves it to where litellm bills from. + + ``cost`` is the field litellm already carries a provider stated cost in, so restating + ``cost_in_usd_ticks`` there is what lets ``llms/xai/cost_calculator.py`` bill the + reported figure. At 10^10 ticks to the dollar, 37756000 ticks is $0.0037756. + """ + + @staticmethod + def _transformed_usage(usage: dict) -> Usage: + raw_response = httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-xai", + "object": "chat.completion", + "created": 0, + "model": "grok-4-latest", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": usage, + }, + ) + + response = XAIChatConfig().transform_response( + model="grok-4-latest", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + return response.usage + + def test_reported_cost_reaches_the_cost_calculator(self): + usage = self._transformed_usage( + { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + } + ) + + assert usage.cost == 0.0037756 + assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0037756) + + def test_usage_without_a_reported_cost_is_left_alone(self): + usage = self._transformed_usage( + {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300} + ) + + assert getattr(usage, "cost", None) is None + + def test_negative_reported_cost_is_not_carried(self): + """A caller who can set api_base must not be able to report negative spend.""" + usage = self._transformed_usage( + { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": -37756000, + } + ) + + assert getattr(usage, "cost", None) is None + + def test_streamed_reported_cost_survives_chunk_aggregation(self): + """Streamed spend only matches if the conversion happens on the chunk. + + Chunk aggregation rebuilds usage from the fields it models plus ``cost``, so a + chunk still carrying only ``cost_in_usd_ticks`` loses the reported amount. + """ + handler = XAIChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + + parsed = handler.chunk_parser( + { + "id": "chatcmpl-xai", + "object": "chat.completion.chunk", + "created": 0, + "model": "grok-4-latest", + "choices": [], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "cost_in_usd_ticks": 37756000, + }, + } + ) + + assert parsed.usage.cost == 0.0037756 + + assembled = litellm.stream_chunk_builder(chunks=[parsed]) + assert assembled.usage.cost == 0.0037756 + assert cost_per_token(model="grok-4-latest", usage=assembled.usage) == ( + 0.0, + 0.0037756, + ) diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 92e76fd18ab..6503e956a51 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -7,7 +7,10 @@ import os import litellm from litellm.types.utils import ( + Choices, CompletionTokensDetailsWrapper, + Message, + ModelResponse, PromptTokensDetailsWrapper, Usage, ) @@ -361,6 +364,145 @@ class TestXAICostCalculator: response_object=object(), usage=usage ) + def test_reported_cost_is_preferred_over_token_math(self): + """The amount xAI reported, carried on usage.cost by the transformation, is billed. + + It lands entirely on completion cost because xAI does not split its total by + direction, the same shape the perplexity calculator returns. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=0.0037756, + ) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost == 0.0 + assert math.isclose(completion_cost, 0.0037756, rel_tol=1e-10) + + def test_reported_cost_suppresses_web_search_surcharge(self): + """The reported total already covers server-side tool calls. + + Without the suppression these 3 searches would be billed a second time on + top of the total xAI already charged. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=100, + web_search_requests=3, + ), + cost=0.0037756, + ) + + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_web_search_surcharge_suppressed_through_the_dispatcher(self): + """The suppression has to hold on the path cost tracking actually uses. + + Legacy behaviour stays intact when xAI reports no cost. + """ + from litellm.llms import get_cost_for_web_search_request + + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + assert get_cost_for_web_search_request("xai", usage, {}) > 0.0 + + reported = Usage( + prompt_tokens=100, completion_tokens=50, total_tokens=150, cost=0.0037756 + ) + setattr(reported, "server_side_tool_usage_details", {"web_search_calls": 3}) + assert get_cost_for_web_search_request("xai", reported, {}) == 0.0 + + def test_no_reported_cost_falls_back_to_token_math(self): + """Absent the provider figure, nothing changes for existing callers.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + + def test_malformed_reported_cost_falls_back_to_token_math(self): + """A junk value must not fail the request, fall back to calculating.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + setattr(usage, "cost", "not-a-number") + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + + def test_boolean_reported_cost_falls_back_to_token_math(self): + """True is an int in python and would otherwise be billed as $1.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + setattr(usage, "cost", True) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + assert completion_cost != 1.0 + + def test_negative_reported_cost_is_rejected(self): + """A negative amount must never reach spend tracking. + + A caller who can set api_base controls the response body, so trusting a + negative figure would let them subtract from their own recorded spend and + slip past a budget. Fall back to token pricing instead, and keep charging + the web search surcharge, since no trustworthy total was reported. + """ + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=-0.0037756, + ) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert prompt_cost > 0.0 + assert completion_cost > 0.0 + assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0 + + def test_non_finite_reported_cost_is_rejected(self): + """NaN compares false against every budget threshold. + + Usage stores a provider supplied cost without validating it, so a caller who + controls the response body could report NaN and leave spend >= max_budget + false for the life of the key rather than mispricing one request. The + infinities are refused alongside it. Fall back to token pricing and keep + charging the web search surcharge, since no trustworthy total was reported. + """ + for reported_cost in (float("nan"), float("inf"), float("-inf")): + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + total_tokens=300, + cost=reported_cost, + ) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 3}) + + prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + + assert math.isfinite(prompt_cost), reported_cost + assert math.isfinite(completion_cost), reported_cost + assert prompt_cost > 0.0, reported_cost + assert completion_cost > 0.0, reported_cost + assert cost_per_web_search_request(usage=usage, model_info={}) > 0.0, reported_cost + + def test_zero_reported_cost_is_honoured(self): + """A reported zero is a real answer, not a missing value.""" + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300, cost=0.0) + + assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0) + def test_grok_4_20_beta_reasoning_cost_calculation(self): """Test cost calculation for grok-4.20-beta-0309-reasoning model.""" usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) @@ -437,6 +579,48 @@ class TestXAICostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + def test_custom_pricing_beats_the_reported_cost(self): + response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")], + usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956), + ) + + billed = litellm.completion_cost( + completion_response=response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.001}, + custom_pricing=True, + ) + + assert math.isclose(billed, 0.551, rel_tol=1e-10) + + def test_deployment_custom_pricing_beats_the_reported_cost(self, monkeypatch): + deployment_id = "xai-deployment-priced-by-the-operator" + monkeypatch.setitem( + litellm.model_cost, + deployment_id, + {"input_cost_per_token": 0.001, "output_cost_per_token": 0.001, "litellm_provider": "xai", "mode": "chat"}, + ) + response = ModelResponse( + id="chatcmpl-xai", + model="grok-4-latest", + choices=[Choices(index=0, message=Message(role="assistant", content="x"), finish_reason="stop")], + usage=Usage(prompt_tokens=198, completion_tokens=353, total_tokens=551, cost=0.0009956), + ) + + billed = litellm.completion_cost( + completion_response=response, + model="xai/grok-4-latest", + custom_llm_provider="xai", + custom_pricing=True, + router_model_id=deployment_id, + ) + + assert math.isclose(billed, 0.551, rel_tol=1e-10) + class TestXAIWebSearchCostHelpers: """Focused coverage for web_search / tool-usage helpers in cost_calculator.py.""" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index d85c6da659b..b99dcb062b4 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -670,6 +670,37 @@ def test_get_provider_specific_params(): ) # Literal type should be select +@pytest.mark.asyncio +async def test_provider_specific_params_includes_hide_secrets(): + """hide-secrets lives in the enterprise package so it is not in + guardrail_class_registry; the endpoint must still advertise it or the + Add Guardrail UI dropdown never offers it (LIT-3548).""" + from litellm.proxy.guardrails.guardrail_endpoints import ( + get_provider_specific_params, + ) + + provider_params = await get_provider_specific_params() + + assert "hide-secrets" in provider_params + # populateGuardrailProviders() in the dashboard only lists providers whose + # entry carries a ui_friendly_name. + assert provider_params["hide-secrets"]["ui_friendly_name"] == "Hide Secrets" + assert provider_params["hide-secrets"]["detect_secrets_config"]["required"] is False + + +@pytest.mark.asyncio +async def test_add_guardrail_settings_restricts_hide_secrets_to_pre_call(): + """hide-secrets only implements async_pre_call_hook, so offering the other + modes in the UI would create configs that boot clean and never run.""" + from litellm.proxy.guardrails.guardrail_endpoints import ( + get_guardrail_ui_settings, + ) + + settings = await get_guardrail_ui_settings() + + assert settings.supported_modes_by_provider["hide-secrets"] == ["pre_call"] + + def test_optional_params_not_returned_when_not_overridden(): """Test that optional_params is not returned when the config model doesn't override it""" from typing import Optional diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py index 36d6414ba9f..635e6332958 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -516,6 +516,39 @@ async def test_team_ids_extracted_from_groups_attribute(saml_env_idp_initiated): assert result.team_ids == ["team-a", "team-b"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "roles", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +async def test_multi_valued_role_attribute_resolves_to_highest_privilege(saml_env_idp_initiated, roles): + """An assertion carrying several roles must not depend on the order the IdP emitted them in.""" + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response( + key_pem, + cert_pem, + attributes={ + "email": ["dave@example.com"], + "role": roles, + }, + ) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_assertion_without_role_attribute_has_no_user_role(saml_env_idp_initiated): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, attributes={"email": ["erin@example.com"]}) + + result = await _acs(_b64(resp), _shared_cache()) + assert result.user_role is None + + @pytest.mark.asyncio async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): cache = DualCache() diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index dd8c752a868..5dfff53f7c3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -6598,13 +6598,94 @@ def test_get_litellm_user_role_with_invalid_role(): assert result is None -def test_get_litellm_user_role_with_list_multiple_roles(): - """Test that get_litellm_user_role takes the first element from a multi-element list.""" +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin", "internal_user"], + ["internal_user", "proxy_admin"], + ], +) +def test_get_litellm_user_role_picks_highest_privilege_regardless_of_order(role_claim): + """A multi-valued role claim resolves to the most privileged role, not the first one listed.""" from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.types import get_litellm_user_role - result = get_litellm_user_role(["proxy_admin", "internal_user"]) - assert result == LitellmUserRoles.PROXY_ADMIN + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["proxy_admin_viewer", "internal_user"], + ["internal_user", "proxy_admin_viewer"], + ], +) +def test_get_litellm_user_role_keeps_org_spend_visibility_for_mixed_roles(role_claim): + """ + Regression for LIT-6077: a user holding both proxy_admin_viewer and internal_user kept + losing org-level spend visibility whenever the IdP happened to list internal_user first. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + +def test_get_litellm_user_role_ignores_unrecognised_entries(): + """Roles LiteLLM does not know about are skipped rather than swallowing the whole claim.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["some_idp_group", "internal_user"]) == LitellmUserRoles.INTERNAL_USER + assert get_litellm_user_role(["some_idp_group", "another_group"]) is None + + +def test_get_litellm_user_role_list_lookup_is_case_insensitive(): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(["INTERNAL_USER", "Proxy_Admin"]) == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "team"], + ["team", "org_admin"], + ], +) +def test_get_litellm_user_role_is_deterministic_for_unranked_roles(role_claim): + """Roles outside the privilege hierarchy still resolve the same way in either claim order.""" + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.ORG_ADMIN + + +@pytest.mark.parametrize( + "role_claim", + [ + ["org_admin", "internal_user"], + ["internal_user", "org_admin"], + ], +) +def test_get_litellm_user_role_prefers_a_ranked_role_over_an_unranked_one(role_claim): + """ + org_admin, team and customer sit outside the privilege ladder, so a claim mixing one of + them with a ranked role settles on the ranked role in either order. Same rule the Entra + app_roles and role_mappings paths already follow. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(role_claim) == LitellmUserRoles.INTERNAL_USER + + +def test_get_litellm_user_role_returns_none_for_non_string_claims(): + from litellm.proxy.management_endpoints.types import get_litellm_user_role + + assert get_litellm_user_role(None) is None + assert get_litellm_user_role({"role": "proxy_admin"}) is None # ============================================================================ @@ -6654,6 +6735,46 @@ def test_process_sso_jwt_access_token_extracts_role_from_access_token(): assert result.user_role == LitellmUserRoles.PROXY_ADMIN +@pytest.mark.parametrize( + "role_claim", + [ + ["internal_user", "proxy_admin_viewer"], + ["proxy_admin_viewer", "internal_user"], + ], +) +def test_process_sso_jwt_access_token_resolves_highest_privilege_role(role_claim): + """ + The generic SSO access-token path must land on the same role for a user whose role + claim holds several roles, whichever order the IdP emitted them in. + """ + import jwt as pyjwt + + from litellm.proxy._types import LitellmUserRoles + + access_token_str = pyjwt.encode( + {"sub": "user-123", "email": "mixed@test.com", "litellm_role": role_claim}, + "secret", + algorithm="HS256", + ) + result = CustomOpenID( + id="user-123", + email="mixed@test.com", + display_name="Mixed Role User", + team_ids=[], + user_role=None, + ) + + with patch.dict(os.environ, {"GENERIC_USER_ROLE_ATTRIBUTE": "litellm_role"}): + process_sso_jwt_access_token( + access_token_str=access_token_str, + sso_jwt_handler=None, + result=result, + role_mappings=None, + ) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + + def test_process_sso_jwt_access_token_does_not_override_existing_role(): """ Test that process_sso_jwt_access_token does NOT override a role that was diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 551e27a18f5..41439f28638 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1104,6 +1104,16 @@ def test_get_autorouter_presets_local_mode_serves_bundled_catalog( assert response.status_code == 200 payload = response.json() assert "anthropic_family" in payload + assert payload["1m_context"]["complexity_router_config"]["classifier_type"] == "heuristic_v2" + assert payload["1m_context"]["complexity_router_config"]["tiers"] == { + "SIMPLE": ["gpt-5.6-luna"], + "MEDIUM": ["gpt-5.6-terra"], + "COMPLEX": ["claude-opus-5"], + "REASONING": ["claude-opus-5"], + } + assert payload["1m_context"]["complexity_router_config"]["tier_model_configs"] == { + "REASONING": [{"model_name": "claude-opus-5", "litellm_params": {"reasoning_effort": "high"}}] + } for preset in payload.values(): assert isinstance(preset["label"], str) assert isinstance(preset["description"], str) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 7c2b9d0be05..3c8bf142835 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3131,9 +3131,9 @@ def test_stream_chunk_builder_prices_proxy_alias_via_model_map(): assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) -def _stream_builder_logging_obj() -> LiteLLMLogging: +def _stream_builder_logging_obj(model: str = "gpt-4o", custom_llm_provider: str = "openai") -> LiteLLMLogging: logging_obj: Final = LiteLLMLogging( - model="gpt-4o", + model=model, messages=[{"role": "user", "content": "hi"}], stream=True, call_type="completion", @@ -3142,10 +3142,11 @@ def _stream_builder_logging_obj() -> LiteLLMLogging: function_id="test-function-id", ) logging_obj.update_environment_variables( - model="gpt-4o", + model=model, user=None, optional_params={}, - litellm_params={"custom_llm_provider": "openai"}, + litellm_params={"custom_llm_provider": custom_llm_provider}, + custom_llm_provider=custom_llm_provider, ) return logging_obj @@ -3237,3 +3238,24 @@ def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): assert response.usage.completion_tokens == 60 assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) assert response._hidden_params["response_cost"] == pytest.approx(0.000704) + + +def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": 0.5}) + usage_chunk: Final = _stream_builder_text_chunk("grok-4", "") + usage_chunk.usage = Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7, cost=0.42) + chunks: Final = [ + _stream_builder_text_chunk("grok-4", "Hello "), + _stream_builder_text_chunk("grok-4", "world.", finish_reason="stop"), + usage_chunk, + ] + logging_obj: Final = _stream_builder_logging_obj(model="grok-4", custom_llm_provider="xai") + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(0.42) + assert response._hidden_params.get("response_cost") is None + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.63) diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/test_litellm/vector_stores/test_main.py index d01e696906a..234e0b01094 100644 --- a/tests/test_litellm/vector_stores/test_main.py +++ b/tests/test_litellm/vector_stores/test_main.py @@ -2,14 +2,17 @@ Tests for litellm/vector_stores/main.py. Pins the router threading contract for vector store search: the router is an -explicit named parameter that reaches the HTTP handler, and it must never leak -into litellm_params/kwargs where logging would model_dump() it (the #19550 -serialization trap). +explicit named parameter that reaches the HTTP handler wrapped in the embedding +executor, and it must never leak into litellm_params/kwargs where logging would +model_dump() it (the #19550 serialization trap). """ from unittest.mock import MagicMock, patch import litellm.vector_stores.main as vector_stores_main +from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, +) from litellm.vector_stores.main import search MOCK_SEARCH_RESPONSE = { @@ -19,17 +22,18 @@ MOCK_SEARCH_RESPONSE = { } -def test_search_threads_router_to_handler(): - """search() must pass its router param through to the HTTP handler""" +def test_search_wraps_router_into_the_handler_embedding_executor(): + """search() hands the HTTP handler a Router-backed embedding executor carrying the + request metadata, and no bare router kwarg (LIT-6750)""" mock_router = MagicMock() logger = MagicMock() with ( - patch( # test-quality-ok: stubs provider config resolution; the seam under test is the router kwarg threading + patch( # test-quality-ok: stubs provider config resolution; the seam under test is the executor threading "litellm.vector_stores.main.ProviderConfigManager.get_provider_vector_stores_config", return_value=MagicMock(), ), - patch.object( # test-quality-ok: the handler call is the observable boundary for the router kwarg contract + patch.object( # test-quality-ok: the handler call is the observable boundary for the executor contract vector_stores_main.base_llm_http_handler, "vector_store_search_handler", return_value=MOCK_SEARCH_RESPONSE, @@ -41,11 +45,16 @@ def test_search_threads_router_to_handler(): custom_llm_provider="s3_vectors", router=mock_router, litellm_logging_obj=logger, + litellm_metadata={"user_api_key_team_id": "team-a"}, ) assert response == MOCK_SEARCH_RESPONSE mock_handler.assert_called_once() - assert mock_handler.call_args.kwargs["router"] is mock_router + assert "router" not in mock_handler.call_args.kwargs + executor = mock_handler.call_args.kwargs["embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.router is mock_router + assert dict(executor.metadata) == {"user_api_key_team_id": "team-a"} def test_search_router_not_in_litellm_params(): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 5c25b8722ef..d5bf3883be4 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16478 + "limit": 16477 }, "LIT011": { - "limit": 5520 + "limit": 5519 }, "LIT012": { "limit": 4489 diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7de7373b20b..7d7066addcf 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1399,9 +1399,6 @@ }, "prefer-const": { "count": 2 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/components/TeamsPage/teamTableColumns.tsx": { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index c1f2ddcf51c..f686ff5644a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -201,6 +201,7 @@ export const guardrailLogoMap = { XecGuard: xecguardLogo.src, "LiteLLM Content Filter": litellmLogo.src, "LiteLLM LLM as a Judge": litellmLogo.src, + "Hide Secrets": litellmLogo.src, Akto: aktoLogo.src, "DeepKeep AI Firewall": deepkeepLogo.src, "Qostodian Nexus": qohashLogo.src, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.test.tsx new file mode 100644 index 00000000000..ab63b7637b1 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.test.tsx @@ -0,0 +1,99 @@ +import React from "react"; +import { fireEvent, screen, waitFor } from "@testing-library/react"; +import { useForm } from "react-hook-form"; +import { renderWithProviders } from "@/../tests/test-utils"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import GuardrailProviderFields from "./guardrail_provider_fields"; +import { populateGuardrailProviderMap } from "./guardrail_info_helpers"; +import type { GuardrailFormValues } from "./GuardrailFormField"; + +vi.mock("@/lib/toast", () => ({ toast: { error: vi.fn() } })); + +const HIDE_SECRETS_PARAMS = { + "hide-secrets": { + ui_friendly_name: "Hide Secrets", + detect_secrets_config: { + param: "detect_secrets_config", + description: "Optional detect-secrets configuration", + required: false, + type: "object", + }, + }, +}; + +const Harness: React.FC<{ onValid: (values: GuardrailFormValues) => void }> = ({ onValid }) => { + const form = useForm(); + return ( +
+ + + + ); +}; + +const renderHarness = () => { + populateGuardrailProviderMap(HIDE_SECRETS_PARAMS); + const onValid = vi.fn(); + renderWithProviders(); + const textarea = screen.getByLabelText(/detect_secrets_config/) as HTMLTextAreaElement; + return { onValid, textarea }; +}; + +describe("GuardrailProviderFields object field", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("commits a valid JSON object as a parsed dict", async () => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: '{"plugins_used": [{"name": "AWSKeyDetector"}]}' } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await waitFor(() => expect(onValid).toHaveBeenCalledTimes(1)); + expect(onValid.mock.calls[0][0].detect_secrets_config).toEqual({ + plugins_used: [{ name: "AWSKeyDetector" }], + }); + }); + + it("blocks submission while the field holds malformed JSON", async () => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: "{not json" } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await screen.findByText("detect_secrets_config must be a valid JSON object"); + expect(onValid).not.toHaveBeenCalled(); + expect(textarea.value).toBe("{not json"); + }); + + it.each(['["array"]', '"scalar"', "null", "42"])("blocks non-object JSON %s", async (raw) => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: raw } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await screen.findByText("detect_secrets_config must be a valid JSON object"); + expect(onValid).not.toHaveBeenCalled(); + }); + + it("treats a cleared field as unset and submits", async () => { + const { onValid, textarea } = renderHarness(); + + fireEvent.change(textarea, { target: { value: '{"a": 1}' } }); + fireEvent.blur(textarea); + fireEvent.change(textarea, { target: { value: "" } }); + fireEvent.blur(textarea); + fireEvent.click(screen.getByRole("button", { name: "save" })); + + await waitFor(() => expect(onValid).toHaveBeenCalledTimes(1)); + expect(onValid.mock.calls[0][0].detect_secrets_config).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx index 76c91d4c176..7d48da89ebf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx @@ -13,6 +13,8 @@ import { FieldGroup } from "@/components/ui/field"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Slider } from "@/components/ui/slider"; +import { Textarea } from "@/components/ui/textarea"; +import { toast } from "@/lib/toast"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { asStringArray, @@ -22,6 +24,7 @@ import { readRecord, requiredRule, type GuardrailFieldControlProps, + type GuardrailFieldRules, type GuardrailFormControl, } from "./GuardrailFormField"; @@ -60,6 +63,44 @@ const BOOLEAN_ITEMS = [ const isSecretKey = (fieldKey: string): boolean => fieldKey.includes("password") || fieldKey.includes("secret") || fieldKey.includes("key"); +const isPlainObject = (value: unknown): value is Record => + typeof value === "object" && value !== null && !Array.isArray(value); + +// Object fields hold the raw text while the user types, so submission must be +// blocked until the value parses to a plain JSON object (or is cleared). +const jsonObjectRule = (fieldKey: string): GuardrailFieldRules => ({ + validate: (value: unknown) => + value === undefined || isPlainObject(value) ? true : `${fieldKey} must be a valid JSON object`, +}); + +// Commits a parsed object (or undefined for a cleared field) to the form on +// blur; anything else stays as raw text so jsonObjectRule blocks submission. +const commitObjectField = (raw: string, onChange: (value: unknown) => void): void => { + const next = raw.trim(); + if (next === "") { + onChange(undefined); + return; + } + let parsed: unknown; + try { + parsed = JSON.parse(next); + } catch { + parsed = next; + } + if (isPlainObject(parsed)) { + onChange(parsed); + } else { + toast.error("Enter a valid JSON object for this configuration"); + } +}; + +const fieldRules = (field: ProviderParam, fieldKey: string): GuardrailFieldRules | undefined => { + if (field.type === "object") { + return jsonObjectRule(fieldKey); + } + return field.required ? requiredRule(`${fieldKey} is required`) : undefined; +}; + interface ProviderFieldInputProps { descriptor: ProviderParam; fieldKey: string; @@ -141,6 +182,25 @@ const ProviderFieldInput: React.FC = ({ descriptor, fie ); } + if (descriptor.type === "object") { + const objectValue = typeof value === "object" && value !== null ? JSON.stringify(value, null, 2) : asText(value); + return ( +