mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_registry_audit_2026_09_02
This commit is contained in:
commit
5a3a2f3d0a
99 changed files with 4667 additions and 475 deletions
|
|
@ -932,7 +932,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
|
||||
|
||||
def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]:
|
||||
if role == "user" or role == "system" or role == "tool":
|
||||
if role in ("user", "system", "developer", "tool"):
|
||||
return {"type": "input_text", "text": content}
|
||||
else:
|
||||
return {"type": "output_text", "text": content}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_ran
|
|||
|
||||
DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
|
||||
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
|
||||
AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"})
|
||||
ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
|
||||
ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000
|
||||
RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
|
||||
|
|
@ -39,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
|
|||
"router_general_settings",
|
||||
"ignore_invalid_deployments",
|
||||
"fallback_access_check",
|
||||
"heuristic_v2_router_limit",
|
||||
}
|
||||
)
|
||||
DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
|
||||
|
|
|
|||
|
|
@ -18,9 +18,11 @@ caller's identity metadata, minus two things that must never be forwarded as-is:
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES
|
||||
from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params
|
||||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin
|
||||
|
||||
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
|
||||
|
|
@ -142,6 +144,19 @@ def forwarded_internal_call_metadata(
|
|||
}
|
||||
|
||||
|
||||
def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]:
|
||||
kwargs: Final = request_kwargs or MappingProxyType({})
|
||||
return MappingProxyType(
|
||||
{k: v for k in ("litellm_session_id", "litellm_trace_id") if isinstance(v := kwargs.get(k), str)}
|
||||
)
|
||||
|
||||
|
||||
def effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None:
|
||||
return initialize_standard_callback_dynamic_params(dict(request_kwargs) if request_kwargs else None).get(
|
||||
"turn_off_message_logging"
|
||||
)
|
||||
|
||||
|
||||
def sanitized_forwardable_call_metadata(
|
||||
parent_metadata: Mapping[str, object],
|
||||
call_origin: InternalCallOrigin,
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ Helper utilities for tracking the cost of built-in tools.
|
|||
from collections.abc import Mapping
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
|
|
@ -13,6 +15,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
|||
from litellm.types.llms.openai import (
|
||||
FileSearchTool,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesToolUsage,
|
||||
WebSearchOptions,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -32,6 +35,17 @@ def _output_item_type(output_item: object) -> str | None:
|
|||
return item_type if isinstance(item_type, str) else None
|
||||
|
||||
|
||||
def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None:
|
||||
tool_usage: Final = getattr(response_object, "tool_usage", None)
|
||||
if tool_usage is None:
|
||||
return None
|
||||
try:
|
||||
web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search
|
||||
except ValidationError:
|
||||
return None
|
||||
return None if web_search is None else web_search.num_requests
|
||||
|
||||
|
||||
def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool:
|
||||
details: Final = getattr(usage, "server_side_tool_usage_details", None)
|
||||
if not isinstance(details, Mapping):
|
||||
|
|
@ -182,15 +196,19 @@ class StandardBuiltInToolCostTracking:
|
|||
|
||||
Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by
|
||||
get_cost_for_web_search_request and never reach here. This path prices per call, so it must count
|
||||
the web_search_call items. Chat-completions responses only expose url_citation annotations with no
|
||||
count, so they floor to a single billable search.
|
||||
the web_search_call items, unless the response reports the billable count itself
|
||||
(Bedrock's tool_usage.web_search.num_requests, which excludes open_page fetches). Chat-completions
|
||||
responses only expose url_citation annotations with no count, so they floor to a single billable search.
|
||||
"""
|
||||
if isinstance(response_object, ResponsesAPIResponse):
|
||||
count = sum(
|
||||
1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call"
|
||||
)
|
||||
return max(count, 1)
|
||||
return 1
|
||||
if not isinstance(response_object, ResponsesAPIResponse):
|
||||
return 1
|
||||
reported: Final = _reported_web_search_requests(response_object)
|
||||
if reported is not None:
|
||||
return reported
|
||||
count: Final = sum(
|
||||
1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call"
|
||||
)
|
||||
return max(count, 1)
|
||||
|
||||
@staticmethod
|
||||
def _handle_file_search_cost(
|
||||
|
|
|
|||
|
|
@ -428,7 +428,7 @@ def _coerce_off_peak_rate(value: object, default: float) -> float:
|
|||
return default
|
||||
|
||||
|
||||
def _apply_off_peak_pricing(
|
||||
def apply_off_peak_pricing(
|
||||
model_info: ModelInfo,
|
||||
current_time: datetime | None,
|
||||
prompt_base_cost: float,
|
||||
|
|
@ -462,7 +462,7 @@ def _apply_off_peak_to_base_costs(
|
|||
has no field for them.
|
||||
"""
|
||||
prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
|
||||
off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing(
|
||||
off_peak_prompt, off_peak_completion, off_peak_cache_read = apply_off_peak_pricing(
|
||||
model_info, current_time, prompt, completion, cache_read
|
||||
)
|
||||
return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read)
|
||||
|
|
|
|||
|
|
@ -2337,6 +2337,9 @@ class CustomStreamWrapper:
|
|||
else:
|
||||
self.sent_last_chunk = True
|
||||
processed_chunk: Final = self.finish_reason_handler()
|
||||
if self.stream_options is None:
|
||||
usage: Final = calculate_total_usage(chunks=self.chunks)
|
||||
processed_chunk._hidden_params["usage"] = usage # pyright: ignore[reportPrivateUsage] # sync parity
|
||||
# see sync __next__'s sibling branch: deliberately do NOT restore
|
||||
# here - this chunk is still this call's own data, and restoring
|
||||
# before returning it would corrupt the caller's own log
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from litellm.exceptions import UnsupportedParamsError
|
|||
from litellm.llms.openai.chat.gpt_5_transformation import (
|
||||
OpenAIGPT5Config,
|
||||
_get_effort_level,
|
||||
is_gpt_reasoning_series_name,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
|
|
@ -35,26 +36,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
|
|||
|
||||
@classmethod
|
||||
def is_model_gpt_5_model(cls, model: str) -> bool:
|
||||
"""Check if the Azure model string refers to a gpt-5 variant.
|
||||
|
||||
Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix
|
||||
used for manual routing.
|
||||
"""
|
||||
# The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
|
||||
# …) are regular chat models: they support temperature and tool_choice but NOT
|
||||
# reasoning_effort. They must NOT be routed through the GPT-5 reasoning path.
|
||||
#
|
||||
# Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
|
||||
# models and must stay on the GPT-5 path. The distinguishing feature is that
|
||||
# the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
|
||||
# (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
|
||||
# number (i.e. "gpt-5.<digit>-chat").
|
||||
#
|
||||
# Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
|
||||
# than a substring check) makes this boundary explicit and avoids any ambiguity
|
||||
# if future model names coincidentally contain "gpt-5-chat" as an interior run.
|
||||
_normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "azure/"
|
||||
return ("gpt-5" in model and not _normalized.startswith("gpt-5-chat")) or "gpt5_series" in model
|
||||
return is_gpt_reasoning_series_name(model) or "gpt5_series" in model
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
"""Get supported parameters for Azure OpenAI GPT-5 models.
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
convert_to_azure_openai_messages,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import GPT_REASONING_SERIES_MARKERS
|
||||
from litellm.types.llms.azure import (
|
||||
API_VERSION_MONTH_SUPPORTED_RESPONSE_FORMAT,
|
||||
API_VERSION_YEAR_SUPPORTED_RESPONSE_FORMAT,
|
||||
|
|
@ -139,7 +140,7 @@ class AzureOpenAIConfig(BaseConfig):
|
|||
name family needs the rename, including the ``gpt-5-chat*`` models that are excluded from
|
||||
the reasoning path by https://github.com/BerriAI/litellm/issues/13781.
|
||||
"""
|
||||
return "gpt-5" in model or "gpt5_series" in model
|
||||
return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) or "gpt5_series" in model
|
||||
|
||||
def _is_response_format_supported_model(self, model: str) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
filter_value_from_dict,
|
||||
)
|
||||
from litellm.llms.azure.common_utils import BaseAzureLLM
|
||||
from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
|
||||
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
|
||||
from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error
|
||||
from litellm.llms.openai.openai import OpenAIConfig
|
||||
|
|
@ -207,20 +208,18 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
message["content"] = texts
|
||||
return stripped_messages
|
||||
|
||||
def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
|
||||
try:
|
||||
if "/" in model:
|
||||
model = model.split("/", 1)[1]
|
||||
if (
|
||||
model in litellm.open_ai_chat_completion_models
|
||||
or model in litellm.open_ai_text_completion_models
|
||||
or model in litellm.open_ai_embedding_models
|
||||
):
|
||||
return True
|
||||
def _is_foundry_model_inference_base(self, api_base: str) -> bool:
|
||||
return is_foundry_model_inference_base(api_base)
|
||||
|
||||
except Exception:
|
||||
def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
|
||||
if api_base is None or self._is_foundry_model_inference_base(api_base):
|
||||
return False
|
||||
return False
|
||||
stripped_model: Final = model.split("/", 1)[1] if "/" in model else model
|
||||
return (
|
||||
stripped_model in litellm.open_ai_chat_completion_models
|
||||
or stripped_model in litellm.open_ai_text_completion_models
|
||||
or stripped_model in litellm.open_ai_embedding_models
|
||||
)
|
||||
|
||||
def _get_openai_compatible_provider_info(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
||||
|
|
@ -10,6 +11,14 @@ from litellm.types.router import GenericLiteLLMParams
|
|||
AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
|
||||
|
||||
|
||||
def is_foundry_model_inference_base(api_base: str) -> bool:
|
||||
parsed: Final = urlparse(api_base)
|
||||
host: Final = parsed.hostname
|
||||
if host is None or not host.endswith(".services.ai.azure.com"):
|
||||
return False
|
||||
return "/openai/deployments" not in parsed.path
|
||||
|
||||
|
||||
def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
|
||||
"""
|
||||
Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment.
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
from typing import Final
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
import litellm
|
||||
from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -16,6 +18,16 @@ from litellm.utils import convert_to_model_response_object
|
|||
from .cohere_transformation import AzureAICohereConfig
|
||||
|
||||
|
||||
def _foundry_models_route_base(api_base: str | None) -> str | None:
|
||||
if api_base is None or not is_foundry_model_inference_base(api_base):
|
||||
return api_base
|
||||
parts: Final = urlsplit(api_base)
|
||||
path: Final = parts.path.rstrip("/")
|
||||
if path.endswith("/models"):
|
||||
return api_base
|
||||
return urlunsplit((parts.scheme, parts.netloc, f"{path}/models", parts.query, parts.fragment))
|
||||
|
||||
|
||||
class AzureAIEmbedding(OpenAIChatCompletion):
|
||||
def _process_response(
|
||||
self,
|
||||
|
|
@ -214,6 +226,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
|
|||
|
||||
assemble result in-order, and return
|
||||
"""
|
||||
resolved_api_base: Final = _foundry_models_route_base(api_base)
|
||||
if aembedding is True:
|
||||
return self.async_embedding(
|
||||
model,
|
||||
|
|
@ -223,7 +236,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
|
|||
model_response,
|
||||
optional_params,
|
||||
api_key,
|
||||
api_base,
|
||||
resolved_api_base,
|
||||
client,
|
||||
)
|
||||
|
||||
|
|
@ -245,7 +258,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
|
|||
model_response=model_response,
|
||||
optional_params=optional_params,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
api_base=resolved_api_base,
|
||||
client=client,
|
||||
)
|
||||
|
||||
|
|
@ -262,7 +275,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
|
|||
model_response,
|
||||
optional_params,
|
||||
api_key,
|
||||
api_base,
|
||||
resolved_api_base,
|
||||
client=(client if client is not None and isinstance(client, OpenAI) else None),
|
||||
aembedding=aembedding,
|
||||
shared_session=shared_session,
|
||||
|
|
|
|||
|
|
@ -99,9 +99,12 @@ from litellm.types.containers.main import (
|
|||
)
|
||||
from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
AgenticLoopSafetyError,
|
||||
converted_stream_requested,
|
||||
is_interception_internal_key,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
|
|
@ -2760,6 +2763,7 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
if self._has_agentic_completion_hook(logging_obj):
|
||||
agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place
|
||||
final_response: Final = run_async_function(
|
||||
self._call_agentic_completion_hooks,
|
||||
response=initial_response,
|
||||
|
|
@ -2770,10 +2774,19 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
stream=False,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=dict(litellm_params),
|
||||
kwargs=agentic_kwargs,
|
||||
api_surface="responses",
|
||||
)
|
||||
return final_response if final_response is not None else initial_response
|
||||
result: Final = final_response if final_response is not None else initial_response
|
||||
if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"):
|
||||
return self._wrap_responses_response_as_fake_stream(
|
||||
result=result,
|
||||
model=model,
|
||||
responses_api_provider_config=responses_api_provider_config,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
return result
|
||||
|
||||
return initial_response
|
||||
|
||||
|
|
@ -2939,6 +2952,7 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place
|
||||
final_response: Final = await self._call_agentic_completion_hooks(
|
||||
response=initial_response,
|
||||
model=model,
|
||||
|
|
@ -2948,15 +2962,12 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
stream=False,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=dict(litellm_params),
|
||||
kwargs=agentic_kwargs,
|
||||
api_surface="responses",
|
||||
)
|
||||
|
||||
result: Final = final_response if final_response is not None else initial_response
|
||||
interception_converted_stream: Final = litellm_params.get(
|
||||
"_code_interpreter_interception_converted_stream"
|
||||
) or litellm_params.get("_websearch_interception_converted_stream")
|
||||
if interception_converted_stream and not litellm_params.get("_agentic_loop_depth"):
|
||||
if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"):
|
||||
return self._wrap_responses_response_as_fake_stream(
|
||||
result=result,
|
||||
model=model,
|
||||
|
|
@ -5420,8 +5431,7 @@ class BaseLLMHTTPHandler:
|
|||
kwargs_for_followup: Final = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if not k.startswith("_websearch_interception")
|
||||
and not k.startswith("_compression_interception")
|
||||
if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES)
|
||||
and k != "_code_interpreter_interception_converted_stream"
|
||||
and k not in internal_keys
|
||||
and k not in optional_params
|
||||
|
|
|
|||
|
|
@ -7,11 +7,13 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate.
|
|||
See https://help.aliyun.com/zh/model-studio/billing-for-model-studio
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
apply_off_peak_pricing,
|
||||
parse_completion_tokens_details,
|
||||
parse_prompt_tokens_details,
|
||||
)
|
||||
|
|
@ -32,6 +34,19 @@ class TokenBreakdown:
|
|||
return self.text_tokens + self.cached_tokens + self.cache_creation_tokens
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenRates:
|
||||
input_rate: float
|
||||
cache_read_rate: float
|
||||
cache_creation_rate: float
|
||||
output_rate: float
|
||||
reasoning_rate: float | None
|
||||
|
||||
@property
|
||||
def billed_reasoning_rate(self) -> float:
|
||||
return self.output_rate if self.reasoning_rate is None else self.reasoning_rate
|
||||
|
||||
|
||||
def _extract_token_breakdown(usage: Usage) -> TokenBreakdown:
|
||||
prompt_details: Final = parse_prompt_tokens_details(usage)
|
||||
cached_tokens: Final = prompt_details["cache_hit_tokens"]
|
||||
|
|
@ -57,69 +72,75 @@ def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) ->
|
|||
return float(value)
|
||||
|
||||
|
||||
def _calculate_prompt_cost(
|
||||
breakdown: TokenBreakdown,
|
||||
model_info: ModelInfo,
|
||||
tier: dict | None,
|
||||
) -> float:
|
||||
if tier is not None:
|
||||
return (
|
||||
(breakdown.text_tokens * tier_rate(tier, "input_cost_per_token"))
|
||||
+ (breakdown.cached_tokens * tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"))
|
||||
+ (
|
||||
breakdown.cache_creation_tokens
|
||||
* tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token")
|
||||
)
|
||||
)
|
||||
|
||||
input_cost: Final = float(model_info.get("input_cost_per_token") or 0.0)
|
||||
cache_read_cost: Final = _flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token")
|
||||
cache_creation_cost: Final = _flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token")
|
||||
|
||||
return (
|
||||
(breakdown.text_tokens * input_cost)
|
||||
+ (breakdown.cached_tokens * cache_read_cost)
|
||||
+ (breakdown.cache_creation_tokens * cache_creation_cost)
|
||||
def _flat_rates(model_info: ModelInfo) -> TokenRates:
|
||||
reasoning_rate: Final = model_info.get("output_cost_per_reasoning_token")
|
||||
return TokenRates(
|
||||
input_rate=float(model_info.get("input_cost_per_token") or 0.0),
|
||||
cache_read_rate=_flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token"),
|
||||
cache_creation_rate=_flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token"),
|
||||
output_rate=float(model_info.get("output_cost_per_token") or 0.0),
|
||||
reasoning_rate=None if reasoning_rate is None else float(reasoning_rate),
|
||||
)
|
||||
|
||||
|
||||
def _calculate_completion_cost(
|
||||
breakdown: TokenBreakdown,
|
||||
model_info: ModelInfo,
|
||||
tier: dict | None,
|
||||
) -> float:
|
||||
def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates:
|
||||
# A tier that declares output rates keeps the request on them, all-or-nothing. A tier table
|
||||
# spelling out only input rates would serve every completion for free, so there the model's
|
||||
# own output rates stand in
|
||||
tier_declares_output: Final = tier is not None and "output_cost_per_token" in tier
|
||||
output_cost: Final = (
|
||||
tier_rate(tier, "output_cost_per_token")
|
||||
if tier_declares_output
|
||||
else float(model_info.get("output_cost_per_token") or 0.0)
|
||||
)
|
||||
tier_declares_reasoning: Final = tier is not None and "output_cost_per_reasoning_token" in tier
|
||||
model_reasoning_rate: Final = None if tier_declares_output else model_info.get("output_cost_per_reasoning_token")
|
||||
reasoning_cost: Final = (
|
||||
tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token")
|
||||
if tier_declares_reasoning
|
||||
else float(model_reasoning_rate)
|
||||
if model_reasoning_rate is not None
|
||||
else output_cost
|
||||
flat_rates: Final = _flat_rates(model_info)
|
||||
tier_declares_output: Final = "output_cost_per_token" in tier
|
||||
tier_declares_reasoning: Final = "output_cost_per_reasoning_token" in tier
|
||||
return TokenRates(
|
||||
input_rate=tier_rate(tier, "input_cost_per_token"),
|
||||
cache_read_rate=tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"),
|
||||
cache_creation_rate=tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token"),
|
||||
output_rate=tier_rate(tier, "output_cost_per_token") if tier_declares_output else flat_rates.output_rate,
|
||||
reasoning_rate=(
|
||||
tier_rate(tier, "output_cost_per_reasoning_token")
|
||||
if tier_declares_reasoning
|
||||
else None
|
||||
if tier_declares_output
|
||||
else flat_rates.reasoning_rate
|
||||
),
|
||||
)
|
||||
|
||||
return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost)
|
||||
|
||||
def _off_peak_rates(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates:
|
||||
input_rate, output_rate, cache_read_rate = apply_off_peak_pricing(
|
||||
model_info, current_time, rates.input_rate, rates.output_rate, rates.cache_read_rate
|
||||
)
|
||||
return replace(rates, input_rate=input_rate, output_rate=output_rate, cache_read_rate=cache_read_rate)
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashscope") -> tuple[float, float]:
|
||||
def _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]:
|
||||
prompt_cost: Final = (
|
||||
(breakdown.text_tokens * rates.input_rate)
|
||||
+ (breakdown.cached_tokens * rates.cache_read_rate)
|
||||
+ (breakdown.cache_creation_tokens * rates.cache_creation_rate)
|
||||
)
|
||||
completion_cost: Final = (breakdown.completion_tokens * rates.output_rate) + (
|
||||
breakdown.reasoning_tokens * rates.billed_reasoning_rate
|
||||
)
|
||||
return prompt_cost, completion_cost
|
||||
|
||||
|
||||
def cost_per_token(
|
||||
model: str,
|
||||
usage: Usage,
|
||||
custom_llm_provider: str = "dashscope",
|
||||
current_time: datetime | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculate cost per token for Dashscope models.
|
||||
|
||||
Supports both tiered and flat pricing with cached and reasoning tokens.
|
||||
Supports both tiered and flat pricing with cached and reasoning tokens, and swaps in the
|
||||
model's off_peak_pricing rates while one of its windows is open.
|
||||
|
||||
Args:
|
||||
model: Model name without provider prefix
|
||||
usage: LiteLLM Usage block
|
||||
custom_llm_provider: The provider id the request resolved to; dashscope or one of its brand aliases
|
||||
current_time: The moment the request is billed at; defaults to now, UTC
|
||||
|
||||
Returns:
|
||||
Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd)
|
||||
|
|
@ -133,8 +154,7 @@ def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashsco
|
|||
if tiered_pricing
|
||||
else None
|
||||
)
|
||||
standard_rates: Final = _flat_rates(model_info) if tier is None else _tier_rates(model_info, tier)
|
||||
rates: Final = _off_peak_rates(model_info, current_time, standard_rates)
|
||||
|
||||
prompt_cost: Final = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tier=tier)
|
||||
completion_cost: Final = _calculate_completion_cost(breakdown=breakdown, model_info=model_info, tier=tier)
|
||||
|
||||
return prompt_cost, completion_cost
|
||||
return _bill(breakdown, rates)
|
||||
|
|
|
|||
|
|
@ -61,6 +61,14 @@ def _get_effort_level(value: str | dict | None) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
GPT_REASONING_SERIES_MARKERS: Final = ("gpt-5", "gpt-6")
|
||||
|
||||
|
||||
def is_gpt_reasoning_series_name(model: str) -> bool:
|
||||
normalized: Final = model.split("/")[-1]
|
||||
return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) and not normalized.startswith("gpt-5-chat")
|
||||
|
||||
|
||||
class OpenAIGPT5Config(OpenAIGPTConfig):
|
||||
"""Configuration for gpt-5 models including GPT-5-Codex variants.
|
||||
|
||||
|
|
@ -73,21 +81,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
|
||||
@classmethod
|
||||
def is_model_gpt_5_model(cls, model: str) -> bool:
|
||||
# The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
|
||||
# …) are regular chat models: they support temperature and tool_choice but NOT
|
||||
# reasoning_effort. They must NOT be routed through the GPT-5 reasoning path.
|
||||
#
|
||||
# Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
|
||||
# models and must stay on the GPT-5 path. The distinguishing feature is that
|
||||
# the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
|
||||
# (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
|
||||
# number (i.e. "gpt-5.<digit>-chat").
|
||||
#
|
||||
# Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
|
||||
# than a substring check) makes this boundary explicit and avoids any ambiguity
|
||||
# if future model names coincidentally contain "gpt-5-chat" as an interior run.
|
||||
_normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "openai/"
|
||||
return "gpt-5" in model and not _normalized.startswith("gpt-5-chat")
|
||||
return is_gpt_reasoning_series_name(model)
|
||||
|
||||
@classmethod
|
||||
def is_model_gpt_5_search_model(cls, model: str) -> bool:
|
||||
|
|
@ -122,6 +116,8 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
def is_model_gpt_5_4_plus_model(cls, model: str) -> bool:
|
||||
"""Check if the model is gpt-5.4 or newer (5.4, 5.5, 5.6, etc., including pro)."""
|
||||
model_name: Final = model.split("/")[-1]
|
||||
if model_name.startswith("gpt-6"):
|
||||
return True
|
||||
if not model_name.startswith("gpt-5."):
|
||||
return False
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import time
|
|||
import uuid
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
|
|
@ -43,6 +44,14 @@ _OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(OpenAI)
|
|||
_AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(AzureOpenAI)
|
||||
|
||||
|
||||
_OPENAI_API_HOST: Final[str] = "api.openai.com"
|
||||
|
||||
|
||||
def is_openai_backed_api_base(api_base: str) -> bool:
|
||||
hostname: Final = urlsplit(api_base).hostname
|
||||
return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}"))
|
||||
|
||||
|
||||
class OpenAIError(BaseLLMException):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ import time
|
|||
import types
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -55,6 +54,7 @@ from .common_utils import (
|
|||
OpenAIError,
|
||||
build_output_token_limit_response,
|
||||
drop_params_from_unprocessable_entity_error,
|
||||
is_openai_backed_api_base,
|
||||
is_output_token_limit_error,
|
||||
)
|
||||
from .workload_identity import resolve_openai_workload_identity_config
|
||||
|
|
@ -1190,10 +1190,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
"""
|
||||
if stream_options is not None:
|
||||
return {"stream_options": stream_options}
|
||||
else:
|
||||
# by default litellm will include usage for openai endpoints
|
||||
if api_base is None or urlparse(api_base).hostname == "api.openai.com":
|
||||
return {"stream_options": {"include_usage": True}}
|
||||
if api_base is None or is_openai_backed_api_base(api_base):
|
||||
return {"stream_options": {"include_usage": True}}
|
||||
return {}
|
||||
|
||||
# Embedding
|
||||
|
|
|
|||
|
|
@ -33,8 +33,9 @@ import time
|
|||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
|
||||
|
||||
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
|
@ -42,6 +43,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
OpenAiResponsesToChatCompletionStreamIterator,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
|
|
@ -74,6 +76,7 @@ from litellm.types.llms.openai import (
|
|||
OutputTextDoneEvent,
|
||||
ResponseAPIUsage,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIOptionalRequestParams,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
ResponsesAPIStreamingResponse,
|
||||
|
|
@ -115,6 +118,199 @@ class ResponsesStreamChunk(TypedDict, total=False):
|
|||
content_index: ReadOnly[int]
|
||||
|
||||
|
||||
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"function_call_output": "output", "message": "content"}
|
||||
)
|
||||
|
||||
_EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {}
|
||||
|
||||
|
||||
def _item_rewrite_field(item: Mapping[str, object]) -> str | None:
|
||||
item_type: Final = item.get("type")
|
||||
if item_type is None:
|
||||
return "content" if "content" in item else None
|
||||
if not isinstance(item_type, str):
|
||||
return None
|
||||
return _PATCHABLE_ITEM_FIELDS.get(item_type)
|
||||
|
||||
|
||||
def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None:
|
||||
field: Final = _item_rewrite_field(item)
|
||||
if field is None or not isinstance(rewritten, Mapping):
|
||||
return None
|
||||
rewritten_content: Final = rewritten.get("content")
|
||||
if isinstance(item.get(field), str) and isinstance(rewritten_content, str):
|
||||
return {**item, field: rewritten_content} # mutable-ok: request input items must stay JSON-plain dicts
|
||||
rewritten_row: Final = cast("AllMessageValues", rewritten) # cast-ok: guardrails hand back chat-shaped rows
|
||||
converted_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(
|
||||
[rewritten_row] # mutable-ok: converter signature takes a list
|
||||
)
|
||||
if len(converted_items) != 1 or not isinstance(converted_items[0], Mapping):
|
||||
return None
|
||||
first_converted: Final = cast("Mapping[str, object]", converted_items[0]) # cast-ok: isinstance-checked above
|
||||
converted_value: Final = first_converted.get(field)
|
||||
if converted_value is None:
|
||||
return None
|
||||
return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts
|
||||
|
||||
|
||||
def _is_function_call_item(item: object) -> bool:
|
||||
return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call")
|
||||
|
||||
|
||||
def _last_message_role(messages: Sequence[object]) -> str | None:
|
||||
if not messages:
|
||||
return None
|
||||
last: Final = messages[-1]
|
||||
role: Final = last.get("role") if isinstance(last, Mapping) else getattr(last, "role", None)
|
||||
return role if isinstance(role, str) else None
|
||||
|
||||
|
||||
def _provenance_unit_bounds(
|
||||
raw_input: Sequence[object],
|
||||
solo_conversions: Sequence[Sequence[object]],
|
||||
) -> tuple[tuple[int, int], ...]:
|
||||
trailing_roles: Final = tuple(
|
||||
accumulate(
|
||||
(_last_message_role(messages) for messages in solo_conversions),
|
||||
lambda previous, current: current if current is not None else previous,
|
||||
)
|
||||
)
|
||||
start_indexes: Final = tuple(
|
||||
index
|
||||
for index in range(len(raw_input))
|
||||
if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
|
||||
)
|
||||
return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input))))
|
||||
|
||||
|
||||
def _input_item_provenance(
|
||||
raw_input: Sequence[object],
|
||||
expected_messages: Sequence[object],
|
||||
) -> tuple[Mapping[int, int], frozenset[int]] | None:
|
||||
if not all(isinstance(item, Mapping) for item in raw_input):
|
||||
return None
|
||||
solo_conversions: Final = tuple(
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=cast("ResponseInputParam", [item]), # cast-ok: items checked as Mappings above
|
||||
responses_api_request=_EMPTY_RESPONSES_REQUEST,
|
||||
)
|
||||
for item in raw_input
|
||||
)
|
||||
full_conversion: Final = tuple(
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=cast("ResponseInputParam", list(raw_input)), # cast-ok: items checked as Mappings above
|
||||
responses_api_request=_EMPTY_RESPONSES_REQUEST,
|
||||
)
|
||||
)
|
||||
if full_conversion != tuple(expected_messages):
|
||||
return None
|
||||
units: Final = _provenance_unit_bounds(raw_input, solo_conversions)
|
||||
unit_messages: Final = tuple(
|
||||
tuple(solo_conversions[start])
|
||||
if end - start == 1
|
||||
else tuple(
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=cast("ResponseInputParam", list(raw_input[start:end])), # cast-ok: checked as Mappings above
|
||||
responses_api_request=_EMPTY_RESPONSES_REQUEST,
|
||||
)
|
||||
)
|
||||
for start, end in units
|
||||
)
|
||||
if tuple(message for messages in unit_messages for message in messages) != full_conversion:
|
||||
return None
|
||||
boundaries: Final = tuple(accumulate((len(messages) for messages in unit_messages), initial=0))
|
||||
item_for_message: Final = MappingProxyType(
|
||||
{
|
||||
message_index: start
|
||||
for unit_index, (start, end) in enumerate(units)
|
||||
if end - start == 1
|
||||
for message_index in range(boundaries[unit_index], boundaries[unit_index + 1])
|
||||
}
|
||||
)
|
||||
tainted: Final = frozenset(
|
||||
message_index
|
||||
for unit_index, (start, end) in enumerate(units)
|
||||
if end - start > 1
|
||||
for message_index in range(boundaries[unit_index], boundaries[unit_index + 1])
|
||||
)
|
||||
return item_for_message, tainted
|
||||
|
||||
|
||||
class _RequestFields(NamedTuple):
|
||||
input: tuple[object, ...]
|
||||
instructions: str | None
|
||||
|
||||
|
||||
class _ExtractedInputs(NamedTuple):
|
||||
inputs: GenericGuardrailAPIInputs
|
||||
task_mappings: tuple[tuple[int, int | None], ...]
|
||||
|
||||
|
||||
def _patched_request_fields(
|
||||
raw_input: object,
|
||||
instructions: object,
|
||||
original_messages: Sequence[object],
|
||||
structured_messages: Sequence[object],
|
||||
) -> _RequestFields | None:
|
||||
if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages):
|
||||
return None
|
||||
offset: Final = 1 if instructions else 0
|
||||
provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:])
|
||||
if provenance is None:
|
||||
return None
|
||||
item_for_message, tainted = provenance
|
||||
changed: Final = tuple(
|
||||
(index, rewritten)
|
||||
for index, (original, rewritten) in enumerate(zip(original_messages, structured_messages))
|
||||
if original != rewritten
|
||||
)
|
||||
instruction_rewrites: Final = tuple(rewritten for index, rewritten in changed if index < offset)
|
||||
rewritten_instructions: Final = (
|
||||
instruction_rewrites[0].get("content")
|
||||
if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping)
|
||||
else instructions
|
||||
)
|
||||
instructions_value: Final = rewritten_instructions if isinstance(rewritten_instructions, str) else None
|
||||
if rewritten_instructions is not None and instructions_value is None:
|
||||
return None
|
||||
body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset)
|
||||
if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes):
|
||||
return None
|
||||
replacements: Final = MappingProxyType(
|
||||
{
|
||||
item_for_message[message_index]: _rewritten_input_item(
|
||||
cast("Mapping[str, object]", raw_input[item_for_message[message_index]]), # cast-ok: checked Mappings
|
||||
rewritten,
|
||||
)
|
||||
for message_index, rewritten in body_changes
|
||||
}
|
||||
)
|
||||
if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()):
|
||||
return None
|
||||
return _RequestFields(
|
||||
input=tuple(replacements.get(index, item) for index, item in enumerate(raw_input)),
|
||||
instructions=instructions_value,
|
||||
)
|
||||
|
||||
|
||||
def _patch_or_convert_request_fields(
|
||||
raw_input: object,
|
||||
instructions: object,
|
||||
original_messages: Sequence[object],
|
||||
structured_messages: Sequence[AllMessageValues],
|
||||
) -> _RequestFields | None:
|
||||
if not isinstance(structured_messages, list):
|
||||
return None
|
||||
patched: Final = _patched_request_fields(raw_input, instructions, original_messages, structured_messages)
|
||||
if patched is not None:
|
||||
return patched
|
||||
input_items, converted_instructions = (
|
||||
LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages)
|
||||
)
|
||||
return _RequestFields(input=tuple(input_items), instructions=converted_instructions)
|
||||
|
||||
|
||||
def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
|
||||
sequence_numbers: Final = (
|
||||
item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None)
|
||||
|
|
@ -162,9 +358,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
Handles both string input and list of message objects.
|
||||
"""
|
||||
input_data: Final[str | ResponseInputParam | None] = data.get("input")
|
||||
if input_data is None:
|
||||
if not isinstance(input_data, (str, list)):
|
||||
return data
|
||||
|
||||
structured_messages: Final = self.get_structured_messages(data)
|
||||
raw_tools: Final = data.get("tools")
|
||||
original_tools: Final[tuple[Mapping[str, object], ...]] = (
|
||||
|
|
@ -173,94 +368,93 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
flattened_tool_groups: Final = tuple(
|
||||
form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools)
|
||||
)
|
||||
flattened_tools: Final = tuple(
|
||||
cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list
|
||||
for group in flattened_tool_groups
|
||||
for tool in group
|
||||
)
|
||||
tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list
|
||||
copy.deepcopy(flattened_tools)
|
||||
)
|
||||
|
||||
# Handle simple string input
|
||||
if isinstance(input_data, str):
|
||||
inputs = GenericGuardrailAPIInputs(texts=[input_data])
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
if structured_messages:
|
||||
inputs["structured_messages"] = structured_messages
|
||||
# Include model information if available
|
||||
model = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
|
||||
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
guardrailed_texts = guardrailed_inputs.get("texts", [])
|
||||
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
|
||||
self._apply_guardrailed_tools_to_data(
|
||||
data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")
|
||||
)
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed string input")
|
||||
return data
|
||||
|
||||
# Handle list input (ResponseInputParam)
|
||||
if not isinstance(input_data, list):
|
||||
extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups)
|
||||
if not extracted.inputs.get("texts"):
|
||||
return data
|
||||
if structured_messages:
|
||||
extracted.inputs["structured_messages"] = structured_messages
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=extracted.inputs,
|
||||
request_data=data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
self._apply_guardrailed_tools_to_data(
|
||||
data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")
|
||||
)
|
||||
written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs)
|
||||
if written_back is not None:
|
||||
data["input"] = list(written_back.input) # mutable-ok: JSON body
|
||||
if written_back.instructions is None:
|
||||
data.pop("instructions", None)
|
||||
else:
|
||||
data["instructions"] = written_back.instructions # rebind-ok: data is an out-param
|
||||
elif isinstance(input_data, str):
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
|
||||
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param
|
||||
else:
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=input_data,
|
||||
responses=guardrailed_inputs.get("texts") or (),
|
||||
task_mappings=extracted.task_mappings,
|
||||
)
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input"))
|
||||
return data
|
||||
|
||||
def _extract_guardrail_inputs(
|
||||
self,
|
||||
data: Mapping[str, object],
|
||||
input_data: "str | ResponseInputParam",
|
||||
flattened_tool_groups: Sequence[Sequence[Mapping[str, object]]],
|
||||
) -> _ExtractedInputs:
|
||||
texts_to_check: Final[list[str]] = []
|
||||
images_to_check: Final[list[str]] = []
|
||||
task_mappings: Final[list[tuple[int, int | None]]] = []
|
||||
|
||||
# Step 1: Extract all text content, images, and tools
|
||||
for msg_idx, message in enumerate(input_data):
|
||||
self._extract_input_text_and_images(
|
||||
message=message,
|
||||
msg_idx=msg_idx,
|
||||
texts_to_check=texts_to_check,
|
||||
images_to_check=images_to_check,
|
||||
task_mappings=task_mappings,
|
||||
tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list
|
||||
copy.deepcopy(
|
||||
tuple(
|
||||
cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list
|
||||
for group in flattened_tool_groups
|
||||
for tool in group
|
||||
)
|
||||
)
|
||||
)
|
||||
if isinstance(input_data, str):
|
||||
texts_to_check.append(input_data)
|
||||
else:
|
||||
for msg_idx, message in enumerate(input_data):
|
||||
self._extract_input_text_and_images(
|
||||
message=message,
|
||||
msg_idx=msg_idx,
|
||||
texts_to_check=texts_to_check,
|
||||
images_to_check=images_to_check,
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
inputs["images"] = images_to_check
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
model: Final = data.get("model")
|
||||
if isinstance(model, str):
|
||||
inputs["model"] = model
|
||||
return _ExtractedInputs(inputs=inputs, task_mappings=tuple(task_mappings))
|
||||
|
||||
# Step 2: Apply guardrail to all texts in batch
|
||||
if texts_to_check:
|
||||
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if images_to_check:
|
||||
inputs["images"] = images_to_check
|
||||
if tools_to_check:
|
||||
inputs["tools"] = tools_to_check
|
||||
if structured_messages:
|
||||
inputs["structured_messages"] = structured_messages
|
||||
# Include model information if available
|
||||
model = data.get("model")
|
||||
if model:
|
||||
inputs["model"] = model
|
||||
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=data,
|
||||
input_type="request",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
guardrailed_texts = guardrailed_inputs.get("texts", [])
|
||||
self._apply_guardrailed_tools_to_data(
|
||||
data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")
|
||||
)
|
||||
|
||||
# Step 3: Map guardrail responses back to original input structure
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=input_data,
|
||||
responses=guardrailed_texts,
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", input_data)
|
||||
|
||||
return data
|
||||
@staticmethod
|
||||
def _written_back_request_fields(
|
||||
data: Mapping[str, object],
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
guardrailed_inputs: GenericGuardrailAPIInputs,
|
||||
) -> _RequestFields | None:
|
||||
guardrailed: Final = guardrailed_inputs.get("structured_messages")
|
||||
if guardrailed is None or guardrailed is structured_messages:
|
||||
return None
|
||||
return _patch_or_convert_request_fields(
|
||||
data.get("input"),
|
||||
data.get("instructions"),
|
||||
structured_messages or (),
|
||||
guardrailed,
|
||||
)
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""Extract tool names from Responses API request (tools[].name for function
|
||||
|
|
@ -331,8 +525,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
async def _apply_guardrail_responses_to_input(
|
||||
self,
|
||||
messages: Any, # Can be List[Dict[str, Any]] or ResponseInputParam
|
||||
responses: list[str],
|
||||
task_mappings: list[tuple[int, int | None]],
|
||||
responses: Sequence[str],
|
||||
task_mappings: Sequence[tuple[int, int | None]],
|
||||
) -> None:
|
||||
"""
|
||||
Apply guardrail responses back to input messages.
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
|
|||
)
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
|
||||
from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import *
|
||||
|
|
@ -88,7 +89,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
parts: Final = model.split("/")
|
||||
if len(parts) > 1 and parts[0] not in ("openai",):
|
||||
return False
|
||||
return "gpt-5" in model and "gpt-5-chat" not in model
|
||||
return is_gpt_reasoning_series_name(model)
|
||||
|
||||
@staticmethod
|
||||
def _supports_reasoning_effort_none(model: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from copy import deepcopy
|
|||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from litellm._logging import _redact_string
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -60,6 +61,7 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
from litellm.constants import (
|
||||
AZURE_OPENAI_AUDIO_PROVIDERS,
|
||||
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
|
||||
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
|
||||
)
|
||||
|
|
@ -984,6 +986,12 @@ def mock_completion(
|
|||
|
||||
|
||||
_OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
|
||||
_OPENAI_API_HOST: Final = "api.openai.com"
|
||||
|
||||
|
||||
def _is_openai_backed_api_base(api_base: str) -> bool:
|
||||
hostname: Final = urlsplit(api_base).hostname
|
||||
return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}"))
|
||||
|
||||
|
||||
def _resolve_openai_api_base(api_base: str | None) -> str:
|
||||
|
|
@ -1053,7 +1061,7 @@ def responses_api_bridge_check(
|
|||
# natively by Chat Completions with reasoning on, so custom-only requests stay on
|
||||
# chat and keep their native custom tool_call response shape.
|
||||
# - The UNSET-effort arm only fires against endpoints known to enforce that
|
||||
# constraint (the default OpenAI endpoint, or Azure OpenAI where api_base is
|
||||
# constraint (any api.openai.com host, or Azure OpenAI where api_base is
|
||||
# always set): chat-only OpenAI-compatible backends registered under the openai
|
||||
# provider with a custom api_base and gpt-5.4+ model names serve tools without
|
||||
# reasoning fine and have no /responses route, so they keep pre-existing
|
||||
|
|
@ -1068,14 +1076,15 @@ def responses_api_bridge_check(
|
|||
reasoning_active = reasoning_effort.get("effort") != "none" or reasoning_effort.get("summary") is not None
|
||||
else:
|
||||
reasoning_active = reasoning_effort != "none"
|
||||
# The reasoning+tools constraint is enforced only by the real OpenAI endpoint (and Azure OpenAI).
|
||||
# Resolve the effective base arg>global>env>default exactly as the chat handler does, so a custom
|
||||
# base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread as the default and
|
||||
# bridged to a /responses route it lacks. A whitespace-only base collapses to the default too.
|
||||
resolved_api_base: Final = _resolve_openai_api_base(api_base)
|
||||
on_constraint_enforcing_endpoint: Final = custom_llm_provider == "azure" or resolved_api_base.strip() in (
|
||||
"",
|
||||
_OPENAI_DEFAULT_API_BASE,
|
||||
# The reasoning+tools constraint is enforced by the real OpenAI backend behind any api.openai.com
|
||||
# host (the default URL or a PrivateLink hostname such as <region>.privatelink.api.openai.com) and
|
||||
# by Azure OpenAI. Resolve the effective base arg>global>env>default exactly as the chat handler
|
||||
# does, so a custom base set via litellm.api_base or OPENAI_BASE_URL/OPENAI_API_BASE isn't misread
|
||||
# as the default and bridged to a /responses route it lacks. A whitespace-only base collapses to
|
||||
# the default too.
|
||||
resolved_api_base: Final = _resolve_openai_api_base(api_base).strip()
|
||||
on_constraint_enforcing_endpoint: Final = (
|
||||
custom_llm_provider == "azure" or resolved_api_base == "" or _is_openai_backed_api_base(resolved_api_base)
|
||||
)
|
||||
if (
|
||||
custom_llm_provider in ("openai", "azure")
|
||||
|
|
@ -7769,7 +7778,7 @@ def transcription(
|
|||
provider=LlmProviders(custom_llm_provider),
|
||||
)
|
||||
|
||||
if custom_llm_provider == "azure" and provider_config is None:
|
||||
if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None:
|
||||
# azure configs
|
||||
api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
|
||||
|
||||
|
|
@ -8056,7 +8065,10 @@ def speech(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None
|
||||
if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers:
|
||||
if custom_llm_provider == "openai" or (
|
||||
custom_llm_provider in litellm.openai_compatible_providers
|
||||
and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS
|
||||
):
|
||||
if voice is None or not (isinstance(voice, str)):
|
||||
raise litellm.BadRequestError(
|
||||
message="'voice' is required to be passed as a string for OpenAI TTS",
|
||||
|
|
@ -8110,7 +8122,7 @@ def speech(
|
|||
aspeech=aspeech,
|
||||
shared_session=shared_session,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
|
||||
# Check if this is Azure Speech Service (Cognitive Services TTS)
|
||||
if model.startswith("speech/"):
|
||||
from litellm.llms.azure.text_to_speech.transformation import (
|
||||
|
|
|
|||
|
|
@ -29354,6 +29354,75 @@
|
|||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"gpt-6-astra": {
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
|
||||
"cache_creation_input_token_cost_flex": 6.25e-06,
|
||||
"cache_creation_input_token_cost_priority": 2.5e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 2e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_flex": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 4e-06,
|
||||
"cache_read_input_token_cost_flex": 5e-07,
|
||||
"cache_read_input_token_cost_priority": 2e-06,
|
||||
"input_cost_per_token": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-05,
|
||||
"input_cost_per_token_above_272k_tokens_flex": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4e-05,
|
||||
"input_cost_per_token_batches": 5e-06,
|
||||
"input_cost_per_token_flex": 5e-06,
|
||||
"input_cost_per_token_priority": 2e-05,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-05,
|
||||
"output_cost_per_token_above_272k_tokens_flex": 3.75e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 0.00015,
|
||||
"output_cost_per_token_batches": 2.5e-05,
|
||||
"output_cost_per_token_flex": 2.5e-05,
|
||||
"output_cost_per_token_priority": 0.0001,
|
||||
"regional_processing_uplift_multiplier_eu": 1.1,
|
||||
"regional_processing_uplift_multiplier_us": 1.1,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_native_streaming": true,
|
||||
"supports_none_reasoning_effort": false,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_cache_breakpoint": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"gpt-5.6": {
|
||||
"cache_creation_input_token_cost": 5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1e-05,
|
||||
|
|
@ -53485,6 +53554,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.95e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53519,6 +53593,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 4.4e-07,
|
||||
"output_cost_per_token": 1.32e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 1.98e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53581,6 +53660,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 4.4e-08,
|
||||
"output_cost_per_token": 1.32e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 1.98e-06,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53769,6 +53853,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.95e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53800,6 +53889,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 5.5e-07,
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 2.475e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
|
|||
|
|
@ -1216,7 +1216,7 @@ class GenerateKeyRequest(KeyRequestBase):
|
|||
organization_id: str | None = None
|
||||
project_id: str | None = None
|
||||
|
||||
@field_validator("team_id", "organization_id", mode="before")
|
||||
@field_validator("team_id", "organization_id", "project_id", mode="before")
|
||||
@classmethod
|
||||
def treat_cleared_id_as_unset(cls, v: object) -> object:
|
||||
if v == "":
|
||||
|
|
@ -4239,6 +4239,8 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase):
|
|||
access_group_id: str
|
||||
access_group_name: str
|
||||
models: tuple[str, ...]
|
||||
mcp_server_ids: tuple[str, ...] = ()
|
||||
agent_ids: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
||||
|
|
|
|||
|
|
@ -16,6 +16,10 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy._types import EnterpriseLicenseData
|
||||
|
||||
|
||||
AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router"
|
||||
HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit."
|
||||
|
||||
|
||||
class LicenseCheck:
|
||||
"""
|
||||
- Check if license in env
|
||||
|
|
@ -149,6 +153,19 @@ class LicenseCheck:
|
|||
return False
|
||||
return team_count > _max_teams_in_license
|
||||
|
||||
def heuristic_v2_router_limit(self) -> int | None:
|
||||
"""
|
||||
How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the
|
||||
signed license lists the auto_router feature, otherwise one. A license verified through
|
||||
the API carries no feature list, so it does not lift the limit either.
|
||||
"""
|
||||
if self.airgapped_license_data is None:
|
||||
return 1
|
||||
allowed_features: Final = self.airgapped_license_data.get("allowed_features")
|
||||
if isinstance(allowed_features, list) and AUTO_ROUTER_LICENSE_FEATURE in allowed_features:
|
||||
return None
|
||||
return 1
|
||||
|
||||
def verify_license_without_api_request(self, public_key, license_key):
|
||||
try:
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
|
|
@ -179,19 +196,21 @@ class LicenseCheck:
|
|||
# Decode and parse the data
|
||||
license_data: Final = json.loads(message.decode())
|
||||
|
||||
self.airgapped_license_data = EnterpriseLicenseData(**license_data)
|
||||
|
||||
# debug information provided in license data
|
||||
verbose_proxy_logger.debug("License data: %s", license_data)
|
||||
|
||||
# Check expiration date
|
||||
expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d")
|
||||
if expiration_date < datetime.now():
|
||||
self.airgapped_license_data = None
|
||||
return False, "License has expired"
|
||||
|
||||
self.airgapped_license_data = EnterpriseLicenseData(**license_data)
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.airgapped_license_data = None
|
||||
verbose_proxy_logger.debug(
|
||||
"litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - %s",
|
||||
e,
|
||||
|
|
|
|||
|
|
@ -916,9 +916,10 @@ class CompresrGuardrail(CustomGuardrail):
|
|||
def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None:
|
||||
"""Compressed content mirrored into the Responses `texts` channel.
|
||||
|
||||
The chat/Anthropic handlers round-trip ``structured_messages``; the
|
||||
Responses translation cannot rebuild its input from chat messages and
|
||||
instead writes back through ``texts``. This matches by value, so a
|
||||
The chat/Anthropic/Responses handlers round-trip
|
||||
``structured_messages``; translations without that round-trip write
|
||||
back through ``texts``, so the compressed content is mirrored there
|
||||
too. This matches by value, so a
|
||||
replacement is applied only when it is unambiguous: one compression per
|
||||
text, and every occurrence in ``texts`` accounted for by a compressed
|
||||
target. Anything else is left uncompressed rather than risk a wrong or
|
||||
|
|
|
|||
|
|
@ -50,6 +50,9 @@ if TYPE_CHECKING:
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
BYPASS_HEADER: Final = "x-headroom-bypass"
|
||||
_STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset(
|
||||
(CallTypes.completion, CallTypes.acompletion, CallTypes.responses, CallTypes.aresponses)
|
||||
)
|
||||
HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve"
|
||||
_HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})")
|
||||
_HASH_CACHE_TTL_SECONDS: Final = 15 * 60
|
||||
|
|
@ -725,6 +728,10 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER)
|
||||
return inputs
|
||||
|
||||
if request_data.get("background"):
|
||||
verbose_proxy_logger.debug("Headroom: background request; skipping compression")
|
||||
return inputs
|
||||
|
||||
structured_messages: Final = inputs.get("structured_messages")
|
||||
if not _is_object_list(structured_messages) or not structured_messages:
|
||||
return inputs
|
||||
|
|
@ -826,9 +833,9 @@ class HeadroomGuardrail(CustomGuardrail):
|
|||
) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
|
||||
base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type)
|
||||
effective: Final = base_result if base_result is not None else kwargs
|
||||
if call_type not in (CallTypes.completion, CallTypes.acompletion):
|
||||
if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES:
|
||||
return base_result
|
||||
if not effective.get("stream"):
|
||||
if not effective.get("stream") or effective.get("background"):
|
||||
return base_result
|
||||
if not has_headroom_retrieve_tool(effective.get("tools")):
|
||||
return base_result
|
||||
|
|
|
|||
|
|
@ -13,10 +13,11 @@ model/{model_id}/update - PATCH endpoint for model update.
|
|||
import asyncio
|
||||
import datetime
|
||||
import json
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from json import JSONDecodeError
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
|
||||
|
|
@ -49,6 +50,7 @@ from litellm.proxy._types import (
|
|||
TeamModelDeleteRequest,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import (
|
||||
coordination_redis_cache,
|
||||
|
|
@ -96,6 +98,9 @@ from litellm.router_strategy.complexity_router import (
|
|||
from litellm.router_utils.auto_router_model_naming import (
|
||||
STRATEGY_ROUTER_PARAM_FIELDS,
|
||||
carries_complexity_router_settings,
|
||||
count_heuristic_v2_routers,
|
||||
heuristic_v2_limit_violation,
|
||||
uses_heuristic_v2_classifier,
|
||||
validate_complexity_router_config_placement,
|
||||
validate_complexity_router_config_write,
|
||||
validate_strategy_router_model_write,
|
||||
|
|
@ -153,6 +158,8 @@ class _ProxyModelTable(Protocol):
|
|||
|
||||
def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
|
||||
|
||||
def create(self, *, data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ...
|
||||
|
||||
def update(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, object]
|
||||
) -> Awaitable[_ProxyModelRow | None]: ...
|
||||
|
|
@ -166,6 +173,9 @@ class _TxModelTables(Protocol):
|
|||
litellm_proxymodeltable: _ProxyModelTable
|
||||
|
||||
|
||||
_RowT = TypeVar("_RowT")
|
||||
|
||||
|
||||
class _ExistingModelRow(Protocol):
|
||||
@property
|
||||
def litellm_params(self) -> Mapping[str, object]: ...
|
||||
|
|
@ -269,6 +279,66 @@ def _raise_on_strategy_router_write_violation(
|
|||
)
|
||||
|
||||
|
||||
HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301
|
||||
_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
|
||||
_HEURISTIC_V2_DB_ROWS_SQL: Final = """
|
||||
SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable"
|
||||
WHERE model_id <> $1
|
||||
AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)
|
||||
-> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2'
|
||||
"""
|
||||
|
||||
|
||||
def _effective_complexity_router_config(
|
||||
incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
|
||||
) -> object:
|
||||
"""The complexity config a write leaves on the row: the incoming one when the write carries it, else the stored one."""
|
||||
incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config
|
||||
if incoming is not None or existing_params is None:
|
||||
return incoming
|
||||
return existing_params.complexity_router_config
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _heuristic_v2_slot(
|
||||
prisma_client: PrismaClient, *, effective_config: object, model_id: str | None
|
||||
) -> AsyncGenerator[_ProxyModelTable, None]:
|
||||
"""Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled.
|
||||
|
||||
A write that leaves the row on classifier_type heuristic_v2 under a limited license runs
|
||||
inside one transaction that takes an advisory lock in its own statement before counting
|
||||
(a statement's snapshot predates anything it locks), so pods cannot both pass the count:
|
||||
the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged
|
||||
against the license limit and the write is refused with a 403 before it happens. The row
|
||||
being edited keeps its own slot through ``model_id``. Every other write, and every write on
|
||||
an unlimited license, goes through the repository table with no lock. Only the row write
|
||||
itself may run inside: anything that needs a second connection (the team model bookkeeping)
|
||||
must wait until the transaction has committed and the lock is released. The transaction
|
||||
writes bypass the repository's publish-on-write, so the config change is published once
|
||||
after commit, the way delete_team_models does.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import _license_check, llm_router
|
||||
|
||||
limit: Final = _license_check.heuristic_v2_router_limit()
|
||||
if limit is None or not uses_heuristic_v2_classifier(effective_config):
|
||||
yield _proxy_model_table(prisma_client)
|
||||
return
|
||||
async with prisma_client.db.tx() as tx_ctx:
|
||||
tables: Final[_TxModelTables] = tx_ctx
|
||||
await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY)
|
||||
rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "")
|
||||
db_held: Final = rows[0].get("held") if rows else 0
|
||||
config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments())
|
||||
held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows)
|
||||
violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit)
|
||||
if violation is not None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}"
|
||||
)
|
||||
yield tables.litellm_proxymodeltable
|
||||
await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
|
||||
|
||||
|
||||
ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add"
|
||||
_REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm")
|
||||
|
||||
|
|
@ -720,22 +790,29 @@ async def patch_model(
|
|||
)
|
||||
|
||||
requested_model_name: Final = patch_data.model_name
|
||||
stored_model_name: str | None = None
|
||||
|
||||
async def write_row(update_data: PrismaCompatibleUpdateDBModel) -> _ProxyModelRow | None:
|
||||
nonlocal stored_model_name
|
||||
stored_model_name = update_data.get("model_name")
|
||||
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
update_data["updated_at"] = cast(str, get_utc_datetime())
|
||||
async with _heuristic_v2_slot(
|
||||
prisma_client,
|
||||
effective_config=_effective_complexity_router_config(
|
||||
patch_data.litellm_params, db_model.litellm_params
|
||||
),
|
||||
model_id=model_id,
|
||||
) as table:
|
||||
return await table.update(where={"model_id": model_id}, data=update_data)
|
||||
|
||||
# Handle team model updates with proper alias management
|
||||
update_data: Final = await _update_team_model_in_db(
|
||||
updated_model: Final = await _update_team_model_in_db(
|
||||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
# Add metadata about update
|
||||
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
update_data["updated_at"] = cast(str, get_utc_datetime())
|
||||
|
||||
# Perform partial update
|
||||
updated_model: Final = await _proxy_model_table(prisma_client).update(
|
||||
where={"model_id": model_id},
|
||||
data=update_data,
|
||||
write_row=write_row,
|
||||
)
|
||||
|
||||
if updated_model is None:
|
||||
|
|
@ -746,7 +823,6 @@ async def patch_model(
|
|||
param=None,
|
||||
)
|
||||
|
||||
stored_model_name: Final = update_data.get("model_name")
|
||||
if (
|
||||
stored_model_name is not None
|
||||
and stored_model_name == requested_model_name
|
||||
|
|
@ -980,7 +1056,8 @@ async def _add_model_to_db(
|
|||
prisma_client: PrismaClient,
|
||||
new_encryption_key: str | None = None,
|
||||
should_create_model_in_db: bool = True,
|
||||
) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
|
||||
slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None,
|
||||
) -> "_ProxyModelRow | LiteLLM_ProxyModelTable":
|
||||
# encrypt litellm params #
|
||||
_litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True)
|
||||
_original_litellm_model_name: Final = model_params.litellm_params.model
|
||||
|
|
@ -998,18 +1075,20 @@ async def _add_model_to_db(
|
|||
if model_params.model_info.id is not None:
|
||||
_data["model_id"] = model_params.model_info.id
|
||||
_create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above
|
||||
if should_create_model_in_db:
|
||||
model_response = await ModelRepository(prisma_client).table.create(data=_create_data)
|
||||
else:
|
||||
model_response = LiteLLM_ProxyModelTable(**_data)
|
||||
return model_response
|
||||
if not should_create_model_in_db:
|
||||
return LiteLLM_ProxyModelTable(**_data)
|
||||
if slot is None:
|
||||
return await _proxy_model_table(prisma_client).create(data=_create_data)
|
||||
async with slot as table:
|
||||
return await table.create(data=_create_data)
|
||||
|
||||
|
||||
async def _add_team_model_to_db(
|
||||
model_params: Deployment,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
|
||||
slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None,
|
||||
) -> "_ProxyModelRow | LiteLLM_ProxyModelTable":
|
||||
"""
|
||||
If 'team_id' is provided,
|
||||
|
||||
|
|
@ -1040,6 +1119,7 @@ async def _add_team_model_to_db(
|
|||
model_params=model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
slot=slot,
|
||||
)
|
||||
|
||||
if original_model_name:
|
||||
|
|
@ -1060,7 +1140,8 @@ async def _update_team_model_in_db(
|
|||
patch_data: updateDeployment,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
) -> PrismaCompatibleUpdateDBModel:
|
||||
write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]],
|
||||
) -> _RowT:
|
||||
"""
|
||||
Handle team model updates with proper alias management.
|
||||
|
||||
|
|
@ -1068,6 +1149,9 @@ async def _update_team_model_in_db(
|
|||
- Creates unique internal model_name and team alias
|
||||
- Adds model to team object
|
||||
- Preserves team_public_model_name for external reference
|
||||
|
||||
The row is written through ``write_row`` before the team's model list is touched, so a
|
||||
refused or failed write leaves the team as it was (the create path orders itself the same way).
|
||||
"""
|
||||
# Validate team_id if present in patch_data
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
|
@ -1079,9 +1163,7 @@ async def _update_team_model_in_db(
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
# Validated before any write, beside the premium check the create path already runs
|
||||
# here. The team ACL is updated below and autocommits, so a validator that raises
|
||||
# further down would leave the team mutated and the deployment row never written.
|
||||
# Validated before the row write, beside the premium check the create path already runs here.
|
||||
#
|
||||
# The merged view is what gets stored, so that is what has to satisfy the invariants.
|
||||
# Validating the patch alone rejected a partial edit of an already valid deployment:
|
||||
|
|
@ -1101,7 +1183,7 @@ async def _update_team_model_in_db(
|
|||
|
||||
# No team_id in patch, proceed with standard update
|
||||
if patch_team_id is None:
|
||||
return update_db_model(db_model=db_model, updated_patch=patch_data)
|
||||
return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data))
|
||||
|
||||
# Determine public model name
|
||||
public_model_name: Final = _get_public_model_name(
|
||||
|
|
@ -1120,11 +1202,14 @@ async def _update_team_model_in_db(
|
|||
db_team_id: Final = db_model.model_info.team_id if db_model.model_info else None
|
||||
is_new_team_assignment: Final = db_team_id != patch_team_id
|
||||
|
||||
# Team rows keep their internal UUID-based model_name; the public name lives in model_info
|
||||
patch_data.model_name = f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None
|
||||
row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=patch_data))
|
||||
|
||||
if is_new_team_assignment:
|
||||
await _setup_new_team_model_assignment(
|
||||
team_id=patch_team_id,
|
||||
public_model_name=public_model_name,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
else:
|
||||
|
|
@ -1132,12 +1217,11 @@ async def _update_team_model_in_db(
|
|||
team_id=patch_team_id,
|
||||
public_model_name=public_model_name,
|
||||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
return update_db_model(db_model=db_model, updated_patch=patch_data)
|
||||
return row
|
||||
|
||||
|
||||
def _get_public_model_name(
|
||||
|
|
@ -1189,13 +1273,9 @@ def _get_public_model_name(
|
|||
async def _setup_new_team_model_assignment(
|
||||
team_id: str,
|
||||
public_model_name: str,
|
||||
patch_data: updateDeployment,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""Set up a new team model with unique name and team membership."""
|
||||
unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}"
|
||||
patch_data.model_name = unique_model_name
|
||||
|
||||
"""Register a newly team-assigned model's public name on the team."""
|
||||
await team_model_add(
|
||||
data=TeamModelAddRequest(
|
||||
team_id=team_id,
|
||||
|
|
@ -1385,7 +1465,6 @@ async def _update_existing_team_model_assignment(
|
|||
team_id: str,
|
||||
public_model_name: str,
|
||||
db_model: Deployment,
|
||||
patch_data: updateDeployment,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> None:
|
||||
|
|
@ -1409,9 +1488,6 @@ async def _update_existing_team_model_assignment(
|
|||
old_public_name: Final = db_model.model_info.team_public_model_name if db_model.model_info else None
|
||||
|
||||
if old_public_name and public_model_name != old_public_name:
|
||||
# Clear user-supplied public name from patch before any early return so the
|
||||
# caller does not overwrite the internal UUID-based model_name in the DB.
|
||||
patch_data.model_name = None
|
||||
if prisma_client is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"prisma_client not initialized; skipping public name update entirely to avoid orphaned entries"
|
||||
|
|
@ -1459,10 +1535,6 @@ async def _update_existing_team_model_assignment(
|
|||
# else: old_public_name == public_model_name (no rename needed)
|
||||
# No team_model_add/delete calls required; public name is already registered
|
||||
|
||||
# Always clear patch_data.model_name to prevent caller from overwriting
|
||||
# the internal UUID-based model_name in the DB with the user-supplied public name
|
||||
patch_data.model_name = None
|
||||
|
||||
|
||||
class ModelManagementAuthChecks:
|
||||
"""
|
||||
|
|
@ -1878,18 +1950,19 @@ async def add_new_model(
|
|||
reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None)
|
||||
try:
|
||||
_original_litellm_model_name: Final = model_params.model_name
|
||||
if model_params.model_info.team_id is None:
|
||||
model_response = await _add_model_to_db(
|
||||
model_params=priced_model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
else:
|
||||
model_response = await _add_team_model_to_db(
|
||||
model_params=priced_model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
add_model: Final = (
|
||||
_add_model_to_db if model_params.model_info.team_id is None else _add_team_model_to_db
|
||||
)
|
||||
model_response = await add_model(
|
||||
model_params=priced_model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
slot=_heuristic_v2_slot(
|
||||
prisma_client,
|
||||
effective_config=priced_model_params.litellm_params.complexity_router_config,
|
||||
model_id=priced_model_params.model_info.id,
|
||||
),
|
||||
)
|
||||
reload_outcome = await proxy_config.add_deployment(
|
||||
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
|
||||
)
|
||||
|
|
@ -1903,6 +1976,8 @@ async def add_new_model(
|
|||
passed_model_info=priced_model_params.model_info,
|
||||
)
|
||||
except Exception as e:
|
||||
if isinstance(e, HTTPException):
|
||||
raise
|
||||
verbose_proxy_logger.exception("Exception in add_new_model: %s", e)
|
||||
|
||||
else:
|
||||
|
|
@ -2070,10 +2145,17 @@ async def update_model(
|
|||
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
|
||||
**({} if renamed_to is None else {"model_name": renamed_to}),
|
||||
}
|
||||
model_response: Final = await _proxy_model_table(prisma_client).update(
|
||||
where={"model_id": _model_id},
|
||||
data=_data,
|
||||
)
|
||||
async with _heuristic_v2_slot(
|
||||
prisma_client,
|
||||
effective_config=_effective_complexity_router_config(
|
||||
model_params.litellm_params, deployment.litellm_params
|
||||
),
|
||||
model_id=_model_id,
|
||||
) as table:
|
||||
model_response: Final = await table.update(
|
||||
where={"model_id": _model_id},
|
||||
data=_data,
|
||||
)
|
||||
if renamed_to is not None:
|
||||
await sync_access_groups_for_renamed_model(
|
||||
prisma_client=prisma_client,
|
||||
|
|
|
|||
|
|
@ -4318,6 +4318,8 @@ async def _resolve_team_access_group_resources(
|
|||
access_group_id=group.access_group_id,
|
||||
access_group_name=group.access_group_name,
|
||||
models=tuple(group.access_model_names or ()),
|
||||
mcp_server_ids=tuple(group.access_mcp_server_ids or ()),
|
||||
agent_ids=tuple(group.access_agent_ids or ()),
|
||||
)
|
||||
for group in resolved_groups
|
||||
),
|
||||
|
|
|
|||
|
|
@ -1730,6 +1730,16 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict:
|
|||
return headers
|
||||
|
||||
|
||||
def _is_vertex_anthropic_count_tokens_route(endpoint: str) -> bool:
|
||||
return endpoint.rsplit("/", 1)[-1].split(":", 1)[0] == "count-tokens"
|
||||
|
||||
|
||||
def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]:
|
||||
if not _is_vertex_anthropic_count_tokens_route(endpoint):
|
||||
return headers
|
||||
return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"})
|
||||
|
||||
|
||||
def get_vertex_pass_through_handler(
|
||||
call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here
|
||||
) -> BaseVertexAIPassThroughHandler:
|
||||
|
|
@ -2128,7 +2138,7 @@ async def _base_vertex_proxy_route(
|
|||
endpoint_func: Final = create_pass_through_route(
|
||||
endpoint=endpoint,
|
||||
target=target,
|
||||
custom_headers=headers,
|
||||
custom_headers=_upstream_headers_for_vertex_route(endpoint, headers),
|
||||
is_streaming_request=is_streaming_request,
|
||||
) # dynamically construct pass-through endpoint based on incoming path
|
||||
|
||||
|
|
|
|||
|
|
@ -120,6 +120,8 @@ from litellm.router_utils.add_retry_fallback_headers import (
|
|||
from litellm.router_utils.auto_router_model_naming import (
|
||||
STRATEGY_ROUTER_PARAM_FIELDS,
|
||||
carries_complexity_router_settings,
|
||||
count_heuristic_v2_routers,
|
||||
heuristic_v2_limit_violation,
|
||||
validate_complexity_router_config_placement,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -301,7 +303,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
)
|
||||
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.litellm_license import LicenseCheck
|
||||
from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck
|
||||
from litellm.proxy.auth.model_checks import (
|
||||
expand_wildcard_deployments_for_model_info,
|
||||
get_all_fallbacks,
|
||||
|
|
@ -4316,6 +4318,19 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object])
|
|||
raise ValueError(f"model {model.get('model_name', '')!r}: {violation}")
|
||||
|
||||
|
||||
def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None:
|
||||
"""
|
||||
Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows.
|
||||
|
||||
Checked here rather than left to router registration for the same reason as the two
|
||||
validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so
|
||||
the router's own refusal would turn the extra router into a silently missing model.
|
||||
"""
|
||||
violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit)
|
||||
if violation is not None:
|
||||
raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}")
|
||||
|
||||
|
||||
def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place
|
||||
"""
|
||||
Stamps `model_info.id` from the raw litellm_params before plugin resolution swaps
|
||||
|
|
@ -5721,6 +5736,7 @@ class ProxyConfig:
|
|||
model_list: Final = config.get("model_list", None)
|
||||
if model_list:
|
||||
router_params["model_list"] = model_list
|
||||
validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit())
|
||||
print( # noqa: T201
|
||||
"\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m"
|
||||
)
|
||||
|
|
@ -5810,6 +5826,7 @@ class ProxyConfig:
|
|||
),
|
||||
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit,
|
||||
)
|
||||
|
||||
if redis_usage_cache is not None and router.cache.redis_cache is None:
|
||||
|
|
@ -6270,6 +6287,7 @@ class ProxyConfig:
|
|||
search_tools=search_tools,
|
||||
ignore_invalid_deployments=True,
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit,
|
||||
)
|
||||
verbose_proxy_logger.debug("updated llm_router: %s", llm_router)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -173,6 +173,9 @@ class _SessionSpendRow(TypedDict):
|
|||
session_cache_hit_count: ReadOnly[int]
|
||||
session_llm_count: ReadOnly[int]
|
||||
session_agent_count: ReadOnly[int]
|
||||
session_total_prompt_tokens: ReadOnly[int]
|
||||
session_total_completion_tokens: ReadOnly[int]
|
||||
session_total_tokens: ReadOnly[int]
|
||||
session_models: ReadOnly[Sequence[str]]
|
||||
|
||||
|
||||
|
|
@ -188,6 +191,9 @@ class _SessionSpendStats(NamedTuple):
|
|||
session_cache_hit_count: int
|
||||
session_llm_count: int
|
||||
session_agent_count: int
|
||||
session_total_prompt_tokens: int
|
||||
session_total_completion_tokens: int
|
||||
session_total_tokens: int
|
||||
session_models: Sequence[str]
|
||||
session_models_truncated: bool
|
||||
|
||||
|
|
@ -4287,8 +4293,8 @@ async def _build_ui_spend_logs_response(
|
|||
Build the paginated response for the UI spend-logs endpoint.
|
||||
|
||||
When ``enrich_session_counts`` is ``True`` (the default for the v1/UI
|
||||
endpoint), each row is enriched with ``session_total_count`` plus spend
|
||||
and call-type aggregates so the frontend knows which sessions are
|
||||
endpoint), each row is enriched with ``session_total_count`` plus spend,
|
||||
token and call-type aggregates so the frontend knows which sessions are
|
||||
expandable (multi-call sessions). One ``GROUP BY (session_id, api_key)``
|
||||
query serves every referenced session, keyed per api key so two callers
|
||||
reusing a session id never see each other's totals. Rows without a
|
||||
|
|
@ -4356,7 +4362,10 @@ async def _build_ui_spend_logs_response(
|
|||
COUNT(*) FILTER (
|
||||
WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
|
||||
)::int AS session_llm_count,
|
||||
COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count
|
||||
COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count,
|
||||
COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens,
|
||||
COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens,
|
||||
COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens
|
||||
FROM "LiteLLM_SpendLogs"
|
||||
WHERE session_id = ANY($1::text[])
|
||||
AND api_key = ANY($2::text[])
|
||||
|
|
@ -4389,6 +4398,9 @@ async def _build_ui_spend_logs_response(
|
|||
session_cache_hit_count=int(row.get("session_cache_hit_count") or 0),
|
||||
session_llm_count=int(row.get("session_llm_count") or 0),
|
||||
session_agent_count=int(row.get("session_agent_count") or 0),
|
||||
session_total_prompt_tokens=int(row.get("session_total_prompt_tokens") or 0),
|
||||
session_total_completion_tokens=int(row.get("session_total_completion_tokens") or 0),
|
||||
session_total_tokens=int(row.get("session_total_tokens") or 0),
|
||||
session_models=models[:_SESSION_MODELS_LIMIT],
|
||||
session_models_truncated=len(models) > _SESSION_MODELS_LIMIT,
|
||||
)
|
||||
|
|
@ -4418,6 +4430,9 @@ async def _build_ui_spend_logs_response(
|
|||
row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count
|
||||
row_dict["session_llm_count"] = session_stats.session_llm_count
|
||||
row_dict["session_agent_count"] = session_stats.session_agent_count
|
||||
row_dict["session_total_prompt_tokens"] = session_stats.session_total_prompt_tokens
|
||||
row_dict["session_total_completion_tokens"] = session_stats.session_total_completion_tokens
|
||||
row_dict["session_total_tokens"] = session_stats.session_total_tokens
|
||||
row_dict["session_models"] = session_stats.session_models
|
||||
row_dict["session_models_truncated"] = session_stats.session_models_truncated
|
||||
enriched.append(row_dict)
|
||||
|
|
|
|||
|
|
@ -7531,6 +7531,9 @@ def create_model_info_response(
|
|||
max_input_tokens = configured_input
|
||||
if configured_output is not None:
|
||||
max_output_tokens = configured_output
|
||||
configured_mode: Final = llm_router.get_configured_mode(model_id)
|
||||
if isinstance(configured_mode, str):
|
||||
base["mode"] = configured_mode
|
||||
|
||||
if max_input_tokens is not None:
|
||||
base["max_input_tokens"] = max_input_tokens
|
||||
|
|
|
|||
|
|
@ -30,7 +30,10 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
|
||||
from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store
|
||||
from litellm.proxy.vector_store_endpoints.utils import (
|
||||
can_user_access_vector_store,
|
||||
filter_listable_vector_stores,
|
||||
)
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import ManagedVectorStoresRepository
|
||||
from litellm.types.vector_stores import (
|
||||
|
|
@ -390,11 +393,10 @@ async def list_vector_stores(
|
|||
|
||||
# Filter vector stores based on access control
|
||||
accessible_vector_stores: Final = []
|
||||
for vs in vector_store_map.values():
|
||||
if await _check_vector_store_access(vs, user_api_key_dict):
|
||||
redacted = LiteLLM_ManagedVectorStore(**vs)
|
||||
redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params"))
|
||||
accessible_vector_stores.append(redacted)
|
||||
for vs in await filter_listable_vector_stores(vector_store_map.values(), user_api_key_dict):
|
||||
redacted = LiteLLM_ManagedVectorStore(**vs)
|
||||
redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params"))
|
||||
accessible_vector_stores.append(redacted)
|
||||
|
||||
total_count: Final = len(accessible_vector_stores)
|
||||
total_pages: Final = (total_count + page_size - 1) // page_size
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
|
||||
is_ui_session_credential,
|
||||
resolve_ui_session_team_ids,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LitellmUserRoles,
|
||||
|
|
@ -160,10 +166,16 @@ async def can_user_access_vector_store(
|
|||
if _is_proxy_admin(user_api_key_dict):
|
||||
return True
|
||||
|
||||
vector_store_team_id: Final = vector_store.get("team_id")
|
||||
if vector_store_team_id is None:
|
||||
if vector_store.get("team_id") is None:
|
||||
return True
|
||||
|
||||
return await _is_vector_store_granted(vector_store, user_api_key_dict)
|
||||
|
||||
|
||||
async def _is_vector_store_granted(
|
||||
vector_store: LiteLLM_ManagedVectorStore,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> bool:
|
||||
vector_store_id: Final = vector_store.get("vector_store_id") or ""
|
||||
|
||||
key_object_permission = user_api_key_dict.object_permission
|
||||
|
|
@ -178,12 +190,70 @@ async def can_user_access_vector_store(
|
|||
if _object_permission_allows_vector_store(team_object_permission, vector_store_id):
|
||||
return True
|
||||
|
||||
if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store_team_id:
|
||||
return True
|
||||
return user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store.get("team_id")
|
||||
|
||||
|
||||
async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth:
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
team: Final = await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
return user_api_key_dict.model_copy(
|
||||
update=MappingProxyType(
|
||||
{
|
||||
"team_id": team_id,
|
||||
"team_object_permission": team.object_permission,
|
||||
"team_object_permission_id": team.object_permission_id,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _vector_store_listing_auth_contexts(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> tuple[UserAPIKeyAuth, ...]:
|
||||
if not is_ui_session_credential(user_api_key_dict):
|
||||
return (user_api_key_dict,)
|
||||
session_key_context: Final = user_api_key_dict.model_copy(
|
||||
update=MappingProxyType({"team_id": None, "team_object_permission": None, "team_object_permission_id": None})
|
||||
)
|
||||
team_ids: Final = await resolve_ui_session_team_ids(user_api_key_dict)
|
||||
team_contexts: Final = tuple([await _team_auth_context(team_id, user_api_key_dict) for team_id in team_ids])
|
||||
return (session_key_context, *team_contexts)
|
||||
|
||||
|
||||
async def _is_vector_store_granted_to_any(
|
||||
vector_store: LiteLLM_ManagedVectorStore,
|
||||
auth_contexts: tuple[UserAPIKeyAuth, ...],
|
||||
) -> bool:
|
||||
for auth_context in auth_contexts:
|
||||
if await _is_vector_store_granted(vector_store, auth_context):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def filter_listable_vector_stores(
|
||||
vector_stores: Iterable[LiteLLM_ManagedVectorStore],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> tuple[LiteLLM_ManagedVectorStore, ...]:
|
||||
"""Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants."""
|
||||
if _is_proxy_admin(user_api_key_dict):
|
||||
return tuple(vector_stores)
|
||||
|
||||
auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict)
|
||||
return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)])
|
||||
|
||||
|
||||
async def get_litellm_managed_vector_store(
|
||||
vector_store_id: str,
|
||||
) -> LiteLLM_ManagedVectorStore | None:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from typing import Any, Final, Literal, cast
|
|||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
AZURE_OPENAI_AUDIO_PROVIDERS,
|
||||
REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS,
|
||||
REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
|
||||
request_timeout,
|
||||
|
|
@ -400,7 +401,7 @@ async def _arealtime(
|
|||
litellm_metadata=_build_litellm_metadata(kwargs),
|
||||
query_params=query_params,
|
||||
)
|
||||
elif _custom_llm_provider == "azure":
|
||||
elif _custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
|
||||
api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
|
||||
# set API KEY
|
||||
api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_API_KEY")
|
||||
|
|
|
|||
|
|
@ -562,6 +562,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
hidden_params: Final = getattr(chunk, "_hidden_params", None)
|
||||
if hidden_params is not None:
|
||||
chunk_dict["_hidden_params"] = dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params
|
||||
if (
|
||||
chunk_dict.get("usage") is None
|
||||
and isinstance(hidden_params, dict)
|
||||
and hidden_params.get("usage") is not None
|
||||
):
|
||||
chunk_dict["usage"] = hidden_params["usage"]
|
||||
return chunk_dict
|
||||
|
||||
def create_reasoning_summary_text_done_event(
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ import time
|
|||
import traceback
|
||||
import weakref
|
||||
from collections import defaultdict
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence
|
||||
from functools import lru_cache, partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
|
||||
|
|
@ -117,6 +117,9 @@ from litellm.router_utils.add_retry_fallback_headers import (
|
|||
from litellm.router_utils.auto_router_model_naming import (
|
||||
AUTO_ROUTER_MODEL_PREFIX,
|
||||
classify_strategy_router_model,
|
||||
count_heuristic_v2_routers,
|
||||
heuristic_v2_limit_violation,
|
||||
uses_heuristic_v2_classifier,
|
||||
)
|
||||
from litellm.router_utils.batch_utils import (
|
||||
_get_router_metadata_variable_name,
|
||||
|
|
@ -211,6 +214,7 @@ from litellm.types.router import (
|
|||
DeploymentTypedDict,
|
||||
FallbackAccessCheck,
|
||||
GuardrailTypedDict,
|
||||
HeuristicV2RouterLimit,
|
||||
LiteLLM_Params,
|
||||
MockRouterTestingParams,
|
||||
ModelGroupInfo,
|
||||
|
|
@ -683,6 +687,7 @@ class Router:
|
|||
background_health_check_model_groups: Sequence[str] | None = None,
|
||||
enable_weighted_failover: bool = False,
|
||||
fallback_access_check: FallbackAccessCheck | None = None,
|
||||
heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
|
||||
|
|
@ -759,6 +764,7 @@ class Router:
|
|||
|
||||
self.set_verbose = set_verbose
|
||||
self.ignore_invalid_deployments = ignore_invalid_deployments
|
||||
self.heuristic_v2_router_limit = heuristic_v2_router_limit
|
||||
self.fallback_access_check: Final = fallback_access_check
|
||||
self.debug_level = debug_level
|
||||
self.enable_pre_call_checks = enable_pre_call_checks
|
||||
|
|
@ -8796,6 +8802,30 @@ class Router:
|
|||
"""
|
||||
return classify_strategy_router_model(litellm_params.model) == "complexity"
|
||||
|
||||
def config_deployments(self) -> Iterator[Mapping[str, object]]:
|
||||
"""The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset)."""
|
||||
for deployment in self.model_list:
|
||||
if not isinstance(deployment, Mapping):
|
||||
continue
|
||||
model_info = deployment.get("model_info")
|
||||
if not (isinstance(model_info, Mapping) and model_info.get("db_model")):
|
||||
yield deployment
|
||||
|
||||
def heuristic_v2_router_limit_violation(self) -> str | None:
|
||||
"""
|
||||
Why one more heuristic_v2 router cannot join this router, or None when it can.
|
||||
|
||||
Judged against every deployment currently on the model_list; an upsert pops the row being
|
||||
edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is
|
||||
resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which
|
||||
is the SDK default, and the proxy injects a resolver backed by its license.
|
||||
"""
|
||||
limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None
|
||||
others: Final = count_heuristic_v2_routers(
|
||||
deployment for deployment in self.model_list if isinstance(deployment, Mapping)
|
||||
)
|
||||
return heuristic_v2_limit_violation(held=others + 1, limit=limit)
|
||||
|
||||
def init_complexity_router_deployment(self, deployment: Deployment):
|
||||
"""
|
||||
Initialize the complexity-router deployment.
|
||||
|
|
@ -8813,6 +8843,10 @@ class Router:
|
|||
)
|
||||
|
||||
complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config
|
||||
if uses_heuristic_v2_classifier(complexity_router_config):
|
||||
limit_violation: Final = self.heuristic_v2_router_limit_violation()
|
||||
if limit_violation is not None:
|
||||
raise ValueError(limit_violation)
|
||||
|
||||
default_model: str | None = deployment.litellm_params.complexity_router_default_model
|
||||
|
||||
|
|
@ -9636,8 +9670,16 @@ class Router:
|
|||
raise e
|
||||
|
||||
def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None:
|
||||
"""Put a deployment back the way it was before a failed upsert popped it.
|
||||
|
||||
A rollback re-admits state that was already serving, so it does not go through the
|
||||
heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first
|
||||
registered, judging the rollback would drop a serving router over an unrelated failed edit.
|
||||
"""
|
||||
if previous_deployment is None or self.has_model_id(model_id):
|
||||
return
|
||||
limit_resolver: Final = self.heuristic_v2_router_limit
|
||||
self.heuristic_v2_router_limit = None
|
||||
try:
|
||||
self.add_deployment(deployment=previous_deployment)
|
||||
verbose_router_logger.info(
|
||||
|
|
@ -9652,6 +9694,8 @@ class Router:
|
|||
model_id,
|
||||
restore_error,
|
||||
)
|
||||
finally:
|
||||
self.heuristic_v2_router_limit = limit_resolver
|
||||
|
||||
@staticmethod
|
||||
def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]:
|
||||
|
|
@ -9981,6 +10025,17 @@ class Router:
|
|||
coerce_token_limit(model_info.get("max_output_tokens")),
|
||||
)
|
||||
|
||||
def get_configured_mode(self, model_name: str) -> "str | None":
|
||||
"""Return the mode explicitly configured for a concrete deployment."""
|
||||
deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name)
|
||||
if deployment is None:
|
||||
return None
|
||||
|
||||
mode: Final = deployment.model_info.get("mode")
|
||||
if isinstance(mode, str) and mode.strip():
|
||||
return mode
|
||||
return None
|
||||
|
||||
def get_configured_display_name(self, model_name: str) -> "str | None":
|
||||
"""
|
||||
Return the display_name explicitly configured in a concrete deployment's
|
||||
|
|
|
|||
|
|
@ -2,23 +2,41 @@
|
|||
Auto-Routing Strategy that works with a Semantic Router Config
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import (
|
||||
effective_turn_off_message_logging,
|
||||
forwarded_internal_call_metadata,
|
||||
parent_session_kwargs,
|
||||
)
|
||||
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from semantic_router.routers import SemanticRouter
|
||||
from semantic_router.routers.base import Route
|
||||
|
||||
from litellm.router import Router
|
||||
from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
else:
|
||||
Router = Any
|
||||
PreRoutingHookResponse = Any
|
||||
Route = Any
|
||||
SemanticRouter = Any
|
||||
LiteLLMRouterEncoder = Any
|
||||
|
||||
|
||||
class _CallerMetadata(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
metadata: Mapping[str, object] | None = None
|
||||
litellm_metadata: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
class AutoRouter(CustomLogger):
|
||||
|
|
@ -50,6 +68,8 @@ class AutoRouter(CustomLogger):
|
|||
"""
|
||||
from semantic_router.routers import SemanticRouter
|
||||
|
||||
from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder
|
||||
|
||||
self.auto_router_config_path: str | None = auto_router_config_path
|
||||
self.auto_router_config: str | None = auto_router_config
|
||||
self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE
|
||||
|
|
@ -59,6 +79,11 @@ class AutoRouter(CustomLogger):
|
|||
self.embedding_model: str = embedding_model
|
||||
self.max_input_chars: int = max_input_chars
|
||||
self.litellm_router_instance: Router = litellm_router_instance
|
||||
self.encoder: LiteLLMRouterEncoder = LiteLLMRouterEncoder(
|
||||
litellm_router_instance=litellm_router_instance,
|
||||
model_name=embedding_model,
|
||||
max_input_chars=max_input_chars,
|
||||
)
|
||||
|
||||
def _load_semantic_routing_routes(self) -> list[Route]:
|
||||
from semantic_router.routers import SemanticRouter
|
||||
|
|
@ -129,9 +154,6 @@ class AutoRouter(CustomLogger):
|
|||
from semantic_router.routers import SemanticRouter
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
|
||||
from litellm.router_strategy.auto_router.litellm_encoder import (
|
||||
LiteLLMRouterEncoder,
|
||||
)
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
resolved_messages: Final = (
|
||||
|
|
@ -149,34 +171,47 @@ class AutoRouter(CustomLogger):
|
|||
#######################
|
||||
routelayer = SemanticRouter(
|
||||
routes=self.loaded_routes,
|
||||
encoder=LiteLLMRouterEncoder(
|
||||
litellm_router_instance=self.litellm_router_instance,
|
||||
model_name=self.embedding_model,
|
||||
max_input_chars=self.max_input_chars,
|
||||
),
|
||||
encoder=self.encoder,
|
||||
auto_sync=self.auto_sync_value,
|
||||
)
|
||||
self.routelayer = routelayer
|
||||
|
||||
message_content: Final = self._extract_text_from_messages(resolved_messages)
|
||||
route_name: Final = self._matched_route_name(routelayer, message_content)
|
||||
route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs)
|
||||
|
||||
return PreRoutingHookResponse(
|
||||
model=route_name or self.default_model,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
def _matched_route_name(self, routelayer: "SemanticRouter", text: str) -> str | None:
|
||||
async def _matched_route_name(
|
||||
self, routelayer: "SemanticRouter", text: str, request_kwargs: Mapping[str, object]
|
||||
) -> str | None:
|
||||
"""Name of the route `text` matches, or None when nothing matched or the match failed.
|
||||
|
||||
The route layer embeds `text` to compare it against the routes, and that embedding call can
|
||||
`text` is embedded here rather than by `routelayer(text=...)` so the caller's metadata reaches
|
||||
`aembedding()` and the embedding's spend lands on the key/team that sent the request;
|
||||
SemanticRouter has no way to pass kwargs through to its encoder. That embedding call can
|
||||
fail (context limit, timeout, provider error). Choosing a model is a routing decision, so a
|
||||
failure here falls back to the default model rather than failing the user's request.
|
||||
"""
|
||||
from semantic_router.schema import RouteChoice
|
||||
|
||||
try:
|
||||
route_choice: Final = routelayer(text=text)
|
||||
caller: Final = _CallerMetadata.model_validate(request_kwargs)
|
||||
query_vector: Final = (
|
||||
await self.encoder.aencode_queries(
|
||||
[text],
|
||||
metadata=forwarded_internal_call_metadata(caller.metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN),
|
||||
litellm_metadata=forwarded_internal_call_metadata(
|
||||
caller.litellm_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN
|
||||
),
|
||||
proxy_server_request={"body": {"model": self.embedding_model, "input": [text]}},
|
||||
turn_off_message_logging=effective_turn_off_message_logging(request_kwargs),
|
||||
**parent_session_kwargs(request_kwargs),
|
||||
)
|
||||
)[0]
|
||||
route_choice: Final = await routelayer.acall(vector=query_vector)
|
||||
except Exception as e: # noqa: BLE001 -- the embedding call behind the route layer can fail many ways (context limit, timeout, provider/network error); none of them may fail the request
|
||||
verbose_router_logger.warning(
|
||||
"AutoRouter: semantic routing failed (%s), falling back to default model %s", e, self.default_model
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under
|
|||
``ignore_invalid_deployments``.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
|
@ -163,6 +163,38 @@ def strategy_router_dependencies(
|
|||
)
|
||||
|
||||
|
||||
def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool:
|
||||
"""Whether this complexity config classifies with the bundled heuristic_v2 model."""
|
||||
return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2"
|
||||
|
||||
|
||||
def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool:
|
||||
"""Whether this deployment is a complexity router that classifies with heuristic_v2."""
|
||||
return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and (
|
||||
uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config"))
|
||||
)
|
||||
|
||||
|
||||
def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int:
|
||||
"""How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers."""
|
||||
return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params"))))
|
||||
|
||||
|
||||
def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None:
|
||||
"""Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits.
|
||||
|
||||
``limit`` None means unlimited. The message is shared by every enforcement point (config
|
||||
load, model writes, router registration) and stays SDK-neutral: it names the cap and what
|
||||
the caller can change; the proxy appends how its license lifts the cap.
|
||||
"""
|
||||
if limit is None or held <= limit:
|
||||
return None
|
||||
return (
|
||||
f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make "
|
||||
f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router."
|
||||
)
|
||||
|
||||
|
||||
def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None:
|
||||
"""Reject a complexity config the router would refuse to build a deployment from.
|
||||
|
||||
|
|
|
|||
|
|
@ -82,6 +82,7 @@ class DeploymentAffinityCheck(CustomLogger):
|
|||
"""
|
||||
|
||||
CACHE_KEY_PREFIX = "deployment_affinity:v1"
|
||||
USER_ID_AFFINITY_PREFIX: Final = "user_id:"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -253,15 +254,6 @@ class DeploymentAffinityCheck(CustomLogger):
|
|||
hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped"
|
||||
return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}"
|
||||
|
||||
@staticmethod
|
||||
def _get_user_key_from_metadata_dict(metadata: dict) -> str | None:
|
||||
# NOTE: affinity is keyed on the *API key hash* provided by the proxy (not the
|
||||
# OpenAI `user` parameter, which is an end-user identifier).
|
||||
user_key: Final = metadata.get("user_api_key_hash")
|
||||
if user_key is None:
|
||||
return None
|
||||
return str(user_key)
|
||||
|
||||
@staticmethod
|
||||
def _get_session_id_from_metadata_dict(metadata: dict) -> str | None:
|
||||
session_id: Final = metadata.get("session_id")
|
||||
|
|
@ -285,22 +277,30 @@ class DeploymentAffinityCheck(CustomLogger):
|
|||
return metadata_dicts
|
||||
|
||||
@staticmethod
|
||||
def _get_user_key_from_request_kwargs(request_kwargs: dict) -> str | None:
|
||||
def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None:
|
||||
value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None)
|
||||
return None if value is None else str(value)
|
||||
|
||||
@classmethod
|
||||
def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None:
|
||||
"""
|
||||
Extract a stable affinity key from request kwargs.
|
||||
|
||||
Source (proxy): `metadata.user_api_key_hash`
|
||||
Source (proxy): `metadata.user_api_key_hash` for virtual-key callers. JWT-authenticated
|
||||
callers carry no key hash, so their `metadata.user_api_key_user_id` stands in for it,
|
||||
namespaced under `USER_ID_AFFINITY_PREFIX` so a user id can never alias a key hash.
|
||||
|
||||
Note: the OpenAI `user` parameter is an end-user identifier and is intentionally
|
||||
not used for deployment affinity.
|
||||
"""
|
||||
# Check metadata dicts (Proxy usage)
|
||||
for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs):
|
||||
user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(metadata=metadata)
|
||||
if user_key is not None:
|
||||
return user_key
|
||||
|
||||
return None
|
||||
metadata_dicts: Final = cls._iter_metadata_dicts(request_kwargs)
|
||||
user_api_key_hash: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_hash")
|
||||
if user_api_key_hash is not None:
|
||||
return user_api_key_hash
|
||||
user_id: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_user_id")
|
||||
if user_id is None:
|
||||
return None
|
||||
return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}"
|
||||
|
||||
@staticmethod
|
||||
def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None:
|
||||
|
|
@ -533,9 +533,9 @@ class DeploymentAffinityCheck(CustomLogger):
|
|||
return typed_healthy_deployments
|
||||
|
||||
verbose_router_logger.debug(
|
||||
"DeploymentAffinityCheck: api-key affinity hit -> deployment=%s user_key=%s",
|
||||
"DeploymentAffinityCheck: caller affinity hit -> deployment=%s user_key=%s",
|
||||
model_id,
|
||||
self._shorten_for_logs(user_key),
|
||||
self._shorten_for_logs(self._hash_user_key(user_key)),
|
||||
)
|
||||
return [deployment]
|
||||
|
||||
|
|
@ -626,7 +626,7 @@ class DeploymentAffinityCheck(CustomLogger):
|
|||
deployment_model_name,
|
||||
model_id,
|
||||
self.ttl_seconds,
|
||||
self._shorten_for_logs(user_key),
|
||||
self._shorten_for_logs(self._hash_user_key(user_key)),
|
||||
)
|
||||
else:
|
||||
verbose_router_logger.debug(
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
|
@ -29,6 +30,13 @@ def is_interception_internal_key(
|
|||
return any(key.startswith(prefix) for prefix in prefixes)
|
||||
|
||||
|
||||
CONVERTED_STREAM_KEYS: Final = frozenset(f"{prefix}_converted_stream" for prefix in INTERCEPTION_INTERNAL_PREFIXES)
|
||||
|
||||
|
||||
def converted_stream_requested(params: Mapping[str, object]) -> bool:
|
||||
return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS)
|
||||
|
||||
|
||||
class AgenticLoopSafetyError(ValueError):
|
||||
"""
|
||||
Raised when an agentic-loop safety rail refuses a rerun.
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ from pydantic import (
|
|||
ConfigDict,
|
||||
Discriminator,
|
||||
Field,
|
||||
NonNegativeInt,
|
||||
PrivateAttr,
|
||||
SerializerFunctionWrapHandler,
|
||||
field_serializer,
|
||||
|
|
@ -1321,6 +1322,18 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject):
|
|||
model_config = {"extra": "allow"}
|
||||
|
||||
|
||||
class WebSearchToolUsage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
num_requests: NonNegativeInt
|
||||
|
||||
|
||||
class ResponsesToolUsage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
web_search: WebSearchToolUsage | None = None
|
||||
|
||||
|
||||
ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"]
|
||||
"""
|
||||
The status of the response generation.
|
||||
|
|
|
|||
|
|
@ -11,8 +11,8 @@ class ModelInfoMetadata(TypedDict):
|
|||
|
||||
class ModelInfoResponse(TypedDict):
|
||||
"""OpenAI-compatible model object. `mode`, `max_input_tokens`, and
|
||||
`max_output_tokens` are attached when the cost map knows them; `metadata`
|
||||
is present only when the endpoint is called with include_metadata=true.
|
||||
`max_output_tokens` are attached when the cost map or deployment config
|
||||
knows them; `metadata` is present only with include_metadata=true.
|
||||
"""
|
||||
|
||||
id: str
|
||||
|
|
|
|||
|
|
@ -885,6 +885,18 @@ class FallbackAccessCheck(Protocol):
|
|||
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ...
|
||||
|
||||
|
||||
class HeuristicV2RouterLimit(Protocol):
|
||||
"""
|
||||
Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited.
|
||||
|
||||
The Router calls it on every registration and limit query instead of caching the answer, so the
|
||||
proxy can keep the limit on its license object (re-verified on config load) rather than hand
|
||||
over a snapshot.
|
||||
"""
|
||||
|
||||
def __call__(self) -> int | None: ...
|
||||
|
||||
|
||||
class LiteLLM_RouterFileObject(TypedDict, total=False):
|
||||
"""
|
||||
Tracking the litellm params hash, used for mapping the file id to the right model
|
||||
|
|
|
|||
|
|
@ -29354,6 +29354,75 @@
|
|||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"gpt-6-astra": {
|
||||
"cache_creation_input_token_cost": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
|
||||
"cache_creation_input_token_cost_flex": 6.25e-06,
|
||||
"cache_creation_input_token_cost_priority": 2.5e-05,
|
||||
"cache_read_input_token_cost": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 2e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_flex": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 4e-06,
|
||||
"cache_read_input_token_cost_flex": 5e-07,
|
||||
"cache_read_input_token_cost_priority": 2e-06,
|
||||
"input_cost_per_token": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 2e-05,
|
||||
"input_cost_per_token_above_272k_tokens_flex": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4e-05,
|
||||
"input_cost_per_token_batches": 5e-06,
|
||||
"input_cost_per_token_flex": 5e-06,
|
||||
"input_cost_per_token_priority": 2e-05,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 922000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 5e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 7.5e-05,
|
||||
"output_cost_per_token_above_272k_tokens_flex": 3.75e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 0.00015,
|
||||
"output_cost_per_token_batches": 2.5e-05,
|
||||
"output_cost_per_token_flex": 2.5e-05,
|
||||
"output_cost_per_token_priority": 0.0001,
|
||||
"regional_processing_uplift_multiplier_eu": 1.1,
|
||||
"regional_processing_uplift_multiplier_us": 1.1,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_native_streaming": true,
|
||||
"supports_none_reasoning_effort": false,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_cache_breakpoint": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_xhigh_reasoning_effort": true
|
||||
},
|
||||
"gpt-5.6": {
|
||||
"cache_creation_input_token_cost": 5e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens": 1e-05,
|
||||
|
|
@ -53485,6 +53554,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.95e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53519,6 +53593,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 4.4e-07,
|
||||
"output_cost_per_token": 1.32e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 1.98e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53581,6 +53660,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 4.4e-08,
|
||||
"output_cost_per_token": 1.32e-06,
|
||||
"output_cost_per_token_above_272k_tokens": 1.98e-06,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53769,6 +53853,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.95e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
@ -53800,6 +53889,11 @@
|
|||
"cache_read_input_token_cost_above_272k_tokens": 5.5e-07,
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 2.475e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.012,
|
||||
"search_context_size_low": 0.012,
|
||||
"search_context_size_medium": 0.012
|
||||
},
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
|
|
|
|||
|
|
@ -2879,7 +2879,6 @@ def response_format_tests(response: litellm.ModelResponse):
|
|||
"model",
|
||||
[
|
||||
"bedrock/mistral.mistral-large-2407-v1:0",
|
||||
"bedrock/cohere.command-r-plus-v1:0",
|
||||
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"mistral.mistral-7b-instruct-v0:2",
|
||||
"meta.llama3-8b-instruct-v1:0",
|
||||
|
|
|
|||
|
|
@ -1168,7 +1168,6 @@ async def test_completion_replicate_llama3_streaming(sync_mode):
|
|||
"model, region",
|
||||
[
|
||||
# ["bedrock/ai21.jamba-instruct-v1:0", "us-east-1"],
|
||||
# ["bedrock/cohere.command-r-plus-v1:0", None],
|
||||
["us.anthropic.claude-sonnet-4-5-20250929-v1:0", None],
|
||||
# ["mistral.mistral-7b-instruct-v0:2", None],
|
||||
# ["meta.llama3-8b-instruct-v1:0", None],
|
||||
|
|
@ -1271,7 +1270,7 @@ def test_bedrock_claude_3_streaming():
|
|||
"model",
|
||||
[
|
||||
"claude-haiku-4-5-20251001",
|
||||
"cohere.command-r-plus-v1:0", # bedrock
|
||||
"bedrock/mistral.mistral-7b-instruct-v0:2",
|
||||
"gpt-3.5-turbo",
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1295,6 +1295,19 @@ def test_text_plus_tool_calls_sequence():
|
|||
# =============================================================================
|
||||
|
||||
|
||||
def test_developer_message_content_uses_input_text():
|
||||
handler = LiteLLMResponsesTransformationHandler()
|
||||
|
||||
input_items, instructions = handler.convert_chat_completion_messages_to_responses_api(
|
||||
[{"role": "developer", "content": "Always answer in French."}]
|
||||
)
|
||||
|
||||
assert instructions is None
|
||||
assert input_items == [
|
||||
{"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "Always answer in French."}]}
|
||||
]
|
||||
|
||||
|
||||
def test_tool_message_output_uses_input_text_not_output_text():
|
||||
"""
|
||||
Test that tool message content uses input_text type, not output_text.
|
||||
|
|
|
|||
|
|
@ -1662,6 +1662,54 @@ def test_generic_cost_per_token_gpt56_cyber(
|
|||
assert completion_cost == pytest.approx(completion_tokens * output_rate)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"service_tier,tier_multiplier",
|
||||
[(None, 1.0), ("flex", 0.5), ("priority", 2.0), ("fast", 2.0)],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"prompt_tokens,input_side_multiplier,output_multiplier",
|
||||
[(100000, 1.0, 1.0), (300000, 2.0, 1.5)],
|
||||
)
|
||||
def test_generic_cost_per_token_gpt_6_astra_price_sheet(
|
||||
_local_model_cost_map,
|
||||
service_tier,
|
||||
tier_multiplier,
|
||||
prompt_tokens,
|
||||
input_side_multiplier,
|
||||
output_multiplier,
|
||||
):
|
||||
"""gpt-6-astra launch price sheet: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens.
|
||||
|
||||
Above 272K prompt tokens the input-side rates double and the output rate is 1.5x on the whole
|
||||
request. Flex is half the applicable rate and fast mode, billed as priority, is double it.
|
||||
"""
|
||||
cached_tokens = 50000
|
||||
cache_write_tokens = 40000
|
||||
text_tokens = prompt_tokens - cached_tokens - cache_write_tokens
|
||||
completion_tokens = 1000
|
||||
usage = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="gpt-6-astra",
|
||||
usage=usage,
|
||||
custom_llm_provider="openai",
|
||||
service_tier=service_tier,
|
||||
)
|
||||
|
||||
input_side = tier_multiplier * input_side_multiplier
|
||||
assert prompt_cost == pytest.approx(
|
||||
input_side * (text_tokens * 1e-5 + cached_tokens * 1e-6 + cache_write_tokens * 1.25e-5)
|
||||
)
|
||||
assert completion_cost == pytest.approx(tier_multiplier * output_multiplier * completion_tokens * 5e-5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,input_cost,output_cost,cache_read_cost",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -6,7 +7,7 @@ import litellm
|
|||
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
|
||||
StandardBuiltInToolCostTracking,
|
||||
)
|
||||
from litellm.types.llms.openai import FileSearchTool, WebSearchOptions
|
||||
from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebSearchOptions
|
||||
from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams
|
||||
|
||||
|
||||
|
|
@ -928,3 +929,125 @@ def test_web_search_gate_reads_server_side_tool_usage_details_without_citations(
|
|||
standard_built_in_tools_params=None,
|
||||
)
|
||||
assert cost == 3 * _DEFAULT_WEB_SEARCH_COST_PER_CALL
|
||||
|
||||
|
||||
_BEDROCK_MANTLE_WEB_SEARCH_MODELS = (
|
||||
"bedrock_mantle/openai.gpt-5.6-sol",
|
||||
"bedrock_mantle/openai.gpt-5.6-terra",
|
||||
"bedrock_mantle/openai.gpt-5.6-luna",
|
||||
"bedrock_mantle/openai.gpt-5.5",
|
||||
"bedrock_mantle/openai.gpt-5.4",
|
||||
)
|
||||
|
||||
_BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012
|
||||
|
||||
|
||||
def _responses_with_web_search(
|
||||
model: str, actions: Sequence[Mapping[str, str]], tool_usage: Mapping[str, object] | None = None
|
||||
) -> ResponsesAPIResponse:
|
||||
payload = {
|
||||
"id": "resp_1",
|
||||
"created_at": 1756900000,
|
||||
"model": model.split("/", 1)[-1],
|
||||
"object": "response",
|
||||
"status": "completed",
|
||||
"output": [
|
||||
{"type": "web_search_call", "id": f"ws_{i}", "status": "completed", "action": action}
|
||||
for i, action in enumerate(actions)
|
||||
],
|
||||
}
|
||||
return ResponsesAPIResponse.model_validate(
|
||||
payload if tool_usage is None else {**payload, "tool_usage": tool_usage}
|
||||
)
|
||||
|
||||
|
||||
def _web_search_cost(model: str, response: ResponsesAPIResponse, custom_llm_provider: str) -> float:
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
|
||||
model=model,
|
||||
response_object=response,
|
||||
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
standard_built_in_tools_params=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS)
|
||||
def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model):
|
||||
"""Two Bedrock-reported web searches bill 2 x $0.012 under the prefixed and the bare model id alike."""
|
||||
pricing = litellm.get_model_info(model)["search_context_cost_per_query"]
|
||||
assert pricing == {
|
||||
"search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
|
||||
"search_context_size_medium": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
|
||||
"search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
|
||||
}
|
||||
|
||||
response = _responses_with_web_search(
|
||||
model,
|
||||
actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
|
||||
tool_usage={"web_search": {"num_requests": 2}},
|
||||
)
|
||||
for cost_model in (model, model.split("/", 1)[1]):
|
||||
cost = _web_search_cost(cost_model, response, "bedrock_mantle")
|
||||
assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
|
||||
f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_requests", [1, 0])
|
||||
def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests):
|
||||
"""A search plus an open_page fetch bills tool_usage.web_search.num_requests, never the two items."""
|
||||
model = "bedrock_mantle/openai.gpt-5.6-sol"
|
||||
response = _responses_with_web_search(
|
||||
model,
|
||||
actions=[
|
||||
{"type": "search", "query": "litellm"},
|
||||
{"type": "open_page", "url": "https://docs.litellm.ai/"},
|
||||
],
|
||||
tool_usage={"web_search": {"num_requests": num_requests}},
|
||||
)
|
||||
|
||||
cost = _web_search_cost(model, response, "bedrock_mantle")
|
||||
|
||||
assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
|
||||
f"{num_requests} reported web search requests must bill {num_requests} x "
|
||||
f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_usage",
|
||||
[None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}],
|
||||
)
|
||||
def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage):
|
||||
"""Without a usable reported count the per-call path keeps counting web_search_call items."""
|
||||
model = "bedrock_mantle/openai.gpt-5.6-sol"
|
||||
response = _responses_with_web_search(
|
||||
model,
|
||||
actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
|
||||
tool_usage=tool_usage,
|
||||
)
|
||||
|
||||
cost = _web_search_cost(model, response, "bedrock_mantle")
|
||||
|
||||
assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), (
|
||||
f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x "
|
||||
f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
|
||||
)
|
||||
|
||||
|
||||
def test_web_search_call_count_reads_reported_count_beside_other_tool_usage_entries(local_model_cost_map):
|
||||
"""OpenAI reports web_search.num_requests next to other tool entries, which must not disable the reported count."""
|
||||
response = _responses_with_web_search(
|
||||
"gpt-5.6",
|
||||
actions=[{"type": "search", "query": "S&P 500 close"}, {"type": "open_page", "url": "https://example.com/"}],
|
||||
tool_usage={
|
||||
"image_gen": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
|
||||
"web_search": {"num_requests": 1},
|
||||
},
|
||||
)
|
||||
|
||||
cost = _web_search_cost("gpt-5.6", response, "openai")
|
||||
|
||||
assert cost == pytest.approx(0.01), f"1 reported OpenAI web search must bill 1 x $0.01, not the 2 items, got ${cost}"
|
||||
|
|
|
|||
|
|
@ -4761,6 +4761,43 @@ async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging
|
|||
assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_fake_stream_final_chunk_carries_hidden_usage(logging_obj: Logging):
|
||||
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
model_response = ModelResponse(
|
||||
id="chatcmpl-fake-stream",
|
||||
model="my-random-model",
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "hello world"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
)
|
||||
model_response.usage = Usage(prompt_tokens=1234, completion_tokens=7, total_tokens=1241)
|
||||
|
||||
wrapper = CustomStreamWrapper(
|
||||
completion_stream=MockResponseIterator(model_response=model_response),
|
||||
model="my-random-model",
|
||||
custom_llm_provider="anthropic",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
final_chunk = None
|
||||
async for chunk in wrapper:
|
||||
final_chunk = chunk
|
||||
|
||||
assert final_chunk is not None
|
||||
hidden_usage = final_chunk._hidden_params.get("usage")
|
||||
assert hidden_usage is not None
|
||||
assert hidden_usage.prompt_tokens == 1234
|
||||
assert hidden_usage.completion_tokens == 7
|
||||
assert hidden_usage.total_tokens == 1241
|
||||
|
||||
|
||||
class TestStableStreamingResponseId:
|
||||
"""
|
||||
All chunks of one streamed response must share the same top-level id
|
||||
|
|
|
|||
|
|
@ -336,3 +336,15 @@ class TestAzureResolvesTheDeclaredDefaultEffort:
|
|||
drop_params=True,
|
||||
)
|
||||
assert ("temperature" in mapped) is temperature_survives
|
||||
|
||||
|
||||
def test_azure_gpt_6_astra_takes_the_reasoning_series_request_shape():
|
||||
params = litellm.get_optional_params(
|
||||
model="gpt-6-astra",
|
||||
custom_llm_provider="azure",
|
||||
max_tokens=100,
|
||||
reasoning_effort="max",
|
||||
)
|
||||
assert params["max_completion_tokens"] == 100
|
||||
assert "max_tokens" not in params
|
||||
assert params["reasoning_effort"] == "max"
|
||||
|
|
|
|||
|
|
@ -31,6 +31,46 @@ async def test_get_openai_compatible_provider_info():
|
|||
assert custom_llm_provider == "azure"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, api_base, expected_provider",
|
||||
[
|
||||
("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com", "azure_ai"),
|
||||
("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com/models", "azure_ai"),
|
||||
("azure_ai/gpt-5.4-nano", "https://my-resource.services.ai.azure.com", "azure_ai"),
|
||||
("azure_ai/gpt-4o", "https://my-resource.openai.azure.com", "azure"),
|
||||
(
|
||||
"azure_ai/gpt-4o",
|
||||
"https://my-resource.services.ai.azure.com/openai/deployments/gpt-4o/chat/completions"
|
||||
"?api-version=2024-08-01-preview",
|
||||
"azure",
|
||||
),
|
||||
("azure_ai/mistral-large-latest", "https://my-resource.services.ai.azure.com", "azure_ai"),
|
||||
("azure_ai/mistral-large-latest", "https://my-resource.openai.azure.com", "azure_ai"),
|
||||
],
|
||||
)
|
||||
def test_foundry_base_keeps_azure_ai_provider(model: str, api_base: str, expected_provider: str):
|
||||
"""Regression for #38276: a Foundry .services.ai.azure.com base must not be reclassified as azure."""
|
||||
config = AzureAIStudioConfig()
|
||||
(
|
||||
_,
|
||||
_,
|
||||
custom_llm_provider,
|
||||
) = config._get_openai_compatible_provider_info(
|
||||
model=model,
|
||||
api_base=api_base,
|
||||
api_key="my-key",
|
||||
custom_llm_provider="azure_ai",
|
||||
)
|
||||
assert custom_llm_provider == expected_provider
|
||||
|
||||
|
||||
def test_is_azure_openai_model_without_api_base_keeps_azure_ai():
|
||||
"""Metadata lookups (get_model_info, supports_* checks) carry no api_base and must not flip the provider."""
|
||||
config = AzureAIStudioConfig()
|
||||
assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base=None) is False
|
||||
assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base="https://my-res.openai.azure.com") is True
|
||||
|
||||
|
||||
def test_azure_ai_validate_environment():
|
||||
config = AzureAIStudioConfig()
|
||||
headers = config.validate_environment(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,69 @@
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from litellm import embedding
|
||||
from litellm.llms.azure_ai.embed.handler import _foundry_models_route_base
|
||||
|
||||
EMBEDDING_PAYLOAD = {
|
||||
"object": "list",
|
||||
"data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}],
|
||||
"model": "text-embedding-3-small",
|
||||
"usage": {"prompt_tokens": 2, "total_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "expected"),
|
||||
[
|
||||
(
|
||||
"https://my-foundry.services.ai.azure.com",
|
||||
"https://my-foundry.services.ai.azure.com/models",
|
||||
),
|
||||
(
|
||||
"https://my-foundry.services.ai.azure.com/",
|
||||
"https://my-foundry.services.ai.azure.com/models",
|
||||
),
|
||||
(
|
||||
"https://my-foundry.services.ai.azure.com?api-version=2024-05-01-preview",
|
||||
"https://my-foundry.services.ai.azure.com/models?api-version=2024-05-01-preview",
|
||||
),
|
||||
(
|
||||
"https://my-foundry.services.ai.azure.com/models",
|
||||
"https://my-foundry.services.ai.azure.com/models",
|
||||
),
|
||||
(
|
||||
"https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small",
|
||||
"https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small",
|
||||
),
|
||||
(
|
||||
"https://my-resource.openai.azure.com",
|
||||
"https://my-resource.openai.azure.com",
|
||||
),
|
||||
(
|
||||
"https://Mistral-serverless.eastus2.models.ai.azure.com",
|
||||
"https://Mistral-serverless.eastus2.models.ai.azure.com",
|
||||
),
|
||||
(None, None),
|
||||
],
|
||||
)
|
||||
def test_foundry_models_route_base(api_base, expected):
|
||||
assert _foundry_models_route_base(api_base) == expected
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_azure_ai_embedding_calls_foundry_models_route():
|
||||
route = respx.post("https://my-foundry.services.ai.azure.com/models/embeddings").mock(
|
||||
return_value=httpx.Response(200, json=EMBEDDING_PAYLOAD)
|
||||
)
|
||||
|
||||
response = embedding(
|
||||
model="azure_ai/text-embedding-3-small",
|
||||
input=["hello world"],
|
||||
api_base="https://my-foundry.services.ai.azure.com",
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.data is not None
|
||||
assert len(response.data) == 1
|
||||
|
|
@ -10,11 +10,11 @@ Tests the cost calculation for Dashscope models including:
|
|||
|
||||
import math
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
# Add the project root to Python path
|
||||
|
||||
import litellm
|
||||
from litellm.llms.dashscope.cost_calculator import (
|
||||
cost_per_token as dashscope_cost_per_token,
|
||||
|
|
@ -526,3 +526,139 @@ class TestDashscopeCostCalculator:
|
|||
|
||||
assert prompt_cost == 0.0
|
||||
assert math.isclose(completion_cost, 500 * 1.6e-06, rel_tol=1e-10)
|
||||
|
||||
OFF_PEAK_WINDOW = "14:00-00:00"
|
||||
INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc)
|
||||
OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc)
|
||||
|
||||
def _register_off_peak_flat_model(self, model_key: str, off_peak_pricing: dict) -> None:
|
||||
litellm.model_cost[model_key] = {
|
||||
"litellm_provider": "dashscope",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 2.4e-06,
|
||||
"output_cost_per_token": 4.8e-06,
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"cache_creation_input_token_cost": 3e-06,
|
||||
"off_peak_pricing": off_peak_pricing,
|
||||
}
|
||||
|
||||
def test_dashscope_off_peak_window_swaps_in_the_off_peak_rates(self):
|
||||
"""
|
||||
Regression (LIT-6782): a deployment configured with off_peak_pricing kept billing the
|
||||
standard dashscope rates inside its window, while the same block on a deepseek
|
||||
deployment billed the off-peak rates.
|
||||
"""
|
||||
self._register_off_peak_flat_model(
|
||||
"dashscope/deepseek-off-peak-test",
|
||||
{
|
||||
"hours_utc": self.OFF_PEAK_WINDOW,
|
||||
"input_cost_per_token": 1.2e-06,
|
||||
"output_cost_per_token": 2.4e-06,
|
||||
"cache_read_input_token_cost": 1e-07,
|
||||
},
|
||||
)
|
||||
usage = Usage(
|
||||
prompt_tokens=1000,
|
||||
completion_tokens=200,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = dashscope_cost_per_token(
|
||||
model="deepseek-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
|
||||
)
|
||||
|
||||
assert math.isclose(prompt_cost, (600 * 1.2e-06) + (300 * 1e-07) + (100 * 3e-06), rel_tol=1e-10)
|
||||
assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)
|
||||
|
||||
peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token(
|
||||
model="deepseek-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW
|
||||
)
|
||||
|
||||
assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10)
|
||||
assert math.isclose(peak_completion_cost, 200 * 4.8e-06, rel_tol=1e-10)
|
||||
|
||||
def test_dashscope_off_peak_window_overrides_the_selected_tier(self):
|
||||
"""An open off-peak window bills the whole request at the flat off-peak rates, whichever tier
|
||||
the input volume selected."""
|
||||
self._register_tiered_model(
|
||||
"dashscope/qwen-tiered-off-peak-test",
|
||||
[
|
||||
{"range": [0, 1000], "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06},
|
||||
{"range": [1000, 2000], "input_cost_per_token": 8e-07, "output_cost_per_token": 3.2e-06},
|
||||
],
|
||||
)
|
||||
litellm.model_cost["dashscope/qwen-tiered-off-peak-test"]["off_peak_pricing"] = {
|
||||
"hours_utc": self.OFF_PEAK_WINDOW,
|
||||
"input_cost_per_token": 1e-07,
|
||||
"output_cost_per_token": 4e-07,
|
||||
}
|
||||
usage = Usage(prompt_tokens=1500, completion_tokens=300)
|
||||
|
||||
prompt_cost, completion_cost = dashscope_cost_per_token(
|
||||
model="qwen-tiered-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
|
||||
)
|
||||
|
||||
assert math.isclose(prompt_cost, 1500 * 1e-07, rel_tol=1e-10)
|
||||
assert math.isclose(completion_cost, 300 * 4e-07, rel_tol=1e-10)
|
||||
|
||||
peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token(
|
||||
model="qwen-tiered-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW
|
||||
)
|
||||
|
||||
assert math.isclose(peak_prompt_cost, 1500 * 8e-07, rel_tol=1e-10)
|
||||
assert math.isclose(peak_completion_cost, 300 * 3.2e-06, rel_tol=1e-10)
|
||||
|
||||
def test_dashscope_off_peak_rates_left_unset_keep_the_standard_rates(self):
|
||||
"""A block that only overrides the input rate leaves output and cache reads on the standard
|
||||
rates, and an explicit reasoning rate is never swapped out."""
|
||||
self._register_off_peak_flat_model(
|
||||
"dashscope/qwen-partial-off-peak-test",
|
||||
{"hours_utc": self.OFF_PEAK_WINDOW, "input_cost_per_token": 1.2e-06},
|
||||
)
|
||||
litellm.model_cost["dashscope/qwen-partial-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-06
|
||||
usage = Usage(
|
||||
prompt_tokens=1000,
|
||||
completion_tokens=200,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = dashscope_cost_per_token(
|
||||
model="qwen-partial-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
|
||||
)
|
||||
|
||||
assert math.isclose(prompt_cost, (700 * 1.2e-06) + (300 * 2e-07), rel_tol=1e-10)
|
||||
assert math.isclose(completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10)
|
||||
|
||||
def test_dashscope_off_peak_output_rate_covers_reasoning_without_a_dedicated_rate(self):
|
||||
"""Reasoning tokens on a model with no dedicated reasoning rate follow the off-peak output
|
||||
rate, the same way they follow the standard output rate outside the window."""
|
||||
self._register_off_peak_flat_model(
|
||||
"dashscope/qwen-reasoning-off-peak-test",
|
||||
{"hours_utc": self.OFF_PEAK_WINDOW, "output_cost_per_token": 2.4e-06},
|
||||
)
|
||||
usage = Usage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=200,
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50),
|
||||
)
|
||||
|
||||
_, completion_cost = dashscope_cost_per_token(
|
||||
model="qwen-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
|
||||
)
|
||||
|
||||
assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)
|
||||
|
||||
def test_dashscope_off_peak_defaults_to_the_current_time(self):
|
||||
"""The proxy's cost dispatch passes no clock, so an all-day window has to apply on the
|
||||
default current time."""
|
||||
self._register_off_peak_flat_model(
|
||||
"dashscope/qwen-all-day-off-peak-test",
|
||||
{"hours_utc": "00:00-00:00", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 2.4e-06},
|
||||
)
|
||||
usage = Usage(prompt_tokens=1000, completion_tokens=200)
|
||||
|
||||
prompt_cost, completion_cost = dashscope_cost_per_token(model="qwen-all-day-off-peak-test", usage=usage)
|
||||
|
||||
assert math.isclose(prompt_cost, 1000 * 1.2e-06, rel_tol=1e-10)
|
||||
assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)
|
||||
|
|
|
|||
|
|
@ -1329,6 +1329,523 @@ class TestOpenAIResponsesHandlerToolInjection:
|
|||
assert "injected_tool" in names
|
||||
|
||||
|
||||
COMPRESSED_MARKER = "[compressed document; retrieve the full text with hash=b573993006976af767214fac]"
|
||||
|
||||
|
||||
class StructuredRewriteGuardrail(CustomGuardrail):
|
||||
"""Guardrail that rewrites whole messages via structured_messages and leaves
|
||||
texts untouched, the way message-compressing guardrails do."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
messages = list(inputs.get("structured_messages") or [])
|
||||
first_user = next(i for i, m in enumerate(messages) if m.get("role") == "user")
|
||||
rewritten = [
|
||||
{**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages)
|
||||
]
|
||||
return {**inputs, "structured_messages": rewritten}
|
||||
|
||||
|
||||
class ToolOutputRewriteGuardrail(CustomGuardrail):
|
||||
"""Guardrail that compresses the first tool-result row, the way Headroom does."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
messages = list(inputs.get("structured_messages") or [])
|
||||
first_tool = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "tool")
|
||||
rewritten = [
|
||||
{**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages)
|
||||
]
|
||||
return {**inputs, "structured_messages": rewritten}
|
||||
|
||||
|
||||
class DroppingRewriteGuardrail(CustomGuardrail):
|
||||
"""Guardrail that rewrites the first user row and drops the last row, so the
|
||||
rewrite can only land through the full-conversion fallback."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
messages = list(inputs.get("structured_messages") or [])
|
||||
first_user = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "user")
|
||||
rewritten = [
|
||||
{**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages)
|
||||
]
|
||||
return {**inputs, "structured_messages": rewritten[:-1]}
|
||||
|
||||
|
||||
def _texts(item: dict) -> list[str]:
|
||||
content = item.get("content")
|
||||
if isinstance(content, str):
|
||||
return [content]
|
||||
return [part["text"] for part in content]
|
||||
|
||||
|
||||
class TestStructuredMessagesWriteBack:
|
||||
"""A guardrail's structured_messages rewrite must land in the Responses request,
|
||||
not only the per-text mapping the chat handler shares with it."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_input_gets_rewritten_messages_and_keeps_instructions(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Answer from the memo only.",
|
||||
"input": [
|
||||
{"role": "user", "content": "memo " * 400},
|
||||
{"role": "assistant", "content": "Understood."},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
|
||||
|
||||
assert result["instructions"] == "Answer from the memo only."
|
||||
user_items = [item for item in result["input"] if item.get("role") == "user"]
|
||||
assert [_texts(item) for item in user_items] == [[COMPRESSED_MARKER], ["What is the codename?"]]
|
||||
assert not any(item.get("role") == "system" for item in result["input"])
|
||||
assert _texts(next(item for item in result["input"] if item.get("role") == "assistant")) == ["Understood."]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_string_input_becomes_rewritten_message_list(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
data = {"model": "gpt-5.6", "input": "memo " * 400}
|
||||
|
||||
result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
|
||||
|
||||
assert [_texts(item) for item in result["input"]] == [[COMPRESSED_MARKER]]
|
||||
assert "instructions" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_developer_item_preserved_verbatim_by_row_patch(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
developer_item = {"role": "developer", "content": "Always answer in French."}
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"input": [
|
||||
developer_item,
|
||||
{"role": "user", "content": "memo " * 400},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
|
||||
|
||||
assert result["input"][0] is developer_item
|
||||
assert developer_item["content"] == "Always answer in French."
|
||||
assert _texts(result["input"][1]) == [COMPRESSED_MARKER]
|
||||
assert _texts(result["input"][2]) == ["What is the codename?"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reasoning_and_function_call_items_survive_tool_output_compression(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
reasoning_item = {
|
||||
"id": "rs_123",
|
||||
"type": "reasoning",
|
||||
"summary": [],
|
||||
"encrypted_content": "gAAAAA-signed-reasoning",
|
||||
}
|
||||
function_call_item = {
|
||||
"id": "fc_123",
|
||||
"type": "function_call",
|
||||
"call_id": "call_abc",
|
||||
"name": "read_document",
|
||||
"arguments": '{"path": "memo.txt"}',
|
||||
"status": "completed",
|
||||
}
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Answer from the memo only.",
|
||||
"input": [
|
||||
reasoning_item,
|
||||
function_call_item,
|
||||
{"type": "function_call_output", "call_id": "call_abc", "output": "memo " * 400},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail())
|
||||
|
||||
assert result["instructions"] == "Answer from the memo only."
|
||||
assert result["input"][0] is reasoning_item
|
||||
assert reasoning_item["encrypted_content"] == "gAAAAA-signed-reasoning"
|
||||
assert result["input"][1] is function_call_item
|
||||
assert function_call_item["id"] == "fc_123"
|
||||
assert result["input"][2] == {
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_abc",
|
||||
"output": COMPRESSED_MARKER,
|
||||
}
|
||||
assert result["input"][3] == {"role": "user", "content": "What is the codename?"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_web_search_call_item_preserved_verbatim(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
web_search_item = {
|
||||
"id": "ws_123",
|
||||
"type": "web_search_call",
|
||||
"status": "completed",
|
||||
"action": {"type": "search", "query": "codename memo"},
|
||||
}
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"input": [
|
||||
web_search_item,
|
||||
{"role": "user", "content": "memo " * 400},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
|
||||
|
||||
assert result["input"][0] is web_search_item
|
||||
assert _texts(result["input"][1]) == [COMPRESSED_MARKER]
|
||||
assert _texts(result["input"][2]) == ["What is the codename?"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_row_count_change_falls_back_to_full_conversion(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"input": [
|
||||
{"role": "developer", "content": "Always answer in French."},
|
||||
{"role": "user", "content": "memo " * 400},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, DroppingRewriteGuardrail())
|
||||
|
||||
assert len(result["input"]) == 2
|
||||
developer = next(item for item in result["input"] if item.get("role") == "developer")
|
||||
assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}]
|
||||
assert _texts(next(item for item in result["input"] if item.get("role") == "user")) == [COMPRESSED_MARKER]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_inputs_object_back_keeps_the_text_mapping(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
original_input = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "Again"}]},
|
||||
]
|
||||
data = {"model": "gpt-5.6", "input": original_input}
|
||||
|
||||
result = await handler.process_input_messages(data, MockGuardrail())
|
||||
|
||||
assert result["input"] is original_input
|
||||
assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]]
|
||||
|
||||
|
||||
class AllToolOutputsRewriteGuardrail(CustomGuardrail):
|
||||
"""Guardrail that compresses every tool-result row."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
messages = list(inputs.get("structured_messages") or [])
|
||||
rewritten = [
|
||||
{**m, "content": COMPRESSED_MARKER} if isinstance(m, dict) and m.get("role") == "tool" else m
|
||||
for m in messages
|
||||
]
|
||||
return {**inputs, "structured_messages": rewritten}
|
||||
|
||||
|
||||
class AssistantRewriteGuardrail(CustomGuardrail):
|
||||
"""Guardrail that rewrites the first assistant row's content."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
messages = list(inputs.get("structured_messages") or [])
|
||||
first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "assistant")
|
||||
rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first else m for i, m in enumerate(messages)]
|
||||
return {**inputs, "structured_messages": rewritten}
|
||||
|
||||
|
||||
class DictStructuredMessagesGuardrail(CustomGuardrail):
|
||||
"""Guardrail that hands back a raw evaluation dict instead of a message list,
|
||||
the way HiddenLayer v2 does."""
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}}
|
||||
|
||||
|
||||
def _parallel_tool_call_input() -> list:
|
||||
return [
|
||||
{"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
|
||||
{"id": "fc_2", "type": "function_call", "call_id": "call_2", "name": "read_b", "arguments": "{}"},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "memo " * 400},
|
||||
{"type": "function_call_output", "call_id": "call_2", "output": "note " * 400},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
]
|
||||
|
||||
|
||||
class TestProvenancePatching:
|
||||
"""The O(n) provenance pass must keep patching rewritten rows in place for the
|
||||
shapes real agent loops produce, and fall back safely everywhere else."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_tool_call_outputs_both_patched(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
raw_input = _parallel_tool_call_input()
|
||||
fc_1, fc_2 = raw_input[0], raw_input[1]
|
||||
data = {"model": "gpt-5.6", "input": raw_input}
|
||||
|
||||
result = await handler.process_input_messages(data, AllToolOutputsRewriteGuardrail())
|
||||
|
||||
assert result["input"][0] is fc_1
|
||||
assert result["input"][1] is fc_2
|
||||
assert result["input"][2] == {"type": "function_call_output", "call_id": "call_1", "output": COMPRESSED_MARKER}
|
||||
assert result["input"][3] == {"type": "function_call_output", "call_id": "call_2", "output": COMPRESSED_MARKER}
|
||||
assert result["input"][4] == {"role": "user", "content": "What is the codename?"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assistant_turn_with_tool_call_keeps_items_verbatim(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
assistant_item = {"role": "assistant", "content": "Let me read the memo."}
|
||||
function_call_item = {
|
||||
"id": "fc_9",
|
||||
"type": "function_call",
|
||||
"call_id": "call_9",
|
||||
"name": "read_document",
|
||||
"arguments": '{"path": "memo.txt"}',
|
||||
}
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"input": [
|
||||
assistant_item,
|
||||
function_call_item,
|
||||
{"type": "function_call_output", "call_id": "call_9", "output": "memo " * 400},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail())
|
||||
|
||||
assert result["input"][0] is assistant_item
|
||||
assert result["input"][1] is function_call_item
|
||||
assert result["input"][2] == {"type": "function_call_output", "call_id": "call_9", "output": COMPRESSED_MARKER}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewrite_of_merged_tool_call_message_falls_back(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
raw_input = _parallel_tool_call_input()
|
||||
data = {"model": "gpt-5.6", "input": raw_input}
|
||||
|
||||
result = await handler.process_input_messages(data, AssistantRewriteGuardrail())
|
||||
|
||||
assert not any(item is original for item in result["input"] for original in raw_input)
|
||||
assistant_items = [item for item in result["input"] if item.get("role") == "assistant"]
|
||||
assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewrite_of_lone_function_call_message_falls_back(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"input": [
|
||||
{"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "memo memo"},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
raw_input = data["input"]
|
||||
result = await handler.process_input_messages(data, AssistantRewriteGuardrail())
|
||||
|
||||
assert not any(item is original for item in result["input"] for original in raw_input)
|
||||
assistant_items = [item for item in result["input"] if item.get("role") == "assistant"]
|
||||
assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]]
|
||||
|
||||
def test_provenance_bails_on_non_mapping_item(self):
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
|
||||
|
||||
assert _input_item_provenance(["not a mapping"], []) is None
|
||||
|
||||
def test_provenance_bails_when_expected_messages_disagree(self):
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
|
||||
|
||||
assert _input_item_provenance([{"role": "user", "content": "hi"}], [{"role": "user", "content": "bye"}]) is None
|
||||
|
||||
def test_provenance_bails_on_unpredicted_merge(self):
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
raw_input = [
|
||||
{"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
|
||||
{"role": "assistant", "content": "Reading the memo now."},
|
||||
]
|
||||
expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=raw_input, responses_api_request={}
|
||||
)
|
||||
assert len(expected) == 1
|
||||
assert _input_item_provenance(raw_input, expected) is None
|
||||
|
||||
def test_provenance_maps_and_taints_parallel_tool_calls(self):
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
raw_input = _parallel_tool_call_input()
|
||||
expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=raw_input, responses_api_request={}
|
||||
)
|
||||
provenance = _input_item_provenance(raw_input, expected)
|
||||
assert provenance is not None
|
||||
item_for_message, tainted = provenance
|
||||
assert tainted == {0}
|
||||
assert dict(item_for_message) == {1: 2, 2: 3, 3: 4}
|
||||
|
||||
|
||||
class TestDictStructuredMessagesGuard:
|
||||
"""A guardrail handing back a non-list structured_messages payload must not
|
||||
blow up the request; the write-back is skipped instead."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_input_survives_dict_structured_messages(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
original_input = [{"role": "user", "content": "Hello"}]
|
||||
data = {"model": "gpt-5.6", "input": original_input}
|
||||
|
||||
result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail())
|
||||
|
||||
assert result["input"] is original_input
|
||||
assert result["input"] == [{"role": "user", "content": "Hello"}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_string_input_survives_dict_structured_messages(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
data = {"model": "gpt-5.6", "input": "Hello there"}
|
||||
|
||||
result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail())
|
||||
|
||||
assert result["input"] == "Hello there"
|
||||
|
||||
|
||||
class SystemRewriteGuardrail(CustomGuardrail):
|
||||
"""Guardrail that rewrites the system row, the way prompt-hardening guardrails do."""
|
||||
|
||||
def __init__(self, rewritten_content: Any = COMPRESSED_MARKER):
|
||||
super().__init__()
|
||||
self.rewritten_content = rewritten_content
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
messages = list(inputs.get("structured_messages") or [])
|
||||
first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "system")
|
||||
rewritten = [
|
||||
{**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages)
|
||||
]
|
||||
return {**inputs, "structured_messages": rewritten}
|
||||
|
||||
|
||||
class TestPatchEdgeBranches:
|
||||
@pytest.mark.asyncio
|
||||
async def test_multimodal_user_item_rewritten_through_conversion(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"input": [
|
||||
{"role": "user", "content": [{"type": "input_text", "text": "memo " * 400}]},
|
||||
{"role": "user", "content": "What is the codename?"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
|
||||
|
||||
assert _texts(result["input"][0]) == [COMPRESSED_MARKER]
|
||||
assert result["input"][1] == {"role": "user", "content": "What is the codename?"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instructions_rewrite_lands_in_instructions_field(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
user_item = {"role": "user", "content": "What is the codename?"}
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Answer from the memo only.",
|
||||
"input": [user_item],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(data, SystemRewriteGuardrail())
|
||||
|
||||
assert result["instructions"] == COMPRESSED_MARKER
|
||||
assert result["input"][0] is user_item
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_string_instructions_rewrite_falls_back(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
user_item = {"role": "user", "content": "What is the codename?"}
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": "Answer from the memo only.",
|
||||
"input": [user_item],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(
|
||||
data, SystemRewriteGuardrail(rewritten_content=[{"type": "text", "text": COMPRESSED_MARKER}])
|
||||
)
|
||||
|
||||
assert result["input"][0] is not user_item
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unpredicted_merge_falls_back_through_patch(self):
|
||||
handler = OpenAIResponsesHandler()
|
||||
raw_input = [
|
||||
{"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"},
|
||||
{"role": "assistant", "content": "Reading the memo now."},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "memo memo"},
|
||||
{"role": "user", "content": "memo " * 400},
|
||||
]
|
||||
data = {"model": "gpt-5.6", "input": raw_input}
|
||||
|
||||
result = await handler.process_input_messages(data, StructuredRewriteGuardrail())
|
||||
|
||||
assert not any(item is original for item in result["input"] for original in raw_input)
|
||||
user_items = [item for item in result["input"] if item.get("role") == "user"]
|
||||
assert _texts(user_items[0]) == [COMPRESSED_MARKER]
|
||||
|
||||
def test_item_rewrite_field_ignores_non_string_type(self):
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import _item_rewrite_field
|
||||
|
||||
assert _item_rewrite_field({"type": 123, "content": "hello"}) is None
|
||||
|
||||
|
||||
class ToolEditingGuardrail(CustomGuardrail):
|
||||
"""Guardrail that rewrites the flattened chat tools it was handed through ``edit``"""
|
||||
|
||||
|
|
|
|||
|
|
@ -1718,6 +1718,8 @@ class TestResponsesSurfaceSharesTheEffortRule:
|
|||
("gpt-5.6-sol", None, False),
|
||||
("gpt-5.6-terra", "none", True),
|
||||
("gpt-5.6-terra", "medium", False),
|
||||
("gpt-6-astra", None, False),
|
||||
("gpt-6-astra", "low", False),
|
||||
],
|
||||
)
|
||||
def test_temperature_follows_the_resolved_effort(
|
||||
|
|
|
|||
|
|
@ -1505,3 +1505,17 @@ class TestACatalogueOlderThanTheCodeDoesNotStripTemperature:
|
|||
drop_params=True,
|
||||
)
|
||||
assert "temperature" not in mapped
|
||||
|
||||
|
||||
def test_gpt_6_astra_takes_the_reasoning_series_request_shape():
|
||||
params = litellm.get_optional_params(
|
||||
model="gpt-6-astra",
|
||||
custom_llm_provider="openai",
|
||||
max_tokens=100,
|
||||
reasoning_effort="max",
|
||||
verbosity="low",
|
||||
)
|
||||
assert params["max_completion_tokens"] == 100
|
||||
assert "max_tokens" not in params
|
||||
assert params["reasoning_effort"] == "max"
|
||||
assert params["verbosity"] == "low"
|
||||
|
|
|
|||
|
|
@ -41,6 +41,8 @@ from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config
|
|||
|
||||
# Models that MUST be classified as GPT-5 (routed through GPT-5 reasoning path)
|
||||
GPT5_MODELS = [
|
||||
"gpt-6-astra",
|
||||
"openai/gpt-6-astra",
|
||||
"gpt-5",
|
||||
"gpt-5.1",
|
||||
"gpt-5.2",
|
||||
|
|
@ -120,6 +122,8 @@ class TestOpenAIGPT5ConfigIsModelGpt5Model:
|
|||
# /v1/responses bridge (when reasoning_effort is set and tools are passed) on
|
||||
# is_model_gpt_5_4_plus_model, so the gpt-5.6 family must land on the True side.
|
||||
GPT5_4_PLUS_MODELS = [
|
||||
"gpt-6-astra",
|
||||
"openai/gpt-6-astra",
|
||||
"gpt-5.4",
|
||||
"gpt-5.5",
|
||||
"gpt-5.5-pro",
|
||||
|
|
|
|||
52
tests/test_litellm/llms/openai/test_openai.py
Normal file
52
tests/test_litellm/llms/openai/test_openai.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
import pytest
|
||||
|
||||
from litellm.llms.openai.openai import OpenAIChatCompletion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
None,
|
||||
"https://api.openai.com/v1",
|
||||
"https://api.openai.com:443/v1",
|
||||
"https://southcentralus.privatelink.api.openai.com/v1",
|
||||
"https://eu.api.openai.com/v1",
|
||||
"https://us.api.openai.com/v1",
|
||||
"HTTPS://API.OPENAI.COM/v1/",
|
||||
],
|
||||
)
|
||||
def test_get_stream_options_defaults_include_usage_on_every_openai_backed_host(api_base):
|
||||
"""
|
||||
PrivateLink and regional hostnames reach the real OpenAI backend, so a stream with no caller
|
||||
stream_options must ask for the usage chunk exactly as the default base does. Regression guard
|
||||
for LIT-6875: spend for those deployments fell back to local token counting.
|
||||
"""
|
||||
assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {
|
||||
"stream_options": {"include_usage": True}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
"https://my-gateway.example/v1",
|
||||
"https://api.openai.com.evil.example/v1",
|
||||
"https://notapi.openai.com/v1",
|
||||
"https://gateway.example/v1?upstream=api.openai.com",
|
||||
"https://openai.internal.example/api.openai.com/v1",
|
||||
],
|
||||
)
|
||||
def test_get_stream_options_leaves_foreign_hosts_without_a_usage_default(api_base):
|
||||
"""Only the host decides: an OpenAI-compatible backend elsewhere may not support stream_options at all."""
|
||||
assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
["https://southcentralus.privatelink.api.openai.com/v1", "https://my-gateway.example/v1"],
|
||||
)
|
||||
def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api_base):
|
||||
caller_options = {"include_usage": False}
|
||||
assert OpenAIChatCompletion().get_stream_options(stream_options=caller_options, api_base=api_base) == {
|
||||
"stream_options": caller_options
|
||||
}
|
||||
|
|
@ -7,7 +7,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.token_counter import token_counter
|
||||
from litellm.llms.openai.common_utils import BaseOpenAILLM
|
||||
from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base
|
||||
|
||||
# Test parameters for different API functions
|
||||
API_FUNCTION_PARAMS = [
|
||||
|
|
@ -392,3 +392,22 @@ async def test_async_genuine_bad_request_still_raises(provider, stream):
|
|||
|
||||
with pytest.raises(litellm.BadRequestError):
|
||||
await _call_and_drain()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "expected"),
|
||||
[
|
||||
("https://api.openai.com/v1", True),
|
||||
("https://api.openai.com:443/v1/", True),
|
||||
("https://southcentralus.privatelink.api.openai.com/v1", True),
|
||||
("https://eu.api.openai.com/v1", True),
|
||||
("HTTPS://API.OPENAI.COM/v1", True),
|
||||
("https://my-gateway.example/v1", False),
|
||||
("https://api.openai.com.evil.example/v1", False),
|
||||
("https://notapi.openai.com/v1", False),
|
||||
("https://gateway.example/v1?upstream=api.openai.com", False),
|
||||
("not a url", False),
|
||||
],
|
||||
)
|
||||
def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected):
|
||||
assert is_openai_backed_api_base(api_base) is expected
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@ import asyncio
|
|||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
|
||||
|
||||
|
||||
from litellm.proxy.auth.litellm_license import LicenseCheck
|
||||
|
||||
|
|
@ -30,3 +32,70 @@ def test_is_over_limit():
|
|||
assert license_check.is_over_limit(101) is False
|
||||
assert license_check.is_over_limit(100) is False
|
||||
assert license_check.is_over_limit(99) is False
|
||||
|
||||
|
||||
def test_heuristic_v2_router_limit() -> None:
|
||||
"""Only the signed license's auto_router feature lifts the one-router limit; an API-verified
|
||||
license (no airgapped data) and an airgapped license without the feature keep it."""
|
||||
license_check = LicenseCheck()
|
||||
license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]}
|
||||
assert license_check.heuristic_v2_router_limit() is None
|
||||
|
||||
license_check.airgapped_license_data = {
|
||||
"expiration_date": "2999-01-01",
|
||||
"allowed_features": ["sso", "auto_router", "audit_logs"],
|
||||
}
|
||||
assert license_check.heuristic_v2_router_limit() is None
|
||||
|
||||
license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]}
|
||||
assert license_check.heuristic_v2_router_limit() == 1
|
||||
|
||||
license_check.airgapped_license_data = {"expiration_date": "2999-01-01"}
|
||||
assert license_check.heuristic_v2_router_limit() == 1
|
||||
|
||||
license_check.airgapped_license_data = None
|
||||
assert license_check.heuristic_v2_router_limit() == 1
|
||||
|
||||
|
||||
def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]:
|
||||
import base64
|
||||
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
||||
|
||||
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
message = json.dumps(
|
||||
{"expiration_date": expiration_date, "user_id": "u", "allowed_features": ["auto_router"]}
|
||||
).encode()
|
||||
signature = private_key.sign(
|
||||
message,
|
||||
padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
|
||||
hashes.SHA256(),
|
||||
)
|
||||
return private_key.public_key(), base64.b64encode(message + b"." + signature).decode()
|
||||
|
||||
|
||||
def test_expired_or_unreadable_license_grants_no_features() -> None:
|
||||
"""The verifier stores the signed payload only after the expiry check passes and clears it when a
|
||||
later verify rejects the license, so a stale payload cannot keep lifting the heuristic_v2 limit."""
|
||||
license_check = LicenseCheck()
|
||||
public_key, valid_key = _signed_license("2999-01-01")
|
||||
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True
|
||||
assert license_check.heuristic_v2_router_limit() is None
|
||||
|
||||
_, expired_key = _signed_license("2000-01-01")
|
||||
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True
|
||||
assert license_check.airgapped_license_data is None
|
||||
assert license_check.heuristic_v2_router_limit() == 1
|
||||
|
||||
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True
|
||||
assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True
|
||||
assert license_check.airgapped_license_data is None
|
||||
|
||||
|
||||
def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None:
|
||||
license_check = LicenseCheck()
|
||||
public_key, license_key = _signed_license("2999-01-01")
|
||||
|
||||
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True
|
||||
assert license_check.heuristic_v2_router_limit() is None
|
||||
|
|
|
|||
|
|
@ -193,6 +193,33 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages(
|
|||
assert "headroom" in _applied_guardrails(request_data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_leaves_background_requests_uncompressed(
|
||||
guardrail: HeadroomGuardrail,
|
||||
):
|
||||
inputs = GenericGuardrailAPIInputs(
|
||||
texts=["A" * 5000],
|
||||
structured_messages=ORIGINAL_MESSAGES,
|
||||
)
|
||||
request_data = {"model": "gpt-4o", "background": True}
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"post",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_make_compress_response(COMPRESSED_MESSAGES),
|
||||
) as post:
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert result is inputs
|
||||
post.assert_not_awaited()
|
||||
assert _recorded_guardrail_entries(request_data) == []
|
||||
|
||||
|
||||
def _recorded_guardrail_response(request_data: dict) -> dict:
|
||||
entries = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(entries) == 1
|
||||
|
|
@ -954,6 +981,40 @@ async def test_passthrough_handler_does_not_log_headroom_as_run(
|
|||
assert "headroom" not in _applied_guardrails(data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstream(
|
||||
guardrail: HeadroomGuardrail,
|
||||
):
|
||||
"""Regression for LIT-6494: on /v1/responses the compressed messages must be
|
||||
written back into `input`, not only the retrieve tool into `tools`, or the
|
||||
model keeps reading the full document and never calls headroom_retrieve."""
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler
|
||||
|
||||
data = {
|
||||
"model": "gpt-5.6",
|
||||
"instructions": ORIGINAL_MESSAGES[0]["content"],
|
||||
"input": [{"role": m["role"], "content": m["content"]} for m in ORIGINAL_MESSAGES[1:]],
|
||||
"tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}],
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"post",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH),
|
||||
):
|
||||
result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"]
|
||||
assert [item["content"] for item in result["input"]] == [
|
||||
COMPRESSED_MESSAGES_WITH_HASH[0]["content"],
|
||||
ORIGINAL_MESSAGES[2]["content"],
|
||||
ORIGINAL_MESSAGES[3]["content"],
|
||||
]
|
||||
assert "A" * 5000 not in json.dumps(result["input"])
|
||||
assert [tool["name"] for tool in result["tools"]] == ["get_weather", HEADROOM_RETRIEVE_TOOL_NAME]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_http_error_raises():
|
||||
guardrail = _make_guardrail()
|
||||
|
|
@ -1950,6 +2011,58 @@ def _openai_text_payload(content: str) -> dict:
|
|||
return _openai_completion_payload({"role": "assistant", "content": content}, "stop")
|
||||
|
||||
|
||||
def _responses_retrieve_tool_definition() -> dict:
|
||||
return {"type": "function", **_retrieve_tool_definition()["function"]}
|
||||
|
||||
|
||||
def _openai_responses_payload(output_item: dict) -> dict:
|
||||
return {
|
||||
"id": "resp_ccr",
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o",
|
||||
"output": [output_item],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"error": None,
|
||||
"incomplete_details": None,
|
||||
"instructions": None,
|
||||
"metadata": {},
|
||||
"temperature": 1.0,
|
||||
"top_p": 1.0,
|
||||
"text": {"format": {"type": "text"}},
|
||||
"truncation": "disabled",
|
||||
}
|
||||
|
||||
|
||||
def _openai_responses_retrieve_call_payload() -> dict:
|
||||
return _openai_responses_payload(
|
||||
{
|
||||
"type": "function_call",
|
||||
"id": "fc_ccr",
|
||||
"call_id": "call_ccr",
|
||||
"name": HEADROOM_RETRIEVE_TOOL_NAME,
|
||||
"arguments": json.dumps({"hash": CCR_HASH}),
|
||||
"status": "completed",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _openai_responses_text_payload(text: str) -> dict:
|
||||
return _openai_responses_payload(
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_ccr",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": text, "annotations": []}],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"call_type, stream, tools, expect_conversion",
|
||||
[
|
||||
|
|
@ -1958,12 +2071,14 @@ def _openai_text_payload(content: str) -> dict:
|
|||
(CallTypes.acompletion, False, [_retrieve_tool_definition()], False),
|
||||
(CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False),
|
||||
(CallTypes.acompletion, True, None, False),
|
||||
(CallTypes.aresponses, True, [_retrieve_tool_definition()], False),
|
||||
(CallTypes.aresponses, True, [_retrieve_tool_definition()], True),
|
||||
(CallTypes.responses, True, [_responses_retrieve_tool_definition()], True),
|
||||
(CallTypes.aresponses, False, [_retrieve_tool_definition()], False),
|
||||
(CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions(
|
||||
async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses(
|
||||
guardrail: HeadroomGuardrail,
|
||||
call_type: CallTypes,
|
||||
stream: bool,
|
||||
|
|
@ -1986,6 +2101,22 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple
|
|||
assert kwargs["stream"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail):
|
||||
kwargs = {
|
||||
"model": "gpt-4o",
|
||||
"stream": True,
|
||||
"background": True,
|
||||
"tools": [_responses_retrieve_tool_definition()],
|
||||
}
|
||||
|
||||
result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.aresponses)
|
||||
|
||||
assert result is kwargs
|
||||
assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs
|
||||
assert kwargs["stream"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs(
|
||||
guardrail: HeadroomGuardrail,
|
||||
|
|
@ -2094,6 +2225,117 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end(
|
|||
assert not any(key.startswith("_headroom_interception") for key in followup_body)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_responses_resolves_ccr_retrieval_end_to_end(
|
||||
guardrail: HeadroomGuardrail,
|
||||
respx_mock: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Regression test for LIT-6481: streaming /v1/responses must resolve the
|
||||
retrieve tool call server-side exactly like streaming /chat/completions does,
|
||||
instead of streaming a headroom_retrieve function_call to the client."""
|
||||
original_content = "the full uncompressed document"
|
||||
final_answer = "the document says hello"
|
||||
guardrail._issued_hashes_by_call_id["ccr-call-id"] = (
|
||||
frozenset({CCR_HASH}),
|
||||
time.monotonic() + 999,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
upstream = respx_mock.post("https://api.openai.com/v1/responses").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json=_openai_responses_retrieve_call_payload()),
|
||||
httpx.Response(200, json=_openai_responses_text_payload(final_answer)),
|
||||
]
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"get",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_make_retrieve_response(original_content),
|
||||
) as mock_get:
|
||||
response = await litellm.aresponses(
|
||||
model="openai/gpt-4o",
|
||||
input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}],
|
||||
tools=[_responses_retrieve_tool_definition()],
|
||||
stream=True,
|
||||
litellm_call_id="ccr-call-id",
|
||||
)
|
||||
events = [event async for event in response]
|
||||
|
||||
streamed_text = "".join(
|
||||
getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta"
|
||||
)
|
||||
assert streamed_text == final_answer
|
||||
assert not any("function_call" in str(getattr(event, "type", "")) for event in events)
|
||||
assert not any(
|
||||
getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
|
||||
)
|
||||
mock_get.assert_called_once()
|
||||
assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0])
|
||||
|
||||
assert len(upstream.calls) == 2
|
||||
followup_body = json.loads(upstream.calls[1].request.content)
|
||||
assert not followup_body.get("stream")
|
||||
assert original_content in json.dumps(followup_body["input"])
|
||||
assert not any(key.startswith("_headroom_interception") for key in followup_body)
|
||||
|
||||
|
||||
def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end(
|
||||
guardrail: HeadroomGuardrail,
|
||||
respx_mock: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""The synchronous responses() path converts the stream the same way, so it
|
||||
must hand back a stream iterator with the resolved answer rather than the
|
||||
completed response object."""
|
||||
original_content = "the full uncompressed document"
|
||||
final_answer = "the document says hello"
|
||||
guardrail._issued_hashes_by_call_id["ccr-call-id"] = (
|
||||
frozenset({CCR_HASH}),
|
||||
time.monotonic() + 999,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
upstream = respx_mock.post("https://api.openai.com/v1/responses").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json=_openai_responses_retrieve_call_payload()),
|
||||
httpx.Response(200, json=_openai_responses_text_payload(final_answer)),
|
||||
]
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler,
|
||||
"get",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_make_retrieve_response(original_content),
|
||||
) as mock_get:
|
||||
response = litellm.responses(
|
||||
model="openai/gpt-4o",
|
||||
input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}],
|
||||
tools=[_responses_retrieve_tool_definition()],
|
||||
stream=True,
|
||||
litellm_call_id="ccr-call-id",
|
||||
)
|
||||
events = list(response)
|
||||
|
||||
streamed_text = "".join(
|
||||
getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta"
|
||||
)
|
||||
assert streamed_text == final_answer
|
||||
assert not any(
|
||||
getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
|
||||
)
|
||||
mock_get.assert_called_once()
|
||||
assert len(upstream.calls) == 2
|
||||
assert not json.loads(upstream.calls[1].request.content).get("stream")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LIT-5018: the turn the model is being asked to act on is never compressed.
|
||||
#
|
||||
|
|
|
|||
|
|
@ -17505,6 +17505,18 @@ def test_generate_key_request_blank_team_id_is_personal():
|
|||
assert GenerateKeyRequest(team_id="team-1").team_id == "team-1"
|
||||
|
||||
|
||||
def test_generate_key_request_blank_organization_and_project_id_are_unset():
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
|
||||
cleared = GenerateKeyRequest(organization_id="", project_id="")
|
||||
assert cleared.organization_id is None
|
||||
assert cleared.project_id is None
|
||||
assert "organization_id" not in cleared.model_dump(exclude_none=True)
|
||||
assert RegenerateKeyRequest(organization_id="").organization_id is None
|
||||
assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").organization_id == "org-1"
|
||||
assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").project_id == "proj-1"
|
||||
|
||||
|
||||
def test_key_request_blank_organization_id_is_unset():
|
||||
from litellm.proxy._types import RegenerateKeyRequest, UpdateKeyRequest
|
||||
|
||||
|
|
|
|||
|
|
@ -28,9 +28,18 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|||
delete_team_models,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.router import Router
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment
|
||||
|
||||
|
||||
async def _passthrough_row(update_data):
|
||||
return update_data
|
||||
|
||||
|
||||
async def _write_empty_row(**kwargs):
|
||||
return await kwargs["write_row"]({})
|
||||
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -1191,7 +1200,7 @@ class TestTeamModelSiblingRouting:
|
|||
team_id = "team_no_alias"
|
||||
public_name = "gpt-4.1-mini"
|
||||
|
||||
async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client):
|
||||
async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client, slot=None):
|
||||
return MagicMock(model_id=str(uuid.uuid4()))
|
||||
|
||||
mock_team_model_add = AsyncMock()
|
||||
|
|
@ -1372,7 +1381,8 @@ class TestTeamModelUpdate:
|
|||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client, # type: ignore
|
||||
prisma_client=prisma_client, # type: ignore,
|
||||
write_row=_passthrough_row,
|
||||
)
|
||||
|
||||
assert result.get("model_name", "").startswith("model_name_test_team_123_")
|
||||
|
|
@ -1435,7 +1445,6 @@ class TestTeamModelUpdate:
|
|||
team_id="team_123",
|
||||
public_model_name="new-public-name",
|
||||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client, # type: ignore
|
||||
)
|
||||
|
|
@ -1481,7 +1490,6 @@ class TestTeamModelUpdate:
|
|||
team_id="team_123",
|
||||
public_model_name="new-public-name",
|
||||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=None,
|
||||
)
|
||||
|
|
@ -1490,39 +1498,72 @@ class TestTeamModelUpdate:
|
|||
mock_delete.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_with_prisma_none_clears_patch_model_name(self):
|
||||
"""Rename path must clear patch_data.model_name even when prisma is unavailable (P1)."""
|
||||
async def test_a_refused_row_write_leaves_the_team_untouched(self):
|
||||
"""The team's model list autocommits, so it is written only after the row write succeeded: a
|
||||
refused write (the heuristic_v2 slot 403, a DB error) must not leave the team listing a name
|
||||
whose row never changed."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_update_existing_team_model_assignment,
|
||||
_update_team_model_in_db,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
db_model = Deployment(
|
||||
model_name="model_name_team_123_uuid1",
|
||||
model_name="gpt-4o",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"),
|
||||
model_info=ModelInfo(
|
||||
team_id="team_123", team_public_model_name="old-public-name"
|
||||
model_info=ModelInfo(),
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
events: list[str] = []
|
||||
written: dict[str, object] = {}
|
||||
|
||||
def patch_data() -> updateDeployment:
|
||||
return updateDeployment(model_name="team-public", model_info=ModelInfo(team_id="team_123"))
|
||||
|
||||
async def refuse_row(update_data):
|
||||
events.append("row")
|
||||
raise HTTPException(status_code=403, detail="slot held")
|
||||
|
||||
async def accept_row(update_data):
|
||||
events.append("row")
|
||||
written.update(update_data)
|
||||
return update_data
|
||||
|
||||
async def team_add(**_):
|
||||
events.append("team_model_add")
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the team auth check needs a live DB; the write order is what is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.allow_team_model_action",
|
||||
AsyncMock(return_value=True),
|
||||
),
|
||||
)
|
||||
patch_data = updateDeployment(
|
||||
model_name="new-public-name",
|
||||
model_info=ModelInfo(team_id="team_123"),
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="test_user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: team models are premium-gated through a proxy global with no injection seam
|
||||
patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
|
||||
side_effect=team_add,
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException):
|
||||
await _update_team_model_in_db(
|
||||
db_model=db_model,
|
||||
patch_data=patch_data(),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=MockPrismaClient(team_exists=True), # type: ignore
|
||||
write_row=refuse_row,
|
||||
)
|
||||
assert events == ["row"]
|
||||
|
||||
await _update_existing_team_model_assignment(
|
||||
team_id="team_123",
|
||||
public_model_name="new-public-name",
|
||||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=None,
|
||||
)
|
||||
|
||||
assert patch_data.model_name is None
|
||||
await _update_team_model_in_db(
|
||||
db_model=db_model,
|
||||
patch_data=patch_data(),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=MockPrismaClient(team_exists=True), # type: ignore
|
||||
write_row=accept_row,
|
||||
)
|
||||
assert events == ["row", "row", "team_model_add"]
|
||||
assert str(written["model_name"]).startswith("model_name_team_123_")
|
||||
assert "team-public" in str(written["model_info"])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_handles_legacy_string_model_info(self):
|
||||
|
|
@ -1574,7 +1615,6 @@ class TestTeamModelUpdate:
|
|||
team_id="team_123",
|
||||
public_model_name="new-public-name",
|
||||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client, # type: ignore
|
||||
)
|
||||
|
|
@ -1614,7 +1654,8 @@ class TestTeamModelUpdate:
|
|||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client, # type: ignore
|
||||
prisma_client=prisma_client, # type: ignore,
|
||||
write_row=_passthrough_row,
|
||||
)
|
||||
assert "403" in str(exc_info.value)
|
||||
|
||||
|
|
@ -1900,7 +1941,8 @@ class TestTeamModelUpdate:
|
|||
db_model=db_model,
|
||||
patch_data=patch_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client, # type: ignore
|
||||
prisma_client=prisma_client, # type: ignore,
|
||||
write_row=_passthrough_row,
|
||||
)
|
||||
|
||||
# team ACL must not be touched on a no-op edit
|
||||
|
|
@ -4311,6 +4353,321 @@ class TestStrategyRouterWriteValidation:
|
|||
is None
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}},
|
||||
{
|
||||
"model_name": "held-v2",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}},
|
||||
},
|
||||
"model_info": {"id": "held-id"},
|
||||
},
|
||||
],
|
||||
heuristic_v2_router_limit=lambda: limit,
|
||||
)
|
||||
|
||||
class _FakeTx:
|
||||
"""Stands in for a prisma transaction: records the raw statements and exposes the model table."""
|
||||
|
||||
def __init__(self, db_held: int) -> None:
|
||||
self.db_held = db_held
|
||||
self.raw_calls: list[tuple[str, tuple[object, ...]]] = []
|
||||
self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock())
|
||||
|
||||
async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]:
|
||||
self.raw_calls.append((sql, args))
|
||||
return [{"held": self.db_held}] if "count(*)" in sql else []
|
||||
|
||||
async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> None:
|
||||
return None
|
||||
|
||||
class _FakeDb:
|
||||
"""Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity."""
|
||||
|
||||
def __init__(self, db_held: int, existing_row: object = None) -> None:
|
||||
self.db = self
|
||||
self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held)
|
||||
self.litellm_proxymodeltable = MagicMock(
|
||||
create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row)
|
||||
)
|
||||
|
||||
def tx(self) -> "TestStrategyRouterWriteValidation._FakeTx":
|
||||
return self.tx_obj
|
||||
|
||||
_V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}
|
||||
_V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"incoming,existing,expected",
|
||||
[
|
||||
(_V2, None, _V2),
|
||||
(_V2, _V1, _V2),
|
||||
(None, _V1, _V1),
|
||||
(None, None, None),
|
||||
("no-config", _V2, _V2),
|
||||
],
|
||||
)
|
||||
def test_effective_complexity_router_config(
|
||||
self, incoming: object, existing: object, expected: object
|
||||
) -> None:
|
||||
"""A write is judged on the config it leaves on the row: the incoming one when it carries one, else the stored one."""
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_effective_complexity_router_config,
|
||||
)
|
||||
from litellm.types.router import updateLiteLLMParams
|
||||
|
||||
incoming_params = None if incoming is None else updateLiteLLMParams(
|
||||
complexity_router_config=None if incoming == "no-config" else incoming
|
||||
)
|
||||
existing_params = None if existing is None else updateLiteLLMParams(complexity_router_config=existing)
|
||||
assert _effective_complexity_router_config(incoming_params, existing_params) == expected
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"limit,effective_config,db_held,config_holds_one,model_id,expected",
|
||||
[
|
||||
(1, _V2, 1, False, None, "refused"),
|
||||
(1, _V2, 0, True, None, "refused"),
|
||||
(1, _V2, 0, False, None, "reserved"),
|
||||
(1, _V2, 0, False, "held-id", "reserved"),
|
||||
(2, _V2, 1, False, None, "reserved"),
|
||||
(1, _V1, 5, True, None, "plain"),
|
||||
(1, None, 5, True, None, "plain"),
|
||||
(None, _V2, 5, True, None, "plain"),
|
||||
],
|
||||
)
|
||||
async def test_heuristic_v2_slot_matrix(
|
||||
self,
|
||||
limit: int | None,
|
||||
effective_config: object,
|
||||
db_held: int,
|
||||
config_holds_one: bool,
|
||||
model_id: str | None,
|
||||
expected: str,
|
||||
) -> None:
|
||||
"""The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows
|
||||
(other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL
|
||||
parameter, and every other write runs on the plain client with no lock."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
HEURISTIC_V2_SLOT_LOCK_KEY,
|
||||
_heuristic_v2_slot,
|
||||
)
|
||||
|
||||
fake = self._FakeDb(db_held)
|
||||
live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam
|
||||
patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change",
|
||||
new=AsyncMock(),
|
||||
) as published,
|
||||
):
|
||||
if expected == "refused":
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id):
|
||||
pass
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "At most 1 auto-router" in str(exc_info.value.detail)
|
||||
assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail)
|
||||
return
|
||||
async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables:
|
||||
handle = tables
|
||||
if expected == "plain":
|
||||
await handle.create(data={})
|
||||
fake.litellm_proxymodeltable.create.assert_awaited_once_with(data={})
|
||||
assert fake.tx_obj.raw_calls == []
|
||||
return
|
||||
assert handle is fake.tx_obj.litellm_proxymodeltable
|
||||
published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable")
|
||||
(lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls
|
||||
assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql
|
||||
assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,)
|
||||
assert count_params == (model_id or "",)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None:
|
||||
"""team_model_add needs a second pool connection, so it must run only after the slot transaction
|
||||
(and its advisory lock) has closed; a pool-sized burst of team creates would otherwise stall on the
|
||||
lock holder waiting for a connection the waiters are occupying."""
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import _add_team_model_to_db
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
events: list[str] = []
|
||||
created = MagicMock(model_id="row-1")
|
||||
|
||||
@asynccontextmanager
|
||||
async def slot():
|
||||
events.append("slot-enter")
|
||||
yield MagicMock(create=AsyncMock(return_value=created))
|
||||
events.append("slot-exit")
|
||||
|
||||
async def team_model_add(**_: object) -> None:
|
||||
events.append("team_model_add")
|
||||
|
||||
deployment = Deployment(
|
||||
model_name="public-v2",
|
||||
litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2),
|
||||
model_info=ModelInfo(id="row-1", team_id="team-1"),
|
||||
)
|
||||
with (
|
||||
patch( # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
lambda value, new_encryption_key=None: value,
|
||||
),
|
||||
patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
|
||||
side_effect=team_model_add,
|
||||
),
|
||||
):
|
||||
result = await _add_team_model_to_db(
|
||||
model_params=deployment,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
prisma_client=MagicMock(),
|
||||
slot=slot(),
|
||||
)
|
||||
|
||||
assert result is created
|
||||
assert events == ["slot-enter", "slot-exit", "team_model_add"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_model_refuses_a_second_heuristic_v2_router_before_the_db_write(self) -> None:
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
add_new_model,
|
||||
)
|
||||
|
||||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
fake = self._FakeDb(db_held=1)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
lambda value, new_encryption_key=None: value,
|
||||
),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_new_model(
|
||||
model_params=Deployment(
|
||||
model_name="second-v2",
|
||||
litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2),
|
||||
),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
assert exc_info.value.code == "403"
|
||||
assert "At most 1 auto-router" in str(exc_info.value.message)
|
||||
fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited()
|
||||
fake.litellm_proxymodeltable.create.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None:
|
||||
"""patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
patch_model,
|
||||
)
|
||||
from litellm.types.router import updateLiteLLMParams
|
||||
|
||||
model_id = "other-id"
|
||||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
fake = self._FakeDb(db_held=1)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch( # test-quality-ok: the write must be refused before this DB step runs
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.get_db_model",
|
||||
new=AsyncMock(return_value=self._db_complexity_router(model_id)),
|
||||
),
|
||||
patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch( # test-quality-ok: the helper's team bookkeeping needs a live DB; the row writer it is handed is what is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db",
|
||||
new=AsyncMock(side_effect=_write_empty_row),
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await patch_model(
|
||||
model_id=model_id,
|
||||
patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=self._V2)),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
fake.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_model_refuses_switching_another_router_to_heuristic_v2(self) -> None:
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
update_model,
|
||||
)
|
||||
from litellm.types.router import ModelInfo, updateLiteLLMParams
|
||||
|
||||
model_id = "other-id"
|
||||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
existing_row = MagicMock()
|
||||
existing_row.model_dump.return_value = {
|
||||
"model_name": "my-auto-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}},
|
||||
},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"]
|
||||
fake = self._FakeDb(db_held=1, existing_row=existing_row)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await update_model(
|
||||
model_params=updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(complexity_router_config=self._V2),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
assert exc_info.value.code == "403"
|
||||
fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
fake.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_model_rejects_prefix_strip(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
|
|
|||
|
|
@ -45,6 +45,10 @@ from litellm.types.router import (
|
|||
from litellm.types.utils import Usage
|
||||
|
||||
|
||||
async def _passthrough_row(update_data):
|
||||
return update_data
|
||||
|
||||
|
||||
def test_model_info_accepts_valid_ptu_fields():
|
||||
info = ModelInfo(
|
||||
id="x",
|
||||
|
|
@ -385,6 +389,7 @@ class TestTeamModelUpdateValidatesBeforeWriting:
|
|||
patch_data=patch_data,
|
||||
user_api_key_dict=MagicMock(),
|
||||
prisma_client=MagicMock(),
|
||||
write_row=_passthrough_row,
|
||||
)
|
||||
return result, touched
|
||||
|
||||
|
|
@ -914,6 +919,7 @@ class TestPtuDeploymentsAreNotBilledPerToken:
|
|||
patch_data=patch,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
prisma_client=MagicMock(),
|
||||
write_row=_passthrough_row,
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 400
|
||||
|
|
|
|||
|
|
@ -9897,11 +9897,11 @@ class TestResolveTeamAccessGroupResources:
|
|||
assert resolved.access_group_mcp_server_ids == ["mcp-1"]
|
||||
assert resolved.access_group_agent_ids == ["agent-1"]
|
||||
assert [
|
||||
(d.access_group_id, d.access_group_name, d.models)
|
||||
(d.access_group_id, d.access_group_name, d.models, d.mcp_server_ids, d.agent_ids)
|
||||
for d in (resolved.access_group_details or [])
|
||||
] == [
|
||||
("ag-1", "shared-models", ("gpt-4", "claude-3")),
|
||||
("ag-2", "extra-models", ("claude-3", "gemini")),
|
||||
("ag-1", "shared-models", ("gpt-4", "claude-3"), ("mcp-1",), ()),
|
||||
("ag-2", "extra-models", ("claude-3", "gemini"), (), ("agent-1",)),
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import pytest
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_base_vertex_proxy_route,
|
||||
_upstream_headers_for_vertex_route,
|
||||
)
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
|
||||
|
|
@ -348,6 +349,93 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
|
|||
assert headers_passed_through is False
|
||||
|
||||
|
||||
VERTEX_ANTHROPIC_MODELS_PREFIX = "v1/projects/test-project/locations/global/publishers/anthropic/models/"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("model_segment", "expects_anthropic_beta"),
|
||||
[
|
||||
("count-tokens:rawPredict", False),
|
||||
("claude-sonnet-4-6:streamRawPredict", True),
|
||||
],
|
||||
)
|
||||
async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
|
||||
model_segment: str, expects_anthropic_beta: bool
|
||||
):
|
||||
with (
|
||||
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
|
||||
"litellm.proxy.proxy_server.llm_router", None
|
||||
),
|
||||
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch( # test-quality-ok: the route offers no injection point for its header preparation
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_prep_headers,
|
||||
patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
|
||||
) as mock_create_route,
|
||||
patch( # test-quality-ok: the route calls auth directly rather than through Depends
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_auth,
|
||||
patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
mock_prep_headers.return_value = (
|
||||
{
|
||||
"anthropic-beta": "tool-search-tool-2025-10-19,web-search-2025-03-05",
|
||||
"content-type": "application/json",
|
||||
"Authorization": "Bearer vertex-access-token",
|
||||
},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
"test-project",
|
||||
"global",
|
||||
)
|
||||
mock_create_route.return_value = AsyncMock()
|
||||
mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key")
|
||||
|
||||
await _base_vertex_proxy_route(
|
||||
endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
|
||||
request=MagicMock(),
|
||||
fastapi_response=MagicMock(),
|
||||
get_vertex_pass_through_handler=MagicMock(),
|
||||
)
|
||||
|
||||
upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
|
||||
assert ("anthropic-beta" in upstream_headers) is expects_anthropic_beta
|
||||
assert upstream_headers["Authorization"] == "Bearer vertex-access-token"
|
||||
assert upstream_headers["content-type"] == "application/json"
|
||||
|
||||
|
||||
def test_upstream_headers_for_vertex_route_filters_anthropic_beta_by_route():
|
||||
headers = {
|
||||
"Anthropic-Beta": "effort-2025-11-24",
|
||||
"content-type": "application/json",
|
||||
"Authorization": "Bearer vertex-access-token",
|
||||
}
|
||||
|
||||
count_tokens_headers = _upstream_headers_for_vertex_route(
|
||||
f"{VERTEX_ANTHROPIC_MODELS_PREFIX}count-tokens:rawPredict", headers
|
||||
)
|
||||
model_headers = _upstream_headers_for_vertex_route(
|
||||
f"{VERTEX_ANTHROPIC_MODELS_PREFIX}claude-sonnet-4-6:rawPredict", headers
|
||||
)
|
||||
|
||||
assert dict(count_tokens_headers) == {
|
||||
"content-type": "application/json",
|
||||
"Authorization": "Bearer vertex-access-token",
|
||||
}
|
||||
assert dict(model_headers) == headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.proxy.proxy_server import (
|
|||
resolve_routing_plugins,
|
||||
validate_deployment_complexity_router_placement,
|
||||
validate_deployment_max_agentic_loops,
|
||||
validate_heuristic_v2_router_limit,
|
||||
)
|
||||
|
||||
from .conftest import normalize
|
||||
|
|
@ -193,6 +194,120 @@ def test_validate_deployment_complexity_router_placement_leaves_valid_deployment
|
|||
assert model["litellm_params"] == litellm_params
|
||||
|
||||
|
||||
def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> dict[str, object]:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {"classifier_type": classifier_type, "tiers": {"SIMPLE": "gpt-4o-mini"}},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None:
|
||||
"""Same reason as the two validators above: the proxy router swallows registration errors, so
|
||||
an over-limit config.yaml must fail here instead of booting with a silently missing router."""
|
||||
with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info:
|
||||
validate_heuristic_v2_router_limit(
|
||||
[_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1
|
||||
)
|
||||
assert "'auto_router' feature lifts the limit" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_list,limit",
|
||||
[
|
||||
([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None),
|
||||
([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1),
|
||||
([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1),
|
||||
],
|
||||
)
|
||||
def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone(
|
||||
model_list: list[dict[str, object]], limit: int | None
|
||||
) -> None:
|
||||
assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None
|
||||
|
||||
|
||||
_TWO_HEURISTIC_V2_ROUTERS_YAML = (
|
||||
"model_list:\n"
|
||||
" - model_name: gpt-4o-mini\n"
|
||||
" litellm_params:\n"
|
||||
" model: openai/gpt-4o-mini\n"
|
||||
" api_key: k\n"
|
||||
" - model_name: v2-a\n"
|
||||
" litellm_params:\n"
|
||||
" model: auto_router/complexity_router\n"
|
||||
" complexity_router_config:\n"
|
||||
" classifier_type: heuristic_v2\n"
|
||||
" tiers: {SIMPLE: gpt-4o-mini}\n"
|
||||
" - model_name: v2-b\n"
|
||||
" litellm_params:\n"
|
||||
" model: auto_router/complexity_router\n"
|
||||
" complexity_router_config:\n"
|
||||
" classifier_type: heuristic_v2\n"
|
||||
" tiers: {SIMPLE: gpt-4o-mini}\n"
|
||||
"router_settings:\n"
|
||||
" heuristic_v2_router_limit: 99\n"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("license_limit", [1, None])
|
||||
async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only(
|
||||
tmp_path, monkeypatch, license_limit: int | None
|
||||
) -> None:
|
||||
"""`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator
|
||||
cannot grant the entitlement by editing the config, and a licensed proxy boots both routers."""
|
||||
f = tmp_path / "c.yaml"
|
||||
f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit
|
||||
)
|
||||
|
||||
if license_limit is None:
|
||||
router, _model_list, _general_settings = await ProxyConfig().load_config(
|
||||
router=None, config_file_path=str(f)
|
||||
)
|
||||
assert router.heuristic_v2_router_limit is not None
|
||||
assert router.heuristic_v2_router_limit() is None
|
||||
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
|
||||
return
|
||||
|
||||
with pytest.raises(ValueError, match=re.escape("config.yaml model_list: At most 1 auto-router")):
|
||||
await ProxyConfig().load_config(router=None, config_file_path=str(f))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_beyond_the_license(
|
||||
tmp_path, monkeypatch
|
||||
) -> None:
|
||||
"""config.yaml holds the one allowed heuristic_v2 router; a second one arriving later from the DB
|
||||
is refused at registration because the router was built with the license's ceiling."""
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
f = tmp_path / "c.yaml"
|
||||
f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML.replace(" - model_name: v2-b\n", " - model_name: v1-b\n", 1).replace(
|
||||
"classifier_type: heuristic_v2\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings",
|
||||
"classifier_type: heuristic\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings",
|
||||
))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1)
|
||||
|
||||
router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f))
|
||||
|
||||
assert router.heuristic_v2_router_limit is not None
|
||||
assert router.heuristic_v2_router_limit() == 1
|
||||
assert sorted(router.complexity_routers) == ["v1-b", "v2-a"]
|
||||
db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"})
|
||||
assert router.upsert_deployment(db_row) is None
|
||||
assert sorted(router.complexity_routers) == ["v1-b", "v2-a"]
|
||||
|
||||
|
||||
def test_validate_deployment_max_agentic_loops_allows_a_deployment_without_the_key():
|
||||
model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}
|
||||
|
||||
|
|
|
|||
|
|
@ -4234,6 +4234,91 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend():
|
|||
assert call_args[2] == [api_key]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens():
|
||||
"""
|
||||
Regression test for LIT-4929: the logs table showed the summed session cost but
|
||||
only the last call's token usage. Every row of a multi-round session must carry
|
||||
the session-wide prompt, completion and total token sums from the aggregate
|
||||
query, while rows outside a session carry none of them.
|
||||
"""
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
_build_ui_spend_logs_response,
|
||||
)
|
||||
|
||||
session_id = "sess-multi-round-tokens"
|
||||
api_key = "hashed-key-xyz"
|
||||
dict_rows = [
|
||||
{
|
||||
"request_id": "req-1",
|
||||
"session_id": session_id,
|
||||
"call_type": "completion",
|
||||
"api_key": api_key,
|
||||
"total_tokens": 10,
|
||||
"prompt_tokens": 7,
|
||||
"completion_tokens": 3,
|
||||
},
|
||||
{
|
||||
"request_id": "req-2",
|
||||
"session_id": session_id,
|
||||
"call_type": "completion",
|
||||
"api_key": api_key,
|
||||
"total_tokens": 50,
|
||||
"prompt_tokens": 35,
|
||||
"completion_tokens": 15,
|
||||
},
|
||||
{
|
||||
"request_id": "req-3",
|
||||
"session_id": None,
|
||||
"call_type": "completion",
|
||||
"api_key": api_key,
|
||||
"total_tokens": 5,
|
||||
"prompt_tokens": 4,
|
||||
"completion_tokens": 1,
|
||||
},
|
||||
]
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"session_id": session_id,
|
||||
"api_key": api_key,
|
||||
"session_total_count": 2,
|
||||
"session_total_spend": 0.06,
|
||||
"mcp_tool_call_count": 0,
|
||||
"mcp_tool_call_spend": 0.0,
|
||||
"session_total_prompt_tokens": 42,
|
||||
"session_total_completion_tokens": 18,
|
||||
"session_total_tokens": 60,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
result = await _build_ui_spend_logs_response(
|
||||
prisma_client=mock_prisma,
|
||||
data=dict_rows,
|
||||
total_records=3,
|
||||
page=1,
|
||||
page_size=50,
|
||||
total_pages=1,
|
||||
enrich_session_counts=True,
|
||||
)
|
||||
|
||||
rows = result["data"]
|
||||
session_rows = rows[:2]
|
||||
assert [row["session_total_tokens"] for row in session_rows] == [60, 60]
|
||||
assert [row["session_total_prompt_tokens"] for row in session_rows] == [42, 42]
|
||||
assert [row["session_total_completion_tokens"] for row in session_rows] == [18, 18]
|
||||
assert [(row["total_tokens"], row["prompt_tokens"], row["completion_tokens"]) for row in session_rows] == [
|
||||
(10, 7, 3),
|
||||
(50, 35, 15),
|
||||
]
|
||||
|
||||
token_keys = ("session_total_tokens", "session_total_prompt_tokens", "session_total_completion_tokens")
|
||||
assert all(key not in rows[2] for key in token_keys)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_ui_spend_logs_response_session_cache_hit_count():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -942,6 +942,47 @@ def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map(
|
|||
assert response["max_output_tokens"] == 8000
|
||||
|
||||
|
||||
def test_create_model_info_response_uses_deployment_mode_for_auto_router():
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "claude-sonnet",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-auto",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {
|
||||
"SIMPLE": "claude-sonnet",
|
||||
"MEDIUM": "claude-sonnet",
|
||||
"COMPLEX": "claude-sonnet",
|
||||
}
|
||||
},
|
||||
"complexity_router_default_model": "claude-sonnet",
|
||||
},
|
||||
"model_info": {
|
||||
"mode": "chat",
|
||||
"max_input_tokens": 1_000_000,
|
||||
"max_output_tokens": 128_000,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
response = create_model_info_response(
|
||||
model_id="claude-auto",
|
||||
provider="openai",
|
||||
llm_router=router,
|
||||
get_model_info=_raise_unmapped,
|
||||
)
|
||||
|
||||
assert response["mode"] == "chat"
|
||||
assert response["max_input_tokens"] == 1_000_000
|
||||
assert response["max_output_tokens"] == 128_000
|
||||
|
||||
|
||||
def test_create_model_info_response_deployment_limits_override_cost_map():
|
||||
router = MagicMock()
|
||||
router.get_configured_token_limits.return_value = (200000, None)
|
||||
|
|
|
|||
|
|
@ -120,9 +120,7 @@ async def test_delete_vector_store_checks_access():
|
|||
"team_id": "team_456",
|
||||
}
|
||||
)
|
||||
mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(
|
||||
return_value=mock_vector_store
|
||||
)
|
||||
mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=mock_vector_store)
|
||||
|
||||
# User from different team should get 403
|
||||
user_api_key_dict = UserAPIKeyAuth(team_id="team_789")
|
||||
|
|
@ -134,9 +132,115 @@ async def test_delete_vector_store_checks_access():
|
|||
):
|
||||
with patch("litellm.vector_store_registry", None):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await delete_vector_store(
|
||||
data=request, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
await delete_vector_store(data=request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "Access denied" in exc_info.value.detail
|
||||
|
||||
|
||||
_UNSCOPED: LiteLLM_ManagedVectorStore = {
|
||||
"vector_store_id": "vs_unscoped",
|
||||
"custom_llm_provider": "openai",
|
||||
"team_id": None,
|
||||
}
|
||||
_TEAM_A_OWNED: LiteLLM_ManagedVectorStore = {
|
||||
"vector_store_id": "vs_team_a",
|
||||
"custom_llm_provider": "openai",
|
||||
"team_id": "team_a",
|
||||
}
|
||||
_UI_CREATED: LiteLLM_ManagedVectorStore = {
|
||||
"vector_store_id": "vs_ui_created",
|
||||
"custom_llm_provider": "openai",
|
||||
"team_id": "litellm-dashboard",
|
||||
}
|
||||
|
||||
|
||||
async def _listed_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str]:
|
||||
from litellm.proxy.vector_store_endpoints.management_endpoints import (
|
||||
list_vector_stores,
|
||||
)
|
||||
|
||||
with patch( # test-quality-ok: the list route reads rows through this module-level DB helper, no injection seam
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db",
|
||||
new=AsyncMock(return_value=[_UNSCOPED, _TEAM_A_OWNED, _UI_CREATED]),
|
||||
):
|
||||
response = await list_vector_stores(user_api_key_dict=user_api_key_dict)
|
||||
return sorted(vs["vector_store_id"] for vs in response["data"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_vector_stores_hides_ungranted_stores_from_non_admin_keys():
|
||||
"""A store with no team_id and no allowlist entry is not listed for a key it was never granted to;
|
||||
only team ownership or an explicit object_permission grant makes a store visible."""
|
||||
assert await _listed_ids(UserAPIKeyAuth()) == []
|
||||
assert await _listed_ids(UserAPIKeyAuth(team_id="team_a")) == ["vs_team_a"]
|
||||
assert await _listed_ids(
|
||||
UserAPIKeyAuth(
|
||||
team_id="team_b",
|
||||
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", vector_stores=["vs_unscoped"]),
|
||||
)
|
||||
) == ["vs_unscoped"]
|
||||
assert await _listed_ids(
|
||||
UserAPIKeyAuth(
|
||||
team_id="team_b",
|
||||
team_object_permission=LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="op-2", vector_stores=["vs_unscoped"]
|
||||
),
|
||||
)
|
||||
) == ["vs_unscoped"]
|
||||
assert await _listed_ids(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) == [
|
||||
"vs_team_a",
|
||||
"vs_ui_created",
|
||||
"vs_unscoped",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("user_team_ids", "session_key_grants", "expected"),
|
||||
[
|
||||
([], None, []),
|
||||
([], ["vs_unscoped"], ["vs_unscoped"]),
|
||||
(["team_a"], None, ["vs_team_a"]),
|
||||
(["team_a", "team_granted"], None, ["vs_team_a", "vs_unscoped"]),
|
||||
],
|
||||
)
|
||||
async def test_list_vector_stores_dashboard_session_resolves_real_teams(
|
||||
user_team_ids: list[str], session_key_grants: list[str] | None, expected: list[str]
|
||||
):
|
||||
"""A dashboard session lists through the user's real teams plus the session key's own grants: stores created
|
||||
from the dashboard (team_id litellm-dashboard) are not visible just because every session shares that team id,
|
||||
while stores owned by or granted to one of the user's teams, or granted to the session key itself, are."""
|
||||
from litellm.models.team import LiteLLM_TeamTableCachedObj
|
||||
|
||||
alice = UserAPIKeyAuth(
|
||||
team_id="litellm-dashboard",
|
||||
user_id="alice",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
object_permission=(
|
||||
LiteLLM_ObjectPermissionTable(object_permission_id="op-4", vector_stores=session_key_grants)
|
||||
if session_key_grants is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
teams = {
|
||||
"team_a": LiteLLM_TeamTableCachedObj(team_id="team_a"),
|
||||
"team_granted": LiteLLM_TeamTableCachedObj(
|
||||
team_id="team_granted",
|
||||
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-3", vector_stores=["vs_unscoped"]),
|
||||
),
|
||||
}
|
||||
|
||||
async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj:
|
||||
return teams[team_id]
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: team rows come from the module-level prisma client, no injection seam
|
||||
"litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object
|
||||
),
|
||||
patch( # test-quality-ok: the user row comes from the module-level prisma client, no injection seam
|
||||
"litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids",
|
||||
new=AsyncMock(return_value=user_team_ids),
|
||||
),
|
||||
):
|
||||
assert await _listed_ids(alice) == expected
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import time
|
||||
from types import TracebackType
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
|
|
@ -294,3 +295,39 @@ async def test_azure_health_check_honors_deployment_realtime_protocol():
|
|||
model_params={"realtime_protocol": "GA"},
|
||||
)
|
||||
assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
|
||||
|
||||
class _ConnectThatStopsAfterCapturingTheUrl:
|
||||
url: str | None = None
|
||||
|
||||
def __call__(self, url: str, **kwargs: object) -> "_ConnectThatStopsAfterCapturingTheUrl":
|
||||
self.url = url
|
||||
return self
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
raise RuntimeError("backend url captured, nothing to bridge")
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai_realtime_route():
|
||||
connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
|
||||
with patch("websockets.connect", connect):
|
||||
await realtime_main._arealtime.__wrapped__(
|
||||
model="azure_ai/gpt-realtime-mini",
|
||||
websocket=MagicMock(),
|
||||
api_base="https://my-project.services.ai.azure.com",
|
||||
api_key="fake-key",
|
||||
litellm_logging_obj=FakeLogging(),
|
||||
)
|
||||
assert connect.url == (
|
||||
"wss://my-project.services.ai.azure.com/openai/realtime"
|
||||
"?api-version=2024-10-01-preview&deployment=gpt-realtime-mini"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.types.utils import (
|
|||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Usage,
|
||||
)
|
||||
|
||||
CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256"
|
||||
|
|
@ -527,6 +528,20 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing():
|
|||
assert response_ids[0].startswith("resp_")
|
||||
|
||||
|
||||
def test_completed_event_restores_usage_hidden_by_stream_options_none():
|
||||
final_chunk = _chunk("", finish_reason="stop")
|
||||
final_chunk._hidden_params = {"usage": Usage(prompt_tokens=117, completion_tokens=5, total_tokens=122)}
|
||||
iterator = _build_iterator([_chunk("the document says hello"), final_chunk])
|
||||
|
||||
events = list(iterator)
|
||||
|
||||
completed = next(
|
||||
event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
||||
)
|
||||
assert completed.response.usage.input_tokens == 117
|
||||
assert completed.response.usage.output_tokens == 5
|
||||
|
||||
|
||||
def test_object_tool_call_arguments_stream_as_valid_json():
|
||||
"""A provider that sends decoded object arguments must still stream valid JSON.
|
||||
|
||||
|
|
|
|||
|
|
@ -330,36 +330,46 @@ ROUTER_CONFIG: Final = json.dumps(
|
|||
)
|
||||
|
||||
|
||||
class FailingRouteLayer:
|
||||
"""Route layer whose embedding call fails, as it does when the prompt exceeds the encoder's window."""
|
||||
|
||||
def __call__(self, text: str) -> Any:
|
||||
raise ValueError(
|
||||
"Internal_litellm_router API call failed. Error: litellm.InternalServerError: "
|
||||
"input is too large to process. increase the physical batch size"
|
||||
)
|
||||
|
||||
|
||||
class FixedRouteLayer:
|
||||
"""Route layer that returns whatever the test tells it to, recording the text it was asked about."""
|
||||
"""Route layer that returns whatever the test tells it to for the query vector it is handed."""
|
||||
|
||||
def __init__(self, route_choice: Any) -> None:
|
||||
self.route_choice = route_choice
|
||||
self.seen_text: str | None = None
|
||||
|
||||
def __call__(self, text: str) -> Any:
|
||||
self.seen_text = text
|
||||
async def acall(self, vector: Any) -> Any:
|
||||
return self.route_choice
|
||||
|
||||
|
||||
def _embedding_response(input: List[str]) -> Any:
|
||||
import litellm
|
||||
|
||||
return litellm.EmbeddingResponse(
|
||||
data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))]
|
||||
)
|
||||
|
||||
|
||||
class StubEmbeddingRouter:
|
||||
"""Stands in for the LiteLLM Router when the route index has to be built for real."""
|
||||
"""Stands in for the LiteLLM Router, recording the text and kwargs each query embedding was made with."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.seen_text: str | None = None
|
||||
self.aembedding_kwargs: Dict[str, Any] | None = None
|
||||
|
||||
def embedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
|
||||
import litellm
|
||||
return _embedding_response(input)
|
||||
|
||||
return litellm.EmbeddingResponse(
|
||||
data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))]
|
||||
async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
|
||||
self.seen_text = input[0]
|
||||
self.aembedding_kwargs = kwargs
|
||||
return _embedding_response(input)
|
||||
|
||||
|
||||
class FailingEmbeddingRouter(StubEmbeddingRouter):
|
||||
"""Router whose query embedding fails, as it does when the prompt exceeds the encoder's window."""
|
||||
|
||||
async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any:
|
||||
raise ValueError(
|
||||
"litellm.InternalServerError: input is too large to process. increase the physical batch size"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -369,7 +379,7 @@ def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs:
|
|||
auto_router_config=ROUTER_CONFIG,
|
||||
default_model="fallback-model",
|
||||
embedding_model="text-embedding-3-small",
|
||||
litellm_router_instance=litellm_router_instance or MagicMock(),
|
||||
litellm_router_instance=litellm_router_instance or StubEmbeddingRouter(),
|
||||
**kwargs,
|
||||
)
|
||||
auto_router.routelayer = routelayer
|
||||
|
|
@ -381,7 +391,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_fall_back_to_default_model_when_the_embedding_call_fails(self):
|
||||
auto_router: Final = _auto_router(FailingRouteLayer())
|
||||
auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=FailingEmbeddingRouter())
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
|
|
@ -440,8 +450,8 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
|
|||
async def test_should_still_route_to_the_matched_route_when_one_matches(self):
|
||||
from semantic_router.schema import RouteChoice
|
||||
|
||||
layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
|
||||
auto_router: Final = _auto_router(layer)
|
||||
router: Final = StubEmbeddingRouter()
|
||||
auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
|
|
@ -451,7 +461,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel:
|
|||
|
||||
assert result is not None
|
||||
assert result.model == "code-model"
|
||||
assert layer.seen_text == "fix this stack trace"
|
||||
assert router.seen_text == "fix this stack trace"
|
||||
|
||||
|
||||
class TestAutoRouterEmbeddingInputCap:
|
||||
|
|
@ -483,8 +493,8 @@ class TestAutoRouterRoutesResponsesApiInput:
|
|||
async def test_should_route_a_string_input_when_messages_is_none(self):
|
||||
from semantic_router.schema import RouteChoice
|
||||
|
||||
layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
|
||||
auto_router: Final = _auto_router(layer)
|
||||
router: Final = StubEmbeddingRouter()
|
||||
auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
|
|
@ -498,14 +508,14 @@ class TestAutoRouterRoutesResponsesApiInput:
|
|||
assert result is not None
|
||||
assert result.model == "code-model"
|
||||
assert result.messages is None
|
||||
assert layer.seen_text == "fix this stack trace"
|
||||
assert router.seen_text == "fix this stack trace"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_route_a_list_input_with_instructions_when_messages_is_none(self):
|
||||
from semantic_router.schema import RouteChoice
|
||||
|
||||
layer: Final = FixedRouteLayer(RouteChoice(name="code-model"))
|
||||
auto_router: Final = _auto_router(layer)
|
||||
router: Final = StubEmbeddingRouter()
|
||||
auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router)
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
|
|
@ -525,13 +535,13 @@ class TestAutoRouterRoutesResponsesApiInput:
|
|||
|
||||
assert result is not None
|
||||
assert result.model == "code-model"
|
||||
assert layer.seen_text is not None
|
||||
assert "fix this stack trace" in layer.seen_text
|
||||
assert router.seen_text is not None
|
||||
assert "fix this stack trace" in router.seen_text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_routing_when_neither_messages_nor_input_is_present(self):
|
||||
layer: Final = FixedRouteLayer(None)
|
||||
auto_router: Final = _auto_router(layer)
|
||||
router: Final = StubEmbeddingRouter()
|
||||
auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router)
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
|
|
@ -540,12 +550,12 @@ class TestAutoRouterRoutesResponsesApiInput:
|
|||
)
|
||||
|
||||
assert result is None
|
||||
assert layer.seen_text is None
|
||||
assert router.seen_text is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_keep_routing_an_empty_messages_list_to_the_default_model(self):
|
||||
layer: Final = FixedRouteLayer(None)
|
||||
auto_router: Final = _auto_router(layer)
|
||||
router: Final = StubEmbeddingRouter()
|
||||
auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router)
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
|
|
@ -555,4 +565,42 @@ class TestAutoRouterRoutesResponsesApiInput:
|
|||
|
||||
assert result is not None
|
||||
assert result.model == "fallback-model"
|
||||
assert layer.seen_text == ""
|
||||
assert router.seen_text == ""
|
||||
|
||||
|
||||
class TestAutoRouterAttributesItsEmbeddingSpend:
|
||||
"""The query embedding is billed to the key that sent the request, like any other call it made."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_forward_the_callers_identity_to_the_query_embedding_minus_its_budget_reservation(self):
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
|
||||
router: Final = StubEmbeddingRouter()
|
||||
auto_router: Final = _auto_router(None, litellm_router_instance=router)
|
||||
request_kwargs: Final = {
|
||||
"metadata": {
|
||||
"user_api_key": "hashed-key",
|
||||
"user_api_key_team_id": "team-1",
|
||||
"user_api_key_budget_reservation": {"reservation_id": "r-1"},
|
||||
},
|
||||
"litellm_session_id": "session-1",
|
||||
}
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "fix this stack trace"}],
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert router.seen_text == "fix this stack trace"
|
||||
assert router.aembedding_kwargs is not None
|
||||
forwarded: Final = router.aembedding_kwargs["metadata"]
|
||||
assert forwarded["user_api_key"] == "hashed-key"
|
||||
assert forwarded["user_api_key_team_id"] == "team-1"
|
||||
assert forwarded[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "autorouter_classifier"
|
||||
assert "user_api_key_budget_reservation" not in forwarded
|
||||
assert router.aembedding_kwargs["litellm_session_id"] == "session-1"
|
||||
assert router.aembedding_kwargs["proxy_server_request"] == {
|
||||
"body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from pydantic import ValidationError
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
|
||||
|
|
@ -1085,6 +1086,165 @@ class TestRouterComplexityDeploymentMethods:
|
|||
router.init_complexity_router_deployment(deployment)
|
||||
assert "auto_router/complexity_router/test-router" in router.complexity_routers
|
||||
|
||||
@staticmethod
|
||||
def _router_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"classifier_type": classifier_type,
|
||||
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
|
||||
},
|
||||
},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
|
||||
_POOL: dict[str, object] = {
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"},
|
||||
}
|
||||
|
||||
def test_heuristic_v2_ceiling_keeps_the_first_router_and_drops_the_rest(self) -> None:
|
||||
"""The proxy runs with ignore_invalid_deployments, so the second heuristic_v2 router is dropped
|
||||
at registration while a heuristic (v1) sibling and the first v2 router stay routable."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
self._POOL,
|
||||
self._router_row("v2-a", "id-a", "heuristic_v2"),
|
||||
self._router_row("v2-b", "id-b", "heuristic_v2"),
|
||||
self._router_row("v1-c", "id-c", "heuristic"),
|
||||
],
|
||||
heuristic_v2_router_limit=lambda: 1,
|
||||
ignore_invalid_deployments=True,
|
||||
)
|
||||
|
||||
assert sorted(router.complexity_routers) == ["v1-c", "v2-a"]
|
||||
assert router.get_deployment(model_id="id-b") is None
|
||||
|
||||
def test_heuristic_v2_ceiling_raises_without_ignore_invalid_deployments(self) -> None:
|
||||
with pytest.raises(ValueError, match="At most 1 auto-router"):
|
||||
Router(
|
||||
model_list=[
|
||||
self._POOL,
|
||||
self._router_row("v2-a", "id-a", "heuristic_v2"),
|
||||
self._router_row("v2-b", "id-b", "heuristic_v2"),
|
||||
],
|
||||
heuristic_v2_router_limit=lambda: 1,
|
||||
)
|
||||
|
||||
def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None:
|
||||
"""The Router never caches the limit: when the resolver's answer moves (the proxy re-verified
|
||||
its license), the next registration and the next limit query see the new value."""
|
||||
limits = {"value": None}
|
||||
router = Router(
|
||||
model_list=[
|
||||
self._POOL,
|
||||
self._router_row("v2-a", "id-a", "heuristic_v2"),
|
||||
self._router_row("v2-b", "id-b", "heuristic_v2"),
|
||||
],
|
||||
heuristic_v2_router_limit=lambda: limits["value"],
|
||||
ignore_invalid_deployments=True,
|
||||
)
|
||||
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
|
||||
assert router.heuristic_v2_router_limit_violation() is None
|
||||
|
||||
limits["value"] = 1
|
||||
assert router.heuristic_v2_router_limit_violation() is not None
|
||||
assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None
|
||||
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
|
||||
|
||||
def test_heuristic_v2_ceiling_tightening_refuses_the_edit_and_keeps_the_live_router(self) -> None:
|
||||
"""Two heuristic_v2 routers registered under an unlimited ceiling, then the ceiling drops to one:
|
||||
an edit to either must be refused before its live row is popped, or the failed re-add and
|
||||
the failed restore would drop a serving router while the write reports success."""
|
||||
limits = {"value": None}
|
||||
router = Router(
|
||||
model_list=[
|
||||
self._POOL,
|
||||
self._router_row("v2-a", "id-a", "heuristic_v2"),
|
||||
self._router_row("v2-b", "id-b", "heuristic_v2"),
|
||||
],
|
||||
heuristic_v2_router_limit=lambda: limits["value"],
|
||||
ignore_invalid_deployments=True,
|
||||
)
|
||||
limits["value"] = 1
|
||||
|
||||
assert router.upsert_deployment(Deployment(**self._router_row("v2-a-renamed", "id-a", "heuristic_v2"))) is None
|
||||
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
|
||||
assert router.get_deployment(model_id="id-a") is not None
|
||||
|
||||
assert router.upsert_deployment(Deployment(**self._router_row("v1-a", "id-a", "heuristic"))) is not None
|
||||
assert sorted(router.complexity_routers) == ["v1-a", "v2-b"]
|
||||
|
||||
def test_config_deployments_excludes_db_rows(self) -> None:
|
||||
"""The proxy counts config.yaml routers from here and DB rows from the database, so a DB-loaded
|
||||
row (``model_info.db_model``) must not show up twice."""
|
||||
router = Router(model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")])
|
||||
db_row = self._router_row("v2-db", "id-db", "heuristic_v2")
|
||||
db_row["model_info"] = {"id": "id-db", "db_model": True}
|
||||
assert router.upsert_deployment(Deployment(**db_row)) is not None
|
||||
|
||||
assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"]
|
||||
assert count_heuristic_v2_routers(router.config_deployments()) == 1
|
||||
|
||||
def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None:
|
||||
"""A rollback after a failed upsert re-admits state that was already serving, so it must not be
|
||||
judged by a ceiling that tightened since: converting one of two live heuristic_v2 routers to a
|
||||
config whose registration fails must leave it serving its previous v2 configuration."""
|
||||
limits = {"value": None}
|
||||
router = Router(
|
||||
model_list=[
|
||||
self._POOL,
|
||||
self._router_row("v2-a", "id-a", "heuristic_v2"),
|
||||
self._router_row("v2-b", "id-b", "heuristic_v2"),
|
||||
],
|
||||
heuristic_v2_router_limit=lambda: limits["value"],
|
||||
ignore_invalid_deployments=True,
|
||||
)
|
||||
limits["value"] = 1
|
||||
|
||||
broken = self._router_row("v1-a", "id-a", "heuristic")
|
||||
broken["litellm_params"]["complexity_router_config"]["tiers"] = {}
|
||||
assert router.upsert_deployment(Deployment(**broken)) is None
|
||||
|
||||
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
|
||||
live = router.get_deployment(model_id="id-a")
|
||||
assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2"
|
||||
assert router.heuristic_v2_router_limit_violation() is not None
|
||||
|
||||
def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None:
|
||||
router = Router(
|
||||
model_list=[
|
||||
self._POOL,
|
||||
self._router_row("v2-a", "id-a", "heuristic_v2"),
|
||||
self._router_row("v2-b", "id-b", "heuristic_v2"),
|
||||
]
|
||||
)
|
||||
|
||||
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
|
||||
assert router.heuristic_v2_router_limit_violation() is None
|
||||
|
||||
def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None:
|
||||
"""A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot
|
||||
while a different deployment switching to heuristic_v2 is refused."""
|
||||
router = Router(
|
||||
model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")],
|
||||
heuristic_v2_router_limit=lambda: 1,
|
||||
ignore_invalid_deployments=True,
|
||||
)
|
||||
|
||||
assert router.heuristic_v2_router_limit_violation() is not None
|
||||
|
||||
edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2")
|
||||
assert router.upsert_deployment(Deployment(**edited)) is not None
|
||||
assert sorted(router.complexity_routers) == ["v2-a-renamed"]
|
||||
|
||||
assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None
|
||||
assert sorted(router.complexity_routers) == ["v2-a-renamed"]
|
||||
assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None
|
||||
assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"]
|
||||
|
||||
def test_hybrid_initialization_waits_for_later_pool_deployments(self):
|
||||
router = Router(
|
||||
model_list=[
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
import asyncio
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import json
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
||||
|
|
@ -102,11 +103,10 @@ async def test_async_user_key_affinity_routes_to_same_deployment():
|
|||
|
||||
# Deterministic routing: first selection uses seq[0], second selection attempts seq[1]
|
||||
# unless the list has been filtered to length=1 by deployment affinity.
|
||||
choice_calls = {"count": 0}
|
||||
choice_calls: Final = itertools.count(1)
|
||||
|
||||
def deterministic_choice(seq):
|
||||
choice_calls["count"] += 1
|
||||
if choice_calls["count"] == 1:
|
||||
def deterministic_choice(seq: Sequence[dict[str, object]]) -> dict[str, object]:
|
||||
if next(choice_calls) == 1:
|
||||
return seq[0]
|
||||
return seq[1] if len(seq) > 1 else seq[0]
|
||||
|
||||
|
|
@ -998,3 +998,212 @@ async def test_model_group_affinity_config_overrides_global():
|
|||
)
|
||||
# All deployments returned (user-key affinity disabled for this group)
|
||||
assert len(filtered) == 2
|
||||
|
||||
|
||||
def _jwt_metadata(user_id: str) -> dict[str, str | None]:
|
||||
return {"user_api_key_hash": None, "user_api_key_user_id": user_id}
|
||||
|
||||
|
||||
def _two_deployments(model_group: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"model_name": model_group,
|
||||
"litellm_params": {"model": "openai/gpt-5.4-mini"},
|
||||
"model_info": {"id": "openai-deployment-a"},
|
||||
},
|
||||
{
|
||||
"model_name": model_group,
|
||||
"litellm_params": {"model": "openai/gpt-5.4-mini"},
|
||||
"model_info": {"id": "openai-deployment-b"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_jwt_user_affinity_routes_to_same_deployment():
|
||||
"""
|
||||
JWT-authenticated proxy requests carry no `user_api_key_hash`, only `user_api_key_user_id`.
|
||||
They must still pin to one deployment per user.
|
||||
"""
|
||||
model_group = "gpt-5.4-mini"
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": model_group,
|
||||
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-a"},
|
||||
"model_info": {"id": "openai-deployment-a"},
|
||||
},
|
||||
{
|
||||
"model_name": model_group,
|
||||
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-b"},
|
||||
"model_info": {"id": "openai-deployment-b"},
|
||||
},
|
||||
],
|
||||
optional_pre_call_checks=["deployment_affinity"],
|
||||
)
|
||||
|
||||
choice_calls = {"count": 0}
|
||||
|
||||
def deterministic_choice(seq):
|
||||
choice_calls["count"] += 1
|
||||
if choice_calls["count"] == 1:
|
||||
return seq[0]
|
||||
return seq[1] if len(seq) > 1 else seq[0]
|
||||
|
||||
with patch( # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy
|
||||
"litellm.router_strategy.simple_shuffle.random.choice",
|
||||
side_effect=deterministic_choice,
|
||||
):
|
||||
first_response = await router.acompletion(
|
||||
model=model_group,
|
||||
messages=[{"role": "user", "content": "Reply with the single word ok"}],
|
||||
mock_response="ok",
|
||||
metadata=_jwt_metadata("jwt-user-alice"),
|
||||
)
|
||||
second_response = await router.acompletion(
|
||||
model=model_group,
|
||||
messages=[{"role": "user", "content": "Reply with the single word ok"}],
|
||||
mock_response="ok",
|
||||
metadata=_jwt_metadata("jwt-user-alice"),
|
||||
)
|
||||
|
||||
first_model_id = first_response._hidden_params["model_id"]
|
||||
assert first_model_id in ("openai-deployment-a", "openai-deployment-b")
|
||||
assert second_response._hidden_params["model_id"] == first_model_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_jwt_auth_metadata_pins_per_user():
|
||||
"""
|
||||
The metadata the proxy stamps for a JWT caller (`UserAPIKeyAuth(api_key=None, user_id=<sub>)`)
|
||||
must claim a pin and be read back by the filter, and another JWT user must not inherit it.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
|
||||
model_group = "gpt-5.4-mini"
|
||||
healthy_deployments = _two_deployments(model_group)
|
||||
callback = DeploymentAffinityCheck(
|
||||
cache=DualCache(),
|
||||
ttl_seconds=60,
|
||||
enable_user_key_affinity=True,
|
||||
enable_responses_api_affinity=False,
|
||||
)
|
||||
|
||||
def proxy_request(user_id: str) -> dict[str, object]:
|
||||
return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
|
||||
data={"model": model_group, "messages": [{"role": "user", "content": "hi"}], "metadata": {}},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key=None, user_id=user_id),
|
||||
_metadata_variable_name="metadata",
|
||||
)
|
||||
|
||||
alice_request = proxy_request("jwt-user-alice")
|
||||
alice_metadata = alice_request["metadata"]
|
||||
assert isinstance(alice_metadata, dict)
|
||||
assert alice_metadata["user_api_key_hash"] is None
|
||||
|
||||
await callback.async_pre_call_deployment_hook(
|
||||
kwargs={
|
||||
**alice_request,
|
||||
"metadata": {**alice_metadata, "deployment_model_name": model_group},
|
||||
"model_info": {"id": "openai-deployment-b"},
|
||||
},
|
||||
call_type=None,
|
||||
)
|
||||
|
||||
alice_pinned = await callback.async_filter_deployments(
|
||||
model=model_group,
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=None,
|
||||
request_kwargs=alice_request,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
assert [deployment["model_info"]["id"] for deployment in alice_pinned] == ["openai-deployment-b"]
|
||||
|
||||
bob_filtered = await callback.async_filter_deployments(
|
||||
model=model_group,
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=None,
|
||||
request_kwargs=proxy_request("jwt-user-bob"),
|
||||
parent_otel_span=None,
|
||||
)
|
||||
assert bob_filtered == healthy_deployments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jwt_user_id_never_reads_a_virtual_key_pin():
|
||||
"""
|
||||
A JWT user id that happens to equal a virtual key's 64-hex hash must not read that key's pin.
|
||||
"""
|
||||
model_group = "gpt-5.4-mini"
|
||||
healthy_deployments = _two_deployments(model_group)
|
||||
callback = DeploymentAffinityCheck(
|
||||
cache=DualCache(),
|
||||
ttl_seconds=60,
|
||||
enable_user_key_affinity=True,
|
||||
enable_responses_api_affinity=False,
|
||||
)
|
||||
key_hash = "a" * 64
|
||||
|
||||
await callback.async_pre_call_deployment_hook(
|
||||
kwargs={
|
||||
"metadata": {"user_api_key_hash": key_hash, "deployment_model_name": model_group},
|
||||
"model_info": {"id": "openai-deployment-b"},
|
||||
},
|
||||
call_type=None,
|
||||
)
|
||||
|
||||
key_pinned = await callback.async_filter_deployments(
|
||||
model=model_group,
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=None,
|
||||
request_kwargs={"metadata": {"user_api_key_hash": key_hash}},
|
||||
parent_otel_span=None,
|
||||
)
|
||||
assert [deployment["model_info"]["id"] for deployment in key_pinned] == ["openai-deployment-b"]
|
||||
|
||||
lookalike_jwt_user = await callback.async_filter_deployments(
|
||||
model=model_group,
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=None,
|
||||
request_kwargs={"metadata": _jwt_metadata(key_hash)},
|
||||
parent_otel_span=None,
|
||||
)
|
||||
assert lookalike_jwt_user == healthy_deployments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_virtual_key_hash_wins_over_user_id_for_affinity():
|
||||
"""
|
||||
A virtual-key caller with a user id pins on the key hash, so two keys owned by one user
|
||||
keep independent pins.
|
||||
"""
|
||||
model_group = "gpt-5.4-mini"
|
||||
healthy_deployments = _two_deployments(model_group)
|
||||
callback = DeploymentAffinityCheck(
|
||||
cache=DualCache(),
|
||||
ttl_seconds=60,
|
||||
enable_user_key_affinity=True,
|
||||
enable_responses_api_affinity=False,
|
||||
)
|
||||
|
||||
await callback.async_pre_call_deployment_hook(
|
||||
kwargs={
|
||||
"metadata": {
|
||||
"user_api_key_hash": "key-one",
|
||||
"user_api_key_user_id": "shared-user",
|
||||
"deployment_model_name": model_group,
|
||||
},
|
||||
"model_info": {"id": "openai-deployment-b"},
|
||||
},
|
||||
call_type=None,
|
||||
)
|
||||
|
||||
other_key_same_user = await callback.async_filter_deployments(
|
||||
model=model_group,
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=None,
|
||||
request_kwargs={"metadata": {"user_api_key_hash": "key-two", "user_api_key_user_id": "shared-user"}},
|
||||
parent_otel_span=None,
|
||||
)
|
||||
assert other_key_same_user == healthy_deployments
|
||||
|
|
|
|||
|
|
@ -1,8 +1,13 @@
|
|||
from collections.abc import Mapping
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_utils.auto_router_model_naming import (
|
||||
carries_complexity_router_settings,
|
||||
classify_strategy_router_model,
|
||||
count_heuristic_v2_routers,
|
||||
heuristic_v2_limit_violation,
|
||||
is_heuristic_v2_router,
|
||||
strategy_router_dependencies,
|
||||
validate_complexity_router_config_placement,
|
||||
validate_complexity_router_config_write,
|
||||
|
|
@ -369,3 +374,55 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie
|
|||
flat param on an s3_vectors vector store, so an unscoped gate would reject a valid deployment.
|
||||
Either complexity field names one on its own, which is what the load itself requires."""
|
||||
assert carries_complexity_router_settings(model, present_fields) is scoped
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_params,expected",
|
||||
[
|
||||
({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True),
|
||||
({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True),
|
||||
({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False),
|
||||
({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False),
|
||||
({"model": "auto_router/complexity_router"}, False),
|
||||
({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False),
|
||||
({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False),
|
||||
({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False),
|
||||
({}, False),
|
||||
],
|
||||
)
|
||||
def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None:
|
||||
"""Only a complexity router whose config selects heuristic_v2 counts toward the license limit."""
|
||||
assert is_heuristic_v2_router(litellm_params) is expected
|
||||
|
||||
|
||||
def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None:
|
||||
v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}
|
||||
rows: list[Mapping[str, object]] = [
|
||||
{"model_name": "a", "litellm_params": v2},
|
||||
{"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}},
|
||||
{"model_name": "c", "litellm_params": v2},
|
||||
{"model_name": "d"},
|
||||
{"model_name": "e", "litellm_params": "not a mapping"},
|
||||
]
|
||||
assert count_heuristic_v2_routers(rows) == 2
|
||||
assert count_heuristic_v2_routers(()) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"held,limit,violates",
|
||||
[
|
||||
(1, 1, False),
|
||||
(2, 1, True),
|
||||
(0, 1, False),
|
||||
(5, None, False),
|
||||
(3, 3, False),
|
||||
(4, 3, True),
|
||||
],
|
||||
)
|
||||
def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None:
|
||||
violation = heuristic_v2_limit_violation(held=held, limit=limit)
|
||||
assert (violation is not None) is violates
|
||||
if violation is not None:
|
||||
assert f"At most {limit} auto-router" in violation
|
||||
assert f"would make {held}" in violation
|
||||
assert "license" not in violation
|
||||
|
|
|
|||
|
|
@ -371,3 +371,20 @@ class TestKimiK3AdvertisesItsDocumentedLevels:
|
|||
"low",
|
||||
"high",
|
||||
)
|
||||
|
||||
|
||||
class TestGpt6AstraAdvertisesItsDocumentedLevels:
|
||||
def test_the_entry_advertises_low_through_max_without_none(self, local_model_cost_map):
|
||||
"""OpenAI documents low, medium, high, xhigh and max for gpt-6-astra. Unlike gpt-5.6-sol it
|
||||
does not take none, so a group must not offer none and must offer max."""
|
||||
from litellm.utils import _get_model_info_helper
|
||||
|
||||
model_info = dict(_get_model_info_helper(model="gpt-6-astra", custom_llm_provider="openai"))
|
||||
|
||||
assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == (
|
||||
"low",
|
||||
"medium",
|
||||
"high",
|
||||
"xhigh",
|
||||
"max",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4478,3 +4478,17 @@ def test_explicit_pricing_precedes_private_provider_response_model(
|
|||
)
|
||||
|
||||
assert selected == expected
|
||||
|
||||
|
||||
def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_model_cost_map):
|
||||
"""gpt-6-astra batch pricing is 50% off the standard $10 input and $50 output rates per 1M tokens."""
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
|
||||
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage, model="gpt-6-astra", custom_llm_provider="openai"
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(1000 * 5e-6)
|
||||
assert completion_cost == pytest.approx(500 * 2.5e-5)
|
||||
|
|
|
|||
|
|
@ -47,8 +47,7 @@ def test_azure_ai_gpt_5_5_model_info(model):
|
|||
|
||||
routed_model, provider, _, _ = get_llm_provider(model=model)
|
||||
assert routed_model == model.split("/", 1)[1]
|
||||
# azure_ai/* models resolve under the azure provider in get_llm_provider
|
||||
assert provider == "azure"
|
||||
assert provider == "azure_ai"
|
||||
|
||||
|
||||
def test_azure_ai_gpt_5_5_backup_matches_main():
|
||||
|
|
|
|||
|
|
@ -784,6 +784,19 @@ def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_respo
|
|||
assert model_info.get("mode") == "responses"
|
||||
|
||||
|
||||
def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses():
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
||||
model_info, model = responses_api_bridge_check(
|
||||
model="gpt-6-astra",
|
||||
custom_llm_provider="openai",
|
||||
tools=[{"type": "function", "function": {"name": "get_capital"}}],
|
||||
)
|
||||
|
||||
assert model == "gpt-6-astra"
|
||||
assert model_info.get("mode") == "responses"
|
||||
|
||||
|
||||
def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses():
|
||||
"""gpt-5.5+ with both tools and reasoning_effort should route to Responses API."""
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
|
@ -1133,6 +1146,82 @@ def test_responses_api_bridge_check_custom_api_base_via_env_with_unset_effort_st
|
|||
assert model_info.get("mode") != "responses"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
"https://southcentralus.privatelink.api.openai.com/v1",
|
||||
"https://privatelink.corp.api.openai.com/v1",
|
||||
"https://api.openai.com:443/v1",
|
||||
"https://api.openai.com/v1/",
|
||||
"HTTPS://API.OPENAI.COM/v1",
|
||||
],
|
||||
)
|
||||
def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_effort_routes_to_responses(api_base):
|
||||
"""
|
||||
A custom api_base whose host is api.openai.com or a subdomain of it (a PrivateLink hostname, a
|
||||
port-qualified or trailing-slash default) still reaches the real OpenAI backend, which rejects
|
||||
function tools with reasoning on Chat Completions, so the unset-effort arm must bridge exactly as
|
||||
it does for the literal default URL. Regression guard for GH #39353.
|
||||
"""
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
||||
model_info, model = responses_api_bridge_check(
|
||||
model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
tools=[{"type": "function", "function": {"name": "get_capital"}}],
|
||||
reasoning_effort=None,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
assert model == "gpt-5.6"
|
||||
assert model_info.get("mode") == "responses"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
"https://api.openai.com.evil.example/v1",
|
||||
"https://notapi.openai.com/v1",
|
||||
"https://gateway.example/v1?upstream=api.openai.com",
|
||||
"https://openai.internal.example/api.openai.com/v1",
|
||||
],
|
||||
)
|
||||
def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_stays_chat(api_base):
|
||||
"""Only the host decides: api.openai.com appearing elsewhere in the URL is still a foreign backend."""
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
||||
model_info, model = responses_api_bridge_check(
|
||||
model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
tools=[{"type": "function", "function": {"name": "get_capital"}}],
|
||||
reasoning_effort=None,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
assert model == "gpt-5.6"
|
||||
assert model_info.get("mode") != "responses"
|
||||
|
||||
|
||||
def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effort_routes_to_responses(monkeypatch):
|
||||
"""A PrivateLink base set through OPENAI_BASE_URL resolves the way the chat handler's does and still bridges."""
|
||||
import litellm
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
||||
monkeypatch.setattr(litellm, "api_base", None)
|
||||
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
|
||||
monkeypatch.setenv("OPENAI_BASE_URL", "https://southcentralus.privatelink.api.openai.com/v1")
|
||||
model_info, model = responses_api_bridge_check(
|
||||
model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
tools=[{"type": "function", "function": {"name": "get_capital"}}],
|
||||
reasoning_effort=None,
|
||||
api_base=None,
|
||||
)
|
||||
|
||||
assert model == "gpt-5.6"
|
||||
assert model_info.get("mode") == "responses"
|
||||
|
||||
|
||||
def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_routes():
|
||||
"""Explicit reasoning_effort keeps its pre-existing bridging behavior on any api_base."""
|
||||
from litellm.main import responses_api_bridge_check
|
||||
|
|
@ -3259,3 +3348,43 @@ def test_stream_chunk_builder_leaves_xai_reported_cost_to_the_calculator(monkeyp
|
|||
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)
|
||||
|
||||
|
||||
FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com"
|
||||
|
||||
|
||||
def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route(
|
||||
respx_mock: respx.MockRouter,
|
||||
):
|
||||
route: Final = respx_mock.post(
|
||||
url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+"
|
||||
).mock(return_value=httpx.Response(200, json={"text": "hello"}))
|
||||
|
||||
response: Final = litellm.transcription(
|
||||
model="azure_ai/whisper-1",
|
||||
file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"),
|
||||
api_base=FOUNDRY_HOST,
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.text == "hello"
|
||||
|
||||
|
||||
def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route(
|
||||
respx_mock: respx.MockRouter,
|
||||
):
|
||||
route: Final = respx_mock.post(
|
||||
url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+"
|
||||
).mock(return_value=httpx.Response(200, content=b"mp3-bytes"))
|
||||
|
||||
response: Final = litellm.speech(
|
||||
model="azure_ai/tts-1",
|
||||
input="hello",
|
||||
voice="alloy",
|
||||
api_base=FOUNDRY_HOST,
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.content == b"mp3-bytes"
|
||||
|
|
|
|||
|
|
@ -52,6 +52,12 @@ PRIORITY_LONG_CONTEXT = {
|
|||
"cache_read_input_token_cost_above_272k_tokens_priority": 8e-08,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06,
|
||||
},
|
||||
"gpt-6-astra": {
|
||||
"input_cost_per_token_above_272k_tokens_priority": 4e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 0.00015,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 4e-06,
|
||||
"cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05,
|
||||
},
|
||||
}
|
||||
|
||||
EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT}
|
||||
|
|
@ -63,6 +69,7 @@ NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5")
|
|||
def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
||||
litellm.add_known_models()
|
||||
|
||||
|
||||
@lru_cache(maxsize=2)
|
||||
|
|
@ -114,6 +121,7 @@ TIERED_COST_CASES = [
|
|||
("gpt-5.6-sol", "priority", 1.6e-05, 6e-05),
|
||||
("gpt-5.6-terra", "priority", 8e-06, 3.6e-05),
|
||||
("gpt-5.6-luna", "priority", 8e-07, 3.6e-06),
|
||||
("gpt-6-astra", "priority", 4e-05, 0.00015),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12553,3 +12553,33 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo
|
|||
bucket = captured.get("litellm_metadata") or captured["metadata"]
|
||||
assert captured["model_info"]["id"] == "provisional-dep"
|
||||
assert bucket["litellm_gateway_injected_cache"] == ""
|
||||
|
||||
|
||||
def test_get_configured_mode_reads_deployment_model_info():
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "my-tts",
|
||||
"litellm_params": {"model": "openai/some-unmapped-mode-model"},
|
||||
"model_info": {"mode": "audio_speech"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert router.get_configured_mode("my-tts") == "audio_speech"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_info", [{}, {"mode": ""}, {"mode": " "}, {"mode": 123}])
|
||||
def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info):
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "plain-model",
|
||||
"litellm_params": {"model": "openai/some-unmapped-mode-model"},
|
||||
"model_info": model_info,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert router.get_configured_mode("plain-model") is None
|
||||
assert router.get_configured_mode("unknown-model") is None
|
||||
|
|
|
|||
|
|
@ -58,6 +58,17 @@ describe("OrganizationDropdown", () => {
|
|||
expect(onChange.mock.calls[0][0]).toBe("org-1");
|
||||
});
|
||||
|
||||
it("emits null, never the empty string, when the selection is cleared", async () => {
|
||||
const onChange = vi.fn();
|
||||
const user = userEvent.setup();
|
||||
render(<OrganizationDropdown organizations={MOCK_ORGS} value="org-1" onChange={onChange} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Clear" }));
|
||||
|
||||
expect(onChange).toHaveBeenCalledTimes(1);
|
||||
expect(onChange).toHaveBeenCalledWith(null);
|
||||
});
|
||||
|
||||
it("should filter options by organization id", async () => {
|
||||
const user = userEvent.setup();
|
||||
render(<OrganizationDropdown organizations={MOCK_ORGS} />);
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import { Organization } from "../networking";
|
|||
interface OrganizationDropdownProps {
|
||||
organizations?: Organization[] | null;
|
||||
value?: string;
|
||||
onChange?: (value: string) => void;
|
||||
onChange?: (value: string | null) => void;
|
||||
disabled?: boolean;
|
||||
loading?: boolean;
|
||||
style?: React.CSSProperties;
|
||||
|
|
@ -32,7 +32,7 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
|
|||
sublabel: org.organization_id,
|
||||
}))}
|
||||
value={value}
|
||||
onValueChange={(organizationId) => onChange?.(organizationId)}
|
||||
onValueChange={(organizationId) => onChange?.(organizationId || null)}
|
||||
placeholder={placeholder}
|
||||
emptyText={loading ? "Loading organizations…" : "No organizations found"}
|
||||
disabled={disabled}
|
||||
|
|
|
|||
|
|
@ -3,9 +3,12 @@ import VectorStorePermissions from "./permissions/VectorStorePermissions";
|
|||
import MCPServerPermissions from "./permissions/MCPServerPermissions";
|
||||
import AgentPermissions from "./permissions/AgentPermissions";
|
||||
import type { ObjectPermission } from "./object_permission_types";
|
||||
import type { InheritedGrant } from "./permissions/inheritedGrants";
|
||||
|
||||
interface ObjectPermissionsViewProps {
|
||||
objectPermission?: ObjectPermission | null;
|
||||
inheritedMcpServers?: InheritedGrant[];
|
||||
inheritedAgents?: InheritedGrant[];
|
||||
variant?: "card" | "inline";
|
||||
className?: string;
|
||||
accessToken?: string | null;
|
||||
|
|
@ -13,6 +16,8 @@ interface ObjectPermissionsViewProps {
|
|||
|
||||
export function ObjectPermissionsView({
|
||||
objectPermission,
|
||||
inheritedMcpServers = [],
|
||||
inheritedAgents = [],
|
||||
variant = "card",
|
||||
className = "",
|
||||
accessToken,
|
||||
|
|
@ -34,9 +39,15 @@ export function ObjectPermissionsView({
|
|||
mcpAccessGroups={mcpAccessGroups}
|
||||
mcpToolPermissions={mcpToolPermissions}
|
||||
mcpToolsets={mcpToolsets}
|
||||
inheritedMcpServers={inheritedMcpServers}
|
||||
accessToken={accessToken}
|
||||
/>
|
||||
<AgentPermissions
|
||||
agents={agents}
|
||||
agentAccessGroups={agentAccessGroups}
|
||||
inheritedAgents={inheritedAgents}
|
||||
accessToken={accessToken}
|
||||
/>
|
||||
<AgentPermissions agents={agents} agentAccessGroups={agentAccessGroups} accessToken={accessToken} />
|
||||
<div className="min-w-0 rounded-md border border-border p-4">
|
||||
<p className="text-sm font-medium text-foreground">Search tools</p>
|
||||
{searchTools.length === 0 ? (
|
||||
|
|
|
|||
|
|
@ -587,9 +587,9 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
}
|
||||
};
|
||||
|
||||
const changeOrganization = (write: FieldWrite) => (orgId: string) => {
|
||||
write(orgId || undefined);
|
||||
setSelectedOrganizationId(orgId || null);
|
||||
const changeOrganization = (write: FieldWrite) => (orgId: string | null) => {
|
||||
write(orgId ?? undefined);
|
||||
setSelectedOrganizationId(orgId);
|
||||
// Clear team and project when org changes
|
||||
setSelectedCreateKeyTeam(null);
|
||||
setSelectedProjectId(null);
|
||||
|
|
|
|||
|
|
@ -0,0 +1,71 @@
|
|||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import AgentPermissions from "./AgentPermissions";
|
||||
import * as networking from "../networking";
|
||||
|
||||
vi.mock("../networking");
|
||||
|
||||
describe("AgentPermissions", () => {
|
||||
const accessToken = "test-token";
|
||||
const agentId = "90337622-756e-4f25-98f0-01fc8174aa24";
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("lists agents inherited from access groups, counts them, and names the groups on hover", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(networking.getAgentsList).mockResolvedValue({
|
||||
agents: [{ agent_id: agentId, agent_name: "support_agent" }],
|
||||
});
|
||||
|
||||
render(
|
||||
<AgentPermissions
|
||||
agents={[]}
|
||||
inheritedAgents={[{ id: agentId, accessGroupNames: ["platform-tools", "support"] }]}
|
||||
accessToken={accessToken}
|
||||
/>,
|
||||
);
|
||||
|
||||
const row = await screen.findByText(/support_agent/);
|
||||
expect(screen.getByText("1")).toBeInTheDocument();
|
||||
expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument();
|
||||
expect(networking.getAgentsList).toHaveBeenCalledWith(accessToken);
|
||||
|
||||
await user.hover(row);
|
||||
expect(
|
||||
await screen.findByText(`Granted via access groups platform-tools, support. Full ID: ${agentId}`),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not double-list an agent that is both granted directly and inherited", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(networking.getAgentsList).mockResolvedValue({
|
||||
agents: [{ agent_id: agentId, agent_name: "support_agent" }],
|
||||
});
|
||||
|
||||
render(
|
||||
<AgentPermissions
|
||||
agents={[agentId]}
|
||||
inheritedAgents={[{ id: agentId, accessGroupNames: ["support"] }]}
|
||||
accessToken={accessToken}
|
||||
/>,
|
||||
);
|
||||
|
||||
const row = await screen.findByText(/support_agent/);
|
||||
expect(screen.getAllByText(/support_agent/)).toHaveLength(1);
|
||||
expect(screen.getByText("1")).toBeInTheDocument();
|
||||
|
||||
await user.hover(row);
|
||||
expect(await screen.findByText(`Full ID: ${agentId}`)).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows the empty state when nothing is granted directly or inherited", () => {
|
||||
render(<AgentPermissions agents={[]} inheritedAgents={[]} accessToken={accessToken} />);
|
||||
|
||||
expect(screen.getByText("No agents or access groups configured")).toBeInTheDocument();
|
||||
expect(screen.getByText("0")).toBeInTheDocument();
|
||||
expect(networking.getAgentsList).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -3,6 +3,7 @@ import { UserGroupIcon } from "@heroicons/react/outline";
|
|||
import { Badge } from "@/components/ui/badge";
|
||||
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
|
||||
import { getAgentsList } from "../networking";
|
||||
import { InheritedGrant, inheritedGrantTooltip } from "./inheritedGrants";
|
||||
|
||||
interface Agent {
|
||||
agent_id: string;
|
||||
|
|
@ -14,16 +15,24 @@ interface Agent {
|
|||
interface AgentPermissionsProps {
|
||||
agents: string[];
|
||||
agentAccessGroups?: string[];
|
||||
inheritedAgents?: InheritedGrant[];
|
||||
accessToken?: string | null;
|
||||
}
|
||||
|
||||
export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }: AgentPermissionsProps) {
|
||||
export function AgentPermissions({
|
||||
agents,
|
||||
agentAccessGroups = [],
|
||||
inheritedAgents = [],
|
||||
accessToken,
|
||||
}: AgentPermissionsProps) {
|
||||
const [agentDetails, setAgentDetails] = useState<Agent[]>([]);
|
||||
const inheritedOnlyAgents = inheritedAgents.filter((grant) => !agents.includes(grant.id));
|
||||
const agentIdCount = agents.length + inheritedOnlyAgents.length;
|
||||
|
||||
// Fetch agent details when component mounts
|
||||
useEffect(() => {
|
||||
const fetchAgentDetails = async () => {
|
||||
if (accessToken && agents.length > 0) {
|
||||
if (accessToken && agentIdCount > 0) {
|
||||
try {
|
||||
const response = await getAgentsList(accessToken);
|
||||
if (response && response.agents && Array.isArray(response.agents)) {
|
||||
|
|
@ -35,7 +44,7 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
|
|||
}
|
||||
};
|
||||
fetchAgentDetails();
|
||||
}, [accessToken, agents.length]);
|
||||
}, [accessToken, agentIdCount]);
|
||||
|
||||
// Function to get display name for agent
|
||||
const getAgentDisplayName = (agentId: string) => {
|
||||
|
|
@ -47,10 +56,10 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
|
|||
return agentId;
|
||||
};
|
||||
|
||||
// Merge agents and access groups into one list
|
||||
const mergedItems = [
|
||||
...agents.map((agent) => ({ type: "agent", value: agent })),
|
||||
...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group })),
|
||||
...agents.map((agent) => ({ type: "agent", value: agent, tooltip: `Full ID: ${agent}` })),
|
||||
...inheritedOnlyAgents.map((grant) => ({ type: "agent", value: grant.id, tooltip: inheritedGrantTooltip(grant) })),
|
||||
...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group, tooltip: "" })),
|
||||
];
|
||||
const totalCount = mergedItems.length;
|
||||
|
||||
|
|
@ -77,7 +86,7 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }
|
|||
{getAgentDisplayName(item.value)}
|
||||
</span>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>{`Full ID: ${item.value}`}</TooltipContent>
|
||||
<TooltipContent>{item.tooltip}</TooltipContent>
|
||||
</Tooltip>
|
||||
</TooltipProvider>
|
||||
) : (
|
||||
|
|
|
|||
|
|
@ -406,4 +406,55 @@ describe("MCPServerPermissions", () => {
|
|||
);
|
||||
await waitFor(() => expect(screen.getByText("Blocked")).toHaveAttribute("data-variant", "destructive"));
|
||||
});
|
||||
|
||||
it("lists servers inherited from access groups, counts them, and names the group on hover", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(networking.fetchMCPServers).mockResolvedValue([
|
||||
{ server_id: mockServerId1, server_name: mockServerName1, alias: mockServerName1 },
|
||||
]);
|
||||
|
||||
render(
|
||||
<MCPServerPermissions
|
||||
mcpServers={[]}
|
||||
mcpAccessGroups={[]}
|
||||
mcpToolPermissions={{}}
|
||||
inheritedMcpServers={[{ id: mockServerId1, accessGroupNames: ["platform-tools"] }]}
|
||||
accessToken={mockAccessToken}
|
||||
/>,
|
||||
);
|
||||
|
||||
const row = await screen.findByText(/DW_MCP/);
|
||||
expect(screen.getByText("1")).toBeInTheDocument();
|
||||
expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument();
|
||||
expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken);
|
||||
|
||||
await user.hover(row);
|
||||
expect(
|
||||
await screen.findByText(`Granted via access group platform-tools. Full ID: ${mockServerId1}`),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not double-list a server that is both granted directly and inherited", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(networking.fetchMCPServers).mockResolvedValue([
|
||||
{ server_id: mockServerId2, server_name: mockServerName2, alias: mockServerName2 },
|
||||
]);
|
||||
|
||||
render(
|
||||
<MCPServerPermissions
|
||||
mcpServers={[mockServerId2]}
|
||||
mcpAccessGroups={[]}
|
||||
mcpToolPermissions={{}}
|
||||
inheritedMcpServers={[{ id: mockServerId2, accessGroupNames: ["platform-tools"] }]}
|
||||
accessToken={mockAccessToken}
|
||||
/>,
|
||||
);
|
||||
|
||||
const row = await screen.findByText(/Test Server/);
|
||||
expect(screen.getAllByText(/Test Server/)).toHaveLength(1);
|
||||
expect(screen.getByText("1")).toBeInTheDocument();
|
||||
|
||||
await user.hover(row);
|
||||
expect(await screen.findByText(`Full ID: ${mockServerId2}`)).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -5,12 +5,14 @@ import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip
|
|||
import { fetchMCPServers, fetchMCPToolsets } from "../networking";
|
||||
import { MCPServer, MCPToolset } from "../mcp_tools/types";
|
||||
import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants";
|
||||
import { InheritedGrant, inheritedGrantTooltip } from "./inheritedGrants";
|
||||
|
||||
interface MCPServerPermissionsProps {
|
||||
mcpServers: string[];
|
||||
mcpAccessGroups?: string[];
|
||||
mcpToolPermissions?: Record<string, string[]>;
|
||||
mcpToolsets?: string[];
|
||||
inheritedMcpServers?: InheritedGrant[];
|
||||
accessToken?: string | null;
|
||||
}
|
||||
|
||||
|
|
@ -19,6 +21,7 @@ export function MCPServerPermissions({
|
|||
mcpAccessGroups = [],
|
||||
mcpToolPermissions = {},
|
||||
mcpToolsets = [],
|
||||
inheritedMcpServers = [],
|
||||
accessToken,
|
||||
}: MCPServerPermissionsProps) {
|
||||
const [mcpServerDetails, setMCPServerDetails] = useState<MCPServer[]>([]);
|
||||
|
|
@ -50,10 +53,16 @@ export function MCPServerPermissions({
|
|||
});
|
||||
};
|
||||
|
||||
const directServerIds = mcpServers.filter(
|
||||
(server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL,
|
||||
);
|
||||
const inheritedOnlyServers = inheritedMcpServers.filter((grant) => !mcpServers.includes(grant.id));
|
||||
const serverIdCount = directServerIds.length + inheritedOnlyServers.length;
|
||||
|
||||
// Fetch MCP server details when component mounts
|
||||
useEffect(() => {
|
||||
const fetchMCPServerDetails = async () => {
|
||||
if (accessToken && mcpServers.length > 0) {
|
||||
if (accessToken && serverIdCount > 0) {
|
||||
try {
|
||||
const response = await fetchMCPServers(accessToken);
|
||||
if (response && Array.isArray(response)) {
|
||||
|
|
@ -67,7 +76,7 @@ export function MCPServerPermissions({
|
|||
}
|
||||
};
|
||||
fetchMCPServerDetails();
|
||||
}, [accessToken, mcpServers.length]);
|
||||
}, [accessToken, serverIdCount]);
|
||||
|
||||
// Fetch toolset details
|
||||
useEffect(() => {
|
||||
|
|
@ -98,12 +107,14 @@ export function MCPServerPermissions({
|
|||
const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL);
|
||||
const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL);
|
||||
|
||||
// Merge servers and access groups into one list
|
||||
const mergedItems = [
|
||||
...mcpServers
|
||||
.filter((server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL)
|
||||
.map((server) => ({ type: "server", value: server })),
|
||||
...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })),
|
||||
...directServerIds.map((server) => ({ type: "server", value: server, tooltip: `Full ID: ${server}` })),
|
||||
...inheritedOnlyServers.map((grant) => ({
|
||||
type: "server",
|
||||
value: grant.id,
|
||||
tooltip: inheritedGrantTooltip(grant),
|
||||
})),
|
||||
...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group, tooltip: "" })),
|
||||
];
|
||||
const totalCount = mergedItems.length + mcpToolsets.length;
|
||||
|
||||
|
|
@ -153,7 +164,7 @@ export function MCPServerPermissions({
|
|||
{getMCPServerDisplayName(item.value)}
|
||||
</span>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>{`Full ID: ${item.value}`}</TooltipContent>
|
||||
<TooltipContent>{item.tooltip}</TooltipContent>
|
||||
</Tooltip>
|
||||
) : (
|
||||
<div className="inline-flex items-center gap-2 min-w-0">
|
||||
|
|
|
|||
|
|
@ -0,0 +1,61 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import { computeInheritedGrants, inheritedGrantTooltip } from "./inheritedGrants";
|
||||
import { TeamAccessGroupModelGrant } from "../team/teamModelAccess";
|
||||
|
||||
const GRANTS: TeamAccessGroupModelGrant[] = [
|
||||
{ access_group_id: "ag-1", access_group_name: "platform-tools", models: [], mcp_server_ids: ["mcp-1", "mcp-2"] },
|
||||
{
|
||||
access_group_id: "ag-2",
|
||||
access_group_name: "support",
|
||||
models: [],
|
||||
mcp_server_ids: ["mcp-2"],
|
||||
agent_ids: ["agent-1"],
|
||||
},
|
||||
];
|
||||
|
||||
describe("computeInheritedGrants", () => {
|
||||
it("attributes each id to every group that grants it, in group order", () => {
|
||||
expect(computeInheritedGrants(["mcp-1", "mcp-2"], GRANTS, (g) => g.mcp_server_ids)).toEqual([
|
||||
{ id: "mcp-1", accessGroupNames: ["platform-tools"] },
|
||||
{ id: "mcp-2", accessGroupNames: ["platform-tools", "support"] },
|
||||
]);
|
||||
});
|
||||
|
||||
it("keeps ids the flat list carries but no group detail explains, with no group names", () => {
|
||||
expect(computeInheritedGrants(["agent-1", "agent-legacy"], GRANTS, (g) => g.agent_ids)).toEqual([
|
||||
{ id: "agent-1", accessGroupNames: ["support"] },
|
||||
{ id: "agent-legacy", accessGroupNames: [] },
|
||||
]);
|
||||
});
|
||||
|
||||
it("falls back to the group details when the flat list is missing, without duplicates", () => {
|
||||
expect(computeInheritedGrants(undefined, GRANTS, (g) => g.mcp_server_ids).map((g) => g.id)).toEqual([
|
||||
"mcp-1",
|
||||
"mcp-2",
|
||||
]);
|
||||
});
|
||||
|
||||
it("returns nothing when neither source has ids", () => {
|
||||
expect(computeInheritedGrants(undefined, undefined, (g) => g.agent_ids)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("inheritedGrantTooltip", () => {
|
||||
it("names a single group", () => {
|
||||
expect(inheritedGrantTooltip({ id: "mcp-1", accessGroupNames: ["platform-tools"] })).toBe(
|
||||
"Granted via access group platform-tools. Full ID: mcp-1",
|
||||
);
|
||||
});
|
||||
|
||||
it("lists several groups", () => {
|
||||
expect(inheritedGrantTooltip({ id: "mcp-2", accessGroupNames: ["platform-tools", "support"] })).toBe(
|
||||
"Granted via access groups platform-tools, support. Full ID: mcp-2",
|
||||
);
|
||||
});
|
||||
|
||||
it("stays generic when the proxy did not say which group granted it", () => {
|
||||
expect(inheritedGrantTooltip({ id: "agent-legacy", accessGroupNames: [] })).toBe(
|
||||
"Granted via an access group. Full ID: agent-legacy",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
import { describeGroups, TeamAccessGroupModelGrant } from "../team/teamModelAccess";
|
||||
|
||||
export interface InheritedGrant {
|
||||
id: string;
|
||||
accessGroupNames: string[];
|
||||
}
|
||||
|
||||
export function computeInheritedGrants(
|
||||
ids: string[] | undefined,
|
||||
grants: TeamAccessGroupModelGrant[] | undefined,
|
||||
idsOf: (grant: TeamAccessGroupModelGrant) => string[] | undefined,
|
||||
): InheritedGrant[] {
|
||||
const known = grants ?? [];
|
||||
const allIds = [...new Set([...(ids ?? []), ...known.flatMap((grant) => idsOf(grant) ?? [])])];
|
||||
return allIds.map((id) => ({
|
||||
id,
|
||||
accessGroupNames: known
|
||||
.filter((grant) => (idsOf(grant) ?? []).includes(id))
|
||||
.map((grant) => grant.access_group_name),
|
||||
}));
|
||||
}
|
||||
|
||||
export const inheritedGrantTooltip = (grant: InheritedGrant): string => {
|
||||
const source = grant.accessGroupNames.length > 0 ? describeGroups(grant.accessGroupNames) : "an access group";
|
||||
return `Granted via ${source}. Full ID: ${grant.id}`;
|
||||
};
|
||||
|
|
@ -38,6 +38,9 @@ vi.mock("@/components/networking", () => ({
|
|||
organizationInfoCall: vi.fn(),
|
||||
getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }),
|
||||
getPassThroughEndpointsCall: vi.fn(),
|
||||
fetchMCPServers: vi.fn().mockResolvedValue([]),
|
||||
fetchMCPToolsets: vi.fn().mockResolvedValue([]),
|
||||
getAgentsList: vi.fn().mockResolvedValue({ agents: [] }),
|
||||
}));
|
||||
|
||||
const can = vi.fn();
|
||||
|
|
@ -302,6 +305,47 @@ describe("TeamInfoView", () => {
|
|||
);
|
||||
});
|
||||
|
||||
it("shows MCP servers and agents inherited from access groups in the Object Permissions card, naming the group on hover", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(networking.fetchMCPServers).mockResolvedValue([
|
||||
{ server_id: "mcp-github-1234", server_name: "github", alias: "github" },
|
||||
]);
|
||||
vi.mocked(networking.getAgentsList).mockResolvedValue({
|
||||
agents: [{ agent_id: "agent-support-5678", agent_name: "support_agent" }],
|
||||
});
|
||||
const platformToolsGroup = {
|
||||
access_group_id: "ag-1",
|
||||
access_group_name: "platform-tools",
|
||||
models: [],
|
||||
mcp_server_ids: ["mcp-github-1234"],
|
||||
agent_ids: ["agent-support-5678"],
|
||||
};
|
||||
const inheritedGrants = {
|
||||
object_permission: null,
|
||||
access_group_ids: ["ag-1"],
|
||||
access_group_mcp_server_ids: ["mcp-github-1234"],
|
||||
access_group_agent_ids: ["agent-support-5678"],
|
||||
access_group_details: [platformToolsGroup],
|
||||
};
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData(inheritedGrants));
|
||||
|
||||
renderWithProviders(<TeamInfoView {...defaultProps} />);
|
||||
|
||||
const serverRow = await screen.findByText(/github \(mcp\.\.\.1234\)/);
|
||||
const agentRow = await screen.findByText(/support_agent \(age\.\.\.5678\)/);
|
||||
expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument();
|
||||
|
||||
await user.hover(serverRow);
|
||||
expect(
|
||||
await screen.findByText("Granted via access group platform-tools. Full ID: mcp-github-1234"),
|
||||
).toBeInTheDocument();
|
||||
await user.hover(agentRow);
|
||||
expect(
|
||||
await screen.findByText("Granted via access group platform-tools. Full ID: agent-support-5678"),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the all-proxy-models badge non-clickable", async () => {
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["all-proxy-models"] }));
|
||||
|
||||
|
|
@ -1844,6 +1888,8 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
|
|||
mcp_access_groups: [],
|
||||
mcp_tool_permissions: {},
|
||||
mcp_toolsets: [],
|
||||
agents: [],
|
||||
agent_access_groups: [],
|
||||
vector_stores: ["vs-1"],
|
||||
};
|
||||
|
||||
|
|
@ -1864,6 +1910,49 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
|
|||
});
|
||||
});
|
||||
|
||||
const openEditorWithAgents = async (user: ReturnType<typeof userEvent.setup>) => {
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(
|
||||
createMockTeamData({
|
||||
models: ["gpt-4"],
|
||||
object_permission: { agents: ["agent-1"], agent_access_groups: ["group-a"] },
|
||||
}),
|
||||
);
|
||||
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...props} />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await screen.findByLabelText("Team Name");
|
||||
};
|
||||
|
||||
it("resends the stored agents and agent_access_groups when the selector is left untouched", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
await openEditorWithAgents(user);
|
||||
|
||||
const payload = await save(user);
|
||||
|
||||
const objectPermission = wireBody(payload).object_permission as Record<string, unknown>;
|
||||
expect(objectPermission.agents).toStrictEqual(["agent-1"]);
|
||||
expect(objectPermission.agent_access_groups).toStrictEqual(["group-a"]);
|
||||
});
|
||||
|
||||
it("sends empty agents and agent_access_groups arrays after the last agent chip is removed", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
await openEditorWithAgents(user);
|
||||
|
||||
await user.click(within(screen.getByLabelText("agent-1")).getByRole("button"));
|
||||
await user.click(within(screen.getByLabelText("group:group-a")).getByRole("button"));
|
||||
expect(screen.queryByLabelText("agent-1")).not.toBeInTheDocument();
|
||||
expect(screen.queryByLabelText("group:group-a")).not.toBeInTheDocument();
|
||||
|
||||
const payload = await save(user);
|
||||
|
||||
const objectPermission = wireBody(payload).object_permission as Record<string, unknown>;
|
||||
expect(objectPermission.agents).toStrictEqual([]);
|
||||
expect(objectPermission.agent_access_groups).toStrictEqual([]);
|
||||
});
|
||||
|
||||
it("resends every stored value once both sections are opened", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
await openEditor(user);
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ import {
|
|||
TeamModelBadge,
|
||||
TeamModelBadgeKind,
|
||||
} from "./teamModelAccess";
|
||||
import { computeInheritedGrants } from "../permissions/inheritedGrants";
|
||||
import MetadataKeyValueFields, {
|
||||
metadataObjectToPairs,
|
||||
metadataPairsSchema,
|
||||
|
|
@ -862,12 +863,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
agents: [],
|
||||
accessGroups: [],
|
||||
};
|
||||
if (agents && agents.length > 0) {
|
||||
updateData.object_permission.agents = agents;
|
||||
}
|
||||
if (agentAccessGroups && agentAccessGroups.length > 0) {
|
||||
updateData.object_permission.agent_access_groups = agentAccessGroups;
|
||||
}
|
||||
updateData.object_permission.agents = agents;
|
||||
updateData.object_permission.agent_access_groups = agentAccessGroups;
|
||||
delete values.agents_and_groups;
|
||||
|
||||
// Handle vector stores permissions
|
||||
|
|
@ -936,6 +933,17 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
|
||||
const { team_info: info } = teamData;
|
||||
|
||||
const inheritedMcpServers = computeInheritedGrants(
|
||||
info.access_group_mcp_server_ids,
|
||||
info.access_group_details,
|
||||
(grant) => grant.mcp_server_ids,
|
||||
);
|
||||
const inheritedAgents = computeInheritedGrants(
|
||||
info.access_group_agent_ids,
|
||||
info.access_group_details,
|
||||
(grant) => grant.agent_ids,
|
||||
);
|
||||
|
||||
const initialKillSwitchOn = info.metadata?.disable_global_guardrails === true;
|
||||
|
||||
const allGuardrails: GuardrailListItem[] = guardrailsData?.guardrails ?? [];
|
||||
|
|
@ -1033,7 +1041,13 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
</div>
|
||||
</Card>
|
||||
|
||||
<ObjectPermissionsView objectPermission={info.object_permission} variant="card" accessToken={accessToken} />
|
||||
<ObjectPermissionsView
|
||||
objectPermission={info.object_permission}
|
||||
inheritedMcpServers={inheritedMcpServers}
|
||||
inheritedAgents={inheritedAgents}
|
||||
variant="card"
|
||||
accessToken={accessToken}
|
||||
/>
|
||||
|
||||
<Card className="block p-6">
|
||||
<GuardrailSettingsView
|
||||
|
|
@ -1883,6 +1897,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
|
||||
<ObjectPermissionsView
|
||||
objectPermission={info.object_permission}
|
||||
inheritedMcpServers={inheritedMcpServers}
|
||||
inheritedAgents={inheritedAgents}
|
||||
variant="inline"
|
||||
className="pt-4 border-t border-border"
|
||||
accessToken={accessToken}
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ export interface TeamAccessGroupModelGrant {
|
|||
access_group_id: string;
|
||||
access_group_name: string;
|
||||
models: string[];
|
||||
mcp_server_ids?: string[];
|
||||
agent_ids?: string[];
|
||||
}
|
||||
|
||||
export type TeamModelBadgeKind = "all-proxy" | "no-default" | "direct" | "access-group";
|
||||
|
|
@ -19,7 +21,7 @@ export function normalizeTeamModelSelection(models: string[] | undefined): strin
|
|||
return models && models.length > 0 ? models : [NO_DEFAULT_MODELS];
|
||||
}
|
||||
|
||||
const describeGroups = (names: string[]): string =>
|
||||
export const describeGroups = (names: string[]): string =>
|
||||
names.length > 1 ? `access groups ${names.join(", ")}` : `access group ${names[0]}`;
|
||||
|
||||
export function computeTeamModelBadges(
|
||||
|
|
|
|||
|
|
@ -303,9 +303,9 @@ export function KeyEditView({
|
|||
}
|
||||
};
|
||||
|
||||
const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | undefined) => {
|
||||
setField(orgId || null);
|
||||
setSelectedOrganizationId(orgId || null);
|
||||
const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => {
|
||||
setField(orgId);
|
||||
setSelectedOrganizationId(orgId);
|
||||
form.setValue("team_id", undefined);
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -75,6 +75,41 @@ describe("Cost column", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("Tokens column", () => {
|
||||
const sessionRow: Partial<LogEntry> = {
|
||||
request_id: "req-session-tokens",
|
||||
total_tokens: 10,
|
||||
prompt_tokens: 7,
|
||||
completion_tokens: 3,
|
||||
session_id: "sess-1",
|
||||
session_total_count: 3,
|
||||
};
|
||||
|
||||
it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => {
|
||||
const aggregatedRow: Partial<LogEntry> = {
|
||||
...sessionRow,
|
||||
session_total_tokens: 60,
|
||||
session_total_prompt_tokens: 42,
|
||||
session_total_completion_tokens: 18,
|
||||
};
|
||||
renderRows([logEntry(aggregatedRow)]);
|
||||
|
||||
const tokensCell = screen.getByRole("cell", { name: /\(42\+18\)/ });
|
||||
expect(tokensCell).toHaveTextContent("60");
|
||||
expect(tokensCell).toHaveTextContent("session total");
|
||||
expect(screen.queryByText("10")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("(7+3)")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => {
|
||||
renderRows([logEntry(sessionRow)]);
|
||||
|
||||
const tokensCell = screen.getByRole("cell", { name: /\(7\+3\)/ });
|
||||
expect(tokensCell).toHaveTextContent("10");
|
||||
expect(tokensCell).not.toHaveTextContent("session total");
|
||||
});
|
||||
});
|
||||
|
||||
describe("Type column", () => {
|
||||
it("shows the conversation badge and composition even when an MCP call represents the conversation", async () => {
|
||||
const user = userEvent.setup();
|
||||
|
|
|
|||
|
|
@ -263,13 +263,20 @@ export const getRequestLogsTableColumns = ({
|
|||
meta: { numeric: true },
|
||||
cell: ({ row }) => {
|
||||
const log = row.original;
|
||||
const showSessionTotal = (log.session_total_count || 1) > 1 && log.session_total_tokens != null;
|
||||
const total = showSessionTotal ? log.session_total_tokens : log.total_tokens;
|
||||
const prompt = showSessionTotal ? log.session_total_prompt_tokens : log.prompt_tokens;
|
||||
const completion = showSessionTotal ? log.session_total_completion_tokens : log.completion_tokens;
|
||||
return (
|
||||
<span className="text-sm">
|
||||
{String(log.total_tokens || "0")}
|
||||
<span className="text-muted-foreground text-xs ml-1">
|
||||
({String(log.prompt_tokens || "0")}+{String(log.completion_tokens || "0")})
|
||||
<div className="flex flex-col items-end">
|
||||
<span className="text-sm">
|
||||
{String(total || "0")}
|
||||
<span className="text-muted-foreground text-xs ml-1">
|
||||
({String(prompt || "0")}+{String(completion || "0")})
|
||||
</span>
|
||||
</span>
|
||||
</span>
|
||||
{showSessionTotal && <span className="text-[10px] text-muted-foreground">session total</span>}
|
||||
</div>
|
||||
);
|
||||
},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -42,6 +42,9 @@ export type LogEntry = {
|
|||
request_duration_ms?: number;
|
||||
session_total_count?: number;
|
||||
session_total_spend?: number;
|
||||
session_total_tokens?: number;
|
||||
session_total_prompt_tokens?: number;
|
||||
session_total_completion_tokens?: number;
|
||||
session_cache_hit_count?: number;
|
||||
mcp_tool_call_count?: number;
|
||||
mcp_tool_call_spend?: number;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue