scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_agent_mcp_grants

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-03 05:51:11 +00:00
commit 342b21266d
62 changed files with 3119 additions and 723 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
basedpyright-code-budget.json
View file

@ -54,7 +54,7 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5595
"limit": 5593
},
"reportMissingTypeArgument": {
"limit": 15288
@ -99,7 +99,7 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44352
"limit": 44349
},
"reportUnknownLambdaType": {
"limit": 109
@ -108,10 +108,10 @@
"limit": 38324
},
"reportUnknownParameterType": {
"limit": 19619
"limit": 19617
},
"reportUnknownVariableType": {
"limit": 29852
"limit": 29849
},
"reportUnnecessaryCast": {
"limit": 111

188
litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs
View file

@ -186,68 +186,180 @@ mod tests {
use serde_json::json;
#[test]
fn supported_params_match_python_mistral_ocr_config() {
fn extract_header_is_a_supported_ocr_param() {
assert!(supported_ocr_params().contains(&"extract_header"));
}
#[test]
fn extract_footer_is_a_supported_ocr_param() {
assert!(supported_ocr_params().contains(&"extract_footer"));
}
#[test]
fn existing_ocr_params_remain_supported() {
for param in [
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
] {
assert!(supported_ocr_params().contains(&param));
}
}
#[test]
fn map_ocr_params_forwards_extract_header() {
let params = json!({"extract_header": true});
assert_eq!(
supported_ocr_params(),
&[
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
"document_annotation_prompt",
"extract_header",
"extract_footer",
"table_format",
"confidence_scores_granularity",
"include_blocks",
"id",
]
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_forwards_extract_footer() {
let params = json!({"extract_footer": true});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_forwards_extract_header_and_footer() {
let params = json!({"extract_header": true, "extract_footer": false});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_drops_unknown_params() {
let params = json!({
"extract_header": true,
"unsupported_param": "value",
"pages": [0, 1]
});
let params = json!({"extract_header": true, "unsupported_param": "value"});
let mapped = map_ocr_params(params.as_object().unwrap());
assert_eq!(mapped.get("extract_header"), Some(&json!(true)));
assert_eq!(mapped.get("pages"), Some(&json!([0, 1])));
assert!(!mapped.contains_key("unsupported_param"));
}
#[test]
fn transform_ocr_request_builds_mistral_body() {
fn new_ocr_params_are_supported() {
for param in [
"table_format",
"confidence_scores_granularity",
"document_annotation_prompt",
"include_blocks",
"id",
] {
assert!(supported_ocr_params().contains(&param));
}
}
#[test]
fn map_ocr_params_forwards_new_ocr_params() {
for (param, value) in [
("table_format", json!("html")),
("confidence_scores_granularity", json!("word")),
(
"document_annotation_prompt",
json!("Extract all invoice line items"),
),
("include_blocks", json!(true)),
("id", json!("req-123")),
] {
let params = json!({param: value});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
}
#[test]
fn transform_ocr_request_includes_each_optional_param() {
let document = json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
});
for (param, value) in [
("table_format", json!("html")),
("confidence_scores_granularity", json!("word")),
(
"document_annotation_prompt",
json!("Extract all invoice line items"),
),
("id", json!("req-123")),
("extract_header", json!(true)),
("include_blocks", json!(true)),
("pages", json!([0, 1])),
] {
let result = transform_ocr_request(
"mistral-ocr-latest",
document.clone(),
json!({param: value}).as_object().unwrap().clone(),
)
.expect("request should transform");
assert_eq!(result.data.get(param), Some(&value));
assert_eq!(result.data.get("model"), Some(&json!("mistral-ocr-latest")));
assert_eq!(result.data.get("document"), Some(&document));
assert_eq!(result.files, None);
}
}
#[test]
fn transform_ocr_request_includes_multiple_new_params() {
let document = json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
});
let optional_params = json!({
"include_image_base64": true,
"table_format": "html"
"table_format": "html",
"confidence_scores_granularity": "page",
"extract_header": true
})
.as_object()
.unwrap()
.clone();
let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params)
let result = transform_ocr_request("mistral-ocr-latest", document, optional_params)
.expect("request should transform");
assert_eq!(result.data.get("table_format"), Some(&json!("html")));
assert_eq!(
result.data,
json!({
"model": "mistral-ocr-latest",
"document": document,
"include_image_base64": true,
"table_format": "html"
})
result.data.get("confidence_scores_granularity"),
Some(&json!("page"))
);
assert_eq!(result.files, None);
assert_eq!(result.data.get("extract_header"), Some(&json!(true)));
}
#[test]
fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
let blocks = json!([{"type": "title", "content": "Invoice"}]);
let confidence_scores = json!({"page": 0.98});
let response = json!({
"pages": [{"index": 0, "markdown": "# Invoice", "blocks": blocks, "confidence_scores": confidence_scores}],
"model": "mistral-ocr-4-0",
"usage_info": {"pages_processed": 1}
});
let result =
transform_ocr_response("mistral-ocr-4-0", response).expect("response should transform");
assert_eq!(result.pages[0].get("blocks"), Some(&blocks));
assert_eq!(
result.pages[0].get("confidence_scores"),
Some(&confidence_scores)
);
}
#[test]
fn transform_ocr_response_preserves_ocr4_page_fields() {
let response = json!({
"pages": [{"index": 0, "markdown": "table page", "tables": [{"rows": 2, "cols": 3}], "hyperlinks": ["https://example.com"], "header": "Acme Corp", "footer": "Page 1"}],
"model": "mistral-ocr-4-0",
"usage_info": {"pages_processed": 1}
});
let result = transform_ocr_response("mistral-ocr-4-0", response.clone())
.expect("response should transform");
assert_eq!(result.pages[0], response["pages"][0]);
}
#[test]

12
litellm/cost_calculator.py
View file

@ -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

17
litellm/litellm_core_utils/streaming_handler.py
View file

@ -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,

7
litellm/llms/azure_ai/vector_stores/transformation.py
View file

@ -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,

20
litellm/llms/base_llm/vector_store/transformation.py
View file

@ -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")

2
litellm/llms/bedrock/vector_stores/transformation.py
View file

@ -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)

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

@ -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)

2
litellm/llms/gemini/vector_stores/transformation.py
View file

@ -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.

7
litellm/llms/milvus/vector_stores/transformation.py
View file

@ -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,

2
litellm/llms/openai/vector_stores/transformation.py
View file

@ -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"

2
litellm/llms/pg_vector/vector_stores/transformation.py
View file

@ -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"

2
litellm/llms/ragflow/vector_stores/transformation.py
View file

@ -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")

209
litellm/llms/s3_vectors/vector_stores/transformation.py
View file

@ -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,

2
litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py
View file

@ -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

2
litellm/llms/vertex_ai/vector_stores/search_api/transformation.py
View file

@ -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.

19
litellm/llms/xai/chat/transformation.py
View file

@ -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

11
litellm/llms/xai/common_utils.py
View file

@ -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(

20
litellm/llms/xai/cost_calculator.py
View file

@ -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

68
litellm/llms/xai/responses/transformation.py
View file

@ -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

12
litellm/main.py
View file

@ -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

97
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -820,20 +820,46 @@ def _should_strip_caller_authorization(
if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate):
return False
normalized_raw_headers: Final = {str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str)}
has_explicit_litellm_admission_header: Final = normalized_raw_headers.get("x-litellm-api-key") is not None
has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers)
if mcp_server.is_oauth_delegate:
return not has_explicit_litellm_admission_header
admission_consumed_authorization_as_litellm_key: Final = (
user_api_key_auth is not None
and bool(getattr(user_api_key_auth, "api_key", None))
and not has_explicit_litellm_admission_header
)
return admission_consumed_authorization_as_litellm_key or (
return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or (
user_api_key_auth is None and not has_explicit_litellm_admission_header
)
LITELLM_VIRTUAL_KEY_PREFIX: Final = "sk-"
def _raw_header_value(raw_headers: Mapping[str, str] | None, name: str) -> str | None:
return next((v for k, v in (raw_headers or {}).items() if isinstance(k, str) and k.lower() == name), None)
def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None) -> bool:
"""Admission only consumes a non-empty ``x-litellm-api-key``; an empty one falls back to ``Authorization``."""
return bool(_raw_header_value(raw_headers, "x-litellm-api-key"))
def _authorization_is_litellm_admission_credential(
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> bool:
"""True when ``Authorization`` carries the LiteLLM key admission validated.
That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the
same key in both headers.
"""
if user_api_key_auth is None or not user_api_key_auth.api_key:
return False
admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key")
if not admission_header:
return True
authorization: Final = _raw_header_value(raw_headers, "authorization")
return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(
admission_header, "Bearer"
)
def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str:
"""Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection.
@ -3277,8 +3303,8 @@ class MCPServerManager:
#########################################################
@staticmethod
def _extract_bearer_token(
oauth2_headers: dict[str, str] | None,
raw_headers: dict[str, str] | None,
oauth2_headers: Mapping[str, str] | None,
raw_headers: Mapping[str, str] | None,
) -> str | None:
"""Extract the bare Bearer token from oauth2_headers or raw_headers.
@ -3298,10 +3324,29 @@ class MCPServerManager:
return auth_value
return None
@staticmethod
def _extract_subject_token(
oauth2_headers: Mapping[str, str] | None,
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""The caller's upstream identity token, or ``None`` when the bearer is a LiteLLM key.
Rejects the key admission validated and, because virtual keys always carry the ``sk-`` prefix,
any other LiteLLM key a client puts in ``Authorization`` next to ``x-litellm-api-key``.
"""
if _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth):
return None
bearer: Final = MCPServerManager._extract_bearer_token(oauth2_headers, raw_headers)
if bearer is not None and bearer.startswith(LITELLM_VIRTUAL_KEY_PREFIX):
return None
return bearer
def _obo_subject_token(
self,
server: MCPServer,
raw_headers: dict[str, str] | None,
raw_headers: Mapping[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""The caller's bearer as the token_exchange (OBO) subject token, for that mode only.
@ -3311,7 +3356,7 @@ class MCPServerManager:
"""
if server.auth_type != MCPAuth.oauth2_token_exchange:
return None
return self._extract_bearer_token(None, raw_headers)
return self._extract_subject_token(None, raw_headers, user_api_key_auth)
def _build_stdio_env(
self,
@ -3566,6 +3611,7 @@ class MCPServerManager:
server: MCPServer,
oauth2_headers: dict[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
raw_headers: Mapping[str, str] | None = None,
) -> None:
"""Run the OBO exchange for a caller-supplied subject at the transport edge.
@ -3577,13 +3623,15 @@ class MCPServerManager:
"""
if server.auth_type != MCPAuth.oauth2_token_exchange:
return
subject_token: Final = self._extract_bearer_token(oauth2_headers, None)
if not subject_token:
if not self._extract_bearer_token(oauth2_headers, None):
return
resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
spec: Final = to_server_spec(resolved_server)
if spec is None or not isinstance(spec.config, TokenExchangeConfig):
return
subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
if subject_token is None:
raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
case Ok(_):
return
@ -3851,7 +3899,7 @@ class MCPServerManager:
# token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes
# never read the inbound bearer, so leave subject_token None to avoid forwarding it.
subject_token: Final = (
self._extract_bearer_token(oauth2_headers, raw_headers)
self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
if server.auth_type == MCPAuth.oauth2_token_exchange
else None
)
@ -3931,6 +3979,7 @@ class MCPServerManager:
async def get_prompts_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -3959,7 +4008,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -3982,6 +4031,7 @@ class MCPServerManager:
async def get_resources_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -4001,7 +4051,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -4024,6 +4074,7 @@ class MCPServerManager:
async def get_resource_templates_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
add_prefix: bool = True,
@ -4043,7 +4094,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client = await self._create_mcp_client(
server=server,
@ -4068,6 +4119,7 @@ class MCPServerManager:
async def read_resource_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
url: AnyUrl,
mcp_auth_header: str | dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
@ -4084,7 +4136,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client: Final = await self._create_mcp_client(
server=server,
@ -4099,6 +4151,7 @@ class MCPServerManager:
async def get_prompt_from_server(
self,
server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
prompt_name: str,
arguments: dict[str, str] | None = None,
mcp_auth_header: str | dict[str, str] | None = None,
@ -4116,7 +4169,7 @@ class MCPServerManager:
extra_headers.update(server.static_headers)
stdio_env: Final = self._build_stdio_env(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers)
subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
client: Final = await self._create_mcp_client(
server=server,
@ -5290,7 +5343,7 @@ class MCPServerManager:
MCPAuth.oauth2_token_exchange,
MCPAuth.oauth2_id_jag,
):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
elif mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
@ -5638,7 +5691,7 @@ class MCPServerManager:
subject_token: str | None = None
if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
elif isinstance(spec.config, PassthroughConfig):
inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers)
per_server_token: Final = _passthrough_token_from_mcp_auth_header(mcp_auth_header)

9
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -2265,6 +2265,7 @@ if MCP_AVAILABLE:
try:
prompts = await global_mcp_server_manager.get_prompts_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -2318,6 +2319,7 @@ if MCP_AVAILABLE:
try:
resources = await global_mcp_server_manager.get_resources_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -2369,6 +2371,7 @@ if MCP_AVAILABLE:
try:
resource_templates = await global_mcp_server_manager.get_resource_templates_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
@ -3296,6 +3299,7 @@ if MCP_AVAILABLE:
return await global_mcp_server_manager.get_prompt_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
prompt_name=original_prompt_name,
arguments=arguments,
mcp_auth_header=server_auth_header,
@ -3346,6 +3350,7 @@ if MCP_AVAILABLE:
return await global_mcp_server_manager.read_resource_from_server(
server=server,
user_api_key_auth=user_api_key_auth,
url=url,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
@ -3808,6 +3813,7 @@ if MCP_AVAILABLE:
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None,
allowed_server_ids: set[str] | None = None,
raw_headers: Mapping[str, str] | None = None,
) -> None:
"""Fail fast with HTTP 401 for MCP servers that need user auth but
didn't receive it on this request. Covers both gateway-managed OAuth2
@ -3952,6 +3958,7 @@ if MCP_AVAILABLE:
server=server,
oauth2_headers=oauth2_headers,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
)
# Pass-through OAuth: when the admin has opted a server into
@ -4280,6 +4287,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers. Must run after
@ -4603,6 +4611,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
client_ip=_sse_client_ip,
allowed_server_ids=toolset_allowed_server_ids,
raw_headers=raw_headers,
)
# Pre-flight auth check for pass-through servers: surface upstream

6
litellm/proxy/management_endpoints/sso/saml_sso.py
View file

@ -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(

62
litellm/proxy/management_endpoints/types.py
View file

@ -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):

19
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -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]:

22
litellm/proxy/public_endpoints/autorouter_presets.json
View file

@ -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.",

1
litellm/vector_stores/main.py
View file

@ -482,7 +482,6 @@ def search(
timeout=timeout or request_timeout,
_is_async=_is_async,
client=kwargs.get("client"),
router=router,
)
return response

4
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 2977
"limit": 2982
},
"ANN002": {
"limit": 71
@ -57,7 +57,7 @@
"limit": 3
},
"BLE001": {
"limit": 2913
"limit": 2915
},
"C401": {
"limit": 8

3
tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
View file

@ -370,12 +370,13 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
custom_llm_provider,
litellm_params,
logging_obj,
embedding_executor=None,
extra_headers=None,
extra_body=None,
timeout=None,
client=None,
_is_async=False,
router: "litellm.Router | None" = None,
embedding_executor=None,
):
litellm_params_dict = (
litellm_params.model_dump(exclude_none=False)

13
tests/router_unit_tests/test_router_embedding_integration.py
View file

@ -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

9
tests/rust-python-harness/README.md
View file

@ -8,14 +8,17 @@ The matrix always has these SDK columns:
- `messages / amessages`
- `responses / aresponses`
- `count_tokens`
- `chat_completions / acompletion`
- `transcription / atranscription`
The harness has three deliberately broad test-strategy folders:
The harness has four deliberately broad test-strategy folders:
| Strategy | Folder |
| --- | --- |
| Public SDK parity over generated and recorded inputs | [`e2e_fuzz_tests/`](e2e_fuzz_tests/) |
| Focused tests of Rust-owned behavior | [`unit_tests_rust/`](unit_tests_rust/) |
| Isolated transform and Python-to-Rust helper coverage | [`validate_sub_methods/`](validate_sub_methods/) |
| Already-existing live-API SDK tests | [`existing_e2e_test_sdk/`](existing_e2e_test_sdk/) |
## Run it
@ -112,7 +115,7 @@ The initial end-to-end entries deliberately show `◐`: the repository has Rust
## Attach parity tests
Each of the three folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list:
Each of the four folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list:
```json
{
@ -123,7 +126,7 @@ Each of the three folders contains a concise `README.md` and a `strategy.json`.
}
```
Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice.
Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family; a selector ending in `/` aggregates every test in that folder, recursively. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice.
Use these coverage values:

9
tests/rust-python-harness/cli.py
View file

@ -6,14 +6,13 @@ from collections.abc import Sequence
from pathlib import Path
from .catalog import load_catalog
from .models import HarnessCase, Strategy
from .models import SDK_FUNCTIONS, HarnessCase, Strategy
from .runner import run_pytest
from .ui import make_dashboard
from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report
REPO_ROOT = Path(__file__).resolve().parents[2]
COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness"
SDK_FUNCTION_CHOICES = ("ocr", "messages", "responses", "count_tokens")
def _parser() -> argparse.ArgumentParser:
@ -42,7 +41,7 @@ def _parser() -> argparse.ArgumentParser:
action="append",
default=[],
dest="sdk_functions",
choices=SDK_FUNCTION_CHOICES,
choices=SDK_FUNCTIONS,
help="run only this SDK function",
)
parser.add_argument(
@ -110,7 +109,7 @@ def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[
)
sdk_functions = _pick_values(
"SDK functions",
[(name, name) for name in SDK_FUNCTION_CHOICES],
[(name, name) for name in SDK_FUNCTIONS],
)
return strategy_ids, sdk_functions
@ -166,7 +165,7 @@ def _print_function_report(report: FunctionReport) -> None:
def _validate_ledger(sdk_functions: set[str]) -> int:
functions = sdk_functions or set(SDK_FUNCTION_CHOICES)
functions = sdk_functions or set(SDK_FUNCTIONS)
reports = tuple(build_function_report(function) for function in sorted(functions))
for report in reports:
_print_function_report(report)

4
tests/rust-python-harness/e2e_fuzz_tests/strategy.json
View file

@ -7,6 +7,8 @@
"ocr": {"coverage": "partial", "selectors": ["tests/test_litellm/ocr/test_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
"messages": {"coverage": "partial", "selectors": ["tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
"responses": {"coverage": "partial", "selectors": ["tests/test_litellm/responses/test_rust_bridge_websocket.py"], "note": "Covers the websocket bridge; full responses parity is still being added."},
"count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."}
"count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."},
"chat_completions": {"coverage": "partial", "selectors": ["tests/test_litellm/rust_bridge/test_chat_completions.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."},
"transcription": {"coverage": "partial", "selectors": ["tests/test_litellm/test_audio_transcription_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}
}
}

3
tests/rust-python-harness/existing_e2e_test_sdk/README.md Normal file
View file

@ -0,0 +1,3 @@
# Existing e2e SDK tests
Wires already-existing live-API SDK tests into the matrix instead of writing new parity tests. Selectors point at real test files and folders, such as `tests/ocr_tests/`, rather than individual node IDs, so future tests added to those folders are picked up automatically.

14
tests/rust-python-harness/existing_e2e_test_sdk/strategy.json Normal file
View file

@ -0,0 +1,14 @@
{
"order": 40,
"id": "existing_e2e_test_sdk",
"label": "Existing e2e SDK tests",
"description": "Wire already-existing live-API SDK tests into the matrix instead of writing new parity tests.",
"functions": {
"ocr": {"coverage": "partial", "selectors": ["tests/ocr_tests/"], "note": "Existing live OCR provider tests; not yet a frozen Rust/Python oracle comparison."},
"messages": {"coverage": "planned", "selectors": []},
"responses": {"coverage": "planned", "selectors": []},
"count_tokens": {"coverage": "planned", "selectors": []},
"chat_completions": {"coverage": "partial", "selectors": ["tests/llm_translation/test_anthropic_completion.py", "tests/llm_translation/test_bedrock_completion.py"], "note": "Existing live chat completion tests for providers with confirmed Rust bridge regressions."},
"transcription": {"coverage": "partial", "selectors": ["tests/audio_tests/test_whisper.py"], "note": "Existing live Whisper transcription test."}
}
}

2
tests/rust-python-harness/models.py
View file

@ -33,7 +33,7 @@ class ConfidenceLevel(str, Enum):
LOW = "LOW"
SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens")
SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription")
@dataclass(frozen=True)

2
tests/rust-python-harness/runner.py
View file

@ -15,6 +15,8 @@ UpdateCallback = Callable[[HarnessRun], None]
def selector_matches_node(selector: str, nodeid: str) -> bool:
normalized_selector = selector.replace("\\", "/")
normalized_nodeid = nodeid.replace("\\", "/")
if normalized_selector.endswith("/"):
return normalized_nodeid.startswith(normalized_selector)
if "::" in normalized_selector:
return normalized_nodeid == normalized_selector or normalized_nodeid.startswith(
f"{normalized_selector}["

1240
tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json
View file

File diff suppressed because it is too large Load diff

2
tests/rust-python-harness/ui.py
View file

@ -119,7 +119,7 @@ class RichDashboard(AbstractContextManager["RichDashboard"]):
table = Table(box=box.ROUNDED, expand=True, title="Strategy × SDK function")
table.add_column("Strategy", ratio=3)
for label in ("ocr/aocr", "messages", "responses", "count_tokens"):
for label in SDK_FUNCTIONS:
table.add_column(label, justify="center", ratio=1)
for strategy in self.strategies:
cells = []

4
tests/rust-python-harness/unit_tests_rust/strategy.json
View file

@ -7,6 +7,8 @@
"ocr": {"coverage": "planned", "selectors": []},
"messages": {"coverage": "planned", "selectors": []},
"responses": {"coverage": "planned", "selectors": []},
"count_tokens": {"coverage": "planned", "selectors": []}
"count_tokens": {"coverage": "planned", "selectors": []},
"chat_completions": {"coverage": "planned", "selectors": []},
"transcription": {"coverage": "planned", "selectors": []}
}
}

4
tests/rust-python-harness/validate_sub_methods/strategy.json
View file

@ -7,6 +7,8 @@
"ocr": {"coverage": "planned", "selectors": []},
"messages": {"coverage": "planned", "selectors": []},
"responses": {"coverage": "planned", "selectors": []},
"count_tokens": {"coverage": "planned", "selectors": []}
"count_tokens": {"coverage": "planned", "selectors": []},
"chat_completions": {"coverage": "planned", "selectors": []},
"transcription": {"coverage": "planned", "selectors": []}
}
}

91
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -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

301
tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py
View file

@ -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 = {}

96
tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py
View file

@ -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

119
tests/test_litellm/llms/xai/test_xai_chat_transformation.py
View file

@ -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,
)

184
tests/test_litellm/llms/xai/test_xai_cost_calculator.py
View file

@ -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."""

6
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -915,6 +915,7 @@ async def test_mcp_get_prompt_success():
)
mock_manager.get_prompt_from_server.assert_awaited_once_with(
server=server,
user_api_key_auth=user_api_key_auth,
prompt_name="hello",
arguments={"foo": "bar"},
mcp_auth_header={"Authorization": "token"},
@ -977,6 +978,7 @@ async def test_mcp_read_resource_success():
)
mock_manager.read_resource_from_server.assert_awaited_once_with(
server=server,
user_api_key_auth=user_api_key_auth,
url="https://example.com/resource",
mcp_auth_header={"Authorization": "token"},
extra_headers={"X-Test": "1"},
@ -8587,7 +8589,9 @@ class TestOboPreflightScopedToAllowedServers:
_, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key)
preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key)
preflight.assert_awaited_once_with(
server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None
)
@pytest.mark.asyncio

311
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -30,6 +30,7 @@ from mcp.types import (
TextResourceContents,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm.constants import MCP_METADATA_TIMEOUT
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
@ -2270,7 +2271,9 @@ class TestMCPServerManager:
"""prompts/list on an OBO server must exchange the caller's bearer, not connect with none."""
server = self._token_exchange_server("te-prompts")
st = await self._capture_subject_token(
lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_prompts_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st == "subj-jwt"
@ -2279,7 +2282,9 @@ class TestMCPServerManager:
"""resources/list on an OBO server must exchange the caller's bearer."""
server = self._token_exchange_server("te-resources")
st = await self._capture_subject_token(
lambda m: m.get_resources_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_resources_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st == "subj-jwt"
@ -2290,6 +2295,7 @@ class TestMCPServerManager:
st = await self._capture_subject_token(
lambda m: m.read_resource_from_server(
server=server,
user_api_key_auth=None,
url="https://up.example.com/r",
raw_headers={"authorization": "Bearer subj-jwt"},
)
@ -2307,7 +2313,9 @@ class TestMCPServerManager:
auth_type=MCPAuth.none,
)
st = await self._capture_subject_token(
lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"})
lambda m: m.get_prompts_from_server(
server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"}
)
)
assert st is None
@ -3254,7 +3262,7 @@ class TestMCPServerManager:
new_callable=AsyncMock,
return_value=mock_client,
):
prompts = await manager.get_prompts_from_server(server, add_prefix=True)
prompts = await manager.get_prompts_from_server(server, user_api_key_auth=None, add_prefix=True)
mock_client.list_prompts.assert_awaited_once()
assert len(prompts) == 1
@ -3289,6 +3297,7 @@ class TestMCPServerManager:
):
result = await manager.get_prompt_from_server(
server=server,
user_api_key_auth=None,
prompt_name="hello",
arguments={"tone": "casual"},
)
@ -3334,6 +3343,7 @@ class TestMCPServerManager:
):
result = await manager.get_resources_from_server(
server=server,
user_api_key_auth=None,
mcp_auth_header="auth",
extra_headers={"X-Test": "1"},
add_prefix=True,
@ -3391,6 +3401,7 @@ class TestMCPServerManager:
):
result = await manager.get_resource_templates_from_server(
server=server,
user_api_key_auth=None,
mcp_auth_header="auth",
extra_headers=None,
add_prefix=False,
@ -3441,6 +3452,7 @@ class TestMCPServerManager:
) as mock_create_client:
result = await manager.read_resource_from_server(
server=server,
user_api_key_auth=None,
url="https://example.com/resource",
mcp_auth_header="auth",
extra_headers={"X-Test": "1"},
@ -11006,3 +11018,294 @@ class TestOpenApiHandlerRelaysUpstreamAuth:
assert result.isError is True
assert "upstream returned HTTP 503" in result.content[0].text
class TestLitellmAdmissionKeyIsNeverTheSubjectToken:
"""The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the
RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it
present, ``Authorization`` is the caller's own identity token and is exchanged as before."""
_ADMISSION_KEY: Final = "sk-litellm-virtual-key"
_USER_TOKEN: Final = "user-idp-jwt"
@staticmethod
def _token_exchange_server(server_id: str) -> MCPServer:
return MCPServer(
server_id=server_id,
name=f"{server_id}-server",
url="https://up.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
)
@staticmethod
def _id_jag_server(server_id: str) -> MCPServer:
return MCPServer(
server_id=server_id,
name=f"{server_id}-server",
url="https://up.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_id_jag,
client_id="cid",
client_secret="csec",
token_exchange_endpoint="https://idp.example.com/token",
id_jag_resource_token_endpoint="https://resource-as.example.com/token",
)
@staticmethod
def _recording_provider() -> MagicMock:
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
provider: Final = MagicMock()
provider.resolve_credentials = AsyncMock(
return_value=Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization"))
)
return provider
@staticmethod
def _subjects_seen_by(provider: MagicMock) -> list[str | None]:
return [
call.args[0].inbound_token.get_secret_value() if call.args[0].inbound_token else None
for call in provider.resolve_credentials.call_args_list
]
@staticmethod
def _manager_with_recording_client() -> MCPServerManager:
manager: Final = MCPServerManager()
client: Final = AsyncMock()
client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False))
client.list_prompts = AsyncMock(return_value=[])
client.read_resource = AsyncMock(return_value=ReadResourceResult(contents=[]))
manager._create_mcp_client = AsyncMock(return_value=client)
return manager
@staticmethod
def _subject_token_given_to_client(manager: MCPServerManager) -> str | None:
return manager._create_mcp_client.call_args.kwargs["subject_token"]
async def _call_tool_subject(self, server: MCPServer, oauth2_headers, raw_headers, user_api_key_auth):
manager: Final = self._manager_with_recording_client()
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="tool",
arguments={},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
proxy_logging_obj=None,
user_api_key_auth=user_api_key_auth,
)
return self._subject_token_given_to_client(manager)
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag])
async def test_tools_call_with_only_the_litellm_key_has_no_subject(self, auth_type):
server = (
self._token_exchange_server("te-call")
if auth_type == MCPAuth.oauth2_token_exchange
else self._id_jag_server("jag-call")
)
subject_token = await self._call_tool_subject(
server,
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_rest_tools_call_with_only_the_litellm_key_has_no_subject(self):
"""The REST facade passes no oauth2_headers; the bearer is reached through raw_headers only."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-rest"),
oauth2_headers=None,
raw_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-split"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={
"X-LiteLLM-API-Key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token == self._USER_TOKEN
@pytest.mark.asyncio
async def test_tools_call_with_an_empty_x_litellm_api_key_has_no_subject(self):
"""Admission ignores an empty ``x-litellm-api-key`` and validates ``Authorization`` instead."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-empty-header"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"x-litellm-api-key": "", "authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_with_the_same_litellm_key_in_both_headers_has_no_subject(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-same-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={
"x-litellm-api-key": self._ADMISSION_KEY,
"authorization": f"Bearer {self._ADMISSION_KEY}",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_with_a_different_litellm_key_in_authorization_has_no_subject(self):
"""A second ``sk-`` virtual key next to ``x-litellm-api-key`` is still a gateway credential."""
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-second-key"),
oauth2_headers={"Authorization": "Bearer sk-another-virtual-key"},
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": "Bearer sk-another-virtual-key",
},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert subject_token is None
@pytest.mark.asyncio
async def test_tools_call_exchanges_the_bearer_when_jwt_admission_left_api_key_unset(self):
subject_token = await self._call_tool_subject(
self._token_exchange_server("te-jwt"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={"authorization": f"Bearer {self._USER_TOKEN}"},
user_api_key_auth=UserAPIKeyAuth(api_key=None, user_id="alice"),
)
assert subject_token == self._USER_TOKEN
@pytest.mark.asyncio
async def test_tools_list_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
manager._fetch_tools_with_timeout = AsyncMock(return_value=[])
await manager._get_tools_from_server(
server=self._token_exchange_server("te-list-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_prompts_list_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
await manager.get_prompts_from_server(
server=self._token_exchange_server("te-prompts-key"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_resource_read_with_only_the_litellm_key_has_no_subject(self):
manager: Final = self._manager_with_recording_client()
await manager.read_resource_from_server(
server=self._token_exchange_server("te-read-key"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
url=AnyUrl("file:///notes.txt"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert self._subject_token_given_to_client(manager) is None
@pytest.mark.asyncio
async def test_resource_read_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
manager: Final = self._manager_with_recording_client()
await manager.read_resource_from_server(
server=self._token_exchange_server("te-read-split"),
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
url=AnyUrl("file:///notes.txt"),
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
)
assert self._subject_token_given_to_client(manager) == self._USER_TOKEN
@pytest.mark.asyncio
async def test_openapi_call_never_hands_the_litellm_key_to_the_exchanger(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
server = MCPServer(
server_id="te-openapi",
name="te_openapi",
server_name="te_openapi",
url=None,
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
spec_path="https://api.example.com/openapi.json",
)
user_auth = UserAPIKeyAuth(api_key="hashed-key", user_id="alice")
await manager.resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
mcp_auth_header=None,
user_api_key_auth=user_auth,
forwarded_headers=None,
)
await manager.resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
mcp_auth_header=None,
user_api_key_auth=user_auth,
forwarded_headers=None,
)
assert self._subjects_seen_by(provider) == [None, self._USER_TOKEN]
@pytest.mark.asyncio
async def test_preflight_challenges_instead_of_exchanging_the_litellm_key(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
with pytest.raises(HTTPException) as exc_info:
await manager.preflight_token_exchange(
server=self._token_exchange_server("te-preflight-key"),
oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"},
)
assert exc_info.value.status_code == 401
headers = exc_info.value.headers or {}
assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "")
assert self._subjects_seen_by(provider) == []
@pytest.mark.asyncio
async def test_preflight_exchanges_the_user_token_when_x_litellm_api_key_admits(self):
provider: Final = self._recording_provider()
manager = MCPServerManager(cred_provider=provider)
await manager.preflight_token_exchange(
server=self._token_exchange_server("te-preflight-split"),
oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"},
user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"),
raw_headers={
"x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}",
"authorization": f"Bearer {self._USER_TOKEN}",
},
)
assert self._subjects_seen_by(provider) == [self._USER_TOKEN]

33
tests/test_litellm/proxy/management_endpoints/test_saml_sso.py
View file

@ -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()

129
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -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

10
tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
View file

@ -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)

30
tests/test_litellm/test_main.py
View file

@ -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)

25
tests/test_litellm/vector_stores/test_main.py
View file

@ -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():

16
tests/test_rust_python_harness.py
View file

@ -66,13 +66,14 @@ def _manifest() -> dict[str, object]:
}
def test_should_load_the_three_harness_strategies_in_order() -> None:
def test_should_load_the_four_harness_strategies_in_order() -> None:
strategies = load_catalog()
assert [strategy.id for strategy in strategies] == [
"e2e_fuzz_tests",
"unit_tests_rust",
"validate_sub_methods",
"existing_e2e_test_sdk",
]
assert all(
tuple(case.sdk_function for case in strategy.cases) == SDK_FUNCTIONS
@ -104,6 +105,8 @@ def test_should_reject_a_manifest_missing_an_sdk_function(tmp_path: Path) -> Non
True,
),
("tests/test_parity.py::test_one", "tests/test_parity.py::test_two", False),
("tests/ocr_tests/", "tests/ocr_tests/test_ocr_mistral.py::test_one", True),
("tests/ocr_tests/", "tests/other_tests/test_ocr_mistral.py::test_one", False),
],
)
def test_should_match_pytest_file_and_node_selectors(
@ -123,6 +126,13 @@ def test_should_only_return_selectors_whose_files_exist(tmp_path: Path) -> None:
assert runnable_selectors((case,), tmp_path) == ("tests/test_parity.py",)
def test_should_treat_an_existing_folder_selector_as_runnable(tmp_path: Path) -> None:
(tmp_path / "tests" / "ocr_tests").mkdir(parents=True)
case = _case(selectors=("tests/ocr_tests/",))
assert runnable_selectors((case,), tmp_path) == ("tests/ocr_tests/",)
def test_should_mark_planned_and_not_applicable_cases_without_running() -> None:
planned = CaseResult(case=_case(coverage=Coverage.PLANNED))
not_applicable = CaseResult(case=_case(coverage=Coverage.NOT_APPLICABLE))
@ -239,8 +249,8 @@ def test_should_report_confidence_for_each_sdk_section() -> None:
}
assert scores["responses"].verified_strategies == 1
assert scores["responses"].required_strategies == 3
assert scores["responses"].percentage == 33
assert scores["responses"].required_strategies == 4
assert scores["responses"].percentage == 25
assert scores["responses"].level.value == "MEDIUM"
assert scores["count_tokens"].percentage == 0
assert scores["count_tokens"].level.value == "LOW"

4
type-discipline-budget.json
View file

@ -27,10 +27,10 @@
"limit": 0
},
"LIT010": {
"limit": 16470
"limit": 16469
},
"LIT011": {
"limit": 5520
"limit": 5519
},
"LIT012": {
"limit": 4489

3
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -1399,9 +1399,6 @@
},
"prefer-const": {
"count": 2
},
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/components/TeamsPage/teamTableColumns.tsx": {

202
ui/litellm-dashboard/src/components/Teams.test.tsx
View file

@ -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(<Teams accessToken="test-token" userID="user-123" userRole="Admin" />);
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<string[]>((resolve) => {
resolveModels = resolve;
}),
);
renderWithQueryClient(<Teams accessToken="test-token" userID="user-123" userRole="Admin" />);
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(<Teams accessToken="test-token" userID="user-123" userRole="Admin" />);
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(<Teams accessToken="test-token" userID="user-123" userRole="Admin" />);
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(<Teams accessToken="test-token" userID="user-123" userRole="Internal User" />);
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(<Teams accessToken="test-token" userID="user-123" userRole="Internal User" />);
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(<Teams accessToken="test-token" userID="user-123" userRole="Internal User" />);
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(<Teams accessToken="test-token" userID="user-123" userRole="Internal User" />);
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(<Teams accessToken="test-token" userID="user-123" userRole="Admin" />);
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("");
});
});

