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/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/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/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..9faaaf492e8 100644
--- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
+++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
@@ -376,7 +376,6 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
timeout=None,
client=None,
_is_async=False,
- router: "litellm.Router | None" = 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/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/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/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx
index bee35ac5c5d..2bceb00aae1 100644
--- a/ui/litellm-dashboard/src/components/Teams.test.tsx
+++ b/ui/litellm-dashboard/src/components/Teams.test.tsx
@@ -15,6 +15,7 @@ import {
teamCreateCall,
} from "./networking";
import Teams from "./Teams";
+import { chooseSelectOption } from "../../tests/test-utils";
const can = vi.fn();
vi.mock("@/app/(dashboard)/hooks/useCan", () => ({
@@ -1488,3 +1489,204 @@ describe("Teams - the exact bytes the create call sends", () => {
expect(teamCreateCall).not.toHaveBeenCalled();
});
});
+
+describe("Teams - the create form keeps the organization and models picks while it is open", () => {
+ const ORGS = [
+ { organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] },
+ { organization_id: "org-2", organization_alias: "Org 2", models: [], members: [] },
+ ];
+
+ const orgField = () => screen.getByRole("combobox", { name: /organization/i });
+ const modelsField = () => screen.getByTestId("create-team-models-select");
+
+ const openCreateModal = async () => {
+ act(() => {
+ fireEvent.click(screen.getAllByRole("button", { name: /create team/i })[0]);
+ });
+ await screen.findByLabelText(/team name/i);
+ };
+
+ beforeEach(() => {
+ vi.clearAllMocks();
+ mockTeamInfoView.mockClear();
+ vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]);
+ vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]);
+ vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] });
+ vi.mocked(getDefaultTeamSettings).mockResolvedValue({ values: {} });
+ mockUseOrganizations.mockReturnValue({ data: ORGS });
+ });
+
+ it("keeps both picks when the organizations list comes back changed from a refetch", async () => {
+ const user = userEvent.setup();
+ renderWithQueryClient();
+ await openCreateModal();
+
+ await chooseSelectOption(user, orgField(), /Org 1/);
+ fireEvent.change(modelsField(), { target: { value: "gpt-4" } });
+
+ mockUseOrganizations.mockReturnValue({ data: ORGS.map((org) => ({ ...org, spend: 1 })) });
+ fireEvent.click(screen.getByText("Additional Settings"));
+
+ expect(orgField()).toHaveValue("Org 1");
+ expect(modelsField()).toHaveValue("gpt-4");
+ });
+
+ it("keeps models picked before the available models finish loading", async () => {
+ let resolveModels: (models: string[]) => void = () => {};
+ vi.mocked(fetchAvailableModelsForTeamOrKey).mockReturnValue(
+ new Promise((resolve) => {
+ resolveModels = resolve;
+ }),
+ );
+ renderWithQueryClient();
+ await openCreateModal();
+
+ fireEvent.change(modelsField(), { target: { value: "gpt-4" } });
+ await act(async () => {
+ resolveModels(["gpt-4", "gpt-3.5-turbo"]);
+ });
+
+ expect(modelsField()).toHaveValue("gpt-4");
+ });
+
+ it("clears the models pick when the organization is changed, since models are org scoped", async () => {
+ const user = userEvent.setup();
+ renderWithQueryClient();
+ await openCreateModal();
+
+ await chooseSelectOption(user, orgField(), /Org 1/);
+ fireEvent.change(modelsField(), { target: { value: "gpt-4" } });
+ await chooseSelectOption(user, orgField(), /Org 2/);
+
+ await waitFor(() => expect(orgField()).toHaveValue("Org 2"));
+ expect(modelsField()).toHaveValue("");
+ });
+
+ it("keeps the models pick when the same organization is chosen again", async () => {
+ const user = userEvent.setup();
+ renderWithQueryClient();
+ await openCreateModal();
+
+ await chooseSelectOption(user, orgField(), /Org 1/);
+ fireEvent.change(modelsField(), { target: { value: "gpt-4" } });
+ await chooseSelectOption(user, orgField(), /Org 1/);
+
+ expect(orgField()).toHaveValue("Org 1");
+ expect(modelsField()).toHaveValue("gpt-4");
+ });
+
+ it("still preselects the only organization an org admin can create teams in", async () => {
+ mockUseOrganizations.mockReturnValue({
+ data: [
+ {
+ organization_id: "org-1",
+ organization_alias: "Org 1",
+ models: [],
+ members: [{ user_id: "user-123", user_role: "org_admin" }],
+ },
+ ],
+ });
+ renderWithQueryClient();
+ await openCreateModal();
+
+ expect(orgField()).toHaveValue("Org 1");
+ expect(orgField()).toBeDisabled();
+ });
+
+ it("leaves an org admin able to pick when their admin orgs narrow to one while the form is open", async () => {
+ const orgAdminOrgs = [
+ {
+ organization_id: "org-1",
+ organization_alias: "Org 1",
+ models: [],
+ members: [{ user_id: "user-123", user_role: "org_admin" }],
+ },
+ {
+ organization_id: "org-2",
+ organization_alias: "Org 2",
+ models: [],
+ members: [{ user_id: "user-123", user_role: "org_admin" }],
+ },
+ ];
+ mockUseOrganizations.mockReturnValue({ data: orgAdminOrgs });
+ renderWithQueryClient();
+ await openCreateModal();
+ expect(orgField()).toHaveValue("");
+
+ mockUseOrganizations.mockReturnValue({ data: [orgAdminOrgs[0]] });
+ fireEvent.click(screen.getByText("Additional Settings"));
+
+ expect(orgField()).toBeEnabled();
+ });
+
+ it("refuses to create the team in an organization the admin has lost access to", async () => {
+ const user = userEvent.setup();
+ const orgAdminOrgs = ORGS.map((org) => ({ ...org, members: [{ user_id: "user-123", user_role: "org_admin" }] }));
+ mockUseOrganizations.mockReturnValue({ data: orgAdminOrgs });
+ renderWithQueryClient();
+ await openCreateModal();
+
+ fireEvent.change(screen.getByTestId("team-name-input"), { target: { value: "Revoked Team" } });
+ await chooseSelectOption(user, orgField(), /Org 1/);
+
+ mockUseOrganizations.mockReturnValue({ data: [orgAdminOrgs[1]] });
+ fireEvent.click(screen.getByText("Additional Settings"));
+
+ const submitButtons = screen.getAllByRole("button", { name: /create team/i });
+ fireEvent.click(submitButtons[submitButtons.length - 1]);
+
+ await screen.findByText(/no longer create teams in this organization/i);
+ expect(teamCreateCall).not.toHaveBeenCalled();
+ });
+
+ it("lets the admin switch to the one organization left after losing access to their pick", async () => {
+ const user = userEvent.setup();
+ const orgAdminOrgs = ORGS.map((org) => ({ ...org, members: [{ user_id: "user-123", user_role: "org_admin" }] }));
+ mockUseOrganizations.mockReturnValue({ data: orgAdminOrgs });
+ const createdTeam = {
+ team_id: "new-team-1",
+ team_alias: "Recovered Team",
+ models: [],
+ organization_id: "org-2",
+ keys: [],
+ members_with_roles: [],
+ spend: 0,
+ };
+ vi.mocked(teamCreateCall).mockResolvedValue(createdTeam);
+ renderWithQueryClient();
+ await openCreateModal();
+
+ fireEvent.change(screen.getByTestId("team-name-input"), { target: { value: "Recovered Team" } });
+ await chooseSelectOption(user, orgField(), /Org 1/);
+
+ mockUseOrganizations.mockReturnValue({ data: [orgAdminOrgs[1]] });
+ fireEvent.click(screen.getByText("Additional Settings"));
+
+ expect(orgField()).toBeEnabled();
+ await chooseSelectOption(user, orgField(), /Org 2/);
+ const submitButtons = screen.getAllByRole("button", { name: /create team/i });
+ fireEvent.click(submitButtons[submitButtons.length - 1]);
+
+ await waitFor(() =>
+ expect(teamCreateCall).toHaveBeenCalledWith(
+ "test-token",
+ expect.objectContaining({ team_alias: "Recovered Team", organization_id: "org-2" }),
+ ),
+ );
+ });
+
+ it("starts the form clean again when the modal is closed and reopened", async () => {
+ const user = userEvent.setup();
+ renderWithQueryClient();
+ await openCreateModal();
+
+ await chooseSelectOption(user, orgField(), /Org 1/);
+ fireEvent.change(modelsField(), { target: { value: "gpt-4" } });
+ fireEvent.click(screen.getByRole("button", { name: /^close$/i }));
+ await waitFor(() => expect(screen.queryByLabelText(/team name/i)).not.toBeInTheDocument());
+
+ await openCreateModal();
+ expect(orgField()).toHaveValue("");
+ expect(modelsField()).toHaveValue("");
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx
index e2eda4adb23..ef58237a6aa 100644
--- a/ui/litellm-dashboard/src/components/Teams.tsx
+++ b/ui/litellm-dashboard/src/components/Teams.tsx
@@ -208,7 +208,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
const queryClient = useQueryClient();
const refreshTeams = () => queryClient.invalidateQueries({ queryKey: teamsTableKeys.all });
const [currentOrg] = useState(null);
- const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState(null);
const isOrgAdmin = userRole !== "Admin";
const [additionalSettingsOpen, setAdditionalSettingsOpen] = useState(false);
@@ -216,17 +215,33 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
const [agentSettingsOpen, setAgentSettingsOpen] = useState(false);
const [searchToolSettingsOpen, setSearchToolSettingsOpen] = useState(false);
+ const adminOrgs = useMemo(
+ () => getAdminOrganizations(userRole, userID, organizations),
+ [userRole, userID, organizations],
+ );
+
const teamCreateSchema = useMemo(
() =>
teamCreateFieldsSchema.superRefine((values, ctx) => {
if (isOrgAdmin && !values.organization_id) {
ctx.addIssue({ code: "custom", message: SUPPRESSED_BY_DESCRIPTION, path: ["organization_id"] });
}
+ const organizationIsStillPickable =
+ values.organization_id == null ||
+ organizations == null ||
+ adminOrgs.some((org) => org.organization_id === values.organization_id);
+ if (!organizationIsStillPickable) {
+ ctx.addIssue({
+ code: "custom",
+ message: "You can no longer create teams in this organization",
+ path: ["organization_id"],
+ });
+ }
if (additionalSettingsOpen && !isParsableJson(values.secret_manager_settings)) {
ctx.addIssue({ code: "custom", message: SUPPRESSED_BY_DESCRIPTION, path: ["secret_manager_settings"] });
}
}),
- [isOrgAdmin, additionalSettingsOpen],
+ [isOrgAdmin, additionalSettingsOpen, adminOrgs, organizations],
);
const form = useZodForm(teamCreateSchema, { defaultValues: EMPTY_TEAM_CREATE_VALUES });
@@ -264,28 +279,6 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
? `Default: ${getBudgetDurationLabel(defaultBudgetDuration)} (${defaultBudgetDuration})`
: "n/a";
- useEffect(() => {
- form.setValue("models", []);
- }, [currentOrgForCreateTeam, userModels]);
-
- // Handle organization preselection when modal opens
- useEffect(() => {
- if (isTeamModalVisible) {
- const adminOrgs = getAdminOrganizations(userRole, userID, organizations);
-
- // Org admins must scope a team to an org, so with exactly one we preselect it.
- // Proxy admins can create org-less teams, so the field stays optional regardless of org count.
- if (isOrgAdmin && adminOrgs.length === 1) {
- const org = adminOrgs[0];
- form.setValue("organization_id", org.organization_id);
- setCurrentOrgForCreateTeam(org);
- } else {
- form.setValue("organization_id", currentOrg?.organization_id || null);
- setCurrentOrgForCreateTeam(currentOrg);
- }
- }
- }, [isTeamModalVisible, isOrgAdmin, userRole, userID, organizations, currentOrg]);
-
// Add this useEffect to fetch guardrails
useEffect(() => {
const fetchGuardrails = async () => {
@@ -320,6 +313,26 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
if (canViewPolicies) fetchPolicies();
}, [accessToken, canViewPolicies]);
+ const openCreateTeamModal = () => {
+ // Org admins must scope a team to an org, so with exactly one we preselect it.
+ // Proxy admins can create org-less teams, so the field stays optional regardless of org count.
+ if (isOrgAdmin && adminOrgs.length === 1) {
+ form.setValue("organization_id", adminOrgs[0].organization_id);
+ }
+ setIsTeamModalVisible(true);
+ };
+
+ const selectCreateTeamOrganization = (
+ next: string,
+ currentOrganizationId: string | null,
+ onChange: (organizationId: string | null) => void,
+ ) => {
+ const nextOrganizationId = next === "" ? null : next;
+ if (nextOrganizationId === currentOrganizationId) return;
+ onChange(nextOrganizationId);
+ form.setValue("models", []);
+ };
+
const resetCreateForm = () => {
form.reset(EMPTY_TEAM_CREATE_VALUES);
setAdditionalSettingsOpen(false);
@@ -636,7 +649,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
subtitle="Manage teams, members, and their access to models and budgets"
primaryAction={
canCreateOrManageTeams(userRole, userID, organizations) ? (
- setIsTeamModalVisible(true)} data-testid="create-team-button">
+
Create Team
@@ -683,9 +696,9 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
)}
{(() => {
- const adminOrgs = getAdminOrganizations(userRole, userID, organizations);
const isSingleOrg = adminOrgs.length === 1;
const hasNoOrgs = adminOrgs.length === 0;
+ const soleOrganizationId = isSingleOrg ? adminOrgs[0].organization_id ?? null : null;
return (
<>
@@ -715,18 +728,13 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser
label: org.organization_alias ?? "",
sublabel: org.organization_id ?? "",
}))}
- disabled={isOrgAdmin && isSingleOrg}
+ disabled={isOrgAdmin && soleOrganizationId !== null && value === soleOrganizationId}
allowClear={!isOrgAdmin}
placeholder={
hasNoOrgs ? "No organizations available" : "Search or select an Organization"
}
emptyText="No organizations available"
- onValueChange={(next) => {
- onChange(next === "" ? null : next);
- setCurrentOrgForCreateTeam(
- adminOrgs.find((org) => org.organization_id === next) ?? null,
- );
- }}
+ onValueChange={(next) => selectCreateTeamOrganization(next, value ?? null, onChange)}
/>
)}
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
index 27a169524b9..836b4c6ff3d 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
@@ -636,7 +636,14 @@ describe("AddAutoRouterTab", () => {
const labels = visibleOptions().map((option) => option.querySelector(".font-medium")?.textContent);
- expect(labels).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family", "Custom Configuration"]);
+ expect(labels).toEqual([
+ "1M Context",
+ "Anthropic Family",
+ "Gemini Family",
+ "Lite",
+ "OpenAI Family",
+ "Custom Configuration",
+ ]);
});
describe("routing test", () => {
@@ -1060,7 +1067,14 @@ describe("AddAutoRouterTab", () => {
expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false);
});
const labels = visibleOptions().map((option) => option.querySelector(".font-medium")?.textContent);
- expect(labels).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family", "Custom Configuration"]);
+ expect(labels).toEqual([
+ "Anthropic Family",
+ "1M Context",
+ "Gemini Family",
+ "Lite",
+ "OpenAI Family",
+ "Custom Configuration",
+ ]);
});
it.each([
diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
index 44ac6973a1c..3dd9911c794 100644
--- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
+++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
@@ -1,9 +1,8 @@
import { describe, it, expect } from "vitest";
-import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json";
+import { BUNDLED_PRESETS_RESPONSE } from "../../tests/mocks/autoRouterPresets";
import {
hydratePresets,
AutoRouterPreset,
- AutoRouterPresetsResponse,
getRequiredModelsInPreset,
getMissingModelsInPreset,
getRequiredModels,
@@ -21,14 +20,20 @@ import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKe
const groupsOnly = (models: Iterable) => buildModelAvailability(models, []);
// Hydrated from the real bundled catalog so a catalog edit flows into these expectations.
-const PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse);
+const PRESETS = hydratePresets(BUNDLED_PRESETS_RESPONSE);
const getAllPresets = (): AutoRouterPreset[] => PRESETS;
const getPresetByKey = (key: string): AutoRouterPreset | undefined => PRESETS.find((p) => p.key === key);
describe("autorouter_presets", () => {
it("hydrates exactly the bundled presets", () => {
const presets = getAllPresets();
- expect(presets.map((p) => p.label).sort()).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family"]);
+ expect(presets.map((p) => p.label).sort()).toEqual([
+ "1M Context",
+ "Anthropic Family",
+ "Gemini Family",
+ "Lite",
+ "OpenAI Family",
+ ]);
// Every preset carries all four fields the UI relies on; a JSON typo dropping one fails here.
for (const p of presets) {
expect(p).toMatchObject({ key: expect.any(String), label: expect.any(String), description: expect.any(String) });
@@ -202,6 +207,25 @@ describe("autorouter_presets", () => {
});
});
+ it("pins the 1M context preset to Luna, Terra, and Opus at high thinking", () => {
+ const preset = getPresetByKey("1m_context")!;
+ const expectedTiers = {
+ SIMPLE: ["gpt-5.6-luna"],
+ MEDIUM: ["gpt-5.6-terra"],
+ COMPLEX: ["claude-opus-5"],
+ REASONING: ["claude-opus-5"],
+ };
+ expect(preset.complexity_router_config.classifier_type).toBe("heuristic_v2");
+ expect(preset.complexity_router_config.tiers).toEqual(expectedTiers);
+ expect(preset.complexity_router_config.tier_model_configs).toEqual({
+ REASONING: [{ model_name: "claude-opus-5", litellm_params: { reasoning_effort: "high" } }],
+ });
+ const prefill = buildPresetPrefill(preset.complexity_router_config, groupsOnly(getRequiredModelsInPreset(preset)));
+ expect(prefill.complexityRouterConfig.tier_model_params).toEqual({
+ REASONING: { "claude-opus-5": { reasoning_effort: "high" } },
+ });
+ });
+
it("pins the gemini preset to concrete model ids, never Google's hot-swapping -latest aliases", () => {
const gemini = getPresetByKey("gemini_family")!;
const config = gemini.complexity_router_config;
diff --git a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts
index f73faaa70fe..ac6da9bceba 100644
--- a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts
+++ b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts
@@ -1,10 +1,15 @@
+import { readFileSync } from "fs";
+import { resolve } from "path";
import { vi } from "vitest";
-import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json";
import { hydratePresets, type AutoRouterPresetsResponse } from "@/lib/autorouter_presets";
// Derived from the real bundled catalog so a preset edit there flows into test expectations
// instead of redding on a stale copy. Exported as vi.fn so a test can override the query state.
-export const BUNDLED_PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse);
+const CATALOG_PATH = resolve(__dirname, "../../../../litellm/proxy/public_endpoints/autorouter_presets.json");
+
+export const BUNDLED_PRESETS_RESPONSE = JSON.parse(readFileSync(CATALOG_PATH, "utf8")) as AutoRouterPresetsResponse;
+
+export const BUNDLED_PRESETS = hydratePresets(BUNDLED_PRESETS_RESPONSE);
export const LOADED_PRESETS_QUERY = {
data: BUNDLED_PRESETS,