74
ui/litellm-dashboard/src/components/Teams.tsx
View file

@ -208,7 +208,6 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
const queryClient = useQueryClient();
const refreshTeams = () => queryClient.invalidateQueries({ queryKey: teamsTableKeys.all });
const [currentOrg] = useState<Organization | null>(null);
const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState<Organization | null>(null);
const isOrgAdmin = userRole !== "Admin";
const [additionalSettingsOpen, setAdditionalSettingsOpen] = useState(false);
@ -216,17 +215,33 @@ const Teams: React.FC<TeamProps> = ({ 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<TeamProps> = ({ 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<TeamProps> = ({ 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<TeamProps> = ({ accessToken, userID, userRole, premiumUser
subtitle="Manage teams, members, and their access to models and budgets"
primaryAction={
canCreateOrManageTeams(userRole, userID, organizations) ? (
<UIButton onClick={() => setIsTeamModalVisible(true)} data-testid="create-team-button">
<UIButton onClick={openCreateTeamModal} data-testid="create-team-button">
<Plus className="size-4" />
Create Team
</UIButton>
@ -683,9 +696,9 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
)}
</FormField>
{(() => {
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<TeamProps> = ({ 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)}
/>
)}
</FormField>

18
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
View file

@ -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([

32
ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
View file

@ -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<string>) => 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;

9
ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts
View file

@ -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,