diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs index 23b520e2ad5..7e705ec4f8f 100644 --- a/.git-blame-ignore-revs +++ b/.git-blame-ignore-revs @@ -11,3 +11,9 @@ # style(ui): run prettier --write across the dashboard (#29622) 7edf3a9cb55548b143df1692f4ed7c4681d7fcf7 + +# style: reformat litellm/ with ruff format (#31317) +430b5b8f1b12dc261a49fda99ac5d1b22381a428 + +# style: unify ruff format width on 120 (#31518) +3dfbeabe626d203ac9de86024519d9a96c484ce4 diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index ff6c40ac9ae..f093caef073 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -59,7 +59,7 @@ jobs: echo "No changed litellm Python files to check with ruff format." exit 0 fi - xargs uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt" + xargs uv run --no-sync ruff format --check --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt" - name: Debug - Check file state run: | diff --git a/Makefile b/Makefile index 7c74526e130..7701f54e15c 100644 --- a/Makefile +++ b/Makefile @@ -82,13 +82,13 @@ install-hooks: ./scripts/install_git_hooks.sh # Formatting -# 88-column wrap matches the Black width the whole repo is formatted to; ruff.toml's -# global line-length is 120 (for E501/isort), so 88 is forced here. +# Wrap width is ruff.toml's single source of truth (line-length = 120), shared by the +# formatter, E501, and the import sorter so there's no 88-vs-120 split to reconcile. format: install-dev - cd litellm && $(UV_RUN) ruff format --line-length 88 --exclude '/enterprise/' . && cd .. + cd litellm && $(UV_RUN) ruff format --exclude '/enterprise/' . && cd .. format-check: install-dev - cd litellm && $(UV_RUN) ruff format --check --line-length 88 --exclude '/enterprise/' . && cd .. + cd litellm && $(UV_RUN) ruff format --check --exclude '/enterprise/' . && cd .. # Linting targets lint-ruff: install-dev diff --git a/litellm/__init__.py b/litellm/__init__.py index 5ae5942f32f..15e95ded906 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -6,9 +6,7 @@ import warnings warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*") # Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances # This warning can accumulate during streaming and cause memory leaks -warnings.filterwarnings( - "ignore", message=".*Accessing the.*attribute on the instance is deprecated.*" -) +warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*") ### INIT VARIABLES ######################### import threading import os @@ -166,13 +164,9 @@ _custom_logger_compatible_callbacks_literal = Literal[ ] cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None -_known_custom_logger_compatible_callbacks: List = list( - get_args(_custom_logger_compatible_callbacks_literal) -) +_known_custom_logger_compatible_callbacks: List = list(get_args(_custom_logger_compatible_callbacks_literal)) callbacks: List[ - Union[ - Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger" - ] # CustomLogger is lazy-loaded + Union[Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger"] # CustomLogger is lazy-loaded ] = [] callback_settings: Dict[str, Dict[str, Any]] = {} initialized_langfuse_clients: int = 0 @@ -183,26 +177,16 @@ prometheus_latency_buckets: Optional[List[float]] = None require_auth_for_metrics_endpoint: Optional[bool] = True argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload. -gcs_pub_sub_use_v1: Optional[bool] = ( - False # if you want to use v1 gcs pubsub logged payload -) -generic_api_use_v1: Optional[bool] = ( - False # if you want to use v1 generic api logged payload -) +gcs_pub_sub_use_v1: Optional[bool] = False # if you want to use v1 gcs pubsub logged payload +generic_api_use_v1: Optional[bool] = False # if you want to use v1 generic api logged payload argilla_transformation_object: Optional[Dict[str, Any]] = None -_async_input_callback: List[ - Union[str, Callable, "CustomLogger"] -] = ( # CustomLogger is lazy-loaded +_async_input_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. -_async_success_callback: List[ - Union[str, Callable, "CustomLogger"] -] = ( # CustomLogger is lazy-loaded +_async_success_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. -_async_failure_callback: List[ - Union[str, Callable, "CustomLogger"] -] = ( # CustomLogger is lazy-loaded +_async_failure_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. pre_call_rules: List[Callable] = [] @@ -261,9 +245,7 @@ route_all_chat_openai_to_responses: bool = ( ) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge # When True, Gemini/Vertex Live setup is deferred until client `session.update`. # Default False preserves historical behavior (auto-send setup on connect). -gemini_live_defer_setup: bool = ( - os.getenv("LITELLM_GEMINI_LIVE_DEFER_SETUP", "false").lower() == "true" -) +gemini_live_defer_setup: bool = os.getenv("LITELLM_GEMINI_LIVE_DEFER_SETUP", "false").lower() == "true" use_legacy_interactions_schema: bool = ( os.getenv("LITELLM_USE_LEGACY_INTERACTIONS_SCHEMA", "false").lower() == "true" ) # When True, sends Api-Revision: 2026-05-07 to Google so responses use the legacy `outputs` @@ -317,9 +299,7 @@ common_cloud_provider_auth_params: dict = { "params": ["project", "region_name", "token"], "providers": ["vertex_ai", "bedrock", "watsonx", "azure", "vertex_ai_beta"], } -use_litellm_proxy: bool = ( - False # when True, requests will be sent to the specified litellm proxy endpoint -) +use_litellm_proxy: bool = False # when True, requests will be sent to the specified litellm proxy endpoint use_client: bool = False ssl_verify: Union[str, bool] = True ssl_security_level: Optional[str] = None @@ -327,9 +307,7 @@ ssl_certificate: Optional[str] = None user_url_validation: bool = True user_url_allowed_hosts: List[str] = [] provider_url_destination_allowed_hosts: List[str] = [] -ssl_ecdh_curve: Optional[str] = ( - None # Set to 'X25519' to disable PQC and improve performance -) +ssl_ecdh_curve: Optional[str] = None # Set to 'X25519' to disable PQC and improve performance disable_streaming_logging: bool = False disable_token_counter: bool = False disable_add_transform_inline_image_block: bool = False @@ -370,9 +348,7 @@ prompt_name_config_map: Dict[str, PromptSpec] = {} ################## ### PREVIEW FEATURES ### enable_preview_features: bool = False -return_response_headers: bool = ( - False # get response headers from LLM Api providers - example x-remaining-requests, -) +return_response_headers: bool = False # get response headers from LLM Api providers - example x-remaining-requests, enable_json_schema_validation: bool = False enable_model_config_credential_overrides: bool = False enable_key_alias_format_validation: bool = ( @@ -384,17 +360,13 @@ enable_gemini_default_thinking_level_low: bool = ( #################### logging: bool = True enable_loadbalancing_on_batch_endpoints: Optional[bool] = None -require_managed_files: bool = ( - False # proxy only - require target_model_names on POST /v1/files -) +require_managed_files: bool = False # proxy only - require target_model_names on POST /v1/files enable_caching_on_provider_specific_optional_params: bool = ( False # feature-flag for caching on optional params - e.g. 'top_k' ) caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -cache: Optional["Cache"] = ( - None # cache object <- use this - https://docs.litellm.ai/docs/caching -) +cache: Optional["Cache"] = None # cache object <- use this - https://docs.litellm.ai/docs/caching default_in_memory_ttl: Optional[float] = None default_redis_ttl: Optional[float] = None default_redis_batch_cache_expiry: Optional[float] = None @@ -404,15 +376,15 @@ max_budget: float = 0.0 # set the max budget across all providers budget_duration: Optional[str] = ( None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). ) -default_soft_budget: float = ( - DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 -) +default_soft_budget: float = DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 forward_traceparent_to_llm_provider: bool = False _current_cost = 0.0 # private variable, used if max budget is set error_logs: Dict = {} -add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt +add_function_to_prompt: bool = ( + False # if function calling not supported by api, append function call details to system prompt +) client_session: Optional[httpx.Client] = None aclient_session: Optional[httpx.AsyncClient] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' @@ -476,9 +448,7 @@ prometheus_user_budget_label_include_email_alias: bool = False prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000 prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0 prometheus_end_user_metrics_cleanup_interval_seconds: Optional[float] = 60.0 -disable_add_prefix_to_prompt: bool = ( - False # used by anthropic, to disable adding prefix to prompt -) +disable_add_prefix_to_prompt: bool = False # used by anthropic, to disable adding prefix to prompt disable_copilot_system_to_assistant: bool = False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. public_mcp_servers: Optional[List[str]] = None public_mcp_hub_strict_whitelist: bool = True @@ -489,9 +459,7 @@ public_agent_groups: Optional[List[str]] = None # Old format: { "displayName": "url" } (for backward compatibility) public_model_groups_links: Dict[str, Union[str, Dict[str, Any]]] = {} #### REQUEST PRIORITIZATION ####### -priority_reservation: Optional[Dict[str, Union[float, "PriorityReservationDict"]]] = ( - None -) +priority_reservation: Optional[Dict[str, Union[float, "PriorityReservationDict"]]] = None # priority_reservation_settings is lazy-loaded via __getattr__ # Only declare for type checking - at runtime __getattr__ handles it if TYPE_CHECKING: @@ -502,9 +470,7 @@ if TYPE_CHECKING: use_aiohttp_transport: bool = True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings disable_aiohttp_transport: bool = False # Set this to true to use httpx instead -disable_aiohttp_trust_env: bool = ( - False # When False, aiohttp will respect HTTP(S)_PROXY env vars -) +disable_aiohttp_trust_env: bool = False # When False, aiohttp will respect HTTP(S)_PROXY env vars force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. network_mock: bool = False # When True, use mock transport — no real network calls @@ -520,9 +486,7 @@ context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True -num_retries_per_request: Optional[int] = ( - None # for the request overall (incl. fallbacks + model retries) -) +num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries) ####### SECRET MANAGERS ##################### secret_manager_client: Optional[Any] = ( None # list of instantiated key management clients - e.g. azure kv, infisical, etc. @@ -539,9 +503,7 @@ output_parse_pii: bool = False from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map model_cost = get_model_cost_map(url=model_cost_map_url) -cost_discount_config: Dict[ - str, float -] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount +cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount cost_margin_config: Dict[ str, Union[float, Dict[str, float]] ] = {} # Provider-specific or global cost margins. Examples: @@ -727,9 +689,7 @@ def is_openai_finetune_model(key: str) -> bool: def add_known_models(model_cost_map: Optional[Dict] = None): _map = model_cost_map if model_cost_map is not None else model_cost for key, value in _map.items(): - if value.get("litellm_provider") == "openai" and not is_openai_finetune_model( - key - ): + if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): open_ai_chat_completion_models.add(key) elif value.get("litellm_provider") == "text-completion-openai": open_ai_text_completion_models.add(key) @@ -807,9 +767,7 @@ def add_known_models(model_cost_map: Optional[Dict] = None): nlp_cloud_models.add(key) elif value.get("litellm_provider") == "aleph_alpha": aleph_alpha_models.add(key) - elif value.get( - "litellm_provider" - ) == "bedrock" and not is_bedrock_pricing_only_model(key): + elif value.get("litellm_provider") == "bedrock" and not is_bedrock_pricing_only_model(key): bedrock_models.add(key) elif value.get("litellm_provider") == "bedrock_converse": bedrock_converse_models.add(key) @@ -1445,9 +1403,7 @@ from . import rag from .types.llms.custom_llm import CustomLLMItem custom_provider_map: List[CustomLLMItem] = [] -_custom_providers: List[ - str -] = [] # internal helper util, used to track names of custom providers +_custom_providers: List[str] = [] # internal helper util, used to track names of custom providers disable_hf_tokenizer_download: Optional[bool] = ( None # disable huggingface tokenizer download. Defaults to openai clk100 ) diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 4d811c3d7d9..b04fae86e47 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -205,9 +205,7 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: return _LAZY_IMPORT_REGISTRY -def _generic_lazy_import( - name: str, import_map: dict[str, tuple[str, str]], category: str -) -> Any: +def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any: """ Generic function that handles lazy importing for most attributes. @@ -325,9 +323,7 @@ def _lazy_import_litellm_logging(name: str) -> Any: def _lazy_import_llm_provider_logic(name: str) -> Any: """Handler for LLM provider logic functions (get_llm_provider, etc.)""" - return _generic_lazy_import( - name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic" - ) + return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") def _lazy_import_utils_module(name: str) -> Any: diff --git a/litellm/_logging.py b/litellm/_logging.py index bb743c32878..5f3c483869d 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -17,9 +17,7 @@ if set_verbose is True: "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs." ) -_ENABLE_SECRET_REDACTION = ( - os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true" -) +_ENABLE_SECRET_REDACTION = os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true" def _redact_string(value: str) -> str: @@ -64,9 +62,7 @@ class SecretRedactionFilter(logging.Filter): # Redact exception tracebacks if record.exc_info and record.exc_info[1] is not None: try: - record.exc_text = _redact_string( - self._formatter.formatException(record.exc_info) - ) + record.exc_text = _redact_string(self._formatter.formatException(record.exc_info)) except Exception: pass @@ -189,9 +185,7 @@ class JsonFormatter(Formatter): json_record["logger"] = f"{record.filename}:{record.lineno}" if record.exc_info: - json_record["stacktrace"] = record.exc_text or self.formatException( - record.exc_info - ) + json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info) return safe_dumps(json_record) diff --git a/litellm/_redis.py b/litellm/_redis.py index fdae674d55d..2bcce0e1083 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -187,8 +187,7 @@ def _build_azure_credential( ) except ImportError: raise ImportError( - "azure-identity is required for Azure AD Redis authentication. " - "Install it with: pip install azure-identity" + "azure-identity is required for Azure AD Redis authentication. Install it with: pip install azure-identity" ) _client_id = azure_client_id or os.environ.get("AZURE_CLIENT_ID") @@ -292,9 +291,7 @@ def get_redis_url_from_environment(): return os.environ["REDIS_URL"] if "REDIS_HOST" not in os.environ or "REDIS_PORT" not in os.environ: - raise ValueError( - "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified for Redis." - ) + raise ValueError("Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified for Redis.") if "REDIS_SSL" in os.environ and os.environ["REDIS_SSL"].lower() == "true": redis_protocol = "rediss" @@ -327,9 +324,7 @@ def _get_redis_client_logic(**env_overrides): **env_overrides, } - _startup_nodes: Optional[Union[str, list]] = redis_kwargs.get( - "startup_nodes", None - ) or get_secret( # type: ignore + _startup_nodes: Optional[Union[str, list]] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore "REDIS_CLUSTER_NODES" ) @@ -340,18 +335,16 @@ def _get_redis_client_logic(**env_overrides): elif _startup_nodes is None: redis_kwargs.pop("startup_nodes", None) - _sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get( - "sentinel_nodes", None - ) or get_secret( # type: ignore + _sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore "REDIS_SENTINEL_NODES" ) if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str): redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes) - _sentinel_password: Optional[str] = redis_kwargs.get( - "sentinel_password", None - ) or get_secret_str("REDIS_SENTINEL_PASSWORD") + _sentinel_password: Optional[str] = redis_kwargs.get("sentinel_password", None) or get_secret_str( + "REDIS_SENTINEL_PASSWORD" + ) if _sentinel_password is not None: redis_kwargs["sentinel_password"] = _sentinel_password @@ -364,17 +357,11 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs["service_name"] = _service_name # Handle GCP IAM authentication - _gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str( - "REDIS_GCP_SERVICE_ACCOUNT" - ) - _gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str( - "REDIS_GCP_SSL_CA_CERTS" - ) + _gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT") + _gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS") if _gcp_service_account is not None: - verbose_logger.debug( - "Setting up GCP IAM authentication for Redis with service account." - ) + verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.") redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func( service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs ) @@ -390,14 +377,9 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs # Handle Azure AD authentication (after GCP IAM block) - _azure_redis_ad_token = redis_kwargs.get("azure_redis_ad_token") or get_secret( - "REDIS_AZURE_AD_TOKEN" - ) + _azure_redis_ad_token = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN") - _azure_ad_enabled = ( - _azure_redis_ad_token is not None - and str(_azure_redis_ad_token).lower() == "true" - ) + _azure_ad_enabled = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true" if _azure_ad_enabled and _gcp_service_account is not None: verbose_logger.warning( @@ -406,15 +388,9 @@ def _get_redis_client_logic(**env_overrides): ) if _azure_ad_enabled and _gcp_service_account is None: - _azure_client_id = redis_kwargs.get("azure_client_id") or get_secret_str( - "AZURE_CLIENT_ID" - ) - _azure_tenant_id = redis_kwargs.get("azure_tenant_id") or get_secret_str( - "AZURE_TENANT_ID" - ) - _azure_client_secret = redis_kwargs.get( - "azure_client_secret" - ) or get_secret_str("AZURE_CLIENT_SECRET") + _azure_client_id = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID") + _azure_tenant_id = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID") + _azure_client_secret = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET") verbose_logger.debug("Setting up Azure AD authentication for Redis.") redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func( @@ -446,9 +422,7 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs.pop("password", None) elif "startup_nodes" in redis_kwargs and redis_kwargs["startup_nodes"] is not None: pass - elif ( - "sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None - ): + elif "sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None: pass elif "host" not in redis_kwargs or redis_kwargs["host"] is None: raise ValueError("Either 'host' or 'url' must be specified for redis.") @@ -505,9 +479,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis: sentinel_kwargs["password"] = sentinel_password if not sentinel_nodes or not service_name: - raise ValueError( - "Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel." - ) + raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.") verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.") @@ -532,9 +504,7 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: sentinel_kwargs["password"] = sentinel_password if not sentinel_nodes or not service_name: - raise ValueError( - "Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel." - ) + raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.") verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.") @@ -593,9 +563,7 @@ def get_redis_async_client( # connection — mirrors the sync path where redis_connect_func is invoked # per connection. Without this, the token would expire after ~1 hour. if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): - cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider( - redis_connect_func._gcp_service_account - ) + cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) # Handle Azure AD authentication for async clusters via CredentialProvider # so the credential's internal cache + silent refresh runs per connection # (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry). @@ -629,9 +597,7 @@ def get_redis_async_client( url_kwargs[arg] = redis_kwargs[arg] else: verbose_logger.debug( - "REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format( - arg - ) + "REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(arg) ) return async_redis.Redis.from_url(**url_kwargs) @@ -650,9 +616,7 @@ def get_redis_async_client( username=os.environ.get("REDIS_USERNAME") or None, ) elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): - redis_kwargs["credential_provider"] = GCPIAMCredentialProvider( - redis_connect_func._gcp_service_account - ) + redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) _pretty_print_redis_config(redis_kwargs=redis_kwargs) @@ -698,18 +662,14 @@ def get_redis_connection_pool( username=os.environ.get("REDIS_USERNAME") or None, ) elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): - redis_kwargs["credential_provider"] = GCPIAMCredentialProvider( - redis_connect_func._gcp_service_account - ) + redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) connection_class = async_redis.Connection if "ssl" in redis_kwargs: connection_class = async_redis.SSLConnection redis_kwargs.pop("ssl", None) redis_kwargs["connection_class"] = connection_class - return async_redis.BlockingConnectionPool( - timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs - ) + return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs) def _pretty_print_redis_config(redis_kwargs: dict) -> None: diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index 586b1c7716c..b973e292a17 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -100,9 +100,7 @@ class GCPIAMCredentialProvider(CredentialProvider): return (token,) async def get_credentials_async(self) -> Tuple[str]: - token = await asyncio.to_thread( - _get_cached_gcp_iam_token, self._gcp_service_account - ) + token = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account) return (token,) @@ -128,9 +126,7 @@ class AzureADCredentialProvider(CredentialProvider): return (token,) async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]: - token_obj = await asyncio.to_thread( - self._credential.get_token, AZURE_REDIS_SCOPE - ) + token_obj = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE) if self._username: return (self._username, token_obj.token) return (token_obj.token,) diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index b290b4340e7..b1bd0a3bba2 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -79,9 +79,7 @@ class ServiceLogging(CustomLogger): if callback == "otel": from litellm.proxy.proxy_server import open_telemetry_logger - if open_telemetry_logger is not None and _is_otel_logger( - open_telemetry_logger - ): + if open_telemetry_logger is not None and _is_otel_logger(open_telemetry_logger): return open_telemetry_logger return None @@ -142,9 +140,7 @@ class ServiceLogging(CustomLogger): ) ) - def service_failure_hook( - self, service: ServiceTypes, duration: float, error: Exception, call_type: str - ): + def service_failure_hook(self, service: ServiceTypes, duration: float, error: Exception, call_type: str): """ [TODO] Not implemented for sync calls yet. V0 is focused on async monitoring (used by proxy). """ @@ -186,9 +182,7 @@ class ServiceLogging(CustomLogger): for callback in litellm.service_callback: if callback == "prometheus_system": await self.init_prometheus_services_logger_if_none() - await self.prometheusServicesLogger.async_service_success_hook( - payload=payload - ) + await self.prometheusServicesLogger.async_service_success_hook(payload=payload) elif callback == "datadog" or isinstance(callback, DataDogLogger): await self.init_datadog_logger_if_none() await self.dd_logger.async_service_success_hook( @@ -205,10 +199,7 @@ class ServiceLogging(CustomLogger): # here is what hid those calls from traces entirely. The OTel # logger decides what to do with a missing parent — legacy V1 # no-ops, V2 emits a root span (and skips metrics-only pings). - if ( - _otel_logger_to_use is not None - and id(_otel_logger_to_use) not in emitted_otel_logger_ids - ): + if _otel_logger_to_use is not None and id(_otel_logger_to_use) not in emitted_otel_logger_ids: emitted_otel_logger_ids.add(id(_otel_logger_to_use)) await _otel_logger_to_use.async_service_success_hook( payload=payload, @@ -249,9 +240,7 @@ class ServiceLogging(CustomLogger): from litellm.proxy.proxy_server import open_telemetry_logger if not hasattr(self, "otel_logger"): - if open_telemetry_logger is not None and isinstance( - open_telemetry_logger, OpenTelemetry - ): + if open_telemetry_logger is not None and isinstance(open_telemetry_logger, OpenTelemetry): self.otel_logger: OpenTelemetry = open_telemetry_logger else: verbose_logger.warning( @@ -319,10 +308,7 @@ class ServiceLogging(CustomLogger): # See the success hook: no parent gate, so background failures # are traced too. V1 no-ops without a parent; V2 emits a root. - if ( - _otel_logger_to_use is not None - and id(_otel_logger_to_use) not in emitted_otel_logger_ids - ): + if _otel_logger_to_use is not None and id(_otel_logger_to_use) not in emitted_otel_logger_ids: emitted_otel_logger_ids.add(id(_otel_logger_to_use)) await _otel_logger_to_use.async_service_failure_hook( payload=payload, @@ -361,9 +347,7 @@ class ServiceLogging(CustomLogger): pass else: raise Exception( - "Duration={} is not a float or timedelta object. type={}".format( - _duration, type(_duration) - ) + "Duration={} is not a float or timedelta object. type={}".format(_duration, type(_duration)) ) # invalid _duration value # Batch polling callbacks (check_batch_cost) don't include call_type in kwargs. # Use .get() to avoid KeyError. diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 4c5dd3e3ba6..1955b5268e1 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -119,17 +119,13 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] last_error = None for path in paths: try: - verbose_logger.debug( - f"Attempting to fetch agent card from {self.base_url}{path}" - ) + verbose_logger.debug(f"Attempting to fetch agent card from {self.base_url}{path}") return await super().get_agent_card( relative_card_path=path, http_kwargs=http_kwargs, ) except Exception as e: - verbose_logger.debug( - f"Failed to fetch agent card from {self.base_url}{path}: {e}" - ) + verbose_logger.debug(f"Failed to fetch agent card from {self.base_url}{path}: {e}") last_error = e continue @@ -138,7 +134,4 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] raise last_error # This shouldn't happen, but just in case - raise Exception( - f"Failed to fetch agent card from {self.base_url}. " - f"Tried paths: {', '.join(paths)}" - ) + raise Exception(f"Failed to fetch agent card from {self.base_url}. Tried paths: {', '.join(paths)}") diff --git a/litellm/a2a_protocol/client.py b/litellm/a2a_protocol/client.py index 05e21284af1..a05f8dc390c 100644 --- a/litellm/a2a_protocol/client.py +++ b/litellm/a2a_protocol/client.py @@ -87,9 +87,7 @@ class A2AClient: extra_headers=self.extra_headers, ) - async def send_message( - self, request: "SendMessageRequest" - ) -> LiteLLMSendMessageResponse: + async def send_message(self, request: "SendMessageRequest") -> LiteLLMSendMessageResponse: """Send a message to the A2A agent.""" from litellm.a2a_protocol.main import asend_message @@ -103,7 +101,5 @@ class A2AClient: from litellm.a2a_protocol.main import asend_message_streaming a2a_client = await self._get_client() - async for chunk in asend_message_streaming( - a2a_client=a2a_client, request=request - ): + async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request): yield chunk diff --git a/litellm/a2a_protocol/cost_calculator.py b/litellm/a2a_protocol/cost_calculator.py index f64174f8be5..f3e84c5b84d 100644 --- a/litellm/a2a_protocol/cost_calculator.py +++ b/litellm/a2a_protocol/cost_calculator.py @@ -97,11 +97,7 @@ class A2ACostCalculator: completion_tokens = getattr(usage, "completion_tokens", 0) or 0 # Calculate costs - input_cost = prompt_tokens * ( - float(input_cost_per_token) if input_cost_per_token else 0.0 - ) - output_cost = completion_tokens * ( - float(output_cost_per_token) if output_cost_per_token else 0.0 - ) + input_cost = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0) + output_cost = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0) return input_cost + output_cost diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py index 49dbb22b158..99706e15cee 100644 --- a/litellm/a2a_protocol/exception_mapping_utils.py +++ b/litellm/a2a_protocol/exception_mapping_utils.py @@ -181,10 +181,7 @@ def handle_a2a_localhost_retry( ImportError: If the A2A SDK is not installed """ if not A2A_SDK_AVAILABLE or _A2AClient is None: - raise ImportError( - "A2A SDK is required for localhost retry handling. " - "Install it with: pip install a2a" - ) + raise ImportError("A2A SDK is required for localhost retry handling. Install it with: pip install a2a") request_type = "streaming " if is_streaming else "" verbose_logger.warning( diff --git a/litellm/a2a_protocol/exceptions.py b/litellm/a2a_protocol/exceptions.py index 546b23105be..b672971e727 100644 --- a/litellm/a2a_protocol/exceptions.py +++ b/litellm/a2a_protocol/exceptions.py @@ -139,10 +139,7 @@ class A2ALocalhostURLError(A2AConnectionError): self.base_url = base_url self.original_error = original_error - message = ( - f"Agent card contains localhost/internal URL '{localhost_url}'. " - f"Retrying with base URL '{base_url}'." - ) + message = f"Agent card contains localhost/internal URL '{localhost_url}'. Retrying with base URL '{base_url}'." super().__init__( message=message, url=localhost_url, diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index a3502f21f95..a84b23a2170 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -75,9 +75,7 @@ class A2ACompletionBridgeHandler: ) if a2a_provider_config is not None: - verbose_logger.info( - f"A2A: Using provider config for {custom_llm_provider}" - ) + verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}") return await a2a_provider_config.handle_non_streaming( request_id=request_id, @@ -91,9 +89,7 @@ class A2ACompletionBridgeHandler: message = params.get("message", {}) # Transform A2A message to OpenAI format - openai_messages = ( - A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) - ) + openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) # Get completion params custom_llm_provider = litellm_params.get("custom_llm_provider") @@ -106,9 +102,7 @@ class A2ACompletionBridgeHandler: else: full_model = model - verbose_logger.info( - f"A2A completion bridge: model={full_model}, api_base={api_base}" - ) + verbose_logger.info(f"A2A completion bridge: model={full_model}, api_base={api_base}") # Build completion params dict completion_params: Dict[str, Any] = { @@ -143,11 +137,9 @@ class A2ACompletionBridgeHandler: response = await litellm.acompletion(**completion_params) # Transform response to A2A format - a2a_response = ( - A2ACompletionBridgeTransformation.openai_response_to_a2a_response( - response=response, - request_id=request_id, - ) + a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response( + response=response, + request_id=request_id, ) verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}") @@ -192,9 +184,7 @@ class A2ACompletionBridgeHandler: ) if a2a_provider_config is not None: - verbose_logger.info( - f"A2A: Using provider config for {custom_llm_provider} (streaming)" - ) + verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider} (streaming)") async for chunk in a2a_provider_config.handle_streaming( request_id=request_id, @@ -217,9 +207,7 @@ class A2ACompletionBridgeHandler: ) # Transform A2A message to OpenAI format - openai_messages = ( - A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) - ) + openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) # Get completion params custom_llm_provider = litellm_params.get("custom_llm_provider") @@ -232,9 +220,7 @@ class A2ACompletionBridgeHandler: else: full_model = model - verbose_logger.info( - f"A2A completion bridge streaming: model={full_model}, api_base={api_base}" - ) + verbose_logger.info(f"A2A completion bridge streaming: model={full_model}, api_base={api_base}") # Build completion params dict completion_params: Dict[str, Any] = { @@ -299,11 +285,9 @@ class A2ACompletionBridgeHandler: # Emit artifact update with accumulated content if accumulated_text: - artifact_event = ( - A2ACompletionBridgeTransformation.create_artifact_update_event( - ctx=ctx, - text=accumulated_text, - ) + artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=accumulated_text, ) yield artifact_event @@ -315,9 +299,7 @@ class A2ACompletionBridgeHandler: ) yield completed_event - verbose_logger.info( - f"A2A completion bridge streaming completed: request_id={request_id}, chunks={chunk_count}" - ) + verbose_logger.info(f"A2A completion bridge streaming completed: request_id={request_id}, chunks={chunk_count}") # Convenience functions that delegate to the class methods diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 06c0a8fc82f..b32963dd6fb 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -104,16 +104,12 @@ class A2ACompletionBridgeTransformation: # ``extra_body.metadata`` so the configured keys remain authoritative # and an A2A caller cannot overwrite server-set run metadata. existing_metadata = extra_body.get("metadata") - existing_dict: Dict[str, Any] = ( - existing_metadata if isinstance(existing_metadata, dict) else {} - ) + existing_dict: Dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {} merged_metadata: Dict[str, Any] = {**forward_metadata, **existing_dict} extra_body = {**extra_body, "metadata": merged_metadata} completion_params["extra_body"] = extra_body - verbose_logger.debug( - f"A2A -> completion forward metadata keys={list(forward_metadata.keys())}" - ) + verbose_logger.debug(f"A2A -> completion forward metadata keys={list(forward_metadata.keys())}") @staticmethod def a2a_message_to_openai_messages( @@ -149,9 +145,7 @@ class A2ACompletionBridgeTransformation: # once at run level via extra_body.metadata (LangGraph POST /runs/wait shape). openai_message: Dict[str, Any] = {"role": openai_role, "content": content} - verbose_logger.debug( - f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}" - ) + verbose_logger.debug(f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}") return [openai_message] diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 2b6f2cd12b4..6694c5c4af3 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -120,9 +120,7 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: litellm_logging_obj.model = model litellm_logging_obj.custom_llm_provider = custom_llm_provider litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details["custom_llm_provider"] = ( - custom_llm_provider - ) + litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider return agent_name @@ -139,19 +137,13 @@ async def _send_message_via_completion_bridge( Requires request; api_base is optional for providers that derive endpoint from model. """ - verbose_logger.info( - f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}" - ) + verbose_logger.info(f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}") from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, ) - params = ( - request.params.model_dump(mode="json") - if hasattr(request.params, "model_dump") - else dict(request.params) - ) + params = request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) response_dict = await A2ACompletionBridgeHandler.handle_non_streaming( request_id=str(request.id), @@ -161,9 +153,7 @@ async def _send_message_via_completion_bridge( agent_extra_headers=agent_extra_headers, ) - return LiteLLMSendMessageResponse.from_dict( - response_dict, request_id=str(request.id) - ) + return LiteLLMSendMessageResponse.from_dict(response_dict, request_id=str(request.id)) async def _execute_a2a_send_with_retry( @@ -203,9 +193,7 @@ async def _execute_a2a_send_with_retry( except Exception: raise if a2a_response is None: - raise RuntimeError( - "A2A send_message failed: no response received after retry attempts." - ) + raise RuntimeError("A2A send_message failed: no response received after retry attempts.") return a2a_response @@ -295,9 +283,7 @@ async def asend_message( # Create A2A client if not provided but api_base is available if a2a_client is None: if api_base is None: - raise ValueError( - "Either a2a_client or api_base is required for standard A2A flow" - ) + raise ValueError("Either a2a_client or api_base is required for standard A2A flow") trace_id = trace_id or str(uuid.uuid4()) extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: @@ -305,9 +291,7 @@ async def asend_message( # Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones) if agent_extra_headers: extra_headers.update(agent_extra_headers) - a2a_client = await create_a2a_client( - base_url=api_base, extra_headers=extra_headers - ) + a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None @@ -317,9 +301,7 @@ async def asend_message( verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}") # Get agent card URL for localhost retry logic - agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr( - a2a_client, "agent_card", None - ) + agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(a2a_client, "agent_card", None) card_url = getattr(agent_card, "url", None) if agent_card else None a2a_response = await _execute_a2a_send_with_retry( @@ -334,9 +316,7 @@ async def asend_message( verbose_logger.info(f"A2A send_message completed, request_id={request.id}") # Wrap in LiteLLM response type for _hidden_params support - response = LiteLLMSendMessageResponse.from_a2a_response( - a2a_response, request_id=str(request.id) - ) + response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id)) # Calculate token usage from request and response response_dict = a2a_response.model_dump(mode="json", exclude_none=True) @@ -389,9 +369,7 @@ def send_message( if loop is not None: return asend_message(a2a_client=a2a_client, request=request, **kwargs) else: - return asyncio.run( - asend_message(a2a_client=a2a_client, request=request, **kwargs) - ) + return asyncio.run(asend_message(a2a_client=a2a_client, request=request, **kwargs)) def _build_streaming_logging_obj( @@ -492,9 +470,7 @@ async def asend_message_streaming( raise ValueError("request is required for completion bridge") # api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore) - verbose_logger.info( - f"A2A streaming using completion bridge: provider={custom_llm_provider}" - ) + verbose_logger.info(f"A2A streaming using completion bridge: provider={custom_llm_provider}") from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, @@ -502,9 +478,7 @@ async def asend_message_streaming( # Extract params from request params = ( - request.params.model_dump(mode="json") - if hasattr(request.params, "model_dump") - else dict(request.params) + request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) ) async for chunk in A2ACompletionBridgeHandler.handle_streaming( @@ -524,9 +498,7 @@ async def asend_message_streaming( # Create A2A client if not provided but api_base is available if a2a_client is None: if api_base is None: - raise ValueError( - "Either a2a_client or api_base is required for standard A2A flow" - ) + raise ValueError("Either a2a_client or api_base is required for standard A2A flow") # Mirror the non-streaming path: always include trace and agent-id headers streaming_extra_headers: Dict[str, str] = { "X-LiteLLM-Trace-Id": str(request.id), @@ -535,9 +507,7 @@ async def asend_message_streaming( streaming_extra_headers["X-LiteLLM-Agent-Id"] = agent_id if agent_extra_headers: streaming_extra_headers.update(agent_extra_headers) - a2a_client = await create_a2a_client( - base_url=api_base, extra_headers=streaming_extra_headers - ) + a2a_client = await create_a2a_client(base_url=api_base, extra_headers=streaming_extra_headers) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None @@ -545,9 +515,7 @@ async def asend_message_streaming( verbose_logger.info(f"A2A send_message_streaming request_id={request.id}") # Build logging object for streaming completion callbacks - agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr( - a2a_client, "agent_card", None - ) + agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(a2a_client, "agent_card", None) card_url = getattr(agent_card, "url", None) if agent_card else None agent_name = getattr(agent_card, "name", "unknown") if agent_card else "unknown" @@ -645,8 +613,7 @@ async def create_a2a_client( """ if not A2A_SDK_AVAILABLE: raise ImportError( - "The 'a2a' package is required for A2A agent invocation. " - "Install it with: pip install a2a-sdk" + "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) verbose_logger.info(f"Creating A2A client for {base_url}") @@ -671,9 +638,7 @@ async def create_a2a_client( httpx_client = _async_handler.client if extra_headers: httpx_client.headers.update(extra_headers) - verbose_proxy_logger.debug( - f"A2A client created with extra_headers={list(extra_headers.keys())}" - ) + verbose_proxy_logger.debug(f"A2A client created with extra_headers={list(extra_headers.keys())}") # Resolve agent card resolver = A2ACardResolver( @@ -682,9 +647,7 @@ async def create_a2a_client( ) agent_card = await resolver.get_agent_card() - verbose_logger.debug( - f"Resolved agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}" - ) + verbose_logger.debug(f"Resolved agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}") # Create A2A client a2a_client = _A2AClient( @@ -718,8 +681,7 @@ async def aget_agent_card( """ if not A2A_SDK_AVAILABLE: raise ImportError( - "The 'a2a' package is required for A2A agent invocation. " - "Install it with: pip install a2a-sdk" + "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) verbose_logger.info(f"Fetching agent card from {base_url}") @@ -737,7 +699,5 @@ async def aget_agent_card( ) agent_card = await resolver.get_agent_card() - verbose_logger.info( - f"Fetched agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}" - ) + verbose_logger.info(f"Fetched agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}") return agent_card diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py index e7f38c6488c..f624aa393ed 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py @@ -30,8 +30,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( - "litellm_params is required for BedrockAgentCoreA2AConfig " - "(must contain model with AgentCore ARN)" + "litellm_params is required for BedrockAgentCoreA2AConfig (must contain model with AgentCore ARN)" ) return await BedrockAgentCoreA2AHandler.handle_non_streaming( request_id=request_id, @@ -51,8 +50,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( - "litellm_params is required for BedrockAgentCoreA2AConfig " - "(must contain model with AgentCore ARN)" + "litellm_params is required for BedrockAgentCoreA2AConfig (must contain model with AgentCore ARN)" ) async for chunk in BedrockAgentCoreA2AHandler.handle_streaming( request_id=request_id, diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index 2f93895099b..c613b68668f 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -44,19 +44,15 @@ class BedrockAgentCoreA2AHandler: Returns: A2A JSON-RPC response dict from the AgentCore agent """ - url, headers, body = ( - BedrockAgentCoreA2ATransformation.get_url_and_signed_request( - request_id=request_id, - params=params, - litellm_params=litellm_params, - method="message/send", - agent_extra_headers=agent_extra_headers, - ) + url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + request_id=request_id, + params=params, + litellm_params=litellm_params, + method="message/send", + agent_extra_headers=agent_extra_headers, ) - verbose_logger.info( - f"BedrockAgentCore A2A: Sending non-streaming request to {url}" - ) + verbose_logger.info(f"BedrockAgentCore A2A: Sending non-streaming request to {url}") client = get_async_httpx_client( llm_provider=cast(Any, httpxSpecialProvider.A2AProvider), @@ -70,9 +66,7 @@ class BedrockAgentCoreA2AHandler: response_data = response.json() if "error" in response_data: - verbose_logger.warning( - f"BedrockAgentCore A2A: Agent returned error: {response_data['error']}" - ) + verbose_logger.warning(f"BedrockAgentCore A2A: Agent returned error: {response_data['error']}") return response_data @@ -96,15 +90,13 @@ class BedrockAgentCoreA2AHandler: Yields: A2A streaming response events from the AgentCore agent """ - url, headers, body = ( - BedrockAgentCoreA2ATransformation.get_url_and_signed_request( - request_id=request_id, - params=params, - litellm_params=litellm_params, - method="message/send", - stream=True, - agent_extra_headers=agent_extra_headers, - ) + url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + request_id=request_id, + params=params, + litellm_params=litellm_params, + method="message/send", + stream=True, + agent_extra_headers=agent_extra_headers, ) verbose_logger.info(f"BedrockAgentCore A2A: Sending streaming request to {url}") @@ -126,15 +118,12 @@ class BedrockAgentCoreA2AHandler: if "application/json" in content_type: # Single JSON response fallback (not SSE) verbose_logger.debug( - "BedrockAgentCore A2A streaming: received JSON instead of SSE, " - "yielding as single event" + "BedrockAgentCore A2A streaming: received JSON instead of SSE, yielding as single event" ) response_body = await response.aread() response_data = json.loads(response_body) yield response_data else: # SSE stream — parse data: lines - async for event in BedrockAgentCoreA2ATransformation.parse_sse_events( - response - ): + async for event in BedrockAgentCoreA2ATransformation.parse_sse_events(response): yield event diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index f868845bb58..091a13ccea5 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -50,9 +50,7 @@ def _filter_reserved_headers( dropped: list = [] for k, v in agent_extra_headers.items(): k_lower = k.lower() - if k_lower in _RESERVED_EXACT_HEADERS or any( - k_lower.startswith(prefix) for prefix in _RESERVED_PREFIX_HEADERS - ): + if k_lower in _RESERVED_EXACT_HEADERS or any(k_lower.startswith(prefix) for prefix in _RESERVED_PREFIX_HEADERS): dropped.append(k) continue filtered[k] = v @@ -115,11 +113,7 @@ class BedrockAgentCoreA2ATransformation: agentcore_model = model # Build optional_params from litellm_params (everything except model and custom_llm_provider) - optional_params = { - k: v - for k, v in litellm_params.items() - if k not in ("model", "custom_llm_provider") - } + optional_params = {k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider")} agentcore_config = AmazonAgentCoreConfig() @@ -200,7 +194,5 @@ class BedrockAgentCoreA2ATransformation: event = json.loads(data_str) yield event except json.JSONDecodeError: - verbose_logger.debug( - f"BedrockAgentCore A2A: Skipping non-JSON SSE line: {data_str[:100]}" - ) + verbose_logger.debug(f"BedrockAgentCore A2A: Skipping non-JSON SSE line: {data_str[:100]}") continue diff --git a/litellm/a2a_protocol/providers/langflow/config.py b/litellm/a2a_protocol/providers/langflow/config.py index 9302c38126b..9edaf151c71 100644 --- a/litellm/a2a_protocol/providers/langflow/config.py +++ b/litellm/a2a_protocol/providers/langflow/config.py @@ -22,8 +22,7 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( - "litellm_params is required for LangFlowA2AConfig " - "(must contain custom_llm_provider and model)" + "litellm_params is required for LangFlowA2AConfig (must contain custom_llm_provider and model)" ) litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) @@ -46,8 +45,7 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( - "litellm_params is required for LangFlowA2AConfig " - "(must contain custom_llm_provider and model)" + "litellm_params is required for LangFlowA2AConfig (must contain custom_llm_provider and model)" ) litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py index b5d3f262a63..352005ff549 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py @@ -91,9 +91,7 @@ class PydanticAIHandler: """ if api_base is None: raise ValueError("api_base is required for Pydantic AI agents") - verbose_logger.info( - f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}" - ) + verbose_logger.info(f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}") # Get raw task response first (not the transformed A2A format) raw_response = await PydanticAITransformation.send_and_get_raw_response( diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 8fac43e7ae1..b9943d83c8a 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -41,17 +41,9 @@ class PydanticAITransformation: Cleaned object with None values removed """ if isinstance(obj, dict): - return { - k: PydanticAITransformation._remove_none_values(v) - for k, v in obj.items() - if v is not None - } + return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None} elif isinstance(obj, list): - return [ - PydanticAITransformation._remove_none_values(item) - for item in obj - if item is not None - ] + return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None] else: return obj @@ -125,9 +117,7 @@ class PydanticAITransformation: status = result.get("status", {}) state = status.get("state", "") - verbose_logger.debug( - f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}" - ) + verbose_logger.debug(f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}") if state == "completed": return poll_data @@ -136,9 +126,7 @@ class PydanticAITransformation: await asyncio.sleep(poll_interval) - raise TimeoutError( - f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds" - ) + raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds") @staticmethod async def _send_and_poll_raw( @@ -211,9 +199,7 @@ class PydanticAITransformation: # Need to poll for completion task_id = result.get("id") if task_id: - verbose_logger.info( - f"Pydantic AI: Task {task_id} submitted, polling for completion..." - ) + verbose_logger.info(f"Pydantic AI: Task {task_id} submitted, polling for completion...") response_data = await PydanticAITransformation._poll_for_completion( client=client, endpoint=endpoint, @@ -222,9 +208,7 @@ class PydanticAITransformation: agent_extra_headers=agent_extra_headers, ) - verbose_logger.info( - f"Pydantic AI: Received completed response for request_id={request_id}" - ) + verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}") return response_data @@ -325,9 +309,7 @@ class PydanticAITransformation: Standard A2A non-streaming response format """ # Extract the agent response text - full_text, message_id, parts = PydanticAITransformation._extract_response_text( - response_data - ) + full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Build standard A2A message a2a_message = { @@ -424,9 +406,7 @@ class PydanticAITransformation: A2A streaming response events """ # Extract the response text from completed task - full_text, message_id, parts = PydanticAITransformation._extract_response_text( - response_data - ) + full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history result = response_data.get("result", {}) @@ -455,9 +435,7 @@ class PydanticAITransformation: "contextId": context_id, "kind": "message", "messageId": input_message_id, - "parts": input_message.get( - "parts", [{"kind": "text", "text": ""}] - ), + "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), "role": "user", "taskId": task_id, } @@ -539,6 +517,4 @@ class PydanticAITransformation: } yield completed_event - verbose_logger.info( - f"Pydantic AI: Fake streaming completed for request_id={request_id}" - ) + verbose_logger.info(f"Pydantic AI: Fake streaming completed for request_id={request_id}") diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index dbc0247618e..07235c1118c 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -56,9 +56,7 @@ class WatsonxOrchestrateHandler: return hashlib.sha256(material.encode()).hexdigest() @staticmethod - def _cp4d_token_ttl_seconds( - expiration: Any, now_wall: Optional[float] = None - ) -> int: + def _cp4d_token_ttl_seconds(expiration: Any, now_wall: Optional[float] = None) -> int: # CP4D returns expiration as absolute Unix epoch seconds, not a duration. expires_at = int(expiration) wall = now_wall if now_wall is not None else time.time() @@ -72,9 +70,7 @@ class WatsonxOrchestrateHandler: username: Optional[str] = None, client: Optional[AsyncHTTPHandler] = None, ) -> str: - cache_key = WatsonxOrchestrateHandler._token_cache_key( - auth_mode, cp4d_host, api_key, username - ) + cache_key = WatsonxOrchestrateHandler._token_cache_key(auth_mode, cp4d_host, api_key, username) now = time.monotonic() cached = _token_cache.get(cache_key) if cached and cached[1] > now: @@ -98,9 +94,7 @@ class WatsonxOrchestrateHandler: ttl_s = int(payload.get("expires_in", 3600)) else: if not username: - raise ValueError( - "'username' is required in litellm_params when auth_mode='cp4d'" - ) + raise ValueError("'username' is required in litellm_params when auth_mode='cp4d'") token_url = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize" response = await client.post( token_url, @@ -140,15 +134,12 @@ class WatsonxOrchestrateHandler: response.raise_for_status() result: Dict[str, Any] = response.json() status = result.get("status", "") - verbose_logger.debug( - f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'" - ) + verbose_logger.debug(f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'") if status in WatsonxOrchestrateTransformation.TERMINAL_STATES: return result raise asyncio.TimeoutError( - f"WXO run '{run_id}' did not reach a terminal state after " - f"{max_attempts * interval_s:.0f}s" + f"WXO run '{run_id}' did not reach a terminal state after {max_attempts * interval_s:.0f}s" ) @staticmethod @@ -172,9 +163,7 @@ class WatsonxOrchestrateHandler: status = run_data.get("status", "") if status not in WatsonxOrchestrateTransformation.SUCCESS_STATES: - raise RuntimeError( - f"WXO run ended with non-success status '{status}': {run_data}" - ) + raise RuntimeError(f"WXO run ended with non-success status '{status}': {run_data}") return run_data @@ -191,9 +180,7 @@ class WatsonxOrchestrateHandler: event = json.loads(data_str) except json.JSONDecodeError: continue - chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result( - event - ) + chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event) if chunk_text: accumulated_text += chunk_text return accumulated_text @@ -208,13 +195,9 @@ class WatsonxOrchestrateHandler: if not cp4d_host: raise ValueError("'cp4d_host' is required in litellm_params for WXO agents") if not instance_id: - raise ValueError( - "'instance_id' is required in litellm_params for WXO agents" - ) + raise ValueError("'instance_id' is required in litellm_params for WXO agents") if not wxo_agent_id: - raise ValueError( - "'wxo_agent_id' is required in litellm_params for WXO agents" - ) + raise ValueError("'wxo_agent_id' is required in litellm_params for WXO agents") if not api_key: raise ValueError("'api_key' is required in litellm_params for WXO agents") @@ -244,9 +227,7 @@ class WatsonxOrchestrateHandler: username=wxo.username, client=client, ) - base_url = WatsonxOrchestrateTransformation.get_api_base_url( - wxo.cp4d_host, wxo.instance_id - ) + base_url = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) auth_headers = { "Authorization": f"Bearer {token}", "Content-Type": "application/json", @@ -273,12 +254,8 @@ class WatsonxOrchestrateHandler: client=client, ) - response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result( - run_data - ) - return WatsonxOrchestrateTransformation.build_a2a_message_response( - request_id=request_id, text=response_text - ) + response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data) + return WatsonxOrchestrateTransformation.build_a2a_message_response(request_id=request_id, text=response_text) @staticmethod async def handle_streaming( @@ -298,9 +275,7 @@ class WatsonxOrchestrateHandler: username=wxo.username, client=client, ) - base_url = WatsonxOrchestrateTransformation.get_api_base_url( - wxo.cp4d_host, wxo.instance_id - ) + base_url = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) auth_headers = { "Authorization": f"Bearer {token}", "Content-Type": "application/json", @@ -330,14 +305,8 @@ class WatsonxOrchestrateHandler: params=params, litellm_params=litellm_params, ) - response_text = ( - WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response( - result - ) - ) - async for ( - chunk - ) in WatsonxOrchestrateTransformation.fake_streaming_from_text( + response_text = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result) + async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( text=response_text, request_id=request_id, chunk_size=chunk_size, @@ -356,13 +325,9 @@ class WatsonxOrchestrateHandler: auth_headers=auth_headers, client=client, ) - accumulated_text = ( - WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) - ) + accumulated_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) else: - accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text( - response - ) + accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( text=accumulated_text, diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py index 824e9dbcdd2..c9bda822aae 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -19,9 +19,7 @@ class WatsonxOrchestrateTransformation: Handles request/response transformation between A2A and the WXO REST API. """ - TERMINAL_STATES = frozenset( - {"completed", "succeeded", "failed", "error", "cancelled"} - ) + TERMINAL_STATES = frozenset({"completed", "succeeded", "failed", "error", "cancelled"}) SUCCESS_STATES = frozenset({"completed", "succeeded"}) @staticmethod @@ -114,11 +112,7 @@ class WatsonxOrchestrateTransformation: verbose_logger.warning("WXO: A2A result has no parts list") return "" for part in parts: - if ( - isinstance(part, dict) - and part.get("kind") == "text" - and part.get("text") - ): + if isinstance(part, dict) and part.get("kind") == "text" and part.get("text"): return str(part["text"]) verbose_logger.warning("WXO: A2A result parts contained no text") return "" @@ -219,6 +213,4 @@ class WatsonxOrchestrateTransformation: }, } - verbose_logger.debug( - f"WXO: Fake streaming completed for request_id={request_id}" - ) + verbose_logger.debug(f"WXO: Fake streaming completed for request_id={request_id}") diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index c5ae9bcdc3c..529154919f3 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -71,11 +71,7 @@ class A2AStreamingIterator: def _collect_text_from_chunk(self, chunk: Any) -> None: """Extract text from a streaming chunk and add to collected parts.""" try: - chunk_dict = ( - chunk.model_dump(mode="json", exclude_none=True) - if hasattr(chunk, "model_dump") - else {} - ) + chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} text = A2ARequestUtils.extract_text_from_response(chunk_dict) if text: self.collected_text_parts.append(text) @@ -85,11 +81,7 @@ class A2AStreamingIterator: def _is_completed_chunk(self, chunk: Any) -> bool: """Check if chunk indicates stream completion.""" try: - chunk_dict = ( - chunk.model_dump(mode="json", exclude_none=True) - if hasattr(chunk, "model_dump") - else {} - ) + chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} result = chunk_dict.get("result", {}) if isinstance(result, dict): status = result.get("status", {}) @@ -110,9 +102,7 @@ class A2AStreamingIterator: prompt_tokens = A2ARequestUtils.count_tokens(input_text) # Use the last (most complete) text from chunks - output_text = ( - self.collected_text_parts[-1] if self.collected_text_parts else "" - ) + output_text = self.collected_text_parts[-1] if self.collected_text_parts else "" completion_tokens = A2ARequestUtils.count_tokens(output_text) total_tokens = prompt_tokens + completion_tokens @@ -168,9 +158,7 @@ class A2AStreamingIterator: result: Dict[str, Any] = { "id": getattr(self.request, "id", "unknown"), "jsonrpc": "2.0", - "usage": ( - usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) - ), + "usage": (usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)), } # Add final chunk result if available diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 97d223088fa..d0082498b09 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -48,9 +48,7 @@ class GetAnthropicBetaHeadersConfig: """Load the local backup beta headers config bundled with the package.""" try: content = json.loads( - files("litellm") - .joinpath("anthropic_beta_headers_config.json") - .read_text(encoding="utf-8") + files("litellm").joinpath("anthropic_beta_headers_config.json").read_text(encoding="utf-8") ) return content except Exception as e: @@ -70,16 +68,14 @@ class GetAnthropicBetaHeadersConfig: """Check if fetched config is a non-empty dict with expected structure.""" if not isinstance(fetched_config, dict): verbose_logger.warning( - "LiteLLM: Fetched beta headers config is not a dict (type=%s). " - "Falling back to local backup.", + "LiteLLM: Fetched beta headers config is not a dict (type=%s). Falling back to local backup.", type(fetched_config).__name__, ) return False if len(fetched_config) == 0: verbose_logger.warning( - "LiteLLM: Fetched beta headers config is empty. " - "Falling back to local backup.", + "LiteLLM: Fetched beta headers config is empty. Falling back to local backup.", ) return False @@ -95,8 +91,7 @@ class GetAnthropicBetaHeadersConfig: if not has_provider: verbose_logger.warning( - "LiteLLM: Fetched beta headers config missing provider keys. " - "Falling back to local backup.", + "LiteLLM: Fetched beta headers config missing provider keys. Falling back to local backup.", ) return False @@ -147,20 +142,16 @@ def get_beta_headers_config(url: str) -> dict: content = GetAnthropicBetaHeadersConfig.fetch_remote_beta_headers_config(url) except Exception as e: verbose_logger.warning( - "LiteLLM: Failed to fetch remote beta headers config from %s: %s. " - "Falling back to local backup.", + "LiteLLM: Failed to fetch remote beta headers config from %s: %s. Falling back to local backup.", url, str(e), ) return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() # Validate the fetched config - if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config( - fetched_config=content - ): + if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(fetched_config=content): verbose_logger.warning( - "LiteLLM: Fetched beta headers config failed integrity check. " - "Using local backup instead. url=%s", + "LiteLLM: Fetched beta headers config failed integrity check. Using local backup instead. url=%s", url, ) return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() @@ -256,9 +247,7 @@ def filter_and_transform_beta_headers( # Check if header is in the mapping if header not in provider_mapping: - verbose_logger.debug( - f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)" - ) + verbose_logger.debug(f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)") continue # Get the mapped header value @@ -266,9 +255,7 @@ def filter_and_transform_beta_headers( # Skip if header is unsupported (null value) if mapped_header is None: - verbose_logger.debug( - f"Dropping unsupported beta header '{header}' for provider '{provider}'" - ) + verbose_logger.debug(f"Dropping unsupported beta header '{header}' for provider '{provider}'") continue # Add the mapped header diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index 4548185bbdc..b4ec83517ee 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -148,9 +148,7 @@ class AnthropicExceptionMapping: parsed = None # If parsed and already in Anthropic format - passthrough - if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict( - parsed - ): + if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(parsed): # Optionally add request_id if provided and not present if request_id and "request_id" not in parsed: parsed["request_id"] = request_id @@ -158,9 +156,7 @@ class AnthropicExceptionMapping: # Extract message - use parsed dict if available, otherwise raw string if parsed is not None: - message = AnthropicExceptionMapping._extract_message_from_dict( - parsed, raw_message - ) + message = AnthropicExceptionMapping._extract_message_from_dict(parsed, raw_message) else: message = raw_message diff --git a/litellm/anthropic_interface/messages/__init__.py b/litellm/anthropic_interface/messages/__init__.py index f71279b226d..52c9ecd5aa4 100644 --- a/litellm/anthropic_interface/messages/__init__.py +++ b/litellm/anthropic_interface/messages/__init__.py @@ -102,9 +102,7 @@ def create( AnthropicMessagesResponse, Iterator[bytes], AsyncIterator[Any], - Coroutine[ - Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]] - ], + Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]], ]: """ Async wrapper for Anthropic's messages API diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index b7dfed6b169..d515cb278bc 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -81,12 +81,8 @@ def get_assistants( ) -> SyncCursorPage[Assistant]: aget_assistants: Optional[bool] = kwargs.pop("aget_assistants", None) if aget_assistants is not None and not isinstance(aget_assistants, bool): - raise Exception( - "Invalid value passed in for aget_assistants. Only bool or None allowed" - ) - optional_params = GenericLiteLLMParams( - api_key=api_key, api_base=api_base, api_version=api_version, **kwargs - ) + raise Exception("Invalid value passed in for aget_assistants. Only bool or None allowed") + optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) litellm_params_dict = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### @@ -138,15 +134,9 @@ def get_assistants( aget_assistants=aget_assistants, # type: ignore ) # type: ignore elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -184,9 +174,7 @@ def get_assistants( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) @@ -200,9 +188,7 @@ def get_assistants( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) @@ -266,18 +252,10 @@ def create_assistants( api_version: Optional[str] = None, **kwargs, ) -> Union[Assistant, Coroutine[Any, Any, Assistant]]: - async_create_assistants: Optional[bool] = kwargs.pop( - "async_create_assistants", None - ) - if async_create_assistants is not None and not isinstance( - async_create_assistants, bool - ): - raise ValueError( - "Invalid value passed in for async_create_assistants. Only bool or None allowed" - ) - optional_params = GenericLiteLLMParams( - api_key=api_key, api_base=api_base, api_version=api_version, **kwargs - ) + async_create_assistants: Optional[bool] = kwargs.pop("async_create_assistants", None) + if async_create_assistants is not None and not isinstance(async_create_assistants, bool): + raise ValueError("Invalid value passed in for async_create_assistants. Only bool or None allowed") + optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) litellm_params_dict = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### @@ -310,9 +288,7 @@ def create_assistants( } # only send params that are not None - create_assistant_data = { - k: v for k, v in create_assistant_data.items() if v is not None - } + create_assistant_data = {k: v for k, v in create_assistant_data.items() if v is not None} response: Optional[Union[Coroutine[Any, Any, Assistant], Assistant]] = None if custom_llm_provider == "openai": @@ -348,15 +324,9 @@ def create_assistants( async_create_assistants=async_create_assistants, # type: ignore ) # type: ignore elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -398,9 +368,7 @@ def create_assistants( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) if response is None: @@ -459,21 +427,13 @@ def delete_assistant( api_version: Optional[str] = None, **kwargs, ) -> Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]: - optional_params = GenericLiteLLMParams( - api_key=api_key, api_base=api_base, api_version=api_version, **kwargs - ) + optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) litellm_params_dict = get_litellm_params(**kwargs) - async_delete_assistants: Optional[bool] = kwargs.pop( - "async_delete_assistants", None - ) - if async_delete_assistants is not None and not isinstance( - async_delete_assistants, bool - ): - raise ValueError( - "Invalid value passed in for async_delete_assistants. Only bool or None allowed" - ) + async_delete_assistants: Optional[bool] = kwargs.pop("async_delete_assistants", None) + if async_delete_assistants is not None and not isinstance(async_delete_assistants, bool): + raise ValueError("Invalid value passed in for async_delete_assistants. Only bool or None allowed") ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 @@ -491,9 +451,7 @@ def delete_assistant( elif timeout is None: timeout = 600.0 - response: Optional[ - Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]] - ] = None + response: Optional[Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]] = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base @@ -503,18 +461,10 @@ def delete_assistant( or "https://api.openai.com/v1" ) organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None + optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None ) # set API KEY - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") - ) + api_key = optional_params.api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY") response = openai_assistants_api.delete_assistant( api_base=api_base, @@ -527,15 +477,9 @@ def delete_assistant( async_delete_assistants=async_delete_assistants, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -577,9 +521,7 @@ def delete_assistant( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="delete_assistant", url="https://github.com/BerriAI/litellm" - ), + request=httpx.Request(method="delete_assistant", url="https://github.com/BerriAI/litellm"), ), ) if response is None: @@ -594,9 +536,7 @@ def delete_assistant( ### THREADS ### -async def acreate_thread( - custom_llm_provider: Literal["openai", "azure"], **kwargs -) -> Thread: +async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwargs) -> Thread: loop = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["acreate_thread"] = True @@ -716,9 +656,7 @@ def create_thread( acreate_thread=acreate_thread, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore api_key = ( optional_params.api_key @@ -729,9 +667,7 @@ def create_thread( ) # type: ignore api_version: Optional[str] = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") + optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") ) # type: ignore extra_body = optional_params.get("extra_body", {}) @@ -767,9 +703,7 @@ def create_thread( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response # type: ignore @@ -874,14 +808,10 @@ def get_thread( aget_thread=aget_thread, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore api_version: Optional[str] = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") + optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") ) # type: ignore api_key = ( @@ -924,9 +854,7 @@ def get_thread( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response # type: ignore @@ -1000,9 +928,7 @@ def add_message( ) -> OpenAIMessage: ### COMMON OBJECTS ### a_add_message = kwargs.pop("a_add_message", None) - _message_data = MessageData( - role=role, content=content, attachments=attachments, metadata=metadata - ) + _message_data = MessageData(role=role, content=content, attachments=attachments, metadata=metadata) litellm_params_dict = get_litellm_params(**kwargs) optional_params = GenericLiteLLMParams(**kwargs) @@ -1065,14 +991,10 @@ def add_message( a_add_message=a_add_message, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore api_version: Optional[str] = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") + optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") ) # type: ignore api_key = ( @@ -1113,9 +1035,7 @@ def add_message( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) @@ -1228,14 +1148,10 @@ def get_messages( aget_messages=aget_messages, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore api_version: Optional[str] = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") + optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") ) # type: ignore api_key = ( @@ -1275,9 +1191,7 @@ def get_messages( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) @@ -1438,15 +1352,9 @@ def run_thread( event_handler=event_handler, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -1492,9 +1400,7 @@ def run_thread( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response # type: ignore diff --git a/litellm/assistants/utils.py b/litellm/assistants/utils.py index bde279602ad..f775c1b6508 100644 --- a/litellm/assistants/utils.py +++ b/litellm/assistants/utils.py @@ -43,11 +43,7 @@ def get_optional_params_add_message( "metadata": None, } - non_default_params = { - k: v - for k, v in passed_params.items() - if (k in default_params and v != default_params[k]) - } + non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])} optional_params = {} ## raise exception if non-default value passed for non-openai/azure embedding calls @@ -55,9 +51,7 @@ def get_optional_params_add_message( if len(non_default_params.keys()) > 0: keys = list(non_default_params.keys()) for k in keys: - if ( - litellm.drop_params is True and k not in supported_params - ): # drop the unsupported non-default values + if litellm.drop_params is True and k not in supported_params: # drop the unsupported non-default values non_default_params.pop(k, None) elif k not in supported_params: raise litellm.utils.UnsupportedParamsError( @@ -108,11 +102,7 @@ def get_optional_params_image_gen( "user": None, } - non_default_params = { - k: v - for k, v in passed_params.items() - if (k in default_params and v != default_params[k]) - } + non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])} optional_params = {} ## raise exception if non-default value passed for non-openai/azure embedding calls @@ -120,9 +110,7 @@ def get_optional_params_image_gen( if len(non_default_params.keys()) > 0: keys = list(non_default_params.keys()) for k in keys: - if ( - litellm.drop_params is True and k not in supported_params - ): # drop the unsupported non-default values + if litellm.drop_params is True and k not in supported_params: # drop the unsupported non-default values non_default_params.pop(k, None) elif k not in supported_params: raise UnsupportedParamsError( diff --git a/litellm/batch_completion/main.py b/litellm/batch_completion/main.py index 446e3f2f990..664977dc8d6 100644 --- a/litellm/batch_completion/main.py +++ b/litellm/batch_completion/main.py @@ -106,9 +106,7 @@ def batch_completion( original_kwargs = {} if "kwargs" in kwargs_modified: original_kwargs = kwargs_modified.pop("kwargs") - future = executor.submit( - litellm.completion, **kwargs_modified, **original_kwargs - ) + future = executor.submit(litellm.completion, **kwargs_modified, **original_kwargs) completions.append(future) # Retrieve the results from the futures @@ -153,13 +151,9 @@ def batch_completion_models(*args, **kwargs): futures = {} with ThreadPoolExecutor(max_workers=len(models)) as executor: for model in models: - futures[model] = executor.submit( - litellm.completion, *args, model=model, **kwargs - ) + futures[model] = executor.submit(litellm.completion, *args, model=model, **kwargs) - for model, future in sorted( - futures.items(), key=lambda x: models.index(x[0]) - ): + for model, future in sorted(futures.items(), key=lambda x: models.index(x[0])): if future.result() is not None: return future.result() elif "deployments" in kwargs: @@ -171,14 +165,10 @@ def batch_completion_models(*args, **kwargs): with ThreadPoolExecutor(max_workers=len(deployments)) as executor: for deployment in deployments: for key in kwargs.keys(): - if ( - key not in deployment - ): # don't override deployment values e.g. model name, api base, etc. + if key not in deployment: # don't override deployment values e.g. model name, api base, etc. deployment[key] = kwargs[key] kwargs = {**deployment, **nested_kwargs} - futures[deployment["model"]] = executor.submit( - litellm.completion, **kwargs - ) + futures[deployment["model"]] = executor.submit(litellm.completion, **kwargs) while futures: # wait for the first returned future @@ -191,9 +181,7 @@ def batch_completion_models(*args, **kwargs): return result except Exception: # if model 1 fails, continue with response from model 2, model3 - print_verbose( - "\n\ngot an exception, ignoring, removing from futures" - ) + print_verbose("\n\ngot an exception, ignoring, removing from futures") print_verbose(futures) new_futures = {} for key, value in futures.items(): @@ -254,10 +242,7 @@ def batch_completion_models_all_responses(*args, **kwargs): responses = [] with concurrent.futures.ThreadPoolExecutor(max_workers=len(models)) as executor: - futures = [ - executor.submit(litellm.completion, *args, model=model, **kwargs) - for model in models - ] + futures = [executor.submit(litellm.completion, *args, model=model, **kwargs) for model in models] for future in futures: try: @@ -265,9 +250,7 @@ def batch_completion_models_all_responses(*args, **kwargs): if result is not None: responses.append(result) except Exception as e: - print_verbose( - f"batch_completion_models_all_responses: model request failed: {str(e)}" - ) + print_verbose(f"batch_completion_models_all_responses: model request failed: {str(e)}") continue return responses diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index aeec58f1dfc..985198ce7ce 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -10,9 +10,7 @@ from litellm.utils import token_counter async def calculate_batch_cost_and_usage( file_content_dictionary: List[dict], - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" - ], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: Optional[str] = None, model_info: Optional[ModelInfo] = None, ) -> Tuple[float, Usage, List[str]]: @@ -36,18 +34,14 @@ async def calculate_batch_cost_and_usage( custom_llm_provider=custom_llm_provider, model_name=model_name, ) - batch_models = _get_batch_models_from_file_content( - file_content_dictionary, model_name - ) + batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name) return batch_cost, batch_usage, batch_models async def _handle_completed_batch( batch: Batch, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" - ], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: Optional[str] = None, litellm_params: Optional[dict] = None, ) -> Tuple[float, Usage, List[str]]: @@ -76,9 +70,7 @@ async def _handle_completed_batch( model_name=model_name, ) - batch_models = _get_batch_models_from_file_content( - file_content_dictionary, model_name - ) + batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name) return batch_cost, batch_usage, batch_models @@ -104,9 +96,7 @@ def _get_batch_models_from_file_content( def _batch_cost_calculator( file_content_dictionary: List[dict], - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", model_name: Optional[str] = None, model_info: Optional[ModelInfo] = None, ) -> float: @@ -118,9 +108,7 @@ def _batch_cost_calculator( and model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False) ): - batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage( - file_content_dictionary, model_name - ) + batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) verbose_logger.debug("vertex_ai_total_cost=%s", batch_cost) return batch_cost @@ -181,9 +169,7 @@ def calculate_vertex_ai_batch_cost_and_usage( ) total_cost += p_cost + c_cost except Exception as e: - verbose_logger.debug( - "vertex_ai batch cost calculation error for line: %s", str(e) - ) + verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e)) prompt_tokens += _prompt completion_tokens += _completion @@ -206,9 +192,7 @@ def calculate_vertex_ai_batch_cost_and_usage( async def _get_batch_output_file_content_as_dictionary( batch: Batch, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", litellm_params: Optional[dict] = None, ) -> List[dict]: """ @@ -235,12 +219,8 @@ async def _get_batch_output_file_content_as_dictionary( is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id) if is_base64_unified_file_id: try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split( - ";" - )[0] - verbose_logger.debug( - f"Extracted LLM output file ID from unified file ID: {file_id}" - ) + file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + verbose_logger.debug(f"Extracted LLM output file ID from unified file ID: {file_id}") except (IndexError, AttributeError) as e: verbose_logger.error( f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}" @@ -380,9 +360,7 @@ def _count_entry_tokens( def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", model_info: Optional[ModelInfo] = None, ) -> float: """ @@ -393,9 +371,7 @@ def _get_batch_job_cost_from_file_content( try: total_cost: float = 0.0 # parse the file content as json - verbose_logger.debug( - "file_content_dictionary=%s", json.dumps(file_content_dictionary, indent=4) - ) + verbose_logger.debug("file_content_dictionary=%s", json.dumps(file_content_dictionary, indent=4)) for _item in file_content_dictionary: if _batch_response_was_successful(_item): _response_body = _get_response_from_batch_job_output_file(_item) @@ -424,9 +400,7 @@ def _get_batch_job_cost_from_file_content( def _get_batch_job_total_usage_from_file_content( file_content_dictionary: List[dict], - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", model_name: Optional[str] = None, ) -> Usage: """ @@ -437,9 +411,7 @@ def _get_batch_job_total_usage_from_file_content( and model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False) ): - _, batch_usage = calculate_vertex_ai_batch_cost_and_usage( - file_content_dictionary, model_name - ) + _, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) return batch_usage # For other providers, use the existing logic @@ -488,11 +460,7 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int: # Nested pre-tokenized prompt: every int contributes a # token. Mixed string/int items still count. total += sum(1 if isinstance(t, int) else 0 for t in chunk) - total += sum( - token_counter(model=model, text=t) - for t in chunk - if isinstance(t, str) - ) + total += sum(token_counter(model=model, text=t) for t in chunk if isinstance(t, str)) return total return 0 diff --git a/litellm/batches/main.py b/litellm/batches/main.py index c5f1c86a4e1..3a2d9e13f77 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -79,11 +79,7 @@ def _resolve_timeout( Returns: Resolved timeout as float """ - timeout = ( - optional_params.timeout - or kwargs.get("request_timeout", default_timeout) - or default_timeout - ) + timeout = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout # Handle httpx.Timeout objects if isinstance(timeout, httpx.Timeout): @@ -109,9 +105,7 @@ async def acreate_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -161,9 +155,7 @@ def create_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -194,9 +186,7 @@ def create_batch( _is_async = kwargs.pop("acreate_batch", False) is True litellm_params = dict(GenericLiteLLMParams(**kwargs)) - litellm_logging_obj: LiteLLMLoggingObj = cast( - LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None) - ) + litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None)) ### TIMEOUT LOGIC ### timeout = _resolve_timeout(optional_params, kwargs, custom_llm_provider) litellm_logging_obj.update_from_kwargs( @@ -224,9 +214,7 @@ def create_batch( extra_body=extra_body, ) if output_expires_after is not None: - _create_batch_request["output_expires_after"] = cast( - FileExpiresAfter, output_expires_after - ) + _create_batch_request["output_expires_after"] = cast(FileExpiresAfter, output_expires_after) if model is not None: provider_config = ProviderConfigManager.get_provider_batches_config( model=model, @@ -244,12 +232,7 @@ def create_batch( api_key=optional_params.api_key, logging_obj=litellm_logging_obj, _is_async=_is_async, - client=( - client - if client is not None - and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) - else None - ), + client=(client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None), timeout=timeout, model=model, ) @@ -288,16 +271,8 @@ def create_batch( _is_async=_is_async, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) + api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") api_key = ( optional_params.api_key @@ -326,18 +301,12 @@ def create_batch( elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") + optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" + optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) + vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.create_batch( _is_async=_is_async, @@ -351,17 +320,13 @@ def create_batch( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support custom_llm_provider={} for 'create_batch'".format( - custom_llm_provider - ), + message="LiteLLM doesn't support custom_llm_provider={} for 'create_batch'".format(custom_llm_provider), model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_batch", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -372,9 +337,7 @@ def create_batch( @client async def aretrieve_batch( batch_id: str, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -420,9 +383,7 @@ def _handle_retrieve_batch_providers_without_provider_config( litellm_params: dict, _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", logging_obj: Optional[Any] = None, ): api_base: Optional[str] = None @@ -459,16 +420,8 @@ def _handle_retrieve_batch_providers_without_provider_config( max_retries=optional_params.max_retries, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) + api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") api_key = ( optional_params.api_key @@ -497,18 +450,12 @@ def _handle_retrieve_batch_providers_without_provider_config( elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") + optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" + optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) + vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.retrieve_batch( _is_async=_is_async, @@ -528,12 +475,7 @@ def _handle_retrieve_batch_providers_without_provider_config( or get_secret_str("ANTHROPIC_API_BASE") or get_secret_str("ANTHROPIC_BASE_URL") ) - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("ANTHROPIC_API_KEY") - ) + api_key = optional_params.api_key or litellm.api_key or litellm.azure_key or get_secret_str("ANTHROPIC_API_KEY") response = anthropic_batches_instance.retrieve_batch( _is_async=_is_async, @@ -555,9 +497,7 @@ def _handle_retrieve_batch_providers_without_provider_config( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="retrieve_batch", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -566,9 +506,7 @@ def _handle_retrieve_batch_providers_without_provider_config( @client def retrieve_batch( batch_id: str, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic" - ] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -581,9 +519,7 @@ def retrieve_batch( """ try: optional_params = GenericLiteLLMParams(**kwargs) - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj", None - ) + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 litellm_params = get_litellm_params( @@ -680,12 +616,7 @@ def retrieve_batch( function_id="batch_retrieve", ), _is_async=_is_async, - client=( - client - if client is not None - and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) - else None - ), + client=(client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None), timeout=timeout, model=model, ) @@ -823,16 +754,8 @@ def list_batches( max_retries=optional_params.max_retries, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) + api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") api_key = ( optional_params.api_key @@ -860,18 +783,12 @@ def list_batches( elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") + optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" + optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) + vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.list_batches( _is_async=_is_async, @@ -895,9 +812,7 @@ def list_batches( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -1014,17 +929,9 @@ def cancel_batch( or "https://api.openai.com/v1" ) organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None - ) - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") + optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None ) + api_key = optional_params.api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY") response = openai_batches_instance.cancel_batch( _is_async=_is_async, @@ -1036,16 +943,8 @@ def cancel_batch( max_retries=optional_params.max_retries, ) elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) + api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") api_key = ( optional_params.api_key @@ -1074,18 +973,12 @@ def cancel_batch( elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or None vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") + optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" + optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) + vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.cancel_batch( _is_async=_is_async, @@ -1107,9 +1000,7 @@ def cancel_batch( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="cancel_batch", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="cancel_batch", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -1117,9 +1008,7 @@ def cancel_batch( raise e -def _handle_async_invoke_status( - batch_id: str, aws_region_name: str, logging_obj=None, **kwargs -) -> "LiteLLMBatch": +def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj=None, **kwargs) -> "LiteLLMBatch": """ Handle async invoke status check for AWS Bedrock. @@ -1168,9 +1057,7 @@ def _handle_async_invoke_status( # Get output S3 URI safely output_s3_uri = "" try: - output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"][ - "s3Uri" - ] + output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] except (KeyError, TypeError): pass @@ -1186,15 +1073,12 @@ def _handle_async_invoke_status( failed_at, _, _, - ) = BedrockBatchesConfig()._parse_timestamps_and_status( - status_response, aws_status_raw - ) + ) = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw) result = LiteLLMBatch( id=status_response["invocationArn"], object="batch", status=normalized_status, - created_at=created_at - or int(time.time()), # Provide default timestamp if None + created_at=created_at or int(time.time()), # Provide default timestamp if None in_progress_at=in_progress_at, completed_at=completed_at, failed_at=failed_at, diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index 915c22b90ec..26f888c8077 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -62,9 +62,7 @@ class BudgetManager: # Load the user_dict from hosted db url = self.api_base + "/get_budget" data = {"project_name": self.project_name} - response = litellm.module_level_client.post( - url, headers=self.headers, json=data - ) + response = litellm.module_level_client.post(url, headers=self.headers, json=data) response = response.json() if response["status"] == "error": self.user_dict = {} # assume this means the user dict hasn't been stored yet @@ -91,9 +89,7 @@ class BudgetManager: elif duration == "yearly": duration_in_days = DAYS_IN_A_YEAR else: - raise ValueError( - """duration needs to be one of ["daily", "weekly", "monthly", "yearly"]""" - ) + raise ValueError("""duration needs to be one of ["daily", "weekly", "monthly", "yearly"]""") self.user_dict[user] = { "total_budget": total_budget, "duration": duration_in_days, @@ -106,9 +102,7 @@ class BudgetManager: def projected_cost(self, model: str, messages: list, user: str): text = "".join(message["content"] for message in messages) prompt_tokens = litellm.token_counter(model=model, text=text) - prompt_cost, _ = litellm.cost_per_token( - model=model, prompt_tokens=prompt_tokens, completion_tokens=0 - ) + prompt_cost, _ = litellm.cost_per_token(model=model, prompt_tokens=prompt_tokens, completion_tokens=0) current_cost = self.user_dict[user].get("current_cost", 0) projected_cost = prompt_cost + current_cost return projected_cost @@ -125,12 +119,8 @@ class BudgetManager: output_text: Optional[str] = None, ): if model and input_text and output_text: - prompt_tokens = litellm.token_counter( - model=model, messages=[{"role": "user", "content": input_text}] - ) - completion_tokens = litellm.token_counter( - model=model, messages=[{"role": "user", "content": output_text}] - ) + prompt_tokens = litellm.token_counter(model=model, messages=[{"role": "user", "content": input_text}]) + completion_tokens = litellm.token_counter(model=model, messages=[{"role": "user", "content": output_text}]) ( prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar, @@ -142,21 +132,15 @@ class BudgetManager: cost = prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar elif completion_obj: cost = litellm.completion_cost(completion_response=completion_obj) - model = completion_obj[ - "model" - ] # if this throws an error try, model = completion_obj['model'] + model = completion_obj["model"] # if this throws an error try, model = completion_obj['model'] else: raise ValueError( "Either a chat completion object or the text response needs to be passed in. Learn more - https://docs.litellm.ai/docs/budget_manager" ) - self.user_dict[user]["current_cost"] = cost + self.user_dict[user].get( - "current_cost", 0 - ) + self.user_dict[user]["current_cost"] = cost + self.user_dict[user].get("current_cost", 0) if "model_cost" in self.user_dict[user]: - self.user_dict[user]["model_cost"][model] = cost + self.user_dict[user][ - "model_cost" - ].get(model, 0) + self.user_dict[user]["model_cost"][model] = cost + self.user_dict[user]["model_cost"].get(model, 0) else: self.user_dict[user]["model_cost"] = {model: cost} @@ -198,9 +182,7 @@ class BudgetManager: current_time = time.time() # Convert duration from days to seconds - duration_in_seconds = ( - self.user_dict[user]["duration"] * HOURS_IN_A_DAY * 60 * 60 - ) + duration_in_seconds = self.user_dict[user]["duration"] * HOURS_IN_A_DAY * 60 * 60 # Check if duration has elapsed if current_time - last_updated_at >= duration_in_seconds: @@ -215,9 +197,7 @@ class BudgetManager: self.reset_on_duration(user) def _save_data_thread(self): - thread = threading.Thread( - target=self.save_data - ) # [Non-Blocking]: saves data without blocking execution + thread = threading.Thread(target=self.save_data) # [Non-Blocking]: saves data without blocking execution thread.start() def save_data(self): @@ -226,15 +206,11 @@ class BudgetManager: # save the user dict with open("user_cost.json", "w") as json_file: - json.dump( - self.user_dict, json_file, indent=4 - ) # Indent for pretty formatting + json.dump(self.user_dict, json_file, indent=4) # Indent for pretty formatting return {"status": "success"} elif self.client_type == "hosted": url = self.api_base + "/set_budget" data = {"project_name": self.project_name, "user_dict": self.user_dict} - response = litellm.module_level_client.post( - url, headers=self.headers, json=data - ) + response = litellm.module_level_client.post(url, headers=self.headers, json=data) response = response.json() return response diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py index 1ec898012e9..ec886b14020 100644 --- a/litellm/caching/_embedding_router.py +++ b/litellm/caching/_embedding_router.py @@ -27,9 +27,7 @@ def resolve_embedding_router( if llm_router is None: return None router_model_names: list[str] = ( - [m["model_name"] for m in llm_model_list if "model_name" in m] - if llm_model_list is not None - else [] + [m["model_name"] for m in llm_model_list if "model_name" in m] if llm_model_list is not None else [] ) if embedding_model in router_model_names: return llm_router diff --git a/litellm/caching/azure_blob_cache.py b/litellm/caching/azure_blob_cache.py index a2246640c30..fca7cf20313 100644 --- a/litellm/caching/azure_blob_cache.py +++ b/litellm/caching/azure_blob_cache.py @@ -52,9 +52,7 @@ class AzureBlobCache(BaseCache): print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}") serialized_value = json.dumps(value) try: - await self.async_container_client.upload_blob( - key, serialized_value, overwrite=True - ) + await self.async_container_client.upload_blob(key, serialized_value, overwrite=True) except Exception as e: # NON blocking - notify users Azure Blob is throwing an exception print_verbose(f"LiteLLM set_cache() - Got exception from Azure Blob: {e}") diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f2269d1945..34badaa3e8a 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -171,9 +171,7 @@ class Cache: # Check REDIS_CLUSTER_NODES env var if no explicit startup nodes if not redis_startup_nodes: _env_cluster_nodes = litellm.get_secret("REDIS_CLUSTER_NODES") - if _env_cluster_nodes is not None and isinstance( - _env_cluster_nodes, str - ): + if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str): redis_startup_nodes = json.loads(_env_cluster_nodes) if redis_startup_nodes: @@ -271,7 +269,9 @@ class Cache: litellm.logging_callback_manager.add_litellm_success_callback("cache") if "cache" not in litellm._async_success_callback: litellm.logging_callback_manager.add_litellm_async_success_callback("cache") - self.supported_call_types = supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"] + self.supported_call_types = ( + supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"] + ) self.type = type self.namespace = namespace self.redis_flush_size = redis_flush_size @@ -294,9 +294,7 @@ class Cache: # Params whose values carry prompt content. Excluded from semantic-cache # scope keys so differently worded prompts share a bucket and match via # vector similarity rather than being split into per-wording buckets. - _SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS: frozenset = frozenset( - {"messages", "prompt", "input"} - ) + _SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS: frozenset = frozenset({"messages", "prompt", "input"}) # Server-set identity (from proxy auth) used to isolate semantic-cache # buckets per tenant. Required once the prompt is out of the scope key, so a @@ -349,11 +347,7 @@ class Cache: combined_kwargs = ModelParamHelper._get_all_llm_api_params() litellm_param_kwargs = all_litellm_params is_semantic_cache = self._is_semantic_cache() - scope_excluded_params = ( - self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS - if is_semantic_cache - else frozenset() - ) + scope_excluded_params = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset() for param in kwargs: if param in scope_excluded_params: continue @@ -361,12 +355,8 @@ class Cache: param_value: Optional[str] = self._get_param_value(param, kwargs) if param_value is not None: cache_key += f"{str(param)}: {str(param_value)}" - elif ( - param not in litellm_param_kwargs - ): # check if user passed in optional param - e.g. top_k - if ( - litellm.enable_caching_on_provider_specific_optional_params is True - ): # feature flagged for now + elif param not in litellm_param_kwargs: # check if user passed in optional param - e.g. top_k + if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now if kwargs[param] is None: continue # ignore None params param_value = kwargs[param] @@ -385,9 +375,7 @@ class Cache: # Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError # when kwargs already contains preset_cache_key from upstream callers kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} - self._set_preset_cache_key_in_kwargs( - preset_cache_key=hashed_cache_key, **kwargs_for_preset - ) + self._set_preset_cache_key_in_kwargs(preset_cache_key=hashed_cache_key, **kwargs_for_preset) return hashed_cache_key def _get_param_value( @@ -415,15 +403,11 @@ class Cache: metadata: Dict = kwargs.get("metadata", {}) or {} litellm_params: Dict = kwargs.get("litellm_params", {}) or {} metadata_in_litellm_params: Dict = litellm_params.get("metadata", {}) or {} - model_group: Optional[str] = metadata.get( - "model_group" - ) or metadata_in_litellm_params.get("model_group") + model_group: Optional[str] = metadata.get("model_group") or metadata_in_litellm_params.get("model_group") caching_group = self._get_caching_group(metadata, model_group) return caching_group or model_group or kwargs["model"] - def _get_caching_group( - self, metadata: dict, model_group: Optional[str] - ) -> Optional[str]: + def _get_caching_group(self, metadata: dict, model_group: Optional[str]) -> Optional[str]: caching_groups: Optional[List] = metadata.get("caching_groups", []) if caching_groups: for group in caching_groups: @@ -503,11 +487,7 @@ class Cache: """ dynamic_cache_control: DynamicCacheControl = kwargs.get("cache", {}) metadata = kwargs.get("metadata") or {} - namespace = ( - dynamic_cache_control.get("namespace") - or metadata.get("redis_namespace") - or self.namespace - ) + namespace = dynamic_cache_control.get("namespace") or metadata.get("redis_namespace") or self.namespace if namespace: hash_hex = f"{namespace}:{hash_hex}" verbose_logger.debug("Final hashed key: %s", hash_hex) @@ -537,11 +517,7 @@ class Cache: Common get cache logic across sync + async implementations """ # Check if a timestamp was stored with the cached response - if ( - cached_result is not None - and isinstance(cached_result, dict) - and "timestamp" in cached_result - ): + if cached_result is not None and isinstance(cached_result, dict) and "timestamp" in cached_result: timestamp = cached_result["timestamp"] current_time = time.time() @@ -586,15 +562,11 @@ class Cache: ) -> None: original_metadata = original_kwargs.get("metadata") cache_lookup_metadata = cache_lookup_kwargs.get("metadata") - if not isinstance(original_metadata, dict) or not isinstance( - cache_lookup_metadata, dict - ): + if not isinstance(original_metadata, dict) or not isinstance(cache_lookup_metadata, dict): return if "semantic-similarity" in cache_lookup_metadata: - original_metadata["semantic-similarity"] = cache_lookup_metadata[ - "semantic-similarity" - ] + original_metadata["semantic-similarity"] = cache_lookup_metadata["semantic-similarity"] def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): """ @@ -616,34 +588,22 @@ class Cache: cache_key = self.get_cache_key(**kwargs) if cache_key is not None: cache_control_args: DynamicCacheControl = kwargs.get("cache", {}) - max_age = ( - cache_control_args.get("s-maxage") - or cache_control_args.get("s-max-age") - or float("inf") - ) + max_age = cache_control_args.get("s-maxage") or cache_control_args.get("s-max-age") or float("inf") cache_lookup_kwargs = self._get_safe_cache_lookup_kwargs(kwargs) if dynamic_cache_object is not None: - cached_result = dynamic_cache_object.get_cache( - cache_key, **cache_lookup_kwargs - ) + cached_result = dynamic_cache_object.get_cache(cache_key, **cache_lookup_kwargs) else: - cached_result = self.cache.get_cache( - cache_key, **cache_lookup_kwargs - ) + cached_result = self.cache.get_cache(cache_key, **cache_lookup_kwargs) self._update_metadata_from_cache_lookup_kwargs( original_kwargs=kwargs, cache_lookup_kwargs=cache_lookup_kwargs, ) - return self._get_cache_logic( - cached_result=cached_result, max_age=max_age - ) + return self._get_cache_logic(cached_result=cached_result, max_age=max_age) except Exception: print_verbose(f"An exception occurred: {traceback.format_exc()}") return None - async def async_get_cache( - self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs - ): + async def async_get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): """ Async get cache implementation. @@ -660,20 +620,12 @@ class Cache: cache_key = self.get_cache_key(**kwargs) if cache_key is not None: cache_control_args = kwargs.get("cache", {}) - max_age = cache_control_args.get( - "s-max-age", cache_control_args.get("s-maxage", float("inf")) - ) + max_age = cache_control_args.get("s-max-age", cache_control_args.get("s-maxage", float("inf"))) if dynamic_cache_object is not None: - cached_result = await dynamic_cache_object.async_get_cache( - cache_key, **kwargs - ) + cached_result = await dynamic_cache_object.async_get_cache(cache_key, **kwargs) else: - cached_result = await self.cache.async_get_cache( - cache_key, **kwargs - ) - return self._get_cache_logic( - cached_result=cached_result, max_age=max_age - ) + cached_result = await self.cache.async_get_cache(cache_key, **kwargs) + return self._get_cache_logic(cached_result=cached_result, max_age=max_age) except Exception: print_verbose(f"An exception occurred: {traceback.format_exc()}") return None @@ -722,16 +674,12 @@ class Cache: try: if self.should_use_cache(**kwargs) is not True: return - cache_key, cached_data, kwargs = self._add_cache_logic( - result=result, **kwargs - ) + cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) self.cache.set_cache(cache_key, cached_data, **kwargs) except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") - async def async_add_cache( - self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs - ): + async def async_add_cache(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): """ Async implementation of add_cache """ @@ -742,13 +690,9 @@ class Cache: # high traffic - fill in results in memory and then flush await self.batch_cache_write(result, **kwargs) else: - cache_key, cached_data, kwargs = self._add_cache_logic( - result=result, **kwargs - ) + cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) if dynamic_cache_object is not None: - await dynamic_cache_object.async_set_cache( - cache_key, cached_data, **kwargs - ) + await dynamic_cache_object.async_set_cache(cache_key, cached_data, **kwargs) else: await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: @@ -899,9 +843,7 @@ class Cache: ) return cache_key, cached_data, kwargs - async def async_add_cache_pipeline( - self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs - ): + async def async_add_cache_pipeline(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): """ Async implementation of add_cache for Embedding calls @@ -925,19 +867,13 @@ class Cache: ) = self.add_embedding_response_to_cache(result, i, kwargs, idx) cache_list.append((cache_key, cached_data)) elif isinstance(kwargs["input"], str): - cache_key, cached_data, kwargs = self.add_embedding_response_to_cache( - result, kwargs["input"], kwargs - ) + cache_key, cached_data, kwargs = self.add_embedding_response_to_cache(result, kwargs["input"], kwargs) cache_list.append((cache_key, cached_data)) if dynamic_cache_object is not None: - await dynamic_cache_object.async_set_cache_pipeline( - cache_list=cache_list, **kwargs - ) + await dynamic_cache_object.async_set_cache_pipeline(cache_list=cache_list, **kwargs) else: - await self.cache.async_set_cache_pipeline( - cache_list=cache_list, **kwargs - ) + await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 1ff4ee04080..c860f8e540d 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -163,11 +163,9 @@ class LLMCachingHandler: """ # Check if caching should be performed BEFORE doing expensive operations if ( - ( - (kwargs.get("caching", None) is None and litellm.cache is not None) - or kwargs.get("caching", False) is True - ) - and (kwargs.get("cache", {}).get("no-cache", False) is not True) + (kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True + ) and ( + kwargs.get("cache", {}).get("no-cache", False) is not True ): # allow users to control returning cached responses from the completion function args = args or () final_embedding_cached_response: Optional[EmbeddingResponse] = None @@ -183,9 +181,7 @@ class LLMCachingHandler: parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) kwargs["parent_otel_span"] = parent_otel_span - if litellm.cache is not None and self._is_call_type_supported_by_cache( - original_function=original_function - ): + if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): verbose_logger.debug("Checking Async Cache") cached_result = await self._retrieve_from_cache( call_type=call_type, @@ -204,9 +200,7 @@ class LLMCachingHandler: api_base=kwargs.get("api_base", None), api_key=kwargs.get("api_key", None), ) - cache_duration_ms = ( - cache_check_end_time - cache_check_start_time - ) * 1000 + cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000 self._update_litellm_logging_obj_environment( logging_obj=logging_obj, model=model, @@ -250,9 +244,7 @@ class LLMCachingHandler: and cached_result is not None and isinstance(cached_result, list) and litellm.cache is not None - and not isinstance( - litellm.cache.cache, S3Cache - ) # s3 doesn't support bulk writing. Exclude. + and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude. ): ( final_embedding_cached_response, @@ -291,9 +283,7 @@ class LLMCachingHandler: cached_result: Optional[Any] = None # Check if caching should be performed BEFORE doing expensive kwargs copy - if litellm.cache is not None and self._is_call_type_supported_by_cache( - original_function=original_function - ): + if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): args = args or () # Now that we confirmed caching will happen, prepare kwargs new_kwargs = kwargs.copy() @@ -376,9 +366,7 @@ class LLMCachingHandler: else: raise ValueError("input must be a string or a list") - def _extract_model_from_cached_results( - self, non_null_list: List[Tuple[int, CachedEmbedding]] - ) -> Optional[str]: + def _extract_model_from_cached_results(self, non_null_list: List[Tuple[int, CachedEmbedding]]) -> Optional[str]: """ Helper method to extract the model name from cached results. @@ -461,9 +449,7 @@ class LLMCachingHandler: elif isinstance(kwargs_input_as_list[idx], str): from litellm.utils import token_counter - prompt_tokens += token_counter( - text=kwargs_input_as_list[idx], count_response_tokens=True - ) + prompt_tokens += token_counter(text=kwargs_input_as_list[idx], count_response_tokens=True) # Aggregate prompt_tokens_details from cached items item_details = cr.get("prompt_tokens_details") if item_details: @@ -471,9 +457,7 @@ class LLMCachingHandler: aggregated_details = {} for key, value in item_details.items(): if isinstance(value, (int, float)): - aggregated_details[key] = ( - aggregated_details.get(key, 0) + value - ) + aggregated_details[key] = aggregated_details.get(key, 0) + value else: aggregated_details[key] = value @@ -483,9 +467,7 @@ class LLMCachingHandler: from litellm.types.utils import PromptTokensDetailsWrapper try: - prompt_tokens_details = PromptTokensDetailsWrapper( - **aggregated_details - ) + prompt_tokens_details = PromptTokensDetailsWrapper(**aggregated_details) except Exception: prompt_tokens_details = None usage = Usage( @@ -554,16 +536,8 @@ class LLMCachingHandler: if details2 is None: return details1 - dict1 = ( - details1.model_dump(exclude_none=True) - if hasattr(details1, "model_dump") - else {} - ) - dict2 = ( - details2.model_dump(exclude_none=True) - if hasattr(details2, "model_dump") - else {} - ) + dict1 = details1.model_dump(exclude_none=True) if hasattr(details1, "model_dump") else {} + dict2 = details2.model_dump(exclude_none=True) if hasattr(details2, "model_dump") else {} merged: dict = {} for key in set(dict1.keys()) | set(dict2.keys()): @@ -632,9 +606,7 @@ class LLMCachingHandler: final_data_list.append(item) _caching_handler_response.final_embedding_cached_response.data = final_data_list - _caching_handler_response.final_embedding_cached_response._hidden_params[ - "cache_hit" - ] = True + _caching_handler_response.final_embedding_cached_response._hidden_params["cache_hit"] = True _caching_handler_response.final_embedding_cached_response._response_ms = ( end_time - start_time ).total_seconds() * 1000 @@ -730,9 +702,7 @@ class LLMCachingHandler: raise ValueError("input must be a string or a list") tasks = [] for idx, i in enumerate(new_kwargs["input"]): - preset_cache_key = litellm.cache.get_cache_key( - **{**new_kwargs, "input": i} - ) + preset_cache_key = litellm.cache.get_cache_key(**{**new_kwargs, "input": i}) tasks.append( litellm.cache.async_get_cache( cache_key=preset_cache_key, @@ -750,18 +720,14 @@ class LLMCachingHandler: request_cache_key = request_kwargs.pop("cache_key", None) if litellm.cache._supports_async() is True: ## check if dual cache is supported ## - self.preset_cache_key = ( - request_cache_key or litellm.cache.get_cache_key(**request_kwargs) - ) + self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs) cached_result = await litellm.cache.async_get_cache( dynamic_cache_object=self.dual_cache, cache_key=self.preset_cache_key, **request_kwargs, ) else: # fallback for caches that don't support async - self.preset_cache_key = ( - request_cache_key or litellm.cache.get_cache_key(**request_kwargs) - ) + self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs) cached_result = litellm.cache.get_cache( dynamic_cache_object=self.dual_cache, cache_key=self.preset_cache_key, @@ -808,10 +774,9 @@ class LLMCachingHandler: """ from litellm.utils import convert_to_model_response_object - if ( - call_type == CallTypes.acompletion.value - or call_type == CallTypes.completion.value - ) and isinstance(cached_result, dict): + if (call_type == CallTypes.acompletion.value or call_type == CallTypes.completion.value) and isinstance( + cached_result, dict + ): if kwargs.get("stream", False) is True: cached_result = self._convert_cached_stream_response( cached_result=cached_result, @@ -825,8 +790,7 @@ class LLMCachingHandler: model_response_object=ModelResponse(), ) if ( - call_type == CallTypes.atext_completion.value - or call_type == CallTypes.text_completion.value + call_type == CallTypes.atext_completion.value or call_type == CallTypes.text_completion.value ) and isinstance(cached_result, dict): if kwargs.get("stream", False) is True: cached_result = self._convert_cached_stream_response( @@ -837,28 +801,26 @@ class LLMCachingHandler: ) else: cached_result = TextCompletionResponse(**cached_result) - elif ( - call_type == CallTypes.aembedding.value - or call_type == CallTypes.embedding.value - ) and isinstance(cached_result, dict): + elif (call_type == CallTypes.aembedding.value or call_type == CallTypes.embedding.value) and isinstance( + cached_result, dict + ): cached_result = convert_to_model_response_object( response_object=cached_result, model_response_object=EmbeddingResponse(), response_type="embedding", ) - elif ( - call_type == CallTypes.arerank.value or call_type == CallTypes.rerank.value - ) and isinstance(cached_result, dict): + elif (call_type == CallTypes.arerank.value or call_type == CallTypes.rerank.value) and isinstance( + cached_result, dict + ): cached_result = convert_to_model_response_object( response_object=cached_result, model_response_object=None, response_type="rerank", ) - elif ( - call_type == CallTypes.atranscription.value - or call_type == CallTypes.transcription.value - ) and isinstance(cached_result, dict): + elif (call_type == CallTypes.atranscription.value or call_type == CallTypes.transcription.value) and isinstance( + cached_result, dict + ): hidden_params = { "model": "whisper-1", "custom_llm_provider": custom_llm_provider, @@ -870,16 +832,12 @@ class LLMCachingHandler: response_type="audio_transcription", hidden_params=hidden_params, ) - elif (call_type == "aresponses" or call_type == "responses") and isinstance( - cached_result, dict - ): + elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict): use_chat_completion_cache = _is_chat_completion_cached_dict(cached_result) if use_chat_completion_cache: if kwargs.get("stream", False) is True: bridge_call_type = ( - CallTypes.acompletion.value - if call_type == "aresponses" - else CallTypes.completion.value + CallTypes.acompletion.value if call_type == "aresponses" else CallTypes.completion.value ) cached_result = self._convert_cached_stream_response( cached_result=cached_result, @@ -949,10 +907,7 @@ class LLMCachingHandler: ) _stream_cached_result: Union[AsyncGenerator, Generator] - if ( - call_type == CallTypes.acompletion.value - or call_type == CallTypes.atext_completion.value - ): + if call_type == CallTypes.acompletion.value or call_type == CallTypes.atext_completion.value: _stream_cached_result = convert_to_streaming_response_async( response_object=cached_result, ) @@ -1005,9 +960,7 @@ class LLMCachingHandler: parent_otel_span = _get_parent_otel_span_from_kwargs(new_kwargs) new_kwargs["parent_otel_span"] = parent_otel_span # [OPTIONAL] ADD TO CACHE - if self._should_store_result_in_cache( - original_function=original_function, kwargs=new_kwargs - ): + if self._should_store_result_in_cache(original_function=original_function, kwargs=new_kwargs): if ( isinstance(result, litellm.ModelResponse) or isinstance(result, litellm.EmbeddingResponse) @@ -1018,9 +971,7 @@ class LLMCachingHandler: if ( isinstance(result, EmbeddingResponse) and litellm.cache is not None - and not isinstance( - litellm.cache.cache, S3Cache - ) # s3 doesn't support bulk writing. Exclude. + and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude. ): asyncio.create_task( litellm.cache.async_add_cache_pipeline( @@ -1058,16 +1009,12 @@ class LLMCachingHandler: if litellm.cache is None: return - if self._should_store_result_in_cache( - original_function=self.original_function, kwargs=new_kwargs - ): + if self._should_store_result_in_cache(original_function=self.original_function, kwargs=new_kwargs): litellm.cache.add_cache(result, **new_kwargs) return - def _should_store_result_in_cache( - self, original_function: Callable, kwargs: Dict[str, Any] - ) -> bool: + def _should_store_result_in_cache(self, original_function: Callable, kwargs: Dict[str, Any]) -> bool: """ Helper function to determine if the result should be stored in the cache. @@ -1113,15 +1060,15 @@ class LLMCachingHandler: """ - complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = _assemble_complete_response_from_streaming_chunks( - result=processed_chunk, - start_time=self.start_time, - end_time=datetime.datetime.now(), - request_kwargs=self.request_kwargs, - streaming_chunks=self.async_streaming_chunks, - is_async=True, + complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = ( + _assemble_complete_response_from_streaming_chunks( + result=processed_chunk, + start_time=self.start_time, + end_time=datetime.datetime.now(), + request_kwargs=self.request_kwargs, + streaming_chunks=self.async_streaming_chunks, + is_async=True, + ) ) # if a complete_streaming_response is assembled, add it to the cache if complete_streaming_response is not None: @@ -1135,15 +1082,15 @@ class LLMCachingHandler: """ Sync internal method to add the streaming response to the cache """ - complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = _assemble_complete_response_from_streaming_chunks( - result=processed_chunk, - start_time=self.start_time, - end_time=datetime.datetime.now(), - request_kwargs=self.request_kwargs, - streaming_chunks=self.sync_streaming_chunks, - is_async=False, + complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = ( + _assemble_complete_response_from_streaming_chunks( + result=processed_chunk, + start_time=self.start_time, + end_time=datetime.datetime.now(), + request_kwargs=self.request_kwargs, + streaming_chunks=self.sync_streaming_chunks, + is_async=False, + ) ) # if a complete_streaming_response is assembled, add it to the cache @@ -1191,9 +1138,7 @@ class LLMCachingHandler: } if litellm.cache is not None: - litellm_params["preset_cache_key"] = ( - litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) - ) + litellm_params["preset_cache_key"] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) else: litellm_params["preset_cache_key"] = None @@ -1202,11 +1147,7 @@ class LLMCachingHandler: user=kwargs.get("user", None), optional_params={}, litellm_params=litellm_params, - input=( - kwargs.get("messages", "") - if not is_embedding - else kwargs.get("input", "") - ), + input=(kwargs.get("messages", "") if not is_embedding else kwargs.get("input", "")), api_key=kwargs.get("api_key", None), original_response=str(cached_result), additional_args=None, diff --git a/litellm/caching/disk_cache.py b/litellm/caching/disk_cache.py index e32c29b3bc6..b51acbe9cfd 100644 --- a/litellm/caching/disk_cache.py +++ b/litellm/caching/disk_cache.py @@ -16,9 +16,7 @@ class DiskCache(BaseCache): try: import diskcache as dc except ModuleNotFoundError as e: - raise ModuleNotFoundError( - "Please install litellm with `litellm[caching]` to use disk caching." - ) from e + raise ModuleNotFoundError("Please install litellm with `litellm[caching]` to use disk caching.") from e # if users don't provider one, use the default litellm cache if disk_cache_dir is None: diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 8060a65b78d..be618815a53 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -69,23 +69,15 @@ class DualCache(BaseCache): self.in_memory_cache = in_memory_cache or InMemoryCache() # If redis_cache is not provided, use the default RedisCache self.redis_cache = redis_cache - self.last_redis_batch_access_time = LimitedSizeOrderedDict( - max_size=default_max_redis_batch_cache_size - ) + self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size) self._last_redis_batch_access_time_lock = Lock() self.redis_batch_cache_expiry = ( - default_redis_batch_cache_expiry - or litellm.default_redis_batch_cache_expiry - or 10 - ) - self.default_in_memory_ttl = ( - default_in_memory_ttl or litellm.default_in_memory_ttl + default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10 ) + self.default_in_memory_ttl = default_in_memory_ttl or litellm.default_in_memory_ttl self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl - def update_cache_ttl( - self, default_in_memory_ttl: Optional[float], default_redis_ttl: Optional[float] - ): + def update_cache_ttl(self, default_in_memory_ttl: Optional[float], default_redis_ttl: Optional[float]): if default_in_memory_ttl is not None: self.default_in_memory_ttl = default_in_memory_ttl @@ -125,9 +117,7 @@ class DualCache(BaseCache): except Exception as e: print_verbose(e) - def increment_cache( - self, key, value: int, local_only: bool = False, **kwargs - ) -> int: + def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> int: """ Key - the key in cache @@ -166,9 +156,7 @@ class DualCache(BaseCache): if result is None and self.redis_cache is not None and local_only is False: # If not found in in-memory cache, try fetching from Redis - redis_result = self.redis_cache.get_cache( - key, parent_otel_span=parent_otel_span - ) + redis_result = self.redis_cache.get_cache(key, parent_otel_span=parent_otel_span) if redis_result is not None: # Update in-memory cache with the value from Redis @@ -196,9 +184,7 @@ class DualCache(BaseCache): new_loop = asyncio.new_event_loop() try: asyncio.set_event_loop(new_loop) - return new_loop.run_until_complete( - self.async_batch_get_cache(**received_args) - ) + return new_loop.run_until_complete(self.async_batch_get_cache(**received_args)) finally: new_loop.close() asyncio.set_event_loop(None) @@ -225,14 +211,10 @@ class DualCache(BaseCache): ): # Try to fetch from in-memory cache first try: - print_verbose( - f"async get cache: cache key: {key}; local_only: {local_only}" - ) + print_verbose(f"async get cache: cache key: {key}; local_only: {local_only}") result = None if self.in_memory_cache is not None: - in_memory_result = await self.in_memory_cache.async_get_cache( - key, **kwargs - ) + in_memory_result = await self.in_memory_cache.async_get_cache(key, **kwargs) print_verbose(f"in_memory_result: {in_memory_result}") if in_memory_result is not None: @@ -240,15 +222,11 @@ class DualCache(BaseCache): if result is None and self.redis_cache is not None and local_only is False: # If not found in in-memory cache, try fetching from Redis - redis_result = await self.redis_cache.async_get_cache( - key, parent_otel_span=parent_otel_span - ) + redis_result = await self.redis_cache.async_get_cache(key, parent_otel_span=parent_otel_span) if redis_result is not None: # Update in-memory cache with the value from Redis - await self.in_memory_cache.async_set_cache( - key, redis_result, **kwargs - ) + await self.in_memory_cache.async_set_cache(key, redis_result, **kwargs) result = redis_result @@ -277,20 +255,15 @@ class DualCache(BaseCache): if ( key not in self.last_redis_batch_access_time - or current_time - self.last_redis_batch_access_time[key] - >= self.redis_batch_cache_expiry + or current_time - self.last_redis_batch_access_time[key] >= self.redis_batch_cache_expiry ): sublist_keys.append(key) - previous_access_times[key] = self.last_redis_batch_access_time.get( - key - ) + previous_access_times[key] = self.last_redis_batch_access_time.get(key) self.last_redis_batch_access_time[key] = current_time return sublist_keys, previous_access_times - def _rollback_redis_batch_key_reservations( - self, previous_access_times: Dict[str, Optional[float]] - ) -> None: + def _rollback_redis_batch_key_reservations(self, previous_access_times: Dict[str, Optional[float]]) -> None: with self._last_redis_batch_access_time_lock: for key, previous_time in previous_access_times.items(): if previous_time is None: @@ -308,9 +281,7 @@ class DualCache(BaseCache): try: result = [None] * len(keys) if self.in_memory_cache is not None: - in_memory_result = await self.in_memory_cache.async_batch_get_cache( - keys, **kwargs - ) + in_memory_result = await self.in_memory_cache.async_batch_get_cache(keys, **kwargs) if in_memory_result is not None: result = in_memory_result @@ -321,9 +292,7 @@ class DualCache(BaseCache): - check the redis cache """ current_time = time.time() - sublist_keys, previous_access_times = self._reserve_redis_batch_keys( - current_time, keys, result - ) + sublist_keys, previous_access_times = self._reserve_redis_batch_keys(current_time, keys, result) # Only hit Redis if enough time has passed since last access. if len(sublist_keys) > 0: @@ -334,15 +303,11 @@ class DualCache(BaseCache): ) except Exception: # Do not throttle subsequent callers if the Redis read fails. - self._rollback_redis_batch_key_reservations( - previous_access_times - ) + self._rollback_redis_batch_key_reservations(previous_access_times) raise # Short-circuit if redis_result is None or contains only None values - if redis_result is None or all( - v is None for v in redis_result.values() - ): + if redis_result is None or all(v is None for v in redis_result.values()): return result # Pre-compute key-to-index mapping for O(1) lookup @@ -353,18 +318,14 @@ class DualCache(BaseCache): result[key_to_index[key]] = value if value is not None and self.in_memory_cache is not None: - await self.in_memory_cache.async_set_cache( - key, value, **kwargs - ) + await self.in_memory_cache.async_set_cache(key, value, **kwargs) return result except Exception: verbose_logger.error(traceback.format_exc()) async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): - print_verbose( - f"async set cache: cache key: {key}; local_only: {local_only}; value: {value}" - ) + print_verbose(f"async set cache: cache key: {key}; local_only: {local_only}; value: {value}") try: if self.in_memory_cache is not None: if "ttl" not in kwargs and self.default_in_memory_ttl is not None: @@ -374,36 +335,26 @@ class DualCache(BaseCache): if self.redis_cache is not None and local_only is False: await self.redis_cache.async_set_cache(key, value, **kwargs) except Exception as e: - verbose_logger.exception( - f"LiteLLM Cache: Excepton async add_cache: {str(e)}" - ) + verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {str(e)}") # async_batch_set_cache - async def async_set_cache_pipeline( - self, cache_list: list, local_only: bool = False, **kwargs - ): + async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs): """ Batch write values to the cache """ - print_verbose( - f"async batch set cache: cache keys: {cache_list}; local_only: {local_only}" - ) + print_verbose(f"async batch set cache: cache keys: {cache_list}; local_only: {local_only}") try: if self.in_memory_cache is not None: if "ttl" not in kwargs and self.default_in_memory_ttl is not None: kwargs["ttl"] = self.default_in_memory_ttl - await self.in_memory_cache.async_set_cache_pipeline( - cache_list=cache_list, **kwargs - ) + await self.in_memory_cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) if self.redis_cache is not None and local_only is False: await self.redis_cache.async_set_cache_pipeline( cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs ) except Exception as e: - verbose_logger.exception( - f"LiteLLM Cache: Excepton async add_cache: {str(e)}" - ) + verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {str(e)}") async def async_increment_cache( self, @@ -428,9 +379,7 @@ class DualCache(BaseCache): result: Optional[float] = None try: if self.in_memory_cache is not None: - result = await self.in_memory_cache.async_increment( - key, value, **kwargs - ) + result = await self.in_memory_cache.async_increment(key, value, **kwargs) if self.redis_cache is not None and local_only is False: result = await self.redis_cache.async_increment( @@ -478,9 +427,7 @@ class DualCache(BaseCache): ) return result - async def async_set_cache_sadd( - self, key, value: List, local_only: bool = False, **kwargs - ) -> None: + async def async_set_cache_sadd(self, key, value: List, local_only: bool = False, **kwargs) -> None: """ Add value to a set @@ -492,14 +439,10 @@ class DualCache(BaseCache): """ try: if self.in_memory_cache is not None: - _ = await self.in_memory_cache.async_set_cache_sadd( - key, value, ttl=kwargs.get("ttl", None) - ) + _ = await self.in_memory_cache.async_set_cache_sadd(key, value, ttl=kwargs.get("ttl", None)) if self.redis_cache is not None and local_only is False: - _ = await self.redis_cache.async_set_cache_sadd( - key, value, ttl=kwargs.get("ttl", None) - ) + _ = await self.redis_cache.async_set_cache_sadd(key, value, ttl=kwargs.get("ttl", None)) return None except Exception as e: diff --git a/litellm/caching/gcs_cache.py b/litellm/caching/gcs_cache.py index 0e6a111eb2b..3345f8fc5eb 100644 --- a/litellm/caching/gcs_cache.py +++ b/litellm/caching/gcs_cache.py @@ -26,15 +26,10 @@ class GCSCache(BaseCache): ) -> None: super().__init__() self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME - self.path_service_account = ( - path_service_account - or GCSBucketBase(bucket_name=None).path_service_account_json - ) + self.path_service_account = path_service_account or GCSBucketBase(bucket_name=None).path_service_account_json self.key_prefix = gcs_path.rstrip("/") + "/" if gcs_path else "" # create httpx clients - self.async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_client = _get_httpx_client() def _construct_headers(self) -> dict: @@ -64,9 +59,7 @@ class GCSCache(BaseCache): data = json.dumps(value) await self.async_client.post(url=url, data=data, headers=headers) except Exception as e: - print_verbose( - f"GCS Caching: async_set_cache() - Got exception from GCS: {e}" - ) + print_verbose(f"GCS Caching: async_set_cache() - Got exception from GCS: {e}") def get_cache(self, key, **kwargs): try: @@ -83,9 +76,7 @@ class GCSCache(BaseCache): return cached_response return None except Exception as e: - verbose_logger.error( - f"GCS Caching: get_cache() - Got exception from GCS: {e}" - ) + verbose_logger.error(f"GCS Caching: get_cache() - Got exception from GCS: {e}") async def async_get_cache(self, key, **kwargs): try: @@ -98,9 +89,7 @@ class GCSCache(BaseCache): return json.loads(response.text) return None except Exception as e: - verbose_logger.error( - f"GCS Caching: async_get_cache() - Got exception from GCS: {e}" - ) + verbose_logger.error(f"GCS Caching: async_get_cache() - Got exception from GCS: {e}") def flush_cache(self): pass diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index ba446dd4f60..2ad3f3f11b7 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -40,9 +40,7 @@ class InMemoryCache(BaseCache): max_size_in_memory if max_size_in_memory is not None else 200 ) # set an upper bound of 200 items in-memory self.default_ttl = default_ttl or 600 - self.max_size_per_item = ( - max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB - ) # 1MB = 1024KB + self.max_size_per_item = max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB # 1MB = 1024KB # in-memory cache self.cache_dict: dict = {} @@ -58,8 +56,7 @@ class InMemoryCache(BaseCache): # Fast path for common primitive types that are typically small if ( isinstance(value, (bool, int, float, str)) - and len(str(value)) - < self.max_size_per_item * MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB + and len(str(value)) < self.max_size_per_item * MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB ): # Conservative estimate return True @@ -73,9 +70,7 @@ class InMemoryCache(BaseCache): return size <= self.max_size_per_item # Fallback for complex types - if isinstance(value, BaseModel) and hasattr( - value, "model_dump" - ): # Pydantic v2 + if isinstance(value, BaseModel) and hasattr(value, "model_dump"): # Pydantic v2 value = value.model_dump() elif hasattr(value, "isoformat"): # datetime objects return True # datetime strings are always small @@ -257,9 +252,7 @@ class InMemoryCache(BaseCache): ) -> Optional[List[float]]: results = [] for increment in increment_list: - result = await self.async_increment( - increment["key"], increment["increment_value"], **kwargs - ) + result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs) results.append(result) return results diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 504ef8a54eb..5ed1bb47eba 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -51,34 +51,24 @@ class QdrantSemanticCache(BaseCache): raise Exception("collection_name must be provided, passed None") self.collection_name = collection_name - print_verbose( - f"qdrant semantic-cache initializing COLLECTION - {self.collection_name}" - ) + print_verbose(f"qdrant semantic-cache initializing COLLECTION - {self.collection_name}") if similarity_threshold is None: raise Exception("similarity_threshold must be provided, passed None") self.similarity_threshold = similarity_threshold self.embedding_model = embedding_model - self.vector_size = ( - vector_size if vector_size is not None else QDRANT_VECTOR_SIZE - ) + self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE headers = {} # check if defined as os.environ/ variable if qdrant_api_base: - if isinstance(qdrant_api_base, str) and qdrant_api_base.startswith( - "os.environ/" - ): + if isinstance(qdrant_api_base, str) and qdrant_api_base.startswith("os.environ/"): qdrant_api_base = get_secret_str(qdrant_api_base) if qdrant_api_key: - if isinstance(qdrant_api_key, str) and qdrant_api_key.startswith( - "os.environ/" - ): + if isinstance(qdrant_api_key, str) and qdrant_api_key.startswith("os.environ/"): qdrant_api_key = get_secret_str(qdrant_api_key) - qdrant_api_base = ( - qdrant_api_base or os.getenv("QDRANT_URL") or os.getenv("QDRANT_API_BASE") - ) + qdrant_api_base = qdrant_api_base or os.getenv("QDRANT_URL") or os.getenv("QDRANT_API_BASE") qdrant_api_key = qdrant_api_key or os.getenv("QDRANT_API_KEY") headers = {"Content-Type": "application/json"} if qdrant_api_key: @@ -94,22 +84,16 @@ class QdrantSemanticCache(BaseCache): self.headers = headers self.sync_client = _get_httpx_client() - self.async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.Caching - ) + self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Caching) if quantization_config is None: - print_verbose( - "Quantization config is not provided. Default binary quantization will be used." - ) + print_verbose("Quantization config is not provided. Default binary quantization will be used.") collection_exists = self.sync_client.get( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/exists", headers=self.headers, ) if collection_exists.status_code != 200: - raise ValueError( - f"Error from qdrant checking if /collections exist {collection_exists.text}" - ) + raise ValueError(f"Error from qdrant checking if /collections exist {collection_exists.text}") if collection_exists.json()["result"]["exists"]: collection_details = self.sync_client.get( @@ -117,9 +101,7 @@ class QdrantSemanticCache(BaseCache): headers=self.headers, ) self.collection_info = collection_details.json() - print_verbose( - f"Collection already exists.\nCollection details:{self.collection_info}" - ) + print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}") self._ensure_cache_key_payload_index() else: quantization_params: Dict[str, Any] @@ -138,13 +120,9 @@ class QdrantSemanticCache(BaseCache): } } elif quantization_config == "product": - quantization_params = { - "product": {"compression": "x16", "always_ram": False} - } + quantization_params = {"product": {"compression": "x16", "always_ram": False}} else: - raise Exception( - "Quantization config must be one of 'scalar', 'binary' or 'product'" - ) + raise Exception("Quantization config must be one of 'scalar', 'binary' or 'product'") new_collection_status = self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}", @@ -160,9 +138,7 @@ class QdrantSemanticCache(BaseCache): headers=self.headers, ) self.collection_info = collection_details.json() - print_verbose( - f"New collection created.\nCollection details:{self.collection_info}" - ) + print_verbose(f"New collection created.\nCollection details:{self.collection_info}") self._ensure_cache_key_payload_index() else: raise Exception("Error while creating new collection") @@ -171,9 +147,7 @@ class QdrantSemanticCache(BaseCache): if cached_response is None: return cached_response try: - cached_response = json.loads( - cached_response - ) # Convert string to dictionary + cached_response = json.loads(cached_response) # Convert string to dictionary except Exception: cached_response = ast.literal_eval(cached_response) return cached_response @@ -202,15 +176,9 @@ class QdrantSemanticCache(BaseCache): }, ) if response.status_code not in (200, 201): - print_verbose( - "Qdrant semantic-cache could not create cache-key payload index: " - f"{response.text}" - ) + print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {response.text}") except Exception as exc: - print_verbose( - "Qdrant semantic-cache could not create cache-key payload index: " - f"{str(exc)}" - ) + print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {str(exc)}") def _payload_matches_cache_key(self, payload: dict, key: str) -> bool: # Pre-isolation points stored only prompt + response with no cache-key @@ -220,9 +188,7 @@ class QdrantSemanticCache(BaseCache): cached_key = payload.get(self.CACHE_KEY_FIELD_NAME) return cached_key is not None and str(cached_key) == str(key) - def _get_embedding( - self, prompt: str, metadata: Dict[str, Any] | None = None - ) -> EmbeddingResponse: + def _get_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse: """Embed via the proxy Router when it serves the model, else direct.""" try: from litellm.proxy.proxy_server import llm_model_list, llm_router @@ -230,9 +196,7 @@ class QdrantSemanticCache(BaseCache): llm_model_list = None llm_router = None - router = resolve_embedding_router( - self.embedding_model, llm_router, llm_model_list - ) + router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) if router is not None: return router.embedding( model=self.embedding_model, @@ -246,18 +210,14 @@ class QdrantSemanticCache(BaseCache): cache={"no-store": True, "no-cache": True}, ) - async def _get_async_embedding( - self, prompt: str, metadata: Dict[str, Any] | None = None - ) -> EmbeddingResponse: + async def _get_async_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse: try: from litellm.proxy.proxy_server import llm_model_list, llm_router except ImportError: llm_model_list = None llm_router = None - router = resolve_embedding_router( - self.embedding_model, llm_router, llm_model_list - ) + router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) if router is not None: return await router.aembedding( model=self.embedding_model, @@ -394,9 +354,7 @@ class QdrantSemanticCache(BaseCache): # get the prompt messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding( - prompt, metadata=kwargs.get("metadata") - ) + embedding_response = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) # get the embedding embedding = embedding_response["data"][0]["embedding"] @@ -432,9 +390,7 @@ class QdrantSemanticCache(BaseCache): messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding( - prompt, metadata=kwargs.get("metadata") - ) + embedding_response = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) # get the embedding embedding = embedding_response["data"][0]["embedding"] diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 594417173ed..dd1c152a421 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -153,8 +153,7 @@ class RedisCircuitBreaker: if self._failure_count >= self.failure_threshold: if self._state != self.OPEN: verbose_logger.warning( - "Redis circuit breaker OPENED after %d consecutive failures — " - "fast-failing Redis calls for %ds", + "Redis circuit breaker OPENED after %d consecutive failures — fast-failing Redis calls for %ds", self._failure_count, self.recovery_timeout, ) @@ -179,9 +178,7 @@ def _redis_circuit_breaker_guard(method): # type: ignore @functools.wraps(method) async def wrapper(self, *args, **kwargs): # type: ignore if self._circuit_breaker.is_open(): - raise Exception( - f"Redis circuit breaker is open — skipping {method.__name__}" - ) + raise Exception(f"Redis circuit breaker is open — skipping {method.__name__}") try: result = await method(self, *args, **kwargs) self._circuit_breaker.record_success() @@ -233,9 +230,7 @@ class RedisCache(BaseCache): redis_kwargs.update(kwargs) self.redis_client = get_redis_client(**redis_kwargs) - self.redis_async_client: Optional[ - Union[async_redis_client, async_redis_cluster_client] - ] = None + self.redis_async_client: Optional[Union[async_redis_client, async_redis_cluster_client]] = None self.redis_kwargs = redis_kwargs self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs) @@ -274,9 +269,7 @@ class RedisCache(BaseCache): _ = asyncio.get_running_loop().create_task(self.ping()) except Exception as e: if "no running event loop" in str(e): - verbose_logger.debug( - "Ignoring async redis ping. No running event loop." - ) + verbose_logger.debug("Ignoring async redis ping. No running event loop.") else: verbose_logger.error( "Error connecting to Async Redis client - {}".format(str(e)), @@ -289,9 +282,7 @@ class RedisCache(BaseCache): if hasattr(self.redis_client, "ping"): self.redis_client.ping() # type: ignore except Exception as e: - verbose_logger.error( - "Error connecting to Sync Redis client", extra={"error": str(e)} - ) + verbose_logger.error("Error connecting to Sync Redis client", extra={"error": str(e)}) self._handle_sync_ping_error(e) def _handle_async_ping_error(self, e: Exception): @@ -350,18 +341,12 @@ class RedisCache(BaseCache): cache_key = self._get_async_client_cache_key() cached_client = in_memory_llm_clients_cache.get_cache(key=cache_key) if cached_client is not None: - redis_async_client = cast( - Union[async_redis_client, async_redis_cluster_client], cached_client - ) + redis_async_client = cast(Union[async_redis_client, async_redis_cluster_client], cached_client) else: # Create new connection pool and client for current event loop self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) - redis_async_client = get_redis_async_client( - connection_pool=self.async_redis_conn_pool, **self.redis_kwargs - ) - in_memory_llm_clients_cache.set_cache( - key=cache_key, value=redis_async_client - ) + redis_async_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs) + in_memory_llm_clients_cache.set_cache(key=cache_key, value=redis_async_client) self.redis_async_client = redis_async_client # type: ignore return redis_async_client @@ -408,9 +393,7 @@ class RedisCache(BaseCache): def set_cache(self, key, value, **kwargs): ttl = self.get_ttl(**kwargs) - print_verbose( - f"Set Redis Cache: key: {key}\nValue {value}\nttl={ttl}, redis_version={self.redis_version}" - ) + print_verbose(f"Set Redis Cache: key: {key}\nValue {value}\nttl={ttl}, redis_version={self.redis_version}") key = self.check_and_fix_namespace(key=key) try: start_time = time.time() @@ -426,13 +409,9 @@ class RedisCache(BaseCache): ) except Exception as e: # NON blocking - notify users Redis is throwing an exception - print_verbose( - f"litellm.caching.caching: set() - Got exception from REDIS : {str(e)}" - ) + print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {str(e)}") - def increment_cache( - self, key, value: int, ttl: Optional[float] = None, **kwargs - ) -> int: + def increment_cache(self, key, value: int, ttl: Optional[float] = None, **kwargs) -> int: _redis_client = self.redis_client start_time = time.time() set_ttl = self.get_ttl(ttl=ttl) @@ -571,24 +550,18 @@ class RedisCache(BaseCache): f"{self.namespace}-{hashlib.sha256(script.encode()).hexdigest()[:16]}" ) - async def run_script( - keys: Sequence[str], args: Sequence[Any], client: Any = None - ) -> Any: - executor: Optional[Callable[..., Awaitable[Any]]] = ( - litellm.in_memory_llm_clients_cache.get_cache(key=script_cache_key) + async def run_script(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any: + executor: Optional[Callable[..., Awaitable[Any]]] = litellm.in_memory_llm_clients_cache.get_cache( + key=script_cache_key ) if executor is None: executor = self._register_script_for_current_loop(script) - litellm.in_memory_llm_clients_cache.set_cache( - key=script_cache_key, value=executor - ) + litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor) return await executor(keys=keys, args=args, client=client) return run_script - def _register_script_for_current_loop( - self, script: str - ) -> Callable[..., Awaitable[Any]]: + def _register_script_for_current_loop(self, script: str) -> Callable[..., Awaitable[Any]]: """ Register the script against the current event loop's Redis client. @@ -599,30 +572,18 @@ class RedisCache(BaseCache): if hasattr(_redis_client, "register_script"): registered_script = _redis_client.register_script(script) - async def standalone_executor( - keys: Sequence[str], args: Sequence[Any], client: Any = None - ) -> Any: - namespaced_keys = tuple( - self.check_and_fix_namespace(key=key) for key in keys - ) - return await registered_script( - keys=namespaced_keys, args=args, client=client - ) + async def standalone_executor(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any: + namespaced_keys = tuple(self.check_and_fix_namespace(key=key) for key in keys) + return await registered_script(keys=namespaced_keys, args=args, client=client) return standalone_executor if hasattr(_redis_client, "script_load"): script_sha = _redis_client.script_load(script) - async def cluster_executor( - keys: Sequence[str], args: Sequence[Any], client: Any = None - ) -> Any: - namespaced_keys = tuple( - self.check_and_fix_namespace(key=key) for key in keys - ) - return await _redis_client.evalsha( - script_sha, len(namespaced_keys), *namespaced_keys, *args - ) + async def cluster_executor(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any: + namespaced_keys = tuple(self.check_and_fix_namespace(key=key) for key in keys) + return await _redis_client.evalsha(script_sha, len(namespaced_keys), *namespaced_keys, *args) return cluster_executor @@ -678,9 +639,7 @@ class RedisCache(BaseCache): nx=nx, ex=ttl, ) - print_verbose( - f"Successfully Set ASYNC Redis Cache: key: {key}\nValue {value}\nttl={ttl}" - ) + print_verbose(f"Successfully Set ASYNC Redis Cache: key: {key}\nValue {value}\nttl={ttl}") end_time = time.time() _duration = end_time - start_time asyncio.create_task( @@ -729,9 +688,7 @@ class RedisCache(BaseCache): # Iterate through each key-value pair in the cache_list and set them in the pipeline. for cache_key, cache_value in cache_list: cache_key = self.check_and_fix_namespace(key=cache_key) - print_verbose( - f"Set ASYNC Redis Cache PIPELINE: key: {cache_key}\nValue {cache_value}\nttl={ttl}" - ) + print_verbose(f"Set ASYNC Redis Cache PIPELINE: key: {cache_key}\nValue {cache_value}\nttl={ttl}") json_cache_value = json.dumps(cache_value) # Set the value with a TTL if it's provided. _td: Optional[timedelta] = None @@ -747,9 +704,7 @@ class RedisCache(BaseCache): return results @_redis_circuit_breaker_guard - async def async_set_cache_pipeline( - self, cache_list: List[Tuple[Any, Any]], ttl: Optional[float] = None, **kwargs - ): + async def async_set_cache_pipeline(self, cache_list: List[Tuple[Any, Any]], ttl: Optional[float] = None, **kwargs): """ Use Redis Pipelines for bulk write operations """ @@ -760,9 +715,7 @@ class RedisCache(BaseCache): _redis_client = self.init_async_client() start_time = time.time() - print_verbose( - f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}" - ) + print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}") cache_value: Any = None try: async with _redis_client.pipeline(transaction=False) as pipe: @@ -824,9 +777,7 @@ class RedisCache(BaseCache): raise @_redis_circuit_breaker_guard - async def async_set_cache_sadd( - self, key, value: List, ttl: Optional[float], **kwargs - ): + async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float], **kwargs): from redis.asyncio import Redis start_time = time.time() @@ -857,12 +808,8 @@ class RedisCache(BaseCache): key = self.check_and_fix_namespace(key=key) print_verbose(f"Set ASYNC Redis Cache: key: {key}\nValue {value}\nttl={ttl}") try: - await self._set_cache_sadd_helper( - redis_client=_redis_client, key=key, value=value, ttl=ttl - ) - print_verbose( - f"Successfully Set ASYNC Redis Cache SADD: key: {key}\nValue {value}\nttl={ttl}" - ) + await self._set_cache_sadd_helper(redis_client=_redis_client, key=key, value=value, ttl=ttl) + print_verbose(f"Successfully Set ASYNC Redis Cache SADD: key: {key}\nValue {value}\nttl={ttl}") end_time = time.time() _duration = end_time - start_time asyncio.create_task( @@ -1006,9 +953,7 @@ class RedisCache(BaseCache): return float(result) async def flush_cache_buffer(self): - print_verbose( - f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}" - ) + print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) self.redis_batch_writing_buffer = [] @@ -1021,9 +966,7 @@ class RedisCache(BaseCache): # cached_response is in `b{} convert it to ModelResponse cached_response = cached_response.decode("utf-8") # Convert bytes to string try: - cached_response = json.loads( - cached_response - ) # Convert string to dictionary + cached_response = json.loads(cached_response) # Convert string to dictionary except Exception: cached_response = ast.literal_eval(cached_response) return cached_response @@ -1044,15 +987,11 @@ class RedisCache(BaseCache): end_time=end_time, parent_otel_span=parent_otel_span, ) - print_verbose( - f"Got Redis Cache: key: {key}, cached_response {cached_response}" - ) + print_verbose(f"Got Redis Cache: key: {key}, cached_response {cached_response}") return self._get_cache_logic(cached_response=cached_response) except Exception as e: # NON blocking - notify users Redis is throwing an exception - verbose_logger.error( - "litellm.caching.caching: get() - Got exception from REDIS: ", e - ) + verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e) def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]: """ @@ -1124,9 +1063,7 @@ class RedisCache(BaseCache): return key_value_dict @_redis_circuit_breaker_guard - async def async_get_cache( - self, key, parent_otel_span: Optional[Span] = None, **kwargs - ): + async def async_get_cache(self, key, parent_otel_span: Optional[Span] = None, **kwargs): from redis.asyncio import Redis _redis_client: Redis = self.init_async_client() # type: ignore @@ -1136,9 +1073,7 @@ class RedisCache(BaseCache): try: print_verbose(f"Get Async Redis Cache: key: {key}") cached_response = await _redis_client.get(key) - print_verbose( - f"Got Async Redis Cache: key: {key}, cached_response {cached_response}" - ) + print_verbose(f"Got Async Redis Cache: key: {key}, cached_response {cached_response}") response = self._get_cache_logic(cached_response=cached_response) end_time = time.time() @@ -1170,9 +1105,7 @@ class RedisCache(BaseCache): event_metadata={"key": key}, ) ) - print_verbose( - f"litellm.caching.caching: async get() - Got exception from REDIS: {str(e)}" - ) + print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {str(e)}") @_redis_circuit_breaker_guard async def async_batch_get_cache( @@ -1277,9 +1210,7 @@ class RedisCache(BaseCache): error=e, call_type=f"sync_ping <- {_get_call_stack_info()}", ) - verbose_logger.error( - f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}" - ) + verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}") raise e async def ping(self) -> bool: @@ -1313,9 +1244,7 @@ class RedisCache(BaseCache): call_type=f"async_ping <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}" - ) + verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}") raise e @_redis_circuit_breaker_guard @@ -1415,9 +1344,7 @@ class RedisCache(BaseCache): # Execute the pipeline and return results results = await pipe.execute() # only return float values - verbose_logger.debug( - f"Increment ASYNC Redis Cache PIPELINE: results: {results}" - ) + verbose_logger.debug(f"Increment ASYNC Redis Cache PIPELINE: results: {results}") return [r for r in results if isinstance(r, float)] @_redis_circuit_breaker_guard @@ -1441,9 +1368,7 @@ class RedisCache(BaseCache): _redis_client: Redis = self.init_async_client() # type: ignore start_time = time.time() - print_verbose( - f"Increment Async Redis Cache Pipeline: increment list: {increment_list}" - ) + print_verbose(f"Increment Async Redis Cache Pipeline: increment list: {increment_list}") try: async with _redis_client.pipeline(transaction=False) as pipe: @@ -1557,9 +1482,7 @@ class RedisCache(BaseCache): call_type=f"async_rpush <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {str(e)}" - ) + verbose_logger.error(f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {str(e)}") raise e async def _pipeline_rpush_helper( @@ -1633,9 +1556,7 @@ class RedisCache(BaseCache): ) raise e - async def handle_lpop_count_for_older_redis_versions( - self, pipe: pipeline, key: str, count: int - ) -> List[bytes]: + async def handle_lpop_count_for_older_redis_versions(self, pipe: pipeline, key: str, count: int) -> List[bytes]: result: List[bytes] = [] for _ in range(count): pipe.lpop(key) @@ -1666,9 +1587,7 @@ class RedisCache(BaseCache): if count is not None and major_version < 7: # For Redis < 7.0, use pipeline to execute multiple LPOP commands async with _redis_client.pipeline(transaction=False) as pipe: - result = await self.handle_lpop_count_for_older_redis_versions( - pipe, key, count - ) + result = await self.handle_lpop_count_for_older_redis_versions(pipe, key, count) else: # For Redis >= 7.0 or when count is None, use native LPOP with count result = await _redis_client.lpop(key, count) @@ -1690,9 +1609,7 @@ class RedisCache(BaseCache): return result.decode("utf-8") except Exception: return result - elif isinstance(result, list) and all( - isinstance(item, bytes) for item in result - ): + elif isinstance(result, list) and all(isinstance(item, bytes) for item in result): try: return [item.decode("utf-8") for item in result] except Exception: @@ -1711,9 +1628,7 @@ class RedisCache(BaseCache): call_type=f"async_lpop <- {_get_call_stack_info()}", ) ) - verbose_logger.error( - f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}" - ) + verbose_logger.error(f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}") raise e async def _pipeline_lpop_helper( @@ -1749,9 +1664,7 @@ class RedisCache(BaseCache): raw_results = [] offset = 0 for count in counts: - key_results = [ - r for r in flat_results[offset : offset + count] if r is not None - ] + key_results = [r for r in flat_results[offset : offset + count] if r is not None] raw_results.append(key_results if key_results else None) offset += count @@ -1768,11 +1681,7 @@ class RedisCache(BaseCache): elif isinstance(r, list): try: decoded_results.append( - [ - item.decode("utf-8") if isinstance(item, bytes) else item - for item in r - if item is not None - ] + [item.decode("utf-8") if isinstance(item, bytes) else item for item in r if item is not None] or None ) except Exception: diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 2dc9224e715..0698ebdcf2a 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -37,9 +37,7 @@ class RedisClusterCache(RedisCache): if self.redis_async_redis_cluster_client: return self.redis_async_redis_cluster_client - _redis_client = get_redis_async_client( - connection_pool=self.async_redis_conn_pool, **self.redis_kwargs - ) + _redis_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs) if isinstance(_redis_client, RedisCluster): self.redis_async_redis_cluster_client = _redis_client diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index e79392bb7f0..d4288cc777c 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -97,8 +97,7 @@ class RedisSemanticCache(BaseCache): # Raise a more informative exception if any of the required keys are missing missing_var = e.args[0] raise ValueError( - f"Missing required Redis configuration: {missing_var}. " - f"Provide {missing_var} or redis_url." + f"Missing required Redis configuration: {missing_var}. Provide {missing_var} or redis_url." ) from e redis_url = f"redis://:{password}@{host}:{port}" @@ -158,10 +157,7 @@ class RedisSemanticCache(BaseCache): ) -> Any: def _is_schema_mismatch(exc: ValueError) -> bool: error_message = str(exc).lower() - return any( - phrase in error_message - for phrase in ("schema does not match", "index schema") - ) + return any(phrase in error_message for phrase in ("schema does not match", "index schema")) try: return semantic_cache_cls( @@ -310,9 +306,7 @@ class RedisSemanticCache(BaseCache): return dict_method() return value - def _get_embedding( - self, prompt: str, metadata: Dict[str, Any] | None = None - ) -> List[float]: + def _get_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> List[float]: """ Routes through the proxy Router when the embedding model is a Router deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies, @@ -324,9 +318,7 @@ class RedisSemanticCache(BaseCache): llm_model_list = None llm_router = None - router = resolve_embedding_router( - self.embedding_model, llm_router, llm_model_list - ) + router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) if router is not None: embedding_response = cast( EmbeddingResponse, @@ -398,9 +390,7 @@ class RedisSemanticCache(BaseCache): value_str = str(value) - prompt_embedding = self._get_embedding( - prompt, metadata=kwargs.get("metadata") - ) + prompt_embedding = self._get_embedding(prompt, metadata=kwargs.get("metadata")) store_kwargs: dict[str, Any] = { "vector": prompt_embedding, @@ -413,9 +403,7 @@ class RedisSemanticCache(BaseCache): store_kwargs["ttl"] = int(ttl) self.llmcache.store(prompt, value_str, **store_kwargs) except Exception as e: - print_verbose( - f"Error setting {value_str or value} in the Redis semantic cache: {str(e)}" - ) + print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {str(e)}") def get_cache(self, key: str, **kwargs) -> Any: """ @@ -439,9 +427,7 @@ class RedisSemanticCache(BaseCache): # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - prompt_embedding = self._get_embedding( - prompt, metadata=kwargs.get("metadata") - ) + prompt_embedding = self._get_embedding(prompt, metadata=kwargs.get("metadata")) check_kwargs: dict[str, Any] = { "prompt": prompt, "vector": prompt_embedding, @@ -485,9 +471,7 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _get_async_embedding( - self, prompt: str, metadata: Dict[str, Any] | None = None - ) -> List[float]: + async def _get_async_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> List[float]: """ Asynchronously generate an embedding for the given prompt. @@ -504,9 +488,7 @@ class RedisSemanticCache(BaseCache): llm_model_list = None llm_router = None - router = resolve_embedding_router( - self.embedding_model, llm_router, llm_model_list - ) + router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) try: if router is not None: embedding_response = await router.aembedding( @@ -547,9 +529,7 @@ class RedisSemanticCache(BaseCache): value_str = str(value) # Generate embedding for the value (response) to cache - prompt_embedding = await self._get_async_embedding( - prompt, metadata=kwargs.get("metadata") - ) + prompt_embedding = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) store_kwargs: dict[str, Any] = { "vector": prompt_embedding, @@ -589,9 +569,7 @@ class RedisSemanticCache(BaseCache): return None # Generate embedding for the prompt - prompt_embedding = await self._get_async_embedding( - prompt, metadata=kwargs.get("metadata") - ) + prompt_embedding = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. @@ -647,9 +625,7 @@ class RedisSemanticCache(BaseCache): aindex = await self.llmcache._get_async_index() return await aindex.info() - async def async_set_cache_pipeline( - self, cache_list: List[Tuple[str, Any]], **kwargs - ) -> None: + async def async_set_cache_pipeline(self, cache_list: List[Tuple[str, Any]], **kwargs) -> None: """ Asynchronously store multiple values in the semantic cache. diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index e26fbe8981c..1ada940a9c9 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -110,9 +110,7 @@ class S3Cache(BaseCache): func = partial(self.set_cache, key, value, **kwargs) await loop.run_in_executor(None, func) except Exception as e: - verbose_logger.error( - f"S3 Caching: async_set_cache() - Got exception from S3: {e}" - ) + verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}") def get_cache(self, key, **kwargs): import botocore @@ -122,9 +120,7 @@ class S3Cache(BaseCache): print_verbose(f"Get S3 Cache: key: {key}") # Download the data from S3 - cached_response = self.s3_client.get_object( - Bucket=self.bucket_name, Key=key - ) + cached_response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key) if cached_response is not None: if "Expires" in cached_response: @@ -135,13 +131,9 @@ class S3Cache(BaseCache): return None # cached_response is in `b{} convert it to ModelResponse - cached_response = ( - cached_response["Body"].read().decode("utf-8") - ) # Convert bytes to string + cached_response = cached_response["Body"].read().decode("utf-8") # Convert bytes to string try: - cached_response = json.loads( - cached_response - ) # Convert string to dictionary + cached_response = json.loads(cached_response) # Convert string to dictionary except Exception: cached_response = ast.literal_eval(cached_response) if not isinstance(cached_response, dict): @@ -153,15 +145,11 @@ class S3Cache(BaseCache): return cached_response except botocore.exceptions.ClientError as e: # type: ignore if e.response["Error"]["Code"] == "NoSuchKey": - verbose_logger.debug( - f"S3 Cache: The specified key '{key}' does not exist in the S3 bucket." - ) + verbose_logger.debug(f"S3 Cache: The specified key '{key}' does not exist in the S3 bucket.") return None except Exception as e: - verbose_logger.error( - f"S3 Caching: get_cache() - Got exception from S3: {e}" - ) + verbose_logger.error(f"S3 Caching: get_cache() - Got exception from S3: {e}") async def async_get_cache(self, key, **kwargs): """ @@ -175,9 +163,7 @@ class S3Cache(BaseCache): result = await loop.run_in_executor(None, func) return result except Exception as e: - verbose_logger.error( - f"S3 Caching: async_get_cache() - Got exception from S3: {e}" - ) + verbose_logger.error(f"S3 Caching: async_get_cache() - Got exception from S3: {e}") return None def flush_cache(self): diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index bf368b74d07..746e91207d8 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -84,31 +84,21 @@ class ValkeySemanticCache(RedisSemanticCache): resolved_url = None if sync_client is None or async_client is None: - resolved_url = redis_url or self._build_valkey_url( - host, port, password, ssl - ) + resolved_url = redis_url or self._build_valkey_url(host, port, password, ssl) self.sync_client = ( sync_client if sync_client is not None else Redis.from_url(resolved_url) # type: ignore[arg-type] ) self.async_client = ( - async_client - if async_client is not None - else AsyncRedis.from_url(resolved_url) # type: ignore[arg-type] + async_client if async_client is not None else AsyncRedis.from_url(resolved_url) # type: ignore[arg-type] ) print_verbose(f"Valkey semantic-cache initializing index - {self.index_name}") @staticmethod - def _build_valkey_url( - host: str | None, port: str | None, password: str | None, ssl: bool = False - ) -> str: + def _build_valkey_url(host: str | None, port: str | None, password: str | None, ssl: bool = False) -> str: host = host or os.environ.get("VALKEY_HOST") or os.environ.get("REDIS_HOST") port = port or os.environ.get("VALKEY_PORT") or os.environ.get("REDIS_PORT") - password = ( - password - or os.environ.get("VALKEY_PASSWORD") - or os.environ.get("REDIS_PASSWORD") - ) + password = password or os.environ.get("VALKEY_PASSWORD") or os.environ.get("REDIS_PASSWORD") if not host or not port: raise ValueError( @@ -157,11 +147,7 @@ class ValkeySemanticCache(RedisSemanticCache): for field in info.get("attributes") or []: if not isinstance(field, (list, tuple)): continue - flat = [ - sub - for item in field - for sub in (item if isinstance(item, (list, tuple)) else [item]) - ] + flat = [sub for item in field for sub in (item if isinstance(item, (list, tuple)) else [item])] for i, marker in enumerate(flat): if marker in (b"dimensions", "dimensions") and i + 1 < len(flat): return int(flat[i + 1]) @@ -207,9 +193,7 @@ class ValkeySemanticCache(RedisSemanticCache): def _doc_key(self, key: str) -> str: return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}" - def _doc_mapping( - self, key: str, prompt: str, value_str: str, embedding: list[float] - ) -> dict: + def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict: return { self.CACHE_KEY_FIELD_NAME: self._scope_tag(key), self.PROMPT_FIELD_NAME: prompt, @@ -223,11 +207,7 @@ class ValkeySemanticCache(RedisSemanticCache): f"(@{self.CACHE_KEY_FIELD_NAME}:{{{scope}}})" f"=>[KNN 1 @{self.EMBEDDING_FIELD_NAME} $vec AS {self.DISTANCE_FIELD_NAME}]" ) - return ( - Query(query_string) - .return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME) - .dialect(2) - ) + return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2) @classmethod def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None: @@ -264,9 +244,7 @@ class ValkeySemanticCache(RedisSemanticCache): self._ensure_index_sync(len(embedding)) doc_key = self._doc_key(key) - self.sync_client.hset( - doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding) - ) + self.sync_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding)) ttl = self._get_ttl(**kwargs) if ttl is not None: self.sync_client.expire(doc_key, ttl) @@ -305,9 +283,7 @@ class ValkeySemanticCache(RedisSemanticCache): await self._ensure_index_async(len(embedding)) doc_key = self._doc_key(key) - await self.async_client.hset( - doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding) - ) + await self.async_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding)) ttl = self._get_ttl(**kwargs) if ttl is not None: await self.async_client.expire(doc_key, ttl) @@ -334,20 +310,11 @@ class ValkeySemanticCache(RedisSemanticCache): print_verbose(f"Error in async Valkey semantic-cache get_cache: {str(e)}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def async_set_cache_pipeline( - self, cache_list: list[tuple[str, Any]], **kwargs: Any - ) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None: try: - await asyncio.gather( - *[ - self.async_set_cache(key, value, **kwargs) - for key, value in cache_list - ] - ) + await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list]) except Exception as e: - print_verbose( - f"Error in Valkey semantic-cache async_set_cache_pipeline: {str(e)}" - ) + print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {str(e)}") async def _index_info(self) -> dict: return await self.async_client.ft(self.index_name).info() diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 87ac5d132d7..8f12d855880 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -41,10 +41,7 @@ class ResponsesToCompletionBridgeHandler: def _is_preformatted_cached_chat_stream(result: Any) -> bool: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - return ( - isinstance(result, CustomStreamWrapper) - and result.custom_llm_provider == "cached_response" - ) + return isinstance(result, CustomStreamWrapper) and result.custom_llm_provider == "cached_response" @staticmethod def _coerce_response_object( @@ -85,9 +82,7 @@ class ResponsesToCompletionBridgeHandler: raise ValueError("Stream completed response is invalid") return response - async def _collect_response_from_stream_async( - self, stream_iter: Any - ) -> "ResponsesAPIResponse": + async def _collect_response_from_stream_async(self, stream_iter: Any) -> "ResponsesAPIResponse": async for _ in stream_iter: pass @@ -102,9 +97,7 @@ class ResponsesToCompletionBridgeHandler: raise ValueError("Stream completed response is invalid") return response - def validate_input_kwargs( - self, kwargs: dict - ) -> ResponsesToCompletionBridgeHandlerInputKwargs: + def validate_input_kwargs(self, kwargs: dict) -> ResponsesToCompletionBridgeHandlerInputKwargs: from litellm import LiteLLMLoggingObj from litellm.types.utils import ModelResponse @@ -234,9 +227,7 @@ class ResponsesToCompletionBridgeHandler: ) else: if self._is_preformatted_cached_chat_stream(result): - return self._apply_post_stream_processing( - result, model, custom_llm_provider - ) + return self._apply_post_stream_processing(result, model, custom_llm_provider) completion_stream = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore sync_stream=True, @@ -248,13 +239,9 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) - return self._apply_post_stream_processing( - streamwrapper, model, custom_llm_provider - ) + return self._apply_post_stream_processing(streamwrapper, model, custom_llm_provider) - async def acompletion( - self, *args, **kwargs - ) -> Union["ModelResponse", "CustomStreamWrapper"]: + async def acompletion(self, *args, **kwargs) -> Union["ModelResponse", "CustomStreamWrapper"]: from litellm import aresponses from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -314,9 +301,7 @@ class ResponsesToCompletionBridgeHandler: elif isinstance(result, ModelResponse): return result elif not stream: - responses_api_response = await self._collect_response_from_stream_async( - result - ) + responses_api_response = await self._collect_response_from_stream_async(result) return self.transformation_handler.transform_response( model=model, raw_response=responses_api_response, @@ -332,9 +317,7 @@ class ResponsesToCompletionBridgeHandler: ) else: if self._is_preformatted_cached_chat_stream(result): - return self._apply_post_stream_processing( - result, model, custom_llm_provider - ) + return self._apply_post_stream_processing(result, model, custom_llm_provider) completion_stream = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore sync_stream=False, @@ -346,9 +329,7 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) - return self._apply_post_stream_processing( - streamwrapper, model, custom_llm_provider - ) + return self._apply_post_stream_processing(streamwrapper, model, custom_llm_provider) @staticmethod def _apply_post_stream_processing( diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index f7ac67927b6..aecb2552b53 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -83,9 +83,7 @@ def _build_reasoning_item( summary: List[Dict[str, Any]] = [] for s in summary_raw or []: if isinstance(s, dict): - summary.append( - {"type": s.get("type", "summary_text"), "text": s.get("text", "")} - ) + summary.append({"type": s.get("type", "summary_text"), "text": s.get("text", "")}) else: summary.append( { @@ -138,9 +136,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return {"type": "function", "name": fn_name} return tool_choice - def _handle_raw_dict_response_item( - self, item: Dict[str, Any], index: int - ) -> Tuple[Optional[Any], int]: + def _handle_raw_dict_response_item(self, item: Dict[str, Any], index: int) -> Tuple[Optional[Any], int]: """ Handle raw dict response items from Responses API (e.g., GPT-5 Codex format). @@ -183,13 +179,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if item_type == "function_call": # Extract provider_specific_fields if present and pass through as-is provider_specific_fields = item.get("provider_specific_fields") - if provider_specific_fields and not isinstance( - provider_specific_fields, dict - ): + if provider_specific_fields and not isinstance(provider_specific_fields, dict): provider_specific_fields = ( - dict(provider_specific_fields) - if hasattr(provider_specific_fields, "__dict__") - else {} + dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} ) tool_call_dict = { @@ -205,9 +197,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if provider_specific_fields: tool_call_dict["provider_specific_fields"] = provider_specific_fields # Also add to function's provider_specific_fields for consistency - tool_call_dict["function"]["provider_specific_fields"] = ( - provider_specific_fields - ) + tool_call_dict["function"]["provider_specific_fields"] = provider_specific_fields msg = Message( content=None, @@ -301,9 +291,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): { "type": "message", "role": role, - "content": self._convert_content_to_responses_format( - content, cast(str, role) - ), # type: ignore[arg-type] + "content": self._convert_content_to_responses_format(content, cast(str, role)), # type: ignore[arg-type] } ) @@ -321,10 +309,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if key in ("max_tokens", "max_completion_tokens"): responses_api_request["max_output_tokens"] = value elif key == "tools" and value is not None: - responses_api_request["tools"] = ( - self._convert_tools_to_responses_format( - cast(List[Dict[str, Any]], value) - ) + responses_api_request["tools"] = self._convert_tools_to_responses_format( + cast(List[Dict[str, Any]], value) ) elif key == "response_format": text_format = self._transform_response_format_to_text_format(value) @@ -345,13 +331,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _build_sanitized_litellm_params(self, litellm_params: dict) -> Dict[str, Any]: """Build sanitized litellm_params with merged metadata.""" - responses_optional_param_keys = set( - ResponsesAPIOptionalRequestParams.__annotations__.keys() - ) + responses_optional_param_keys = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) sanitized: Dict[str, Any] = { - key: value - for key, value in litellm_params.items() - if key not in responses_optional_param_keys + key: value for key, value in litellm_params.items() if key not in responses_optional_param_keys } legacy_metadata = litellm_params.get("metadata") existing_litellm_metadata = litellm_params.get("litellm_metadata") @@ -427,9 +409,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if instructions: responses_api_request["instructions"] = instructions - self._map_optional_params_to_responses_api_request( - optional_params, responses_api_request - ) + self._map_optional_params_to_responses_api_request(optional_params, responses_api_request) stream = optional_params.get("stream") or litellm_params.get("stream", False) verbose_logger.debug(f"Chat provider: Stream parameter: {stream}") @@ -442,9 +422,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): previous_response_id = optional_params.get("previous_response_id") if previous_response_id: # Use the existing session handler for responses API - verbose_logger.debug( - f"Chat provider: Warning ignoring previous response ID: {previous_response_id}" - ) + verbose_logger.debug(f"Chat provider: Warning ignoring previous response ID: {previous_response_id}") # Convert back to responses API format for the actual request @@ -464,13 +442,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "client": client, } - verbose_logger.debug( - f"Chat provider: Final request model={api_model}, input_items={len(input_items)}" - ) + verbose_logger.debug(f"Chat provider: Final request model={api_model}, input_items={len(input_items)}") - self._merge_responses_api_request_into_request_data( - request_data, responses_api_request, instructions - ) + self._merge_responses_api_request_into_request_data(request_data, responses_api_request, instructions) if headers: request_data["extra_headers"] = headers @@ -524,11 +498,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): encrypted_content=getattr(item, "encrypted_content", None), summary_raw=item.summary, ) - reasoning_content = " ".join( - s["text"] - for s in pending_reasoning_item["summary"] - if s.get("text") - ) + reasoning_content = " ".join(s["text"] for s in pending_reasoning_item["summary"] if s.get("text")) elif isinstance(item, ResponseOutputMessage): for content in item.content: @@ -545,11 +515,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): annotations=annotations, reasoning_items=cast( Optional[List[ChatCompletionReasoningItem]], - ( - [pending_reasoning_item] - if pending_reasoning_item is not None - else None - ), + ([pending_reasoning_item] if pending_reasoning_item is not None else None), ), ) @@ -570,23 +536,25 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): LiteLLMCompletionResponsesConfig, ) - tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=item, - index=tool_call_index, + tool_call_dict = ( + LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( + tool_call_item=item, + index=tool_call_index, + ) ) accumulated_tool_calls.append(tool_call_dict) tool_call_index += 1 - elif ResponseApplyPatchToolCall is not None and isinstance( - item, ResponseApplyPatchToolCall - ): + elif ResponseApplyPatchToolCall is not None and isinstance(item, ResponseApplyPatchToolCall): from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) - tool_call_dict = LiteLLMCompletionResponsesConfig.convert_apply_patch_tool_call_to_chat_completion_tool_call( - tool_call_item=item, - index=tool_call_index, + tool_call_dict = ( + LiteLLMCompletionResponsesConfig.convert_apply_patch_tool_call_to_chat_completion_tool_call( + tool_call_item=item, + index=tool_call_index, + ) ) accumulated_tool_calls.append(tool_call_dict) tool_call_index += 1 @@ -607,25 +575,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): reasoning_content=reasoning_content, reasoning_items=cast( Optional[List[ChatCompletionReasoningItem]], - ( - [pending_reasoning_item] - if pending_reasoning_item is not None - else None - ), + ([pending_reasoning_item] if pending_reasoning_item is not None else None), ), ) - choices.append( - Choices(message=msg, finish_reason="tool_calls", index=index) - ) + choices.append(Choices(message=msg, finish_reason="tool_calls", index=index)) reasoning_content = None pending_reasoning_item = None return choices @classmethod - def _extract_output_from_completed_event( - cls, parsed_chunk: Dict[str, Any] - ) -> Optional[List[Dict[str, Any]]]: + def _extract_output_from_completed_event(cls, parsed_chunk: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]: response_payload = parsed_chunk.get("response") if not isinstance(response_payload, dict): return None @@ -635,9 +595,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return cast(List[Dict[str, Any]], response_output) @classmethod - def _recover_output_items_from_raw_sse( - cls, raw_sse: Optional[str] - ) -> List[Dict[str, Any]]: + def _recover_output_items_from_raw_sse(cls, raw_sse: Optional[str]) -> List[Dict[str, Any]]: if not raw_sse or not isinstance(raw_sse, str): return [] @@ -652,9 +610,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): event_type = parsed_chunk.get("type") if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: - recovered_output = cls._extract_output_from_completed_event( - parsed_chunk - ) + recovered_output = cls._extract_output_from_completed_event(parsed_chunk) if recovered_output is not None: return recovered_output continue @@ -688,9 +644,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return [] @classmethod - def _recover_output_items_from_logging( - cls, logging_obj: "LiteLLMLoggingObj" - ) -> List[Dict[str, Any]]: + def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> List[Dict[str, Any]]: model_call_details = getattr(logging_obj, "model_call_details", {}) or {} original_response = model_call_details.get("original_response") return cls._recover_output_items_from_raw_sse(original_response) @@ -721,9 +675,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): output_items = raw_response.output if len(output_items) == 0: - recovered_output_items = self._recover_output_items_from_logging( - logging_obj - ) + recovered_output_items = self._recover_output_items_from_logging(logging_obj) if recovered_output_items: output_items = cast(Any, recovered_output_items) raw_response.output = cast(Any, recovered_output_items) @@ -739,17 +691,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) if len(choices) == 0: - if ( - raw_response.incomplete_details is not None - and raw_response.incomplete_details.reason is not None - ): - raise ValueError( - f"{model} unable to complete request: {raw_response.incomplete_details.reason}" - ) + if raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None: + raise ValueError(f"{model} unable to complete request: {raw_response.incomplete_details.reason}") else: - raise ValueError( - f"Unknown items in responses API response: {output_items}" - ) + raise ValueError(f"Unknown items in responses API response: {output_items}") setattr(model_response, "choices", choices) @@ -758,28 +703,21 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): setattr( model_response, "usage", - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - raw_response.usage - ), + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage), ) # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id raw_response_hidden_params = getattr(raw_response, "_hidden_params", {}) if raw_response_hidden_params: - if ( - not hasattr(model_response, "_hidden_params") - or model_response._hidden_params is None - ): + if not hasattr(model_response, "_hidden_params") or model_response._hidden_params is None: model_response._hidden_params = {} # Merge the raw_response hidden params with model_response hidden params # Preserve existing keys in model_response but add/override with raw_response params for key, value in raw_response_hidden_params.items(): if key == "additional_headers" and key in model_response._hidden_params: # Merge additional_headers to preserve both sets - existing_additional_headers = model_response._hidden_params.get( - "additional_headers", {} - ) + existing_additional_headers = model_response._hidden_params.get("additional_headers", {}) merged_headers = {**value, **existing_additional_headers} model_response._hidden_params[key] = merged_headers else: @@ -789,19 +727,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def get_model_response_iterator( self, - streaming_response: Union[ - Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel" - ], + streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"], sync_stream: bool, json_mode: Optional[bool] = False, ) -> BaseModelResponseIterator: - return OpenAiResponsesToChatCompletionStreamIterator( - streaming_response, sync_stream, json_mode - ) + return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode) - def _convert_content_str_to_input_text( - self, content: str, role: str - ) -> Dict[str, Any]: + def _convert_content_str_to_input_text(self, content: str, role: str) -> Dict[str, Any]: if role == "user" or role == "system" or role == "tool": return {"type": "input_text", "text": content} else: @@ -828,9 +760,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if actual_image_url is None: raise ValueError(f"Invalid image URL: {content_image_url}") - image_param = ResponseInputImageParam( - image_url=actual_image_url, detail="auto", type="input_image" - ) + image_param = ResponseInputImageParam(image_url=actual_image_url, detail="auto", type="input_image") if detail: image_param["detail"] = detail @@ -857,9 +787,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): """Convert chat completion content to responses API format""" from litellm.types.llms.openai import ChatCompletionImageObject - verbose_logger.debug( - f"Chat provider: Converting content to responses format - input type: {type(content)}" - ) + verbose_logger.debug(f"Chat provider: Converting content to responses format - input type: {type(content)}") if content is None: return [self._convert_content_str_to_input_text("", role)] @@ -870,9 +798,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif isinstance(content, list): result = [] for i, item in enumerate(content): - verbose_logger.debug( - f"Chat provider: Processing content item {i}: {type(item)} = {item}" - ) + verbose_logger.debug(f"Chat provider: Processing content item {i}: {type(item)} = {item}") if isinstance(item, str): converted = self._convert_content_str_to_input_text(item, role) result.append(converted) @@ -881,9 +807,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Handle multimodal content original_type = item.get("type") if original_type == "text": - converted = self._convert_content_str_to_input_text( - item.get("text", ""), role - ) + converted = self._convert_content_str_to_input_text(item.get("text", ""), role) result.append(converted) verbose_logger.debug(f"Chat provider: text -> {converted}") elif original_type == "image_url": @@ -895,18 +819,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ), ) result.append(converted) - verbose_logger.debug( - f"Chat provider: image_url -> {converted}" - ) + verbose_logger.debug(f"Chat provider: image_url -> {converted}") else: # Try to map other types to responses API format item_type = original_type or "input_text" if item_type == "image": converted = {"type": "input_image", **item} result.append(converted) - verbose_logger.debug( - f"Chat provider: image -> {converted}" - ) + verbose_logger.debug(f"Chat provider: image -> {converted}") elif item_type == "file": # Map Chat Completion file to Responses API input_file # {"type": "file", "file": {"file_data": "...", "filename": "..."}} @@ -918,9 +838,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if key in file_data: converted[key] = file_data[key] result.append(converted) - verbose_logger.debug( - f"Chat provider: file -> {converted}" - ) + verbose_logger.debug(f"Chat provider: file -> {converted}") elif item_type in [ "input_text", "input_image", @@ -932,18 +850,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ]: # Already in responses API format result.append(item) - verbose_logger.debug( - f"Chat provider: passthrough -> {item}" - ) + verbose_logger.debug(f"Chat provider: passthrough -> {item}") else: # Default to input_text for unknown types - converted = self._convert_content_str_to_input_text( - str(item.get("text", item)), role - ) + converted = self._convert_content_str_to_input_text(str(item.get("text", item)), role) result.append(converted) - verbose_logger.debug( - f"Chat provider: unknown({original_type}) -> {converted}" - ) + verbose_logger.debug(f"Chat provider: unknown({original_type}) -> {converted}") verbose_logger.debug(f"Chat provider: Final converted content: {result}") return result else: @@ -951,17 +863,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): verbose_logger.debug(f"Chat provider: Other content type -> {result}") return result - def _convert_tools_to_responses_format( - self, tools: List[Dict[str, Any]] - ) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]: + def _convert_tools_to_responses_format(self, tools: List[Dict[str, Any]]) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]: """Convert chat completion tools to responses API tools format""" responses_tools: List["ALL_RESPONSES_API_TOOL_PARAMS"] = [] for tool in tools: # convert function tool from chat completion to responses API format if tool.get("type") == "function": - function_tool = cast( - ChatCompletionToolParamFunctionChunk, tool.get("function") - ) + function_tool = cast(ChatCompletionToolParamFunctionChunk, tool.get("function")) responses_tools.append( FunctionToolParam( name=function_tool["name"], @@ -987,9 +895,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if not extra_body: return optional_params - supported_responses_api_params = set( - ResponsesAPIOptionalRequestParams.__annotations__.keys() - ) + supported_responses_api_params = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) # Also include params we handle specially supported_responses_api_params.update( { @@ -1007,9 +913,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return optional_params - def _map_reasoning_effort( - self, reasoning_effort: Union[str, Dict[str, Any]] - ) -> Optional[Reasoning]: + def _map_reasoning_effort(self, reasoning_effort: Union[str, Dict[str, Any]]) -> Optional[Reasoning]: # If dict is passed, convert it directly to Reasoning object if isinstance(reasoning_effort, dict): return Reasoning(**reasoning_effort) # type: ignore[typeddict-item] @@ -1017,46 +921,25 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Check if auto-summary is enabled via flag or environment variable # Priority: litellm.reasoning_auto_summary flag > LITELLM_REASONING_AUTO_SUMMARY env var auto_summary_enabled = ( - litellm.reasoning_auto_summary - or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" + litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" ) # If string is passed, map with optional summary based on flag/env var if reasoning_effort == "none": - return ( - Reasoning(effort="none", summary="detailed") - if auto_summary_enabled - else Reasoning(effort="none") - ) # type: ignore + return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") # type: ignore elif reasoning_effort == "high": - return ( - Reasoning(effort="high", summary="detailed") - if auto_summary_enabled - else Reasoning(effort="high") - ) + return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high") elif reasoning_effort == "xhigh": - return ( - Reasoning(effort="xhigh", summary="detailed") - if auto_summary_enabled - else Reasoning(effort="xhigh") - ) # type: ignore[typeddict-item] + return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") # type: ignore[typeddict-item] elif reasoning_effort == "medium": return ( - Reasoning(effort="medium", summary="detailed") - if auto_summary_enabled - else Reasoning(effort="medium") + Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium") ) elif reasoning_effort == "low": - return ( - Reasoning(effort="low", summary="detailed") - if auto_summary_enabled - else Reasoning(effort="low") - ) + return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low") elif reasoning_effort == "minimal": return ( - Reasoning(effort="minimal", summary="detailed") - if auto_summary_enabled - else Reasoning(effort="minimal") + Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal") ) return None @@ -1072,10 +955,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): responses_api_request: The responses API request dict to modify web_search_options: Web search configuration (dict or other value) """ - if ( - "tools" not in responses_api_request - or responses_api_request["tools"] is None - ): + if "tools" not in responses_api_request or responses_api_request["tools"] is None: responses_api_request["tools"] = [] # Get the tools list with proper type narrowing @@ -1165,17 +1045,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): annotation_dict = annotation else: # Skip unsupported annotation types - verbose_logger.debug( - f"Skipping unsupported annotation type: {type(annotation)}" - ) + verbose_logger.debug(f"Skipping unsupported annotation type: {type(annotation)}") continue result.append(annotation_dict) # type: ignore except Exception as e: # Skip malformed annotations - verbose_logger.debug( - f"Skipping malformed annotation: {annotation}, error: {e}" - ) + verbose_logger.debug(f"Skipping malformed annotation: {annotation}, error: {e}") continue return result if result else None @@ -1196,9 +1072,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): - def __init__( - self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False - ): + def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): super().__init__(streaming_response, sync_stream, json_mode) def _handle_string_chunk( @@ -1211,9 +1085,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if not str_line or str_line.startswith("event:"): # ignore. - return GenericStreamingChunk( - text="", tool_use=None, is_finished=False, finish_reason="", usage=None - ) + return GenericStreamingChunk(text="", tool_use=None, is_finished=False, finish_reason="", usage=None) index = str_line.find("data:") if index != -1: str_line = str_line[index + 5 :] @@ -1258,9 +1130,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): event_type = event_type.value if parsed_chunk.get("object") == "chat.completion.chunk" or ( - event_type is None - and isinstance(parsed_chunk.get("choices"), list) - and parsed_chunk.get("choices") + event_type is None and isinstance(parsed_chunk.get("choices"), list) and parsed_chunk.get("choices") ): return ModelResponseStream(**parsed_chunk) @@ -1284,13 +1154,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if output_item.get("type") == "function_call": # Extract provider_specific_fields if present provider_specific_fields = output_item.get("provider_specific_fields") - if provider_specific_fields and not isinstance( - provider_specific_fields, dict - ): + if provider_specific_fields and not isinstance(provider_specific_fields, dict): provider_specific_fields = ( - dict(provider_specific_fields) - if hasattr(provider_specific_fields, "__dict__") - else {} + dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} ) function_chunk = ChatCompletionToolCallFunctionChunk( @@ -1299,9 +1165,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) if provider_specific_fields: - function_chunk["provider_specific_fields"] = ( - provider_specific_fields - ) + function_chunk["provider_specific_fields"] = provider_specific_fields from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -1344,9 +1208,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): id=None, index=tool_call_index, type="function", - function=ChatCompletionToolCallFunctionChunk( - name=None, arguments=content_part - ), + function=ChatCompletionToolCallFunctionChunk(name=None, arguments=content_part), ) ] ), @@ -1355,22 +1217,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ] ) else: - raise ValueError( - f"Chat provider: Invalid function argument delta {parsed_chunk}" - ) + raise ValueError(f"Chat provider: Invalid function argument delta {parsed_chunk}") elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") == "function_call": # Extract provider_specific_fields if present provider_specific_fields = output_item.get("provider_specific_fields") - if provider_specific_fields and not isinstance( - provider_specific_fields, dict - ): + if provider_specific_fields and not isinstance(provider_specific_fields, dict): provider_specific_fields = ( - dict(provider_specific_fields) - if hasattr(provider_specific_fields, "__dict__") - else {} + dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} ) function_chunk = ChatCompletionToolCallFunctionChunk( @@ -1380,9 +1236,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): # Add provider_specific_fields to function if present if provider_specific_fields: - function_chunk["provider_specific_fields"] = ( - provider_specific_fields - ) + function_chunk["provider_specific_fields"] = provider_specific_fields tool_call_index = parsed_chunk.get("output_index", 0) tool_call_chunk = ChatCompletionToolCallChunk( @@ -1458,9 +1312,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): output_items = response_data.get("output", []) if response_data else [] has_function_calls = any( - item.get("type") == "function_call" - for item in output_items - if isinstance(item, dict) + item.get("type") == "function_call" for item in output_items if isinstance(item, dict) ) finish_reason = "tool_calls" if has_function_calls else "stop" @@ -1488,11 +1340,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if response_data.get("usage"): from litellm.responses.utils import ResponseAPILoggingUtils - usage = ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - response_data.get("usage") - ) - ) + usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage")) return ModelResponseStream( choices=[ StreamingChoices( @@ -1509,9 +1357,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): else: pass # For any unhandled event types, create a minimal valid chunk or skip - verbose_logger.debug( - f"Chat provider: Unhandled event type '{event_type}', creating empty chunk" - ) + verbose_logger.debug(f"Chat provider: Unhandled event type '{event_type}', creating empty chunk") # Return a minimal valid chunk for unknown events return ModelResponseStream( @@ -1534,9 +1380,5 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): Returns: ModelResponseStream: OpenAI-formatted streaming chunk """ - verbose_logger.debug( - f"Chat provider: transform_streaming_response called with chunk: {chunk}" - ) - return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( - chunk - ) + verbose_logger.debug(f"Chat provider: transform_streaming_response called with chunk: {chunk}") + return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(chunk) diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 45795c9ca15..004dd82cbaa 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -107,8 +107,7 @@ def _normalize_messages_for_compression( """ if call_type not in _SUPPORTED_CALL_TYPES: raise ValueError( - f"Unsupported call_type={call_type!r} for compression. " - f"Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}." + f"Unsupported call_type={call_type!r} for compression. Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}." ) original_messages: List[Dict[str, Any]] = [dict(m) for m in messages] @@ -334,9 +333,7 @@ def _select_kept_indices_for_budget( return kept_indices, truncated_overrides -def _get_dropped_tool_span_indices( - kept_indices: Set[int], tool_exchange_spans: List[Set[int]] -) -> Set[int]: +def _get_dropped_tool_span_indices(kept_indices: Set[int], tool_exchange_spans: List[Set[int]]) -> Set[int]: dropped_tool_span_indices: Set[int] = set() for span in tool_exchange_spans: if not any(idx in kept_indices for idx in span): @@ -440,9 +437,7 @@ def compress( tool_exchange_spans: List[Set[int]] = [] if _is_anthropic_call_type(call_type_str): - tool_exchange_spans, tool_sequence_error = ( - _extract_anthropic_tool_exchange_spans(original_messages) - ) + tool_exchange_spans, tool_sequence_error = _extract_anthropic_tool_exchange_spans(original_messages) if tool_sequence_error is not None: return CompressedResult( messages=original_messages, @@ -484,9 +479,7 @@ def compress( # Use the truncated version if we made one, otherwise the original compressed_messages.append(truncated_overrides.get(i, msg)) else: - key = extract_key( - normalized_messages[i], fallback_index=i, used_keys=used_keys - ) + key = extract_key(normalized_messages[i], fallback_index=i, used_keys=used_keys) content = _content_to_text(msg.get("content", "")) cache[key] = content compressed_messages.append(stub_message(msg, key)) @@ -503,11 +496,7 @@ def compress( messages=compressed_messages, original_tokens=original_tokens, compressed_tokens=compressed_tokens, - compression_ratio=( - round(1 - (compressed_tokens / original_tokens), 4) - if original_tokens > 0 - else 0.0 - ), + compression_ratio=(round(1 - (compressed_tokens / original_tokens), 4) if original_tokens > 0 else 0.0), cache=cache, tools=tools, ) diff --git a/litellm/compression/content_detection.py b/litellm/compression/content_detection.py index 975117eb608..4a072b63f2c 100644 --- a/litellm/compression/content_detection.py +++ b/litellm/compression/content_detection.py @@ -33,9 +33,7 @@ def detect_content_type(content: str) -> str: sample = stripped[:5000] keyword_matches = len(_CODE_KEYWORDS.findall(sample)) lines = sample.split("\n") - indented_lines = sum( - 1 for line in lines if line.startswith((" ", "\t")) and line.strip() - ) + indented_lines = sum(1 for line in lines if line.startswith((" ", "\t")) and line.strip()) # If we see multiple code keywords or significant indentation, it's likely code if keyword_matches >= 3 or (indented_lines > len(lines) * 0.3 and len(lines) > 5): diff --git a/litellm/compression/message_stubbing.py b/litellm/compression/message_stubbing.py index 2330f1bbc9e..8d4e65752c1 100644 --- a/litellm/compression/message_stubbing.py +++ b/litellm/compression/message_stubbing.py @@ -26,9 +26,7 @@ def extract_key(message: dict, fallback_index: int, used_keys: Set[str]) -> str: """ content = message.get("content", "") if isinstance(content, list): - content = " ".join( - p.get("text", "") if isinstance(p, dict) else str(p) for p in content - ) + content = " ".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in content) key = None for pattern in _FILE_PATH_PATTERNS: @@ -62,9 +60,7 @@ def stub_message(message: dict, key: str) -> dict: """ content = message.get("content", "") if isinstance(content, list): - content = " ".join( - p.get("text", "") if isinstance(p, dict) else str(p) for p in content - ) + content = " ".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in content) line_count = content.count("\n") + 1 content_type = detect_content_type(content) @@ -91,9 +87,7 @@ def truncate_message(message: dict, max_tokens: int) -> dict: """ content = message.get("content", "") if isinstance(content, list): - content = " ".join( - p.get("text", "") if isinstance(p, dict) else str(p) for p in content - ) + content = " ".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in content) # Rough conversion: 1 token ≈ 3 characters target_chars = max(100, max_tokens * 3) @@ -113,8 +107,6 @@ def truncate_message(message: dict, max_tokens: int) -> dict: first_count = (target_lines * 7) // 10 last_count = target_lines - first_count truncated = ( - "\n".join(lines[:first_count]) - + "\n...[truncated for context window]...\n" - + "\n".join(lines[-last_count:]) + "\n".join(lines[:first_count]) + "\n...[truncated for context window]...\n" + "\n".join(lines[-last_count:]) ) return {**message, "content": truncated} diff --git a/litellm/compression/retrieval_tool.py b/litellm/compression/retrieval_tool.py index 1ee24784a63..99431a2a15d 100644 --- a/litellm/compression/retrieval_tool.py +++ b/litellm/compression/retrieval_tool.py @@ -17,8 +17,7 @@ def build_retrieval_tool(available_keys: List[str]) -> dict: "description": ( "Retrieve the full content of a file or message that was " "compressed to save tokens. Use this when you need the complete " - "content to answer accurately. Available keys: " - + ", ".join(available_keys) + "content to answer accurately. Available keys: " + ", ".join(available_keys) ), "parameters": { "type": "object", diff --git a/litellm/compression/scoring/bm25.py b/litellm/compression/scoring/bm25.py index e8e1bf631eb..7f919ef16fb 100644 --- a/litellm/compression/scoring/bm25.py +++ b/litellm/compression/scoring/bm25.py @@ -91,11 +91,7 @@ def bm25_score_messages( return exact if len(query_term) < 4: return 0 - return sum( - count - for token, count in tf_counts.items() - if token != query_term and token.startswith(query_term) - ) + return sum(count for token, count in tf_counts.items() if token != query_term and token.startswith(query_term)) # Score each document scores: List[float] = [] diff --git a/litellm/constants.py b/litellm/constants.py index 9fd3237ce09..aeb74a65839 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -4,28 +4,16 @@ from typing import List, Literal, Optional from litellm.litellm_core_utils.env_utils import get_env_int -DEFAULT_HEALTH_CHECK_PROMPT = str( - os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm") -) -AZURE_DEFAULT_RESPONSES_API_VERSION = str( - os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview") -) +DEFAULT_HEALTH_CHECK_PROMPT = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) +AZURE_DEFAULT_RESPONSES_API_VERSION = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) ROUTER_MAX_FALLBACKS = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) DEFAULT_BATCH_SIZE = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) -DEFAULT_S3_FLUSH_INTERVAL_SECONDS = int( - os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10) -) +DEFAULT_S3_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_S3_BATCH_SIZE = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) -DEFAULT_SQS_FLUSH_INTERVAL_SECONDS = int( - os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10) -) -DEFAULT_NUM_WORKERS_LITELLM_PROXY = int( - os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1) -) -DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int( - os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1) -) +DEFAULT_SQS_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) +DEFAULT_NUM_WORKERS_LITELLM_PROXY = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) +DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) DEFAULT_SQS_BATCH_SIZE = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) SQS_SEND_MESSAGE_ACTION = "SendMessage" SQS_API_VERSION = "2012-11-05" @@ -34,9 +22,7 @@ DEFAULT_MAX_RETRIES = int(os.getenv("DEFAULT_MAX_RETRIES", 2)) # radius: each record fans out to spend logs + every callback integration. MAX_CALLBACK_LOG_RECORDS = 1000 DEFAULT_MAX_RECURSE_DEPTH = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100)) -DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int( - os.getenv("DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER", 10) -) +DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER", 10)) DEFAULT_FAILURE_THRESHOLD_PERCENT = float( os.getenv("DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5) ) # default cooldown a deployment if 50% of requests fail in a given minute @@ -44,12 +30,8 @@ DEFAULT_MAX_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) DEFAULT_ALLOWED_FAILS = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) DEFAULT_REDIS_SYNC_INTERVAL = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) DEFAULT_COOLDOWN_TIME_SECONDS = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) -DEFAULT_REPLICATE_POLLING_RETRIES = int( - os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5) -) -DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int( - os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1) -) +DEFAULT_REPLICATE_POLLING_RETRIES = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) +DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) # Maximum wall-clock seconds a streaming response is allowed to run. @@ -67,9 +49,7 @@ MAX_BASE64_LENGTH_FOR_LOGGING = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 6 # When true, adds detailed per-phase timing breakdown headers to responses. # Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms -LITELLM_DETAILED_TIMING = ( - os.getenv("LITELLM_DETAILED_TIMING", "false").lower() == "true" -) +LITELLM_DETAILED_TIMING = os.getenv("LITELLM_DETAILED_TIMING", "false").lower() == "true" # Model cost map validation constants MODEL_COST_MAP_MIN_MODEL_COUNT = int( @@ -91,9 +71,7 @@ MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int( # Surrogate-repair fallback in _read_request_body runs two full-body re.sub passes # that block the event loop on multi-MB malformed bodies. Skip the repair above this # size and raise the existing 400 immediately. Set to 0 to disable the cap. -MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB = get_env_int( - "MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 1 -) +MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB = get_env_int("MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 1) SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD = int( os.getenv("SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD", 1000) ) # Minimum number of requests to consider "reasonable traffic". Used for single-deployment cooldown logic. @@ -101,42 +79,28 @@ DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS = int( os.getenv("DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS", 5) ) # Minimum number of requests before applying error rate cooldown. Prevents cooldown from triggering on first failure. -DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET = int( - os.getenv("DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET", 0) -) +DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET", 0)) # MCP Semantic Tool Filter Defaults DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL = str( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL", "text-embedding-3-small") ) -DEFAULT_MCP_SEMANTIC_FILTER_TOP_K = int( - os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_TOP_K", 10) -) +DEFAULT_MCP_SEMANTIC_FILTER_TOP_K = int(os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_TOP_K", 10)) DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD = float( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3) ) -MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int( - os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150) -) +MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)) # Semantic Guard Defaults DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL = str( os.getenv("DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL", "text-embedding-3-small") ) -DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float( - os.getenv("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75) -) +DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75)) # MCP OAuth2 Client Credentials Defaults -MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS = int( - os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60") -) -MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE = int( - os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200") -) -MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int( - os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600") -) +MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) +MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200")) +MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600")) # Default npm cache directory for STDIO MCP servers. # npm/npx needs a writable cache dir; in containers the default (~/.npm) @@ -149,9 +113,7 @@ MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX = "mcp:per_user_token" MCP_PER_USER_TOKEN_DEFAULT_TTL = int( os.getenv("MCP_PER_USER_TOKEN_DEFAULT_TTL", "43200") # 12 hours ) -MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int( - os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60") -) +MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int(os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0")) @@ -166,14 +128,11 @@ MCP_HEALTH_CHECK_TIMEOUT = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", " # Extend via LITELLM_MCP_STDIO_EXTRA_COMMANDS env var (comma-separated). _MCP_STDIO_EXTRA_COMMANDS = os.getenv("LITELLM_MCP_STDIO_EXTRA_COMMANDS", "") MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset( - {"npx", "uvx", "python", "python3", "node", "docker", "deno"} - | (set(_MCP_STDIO_EXTRA_COMMANDS.split(",")) - {""}) + {"npx", "uvx", "python", "python3", "node", "docker", "deno"} | (set(_MCP_STDIO_EXTRA_COMMANDS.split(",")) - {""}) ) # MCP OAuth2 Token Exchange (OBO) Defaults -MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE = int( - os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500") -) +MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE = int(os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500")) LITELLM_UI_ALLOW_HEADERS = [ "x-litellm-semantic-filter", @@ -189,9 +148,7 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO", 128) ) DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int( - os.getenv( - "DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE", 512 - ) + os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE", 512) ) # Maximum number of callbacks that can be registered @@ -223,31 +180,19 @@ OPEN_SANDBOX_DEFAULT_TIMEOUT = 300 OPEN_SANDBOX_READY_TIMEOUT = 30.0 OPEN_SANDBOX_POLL_INTERVAL = 0.2 -DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int( - os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024) -) +DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024)) DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET = int( os.getenv("DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET", 2048) ) -DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET = int( - os.getenv("DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET", 4096) -) -DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET = int( - os.getenv("DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET", 8192) -) -DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET = int( - os.getenv("DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET", 16384) -) +DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET", 4096)) +DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET", 8192)) +DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET", 16384)) MAX_TOKEN_TRIMMING_ATTEMPTS = int( os.getenv("MAX_TOKEN_TRIMMING_ATTEMPTS", 10) ) # Maximum number of attempts to trim the message -RUNWAYML_DEFAULT_API_VERSION = str( - os.getenv("RUNWAYML_DEFAULT_API_VERSION", "2024-11-06") -) -RUNWAYML_POLLING_TIMEOUT = int( - os.getenv("RUNWAYML_POLLING_TIMEOUT", 600) -) # 10 minutes default for image generation +RUNWAYML_DEFAULT_API_VERSION = str(os.getenv("RUNWAYML_DEFAULT_API_VERSION", "2024-11-06")) +RUNWAYML_POLLING_TIMEOUT = int(os.getenv("RUNWAYML_POLLING_TIMEOUT", 600)) # 10 minutes default for image generation ########## Networking constants ############################################################## _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour @@ -255,9 +200,7 @@ _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client fo # Aiohttp connection pooling - prevents memory leaks from unbounded connection growth # Set to 0 for unlimited (not recommended for production) AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000)) -AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int( - os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500) -) +AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500)) AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) # TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs @@ -283,9 +226,7 @@ AIOHTTP_NEEDS_CLEANUP_CLOSED = (3, 13, 0) <= sys.version_info < ( # Default to None (unlimited) to match OpenAI's official agents SDK behavior # https://github.com/openai/openai-agents-python/blob/cf1b933660e44fd37b4350c41febab8221801409/src/agents/realtime/openai_realtime.py#L235 _max_size_env = os.getenv("REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES") -REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES = ( - int(_max_size_env) if _max_size_env is not None else None -) +REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES = int(_max_size_env) if _max_size_env is not None else None # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones @@ -316,9 +257,7 @@ REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer" -REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = ( - "litellm_daily_end_user_spend_update_buffer" -) +REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer" REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) @@ -327,12 +266,8 @@ LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1 TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. -MAX_SIZE_IN_MEMORY_QUEUE = int( - os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)) -) -MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int( - os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000) -) +MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) +MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int(os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000)) ############################################################################################### MINIMUM_PROMPT_CACHE_TOKEN_COUNT = int( os.getenv("MINIMUM_PROMPT_CACHE_TOKEN_COUNT", 1024) @@ -344,49 +279,31 @@ HOURS_IN_A_DAY = int(os.getenv("HOURS_IN_A_DAY", 24)) DAYS_IN_A_WEEK = int(os.getenv("DAYS_IN_A_WEEK", 7)) DAYS_IN_A_MONTH = int(os.getenv("DAYS_IN_A_MONTH", 28)) DAYS_IN_A_YEAR = int(os.getenv("DAYS_IN_A_YEAR", 365)) -REPLICATE_MODEL_NAME_WITH_ID_LENGTH = int( - os.getenv("REPLICATE_MODEL_NAME_WITH_ID_LENGTH", 64) -) +REPLICATE_MODEL_NAME_WITH_ID_LENGTH = int(os.getenv("REPLICATE_MODEL_NAME_WITH_ID_LENGTH", 64)) #### TOKEN COUNTING #### FUNCTION_DEFINITION_TOKEN_COUNT = int(os.getenv("FUNCTION_DEFINITION_TOKEN_COUNT", 9)) SYSTEM_MESSAGE_TOKEN_COUNT = int(os.getenv("SYSTEM_MESSAGE_TOKEN_COUNT", 4)) TOOL_CHOICE_OBJECT_TOKEN_COUNT = int(os.getenv("TOOL_CHOICE_OBJECT_TOKEN_COUNT", 4)) -DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT = int( - os.getenv("DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT", 10) -) -DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT = int( - os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20) -) -MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES = int( - os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768) -) -MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES = int( - os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000) -) +DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT = int(os.getenv("DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT", 10)) +DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT = int(os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20)) +MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES = int(os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768)) +MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES = int(os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000)) MAX_TILE_WIDTH = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT = int(os.getenv("MAX_TILE_HEIGHT", 512)) -OPENAI_FILE_SEARCH_COST_PER_1K_CALLS = float( - os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000) -) +OPENAI_FILE_SEARCH_COST_PER_1K_CALLS = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) # Azure OpenAI Assistants feature costs # Source: https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/ AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY = float( os.getenv("AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY", 0.1) # $0.1 USD per 1 GB/Day ) AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS = float( - os.getenv( - "AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS", 3.0 - ) # $0.003 USD per 1K Tokens + os.getenv("AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS", 3.0) # $0.003 USD per 1K Tokens ) AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS = float( - os.getenv( - "AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS", 12.0 - ) # $0.012 USD per 1K Tokens + os.getenv("AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS", 12.0) # $0.012 USD per 1K Tokens ) AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY = float( - os.getenv( - "AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY", 0.1 - ) # $0.1 USD per 1 GB/Day (same as file search) + os.getenv("AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY", 0.1) # $0.1 USD per 1 GB/Day (same as file search) ) MIN_NON_ZERO_TEMPERATURE = float(os.getenv("MIN_NON_ZERO_TEMPERATURE", 0.0001)) #### RELIABILITY #### @@ -400,9 +317,7 @@ _REALTIME_BODY_CACHE_SIZE = 1000 # Keep realtime helper caches bounded; workloa INITIAL_RETRY_DELAY = float(os.getenv("INITIAL_RETRY_DELAY", 0.5)) MAX_RETRY_DELAY = float(os.getenv("MAX_RETRY_DELAY", 8.0)) JITTER = float(os.getenv("JITTER", 0.75)) -DEFAULT_IN_MEMORY_TTL = int( - os.getenv("DEFAULT_IN_MEMORY_TTL", 5) -) # default time to live for the in-memory cache +DEFAULT_IN_MEMORY_TTL = int(os.getenv("DEFAULT_IN_MEMORY_TTL", 5)) # default time to live for the in-memory cache DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE = int( os.getenv("DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 1000) ) # default max size for redis batch cache @@ -410,23 +325,13 @@ DEFAULT_POLLING_INTERVAL = float( os.getenv("DEFAULT_POLLING_INTERVAL", 0.03) ) # default polling interval for the scheduler AZURE_OPERATION_POLLING_TIMEOUT = int(os.getenv("AZURE_OPERATION_POLLING_TIMEOUT", 120)) -AZURE_DOCUMENT_INTELLIGENCE_API_VERSION = str( - os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30") -) -AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI = int( - os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96) -) +AZURE_DOCUMENT_INTELLIGENCE_API_VERSION = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30")) +AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96)) REDIS_SOCKET_TIMEOUT = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1)) REDIS_CONNECTION_POOL_TIMEOUT = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5)) -REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD = int( - os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5) -) -REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT = int( - os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60) -) -REDIS_CIRCUIT_BREAKER_ENABLED = ( - os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true" -) +REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5)) +REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)) +REDIS_CIRCUIT_BREAKER_ENABLED = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true" # Default Redis major version to assume when version cannot be determined # Using 7 as it's the modern version that supports LPOP with count parameter DEFAULT_REDIS_MAJOR_VERSION = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7)) @@ -436,17 +341,11 @@ NON_LLM_CONNECTION_TIMEOUT = int( MAX_EXCEPTION_MESSAGE_LENGTH = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000)) MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048)) BEDROCK_MAX_POLICY_SIZE = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75)) -BEDROCK_MIN_THINKING_BUDGET_TOKENS = int( - os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024) -) +BEDROCK_MIN_THINKING_BUDGET_TOKENS = int(os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024)) # Anthropic's Messages API rejects thinking.budget_tokens < 1024. ANTHROPIC_MIN_THINKING_BUDGET_TOKENS = 1024 -REPLICATE_POLLING_DELAY_SECONDS = float( - os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5) -) -DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS = int( - os.getenv("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096) -) +REPLICATE_POLLING_DELAY_SECONDS = float(os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5)) +DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS = int(os.getenv("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096)) DEFAULT_OCI_CHAT_MAX_TOKENS = 4096 TOGETHER_AI_4_B = int(os.getenv("TOGETHER_AI_4_B", 4)) TOGETHER_AI_8_B = int(os.getenv("TOGETHER_AI_8_B", 8)) @@ -475,13 +374,9 @@ DEFAULT_REQUEST_TIMEOUT_SECONDS: float = 6000.0 # deadline and connect handshake (see ``http_handler`` cached handler paths). COMPLETION_HTTP_FALLBACK_SECONDS: float = 600.0 HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: float = 5.0 -request_timeout: float = float( - os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS))) -) +request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS)))) request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ -DEFAULT_A2A_AGENT_TIMEOUT: float = float( - os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000) -) # 10 minutes +DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes # Patterns that indicate a localhost/internal URL in A2A agent cards that should be # replaced with the original base_url. This is a common misconfiguration where # developers deploy agents with development URLs in their agent cards. @@ -511,16 +406,10 @@ FIREWORKS_AI_16_B = int(os.getenv("FIREWORKS_AI_16_B", 16)) FIREWORKS_AI_80_B = int(os.getenv("FIREWORKS_AI_80_B", 80)) #### Logging callback constants #### REDACTED_BY_LITELM_STRING = "REDACTED_BY_LITELM" -MAX_LANGFUSE_INITIALIZED_CLIENTS = int( - os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50) -) -LOGGING_WORKER_CONCURRENCY = int( - os.getenv("LOGGING_WORKER_CONCURRENCY", 100) -) # Must be above 0 +MAX_LANGFUSE_INITIALIZED_CLIENTS = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)) +LOGGING_WORKER_CONCURRENCY = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 LOGGING_WORKER_MAX_QUEUE_SIZE = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) -LOGGING_WORKER_MAX_TIME_PER_COROUTINE = float( - os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0) -) +LOGGING_WORKER_MAX_TIME_PER_COROUTINE = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) LOGGING_WORKER_CLEAR_PERCENTAGE = int( os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50) ) # Percentage of queue to clear (default: 50%) @@ -535,17 +424,13 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv( LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED = 499 -EMAIL_BUDGET_ALERT_TTL = int( - os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60) -) # 24 hours in seconds +EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float( os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8) ) # 80% of max budget ############### LLM Provider Constants ############### ### ANTHROPIC CONSTANTS ### -ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv( - "ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01" -) +ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01") ANTHROPIC_SKILLS_API_BETA_VERSION = "skills-2025-10-02" ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES = { "low": 1, @@ -560,9 +445,7 @@ LITELLM_WEB_SEARCH_TOOL_NAME = "litellm_web_search" DEFAULT_IMAGE_ENDPOINT_MODEL = "dall-e-2" DEFAULT_VIDEO_ENDPOINT_MODEL = "sora-2" -DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS = int( - os.getenv("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8) -) +DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS = int(os.getenv("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8)) ### DATAFORSEO CONSTANTS ### DEFAULT_DATAFORSEO_LOCATION_CODE = int( @@ -945,8 +828,7 @@ clarifai_models: set = set( "clarifai/qwen.qwenLM.Qwen3-30B-A3B-Instruct-2507", "clarifai/qwen.qwen3.qwen3-next-80B-A3B-Thinking", "clarifai/openai.chat-completion.gpt-oss-120b", - "clarifai/qwen.qwenLM.Qwen3-30B-A3B-Thinking-2507" - "clarifai/openai.chat-completion.gpt-5-nano", + "clarifai/qwen.qwenLM.Qwen3-30B-A3B-Thinking-2507clarifai/openai.chat-completion.gpt-5-nano", "clarifai/openai.chat-completion.gpt-4o", "clarifai/gcp.generate.gemini-2_5-pro", "clarifai/anthropic.completion.claude-sonnet-4", @@ -1375,9 +1257,7 @@ OPENAI_FINISH_REASONS = [ "tool_calls", "content_filter", ] -HUMANLOOP_PROMPT_CACHE_TTL_SECONDS = int( - os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60) -) # 1 minute +HUMANLOOP_PROMPT_CACHE_TTL_SECONDS = int(os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60)) # 1 minute RESPONSE_FORMAT_TOOL_NAME = "json_tool_call" # default tool name used when converting response format to tool call ########################### Logging Callback Constants ########################### @@ -1385,9 +1265,7 @@ AZURE_STORAGE_MSFT_VERSION = "2019-07-07" PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES = int( os.getenv("PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5) ) -CLOUDZERO_EXPORT_INTERVAL_MINUTES = int( - os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60) -) +CLOUDZERO_EXPORT_INTERVAL_MINUTES = int(os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60)) MCP_TOOL_NAME_PREFIX = "mcp_tool" MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)) @@ -1450,37 +1328,23 @@ PASS_THROUGH_HEADER_PREFIX = "x-pass-" BASE_MCP_ROUTE = "/mcp" -BATCH_STATUS_POLL_INTERVAL_SECONDS = int( - os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600) -) # 1 hour -BATCH_STATUS_POLL_MAX_ATTEMPTS = int( - os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24) -) # for 24 hours +BATCH_STATUS_POLL_INTERVAL_SECONDS = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour +BATCH_STATUS_POLL_MAX_ATTEMPTS = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours -HEALTH_CHECK_TIMEOUT_SECONDS = int( - os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60) -) # 60 seconds -_background_health_check_max_tokens_env = os.getenv( - "BACKGROUND_HEALTH_CHECK_MAX_TOKENS" -) +HEALTH_CHECK_TIMEOUT_SECONDS = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds +_background_health_check_max_tokens_env = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS") try: _raw_background_health_check_max_tokens = ( - _background_health_check_max_tokens_env.strip() - if _background_health_check_max_tokens_env is not None - else "" + _background_health_check_max_tokens_env.strip() if _background_health_check_max_tokens_env is not None else "" ) BACKGROUND_HEALTH_CHECK_MAX_TOKENS: Optional[int] = ( - int(_raw_background_health_check_max_tokens) - if _raw_background_health_check_max_tokens - else None + int(_raw_background_health_check_max_tokens) if _raw_background_health_check_max_tokens else None ) except (ValueError, TypeError): BACKGROUND_HEALTH_CHECK_MAX_TOKENS = None -_background_health_check_max_tokens_reasoning_env = os.getenv( - "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING" -) +_background_health_check_max_tokens_reasoning_env = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING") try: _raw_background_health_check_max_tokens_reasoning = ( _background_health_check_max_tokens_reasoning_env.strip() @@ -1522,9 +1386,7 @@ LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int( os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600) ) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" -LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv( - "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false" -) +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false") LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS = int( os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400) ) # 24 hours default @@ -1541,15 +1403,11 @@ CLI_SSO_SESSION_TTL_SECONDS = 600 CLI_SESSION_KEY_PREFIX = "cli-session" # Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility CLI_JWT_EXPIRATION_HOURS = int( - os.getenv("CLI_JWT_EXPIRATION_HOURS") - or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS") - or 24 + os.getenv("CLI_JWT_EXPIRATION_HOURS") or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS") or 24 ) # Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g. # "employment_type->acme_employment_type,org_info.department->department" -CLI_SSO_CLAIM_MAP = ( - os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or "" -) +CLI_SSO_CLAIM_MAP = os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or "" CLI_SSO_CLAIM_MAX_SCALAR_LENGTH = 1024 ########################### UI SESSION DURATION ########################### @@ -1563,54 +1421,34 @@ DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics" CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data" MAVVRIK_FOCUS_EXPORT_JOB_NAME = "mavvrik_focus_export_usage_data" -CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( - os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000) -) +CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)) SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) -SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int( - os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) -) +SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)) SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float( os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) ) SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") -SPEND_LOG_PARTITION_PRECREATE_AHEAD = int( - os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7) -) +SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) -SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int( - os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000) -) -DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int( - os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60) -) # 1 minute -PROXY_BUDGET_RESCHEDULER_MIN_TIME = int( - os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597) -) +SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) +DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute +PROXY_BUDGET_RESCHEDULER_MIN_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) -MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max( - 1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) -) -STALE_OBJECT_CLEANUP_BATCH_SIZE = max( - 1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000)) -) +MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) +STALE_OBJECT_CLEANUP_BATCH_SIZE = max(1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))) # Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and # CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on # installations with large numbers of stale managed objects). _batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true" -PROXY_BUDGET_RESCHEDULER_MAX_TIME = int( - os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605) -) -PROXY_BATCH_WRITE_AT = int( - os.getenv("PROXY_BATCH_WRITE_AT", 10) -) # in seconds, increased from 10 +PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) +PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 # APScheduler Configuration - MEMORY LEAK FIX # These settings prevent memory leaks in APScheduler's normalize() and _apply_jitter() functions @@ -1621,12 +1459,8 @@ APSCHEDULER_COALESCE = os.getenv("APSCHEDULER_COALESCE", "True").lower() in [ APSCHEDULER_MISFIRE_GRACE_TIME = int( os.getenv("APSCHEDULER_MISFIRE_GRACE_TIME", 3600) ) # ignore runs older than 1 hour (was 120) -APSCHEDULER_MAX_INSTANCES = int( - os.getenv("APSCHEDULER_MAX_INSTANCES", 1) -) # prevent concurrent job instances -APSCHEDULER_REPLACE_EXISTING = os.getenv( - "APSCHEDULER_REPLACE_EXISTING", "True" -).lower() in [ +APSCHEDULER_MAX_INSTANCES = int(os.getenv("APSCHEDULER_MAX_INSTANCES", 1)) # prevent concurrent job instances +APSCHEDULER_REPLACE_EXISTING = os.getenv("APSCHEDULER_REPLACE_EXISTING", "True").lower() in [ "true", "1", ] # always replace existing jobs @@ -1635,38 +1469,24 @@ APSCHEDULER_REPLACE_EXISTING = os.getenv( # This will run tag spcific tasks at a later time to smooth QPS DAILY_TAG_SPEND_BATCH_MULTIPLIER = 2.3 -DEFAULT_HEALTH_CHECK_INTERVAL = int( - os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300) -) # 5 minutes +DEFAULT_HEALTH_CHECK_INTERVAL = int(os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300)) # 5 minutes DEFAULT_SHARED_HEALTH_CHECK_TTL = int( os.getenv("DEFAULT_SHARED_HEALTH_CHECK_TTL", 300) ) # 5 minutes - TTL for cached health check results DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL = int( os.getenv("DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL", 60) ) # 1 minute - TTL for health check lock -DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER = ( - 2 # health state is stale after interval * this -) -PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS = int( - os.getenv("PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS", 9) -) +DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER = 2 # health state is stale after interval * this +PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS = int(os.getenv("PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS", 9)) DEFAULT_MODEL_CREATED_AT_TIME = int( os.getenv("DEFAULT_MODEL_CREATED_AT_TIME", 1677610602) ) # returns on `/models` endpoint -DEFAULT_SLACK_ALERTING_THRESHOLD = int( - os.getenv("DEFAULT_SLACK_ALERTING_THRESHOLD", 300) -) +DEFAULT_SLACK_ALERTING_THRESHOLD = int(os.getenv("DEFAULT_SLACK_ALERTING_THRESHOLD", 300)) MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20)) -MAX_POLICY_ESTIMATE_IMPACT_ROWS = int( - os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000) -) -DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float( - os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7) -) +MAX_POLICY_ESTIMATE_IMPACT_ROWS = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000)) +DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7)) LENGTH_OF_LITELLM_GENERATED_KEY = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16)) -SECRET_MANAGER_REFRESH_INTERVAL = int( - os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400) -) +SECRET_MANAGER_REFRESH_INTERVAL = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400)) LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "default_internal_user_params", "default_team_params", @@ -1678,9 +1498,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "cost_margin_config", ] SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] -DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int( - os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60) -) +DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding @@ -1760,9 +1578,7 @@ SENTRY_PII_DENYLIST = [ ] # CoroutineChecker cache configuration -COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int( - os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000) -) +COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int(os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000)) ########################### RAG Text Splitter Constants ########################### DEFAULT_CHUNK_SIZE = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000)) @@ -1770,31 +1586,19 @@ DEFAULT_CHUNK_OVERLAP = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200)) ########################### S3 Vectors RAG Constants ########################### S3_VECTORS_DEFAULT_DIMENSION = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 1024)) -S3_VECTORS_DEFAULT_DISTANCE_METRIC = str( - os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine") -) +S3_VECTORS_DEFAULT_DISTANCE_METRIC = str(os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine")) S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS = ["source_text"] ########################### Microsoft SSO Constants ########################### -MICROSOFT_USER_EMAIL_ATTRIBUTE = str( - os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName") -) -MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE = str( - os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName") -) +MICROSOFT_USER_EMAIL_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName")) +MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName")) MICROSOFT_USER_ID_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id")) -MICROSOFT_USER_FIRST_NAME_ATTRIBUTE = str( - os.getenv("MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "givenName") -) -MICROSOFT_USER_LAST_NAME_ATTRIBUTE = str( - os.getenv("MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "surname") -) +MICROSOFT_USER_FIRST_NAME_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "givenName")) +MICROSOFT_USER_LAST_NAME_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "surname")) # Maximum payload size (in bytes) to fully serialize for DEBUG logging. # Payloads larger than this are truncated to avoid multi-second json.dumps blocking the response. -MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG = int( - os.getenv("MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG", 102400) -) # 100 KB +MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG = int(os.getenv("MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG", 102400)) # 100 KB # Policy template enrichment MAX_COMPETITOR_NAMES = int(os.getenv("MAX_COMPETITOR_NAMES", 100)) diff --git a/litellm/containers/endpoint_factory.py b/litellm/containers/endpoint_factory.py index a5f6951862f..bebdfa2f9e6 100644 --- a/litellm/containers/endpoint_factory.py +++ b/litellm/containers/endpoint_factory.py @@ -97,9 +97,7 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable: ) if container_provider_config is None: - raise ValueError( - f"Container provider config not found for: {resolved_custom_llm_provider}" - ) + raise ValueError(f"Container provider config not found for: {resolved_custom_llm_provider}") # Build optional params for logging optional_params = {k: kwargs.get(k) for k in path_params if k in kwargs} @@ -239,9 +237,5 @@ retrieve_container_file = _generated_endpoints.get("retrieve_container_file") aretrieve_container_file = _generated_endpoints.get("aretrieve_container_file") delete_container_file = _generated_endpoints.get("delete_container_file") adelete_container_file = _generated_endpoints.get("adelete_container_file") -retrieve_container_file_content = _generated_endpoints.get( - "retrieve_container_file_content" -) -aretrieve_container_file_content = _generated_endpoints.get( - "aretrieve_container_file_content" -) +retrieve_container_file_content = _generated_endpoints.get("retrieve_container_file_content") +aretrieve_container_file_content = _generated_endpoints.get("aretrieve_container_file_content") diff --git a/litellm/containers/main.py b/litellm/containers/main.py index c0ca550c9a9..caf6c684844 100644 --- a/litellm/containers/main.py +++ b/litellm/containers/main.py @@ -211,31 +211,23 @@ def create_container( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError( - f"container operations are not supported for {custom_llm_provider}" - ) + raise ValueError(f"container operations are not supported for {custom_llm_provider}") local_vars.update(kwargs) # Get ContainerCreateOptionalRequestParams with only valid parameters container_create_optional_params: ContainerCreateOptionalRequestParams = ( - ContainerRequestUtils.get_requested_container_create_optional_param( - local_vars - ) + ContainerRequestUtils.get_requested_container_create_optional_param(local_vars) ) # Get optional parameters for the container API - container_create_request_params: Dict = ( - ContainerRequestUtils.get_optional_params_container_create( - container_provider_config=container_provider_config, - container_create_optional_params=container_create_optional_params, - ) + container_create_request_params: Dict = ContainerRequestUtils.get_optional_params_container_create( + container_provider_config=container_provider_config, + container_create_optional_params=container_create_optional_params, ) # Pre Call logging @@ -440,22 +432,16 @@ def list_containers( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError( - f"Container provider config not found for provider: {custom_llm_provider}" - ) + raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}") # Get container list request parameters container_list_optional_params: ContainerListOptionalRequestParams = ( - ContainerRequestUtils.get_requested_container_list_optional_param( - local_vars - ) + ContainerRequestUtils.get_requested_container_list_optional_param(local_vars) ) # Pre Call logging @@ -641,27 +627,21 @@ def retrieve_container( ) # Decode container ID and extract provider info - original_container_id, resolved_custom_llm_provider, litellm_params = ( - decode_managed_container_id_for_request( - container_id=container_id, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - ) + original_container_id, resolved_custom_llm_provider, litellm_params = decode_managed_container_id_for_request( + container_id=container_id, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, ) # True when input was a LiteLLM-managed ID (any length); needed to re-encode output for routing affinity was_encoded = original_container_id != container_id # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(resolved_custom_llm_provider), - ) + container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(resolved_custom_llm_provider), ) if container_provider_config is None: - raise ValueError( - f"Container provider config not found for provider: {resolved_custom_llm_provider}" - ) + raise ValueError(f"Container provider config not found for provider: {resolved_custom_llm_provider}") # Pre Call logging litellm_logging_obj.update_from_kwargs( @@ -865,27 +845,21 @@ def delete_container( ) # Decode container ID and extract provider info - original_container_id, resolved_custom_llm_provider, litellm_params = ( - decode_managed_container_id_for_request( - container_id=container_id, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - ) + original_container_id, resolved_custom_llm_provider, litellm_params = decode_managed_container_id_for_request( + container_id=container_id, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, ) # True when input was a LiteLLM-managed ID (any length); needed to re-encode output for routing affinity was_encoded = original_container_id != container_id # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(resolved_custom_llm_provider), - ) + container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(resolved_custom_llm_provider), ) if container_provider_config is None: - raise ValueError( - f"Container provider config not found for provider: {resolved_custom_llm_provider}" - ) + raise ValueError(f"Container provider config not found for provider: {resolved_custom_llm_provider}") # Pre Call logging litellm_logging_obj.update_from_kwargs( @@ -1103,25 +1077,19 @@ def list_container_files( ) # Decode container ID and extract provider info - original_container_id, resolved_custom_llm_provider, litellm_params = ( - decode_managed_container_id_for_request( - container_id=container_id, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - ) + original_container_id, resolved_custom_llm_provider, litellm_params = decode_managed_container_id_for_request( + container_id=container_id, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(resolved_custom_llm_provider), - ) + container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(resolved_custom_llm_provider), ) if container_provider_config is None: - raise ValueError( - f"Container provider config not found for provider: {resolved_custom_llm_provider}" - ) + raise ValueError(f"Container provider config not found for provider: {resolved_custom_llm_provider}") # Pre Call logging litellm_logging_obj.update_from_kwargs( @@ -1363,25 +1331,19 @@ def upload_container_file( ) # Decode container ID and extract provider info - original_container_id, resolved_custom_llm_provider, litellm_params = ( - decode_managed_container_id_for_request( - container_id=container_id, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - ) + original_container_id, resolved_custom_llm_provider, litellm_params = decode_managed_container_id_for_request( + container_id=container_id, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(resolved_custom_llm_provider), - ) + container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(resolved_custom_llm_provider), ) if container_provider_config is None: - raise ValueError( - f"Container provider config not found for provider: {resolved_custom_llm_provider}" - ) + raise ValueError(f"Container provider config not found for provider: {resolved_custom_llm_provider}") # Pre Call logging litellm_logging_obj.update_from_kwargs( diff --git a/litellm/containers/utils.py b/litellm/containers/utils.py index 7c66eb70eb5..2b115c6b3c4 100644 --- a/litellm/containers/utils.py +++ b/litellm/containers/utils.py @@ -66,11 +66,7 @@ class ContainerRequestUtils: supported_params = container_provider_config.get_supported_openai_params() # Filter out unsupported parameters - filtered_params = { - k: v - for k, v in container_create_optional_params.items() - if k in supported_params - } + filtered_params = {k: v for k, v in container_create_optional_params.items() if k in supported_params} return container_provider_config.map_openai_params( container_create_optional_params=filtered_params, # type: ignore diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f95deda0f81..8ddc69f5396 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -317,9 +317,7 @@ def cost_per_token( ### SERVICE TIER ### service_tier: Optional[str] = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### - data_residency: Optional[ - str - ] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") response: Optional[Any] = None, ### REQUEST MODEL ### request_model: Optional[str] = None, # original request model for router detection @@ -376,9 +374,7 @@ def cost_per_token( # either `cache_write_tokens` (kimi-k2) or `cache_creation_tokens`. # Mirror db_spend_update_writer to stay symmetric. _cache_creation_tokens = float( - getattr(_pt_details, "cache_write_tokens", 0) - or getattr(_pt_details, "cache_creation_tokens", 0) - or 0 + getattr(_pt_details, "cache_write_tokens", 0) or getattr(_pt_details, "cache_creation_tokens", 0) or 0 ) _anthropic_read = getattr(usage_object, "cache_read_input_tokens", None) @@ -451,12 +447,8 @@ def cost_per_token( else: model_with_provider = f"{custom_llm_provider}/{model}" if region_name is not None: - model_with_provider_and_region = ( - f"{custom_llm_provider}/{region_name}/{model}" - ) - if ( - model_with_provider_and_region in model_cost_ref - ): # use region based pricing, if it's available + model_with_provider_and_region = f"{custom_llm_provider}/{region_name}/{model}" + if model_with_provider_and_region in model_cost_ref: # use region based pricing, if it's available model_with_provider = model_with_provider_and_region else: _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) @@ -475,9 +467,7 @@ def cost_per_token( Option2. model = "openai/gpt-4" - model = provider/model Option3. model = "anthropic.claude-3" - model = model """ - if ( - model_with_provider in model_cost_ref - ): # Option 2. use model with provider, model = "openai/gpt-4" + if model_with_provider in model_cost_ref: # Option 2. use model with provider, model = "openai/gpt-4" model = model_with_provider elif model in model_cost_ref: # Option 1. use model passed, model="gpt-4" model = model @@ -488,9 +478,7 @@ def cost_per_token( # see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models if call_type == "speech" or call_type == "aspeech": - speech_model_info = litellm.get_model_info( - model=model_without_prefix, custom_llm_provider=custom_llm_provider - ) + speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) cost_metric = select_cost_metric_for_model(speech_model_info) prompt_cost: float = 0.0 completion_cost: float = 0.0 @@ -587,11 +575,7 @@ def cost_per_token( model=model, custom_llm_provider=custom_llm_provider, number_of_queries=number_of_queries or 1, - optional_params=( - response._hidden_params - if response and hasattr(response, "_hidden_params") - else None - ), + optional_params=(response._hidden_params if response and hasattr(response, "_hidden_params") else None), ) elif custom_llm_provider == "vertex_ai": cost_router = google_cost_router( @@ -615,13 +599,9 @@ def cost_per_token( service_tier=service_tier, ) elif custom_llm_provider == "anthropic": - return anthropic_cost_per_token( - model=model, usage=usage_block, service_tier=service_tier - ) + return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier) elif custom_llm_provider == "bedrock": - return bedrock_cost_per_token( - model=model, usage=usage_block, service_tier=service_tier - ) + return bedrock_cost_per_token(model=model, usage=usage_block, service_tier=service_tier) elif custom_llm_provider == "openai": return openai_cost_per_token( model=model, @@ -641,9 +621,7 @@ def cost_per_token( service_tier=service_tier, ) elif custom_llm_provider == "gemini": - return gemini_cost_per_token( - model=model, usage=usage_block, service_tier=service_tier - ) + return gemini_cost_per_token(model=model, usage=usage_block, service_tier=service_tier) elif custom_llm_provider == "deepseek": return deepseek_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "perplexity": @@ -667,13 +645,9 @@ def cost_per_token( service_tier=service_tier, ) else: - model_info = _cached_get_model_info_helper( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) - if (model_info.get("input_cost_per_token") or 0.0) > 0 or ( - model_info.get("output_cost_per_token") or 0.0 - ) > 0: + if (model_info.get("input_cost_per_token") or 0.0) > 0 or (model_info.get("output_cost_per_token") or 0.0) > 0: return generic_cost_per_token( model=model, usage=usage_block, @@ -682,10 +656,7 @@ def cost_per_token( data_residency=data_residency, ) - if ( - model_info.get("input_cost_per_second", None) is not None - and response_time_ms is not None - ): + if model_info.get("input_cost_per_second", None) is not None and response_time_ms is not None: verbose_logger.debug( "For model=%s - input_cost_per_second: %s; response time: %s", model, @@ -697,10 +668,7 @@ def cost_per_token( model_info["input_cost_per_second"] * response_time_ms / 1000 # type: ignore ) - if ( - model_info.get("output_cost_per_second", None) is not None - and response_time_ms is not None - ): + if model_info.get("output_cost_per_second", None) is not None and response_time_ms is not None: verbose_logger.debug( "For model=%s - output_cost_per_second: %s; response time: %s", model, @@ -724,7 +692,9 @@ def cost_per_token( def get_replicate_completion_pricing(completion_response: dict, total_time=0.0): # see https://replicate.com/pricing # for all litellm currently supported LLMs, almost all requests go to a100_80gb - a100_80gb_price_per_second_public = DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND # assume all calls sent to A100 80GB for now + a100_80gb_price_per_second_public = ( + DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND # assume all calls sent to A100 80GB for now + ) if total_time == 0.0: # total time is in ms start_time = completion_response.get("created", time.time()) end_time = getattr(completion_response, "ended", time.time()) @@ -773,9 +743,7 @@ def _select_model_name_for_cost_calc( return_model: Optional[str] = None region_name: Optional[str] = None - custom_llm_provider = _get_provider_for_cost_calc( - model=model, custom_llm_provider=custom_llm_provider - ) + custom_llm_provider = _get_provider_for_cost_calc(model=model, custom_llm_provider=custom_llm_provider) completion_response_model: Optional[str] = None if completion_response is not None: @@ -788,10 +756,7 @@ def _select_model_name_for_cost_calc( if custom_pricing is True: if router_model_id is not None and router_model_id in litellm.model_cost: entry = litellm.model_cost[router_model_id] - if ( - entry.get("input_cost_per_token") is not None - or entry.get("input_cost_per_second") is not None - ): + if entry.get("input_cost_per_token") is not None or entry.get("input_cost_per_second") is not None: return_model = router_model_id else: return_model = model @@ -802,14 +767,9 @@ def _select_model_name_for_cost_calc( return_model = base_model elif completion_response_model is None and hidden_params is not None: - if ( - hidden_params.get("model", None) is not None - and len(hidden_params["model"]) > 0 - ): + if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0: return_model = hidden_params.get("model", model) - elif ( - hidden_params is not None and hidden_params.get("region_name", None) is not None - ): + elif hidden_params is not None and hidden_params.get("region_name", None) is not None: region_name = hidden_params.get("region_name", None) if return_model is None and completion_response_model is not None: @@ -897,10 +857,7 @@ def _normalize_service_tier(service_tier: object) -> str | None: on the response usage) instead of crashing the downstream cost-key lookup, which calls service_tier.lower() """ - if ( - not isinstance(service_tier, str) - or service_tier.lower() == ServiceTier.AUTO.value - ): + if not isinstance(service_tier, str) or service_tier.lower() == ServiceTier.AUTO.value: return None return service_tier @@ -926,20 +883,12 @@ def _get_usage_object( and (isinstance(usage_obj, dict) or isinstance(usage_obj, ResponseAPIUsage)) and ResponseAPILoggingUtils._is_response_api_usage(usage_obj) ): - return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage_obj - ) - elif TranscriptionUsageObjectTransformation.is_transcription_usage_object( - usage_obj - ): - return ( - TranscriptionUsageObjectTransformation.transform_transcription_usage_object( - cast( - Union[ - TranscriptionUsageDurationObject, TranscriptionUsageTokensObject - ], - usage_obj, - ) + return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage_obj) + elif TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj): + return TranscriptionUsageObjectTransformation.transform_transcription_usage_object( + cast( + Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject], + usage_obj, ) ) elif isinstance(usage_obj, dict): @@ -947,9 +896,7 @@ def _get_usage_object( elif isinstance(usage_obj, BaseModel): return Usage(**usage_obj.model_dump()) else: - verbose_logger.debug( - f"Unknown usage object type: {type(usage_obj)}, usage_obj: {usage_obj}" - ) + verbose_logger.debug(f"Unknown usage object type: {type(usage_obj)}, usage_obj: {usage_obj}") return None @@ -958,24 +905,18 @@ def _is_known_usage_objects(usage_obj): return ( isinstance(usage_obj, litellm.Usage) or isinstance(usage_obj, ResponseAPIUsage) - or TranscriptionUsageObjectTransformation.is_transcription_usage_object( - usage_obj - ) + or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj) ) -def _infer_call_type( - call_type: Optional[CallTypesLiteral], completion_response: Any -) -> Optional[CallTypesLiteral]: +def _infer_call_type(call_type: Optional[CallTypesLiteral], completion_response: Any) -> Optional[CallTypesLiteral]: if call_type is not None: return call_type if completion_response is None: return None - if isinstance(completion_response, ModelResponse) or isinstance( - completion_response, ModelResponseStream - ): + if isinstance(completion_response, ModelResponse) or isinstance(completion_response, ModelResponseStream): return "completion" elif isinstance(completion_response, EmbeddingResponse): return "embedding" @@ -1053,9 +994,7 @@ def _apply_cost_margin( if custom_llm_provider and custom_llm_provider in litellm.cost_margin_config: margin_config = litellm.cost_margin_config[custom_llm_provider] if verbose_logger.isEnabledFor(logging.DEBUG): - verbose_logger.debug( - f"Found provider-specific margin config for {custom_llm_provider}: {margin_config}" - ) + verbose_logger.debug(f"Found provider-specific margin config for {custom_llm_provider}: {margin_config}") elif "global" in litellm.cost_margin_config: margin_config = litellm.cost_margin_config["global"] if verbose_logger.isEnabledFor(logging.DEBUG): @@ -1184,9 +1123,7 @@ def completion_cost( ### SERVICE TIER ### service_tier: Optional[str] = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### - data_residency: Optional[ - str - ] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ) -> float: """ Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. @@ -1235,9 +1172,7 @@ def completion_cost( cache_creation_input_tokens: Optional[int] = None cache_read_input_tokens: Optional[int] = None audio_transcription_file_duration: float = 0.0 - cost_per_token_usage_object: Optional[Usage] = _get_usage_object( - completion_response=completion_response - ) + cost_per_token_usage_object: Optional[Usage] = _get_usage_object(completion_response=completion_response) rerank_billed_units: Optional[RerankBilledUnits] = None # Extract service_tier from optional_params if not provided directly @@ -1258,9 +1193,7 @@ def completion_cost( # Extract service_tier from usage object if not provided if service_tier is None and cost_per_token_usage_object is not None: if isinstance(cost_per_token_usage_object, BaseModel): - service_tier = getattr( - cost_per_token_usage_object, "service_tier", None - ) + service_tier = getattr(cost_per_token_usage_object, "service_tier", None) elif isinstance(cost_per_token_usage_object, dict): service_tier = cost_per_token_usage_object.get("service_tier") @@ -1285,23 +1218,16 @@ def completion_cost( for idx, model in enumerate(potential_model_names): try: if verbose_logger.isEnabledFor(logging.DEBUG): - verbose_logger.debug( - f"selected model name for cost calculation: {model}" - ) + verbose_logger.debug(f"selected model name for cost calculation: {model}") if completion_response is not None and ( - isinstance(completion_response, BaseModel) - or isinstance(completion_response, dict) + isinstance(completion_response, BaseModel) or isinstance(completion_response, dict) ): # tts returns a custom class if isinstance(completion_response, dict): - usage_obj: Optional[Union[dict, Usage]] = ( - completion_response.get("usage", {}) - ) + usage_obj: Optional[Union[dict, Usage]] = completion_response.get("usage", {}) else: usage_obj = getattr(completion_response, "usage", {}) - if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects( - usage_obj=usage_obj - ): + if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects(usage_obj=usage_obj): _usage_for_dump = cast(BaseModel, usage_obj) setattr( completion_response, @@ -1319,9 +1245,7 @@ def completion_cost( _usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( _usage ).model_dump() - elif TranscriptionUsageObjectTransformation.is_transcription_usage_object( - _usage - ): + elif TranscriptionUsageObjectTransformation.is_transcription_usage_object(_usage): tr_usage = TranscriptionUsageObjectTransformation.transform_transcription_usage_object( cast( Union[ @@ -1339,29 +1263,21 @@ def completion_cost( # get input/output tokens from completion_response prompt_tokens = _usage.get("prompt_tokens", 0) completion_tokens = _usage.get("completion_tokens", 0) - cache_creation_input_tokens = _usage.get( - "cache_creation_input_tokens", 0 - ) + cache_creation_input_tokens = _usage.get("cache_creation_input_tokens", 0) cache_read_input_tokens = _usage.get("cache_read_input_tokens", 0) if ( "prompt_tokens_details" in _usage and _usage["prompt_tokens_details"] != {} and _usage["prompt_tokens_details"] ): - prompt_tokens_details = ( - _usage.get("prompt_tokens_details") or {} - ) - cache_read_input_tokens = prompt_tokens_details.get( - "cached_tokens", 0 - ) + prompt_tokens_details = _usage.get("prompt_tokens_details") or {} + cache_read_input_tokens = prompt_tokens_details.get("cached_tokens", 0) total_time = getattr(completion_response, "_response_ms", 0) hidden_params = getattr(completion_response, "_hidden_params", None) if hidden_params is not None: - custom_llm_provider = hidden_params.get( - "custom_llm_provider", custom_llm_provider or None - ) + custom_llm_provider = hidden_params.get("custom_llm_provider", custom_llm_provider or None) region_name = hidden_params.get("region_name", region_name) # For Gemini/Vertex AI responses, trafficType is stored in @@ -1369,14 +1285,10 @@ def completion_cost( # by the cost key lookup (_priority / _flex suffixes) so that # ON_DEMAND_PRIORITY requests are billed at priority prices. if service_tier is None: - provider_specific = ( - hidden_params.get("provider_specific_fields") or {} - ) + provider_specific = hidden_params.get("provider_specific_fields") or {} raw_traffic_type = provider_specific.get("traffic_type") if raw_traffic_type: - service_tier = _map_traffic_type_to_service_tier( - raw_traffic_type - ) + service_tier = _map_traffic_type_to_service_tier(raw_traffic_type) else: if model is None: raise ValueError( @@ -1392,9 +1304,7 @@ def completion_cost( if call_type in _A2A_CALL_TYPES: from litellm.a2a_protocol.cost_calculator import A2ACostCalculator - return A2ACostCalculator.calculate_a2a_cost( - litellm_logging_obj=litellm_logging_obj - ) + return A2ACostCalculator.calculate_a2a_cost(litellm_logging_obj=litellm_logging_obj) if model is None: raise ValueError( @@ -1411,9 +1321,9 @@ def completion_cost( str(e) ) ) - if CostCalculatorUtils._call_type_has_image_response( - call_type - ) and isinstance(completion_response, ImageResponse): + if CostCalculatorUtils._call_type_has_image_response(call_type) and isinstance( + completion_response, ImageResponse + ): ### IMAGE GENERATION COST CALCULATION ### return CostCalculatorUtils.route_image_generation_cost_calculator( model=model, @@ -1430,9 +1340,7 @@ def completion_cost( # Extract custom model_info for deployment-specific pricing _video_model_info: Optional[ModelInfo] = None if custom_pricing and litellm_logging_obj is not None: - _litellm_params = getattr( - litellm_logging_obj, "litellm_params", None - ) + _litellm_params = getattr(litellm_logging_obj, "litellm_params", None) if _litellm_params is not None: _metadata = _litellm_params.get("metadata", {}) or {} _video_model_info = _metadata.get("model_info", None) @@ -1446,9 +1354,7 @@ def completion_cost( duration_seconds = usage_obj.get("duration_seconds", None) _vr = usage_obj.get("video_resolution", None) else: - duration_seconds = getattr( - usage_obj, "duration_seconds", None - ) + duration_seconds = getattr(usage_obj, "duration_seconds", None) _vr = getattr(usage_obj, "video_resolution", None) if _vr is not None: video_resolution = str(_vr).strip().lower() @@ -1487,9 +1393,7 @@ def completion_cost( getattr(completion_response, "duration", 0.0), ) elif call_type in _RERANK_CALL_TYPES: - if completion_response is not None and isinstance( - completion_response, RerankResponse - ): + if completion_response is not None and isinstance(completion_response, RerankResponse): meta_obj = completion_response.meta if meta_obj is not None: billed_units = meta_obj.get("billed_units", {}) or {} @@ -1501,9 +1405,7 @@ def completion_cost( total_tokens=billed_units.get("total_tokens"), ) - search_units = ( - billed_units.get("search_units") or 1 - ) # cohere charges per request by default. + search_units = billed_units.get("search_units") or 1 # cohere charges per request by default. completion_tokens = search_units elif call_type in _SEARCH_CALL_TYPES: from litellm.search import search_provider_cost_per_query @@ -1577,10 +1479,7 @@ def completion_cost( elif call_type == _AREALTIME_CALL_TYPE and isinstance( completion_response, LiteLLMRealtimeStreamLoggingObject ): - if ( - cost_per_token_usage_object is None - or custom_llm_provider is None - ): + if cost_per_token_usage_object is None or custom_llm_provider is None: raise ValueError( "usage object and custom_llm_provider must be provided for realtime stream cost calculation. Got cost_per_token_usage_object={}, custom_llm_provider={}".format( cost_per_token_usage_object, @@ -1599,59 +1498,36 @@ def completion_cost( MCPCostCalculator, ) - return MCPCostCalculator.calculate_mcp_tool_call_cost( - litellm_logging_obj=litellm_logging_obj - ) + return MCPCostCalculator.calculate_mcp_tool_call_cost(litellm_logging_obj=litellm_logging_obj) # Calculate cost based on prompt_tokens, completion_tokens - if ( - "togethercomputer" in model - or "together_ai" in model - or custom_llm_provider == "together_ai" - ): + if "togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai": # together ai prices based on size of llm # get_model_params_and_category takes a model name and returns the category of LLM size it is in model_prices_and_context_window.json - model = get_model_params_and_category( - model, call_type=CallTypes(call_type) - ) + model = get_model_params_and_category(model, call_type=CallTypes(call_type)) # replicate llms are calculate based on time for request running # see https://replicate.com/pricing - elif ( - model in litellm.replicate_models or "replicate" in model - ) and model not in litellm.model_cost: + elif (model in litellm.replicate_models or "replicate" in model) and model not in litellm.model_cost: # for unmapped replicate model, default to replicate's time tracking logic - return get_replicate_completion_pricing( - completion_response, total_time - ) # type: ignore + return get_replicate_completion_pricing(completion_response, total_time) # type: ignore if model is None: raise ValueError( f"Model is None and does not exist in passed completion_response. Passed completion_response={completion_response}, model={model}" ) - if ( - custom_llm_provider is not None - and custom_llm_provider == "vertex_ai" - ): + if custom_llm_provider is not None and custom_llm_provider == "vertex_ai": # Calculate the prompt characters + response characters if len(messages) > 0: prompt_string = litellm.utils.get_formatted_prompt( data={"messages": messages}, call_type="completion" ) - prompt_characters = litellm.utils._count_characters( - text=prompt_string - ) - if completion_response is not None and isinstance( - completion_response, ModelResponse - ): - completion_string = litellm.utils.get_response_string( - response_obj=completion_response - ) - completion_characters = litellm.utils._count_characters( - text=completion_string - ) + prompt_characters = litellm.utils._count_characters(text=prompt_string) + if completion_response is not None and isinstance(completion_response, ModelResponse): + completion_string = litellm.utils.get_response_string(response_obj=completion_response) + completion_characters = litellm.utils._count_characters(text=completion_string) # Get the original request model for router detection request_model_for_cost = None @@ -1688,12 +1564,8 @@ def completion_cost( if custom_llm_provider == "azure_ai": model_for_additional_costs = request_model_for_cost if completion_response is not None: - hidden_params = ( - getattr(completion_response, "_hidden_params", None) or {} - ) - hidden_model = hidden_params.get("model") or hidden_params.get( - "litellm_model_name" - ) + hidden_params = getattr(completion_response, "_hidden_params", None) or {} + hidden_model = hidden_params.get("model") or hidden_params.get("litellm_model_name") if hidden_model and ( "model_router" in (hidden_model or "").lower() or "model-router" in (hidden_model or "").lower() @@ -1712,17 +1584,13 @@ def completion_cost( else: additional_costs = None - _final_cost = ( - prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar - ) - cost_for_built_in_tools = ( - StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( - model=model, - response_object=completion_response, - usage=cost_per_token_usage_object, - standard_built_in_tools_params=standard_built_in_tools_params, - custom_llm_provider=custom_llm_provider, - ) + _final_cost = prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar + cost_for_built_in_tools = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=completion_response, + usage=cost_per_token_usage_object, + standard_built_in_tools_params=standard_built_in_tools_params, + custom_llm_provider=custom_llm_provider, ) _final_cost += cost_for_built_in_tools if additional_costs: @@ -1763,23 +1631,17 @@ def completion_cost( _cache_read_cost: Optional[float] = None _cache_creation_cost: Optional[float] = None if cost_per_token_usage_object is not None: - _cr = getattr( - cost_per_token_usage_object, "cache_read_input_tokens", None - ) or (cost_per_token_usage_object.model_extra or {}).get( - "cache_read_input_tokens" - ) + _cr = getattr(cost_per_token_usage_object, "cache_read_input_tokens", None) or ( + cost_per_token_usage_object.model_extra or {} + ).get("cache_read_input_tokens") _cc = getattr( cost_per_token_usage_object, "cache_creation_input_tokens", None, - ) or (cost_per_token_usage_object.model_extra or {}).get( - "cache_creation_input_tokens" - ) + ) or (cost_per_token_usage_object.model_extra or {}).get("cache_creation_input_tokens") if (_cr or _cc) and model: try: - _mi = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + _mi = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) _cr_rate = _mi.get("cache_read_input_token_cost") if _cr and _cr_rate is not None: _cache_read_cost = float(_cr) * float(_cr_rate) @@ -1814,11 +1676,7 @@ def completion_cost( ) if idx == len(potential_model_names) - 1: raise e - raise Exception( - "Unable to calculat cost for received potential model names - {}".format( - potential_model_names - ) - ) + raise Exception("Unable to calculat cost for received potential model names - {}".format(potential_model_names)) except Exception as e: raise e @@ -1832,10 +1690,7 @@ def get_response_cost_from_hidden_params( _hidden_params_dict = hidden_params additional_headers = _hidden_params_dict.get("additional_headers", {}) - if ( - additional_headers - and "llm_provider-x-litellm-response-cost" in additional_headers - ): + if additional_headers and "llm_provider-x-litellm-response-cost" in additional_headers: response_cost = additional_headers["llm_provider-x-litellm-response-cost"] if response_cost is None: return None @@ -1892,9 +1747,7 @@ def response_cost_calculator( ### SERVICE TIER ### service_tier: Optional[str] = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### - data_residency: Optional[ - str - ] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ) -> float: """ Returns @@ -1908,9 +1761,7 @@ def response_cost_calculator( if isinstance(response_object, BaseModel): if hasattr(response_object, "_hidden_params"): response_object._hidden_params["optional_params"] = optional_params - provider_response_cost = get_response_cost_from_hidden_params( - response_object._hidden_params - ) + provider_response_cost = get_response_cost_from_hidden_params(response_object._hidden_params) if provider_response_cost is not None: return provider_response_cost @@ -1957,17 +1808,13 @@ def ocr_cost( # validate it's an OCR response ######################################################### if response is None or not isinstance(response, OCRResponse): - raise ValueError( - f"response must be of type OCRResponse got type={type(response)}" - ) + raise ValueError(f"response must be of type OCRResponse got type={type(response)}") if response.usage_info is None: raise ValueError("OCR response usage_info is None") try: - model_info: Optional[ModelInfo] = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info: Optional[ModelInfo] = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: model_info = None @@ -2043,9 +1890,7 @@ def vector_store_search_cost( ) if config is None: - verbose_logger.debug( - f"Vector store search is not supported for {custom_llm_provider}" - ) + verbose_logger.debug(f"Vector store search is not supported for {custom_llm_provider}") return 0.0, 0.0 return config.calculate_vector_store_cost( @@ -2062,9 +1907,7 @@ def rerank_cost( Returns - float or None: cost of response OR none if error. """ - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, custom_llm_provider=custom_llm_provider - ) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) try: config = ProviderConfigManager.get_provider_rerank_config( @@ -2091,12 +1934,8 @@ def rerank_cost( raise e -def transcription_cost( - model: str, custom_llm_provider: Optional[str], duration: float -) -> Tuple[float, float]: - return openai_cost_per_second( - model=model, custom_llm_provider=custom_llm_provider, duration=duration - ) +def transcription_cost(model: str, custom_llm_provider: Optional[str], duration: float) -> Tuple[float, float]: + return openai_cost_per_second(model=model, custom_llm_provider=custom_llm_provider, duration=duration) def default_image_cost_calculator( @@ -2125,11 +1964,7 @@ def default_image_cost_calculator( """ # Standardize size format to use "-x-" size_str: str = size or "1024-x-1024" - size_str = ( - size_str.replace("x", "-x-") - if "x" in size_str and "-x-" not in size_str - else size_str - ) + size_str = size_str.replace("x", "-x-") if "x" in size_str and "-x-" not in size_str else size_str # Parse dimensions height, width = map(int, size_str.split("-x-")) @@ -2138,29 +1973,17 @@ def default_image_cost_calculator( base_model_name = f"{size_str}/{model}" model_name_without_custom_llm_provider: Optional[str] = None if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): - model_name_without_custom_llm_provider = model.replace( - f"{custom_llm_provider}/", "" - ) - base_model_name = ( - f"{custom_llm_provider}/{size_str}/{model_name_without_custom_llm_provider}" - ) - model_name_with_quality = ( - f"{quality}/{base_model_name}" if quality else base_model_name - ) + model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "") + base_model_name = f"{custom_llm_provider}/{size_str}/{model_name_without_custom_llm_provider}" + model_name_with_quality = f"{quality}/{base_model_name}" if quality else base_model_name # gpt-image-1 models use low, medium, high quality. If user did not specify quality, use medium fot gpt-image-1 model family - model_name_with_v2_quality = ( - f"{ImageGenerationRequestQuality.HIGH.value}/{base_model_name}" - ) + model_name_with_v2_quality = f"{ImageGenerationRequestQuality.HIGH.value}/{base_model_name}" - verbose_logger.debug( - f"Looking up cost for models: {model_name_with_quality}, {base_model_name}" - ) + verbose_logger.debug(f"Looking up cost for models: {model_name_with_quality}, {base_model_name}") model_without_provider = f"{size_str}/{model.split('/')[-1]}" - model_with_quality_without_provider = ( - f"{quality}/{model_without_provider}" if quality else model_without_provider - ) + model_with_quality_without_provider = f"{quality}/{model_without_provider}" if quality else model_without_provider # Try model with quality first, fall back to base model name cost_info: Optional[dict] = None @@ -2178,26 +2001,16 @@ def default_image_cost_calculator( cost_info = litellm.model_cost[_model] break if cost_info is None: - raise Exception( - f"Model not found in cost map. Tried checking {models_to_check}" - ) + raise Exception(f"Model not found in cost map. Tried checking {models_to_check}") # Priority 1: Use per-image pricing if available (for gpt-image-1 and similar models) - if ( - "input_cost_per_image" in cost_info - and cost_info["input_cost_per_image"] is not None - ): + if "input_cost_per_image" in cost_info and cost_info["input_cost_per_image"] is not None: return cost_info["input_cost_per_image"] * n # Priority 2: Fall back to per-pixel pricing for backward compatibility - elif ( - "input_cost_per_pixel" in cost_info - and cost_info["input_cost_per_pixel"] is not None - ): + elif "input_cost_per_pixel" in cost_info and cost_info["input_cost_per_pixel"] is not None: return cost_info["input_cost_per_pixel"] * height * width * n else: - raise Exception( - f"No pricing information found for model {model}. Tried checking {models_to_check}" - ) + raise Exception(f"No pricing information found for model {model}. Tried checking {models_to_check}") def default_video_cost_calculator( @@ -2234,12 +2047,8 @@ def default_video_cost_calculator( base_model_name = model model_name_without_custom_llm_provider: Optional[str] = None if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): - model_name_without_custom_llm_provider = model.replace( - f"{custom_llm_provider}/", "" - ) - base_model_name = ( - f"{custom_llm_provider}/{model_name_without_custom_llm_provider}" - ) + model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "") + base_model_name = f"{custom_llm_provider}/{model_name_without_custom_llm_provider}" verbose_logger.debug(f"Looking up cost for video model: {base_model_name}") @@ -2299,9 +2108,7 @@ def batch_cost_calculator( deployment-specific pricing is used. """ - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, custom_llm_provider=custom_llm_provider - ) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) verbose_logger.debug( "Calculating batch cost per token. model=%s, custom_llm_provider=%s", @@ -2311,9 +2118,7 @@ def batch_cost_calculator( if model_info is None: try: - model_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: model_info = None elif not any( @@ -2329,9 +2134,7 @@ def batch_cost_calculator( # but carries no pricing fields. Fall back to the global pricing table so # that standard model pricing is used instead of silently returning $0. try: - global_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + global_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) if global_info: model_info = global_info except Exception: @@ -2358,13 +2161,8 @@ def batch_cost_calculator( # Add cache read cost if applicable details = _parse_prompt_tokens_details(usage) cache_read_tokens = details["cache_hit_tokens"] - cache_read_cost_key = _get_service_tier_cost_key( - "cache_read_input_token_cost", None - ) - total_prompt_cost += ( - calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) - / 2 - ) + cache_read_cost_key = _get_service_tier_cost_key("cache_read_input_token_cost", None) + total_prompt_cost += calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) / 2 if output_cost_per_token_batches: total_completion_cost = usage.completion_tokens * output_cost_per_token_batches elif output_cost_per_token: @@ -2409,10 +2207,7 @@ class BaseTokenUsageProcessor: setattr(combined, attr, current_val + new_val) # Handle nested prompt_tokens_details if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - if ( - not hasattr(combined, "prompt_tokens_details") - or not combined.prompt_tokens_details - ): + if not hasattr(combined, "prompt_tokens_details") or not combined.prompt_tokens_details: combined.prompt_tokens_details = PromptTokensDetailsWrapper() # Check what keys exist in the model's prompt_tokens_details @@ -2423,9 +2218,7 @@ class BaseTokenUsageProcessor: and not attr.startswith("_") and not callable(getattr(usage.prompt_tokens_details, attr)) ): - current_val = ( - getattr(combined.prompt_tokens_details, attr, 0) or 0 - ) + current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 if new_val is not None and isinstance(new_val, (int, float)): setattr( @@ -2435,27 +2228,15 @@ class BaseTokenUsageProcessor: ) # Handle nested completion_tokens_details - if ( - hasattr(usage, "completion_tokens_details") - and usage.completion_tokens_details - ): - if ( - not hasattr(combined, "completion_tokens_details") - or not combined.completion_tokens_details - ): - combined.completion_tokens_details = ( - CompletionTokensDetailsWrapper() - ) + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: + if not hasattr(combined, "completion_tokens_details") or not combined.completion_tokens_details: + combined.completion_tokens_details = CompletionTokensDetailsWrapper() # Check what keys exist in the model's completion_tokens_details # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings for attr in type(usage.completion_tokens_details).model_fields: - if not attr.startswith("_") and not callable( - getattr(usage.completion_tokens_details, attr) - ): - current_val = ( - getattr(combined.completion_tokens_details, attr, 0) or 0 - ) + if not attr.startswith("_") and not callable(getattr(usage.completion_tokens_details, attr)): + current_val = getattr(combined.completion_tokens_details, attr, 0) or 0 new_val = getattr(usage.completion_tokens_details, attr, 0) or 0 if isinstance(new_val, (int, float)): setattr( @@ -2481,10 +2262,8 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): ) usage_objects: List[Usage] = [] for result in response_done_events: - usage_object = ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - result["response"].get("usage", {}) - ) + usage_object = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + result["response"].get("usage", {}) ) usage_objects.append(usage_object) return usage_objects @@ -2496,14 +2275,8 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): """ Collect and combine usage from realtime stream results """ - collected_usage_objects = ( - RealtimeAPITokenUsageProcessor.collect_usage_from_realtime_stream_results( - results - ) - ) - combined_usage_object = RealtimeAPITokenUsageProcessor.combine_usage_objects( - collected_usage_objects - ) + collected_usage_objects = RealtimeAPITokenUsageProcessor.collect_usage_from_realtime_stream_results(results) + combined_usage_object = RealtimeAPITokenUsageProcessor.combine_usage_objects(collected_usage_objects) return combined_usage_object @staticmethod @@ -2516,9 +2289,7 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): ) -_TRANSCRIPTION_COMPLETED_EVENT_TYPE = ( - "conversation.item.input_audio_transcription.completed" -) +_TRANSCRIPTION_COMPLETED_EVENT_TYPE = "conversation.item.input_audio_transcription.completed" def handle_realtime_stream_cost_calculation( @@ -2540,9 +2311,7 @@ def handle_realtime_stream_cost_calculation( potential_model_names = [] for result in results: if result["type"] == "session.created": - received_model = cast(OpenAIRealtimeStreamSessionEvents, result)[ - "session" - ].get("model", None) + received_model = cast(OpenAIRealtimeStreamSessionEvents, result)["session"].get("model", None) potential_model_names.append(received_model) potential_model_names.append(litellm_model_name) @@ -2591,20 +2360,14 @@ def handle_realtime_transcription_cost_calculation( - {"type": "tokens", "input_tokens": ...} → priced via input/audio token cost """ completed_events = [ - cast(dict, result) - for result in results - if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE + cast(dict, result) for result in results if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE ] if not completed_events: return 0.0 - model_name = ( - _get_transcription_model_name_from_results(results) or litellm_model_name - ) + model_name = _get_transcription_model_name_from_results(results) or litellm_model_name try: - model_info = litellm.get_model_info( - model=model_name, custom_llm_provider=custom_llm_provider - ) + model_info = litellm.get_model_info(model=model_name, custom_llm_provider=custom_llm_provider) except Exception: model_info = None @@ -2627,9 +2390,9 @@ def _get_transcription_model_name_from_results( "session.updated", ): session = cast(dict, result).get("session", {}) or {} - transcription = ( - (session.get("audio", {}) or {}).get("input", {}) or {} - ).get("transcription", {}) or session.get("input_audio_transcription", {}) + transcription = ((session.get("audio", {}) or {}).get("input", {}) or {}).get( + "transcription", {} + ) or session.get("input_audio_transcription", {}) model = (transcription or {}).get("model") or session.get("model") if model: return model @@ -2650,15 +2413,9 @@ def _transcription_usage_cost(usage: dict, model_info: Optional[ModelInfo]) -> f text_tokens = input_token_details.get("text_tokens") or 0 output_tokens = usage.get("output_tokens") or 0 audio_cost = float(audio_tokens) * float( - model_info.get("input_cost_per_audio_token") - or model_info.get("input_cost_per_token") - or 0.0 - ) - text_cost = float(text_tokens) * float( - model_info.get("input_cost_per_token") or 0.0 - ) - output_cost = float(output_tokens) * float( - model_info.get("output_cost_per_token") or 0.0 + model_info.get("input_cost_per_audio_token") or model_info.get("input_cost_per_token") or 0.0 ) + text_cost = float(text_tokens) * float(model_info.get("input_cost_per_token") or 0.0) + output_cost = float(output_tokens) * float(model_info.get("output_cost_per_token") or 0.0) return audio_cost + text_cost + output_cost return 0.0 diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 13af0a30fe0..f2b443eb7bf 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -29,9 +29,7 @@ class SpeechToCompletionBridgeHandler: super().__init__() self.transformation_handler = SpeechToCompletionBridgeTransformationHandler() - def validate_input_kwargs( - self, kwargs: dict - ) -> SpeechToCompletionBridgeHandlerInputKwargs: + def validate_input_kwargs(self, kwargs: dict) -> SpeechToCompletionBridgeHandlerInputKwargs: from litellm import LiteLLMLoggingObj model = kwargs.get("model") diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 5dce467d443..94de4878b65 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -29,9 +29,7 @@ class SpeechToCompletionBridgeTransformationHandler: if isinstance(voice, str): passed_optional_params["audio"] = {"voice": voice} if "response_format" in optional_params: - passed_optional_params["audio"]["format"] = optional_params[ - "response_format" - ] + passed_optional_params["audio"]["format"] = optional_params["response_format"] return_kwargs = { "model": model, @@ -53,9 +51,7 @@ class SpeechToCompletionBridgeTransformationHandler: return_kwargs = {k: v for k, v in return_kwargs.items() if v is not None} return return_kwargs - def _convert_pcm16_to_wav( - self, pcm_data: bytes, sample_rate: int = 24000, channels: int = 1 - ) -> bytes: + def _convert_pcm16_to_wav(self, pcm_data: bytes, sample_rate: int = 24000, channels: int = 1) -> bytes: """ Convert raw PCM16 data to WAV format. @@ -97,13 +93,9 @@ class SpeechToCompletionBridgeTransformationHandler: def _is_gemini_tts_model(self, model: str) -> bool: """Check if the model is a Gemini TTS model that returns PCM16 data.""" - return "gemini" in model.lower() and ( - "tts" in model.lower() or "preview-tts" in model.lower() - ) + return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower()) - def transform_response( - self, model_response: "ModelResponse" - ) -> "HttpxBinaryResponseContent": + def transform_response(self, model_response: "ModelResponse") -> "HttpxBinaryResponseContent": import base64 import httpx diff --git a/litellm/evals/main.py b/litellm/evals/main.py index df6d3accb82..d4e9d638583 100644 --- a/litellm/evals/main.py +++ b/litellm/evals/main.py @@ -152,10 +152,8 @@ def create_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -175,9 +173,7 @@ def create_eval( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request request_body = evals_api_provider_config.transform_create_eval_request( @@ -188,9 +184,7 @@ def create_eval( # Get API base and URL api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE - url = evals_api_provider_config.get_complete_url( - api_base=api_base, endpoint="evals" - ) + url = evals_api_provider_config.get_complete_url(api_base=api_base, endpoint="evals") # Pre-call logging litellm_logging_obj.update_from_kwargs( @@ -343,10 +337,8 @@ def list_evals( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -371,9 +363,7 @@ def list_evals( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request url, query_params = evals_api_provider_config.transform_list_evals_request( @@ -513,10 +503,8 @@ def get_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -524,9 +512,7 @@ def get_eval( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE @@ -682,10 +668,8 @@ def update_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -732,9 +716,7 @@ def update_eval( "user_agent", } # Only include user-provided metadata keys - filtered_metadata = { - k: v for k, v in metadata.items() if k not in internal_keys - } + filtered_metadata = {k: v for k, v in metadata.items() if k not in internal_keys} if filtered_metadata: # Only add if there's user metadata update_request["metadata"] = filtered_metadata @@ -744,9 +726,7 @@ def update_eval( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE @@ -893,10 +873,8 @@ def delete_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -904,9 +882,7 @@ def delete_eval( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE @@ -1047,10 +1023,8 @@ def cancel_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1058,9 +1032,7 @@ def cancel_eval( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE @@ -1230,10 +1202,8 @@ def create_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1254,9 +1224,7 @@ def create_run( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE @@ -1418,10 +1386,8 @@ def list_runs( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1444,9 +1410,7 @@ def list_runs( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request url, query_params = evals_api_provider_config.transform_list_runs_request( @@ -1592,10 +1556,8 @@ def get_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1603,9 +1565,7 @@ def get_run( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE @@ -1752,10 +1712,8 @@ def cancel_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1763,9 +1721,7 @@ def cancel_run( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE @@ -1921,10 +1877,8 @@ def delete_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1932,9 +1886,7 @@ def delete_run( # Validate environment and get headers headers = extra_headers or {} - headers = evals_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 1cbef6b0b49..d97ba347b07 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -146,9 +146,7 @@ class AuthenticationError(openai.AuthenticationError): # type: ignore self.num_retries = num_retries self.response = response or httpx.Response( status_code=self.status_code, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), # mock request object + request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ) super().__init__( self.message, response=self.response, body=None @@ -192,9 +190,7 @@ class NotFoundError(openai.NotFoundError): # type: ignore self.num_retries = num_retries self.response = response or httpx.Response( status_code=self.status_code, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), # mock request object + request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ) super().__init__( self.message, response=self.response, body=None @@ -347,9 +343,7 @@ class Timeout(openai.APITimeoutError): # type: ignore method="POST", url="https://api.openai.com/v1", ) - super().__init__( - request=request - ) # Call the base class constructor with the parameters it needs + super().__init__(request=request) # Call the base class constructor with the parameters it needs self.status_code = exception_status_code or 408 self.message = "litellm.Timeout: {}".format(message) self.model = model @@ -438,9 +432,7 @@ class RateLimitError(openai.RateLimitError): # type: ignore litellm_debug_info: Optional[str] = None, max_retries: Optional[int] = None, num_retries: Optional[int] = None, - category: Union[str, RateLimitErrorCategory] = ( - RateLimitErrorCategory.VENDOR_RATE_LIMIT - ), + category: Union[str, RateLimitErrorCategory] = (RateLimitErrorCategory.VENDOR_RATE_LIMIT), rate_limit_type: Optional[Union[str, RateLimitType]] = None, headers: Optional[Dict[str, str]] = None, detail: Any = None, @@ -452,16 +444,12 @@ class RateLimitError(openai.RateLimitError): # type: ignore self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - self.category = ( - category.value if isinstance(category, RateLimitErrorCategory) else category - ) + self.category = category.value if isinstance(category, RateLimitErrorCategory) else category # Which dimension was exceeded — request count, token count, parallel # requests, budget, max iterations. None when the source didn't # classify the failure (e.g. legacy vendor 429 with no header hints). self.rate_limit_type: Optional[str] = ( - rate_limit_type.value - if isinstance(rate_limit_type, RateLimitType) - else rate_limit_type + rate_limit_type.value if isinstance(rate_limit_type, RateLimitType) else rate_limit_type ) # Headers explicitly attached to the error (e.g. retry-after, # rate_limit_type, reset_at). Preserved across the proxy boundary so @@ -476,12 +464,8 @@ class RateLimitError(openai.RateLimitError): # type: ignore # headers stay reachable on `e.response.headers` for callers that # explicitly want them; only the proxy-supplied `headers=` kwarg # makes it onto `self.headers`. - _response_headers = ( - getattr(response, "headers", None) if response is not None else None - ) - self.headers: Optional[Dict[str, str]] = ( - {k: str(v) for k, v in headers.items()} if headers else None - ) + _response_headers = getattr(response, "headers", None) if response is not None else None + self.headers: Optional[Dict[str, str]] = {k: str(v) for k, v in headers.items()} if headers else None # Mirrors FastAPI HTTPException.detail so the same instance can be # serialized through both the ProxyException and HTTPException paths. self.detail = detail if detail is not None else self.message @@ -664,9 +648,7 @@ class ServiceUnavailableError(openai.APIStatusError): # type: ignore self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers = ( - getattr(response, "headers", None) if response is not None else None - ) + _response_headers = getattr(response, "headers", None) if response is not None else None self.response = httpx.Response( status_code=self.status_code, headers=_response_headers, @@ -714,9 +696,7 @@ class BadGatewayError(openai.APIStatusError): # type: ignore self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers = ( - getattr(response, "headers", None) if response is not None else None - ) + _response_headers = getattr(response, "headers", None) if response is not None else None self.response = httpx.Response( status_code=self.status_code, headers=_response_headers, @@ -764,9 +744,7 @@ class InternalServerError(openai.InternalServerError): # type: ignore self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers = ( - getattr(response, "headers", None) if response is not None else None - ) + _response_headers = getattr(response, "headers", None) if response is not None else None self.response = httpx.Response( status_code=self.status_code, headers=_response_headers, @@ -915,9 +893,7 @@ class APIResponseValidationError(openai.APIResponseValidationError): # type: ig class JSONSchemaValidationError(APIResponseValidationError): - def __init__( - self, model: str, llm_provider: str, raw_response: str, schema: str - ) -> None: + def __init__(self, model: str, llm_provider: str, raw_response: str, schema: str) -> None: self.raw_response = raw_response self.schema = schema self.model = model @@ -953,9 +929,7 @@ class UnsupportedParamsError(BadRequestError): self.litellm_debug_info = litellm_debug_info response = response or httpx.Response( status_code=self.status_code, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), # mock request object + request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ) self.max_retries = max_retries self.num_retries = num_retries @@ -1005,10 +979,7 @@ class BudgetExceededError(Exception): # to match the normalization RateLimitError.__init__ performs. self.category: str = RateLimitErrorCategory.LITELLM_RATE_LIMIT.value self.rate_limit_type: str = RateLimitType.BUDGET.value - message = ( - message - or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" - ) + message = message or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" self.message = message super().__init__(message) @@ -1022,9 +993,7 @@ class InvalidRequestError(openai.BadRequestError): # type: ignore self.llm_provider = llm_provider self.response = httpx.Response( status_code=400, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), # mock request object + request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ) super().__init__( message=self.message, response=self.response, body=None @@ -1061,9 +1030,7 @@ class LiteLLMUnknownProvider(BadRequestError): self.message = LiteLLMCommonStrings.llm_provider_not_provided.value.format( model=model, custom_llm_provider=custom_llm_provider ) - super().__init__( - self.message, model=model, llm_provider=custom_llm_provider, response=None - ) + super().__init__(self.message, model=model, llm_provider=custom_llm_provider, response=None) def __str__(self): return self.message @@ -1248,8 +1215,5 @@ class SensitiveDataRouteException(Exception): self.guardrail_name = guardrail_name self.detection_info = detection_info or {} self.sticky_session_routing = sticky_session_routing - self.message = ( - message - or f"Sensitive data detected by {guardrail_name}. Routing to model: {route_to_model}" - ) + self.message = message or f"Sensitive data detected by {guardrail_name}. Routing to model: {route_to_model}" super().__init__(self.message) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ba07d116e26..831e588e5ba 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -26,9 +26,7 @@ streamable_http_client: Optional[Any] = None try: import mcp.client.streamable_http as streamable_http_module # type: ignore - streamable_http_client = getattr( - streamable_http_module, "streamable_http_client", None - ) + streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None) except ImportError: pass from mcp.types import CallToolRequestParams as MCPCallToolRequestParams @@ -62,9 +60,7 @@ def to_basic_auth(auth_value: str) -> str: def _strip_header_whitespace(headers: Dict[str, str]) -> Dict[str, str]: return { - (key.strip() if isinstance(key, str) else key): ( - value.strip() if isinstance(value, str) else value - ) + (key.strip() if isinstance(key, str) else key): (value.strip() if isinstance(value, str) else value) for key, value in headers.items() } @@ -107,10 +103,7 @@ class MCPSigV4Auth(httpx.Auth): try: from botocore.credentials import Credentials except ImportError: - raise ImportError( - "Missing botocore to use AWS SigV4 authentication. " - "Run 'pip install boto3'." - ) + raise ImportError("Missing botocore to use AWS SigV4 authentication. Run 'pip install boto3'.") self.service_name = aws_service_name or "bedrock-agentcore" self.region_name = aws_region_name or "us-east-1" # Note: os.environ/ prefixed values are already resolved by @@ -157,9 +150,7 @@ class MCPSigV4Auth(httpx.Auth): import boto3 from botocore.credentials import Credentials - session_name = ( - aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" - ) + session_name = aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" sts_kwargs: dict = {"region_name": aws_region_name} if aws_access_key_id and aws_secret_access_key: sts_kwargs["aws_access_key_id"] = aws_access_key_id @@ -178,9 +169,7 @@ class MCPSigV4Auth(httpx.Auth): token=sts_creds["SessionToken"], ) - def auth_flow( - self, request: httpx.Request - ) -> Generator[httpx.Request, httpx.Response, None]: + def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest @@ -282,10 +271,7 @@ class MCPClient: ) # HTTP transport (default) if streamable_http_client is None: - raise ImportError( - "streamable_http_client is not available. " - "Please install mcp with HTTP support." - ) + raise ImportError("streamable_http_client is not available. Please install mcp with HTTP support.") headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() verbose_logger.debug("litellm headers for streamable_http_client: %s", headers) @@ -299,9 +285,7 @@ class MCPClient: ) return transport_ctx, http_client - def _get_safe_stdio_env( - self, provided_env: Optional[Dict[str, str]] - ) -> Optional[Dict[str, str]]: + def _get_safe_stdio_env(self, provided_env: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]: """ Return a safe environment for the stdio subprocess. @@ -393,18 +377,12 @@ class MCPClient: try: await transport_ctx.__aexit__(None, None, None) except BaseException as exit_error: - verbose_logger.debug( - f"Error during transport context exit: {exit_error}" - ) + verbose_logger.debug(f"Error during transport context exit: {exit_error}") root_cause = _first_non_cancelled_cause(exit_error) - if root_cause is not None and isinstance( - in_flight_error, asyncio.CancelledError - ): + if root_cause is not None and isinstance(in_flight_error, asyncio.CancelledError): raise root_cause from in_flight_error - async def run_with_session( - self, operation: Callable[[ClientSession], Awaitable[TSessionResult]] - ) -> TSessionResult: + async def run_with_session(self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]) -> TSessionResult: """Open a session, run the provided coroutine, and clean up.""" http_client: Optional[httpx.AsyncClient] = None try: @@ -412,9 +390,7 @@ class MCPClient: transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) except Exception: - verbose_logger.warning( - "MCP client run_with_session failed for %s", self.server_url or "stdio" - ) + verbose_logger.warning("MCP client run_with_session failed for %s", self.server_url or "stdio") raise finally: if http_client is not None: @@ -483,17 +459,11 @@ class MCPClient: """Create an httpx.AsyncClient with LiteLLM's SSL configuration.""" # Get unified SSL configuration using the same logic as http_handler.py ssl_config = get_ssl_configuration(self.ssl_verify) - verbose_logger.debug( - f"MCP client using SSL configuration: {type(ssl_config).__name__}" - ) + verbose_logger.debug(f"MCP client using SSL configuration: {type(ssl_config).__name__}") # The MCP SDK's sse_client and streamable_http_client call this factory without # passing auth=, so the fallback is used: a v2-resolved auth if present, else the # SigV4 aws_auth. Both are None for the common case — no behavior change. - fallback_auth = ( - self._resolved_auth - if self._resolved_auth is not None - else self._aws_auth - ) + fallback_auth = self._resolved_auth if self._resolved_auth is not None else self._aws_auth effective_auth = auth if auth is not None else fallback_auth return httpx.AsyncClient( headers=headers, @@ -515,9 +485,7 @@ class MCPClient: MCP client (triggering the upstream OAuth flow) rather than masking them as "connected, no tools". """ - verbose_logger.debug( - f"MCP client listing tools from {self.server_url or 'stdio'}" - ) + verbose_logger.debug(f"MCP client listing tools from {self.server_url or 'stdio'}") async def _list_tools_operation(session: ClientSession): return await session.list_tools() @@ -526,9 +494,7 @@ class MCPClient: result = await self.run_with_session(_list_tools_operation) tool_count = len(result.tools) tool_names = [tool.name for tool in result.tools] - verbose_logger.info( - f"MCP client listed {tool_count} tools from {self.server_url or 'stdio'}: {tool_names}" - ) + verbose_logger.info(f"MCP client listed {tool_count} tools from {self.server_url or 'stdio'}: {tool_names}") return result.tools except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") @@ -562,13 +528,9 @@ class MCPClient: """ Call an MCP Tool. """ - verbose_logger.info( - f"MCP client calling tool '{call_tool_request_params.name}'" - ) + verbose_logger.info(f"MCP client calling tool '{call_tool_request_params.name}'") - async def on_progress( - progress: float, total: float | None, message: str | None - ): + async def on_progress(progress: float, total: float | None, message: str | None): percentage = (progress / total * 100) if total else 0 verbose_logger.info( f"MCP Tool '{call_tool_request_params.name}' progress: " @@ -591,14 +553,10 @@ class MCPClient: try: tool_result = await self.run_with_session(_call_tool_operation) - verbose_logger.info( - f"MCP client tool call '{call_tool_request_params.name}' completed successfully" - ) + verbose_logger.info(f"MCP client tool call '{call_tool_request_params.name}' completed successfully") return tool_result except asyncio.CancelledError: - verbose_logger.warning( - f"MCP client tool call timed out after {self.timeout}s for {self.server_url}" - ) + verbose_logger.warning(f"MCP client tool call timed out after {self.timeout}s for {self.server_url}") raise except Exception as e: import traceback @@ -623,17 +581,13 @@ class MCPClient: ) # Return a default error result instead of raising return MCPCallToolResult( - content=[ - TextContent(type="text", text=f"{error_type}: {str(e)}") - ], # Empty content for error case + content=[TextContent(type="text", text=f"{error_type}: {str(e)}")], # Empty content for error case isError=True, ) async def list_prompts(self) -> List[Prompt]: """List available prompts from the server.""" - verbose_logger.debug( - f"MCP client listing tools from {self.server_url or 'stdio'}" - ) + verbose_logger.debug(f"MCP client listing tools from {self.server_url or 'stdio'}") async def _list_prompts_operation(session: ClientSession): return await session.list_prompts() @@ -667,13 +621,9 @@ class MCPClient: # Return empty list instead of raising to allow graceful degradation return [] - async def get_prompt( - self, get_prompt_request_params: GetPromptRequestParams - ) -> GetPromptResult: + async def get_prompt(self, get_prompt_request_params: GetPromptRequestParams) -> GetPromptResult: """Fetch a prompt definition from the MCP server.""" - verbose_logger.info( - f"MCP client fetching prompt '{get_prompt_request_params.name}'" - ) + verbose_logger.info(f"MCP client fetching prompt '{get_prompt_request_params.name}'") async def _get_prompt_operation(session: ClientSession): verbose_logger.debug("MCP client sending get_prompt request to session") @@ -684,9 +634,7 @@ class MCPClient: try: get_prompt_result = await self.run_with_session(_get_prompt_operation) - verbose_logger.info( - f"MCP client get_prompt '{get_prompt_request_params.name}' completed successfully" - ) + verbose_logger.info(f"MCP client get_prompt '{get_prompt_request_params.name}' completed successfully") return get_prompt_result except asyncio.CancelledError: verbose_logger.warning("MCP client get_prompt was cancelled") @@ -716,9 +664,7 @@ class MCPClient: async def list_resources(self) -> list[Resource]: """List available resources from the server.""" - verbose_logger.debug( - f"MCP client listing resources from {self.server_url or 'stdio'}" - ) + verbose_logger.debug(f"MCP client listing resources from {self.server_url or 'stdio'}") async def _list_resources_operation(session: ClientSession): return await session.list_resources() @@ -754,9 +700,7 @@ class MCPClient: async def list_resource_templates(self) -> list[ResourceTemplate]: """List available resource templates from the server.""" - verbose_logger.debug( - f"MCP client listing resource templates from {self.server_url or 'stdio'}" - ) + verbose_logger.debug(f"MCP client listing resource templates from {self.server_url or 'stdio'}") async def _list_resource_templates_operation(session: ClientSession): return await session.list_resource_templates() @@ -764,9 +708,7 @@ class MCPClient: try: result = await self.run_with_session(_list_resource_templates_operation) resource_template_count = len(result.resourceTemplates) - resource_template_names = [ - resourceTemplate.name for resourceTemplate in result.resourceTemplates - ] + resource_template_names = [resourceTemplate.name for resourceTemplate in result.resourceTemplates] verbose_logger.info( f"MCP client listed {resource_template_count} resource templates from {self.server_url or 'stdio'}: {resource_template_names}" ) @@ -802,9 +744,7 @@ class MCPClient: try: read_resource_result = await self.run_with_session(_read_resource_operation) - verbose_logger.info( - f"MCP client read_resource '{url}' completed successfully" - ) + verbose_logger.info(f"MCP client read_resource '{url}' completed successfully") return read_resource_result except asyncio.CancelledError: verbose_logger.warning("MCP client read_resource was cancelled") diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index bd42f7e7111..c65b266bd02 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -90,9 +90,7 @@ async def load_mcp_tools( """ tools = await session.list_tools() if format == "openai": - return [ - transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools - ] + return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools] return tools.tools @@ -148,10 +146,8 @@ async def call_openai_tool( Returns: The result of the MCP tool call. """ - mcp_tool_call_request_params = ( - transform_openai_tool_call_request_to_mcp_tool_call_request( - openai_tool=openai_tool, - ) + mcp_tool_call_request_params = transform_openai_tool_call_request_to_mcp_tool_call_request( + openai_tool=openai_tool, ) return await call_mcp_tool( session=session, diff --git a/litellm/files/main.py b/litellm/files/main.py index 582d9d5cdd0..3b359b55fe3 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -26,9 +26,7 @@ FileCreateProvider = Literal[ "manus", "anthropic", ] -FileRetrieveProvider = Literal[ - "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus", "anthropic" -] +FileRetrieveProvider = Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus", "anthropic"] FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"] FileListProvider = Literal["openai", "azure", "manus", "anthropic"] import litellm @@ -91,9 +89,7 @@ def _add_trusted_model_credentials_to_litellm_params( ) -> None: trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials") if isinstance(trusted_model_credentials, type(MappingProxyType({}))): - litellm_params_dict["_litellm_internal_model_credentials"] = ( - trusted_model_credentials - ) + litellm_params_dict["_litellm_internal_model_credentials"] = trusted_model_credentials @client @@ -162,9 +158,7 @@ def create_file( _is_async = kwargs.pop("acreate_file", False) is True optional_params = GenericLiteLLMParams(**kwargs) litellm_params_dict = dict(**kwargs) - logging_obj = cast( - Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj") - ) + logging_obj = cast(Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")) if logging_obj is None: raise ValueError("logging_obj is required") client = kwargs.get("client") @@ -215,12 +209,7 @@ def create_file( api_key=optional_params.api_key, logging_obj=logging_obj, _is_async=_is_async, - client=( - client - if client is not None - and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) - else None - ), + client=(client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None), timeout=timeout, ) elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: @@ -264,9 +253,7 @@ def create_file( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_file", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_file", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -405,9 +392,7 @@ def file_retrieve( stream=False, call_type="afile_retrieve" if _is_async else "file_retrieve", start_time=time.time(), - litellm_call_id=kwargs.get( - "litellm_call_id", str(uuid_module.uuid4()) - ), + litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())), function_id=str(kwargs.get("id") or ""), ) @@ -420,10 +405,7 @@ def file_retrieve( logging_obj=logging_obj, _is_async=_is_async, client=( - client - if client is not None - and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) - else None + client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None ), timeout=timeout, ) @@ -510,9 +492,7 @@ def file_delete( try: try: if model is not None: - _, custom_llm_provider, _, _ = get_llm_provider( - model, custom_llm_provider - ) + _, custom_llm_provider, _, _ = get_llm_provider(model, custom_llm_provider) except Exception: pass optional_params = GenericLiteLLMParams(**kwargs) @@ -592,9 +572,7 @@ def file_delete( stream=False, call_type="afile_delete" if _is_async else "file_delete", start_time=time.time(), - litellm_call_id=kwargs.get( - "litellm_call_id", str(uuid_module.uuid4()) - ), + litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())), function_id=str(kwargs.get("id") or ""), ) @@ -606,10 +584,7 @@ def file_delete( logging_obj=logging_obj, _is_async=_is_async, client=( - client - if client is not None - and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) - else None + client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None ), timeout=timeout, ) @@ -731,9 +706,7 @@ def file_list( stream=False, call_type="afile_list" if _is_async else "file_list", start_time=time.time(), - litellm_call_id=kwargs.get( - "litellm_call_id", str(uuid_module.uuid4()) - ), + litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())), function_id=str(kwargs.get("id", "")), ) @@ -745,12 +718,7 @@ def file_list( headers=extra_headers or {}, logging_obj=logging_obj, _is_async=_is_async, - client=( - client - if client is not None - and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) - else None - ), + client=(client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None), timeout=timeout, ) return response @@ -794,9 +762,7 @@ def file_list( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="file_list", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="file_list", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -886,9 +852,7 @@ def file_content( try: if model is not None: - _, custom_llm_provider, _, _ = get_llm_provider( - model, custom_llm_provider - ) + _, custom_llm_provider, _, _ = get_llm_provider(model, custom_llm_provider) except Exception: pass @@ -922,9 +886,7 @@ def file_content( chunk_size=chunk_size, optional_params=optional_params, timeout=timeout, - logging_obj=cast( - Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj") - ), + logging_obj=cast(Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")), _is_async=_is_async, client=client, ) @@ -946,9 +908,7 @@ def file_content( stream=False, call_type="afile_content" if _is_async else "file_content", start_time=time.time(), - litellm_call_id=kwargs.get( - "litellm_call_id", str(uuid_module.uuid4()) - ), + litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())), function_id=str(kwargs.get("id") or ""), ) @@ -959,12 +919,7 @@ def file_content( headers=extra_headers or {}, logging_obj=logging_obj, _is_async=_is_async, - client=( - client - if client is not None - and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) - else None - ), + client=(client if client is not None and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) else None), timeout=timeout, ) return response @@ -1004,18 +959,12 @@ def file_content( elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") + optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" + optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) + vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_files_instance.file_content( _is_async=_is_async, @@ -1047,9 +996,7 @@ def file_content( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -1095,9 +1042,9 @@ def file_content_streaming( headers=response.headers, ) - response: Union[ - FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult] - ] = FileContentStreamingResult(stream_iterator=iter(()), headers={}) + response: Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]] = ( + FileContentStreamingResult(stream_iterator=iter(()), headers={}) + ) if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: openai_creds = get_openai_credentials( api_base=optional_params.api_base, @@ -1130,9 +1077,7 @@ def file_content_streaming( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index b7095ce7e2b..6d84f73dcfe 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -94,9 +94,7 @@ class FileContentStreamingResponse: self._close_completed = True self._logging_completed = True stream_to_close = self.stream_iterator - self.stream_iterator = cast( - Union[Iterator[bytes], AsyncIterator[bytes]], iter(()) - ) + self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(())) # Shield cleanup from request cancellation so upstream HTTP connections # are released promptly on client disconnects. @@ -115,9 +113,7 @@ class FileContentStreamingResponse: self._close_completed = True self._logging_completed = True stream_to_close = self.stream_iterator - self.stream_iterator = cast( - Union[Iterator[bytes], AsyncIterator[bytes]], iter(()) - ) + self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(())) if hasattr(stream_to_close, "close"): cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined] @@ -134,9 +130,7 @@ class FileContentStreamingResponse: def _sync_hidden_params(self) -> None: litellm_params: dict[str, Any] = {} if self.logging_obj is not None: - litellm_params = ( - self.logging_obj.model_call_details.get("litellm_params", {}) or {} - ) + litellm_params = self.logging_obj.model_call_details.get("litellm_params", {}) or {} if "api_base" not in self._hidden_params and litellm_params.get("api_base"): self._hidden_params["api_base"] = litellm_params["api_base"] @@ -232,12 +226,8 @@ class FileContentStreamingResponse: self._logging_completed = True end_time = datetime.datetime.now() traceback_str = traceback.format_exc() - self.logging_obj.failure_handler( - error, traceback_str, self._start_time, end_time - ) - await self.logging_obj.async_failure_handler( - error, traceback_str, self._start_time, end_time - ) + self.logging_obj.failure_handler(error, traceback_str, self._start_time, end_time) + await self.logging_obj.async_failure_handler(error, traceback_str, self._start_time, end_time) def _log_failure_sync(self, error: Exception) -> None: if self._logging_completed or self.logging_obj is None: @@ -245,6 +235,4 @@ class FileContentStreamingResponse: self._logging_completed = True end_time = datetime.datetime.now() - self.logging_obj.failure_handler( - error, traceback.format_exc(), self._start_time, end_time - ) + self.logging_obj.failure_handler(error, traceback.format_exc(), self._start_time, end_time) diff --git a/litellm/files/types.py b/litellm/files/types.py index ba42a39f666..6bf7b1a1cc2 100644 --- a/litellm/files/types.py +++ b/litellm/files/types.py @@ -1,8 +1,6 @@ from typing import AsyncIterator, Dict, Iterator, Literal, NamedTuple, Union -FileContentProvider = Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus" -] +FileContentProvider = Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"] class FileContentStreamingResult(NamedTuple): diff --git a/litellm/files/utils.py b/litellm/files/utils.py index a0df7a89b0f..3ee4953bfef 100644 --- a/litellm/files/utils.py +++ b/litellm/files/utils.py @@ -26,24 +26,18 @@ class FilesAPIUtils: """ @staticmethod - def is_batch_jsonl_file( - create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData - ) -> bool: + def is_batch_jsonl_file(create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData) -> bool: """ Check if the file is a batch jsonl file """ return ( create_file_data.get("purpose") == "batch" - and FilesAPIUtils.valid_content_type( - extracted_file_data.get("content_type") - ) + and FilesAPIUtils.valid_content_type(extracted_file_data.get("content_type")) and extracted_file_data.get("content") is not None ) @staticmethod - def is_batch_jsonl_request( - create_file_data: CreateFileRequest, content_type: Optional[str] - ) -> bool: + def is_batch_jsonl_request(create_file_data: CreateFileRequest, content_type: Optional[str]) -> bool: """ Batch-jsonl check from metadata only, so the body can stay a streamable Path/handle instead of being read into memory. diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index 846a8a504a8..8a8a916fa9c 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -90,9 +90,7 @@ async def acreate_fine_tuning_job( Async: Creates and executes a batch from an uploaded file of request """ - verbose_logger.debug( - "inside acreate_fine_tuning_job model=%s and kwargs=%s", model, kwargs - ) + verbose_logger.debug("inside acreate_fine_tuning_job model=%s and kwargs=%s", model, kwargs) try: loop = asyncio.get_event_loop() kwargs["acreate_fine_tuning_job"] = True @@ -126,9 +124,7 @@ async def acreate_fine_tuning_job( raise e -def _build_fine_tuning_job_data( - model, training_file, hyperparameters, suffix, validation_file, integrations, seed -): +def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, validation_file, integrations, seed): return FineTuningJobCreate( model=model, training_file=training_file, @@ -245,17 +241,9 @@ def create_fine_tuning_job( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -272,9 +260,7 @@ def create_fine_tuning_job( get_secret_str("AZURE_AD_TOKEN") # type: ignore # Prepare Azure-specific parameters for extra_body - extra_body = _prepare_azure_extra_body( - extra_body, kwargs, azure_specific_hyperparams - ) + extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams) create_fine_tuning_job_data_dict = _build_fine_tuning_job_data( model, @@ -303,18 +289,12 @@ def create_fine_tuning_job( elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" vertex_ai_project = ( - optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") + optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) vertex_ai_location = ( - optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str( - "VERTEXAI_CREDENTIALS" + optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) + vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_fine_tuning_apis_instance.create_fine_tuning_job( _is_async=_is_async, create_fine_tuning_job_data=_build_fine_tuning_job_data( @@ -344,9 +324,7 @@ def create_fine_tuning_job( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -464,17 +442,9 @@ def cancel_fine_tuning_job( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -510,9 +480,7 @@ def cancel_fine_tuning_job( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -633,17 +601,9 @@ def list_fine_tuning_jobs( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -680,9 +640,7 @@ def list_fine_tuning_jobs( response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request( - method="create_thread", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response @@ -769,17 +727,9 @@ def retrieve_fine_tuning_job( or "https://api.openai.com/v1" ) organization = ( - optional_params.organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None - ) - api_key = ( - optional_params.api_key - or litellm.api_key - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") + optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None ) + api_key = optional_params.api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY") response = openai_fine_tuning_apis_instance.retrieve_fine_tuning_job( api_base=api_base, @@ -794,17 +744,9 @@ def retrieve_fine_tuning_job( ) # Azure OpenAI elif custom_llm_provider == "azure": - api_base = ( - optional_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) # type: ignore + api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore - api_version = ( - optional_params.api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) # type: ignore + api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 209e03d2bda..82777fb1378 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -25,14 +25,12 @@ class GenerateContentToCompletionHandler: """Prepare kwargs for litellm.completion/acompletion""" # Transform generate_content request to completion format - completion_request = ( - GOOGLE_GENAI_ADAPTER.translate_generate_content_to_completion( - model=model, - contents=contents, - config=config, - litellm_params=litellm_params, - **(extra_kwargs or {}), - ) + completion_request = GOOGLE_GENAI_ADAPTER.translate_generate_content_to_completion( + model=model, + contents=contents, + config=config, + litellm_params=litellm_params, + **(extra_kwargs or {}), ) completion_kwargs: Dict[str, Any] = dict(completion_request) @@ -62,15 +60,13 @@ class GenerateContentToCompletionHandler: ) -> Union[Dict[str, Any], AsyncIterator[bytes]]: """Handle generate_content call asynchronously using completion adapter""" - completion_kwargs = ( - GenerateContentToCompletionHandler._prepare_completion_kwargs( - model=model, - contents=contents, - config=config, - stream=stream, - litellm_params=litellm_params, - extra_kwargs=kwargs, - ) + completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs( + model=model, + contents=contents, + config=config, + stream=stream, + litellm_params=litellm_params, + extra_kwargs=kwargs, ) try: @@ -81,10 +77,8 @@ class GenerateContentToCompletionHandler: # This can happen in error cases or when stream is not properly supported if not hasattr(completion_response, "__aiter__"): # If it's not a stream, treat it as a regular response - generate_content_response = ( - GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( - cast(ModelResponse, completion_response) - ) + generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( + cast(ModelResponse, completion_response) ) return generate_content_response else: @@ -97,17 +91,13 @@ class GenerateContentToCompletionHandler: raise ValueError("Failed to transform streaming response") else: # Transform completion response back to generate_content format - generate_content_response = ( - GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( - cast(ModelResponse, completion_response) - ) + generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( + cast(ModelResponse, completion_response) ) return generate_content_response except Exception as e: - raise ValueError( - f"Error calling litellm.acompletion for generate_content: {str(e)}" - ) + raise ValueError(f"Error calling litellm.acompletion for generate_content: {str(e)}") @staticmethod def generate_content_handler( @@ -135,15 +125,13 @@ class GenerateContentToCompletionHandler: **kwargs, ) - completion_kwargs = ( - GenerateContentToCompletionHandler._prepare_completion_kwargs( - model=model, - contents=contents, - config=config, - stream=stream, - litellm_params=litellm_params, - extra_kwargs=kwargs, - ) + completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs( + model=model, + contents=contents, + config=config, + stream=stream, + litellm_params=litellm_params, + extra_kwargs=kwargs, ) try: @@ -154,10 +142,8 @@ class GenerateContentToCompletionHandler: # This can happen in error cases or when stream is not properly supported if not hasattr(completion_response, "__iter__"): # If it's not a stream, treat it as a regular response - generate_content_response = ( - GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( - cast(ModelResponse, completion_response) - ) + generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( + cast(ModelResponse, completion_response) ) return generate_content_response else: @@ -170,14 +156,10 @@ class GenerateContentToCompletionHandler: raise ValueError("Failed to transform streaming response") else: # Transform completion response back to generate_content format - generate_content_response = ( - GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( - cast(ModelResponse, completion_response) - ) + generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( + cast(ModelResponse, completion_response) ) return generate_content_response except Exception as e: - raise ValueError( - f"Error calling litellm.completion for generate_content: {str(e)}" - ) + raise ValueError(f"Error calling litellm.completion for generate_content: {str(e)}") diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index d8f9f1feb0b..02dde12a30d 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -49,17 +49,13 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): if self._returned_response: raise StopIteration self._returned_response = True - return GoogleGenAIAdapter().translate_completion_to_generate_content( - self.completion_stream - ) + return GoogleGenAIAdapter().translate_completion_to_generate_content(self.completion_stream) for chunk in self.completion_stream: if chunk == "None" or chunk is None: continue - transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content( - chunk, self - ) + transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content(chunk, self) if transformed_chunk: return transformed_chunk @@ -75,17 +71,13 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): if self._returned_response: raise StopAsyncIteration self._returned_response = True - return GoogleGenAIAdapter().translate_completion_to_generate_content( - self.completion_stream - ) + return GoogleGenAIAdapter().translate_completion_to_generate_content(self.completion_stream) async for chunk in self.completion_stream: if chunk == "None" or chunk is None: continue - transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content( - chunk, self - ) + transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content(chunk, self) if transformed_chunk: return transformed_chunk @@ -100,13 +92,10 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): try: # For tool calls with no arguments, accumulated_args will be "", which is not valid JSON. # We default to an empty JSON object in this case. - parsed_args = json.loads( - tool_call_data["arguments"] or "{}" - ) + parsed_args = json.loads(tool_call_data["arguments"] or "{}") function_call_part = { "functionCall": { - "name": tool_call_data["name"] - or "undefined_tool_name", + "name": tool_call_data["name"] or "undefined_tool_name", "args": parsed_args, } } @@ -163,9 +152,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): yield payload.encode() elif isinstance(chunk, ModelResponseStream): # Transform OpenAI streaming chunk to Google GenAI format - transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content( - chunk, self - ) + transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content(chunk, self) if isinstance(transformed_chunk, dict): # Only return non-empty chunks payload = f"data: {json.dumps(transformed_chunk)}\n\n" @@ -209,9 +196,7 @@ class GoogleGenAIAdapter: """ # Extract top-level fields from kwargs - system_instruction = kwargs.get("systemInstruction") or kwargs.get( - "system_instruction" - ) + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") tools = kwargs.get("tools") tool_config = kwargs.get("toolConfig") or kwargs.get("tool_config") @@ -222,9 +207,7 @@ class GoogleGenAIAdapter: contents_list = contents # Transform contents to OpenAI messages format - messages = self._transform_contents_to_messages( - contents_list, system_instruction=system_instruction - ) + messages = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction) # Create base request as dict (which is compatible with ChatCompletionRequest) completion_request: ChatCompletionRequest = { @@ -271,9 +254,7 @@ class GoogleGenAIAdapter: # Handle tool_config (tool choice) if tool_config: - tool_choice = self._transform_google_genai_tool_config_to_openai( - tool_config - ) + tool_choice = self._transform_google_genai_tool_config_to_openai(tool_config) if tool_choice: completion_request["tool_choice"] = tool_choice @@ -316,9 +297,7 @@ class GoogleGenAIAdapter: completion_stream: Any, ) -> Union[AsyncIterator[bytes], None]: """Transform streaming completion output to Google GenAI format""" - google_genai_wrapper = GoogleGenAIStreamWrapper( - completion_stream=completion_stream - ) + google_genai_wrapper = GoogleGenAIStreamWrapper(completion_stream=completion_stream) # Return the SSE-wrapped version for proper event formatting return google_genai_wrapper.async_google_genai_sse_wrapper() @@ -374,11 +353,7 @@ class GoogleGenAIAdapter: if system_instruction: system_parts = system_instruction.get("parts", []) if system_parts and "text" in system_parts[0]: - messages.append( - ChatCompletionSystemMessage( - role="system", content=system_parts[0]["text"] - ) - ) + messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"])) for content in contents: role = content.get("role", "user") @@ -386,9 +361,7 @@ class GoogleGenAIAdapter: if role == "user": # Handle user messages with potential function responses - content_parts: List[ - Union[ChatCompletionTextObject, ChatCompletionImageObject] - ] = [] + content_parts: List[Union[ChatCompletionTextObject, ChatCompletionImageObject]] = [] tool_messages: List[ChatCompletionToolMessage] = [] for part in parts: @@ -410,9 +383,7 @@ class GoogleGenAIAdapter: ChatCompletionImageObject, { "type": "image_url", - "image_url": { - "url": f"data:{mime_type};base64,{data}" - }, + "image_url": {"url": f"data:{mime_type};base64,{data}"}, }, ) ) @@ -426,11 +397,7 @@ class GoogleGenAIAdapter: ) tool_messages.append(tool_message) elif isinstance(part, str): - content_parts.append( - cast( - ChatCompletionTextObject, {"type": "text", "text": part} - ) - ) + content_parts.append(cast(ChatCompletionTextObject, {"type": "text", "text": part})) # Add user message if there's content if content_parts: @@ -441,18 +408,10 @@ class GoogleGenAIAdapter: and content_parts[0].get("type") == "text" ): text_part = cast(ChatCompletionTextObject, content_parts[0]) - messages.append( - ChatCompletionUserMessage( - role="user", content=text_part["text"] - ) - ) + messages.append(ChatCompletionUserMessage(role="user", content=text_part["text"])) else: # Use multimodal format (array of content parts) - messages.append( - ChatCompletionUserMessage( - role="user", content=content_parts - ) - ) + messages.append(ChatCompletionUserMessage(role="user", content=content_parts)) # Add tool messages messages.extend(tool_messages) @@ -520,15 +479,13 @@ class GoogleGenAIAdapter: # Handle different choice types (Choices vs StreamingChoices) if isinstance(choice, Choices): if not choice.message: - raise ValueError( - "Invalid completion response: no message found in choice" - ) + raise ValueError("Invalid completion response: no message found in choice") parts = self._transform_openai_message_to_google_genai_parts(choice.message) else: # Fallback for generic choice objects - message_content = getattr(choice, "message", {}).get( + message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get( "content", "" - ) or getattr(choice, "delta", {}).get("content", "") + ) parts = [{"text": message_content}] if message_content else [] # Create Google GenAI format response @@ -536,9 +493,7 @@ class GoogleGenAIAdapter: "candidates": [ { "content": {"parts": parts, "role": "model"}, - "finishReason": self._map_finish_reason( - getattr(choice, "finish_reason", None) - ), + "finishReason": self._map_finish_reason(getattr(choice, "finish_reason", None)), "index": 0, "safetyRatings": [], } @@ -589,9 +544,7 @@ class GoogleGenAIAdapter: # Handle streaming choice if isinstance(choice, StreamingChoices): if choice.delta: - parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation( - choice.delta, wrapper - ) + parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation(choice.delta, wrapper) else: parts = [] finish_reason = getattr(choice, "finish_reason", None) @@ -610,11 +563,7 @@ class GoogleGenAIAdapter: "candidates": [ { "content": {"parts": parts, "role": "model"}, - "finishReason": ( - self._map_finish_reason(finish_reason) - if finish_reason - else None - ), + "finishReason": (self._map_finish_reason(finish_reason) if finish_reason else None), "index": 0, "safetyRatings": [], } @@ -660,11 +609,7 @@ class GoogleGenAIAdapter: for tool_call in message.tool_calls: if hasattr(tool_call, "function") and tool_call.function: try: - args = ( - json.loads(tool_call.function.arguments) - if tool_call.function.arguments - else {} - ) + args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {} except json.JSONDecodeError: args = {} @@ -717,18 +662,14 @@ class GoogleGenAIAdapter: # Optimization: Skip chunks that have no new data if not function_name and not args_chunk: - verbose_logger.debug( - f"Skipping empty tool call chunk for index: {tool_call_index}" - ) + verbose_logger.debug(f"Skipping empty tool call chunk for index: {tool_call_index}") continue if function_name: wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name if args_chunk: - wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += ( - args_chunk - ) + wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += args_chunk # Attempt to parse and emit a complete tool call accumulated_data = wrapper.accumulated_tool_calls[tool_call_index] @@ -744,9 +685,7 @@ class GoogleGenAIAdapter: # The part will be created by a later chunk that brings the name. if accumulated_name: # If successful, create the part and clean up - function_call_part = { - "functionCall": {"name": accumulated_name, "args": parsed_args} - } + function_call_part = {"functionCall": {"name": accumulated_name, "args": parsed_args}} parts.append(function_call_part) # Remove the completed tool call from the accumulator diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index d35601c7e13..8e77c562094 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -103,9 +103,7 @@ class GenerateContentHelper: Returns: GenerateContentSetupResult containing all setup information """ - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj" - ) + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) # get llm provider logic @@ -135,11 +133,11 @@ class GenerateContentHelper: litellm_params.custom_llm_provider = custom_llm_provider # get provider config - generate_content_provider_config: Optional[ - BaseGoogleGenAIGenerateContentConfig - ] = ProviderConfigManager.get_provider_google_genai_generate_content_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), + generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] = ( + ProviderConfigManager.get_provider_google_genai_generate_content_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) ) if generate_content_provider_config is None: @@ -163,30 +161,24 @@ class GenerateContentHelper: # Construct request body ######################################################################################### # Create Google Optional Params Config - generate_content_config_dict = ( - generate_content_provider_config.map_generate_content_optional_params( - generate_content_config_dict=config or {}, - model=model, - ) + generate_content_config_dict = generate_content_provider_config.map_generate_content_optional_params( + generate_content_config_dict=config or {}, + model=model, ) # Extract systemInstruction from kwargs to pass to transform - system_instruction = kwargs.get("systemInstruction") or kwargs.get( - "system_instruction" - ) + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Native top-level REST fields arrive as loose kwargs and are otherwise dropped. native_request_fields: dict[str, object] = { field: kwargs[field] for field in generate_content_provider_config.get_generate_content_request_top_level_fields() if field in kwargs } - request_body = ( - generate_content_provider_config.transform_generate_content_request( - model=model, - contents=contents, - tools=tools, - generate_content_config_dict=generate_content_config_dict, - system_instruction=system_instruction, - ) + request_body = generate_content_provider_config.transform_generate_content_request( + model=model, + contents=contents, + tools=tools, + generate_content_config_dict=generate_content_config_dict, + system_instruction=system_instruction, ) # Pre Call logging @@ -328,12 +320,8 @@ def generate_content( config = kwargs.pop("generationConfig") # Check for mock response first litellm_params = GenericLiteLLMParams(**kwargs) - if litellm_params.mock_response and isinstance( - litellm_params.mock_response, str - ): - return GenerateContentHelper.mock_generate_content_response( - mock_response=litellm_params.mock_response - ) + if litellm_params.mock_response and isinstance(litellm_params.mock_response, str): + return GenerateContentHelper.mock_generate_content_response(mock_response=litellm_params.mock_response) # Setup the call setup_result = GenerateContentHelper.setup_generate_content_call( @@ -346,9 +334,7 @@ def generate_content( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get( - "system_instruction" - ) + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -375,9 +361,7 @@ def generate_content( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=_merge_native_request_fields( - setup_result.native_request_fields, extra_body - ), + extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), @@ -439,9 +423,7 @@ async def agenerate_content_stream( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get( - "system_instruction" - ) + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -449,17 +431,15 @@ async def agenerate_content_stream( kwargs.pop("stream", None) # Use the adapter to convert to completion format - return ( - await GenerateContentToCompletionHandler.async_generate_content_handler( - model=model, - contents=contents, # type: ignore - config=setup_result.generate_content_config_dict, - litellm_params=setup_result.litellm_params, - tools=tools, - stream=True, - extra_headers=extra_headers, - **kwargs, - ) + return await GenerateContentToCompletionHandler.async_generate_content_handler( + model=model, + contents=contents, # type: ignore + config=setup_result.generate_content_config_dict, + litellm_params=setup_result.litellm_params, + tools=tools, + stream=True, + extra_headers=extra_headers, + **kwargs, ) # Call the handler with async enabled and streaming @@ -474,9 +454,7 @@ async def agenerate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=_merge_native_request_fields( - setup_result.native_request_fields, extra_body - ), + extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body), timeout=timeout or request_timeout, _is_async=True, client=kwargs.get("client"), @@ -533,9 +511,7 @@ def generate_content_stream( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get( - "system_instruction" - ) + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -565,9 +541,7 @@ def generate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=_merge_native_request_fields( - setup_result.native_request_fields, extra_body - ), + extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index a8d0e5976f0..900a171640b 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -98,9 +98,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: ) -class GoogleGenAIGenerateContentStreamingIterator( - BaseGoogleGenAIGenerateContentStreamingIterator -): +class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator): """ Streaming iterator specifically for Google GenAI generate content API. """ @@ -148,14 +146,10 @@ class GoogleGenAIGenerateContentStreamingIterator( async def __anext__(self): # This should not be used for sync responses # If you need async iteration, use AsyncGoogleGenAIGenerateContentStreamingIterator - raise NotImplementedError( - "Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration" - ) + raise NotImplementedError("Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration") -class AsyncGoogleGenAIGenerateContentStreamingIterator( - BaseGoogleGenAIGenerateContentStreamingIterator -): +class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator): """ Async streaming iterator specifically for Google GenAI generate content API. """ diff --git a/litellm/images/main.py b/litellm/images/main.py index 34e77ffe8ab..17ea9aa177b 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -111,9 +111,7 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) - _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) - ) + _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None)) # Await normally init_response = await loop.run_in_executor(None, func_with_context) @@ -127,9 +125,7 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: response = await init_response # type: ignore if response is None: - raise ValueError( - "Unable to get Image Response. Please pass a valid llm_provider." - ) + raise ValueError("Unable to get Image Response. Please pass a valid llm_provider.") return response except Exception as e: @@ -272,15 +268,10 @@ def image_generation( } # model-specific params - pass them straight to the model/provider image_generation_config: Optional[BaseImageGenerationConfig] = None - if ( - custom_llm_provider is not None - and custom_llm_provider in LlmProviders._member_map_.values() - ): - image_generation_config = ( - ProviderConfigManager.get_provider_image_generation_config( - model=base_model or model, - provider=LlmProviders(custom_llm_provider), - ) + if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values(): + image_generation_config = ProviderConfigManager.get_provider_image_generation_config( + model=base_model or model, + provider=LlmProviders(custom_llm_provider), ) optional_params = get_optional_params_image_gen( @@ -327,11 +318,7 @@ def image_generation( api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") - api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) + api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") api_key = ( api_key @@ -341,9 +328,7 @@ def image_generation( or get_secret_str("AZURE_API_KEY") ) - azure_ad_token = optional_params.pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") + azure_ad_token = optional_params.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") # Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided if azure_ad_token_provider is None: @@ -355,10 +340,7 @@ def image_generation( tenant_id = litellm_params_dict.get("tenant_id") client_id = litellm_params_dict.get("client_id") client_secret = litellm_params_dict.get("client_secret") - azure_scope = ( - litellm_params_dict.get("azure_scope") - or "https://cognitiveservices.azure.com/.default" - ) + azure_scope = litellm_params_dict.get("azure_scope") or "https://cognitiveservices.azure.com/.default" # Create token provider if credentials are available if tenant_id and client_id and client_secret: @@ -413,9 +395,7 @@ def image_generation( litellm.LlmProviders.DASHSCOPE, ): if image_generation_config is None: - raise ValueError( - f"image generation config is not supported for {custom_llm_provider}" - ) + raise ValueError(f"image generation config is not supported for {custom_llm_provider}") # Resolve api_base from litellm.api_base if not explicitly provided _api_base = api_base or litellm.api_base @@ -524,9 +504,7 @@ def image_generation( api_base=api_base, api_key=api_key, ) - elif ( - custom_llm_provider in litellm._custom_providers - ): # Assume custom LLM provider + elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider # Get the Custom Handler custom_handler: Optional[CustomLLM] = None for item in litellm.custom_provider_map: @@ -534,9 +512,7 @@ def image_generation( custom_handler = item["custom_handler"] if custom_handler is None: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) + raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) ## ROUTE LLM CALL ## if aimg_generation is True: @@ -612,15 +588,11 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse: func_with_context = partial(ctx.run, func) if custom_llm_provider is None and model is not None: - _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) - ) + _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None)) # Await normally init_response = await loop.run_in_executor(None, func_with_context) - if isinstance(init_response, dict) or isinstance( - init_response, ImageResponse - ): ## CACHING SCENARIO + if isinstance(init_response, dict) or isinstance(init_response, ImageResponse): ## CACHING SCENARIO if isinstance(init_response, dict): init_response = ImageResponse(**init_response) response = init_response @@ -793,9 +765,7 @@ def image_edit( _is_async = kwargs.pop("async_call", False) is True # add images / or return a single image - images = ( - image if isinstance(image, list) else ([image] if image is not None else []) - ) + images = image if isinstance(image, list) else ([image] if image is not None else []) headers_from_kwargs = kwargs.get("headers") merged_extra_headers: Dict[str, Any] = {} @@ -822,17 +792,13 @@ def image_edit( custom_handler = item["custom_handler"] if custom_handler is None: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) + raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) model_response = ImageResponse() if _is_async: async_custom_client: Optional[AsyncHTTPHandler] = None - if kwargs.get("client") is not None and isinstance( - kwargs.get("client"), AsyncHTTPHandler - ): + if kwargs.get("client") is not None and isinstance(kwargs.get("client"), AsyncHTTPHandler): async_custom_client = kwargs.get("client") return custom_handler.aimage_edit( @@ -849,9 +815,7 @@ def image_edit( ) else: custom_client: Optional[HTTPHandler] = None - if kwargs.get("client") is not None and isinstance( - kwargs.get("client"), HTTPHandler - ): + if kwargs.get("client") is not None and isinstance(kwargs.get("client"), HTTPHandler): custom_client = kwargs.get("client") return custom_handler.image_edit( @@ -881,19 +845,15 @@ def image_edit( local_vars.update(kwargs) # Get ImageEditOptionalRequestParams with only valid parameters image_edit_optional_params: ImageEditOptionalRequestParams = ( - _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( - local_vars - ) + _get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars) ) # Get optional parameters for the responses API - image_edit_request_params: Dict = ( - _get_ImageEditRequestUtils().get_optional_params_image_edit( - model=model, - image_edit_provider_config=image_edit_provider_config, - image_edit_optional_params=image_edit_optional_params, - drop_params=kwargs.get("drop_params"), - additional_drop_params=kwargs.get("additional_drop_params"), - ) + image_edit_request_params: Dict = _get_ImageEditRequestUtils().get_optional_params_image_edit( + model=model, + image_edit_provider_config=image_edit_provider_config, + image_edit_optional_params=image_edit_optional_params, + drop_params=kwargs.get("drop_params"), + additional_drop_params=kwargs.get("additional_drop_params"), ) # Pre Call logging diff --git a/litellm/images/utils.py b/litellm/images/utils.py index 8d3e96f1433..f0d4c985c01 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -39,9 +39,7 @@ class ImageEditRequestUtils: for param in additional_drop_params: filtered_optional_params.pop(param, None) - unsupported_params = [ - param for param in filtered_optional_params if param not in supported_params - ] + unsupported_params = [param for param in filtered_optional_params if param not in supported_params] if unsupported_params: if should_drop: @@ -54,9 +52,7 @@ class ImageEditRequestUtils: ) mapped_params = image_edit_provider_config.map_openai_params( - image_edit_optional_params=cast( - ImageEditOptionalRequestParams, filtered_optional_params - ), + image_edit_optional_params=cast(ImageEditOptionalRequestParams, filtered_optional_params), model=model, drop_params=should_drop, ) @@ -77,9 +73,7 @@ class ImageEditRequestUtils: ImageEditOptionalRequestParams instance with only the valid parameters """ valid_keys = get_type_hints(ImageEditOptionalRequestParams).keys() - filtered_params = { - k: v for k, v in params.items() if k in valid_keys and v is not None - } + filtered_params = {k: v for k, v in params.items() if k in valid_keys and v is not None} return cast(ImageEditOptionalRequestParams, filtered_params) @staticmethod @@ -99,9 +93,7 @@ class ImageEditRequestUtils: # Save current position current_pos = image_data.tell() image_data.seek(0) - bytes_data = image_data.read( - 100 - ) # First 100 bytes are enough for detection + bytes_data = image_data.read(100) # First 100 bytes are enough for detection # Restore position image_data.seek(current_pos) elif isinstance(image_data, BufferedReader): diff --git a/litellm/integrations/SlackAlerting/batching_handler.py b/litellm/integrations/SlackAlerting/batching_handler.py index 828f3eb4175..42f4f562422 100644 --- a/litellm/integrations/SlackAlerting/batching_handler.py +++ b/litellm/integrations/SlackAlerting/batching_handler.py @@ -40,9 +40,7 @@ def squash_payloads(queue): return squashed -def _print_alerting_payload_warning( - payload: dict, slackAlertingInstance: SlackAlertingType -): +def _print_alerting_payload_warning(payload: dict, slackAlertingInstance: SlackAlertingType): """ Print the payload to the console when slackAlertingInstance.alerting_args.log_to_console is True @@ -70,12 +68,8 @@ async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count) data=json.dumps(payload), ) if response.status_code != 200: - verbose_proxy_logger.debug( - f"Error sending slack alert to url={item['url']}. Error={response.text}" - ) + verbose_proxy_logger.debug(f"Error sending slack alert to url={item['url']}. Error={response.text}") except Exception as e: verbose_proxy_logger.debug(f"Error sending slack alert: {str(e)}") finally: - _print_alerting_payload_warning( - payload, slackAlertingInstance=slackAlertingInstance - ) + _print_alerting_payload_warning(payload, slackAlertingInstance=slackAlertingInstance) diff --git a/litellm/integrations/SlackAlerting/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py index 8c0da9bb4fb..136b6583f38 100644 --- a/litellm/integrations/SlackAlerting/hanging_request_check.py +++ b/litellm/integrations/SlackAlerting/hanging_request_check.py @@ -41,8 +41,7 @@ class AlertingHangingRequestCheck: # stay cached for at least 1.5x the threshold to guarantee a check # happens after they cross it self.hanging_request_cache_ttl = int( - self.slack_alerting_object.alerting_threshold * 1.5 - + HANGING_ALERT_BUFFER_TIME_SECONDS + self.slack_alerting_object.alerting_threshold * 1.5 + HANGING_ALERT_BUFFER_TIME_SECONDS ) self.hanging_request_cache = InMemoryCache( default_ttl=self.hanging_request_cache_ttl, @@ -62,9 +61,7 @@ class AlertingHangingRequestCheck: model = request_data.get("model", "") api_base: Optional[str] = None - if request_data.get("deployment", None) is not None and isinstance( - request_data["deployment"], dict - ): + if request_data.get("deployment", None) is not None and isinstance(request_data["deployment"], dict): api_base = litellm.get_api_base( model=model, optional_params=request_data["deployment"].get("litellm_params", {}), @@ -104,9 +101,7 @@ class AlertingHangingRequestCheck: ) for request_id in hanging_requests: - hanging_request_data: Optional[ - HangingRequestData - ] = await self.hanging_request_cache.async_get_cache( + hanging_request_data: Optional[HangingRequestData] = await self.hanging_request_cache.async_get_cache( key=request_id, ) @@ -116,12 +111,10 @@ class AlertingHangingRequestCheck: if hanging_request_data.alerted: continue - request_status = ( - await proxy_logging_obj.internal_usage_cache.async_get_cache( - key="request_status:{}".format(hanging_request_data.request_id), - litellm_parent_otel_span=None, - local_only=True, - ) + request_status = await proxy_logging_obj.internal_usage_cache.async_get_cache( + key="request_status:{}".format(hanging_request_data.request_id), + litellm_parent_otel_span=None, + local_only=True, ) # this means the request status was either success or fail # and is not hanging @@ -141,9 +134,7 @@ class AlertingHangingRequestCheck: ################ # Send the Alert on Slack ################ - await self.send_hanging_request_alert( - hanging_request_data=hanging_request_data - ) + await self.send_hanging_request_alert(hanging_request_data=hanging_request_data) # flag so the entry is skipped on later ticks; one alert per hang, # with the existing TTL still handling cleanup hanging_request_data.alerted = True diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 35731306b93..e93c650ed97 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -62,9 +62,7 @@ class SlackAlerting(CustomBatchLogger): def __init__( self, internal_usage_cache: Optional[DualCache] = None, - alerting_threshold: Optional[ - float - ] = None, # threshold for slow / hanging llm responses (in seconds) + alerting_threshold: Optional[float] = None, # threshold for slow / hanging llm responses (in seconds) alerting: Optional[List] = [], alert_types: List[AlertType] = DEFAULT_ALERT_TYPES, alert_to_webhook_url: Optional[ @@ -81,12 +79,8 @@ class SlackAlerting(CustomBatchLogger): self.alerting = alerting self.alert_types = alert_types self.internal_usage_cache = internal_usage_cache or DualCache() - self.async_http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) - self.alert_to_webhook_url = process_slack_alerting_variables( - alert_to_webhook_url=alert_to_webhook_url - ) + self.async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) self.is_running = False self.alerting_args = SlackAlertingArgs(**alerting_args) self.default_webhook_url = default_webhook_url @@ -98,9 +92,7 @@ class SlackAlerting(CustomBatchLogger): self.alert_type_config: Dict[str, AlertTypeConfig] = {} if alert_type_config: for key, val in alert_type_config.items(): - self.alert_type_config[key] = ( - AlertTypeConfig(**val) if isinstance(val, dict) else val - ) + self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val self.digest_buckets: Dict[str, DigestEntry] = {} self.digest_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) @@ -130,23 +122,14 @@ class SlackAlerting(CustomBatchLogger): self.periodic_started = True if alert_type_config is not None: for key, val in alert_type_config.items(): - self.alert_type_config[key] = ( - AlertTypeConfig(**val) if isinstance(val, dict) else val - ) + self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val if alert_to_webhook_url is not None: # update the dict if self.alert_to_webhook_url is None: - self.alert_to_webhook_url = process_slack_alerting_variables( - alert_to_webhook_url=alert_to_webhook_url - ) + self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) else: - _new_values = ( - process_slack_alerting_variables( - alert_to_webhook_url=alert_to_webhook_url - ) - or {} - ) + _new_values = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) or {} self.alert_to_webhook_url.update(_new_values) if llm_router is not None: self.llm_router = llm_router @@ -161,15 +144,11 @@ class SlackAlerting(CustomBatchLogger): # Convert to dict for processing cache_value = dict(outage_value) - if "deployment_ids" in cache_value and isinstance( - cache_value["deployment_ids"], set - ): + if "deployment_ids" in cache_value and isinstance(cache_value["deployment_ids"], set): cache_value["deployment_ids"] = list(cache_value["deployment_ids"]) return cache_value - def _restore_outage_value_from_cache( - self, outage_value: Optional[dict] - ) -> Optional[dict]: + def _restore_outage_value_from_cache(self, outage_value: Optional[dict]) -> Optional[dict]: """ Helper method to restore outage value after retrieving from cache. Converts list objects back to sets for proper handling. @@ -234,9 +213,7 @@ class SlackAlerting(CustomBatchLogger): _deployment_latency_map: Optional[dict] = None try: # try sorting deployments by latency - _deployment_latencies = sorted( - _deployment_latencies.items(), key=lambda x: x[1] - ) + _deployment_latencies = sorted(_deployment_latencies.items(), key=lambda x: x[1]) _deployment_latency_map = dict(_deployment_latencies) except Exception: pass @@ -276,23 +253,13 @@ class SlackAlerting(CustomBatchLogger): alerting_metadata: dict = {} if time_difference_float > self.alerting_threshold: # add deployment latencies to alert - if ( - kwargs is not None - and "litellm_params" in kwargs - and "metadata" in kwargs["litellm_params"] - ): + if kwargs is not None and "litellm_params" in kwargs and "metadata" in kwargs["litellm_params"]: _metadata: dict = kwargs["litellm_params"]["metadata"] - request_info = _add_key_name_and_team_to_alert( - request_info=request_info, metadata=_metadata - ) + request_info = _add_key_name_and_team_to_alert(request_info=request_info, metadata=_metadata) - _deployment_latency_map = self._get_deployment_latencies_to_alert( - metadata=_metadata - ) + _deployment_latency_map = self._get_deployment_latencies_to_alert(metadata=_metadata) if _deployment_latency_map is not None: - request_info += ( - f"\nAvailable Deployment Latencies\n{_deployment_latency_map}" - ) + request_info += f"\nAvailable Deployment Latencies\n{_deployment_latency_map}" if "alerting_metadata" in _metadata: alerting_metadata = _metadata["alerting_metadata"] @@ -305,9 +272,7 @@ class SlackAlerting(CustomBatchLogger): api_base=api_base, ) - async def async_update_daily_reports( - self, deployment_metrics: DeploymentMetrics - ) -> int: + async def async_update_daily_reports(self, deployment_metrics: DeploymentMetrics) -> int: """ Store the perf by deployment in cache - Number of failed requests per deployment @@ -338,9 +303,7 @@ class SlackAlerting(CustomBatchLogger): ## LATENCY ## if deployment_metrics.latency_per_output_token is not None: await self.internal_usage_cache.async_increment_cache( - key="{}:{}".format( - deployment_metrics.id, SlackAlertingCacheKeys.latency_key.value - ), + key="{}:{}".format(deployment_metrics.id, SlackAlertingCacheKeys.latency_key.value), value=deployment_metrics.latency_per_output_token, parent_otel_span=None, # no attached request, this is a background operation ) @@ -370,13 +333,8 @@ class SlackAlerting(CustomBatchLogger): ids = router.get_model_ids() # get keys - failed_request_keys = [ - "{}:{}".format(id, SlackAlertingCacheKeys.failed_requests_key.value) - for id in ids - ] - latency_keys = [ - "{}:{}".format(id, SlackAlertingCacheKeys.latency_key.value) for id in ids - ] + failed_request_keys = ["{}:{}".format(id, SlackAlertingCacheKeys.failed_requests_key.value) for id in ids] + latency_keys = ["{}:{}".format(id, SlackAlertingCacheKeys.latency_key.value) for id in ids] combined_metrics_keys = failed_request_keys + latency_keys # reduce cache calls @@ -396,18 +354,13 @@ class SlackAlerting(CustomBatchLogger): if all_none: return False - failed_request_values = combined_metrics_values[ - : len(failed_request_keys) - ] # # [1, 2, None, ..] + failed_request_values = combined_metrics_values[: len(failed_request_keys)] # # [1, 2, None, ..] latency_values = combined_metrics_values[len(failed_request_keys) :] # find top 5 failed ## Replace None values with a placeholder value (-1 in this case) placeholder_value = 0 - replaced_failed_values = [ - value if value is not None else placeholder_value - for value in failed_request_values - ] + replaced_failed_values = [value if value is not None else placeholder_value for value in failed_request_values] ## Get the indices of top 5 keys with the highest numerical values (ignoring None and 0 values) top_5_failed = sorted( @@ -415,17 +368,12 @@ class SlackAlerting(CustomBatchLogger): key=lambda i: replaced_failed_values[i], reverse=True, )[:5] - top_5_failed = [ - index for index in top_5_failed if replaced_failed_values[index] > 0 - ] + top_5_failed = [index for index in top_5_failed if replaced_failed_values[index] > 0] # find top 5 slowest # Replace None values with a placeholder value (-1 in this case) placeholder_value = 0 - replaced_slowest_values = [ - value if value is not None else placeholder_value - for value in latency_values - ] + replaced_slowest_values = [value if value is not None else placeholder_value for value in latency_values] # Get the indices of top 5 values with the highest numerical values (ignoring None and 0 values) top_5_slowest = sorted( @@ -433,9 +381,7 @@ class SlackAlerting(CustomBatchLogger): key=lambda i: replaced_slowest_values[i], reverse=True, )[:5] - top_5_slowest = [ - index for index in top_5_slowest if replaced_slowest_values[index] > 0 - ] + top_5_slowest = [index for index in top_5_slowest if replaced_slowest_values[index] > 0] # format alert -> return the litellm model name + api base message = f"\n\nTime: `{time.time()}`s\nHere are today's key metrics 📈: \n\n" @@ -453,14 +399,14 @@ class SlackAlerting(CustomBatchLogger): api_base = litellm.get_api_base( model=deployment_name, - optional_params=( - _deployment["litellm_params"] if _deployment is not None else {} - ), + optional_params=(_deployment["litellm_params"] if _deployment is not None else {}), ) if api_base is None: api_base = "" value = replaced_failed_values[top_5_failed[i]] - message += f"\t{i + 1}. Deployment: `{deployment_name}`, Failed Requests: `{value}`, API Base: `{api_base}`\n" + message += ( + f"\t{i + 1}. Deployment: `{deployment_name}`, Failed Requests: `{value}`, API Base: `{api_base}`\n" + ) message += "\n\n*😅 Top Slowest Deployments:*\n\n" if not top_5_slowest: @@ -474,9 +420,7 @@ class SlackAlerting(CustomBatchLogger): deployment_name = "" api_base = litellm.get_api_base( model=deployment_name, - optional_params=( - _deployment["litellm_params"] if _deployment is not None else {} - ), + optional_params=(_deployment["litellm_params"] if _deployment is not None else {}), ) value = round(replaced_slowest_values[top_5_slowest[i]], 3) message += f"\t{i + 1}. Deployment: `{deployment_name}`, Latency per output token: `{value}s/token`, API Base: `{api_base}`\n\n" @@ -485,9 +429,7 @@ class SlackAlerting(CustomBatchLogger): latency_cache_keys = [(key, 0) for key in latency_keys] failed_request_cache_keys = [(key, 0) for key in failed_request_keys] combined_metrics_cache_keys = latency_cache_keys + failed_request_cache_keys - await self.internal_usage_cache.async_set_cache_pipeline( - cache_list=combined_metrics_cache_keys - ) + await self.internal_usage_cache.async_set_cache_pipeline(cache_list=combined_metrics_cache_keys) message += f"\n\nNext Run is at: `{time.time() + self.alerting_args.daily_report_frequency}`s" @@ -511,9 +453,7 @@ class SlackAlerting(CustomBatchLogger): if AlertType.llm_requests_hanging not in self.alert_types: return - await self.hanging_request_check.add_request_to_hanging_request_check( - request_data=request_data - ) + await self.hanging_request_check.add_request_to_hanging_request_check(request_data=request_data) async def failed_tracking_alert(self, error_message: str, failing_model: str): """ @@ -686,9 +626,7 @@ class SlackAlerting(CustomBatchLogger): if user_info.max_budget is not None: if user_info.spend >= user_info.max_budget: event = "budget_crossed" - event_message += ( - f"Budget Crossed\n Total Budget:`{user_info.max_budget}`" - ) + event_message += f"Budget Crossed\n Total Budget:`{user_info.max_budget}`" elif percent_left <= SLACK_ALERTING_THRESHOLD_5_PERCENT: event = "threshold_crossed" event_message += "5% Threshold Crossed " @@ -755,9 +693,7 @@ class SlackAlerting(CustomBatchLogger): projected_spend=None, event="spend_tracked", event_group=Litellm_EntityType.END_USER, - event_message="Customer spend tracked. Customer={}, spend={}".format( - end_user_id, response_cost - ), + event_message="Customer spend tracked. Customer={}, spend={}".format(end_user_id, response_cost), ) await self.send_webhook_alert(webhook_event=event) @@ -852,9 +788,9 @@ class SlackAlerting(CustomBatchLogger): ### UNIQUE CACHE KEY ### cache_key = provider + region_name - outage_value: Optional[ - ProviderRegionOutageModel - ] = await self.internal_usage_cache.async_get_cache(key=cache_key) + outage_value: Optional[ProviderRegionOutageModel] = await self.internal_usage_cache.async_get_cache( + key=cache_key + ) # Convert deployment_ids back to set if it was stored as a list if outage_value is not None: @@ -904,8 +840,7 @@ class SlackAlerting(CustomBatchLogger): ## MINOR OUTAGE ALERT SENT ## if ( outage_value["minor_alert_sent"] is False - and len(outage_value["alerts"]) - >= self.alerting_args.minor_outage_alert_threshold + and len(outage_value["alerts"]) >= self.alerting_args.minor_outage_alert_threshold and len(_deployment_set) > 1 # make sure it's not just 1 bad deployment ): msg = self._outage_alert_msg_factory( @@ -929,8 +864,7 @@ class SlackAlerting(CustomBatchLogger): ## MAJOR OUTAGE ALERT SENT ## elif ( outage_value["major_alert_sent"] is False - and len(outage_value["alerts"]) - >= self.alerting_args.major_outage_alert_threshold + and len(outage_value["alerts"]) >= self.alerting_args.major_outage_alert_threshold and len(_deployment_set) > 1 # make sure it's not just 1 bad deployment ): msg = self._outage_alert_msg_factory( @@ -955,9 +889,7 @@ class SlackAlerting(CustomBatchLogger): ## update cache ## # Convert set to list for JSON serialization cache_value = self._prepare_outage_value_for_cache(outage_value) - await self.internal_usage_cache.async_set_cache( - key=cache_key, value=cache_value - ) + await self.internal_usage_cache.async_set_cache(key=cache_key, value=cache_value) async def outage_alerts( self, @@ -979,9 +911,7 @@ class SlackAlerting(CustomBatchLogger): max_alerts_size = 10 """ try: - outage_value: Optional[ - OutageModel - ] = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore + outage_value: Optional[OutageModel] = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore if ( getattr(exception, "status_code", None) is None or ( @@ -1004,9 +934,7 @@ class SlackAlerting(CustomBatchLogger): model, provider, _, _ = litellm.get_llm_provider(model=model) except Exception: provider = "" - api_base = litellm.get_api_base( - model=model, optional_params=deployment.litellm_params - ) + api_base = litellm.get_api_base(model=model, optional_params=deployment.litellm_params) if outage_value is None: outage_value = OutageModel( @@ -1025,10 +953,7 @@ class SlackAlerting(CustomBatchLogger): ) return - if ( - len(outage_value["alerts"]) - < self.alerting_args.max_outage_alert_list_size - ): + if len(outage_value["alerts"]) < self.alerting_args.max_outage_alert_list_size: outage_value["alerts"].append(exception.status_code) # type: ignore else: # prevent memory leaks pass @@ -1038,8 +963,7 @@ class SlackAlerting(CustomBatchLogger): ## MINOR OUTAGE ALERT SENT ## if ( outage_value["minor_alert_sent"] is False - and len(outage_value["alerts"]) - >= self.alerting_args.minor_outage_alert_threshold + and len(outage_value["alerts"]) >= self.alerting_args.minor_outage_alert_threshold ): msg = self._outage_alert_msg_factory( alert_type="Minor", @@ -1060,8 +984,7 @@ class SlackAlerting(CustomBatchLogger): outage_value["minor_alert_sent"] = True elif ( outage_value["major_alert_sent"] is False - and len(outage_value["alerts"]) - >= self.alerting_args.major_outage_alert_threshold + and len(outage_value["alerts"]) >= self.alerting_args.major_outage_alert_threshold ): msg = self._outage_alert_msg_factory( alert_type="Major", @@ -1084,15 +1007,11 @@ class SlackAlerting(CustomBatchLogger): ## update cache ## # Convert set to list for JSON serialization cache_value = self._prepare_outage_value_for_cache(outage_value) - await self.internal_usage_cache.async_set_cache( - key=deployment_id, value=cache_value - ) + await self.internal_usage_cache.async_set_cache(key=deployment_id, value=cache_value) except Exception: pass - async def model_added_alert( - self, model_name: str, litellm_model_name: str, passed_model_info: Any - ): + async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: Any): base_model_from_user = getattr(passed_model_info, "base_model", None) model_info = {} base_model = "" @@ -1193,14 +1112,10 @@ Model Info: if premium_user is not True: if email_logo_url is not None or email_support_contact is not None: - raise ValueError( - f"Trying to Customize Email Alerting\n {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"Trying to Customize Email Alerting\n {CommonProxyErrors.not_premium_user.value}") return - async def send_key_created_or_user_invited_email( - self, webhook_event: WebhookEvent - ) -> bool: + async def send_key_created_or_user_invited_email(self, webhook_event: WebhookEvent) -> bool: try: from litellm.proxy.utils import send_email @@ -1213,13 +1128,9 @@ Model Info: return False from litellm.proxy.proxy_server import premium_user, prisma_client - email_logo_url = os.getenv( - "SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", None) - ) + email_logo_url = os.getenv("SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", None)) email_support_contact = os.getenv("EMAIL_SUPPORT_CONTACT", None) - await self._check_if_using_premium_email_feature( - premium_user, email_logo_url, email_support_contact - ) + await self._check_if_using_premium_email_feature(premium_user, email_logo_url, email_support_contact) if email_logo_url is None: email_logo_url = LITELLM_LOGO_URL if email_support_contact is None: @@ -1228,14 +1139,8 @@ Model Info: event_name = webhook_event.event_message recipient_email = webhook_event.user_email recipient_user_id = webhook_event.user_id - if ( - recipient_email is None - and recipient_user_id is not None - and prisma_client is not None - ): - user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": recipient_user_id} - ) + if recipient_email is None and recipient_user_id is not None and prisma_client is not None: + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": recipient_user_id}) if user_row is not None: recipient_email = user_row.user_email @@ -1265,9 +1170,7 @@ Model Info: team_id = webhook_event.team_id team_name = "Default Team" if team_id is not None and prisma_client is not None: - team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if team_row is not None: team_name = team_row.team_alias or "-" email_html_content = USER_INVITED_EMAIL_TEMPLATE.format( @@ -1302,9 +1205,7 @@ Model Info: verbose_proxy_logger.error("Error sending email alert %s", str(e)) return False - async def send_email_alert_using_smtp( - self, webhook_event: WebhookEvent, alert_type: str - ) -> bool: + async def send_email_alert_using_smtp(self, webhook_event: WebhookEvent, alert_type: str) -> bool: """ Sends structured Email alert to an SMTP server @@ -1315,13 +1216,9 @@ Model Info: from litellm.proxy.proxy_server import premium_user from litellm.proxy.utils import send_email - email_logo_url = os.getenv( - "SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", None) - ) + email_logo_url = os.getenv("SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", None)) email_support_contact = os.getenv("EMAIL_SUPPORT_CONTACT", None) - await self._check_if_using_premium_email_feature( - premium_user, email_logo_url, email_support_contact - ) + await self._check_if_using_premium_email_feature(premium_user, email_logo_url, email_support_contact) if email_logo_url is None: email_logo_url = LITELLM_LOGO_URL @@ -1334,9 +1231,7 @@ Model Info: max_budget = webhook_event.max_budget email_html_content = "Alert from LiteLLM Server" if recipient_email is None: - verbose_proxy_logger.error( - "Trying to send email alert to no recipient", extra=webhook_event.dict() - ) + verbose_proxy_logger.error("Trying to send email alert to no recipient", extra=webhook_event.dict()) if webhook_event.event == "budget_crossed": email_html_content = f""" @@ -1404,30 +1299,16 @@ Model Info: return # Start periodic flush if not already started - if ( - not self.periodic_started - and self.alerting is not None - and len(self.alerting) > 0 - ): + if not self.periodic_started and self.alerting is not None and len(self.alerting) > 0: asyncio.create_task(self.periodic_flush()) self.periodic_started = True - if ( - "webhook" in self.alerting - and alert_type == "budget_alerts" - and user_info is not None - ): + if "webhook" in self.alerting and alert_type == "budget_alerts" and user_info is not None: await self.send_webhook_alert(webhook_event=user_info) - if ( - "email" in self.alerting - and alert_type == "budget_alerts" - and user_info is not None - ): + if "email" in self.alerting and alert_type == "budget_alerts" and user_info is not None: # only send budget alerts over Email - await self.send_email_alert_using_smtp( - webhook_event=user_info, alert_type=alert_type - ) + await self.send_email_alert_using_smtp(webhook_event=user_info, alert_type=alert_type) if "slack" not in self.alerting: return @@ -1441,13 +1322,8 @@ Model Info: _atc = self.alert_type_config.get(alert_type_name_str) if _atc is not None and _atc.digest: # Resolve webhook URL for this alert type (needed for digest entry) - if ( - self.alert_to_webhook_url is not None - and alert_type in self.alert_to_webhook_url - ): - _digest_webhook: Optional[Union[str, List[str]]] = ( - self.alert_to_webhook_url[alert_type] - ) + if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url: + _digest_webhook: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type] elif self.default_webhook_url is not None: _digest_webhook = self.default_webhook_url else: @@ -1485,7 +1361,9 @@ Model Info: if alert_type == "daily_reports" or alert_type == "new_model_added": formatted_message = alert_type_formatted + message else: - formatted_message = f"{alert_type_formatted}\nLevel: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" + formatted_message = ( + f"{alert_type_formatted}\nLevel: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" + ) if kwargs: for key, value in kwargs.items(): @@ -1497,13 +1375,8 @@ Model Info: formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`" # check if we find the slack webhook url in self.alert_to_webhook_url - if ( - self.alert_to_webhook_url is not None - and alert_type in self.alert_to_webhook_url - ): - slack_webhook_url: Optional[Union[str, List[str]]] = ( - self.alert_to_webhook_url[alert_type] - ) + if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url: + slack_webhook_url: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type] elif self.default_webhook_url is not None: slack_webhook_url = self.default_webhook_url else: @@ -1543,9 +1416,7 @@ Model Info: squashed_queue = squash_payloads(self.log_queue) tasks = [ - send_to_webhook( - slackAlertingInstance=self, item=item["item"], count=item["count"] - ) + send_to_webhook(slackAlertingInstance=self, item=item["item"], count=item["count"]) for item in squashed_queue.values() ] await asyncio.gather(*tasks) @@ -1645,9 +1516,7 @@ Model Info: ): completion_tokens = response_obj.usage.completion_tokens # type: ignore if completion_tokens is not None and completion_tokens > 0: - final_value = float( - response_s.total_seconds() / completion_tokens - ) + final_value = float(response_s.total_seconds() / completion_tokens) if isinstance(final_value, timedelta): final_value = final_value.total_seconds() @@ -1692,9 +1561,7 @@ Model Info: ) if "region_outage_alerts" in self.alert_types: - await self.region_outage_alerts( - exception=kwargs["exception"], deployment_id=model_id - ) + await self.region_outage_alerts(exception=kwargs["exception"], deployment_id=model_id) except Exception: pass @@ -1781,7 +1648,9 @@ Model Info: todays_date = datetime.datetime.now().date() start_date = todays_date - datetime.timedelta(days=days) - _event_cache_key = f"weekly_spend_report_sent_{start_date.strftime('%Y-%m-%d')}_{todays_date.strftime('%Y-%m-%d')}" + _event_cache_key = ( + f"weekly_spend_report_sent_{start_date.strftime('%Y-%m-%d')}_{todays_date.strftime('%Y-%m-%d')}" + ) if await self.internal_usage_cache.async_get_cache(key=_event_cache_key): return @@ -1800,9 +1669,7 @@ Model Info: _spend_message += "\n*Team Spend Report:*\n" for spend in spend_per_team: _team_spend = round(float(spend["total_spend"]), 4) - _spend_message += ( - f"Team: `{spend['team_alias']}` | Spend: `${_team_spend}`\n" - ) + _spend_message += f"Team: `{spend['team_alias']}` | Spend: `${_team_spend}`\n" if spend_per_tag is not None: _spend_message += "\n*Tag Spend Report:*\n" @@ -1840,9 +1707,7 @@ Model Info: todays_date = datetime.datetime.now().date() first_day_of_month = todays_date.replace(day=1) _, last_day_of_month = monthrange(todays_date.year, todays_date.month) - last_day_of_month = first_day_of_month + datetime.timedelta( - days=last_day_of_month - 1 - ) + last_day_of_month = first_day_of_month + datetime.timedelta(days=last_day_of_month - 1) _event_cache_key = f"monthly_spend_report_sent_{first_day_of_month.strftime('%Y-%m-%d')}_{last_day_of_month.strftime('%Y-%m-%d')}" if await self.internal_usage_cache.async_get_cache(key=_event_cache_key): @@ -1867,9 +1732,7 @@ Model Info: _team_spend = float(_team_spend) # round to 4 decimal places _team_spend = round(_team_spend, 4) - _spend_message += ( - f"Team: `{spend['team_alias']}` | Spend: `${_team_spend}`\n" - ) + _spend_message += f"Team: `{spend['team_alias']}` | Spend: `${_team_spend}`\n" if monthly_spend_per_tag is not None: _spend_message += "\n*Tag Spend Report:*\n" @@ -1908,13 +1771,9 @@ Model Info: ) # call prometheuslogger. - falllback_success_info_prometheus = ( - await get_fallback_metric_from_prometheus() - ) + falllback_success_info_prometheus = await get_fallback_metric_from_prometheus() - fallback_message = ( - f"*Fallback Statistics:*\n{falllback_success_info_prometheus}" - ) + fallback_message = f"*Fallback Statistics:*\n{falllback_success_info_prometheus}" await self.send_alert( message=fallback_message, @@ -1969,9 +1828,7 @@ Model Info: ) except Exception as e: - verbose_proxy_logger.error( - "Error sending send_virtual_key_event_slack %s", e - ) + verbose_proxy_logger.error("Error sending send_virtual_key_event_slack %s", e) return @@ -1982,10 +1839,7 @@ Model Info: if request_data is None: return False - if ( - request_data.get("litellm_status", "") != "success" - and request_data.get("litellm_status", "") != "fail" - ): + if request_data.get("litellm_status", "") != "success" and request_data.get("litellm_status", "") != "fail": ## CHECK IF CACHE IS UPDATED litellm_call_id = request_data.get("litellm_call_id", "") status: Optional[str] = await self.internal_usage_cache.async_get_cache( diff --git a/litellm/integrations/SlackAlerting/utils.py b/litellm/integrations/SlackAlerting/utils.py index e2580768178..4424bedba81 100644 --- a/litellm/integrations/SlackAlerting/utils.py +++ b/litellm/integrations/SlackAlerting/utils.py @@ -34,9 +34,7 @@ def process_slack_alerting_variables( if "os.environ/" in webhook_url: _env_value = get_secret(secret_name=webhook_url) if not isinstance(_env_value, str): - raise ValueError( - f"Invalid webhook url value for: {webhook_url}. Got type={type(_env_value)}" - ) + raise ValueError(f"Invalid webhook url value for: {webhook_url}. Got type={type(_env_value)}") _webhook_values.append(_env_value) else: _webhook_values.append(webhook_url) @@ -47,9 +45,7 @@ def process_slack_alerting_variables( if "os.environ/" in webhook_urls: _env_value = get_secret(secret_name=webhook_urls) if not isinstance(_env_value, str): - raise ValueError( - f"Invalid webhook url value for: {webhook_urls}. Got type={type(_env_value)}" - ) + raise ValueError(f"Invalid webhook url value for: {webhook_urls}. Got type={type(_env_value)}") _webhook_value_str = _env_value else: _webhook_value_str = webhook_urls @@ -76,10 +72,7 @@ async def _add_langfuse_trace_id_to_alert( # Only run if langfuse is added as a callback ######################################################### - if ( - request_data is not None - and request_data.get("litellm_logging_obj", None) is not None - ): + if request_data is not None and request_data.get("litellm_logging_obj", None) is not None: trace_id: Optional[str] = None litellm_logging_obj: Logging = request_data["litellm_logging_obj"] @@ -89,9 +82,7 @@ async def _add_langfuse_trace_id_to_alert( break await asyncio.sleep(3) # wait 3s before retrying for trace id ######################################################### - langfuse_object = litellm_logging_obj._get_callback_object( - service_name="langfuse" - ) + langfuse_object = litellm_logging_obj._get_callback_object(service_name="langfuse") if langfuse_object is not None: base_url = langfuse_object.Langfuse.base_url return f"{base_url}/trace/{trace_id}" diff --git a/litellm/integrations/_types/open_inference.py b/litellm/integrations/_types/open_inference.py index 3404df7495f..8ce3ec6f492 100644 --- a/litellm/integrations/_types/open_inference.py +++ b/litellm/integrations/_types/open_inference.py @@ -73,15 +73,11 @@ class SpanAttributes: """ Number of tokens in the prompt. """ - LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = ( - "llm.token_count.prompt_details.cache_write" - ) + LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = "llm.token_count.prompt_details.cache_write" """ Number of tokens in the prompt that were written to cache. """ - LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = ( - "llm.token_count.prompt_details.cache_read" - ) + LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = "llm.token_count.prompt_details.cache_read" """ Number of tokens in the prompt that were read from cache. """ @@ -93,15 +89,11 @@ class SpanAttributes: """ Number of tokens in the completion. """ - LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = ( - "llm.token_count.completion_details.reasoning" - ) + LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = "llm.token_count.completion_details.reasoning" """ Number of tokens used for reasoning steps in the completion. """ - LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = ( - "llm.token_count.completion_details.audio" - ) + LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = "llm.token_count.completion_details.audio" """ The number of audio input tokens generated by the model """ diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py index 4f17806a6b7..c60e5cb0e2a 100644 --- a/litellm/integrations/agentops/agentops.py +++ b/litellm/integrations/agentops/agentops.py @@ -65,9 +65,7 @@ class AgentOps(OpenTelemetry): headers = f"Authorization=Bearer {jwt_token}" if jwt_token else None - otel_config = OpenTelemetryConfig( - exporter="otlp_http", endpoint=config.endpoint, headers=headers - ) + otel_config = OpenTelemetryConfig(exporter="otlp_http", endpoint=config.endpoint, headers=headers) # Initialize OpenTelemetry with our config super().__init__(config=otel_config, callback_name="agentops") diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 296bfb6fc85..1314fd82255 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -78,11 +78,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): # provider transform, where each tool_config point appends at most one # cachePoint to the tools. That block also counts toward Anthropic's # limit, so reserve a slot for it here to leave room. - reserved_blocks = ( - 1 - if any(p.get("location") == "tool_config" for p in remaining_points) - else 0 - ) + reserved_blocks = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0 processed_messages = self._apply_message_injections( points=message_points, @@ -111,10 +107,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): ``max_blocks`` is reached. Injection points are honored in config order, so earlier points win when slots are scarce. """ - used_blocks = sum( - AnthropicCacheControlHook._count_cache_control_blocks(msg) - for msg in messages - ) + used_blocks = sum(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages) limit_reached = False for point in points: @@ -122,27 +115,21 @@ class AnthropicCacheControlHook(CustomPromptManagement): limit_reached = True break - control: ChatCompletionCachedContent = point.get( - "control", None - ) or ChatCompletionCachedContent(type="ephemeral") + control: ChatCompletionCachedContent = point.get("control", None) or ChatCompletionCachedContent( + type="ephemeral" + ) - for target_index in AnthropicCacheControlHook._resolve_target_indices( - point=point, messages=messages - ): + for target_index in AnthropicCacheControlHook._resolve_target_indices(point=point, messages=messages): if used_blocks >= max_blocks: limit_reached = True break - if AnthropicCacheControlHook._message_has_cache_control( - messages[target_index] - ): + if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]): # Client already marked this message; don't overwrite it. continue - messages[target_index] = ( - AnthropicCacheControlHook._safe_insert_cache_control_in_message( - messages[target_index], control - ) + messages[target_index] = AnthropicCacheControlHook._safe_insert_cache_control_in_message( + messages[target_index], control ) used_blocks += 1 @@ -190,11 +177,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): # Case 2: Target by role targetted_role = point.get("role", None) if targetted_role is not None: - return [ - idx - for idx, msg in enumerate(messages) - if msg.get("role") == targetted_role - ] + return [idx for idx, msg in enumerate(messages) if msg.get("role") == targetted_role] return [] @@ -338,9 +321,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): _init_custom_logger_compatible_class, ) - if AnthropicCacheControlHook.should_use_anthropic_cache_control_hook( - non_default_params - ): + if AnthropicCacheControlHook.should_use_anthropic_cache_control_hook(non_default_params): return _init_custom_logger_compatible_class( logging_integration="anthropic_cache_control_hook", internal_usage_cache=None, diff --git a/litellm/integrations/argilla.py b/litellm/integrations/argilla.py index a362ce7e4d7..a86b6f9e388 100644 --- a/litellm/integrations/argilla.py +++ b/litellm/integrations/argilla.py @@ -47,12 +47,8 @@ class ArgillaLogger(CustomBatchLogger): **kwargs, ): if litellm.argilla_transformation_object is None: - raise Exception( - "'litellm.argilla_transformation_object' is required, to log your payload to Argilla." - ) - self.validate_argilla_transformation_object( - litellm.argilla_transformation_object - ) + raise Exception("'litellm.argilla_transformation_object' is required, to log your payload to Argilla.") + self.validate_argilla_transformation_object(litellm.argilla_transformation_object) self.argilla_transformation_object = litellm.argilla_transformation_object self.default_credentials = self.get_credentials_from_env( argilla_api_key=argilla_api_key, @@ -61,30 +57,21 @@ class ArgillaLogger(CustomBatchLogger): ) self.sampling_rate: float = ( float(os.getenv("ARGILLA_SAMPLING_RATE")) # type: ignore - if os.getenv("ARGILLA_SAMPLING_RATE") is not None - and os.getenv("ARGILLA_SAMPLING_RATE").strip().isdigit() # type: ignore + if os.getenv("ARGILLA_SAMPLING_RATE") is not None and os.getenv("ARGILLA_SAMPLING_RATE").strip().isdigit() # type: ignore else 1.0 ) - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) - _batch_size = ( - os.getenv("ARGILLA_BATCH_SIZE", None) or litellm.argilla_batch_size - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + _batch_size = os.getenv("ARGILLA_BATCH_SIZE", None) or litellm.argilla_batch_size if _batch_size: self.batch_size = int(_batch_size) asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) - def validate_argilla_transformation_object( - self, argilla_transformation_object: Dict[str, Any] - ): + def validate_argilla_transformation_object(self, argilla_transformation_object: Dict[str, Any]): if not isinstance(argilla_transformation_object, dict): - raise Exception( - "'argilla_transformation_object' must be a dictionary, to log your payload to Argilla." - ) + raise Exception("'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.") for v in argilla_transformation_object.values(): if v not in SUPPORTED_PAYLOAD_FIELDS: @@ -102,21 +89,11 @@ class ArgillaLogger(CustomBatchLogger): if _credentials_api_key is None: raise Exception("Invalid Argilla API Key given. _credentials_api_key=None.") - _credentials_base_url = ( - argilla_base_url - or os.getenv("ARGILLA_BASE_URL") - or "http://localhost:6900/" - ) + _credentials_base_url = argilla_base_url or os.getenv("ARGILLA_BASE_URL") or "http://localhost:6900/" if _credentials_base_url is None: - raise Exception( - "Invalid Argilla Base URL given. _credentials_base_url=None." - ) + raise Exception("Invalid Argilla Base URL given. _credentials_base_url=None.") - _credentials_dataset_name = ( - argilla_dataset_name - or os.getenv("ARGILLA_DATASET_NAME") - or "litellm-completion" - ) + _credentials_dataset_name = argilla_dataset_name or os.getenv("ARGILLA_DATASET_NAME") or "litellm-completion" if _credentials_dataset_name is None: raise Exception("Invalid Argilla Dataset give. Value=None.") else: @@ -138,19 +115,13 @@ class ArgillaLogger(CustomBatchLogger): ARGILLA_DATASET_NAME=_credentials_dataset_name, ) - def get_chat_messages( - self, payload: StandardLoggingPayload - ) -> List[Dict[str, Any]]: + def get_chat_messages(self, payload: StandardLoggingPayload) -> List[Dict[str, Any]]: payload_messages = payload.get("messages", None) if payload_messages is None: raise Exception("No chat messages found in payload.") - if ( - isinstance(payload_messages, list) - and len(payload_messages) > 0 - and isinstance(payload_messages[0], dict) - ): + if isinstance(payload_messages, list) and len(payload_messages) > 0 and isinstance(payload_messages[0], dict): return payload_messages elif isinstance(payload_messages, dict): return [payload_messages] @@ -166,20 +137,14 @@ class ArgillaLogger(CustomBatchLogger): if isinstance(response, str): return response elif isinstance(response, dict): - return ( - response.get("choices", [{}])[0].get("message", {}).get("content", "") - ) + return response.get("choices", [{}])[0].get("message", {}).get("content", "") else: raise Exception(f"Invalid response format: {response}") - def _prepare_log_data( - self, kwargs, response_obj, start_time, end_time - ) -> Optional[ArgillaItem]: + def _prepare_log_data(self, kwargs, response_obj, start_time, end_time) -> Optional[ArgillaItem]: try: # Ensure everything in the payload is converted to str - payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if payload is None: raise Exception("Error logging request payload. Payload=none.") @@ -220,13 +185,9 @@ class ArgillaLogger(CustomBatchLogger): ) if response.status_code >= 300: - verbose_logger.error( - f"Argilla Error: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"Argilla Error: {response.status_code} - {response.text}") else: - verbose_logger.debug( - f"Batch of {len(self.log_queue)} runs successfully created" - ) + verbose_logger.debug(f"Batch of {len(self.log_queue)} runs successfully created") self.log_queue.clear() except Exception: @@ -258,9 +219,7 @@ class ArgillaLogger(CustomBatchLogger): return self.log_queue.append(data) - verbose_logger.debug( - f"Langsmith, event added to queue. Will flush in {self.flush_interval} seconds..." - ) + verbose_logger.debug(f"Langsmith, event added to queue. Will flush in {self.flush_interval} seconds...") if len(self.log_queue) >= self.batch_size: self._send_batch() @@ -284,9 +243,7 @@ class ArgillaLogger(CustomBatchLogger): kwargs, response_obj, ) - payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) @@ -312,18 +269,14 @@ class ArgillaLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception: - verbose_logger.exception( - "Argilla Layer Error - error logging async success event." - ) + verbose_logger.exception("Argilla Layer Error - error logging async success event.") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): sampling_rate = self.sampling_rate random_sample = random.random() if random_sample > sampling_rate: verbose_logger.info( - "Skipping Langsmith logging. Sampling rate={}, random_sample={}".format( - sampling_rate, random_sample - ) + "Skipping Langsmith logging. Sampling rate={}, random_sample={}".format(sampling_rate, random_sample) ) return # Skip logging verbose_logger.info("Langsmith Failure Event Logging!") @@ -338,9 +291,7 @@ class ArgillaLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception: - verbose_logger.exception( - "Langsmith Layer Error - error logging async failure event." - ) + verbose_logger.exception("Langsmith Layer Error - error logging async failure event.") async def async_send_batch(self): """ @@ -378,13 +329,9 @@ class ArgillaLogger(CustomBatchLogger): response.raise_for_status() if response.status_code >= 300: - verbose_logger.error( - f"Argilla Error: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"Argilla Error: {response.status_code} - {response.text}") else: - verbose_logger.debug( - "Batch of %s runs successfully created", len(self.log_queue) - ) + verbose_logger.debug("Batch of %s runs successfully created", len(self.log_queue)) except httpx.HTTPStatusError: verbose_logger.exception("Argilla HTTP Error") except Exception: diff --git a/litellm/integrations/arize/__init__.py b/litellm/integrations/arize/__init__.py index bc06c7a51eb..ab2627801e6 100644 --- a/litellm/integrations/arize/__init__.py +++ b/litellm/integrations/arize/__init__.py @@ -13,22 +13,16 @@ from .arize_phoenix_prompt_manager import ArizePhoenixPromptManager global_arize_config: Optional[dict] = None -def prompt_initializer( - litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec" -) -> "CustomPromptManagement": +def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec") -> "CustomPromptManagement": """ Initialize a prompt from Arize Phoenix. """ - api_key = getattr(litellm_params, "api_key", None) or os.environ.get( - "PHOENIX_API_KEY" - ) + api_key = getattr(litellm_params, "api_key", None) or os.environ.get("PHOENIX_API_KEY") api_base = getattr(litellm_params, "api_base", None) prompt_id = getattr(litellm_params, "prompt_id", None) if not api_key or not api_base: - raise ValueError( - "api_key and api_base are required for Arize Phoenix prompt integration" - ) + raise ValueError("api_key and api_base are required for Arize Phoenix prompt integration") try: arize_prompt_manager = ArizePhoenixPromptManager( @@ -36,9 +30,7 @@ def prompt_initializer( "api_key": api_key, "api_base": api_base, "prompt_id": prompt_id, - **litellm_params.model_dump( - exclude={"api_key", "api_base", "prompt_id"} - ), + **litellm_params.model_dump(exclude={"api_key", "api_base", "prompt_id"}), }, ) diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 75710e10498..44fd7a0d01a 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -48,9 +48,7 @@ class ArizeOTELAttributes(BaseLLMObsOTELAttributes): for idx, msg in enumerate(messages): prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}" # Set the role per message. - safe_set_attribute( - span, f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", msg.get("role") - ) + safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", msg.get("role")) # Set the content per message. safe_set_attribute( span, @@ -164,9 +162,7 @@ def _set_audio_outputs(span: "Span", response_obj, audio_attrs, span_attrs): audio_transcript = audio_item.get("transcript") if audio_transcript: - safe_set_attribute( - span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript - ) + safe_set_attribute(span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript) def _set_embedding_outputs(span: "Span", response_obj, embedding_attrs, span_attrs): @@ -220,9 +216,7 @@ def _set_structured_outputs(span: "Span", response_obj, msg_attrs, span_attrs): message_content = getattr(first_content, "text", "") message_role = getattr(item, "role", "assistant") safe_set_attribute(span, span_attrs.OUTPUT_VALUE, message_content) - safe_set_attribute( - span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content - ) + safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content) safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_ROLE}", message_role) @@ -253,19 +247,11 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): if not usage: return - safe_set_attribute( - span, span_attrs.LLM_TOKEN_COUNT_TOTAL, _safe_get(usage, "total_tokens") - ) - completion_tokens = _safe_get(usage, "completion_tokens") or _safe_get( - usage, "output_tokens" - ) + safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_TOTAL, _safe_get(usage, "total_tokens")) + completion_tokens = _safe_get(usage, "completion_tokens") or _safe_get(usage, "output_tokens") if completion_tokens: - safe_set_attribute( - span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens - ) - prompt_tokens = _safe_get(usage, "prompt_tokens") or _safe_get( - usage, "input_tokens" - ) + safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens) + prompt_tokens = _safe_get(usage, "prompt_tokens") or _safe_get(usage, "input_tokens") if prompt_tokens: safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens) @@ -273,9 +259,7 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): # API (Usage) and in `output_tokens_details` for Responses API # (ResponseAPIUsage). Both nested objects may be plain Pydantic models # without `.get`. - token_details = _safe_get(usage, "completion_tokens_details") or _safe_get( - usage, "output_tokens_details" - ) + token_details = _safe_get(usage, "completion_tokens_details") or _safe_get(usage, "output_tokens_details") reasoning_tokens = _safe_get(token_details, "reasoning_tokens") if reasoning_tokens: safe_set_attribute( @@ -291,12 +275,8 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): # `cache_creation_input_tokens` # All emits are conditional, so when none of these fields exist (the # situation in the existing test fixtures) no extra attributes are set. - prompt_token_details = _safe_get(usage, "prompt_tokens_details") or _safe_get( - usage, "input_tokens_details" - ) - cache_read = _safe_get(prompt_token_details, "cached_tokens") or _safe_get( - usage, "cache_read_input_tokens" - ) + prompt_token_details = _safe_get(usage, "prompt_tokens_details") or _safe_get(usage, "input_tokens_details") + cache_read = _safe_get(prompt_token_details, "cached_tokens") or _safe_get(usage, "cache_read_input_tokens") if cache_read: safe_set_attribute( span, @@ -374,33 +354,24 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: ): return OpenInferenceSpanKindValues.LLM.value - if any( - keyword in lowered - for keyword in ("file", "batch", "container", "fine_tuning_job") - ): + if any(keyword in lowered for keyword in ("file", "batch", "container", "fine_tuning_job")): return OpenInferenceSpanKindValues.CHAIN.value return OpenInferenceSpanKindValues.UNKNOWN.value -def _set_tool_attributes( - span: "Span", optional_tools: Optional[list], metadata_tools: Optional[list] -): +def _set_tool_attributes(span: "Span", optional_tools: Optional[list], metadata_tools: Optional[list]): """set tool attributes on span from optional_params or tool call metadata""" if optional_tools: for idx, tool in enumerate(optional_tools): if not isinstance(tool, dict): continue - function = ( - tool.get("function") if isinstance(tool.get("function"), dict) else None - ) + function = tool.get("function") if isinstance(tool.get("function"), dict) else None if not function: continue tool_name = function.get("name") if tool_name: - safe_set_attribute( - span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name - ) + safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name) tool_description = function.get("description") if tool_description: safe_set_attribute( @@ -437,9 +408,7 @@ def _set_tool_attributes( ) -def set_attributes( - span: "Span", kwargs, response_obj, attributes: Type[BaseLLMObsOTELAttributes] -): +def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMObsOTELAttributes]): """ Populates span with OpenInference-compliant LLM attributes for Arize and Phoenix tracing. """ @@ -458,17 +427,11 @@ def set_attributes( try: optional_params = _sanitize_optional_params(kwargs.get("optional_params")) litellm_params = kwargs.get("litellm_params", {}) or {} - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - metadata = ( - standard_logging_payload.get("metadata") - if standard_logging_payload - else None - ) + metadata = standard_logging_payload.get("metadata") if standard_logging_payload else None _set_metadata_attributes(span, metadata, SpanAttributes) metadata_tools = _extract_metadata_tools(metadata) @@ -492,19 +455,13 @@ def set_attributes( _set_tool_attributes(span, optional_tools, metadata_tools) attributes.set_messages(span, kwargs) - model_params = ( - standard_logging_payload.get("model_parameters") - if standard_logging_payload - else None - ) + model_params = standard_logging_payload.get("model_parameters") if standard_logging_payload else None _set_model_params(span, model_params, SpanAttributes) _set_response_attributes(span=span, response_obj=response_obj_for_attrs) except Exception as e: - verbose_logger.error( - f"[Arize/Phoenix] Failed to set OpenInference span attributes: {e}" - ) + verbose_logger.error(f"[Arize/Phoenix] Failed to set OpenInference span attributes: {e}") if hasattr(span, "record_exception"): span.record_exception(e) @@ -562,9 +519,7 @@ def _set_request_attributes( if kwargs.get("model"): safe_set_attribute(span, span_attrs.LLM_MODEL_NAME, kwargs.get("model")) - safe_set_attribute( - span, "llm.request.type", standard_logging_payload.get("call_type") - ) + safe_set_attribute(span, "llm.request.type", standard_logging_payload.get("call_type")) safe_set_attribute( span, span_attrs.LLM_PROVIDER, @@ -572,19 +527,13 @@ def _set_request_attributes( ) if optional_params.get("max_tokens"): - safe_set_attribute( - span, "llm.request.max_tokens", optional_params.get("max_tokens") - ) + safe_set_attribute(span, "llm.request.max_tokens", optional_params.get("max_tokens")) if optional_params.get("temperature"): - safe_set_attribute( - span, "llm.request.temperature", optional_params.get("temperature") - ) + safe_set_attribute(span, "llm.request.temperature", optional_params.get("temperature")) if optional_params.get("top_p"): safe_set_attribute(span, "llm.request.top_p", optional_params.get("top_p")) - safe_set_attribute( - span, "llm.is_streaming", str(optional_params.get("stream", False)) - ) + safe_set_attribute(span, "llm.is_streaming", str(optional_params.get("stream", False))) if optional_params.get("user"): safe_set_attribute(span, "llm.user", optional_params.get("user")) @@ -599,9 +548,7 @@ def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) -> if not model_params: return - safe_set_attribute( - span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params) - ) + safe_set_attribute(span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params)) if model_params.get("user"): user_id = model_params.get("user") if user_id is not None: @@ -767,9 +714,7 @@ def _emit_message_tool_calls(span: "Span", prefix: str, message) -> None: continue tc_prefix = f"{prefix}.{MessageAttributes.MESSAGE_TOOL_CALLS}.{tc_idx}" if tc["id"]: - safe_set_attribute( - span, f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_ID}", tc["id"] - ) + safe_set_attribute(span, f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_ID}", tc["id"]) fn = tc["function"] if fn["name"]: safe_set_attribute( @@ -862,9 +807,7 @@ def _emit_input_message_extras(span: "Span", prefix: str, message: dict) -> None ) -def _set_session_and_user_attrs( - span: "Span", kwargs: dict, standard_logging_payload -) -> None: +def _set_session_and_user_attrs(span: "Span", kwargs: dict, standard_logging_payload) -> None: """Emit `SESSION_ID` / `USER_ID` / team metadata when source data exists. `SESSION_ID` is emitted only when an explicit end-user identifier exists @@ -970,11 +913,7 @@ def _maybe_normalize_passthrough( passthrough I/O (with central redaction) for free and this helper's `complete_input_dict` fallback can be deleted. See follow-up issue. """ - call_type = ( - standard_logging_payload.get("call_type") - if isinstance(standard_logging_payload, dict) - else None - ) + call_type = standard_logging_payload.get("call_type") if isinstance(standard_logging_payload, dict) else None if not _is_passthrough_call_type(call_type): return @@ -989,18 +928,12 @@ def _maybe_normalize_passthrough( # --- INPUT -------------------------------------------------------------- additional_args = kwargs.get("additional_args") or {} - complete_input_dict = ( - additional_args.get("complete_input_dict") - if isinstance(additional_args, dict) - else None - ) + complete_input_dict = additional_args.get("complete_input_dict") if isinstance(additional_args, dict) else None if isinstance(complete_input_dict, dict): _set_passthrough_input_attributes(span, complete_input_dict.get("messages")) # --- OUTPUT ------------------------------------------------------------- - parsed_response = _parse_passthrough_response( - raw_response_obj, coerced_response_obj, kwargs - ) + parsed_response = _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs) if not isinstance(parsed_response, dict): return @@ -1094,19 +1027,12 @@ def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs): candidates = [] if isinstance(coerced_response_obj, dict): candidates.append(coerced_response_obj) - if ( - isinstance(raw_response_obj, dict) - and raw_response_obj is not coerced_response_obj - ): + if isinstance(raw_response_obj, dict) and raw_response_obj is not coerced_response_obj: candidates.append(raw_response_obj) for candidate in candidates: # StandardPassThroughResponseObject wrapper: {"response": "..."}. - if ( - "response" in candidate - and "content" not in candidate - and "choices" not in candidate - ): + if "response" in candidate and "content" not in candidate and "choices" not in candidate: inner = candidate.get("response") if isinstance(inner, str): try: diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index fe2f9f41f1b..e5fdb231933 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -195,20 +195,14 @@ class ArizeLogger(OpenTelemetry): # the suggested param is `arize_space_key` ######################################################### if standard_callback_dynamic_params.get("arize_space_id"): - dynamic_headers["arize-space-id"] = standard_callback_dynamic_params.get( - "arize_space_id" - ) + dynamic_headers["arize-space-id"] = standard_callback_dynamic_params.get("arize_space_id") if standard_callback_dynamic_params.get("arize_space_key"): - dynamic_headers["arize-space-id"] = standard_callback_dynamic_params.get( - "arize_space_key" - ) + dynamic_headers["arize-space-id"] = standard_callback_dynamic_params.get("arize_space_key") ######################################################### # `api_key` handling ######################################################### if standard_callback_dynamic_params.get("arize_api_key"): - dynamic_headers["api_key"] = standard_callback_dynamic_params.get( - "arize_api_key" - ) + dynamic_headers["api_key"] = standard_callback_dynamic_params.get("arize_api_key") return dynamic_headers diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index d48dba8e7bb..db7aed1a71c 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -118,9 +118,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore try: provider.force_flush() except Exception as e: - verbose_logger.debug( - "ArizePhoenixLogger: TracerProvider force_flush failed: %s", e - ) + verbose_logger.debug("ArizePhoenixLogger: TracerProvider force_flush failed: %s", e) def _get_litellm_resource_for_project(self, project_name: str): """ @@ -149,9 +147,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore """Create a TracerProvider for *project_name* (caller holds no cache lock).""" from opentelemetry.sdk.trace import TracerProvider - provider = TracerProvider( - resource=self._get_litellm_resource_for_project(project_name) - ) + provider = TracerProvider(resource=self._get_litellm_resource_for_project(project_name)) provider.add_span_processor(self._shared_span_processor) return provider @@ -163,9 +159,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore with self._project_providers_lock: if project_name in self._project_providers: self._project_providers.move_to_end(project_name) - return self._project_providers[project_name].get_tracer( - LITELLM_TRACER_NAME - ) + return self._project_providers[project_name].get_tracer(LITELLM_TRACER_NAME) # OTELResourceDetector().detect() is synchronous; build outside the lock so # concurrent requests for other projects are not blocked on cache misses. @@ -174,9 +168,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore with self._project_providers_lock: if project_name in self._project_providers: self._project_providers.move_to_end(project_name) - return self._project_providers[project_name].get_tracer( - LITELLM_TRACER_NAME - ) + return self._project_providers[project_name].get_tracer(LITELLM_TRACER_NAME) if len(self._project_providers) >= _MAX_PROJECT_PROVIDERS: self._project_providers.popitem(last=False) @@ -241,14 +233,10 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore detection to route their telemetry into arbitrary Arize/Phoenix projects. """ litellm_params = kwargs.get("litellm_params") - return isinstance(litellm_params, dict) and bool( - litellm_params.get("proxy_server_request") - ) + return isinstance(litellm_params, dict) and bool(litellm_params.get("proxy_server_request")) @staticmethod - def _project_from_metadata_dict( - metadata: dict, metadata_key: str, *, proxy_mode: bool - ) -> Optional[str]: + def _project_from_metadata_dict(metadata: dict, metadata_key: str, *, proxy_mode: bool) -> Optional[str]: """ Read a Phoenix project field from proxy/SDK metadata. @@ -258,25 +246,19 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore """ auth_metadata = metadata.get("user_api_key_auth_metadata") if isinstance(auth_metadata, dict): - project = ArizePhoenixLogger._normalize_project_name( - auth_metadata.get(metadata_key) - ) + project = ArizePhoenixLogger._normalize_project_name(auth_metadata.get(metadata_key)) if project: return project if not proxy_mode: - return ArizePhoenixLogger._normalize_project_name( - metadata.get(metadata_key) - ) + return ArizePhoenixLogger._normalize_project_name(metadata.get(metadata_key)) return None @staticmethod def _metadata_project_from_kwargs(kwargs: dict, metadata_key: str) -> Optional[str]: proxy_mode = ArizePhoenixLogger._is_proxy_request(kwargs) for metadata in ArizePhoenixLogger._iter_metadata_dicts_from_kwargs(kwargs): - project = ArizePhoenixLogger._project_from_metadata_dict( - metadata, metadata_key, proxy_mode=proxy_mode - ) + project = ArizePhoenixLogger._project_from_metadata_dict(metadata, metadata_key, proxy_mode=proxy_mode) if project: return project return None @@ -290,21 +272,16 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore ``user_api_key_auth_metadata.phoenix_project_name``, env, then ``default``. SDK priority: request metadata fields, then env, then ``default``. """ - override = ArizePhoenixLogger._metadata_project_from_kwargs( - kwargs, "phoenix_project_name_override" - ) + override = ArizePhoenixLogger._metadata_project_from_kwargs(kwargs, "phoenix_project_name_override") if override: return override - phoenix_name = ArizePhoenixLogger._metadata_project_from_kwargs( - kwargs, "phoenix_project_name" - ) + phoenix_name = ArizePhoenixLogger._metadata_project_from_kwargs(kwargs, "phoenix_project_name") if phoenix_name: return phoenix_name env_name = ArizePhoenixLogger._normalize_project_name( - os.environ.get("PHOENIX_PROJECT_NAME") - or os.environ.get("ARIZE_PROJECT_NAME") + os.environ.get("PHOENIX_PROJECT_NAME") or os.environ.get("ARIZE_PROJECT_NAME") ) if env_name: return env_name @@ -335,11 +312,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore proxy_server_request = litellm_params.get("proxy_server_request", {}) or {} headers = proxy_server_request.get("headers", {}) or {} - traceparent_ctx = ( - self.get_traceparent_from_header(headers=headers) - if headers.get("traceparent") - else None - ) + traceparent_ctx = self.get_traceparent_from_header(headers=headers) if headers.get("traceparent") else None is_proxy_mode = bool(proxy_server_request) @@ -347,9 +320,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time")) parent_span = tracer.start_span( name="litellm_proxy_request", - start_time=( - self._to_ns(start_time_val) if start_time_val is not None else None - ), + start_time=(self._to_ns(start_time_val) if start_time_val is not None else None), context=traceparent_ctx, kind=self.span_kind.SERVER, ) @@ -359,14 +330,10 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore return traceparent_ctx, None def _handle_success(self, kwargs, response_obj, start_time, end_time): - self._handle_phoenix_trace( - kwargs, response_obj, start_time, end_time, success=True - ) + self._handle_phoenix_trace(kwargs, response_obj, start_time, end_time, success=True) def _handle_failure(self, kwargs, response_obj, start_time, end_time): - self._handle_phoenix_trace( - kwargs, response_obj, start_time, end_time, success=False - ) + self._handle_phoenix_trace(kwargs, response_obj, start_time, end_time, success=False) def _handle_phoenix_trace( self, @@ -402,9 +369,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore self._record_exception_on_span(span=span, kwargs=kwargs) if success: - self._maybe_log_raw_request( - kwargs, response_obj, start_time, end_time, span - ) + self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) span.end(end_time=self._to_ns(end_time)) self._create_guardrail_span(kwargs=kwargs, context=ctx) @@ -471,9 +436,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore if api_key is not None: otlp_auth_headers = f"Authorization=Bearer {api_key}" elif "app.phoenix.arize.com" in endpoint: - raise ValueError( - "PHOENIX_API_KEY must be set when using Phoenix Cloud (app.phoenix.arize.com)." - ) + raise ValueError("PHOENIX_API_KEY must be set when using Phoenix Cloud (app.phoenix.arize.com).") project_name = os.environ.get("PHOENIX_PROJECT_NAME") or "default" diff --git a/litellm/integrations/arize/arize_phoenix_client.py b/litellm/integrations/arize/arize_phoenix_client.py index 8c3c2a5ff0f..7c0715d2e1e 100644 --- a/litellm/integrations/arize/arize_phoenix_client.py +++ b/litellm/integrations/arize/arize_phoenix_client.py @@ -11,9 +11,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler def _sanitize_id(identifier: str) -> str: """Reject path traversal characters and URL-encode the identifier.""" if any(c in identifier for c in ("/", "\\", "#", "?")): - raise ValueError( - f"Invalid identifier {identifier!r}: contains disallowed characters" - ) + raise ValueError(f"Invalid identifier {identifier!r}: contains disallowed characters") if ".." in identifier: raise ValueError(f"Invalid identifier {identifier!r}: path traversal detected") return urllib.parse.quote(identifier, safe="") @@ -87,17 +85,11 @@ class ArizePhoenixClient: f"Access denied to prompt version '{prompt_version_id}'. Check your Arize Phoenix permissions." ) elif response.status_code == 401: - raise Exception( - "Authentication failed. Check your Arize Phoenix API key and permissions." - ) + raise Exception("Authentication failed. Check your Arize Phoenix API key and permissions.") else: - raise Exception( - f"Failed to fetch prompt version '{prompt_version_id}': {e}" - ) + raise Exception(f"Failed to fetch prompt version '{prompt_version_id}': {e}") else: - raise Exception( - f"Error fetching prompt version '{prompt_version_id}': {e}" - ) + raise Exception(f"Error fetching prompt version '{prompt_version_id}': {e}") def test_connection(self) -> bool: """ diff --git a/litellm/integrations/arize/arize_phoenix_prompt_manager.py b/litellm/integrations/arize/arize_phoenix_prompt_manager.py index df56d7bd391..4053b725a0f 100644 --- a/litellm/integrations/arize/arize_phoenix_prompt_manager.py +++ b/litellm/integrations/arize/arize_phoenix_prompt_manager.py @@ -44,9 +44,7 @@ class ArizePhoenixPromptTemplate: self.template_format = metadata.get("template_format", "MUSTACHE") def __repr__(self): - return ( - f"ArizePhoenixPromptTemplate(id='{self.template_id}', model='{self.model}')" - ) + return f"ArizePhoenixPromptTemplate(id='{self.template_id}', model='{self.model}')" class ArizePhoenixTemplateManager: @@ -71,9 +69,7 @@ class ArizePhoenixTemplateManager: self.api_base = api_base self.prompt_id = prompt_id self.prompts: Dict[str, ArizePhoenixPromptTemplate] = {} - self.arize_client = ArizePhoenixClient( - api_key=self.api_key, api_base=self.api_base - ) + self.arize_client = ArizePhoenixClient(api_key=self.api_key, api_base=self.api_base) # Templates fetched from Arize Phoenix come from external workspace # users; in a plain `Environment()` a malicious template could reach @@ -109,13 +105,9 @@ class ArizePhoenixTemplateManager: else: raise ValueError(f"Prompt version '{prompt_version_id}' not found") except Exception as e: - raise Exception( - f"Failed to load prompt version '{prompt_version_id}' from Arize Phoenix: {e}" - ) + raise Exception(f"Failed to load prompt version '{prompt_version_id}' from Arize Phoenix: {e}") - def _parse_prompt_data( - self, data: Dict[str, Any], prompt_version_id: str - ) -> ArizePhoenixPromptTemplate: + def _parse_prompt_data(self, data: Dict[str, Any], prompt_version_id: str) -> ArizePhoenixPromptTemplate: """Parse Arize Phoenix prompt data and extract messages and metadata.""" template_data = data.get("template", {}) messages = template_data.get("messages", []) @@ -154,9 +146,7 @@ class ArizePhoenixTemplateManager: metadata=metadata, ) - def render_template( - self, template_id: str, variables: Optional[Dict[str, Any]] = None - ) -> List[AllMessageValues]: + def render_template(self, template_id: str, variables: Optional[Dict[str, Any]] = None) -> List[AllMessageValues]: """Render a template with the given variables and return formatted messages.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") @@ -272,9 +262,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): raise ValueError(f"Prompt template '{prompt_id}' not found") # Render the template - rendered_messages = self.prompt_manager.render_template( - prompt_id, prompt_variables or {} - ) + rendered_messages = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) # Extract metadata metadata = { @@ -317,9 +305,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): try: # Get the rendered messages and metadata - rendered_messages, prompt_metadata = self.get_prompt_template( - prompt_id, prompt_variables - ) + rendered_messages, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Merge rendered messages with existing messages if rendered_messages: @@ -353,9 +339,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): # Log error but don't fail the call import litellm - litellm._logging.verbose_proxy_logger.error( - f"Error in Arize Phoenix prompt pre_call_hook: {e}" - ) + litellm._logging.verbose_proxy_logger.error(f"Error in Arize Phoenix prompt pre_call_hook: {e}") return messages, litellm_params def get_available_prompts(self) -> List[str]: @@ -408,9 +392,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): self.prompt_manager._load_prompt_from_arize(prompt_id) # Get the rendered messages and metadata - rendered_messages, prompt_metadata = self.get_prompt_template( - prompt_id, prompt_variables - ) + rendered_messages, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Extract model from metadata (if specified) template_model = prompt_metadata.get("model") diff --git a/litellm/integrations/athina.py b/litellm/integrations/athina.py index 49b9e9e6872..d1bf8e68624 100644 --- a/litellm/integrations/athina.py +++ b/litellm/integrations/athina.py @@ -12,10 +12,7 @@ class AthinaLogger: "athina-api-key": self.athina_api_key, "Content-Type": "application/json", } - self.athina_logging_url = ( - os.getenv("ATHINA_BASE_URL", "https://log.athina.ai") - + "/api/v1/log/inference" - ) + self.athina_logging_url = os.getenv("ATHINA_BASE_URL", "https://log.athina.ai") + "/api/v1/log/inference" self.additional_keys = [ "environment", "prompt_slug", @@ -42,9 +39,7 @@ class AthinaLogger: if "complete_streaming_response" in kwargs: # Log the completion response in streaming mode completion_response = kwargs["complete_streaming_response"] - response_json = ( - completion_response.model_dump() if completion_response else {} - ) + response_json = completion_response.model_dump() if completion_response else {} else: # Skip logging if the completion response is not available return @@ -56,30 +51,19 @@ class AthinaLogger: "request": kwargs, "response": response_json, "prompt_tokens": response_json.get("usage", {}).get("prompt_tokens"), - "completion_tokens": response_json.get("usage", {}).get( - "completion_tokens" - ), + "completion_tokens": response_json.get("usage", {}).get("completion_tokens"), "total_tokens": response_json.get("usage", {}).get("total_tokens"), } - if ( - type(end_time) is datetime.datetime - and type(start_time) is datetime.datetime - ): - data["response_time"] = int( - (end_time - start_time).total_seconds() * 1000 - ) + if type(end_time) is datetime.datetime and type(start_time) is datetime.datetime: + data["response_time"] = int((end_time - start_time).total_seconds() * 1000) if "messages" in kwargs: data["prompt"] = kwargs.get("messages", None) # Directly add tools or functions if present optional_params = kwargs.get("optional_params", {}) - data.update( - (k, v) - for k, v in optional_params.items() - if k in ["tools", "functions"] - ) + data.update((k, v) for k, v in optional_params.items() if k in ["tools", "functions"]) # Add additional metadata keys metadata = kwargs.get("litellm_params", {}).get("metadata", {}) @@ -93,13 +77,9 @@ class AthinaLogger: data=json.dumps(data, default=str), ) if response.status_code != 200: - print_verbose( - f"Athina Logger Error - {response.text}, {response.status_code}" - ) + print_verbose(f"Athina Logger Error - {response.text}, {response.status_code}") else: print_verbose(f"Athina Logger Succeeded - {response.text}") except Exception as e: - print_verbose( - f"Athina Logger Error - {e}, Stack trace: {traceback.format_exc()}" - ) + print_verbose(f"Athina Logger Error - {e}, Stack trace: {traceback.format_exc()}") pass diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 0cfd49cda37..182a2e185ef 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -63,32 +63,16 @@ class AzureSentinelLogger(CustomBatchLogger): audit_stream_name (str, optional): Stream name from DCR for audit logs. If not provided, audit logs use the standard stream name. """ - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - resolved_dcr_immutable_id = dcr_immutable_id or os.getenv( - "AZURE_SENTINEL_DCR_IMMUTABLE_ID" - ) - resolved_stream_name = ( - stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM" - ) + resolved_dcr_immutable_id = dcr_immutable_id or os.getenv("AZURE_SENTINEL_DCR_IMMUTABLE_ID") + resolved_stream_name = stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM" resolved_audit_stream_name = audit_stream_name or resolved_stream_name resolved_endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT") - resolved_tenant_id = ( - tenant_id - or os.getenv("AZURE_SENTINEL_TENANT_ID") - or os.getenv("AZURE_TENANT_ID") - ) - resolved_client_id = ( - client_id - or os.getenv("AZURE_SENTINEL_CLIENT_ID") - or os.getenv("AZURE_CLIENT_ID") - ) + resolved_tenant_id = tenant_id or os.getenv("AZURE_SENTINEL_TENANT_ID") or os.getenv("AZURE_TENANT_ID") + resolved_client_id = client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv("AZURE_CLIENT_ID") resolved_client_secret = ( - client_secret - or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") - or os.getenv("AZURE_CLIENT_SECRET") + client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET") ) if not resolved_dcr_immutable_id: @@ -144,9 +128,7 @@ class AzureSentinelLogger(CustomBatchLogger): self.audit_log_queue: List[StandardAuditLogPayload] = [] @staticmethod - def _build_api_endpoint( - endpoint: str, dcr_immutable_id: str, stream_name: str - ) -> str: + def _build_api_endpoint(endpoint: str, dcr_immutable_id: str, stream_name: str) -> str: return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01" async def _get_oauth_token(self) -> str: @@ -157,9 +139,7 @@ class AzureSentinelLogger(CustomBatchLogger): Bearer token string """ if ( - self.oauth_token - and self.oauth_token_expires_at - and time.time() < self.oauth_token_expires_at - 60 + self.oauth_token and self.oauth_token_expires_at and time.time() < self.oauth_token_expires_at - 60 ): # Refresh 60 seconds before expiry return self.oauth_token @@ -168,9 +148,7 @@ class AzureSentinelLogger(CustomBatchLogger): assert self.client_id is not None, "client_id is required" assert self.client_secret is not None, "client_secret is required" - token_url = ( - f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" - ) + token_url = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" token_data = { "client_id": self.client_id, @@ -186,9 +164,7 @@ class AzureSentinelLogger(CustomBatchLogger): ) if response.status_code != 200: - raise Exception( - f"Failed to get OAuth2 token: {response.status_code} - {response.text}" - ) + raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}") token_response = response.json() self.oauth_token = token_response.get("access_token") @@ -213,15 +189,11 @@ class AzureSentinelLogger(CustomBatchLogger): Raises a NON Blocking verbose_logger.exception if an error occurs """ try: - verbose_logger.debug( - "Azure Sentinel: Logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("Azure Sentinel: Logging - Enters logging function for model %s", kwargs) standard_logging_payload = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: - verbose_logger.warning( - "Azure Sentinel: standard_logging_object not found in kwargs" - ) + verbose_logger.warning("Azure Sentinel: standard_logging_object not found in kwargs") return self.log_queue.append(standard_logging_payload) @@ -230,9 +202,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"Azure Sentinel Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Azure Sentinel Layer Error - {str(e)}\n{traceback.format_exc()}") pass async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -254,9 +224,7 @@ class AzureSentinelLogger(CustomBatchLogger): standard_logging_payload = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: - verbose_logger.warning( - "Azure Sentinel: standard_logging_object not found in kwargs" - ) + verbose_logger.warning("Azure Sentinel: standard_logging_object not found in kwargs") return self.log_queue.append(standard_logging_payload) @@ -265,14 +233,10 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"Azure Sentinel Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Azure Sentinel Layer Error - {str(e)}\n{traceback.format_exc()}") pass - async def async_log_audit_log_event( - self, audit_log: StandardAuditLogPayload - ) -> None: + async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: """ Async log LiteLLM audit log events to Azure Sentinel. @@ -293,9 +257,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_audit_batch() except Exception as e: - verbose_logger.exception( - f"Azure Sentinel Audit Log Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Azure Sentinel Audit Log Layer Error - {str(e)}\n{traceback.format_exc()}") pass async def async_send_batch(self): @@ -331,9 +293,7 @@ class AzureSentinelLogger(CustomBatchLogger): if not log_queue: return - verbose_logger.debug( - "Azure Sentinel - about to flush %s %s", len(log_queue), log_type - ) + verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type) # Get OAuth2 token bearer_token = await self._get_oauth_token() @@ -349,9 +309,7 @@ class AzureSentinelLogger(CustomBatchLogger): } # Send the request - response = await self.async_httpx_client.post( - url=api_endpoint, data=body.encode("utf-8"), headers=headers - ) + response = await self.async_httpx_client.post(url=api_endpoint, data=body.encode("utf-8"), headers=headers) if response.status_code not in [200, 204]: verbose_logger.error( @@ -359,9 +317,7 @@ class AzureSentinelLogger(CustomBatchLogger): response.status_code, response.text, ) - raise Exception( - f"Failed to send logs to Azure Sentinel: {response.status_code} - {response.text}" - ) + raise Exception(f"Failed to send logs to Azure Sentinel: {response.status_code} - {response.text}") verbose_logger.debug( "Azure Sentinel: Response from API status_code: %s", @@ -369,9 +325,7 @@ class AzureSentinelLogger(CustomBatchLogger): ) except Exception as e: - verbose_logger.exception( - f"Azure Sentinel Error sending batch API - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Azure Sentinel Error sending batch API - {str(e)}\n{traceback.format_exc()}") finally: log_queue.clear() diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index b06fa13e918..5ccd1a86bff 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -24,42 +24,30 @@ class AzureBlobStorageLogger(CustomBatchLogger): **kwargs, ): try: - verbose_logger.debug( - "AzureBlobStorageLogger: in init azure blob storage logger" - ) + verbose_logger.debug("AzureBlobStorageLogger: in init azure blob storage logger") # Env Variables used for Azure Storage Authentication self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID") self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID") self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET") - self.azure_storage_account_key: Optional[str] = os.getenv( - "AZURE_STORAGE_ACCOUNT_KEY" - ) + self.azure_storage_account_key: Optional[str] = os.getenv("AZURE_STORAGE_ACCOUNT_KEY") # Required Env Variables for Azure Storage _azure_storage_account_name = os.getenv("AZURE_STORAGE_ACCOUNT_NAME") if not _azure_storage_account_name: - raise ValueError( - "Missing required environment variable: AZURE_STORAGE_ACCOUNT_NAME" - ) + raise ValueError("Missing required environment variable: AZURE_STORAGE_ACCOUNT_NAME") self.azure_storage_account_name: str = _azure_storage_account_name _azure_storage_file_system = os.getenv("AZURE_STORAGE_FILE_SYSTEM") if not _azure_storage_file_system: - raise ValueError( - "Missing required environment variable: AZURE_STORAGE_FILE_SYSTEM" - ) + raise ValueError("Missing required environment variable: AZURE_STORAGE_FILE_SYSTEM") self.azure_storage_file_system: str = _azure_storage_file_system self._service_client = None # Time that the azure service client expires, in order to reset the connection pool and keep it fresh self._service_client_timeout: Optional[float] = None # Internal variables used for Token based authentication - self.azure_auth_token: Optional[str] = ( - None # the Azure AD token to use for Azure Storage API requests - ) - self.token_expiry: Optional[datetime] = ( - None # the expiry time of the currentAzure AD token - ) + self.azure_auth_token: Optional[str] = None # the Azure AD token to use for Azure Storage API requests + self.token_expiry: Optional[datetime] = None # the expiry time of the currentAzure AD token asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() @@ -84,9 +72,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): "AzureBlobStorageLogger: Logging - Enters logging function for model %s", kwargs, ) - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") @@ -110,9 +96,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): "AzureBlobStorageLogger: Logging - Enters logging function for model %s", kwargs, ) - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") @@ -143,13 +127,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): await self.async_upload_payload_to_azure_blob_storage(payload=payload) except Exception as e: - verbose_logger.exception( - f"AzureBlobStorageLogger Error sending batch API - {str(e)}" - ) + verbose_logger.exception(f"AzureBlobStorageLogger Error sending batch API - {str(e)}") - async def async_upload_payload_to_azure_blob_storage( - self, payload: StandardLoggingPayload - ): + async def async_upload_payload_to_azure_blob_storage(self, payload: StandardLoggingPayload): """ Uploads the payload to Azure Blob Storage using a 3-step process: 1. Create file resource @@ -158,18 +138,12 @@ class AzureBlobStorageLogger(CustomBatchLogger): """ try: if self.azure_storage_account_key: - await self.upload_to_azure_data_lake_with_azure_account_key( - payload=payload - ) + await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload) else: # Get a valid token instead of always requesting a new one await self.set_valid_azure_ad_token() - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) - json_payload = ( - safe_dumps(payload) + "\n" - ) # Add newline for each log entry + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + json_payload = safe_dumps(payload) + "\n" # Add newline for each log entry payload_bytes = json_payload.encode("utf-8") filename = f"{payload.get('id') or str(uuid.uuid4())}.json" base_url = f"https://{self.azure_storage_account_name}.dfs.core.windows.net/{self.azure_storage_file_system}/{filename}" @@ -179,9 +153,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): await self._append_data(async_client, base_url, json_payload) await self._flush_data(async_client, base_url, len(payload_bytes)) - verbose_logger.debug( - f"Successfully uploaded log to Azure Blob Storage: {filename}" - ) + verbose_logger.debug(f"Successfully uploaded log to Azure Blob Storage: {filename}") except Exception as e: verbose_logger.exception(f"Error uploading to Azure Blob Storage: {str(e)}") @@ -203,9 +175,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): verbose_logger.exception(f"Error creating file resource: {str(e)}") raise - async def _append_data( - self, client: AsyncHTTPHandler, base_url: str, json_payload: str - ): + async def _append_data(self, client: AsyncHTTPHandler, base_url: str, json_payload: str): """Helper method to append data to the file""" try: verbose_logger.debug(f"Appending data to file: {base_url}") @@ -234,9 +204,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): "Content-Length": "0", "Authorization": f"Bearer {self.azure_auth_token}", } - response = await client.patch( - f"{base_url}?action=flush&position={position}", headers=headers - ) + response = await client.patch(f"{base_url}?action=flush&position={position}", headers=headers) response.raise_for_status() verbose_logger.debug("Successfully flushed data") except Exception as e: @@ -282,17 +250,11 @@ class AzureBlobStorageLogger(CustomBatchLogger): client_secret is not None, ) if tenant_id is None: - raise ValueError( - "Missing required environment variable: AZURE_STORAGE_TENANT_ID" - ) + raise ValueError("Missing required environment variable: AZURE_STORAGE_TENANT_ID") if client_id is None: - raise ValueError( - "Missing required environment variable: AZURE_STORAGE_CLIENT_ID" - ) + raise ValueError("Missing required environment variable: AZURE_STORAGE_CLIENT_ID") if client_secret is None: - raise ValueError( - "Missing required environment variable: AZURE_STORAGE_CLIENT_SECRET" - ) + raise ValueError("Missing required environment variable: AZURE_STORAGE_CLIENT_SECRET") token_provider = get_azure_ad_token_from_entra_id( tenant_id=tenant_id, @@ -331,11 +293,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): from azure.storage.filedatalake.aio import DataLakeServiceClient # expire old clients to recover from connection issues - if ( - self._service_client_timeout - and self._service_client - and self._service_client_timeout > time.time() - ): + if self._service_client_timeout and self._service_client and self._service_client_timeout > time.time(): await self._service_client.close() self._service_client = None if not self._service_client: @@ -346,9 +304,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client_timeout = time.time() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS return self._service_client - async def upload_to_azure_data_lake_with_azure_account_key( - self, payload: StandardLoggingPayload - ): + async def upload_to_azure_data_lake_with_azure_account_key(self, payload: StandardLoggingPayload): """ Uploads the payload to Azure Data Lake using the Azure SDK @@ -359,9 +315,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): service_client = await self.get_service_client() # Get file system client - file_system_client = service_client.get_file_system_client( - file_system=self.azure_storage_file_system - ) + file_system_client = service_client.get_file_system_client(file_system=self.azure_storage_file_system) try: # Create directory with today's date @@ -391,9 +345,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): # Flush the content to finalize the file await file_client.flush_data(position=len(content), offset=0) - verbose_logger.debug( - f"Successfully uploaded and wrote to {today}/{file_name}" - ) + verbose_logger.debug(f"Successfully uploaded and wrote to {today}/{file_name}") except Exception as e: verbose_logger.exception(f"Error occurred: {str(e)}") diff --git a/litellm/integrations/bitbucket/__init__.py b/litellm/integrations/bitbucket/__init__.py index 111d38f78a4..2b9bd568e32 100644 --- a/litellm/integrations/bitbucket/__init__.py +++ b/litellm/integrations/bitbucket/__init__.py @@ -29,9 +29,7 @@ def set_global_bitbucket_config(config: dict) -> None: litellm.global_bitbucket_config = config # type: ignore -def prompt_initializer( - litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec" -) -> "CustomPromptManagement": +def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec") -> "CustomPromptManagement": """ Initialize a prompt from a BitBucket repository. """ @@ -39,9 +37,7 @@ def prompt_initializer( prompt_id = getattr(litellm_params, "prompt_id", None) if not bitbucket_config: - raise ValueError( - "bitbucket_config is required for BitBucket prompt integration" - ) + raise ValueError("bitbucket_config is required for BitBucket prompt integration") try: bitbucket_prompt_manager = BitBucketPromptManager( diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index e742cc14b7d..c02d56811a7 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -12,15 +12,11 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler def _sanitize_file_path(file_path: str) -> str: """Reject path traversal and URL-encode each path segment.""" if "#" in file_path or "?" in file_path: - raise ValueError( - f"Invalid file path {file_path!r}: contains URL special characters" - ) + raise ValueError(f"Invalid file path {file_path!r}: contains URL special characters") parts = file_path.split("/") for part in parts: if part == "..": - raise ValueError( - f"Invalid file path {file_path!r}: path traversal detected" - ) + raise ValueError(f"Invalid file path {file_path!r}: path traversal detected") return "/".join(urllib.parse.quote(part, safe="") for part in parts) @@ -115,17 +111,13 @@ class BitBucketClient: f"Access denied to file '{file_path}'. Check your BitBucket permissions for workspace '{self.workspace}' and repository '{self.repository}'." ) elif e.response.status_code == 401: - raise Exception( - "Authentication failed. Check your BitBucket access token and permissions." - ) + raise Exception("Authentication failed. Check your BitBucket access token and permissions.") else: raise Exception(f"Failed to fetch file '{file_path}': {e}") else: raise Exception(f"Error fetching file '{file_path}': {e}") - def list_files( - self, directory_path: str = "", file_extension: str = ".prompt" - ) -> List[str]: + def list_files(self, directory_path: str = "", file_extension: str = ".prompt") -> List[str]: """ List files in a directory with a specific extension. @@ -164,9 +156,7 @@ class BitBucketClient: f"Access denied to directory '{directory_path}'. Check your BitBucket permissions for workspace '{self.workspace}' and repository '{self.repository}'." ) elif e.response.status_code == 401: - raise Exception( - "Authentication failed. Check your BitBucket access token and permissions." - ) + raise Exception("Authentication failed. Check your BitBucket access token and permissions.") else: raise Exception(f"Failed to list files in '{directory_path}': {e}") else: diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py index 844fa9f38cb..6dca4d76c04 100644 --- a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -44,9 +44,7 @@ class BitBucketPromptTemplate: self.temperature = metadata.get("temperature") self.max_tokens = metadata.get("max_tokens") self.input_schema = metadata.get("input", {}).get("schema", {}) - self.optional_params = { - k: v for k, v in metadata.items() if k not in ["model", "input", "content"] - } + self.optional_params = {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]} def __repr__(self): return f"BitBucketPromptTemplate(id='{self.template_id}', model='{self.model}')" @@ -101,9 +99,7 @@ class BitBucketTemplateManager: """Load a specific .prompt file from BitBucket.""" try: # Fetch the .prompt file from BitBucket - prompt_content = self.bitbucket_client.get_file_content( - f"{prompt_id}.prompt" - ) + prompt_content = self.bitbucket_client.get_file_content(f"{prompt_id}.prompt") if prompt_content: template = self._parse_prompt_file(prompt_content, prompt_id) @@ -111,9 +107,7 @@ class BitBucketTemplateManager: except Exception as e: raise Exception(f"Failed to load prompt '{prompt_id}' from BitBucket: {e}") - def _parse_prompt_file( - self, content: str, prompt_id: str - ) -> BitBucketPromptTemplate: + def _parse_prompt_file(self, content: str, prompt_id: str) -> BitBucketPromptTemplate: """Parse a .prompt file content and extract metadata and template.""" # Split frontmatter and content if content.startswith("---"): @@ -168,9 +162,7 @@ class BitBucketTemplateManager: result[key] = value.strip("\"'") return result - def render_template( - self, template_id: str, variables: Optional[Dict[str, Any]] = None - ) -> str: + def render_template(self, template_id: str, variables: Optional[Dict[str, Any]] = None) -> str: """Render a template with the given variables.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") @@ -259,9 +251,7 @@ class BitBucketPromptManager(CustomPromptManagement): raise ValueError(f"Prompt template '{prompt_id}' not found") # Render the template - rendered_prompt = self.prompt_manager.render_template( - prompt_id, prompt_variables or {} - ) + rendered_prompt = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) # Extract metadata metadata = { @@ -291,9 +281,7 @@ class BitBucketPromptManager(CustomPromptManagement): try: # Get the rendered prompt and metadata - rendered_prompt, prompt_metadata = self.get_prompt_template( - prompt_id, prompt_variables - ) + rendered_prompt, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Parse the rendered prompt into messages parsed_messages = self._parse_prompt_to_messages(rendered_prompt) @@ -332,9 +320,7 @@ class BitBucketPromptManager(CustomPromptManagement): # Log error but don't fail the call import litellm - litellm._logging.verbose_proxy_logger.error( - f"Error in BitBucket prompt pre_call_hook: {e}" - ) + litellm._logging.verbose_proxy_logger.error(f"Error in BitBucket prompt pre_call_hook: {e}") return messages, litellm_params def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]: @@ -389,9 +375,7 @@ class BitBucketPromptManager(CustomPromptManagement): # Add the last message if current_role and current_content: - messages.append( - {"role": current_role, "content": "\n".join(current_content).strip()} - ) + messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # If no role indicators found, treat as a single user message if not messages and prompt_content.strip(): @@ -466,9 +450,7 @@ class BitBucketPromptManager(CustomPromptManagement): self.prompt_manager._load_prompt_from_bitbucket(prompt_id) # Get the rendered prompt and metadata - rendered_prompt, prompt_metadata = self.get_prompt_template( - prompt_id, prompt_variables - ) + rendered_prompt, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Convert rendered content to chat messages messages = self._parse_prompt_to_messages(rendered_prompt) diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index a69785ba2e3..686c37d3e17 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -34,16 +34,12 @@ def get_utc_datetime(): class BraintrustLogger(CustomLogger): - def __init__( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> None: + def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> None: super().__init__() self.is_mock_mode = should_use_braintrust_mock() if self.is_mock_mode: create_mock_braintrust_client() - verbose_logger.info( - "[BRAINTRUST MOCK] Braintrust logger initialized in mock mode" - ) + verbose_logger.info("[BRAINTRUST MOCK] Braintrust logger initialized in mock mode") self.validate_environment(api_key=api_key) self.api_base = api_base or os.getenv("BRAINTRUST_API_BASE") or API_BASE self.default_project_id = None @@ -52,12 +48,8 @@ class BraintrustLogger(CustomLogger): "Authorization": "Bearer " + self.api_key, "Content-Type": "application/json", } - self._project_id_cache: Dict[ - str, str - ] = {} # Cache mapping project names to IDs - self.global_braintrust_http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self._project_id_cache: Dict[str, str] = {} # Cache mapping project names to IDs + self.global_braintrust_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.global_braintrust_sync_http_handler = HTTPHandler() def validate_environment(self, api_key: Optional[str]): @@ -143,23 +135,16 @@ class BraintrustLogger(CustomLogger): output = None choices = [] if response_obj is not None and ( - kwargs.get("call_type", None) == "embedding" - or isinstance(response_obj, litellm.EmbeddingResponse) + kwargs.get("call_type", None) == "embedding" or isinstance(response_obj, litellm.EmbeddingResponse) ): output = None - elif response_obj is not None and isinstance( - response_obj, litellm.ModelResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.ModelResponse): output = response_obj["choices"][0]["message"].json() choices = response_obj["choices"] - elif response_obj is not None and isinstance( - response_obj, litellm.TextCompletionResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.TextCompletionResponse): output = response_obj.choices[0].text choices = response_obj.choices - elif response_obj is not None and isinstance( - response_obj, litellm.ImageResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.ImageResponse): output = response_obj["data"] litellm_params = kwargs.get("litellm_params", {}) or {} @@ -169,9 +154,7 @@ class BraintrustLogger(CustomLogger): project_id = dynamic_metadata.get("project_id") if project_id is None: project_name = dynamic_metadata.get("project_name") - project_id = ( - self.get_project_id_sync(project_name) if project_name else None - ) + project_id = self.get_project_id_sync(project_name) if project_name else None if project_id is None: if self.default_project_id is None: @@ -206,8 +189,7 @@ class BraintrustLogger(CustomLogger): "completion_tokens": usage_obj.completion_tokens, "total_tokens": usage_obj.total_tokens, "total_cost": cost, - "time_to_first_token": end_time.timestamp() - - start_time.timestamp(), + "time_to_first_token": end_time.timestamp() - start_time.timestamp(), "start": start_time.timestamp(), "end": end_time.timestamp(), } @@ -278,23 +260,16 @@ class BraintrustLogger(CustomLogger): output = None choices = [] if response_obj is not None and ( - kwargs.get("call_type", None) == "embedding" - or isinstance(response_obj, litellm.EmbeddingResponse) + kwargs.get("call_type", None) == "embedding" or isinstance(response_obj, litellm.EmbeddingResponse) ): output = None - elif response_obj is not None and isinstance( - response_obj, litellm.ModelResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.ModelResponse): output = response_obj["choices"][0]["message"].json() choices = response_obj["choices"] - elif response_obj is not None and isinstance( - response_obj, litellm.TextCompletionResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.TextCompletionResponse): output = response_obj.choices[0].text choices = response_obj.choices - elif response_obj is not None and isinstance( - response_obj, litellm.ImageResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.ImageResponse): output = response_obj["data"] litellm_params = kwargs.get("litellm_params", {}) @@ -304,11 +279,7 @@ class BraintrustLogger(CustomLogger): project_id = dynamic_metadata.get("project_id") if project_id is None: project_name = dynamic_metadata.get("project_name") - project_id = ( - await self.get_project_id_async(project_name) - if project_name - else None - ) + project_id = await self.get_project_id_async(project_name) if project_name else None if project_id is None: if self.default_project_id is None: @@ -350,14 +321,8 @@ class BraintrustLogger(CustomLogger): api_call_start_time = kwargs.get("api_call_start_time") completion_start_time = kwargs.get("completion_start_time") - if ( - api_call_start_time is not None - and completion_start_time is not None - ): - metrics["time_to_first_token"] = ( - completion_start_time.timestamp() - - api_call_start_time.timestamp() - ) + if api_call_start_time is not None and completion_start_time is not None: + metrics["time_to_first_token"] = completion_start_time.timestamp() - api_call_start_time.timestamp() # Allow metadata override for span name span_name = dynamic_metadata.get("span_name", "Chat Completion") diff --git a/litellm/integrations/braintrust_mock_client.py b/litellm/integrations/braintrust_mock_client.py index 1af14deeab6..e2b732d6e9c 100644 --- a/litellm/integrations/braintrust_mock_client.py +++ b/litellm/integrations/braintrust_mock_client.py @@ -156,11 +156,7 @@ def create_mock_braintrust_client(): # This is required for async calls to be mocked create_mock_braintrust_factory_client() - verbose_logger.debug( - f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms" - ) - verbose_logger.debug( - "[BRAINTRUST MOCK] Braintrust mock client initialization complete" - ) + verbose_logger.debug(f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms") + verbose_logger.debug("[BRAINTRUST MOCK] Braintrust mock client initialization complete") _mocks_initialized = True diff --git a/litellm/integrations/cloudzero/cloudzero.py b/litellm/integrations/cloudzero/cloudzero.py index 8decd4ef23f..121b1dc6967 100644 --- a/litellm/integrations/cloudzero/cloudzero.py +++ b/litellm/integrations/cloudzero/cloudzero.py @@ -60,15 +60,11 @@ class CloudZeroLogger(CustomLogger): # if using redis, ensure only one pod exports the data at a time if pod_lock_manager and pod_lock_manager.redis_cache: - if await pod_lock_manager.acquire_lock( - cronjob_id=CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME - ): + if await pod_lock_manager.acquire_lock(cronjob_id=CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME): try: await self._hourly_usage_data_export() finally: - await pod_lock_manager.release_lock( - cronjob_id=CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME - ) + await pod_lock_manager.release_lock(cronjob_id=CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME) else: # if not using redis, export the data directly await self._hourly_usage_data_export() @@ -86,9 +82,7 @@ class CloudZeroLogger(CustomLogger): current_time_utc = datetime.now(timezone.utc) # Mitigates the possibility of missing spend if an hour is skipped due to a restart in an ephemeral environment - one_hour_ago_utc = current_time_utc - timedelta( - minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES * 2 - ) + one_hour_ago_utc = current_time_utc - timedelta(minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES * 2) await self.export_usage_data( limit=CLOUDZERO_MAX_FETCHED_DATA_RECORDS, operation="replace_hourly", @@ -130,9 +124,7 @@ class CloudZeroLogger(CustomLogger): # Initialize database connection and load data database = LiteLLMDatabase() verbose_logger.debug("CloudZero Logger: Loading usage data from database") - data = await database.get_usage_data( - limit=limit, start_time_utc=start_time_utc, end_time_utc=end_time_utc - ) + data = await database.get_usage_data(limit=limit, start_time_utc=start_time_utc, end_time_utc=end_time_utc) if data.is_empty(): verbose_logger.debug("CloudZero Logger: No usage data found to export") @@ -145,9 +137,7 @@ class CloudZeroLogger(CustomLogger): cbf_data = transformer.transform(data) if cbf_data.is_empty(): - verbose_logger.warning( - "CloudZero Logger: No valid data after transformation" - ) + verbose_logger.warning("CloudZero Logger: No valid data after transformation") return # Send data to CloudZero @@ -157,19 +147,13 @@ class CloudZeroLogger(CustomLogger): user_timezone=self.timezone, ) - verbose_logger.debug( - f"CloudZero Logger: Transmitting {len(cbf_data)} records to CloudZero" - ) + verbose_logger.debug(f"CloudZero Logger: Transmitting {len(cbf_data)} records to CloudZero") streamer.send_batched(cbf_data, operation=operation) - verbose_logger.debug( - f"CloudZero Logger: Successfully exported {len(cbf_data)} records to CloudZero" - ) + verbose_logger.debug(f"CloudZero Logger: Successfully exported {len(cbf_data)} records to CloudZero") except Exception as e: - verbose_logger.error( - f"CloudZero Logger: Error exporting usage data: {str(e)}" - ) + verbose_logger.error(f"CloudZero Logger: Error exporting usage data: {str(e)}") raise async def dry_run_export_usage_data(self, limit: Optional[int] = 10000): @@ -207,9 +191,7 @@ class CloudZeroLogger(CustomLogger): }, } - verbose_logger.debug( - f"CloudZero Dry Run: Processing {len(data)} records..." - ) + verbose_logger.debug(f"CloudZero Dry Run: Processing {len(data)} records...") # Convert usage data to dict format for response usage_data_sample = data.head(50).to_dicts() # Return first 50 rows @@ -219,21 +201,15 @@ class CloudZeroLogger(CustomLogger): cbf_data = transformer.transform(data) if cbf_data.is_empty(): - verbose_logger.warning( - "CloudZero Dry Run: No valid data after transformation" - ) + verbose_logger.warning("CloudZero Dry Run: No valid data after transformation") return { "usage_data": usage_data_sample, "cbf_data": [], "summary": { "total_records": len(usage_data_sample), - "total_cost": sum( - row.get("spend", 0) for row in usage_data_sample - ), + "total_cost": sum(row.get("spend", 0) for row in usage_data_sample), "total_tokens": sum( - row.get("prompt_tokens", 0) - + row.get("completion_tokens", 0) - for row in usage_data_sample + row.get("prompt_tokens", 0) + row.get("completion_tokens", 0) for row in usage_data_sample ), "unique_accounts": 0, "unique_services": 0, @@ -246,26 +222,14 @@ class CloudZeroLogger(CustomLogger): # Calculate summary statistics total_cost = sum(record.get("cost/cost", 0) for record in cbf_data_dict) unique_accounts = len( - set( - record.get("resource/account", "") - for record in cbf_data_dict - if record.get("resource/account") - ) + set(record.get("resource/account", "") for record in cbf_data_dict if record.get("resource/account")) ) unique_services = len( - set( - record.get("resource/service", "") - for record in cbf_data_dict - if record.get("resource/service") - ) - ) - total_tokens = sum( - record.get("usage/amount", 0) for record in cbf_data_dict + set(record.get("resource/service", "") for record in cbf_data_dict if record.get("resource/service")) ) + total_tokens = sum(record.get("usage/amount", 0) for record in cbf_data_dict) - verbose_logger.debug( - f"CloudZero Logger: Dry run completed for {len(cbf_data)} records" - ) + verbose_logger.debug(f"CloudZero Logger: Dry run completed for {len(cbf_data)} records") return { "usage_data": usage_data_sample, @@ -296,32 +260,22 @@ class CloudZeroLogger(CustomLogger): console.print("[yellow]No CBF data to display[/yellow]") return - console.print( - f"\n[bold green]💰 CloudZero CBF Transformed Data ({len(cbf_data)} records)[/bold green]" - ) + console.print(f"\n[bold green]💰 CloudZero CBF Transformed Data ({len(cbf_data)} records)[/bold green]") # Convert to dicts for easier processing records = cbf_data.to_dicts() # Create main CBF table - cbf_table = Table( - show_header=True, header_style="bold cyan", box=SIMPLE, padding=(0, 1) - ) + cbf_table = Table(show_header=True, header_style="bold cyan", box=SIMPLE, padding=(0, 1)) cbf_table.add_column("time/usage_start", style="blue", no_wrap=False) cbf_table.add_column("cost/cost", style="green", justify="right", no_wrap=False) - cbf_table.add_column( - "entity_type", style="magenta", justify="right", no_wrap=False - ) - cbf_table.add_column( - "entity_id", style="magenta", justify="right", no_wrap=False - ) + cbf_table.add_column("entity_type", style="magenta", justify="right", no_wrap=False) + cbf_table.add_column("entity_id", style="magenta", justify="right", no_wrap=False) cbf_table.add_column("team_id", style="cyan", no_wrap=False) cbf_table.add_column("team_alias", style="cyan", no_wrap=False) cbf_table.add_column("user_email", style="cyan", no_wrap=False) cbf_table.add_column("api_key_alias", style="yellow", no_wrap=False) - cbf_table.add_column( - "usage/amount", style="yellow", justify="right", no_wrap=False - ) + cbf_table.add_column("usage/amount", style="yellow", justify="right", no_wrap=False) cbf_table.add_column("resource/id", style="magenta", no_wrap=False) cbf_table.add_column("resource/service", style="cyan", no_wrap=False) cbf_table.add_column("resource/account", style="white", no_wrap=False) @@ -364,18 +318,10 @@ class CloudZeroLogger(CustomLogger): # Show summary statistics total_cost = sum(record.get("cost/cost", 0) for record in records) unique_accounts = len( - set( - record.get("resource/account", "") - for record in records - if record.get("resource/account") - ) + set(record.get("resource/account", "") for record in records if record.get("resource/account")) ) unique_services = len( - set( - record.get("resource/service", "") - for record in records - if record.get("resource/service") - ) + set(record.get("resource/service", "") for record in records if record.get("resource/service")) ) # Count total tokens from usage metrics @@ -388,9 +334,7 @@ class CloudZeroLogger(CustomLogger): console.print(f" Unique Accounts: {unique_accounts}") console.print(f" Unique Services: {unique_services}") - console.print( - "\n[dim]💡 This is the CloudZero CBF format ready for AnyCost ingestion[/dim]" - ) + console.print("\n[dim]💡 This is the CloudZero CBF format ready for AnyCost ingestion[/dim]") @staticmethod async def init_cloudzero_background_job(scheduler: AsyncIOScheduler): @@ -402,10 +346,8 @@ class CloudZeroLogger(CustomLogger): from litellm.constants import CLOUDZERO_EXPORT_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger - prometheus_loggers: List[CustomLogger] = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=CloudZeroLogger - ) + prometheus_loggers: List[CustomLogger] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=CloudZeroLogger ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s cloudzero loggers", len(prometheus_loggers)) diff --git a/litellm/integrations/cloudzero/cz_resource_names.py b/litellm/integrations/cloudzero/cz_resource_names.py index 20862c1c7ec..15cb66002f7 100644 --- a/litellm/integrations/cloudzero/cz_resource_names.py +++ b/litellm/integrations/cloudzero/cz_resource_names.py @@ -30,9 +30,7 @@ class CZEntityType(str, Enum): class CZRNGenerator: """Generate CloudZero Resource Names (CZRNs) for LiteLLM resources.""" - CZRN_REGEX = re.compile( - r"^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$" - ) + CZRN_REGEX = re.compile(r"^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$") def __init__(self): """Initialize CZRN generator.""" @@ -138,9 +136,7 @@ class CZRNGenerator: return normalized return provider_map.get(normalized, normalized) - def _normalize_component( - self, component: str, allow_uppercase: bool = False - ) -> str: + def _normalize_component(self, component: str, allow_uppercase: bool = False) -> str: """Normalize a CZRN component to meet format requirements.""" if not component: return "unknown" diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index d673536e72d..47d6f7474a2 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -30,9 +30,7 @@ from rich.console import Console class CloudZeroStreamer: """Stream CBF data to CloudZero AnyCost API with proper batching and timezone handling.""" - def __init__( - self, api_key: str, connection_id: str, user_timezone: Optional[str] = None - ): + def __init__(self, api_key: str, connection_id: str, user_timezone: Optional[str] = None): """Initialize CloudZero streamer with credentials.""" self.api_key = api_key self.connection_id = connection_id @@ -45,16 +43,12 @@ class CloudZeroStreamer: try: self.user_timezone = zoneinfo.ZoneInfo(user_timezone) except zoneinfo.ZoneInfoNotFoundError: - self.console.print( - f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]" - ) + self.console.print(f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]") self.user_timezone = timezone.utc else: self.user_timezone = timezone.utc - def send_batched( - self, data: pl.DataFrame, operation: str = "replace_hourly" - ) -> None: + def send_batched(self, data: pl.DataFrame, operation: str = "replace_hourly") -> None: """Send CBF data in daily batches to CloudZero AnyCost API.""" if data.is_empty(): self.console.print("[yellow]No data to send to CloudZero[/yellow]") @@ -67,9 +61,7 @@ class CloudZeroStreamer: self.console.print("[yellow]No valid daily batches to send[/yellow]") return - self.console.print( - f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]" - ) + self.console.print(f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]") for batch_date, batch_data in daily_batches.items(): self._send_daily_batch(batch_date, batch_data, operation) @@ -80,9 +72,7 @@ class CloudZeroStreamer: # Ensure we have the required columns if "time/usage_start" not in data.columns: - self.console.print( - "[red]Error: Missing 'time/usage_start' column for date grouping[/red]" - ) + self.console.print("[red]Error: Missing 'time/usage_start' column for date grouping[/red]") return {} timestamp_str: Optional[str] = None @@ -103,17 +93,11 @@ class CloudZeroStreamer: daily_batches[batch_date].append(row) except Exception as e: - self.console.print( - f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]" - ) + self.console.print(f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]") continue # Convert lists back to DataFrames - return { - date_key: pl.DataFrame(records) - for date_key, records in daily_batches.items() - if records - } + return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records} def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime: """Parse timestamp string and convert to UTC.""" @@ -164,9 +148,7 @@ class CloudZeroStreamer: except ValueError as e: raise ValueError(f"Could not parse timestamp '{timestamp_str}': {e}") - def _send_daily_batch( - self, batch_date: str, batch_data: pl.DataFrame, operation: str - ) -> None: + def _send_daily_batch(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> None: """Send a single daily batch to CloudZero API.""" if batch_data.is_empty(): return @@ -184,9 +166,7 @@ class CloudZeroStreamer: try: with httpx.Client(timeout=30.0) as client: - self.console.print( - f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]" - ) + self.console.print(f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]") response = client.post(url, headers=headers, json=payload) response.raise_for_status() @@ -196,9 +176,7 @@ class CloudZeroStreamer: ) except httpx.RequestError as e: - self.console.print( - f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]" - ) + self.console.print(f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]") raise except httpx.HTTPStatusError as e: self.console.print( @@ -206,9 +184,7 @@ class CloudZeroStreamer: ) raise - def _prepare_batch_payload( - self, batch_date: str, batch_data: pl.DataFrame, operation: str - ) -> dict[str, Any]: + def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> dict[str, Any]: """Prepare batch payload according to CloudZero AnyCost API format.""" # Convert batch_date to month for the API (YYYY-MM format) try: @@ -229,9 +205,7 @@ class CloudZeroStreamer: return payload - def _convert_cbf_to_api_format( - self, row: dict[str, Any] - ) -> Optional[dict[str, Any]]: + def _convert_cbf_to_api_format(self, row: dict[str, Any]) -> Optional[dict[str, Any]]: """Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them.""" try: # CloudZero expects CBF format field names directly, not converted names @@ -253,16 +227,12 @@ class CloudZeroStreamer: # Ensure timestamp is in UTC format if "time/usage_start" in api_record: - api_record["time/usage_start"] = self._ensure_utc_timestamp( - api_record["time/usage_start"] - ) + api_record["time/usage_start"] = self._ensure_utc_timestamp(api_record["time/usage_start"]) return api_record except Exception as e: - self.console.print( - f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]" - ) + self.console.print(f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]") return None def _ensure_utc_timestamp(self, timestamp_str: str) -> str: diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index 2d84796150a..c72001aee1a 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -78,9 +78,7 @@ class CBFTransformer: ) if len(cbf_data) > 0: - console.print( - f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]" - ) + console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]") return pl.DataFrame(cbf_data) @@ -100,9 +98,7 @@ class CBFTransformer: # Build dimensions for CloudZero model = str(row.get("model", "")) - api_key_hash = str(row.get("api_key", ""))[ - :8 - ] # First 8 chars for identification + api_key_hash = str(row.get("api_key", ""))[:8] # First 8 chars for identification # Handle team information with fallbacks team_id = row.get("team_id") @@ -110,9 +106,7 @@ class CBFTransformer: user_email = row.get("user_email") # Use team_alias if available, otherwise team_id, otherwise fallback to 'unknown' - entity_id = ( - str(team_alias) if team_alias else (str(team_id) if team_id else "unknown") - ) + entity_id = str(team_alias) if team_alias else (str(team_id) if team_id else "unknown") # Get alias fields if they exist api_key_alias = row.get("api_key_alias") @@ -152,9 +146,7 @@ class CBFTransformer: ) = czrn_components # Build resource/account as concat of api_key_alias and api_key_prefix - resource_account = ( - f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash - ) + resource_account = f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash # CloudZero CBF format with proper column names cbf_record = { @@ -171,9 +163,7 @@ class CBFTransformer: "resource/service": str(row.get("model_group", "")), # Send model_group "resource/account": resource_account, # Send api_key_alias|api_key_prefix "resource/region": region, # Maps to CZRN region (cross-region) - "resource/usage_family": str( - row.get("custom_llm_provider", "") - ), # Send provider + "resource/usage_family": str(row.get("custom_llm_provider", "")), # Send provider # Action field "action/operation": str(team_id) if team_id else "", # Send team_id # Line item details @@ -182,15 +172,11 @@ class CBFTransformer: # Add CZRN components that don't have direct CBF column mappings as resource tags cbf_record["resource/tag:provider"] = provider # CZRN provider component - cbf_record["resource/tag:model"] = ( - cloud_local_id # CZRN cloud-local-id component (model) - ) + cbf_record["resource/tag:model"] = cloud_local_id # CZRN cloud-local-id component (model) # Add resource tags for all dimensions (using resource/tag: format) for key, value in dimensions.items(): - if ( - value and value != "N/A" and value != "unknown" - ): # Only add meaningful tags + if value and value != "N/A" and value != "unknown": # Only add meaningful tags cbf_record[f"resource/tag:{key}"] = str(value) # Add token breakdown as resource tags for analysis (excluding total_tokens per LIT-1907) diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index 362581937d7..cd7b211f1a5 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -104,9 +104,7 @@ class ChatCompletionFunctionToolChoice(TypedDict): function: dict[str, str] -CodeExecutionFunctionToolChoice = ( - ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice -) +CodeExecutionFunctionToolChoice = ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None: @@ -145,9 +143,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): self._container_cache: dict[str, tuple[Any, dict[str, Any] | None, float]] = {} @classmethod - def from_config_yaml( - cls, config: CodeInterpreterInterceptionConfig - ) -> "CodeInterpreterInterceptionLogger": + def from_config_yaml(cls, config: CodeInterpreterInterceptionConfig) -> "CodeInterpreterInterceptionLogger": return cls( enabled=bool(config.get("enabled", True)), enabled_providers=config.get("enabled_providers"), @@ -171,9 +167,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) return CodeInterpreterInterceptionLogger.from_config_yaml(params) - async def async_pre_call_deployment_hook( - self, kwargs: dict[str, Any], call_type: CallTypes | None - ) -> dict | None: + async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: if not kwargs.get("_agentic_loop_depth"): kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None) kwargs.pop(_SANDBOX_KEY, None) @@ -187,19 +181,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger): CallTypes.acompletion, ): return None - if ( - self.enabled_providers is not None - and self._resolve_provider(kwargs) not in self.enabled_providers - ): + if self.enabled_providers is not None and self._resolve_provider(kwargs) not in self.enabled_providers: return None tools = kwargs.get("tools") if not isinstance(tools, list): return None - if not any( - isinstance(tool, dict) and tool.get("type") == "code_interpreter" - for tool in tools - ): + if not any(isinstance(tool, dict) and tool.get("type") == "code_interpreter" for tool in tools): return None kwargs[_INTERCEPTION_ACTIVE_KEY] = True @@ -211,11 +199,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): function_tool = self._get_function_tool(call_type=call_type) kwargs["tools"] = [ - ( - function_tool - if isinstance(tool, dict) and tool.get("type") == "code_interpreter" - else tool - ) + (function_tool if isinstance(tool, dict) and tool.get("type") == "code_interpreter" else tool) for tool in tools ] if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")): @@ -256,9 +240,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): "required": ["code"], } - def _get_function_tool( - self, call_type: CallTypes | None - ) -> CodeExecutionFunctionTool: + def _get_function_tool(self, call_type: CallTypes | None) -> CodeExecutionFunctionTool: description = "Execute python code in a sandbox and return stdout." if call_type in (CallTypes.completion, CallTypes.acompletion): return { @@ -299,10 +281,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): tool_choice.get("type") == "code_interpreter" or tool_choice.get("name") == "code_interpreter" or tool_choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME - or ( - isinstance(function, dict) - and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME - ) + or (isinstance(function, dict) and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME) ) def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None: @@ -331,16 +310,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return False, {} if not kwargs.get(_INTERCEPTION_ACTIVE_KEY): return False, {} - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: return False, {} tool_calls = ( self._extract_chat_completion_code_execution_tool_calls(response=response) - if kwargs.get("_agentic_loop_api_surface") - == CHAT_COMPLETION_AGENTIC_SURFACE + if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE else self._extract_code_execution_tool_calls(response=response) ) if not tool_calls: @@ -381,9 +356,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): for tool_call in tool_calls: arguments = tool_call.get("arguments", "") code = self._parse_code(arguments) - stdout = await self._run_tool_call( - container=container, params=params, arguments=arguments - ) + stdout = await self._run_tool_call(container=container, params=params, arguments=arguments) input_list.append( { "type": "function_call", @@ -406,9 +379,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): "status": "completed", "code": code, "container_id": container_id, - "outputs": ( - [{"type": "logs", "logs": stdout}] if stdout else [] - ), + "outputs": ([{"type": "logs", "logs": stdout}] if stdout else []), } ) except Exception: @@ -469,9 +440,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): request_patch = AgenticLoopRequestPatch( model=model, - messages=list(messages) - + [self._build_chat_completion_assistant_message(tool_calls)] - + tool_messages, + messages=list(messages) + [self._build_chat_completion_assistant_message(tool_calls)] + tool_messages, tools=self._get_followup_tools( tools=optional_params.get("tools"), call_type=CallTypes.completion, @@ -500,12 +469,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]: arguments = tool_call.get("arguments", "") code = self._parse_code(arguments) - stdout = await self._run_tool_call( - container=container, params=params, arguments=arguments - ) - tool_call_id = ( - tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex - ) + stdout = await self._run_tool_call(container=container, params=params, arguments=arguments) + tool_call_id = tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex return ( { "role": "tool", @@ -522,9 +487,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): }, ) - async def async_agentic_loop_cleanup_hook( - self, plan: AgenticLoopPlan, kwargs: dict - ) -> None: + async def async_agentic_loop_cleanup_hook(self, plan: AgenticLoopPlan, kwargs: dict) -> None: metadata = plan.metadata or {} if plan else {} await self._delete_container_for_cache_key(metadata.get("sandbox_key")) @@ -534,14 +497,10 @@ class CodeInterpreterInterceptionLogger(CustomLogger): k: v for k, v in kwargs.items() if k not in {"litellm_logging_obj", "acompletion"} - and not is_interception_internal_key( - k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES - ) + and not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES) } - def _get_followup_tools( - self, tools: object, call_type: CallTypes | None - ) -> list[dict[str, Any]] | None: + def _get_followup_tools(self, tools: object, call_type: CallTypes | None) -> list[dict[str, Any]] | None: if not isinstance(tools, list): return None return [ @@ -553,21 +512,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger): for tool in tools ] - def _get_followup_optional_params( - self, optional_params: dict[str, object] - ) -> dict[str, object]: - drop_tool_choice = self._tool_choice_targets_code_interpreter( - optional_params.get("tool_choice") - ) + def _get_followup_optional_params(self, optional_params: dict[str, object]) -> dict[str, object]: + drop_tool_choice = self._tool_choice_targets_code_interpreter(optional_params.get("tool_choice")) return { - k: v - for k, v in optional_params.items() - if k != "tools" and not (k == "tool_choice" and drop_tool_choice) + k: v for k, v in optional_params.items() if k != "tools" and not (k == "tool_choice" and drop_tool_choice) } - async def async_post_agentic_loop_response_hook( - self, response: Any, plan: AgenticLoopPlan, kwargs: dict - ) -> Any: + async def async_post_agentic_loop_response_hook(self, response: Any, plan: AgenticLoopPlan, kwargs: dict) -> Any: metadata = plan.metadata or {} if plan else {} await self._delete_container_for_cache_key(metadata.get("sandbox_key")) @@ -576,18 +527,12 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return response is_dict = isinstance(response, dict) - output = ( - response.get("output") if is_dict else getattr(response, "output", None) - ) + output = response.get("output") if is_dict else getattr(response, "output", None) if not isinstance(output, list): return response def _item_type(item: Any) -> Any: - return ( - item.get("type") - if isinstance(item, dict) - else getattr(item, "type", None) - ) + return item.get("type") if isinstance(item, dict) else getattr(item, "type", None) insert_at = next( (i for i, item in enumerate(output) if _item_type(item) == "message"), @@ -607,9 +552,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): except (json.JSONDecodeError, TypeError, AttributeError): return "" - async def _run_tool_call( - self, container: Any, params: dict[str, Any] | None, arguments: str - ) -> str: + async def _run_tool_call(self, container: Any, params: dict[str, Any] | None, arguments: str) -> str: try: code = json.loads(arguments).get("code", "") if arguments else "" except (json.JSONDecodeError, TypeError): @@ -618,17 +561,11 @@ class CodeInterpreterInterceptionLogger(CustomLogger): result = await self._run_code(container=container, params=params, code=code) if getattr(result, "error", None): error = result.error - message = ( - error.get("value") or error.get("name") - if isinstance(error, dict) - else str(error) - ) + message = error.get("value") or error.get("name") if isinstance(error, dict) else str(error) return f"[execution error] {message}" return getattr(result, "stdout", "") or "" - async def _get_or_create_container( - self, cache_key: str | None - ) -> tuple[Any, dict[str, Any] | None]: + async def _get_or_create_container(self, cache_key: str | None) -> tuple[Any, dict[str, Any] | None]: if cache_key: cached = self._container_cache.get(cache_key) if cached is not None: @@ -657,15 +594,11 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) return container, params - async def _run_code( - self, container: Any, params: dict[str, Any] | None, code: str - ) -> Any: + async def _run_code(self, container: Any, params: dict[str, Any] | None, code: str) -> Any: if self.sandbox_config is not None: return await self.sandbox_config.arun_code(container=container, code=code) if params is None: - raise ValueError( - "CodeInterpreterInterception: no sandbox available to run code." - ) + raise ValueError("CodeInterpreterInterception: no sandbox available to run code.") return await litellm.arun_code( provider=params["sandbox_provider"], container=container, @@ -673,9 +606,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): api_key=params.get("api_key"), ) - async def _delete_container( - self, container: Any, params: dict[str, Any] | None - ) -> None: + async def _delete_container(self, container: Any, params: dict[str, Any] | None) -> None: try: if self.sandbox_config is not None: await self.sandbox_config.adelete_sandbox(container=container) @@ -689,9 +620,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): api_base=params.get("api_base"), ) except Exception: - verbose_logger.exception( - "CodeInterpreterInterception: failed to delete sandbox container" - ) + verbose_logger.exception("CodeInterpreterInterception: failed to delete sandbox container") async def _delete_container_for_cache_key(self, cache_key: str | None) -> None: if not cache_key: @@ -708,9 +637,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return list(messages) return [] - def _extract_code_execution_tool_calls( - self, response: object - ) -> list[CodeExecutionToolCall]: + def _extract_code_execution_tool_calls(self, response: object) -> list[CodeExecutionToolCall]: if isinstance(response, dict): output = response.get("output", []) else: @@ -720,17 +647,9 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return [ { - "call_id": ( - item.get("call_id") - if isinstance(item, dict) - else getattr(item, "call_id", None) - ), + "call_id": (item.get("call_id") if isinstance(item, dict) else getattr(item, "call_id", None)), "name": LITELLM_CODE_EXECUTION_TOOL_NAME, - "arguments": ( - item.get("arguments") - if isinstance(item, dict) - else getattr(item, "arguments", "") - ), + "arguments": (item.get("arguments") if isinstance(item, dict) else getattr(item, "arguments", "")), } for item in output if self._is_code_execution_call(item) @@ -751,18 +670,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return [ normalized for tool_call in tool_calls - if (normalized := self._normalize_chat_completion_tool_call(tool_call)) - is not None + if (normalized := self._normalize_chat_completion_tool_call(tool_call)) is not None ] @staticmethod def _normalize_chat_completion_tool_call( tool_call: ChatCompletionMessageToolCall, ) -> CodeExecutionToolCall | None: - if ( - tool_call.type != "function" - or tool_call.function.name != LITELLM_CODE_EXECUTION_TOOL_NAME - ): + if tool_call.type != "function" or tool_call.function.name != LITELLM_CODE_EXECUTION_TOOL_NAME: return None arguments = tool_call.function.arguments @@ -814,10 +729,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): def _is_code_execution_call(self, item: Any) -> bool: if isinstance(item, dict): - return ( - item.get("type") == "function_call" - and item.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME - ) + return item.get("type") == "function_call" and item.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME return ( getattr(item, "type", None) == "function_call" and getattr(item, "name", None) == LITELLM_CODE_EXECUTION_TOOL_NAME diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 8899089500d..c82f9ff477f 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -53,9 +53,7 @@ class CompressionInterceptionLogger(CustomLogger): self._compression_cache_by_call_id: Dict[str, Tuple[Dict[str, str], float]] = {} @classmethod - def from_config_yaml( - cls, config: CompressionInterceptionConfig - ) -> "CompressionInterceptionLogger": + def from_config_yaml(cls, config: CompressionInterceptionConfig) -> "CompressionInterceptionLogger": return cls( enabled=bool(config.get("enabled", True)), compression_trigger=int(config.get("compression_trigger", 200_000)), @@ -124,9 +122,7 @@ class CompressionInterceptionLogger(CustomLogger): kwargs["messages"] = compressed["messages"] if compressed_tools: kwargs["tools"] = self._merge_tools( - existing_tools=cast( - Optional[List[Dict[str, Any]]], kwargs.get("tools") - ), + existing_tools=cast(Optional[List[Dict[str, Any]]], kwargs.get("tools")), compressed_tools=compressed_tools, ) call_id = cast(Optional[str], kwargs.get("litellm_call_id")) @@ -166,9 +162,7 @@ class CompressionInterceptionLogger(CustomLogger): if not self._has_retrieval_tool(tools): return False, {} - tool_calls, thinking_blocks = self._extract_retrieval_tool_calls( - response=response - ) + tool_calls, thinking_blocks = self._extract_retrieval_tool_calls(response=response) if not tool_calls: return False, {} @@ -196,9 +190,7 @@ class CompressionInterceptionLogger(CustomLogger): call_id = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) cache = self._get_cache(call_id=call_id) - retrieval_results = [ - self._resolve_retrieval_content(tc, cache) for tc in tool_calls - ] + retrieval_results = [self._resolve_retrieval_content(tc, cache) for tc in tool_calls] assistant_message = { "role": "assistant", @@ -228,20 +220,15 @@ class CompressionInterceptionLogger(CustomLogger): max_tokens = cast( Optional[int], - anthropic_messages_optional_request_params.get("max_tokens") - or kwargs.get("max_tokens"), + anthropic_messages_optional_request_params.get("max_tokens") or kwargs.get("max_tokens"), ) optional_params_without_max_tokens = { - k: v - for k, v in anthropic_messages_optional_request_params.items() - if k != "max_tokens" + k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens" } full_model_name = model if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get( - "agentic_loop_params", {} - ) + agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = cast(str, agentic_params.get("model", model)) request_patch = AgenticLoopRequestPatch( @@ -277,21 +264,15 @@ class CompressionInterceptionLogger(CustomLogger): return {} return cache_entry[0] - def _resolve_call_id( - self, logging_obj: Any, kwargs: Dict[str, Any] - ) -> Optional[str]: + def _resolve_call_id(self, logging_obj: Any, kwargs: Dict[str, Any]) -> Optional[str]: if logging_obj is not None: logging_call_id = getattr(logging_obj, "litellm_call_id", None) if isinstance(logging_call_id, str) and logging_call_id: return logging_call_id kwargs_call_id = kwargs.get("litellm_call_id") - return cast( - Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None - ) + return cast(Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None) - def _resolve_retrieval_content( - self, tool_call: Dict[str, Any], cache: Dict[str, str] - ) -> str: + def _resolve_retrieval_content(self, tool_call: Dict[str, Any], cache: Dict[str, str]) -> str: raw_input = tool_call.get("input", {}) key = "" if isinstance(raw_input, dict): @@ -302,9 +283,7 @@ class CompressionInterceptionLogger(CustomLogger): return cache[key] return f"[compressed content key '{key}' not found]" - def _extract_retrieval_tool_calls( - self, response: Any - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + def _extract_retrieval_tool_calls(self, response: Any) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: if isinstance(response, dict): content = response.get("content", []) else: @@ -322,10 +301,7 @@ class CompressionInterceptionLogger(CustomLogger): block_name = block.get("name") if block_type in ("thinking", "redacted_thinking"): thinking_blocks.append(block) - if ( - block_type == "tool_use" - and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME - ): + if block_type == "tool_use" and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME: tool_calls.append( { "id": block.get("id"), @@ -352,10 +328,7 @@ class CompressionInterceptionLogger(CustomLogger): "data": getattr(block, "data", ""), } ) - if ( - block_type == "tool_use" - and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME - ): + if block_type == "tool_use" and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME: tool_calls.append( { "id": getattr(block, "id", None), @@ -370,9 +343,7 @@ class CompressionInterceptionLogger(CustomLogger): def _prepare_followup_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: internal_keys = {"litellm_logging_obj"} return { - k: v - for k, v in kwargs.items() - if not k.startswith("_compression_interception") and k not in internal_keys + k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys } def _has_retrieval_tool(self, tools: Any) -> bool: @@ -385,10 +356,7 @@ class CompressionInterceptionLogger(CustomLogger): if tool.get("type") == "function" and isinstance(function, dict): if function.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME: return True - if ( - tool.get("type") == "custom" - and tool.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME - ): + if tool.get("type") == "custom" and tool.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME: return True return False diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index 8f4844501c3..aded12fa399 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -41,20 +41,14 @@ class CustomBatchLogger(CustomLogger): self.batch_size: int = batch_size or litellm.DEFAULT_BATCH_SIZE self.last_flush_time = time.time() self.flush_lock = flush_lock - self.max_queue_size: int = ( - max_queue_size - if max_queue_size is not None - else self.DEFAULT_MAX_QUEUE_SIZE - ) + self.max_queue_size: int = max_queue_size if max_queue_size is not None else self.DEFAULT_MAX_QUEUE_SIZE super().__init__(**kwargs) async def periodic_flush(self): while True: await asyncio.sleep(self.flush_interval) - verbose_logger.debug( - f"CustomLogger periodic flush after {self.flush_interval} seconds" - ) + verbose_logger.debug(f"CustomLogger periodic flush after {self.flush_interval} seconds") await self.flush_queue() async def flush_queue(self): @@ -64,9 +58,7 @@ class CustomBatchLogger(CustomLogger): async with self.flush_lock: if self.log_queue: log_queue_length = len(self.log_queue) - verbose_logger.debug( - "CustomLogger: Flushing batch of %s events", len(self.log_queue) - ) + verbose_logger.debug("CustomLogger: Flushing batch of %s events", len(self.log_queue)) try: await self.async_send_batch() except Exception: @@ -76,8 +68,7 @@ class CustomBatchLogger(CustomLogger): # their own errors, so this only affects loggers that opt # in to surfacing failures (e.g. Rubrik). verbose_logger.exception( - "CustomLogger: async_send_batch raised; preserving " - "%s events in queue for retry", + "CustomLogger: async_send_batch raised; preserving %s events in queue for retry", log_queue_length, ) # Guard against unbounded queue growth if the destination @@ -87,8 +78,7 @@ class CustomBatchLogger(CustomLogger): if overflow > 0: del self.log_queue[:overflow] verbose_logger.warning( - "CustomLogger: log queue exceeded max_queue_size=%s; " - "dropped %s oldest events.", + "CustomLogger: log queue exceeded max_queue_size=%s; dropped %s oldest events.", self.max_queue_size, overflow, ) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 158ef0b4356..59d37639098 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -86,9 +86,7 @@ class CustomGuardrail(CustomLogger): self, guardrail_name: Optional[str] = None, supported_event_hooks: Optional[List[GuardrailEventHooks]] = None, - event_hook: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode] - ] = None, + event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = None, default_on: bool = False, mask_request_content: bool = False, mask_response_content: bool = False, @@ -120,9 +118,7 @@ class CustomGuardrail(CustomLogger): """ self.guardrail_name = guardrail_name self.supported_event_hooks = supported_event_hooks - self.event_hook: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode] - ] = event_hook + self.event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = event_hook self.default_on: bool = default_on self.mask_request_content: bool = mask_request_content self.mask_response_content: bool = mask_response_content @@ -131,9 +127,7 @@ class CustomGuardrail(CustomLogger): self.on_violation: Optional[str] = on_violation self.realtime_violation_message: Optional[str] = realtime_violation_message self.on_sensitive_data: Optional[str] = on_sensitive_data - self.sensitive_data_route_to_model: Optional[str] = ( - sensitive_data_route_to_model - ) + self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing if supported_event_hooks: @@ -141,9 +135,7 @@ class CustomGuardrail(CustomLogger): self._validate_event_hook(event_hook, supported_event_hooks) super().__init__(**kwargs) - def render_violation_message( - self, default: str, context: Optional[Dict[str, Any]] = None - ) -> str: + def render_violation_message(self, default: str, context: Optional[Dict[str, Any]] = None) -> str: """Return a custom violation message if template is configured.""" if not self.violation_message_template: @@ -247,9 +239,7 @@ class CustomGuardrail(CustomLogger): sticky_session_routing=self.sticky_session_routing, ) - def _get_session_id_from_request_data( - self, request_data: Dict[str, Any] - ) -> Optional[str]: + def _get_session_id_from_request_data(self, request_data: Dict[str, Any]) -> Optional[str]: """Extract session_id from request data.""" return get_session_id_from_request_data(request_data) @@ -258,10 +248,7 @@ class CustomGuardrail(CustomLogger): Returns True if this guardrail is configured to route requests to a different model when sensitive data is detected. """ - return ( - self.on_sensitive_data == "route" - and self.sensitive_data_route_to_model is not None - ) + return self.on_sensitive_data == "route" and self.sensitive_data_route_to_model is not None def handle_sensitive_data_detection( self, @@ -297,8 +284,7 @@ class CustomGuardrail(CustomLogger): except ValueError: raise GuardrailRaisedException( message=( - f"Sensitive data detected by {self.guardrail_name} " - "(routing skipped: request has no session_id)" + f"Sensitive data detected by {self.guardrail_name} (routing skipped: request has no session_id)" ), guardrail_name=self.guardrail_name, ) @@ -319,9 +305,7 @@ class CustomGuardrail(CustomLogger): def _validate_event_hook( self, - event_hook: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode] - ], + event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]], supported_event_hooks: List[GuardrailEventHooks], ) -> None: def _validate_event_hook_list_is_in_supported_event_hooks( @@ -332,18 +316,14 @@ class CustomGuardrail(CustomLogger): if isinstance(hook, str): hook = GuardrailEventHooks(hook) if hook not in supported_event_hooks: - raise ValueError( - f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}" - ) + raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}") if event_hook is None: return if isinstance(event_hook, str): event_hook = GuardrailEventHooks(event_hook) if isinstance(event_hook, list): - _validate_event_hook_list_is_in_supported_event_hooks( - event_hook, supported_event_hooks - ) + _validate_event_hook_list_is_in_supported_event_hooks(event_hook, supported_event_hooks) elif isinstance(event_hook, Mode): tag_values_flat: list = [] for v in event_hook.tags.values(): @@ -351,23 +331,13 @@ class CustomGuardrail(CustomLogger): tag_values_flat.extend(v) else: tag_values_flat.append(v) - _validate_event_hook_list_is_in_supported_event_hooks( - tag_values_flat, supported_event_hooks - ) + _validate_event_hook_list_is_in_supported_event_hooks(tag_values_flat, supported_event_hooks) if event_hook.default: - default_list = ( - event_hook.default - if isinstance(event_hook.default, list) - else [event_hook.default] - ) - _validate_event_hook_list_is_in_supported_event_hooks( - default_list, supported_event_hooks - ) + default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default] + _validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks) elif isinstance(event_hook, GuardrailEventHooks): if event_hook not in supported_event_hooks: - raise ValueError( - f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}" - ) + raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}") @staticmethod def _get_admin_metadata(data: dict) -> dict: @@ -431,9 +401,7 @@ class CustomGuardrail(CustomLogger): return True raise - def get_guardrail_from_metadata( - self, data: dict - ) -> Union[List[str], List[Dict[str, DynamicGuardrailParams]]]: + def get_guardrail_from_metadata(self, data: dict) -> Union[List[str], List[Dict[str, DynamicGuardrailParams]]]: """ Returns the guardrail(s) to be run from the metadata or root """ @@ -522,12 +490,7 @@ class CustomGuardrail(CustomLogger): if self._pre_call_hook_already_ran(kwargs): return kwargs - if ( - self.should_run_guardrail( - data=kwargs, event_type=GuardrailEventHooks.pre_call - ) - is not True - ): + if self.should_run_guardrail(data=kwargs, event_type=GuardrailEventHooks.pre_call) is not True: return kwargs # CHECK IF GUARDRAIL REJECTS THE REQUEST @@ -568,12 +531,7 @@ class CustomGuardrail(CustomLogger): if litellm_guardrails is None or not isinstance(litellm_guardrails, list): return response - if ( - self.should_run_guardrail( - data=request_data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: return response # CHECK IF GUARDRAIL REJECTS THE REQUEST @@ -604,9 +562,7 @@ class CustomGuardrail(CustomLogger): """ requested_guardrails = self.get_guardrail_from_metadata(data) disable_global_guardrail = self.get_disable_global_guardrail(data) - opted_out_global_guardrails = ( - self.get_opted_out_global_guardrails_from_metadata(data) - ) + opted_out_global_guardrails = self.get_opted_out_global_guardrails_from_metadata(data) verbose_logger.debug( "inside should_run_guardrail for guardrail=%s event_type= %s guardrail_supported_event_hooks= %s requested_guardrails= %s self.default_on= %s", self.guardrail_name, @@ -615,10 +571,7 @@ class CustomGuardrail(CustomLogger): requested_guardrails, self.default_on, ) - if ( - self.default_on is True - and self.guardrail_name in opted_out_global_guardrails - ): + if self.default_on is True and self.guardrail_name in opted_out_global_guardrails: return False if self.default_on is True and disable_global_guardrail is True: @@ -662,9 +615,7 @@ class CustomGuardrail(CustomLogger): raise ImportError( "Setting tag-based guardrails is only available in litellm-enterprise. You must be a premium user to use this feature." ) - result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag( - data, self.event_hook, event_type - ) + result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag(data, self.event_hook, event_type) if result is not None: return result return True @@ -690,9 +641,7 @@ class CustomGuardrail(CustomLogger): return True if self.event_hook.default: default_list = ( - self.event_hook.default - if isinstance(self.event_hook.default, list) - else [self.event_hook.default] + self.event_hook.default if isinstance(self.event_hook.default, list) else [self.event_hook.default] ) return event_type.value in default_list return False @@ -722,9 +671,7 @@ class CustomGuardrail(CustomLogger): for guardrail in requested_guardrails: if isinstance(guardrail, dict) and self.guardrail_name in guardrail: # Get the configuration for this guardrail - guardrail_config: DynamicGuardrailParams = DynamicGuardrailParams( - **guardrail[self.guardrail_name] - ) + guardrail_config: DynamicGuardrailParams = DynamicGuardrailParams(**guardrail[self.guardrail_name]) extra_body = guardrail_config.get("extra_body", {}) if self._validate_premium_user() is not True: if isinstance(extra_body, dict) and extra_body: @@ -779,9 +726,7 @@ class CustomGuardrail(CustomLogger): from litellm.types.utils import GuardrailMode # Use event_type if provided, otherwise fall back to self.event_hook - guardrail_mode: Union[ - GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks] - ] + guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]] if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): @@ -795,9 +740,7 @@ class CustomGuardrail(CustomLogger): # Sanitize the response to ensure it's JSON serializable and free of circular refs # This prevents RecursionErrors in downstream loggers (Langfuse, Datadog, etc.) - clean_guardrail_response = filter_exceptions_from_params( - guardrail_json_response - ) + clean_guardrail_response = filter_exceptions_from_params(guardrail_json_response) # Strip secret_fields to prevent plaintext Authorization headers from # being persisted to spend logs, OTEL traces, or other logging backends. @@ -812,9 +755,7 @@ class CustomGuardrail(CustomLogger): # Default-safe behavior: never persist raw matched spans in standard # guardrail logging payloads (single shared implementation; Bedrock hooks pass # raw provider JSON so redaction is not duplicated upstream). - clean_guardrail_response = redact_nested_match_and_regex_keys( - clean_guardrail_response - ) + clean_guardrail_response = redact_nested_match_and_regex_keys(clean_guardrail_response) slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, @@ -908,9 +849,7 @@ class CustomGuardrail(CustomLogger): This gets logged on downsteam Langfuse, DataDog, etc. """ # Convert None to empty dict to satisfy type requirements - guardrail_response: Union[Dict[str, Any], str] = ( - {} if response is None else response - ) + guardrail_response: Union[Dict[str, Any], str] = {} if response is None else response # For apply_guardrail functions in custom_code_guardrail scenario, # simplify the logged response to "allow", "deny", or "mask" @@ -958,11 +897,7 @@ class CustomGuardrail(CustomLogger): ), ): return True - if ( - HTTPException is not None - and isinstance(e, HTTPException) - and e.status_code == 400 - ): + if HTTPException is not None and isinstance(e, HTTPException) and e.status_code == 400: return True return False @@ -981,9 +916,7 @@ class CustomGuardrail(CustomLogger): This gets logged on downsteam Langfuse, DataDog, etc. """ guardrail_status: GuardrailStatus = ( - "guardrail_intervened" - if self._is_guardrail_intervention(e) - else "guardrail_failed_to_respond" + "guardrail_intervened" if self._is_guardrail_intervention(e) else "guardrail_failed_to_respond" ) # For custom_code_guardrail scenario, log as "deny" instead of full exception # Check if this is from custom_code_guardrail by checking the class name @@ -1071,10 +1004,7 @@ class CustomGuardrail(CustomLogger): # /responses # User/System messages are stored in the "input" key, use litellm transformation to get the messages ######################################################### - if ( - call_type == CallTypes.responses.value - or call_type == CallTypes.aresponses.value - ): + if call_type == CallTypes.responses.value or call_type == CallTypes.aresponses.value: from typing import cast from litellm.responses.litellm_completion_transformation.transformation import ( @@ -1105,9 +1035,7 @@ def _append_slg_to_litellm_params(lp: object, entries: list) -> None: existing.append(entry) -def _sync_guardrail_info_to_logging_obj( - request_data: dict, logging_obj: object -) -> None: +def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object) -> None: """Copy standard_logging_guardrail_information from request_data into logging_obj. The @log_guardrail_information decorator writes guardrail info to @@ -1120,9 +1048,7 @@ def _sync_guardrail_info_to_logging_obj( """ if logging_obj is None: return - meta_src = ( - request_data.get("metadata") or request_data.get("litellm_metadata") or {} - ) + meta_src = request_data.get("metadata") or request_data.get("litellm_metadata") or {} slg_info = meta_src.get("standard_logging_guardrail_information") if not slg_info: return diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 6d65b4ec0d2..108928871b0 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -145,9 +145,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_pre_api_call(self, model, messages, kwargs): pass - async def async_pre_request_hook( - self, model: str, messages: List, kwargs: Dict - ) -> Optional[Dict]: + async def async_pre_request_hook(self, model: str, messages: List, kwargs: Dict) -> Optional[Dict]: """ Hook called before making the API request to allow modifying request parameters. @@ -273,9 +271,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ pass - async def async_pre_call_check( - self, deployment: dict, parent_otel_span: Optional[Span] - ) -> Optional[dict]: + async def async_pre_call_check(self, deployment: dict, parent_otel_span: Optional[Span]) -> Optional[dict]: pass def pre_call_check(self, deployment: dict) -> Optional[dict]: @@ -311,29 +307,21 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ): pass - async def log_success_fallback_event( - self, original_model_group: str, kwargs: dict, original_exception: Exception - ): + async def log_success_fallback_event(self, original_model_group: str, kwargs: dict, original_exception: Exception): pass - async def log_failure_fallback_event( - self, original_model_group: str, kwargs: dict, original_exception: Exception - ): + async def log_failure_fallback_event(self, original_model_group: str, kwargs: dict, original_exception: Exception): pass #### ADAPTERS #### Allow calling 100+ LLMs in custom format - https://github.com/BerriAI/litellm/pulls - def translate_completion_input_params( - self, kwargs - ) -> Optional[ChatCompletionRequest]: + def translate_completion_input_params(self, kwargs) -> Optional[ChatCompletionRequest]: """ Translates the input params, from the provider's native format to the litellm.completion() format. """ pass - def translate_completion_output_params( - self, response: ModelResponse - ) -> Optional[BaseModel]: + def translate_completion_output_params(self, response: ModelResponse) -> Optional[BaseModel]: """ Translates the output params, from the OpenAI format to the custom format. """ @@ -435,15 +423,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> Any: pass - async def async_logging_hook( - self, kwargs: dict, result: Any, call_type: str - ) -> Tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result - def logging_hook( - self, kwargs: dict, result: Any, call_type: str - ) -> Tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result @@ -485,9 +469,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac except Exception: print_verbose(f"Custom Logger Error - {traceback.format_exc()}") - async def async_log_input_event( - self, model, messages, kwargs, print_verbose, callback_func - ): + async def async_log_input_event(self, model, messages, kwargs, print_verbose, callback_func): try: kwargs["model"] = model kwargs["messages"] = messages @@ -499,9 +481,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac except Exception: print_verbose(f"Custom Logger Error - {traceback.format_exc()}") - def log_event( - self, kwargs, response_obj, start_time, end_time, print_verbose, callback_func - ): + def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose, callback_func): # Method definition try: kwargs["log_event_type"] = "post_api_call" @@ -515,9 +495,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac print_verbose(f"Custom Logger Error - {traceback.format_exc()}") pass - async def async_log_event( - self, kwargs, response_obj, start_time, end_time, print_verbose, callback_func - ): + async def async_log_event(self, kwargs, response_obj, start_time, end_time, print_verbose, callback_func): # Method definition try: kwargs["log_event_type"] = "post_api_call" @@ -834,15 +812,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def _truncate_text(self, text: str, max_length: int) -> str: """Truncate text if it exceeds max_length""" return ( - text[:max_length] - + "...truncated by litellm, this logger does not support large content" + text[:max_length] + "...truncated by litellm, this logger does not support large content" if len(text) > max_length else text ) - def _select_metadata_field( - self, request_kwargs: Optional[Dict] = None - ) -> Optional[str]: + def _select_metadata_field(self, request_kwargs: Optional[Dict] = None) -> Optional[str]: """ Select the metadata field to use for logging @@ -857,9 +832,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return LITELLM_METADATA_FIELD return OLD_LITELLM_METADATA_FIELD - def redact_standard_logging_payload_from_model_call_details( - self, model_call_details: Dict - ) -> Dict: + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: Dict) -> Dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. @@ -876,12 +849,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac from litellm import Choices, Message, ModelResponse - turn_off_message_logging: bool = getattr( - self, "turn_off_message_logging", False - ) - excluded_fields: Optional[List[str]] = getattr( - litellm, "standard_logging_payload_excluded_fields", None - ) + turn_off_message_logging: bool = getattr(self, "turn_off_message_logging", False) + excluded_fields: Optional[List[str]] = getattr(litellm, "standard_logging_payload_excluded_fields", None) # Early return if no processing needed if turn_off_message_logging is False and not excluded_fields: @@ -907,18 +876,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if turn_off_message_logging: redacted_str = "redacted-by-litellm" - if ( - "messages" not in (excluded_fields or []) - and standard_logging_object_copy.get("messages") is not None - ): - standard_logging_object_copy["messages"] = [ - Message(content=redacted_str).model_dump() - ] + if "messages" not in (excluded_fields or []) and standard_logging_object_copy.get("messages") is not None: + standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()] - if ( - "response" not in (excluded_fields or []) - and standard_logging_object_copy.get("response") is not None - ): + if "response" not in (excluded_fields or []) and standard_logging_object_copy.get("response") is not None: response = standard_logging_object_copy["response"] # Check if this is a ResponsesAPIResponse (has "output" field) if isinstance(response, dict) and "output" in response: @@ -929,30 +890,20 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac # Redact content in output array if isinstance(response_copy.get("output"), list): for output_item in response_copy["output"]: - if ( - isinstance(output_item, dict) - and "content" in output_item - ): + if isinstance(output_item, dict) and "content" in output_item: if isinstance(output_item["content"], list): # Redact text in content items for content_item in output_item["content"]: - if ( - isinstance(content_item, dict) - and "text" in content_item - ): + if isinstance(content_item, dict) and "text" in content_item: content_item["text"] = redacted_str standard_logging_object_copy["response"] = response_copy else: # Standard ModelResponse format - model_response = ModelResponse( - choices=[Choices(message=Message(content=redacted_str))] - ) + model_response = ModelResponse(choices=[Choices(message=Message(content=redacted_str))]) model_response_dict = model_response.model_dump() standard_logging_object_copy["response"] = model_response_dict - model_call_details_copy["standard_logging_object"] = ( - standard_logging_object_copy - ) + model_call_details_copy["standard_logging_object"] = standard_logging_object_copy return model_call_details_copy async def get_proxy_server_request_from_cold_storage_with_object_key( @@ -978,12 +929,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac for callback_obj in all_callbacks: if hasattr(callback_obj, "increment_callback_logging_failure"): - verbose_logger.debug( - f"Incrementing callback failure metric for {callback_name}" - ) - callback_obj.increment_callback_logging_failure( - callback_name=callback_name - ) # type: ignore + verbose_logger.debug(f"Incrementing callback failure metric for {callback_name}") + callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore return verbose_logger.debug( @@ -994,9 +941,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac except Exception as e: from litellm._logging import verbose_logger - verbose_logger.debug( - f"Error in handle_callback_failure for {callback_name}: {str(e)}" - ) + verbose_logger.debug(f"Error in handle_callback_failure for {callback_name}: {str(e)}") async def _strip_base64_from_messages( self, @@ -1015,14 +960,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ raw_messages: Any = payload.get("messages", []) messages: List[Any] = raw_messages if isinstance(raw_messages, list) else [] - verbose_logger.debug( - f"[CustomLogger] Stripping base64 from {len(messages)} messages" - ) + verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages") if messages: - payload["messages"] = self._process_messages( - messages=messages, max_depth=max_depth - ) + payload["messages"] = self._process_messages(messages=messages, max_depth=max_depth) total_items = 0 for m in payload.get("messages", []) or []: @@ -1031,9 +972,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if isinstance(content, list): total_items += len(content) - verbose_logger.debug( - f"[CustomLogger] Completed base64 strip; retained {total_items} content items" - ) + verbose_logger.debug(f"[CustomLogger] Completed base64 strip; retained {total_items} content items") return payload def _strip_base64_from_messages_sync( @@ -1053,14 +992,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ raw_messages: Any = payload.get("messages", []) messages: List[Any] = raw_messages if isinstance(raw_messages, list) else [] - verbose_logger.debug( - f"[CustomLogger] Stripping base64 from {len(messages)} messages" - ) + verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages") if messages: - payload["messages"] = self._process_messages( - messages=messages, max_depth=max_depth - ) + payload["messages"] = self._process_messages(messages=messages, max_depth=max_depth) total_items = 0 for m in payload.get("messages", []) or []: @@ -1069,9 +1004,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if isinstance(content, list): total_items += len(content) - verbose_logger.debug( - f"[CustomLogger] Completed base64 strip; retained {total_items} content items" - ) + verbose_logger.debug(f"[CustomLogger] Completed base64 strip; retained {total_items} content items") return payload def _redact_base64( @@ -1082,30 +1015,20 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> Any: """Recursively redact inline base64 from any nested structure with a max recursion depth limit.""" if depth > max_depth: - verbose_logger.warning( - f"[CustomLogger] Max recursion depth {max_depth} reached while redacting base64" - ) + verbose_logger.warning(f"[CustomLogger] Max recursion depth {max_depth} reached while redacting base64") return "[MAX_DEPTH_REACHED]" if isinstance(value, str): if _BASE64_INLINE_PATTERN.search(value): - verbose_logger.debug( - f"[CustomLogger] Redacted inline base64 string: {value[:40]}..." - ) + verbose_logger.debug(f"[CustomLogger] Redacted inline base64 string: {value[:40]}...") return _BASE64_INLINE_PATTERN.sub("[BASE64_REDACTED]", value) return value if isinstance(value, list): - return [ - self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth) - for v in value - ] + return [self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth) for v in value] if isinstance(value, dict): - return { - k: self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth) - for k, v in value.items() - } + return {k: self._redact_base64(value=v, depth=depth + 1, max_depth=max_depth) for k, v in value.items()} return value @@ -1132,14 +1055,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac cleaned: List[Any] = [] for c in contents: if self._should_keep_content(content=c): - cleaned.append( - self._redact_base64(value=c, max_depth=max_depth) - ) + cleaned.append(self._redact_base64(value=c, max_depth=max_depth)) msg["content"] = cleaned else: - msg["content"] = self._redact_base64( - value=contents, max_depth=max_depth - ) + msg["content"] = self._redact_base64(value=contents, max_depth=max_depth) for key, val in list(msg.items()): if key != "content": diff --git a/litellm/integrations/custom_prompt_management.py b/litellm/integrations/custom_prompt_management.py index 61e619aba65..fbca1867793 100644 --- a/litellm/integrations/custom_prompt_management.py +++ b/litellm/integrations/custom_prompt_management.py @@ -18,9 +18,7 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase): **kwargs, ): self.ignore_prompt_manager_model = ignore_prompt_manager_model - self.ignore_prompt_manager_optional_params = ( - ignore_prompt_manager_optional_params - ) + self.ignore_prompt_manager_optional_params = ignore_prompt_manager_optional_params def get_chat_completion_prompt( self, @@ -65,9 +63,7 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase): prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ) -> PromptManagementClient: - raise NotImplementedError( - "Custom prompt management does not support compile prompt helper" - ) + raise NotImplementedError("Custom prompt management does not support compile prompt helper") async def async_compile_prompt_helper( self, @@ -78,6 +74,4 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase): prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ) -> PromptManagementClient: - raise NotImplementedError( - "Custom prompt management does not support async compile prompt helper" - ) + raise NotImplementedError("Custom prompt management does not support async compile prompt helper") diff --git a/litellm/integrations/custom_secret_manager.py b/litellm/integrations/custom_secret_manager.py index 45ffa2e08cf..a1bb7b00d92 100644 --- a/litellm/integrations/custom_secret_manager.py +++ b/litellm/integrations/custom_secret_manager.py @@ -224,14 +224,10 @@ class CustomSecretManager(BaseSecretManager): Raises: ValueError: If required configuration is missing """ - verbose_logger.debug( - "No environment validation configured for custom secret manager" - ) + verbose_logger.debug("No environment validation configured for custom secret manager") return True - async def async_health_check( - self, timeout: Optional[Union[float, httpx.Timeout]] = None - ) -> bool: + async def async_health_check(self, timeout: Optional[Union[float, httpx.Timeout]] = None) -> bool: """ Perform a health check on your secret manager. @@ -243,9 +239,7 @@ class CustomSecretManager(BaseSecretManager): Returns: True if the secret manager is healthy, False otherwise """ - verbose_logger.debug( - f"Health check not implemented for {self.secret_manager_name}" - ) + verbose_logger.debug(f"Health check not implemented for {self.secret_manager_name}") return True def __repr__(self) -> str: diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index b0cd0eb1172..6775858c124 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -129,9 +129,7 @@ class DataDogLogger( if self.is_mock_mode: create_mock_datadog_client() - verbose_logger.debug( - "[DATADOG MOCK] Datadog logger initialized in mock mode" - ) + verbose_logger.debug("[DATADOG MOCK] Datadog logger initialized in mock mode") ######################################################### # Handle datadog_params set as litellm.datadog_params @@ -139,9 +137,7 @@ class DataDogLogger( dict_datadog_params = self._get_datadog_params() kwargs.update(dict_datadog_params) - self.async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Configure DataDog endpoint (Agent or Direct API) # Prefer explicit kwargs, then fall back to env vars @@ -173,9 +169,7 @@ class DataDogLogger( batch_size=_resolve_dd_batch_size(), ) except Exception as e: - verbose_logger.exception( - f"Datadog: Got exception on init Datadog client {str(e)}" - ) + verbose_logger.exception(f"Datadog: Got exception on init Datadog client {str(e)}") raise e def _get_datadog_params(self) -> Dict: @@ -190,9 +184,7 @@ class DataDogLogger( dict_datadog_params = litellm.datadog_params.model_dump() elif isinstance(litellm.datadog_params, Dict): # only allow params that are of DatadogInitParams - dict_datadog_params = DatadogInitParams( - **litellm.datadog_params - ).model_dump() + dict_datadog_params = DatadogInitParams(**litellm.datadog_params).model_dump() return dict_datadog_params def _configure_dd_agent( @@ -211,9 +203,7 @@ class DataDogLogger( dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True. Optional when using agent. allow_env_credentials: When False, never read the API key from DD_API_KEY env var. """ - resolved_port = dd_agent_port or os.getenv( - "LITELLM_DD_AGENT_PORT", "10518" - ) # default port for logs + resolved_port = dd_agent_port or os.getenv("LITELLM_DD_AGENT_PORT", "10518") # default port for logs self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs" self.DD_API_KEY = dd_api_key or ( os.getenv("DD_API_KEY") if allow_env_credentials else None @@ -237,9 +227,7 @@ class DataDogLogger( Raises: Exception: If required credentials are not provided via args or env vars """ - resolved_api_key = dd_api_key or ( - os.getenv("DD_API_KEY") if allow_env_credentials else None - ) + resolved_api_key = dd_api_key or (os.getenv("DD_API_KEY") if allow_env_credentials else None) resolved_site = dd_site or os.getenv("DD_SITE") if resolved_api_key is None: @@ -263,28 +251,20 @@ class DataDogLogger( Raises a NON Blocking verbose_logger.exception if an error occurs """ try: - verbose_logger.debug( - "Datadog: Logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("Datadog: Logging - Enters logging function for model %s", kwargs) await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: - verbose_logger.exception( - f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}") pass async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug( - "Datadog: Logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("Datadog: Logging - Enters logging function for model %s", kwargs) await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: - verbose_logger.exception( - f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}") pass async def async_post_call_failure_hook( @@ -323,36 +303,24 @@ class DataDogLogger( LiteLLMProxyRequestSetup, ) - _meta = ( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + _meta = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict ) user_context = dict(_meta) if isinstance(_meta, dict) else _meta except Exception: # Fallback if proxy not available (e.g. SDK-only): minimal safe fields if hasattr(user_api_key_dict, "request_route"): - user_context["request_route"] = getattr( - user_api_key_dict, "request_route", None - ) + user_context["request_route"] = getattr(user_api_key_dict, "request_route", None) if hasattr(user_api_key_dict, "team_id"): - user_context["team_id"] = getattr( - user_api_key_dict, "team_id", None - ) + user_context["team_id"] = getattr(user_api_key_dict, "team_id", None) if hasattr(user_api_key_dict, "user_id"): - user_context["user_id"] = getattr( - user_api_key_dict, "user_id", None - ) + user_context["user_id"] = getattr(user_api_key_dict, "user_id", None) if hasattr(user_api_key_dict, "end_user_id"): - user_context["end_user_id"] = getattr( - user_api_key_dict, "end_user_id", None - ) + user_context["end_user_id"] = getattr(user_api_key_dict, "end_user_id", None) message_payload: DatadogProxyFailureHookJsonMessage = { - "exception": error_information.get("error_message") - or str(original_exception), - "error_class": error_information.get("error_class") - or original_exception.__class__.__name__, + "exception": error_information.get("error_message") or str(original_exception), + "error_class": error_information.get("error_class") or original_exception.__class__.__name__, "status_code": status_code, "traceback": error_information.get("traceback") or "", "user_api_key_dict": user_context, @@ -372,9 +340,7 @@ class DataDogLogger( if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception as e: - verbose_logger.exception( - f"Datadog: async_post_call_failure_hook - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Datadog: async_post_call_failure_hook - {str(e)}\n{traceback.format_exc()}") return None async def async_send_batch(self): @@ -403,24 +369,18 @@ class DataDogLogger( ) if self.is_mock_mode: - verbose_logger.debug( - "[DATADOG MOCK] Mock mode enabled - API calls will be intercepted" - ) + verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") undelivered = await self._send_with_413_split(batch_to_send) if undelivered: self.log_queue = undelivered + self.log_queue if self.is_mock_mode: - verbose_logger.debug( - f"[DATADOG MOCK] Batch of {len(batch_to_send)} events successfully mocked" - ) + verbose_logger.debug(f"[DATADOG MOCK] Batch of {len(batch_to_send)} events successfully mocked") except Exception as e: self.log_queue = batch_to_send + self.log_queue - verbose_logger.exception( - f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}") async def _send_with_413_split(self, batch: List) -> List: """ @@ -444,9 +404,7 @@ class DataDogLogger( if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413: response = e.response else: - verbose_logger.exception( - f"Datadog Error sending batch API - {str(e)}" - ) + verbose_logger.exception(f"Datadog Error sending batch API - {str(e)}") return self._undelivered(chunk, pending) if response.status_code == 413: @@ -484,9 +442,7 @@ class DataDogLogger( async with self.flush_lock: if self.log_queue: - verbose_logger.debug( - "Datadog: Flushing batch of %s events", len(self.log_queue) - ) + verbose_logger.debug("Datadog: Flushing batch of %s events", len(self.log_queue)) await self.async_send_batch() if not self.log_queue: self.last_flush_time = time.time() @@ -528,9 +484,7 @@ class DataDogLogger( response.raise_for_status() if response.status_code != 202: - raise Exception( - f"Response from datadog API status_code: {response.status_code}, text: {response.text}" - ) + raise Exception(f"Response from datadog API status_code: {response.status_code}, text: {response.text}") verbose_logger.debug( "Datadog: Response from datadog API status_code: %s, text: %s", @@ -539,9 +493,7 @@ class DataDogLogger( ) except Exception as e: - verbose_logger.exception( - f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}") pass pass @@ -554,9 +506,7 @@ class DataDogLogger( ) self.log_queue.append(dd_payload) - verbose_logger.debug( - f"Datadog, event added to queue. Will flush in {self.flush_interval} seconds..." - ) + verbose_logger.debug(f"Datadog, event added to queue. Will flush in {self.flush_interval} seconds...") if len(self.log_queue) >= self.batch_size: await self.flush_queue() @@ -572,9 +522,7 @@ class DataDogLogger( verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload) dd_payload = DatadogPayload( ddsource=get_datadog_source(), - ddtags=",".join( - get_datadog_tags(standard_logging_object=standard_logging_object) - ), + ddtags=",".join(get_datadog_tags(standard_logging_object=standard_logging_object)), hostname=get_datadog_hostname(), message=json_payload, service=get_datadog_service(), @@ -603,9 +551,7 @@ class DataDogLogger( DatadogPayload: defined in types.py """ - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object not found in kwargs") @@ -687,9 +633,7 @@ class DataDogLogger( self.log_queue.append(_dd_payload) except Exception as e: - verbose_logger.exception( - f"Datadog: Logger - Exception in async_service_failure_hook: {e}" - ) + verbose_logger.exception(f"Datadog: Logger - Exception in async_service_failure_hook: {e}") pass async def async_service_success_hook( @@ -729,9 +673,7 @@ class DataDogLogger( self.log_queue.append(_dd_payload) except Exception as e: - verbose_logger.exception( - f"Datadog: Logger - Exception in async_service_failure_hook: {e}" - ) + verbose_logger.exception(f"Datadog: Logger - Exception in async_service_failure_hook: {e}") def _create_v0_logging_payload( self, @@ -748,9 +690,7 @@ class DataDogLogger( """ litellm_params = kwargs.get("litellm_params", {}) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None + metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None messages = kwargs.get("messages") optional_params = kwargs.get("optional_params", {}) call_type = kwargs.get("call_type", "litellm.completion") @@ -834,9 +774,7 @@ class DataDogLogger( if span_id is not None: dd_payload["dd.span_id"] = span_id except Exception: - verbose_logger.exception( - "Datadog: Failed to attach trace context to payload" - ) + verbose_logger.exception("Datadog: Failed to attach trace context to payload") def _get_active_trace_context(self) -> Optional[Dict[str, str]]: try: @@ -863,9 +801,7 @@ class DataDogLogger( trace_context["span_id"] = str(span_id) return trace_context except Exception: - verbose_logger.exception( - "Datadog: Failed to retrieve active trace context from tracer" - ) + verbose_logger.exception("Datadog: Failed to retrieve active trace context from tracer") return None async def async_health_check(self) -> IntegrationHealthCheckStatus: diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 0f954eb1ce0..714a50eb2f2 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -57,9 +57,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): self.upload_url = f"https://api.{self.dd_site}/api/v2/cost/custom_costs" - self.async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Initialize lock and start periodic flush task self.flush_lock = asyncio.Lock() @@ -73,9 +71,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: return @@ -88,9 +84,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"Datadog Cost Management: Error in async_log_success_event: {str(e)}" - ) + verbose_logger.exception(f"Datadog Cost Management: Error in async_log_success_event: {str(e)}") async def async_send_batch(self): if not self.log_queue: @@ -103,28 +97,21 @@ class DatadogCostManagementLogger(CustomBatchLogger): aggregated_entries = self._aggregate_costs(batch_to_send) if not aggregated_entries: verbose_logger.debug( - "Datadog Cost Management: batch produced no aggregable entries; " - "dropping %d log(s) from queue.", + "Datadog Cost Management: batch produced no aggregable entries; dropping %d log(s) from queue.", len(batch_to_send), ) return await self._upload_to_datadog(aggregated_entries) except Exception as e: self.log_queue = batch_to_send + self.log_queue - verbose_logger.exception( - f"Datadog Cost Management: Error in async_send_batch: {str(e)}" - ) + verbose_logger.exception(f"Datadog Cost Management: Error in async_send_batch: {str(e)}") - def _aggregate_costs( - self, logs: List[StandardLoggingPayload] - ) -> List[DatadogFOCUSCostEntry]: + def _aggregate_costs(self, logs: List[StandardLoggingPayload]) -> List[DatadogFOCUSCostEntry]: """ Aggregates costs by Provider, Model, and Date. Returns a list of DatadogFOCUSCostEntry. """ - aggregator: Dict[ - Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry - ] = {} + aggregator: Dict[Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry] = {} for log in logs: try: @@ -172,9 +159,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): aggregator[key]["BilledCost"] += cost except Exception as e: - verbose_logger.warning( - f"Error processing log for cost aggregation: {e}" - ) + verbose_logger.warning(f"Error processing log for cost aggregation: {e}") continue return list(aggregator.values()) @@ -229,11 +214,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): nested = metadata.get(nested_key) if isinstance(nested, dict): for k, v in nested.items(): - if ( - k in allow - and v is not None - and not isinstance(v, (dict, list)) - ): + if k in allow and v is not None and not isinstance(v, (dict, list)): self._set_custom_tag(tags, k, str(v)) return tags @@ -268,9 +249,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): # The API endpoint expects a list of objects directly in the body (file content behavior) data_json = safe_dumps(payload) - response = await self.async_client.put( - self.upload_url, content=data_json, headers=headers - ) + response = await self.async_client.put(self.upload_url, content=data_json, headers=headers) response.raise_for_status() diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 201d3fb0a41..1078f05165a 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -53,18 +53,14 @@ class DataDogLLMObsLogger(CustomBatchLogger): if self.is_mock_mode: create_mock_datadog_client() - verbose_logger.debug( - "[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode" - ) + verbose_logger.debug("[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode") # Configure DataDog endpoint (Agent or Direct API) # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST # Check for agent mode FIRST - agent mode doesn't require DD_API_KEY or DD_SITE dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") - self.async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.DD_API_KEY = os.getenv("DD_API_KEY") if dd_agent_host: @@ -74,9 +70,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): if os.getenv("DD_API_KEY", None) is None: raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'") if os.getenv("DD_SITE", None) is None: - raise Exception( - "DD_SITE is not set, set 'DD_SITE=<>', example sit = `us5.datadoghq.com`" - ) + raise Exception("DD_SITE is not set, set 'DD_SITE=<>', example sit = `us5.datadoghq.com`") self._configure_dd_direct_api() # Optional override for testing @@ -108,9 +102,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Use specific port for LLM Obs (Trace Agent) to avoid conflict with Logs Agent (10518) agent_port = os.getenv("LITELLM_DD_LLM_OBS_PORT", "8126") self.DD_SITE = "localhost" # Not used for URL construction in agent mode - self.intake_url = ( - f"http://{dd_agent_host}:{agent_port}/api/intake/llm-obs/v1/trace/spans" - ) + self.intake_url = f"http://{dd_agent_host}:{agent_port}/api/intake/llm-obs/v1/trace/spans" verbose_logger.debug(f"DataDogLLMObs: Using DD Agent at {self.intake_url}") def _configure_dd_direct_api(self): @@ -122,13 +114,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): self.DD_SITE = os.getenv("DD_SITE") if not self.DD_SITE: - raise Exception( - "DD_SITE is not set, set 'DD_SITE=<>', example site = `us5.datadoghq.com`" - ) + raise Exception("DD_SITE is not set, set 'DD_SITE=<>', example site = `us5.datadoghq.com`") - self.intake_url = ( - f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans" - ) + self.intake_url = f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans" def _get_datadog_llm_obs_params(self) -> Dict: """ @@ -138,12 +126,8 @@ class DataDogLLMObsLogger(CustomBatchLogger): """ dict_datadog_llm_obs_params: Dict = {} if litellm.datadog_llm_observability_params is not None: - if isinstance( - litellm.datadog_llm_observability_params, DatadogLLMObsInitParams - ): - dict_datadog_llm_obs_params = ( - litellm.datadog_llm_observability_params.model_dump() - ) + if isinstance(litellm.datadog_llm_observability_params, DatadogLLMObsInitParams): + dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump() elif isinstance(litellm.datadog_llm_observability_params, Dict): # only allow params that are of DatadogLLMObsInitParams dict_datadog_llm_obs_params = DatadogLLMObsInitParams( @@ -153,9 +137,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug( - f"DataDogLLMObs: Logging success event for model {kwargs.get('model', 'unknown')}" - ) + verbose_logger.debug(f"DataDogLLMObs: Logging success event for model {kwargs.get('model', 'unknown')}") payload = self.create_llm_obs_payload(kwargs, start_time, end_time) verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}") self.log_queue.append(payload) @@ -163,15 +145,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"DataDogLLMObs: Error logging success event - {str(e)}" - ) + verbose_logger.exception(f"DataDogLLMObs: Error logging success event - {str(e)}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug( - f"DataDogLLMObs: Logging failure event for model {kwargs.get('model', 'unknown')}" - ) + verbose_logger.debug(f"DataDogLLMObs: Logging failure event for model {kwargs.get('model', 'unknown')}") payload = self.create_llm_obs_payload(kwargs, start_time, end_time) verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}") self.log_queue.append(payload) @@ -179,23 +157,17 @@ class DataDogLLMObsLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"DataDogLLMObs: Error logging failure event - {str(e)}" - ) + verbose_logger.exception(f"DataDogLLMObs: Error logging failure event - {str(e)}") async def async_send_batch(self): try: if not self.log_queue: return - verbose_logger.debug( - f"DataDogLLMObs: Flushing {len(self.log_queue)} events" - ) + verbose_logger.debug(f"DataDogLLMObs: Flushing {len(self.log_queue)} events") if self.is_mock_mode: - verbose_logger.debug( - "[DATADOG MOCK] Mock mode enabled - API calls will be intercepted" - ) + verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") # Prepare the payload payload = { @@ -215,9 +187,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): try: verbose_logger.debug("payload %s", safe_dumps(payload)) except Exception as debug_error: - verbose_logger.debug( - "payload serialization failed: %s", str(debug_error) - ) + verbose_logger.debug("payload serialization failed: %s", str(debug_error)) json_payload = safe_dumps(payload) @@ -237,27 +207,17 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) if self.is_mock_mode: - verbose_logger.debug( - f"[DATADOG MOCK] Batch of {len(self.log_queue)} events successfully mocked" - ) + verbose_logger.debug(f"[DATADOG MOCK] Batch of {len(self.log_queue)} events successfully mocked") else: - verbose_logger.debug( - f"DataDogLLMObs: Successfully sent batch - status_code: {response.status_code}" - ) + verbose_logger.debug(f"DataDogLLMObs: Successfully sent batch - status_code: {response.status_code}") self.log_queue.clear() except httpx.HTTPStatusError as e: - verbose_logger.exception( - f"DataDogLLMObs: Error sending batch - {e.response.text}" - ) + verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {e.response.text}") except Exception as e: verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {str(e)}") - def create_llm_obs_payload( - self, kwargs: Dict, start_time: datetime, end_time: datetime - ) -> LLMObsPayload: - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + def create_llm_obs_payload(self, kwargs: Dict, start_time: datetime, end_time: datetime) -> LLMObsPayload: + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise Exception("DataDogLLMObs: standard_logging_object is not set") @@ -266,11 +226,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): metadata = kwargs.get("litellm_params", {}).get("metadata", {}) - input_meta = InputMeta( - messages=handle_any_messages_to_chat_completion_str_messages_conversion( - messages - ) - ) + input_meta = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages)) output_meta = OutputMeta( messages=self._get_response_messages( standard_logging_payload=standard_logging_payload, @@ -285,9 +241,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): metadata_parent_id = metadata.get("parent_id") meta = Meta( - kind=self._get_datadog_span_kind( - standard_logging_payload.get("call_type"), metadata_parent_id - ), + kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id), input=input_meta, output=output_meta, metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload), @@ -300,9 +254,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): output_tokens=float(standard_logging_payload.get("completion_tokens", 0)), total_tokens=float(standard_logging_payload.get("total_tokens", 0)), total_cost=float(standard_logging_payload.get("response_cost", 0)), - time_to_first_token=self._get_time_to_first_token_seconds( - standard_logging_payload - ), + time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload), ) payload: LLMObsPayload = LLMObsPayload( @@ -338,9 +290,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): pass return None - def _assemble_error_info( - self, standard_logging_payload: StandardLoggingPayload - ) -> Optional[DDLLMObsError]: + def _assemble_error_info(self, standard_logging_payload: StandardLoggingPayload) -> Optional[DDLLMObsError]: """ Assemble error information for failure cases according to DD LLM Obs API spec """ @@ -349,8 +299,8 @@ class DataDogLLMObsLogger(CustomBatchLogger): if standard_logging_payload.get("status") == "failure": # Try to get structured error information first - error_information: Optional[StandardLoggingPayloadErrorInformation] = ( - standard_logging_payload.get("error_information") + error_information: Optional[StandardLoggingPayloadErrorInformation] = standard_logging_payload.get( + "error_information" ) if error_information: @@ -363,9 +313,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) return error_info - def _get_time_to_first_token_seconds( - self, standard_logging_payload: StandardLoggingPayload - ) -> float: + def _get_time_to_first_token_seconds(self, standard_logging_payload: StandardLoggingPayload) -> float: """ Get the time to first token in seconds @@ -374,9 +322,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): For non streaming calls, CompletionStartTime is time we get the response back """ start_time: Optional[float] = standard_logging_payload.get("startTime") - completion_start_time: Optional[float] = standard_logging_payload.get( - "completionStartTime" - ) + completion_start_time: Optional[float] = standard_logging_payload.get("completionStartTime") end_time: Optional[float] = standard_logging_payload.get("endTime") if completion_start_time is not None and start_time is not None: @@ -538,9 +484,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default fallback for unknown or passthrough operations return "llm" - def _ensure_string_content( - self, messages: Optional[Union[str, List[Any], Dict[Any, Any]]] - ) -> List[Any]: + def _ensure_string_content(self, messages: Optional[Union[str, List[Any], Dict[Any, Any]]]) -> List[Any]: if messages is None: return [] if isinstance(messages, str): @@ -551,28 +495,20 @@ class DataDogLLMObsLogger(CustomBatchLogger): return [str(messages.get("content", ""))] return [] - def _get_dd_llm_obs_payload_metadata( - self, standard_logging_payload: StandardLoggingPayload - ) -> Dict[str, Any]: + def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> Dict[str, Any]: """ Fields to track in DD LLM Observability metadata from litellm standard logging payload """ _metadata: Dict[str, Any] = { "model_name": standard_logging_payload.get("model", "unknown"), - "model_provider": standard_logging_payload.get( - "custom_llm_provider", "unknown" - ), + "model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"), "id": standard_logging_payload.get("id", "unknown"), "trace_id": standard_logging_payload.get("trace_id", "unknown"), "cache_hit": standard_logging_payload.get("cache_hit", "unknown"), "cache_key": standard_logging_payload.get("cache_key", "unknown"), "saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0), - "guardrail_information": standard_logging_payload.get( - "guardrail_information", None - ), - "is_streamed_request": self._get_stream_value_from_payload( - standard_logging_payload - ), + "guardrail_information": standard_logging_payload.get("guardrail_information", None), + "is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload), } ######################################################### @@ -591,28 +527,20 @@ class DataDogLLMObsLogger(CustomBatchLogger): tool_call_metadata = self._extract_tool_call_metadata(standard_logging_payload) _metadata.update(tool_call_metadata) - _standard_logging_metadata: dict = ( - dict(standard_logging_payload.get("metadata", {})) or {} - ) + _standard_logging_metadata: dict = dict(standard_logging_payload.get("metadata", {})) or {} _metadata.update(_standard_logging_metadata) return _metadata - def _get_latency_metrics( - self, standard_logging_payload: StandardLoggingPayload - ) -> DDLLMObsLatencyMetrics: + def _get_latency_metrics(self, standard_logging_payload: StandardLoggingPayload) -> DDLLMObsLatencyMetrics: """ Get the latency metrics from the standard logging payload """ latency_metrics: DDLLMObsLatencyMetrics = DDLLMObsLatencyMetrics() # Add latency metrics to metadata # Time to first token (convert from seconds to milliseconds for consistency) - time_to_first_token_seconds = self._get_time_to_first_token_seconds( - standard_logging_payload - ) + time_to_first_token_seconds = self._get_time_to_first_token_seconds(standard_logging_payload) if time_to_first_token_seconds > 0: - latency_metrics["time_to_first_token_ms"] = ( - time_to_first_token_seconds * 1000 - ) + latency_metrics["time_to_first_token_ms"] = time_to_first_token_seconds * 1000 # LiteLLM overhead time hidden_params = standard_logging_payload.get("hidden_params", {}) @@ -621,8 +549,8 @@ class DataDogLLMObsLogger(CustomBatchLogger): latency_metrics["litellm_overhead_time_ms"] = litellm_overhead_ms # Guardrail overhead latency - guardrail_info: Optional[list[StandardLoggingGuardrailInformation]] = ( - standard_logging_payload.get("guardrail_information") + guardrail_info: Optional[list[StandardLoggingGuardrailInformation]] = standard_logging_payload.get( + "guardrail_information" ) if guardrail_info is not None: total_duration = 0.0 @@ -637,9 +565,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): return latency_metrics - def _get_stream_value_from_payload( - self, standard_logging_payload: StandardLoggingPayload - ) -> bool: + def _get_stream_value_from_payload(self, standard_logging_payload: StandardLoggingPayload) -> bool: """ Extract the stream value from standard logging payload. @@ -664,18 +590,14 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default to False for non-streaming requests return False - def _get_spend_metrics( - self, standard_logging_payload: StandardLoggingPayload - ) -> DDLLMObsSpendMetrics: + def _get_spend_metrics(self, standard_logging_payload: StandardLoggingPayload) -> DDLLMObsSpendMetrics: """ Get the spend metrics from the standard logging payload """ spend_metrics: DDLLMObsSpendMetrics = DDLLMObsSpendMetrics() # send response cost - spend_metrics["response_cost"] = standard_logging_payload.get( - "response_cost", 0.0 - ) + spend_metrics["response_cost"] = standard_logging_payload.get("response_cost", 0.0) # Get budget information from metadata metadata = standard_logging_payload.get("metadata", {}) @@ -691,9 +613,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): try: spend_metrics["user_api_key_spend"] = float(user_api_key_spend) except (ValueError, TypeError): - verbose_logger.debug( - f"Invalid user_api_key_spend value: {user_api_key_spend}" - ) + verbose_logger.debug(f"Invalid user_api_key_spend value: {user_api_key_spend}") # API key budget reset datetime user_api_key_budget_reset_at = metadata.get("user_api_key_budget_reset_at") @@ -720,18 +640,14 @@ class DataDogLLMObsLogger(CustomBatchLogger): spend_metrics["user_api_key_budget_reset_at"] = iso_string # Debug logging to verify the conversion - verbose_logger.debug( - f"Converted budget_reset_at to ISO format: {iso_string}" - ) + verbose_logger.debug(f"Converted budget_reset_at to ISO format: {iso_string}") except Exception as e: verbose_logger.debug(f"Error processing budget reset datetime: {e}") verbose_logger.debug(f"Original value: {user_api_key_budget_reset_at}") return spend_metrics - def _process_input_messages_preserving_tool_calls( - self, messages: List[Any] - ) -> List[Dict[str, Any]]: + def _process_input_messages_preserving_tool_calls(self, messages: List[Any]) -> List[Dict[str, Any]]: """ Process input messages while preserving tool_calls and tool message types. @@ -746,19 +662,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): processed.append(msg) else: # For regular messages, still apply string conversion - converted = ( - handle_any_messages_to_chat_completion_str_messages_conversion( - [msg] - ) - ) + converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg]) processed.extend(converted) else: # For non-dict messages, apply string conversion - converted = ( - handle_any_messages_to_chat_completion_str_messages_conversion( - [msg] - ) - ) + converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg]) processed.extend(converted) return processed @@ -793,26 +701,18 @@ class DataDogLLMObsLogger(CustomBatchLogger): if function_arguments: # Store arguments as JSON string for Datadog if isinstance(function_arguments, str): - kv_pairs[f"tool_calls.{idx}.function.arguments"] = ( - function_arguments - ) + kv_pairs[f"tool_calls.{idx}.function.arguments"] = function_arguments else: import json - kv_pairs[f"tool_calls.{idx}.function.arguments"] = ( - json.dumps(function_arguments) - ) + kv_pairs[f"tool_calls.{idx}.function.arguments"] = json.dumps(function_arguments) except (KeyError, TypeError, ValueError) as e: - verbose_logger.debug( - f"DataDogLLMObs: Error processing tool call {idx}: {str(e)}" - ) + verbose_logger.debug(f"DataDogLLMObs: Error processing tool call {idx}: {str(e)}") continue return kv_pairs - def _extract_tool_call_metadata( - self, standard_logging_payload: StandardLoggingPayload - ) -> Dict[str, Any]: + def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> Dict[str, Any]: """ Extract tool call information from both input messages and response for Datadog metadata. """ @@ -841,16 +741,12 @@ class DataDogLLMObsLogger(CustomBatchLogger): if message and isinstance(message, dict): tool_calls = message.get("tool_calls") if tool_calls: - response_tool_calls_kv = self._tool_calls_kv_pair( - tool_calls - ) + response_tool_calls_kv = self._tool_calls_kv_pair(tool_calls) # Prefix with "output_" to distinguish from input tool calls for key, value in response_tool_calls_kv.items(): tool_call_metadata[f"output_{key}"] = value except Exception as e: - verbose_logger.debug( - f"DataDogLLMObs: Error extracting tool call metadata: {str(e)}" - ) + verbose_logger.debug(f"DataDogLLMObs: Error extracting tool call metadata: {str(e)}") return tool_call_metadata diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index bd5c165cba2..b1e4bc73e77 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -34,15 +34,11 @@ class DatadogMetricsLogger(CustomBatchLogger): self.dd_site = os.getenv("DD_SITE", "datadoghq.com") if not self.dd_api_key: - verbose_logger.warning( - "Datadog Metrics: DD_API_KEY is required. Integration will not work." - ) + verbose_logger.warning("Datadog Metrics: DD_API_KEY is required. Integration will not work.") self.upload_url = f"https://api.{self.dd_site}/api/v2/series" - self.async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Initialize lock self.flush_lock = asyncio.Lock() @@ -155,8 +151,7 @@ class DatadogMetricsLogger(CustomBatchLogger): "points": [ { "timestamp": timestamp, - "value": litellm_overhead_time_ms - / 1000, # convert ms → seconds + "value": litellm_overhead_time_ms / 1000, # convert ms → seconds } ], "tags": overhead_tags, @@ -175,54 +170,40 @@ class DatadogMetricsLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: return - self._add_metrics_from_log( - log=standard_logging_object, kwargs=kwargs, status_code="200" - ) + self._add_metrics_from_log(log=standard_logging_object, kwargs=kwargs, status_code="200") if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception as e: - verbose_logger.exception( - f"Datadog Metrics: Error in async_log_success_event: {str(e)}" - ) + verbose_logger.exception(f"Datadog Metrics: Error in async_log_success_event: {str(e)}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: return # Extract status code from error information status_code = "500" # default - error_information = ( - standard_logging_object.get("error_information", {}) or {} - ) + error_information = standard_logging_object.get("error_information", {}) or {} error_code = error_information.get("error_code") # type: ignore if error_code is not None: status_code = str(error_code) - self._add_metrics_from_log( - log=standard_logging_object, kwargs=kwargs, status_code=status_code - ) + self._add_metrics_from_log(log=standard_logging_object, kwargs=kwargs, status_code=status_code) if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception as e: - verbose_logger.exception( - f"Datadog Metrics: Error in async_log_failure_event: {str(e)}" - ) + verbose_logger.exception(f"Datadog Metrics: Error in async_log_failure_event: {str(e)}") async def async_send_batch(self): if not self.log_queue: @@ -234,9 +215,7 @@ class DatadogMetricsLogger(CustomBatchLogger): try: await self._upload_to_datadog(payload_data) except Exception as e: - verbose_logger.exception( - f"Datadog Metrics: Error in async_send_batch: {str(e)}" - ) + verbose_logger.exception(f"Datadog Metrics: Error in async_send_batch: {str(e)}") raise async def _upload_to_datadog(self, payload: DatadogMetricsPayload): diff --git a/litellm/integrations/datadog/datadog_mock_client.py b/litellm/integrations/datadog/datadog_mock_client.py index 7f9beab72cc..c50cdc6a019 100644 --- a/litellm/integrations/datadog/datadog_mock_client.py +++ b/litellm/integrations/datadog/datadog_mock_client.py @@ -28,6 +28,4 @@ _config = MockClientConfig( patch_sync_client=True, ) -create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory( - _config -) +create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory(_config) diff --git a/litellm/integrations/datadog/datadog_team_handler.py b/litellm/integrations/datadog/datadog_team_handler.py index 3a5b73fc005..cae954f753c 100644 --- a/litellm/integrations/datadog/datadog_team_handler.py +++ b/litellm/integrations/datadog/datadog_team_handler.py @@ -54,11 +54,9 @@ class DataDogHandler: # if not cached, create a new datadog logger and cache it if temp_datadog_logger is None: - temp_datadog_logger = ( - DataDogHandler._create_datadog_logger_from_credentials( - credentials=credentials_dict, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) + temp_datadog_logger = DataDogHandler._create_datadog_logger_from_credentials( + credentials=credentials_dict, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, ) return temp_datadog_logger @@ -73,10 +71,7 @@ class DataDogHandler: """ # When the destination is caller-supplied (dd_agent_host/dd_site), never fall back to the # proxy's DD_API_KEY env var, otherwise it would be sent to a team-controlled host. - allow_env_credentials = ( - credentials.get("dd_agent_host") is None - and credentials.get("dd_site") is None - ) + allow_env_credentials = credentials.get("dd_agent_host") is None and credentials.get("dd_site") is None datadog_logger = DataDogLogger( dd_api_key=credentials.get("dd_api_key"), dd_site=credentials.get("dd_site"), @@ -89,9 +84,7 @@ class DataDogHandler: service_name="datadog", logging_obj=datadog_logger, ) - verbose_logger.debug( - "Datadog: Created and cached new DataDogLogger for team-scoped credentials" - ) + verbose_logger.debug("Datadog: Created and cached new DataDogLogger for team-scoped credentials") return datadog_logger @staticmethod diff --git a/litellm/integrations/deepeval/api.py b/litellm/integrations/deepeval/api.py index 5e446e26feb..fccc5970433 100644 --- a/litellm/integrations/deepeval/api.py +++ b/litellm/integrations/deepeval/api.py @@ -58,13 +58,9 @@ class Api: # using the global non-eu variable for base url self.base_api_url = base_url or API_BASE_URL self.sync_http_handler = HTTPHandler() - self.async_http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - def _http_request( - self, method: str, url: str, headers=None, json=None, params=None - ): + def _http_request(self, method: str, url: str, headers=None, json=None, params=None): if method != "POST": raise Exception("Only POST requests are supported") try: @@ -79,9 +75,7 @@ class Api: except Exception as e: raise e - def send_request( - self, method: HttpMethods, endpoint: Endpoints, body=None, params=None - ): + def send_request(self, method: HttpMethods, endpoint: Endpoints, body=None, params=None): url = f"{self.base_api_url}{endpoint.value}" res = self._http_request( method=method.value, @@ -100,9 +94,7 @@ class Api: verbose_logger.debug(res.json()) raise Exception(res.json().get("error", res.text)) - async def a_send_request( - self, method: HttpMethods, endpoint: Endpoints, body=None, params=None - ): + async def a_send_request(self, method: HttpMethods, endpoint: Endpoints, body=None, params=None): if method != HttpMethods.POST: raise Exception("Only POST requests are supported") diff --git a/litellm/integrations/deepeval/deepeval.py b/litellm/integrations/deepeval/deepeval.py index 972843e120a..90c1d8eedce 100644 --- a/litellm/integrations/deepeval/deepeval.py +++ b/litellm/integrations/deepeval/deepeval.py @@ -25,39 +25,27 @@ class DeepEvalLogger(CustomLogger): self.litellm_environment = os.getenv("LITELM_ENVIRONMENT", "development") validate_environment(self.litellm_environment) if not api_key: - raise ValueError( - "Please set 'CONFIDENT_API_KEY=<>' in your environment variables." - ) + raise ValueError("Please set 'CONFIDENT_API_KEY=<>' in your environment variables.") self.api = Api(api_key=api_key) super().__init__(*args, **kwargs) def log_success_event(self, kwargs, response_obj, start_time, end_time): """Logs a success event to DeepEval's platform.""" - self._sync_event_handler( - kwargs, response_obj, start_time, end_time, is_success=True - ) + self._sync_event_handler(kwargs, response_obj, start_time, end_time, is_success=True) def log_failure_event(self, kwargs, response_obj, start_time, end_time): """Logs a failure event to DeepEval's platform.""" - self._sync_event_handler( - kwargs, response_obj, start_time, end_time, is_success=False - ) + self._sync_event_handler(kwargs, response_obj, start_time, end_time, is_success=False) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """Logs a failure event to DeepEval's platform.""" - await self._async_event_handler( - kwargs, response_obj, start_time, end_time, is_success=False - ) + await self._async_event_handler(kwargs, response_obj, start_time, end_time, is_success=False) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Logs a success event to DeepEval's platform.""" - await self._async_event_handler( - kwargs, response_obj, start_time, end_time, is_success=True - ) + await self._async_event_handler(kwargs, response_obj, start_time, end_time, is_success=True) - def _prepare_trace_api( - self, kwargs, response_obj, start_time, end_time, is_success - ): + def _prepare_trace_api(self, kwargs, response_obj, start_time, end_time, is_success): _start_time = to_zod_compatible_iso(start_time) _end_time = to_zod_compatible_iso(end_time) _standard_logging_object = kwargs.get("standard_logging_object", {}) @@ -85,12 +73,8 @@ class DeepEvalLogger(CustomLogger): body = trace_api.dict(by_alias=True, exclude_none=True) return body - def _sync_event_handler( - self, kwargs, response_obj, start_time, end_time, is_success - ): - body = self._prepare_trace_api( - kwargs, response_obj, start_time, end_time, is_success - ) + def _sync_event_handler(self, kwargs, response_obj, start_time, end_time, is_success): + body = self._prepare_trace_api(kwargs, response_obj, start_time, end_time, is_success) try: response = self.api.send_request( method=HttpMethods.POST, @@ -99,29 +83,19 @@ class DeepEvalLogger(CustomLogger): ) except Exception as e: raise e - verbose_logger.debug( - "DeepEvalLogger: sync_log_failure_event: Api response %s", response - ) + verbose_logger.debug("DeepEvalLogger: sync_log_failure_event: Api response %s", response) - async def _async_event_handler( - self, kwargs, response_obj, start_time, end_time, is_success - ): - body = self._prepare_trace_api( - kwargs, response_obj, start_time, end_time, is_success - ) + async def _async_event_handler(self, kwargs, response_obj, start_time, end_time, is_success): + body = self._prepare_trace_api(kwargs, response_obj, start_time, end_time, is_success) response = await self.api.a_send_request( method=HttpMethods.POST, endpoint=Endpoints.TRACING_ENDPOINT, body=body, ) - verbose_logger.debug( - "DeepEvalLogger: async_event_handler: Api response %s", response - ) + verbose_logger.debug("DeepEvalLogger: async_event_handler: Api response %s", response) - def _create_base_api_span( - self, kwargs, standard_logging_object, start_time, end_time, is_success - ): + def _create_base_api_span(self, kwargs, standard_logging_object, start_time, end_time, is_success): # extract usage usage = standard_logging_object.get("response", {}).get("usage", {}) if is_success: @@ -135,12 +109,8 @@ class DeepEvalLogger(CustomLogger): output = str(standard_logging_object.get("error_string", "")) return BaseApiSpan( uuid=standard_logging_object.get("id", uuid.uuid4()), - name=( - "litellm_success_callback" if is_success else "litellm_failure_callback" - ), - status=( - TraceSpanApiStatus.SUCCESS if is_success else TraceSpanApiStatus.ERRORED - ), + name=("litellm_success_callback" if is_success else "litellm_failure_callback"), + status=(TraceSpanApiStatus.SUCCESS if is_success else TraceSpanApiStatus.ERRORED), type=SpanApiType.LLM, traceUuid=standard_logging_object.get("trace_id", uuid.uuid4()), startTime=str(start_time), @@ -149,9 +119,7 @@ class DeepEvalLogger(CustomLogger): output=output, model=standard_logging_object.get("model", None), inputTokenCount=usage.get("prompt_tokens", None) if is_success else None, - outputTokenCount=( - usage.get("completion_tokens", None) if is_success else None - ), + outputTokenCount=(usage.get("completion_tokens", None) if is_success else None), ) def _create_trace_api( diff --git a/litellm/integrations/deepeval/utils.py b/litellm/integrations/deepeval/utils.py index 0beb22db9e3..3df9aceb241 100644 --- a/litellm/integrations/deepeval/utils.py +++ b/litellm/integrations/deepeval/utils.py @@ -3,16 +3,10 @@ from litellm.integrations.deepeval.types import Environment def to_zod_compatible_iso(dt: datetime) -> str: - return ( - dt.astimezone(timezone.utc) - .isoformat(timespec="milliseconds") - .replace("+00:00", "Z") - ) + return dt.astimezone(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z") def validate_environment(environment: str): if environment not in [env.value for env in Environment]: valid_values = ", ".join(f'"{env.value}"' for env in Environment) - raise ValueError( - f"Invalid environment: {environment}. Please use one of the following instead: {valid_values}" - ) + raise ValueError(f"Invalid environment: {environment}. Please use one of the following instead: {valid_values}") diff --git a/litellm/integrations/dotprompt/__init__.py b/litellm/integrations/dotprompt/__init__.py index 394929f4a25..8432d50e32b 100644 --- a/litellm/integrations/dotprompt/__init__.py +++ b/litellm/integrations/dotprompt/__init__.py @@ -42,9 +42,7 @@ def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict: return {"content": content.strip(), "metadata": metadata} -def prompt_initializer( - litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec" -) -> "CustomPromptManagement": +def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec") -> "CustomPromptManagement": """ Initialize a prompt from a .prompt file. """ diff --git a/litellm/integrations/dotprompt/dotprompt_manager.py b/litellm/integrations/dotprompt/dotprompt_manager.py index 37fdf7da693..3ba9efd68b7 100644 --- a/litellm/integrations/dotprompt/dotprompt_manager.py +++ b/litellm/integrations/dotprompt/dotprompt_manager.py @@ -69,11 +69,7 @@ class DotpromptManager(CustomPromptManagement): def prompt_manager(self) -> PromptManager: """Lazy-load the prompt manager.""" if self._prompt_manager is None: - if ( - self.prompt_directory is None - and not self.prompt_data - and not self.prompt_file - ): + if self.prompt_directory is None and not self.prompt_data and not self.prompt_file: raise ValueError( "Either prompt_directory or prompt_data must be set before using dotprompt manager. " "Set litellm.global_prompt_directory, initialize with prompt_directory parameter, or provide prompt_data." @@ -129,14 +125,10 @@ class DotpromptManager(CustomPromptManagement): try: # Get the prompt template (versioned or base) - template = self.prompt_manager.get_prompt( - prompt_id=prompt_id, version=prompt_version - ) + template = self.prompt_manager.get_prompt(prompt_id=prompt_id, version=prompt_version) if template is None: version_str = f" (version {prompt_version})" if prompt_version else "" - raise ValueError( - f"Prompt '{prompt_id}'{version_str} not found in prompt directory" - ) + raise ValueError(f"Prompt '{prompt_id}'{version_str} not found in prompt directory") # Render the template with variables (pass version for proper lookup) rendered_content = self.prompt_manager.render( @@ -282,29 +274,17 @@ class DotpromptManager(CustomPromptManagement): # Check for role prefixes if line.startswith("System:"): if current_role and current_content: - messages.append( - self._create_message( - current_role, "\n".join(current_content).strip() - ) - ) + messages.append(self._create_message(current_role, "\n".join(current_content).strip())) current_role = "system" current_content = [line[7:].strip()] # Remove "System:" prefix elif line.startswith("User:"): if current_role and current_content: - messages.append( - self._create_message( - current_role, "\n".join(current_content).strip() - ) - ) + messages.append(self._create_message(current_role, "\n".join(current_content).strip())) current_role = "user" current_content = [line[5:].strip()] # Remove "User:" prefix elif line.startswith("Assistant:"): if current_role and current_content: - messages.append( - self._create_message( - current_role, "\n".join(current_content).strip() - ) - ) + messages.append(self._create_message(current_role, "\n".join(current_content).strip())) current_role = "assistant" current_content = [line[10:].strip()] # Remove "Assistant:" prefix else: diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index 6407a18d0b3..dd198ba1272 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -93,9 +93,7 @@ class PromptManager: def _load_prompts(self) -> None: """Load all .prompt files from the prompt directory.""" if not self.prompt_directory or not self.prompt_directory.exists(): - raise ValueError( - f"Prompt directory does not exist: {self.prompt_directory}" - ) + raise ValueError(f"Prompt directory does not exist: {self.prompt_directory}") prompt_files = list(self.prompt_directory.glob("*.prompt")) @@ -109,9 +107,7 @@ class PromptManager: # Optional: print(f"Error loading prompt file {prompt_file}") pass - def _load_prompts_from_json( - self, prompt_data: Dict[str, Dict[str, Any]], prompt_id: Optional[str] = None - ) -> None: + def _load_prompts_from_json(self, prompt_data: Dict[str, Dict[str, Any]], prompt_id: Optional[str] = None) -> None: """Load prompts from JSON data structure. Expected format: @@ -147,9 +143,7 @@ class PromptManager: # Optional: print(f"Error loading prompt from JSON: {prompt_id}") pass - def _load_prompt_file( - self, file_path: Union[str, Path], prompt_id: str - ) -> PromptTemplate: + def _load_prompt_file(self, file_path: Union[str, Path], prompt_id: str) -> PromptTemplate: """Load and parse a single .prompt file.""" if isinstance(file_path, str): file_path = Path(file_path) @@ -213,9 +207,7 @@ class PromptManager: if template is None: available_prompts = list(self.prompts.keys()) version_str = f" (version {version})" if version else "" - raise KeyError( - f"Prompt '{prompt_id}'{version_str} not found. Available prompts: {available_prompts}" - ) + raise KeyError(f"Prompt '{prompt_id}'{version_str} not found. Available prompts: {available_prompts}") variables = prompt_variables or {} @@ -231,9 +223,7 @@ class PromptManager: except Exception as e: raise ValueError(f"Error rendering template '{prompt_id}': {e}") - def _validate_input( - self, variables: Dict[str, Any], schema: Dict[str, Any] - ) -> None: + def _validate_input(self, variables: Dict[str, Any], schema: Dict[str, Any]) -> None: """Basic validation of input variables against schema.""" for field_name, field_type in schema.items(): if field_name in variables: @@ -265,9 +255,7 @@ class PromptManager: return type_mapping.get(schema_type.lower(), str) # type: ignore - def get_prompt( - self, prompt_id: str, version: Optional[int] = None - ) -> Optional[PromptTemplate]: + def get_prompt(self, prompt_id: str, version: Optional[int] = None) -> Optional[PromptTemplate]: """ Get a prompt template by ID and optional version. @@ -302,13 +290,9 @@ class PromptManager: if self.prompt_directory: self._load_prompts() - def add_prompt( - self, prompt_id: str, content: str, metadata: Optional[Dict[str, Any]] = None - ) -> None: + def add_prompt(self, prompt_id: str, content: str, metadata: Optional[Dict[str, Any]] = None) -> None: """Add a prompt template programmatically.""" - template = PromptTemplate( - content=content, metadata=metadata or {}, template_id=prompt_id - ) + template = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id) self.prompts[prompt_id] = template def prompt_file_to_json(self, file_path: Union[str, Path]) -> Dict[str, Any]: @@ -365,8 +349,6 @@ class PromptManager: } return result - def load_prompts_from_json_data( - self, prompt_data: Dict[str, Dict[str, Any]] - ) -> None: + def load_prompts_from_json_data(self, prompt_data: Dict[str, Dict[str, Any]]) -> None: """Load additional prompts from JSON data (merges with existing prompts).""" self._load_prompts_from_json(prompt_data) diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index dfc05ae1f32..ab76fa3c8bd 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -16,32 +16,24 @@ class DyanmoDBLogger: # Instance variables import boto3 - self.dynamodb: Any = boto3.resource( - "dynamodb", region_name=os.environ["AWS_REGION_NAME"] - ) + self.dynamodb: Any = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"]) if litellm.dynamodb_table_name is None: raise ValueError( "LiteLLM Error, trying to use DynamoDB but not table name passed. Create a table and set `litellm.dynamodb_table_name=`" ) self.table_name = litellm.dynamodb_table_name - async def _async_log_event( - self, kwargs, response_obj, start_time, end_time, print_verbose - ): + async def _async_log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): self.log_event(kwargs, response_obj, start_time, end_time, print_verbose) def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): try: - print_verbose( - f"DynamoDB Logging - Enters logging function for model {kwargs}" - ) + print_verbose(f"DynamoDB Logging - Enters logging function for model {kwargs}") # construct payload to send to DynamoDB # follows the same params as langfuse.py litellm_params = kwargs.get("litellm_params", {}) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None + metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None messages = kwargs.get("messages") optional_params = kwargs.get("optional_params", {}) call_type = kwargs.get("call_type", "litellm.completion") @@ -80,9 +72,7 @@ class DyanmoDBLogger: print_verbose(f"Response from DynamoDB:{str(response)}") - print_verbose( - f"DynamoDB Layer Logging - final response object: {response_obj}" - ) + print_verbose(f"DynamoDB Layer Logging - final response object: {response_obj}") return response except Exception: print_verbose(f"DynamoDB Layer Error - {traceback.format_exc()}") diff --git a/litellm/integrations/email_alerting.py b/litellm/integrations/email_alerting.py index b721dc50464..35d63a691f9 100644 --- a/litellm/integrations/email_alerting.py +++ b/litellm/integrations/email_alerting.py @@ -15,9 +15,7 @@ LITELLM_SUPPORT_CONTACT = "support@berri.ai" async def get_all_team_member_emails(team_id: Optional[str] = None) -> list: - verbose_logger.debug( - "Email Alerting: Getting all team members for team_id=%s", team_id - ) + verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id) if team_id is None: return [] from litellm.proxy.proxy_server import prisma_client @@ -76,9 +74,7 @@ async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool: _team_id = webhook_event.team_id team_alias = webhook_event.team_alias - verbose_logger.debug( - "Email Alerting: Sending Team Budget Alert for team=%s", team_alias - ) + verbose_logger.debug("Email Alerting: Sending Team Budget Alert for team=%s", team_alias) email_logo_url = os.getenv("SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", None)) email_support_contact = os.getenv("EMAIL_SUPPORT_CONTACT", None) @@ -93,9 +89,7 @@ async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool: email_support_contact = LITELLM_SUPPORT_CONTACT recipient_emails = await get_all_team_member_emails(_team_id) recipient_emails_str: str = ",".join(recipient_emails) - verbose_logger.debug( - "Email Alerting: Sending team budget alert to %s", recipient_emails_str - ) + verbose_logger.debug("Email Alerting: Sending team budget alert to %s", recipient_emails_str) event_name = webhook_event.event_message max_budget = webhook_event.max_budget diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index cd25a87729f..3d79046bf6c 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -24,9 +24,7 @@ class FocusDestinationFactory: ) -> FocusDestination: """Return a destination implementation for the requested provider.""" provider_lower = provider.lower() - normalized_config = FocusDestinationFactory._resolve_config( - provider=provider_lower, overrides=config or {} - ) + normalized_config = FocusDestinationFactory._resolve_config(provider=provider_lower, overrides=config or {}) if provider_lower == "s3": return FocusS3Destination(prefix=prefix, config=normalized_config) if provider_lower == "vantage": @@ -35,9 +33,7 @@ class FocusDestinationFactory: return FocusGCSDestination(prefix=prefix, config=normalized_config) if provider_lower == "mavvrik": return FocusMavvrikDestination(prefix=prefix, config=normalized_config) - raise NotImplementedError( - f"Provider '{provider}' not supported for Focus export" - ) + raise NotImplementedError(f"Provider '{provider}' not supported for Focus export") @staticmethod def _resolve_config( @@ -47,18 +43,12 @@ class FocusDestinationFactory: ) -> Dict[str, Any]: if provider == "s3": resolved = { - "bucket_name": overrides.get("bucket_name") - or os.getenv("FOCUS_S3_BUCKET_NAME"), - "region_name": overrides.get("region_name") - or os.getenv("FOCUS_S3_REGION_NAME"), - "endpoint_url": overrides.get("endpoint_url") - or os.getenv("FOCUS_S3_ENDPOINT_URL"), - "aws_access_key_id": overrides.get("aws_access_key_id") - or os.getenv("FOCUS_S3_ACCESS_KEY"), - "aws_secret_access_key": overrides.get("aws_secret_access_key") - or os.getenv("FOCUS_S3_SECRET_KEY"), - "aws_session_token": overrides.get("aws_session_token") - or os.getenv("FOCUS_S3_SESSION_TOKEN"), + "bucket_name": overrides.get("bucket_name") or os.getenv("FOCUS_S3_BUCKET_NAME"), + "region_name": overrides.get("region_name") or os.getenv("FOCUS_S3_REGION_NAME"), + "endpoint_url": overrides.get("endpoint_url") or os.getenv("FOCUS_S3_ENDPOINT_URL"), + "aws_access_key_id": overrides.get("aws_access_key_id") or os.getenv("FOCUS_S3_ACCESS_KEY"), + "aws_secret_access_key": overrides.get("aws_secret_access_key") or os.getenv("FOCUS_S3_SECRET_KEY"), + "aws_session_token": overrides.get("aws_session_token") or os.getenv("FOCUS_S3_SESSION_TOKEN"), } if not resolved.get("bucket_name"): raise ValueError("FOCUS_S3_BUCKET_NAME must be provided for S3 exports") @@ -66,39 +56,28 @@ class FocusDestinationFactory: if provider == "vantage": resolved = { "api_key": overrides.get("api_key") or os.getenv("VANTAGE_API_KEY"), - "integration_token": overrides.get("integration_token") - or os.getenv("VANTAGE_INTEGRATION_TOKEN"), - "base_url": overrides.get("base_url") - or os.getenv("VANTAGE_BASE_URL", "https://api.vantage.sh"), + "integration_token": overrides.get("integration_token") or os.getenv("VANTAGE_INTEGRATION_TOKEN"), + "base_url": overrides.get("base_url") or os.getenv("VANTAGE_BASE_URL", "https://api.vantage.sh"), } if not resolved.get("api_key"): raise ValueError("VANTAGE_API_KEY must be provided for Vantage exports") if not resolved.get("integration_token"): - raise ValueError( - "VANTAGE_INTEGRATION_TOKEN must be provided for Vantage exports" - ) + raise ValueError("VANTAGE_INTEGRATION_TOKEN must be provided for Vantage exports") return {k: v for k, v in resolved.items() if v is not None} if provider == "gcs": resolved = { - "bucket_name": overrides.get("bucket_name") - or os.getenv("FOCUS_GCS_BUCKET_NAME"), + "bucket_name": overrides.get("bucket_name") or os.getenv("FOCUS_GCS_BUCKET_NAME"), "service_account_json": overrides.get("service_account_json") or os.getenv("FOCUS_GCS_PATH_SERVICE_ACCOUNT"), } if not resolved.get("bucket_name"): - raise ValueError( - "FOCUS_GCS_BUCKET_NAME must be provided for GCS exports" - ) + raise ValueError("FOCUS_GCS_BUCKET_NAME must be provided for GCS exports") return {k: v for k, v in resolved.items() if v is not None} if provider == "mavvrik": resolved = { "api_key": overrides.get("api_key") or os.getenv("MAVVRIK_API_KEY"), - "api_endpoint": overrides.get("api_endpoint") - or os.getenv("MAVVRIK_API_ENDPOINT"), - "connection_id": overrides.get("connection_id") - or os.getenv("MAVVRIK_CONNECTION_ID"), + "api_endpoint": overrides.get("api_endpoint") or os.getenv("MAVVRIK_API_ENDPOINT"), + "connection_id": overrides.get("connection_id") or os.getenv("MAVVRIK_CONNECTION_ID"), } return {k: v for k, v in resolved.items() if v is not None} - raise NotImplementedError( - f"Provider '{provider}' not supported for Focus export configuration" - ) + raise NotImplementedError(f"Provider '{provider}' not supported for Focus export configuration") diff --git a/litellm/integrations/focus/destinations/gcs_destination.py b/litellm/integrations/focus/destinations/gcs_destination.py index b04c16c9d32..e4525ccd267 100644 --- a/litellm/integrations/focus/destinations/gcs_destination.py +++ b/litellm/integrations/focus/destinations/gcs_destination.py @@ -41,22 +41,16 @@ class FocusGCSDestination(GCSBucketBase, FocusDestination): filename: str, ) -> None: object_name = self._build_object_key(time_window=time_window, filename=filename) - headers = await self.construct_request_headers( - service_account_json=self.path_service_account_json - ) + headers = await self.construct_request_headers(service_account_json=self.path_service_account_json) headers["Content-Type"] = "application/octet-stream" encoded_name = encode_gcs_object_name_for_url(object_name) url = ( f"https://storage.googleapis.com/upload/storage/v1/b/" f"{self.BUCKET_NAME}/o?uploadType=media&name={encoded_name}" ) - response = await self.async_httpx_client.post( - url=url, headers=headers, data=content - ) + response = await self.async_httpx_client.post(url=url, headers=headers, data=content) if response.status_code != 200: - raise RuntimeError( - f"GCS upload failed: status={response.status_code} body={response.text}" - ) + raise RuntimeError(f"GCS upload failed: status={response.status_code} body={response.text}") verbose_logger.debug( "Focus GCS: uploaded %d bytes to gs://%s/%s", len(content), diff --git a/litellm/integrations/focus/destinations/mavvrik_destination.py b/litellm/integrations/focus/destinations/mavvrik_destination.py index 81944ba69e2..cf500a71b52 100644 --- a/litellm/integrations/focus/destinations/mavvrik_destination.py +++ b/litellm/integrations/focus/destinations/mavvrik_destination.py @@ -41,14 +41,9 @@ def _validate_api_endpoint(api_endpoint: str) -> None: def _validate_gcs_url(url: str, label: str) -> None: parsed = urlparse(url) if parsed.scheme != "https": - raise ValueError( - f"Mavvrik FOCUS destination: {label} must be HTTPS, got scheme '{parsed.scheme}'" - ) + raise ValueError(f"Mavvrik FOCUS destination: {label} must be HTTPS, got scheme '{parsed.scheme}'") hostname = (parsed.hostname or "").lower() - if not ( - hostname == "storage.googleapis.com" - or hostname.endswith(".storage.googleapis.com") - ): + if not (hostname == "storage.googleapis.com" or hostname.endswith(".storage.googleapis.com")): raise ValueError( f"Mavvrik FOCUS destination: {label} must be a GCS endpoint (storage.googleapis.com), got '{hostname}'" ) @@ -90,9 +85,7 @@ class FocusMavvrikDestination(FocusDestination): self.api_endpoint = api_endpoint.rstrip("/") self.connection_id = connection_id self.prefix = prefix - self._http: AsyncHTTPHandler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self._http: AsyncHTTPHandler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self._registered = False @property @@ -132,9 +125,7 @@ class FocusMavvrikDestination(FocusDestination): "Re-enable the connection in the Mavvrik dashboard." ) if resp.status_code >= 400: - raise RuntimeError( - f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}" - ) + raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}") self._registered = True metrics_marker = resp.json().get("metricsMarker", 0) verbose_logger.debug( @@ -159,13 +150,9 @@ class FocusMavvrikDestination(FocusDestination): ) signed_url = resp.json().get("url") if not signed_url: - raise RuntimeError( - f"Mavvrik FOCUS destination: response missing 'url' field: {resp.json()}" - ) + raise RuntimeError(f"Mavvrik FOCUS destination: response missing 'url' field: {resp.json()}") _validate_gcs_url(signed_url, "signed URL") - verbose_logger.debug( - "Mavvrik FOCUS destination: got signed URL for date %s", date_str - ) + verbose_logger.debug("Mavvrik FOCUS destination: got signed URL for date %s", date_str) return signed_url async def _upload_to_gcs(self, signed_url: str, content: bytes) -> None: @@ -205,9 +192,7 @@ class FocusMavvrikDestination(FocusDestination): session_uri = init_resp.headers.get("Location") if not session_uri: - raise RuntimeError( - "Mavvrik FOCUS destination: GCS session init missing Location header" - ) + raise RuntimeError("Mavvrik FOCUS destination: GCS session init missing Location header") _validate_gcs_url(session_uri, "session URI") verbose_logger.debug( @@ -224,11 +209,7 @@ class FocusMavvrikDestination(FocusDestination): chunk = gzip_bytes[offset : offset + _GCS_CHUNK_SIZE] chunk_end = offset + len(chunk) - 1 is_final = (offset + len(chunk)) >= total - content_range = ( - f"bytes {offset}-{chunk_end}/{total}" - if is_final - else f"bytes {offset}-{chunk_end}/*" - ) + content_range = f"bytes {offset}-{chunk_end}/{total}" if is_final else f"bytes {offset}-{chunk_end}/*" expected_statuses = {200, 201} if is_final else {308} resp = await self._http.client.request( @@ -256,12 +237,8 @@ class FocusMavvrikDestination(FocusDestination): except Exception: # Cancel the open GCS session so it doesn't linger for up to 1 week. try: - await self._http.client.request( - method="DELETE", url=session_uri, timeout=10.0 - ) - verbose_logger.debug( - "Mavvrik FOCUS destination: cancelled GCS session after error" - ) + await self._http.client.request(method="DELETE", url=session_uri, timeout=10.0) + verbose_logger.debug("Mavvrik FOCUS destination: cancelled GCS session after error") except Exception: pass raise @@ -283,12 +260,9 @@ class FocusMavvrikDestination(FocusDestination): ) if resp.status_code >= 400: raise RuntimeError( - f"Mavvrik FOCUS destination: failed to update metricsMarker " - f"({resp.status_code}): {resp.text[:200]}" + f"Mavvrik FOCUS destination: failed to update metricsMarker ({resp.status_code}): {resp.text[:200]}" ) - verbose_logger.debug( - "Mavvrik FOCUS destination: metricsMarker advanced to %s", date_epoch - ) + verbose_logger.debug("Mavvrik FOCUS destination: metricsMarker advanced to %s", date_epoch) async def get_metrics_marker(self) -> Optional[int]: """Register with Mavvrik and return the current metricsMarker. @@ -311,14 +285,10 @@ class FocusMavvrikDestination(FocusDestination): "Re-enable the connection in the Mavvrik dashboard." ) if resp.status_code >= 400: - raise RuntimeError( - f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}" - ) + raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}") self._registered = True metrics_marker = resp.json().get("metricsMarker", 0) - verbose_logger.debug( - "Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker - ) + verbose_logger.debug("Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker) return metrics_marker async def deliver( @@ -356,6 +326,4 @@ class FocusMavvrikDestination(FocusDestination): await self._upload_to_gcs(signed_url, content) await self._update_metrics_marker(date_epoch) - verbose_logger.debug( - "Mavvrik FOCUS destination: upload complete for date=%s", date_str - ) + verbose_logger.debug("Mavvrik FOCUS destination: upload complete for date=%s", date_str) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 4e3dd2b6d8b..ffd37aa195b 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -67,20 +67,14 @@ def _strip_unsupported_columns(csv_bytes: bytes) -> bytes: return csv_bytes header_cols = lines[0].decode("utf-8").split(",") - keep_indices = [ - i - for i, col in enumerate(header_cols) - if col.strip('"') in VANTAGE_SUPPORTED_COLUMNS - ] + keep_indices = [i for i, col in enumerate(header_cols) if col.strip('"') in VANTAGE_SUPPORTED_COLUMNS] # If all columns are supported, return as-is if len(keep_indices) == len(header_cols): return csv_bytes dropped = [col for i, col in enumerate(header_cols) if i not in keep_indices] - verbose_logger.debug( - "Vantage destination: dropping unsupported columns: %s", dropped - ) + verbose_logger.debug("Vantage destination: dropping unsupported columns: %s", dropped) output = io.StringIO() writer = csv.writer(output) @@ -143,10 +137,7 @@ class FocusVantageDestination(FocusDestination): # Check both size and row-count limits before single-shot upload lines = content.split(b"\n") data_line_count = sum(1 for line in lines[1:] if line.strip()) - within_limits = ( - len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD - and data_line_count <= VANTAGE_MAX_ROWS_PER_UPLOAD - ) + within_limits = len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD and data_line_count <= VANTAGE_MAX_ROWS_PER_UPLOAD if within_limits: await self._upload_csv(client, content, filename) return @@ -154,9 +145,7 @@ class FocusVantageDestination(FocusDestination): # Otherwise split into batches respecting both limits await self._upload_batched(client, content, filename) - async def _upload_csv( - self, client: AsyncHTTPHandler, csv_bytes: bytes, filename: str - ) -> None: + async def _upload_csv(self, client: AsyncHTTPHandler, csv_bytes: bytes, filename: str) -> None: url = f"{self.base_url}/v2/integrations/{self.integration_token}/costs.csv" headers = { "Authorization": f"Bearer {self.api_key}", @@ -174,9 +163,7 @@ class FocusVantageDestination(FocusDestination): filename, ) - async def _upload_batched( - self, client: AsyncHTTPHandler, csv_bytes: bytes, filename: str - ) -> None: + async def _upload_batched(self, client: AsyncHTTPHandler, csv_bytes: bytes, filename: str) -> None: """Split the CSV into batches and upload each. Continues uploading remaining batches even if one fails, then raises @@ -195,16 +182,12 @@ class FocusVantageDestination(FocusDestination): try: # If a single batch still exceeds 2 MB, split further by size if len(batch_csv) > VANTAGE_MAX_BYTES_PER_UPLOAD: - await self._upload_size_limited( - client, header, batch_lines, filename, batch_num - ) + await self._upload_size_limited(client, header, batch_lines, filename, batch_num) else: batch_filename = f"{filename}.part{batch_num}" await self._upload_csv(client, batch_csv, batch_filename) except Exception as e: - verbose_logger.error( - "Vantage destination: batch %d failed: %s", batch_num, e - ) + verbose_logger.error("Vantage destination: batch %d failed: %s", batch_num, e) if first_error is None: first_error = e batch_num += 1 @@ -244,10 +227,7 @@ class FocusVantageDestination(FocusDestination): ) continue - if ( - current_size + line_size > VANTAGE_MAX_BYTES_PER_UPLOAD - and current_chunk - ): + if current_size + line_size > VANTAGE_MAX_BYTES_PER_UPLOAD and current_chunk: batch_csv = header + b"\n" + b"\n".join(current_chunk) + b"\n" batch_filename = f"{filename}.part{batch_offset}_{sub_batch}" try: diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 37da18a0eb7..67ae6bcc3d0 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -42,9 +42,7 @@ class FocusExportEngine: return FocusCsvSerializer() if self.export_format == "parquet": return FocusParquetSerializer() - raise NotImplementedError( - f"Export format '{self.export_format}' not supported. Use 'parquet' or 'csv'." - ) + raise NotImplementedError(f"Export format '{self.export_format}' not supported. Use 'parquet' or 'csv'.") async def dry_run_export_usage_data(self, limit: Optional[int]) -> Dict[str, Any]: data = await self._database.get_usage_data(limit=limit) @@ -111,16 +109,12 @@ class FocusExportEngine: normalized = self._transformer.transform(data) if normalized.is_empty(): - verbose_logger.debug( - "Focus export: normalized data empty for window %s", window - ) + verbose_logger.debug("Focus export: normalized data empty for window %s", window) return await self._serialize_and_upload(normalized, window) - async def _serialize_and_upload( - self, frame: pl.DataFrame, window: FocusTimeWindow - ) -> None: + async def _serialize_and_upload(self, frame: pl.DataFrame, window: FocusTimeWindow) -> None: payload = self._serializer.serialize(frame) if not payload: verbose_logger.debug("Focus export: serializer returned empty payload") diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 083b0e1463a..ac6f1f7af1f 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -39,20 +39,12 @@ class FocusLogger(CustomLogger): ) -> None: super().__init__(**kwargs) self.provider = (provider or os.getenv("FOCUS_PROVIDER") or "s3").lower() - self.export_format = ( - export_format or os.getenv("FOCUS_FORMAT") or "parquet" - ).lower() + self.export_format = (export_format or os.getenv("FOCUS_FORMAT") or "parquet").lower() self.frequency = (frequency or os.getenv("FOCUS_FREQUENCY") or "hourly").lower() self.cron_offset_minute = ( - cron_offset_minute - if cron_offset_minute is not None - else int(os.getenv("FOCUS_CRON_OFFSET", "5")) - ) - raw_interval = ( - interval_seconds - if interval_seconds is not None - else os.getenv("FOCUS_INTERVAL_SECONDS") + cron_offset_minute if cron_offset_minute is not None else int(os.getenv("FOCUS_CRON_OFFSET", "5")) ) + raw_interval = interval_seconds if interval_seconds is not None else os.getenv("FOCUS_INTERVAL_SECONDS") self.interval_seconds: Optional[int] = None if raw_interval is not None: try: @@ -63,11 +55,7 @@ class FocusLogger(CustomLogger): raw_interval, ) env_prefix = os.getenv("FOCUS_PREFIX") - self.prefix: str = ( - prefix - if prefix is not None - else (env_prefix if env_prefix else "focus_exports") - ) + self.prefix: str = prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports") self._destination_config = destination_config self._engine: Optional["FocusExportEngine"] = None @@ -100,9 +88,7 @@ class FocusLogger(CustomLogger): automatic scheduler runs. """ if bool(start_time_utc) ^ bool(end_time_utc): - raise ValueError( - "start_time_utc and end_time_utc must be provided together" - ) + raise ValueError("start_time_utc and end_time_utc must be provided together") if start_time_utc and end_time_utc: window = FocusTimeWindow( @@ -115,9 +101,7 @@ class FocusLogger(CustomLogger): # No time bounds → export all available data await self._export_all(limit=limit) - async def dry_run_export_usage_data( - self, limit: Optional[int] = DEFAULT_DRY_RUN_LIMIT - ) -> dict[str, Any]: + async def dry_run_export_usage_data(self, limit: Optional[int] = DEFAULT_DRY_RUN_LIMIT) -> dict[str, Any]: """Return transformed data without uploading.""" engine = self._ensure_engine() return await engine.dry_run_export_usage_data(limit=limit) @@ -133,18 +117,14 @@ class FocusLogger(CustomLogger): pod_lock_manager = getattr(writer, "pod_lock_manager", None) if pod_lock_manager and pod_lock_manager.redis_cache: - acquired = await pod_lock_manager.acquire_lock( - cronjob_id=FOCUS_USAGE_DATA_JOB_NAME - ) + acquired = await pod_lock_manager.acquire_lock(cronjob_id=FOCUS_USAGE_DATA_JOB_NAME) if not acquired: verbose_logger.debug("Focus export: unable to acquire pod lock") return try: await self._run_scheduled_export() finally: - await pod_lock_manager.release_lock( - cronjob_id=FOCUS_USAGE_DATA_JOB_NAME - ) + await pod_lock_manager.release_lock(cronjob_id=FOCUS_USAGE_DATA_JOB_NAME) else: await self._run_scheduled_export() @@ -158,15 +138,11 @@ class FocusLogger(CustomLogger): # which have their own dedicated scheduling method. focus_loggers: List[CustomLogger] = [ cb - for cb in litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=FocusLogger - ) + for cb in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=FocusLogger) if type(cb) is FocusLogger ] if not focus_loggers: - verbose_logger.debug( - "No Focus export logger registered; skipping scheduler" - ) + verbose_logger.debug("No Focus export logger registered; skipping scheduler") return focus_logger = cast(FocusLogger, focus_loggers[0]) diff --git a/litellm/integrations/focus/serializers/csv.py b/litellm/integrations/focus/serializers/csv.py index 8e33c557be2..c0790358179 100644 --- a/litellm/integrations/focus/serializers/csv.py +++ b/litellm/integrations/focus/serializers/csv.py @@ -19,15 +19,9 @@ class FocusCsvSerializer(FocusSerializer): # Cast Decimal columns to Float64 so CSV output uses standard # floating-point notation (e.g. "1.5") instead of fixed-point # strings (e.g. "1.500000") that some parsers may reject. - decimal_cols = [ - col - for col, dtype in zip(frame.columns, frame.dtypes) - if isinstance(dtype, pl.Decimal) - ] + decimal_cols = [col for col, dtype in zip(frame.columns, frame.dtypes) if isinstance(dtype, pl.Decimal)] if decimal_cols: - frame = frame.with_columns( - [pl.col(c).cast(pl.Float64) for c in decimal_cols] - ) + frame = frame.with_columns([pl.col(c).cast(pl.Float64) for c in decimal_cols]) buffer = io.BytesIO() frame.write_csv(buffer) return buffer.getvalue() diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index a17df29b912..0fbefee75de 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -36,11 +36,7 @@ def _build_tags_expr(available_keys: list[str]) -> pl.Expr: tags = {k: str(v) for k, v in row.items() if v is not None} return json.dumps(tags) if tags else "{}" - return ( - pl.struct(available_keys) - .map_elements(_struct_to_json, return_dtype=pl.String) - .alias("Tags") - ) + return pl.struct(available_keys).map_elements(_struct_to_json, return_dtype=pl.String).alias("Tags") class FocusTransformer: @@ -97,9 +93,7 @@ class FocusTransformer: pl.lit("Usage-Based").alias("ChargeFrequency"), fmt(pl.col("ChargePeriodEnd")).alias("ChargePeriodEnd"), fmt(pl.col("ChargePeriodStart")).alias("ChargePeriodStart"), - dec( - pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0) - ).alias("ConsumedQuantity"), + dec(pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0)).alias("ConsumedQuantity"), pl.lit("Requests").alias("ConsumedUnit"), dec(pl.col("spend").fill_null(0.0)).alias("ContractedCost"), none_str.alias("ContractedUnitPrice"), @@ -111,9 +105,7 @@ class FocusTransformer: none_str.alias("AvailabilityZone"), pl.lit("USD").alias("PricingCurrency"), none_str.alias("PricingCategory"), - dec( - pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0) - ).alias("PricingQuantity"), + dec(pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0)).alias("PricingQuantity"), none_dec.alias("PricingCurrencyContractedUnitPrice"), dec(pl.col("spend").fill_null(0.0)).alias("PricingCurrencyEffectiveCost"), none_dec.alias("PricingCurrencyListUnitPrice"), diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index f9ff7e8c7a1..0ec6d496689 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -52,9 +52,7 @@ class LLMResponse(BaseModel): default=None, description="Optional. When available, logprobs are used to compute Uncertainty.", ) - created_at: str = Field( - ..., description='timestamp constructed in "%Y-%m-%dT%H:%M:%S" format' - ) + created_at: str = Field(..., description='timestamp constructed in "%Y-%m-%dT%H:%M:%S" format') tags: Optional[List[str]] = None user_metadata: Optional[Dict[str, Any]] = None @@ -73,9 +71,7 @@ class GalileoObserve(CustomLogger): self.base_url = GALILEO_CLOUD_API_BASE_URL self.use_v2_api = bool(self.api_key) self.headers: Optional[Dict[str, str]] = None - self.async_httpx_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) @staticmethod def _normalize_base_url(base_url: Optional[str]) -> Optional[str]: @@ -108,8 +104,7 @@ class GalileoObserve(CustomLogger): return IntegrationHealthCheckStatus( status="unhealthy", error_message=( - "GALILEO_API_KEY or GALILEO_USERNAME and GALILEO_PASSWORD " - "environment variables must be set" + "GALILEO_API_KEY or GALILEO_USERNAME and GALILEO_PASSWORD environment variables must be set" ), ) @@ -181,9 +176,7 @@ class GalileoObserve(CustomLogger): return False @staticmethod - def _galileo_input_messages( - messages: Optional[Any], input_text: str - ) -> List[Dict[str, str]]: + def _galileo_input_messages(messages: Optional[Any], input_text: str) -> List[Dict[str, str]]: if isinstance(messages, dict): messages = messages.get("messages") if not messages: @@ -201,9 +194,7 @@ class GalileoObserve(CustomLogger): galileo_messages.append( { "role": str(role), - "content": convert_content_list_to_str( - message=cast(AllMessageValues, message) - ), + "content": convert_content_list_to_str(message=cast(AllMessageValues, message)), } ) @@ -267,9 +258,7 @@ class GalileoObserve(CustomLogger): "parent_id": trace_id, "name": record.get("node_type", "litellm"), "created_at": created_at, - "input": GalileoObserve._galileo_input_messages( - record.get("messages"), record.get("input_text", "") - ), + "input": GalileoObserve._galileo_input_messages(record.get("messages"), record.get("input_text", "")), "output": { "role": "assistant", "content": record.get("output_text", ""), @@ -303,11 +292,7 @@ class GalileoObserve(CustomLogger): "duration_ns": int(record.get("latency_ms", 0)) * 1_000_000, **GalileoObserve._token_metrics_from_record(record), }, - "spans": [ - GalileoObserve._record_to_v2_span( - record, trace_id=trace_id, span_id=span_id - ) - ], + "spans": [GalileoObserve._record_to_v2_span(record, trace_id=trace_id, span_id=span_id)], } def _build_traces_payload(self, records: List[dict]) -> Dict[str, Any]: @@ -351,9 +336,7 @@ class GalileoObserve(CustomLogger): redacted: Dict[str, str] = {} for key, value in headers.items(): if key.lower() in {"authorization", "galileo-api-key"} and value: - redacted[key] = ( - f"{value[:8]}...{value[-4:]}" if len(value) > 12 else "***" - ) + redacted[key] = f"{value[:8]}...{value[-4:]}" if len(value) > 12 else "***" else: redacted[key] = value return redacted @@ -391,13 +374,9 @@ class GalileoObserve(CustomLogger): continue for field in ("id", "trace_id", "parent_id"): if field not in span: - missing_fields.append( - f"traces[{trace_index}].spans[{span_index}].{field}" - ) + missing_fields.append(f"traces[{trace_index}].spans[{span_index}].{field}") if trace_id and span.get("trace_id") != trace_id: - missing_fields.append( - f"traces[{trace_index}].spans[{span_index}].trace_id mismatch" - ) + missing_fields.append(f"traces[{trace_index}].spans[{span_index}].trace_id mismatch") if missing_fields: verbose_logger.debug( @@ -516,16 +495,11 @@ class GalileoObserve(CustomLogger): call_type = kwargs.get("call_type") prompt = self._build_prompt(kwargs) - if ( - level == "ERROR" - and status_message is not None - and isinstance(status_message, str) - ): + if level == "ERROR" and status_message is not None and isinstance(status_message, str): return self._prompt_to_input_text(prompt), status_message, prompt if response_obj is not None and ( - call_type in ("embedding", "aembedding") - or isinstance(response_obj, litellm.EmbeddingResponse) + call_type in ("embedding", "aembedding") or isinstance(response_obj, litellm.EmbeddingResponse) ): # Match Langfuse OTEL: log embeddings without serializing vectors. return self._prompt_to_input_text(prompt), "embedding-output", prompt @@ -538,14 +512,10 @@ class GalileoObserve(CustomLogger): kwargs.get("messages") or [], ) - if response_obj is not None and isinstance( - response_obj, HttpxBinaryResponseContent - ): + if response_obj is not None and isinstance(response_obj, HttpxBinaryResponseContent): return self._prompt_to_input_text(prompt), "speech-output", prompt - if response_obj is not None and isinstance( - response_obj, litellm.TextCompletionResponse - ): + if response_obj is not None and isinstance(response_obj, litellm.TextCompletionResponse): output = self._get_text_completion_content_for_galileo(response_obj) return ( self._prompt_to_input_text(prompt), @@ -561,9 +531,7 @@ class GalileoObserve(CustomLogger): prompt, ) - if response_obj is not None and isinstance( - response_obj, litellm.TranscriptionResponse - ): + if response_obj is not None and isinstance(response_obj, litellm.TranscriptionResponse): output = response_obj.get("text", None) return ( self._prompt_to_input_text(prompt), @@ -571,9 +539,7 @@ class GalileoObserve(CustomLogger): prompt, ) - if response_obj is not None and isinstance( - response_obj, litellm.RerankResponse - ): + if response_obj is not None and isinstance(response_obj, litellm.RerankResponse): output = response_obj.results rerank_prompt = self._langfuse_style_rerank_prompt(kwargs) return ( @@ -590,11 +556,7 @@ class GalileoObserve(CustomLogger): kwargs.get("messages") or [], ) - if ( - call_type == "_arealtime" - and response_obj is not None - and isinstance(response_obj, list) - ): + if call_type == "_arealtime" and response_obj is not None and isinstance(response_obj, list): input_val = kwargs.get("input") return ( self._serialize_galileo_output(input_val), @@ -602,11 +564,7 @@ class GalileoObserve(CustomLogger): input_val, ) - if ( - call_type == "pass_through_endpoint" - and response_obj is not None - and isinstance(response_obj, dict) - ): + if call_type == "pass_through_endpoint" and response_obj is not None and isinstance(response_obj, dict): output = response_obj.get("response", "") return ( self._prompt_to_input_text(prompt), @@ -624,12 +582,8 @@ class GalileoObserve(CustomLogger): return self._prompt_to_input_text(prompt), "", kwargs.get("messages") or [] - def get_output_str_from_response( - self, response_obj: Any, kwargs: Dict[str, Any] - ) -> str: - _, output_text, _ = self._get_galileo_input_output_content( - kwargs=kwargs, response_obj=response_obj - ) + def get_output_str_from_response(self, response_obj: Any, kwargs: Dict[str, Any]) -> str: + _, output_text, _ = self._get_galileo_input_output_content(kwargs=kwargs, response_obj=response_obj) return output_text @staticmethod @@ -646,10 +600,7 @@ class GalileoObserve(CustomLogger): if str(msg.get("role", "")).lower() in ("user", "human"): content = msg.get("content") or "" if isinstance(content, list): - content = " ".join( - b.get("text", "") if isinstance(b, dict) else str(b) - for b in content - ) + content = " ".join(b.get("text", "") if isinstance(b, dict) else str(b) for b in content) if content: return str(content) # Fallback: first non-empty content of any role @@ -657,17 +608,12 @@ class GalileoObserve(CustomLogger): if isinstance(msg, dict): content = msg.get("content") or "" if isinstance(content, list): - content = " ".join( - b.get("text", "") if isinstance(b, dict) else str(b) - for b in content - ) + content = " ".join(b.get("text", "") if isinstance(b, dict) else str(b) for b in content) if content: return str(content) return "" - async def async_log_success_event( - self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any - ): + async def async_log_success_event(self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any): verbose_logger.debug("On Async Success") try: await self._async_log_success_event_impl( @@ -677,13 +623,9 @@ class GalileoObserve(CustomLogger): end_time=end_time, ) except Exception: - verbose_logger.exception( - "Galileo Logger: unexpected error in async_log_success_event" - ) + verbose_logger.exception("Galileo Logger: unexpected error in async_log_success_event") - async def _async_log_success_event_impl( - self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any - ): + async def _async_log_success_event_impl(self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any): if not self._is_configured(): verbose_logger.debug( "Galileo Logger: skipping — GALILEO_PROJECT_ID=%s GALILEO_API_KEY=%s GALILEO_BASE_URL=%s", @@ -695,14 +637,10 @@ class GalileoObserve(CustomLogger): slo: Optional[Dict[str, Any]] = kwargs.get("standard_logging_object") if slo is None: - verbose_logger.debug( - "Galileo Logger: no standard_logging_object in kwargs, skipping" - ) + verbose_logger.debug("Galileo Logger: no standard_logging_object in kwargs, skipping") return - _call_type: str = str( - slo.get("call_type") or kwargs.get("call_type") or "litellm" - ) + _call_type: str = str(slo.get("call_type") or kwargs.get("call_type") or "litellm") input_text, output_text, messages = self._get_galileo_input_output_content( kwargs=kwargs, response_obj=response_obj @@ -715,9 +653,7 @@ class GalileoObserve(CustomLogger): "Galileo Logger: standard_logging_object missing startTime/endTime, " "falling back to start_time/end_time params" ) - if not isinstance(start_time, datetime) or not isinstance( - end_time, datetime - ): + if not isinstance(start_time, datetime) or not isinstance(end_time, datetime): return start_ts = start_time end_ts = end_time @@ -757,17 +693,13 @@ class GalileoObserve(CustomLogger): if isinstance(messages, list) and messages: request_dict["messages"] = messages self.in_memory_records.append(request_dict) - verbose_logger.debug( - "Galileo Logger: queued record, in_memory=%d", len(self.in_memory_records) - ) + verbose_logger.debug("Galileo Logger: queued record, in_memory=%d", len(self.in_memory_records)) # Bound the buffer so persistent flush failures cannot grow it # without limit. Drop the oldest records once we exceed the cap. if len(self.in_memory_records) > GALILEO_MAX_IN_MEMORY_RECORDS: dropped = len(self.in_memory_records) - GALILEO_MAX_IN_MEMORY_RECORDS - self.in_memory_records = self.in_memory_records[ - -GALILEO_MAX_IN_MEMORY_RECORDS: - ] + self.in_memory_records = self.in_memory_records[-GALILEO_MAX_IN_MEMORY_RECORDS:] verbose_logger.warning( "Galileo Logger: in-memory buffer exceeded %s records; " "dropped %s oldest record(s). Check Galileo connectivity/credentials.", @@ -789,15 +721,11 @@ class GalileoObserve(CustomLogger): ingest_request = self._get_ingest_request() if ingest_request is None: - verbose_logger.debug( - "Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID — skipping flush" - ) + verbose_logger.debug("Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID — skipping flush") return if not await self._ensure_headers(): - verbose_logger.debug( - "Galileo Logger: could not set request headers — skipping flush" - ) + verbose_logger.debug("Galileo Logger: could not set request headers — skipping flush") return url, payload = ingest_request @@ -817,20 +745,14 @@ class GalileoObserve(CustomLogger): ) except httpx.HTTPStatusError as e: self._log_http_status_error(error=e, url=url) - verbose_logger.debug( - "Galileo Logger: failed to flush in memory records: %s", e - ) + verbose_logger.debug("Galileo Logger: failed to flush in memory records: %s", e) return except Exception as e: - verbose_logger.debug( - "Galileo Logger: failed to flush in memory records: %s", e - ) + verbose_logger.debug("Galileo Logger: failed to flush in memory records: %s", e) return if response.is_success: - verbose_logger.debug( - "Galileo Logger: successfully flushed in memory records" - ) + verbose_logger.debug("Galileo Logger: successfully flushed in memory records") verbose_logger.debug( "Galileo Logger flush response: status=%s body=%s", response.status_code, diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 90057984235..c2e0ad64586 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -32,14 +32,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): super().__init__(bucket_name=bucket_name) self.batch_size = int(os.getenv("GCS_BATCH_SIZE", GCS_DEFAULT_BATCH_SIZE)) - self.flush_interval = int( - os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS) - ) + self.flush_interval = int(os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS)) self.use_batched_logging = ( - os.getenv( - "GCS_USE_BATCHED_LOGGING", str(GCS_DEFAULT_USE_BATCHED_LOGGING).lower() - ).lower() - == "true" + os.getenv("GCS_USE_BATCHED_LOGGING", str(GCS_DEFAULT_USE_BATCHED_LOGGING).lower()).lower() == "true" ) self.flush_lock = asyncio.Lock() super().__init__( @@ -72,19 +67,13 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): kwargs, response_obj, ) - logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) if self.log_queue.full(): await self.flush_queue() - await self.log_queue.put( - GCSLogQueueItem( - payload=logging_payload, kwargs=kwargs, response_obj=response_obj - ) - ) + await self.log_queue.put(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj)) except Exception as e: verbose_logger.exception(f"GCS Bucket logging error: {str(e)}") @@ -97,19 +86,13 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): response_obj, ) - logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") # When queue is at maxsize, flush immediately to make room (no blocking, no data dropped) if self.log_queue.full(): await self.flush_queue() - await self.log_queue.put( - GCSLogQueueItem( - payload=logging_payload, kwargs=kwargs, response_obj=response_obj - ) - ) + await self.log_queue.put(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj)) except Exception as e: verbose_logger.exception(f"GCS Bucket logging error: {str(e)}") @@ -147,15 +130,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): This key may contain sensitive information (bucket names, paths) - use _sanitize_config_key() for logging purposes. """ - standard_callback_dynamic_params = ( - kwargs.get("standard_callback_dynamic_params", None) or {} - ) + standard_callback_dynamic_params = kwargs.get("standard_callback_dynamic_params", None) or {} - bucket_name = ( - standard_callback_dynamic_params.get("gcs_bucket_name", None) - or self.BUCKET_NAME - or "default" - ) + bucket_name = standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME or "default" path_service_account = ( standard_callback_dynamic_params.get("gcs_path_service_account", None) or self.path_service_account_json @@ -174,9 +151,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): hash_obj = hashlib.sha256(config_key.encode("utf-8")) return f"config-{hash_obj.hexdigest()[:8]}" - def _group_items_by_config( - self, items: List[GCSLogQueueItem] - ) -> Dict[str, List[GCSLogQueueItem]]: + def _group_items_by_config(self, items: List[GCSLogQueueItem]) -> Dict[str, List[GCSLogQueueItem]]: """ Group items by their GCS config (bucket + credentials). This ensures items with different configs are processed separately. @@ -203,9 +178,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): lines.append(json_line) return "\n".join(lines) - async def _send_grouped_batch( - self, items: List[GCSLogQueueItem], config_key: str - ) -> Tuple[int, int]: + async def _send_grouped_batch(self, items: List[GCSLogQueueItem], config_key: str) -> Tuple[int, int]: """ Send a batch of items that share the same GCS config. @@ -218,9 +191,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): first_kwargs = items[0]["kwargs"] try: - gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( - first_kwargs - ) + gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(first_kwargs) headers = await self.construct_request_headers( vertex_instance=gcs_logging_config["vertex_instance"], @@ -228,9 +199,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) bucket_name = gcs_logging_config["bucket_name"] - current_date = self._get_object_date_from_datetime( - datetime.now(timezone.utc) - ) + current_date = self._get_object_date_from_datetime(datetime.now(timezone.utc)) batch_id = f"{int(time.time() * 1000)}-{uuid.uuid4().hex[:8]}" object_name = self._generate_batch_object_name(current_date, batch_id) combined_payload = self._combine_payloads_to_ndjson(items) @@ -249,9 +218,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): except Exception as e: success_count = 0 error_count = len(items) - verbose_logger.exception( - f"GCS Bucket error logging batch payload to GCS bucket: {str(e)}" - ) + verbose_logger.exception(f"GCS Bucket error logging batch payload to GCS bucket: {str(e)}") return (success_count, error_count) async def _send_individual_logs(self, items: List[GCSLogQueueItem]) -> None: @@ -267,9 +234,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): Send a single log item to GCS as an individual object. """ try: - gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( - item["kwargs"] - ) + gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(item["kwargs"]) headers = await self.construct_request_headers( vertex_instance=gcs_logging_config["vertex_instance"], @@ -290,9 +255,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): logging_payload=item["payload"], ) except Exception as e: - verbose_logger.exception( - f"GCS Bucket error logging individual payload to GCS bucket: {str(e)}" - ) + verbose_logger.exception(f"GCS Bucket error logging individual payload to GCS bucket: {str(e)}") async def async_send_batch(self): """ @@ -316,9 +279,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): else: await self._send_individual_logs(items_to_process) - def _get_object_name( - self, kwargs: Dict, logging_payload: StandardLoggingPayload, response_obj: Any - ) -> str: + def _get_object_name(self, kwargs: Dict, logging_payload: StandardLoggingPayload, response_obj: Any) -> str: """ Get the object name to use for the current payload """ @@ -337,9 +298,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): _litellm_params = kwargs.get("litellm_params", None) or {} _metadata = _litellm_params.get("metadata", None) or {} if "gcs_log_id" in _metadata: - safe_log_id = sanitize_cloud_object_component( - _metadata.get("gcs_log_id"), fallback="" - ) + safe_log_id = sanitize_cloud_object_component(_metadata.get("gcs_log_id"), fallback="") if safe_log_id: object_name = f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" @@ -356,9 +315,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): Tries current day, next day, and previous day until it finds the payload """ if start_time_utc is None: - raise ValueError( - "start_time_utc is required for getting a payload from GCS Bucket" - ) + raise ValueError("start_time_utc is required for getting a payload from GCS Bucket") dates_to_try = [ start_time_utc, @@ -379,9 +336,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): loaded_response = json.loads(response) return loaded_response except Exception as e: - verbose_logger.debug( - f"Failed to fetch payload for date {date_str}: {str(e)}" - ) + verbose_logger.debug(f"Failed to fetch payload for date {date_str}: {str(e)}") continue return None @@ -415,9 +370,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): """ while True: await asyncio.sleep(self.flush_interval) - verbose_logger.debug( - f"GCS Bucket periodic flush after {self.flush_interval} seconds" - ) + verbose_logger.debug(f"GCS Bucket periodic flush after {self.flush_interval} seconds") await self.flush_queue() async def async_health_check(self) -> IntegrationHealthCheckStatus: diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_base.py b/litellm/integrations/gcs_bucket/gcs_bucket_base.py index 1c5e30777a2..0eabf16cff9 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_base.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_base.py @@ -37,9 +37,7 @@ class GCSBucketBase(CustomBatchLogger): mock_vertex_auth_methods() create_mock_gcs_client() - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) _path_service_account = os.getenv("GCS_PATH_SERVICE_ACCOUNT") _bucket_name = bucket_name or os.getenv("GCS_BUCKET_NAME") self.path_service_account_json: Optional[str] = _path_service_account @@ -74,9 +72,7 @@ class GCSBucketBase(CustomBatchLogger): custom_llm_provider="vertex_ai", api_base=None, ) - verbose_logger.debug( - "constructed auth_header [set=%s]", auth_header is not None - ) + verbose_logger.debug("constructed auth_header [set=%s]", auth_header is not None) headers = { "Authorization": f"Bearer {auth_header}", # auth_header "Content-Type": "application/json", @@ -112,9 +108,7 @@ class GCSBucketBase(CustomBatchLogger): custom_llm_provider="vertex_ai", api_base=None, ) - verbose_logger.debug( - "constructed auth_header [set=%s]", auth_header is not None - ) + verbose_logger.debug("constructed auth_header [set=%s]", auth_header is not None) headers = { "Authorization": f"Bearer {auth_header}", # auth_header "Content-Type": "application/json", @@ -143,9 +137,7 @@ class GCSBucketBase(CustomBatchLogger): return bucket_name, object_name return bucket_name, object_name - async def get_gcs_logging_config( - self, kwargs: Optional[Dict[str, Any]] = {} - ) -> GCSLoggingConfig: + async def get_gcs_logging_config(self, kwargs: Optional[Dict[str, Any]] = {}) -> GCSLoggingConfig: """ This function is used to get the GCS logging config for the GCS Bucket Logger. It checks if the dynamic parameters are provided in the kwargs and uses them to get the GCS logging config. @@ -154,25 +146,21 @@ class GCSBucketBase(CustomBatchLogger): if kwargs is None: kwargs = {} - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params", None) + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get( + "standard_callback_dynamic_params", None ) bucket_name: str path_service_account: Optional[str] if standard_callback_dynamic_params is not None: verbose_logger.debug("Using dynamic GCS logging") - verbose_logger.debug( - "standard_callback_dynamic_params: %s", standard_callback_dynamic_params - ) + verbose_logger.debug("standard_callback_dynamic_params: %s", standard_callback_dynamic_params) _bucket_name: Optional[str] = ( - standard_callback_dynamic_params.get("gcs_bucket_name", None) - or self.BUCKET_NAME + standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME ) _path_service_account: Optional[str] = ( - standard_callback_dynamic_params.get("gcs_path_service_account", None) - or self.path_service_account_json + standard_callback_dynamic_params.get("gcs_path_service_account", None) or self.path_service_account_json ) if _bucket_name is None: @@ -181,9 +169,7 @@ class GCSBucketBase(CustomBatchLogger): ) bucket_name = _bucket_name path_service_account = _path_service_account - vertex_instance = await self.get_or_create_vertex_instance( - credentials=path_service_account - ) + vertex_instance = await self.get_or_create_vertex_instance(credentials=path_service_account) else: # If no dynamic parameters, use the default instance if self.BUCKET_NAME is None: @@ -192,9 +178,7 @@ class GCSBucketBase(CustomBatchLogger): ) bucket_name = self.BUCKET_NAME path_service_account = self.path_service_account_json - vertex_instance = await self.get_or_create_vertex_instance( - credentials=path_service_account - ) + vertex_instance = await self.get_or_create_vertex_instance(credentials=path_service_account) return GCSLoggingConfig( bucket_name=bucket_name, @@ -202,9 +186,7 @@ class GCSBucketBase(CustomBatchLogger): path_service_account=path_service_account, ) - async def get_or_create_vertex_instance( - self, credentials: Optional[str] - ) -> VertexBase: + async def get_or_create_vertex_instance(self, credentials: Optional[str]) -> VertexBase: """ This function is used to get the Vertex instance for the GCS Bucket Logger. It checks if the Vertex instance is already created and cached, if not it creates a new instance and caches it. @@ -240,9 +222,7 @@ class GCSBucketBase(CustomBatchLogger): https://cloud.google.com/storage/docs/downloading-objects#download-object-json """ try: - gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( - kwargs=kwargs - ) + gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(kwargs=kwargs) headers = await self.construct_request_headers( vertex_instance=gcs_logging_config["vertex_instance"], service_account_json=gcs_logging_config["path_service_account"], @@ -260,14 +240,10 @@ class GCSBucketBase(CustomBatchLogger): response = await self.async_httpx_client.get(url=url, headers=headers) if response.status_code != 200: - verbose_logger.error( - "GCS object download error: %s", str(response.text) - ) + verbose_logger.error("GCS object download error: %s", str(response.text)) return None - verbose_logger.debug( - "GCS object download response status code: %s", response.status_code - ) + verbose_logger.debug("GCS object download response status code: %s", response.status_code) # Return the content of the downloaded object return response.content @@ -281,9 +257,7 @@ class GCSBucketBase(CustomBatchLogger): Delete an object from GCS. """ try: - gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( - kwargs=kwargs - ) + gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(kwargs=kwargs) headers = await self.construct_request_headers( vertex_instance=gcs_logging_config["vertex_instance"], service_account_json=gcs_logging_config["path_service_account"], diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py index 9161455a246..fae7ddaf536 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py @@ -38,14 +38,10 @@ _mocks_initialized = False # Default mock latency in seconds (simulates network round-trip) # Typical GCS API calls take 100-300ms for uploads, 50-150ms for GET/DELETE -_MOCK_LATENCY_SECONDS = ( - float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 -) +_MOCK_LATENCY_SECONDS = float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 -async def _mock_async_handler_get( - self, url, params=None, headers=None, follow_redirects=None -): +async def _mock_async_handler_get(self, url, params=None, headers=None, follow_redirects=None): """Monkey-patched AsyncHTTPHandler.get that intercepts GCS calls.""" # Only mock GCS API calls if isinstance(url, str) and "storage.googleapis.com" in url: @@ -154,10 +150,7 @@ def create_mock_gcs_client(): This function is idempotent - it only initializes mocks once, even if called multiple times. """ - global \ - _original_async_handler_get, \ - _original_async_handler_delete, \ - _mocks_initialized + global _original_async_handler_get, _original_async_handler_delete, _mocks_initialized # Use factory for POST handler _create_mock_gcs_post() @@ -181,9 +174,7 @@ def create_mock_gcs_client(): AsyncHTTPHandler.delete = _mock_async_handler_delete # type: ignore verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.delete") - verbose_logger.debug( - f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms" - ) + verbose_logger.debug(f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms") verbose_logger.debug("[GCS MOCK] GCS mock client initialization complete") _mocks_initialized = True @@ -205,29 +196,17 @@ def mock_vertex_auth_methods(): "_original_ensure_access_token_async", VertexBase._ensure_access_token_async, ) - setattr( - VertexBase, "_original_ensure_access_token", VertexBase._ensure_access_token - ) - setattr( - VertexBase, "_original_get_token_and_url", VertexBase._get_token_and_url - ) + setattr(VertexBase, "_original_ensure_access_token", VertexBase._ensure_access_token) + setattr(VertexBase, "_original_get_token_and_url", VertexBase._get_token_and_url) - async def _mock_ensure_access_token_async( - self, credentials, project_id, custom_llm_provider - ): + async def _mock_ensure_access_token_async(self, credentials, project_id, custom_llm_provider): """Mock async auth method - returns fake token.""" - verbose_logger.debug( - "[GCS MOCK] Vertex AI auth: _ensure_access_token_async called" - ) + verbose_logger.debug("[GCS MOCK] Vertex AI auth: _ensure_access_token_async called") return ("mock-gcs-token", "mock-project-id") - def _mock_ensure_access_token( - self, credentials, project_id, custom_llm_provider - ): + def _mock_ensure_access_token(self, credentials, project_id, custom_llm_provider): """Mock sync auth method - returns fake token.""" - verbose_logger.debug( - "[GCS MOCK] Vertex AI auth: _ensure_access_token called" - ) + verbose_logger.debug("[GCS MOCK] Vertex AI auth: _ensure_access_token called") return ("mock-gcs-token", "mock-project-id") def _mock_get_token_and_url( diff --git a/litellm/integrations/gcs_pubsub/pub_sub.py b/litellm/integrations/gcs_pubsub/pub_sub.py index db7f9bb4d0b..c1bccb0b390 100644 --- a/litellm/integrations/gcs_pubsub/pub_sub.py +++ b/litellm/integrations/gcs_pubsub/pub_sub.py @@ -48,15 +48,11 @@ class GcsPubSubLogger(CustomBatchLogger): _premium_user_check() - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.project_id = project_id or os.getenv("GCS_PUBSUB_PROJECT_ID") self.topic_id = topic_id or os.getenv("GCS_PUBSUB_TOPIC_ID") - self.path_service_account_json = credentials_path or os.getenv( - "GCS_PATH_SERVICE_ACCOUNT" - ) + self.path_service_account_json = credentials_path or os.getenv("GCS_PATH_SERVICE_ACCOUNT") if not self.project_id or not self.topic_id: raise ValueError("Both project_id and topic_id must be provided") @@ -116,9 +112,7 @@ class GcsPubSubLogger(CustomBatchLogger): _premium_user_check() try: - verbose_logger.debug( - "PubSub: Logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("PubSub: Logging - Enters logging function for model %s", kwargs) standard_logging_payload = kwargs.get("standard_logging_object", None) # Backwards compatibility with old logging payload @@ -138,9 +132,7 @@ class GcsPubSubLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"PubSub Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"PubSub Layer Error - {str(e)}\n{traceback.format_exc()}") pass async def async_send_batch(self): @@ -151,17 +143,13 @@ class GcsPubSubLogger(CustomBatchLogger): if not self.log_queue: return - verbose_logger.debug( - f"PubSub - about to flush {len(self.log_queue)} events" - ) + verbose_logger.debug(f"PubSub - about to flush {len(self.log_queue)} events") for message in self.log_queue: await self.publish_message(message) except Exception as e: - verbose_logger.exception( - f"PubSub Error sending batch - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"PubSub Error sending batch - {str(e)}\n{traceback.format_exc()}") finally: self.log_queue.clear() @@ -189,18 +177,14 @@ class GcsPubSubLogger(CustomBatchLogger): # Base64 encode the message import base64 - encoded_message = base64.b64encode(message_data.encode("utf-8")).decode( - "utf-8" - ) + encoded_message = base64.b64encode(message_data.encode("utf-8")).decode("utf-8") # Construct request body request_body = {"messages": [{"data": encoded_message}]} url = f"https://pubsub.googleapis.com/v1/projects/{self.project_id}/topics/{self.topic_id}:publish" - response = await self.async_httpx_client.post( - url=url, headers=headers, json=request_body - ) + response = await self.async_httpx_client.post(url=url, headers=headers, json=request_body) if response.status_code not in [200, 202]: verbose_logger.error("Pub/Sub publish error: %s", str(response.text)) diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 2982df8fda2..da6009c3a94 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -37,15 +37,11 @@ def load_compatible_callbacks() -> Dict: Dict: Dictionary of compatible callbacks configuration """ try: - json_path = os.path.join( - os.path.dirname(__file__), "generic_api_compatible_callbacks.json" - ) + json_path = os.path.join(os.path.dirname(__file__), "generic_api_compatible_callbacks.json") with open(json_path, "r") as f: return json.load(f) except Exception as e: - verbose_logger.warning( - f"Error loading generic_api_compatible_callbacks.json: {str(e)}" - ) + verbose_logger.warning(f"Error loading generic_api_compatible_callbacks.json: {str(e)}") return {} @@ -127,9 +123,7 @@ class GenericAPILogger(CustomBatchLogger): ######################################################### if callback_name: if is_callback_compatible(callback_name): - verbose_logger.debug( - f"Loading configuration for callback: {callback_name}" - ) + verbose_logger.debug(f"Loading configuration for callback: {callback_name}") callback_config = get_callback_config(callback_name) # Use config from JSON if not explicitly provided @@ -156,9 +150,7 @@ class GenericAPILogger(CustomBatchLogger): ######################################################### # Init httpx client ######################################################### - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) endpoint = endpoint or os.getenv("GENERIC_LOGGER_ENDPOINT") if endpoint is None: raise ValueError( @@ -180,9 +172,7 @@ class GenericAPILogger(CustomBatchLogger): "ndjson", "single", ]: - raise ValueError( - f"Invalid log_format: {log_format}. Must be one of: 'json_array', 'ndjson', 'single'" - ) + raise ValueError(f"Invalid log_format: {log_format}. Must be one of: 'json_array', 'ndjson', 'single'") self.log_format: LOG_FORMAT_TYPES = log_format or "json_array" verbose_logger.debug( @@ -223,9 +213,7 @@ class GenericAPILogger(CustomBatchLogger): key, value = item.split("=", 1) headers_dict[key.strip()] = value.strip() except Exception as e: - verbose_logger.warning( - f"Error parsing headers from environment variables: {str(e)}" - ) + verbose_logger.warning(f"Error parsing headers from environment variables: {str(e)}") # 2. Update with litellm generic headers if available if litellm.generic_logger_headers: @@ -273,8 +261,7 @@ class GenericAPILogger(CustomBatchLogger): raise verbose_logger.warning( - "Generic API Logger - retrying request to %s after error: %s " - "(attempt %s/%s)", + "Generic API Logger - retrying request to %s after error: %s (attempt %s/%s)", self.endpoint, str(e), attempt + 1, @@ -300,9 +287,7 @@ class GenericAPILogger(CustomBatchLogger): return try: - verbose_logger.debug( - "Generic API Logger - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("Generic API Logger - Enters logging function for model %s", kwargs) standard_logging_payload = kwargs.get("standard_logging_object", None) # Backwards compatibility with old logging payload @@ -322,9 +307,7 @@ class GenericAPILogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"Generic API Logger Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Generic API Logger Error - {str(e)}\n{traceback.format_exc()}") pass async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -338,9 +321,7 @@ class GenericAPILogger(CustomBatchLogger): return try: - verbose_logger.debug( - "Generic API Logger - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("Generic API Logger - Enters logging function for model %s", kwargs) standard_logging_payload = kwargs.get("standard_logging_object", None) if litellm.generic_api_use_v1 is True: @@ -358,9 +339,7 @@ class GenericAPILogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception( - f"Generic API Logger Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Generic API Logger Error - {str(e)}\n{traceback.format_exc()}") async def async_send_batch(self): """ @@ -392,9 +371,7 @@ class GenericAPILogger(CustomBatchLogger): # Log results for idx, result in enumerate(responses): if isinstance(result, Exception): - verbose_logger.exception( - f"Generic API Logger - Error sending log {idx}: {result}" - ) + verbose_logger.exception(f"Generic API Logger - Error sending log {idx}: {result}") else: # result is a Response object verbose_logger.debug( @@ -418,30 +395,22 @@ class GenericAPILogger(CustomBatchLogger): ) except Exception as e: - verbose_logger.exception( - f"Generic API Logger Error sending batch - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Generic API Logger Error sending batch - {str(e)}\n{traceback.format_exc()}") finally: self.log_queue.clear() - def _get_v1_logging_payload( - self, kwargs, response_obj, start_time, end_time - ) -> dict: + def _get_v1_logging_payload(self, kwargs, response_obj, start_time, end_time) -> dict: """ Maintained for backwards compatibility with old logging payload Returns a dict of the payload to send to the Generic API Endpoint """ - verbose_logger.debug( - f"GenericAPILogger Logging - Enters logging function for model {kwargs}" - ) + verbose_logger.debug(f"GenericAPILogger Logging - Enters logging function for model {kwargs}") # construct payload to send custom logger # follows the same params as langfuse.py litellm_params = kwargs.get("litellm_params", {}) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None + metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None messages = kwargs.get("messages") cost = kwargs.get("response_cost", 0.0) optional_params = kwargs.get("optional_params", {}) diff --git a/litellm/integrations/generic_prompt_management/__init__.py b/litellm/integrations/generic_prompt_management/__init__.py index 7466dc9c68d..44c61aa5f50 100644 --- a/litellm/integrations/generic_prompt_management/__init__.py +++ b/litellm/integrations/generic_prompt_management/__init__.py @@ -30,9 +30,7 @@ def set_global_generic_prompt_config(config: dict) -> None: litellm.global_generic_prompt_config = config # type: ignore -def prompt_initializer( - litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec" -) -> "CustomPromptManagement": +def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec") -> "CustomPromptManagement": """ Initialize a prompt from a generic prompt management API. """ diff --git a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py index 858bfd458b6..f9837efdde2 100644 --- a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py +++ b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py @@ -74,9 +74,7 @@ class GenericPromptManager(CustomPromptManagement): self.api_key = api_key self.timeout = timeout self.prompt_id = prompt_id - self.additional_provider_specific_query_params = ( - additional_provider_specific_query_params - ) + self.additional_provider_specific_query_params = additional_provider_specific_query_params self._prompt_cache: Dict[str, PromptManagementClient] = {} @property @@ -94,9 +92,7 @@ class GenericPromptManager(CustomPromptManagement): headers["Authorization"] = f"Bearer {self.api_key}" return headers - def _fetch_prompt_from_api( - self, prompt_id: Optional[str], prompt_spec: Optional[PromptSpec] - ) -> Dict[str, Any]: + def _fetch_prompt_from_api(self, prompt_id: Optional[str], prompt_spec: Optional[PromptSpec]) -> Dict[str, Any]: """ Fetch a prompt from the API. @@ -147,8 +143,7 @@ class GenericPromptManager(CustomPromptManagement): "prompt_id": prompt_id, **( prompt_spec.litellm_params.provider_specific_query_params - if prompt_spec - and prompt_spec.litellm_params.provider_specific_query_params + if prompt_spec and prompt_spec.litellm_params.provider_specific_query_params else {} ), } @@ -204,9 +199,7 @@ class GenericPromptManager(CustomPromptManagement): prompt_id=prompt_id, prompt_template=api_response.get("prompt_template", []), prompt_template_model=api_response.get("prompt_template_model"), - prompt_template_optional_params=api_response.get( - "prompt_template_optional_params" - ), + prompt_template_optional_params=api_response.get("prompt_template_optional_params"), completed_messages=None, ) @@ -223,8 +216,7 @@ class GenericPromptManager(CustomPromptManagement): in the _compile_prompt_helper method. """ if prompt_id is not None or ( - prompt_spec is not None - and prompt_spec.litellm_params.provider_specific_query_params is not None + prompt_spec is not None and prompt_spec.litellm_params.provider_specific_query_params is not None ): return True return False @@ -299,9 +291,7 @@ class GenericPromptManager(CustomPromptManagement): api_response = self._fetch_prompt_from_api(prompt_id, prompt_spec) # Parse the response - prompt_client = self._parse_api_response( - prompt_id, prompt_spec, api_response - ) + prompt_client = self._parse_api_response(prompt_id, prompt_spec, api_response) # Cache the result self._prompt_cache[cache_key] = prompt_client @@ -339,14 +329,10 @@ class GenericPromptManager(CustomPromptManagement): try: # Fetch from API - api_response = await self.async_fetch_prompt_from_api( - prompt_id=prompt_id, prompt_spec=prompt_spec - ) + api_response = await self.async_fetch_prompt_from_api(prompt_id=prompt_id, prompt_spec=prompt_spec) # Parse the response - prompt_client = self._parse_api_response( - prompt_id, prompt_spec, api_response - ) + prompt_client = self._parse_api_response(prompt_id, prompt_spec, api_response) # Cache the result self._prompt_cache[cache_key] = prompt_client @@ -358,9 +344,7 @@ class GenericPromptManager(CustomPromptManagement): return prompt_client except Exception as e: - raise ValueError( - f"Error compiling prompt '{prompt_id}': {e}, prompt_spec: {prompt_spec}" - ) + raise ValueError(f"Error compiling prompt '{prompt_id}': {e}, prompt_spec: {prompt_spec}") def _apply_variables( self, @@ -383,15 +367,11 @@ class GenericPromptManager(CustomPromptManagement): updated_messages: List[AllMessageValues] = [] for message in prompt_client["prompt_template"]: updated_message = dict(message) # type: ignore - if "content" in updated_message and isinstance( - updated_message["content"], str - ): + if "content" in updated_message and isinstance(updated_message["content"], str): content = updated_message["content"] for key, value in variables.items(): content = content.replace(f"{{{key}}}", str(value)) - content = content.replace( - f"{{{{{key}}}}}", str(value) - ) # Also support {{key}} + content = content.replace(f"{{{{{key}}}}}", str(value)) # Also support {{key}} updated_message["content"] = content updated_messages.append(updated_message) # type: ignore @@ -399,9 +379,7 @@ class GenericPromptManager(CustomPromptManagement): prompt_id=prompt_client["prompt_id"], prompt_template=updated_messages, prompt_template_model=prompt_client["prompt_template_model"], - prompt_template_optional_params=prompt_client[ - "prompt_template_optional_params" - ], + prompt_template_optional_params=prompt_client["prompt_template_optional_params"], completed_messages=None, ) @@ -439,8 +417,7 @@ class GenericPromptManager(CustomPromptManagement): prompt_label=prompt_label, prompt_version=prompt_version, ignore_prompt_manager_model=( - ignore_prompt_manager_model - or prompt_spec.litellm_params.ignore_prompt_manager_model + ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model if prompt_spec else False ), @@ -481,8 +458,7 @@ class GenericPromptManager(CustomPromptManagement): prompt_label=prompt_label, prompt_version=prompt_version, ignore_prompt_manager_model=( - ignore_prompt_manager_model - or prompt_spec.litellm_params.ignore_prompt_manager_model + ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model if prompt_spec else False ), diff --git a/litellm/integrations/gitlab/__init__.py b/litellm/integrations/gitlab/__init__.py index 24e7ddea9e8..f06c28c5001 100644 --- a/litellm/integrations/gitlab/__init__.py +++ b/litellm/integrations/gitlab/__init__.py @@ -30,9 +30,7 @@ def set_global_gitlab_config(config: dict) -> None: litellm.global_gitlab_config = config # type: ignore -def prompt_initializer( - litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec" -) -> "CustomPromptManagement": +def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec") -> "CustomPromptManagement": """ Initialize a prompt from a Gitlab repository. """ diff --git a/litellm/integrations/gitlab/gitlab_client.py b/litellm/integrations/gitlab/gitlab_client.py index 60f73256185..ca366274ccd 100644 --- a/litellm/integrations/gitlab/gitlab_client.py +++ b/litellm/integrations/gitlab/gitlab_client.py @@ -108,9 +108,7 @@ class GitLabClient: raise ValueError("ref must be a non-empty string") self.ref = ref - def get_file_content( - self, file_path: str, *, ref: Optional[str] = None - ) -> Optional[str]: + def get_file_content(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]: """ Fetch the content of a file from the GitLab repository at the given ref (tag, branch, or commit SHA). If `ref` is None, uses self.ref. @@ -132,11 +130,7 @@ class GitLabClient: resp.raise_for_status() ctype = (resp.headers.get("content-type") or "").lower() - if ( - ctype.startswith("text/") - or "charset=" in ctype - or ctype.startswith("application/json") - ): + if ctype.startswith("text/") or "charset=" in ctype or ctype.startswith("application/json"): return resp.text try: return resp.content.decode("utf-8") @@ -152,14 +146,10 @@ class GitLabClient: f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'." ) if status == 401: - raise Exception( - "Authentication failed. Check your GitLab token and auth_method." - ) + raise Exception("Authentication failed. Check your GitLab token and auth_method.") raise Exception(f"Failed to fetch file '{file_path}': {e}") - def _get_file_content_via_json( - self, file_path: str, *, ref: Optional[str] = None - ) -> Optional[str]: + def _get_file_content_via_json(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]: """ Fallback for get_file_content(): use the JSON file API which returns base64 content. """ @@ -187,12 +177,8 @@ class GitLabClient: f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'." ) if status == 401: - raise Exception( - "Authentication failed. Check your GitLab token and auth_method." - ) - raise Exception( - f"Failed to fetch file '{file_path}' via JSON endpoint: {e}" - ) + raise Exception("Authentication failed. Check your GitLab token and auth_method.") + raise Exception(f"Failed to fetch file '{file_path}' via JSON endpoint: {e}") def list_files( self, @@ -240,9 +226,7 @@ class GitLabClient: f"Access denied to directory '{directory_path}'. Check your GitLab permissions for project '{self.project}'." ) if status == 401: - raise Exception( - "Authentication failed. Check your GitLab token and auth_method." - ) + raise Exception("Authentication failed. Check your GitLab token and auth_method.") raise Exception(f"Failed to list files in '{directory_path}': {e}") def get_repository_info(self) -> Dict[str, Any]: @@ -274,9 +258,7 @@ class GitLabClient: except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata( - self, file_path: str, *, ref: Optional[str] = None - ) -> Optional[Dict[str, Any]]: + def get_file_metadata(self, file_path: str, *, ref: Optional[str] = None) -> Optional[Dict[str, Any]]: """ Get minimal metadata about a file via RAW endpoint headers at a given ref. diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index 99d9d9b285b..4896f95f398 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -54,9 +54,7 @@ class GitLabPromptTemplate: self.temperature = metadata.get("temperature") self.max_tokens = metadata.get("max_tokens") self.input_schema = metadata.get("input", {}).get("schema", {}) - self.optional_params = { - k: v for k, v in metadata.items() if k not in ["model", "input", "content"] - } + self.optional_params = {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]} def __repr__(self): return f"GitLabPromptTemplate(id='{self.template_id}', model='{self.model}')" @@ -86,9 +84,7 @@ class GitLabTemplateManager: # Folder inside repo to look for prompts (e.g., "prompts" or "prompts/chat") self.prompts_path: str = ( - self.gitlab_config.get("prompts_path") - or self.gitlab_config.get("folder") - or "" + self.gitlab_config.get("prompts_path") or self.gitlab_config.get("folder") or "" ).strip("/") # Templates fetched from a GitLab repo are not trustworthy: @@ -134,9 +130,7 @@ class GitLabTemplateManager: # ---------- loading ---------- - def _load_prompt_from_gitlab( - self, prompt_id: str, *, ref: Optional[str] = None - ) -> None: + def _load_prompt_from_gitlab(self, prompt_id: str, *, ref: Optional[str] = None) -> None: """Load a specific .prompt file from GitLab (scoped under prompts_path if set).""" try: # prompt_id = decode_prompt_id(prompt_id) @@ -146,9 +140,7 @@ class GitLabTemplateManager: template = self._parse_prompt_file(prompt_content, prompt_id) self.prompts[prompt_id] = template except Exception as e: - raise Exception( - f"Failed to load prompt '{encode_prompt_id(prompt_id)}' from GitLab: {e}" - ) + raise Exception(f"Failed to load prompt '{encode_prompt_id(prompt_id)}' from GitLab: {e}") def load_all_prompts(self, *, recursive: bool = True) -> List[str]: """ @@ -215,9 +207,7 @@ class GitLabTemplateManager: result[key] = value.strip("\"'") return result - def render_template( - self, template_id: str, variables: Optional[Dict[str, Any]] = None - ) -> str: + def render_template(self, template_id: str, variables: Optional[Dict[str, Any]] = None) -> str: if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") template = self.prompts[template_id] @@ -335,9 +325,7 @@ class GitLabPromptManager(CustomPromptManagement): if not template: raise ValueError(f"Prompt template '{prompt_id}' not found") - rendered_prompt = self.prompt_manager.render_template( - prompt_id, prompt_variables or {} - ) + rendered_prompt = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) metadata = { "model": template.model, @@ -364,17 +352,13 @@ class GitLabPromptManager(CustomPromptManagement): # Precedence: explicit prompt_version → per-call git_ref kwarg → manager override → config default git_ref = prompt_version or kwargs.get("git_ref") or self._ref_override - rendered_prompt, prompt_metadata = self.get_prompt_template( - prompt_id, prompt_variables, ref=git_ref - ) + rendered_prompt, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables, ref=git_ref) parsed_messages = self._parse_prompt_to_messages(rendered_prompt) if parsed_messages: final_messages: List[AllMessageValues] = parsed_messages else: - final_messages = [ - {"role": "user", "content": rendered_prompt} - ] + messages # type: ignore + final_messages = [{"role": "user", "content": rendered_prompt}] + messages # type: ignore if litellm_params is None: litellm_params = {} @@ -396,9 +380,7 @@ class GitLabPromptManager(CustomPromptManagement): except Exception as e: import litellm - litellm._logging.verbose_proxy_logger.error( - f"Error in GitLab prompt pre_call_hook: {e}" - ) + litellm._logging.verbose_proxy_logger.error(f"Error in GitLab prompt pre_call_hook: {e}") return messages, litellm_params def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]: @@ -446,9 +428,7 @@ class GitLabPromptManager(CustomPromptManagement): current_content.append(line) if current_role and current_content: - messages.append( - {"role": current_role, "content": "\n".join(current_content).strip()} - ) # type: ignore + messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore if not messages and prompt_content.strip(): messages = [{"role": "user", "content": prompt_content.strip()}] # type: ignore return messages @@ -514,9 +494,7 @@ class GitLabPromptManager(CustomPromptManagement): ) self.prompt_manager._load_prompt_from_gitlab(decoded_id, ref=git_ref) - rendered_prompt, prompt_metadata = self.get_prompt_template( - prompt_id, prompt_variables - ) + rendered_prompt, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) messages = self._parse_prompt_to_messages(rendered_prompt) template_model = prompt_metadata.get("model") @@ -678,9 +656,7 @@ class GitLabPromptCache: ref=ref, gitlab_client=gitlab_client, ) - self.template_manager: GitLabTemplateManager = ( - self.prompt_manager.prompt_manager - ) + self.template_manager: GitLabTemplateManager = self.prompt_manager.prompt_manager # In-memory stores self._by_file: Dict[str, Dict[str, Any]] = {} @@ -695,9 +671,7 @@ class GitLabPromptCache: Scan GitLab for all .prompt files under prompts_path, load and parse each, and return the mapping of repo file path -> JSON-like dict. """ - ids = self.template_manager.list_templates( - recursive=recursive - ) # IDs relative to prompts_path + ids = self.template_manager.list_templates(recursive=recursive) # IDs relative to prompts_path for pid in ids: # Ensure template is loaded into TemplateManager if pid not in self.template_manager.prompts: @@ -711,9 +685,7 @@ class GitLabPromptCache: if tmpl is None: continue - file_path = self.template_manager._id_to_repo_path( - pid - ) # "prompts/chat/..../file.prompt" + file_path = self.template_manager._id_to_repo_path(pid) # "prompts/chat/..../file.prompt" entry = self._template_to_json(pid, tmpl) self._by_file[file_path] = entry @@ -757,9 +729,7 @@ class GitLabPromptCache: # Internals # ------------------------- - def _template_to_json( - self, prompt_id: str, tmpl: GitLabPromptTemplate - ) -> Dict[str, Any]: + def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> Dict[str, Any]: """ Normalize a GitLabPromptTemplate into a JSON-like dict that is easy to serialize. """ @@ -774,9 +744,7 @@ class GitLabPromptCache: return { "id": prompt_id, # e.g. "greet/hi" - "path": self.template_manager._id_to_repo_path( - prompt_id - ), # e.g. "prompts/chat/greet/hi.prompt" + "path": self.template_manager._id_to_repo_path(prompt_id), # e.g. "prompts/chat/greet/hi.prompt" "content": tmpl.content, # rendered content (without frontmatter) "metadata": md, # parsed frontmatter "model": model, diff --git a/litellm/integrations/greenscale.py b/litellm/integrations/greenscale.py index 430c3d0abf2..e2aca361010 100644 --- a/litellm/integrations/greenscale.py +++ b/litellm/integrations/greenscale.py @@ -22,18 +22,12 @@ class GreenscaleLogger: data = { "modelId": kwargs.get("model"), "inputTokenCount": response_json.get("usage", {}).get("prompt_tokens"), - "outputTokenCount": response_json.get("usage", {}).get( - "completion_tokens" - ), + "outputTokenCount": response_json.get("usage", {}).get("completion_tokens"), } - data["timestamp"] = datetime.now(timezone.utc).strftime( - "%Y-%m-%dT%H:%M:%SZ" - ) + data["timestamp"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") if type(end_time) is datetime and type(start_time) is datetime: - data["invocationLatency"] = int( - (end_time - start_time).total_seconds() * 1000 - ) + data["invocationLatency"] = int((end_time - start_time).total_seconds() * 1000) # Add additional metadata keys to tags tags = [] @@ -45,9 +39,7 @@ class GreenscaleLogger: elif key == "greenscale_application": data["application"] = value else: - tags.append( - {"key": key.replace("greenscale_", ""), "value": str(value)} - ) + tags.append({"key": key.replace("greenscale_", ""), "value": str(value)}) data["tags"] = tags @@ -60,13 +52,9 @@ class GreenscaleLogger: data=json.dumps(data, default=str), ) if response.status_code != 200: - print_verbose( - f"Greenscale Logger Error - {response.text}, {response.status_code}" - ) + print_verbose(f"Greenscale Logger Error - {response.text}, {response.status_code}") else: print_verbose(f"Greenscale Logger Succeeded - {response.text}") except Exception as e: - print_verbose( - f"Greenscale Logger Error - {e}, Stack trace: {traceback.format_exc()}" - ) + print_verbose(f"Greenscale Logger Error - {e}, Stack trace: {traceback.format_exc()}") pass diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index 376952033a0..21e9479491e 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -31,9 +31,7 @@ class HeliconeLogger: self.is_mock_mode = should_use_helicone_mock() if self.is_mock_mode: create_mock_helicone_client() - verbose_logger.info( - "[HELICONE MOCK] Helicone logger initialized in mock mode" - ) + verbose_logger.info("[HELICONE MOCK] Helicone logger initialized in mock mode") self.provider_url = "https://api.openai.com/v1" self.key = os.getenv("HELICONE_API_KEY") @@ -106,9 +104,7 @@ class HeliconeLogger: if metadata is None: metadata = {} - proxy_headers = ( - litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} - ) + proxy_headers = litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} for header_key in proxy_headers: if header_key.startswith("helicone_"): @@ -121,14 +117,10 @@ class HeliconeLogger: return metadata - def log_success( - self, model, messages, response_obj, start_time, end_time, print_verbose, kwargs - ): + def log_success(self, model, messages, response_obj, start_time, end_time, print_verbose, kwargs): # Method definition try: - print_verbose( - f"Helicone Logging - Enters logging function for model {model}" - ) + print_verbose(f"Helicone Logging - Enters logging function for model {model}") litellm_params = kwargs.get("litellm_params", {}) custom_llm_provider = litellm_params.get("custom_llm_provider", "") kwargs.get("litellm_call_id", None) @@ -136,29 +128,19 @@ class HeliconeLogger: metadata = self.add_metadata_from_header(litellm_params, metadata) # Check if model is a vertex_ai model - is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith( - "vertex_ai/" - ) + is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/") model = ( model - if any( - accepted_model in model - for accepted_model in self.helicone_model_list - ) - or is_vertex_ai + if any(accepted_model in model for accepted_model in self.helicone_model_list) or is_vertex_ai else "gpt-3.5-turbo" ) provider_request = {"model": model, "messages": messages} - if isinstance(response_obj, litellm.EmbeddingResponse) or isinstance( - response_obj, litellm.ModelResponse - ): + if isinstance(response_obj, litellm.EmbeddingResponse) or isinstance(response_obj, litellm.ModelResponse): response_obj = response_obj.json() if "claude" in model and not is_vertex_ai: - response_obj = self.claude_mapping( - model=model, messages=messages, response_obj=response_obj - ) + response_obj = self.claude_mapping(model=model, messages=messages, response_obj=response_obj) providerResponse = { "json": response_obj, @@ -183,13 +165,9 @@ class HeliconeLogger: "Content-Type": "application/json", } start_time_seconds = int(start_time.timestamp()) - start_time_milliseconds = int( - (start_time.timestamp() - start_time_seconds) * 1000 - ) + start_time_milliseconds = int((start_time.timestamp() - start_time_seconds) * 1000) end_time_seconds = int(end_time.timestamp()) - end_time_milliseconds = int( - (end_time.timestamp() - end_time_seconds) * 1000 - ) + end_time_milliseconds = int((end_time.timestamp() - end_time_seconds) * 1000) meta = {"Helicone-Auth": f"Bearer {self.key}"} meta.update(metadata) data = { @@ -213,9 +191,7 @@ class HeliconeLogger: response = litellm.module_level_client.post(url, headers=headers, json=data) if response.status_code == 200: if self.is_mock_mode: - print_verbose( - "[HELICONE MOCK] Helicone Logging - Successfully mocked!" - ) + print_verbose("[HELICONE MOCK] Helicone Logging - Successfully mocked!") else: print_verbose("Helicone Logging - Success!") else: diff --git a/litellm/integrations/helicone_mock_client.py b/litellm/integrations/helicone_mock_client.py index c2d3dfdf5bc..02530692d43 100644 --- a/litellm/integrations/helicone_mock_client.py +++ b/litellm/integrations/helicone_mock_client.py @@ -32,6 +32,4 @@ _config = MockClientConfig( patch_http_handler=True, # Patch HTTPHandler.post directly ) -create_mock_helicone_client, should_use_helicone_mock = create_mock_client_factory( - _config -) +create_mock_helicone_client, should_use_helicone_mock = create_mock_client_factory(_config) diff --git a/litellm/integrations/humanloop.py b/litellm/integrations/humanloop.py index 369df5ee0bd..2a5cb70baee 100644 --- a/litellm/integrations/humanloop.py +++ b/litellm/integrations/humanloop.py @@ -32,12 +32,8 @@ class HumanLoopPromptManager(DualCache): def integration_name(self): return "humanloop" - def _get_prompt_from_id_cache( - self, humanloop_prompt_id: str - ) -> Optional[PromptManagementClient]: - return cast( - Optional[PromptManagementClient], self.get_cache(key=humanloop_prompt_id) - ) + def _get_prompt_from_id_cache(self, humanloop_prompt_id: str) -> Optional[PromptManagementClient]: + return cast(Optional[PromptManagementClient], self.get_cache(key=humanloop_prompt_id)) def _compile_prompt_helper( self, prompt_template: List[AllMessageValues], prompt_variables: Dict[str, Any] @@ -64,9 +60,7 @@ class HumanLoopPromptManager(DualCache): return compiled_prompts - def _get_prompt_from_id_api( - self, humanloop_prompt_id: str, humanloop_api_key: str - ) -> PromptManagementClient: + def _get_prompt_from_id_api(self, humanloop_prompt_id: str, humanloop_api_key: str) -> PromptManagementClient: client = _get_httpx_client() base_url = "https://api.humanloop.com/v5/prompts/{}".format(humanloop_prompt_id) @@ -104,14 +98,10 @@ class HumanLoopPromptManager(DualCache): optional_params=optional_params, ) - def _get_prompt_from_id( - self, humanloop_prompt_id: str, humanloop_api_key: str - ) -> PromptManagementClient: + def _get_prompt_from_id(self, humanloop_prompt_id: str, humanloop_api_key: str) -> PromptManagementClient: prompt = self._get_prompt_from_id_cache(humanloop_prompt_id) if prompt is None: - prompt = self._get_prompt_from_id_api( - humanloop_prompt_id, humanloop_api_key - ) + prompt = self._get_prompt_from_id_api(humanloop_prompt_id, humanloop_api_key) self.set_cache( key=humanloop_prompt_id, value=prompt, @@ -136,9 +126,7 @@ class HumanLoopPromptManager(DualCache): return compiled_prompt - def _get_model_from_prompt( - self, prompt_management_client: PromptManagementClient, model: str - ) -> str: + def _get_model_from_prompt(self, prompt_management_client: PromptManagementClient, model: str) -> str: if prompt_management_client["model"] is not None: return prompt_management_client["model"] else: @@ -167,9 +155,7 @@ class HumanloopLogger(CustomLogger): List[AllMessageValues], dict, ]: - humanloop_api_key = dynamic_callback_params.get( - "humanloop_api_key" - ) or get_secret_str("HUMANLOOP_API_KEY") + humanloop_api_key = dynamic_callback_params.get("humanloop_api_key") or get_secret_str("HUMANLOOP_API_KEY") if prompt_id is None: raise ValueError("prompt_id is required for Humanloop integration") @@ -201,8 +187,6 @@ class HumanloopLogger(CustomLogger): **prompt_template_optional_params, } - model = prompt_manager._get_model_from_prompt( - prompt_management_client=prompt_template, model=model - ) + model = prompt_manager._get_model_from_prompt(prompt_management_client=prompt_template, model=model) return model, updated_messages, updated_non_default_params diff --git a/litellm/integrations/lago.py b/litellm/integrations/lago.py index c7c010f9976..0052e04644d 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -32,9 +32,7 @@ class LagoLogger(CustomLogger): def __init__(self) -> None: super().__init__() self.validate_environment() - self.async_http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_http_handler = HTTPHandler() def validate_environment(self): @@ -70,8 +68,7 @@ class LagoLogger(CustomLogger): usage = {} if ( - isinstance(response_obj, litellm.ModelResponse) - or isinstance(response_obj, litellm.EmbeddingResponse) + isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse) ) and hasattr(response_obj, "usage"): usage = { "prompt_tokens": response_obj["usage"].get("prompt_tokens", 0), @@ -89,9 +86,7 @@ class LagoLogger(CustomLogger): charge_by: Literal["end_user_id", "team_id", "user_id"] = "end_user_id" external_customer_id: Optional[str] = None - if os.getenv("LAGO_API_CHARGE_BY", None) is not None and isinstance( - os.environ["LAGO_API_CHARGE_BY"], str - ): + if os.getenv("LAGO_API_CHARGE_BY", None) is not None and isinstance(os.environ["LAGO_API_CHARGE_BY"], str): if os.environ["LAGO_API_CHARGE_BY"] in [ "end_user_id", "user_id", @@ -124,9 +119,7 @@ class LagoLogger(CustomLogger): } } - verbose_logger.debug( - "\033[91mLogged Lago Object:\n{}\033[0m\n".format(returned_val) - ) + verbose_logger.debug("\033[91mLogged Lago Object:\n{}\033[0m\n".format(returned_val)) return returned_val def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -166,9 +159,7 @@ class LagoLogger(CustomLogger): verbose_logger.debug("ENTERS LAGO CALLBACK") _url = os.getenv("LAGO_API_BASE") assert _url is not None and isinstance(_url, str), ( - "LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={}".format( - _url - ) + "LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={}".format(_url) ) if _url.endswith("/"): _url += "api/v1/events" diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index b1c6956a16c..8068a8c0b0c 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -76,15 +76,9 @@ def _extract_cache_read_input_tokens(usage_obj) -> int: # Check prompt_tokens_details.cached_tokens (used by Gemini and other providers) if hasattr(usage_obj, "prompt_tokens_details"): prompt_tokens_details = getattr(usage_obj, "prompt_tokens_details", None) - if prompt_tokens_details is not None and hasattr( - prompt_tokens_details, "cached_tokens" - ): + if prompt_tokens_details is not None and hasattr(prompt_tokens_details, "cached_tokens"): cached_tokens = getattr(prompt_tokens_details, "cached_tokens", None) - if ( - cached_tokens is not None - and isinstance(cached_tokens, (int, float)) - and cached_tokens > 0 - ): + if cached_tokens is not None and isinstance(cached_tokens, (int, float)) and cached_tokens > 0: cache_read_input_tokens = cached_tokens return cache_read_input_tokens @@ -101,14 +95,10 @@ def resolve_langfuse_credentials( secret_key = langfuse_secret or langfuse_secret_key public_key = langfuse_public_key else: - secret_key = ( - langfuse_secret or langfuse_secret_key or os.getenv("LANGFUSE_SECRET_KEY") - ) + secret_key = langfuse_secret or langfuse_secret_key or os.getenv("LANGFUSE_SECRET_KEY") public_key = langfuse_public_key or os.getenv("LANGFUSE_PUBLIC_KEY") - resolved_host = langfuse_host or os.getenv( - "LANGFUSE_HOST", "https://cloud.langfuse.com" - ) + resolved_host = langfuse_host or os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com") return public_key, secret_key, resolved_host @@ -130,25 +120,18 @@ class LangFuseLogger: raise Exception( f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\n{traceback.format_exc()}\033[0m" ) - self.public_key, self.secret_key, self.langfuse_host = ( - resolve_langfuse_credentials( - langfuse_public_key=langfuse_public_key, - langfuse_secret=langfuse_secret, - langfuse_host=langfuse_host, - allow_env_credentials=allow_env_credentials, - ) + self.public_key, self.secret_key, self.langfuse_host = resolve_langfuse_credentials( + langfuse_public_key=langfuse_public_key, + langfuse_secret=langfuse_secret, + langfuse_host=langfuse_host, + allow_env_credentials=allow_env_credentials, ) - if not ( - self.langfuse_host.startswith("http://") - or self.langfuse_host.startswith("https://") - ): + if not (self.langfuse_host.startswith("http://") or self.langfuse_host.startswith("https://")): # add http:// if unset, assume communicating over private network - e.g. render self.langfuse_host = "http://" + self.langfuse_host self.langfuse_release = os.getenv("LANGFUSE_RELEASE") self.langfuse_debug = os.getenv("LANGFUSE_DEBUG") - self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval( - flush_interval - ) + self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval(flush_interval) if should_use_langfuse_mock(): self.langfuse_client = create_mock_langfuse_client() @@ -188,16 +171,10 @@ class LangFuseLogger: if os.getenv("UPSTREAM_LANGFUSE_SECRET_KEY") is not None: upstream_langfuse_debug_env = os.getenv("UPSTREAM_LANGFUSE_DEBUG") upstream_langfuse_debug = ( - str_to_bool(upstream_langfuse_debug_env) - if upstream_langfuse_debug_env is not None - else None - ) - self.upstream_langfuse_secret_key = os.getenv( - "UPSTREAM_LANGFUSE_SECRET_KEY" - ) - self.upstream_langfuse_public_key = os.getenv( - "UPSTREAM_LANGFUSE_PUBLIC_KEY" + str_to_bool(upstream_langfuse_debug_env) if upstream_langfuse_debug_env is not None else None ) + self.upstream_langfuse_secret_key = os.getenv("UPSTREAM_LANGFUSE_SECRET_KEY") + self.upstream_langfuse_public_key = os.getenv("UPSTREAM_LANGFUSE_PUBLIC_KEY") self.upstream_langfuse_host = os.getenv("UPSTREAM_LANGFUSE_HOST") self.upstream_langfuse_release = os.getenv("UPSTREAM_LANGFUSE_RELEASE") self.upstream_langfuse_debug = upstream_langfuse_debug_env @@ -206,11 +183,7 @@ class LangFuseLogger: secret_key=self.upstream_langfuse_secret_key, host=self.upstream_langfuse_host, release=self.upstream_langfuse_release, - debug=( - upstream_langfuse_debug - if upstream_langfuse_debug is not None - else False - ), + debug=(upstream_langfuse_debug if upstream_langfuse_debug is not None else False), ) else: self.upstream_langfuse = None @@ -231,9 +204,7 @@ class LangFuseLogger: ) langfuse_client = Langfuse(**parameters) litellm.initialized_langfuse_clients += 1 - verbose_logger.debug( - f"Created langfuse client number {litellm.initialized_langfuse_clients}" - ) + verbose_logger.debug(f"Created langfuse client number {litellm.initialized_langfuse_clients}") return langfuse_client @staticmethod @@ -254,21 +225,15 @@ class LangFuseLogger: if metadata is None: metadata = {} - proxy_headers = ( - litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} - ) + proxy_headers = litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} for metadata_param_key in proxy_headers: if metadata_param_key.startswith("langfuse_"): trace_param_key = metadata_param_key.replace("langfuse_", "", 1) if trace_param_key in metadata: - verbose_logger.warning( - f"Overwriting Langfuse `{trace_param_key}` from request header" - ) + verbose_logger.warning(f"Overwriting Langfuse `{trace_param_key}` from request header") else: - verbose_logger.debug( - f"Found Langfuse `{trace_param_key}` in request header" - ) + verbose_logger.debug(f"Found Langfuse `{trace_param_key}` in request header") metadata[trace_param_key] = proxy_headers.get(metadata_param_key) return metadata @@ -298,9 +263,7 @@ class LangFuseLogger: Logs a success or error event on Langfuse """ try: - verbose_logger.debug( - f"Langfuse Logging - Enters logging function for model {kwargs}" - ) + verbose_logger.debug(f"Langfuse Logging - Enters logging function for model {kwargs}") # set default values for input/output for langfuse logging input = None @@ -308,9 +271,7 @@ class LangFuseLogger: litellm_params = kwargs.get("litellm_params", {}) litellm_call_id = kwargs.get("litellm_call_id", None) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None + metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None metadata = self.add_metadata_from_header(litellm_params, metadata) optional_params = safe_deep_copy(kwargs.get("optional_params", {})) @@ -341,9 +302,7 @@ class LangFuseLogger: level=level, status_message=status_message, ) - verbose_logger.debug( - f"OUTPUT IN LANGFUSE: {output}; original: {response_obj}" - ) + verbose_logger.debug(f"OUTPUT IN LANGFUSE: {output}; original: {response_obj}") trace_id = None generation_id = None if self._is_langfuse_v2(): @@ -373,16 +332,12 @@ class LangFuseLogger: input=input, response_obj=response_obj, ) - verbose_logger.debug( - f"Langfuse Layer Logging - final response object: {response_obj}" - ) + verbose_logger.debug(f"Langfuse Layer Logging - final response object: {response_obj}") verbose_logger.info("Langfuse Layer Logging - logging success") return {"trace_id": trace_id, "generation_id": generation_id} except Exception as e: - verbose_logger.exception( - "Langfuse Layer Error(): Exception occured - {}".format(str(e)) - ) + verbose_logger.exception("Langfuse Layer Error(): Exception occured - {}".format(str(e))) return {"trace_id": None, "generation_id": None} def _get_langfuse_input_output_content( @@ -420,52 +375,33 @@ class LangFuseLogger: """ input = None output: Optional[Union[str, dict, List[Any]]] = None - if ( - level == "ERROR" - and status_message is not None - and isinstance(status_message, str) - ): + if level == "ERROR" and status_message is not None and isinstance(status_message, str): input = prompt output = status_message elif response_obj is not None and ( - kwargs.get("call_type", None) == "embedding" - or isinstance(response_obj, litellm.EmbeddingResponse) + kwargs.get("call_type", None) == "embedding" or isinstance(response_obj, litellm.EmbeddingResponse) ): input = prompt output = None - elif response_obj is not None and isinstance( - response_obj, litellm.ModelResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.ModelResponse): input = prompt output = self._get_chat_content_for_langfuse(response_obj) - elif response_obj is not None and isinstance( - response_obj, litellm.HttpxBinaryResponseContent - ): + elif response_obj is not None and isinstance(response_obj, litellm.HttpxBinaryResponseContent): input = prompt output = "speech-output" - elif response_obj is not None and isinstance( - response_obj, litellm.TextCompletionResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.TextCompletionResponse): input = prompt output = self._get_text_completion_content_for_langfuse(response_obj) - elif response_obj is not None and isinstance( - response_obj, litellm.ImageResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.ImageResponse): input = prompt output = response_obj.get("data", None) - elif response_obj is not None and isinstance( - response_obj, litellm.TranscriptionResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.TranscriptionResponse): input = prompt output = response_obj.get("text", None) - elif response_obj is not None and isinstance( - response_obj, litellm.RerankResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.RerankResponse): input = prompt output = response_obj.results - elif response_obj is not None and isinstance( - response_obj, litellm.ResponsesAPIResponse - ): + elif response_obj is not None and isinstance(response_obj, litellm.ResponsesAPIResponse): input = prompt output = self._get_responses_api_content_for_langfuse(response_obj) elif ( @@ -486,9 +422,7 @@ class LangFuseLogger: output = response_obj.get("response", "") return input, output - async def _async_log_event( - self, kwargs, response_obj, start_time, end_time, user_id - ): + async def _async_log_event(self, kwargs, response_obj, start_time, end_time, user_id): """ Langfuse SDK uses a background thread to log events @@ -528,9 +462,7 @@ class LangFuseLogger: ) custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) - model_name = reconstruct_model_name( - kwargs.get("model", ""), custom_llm_provider, metadata - ) + model_name = reconstruct_model_name(kwargs.get("model", ""), custom_llm_provider, metadata) trace.generation( CreateGeneration( @@ -579,19 +511,13 @@ class LangFuseLogger: if standard_logging_object is None: end_user_id = None - prompt_management_metadata: Optional[ - StandardLoggingPromptManagementMetadata - ] = None + prompt_management_metadata: Optional[StandardLoggingPromptManagementMetadata] = None else: - end_user_id = standard_logging_object["metadata"].get( - "user_api_key_end_user_id", None - ) + end_user_id = standard_logging_object["metadata"].get("user_api_key_end_user_id", None) prompt_management_metadata = cast( Optional[StandardLoggingPromptManagementMetadata], - standard_logging_object["metadata"].get( - "prompt_management_metadata", None - ), + standard_logging_object["metadata"].get("prompt_management_metadata", None), ) # Clean Metadata before logging - never log raw metadata @@ -599,9 +525,7 @@ class LangFuseLogger: # we clean out all extra litellm metadata params before logging clean_metadata: Dict[str, Any] = {} if prompt_management_metadata is not None: - clean_metadata["prompt_management_metadata"] = ( - prompt_management_metadata - ) + clean_metadata["prompt_management_metadata"] = prompt_management_metadata if isinstance(metadata, dict): for key, value in metadata.items(): # generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy @@ -624,9 +548,7 @@ class LangFuseLogger: clean_metadata[key] = value # Add default langfuse tags - tags = self.add_default_langfuse_tags( - tags=tags, kwargs=kwargs, metadata=metadata - ) + tags = self.add_default_langfuse_tags(tags=tags, kwargs=kwargs, metadata=metadata) session_id = clean_metadata.pop("session_id", None) trace_name = cast(Optional[str], clean_metadata.pop("trace_name", None)) @@ -649,9 +571,9 @@ class LangFuseLogger: mask_output = clean_metadata.pop("mask_output", False) # Look for masking function in the dedicated location first (set by scrub_sensitive_keys_in_metadata) # Fall back to metadata for backwards compatibility - masking_function = litellm_params.get( - "_langfuse_masking_function" - ) or clean_metadata.pop("langfuse_masking_function", None) + masking_function = litellm_params.get("_langfuse_masking_function") or clean_metadata.pop( + "langfuse_masking_function", None + ) # Apply custom masking function if provided if masking_function is not None and callable(masking_function): @@ -672,27 +594,19 @@ class LangFuseLogger: for metadata_param_key in update_trace_keys: trace_param_key = metadata_param_key.replace("trace_", "") if trace_param_key not in trace_params: - updated_trace_value = clean_metadata.pop( - metadata_param_key, None - ) + updated_trace_value = clean_metadata.pop(metadata_param_key, None) if updated_trace_value is not None: trace_params[trace_param_key] = updated_trace_value # Pop the trace specific keys that would have been popped if there were a new trace - for key in list( - filter(lambda key: key.startswith("trace_"), clean_metadata.keys()) - ): + for key in list(filter(lambda key: key.startswith("trace_"), clean_metadata.keys())): clean_metadata.pop(key, None) # Special keys that are found in the function arguments and not the metadata if "input" in update_trace_keys: - trace_params["input"] = ( - input if not mask_input else "redacted-by-litellm" - ) + trace_params["input"] = input if not mask_input else "redacted-by-litellm" if "output" in update_trace_keys: - trace_params["output"] = ( - output if not mask_output else "redacted-by-litellm" - ) + trace_params["output"] = output if not mask_output else "redacted-by-litellm" else: # don't overwrite an existing trace trace_params = { "id": trace_id, @@ -704,19 +618,13 @@ class LangFuseLogger: ), # If provided just version, it will applied to the trace as well, if applied a trace version it will take precedence "user_id": end_user_id, } - for key in list( - filter(lambda key: key.startswith("trace_"), clean_metadata.keys()) - ): - trace_params[key.replace("trace_", "")] = clean_metadata.pop( - key, None - ) + for key in list(filter(lambda key: key.startswith("trace_"), clean_metadata.keys())): + trace_params[key.replace("trace_", "")] = clean_metadata.pop(key, None) if level == "ERROR": trace_params["status_message"] = output else: - trace_params["output"] = ( - output if not mask_output else "redacted-by-litellm" - ) + trace_params["output"] = output if not mask_output else "redacted-by-litellm" if debug is True or (isinstance(debug, str) and debug.lower() == "true"): if "metadata" in trace_params: @@ -731,9 +639,7 @@ class LangFuseLogger: clean_metadata["litellm_response_cost"] = cost if standard_logging_object is not None: hidden_params = standard_logging_object.get("hidden_params", {}) - clean_metadata["hidden_params"] = filter_exceptions_from_params( - hidden_params - ) + clean_metadata["hidden_params"] = filter_exceptions_from_params(hidden_params) if ( litellm.langfuse_default_tags is not None @@ -791,30 +697,19 @@ class LangFuseLogger: usage = None usage_details = None if response_obj is not None: - if ( - hasattr(response_obj, "id") - and response_obj.get("id", None) is not None - ): - generation_id = litellm.utils.get_logging_id( - start_time, response_obj - ) + if hasattr(response_obj, "id") and response_obj.get("id", None) is not None: + generation_id = litellm.utils.get_logging_id(start_time, response_obj) _usage_obj = getattr(response_obj, "usage", None) if _usage_obj: # Safely get usage values, defaulting None to 0 for Langfuse compatibility. # Some providers may return null for token counts. prompt_tokens = getattr(_usage_obj, "prompt_tokens", None) or 0 - completion_tokens = ( - getattr(_usage_obj, "completion_tokens", None) or 0 - ) + completion_tokens = getattr(_usage_obj, "completion_tokens", None) or 0 total_tokens = getattr(_usage_obj, "total_tokens", None) or 0 - cache_creation_input_tokens = ( - _usage_obj.get("cache_creation_input_tokens") or 0 - ) - cache_read_input_tokens = _extract_cache_read_input_tokens( - _usage_obj - ) + cache_creation_input_tokens = _usage_obj.get("cache_creation_input_tokens") or 0 + cache_read_input_tokens = _extract_cache_read_input_tokens(_usage_obj) usage = { "prompt_tokens": prompt_tokens, @@ -836,12 +731,8 @@ class LangFuseLogger: # if `generation_name` is None, use sensible default values # If using litellm proxy user `key_alias` if not None # If `key_alias` is None, just log `litellm-{call_type}` as the generation name - _user_api_key_alias = cast( - Optional[str], clean_metadata.get("user_api_key_alias", None) - ) - generation_name = ( - f"litellm-{cast(str, kwargs.get('call_type', 'completion'))}" - ) + _user_api_key_alias = cast(Optional[str], clean_metadata.get("user_api_key_alias", None)) + generation_name = f"litellm-{cast(str, kwargs.get('call_type', 'completion'))}" if _user_api_key_alias is not None: generation_name = f"litellm:{_user_api_key_alias}" @@ -854,9 +745,7 @@ class LangFuseLogger: optional_params["system_fingerprint"] = system_fingerprint custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) - model_name = reconstruct_model_name( - kwargs.get("model", ""), custom_llm_provider, metadata - ) + model_name = reconstruct_model_name(kwargs.get("model", ""), custom_llm_provider, metadata) generation_params = { "name": generation_name, @@ -889,9 +778,7 @@ class LangFuseLogger: generation_params["status_message"] = output if self._supports_completion_start_time(): - generation_params["completion_start_time"] = kwargs.get( - "completion_start_time", None - ) + generation_params["completion_start_time"] = kwargs.get("completion_start_time", None) generation_client = trace.generation(**generation_params) @@ -965,9 +852,7 @@ class LangFuseLogger: - cache_key """ - if litellm.langfuse_default_tags is not None and isinstance( - litellm.langfuse_default_tags, list - ): + if litellm.langfuse_default_tags is not None and isinstance(litellm.langfuse_default_tags, list): if "cache_hit" in litellm.langfuse_default_tags: _cache_hit_value = kwargs.get("cache_hit", False) tags.append(f"cache_hit:{_cache_hit_value}") @@ -976,9 +861,7 @@ class LangFuseLogger: _cache_key = _hidden_params.get("cache_key", None) if _cache_key is None and litellm.cache is not None: # fallback to using "preset_cache_key" - _preset_cache_key = litellm.cache._get_preset_cache_key_from_kwargs( - **kwargs - ) + _preset_cache_key = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) _cache_key = _preset_cache_key tags.append(f"cache_key:{_cache_key}") return tags @@ -1000,9 +883,7 @@ class LangFuseLogger: return Version(self.langfuse_sdk_version) >= Version("2.7.3") @staticmethod - def _apply_masking_function( - data: Any, masking_function: Callable[[Any], Any] - ) -> Any: + def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any: """ Apply a masking function to data, handling different data types. @@ -1022,22 +903,15 @@ class LangFuseLogger: elif isinstance(data, dict): masked_dict = {} for key, value in data.items(): - masked_dict[key] = LangFuseLogger._apply_masking_function( - value, masking_function - ) + masked_dict[key] = LangFuseLogger._apply_masking_function(value, masking_function) return masked_dict elif isinstance(data, list): - return [ - LangFuseLogger._apply_masking_function(item, masking_function) - for item in data - ] + return [LangFuseLogger._apply_masking_function(item, masking_function) for item in data] else: # For other types, try to apply the function directly return masking_function(data) except Exception as e: - verbose_logger.warning( - f"Failed to apply masking function: {e}. Returning original data." - ) + verbose_logger.warning(f"Failed to apply masking function: {e}. Returning original data.") return data @staticmethod @@ -1065,18 +939,12 @@ class LangFuseLogger: Log guardrail information as a span """ if standard_logging_object is None: - verbose_logger.debug( - "Not logging guardrail information as span because standard_logging_object is None" - ) + verbose_logger.debug("Not logging guardrail information as span because standard_logging_object is None") return - guardrail_information = standard_logging_object.get( - "guardrail_information", None - ) + guardrail_information = standard_logging_object.get("guardrail_information", None) if not guardrail_information: - verbose_logger.debug( - "Not logging guardrail information as span because guardrail_information is empty" - ) + verbose_logger.debug("Not logging guardrail information as span because guardrail_information is empty") return if not isinstance(guardrail_information, list): @@ -1101,9 +969,7 @@ class LangFuseLogger: metadata={ "guardrail_name": guardrail_entry.get("guardrail_name", None), "guardrail_mode": guardrail_entry.get("guardrail_mode", None), - "guardrail_masked_entity_count": guardrail_entry.get( - "masked_entity_count", None - ), + "guardrail_masked_entity_count": guardrail_entry.get("masked_entity_count", None), }, start_time=guardrail_entry.get("start_time", None), # type: ignore end_time=guardrail_entry.get("end_time", None), # type: ignore @@ -1142,9 +1008,7 @@ def _add_prompt_to_generation_params( elif "version" in user_prompt and "prompt" in user_prompt: # prompts if isinstance(user_prompt["prompt"], str): - prompt_text_params = getattr( - Prompt_Text, "model_fields", Prompt_Text.__fields__ - ) + prompt_text_params = getattr(Prompt_Text, "model_fields", Prompt_Text.__fields__) _data = { "name": user_prompt["name"], "prompt": user_prompt["prompt"], @@ -1158,9 +1022,7 @@ def _add_prompt_to_generation_params( generation_params["prompt"] = TextPromptClient(prompt=_prompt_obj) elif isinstance(user_prompt["prompt"], list): - prompt_chat_params = getattr( - Prompt_Chat, "model_fields", Prompt_Chat.__fields__ - ) + prompt_chat_params = getattr(Prompt_Chat, "model_fields", Prompt_Chat.__fields__) _data = { "name": user_prompt["name"], "prompt": user_prompt["prompt"], @@ -1175,25 +1037,14 @@ def _add_prompt_to_generation_params( generation_params["prompt"] = ChatPromptClient(prompt=_prompt_obj) else: - verbose_logger.error( - "[Non-blocking] Langfuse Logger: Invalid prompt format" - ) + verbose_logger.error("[Non-blocking] Langfuse Logger: Invalid prompt format") else: - verbose_logger.error( - "[Non-blocking] Langfuse Logger: Invalid prompt format. No prompt logged to Langfuse" - ) - elif ( - prompt_management_metadata is not None - and prompt_management_metadata["prompt_integration"] == "langfuse" - ): + verbose_logger.error("[Non-blocking] Langfuse Logger: Invalid prompt format. No prompt logged to Langfuse") + elif prompt_management_metadata is not None and prompt_management_metadata["prompt_integration"] == "langfuse": try: - generation_params["prompt"] = langfuse_client.get_prompt( - prompt_management_metadata["prompt_id"] - ) + generation_params["prompt"] = langfuse_client.get_prompt(prompt_management_metadata["prompt_id"]) except Exception as e: - verbose_logger.debug( - f"[Non-blocking] Langfuse Logger: Error getting prompt client for logging: {e}" - ) + verbose_logger.debug(f"[Non-blocking] Langfuse Logger: Error getting prompt client for logging: {e}") pass else: @@ -1221,9 +1072,7 @@ def log_provider_specific_information_as_span( if _hidden_params is None: return - vertex_ai_grounding_metadata = _hidden_params.get( - "vertex_ai_grounding_metadata", None - ) + vertex_ai_grounding_metadata = _hidden_params.get("vertex_ai_grounding_metadata", None) if vertex_ai_grounding_metadata is not None: if isinstance(vertex_ai_grounding_metadata, list): diff --git a/litellm/integrations/langfuse/langfuse_handler.py b/litellm/integrations/langfuse/langfuse_handler.py index 797c1609f80..b1d083bd7d4 100644 --- a/litellm/integrations/langfuse/langfuse_handler.py +++ b/litellm/integrations/langfuse/langfuse_handler.py @@ -36,12 +36,7 @@ class LangFuseHandler: """ temp_langfuse_logger: Optional[LangFuseLogger] = globalLangfuseLogger - if ( - LangFuseHandler._dynamic_langfuse_credentials_are_passed( - standard_callback_dynamic_params - ) - is False - ): + if LangFuseHandler._dynamic_langfuse_credentials_are_passed(standard_callback_dynamic_params) is False: return LangFuseHandler._return_global_langfuse_logger( globalLangfuseLogger=globalLangfuseLogger, in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, @@ -61,11 +56,9 @@ class LangFuseHandler: # if not cached, create a new langfuse logger and cache it if temp_langfuse_logger is None: - temp_langfuse_logger = ( - LangFuseHandler._create_langfuse_logger_from_credentials( - credentials=credentials_dict, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) + temp_langfuse_logger = LangFuseHandler._create_langfuse_logger_from_credentials( + credentials=credentials_dict, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, ) return temp_langfuse_logger @@ -94,11 +87,9 @@ class LangFuseHandler: service_name="langfuse", ) if globalLangfuseLogger is None: - globalLangfuseLogger = ( - LangFuseHandler._create_langfuse_logger_from_credentials( - credentials=credentials_dict, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) + globalLangfuseLogger = LangFuseHandler._create_langfuse_logger_from_credentials( + credentials=credentials_dict, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, ) return globalLangfuseLogger @@ -115,8 +106,7 @@ class LangFuseHandler: langfuse_logger = LangFuseLogger( langfuse_public_key=credentials.get("langfuse_public_key"), - langfuse_secret=credentials.get("langfuse_secret") - or credentials.get("langfuse_secret_key"), + langfuse_secret=credentials.get("langfuse_secret") or credentials.get("langfuse_secret_key"), langfuse_host=credentials.get("langfuse_host"), allow_env_credentials=credentials.get("langfuse_host") is None, ) @@ -143,9 +133,7 @@ class LangFuseHandler: return LangfuseLoggingConfig( langfuse_secret=standard_callback_dynamic_params.get("langfuse_secret") or standard_callback_dynamic_params.get("langfuse_secret_key"), - langfuse_public_key=standard_callback_dynamic_params.get( - "langfuse_public_key" - ), + langfuse_public_key=standard_callback_dynamic_params.get("langfuse_public_key"), langfuse_host=standard_callback_dynamic_params.get("langfuse_host"), ) diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 7370bcdf934..fc7c1b211c0 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -48,9 +48,7 @@ class LangfuseOtelLogger(OpenTelemetry): ######################################################### # Set Langfuse specific attributes ######################################################### - LangfuseOtelLogger._set_langfuse_specific_attributes( - span=span, kwargs=kwargs, response_obj=response_obj - ) + LangfuseOtelLogger._set_langfuse_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj) return @staticmethod @@ -141,11 +139,7 @@ class LangfuseOtelLogger(OpenTelemetry): function = tool_call.get("function", {}) arguments_str = function.get("arguments", "{}") try: - arguments_obj = ( - json.loads(arguments_str) - if isinstance(arguments_str, str) - else arguments_str - ) + arguments_obj = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str except json.JSONDecodeError: arguments_obj = {} langfuse_tool_call = { @@ -193,18 +187,12 @@ class LangfuseOtelLogger(OpenTelemetry): output_items_data.append( { "role": getattr(item, "role", "assistant"), - "content": getattr( - getattr(item, "content", [{}])[0], "text", "" - ), + "content": getattr(getattr(item, "content", [{}])[0], "text", ""), } ) elif item_type == "function_call": arguments_str = getattr(item, "arguments", "{}") - arguments_obj = ( - json.loads(arguments_str) - if isinstance(arguments_str, str) - else arguments_str - ) + arguments_obj = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str langfuse_tool_call = { "id": getattr(item, "id", ""), "name": getattr(item, "name", ""), @@ -379,12 +367,8 @@ class LangfuseOtelLogger(OpenTelemetry): """ dynamic_headers = {} - dynamic_langfuse_public_key = standard_callback_dynamic_params.get( - "langfuse_public_key" - ) - dynamic_langfuse_secret_key = standard_callback_dynamic_params.get( - "langfuse_secret_key" - ) + dynamic_langfuse_public_key = standard_callback_dynamic_params.get("langfuse_public_key") + dynamic_langfuse_secret_key = standard_callback_dynamic_params.get("langfuse_secret_key") if dynamic_langfuse_public_key and dynamic_langfuse_secret_key: auth_header = LangfuseOtelLogger._get_langfuse_authorization_header( public_key=dynamic_langfuse_public_key, diff --git a/litellm/integrations/langfuse/langfuse_otel_attributes.py b/litellm/integrations/langfuse/langfuse_otel_attributes.py index fb4a0a6a36c..46bfc21968f 100644 --- a/litellm/integrations/langfuse/langfuse_otel_attributes.py +++ b/litellm/integrations/langfuse/langfuse_otel_attributes.py @@ -74,9 +74,7 @@ def get_output_content_by_type( if isinstance(response_obj, BaseModel): return response_obj.model_dump_json() - if response_obj and ( - isinstance(response_obj, dict) or isinstance(response_obj, list) - ): + if response_obj and (isinstance(response_obj, dict) or isinstance(response_obj, list)): return json.dumps(response_obj) else: return "" diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index cae59295634..0e06f516ecd 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -79,9 +79,7 @@ def langfuse_client_init( allow_env_credentials=allow_env_credentials, ) - if not ( - langfuse_host.startswith("http://") or langfuse_host.startswith("https://") - ): + if not (langfuse_host.startswith("http://") or langfuse_host.startswith("https://")): # add http:// if unset, assume communicating over private network - e.g. render langfuse_host = "http://" + langfuse_host @@ -94,9 +92,7 @@ def langfuse_client_init( "host": langfuse_host, "release": langfuse_release, "debug": langfuse_debug, - "flush_interval": LangFuseLogger._get_langfuse_flush_interval( - flush_interval - ), # flush interval in seconds + "flush_interval": LangFuseLogger._get_langfuse_flush_interval(flush_interval), # flush interval in seconds } if Version(langfuse.version.__version__) >= Version("2.6.0"): @@ -148,9 +144,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ) -> PROMPT_CLIENT: - prompt_client = langfuse_client.get_prompt( - langfuse_prompt_id, label=prompt_label, version=prompt_version - ) + prompt_client = langfuse_client.get_prompt(langfuse_prompt_id, label=prompt_label, version=prompt_version) return prompt_client @@ -168,17 +162,13 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge compiled_prompt = langfuse_prompt_client.compile(**langfuse_prompt_variables) if isinstance(compiled_prompt, str): - compiled_prompt = [ - ChatCompletionSystemMessage(role="system", content=compiled_prompt) - ] + compiled_prompt = [ChatCompletionSystemMessage(role="system", content=compiled_prompt)] else: compiled_prompt = cast(List[AllMessageValues], compiled_prompt) return compiled_prompt - def _get_optional_params_from_langfuse( - self, langfuse_prompt_client: PROMPT_CLIENT - ) -> dict: + def _get_optional_params_from_langfuse(self, langfuse_prompt_client: PROMPT_CLIENT) -> dict: config = langfuse_prompt_client.config optional_params = {} for k, v in config.items(): @@ -276,9 +266,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge template_model = langfuse_prompt_client.config.get("model") - template_optional_params = self._get_optional_params_from_langfuse( - langfuse_prompt_client - ) + template_optional_params = self._get_optional_params_from_langfuse(langfuse_prompt_client) return PromptManagementClient( prompt_id=prompt_id, @@ -307,20 +295,14 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge ) def log_success_event(self, kwargs, response_obj, start_time, end_time): - return run_async_function( - self.async_log_success_event, kwargs, response_obj, start_time, end_time - ) + return run_async_function(self.async_log_success_event, kwargs, response_obj, start_time, end_time) def log_failure_event(self, kwargs, response_obj, start_time, end_time): - return run_async_function( - self.async_log_failure_event, kwargs, response_obj, start_time, end_time - ) + return run_async_function(self.async_log_failure_event, kwargs, response_obj, start_time, end_time) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - standard_callback_dynamic_params = kwargs.get( - "standard_callback_dynamic_params" - ) + standard_callback_dynamic_params = kwargs.get("standard_callback_dynamic_params") langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( globalLangfuseLogger=self, standard_callback_dynamic_params=standard_callback_dynamic_params, @@ -336,16 +318,12 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception( - f"Langfuse Layer Error - Exception occurred while logging success event: {str(e)}" - ) + verbose_logger.exception(f"Langfuse Layer Error - Exception occurred while logging success event: {str(e)}") self.handle_callback_failure(callback_name="langfuse") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: - standard_callback_dynamic_params = kwargs.get( - "standard_callback_dynamic_params" - ) + standard_callback_dynamic_params = kwargs.get("standard_callback_dynamic_params") langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( globalLangfuseLogger=self, standard_callback_dynamic_params=standard_callback_dynamic_params, @@ -357,9 +335,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge ) status_message = str(kwargs.get("exception", "Unknown error")) if standard_logging_object is not None: - status_message = ( - standard_logging_object.get("error_str", None) or status_message - ) + status_message = standard_logging_object.get("error_str", None) or status_message langfuse_logger_to_use.log_event_on_langfuse( start_time=start_time, end_time=end_time, @@ -372,7 +348,5 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception( - f"Langfuse Layer Error - Exception occurred while logging failure event: {str(e)}" - ) + verbose_logger.exception(f"Langfuse Layer Error - Exception occurred while logging failure event: {str(e)}") self.handle_callback_failure(callback_name="langfuse") diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 15f92e8b322..18c4baccd51 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -54,9 +54,7 @@ class LangsmithLogger(CustomBatchLogger): if self.is_mock_mode: create_mock_langsmith_client() - verbose_logger.debug( - "[LANGSMITH MOCK] LangSmith logger initialized in mock mode" - ) + verbose_logger.debug("[LANGSMITH MOCK] LangSmith logger initialized in mock mode") self.default_credentials = self.get_credentials_from_env( langsmith_api_key=langsmith_api_key, @@ -70,31 +68,21 @@ class LangsmithLogger(CustomBatchLogger): and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore else 1.0 ) - self.langsmith_default_run_name = os.getenv( - "LANGSMITH_DEFAULT_RUN_NAME", "LLMRun" - ) - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) - _batch_size = ( - os.getenv("LANGSMITH_BATCH_SIZE", None) or litellm.langsmith_batch_size - ) + self.langsmith_default_run_name = os.getenv("LANGSMITH_DEFAULT_RUN_NAME", "LLMRun") + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + _batch_size = os.getenv("LANGSMITH_BATCH_SIZE", None) or litellm.langsmith_batch_size if _batch_size: self.batch_size = int(_batch_size) self.log_queue: List[LangsmithQueueObject] = [] - self._flush_task: Optional[asyncio.Task[Any]] = ( - self._start_periodic_flush_task() - ) + self._flush_task: Optional[asyncio.Task[Any]] = self._start_periodic_flush_task() def _start_periodic_flush_task(self) -> Optional[asyncio.Task[Any]]: """Start the periodic flush task only when an event loop is already running.""" try: loop = asyncio.get_running_loop() except RuntimeError: - verbose_logger.debug( - "Langsmith logger init: no running event loop, skipping periodic flush task startup" - ) + verbose_logger.debug("Langsmith logger init: no running event loop, skipping periodic flush task startup") return None return loop.create_task(self.periodic_flush()) @@ -121,19 +109,11 @@ class LangsmithLogger(CustomBatchLogger): _credentials_tenant_id = langsmith_tenant_id else: _credentials_api_key = langsmith_api_key or os.getenv("LANGSMITH_API_KEY") - _credentials_project = ( - langsmith_project - or os.getenv("LANGSMITH_PROJECT") - or "litellm-completion" - ) + _credentials_project = langsmith_project or os.getenv("LANGSMITH_PROJECT") or "litellm-completion" _credentials_base_url = ( - langsmith_base_url - or os.getenv("LANGSMITH_BASE_URL") - or "https://api.smith.langchain.com" - ) - _credentials_tenant_id = langsmith_tenant_id or os.getenv( - "LANGSMITH_TENANT_ID" + langsmith_base_url or os.getenv("LANGSMITH_BASE_URL") or "https://api.smith.langchain.com" ) + _credentials_tenant_id = langsmith_tenant_id or os.getenv("LANGSMITH_TENANT_ID") return LangsmithCredentialsObject( LANGSMITH_API_KEY=_credentials_api_key, @@ -142,13 +122,9 @@ class LangsmithLogger(CustomBatchLogger): LANGSMITH_TENANT_ID=_credentials_tenant_id, ) - def _extract_metadata_fields( - self, metadata: dict, credentials: LangsmithCredentialsObject - ): + def _extract_metadata_fields(self, metadata: dict, credentials: LangsmithCredentialsObject): return { - "project_name": metadata.get( - "project_name", credentials["LANGSMITH_PROJECT"] - ), + "project_name": metadata.get("project_name", credentials["LANGSMITH_PROJECT"]), "run_name": metadata.get("run_name", self.langsmith_default_run_name), "run_id": metadata.get("id", metadata.get("run_id", None)), "parent_run_id": metadata.get("parent_run_id", None), @@ -170,14 +146,10 @@ class LangsmithLogger(CustomBatchLogger): extra_metadata = redact_user_api_key_info(metadata=extra_metadata) nested = extra_metadata.get("requester_metadata") if isinstance(nested, dict): - extra_metadata["requester_metadata"] = redact_user_api_key_info( - metadata=nested - ) + extra_metadata["requester_metadata"] = redact_user_api_key_info(metadata=nested) return extra_metadata - def _build_outputs_with_usage( - self, payload: StandardLoggingPayload - ) -> Dict[str, Any]: + def _build_outputs_with_usage(self, payload: StandardLoggingPayload) -> Dict[str, Any]: response = payload["response"] outputs: Dict[str, Any] if isinstance(response, dict): @@ -222,9 +194,7 @@ class LangsmithLogger(CustomBatchLogger): f"Langsmith Logging - project_name: {fields['project_name']}, run_name {fields['run_name']}" ) - payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if payload is None: raise Exception("Error logging request payload. Payload=none.") @@ -295,9 +265,7 @@ class LangsmithLogger(CustomBatchLogger): credentials=credentials, ) ) - verbose_logger.debug( - f"Langsmith, event added to queue. Will flush in {self.flush_interval} seconds..." - ) + verbose_logger.debug(f"Langsmith, event added to queue. Will flush in {self.flush_interval} seconds...") if len(self.log_queue) >= self.batch_size: self._send_batch() @@ -344,9 +312,7 @@ class LangsmithLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception: - verbose_logger.exception( - "Langsmith Layer Error - error logging async success event." - ) + verbose_logger.exception("Langsmith Layer Error - error logging async success event.") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -383,9 +349,7 @@ class LangsmithLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception: - verbose_logger.exception( - "Langsmith Layer Error - error logging async failure event." - ) + verbose_logger.exception("Langsmith Layer Error - error logging async failure event.") async def async_send_batch(self): """ @@ -414,9 +378,7 @@ class LangsmithLogger(CustomBatchLogger): queue_objects=batch_group.queue_objects, ) - def _add_endpoint_to_url( - self, url: str, endpoint: str, api_version: str = "/api/v1" - ) -> str: + def _add_endpoint_to_url(self, url: str, endpoint: str, api_version: str = "/api/v1") -> str: if api_version not in url: url = f"{url.rstrip('/')}{api_version}" @@ -451,13 +413,9 @@ class LangsmithLogger(CustomBatchLogger): elements_to_log = [queue_object["data"] for queue_object in queue_objects] try: - verbose_logger.debug( - "Sending batch of %s runs to Langsmith", len(elements_to_log) - ) + verbose_logger.debug("Sending batch of %s runs to Langsmith", len(elements_to_log)) if self.is_mock_mode: - verbose_logger.debug( - "[LANGSMITH MOCK] Mock mode enabled - API calls will be intercepted" - ) + verbose_logger.debug("[LANGSMITH MOCK] Mock mode enabled - API calls will be intercepted") response = await self.async_httpx_client.post( url=url, json={"post": elements_to_log}, @@ -466,26 +424,16 @@ class LangsmithLogger(CustomBatchLogger): response.raise_for_status() if response.status_code >= 300: - verbose_logger.error( - f"Langsmith Error: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"Langsmith Error: {response.status_code} - {response.text}") else: if self.is_mock_mode: - verbose_logger.debug( - f"[LANGSMITH MOCK] Batch of {len(elements_to_log)} runs successfully mocked" - ) + verbose_logger.debug(f"[LANGSMITH MOCK] Batch of {len(elements_to_log)} runs successfully mocked") else: - verbose_logger.debug( - f"Batch of {len(self.log_queue)} runs successfully created" - ) + verbose_logger.debug(f"Batch of {len(self.log_queue)} runs successfully created") except httpx.HTTPStatusError as e: - verbose_logger.exception( - f"Langsmith HTTP Error: {e.response.status_code} - {e.response.text}" - ) + verbose_logger.exception(f"Langsmith HTTP Error: {e.response.status_code} - {e.response.text}") except Exception: - verbose_logger.exception( - f"Langsmith Layer Error - {traceback.format_exc()}" - ) + verbose_logger.exception(f"Langsmith Layer Error - {traceback.format_exc()}") def _group_batches_by_credentials(self) -> Dict[CredentialsKey, BatchGroup]: """Groups queue objects by credentials using a proper key structure""" @@ -494,10 +442,7 @@ class LangsmithLogger(CustomBatchLogger): for queue_object in self.log_queue: credentials = queue_object["credentials"] # if credential missing, skip - log warning - if ( - credentials["LANGSMITH_API_KEY"] is None - or credentials["LANGSMITH_PROJECT"] is None - ): + if credentials["LANGSMITH_API_KEY"] is None or credentials["LANGSMITH_PROJECT"] is None: verbose_logger.warning( "Langsmith Logging - credentials missing - api_key: %s, project: %s", credentials["LANGSMITH_API_KEY"], @@ -512,30 +457,24 @@ class LangsmithLogger(CustomBatchLogger): ) if key not in log_queue_by_credentials: - log_queue_by_credentials[key] = BatchGroup( - credentials=credentials, queue_objects=[] - ) + log_queue_by_credentials[key] = BatchGroup(credentials=credentials, queue_objects=[]) log_queue_by_credentials[key].queue_objects.append(queue_object) return log_queue_by_credentials def _get_sampling_rate_to_use_for_request(self, kwargs: Dict[str, Any]) -> float: - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params", None) + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get( + "standard_callback_dynamic_params", None ) sampling_rate: float = self.sampling_rate if standard_callback_dynamic_params is not None: - _sampling_rate = standard_callback_dynamic_params.get( - "langsmith_sampling_rate" - ) + _sampling_rate = standard_callback_dynamic_params.get("langsmith_sampling_rate") if _sampling_rate is not None: sampling_rate = float(_sampling_rate) return sampling_rate - def _get_credentials_to_use_for_request( - self, kwargs: Dict[str, Any] - ) -> LangsmithCredentialsObject: + def _get_credentials_to_use_for_request(self, kwargs: Dict[str, Any]) -> LangsmithCredentialsObject: """ Handles key/team based logging @@ -543,27 +482,16 @@ class LangsmithLogger(CustomBatchLogger): Otherwise, use the default credentials. """ - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params", None) + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get( + "standard_callback_dynamic_params", None ) if standard_callback_dynamic_params is not None: credentials = self.get_credentials_from_env( - langsmith_api_key=standard_callback_dynamic_params.get( - "langsmith_api_key", None - ), - langsmith_project=standard_callback_dynamic_params.get( - "langsmith_project", None - ), - langsmith_base_url=standard_callback_dynamic_params.get( - "langsmith_base_url", None - ), - langsmith_tenant_id=standard_callback_dynamic_params.get( - "langsmith_tenant_id", None - ), - allow_env_credentials=standard_callback_dynamic_params.get( - "langsmith_base_url", None - ) - is None, + langsmith_api_key=standard_callback_dynamic_params.get("langsmith_api_key", None), + langsmith_project=standard_callback_dynamic_params.get("langsmith_project", None), + langsmith_base_url=standard_callback_dynamic_params.get("langsmith_base_url", None), + langsmith_tenant_id=standard_callback_dynamic_params.get("langsmith_tenant_id", None), + allow_env_credentials=standard_callback_dynamic_params.get("langsmith_base_url", None) is None, ) else: credentials = self.default_credentials diff --git a/litellm/integrations/langsmith_mock_client.py b/litellm/integrations/langsmith_mock_client.py index 0226bdecc27..1e20e1b5ce5 100644 --- a/litellm/integrations/langsmith_mock_client.py +++ b/litellm/integrations/langsmith_mock_client.py @@ -29,6 +29,4 @@ _config = MockClientConfig( patch_sync_client=False, ) -create_mock_langsmith_client, should_use_langsmith_mock = create_mock_client_factory( - _config -) +create_mock_langsmith_client, should_use_langsmith_mock = create_mock_client_factory(_config) diff --git a/litellm/integrations/langtrace.py b/litellm/integrations/langtrace.py index ac1069f440e..a9e580a83b2 100644 --- a/litellm/integrations/langtrace.py +++ b/litellm/integrations/langtrace.py @@ -86,12 +86,8 @@ class LangtraceAttributes: usage = response_obj.get("usage") if usage: usage_attributes = { - SpanAttributes.LLM_USAGE_PROMPT_TOKENS.value: usage.get( - "prompt_tokens" - ), - SpanAttributes.LLM_USAGE_COMPLETION_TOKENS.value: usage.get( - "completion_tokens" - ), + SpanAttributes.LLM_USAGE_PROMPT_TOKENS.value: usage.get("prompt_tokens"), + SpanAttributes.LLM_USAGE_COMPLETION_TOKENS.value: usage.get("completion_tokens"), SpanAttributes.LLM_USAGE_TOTAL_TOKENS.value: usage.get("total_tokens"), } self.set_span_attributes(span, usage_attributes) diff --git a/litellm/integrations/levo/levo.py b/litellm/integrations/levo/levo.py index 4b08ce50f74..a865944485c 100644 --- a/litellm/integrations/levo/levo.py +++ b/litellm/integrations/levo/levo.py @@ -56,17 +56,11 @@ class LevoLogger(OpenTelemetry): # Validate required env vars if not api_key: - raise ValueError( - "LEVOAI_API_KEY environment variable is required for Levo integration." - ) + raise ValueError("LEVOAI_API_KEY environment variable is required for Levo integration.") if not org_id: - raise ValueError( - "LEVOAI_ORG_ID environment variable is required for Levo integration." - ) + raise ValueError("LEVOAI_ORG_ID environment variable is required for Levo integration.") if not workspace_id: - raise ValueError( - "LEVOAI_WORKSPACE_ID environment variable is required for Levo integration." - ) + raise ValueError("LEVOAI_WORKSPACE_ID environment variable is required for Levo integration.") if not collector_url: raise ValueError( "LEVOAI_COLLECTOR_URL environment variable is required for Levo integration. " diff --git a/litellm/integrations/literal_ai.py b/litellm/integrations/literal_ai.py index 042779ba844..c8c931eb667 100644 --- a/litellm/integrations/literal_ai.py +++ b/litellm/integrations/literal_ai.py @@ -33,9 +33,7 @@ class LiteralAILogger(CustomBatchLogger): } if env: self.headers["x-env"] = env - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_http_handler = HTTPHandler() batch_size = os.getenv("LITERAL_BATCH_SIZE", None) self.flush_lock = asyncio.Lock() @@ -62,9 +60,7 @@ class LiteralAILogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: self._send_batch() except Exception: - verbose_logger.exception( - "Literal AI Layer Error - error logging success event." - ) + verbose_logger.exception("Literal AI Layer Error - error logging success event.") def log_failure_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.info("Literal AI Failure Event Logging!") @@ -79,9 +75,7 @@ class LiteralAILogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: self._send_batch() except Exception: - verbose_logger.exception( - "Literal AI Layer Error - error logging failure event." - ) + verbose_logger.exception("Literal AI Layer Error - error logging failure event.") def _send_batch(self): if not self.log_queue: @@ -101,13 +95,9 @@ class LiteralAILogger(CustomBatchLogger): ) if response.status_code >= 300: - verbose_logger.error( - f"Literal AI Error: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"Literal AI Error: {response.status_code} - {response.text}") else: - verbose_logger.debug( - f"Batch of {len(self.log_queue)} runs successfully created" - ) + verbose_logger.debug(f"Batch of {len(self.log_queue)} runs successfully created") except Exception: verbose_logger.exception("Literal AI Layer Error") @@ -128,9 +118,7 @@ class LiteralAILogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception: - verbose_logger.exception( - "Literal AI Layer Error - error logging async success event." - ) + verbose_logger.exception("Literal AI Layer Error - error logging async success event.") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.info("Literal AI Failure Event Logging!") @@ -145,9 +133,7 @@ class LiteralAILogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception: - verbose_logger.exception( - "Literal AI Layer Error - error logging async failure event." - ) + verbose_logger.exception("Literal AI Layer Error - error logging async failure event.") async def async_send_batch(self): if not self.log_queue: @@ -167,24 +153,16 @@ class LiteralAILogger(CustomBatchLogger): headers=self.headers, ) if response.status_code >= 300: - verbose_logger.error( - f"Literal AI Error: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"Literal AI Error: {response.status_code} - {response.text}") else: - verbose_logger.debug( - f"Batch of {len(self.log_queue)} runs successfully created" - ) + verbose_logger.debug(f"Batch of {len(self.log_queue)} runs successfully created") except httpx.HTTPStatusError as e: - verbose_logger.exception( - f"Literal AI HTTP Error: {e.response.status_code} - {e.response.text}" - ) + verbose_logger.exception(f"Literal AI HTTP Error: {e.response.status_code} - {e.response.text}") except Exception: verbose_logger.exception("Literal AI Layer Error") def _prepare_log_data(self, kwargs, response_obj, start_time, end_time) -> dict: - logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") diff --git a/litellm/integrations/logfire_logger.py b/litellm/integrations/logfire_logger.py index 2345dc869c6..c92dfff2934 100644 --- a/litellm/integrations/logfire_logger.py +++ b/litellm/integrations/logfire_logger.py @@ -39,25 +39,17 @@ class LogfireLogger: raise e def _get_span_config(self, payload) -> SpanConfig: - if ( - payload["call_type"] == "completion" - or payload["call_type"] == "acompletion" - ): + if payload["call_type"] == "completion" or payload["call_type"] == "acompletion": return SpanConfig( message_template="Chat Completion with {request_data[model]!r}", span_data={"request_data": payload}, ) - elif ( - payload["call_type"] == "embedding" or payload["call_type"] == "aembedding" - ): + elif payload["call_type"] == "embedding" or payload["call_type"] == "aembedding": return SpanConfig( message_template="Embedding Creation with {request_data[model]!r}", span_data={"request_data": payload}, ) - elif ( - payload["call_type"] == "image_generation" - or payload["call_type"] == "aimage_generation" - ): + elif payload["call_type"] == "image_generation" or payload["call_type"] == "aimage_generation": return SpanConfig( message_template="Image Generation with {request_data[model]!r}", span_data={"request_data": payload}, @@ -98,16 +90,12 @@ class LogfireLogger: try: import logfire - verbose_logger.debug( - f"logfire Logging - Enters logging function for model {kwargs}" - ) + verbose_logger.debug(f"logfire Logging - Enters logging function for model {kwargs}") if not response_obj: response_obj = {} litellm_params = kwargs.get("litellm_params", {}) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None + metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None messages = kwargs.get("messages") optional_params = kwargs.get("optional_params", {}) call_type = kwargs.get("call_type", "completion") @@ -169,11 +157,7 @@ class LogfireLogger: ) print_verbose(f"\ndd Logger - Logging payload = {payload}") - print_verbose( - f"Logfire Layer Logging - final response object: {response_obj}" - ) + print_verbose(f"Logfire Layer Logging - final response object: {response_obj}") except Exception as e: - verbose_logger.debug( - f"Logfire Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.debug(f"Logfire Layer Error - {str(e)}\n{traceback.format_exc()}") pass diff --git a/litellm/integrations/lunary.py b/litellm/integrations/lunary.py index 7b1cbc32d43..aaf5751cb79 100644 --- a/litellm/integrations/lunary.py +++ b/litellm/integrations/lunary.py @@ -130,11 +130,7 @@ class LunaryLogger: pass if response_obj: - usage = ( - parse_usage(response_obj["usage"]) - if "usage" in response_obj - else None - ) + usage = parse_usage(response_obj["usage"]) if "usage" in response_obj else None output = response_obj["choices"] if "choices" in response_obj else None diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 49030ffc2e1..26b2f32f32f 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -66,9 +66,7 @@ def _parse_metrics_marker( continue except Exception: pass - verbose_proxy_logger.warning( - "Mavvrik FOCUS: could not parse metricsMarker %r — skipping catch-up", marker - ) + verbose_proxy_logger.warning("Mavvrik FOCUS: could not parse metricsMarker %r — skipping catch-up", marker) return None @@ -134,9 +132,7 @@ class MavvrikFocusLogger(FocusLogger): ) payload = b"" if data.is_empty(): - verbose_proxy_logger.debug( - "Mavvrik FOCUS export: no usage data for window %s", window - ) + verbose_proxy_logger.debug("Mavvrik FOCUS export: no usage data for window %s", window) else: normalized = engine._transformer.transform(data) if not normalized.is_empty(): @@ -179,9 +175,7 @@ class MavvrikFocusLogger(FocusLogger): marker = await destination.get_metrics_marker() now = datetime.now(timezone.utc) - yesterday = now.replace(hour=0, minute=0, second=0, microsecond=0) - timedelta( - days=1 - ) + yesterday = now.replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=1) last_ingested = _parse_metrics_marker(marker) @@ -189,15 +183,10 @@ class MavvrikFocusLogger(FocusLogger): earliest_catchup = yesterday - timedelta(days=self._MAX_CATCHUP_DAYS - 1) if is_empty_marker or (last_ingested is not None and last_ingested < yesterday): catch_up_date = ( - earliest_catchup - if last_ingested is None - else max(last_ingested + timedelta(days=1), earliest_catchup) + earliest_catchup if last_ingested is None else max(last_ingested + timedelta(days=1), earliest_catchup) ) - if ( - last_ingested is not None - and last_ingested + timedelta(days=1) < earliest_catchup - ): + if last_ingested is not None and last_ingested + timedelta(days=1) < earliest_catchup: verbose_proxy_logger.warning( "Mavvrik FOCUS export: metricsMarker is more than %d days behind " "(%s). Catching up from %s only; earlier data will not be re-exported.", @@ -244,20 +233,14 @@ class MavvrikFocusLogger(FocusLogger): pod_lock_manager = getattr(writer, "pod_lock_manager", None) if pod_lock_manager and pod_lock_manager.redis_cache: - acquired = await pod_lock_manager.acquire_lock( - cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME - ) + acquired = await pod_lock_manager.acquire_lock(cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME) if not acquired: - verbose_proxy_logger.debug( - "Mavvrik FOCUS export: unable to acquire pod lock" - ) + verbose_proxy_logger.debug("Mavvrik FOCUS export: unable to acquire pod lock") return try: await self._run_scheduled_export() finally: - await pod_lock_manager.release_lock( - cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME - ) + await pod_lock_manager.release_lock(cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME) else: await self._run_scheduled_export() @@ -268,9 +251,7 @@ class MavvrikFocusLogger(FocusLogger): """Register the Mavvrik FOCUS export job on the provided scheduler.""" loggers: List[MavvrikFocusLogger] = [ cb - for cb in litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=MavvrikFocusLogger - ) + for cb in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=MavvrikFocusLogger) if type(cb) is MavvrikFocusLogger ] if not loggers and "mavvrik" in litellm.callbacks: @@ -289,9 +270,7 @@ class MavvrikFocusLogger(FocusLogger): if isinstance(instance, MavvrikFocusLogger): loggers = [instance] if not loggers: - verbose_proxy_logger.debug( - "No MavvrikFocusLogger registered; skipping scheduler" - ) + verbose_proxy_logger.debug("No MavvrikFocusLogger registered; skipping scheduler") return logger = loggers[0] @@ -302,6 +281,4 @@ class MavvrikFocusLogger(FocusLogger): replace_existing=True, **trigger_kwargs, ) - verbose_proxy_logger.info( - "mavvrik_focus: background export job scheduled (%s)", trigger_kwargs - ) + verbose_proxy_logger.info("mavvrik_focus: background export job scheduled (%s)", trigger_kwargs) diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index 6378e55f7e1..1952c95eac9 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -60,10 +60,7 @@ class MlflowLogger(CustomLogger): inputs = self._construct_input(kwargs) input_messages = inputs.get("messages", []) - output_messages = [ - c.message.model_dump(exclude_none=True) - for c in getattr(response_obj, "choices", []) - ] + output_messages = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])] if messages := [*input_messages, *output_messages]: set_span_chat_messages(span, messages) if tools := inputs.get("tools"): @@ -130,9 +127,7 @@ class MlflowLogger(CustomLogger): # If this is the final chunk, end the span. The final chunk # has the assembled streaming response (key differs between sync/async paths). - final_response = kwargs.get("complete_streaming_response") or kwargs.get( - "async_complete_streaming_response" - ) + final_response = kwargs.get("complete_streaming_response") or kwargs.get("async_complete_streaming_response") if final_response: end_time_ns = int(end_time.timestamp() * 1e9) @@ -156,9 +151,7 @@ class MlflowLogger(CustomLogger): span.add_event( SpanEvent( name="streaming_chunk", - attributes={ - "delta": json.dumps(choice.delta.model_dump, default=str) - }, + attributes={"delta": json.dumps(choice.delta.model_dump, default=str)}, ) ) except Exception: @@ -192,9 +185,7 @@ class MlflowLogger(CustomLogger): "call_type": kwargs.get("call_type"), "model": kwargs.get("model"), } - standard_obj: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_obj: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_obj: attributes.update( { @@ -267,9 +258,7 @@ class MlflowLogger(CustomLogger): span_type=span_type, inputs=inputs, attributes=attributes, - tags=self._transform_tag_list_to_dict( - attributes.get("request_tags", []) - ), + tags=self._transform_tag_list_to_dict(attributes.get("request_tags", [])), start_time_ns=start_time_ns, ) diff --git a/litellm/integrations/mock_client_factory.py b/litellm/integrations/mock_client_factory.py index 3b013f96a79..76c0ac03b7b 100644 --- a/litellm/integrations/mock_client_factory.py +++ b/litellm/integrations/mock_client_factory.py @@ -25,14 +25,10 @@ class MockClientConfig: default_latency_ms: int = 100 # Default mock latency in milliseconds default_status_code: int = 200 # Default HTTP status code default_json_data: Optional[Dict] = None # Default JSON response data - url_matchers: Optional[List[str]] = ( - None # List of strings to match in URLs (e.g., ["storage.googleapis.com"]) - ) + url_matchers: Optional[List[str]] = None # List of strings to match in URLs (e.g., ["storage.googleapis.com"]) patch_async_handler: bool = True # Whether to patch AsyncHTTPHandler.post patch_sync_client: bool = False # Whether to patch httpx.Client.post - patch_http_handler: bool = ( - False # Whether to patch HTTPHandler.post (for sync calls that use HTTPHandler) - ) + patch_http_handler: bool = False # Whether to patch HTTPHandler.post (for sync calls that use HTTPHandler) def __post_init__(self): """Ensure url_matchers is a list.""" @@ -124,9 +120,7 @@ def create_mock_client_factory(config: MockClientConfig): import os latency_env = f"{config.name.upper()}_MOCK_LATENCY_MS" - _MOCK_LATENCY_SECONDS = ( - float(os.getenv(latency_env, str(config.default_latency_ms))) / 1000.0 - ) + _MOCK_LATENCY_SECONDS = float(os.getenv(latency_env, str(config.default_latency_ms))) / 1000.0 # Create URL matcher function def _is_mock_url(url) -> bool: @@ -241,9 +235,7 @@ def create_mock_client_factory(config: MockClientConfig): if _mocks_initialized: return - verbose_logger.debug( - f"[{config.name} MOCK] Initializing {config.name} mock client..." - ) + verbose_logger.debug(f"[{config.name} MOCK] Initializing {config.name} mock client...") if config.patch_async_handler and _original_async_handler_post is None: from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -264,12 +256,8 @@ def create_mock_client_factory(config: MockClientConfig): HTTPHandler.post = _mock_http_handler_post # type: ignore verbose_logger.debug(f"[{config.name} MOCK] Patched HTTPHandler.post") - verbose_logger.debug( - f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms" - ) - verbose_logger.debug( - f"[{config.name} MOCK] {config.name} mock client initialization complete" - ) + verbose_logger.debug(f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS * 1000:.0f}ms") + verbose_logger.debug(f"[{config.name} MOCK] {config.name} mock client initialization complete") _mocks_initialized = True @@ -284,9 +272,7 @@ def create_mock_client_factory(config: MockClientConfig): result = bool(result) if result is not None else False if result: - verbose_logger.info( - f"{config.name} Mock Mode: ENABLED - API calls will be mocked" - ) + verbose_logger.info(f"{config.name} Mock Mode: ENABLED - API calls will be mocked") return result diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index 753b8520337..3c2bed60ef4 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -120,10 +120,7 @@ class NewRelicLogger(CustomLogger): f"content recording: {self.record_content}" ) except Exception as e: - verbose_logger.error( - f"Failed to initialize New Relic agent: {e}. " - "Integration will be disabled." - ) + verbose_logger.error(f"Failed to initialize New Relic agent: {e}. Integration will be disabled.") self.enabled = False def _get_newrelic_params(self) -> Dict: @@ -138,9 +135,7 @@ class NewRelicLogger(CustomLogger): dict_newrelic_params = litellm.newrelic_params.model_dump() elif isinstance(litellm.newrelic_params, Dict): # only allow params that are of NewRelicInitParams - dict_newrelic_params = NewRelicInitParams( - **litellm.newrelic_params - ).model_dump() + dict_newrelic_params = NewRelicInitParams(**litellm.newrelic_params).model_dump() return dict_newrelic_params @property @@ -221,13 +216,9 @@ class NewRelicLogger(CustomLogger): if app and app.enabled: app.record_custom_metric(metric_name, 1) - verbose_logger.info( - f"Emitted New Relic supportability metric: {metric_name}" - ) + verbose_logger.info(f"Emitted New Relic supportability metric: {metric_name}") else: - verbose_logger.info( - "New Relic application is not enabled; skipping metric recording." - ) + verbose_logger.info("New Relic application is not enabled; skipping metric recording.") except Exception as e: verbose_logger.warning(f"Failed to emit supportability metric: {e}") @@ -241,18 +232,14 @@ class NewRelicLogger(CustomLogger): """ # Quick check without lock to avoid unnecessary locking current_time = time.time() - time_since_last_emission = ( - current_time - NewRelicLogger._last_metric_emission_time - ) + time_since_last_emission = current_time - NewRelicLogger._last_metric_emission_time if time_since_last_emission >= 97200: # 27 hours = 97200 seconds # Acquire lock to ensure only one thread emits with NewRelicLogger._metric_lock: # Double-check inside lock in case another thread just emitted current_time = time.time() - time_since_last_emission = ( - current_time - NewRelicLogger._last_metric_emission_time - ) + time_since_last_emission = current_time - NewRelicLogger._last_metric_emission_time if time_since_last_emission >= 97200: self._emit_supportability_metric() @@ -292,9 +279,7 @@ class NewRelicLogger(CustomLogger): metadata = litellm_params.get("metadata") or {} headers = metadata.get("headers") or {} # Normalize header key lookup to be case-insensitive per W3C spec - traceparent = next( - (v for k, v in headers.items() if k.lower() == "traceparent"), None - ) + traceparent = next((v for k, v in headers.items() if k.lower() == "traceparent"), None) if traceparent: # Extract trace_id from traceparent header if available @@ -309,9 +294,7 @@ class NewRelicLogger(CustomLogger): trace_id = slo_trace_id except Exception as e: - verbose_logger.warning( - f"Unable to parse New Relic trace context from upstream sources: {e}" - ) + verbose_logger.warning(f"Unable to parse New Relic trace context from upstream sources: {e}") if not trace_id: trace_id = uuid.uuid4().hex @@ -439,9 +422,7 @@ class NewRelicLogger(CustomLogger): if standard_logging_object: response_time = standard_logging_object.get("response_time") if response_time is not None: - return ( - float(response_time) * 1000.0 - ) # SLO stores seconds; convert to ms + return float(response_time) * 1000.0 # SLO stores seconds; convert to ms duration_ms = kwargs.get("llm_api_duration_ms") if duration_ms is not None: @@ -552,15 +533,11 @@ class NewRelicLogger(CustomLogger): # callback an unredacted async_complete_streaming_response, so without # this gate generated content would still reach NR even when the user # has globally disabled message logging. - record_content = self.record_content and not should_redact_message_logging( - kwargs - ) + record_content = self.record_content and not should_redact_message_logging(kwargs) # Extract request messages, preferring StandardLoggingPayload. # SLO messages can be a string (serialized/redacted), so only use it when it's a list. - slo_messages = ( - standard_logging_object.get("messages") if standard_logging_object else None - ) + slo_messages = standard_logging_object.get("messages") if standard_logging_object else None if isinstance(slo_messages, list): request_messages = slo_messages else: @@ -658,9 +635,7 @@ class NewRelicLogger(CustomLogger): if app and app.enabled: app.record_custom_event("LlmChatCompletionSummary", event_data) else: - verbose_logger.warning( - "New Relic application is not enabled; skipping summary event recording." - ) + verbose_logger.warning("New Relic application is not enabled; skipping summary event recording.") except Exception as e: verbose_logger.warning(f"Failed to record New Relic summary event: {e}") @@ -685,9 +660,7 @@ class NewRelicLogger(CustomLogger): app = _newrelic_agent.application() if not (app and app.enabled): - verbose_logger.warning( - "New Relic application is not enabled; skipping message event recording." - ) + verbose_logger.warning("New Relic application is not enabled; skipping message event recording.") return for message in messages: @@ -763,9 +736,7 @@ class NewRelicLogger(CustomLogger): self._check_and_emit_periodic_metric() # Use StandardLoggingPayload where available for normalized, pre-computed values - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") # Get trace context trace_id = self._get_trace_context(kwargs, standard_logging_object) @@ -776,22 +747,16 @@ class NewRelicLogger(CustomLogger): # Extract data from response llm_response_id = self._extract_completion_id(kwargs, response_obj) vendor = self._get_vendor(kwargs, standard_logging_object) - request_model, response_model = self._get_model_names( - kwargs, response_obj, standard_logging_object - ) + request_model, response_model = self._get_model_names(kwargs, response_obj, standard_logging_object) usage = self._extract_usage(response_obj, standard_logging_object) finish_reason = self._get_finish_reason(response_obj) # Extract additional summary event fields - duration = self._get_duration( - kwargs, start_time, end_time, standard_logging_object - ) + duration = self._get_duration(kwargs, start_time, end_time, standard_logging_object) request_params = self._get_request_params(kwargs, standard_logging_object) # Extract all messages - messages = self._extract_all_messages( - kwargs, response_obj, response_model, vendor, standard_logging_object - ) + messages = self._extract_all_messages(kwargs, response_obj, response_model, vendor, standard_logging_object) # Record summary event self._record_summary_event( diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index b234ab11ddb..e9cc68a7841 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -29,9 +29,7 @@ class OpenMeterLogger(CustomLogger): def __init__(self) -> None: super().__init__() self.validate_environment() - self.async_http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_http_handler = HTTPHandler() def validate_environment(self): @@ -56,8 +54,7 @@ class OpenMeterLogger(CustomLogger): model = kwargs.get("model") usage = {} if ( - isinstance(response_obj, litellm.ModelResponse) - or isinstance(response_obj, litellm.EmbeddingResponse) + isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse) ) and hasattr(response_obj, "usage"): usage = { "prompt_tokens": response_obj["usage"].get("prompt_tokens", 0), @@ -70,9 +67,7 @@ class OpenMeterLogger(CustomLogger): # resolved solely from the key-bound user_api_key_user_id. Proxies # serving multi-tenant traffic enable this to prevent clients from # forging attribution by setting `user` in the request body. - trust_request_user = ( - os.getenv("OPENMETER_TRUST_REQUEST_USER", "true").lower() != "false" - ) + trust_request_user = os.getenv("OPENMETER_TRUST_REQUEST_USER", "true").lower() != "false" user_param = kwargs.get("user", None) if trust_request_user else None # If no user provided directly, try to get it from token user_id diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index c652141a2b6..de543fa042b 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -153,22 +153,16 @@ def _resolve_metric_attribute_filter( include = attributes.include_list or None exclude = attributes.exclude_list or None if include and exclude: - raise ValueError( - "otel.attributes: include_list and exclude_list are mutually exclusive" - ) + raise ValueError("otel.attributes: include_list and exclude_list are mutually exclusive") requested = include or exclude or [] if TOKEN_TYPE_ATTRIBUTE in requested: raise ValueError( - f"otel.attributes: {TOKEN_TYPE_ATTRIBUTE} is a structural token-usage " - "discriminator and cannot be filtered" + f"otel.attributes: {TOKEN_TYPE_ATTRIBUTE} is a structural token-usage discriminator and cannot be filtered" ) - unknown = sorted( - name for name in requested if name not in VALID_METRIC_ATTRIBUTE_NAMES - ) + unknown = sorted(name for name in requested if name not in VALID_METRIC_ATTRIBUTE_NAMES) if unknown: raise ValueError( - f"otel.attributes: unknown attribute name(s) {unknown}. " - f"Valid names: {sorted(VALID_METRIC_ATTRIBUTE_NAMES)}" + f"otel.attributes: unknown attribute name(s) {unknown}. Valid names: {sorted(VALID_METRIC_ATTRIBUTE_NAMES)}" ) return ( frozenset(include) if include else None, @@ -212,8 +206,7 @@ def _freeze_for_dedupe(value: object, _depth: int = 0) -> HashableScope: return frozenset(_freeze_for_dedupe(item, _depth + 1) for item in value) if isinstance(value, dict): return frozenset( - (_freeze_for_dedupe(key, _depth + 1), _freeze_for_dedupe(item, _depth + 1)) - for key, item in value.items() + (_freeze_for_dedupe(key, _depth + 1), _freeze_for_dedupe(item, _depth + 1)) for key, item in value.items() ) if isinstance(value, (str, int, float, bytes)) or value is None: return value @@ -250,35 +243,23 @@ class OpenTelemetryConfig: # automatically infer "otlp_http" to send traces to the endpoint. # This fixes an issue where UI-configured OTEL settings would default # to console output instead of sending traces to the configured endpoint. - if ( - self.endpoint - and isinstance(self.exporter, str) - and self.exporter == "console" - ): + if self.endpoint and isinstance(self.exporter, str) and self.exporter == "console": self.exporter = "otlp_http" if not self.service_name: self.service_name = os.getenv("OTEL_SERVICE_NAME", "litellm") if not self.deployment_environment: - self.deployment_environment = os.getenv( - "OTEL_ENVIRONMENT_NAME", "production" - ) + self.deployment_environment = os.getenv("OTEL_ENVIRONMENT_NAME", "production") if not self.model_id: self.model_id = os.getenv("OTEL_MODEL_ID", self.service_name) if self.ignore_context_propagation is None: - self.ignore_context_propagation = str_to_bool( - os.getenv("OTEL_IGNORE_CONTEXT_PROPAGATION") - ) + self.ignore_context_propagation = str_to_bool(os.getenv("OTEL_IGNORE_CONTEXT_PROPAGATION")) # Resolve the env opt-in once here so self.semconv_stability_opt_in is the # single source of truth: the union of programmatic and env categories. - self.semconv_stability_opt_in |= parse_semconv_opt_in( - os.getenv(OTEL_SEMCONV_STABILITY_OPT_IN_ENV) - ) + self.semconv_stability_opt_in |= parse_semconv_opt_in(os.getenv(OTEL_SEMCONV_STABILITY_OPT_IN_ENV)) self.baggage_team_metadata_keys = _normalize_team_metadata_keys( self.baggage_team_metadata_keys - ) or _normalize_team_metadata_keys( - os.getenv("LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS") - ) + ) or _normalize_team_metadata_keys(os.getenv("LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS")) @classmethod def from_env(cls): @@ -293,21 +274,13 @@ class OpenTelemetryConfig: InMemorySpanExporter, ) - exporter = os.getenv( - "OTEL_EXPORTER_OTLP_PROTOCOL", os.getenv("OTEL_EXPORTER", "console") - ) + exporter = os.getenv("OTEL_EXPORTER_OTLP_PROTOCOL", os.getenv("OTEL_EXPORTER", "console")) endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", os.getenv("OTEL_ENDPOINT")) headers = os.getenv( "OTEL_EXPORTER_OTLP_HEADERS", os.getenv("OTEL_HEADERS") ) # example: OTEL_HEADERS=x-honeycomb-team=B85YgLm96***" - enable_metrics: bool = ( - os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_METRICS", "false").lower() - == "true" - ) - enable_events: bool = ( - os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS", "false").lower() - == "true" - ) + enable_metrics: bool = os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_METRICS", "false").lower() == "true" + enable_events: bool = os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS", "false").lower() == "true" service_name = os.getenv("OTEL_SERVICE_NAME", "litellm") deployment_environment = os.getenv("OTEL_ENVIRONMENT_NAME", "production") model_id = os.getenv("OTEL_MODEL_ID", service_name) @@ -342,13 +315,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if config is None: config = OpenTelemetryConfig.from_env() if team_metadata_keys_override is not None: - config.baggage_team_metadata_keys = _normalize_team_metadata_keys( - team_metadata_keys_override - ) + config.baggage_team_metadata_keys = _normalize_team_metadata_keys(team_metadata_keys_override) if metric_attributes_override is not None: - config.attributes = _build_metric_attribute_filter( - metric_attributes_override - ) + config.attributes = _build_metric_attribute_filter(metric_attributes_override) self.config = config self.callback_name = callback_name @@ -415,9 +384,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): try: from litellm.proxy import proxy_server except ImportError: - verbose_logger.warning( - "Proxy Server is not installed. Skipping OpenTelemetry initialization." - ) + verbose_logger.warning("Proxy Server is not installed. Skipping OpenTelemetry initialization.") return # Add self as a service callback @@ -514,9 +481,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): def _skip_set_global(self) -> bool: # langfuse_otel relies on the Langfuse SDK's providers; don't overwrite them. - return self.config.skip_set_global or ( - hasattr(self, "callback_name") and self.callback_name == "langfuse_otel" - ) + return self.config.skip_set_global or (hasattr(self, "callback_name") and self.callback_name == "langfuse_otel") def _compute_capture_mode_from_init_state(self) -> Optional[str]: """Sample explicit settings at init. Returns the resolved mode or @@ -556,11 +521,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return CAPTURE_MODE_NO_CONTENT if self._capture_mode_cached is not None: return self._capture_mode_cached - return ( - CAPTURE_MODE_SPAN_AND_EVENT - if self.message_logging - else CAPTURE_MODE_NO_CONTENT - ) + return CAPTURE_MODE_SPAN_AND_EVENT if self.message_logging else CAPTURE_MODE_NO_CONTENT def _capture_in_span(self) -> bool: return self._resolve_capture_mode() in ( @@ -676,9 +637,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): from opentelemetry.sdk._logs.export import BatchLogRecordProcessor def create_logger_provider(): - provider = OTLoggerProvider( - resource=self._get_litellm_resource(self.config) - ) + provider = OTLoggerProvider(resource=self._get_litellm_resource(self.config)) log_exporter = self._get_log_exporter() provider.add_log_record_processor( BatchLogRecordProcessor(log_exporter) # type: ignore[arg-type] @@ -875,9 +834,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # _record_exception_on_span only stamps when error_code is set; # bare TypeError etc. has none, and the span is about to be ended. - error_code = ( - error_information.get("error_code") if error_information else None - ) + error_code = error_information.get("error_code") if error_information else None if not error_code: self.set_response_status_code_attribute(parent_otel_span, 500) @@ -951,9 +908,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): "metadata": metadata, }, } - context = ( - _trace.set_span_in_context(parent_span) if parent_span is not None else None - ) + context = _trace.set_span_in_context(parent_span) if parent_span is not None else None self._create_guardrail_span(kwargs=kwargs, context=context) async def async_post_call_success_hook( @@ -966,9 +921,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): litellm_logging_obj = data.get("litellm_logging_obj") - if litellm_logging_obj is not None and isinstance( - litellm_logging_obj, LiteLLMLogging - ): + if litellm_logging_obj is not None and isinstance(litellm_logging_obj, LiteLLMLogging): kwargs = litellm_logging_obj.model_call_details parent_span = user_api_key_dict.parent_otel_span @@ -1000,43 +953,31 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if dynamic_headers is not None: # Create spans using a temporary tracer with dynamic headers tracer_to_use = self._get_tracer_with_dynamic_headers(dynamic_headers) - verbose_logger.debug( - "[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", dynamic_headers - ) + verbose_logger.debug("[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", dynamic_headers) else: # For langfuse_otel without dynamic headers, create a provider with env var credentials if hasattr(self, "callback_name") and self.callback_name == "langfuse_otel": # Use the headers from config (which were set from env vars during init) - env_var_headers = ( - self._get_headers_dictionary(self.OTEL_HEADERS) - if self.OTEL_HEADERS - else {} - ) + env_var_headers = self._get_headers_dictionary(self.OTEL_HEADERS) if self.OTEL_HEADERS else {} if env_var_headers: - tracer_to_use = self._get_tracer_with_dynamic_headers( - env_var_headers - ) + tracer_to_use = self._get_tracer_with_dynamic_headers(env_var_headers) verbose_logger.debug( "[OTEL DEBUG] Using env var credentials for langfuse_otel (master key request)" ) else: # No env vars set, use global tracer (will be NoOp) tracer_to_use = self.tracer - verbose_logger.debug( - "[OTEL DEBUG] No credentials available for langfuse_otel" - ) + verbose_logger.debug("[OTEL DEBUG] No credentials available for langfuse_otel") else: tracer_to_use = self.tracer - verbose_logger.debug( - "[OTEL DEBUG] Using GLOBAL tracer (no dynamic headers)" - ) + verbose_logger.debug("[OTEL DEBUG] Using GLOBAL tracer (no dynamic headers)") return tracer_to_use def _get_dynamic_otel_headers_from_kwargs(self, kwargs) -> Optional[dict]: """Extract dynamic headers from kwargs if available.""" - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params") + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get( + "standard_callback_dynamic_params" ) if not standard_callback_dynamic_params: @@ -1055,15 +996,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Prevents thread exhaustion by reusing providers for the same credential sets (e.g. per-team keys) cache_key = str(sorted(dynamic_headers.items())) if cache_key in self._tracer_provider_cache: - return self._tracer_provider_cache[cache_key].get_tracer( - LITELLM_TRACER_NAME - ) + return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME) # Create a temporary tracer provider with dynamic headers temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config)) - temp_provider.add_span_processor( - self._get_span_processor(dynamic_headers=dynamic_headers) - ) + temp_provider.add_span_processor(self._get_span_processor(dynamic_headers=dynamic_headers)) # Store in cache for reuse self._tracer_provider_cache[cache_key] = temp_provider @@ -1211,19 +1148,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Decide whether to create a primary span # Always create if no parent span exists (backward compatibility) # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled - should_create_primary_span = parent_span is None or get_secret_bool( - "USE_OTEL_LITELLM_REQUEST_SPAN" - ) + should_create_primary_span = parent_span is None or get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN") if should_create_primary_span: # Create a new litellm_request span - span = self._start_primary_span( - kwargs, response_obj, start_time, end_time, ctx - ) + span = self._start_primary_span(kwargs, response_obj, start_time, end_time, ctx) # Raw-request sub-span (if enabled) - child of litellm_request span - self._maybe_log_raw_request( - kwargs, response_obj, start_time, end_time, span - ) + self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) # Do NOT duplicate attributes onto the parent proxy-request span. # The child litellm_request span already carries all attributes; # copying them to the parent doubles storage and complicates @@ -1239,15 +1170,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): parent_span.set_status(Status(StatusCode.OK)) self.set_attributes(parent_span, kwargs, response_obj) # Raw-request as direct child of parent_span - self._maybe_log_raw_request( - kwargs, response_obj, start_time, end_time, parent_span - ) + self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, parent_span) # 3. Guardrail span — ensure guardrails are always parented to an # existing span so they never become orphaned root spans (Issue #5). - guardrail_ctx = self._resolve_guardrail_context( - span=span, parent_span=parent_span, fallback_ctx=ctx - ) + guardrail_ctx = self._resolve_guardrail_context(span=span, parent_span=parent_span, fallback_ctx=ctx) self._create_guardrail_span(kwargs=kwargs, context=guardrail_ctx) # 4. Metrics & cost recording @@ -1306,9 +1233,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): span.end(end_time=self._to_ns(end_time)) return span - def _maybe_log_raw_request( - self, kwargs, response_obj, start_time, end_time, parent_span - ): + def _maybe_log_raw_request(self, kwargs, response_obj, start_time, end_time, parent_span): from opentelemetry import trace from opentelemetry.trace import Status, StatusCode @@ -1412,42 +1337,28 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): http_route = metadata.get("user_api_key_request_route") if http_route: - self.safe_set_attribute( - span=span, key=HTTP_ROUTE_ATTRIBUTE, value=http_route - ) + self.safe_set_attribute(span=span, key=HTTP_ROUTE_ATTRIBUTE, value=http_route) # ``user_api_key_team_metadata`` is dropped from the standard logging # payload metadata, so read it from the raw request metadata in kwargs. # ``metadata`` and ``litellm_metadata`` are alternate names for the same # full metadata dict (the name varies by endpoint), so first-truthy wins. - raw_metadata = ( - litellm_params.get("metadata") - or litellm_params.get("litellm_metadata") - or {} - ) + raw_metadata = litellm_params.get("metadata") or litellm_params.get("litellm_metadata") or {} team_metadata = self._team_metadata_json( raw_metadata.get("user_api_key_team_metadata"), self.config.baggage_team_metadata_keys, ) if team_metadata: - self.safe_set_attribute( - span=span, key=TEAM_METADATA_ATTRIBUTE, value=team_metadata - ) + self.safe_set_attribute(span=span, key=TEAM_METADATA_ATTRIBUTE, value=team_metadata) model_group = standard_logging_payload.get("model_group") if model_group: - self.safe_set_attribute( - span=span, key=MODEL_GROUP_ATTRIBUTE, value=model_group - ) + self.safe_set_attribute(span=span, key=MODEL_GROUP_ATTRIBUTE, value=model_group) hidden_params = standard_logging_payload.get("hidden_params") or {} - provider_model = hidden_params.get( - "litellm_model_name" - ) or standard_logging_payload.get("model") + provider_model = hidden_params.get("litellm_model_name") or standard_logging_payload.get("model") if provider_model: - self.safe_set_attribute( - span=span, key=PROVIDER_MODEL_ATTRIBUTE, value=provider_model - ) + self.safe_set_attribute(span=span, key=PROVIDER_MODEL_ATTRIBUTE, value=provider_model) @staticmethod def _team_metadata_json(value: Any, allowed_keys: List[str]) -> Optional[str]: @@ -1473,11 +1384,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): attributes = self.config.attributes if attributes is None and self.callback_name in (None, "otel"): otel_settings = (litellm.callback_settings or {}).get("otel") or {} - raw = ( - otel_settings.get("attributes") - if isinstance(otel_settings, dict) - else None - ) + raw = otel_settings.get("attributes") if isinstance(otel_settings, dict) else None if raw is not None: attributes = _build_metric_attribute_filter(raw) ( @@ -1492,9 +1399,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if self._metric_attr_include is not None: return {k: v for k, v in attrs.items() if k in self._metric_attr_include} if self._metric_attr_exclude is not None: - return { - k: v for k, v in attrs.items() if k not in self._metric_attr_exclude - } + return {k: v for k, v in attrs.items() if k not in self._metric_attr_exclude} return attrs def _record_metrics(self, kwargs, response_obj, start_time, end_time): @@ -1504,9 +1409,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): common_attrs = { "gen_ai.operation.name": ( - self._gen_ai_operation_name(kwargs) - if self._gen_ai_semconv_latest_experimental - else "chat" + self._gen_ai_operation_name(kwargs) if self._gen_ai_semconv_latest_experimental else "chat" ), "gen_ai.system": provider, "gen_ai.request.model": kwargs.get("model"), @@ -1525,31 +1428,19 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): common_attrs[f"metadata.{key}"] = str(value) # get hidden params - hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get( - "hidden_params", {} - ) + hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get("hidden_params", {}) if hidden_params: common_attrs["hidden_params"] = safe_dumps(hidden_params) common_attrs = self._filter_metric_attributes(common_attrs) if self._operation_duration_histogram: - self._operation_duration_histogram.record( - duration_s, attributes=common_attrs - ) - if ( - response_obj - and (usage := response_obj.get("usage")) - and self._token_usage_histogram - ): + self._operation_duration_histogram.record(duration_s, attributes=common_attrs) + if response_obj and (usage := response_obj.get("usage")) and self._token_usage_histogram: in_attrs = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "input"} out_attrs = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "output"} - self._token_usage_histogram.record( - usage.get("prompt_tokens", 0), attributes=in_attrs - ) - self._token_usage_histogram.record( - usage.get("completion_tokens", 0), attributes=out_attrs - ) + self._token_usage_histogram.record(usage.get("prompt_tokens", 0), attributes=in_attrs) + self._token_usage_histogram.record(usage.get("completion_tokens", 0), attributes=out_attrs) cost = kwargs.get("response_cost") if self._cost_histogram and cost: @@ -1557,9 +1448,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Record latency metrics (TTFT, TPOT, and Total Generation Time) self._record_time_to_first_token_metric(kwargs, common_attrs) - self._record_time_per_output_token_metric( - kwargs, response_obj, end_time, duration_s, common_attrs - ) + self._record_time_per_output_token_metric(kwargs, response_obj, end_time, duration_s, common_attrs) self._record_response_duration_metric(kwargs, end_time, common_attrs) @staticmethod @@ -1604,9 +1493,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return # Skip recording if conversion failed time_to_first_token_seconds = completion_start_ts - api_call_start_ts - self._time_to_first_token_histogram.record( - time_to_first_token_seconds, attributes=common_attrs - ) + self._time_to_first_token_histogram.record(time_to_first_token_seconds, attributes=common_attrs) def _record_time_per_output_token_metric( self, @@ -1643,12 +1530,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Fallback to duration_s if conversion failed generation_time_seconds = duration_s if generation_time_seconds > 0: - time_per_output_token_seconds = ( - generation_time_seconds / completion_tokens - ) - self._time_per_output_token_histogram.record( - time_per_output_token_seconds, attributes=common_attrs - ) + time_per_output_token_seconds = generation_time_seconds / completion_tokens + self._time_per_output_token_histogram.record(time_per_output_token_seconds, attributes=common_attrs) return if completion_start_time is not None: @@ -1674,9 +1557,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if generation_time_seconds > 0: time_per_output_token_seconds = generation_time_seconds / completion_tokens - self._time_per_output_token_histogram.record( - time_per_output_token_seconds, attributes=common_attrs - ) + self._time_per_output_token_histogram.record(time_per_output_token_seconds, attributes=common_attrs) def _record_response_duration_metric( self, @@ -1717,9 +1598,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): response_duration_seconds = end_time_ts - api_call_start_ts if response_duration_seconds > 0: - self._response_duration_histogram.record( - response_duration_seconds, attributes=common_attrs - ) + self._response_duration_histogram.record(response_duration_seconds, attributes=common_attrs) @staticmethod def _otel_log_types(): @@ -1760,9 +1639,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): otel_logger = self._logger_provider.get_logger(LITELLM_LOGGER_NAME) parent_ctx = span.get_span_context() - provider = (kwargs.get("litellm_params") or {}).get( - "custom_llm_provider", "Unknown" - ) + provider = (kwargs.get("litellm_params") or {}).get("custom_llm_provider", "Unknown") if self._gen_ai_semconv_latest_experimental: self._emit_inference_details_event( @@ -1856,31 +1733,23 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return _trace.set_span_in_context(parent_span) return fallback_ctx - def _create_guardrail_span( - self, kwargs: Optional[dict], context: Optional[Context] - ): + def _create_guardrail_span(self, kwargs: Optional[dict], context: Optional[Context]): """ Creates a span for Guardrail, if any guardrail information is present in standard_logging_object """ # Create span for guardrail information kwargs = kwargs or {} - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: return - guardrail_information_data = standard_logging_payload.get( - "guardrail_information" - ) + guardrail_information_data = standard_logging_payload.get("guardrail_information") if not guardrail_information_data: return guardrail_information_list = [ - information - for information in guardrail_information_data - if isinstance(information, dict) + information for information in guardrail_information_data if isinstance(information, dict) ] if not guardrail_information_list: @@ -1938,15 +1807,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): masked_entity_count = guardrail_information.get("masked_entity_count") if masked_entity_count is not None: - guardrail_span.set_attribute( - "masked_entity_count", safe_dumps(masked_entity_count) - ) + guardrail_span.set_attribute("masked_entity_count", safe_dumps(masked_entity_count)) guardrail_response = guardrail_information.get("guardrail_response") if guardrail_response is not None: - guardrail_span.set_attribute( - "guardrail_response", safe_dumps(guardrail_response) - ) + guardrail_span.set_attribute("guardrail_response", safe_dumps(guardrail_response)) # Surface guardrail_status (success / guardrail_intervened / # guardrail_failed_to_respond / not_run) as a top-level span @@ -1975,9 +1840,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if violation_categories: # OTel sequence attributes must be homogeneous primitives; # serialise to JSON once so set_attribute never coerces. - guardrail_span.set_attribute( - "guardrail_violation_categories", safe_dumps(violation_categories) - ) + guardrail_span.set_attribute("guardrail_violation_categories", safe_dumps(violation_categories)) self._set_team_attributes_from_kwargs(guardrail_span, kwargs) @@ -2015,9 +1878,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Decide whether to create a primary span # Always create if no parent span exists (backward compatibility) # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled - should_create_primary_span = parent_otel_span is None or get_secret_bool( - "USE_OTEL_LITELLM_REQUEST_SPAN" - ) + should_create_primary_span = parent_otel_span is None or get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN") span = None if should_create_primary_span: @@ -2085,9 +1946,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): span.record_exception(exception) # Get StandardLoggingPayload for structured error information - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: return @@ -2156,9 +2015,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) except Exception as e: - verbose_logger.exception( - "OpenTelemetry: Error recording exception on span: %s", str(e) - ) + verbose_logger.exception("OpenTelemetry: Error recording exception on span: %s", str(e)) def set_tools_attributes(self, span: Span, tools): import json @@ -2191,9 +2048,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=json.dumps(function.get("parameters")), ) except Exception as e: - verbose_logger.error( - "OpenTelemetry: Error setting tools attributes: %s", str(e) - ) + verbose_logger.error("OpenTelemetry: Error setting tools attributes: %s", str(e)) pass def cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: @@ -2229,9 +2084,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): for key in keys: _value = _function.get(key) if _value: - kv_pairs[ - f"{SpanAttributes.LLM_COMPLETIONS.value}.{idx}.function_call.{key}" - ] = _value + kv_pairs[f"{SpanAttributes.LLM_COMPLETIONS.value}.{idx}.function_call.{key}"] = _value return kv_pairs @@ -2240,18 +2093,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if self.callback_name == "langtrace": from litellm.integrations.langtrace import LangtraceAttributes - LangtraceAttributes().set_langtrace_attributes( - span, kwargs, response_obj - ) + LangtraceAttributes().set_langtrace_attributes(span, kwargs, response_obj) return elif self.callback_name == "langfuse_otel": from litellm.integrations.langfuse.langfuse_otel import ( LangfuseOtelLogger, ) - LangfuseOtelLogger.set_langfuse_otel_attributes( - span, kwargs, response_obj - ) + LangfuseOtelLogger.set_langfuse_otel_attributes(span, kwargs, response_obj) return elif self.callback_name == "weave_otel": from litellm.integrations.weave.weave_otel import ( @@ -2264,9 +2113,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): optional_params = kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) or {} - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") @@ -2277,14 +2124,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ############################################# metadata = standard_logging_payload["metadata"] for key, value in metadata.items(): - self.safe_set_attribute( - span=span, key="metadata.{}".format(key), value=value - ) + self.safe_set_attribute(span=span, key="metadata.{}".format(key), value=value) # get hidden params - hidden_params = getattr( - standard_logging_payload, "hidden_params", None - ) or (standard_logging_payload or {}).get("hidden_params", {}) + hidden_params = getattr(standard_logging_payload, "hidden_params", None) or ( + standard_logging_payload or {} + ).get("hidden_params", {}) if hidden_params: self.safe_set_attribute( span=span, @@ -2298,9 +2143,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): litellm_params=litellm_params, ) # Cost breakdown tracking - cost_breakdown: Optional[CostBreakdown] = standard_logging_payload.get( - "cost_breakdown" - ) + cost_breakdown: Optional[CostBreakdown] = standard_logging_payload.get("cost_breakdown") if cost_breakdown: for key, value in cost_breakdown.items(): if value is not None: @@ -2393,9 +2236,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # but Embeddings and Image-gen responses do not. Fall back to # the litellm call ID so every call type can be correlated # across LiteLLM UI, Phoenix traces, and provider logs (Issue #8). - response_id = ( - response_obj.get("id") if response_obj else None - ) or standard_logging_payload.get("id") + response_id = (response_obj.get("id") if response_obj else None) or standard_logging_payload.get("id") if response_id: self.safe_set_attribute( span=span, @@ -2453,11 +2294,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): self.set_tools_attributes(span, tools) if kwargs.get("messages"): - transformed_messages = ( - self._transform_messages_to_otel_semantic_conventions( - kwargs.get("messages") - ) - ) + transformed_messages = self._transform_messages_to_otel_semantic_conventions(kwargs.get("messages")) self.safe_set_attribute( span=span, key=SpanAttributes.GEN_AI_INPUT_MESSAGES.value, @@ -2474,11 +2311,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): system_instructions = ( kwargs.get("system_instructions") if kwargs.get("system_instructions") is not None - else ( - kwargs.get("instructions") - if kwargs.get("instructions") is not None - else kwargs.get("system") - ) + else (kwargs.get("instructions") if kwargs.get("instructions") is not None else kwargs.get("system")) ) if system_instructions: if isinstance(system_instructions, str): @@ -2489,10 +2322,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=system_instructions, ) else: - transformed_system_instructions = ( - self._transform_messages_to_otel_semantic_conventions( - system_instructions - ) + transformed_system_instructions = self._transform_messages_to_otel_semantic_conventions( + system_instructions ) self.safe_set_attribute( span=span, @@ -2525,10 +2356,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ############################################# if response_obj is not None: if response_obj.get("choices"): - transformed_choices = ( - self._transform_choices_to_otel_semantic_conventions( - response_obj.get("choices") - ) + transformed_choices = self._transform_choices_to_otel_semantic_conventions( + response_obj.get("choices") ) self.safe_set_attribute( span=span, @@ -2567,9 +2396,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # type="message" contains a "content" list of # OutputText objects (type="output_text"). output_items = response_obj.get("output") - output_messages = self._transform_responses_api_output_to_otel( - output_items - ) + output_messages = self._transform_responses_api_output_to_otel(output_items) if output_messages: self.safe_set_attribute( span=span, @@ -2613,12 +2440,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) except Exception as e: - self.handle_callback_failure( - callback_name=self.callback_name or "opentelemetry" - ) - verbose_logger.exception( - "OpenTelemetry logging error in set_attributes %s", str(e) - ) + self.handle_callback_failure(callback_name=self.callback_name or "opentelemetry") + verbose_logger.exception("OpenTelemetry logging error in set_attributes %s", str(e)) def _cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: """ @@ -2644,9 +2467,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): primitive_value = self._cast_as_primitive_value_type(value) span.set_attribute(key, primitive_value) - def _transform_messages_to_otel_semantic_conventions( - self, messages: Union[List[dict], str] - ) -> List[dict]: + def _transform_messages_to_otel_semantic_conventions(self, messages: Union[List[dict], str]) -> List[dict]: """ Transforms LiteLLM/OpenAI style messages into OTEL GenAI 1.38 compliant format. OTEL expects a 'parts' array instead of a single 'content' string. @@ -2687,9 +2508,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return transformed - def _transform_choices_to_otel_semantic_conventions( - self, choices: List[dict] - ) -> List[dict]: + def _transform_choices_to_otel_semantic_conventions(self, choices: List[dict]) -> List[dict]: """ Transforms choices into OTEL GenAI 1.38 compliant format for output.messages. """ @@ -2698,9 +2517,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): message = choice.get("message") or {} finish_reason = choice.get("finish_reason") - transformed_msg = self._transform_messages_to_otel_semantic_conventions( - [message] - )[0] + transformed_msg = self._transform_messages_to_otel_semantic_conventions([message])[0] if finish_reason: transformed_msg["finish_reason"] = finish_reason @@ -2897,16 +2714,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Priority 1: Explicit parent span from metadata if parent_otel_span is not None: - verbose_logger.debug( - "OpenTelemetry: Using explicit parent span from metadata" - ) + verbose_logger.debug("OpenTelemetry: Using explicit parent span from metadata") return trace.set_span_in_context(parent_otel_span), None # Priority 2: HTTP traceparent header if traceparent is not None: - verbose_logger.debug( - "OpenTelemetry: Using traceparent header for context propagation" - ) + verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation") carrier = {"traceparent": traceparent} return ( TraceContextTextMapPropagator().extract(carrier=carrier), @@ -2928,14 +2741,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) return context.get_current(), current_span except Exception as e: - verbose_logger.debug( - "OpenTelemetry: Error getting current span: %s", str(e) - ) + verbose_logger.debug("OpenTelemetry: Error getting current span: %s", str(e)) # Priority 4: No parent context - verbose_logger.debug( - "OpenTelemetry: No parent context found, creating root span" - ) + verbose_logger.debug("OpenTelemetry: No parent context found, creating root span") return None, None def _get_span_processor(self, dynamic_headers: Optional[dict] = None): @@ -2952,26 +2761,17 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): self.OTEL_ENDPOINT, self.OTEL_HEADERS, ) - _split_otel_headers = OpenTelemetry._get_headers_dictionary( - headers=dynamic_headers or self.OTEL_HEADERS - ) + _split_otel_headers = OpenTelemetry._get_headers_dictionary(headers=dynamic_headers or self.OTEL_HEADERS) if dynamic_headers: verbose_logger.debug( "[OTEL DEBUG] Creating span processor with DYNAMIC headers: %s", - { - k: v[:20] + "..." if len(str(v)) > 20 else v - for k, v in _split_otel_headers.items() - }, + {k: v[:20] + "..." if len(str(v)) > 20 else v for k, v in _split_otel_headers.items()}, ) else: - verbose_logger.debug( - "[OTEL DEBUG] Creating span processor with GLOBAL headers" - ) + verbose_logger.debug("[OTEL DEBUG] Creating span processor with GLOBAL headers") - if hasattr( - self.OTEL_EXPORTER, "export" - ): # Check if it has the export method that SpanExporter requires + if hasattr(self.OTEL_EXPORTER, "export"): # Check if it has the export method that SpanExporter requires verbose_logger.debug( "OpenTelemetry: intiializing SpanExporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, @@ -3003,13 +2803,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): "OpenTelemetry: intiializing http exporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, ) - normalized_endpoint = self._normalize_otel_endpoint( - self.OTEL_ENDPOINT, "traces" - ) + normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces") return BatchSpanProcessor( - OTLPSpanExporterHTTP( - endpoint=normalized_endpoint, headers=_split_otel_headers - ), + OTLPSpanExporterHTTP(endpoint=normalized_endpoint, headers=_split_otel_headers), ) elif self.OTEL_EXPORTER == "otlp_grpc" or self.OTEL_EXPORTER == "grpc": try: @@ -3026,13 +2822,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): "OpenTelemetry: intiializing grpc exporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, ) - normalized_endpoint = self._normalize_otel_endpoint( - self.OTEL_ENDPOINT, "traces" - ) + normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces") return BatchSpanProcessor( - OTLPSpanExporterGRPC( - endpoint=normalized_endpoint, headers=_split_otel_headers - ), + OTLPSpanExporterGRPC(endpoint=normalized_endpoint, headers=_split_otel_headers), ) else: verbose_logger.debug( @@ -3094,9 +2886,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): self.OTEL_EXPORTER, normalized_endpoint, ) - return OTLPLogExporter( - endpoint=normalized_endpoint, headers=_split_otel_headers - ) + return OTLPLogExporter(endpoint=normalized_endpoint, headers=_split_otel_headers) elif self.OTEL_EXPORTER == "otlp_grpc" or self.OTEL_EXPORTER == "grpc": try: from opentelemetry.exporter.otlp.proto.grpc._log_exporter import ( @@ -3113,9 +2903,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): self.OTEL_EXPORTER, normalized_endpoint, ) - return OTLPLogExporter( - endpoint=normalized_endpoint, headers=_split_otel_headers - ) + return OTLPLogExporter(endpoint=normalized_endpoint, headers=_split_otel_headers) else: verbose_logger.warning( "OpenTelemetry: Unknown log exporter '%s', defaulting to console. Supported: console, otlp_http, otlp_grpc", @@ -3144,9 +2932,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) _split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS) - normalized_endpoint = self._normalize_otel_endpoint( - self.OTEL_ENDPOINT, "metrics" - ) + normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "metrics") if self.OTEL_EXPORTER == "console": exporter = ConsoleMetricExporter() @@ -3194,9 +2980,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): exporter = ConsoleMetricExporter() return PeriodicExportingMetricReader(exporter, export_interval_millis=5000) - def _normalize_otel_endpoint( - self, endpoint: Optional[str], signal_type: str - ) -> Optional[str]: + def _normalize_otel_endpoint(self, endpoint: Optional[str], signal_type: str) -> Optional[str]: """ Normalize the endpoint URL for a specific OpenTelemetry signal type. @@ -3445,13 +3229,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if url_path: self.safe_set_attribute(span=span, key=URL_PATH_ATTRIBUTE, value=url_path) if http_route: - self.safe_set_attribute( - span=span, key=HTTP_ROUTE_ATTRIBUTE, value=http_route - ) + self.safe_set_attribute(span=span, key=HTTP_ROUTE_ATTRIBUTE, value=http_route) - def set_response_status_code_attribute( - self, span: Optional[Span], status_code: Optional[int] - ) -> None: + def set_response_status_code_attribute(self, span: Optional[Span], status_code: Optional[int]) -> None: """ Set OTel-standard ``http.response.status_code`` (int) on the proxy SERVER span. The failure path sets this from the error code in @@ -3483,20 +3263,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): StandardLoggingPayloadSetup, ) - error_information = StandardLoggingPayloadSetup.get_error_information( - original_exception=exception - ) + error_information = StandardLoggingPayloadSetup.get_error_information(original_exception=exception) error_information["error_code"] = str(status_code) self._record_exception_on_span( span=span, - kwargs={ - "standard_logging_object": {"error_information": error_information} - }, + kwargs={"standard_logging_object": {"error_information": error_information}}, ) - def set_preprocessing_duration_attribute( - self, span: Optional[Span], container: Any - ) -> None: + def set_preprocessing_duration_attribute(self, span: Optional[Span], container: Any) -> None: """ Set ``litellm.preprocessing.duration_ms`` (proxy-receive -> first provider handoff) on the proxy SERVER span. ``litellm_received_at`` diff --git a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py index e45fe149e13..98d24f1f7cc 100644 --- a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py +++ b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py @@ -55,9 +55,7 @@ class OTELSemconvCategory(Enum): # Reverse lookup: opt-in token string -> OTELSemconvCategory. -_SEMCONV_CATEGORY_BY_VALUE = { - category.value: category for category in OTELSemconvCategory -} +_SEMCONV_CATEGORY_BY_VALUE = {category.value: category for category in OTELSemconvCategory} # LiteLLM optional_params key -> OTEL gen_ai semconv span attribute. @@ -123,13 +121,9 @@ class OTELGenAISemconvMixin: def _capture_in_event(self) -> bool: ... - def _transform_messages_to_otel_semantic_conventions( - self, messages: Union[List[dict], str] - ) -> List[dict]: ... + def _transform_messages_to_otel_semantic_conventions(self, messages: Union[List[dict], str]) -> List[dict]: ... - def _transform_choices_to_otel_semantic_conventions( - self, choices: List[dict] - ) -> List[dict]: ... + def _transform_choices_to_otel_semantic_conventions(self, choices: List[dict]) -> List[dict]: ... def _to_ns(self, dt: datetime) -> int: ... @@ -141,10 +135,7 @@ class OTELGenAISemconvMixin: Every semconv behavior is gated on this; ``False`` => legacy output. """ - return ( - OTELSemconvCategory.GEN_AI_LATEST_EXPERIMENTAL - in self.config.semconv_stability_opt_in - ) + return OTELSemconvCategory.GEN_AI_LATEST_EXPERIMENTAL in self.config.semconv_stability_opt_in @staticmethod def _gen_ai_operation_name(kwargs: dict) -> str: @@ -162,9 +153,7 @@ class OTELGenAISemconvMixin: case _: return "chat" - def _set_semconv_request_attributes( - self, span: Span, optional_params: dict - ) -> None: + def _set_semconv_request_attributes(self, span: Span, optional_params: dict) -> None: """Add ``gen_ai.request.*`` span attributes from ``optional_params``. Covers the sampling params plus the conditionally-required @@ -180,9 +169,7 @@ class OTELGenAISemconvMixin: # Spec types this as string[]. safe_set_attribute coerces to a # primitive, so set the array directly via the span API. stop_list = stop if isinstance(stop, list) else [stop] - span.set_attribute( - "gen_ai.request.stop_sequences", [str(s) for s in stop_list] - ) + span.set_attribute("gen_ai.request.stop_sequences", [str(s) for s in stop_list]) # Conditionally required: set only when the request is streaming. if optional_params.get("stream"): @@ -193,30 +180,22 @@ class OTELGenAISemconvMixin: # suppressing nonsensical values (0, negative, non-int). n = optional_params.get("n") if isinstance(n, int) and n > 1: - self.safe_set_attribute( - span=span, key="gen_ai.request.choice.count", value=n - ) + self.safe_set_attribute(span=span, key="gen_ai.request.choice.count", value=n) - def _set_semconv_cache_token_attributes( - self, span: Span, standard_logging_payload - ) -> None: + def _set_semconv_cache_token_attributes(self, span: Span, standard_logging_payload) -> None: """Add ``gen_ai.usage.cache_*.input_tokens`` from the usage object. No-op when the payload or the usage values are missing/zero. """ if not standard_logging_payload: return - usage = (standard_logging_payload.get("metadata") or {}).get( - "usage_object" - ) or {} + usage = (standard_logging_payload.get("metadata") or {}).get("usage_object") or {} for source_key, semconv_key in _SEMCONV_CACHE_TOKEN_ATTRIBUTES.items(): value = usage.get(source_key) if value: self.safe_set_attribute(span=span, key=semconv_key, value=value) - def _build_inference_details_attrs( - self, kwargs: dict, response_obj: dict, provider: str - ) -> Dict[str, Any]: + def _build_inference_details_attrs(self, kwargs: dict, response_obj: dict, provider: str) -> Dict[str, Any]: """Build the attribute payload for the inference-details event. Always includes provider/operation; input/output messages are added @@ -230,12 +209,8 @@ class OTELGenAISemconvMixin: if not self._capture_in_event(): return attrs - input_messages = self._transform_messages_to_otel_semantic_conventions( - kwargs.get("messages") or [] - ) - output_messages = self._transform_choices_to_otel_semantic_conventions( - response_obj.get("choices", []) - ) + input_messages = self._transform_messages_to_otel_semantic_conventions(kwargs.get("messages") or []) + output_messages = self._transform_choices_to_otel_semantic_conventions(response_obj.get("choices", [])) if input_messages: attrs["gen_ai.input.messages"] = safe_dumps(input_messages) if output_messages: @@ -264,8 +239,6 @@ class OTELGenAISemconvMixin: severity_number=SeverityNumber.INFO, severity_text="INFO", body=None, - attributes=self._build_inference_details_attrs( - kwargs, response_obj, provider - ), + attributes=self._build_inference_details_attrs(kwargs, response_obj, provider), ) otel_logger.emit(log_record) diff --git a/litellm/integrations/opik/opik.py b/litellm/integrations/opik/opik.py index b15b024ff09..fd84ad56247 100644 --- a/litellm/integrations/opik/opik.py +++ b/litellm/integrations/opik/opik.py @@ -26,9 +26,7 @@ except Exception: def _should_skip_event(kwargs: Dict[str, Any]) -> bool: """Check if event should be skipped due to missing standard_logging_object.""" if kwargs.get("standard_logging_object") is None: - verbose_logger.debug( - "OpikLogger skipping event; no standard_logging_object found" - ) + verbose_logger.debug("OpikLogger skipping event; no standard_logging_object found") return True return False @@ -39,9 +37,7 @@ class OpikLogger(CustomBatchLogger): """ def __init__(self, **kwargs: Any) -> None: - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_httpx_client = _get_httpx_client() self.opik_project_name: str = ( @@ -165,13 +161,9 @@ class OpikLogger(CustomBatchLogger): verbose_logger.debug("OpikLogger - Flushing batch") await self.flush_queue() except Exception as e: - verbose_logger.exception( - f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}") - def _sync_send( - self, url: str, headers: Dict[str, str], batch: Dict[str, Any] - ) -> None: + def _sync_send(self, url: str, headers: Dict[str, str], batch: Dict[str, Any]) -> None: try: response = self.sync_httpx_client.post( url=url, @@ -180,13 +172,9 @@ class OpikLogger(CustomBatchLogger): ) response.raise_for_status() if response.status_code != 204: - raise Exception( - f"Response from opik API status_code: {response.status_code}, text: {response.text}" - ) + raise Exception(f"Response from opik API status_code: {response.status_code}, text: {response.text}") except Exception as e: - verbose_logger.exception( - f"OpikLogger failed to send batch - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"OpikLogger failed to send batch - {str(e)}\n{traceback.format_exc()}") def log_success_event( self, @@ -257,13 +245,9 @@ class OpikLogger(CustomBatchLogger): batch={"spans": [span_payload.__dict__]}, ) except Exception as e: - verbose_logger.exception( - f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}") - async def _submit_batch( - self, url: str, headers: Dict[str, str], batch: Dict[str, Any] - ) -> None: + async def _submit_batch(self, url: str, headers: Dict[str, str], batch: Dict[str, Any]) -> None: try: response = await self.async_httpx_client.post( url=url, @@ -273,13 +257,9 @@ class OpikLogger(CustomBatchLogger): response.raise_for_status() if response.status_code >= 300: - verbose_logger.error( - f"OpikLogger - Error: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"OpikLogger - Error: {response.status_code} - {response.text}") else: - verbose_logger.info( - f"OpikLogger - {len(self.log_queue)} Opik events submitted" - ) + verbose_logger.info(f"OpikLogger - {len(self.log_queue)} Opik events submitted") except Exception as e: verbose_logger.exception(f"OpikLogger failed to send batch - {str(e)}") @@ -302,12 +282,8 @@ class OpikLogger(CustomBatchLogger): # Send trace batch if len(traces) > 0: - await self._submit_batch( - url=self.trace_url, headers=self.headers, batch={"traces": traces} - ) + await self._submit_batch(url=self.trace_url, headers=self.headers, batch={"traces": traces}) verbose_logger.info(f"Sent {len(traces)} traces") if len(spans) > 0: - await self._submit_batch( - url=self.span_url, headers=self.headers, batch={"spans": spans} - ) + await self._submit_batch(url=self.span_url, headers=self.headers, batch={"spans": spans}) verbose_logger.info(f"Sent {len(spans)} spans") diff --git a/litellm/integrations/opik/opik_payload_builder/api.py b/litellm/integrations/opik/opik_payload_builder/api.py index e3ffab80ae8..6a5f9bfddc5 100644 --- a/litellm/integrations/opik/opik_payload_builder/api.py +++ b/litellm/integrations/opik/opik_payload_builder/api.py @@ -44,9 +44,7 @@ def build_opik_payload( standard_logging_metadata = standard_logging_object.get("metadata", {}) or {} # Extract and merge Opik metadata - opik_metadata = extractors.extract_opik_metadata( - litellm_metadata, standard_logging_metadata - ) + opik_metadata = extractors.extract_opik_metadata(litellm_metadata, standard_logging_metadata) # Extract project name current_project_name = opik_metadata.get("project_name", project_name) diff --git a/litellm/integrations/opik/opik_payload_builder/extractors.py b/litellm/integrations/opik/opik_payload_builder/extractors.py index 1e3a664acc1..73058b2a524 100644 --- a/litellm/integrations/opik/opik_payload_builder/extractors.py +++ b/litellm/integrations/opik/opik_payload_builder/extractors.py @@ -66,9 +66,7 @@ def extract_opik_metadata( if requester_opik: opik_meta.update(requester_opik) - _logging.verbose_logger.debug( - f"litellm_opik_metadata - {json.dumps(opik_meta, default=str)}" - ) + _logging.verbose_logger.debug(f"litellm_opik_metadata - {json.dumps(opik_meta, default=str)}") return opik_meta @@ -94,9 +92,7 @@ def extract_span_identifiers( try: return current_span_data.trace_id, current_span_data.id except AttributeError: - _logging.verbose_logger.warning( - f"Unexpected current_span_data format: {type(current_span_data)}" - ) + _logging.verbose_logger.warning(f"Unexpected current_span_data format: {type(current_span_data)}") return None, None @@ -156,9 +152,7 @@ def apply_proxy_header_overrides( if isinstance(parsed_tags, list): tags.extend(parsed_tags) except (json.JSONDecodeError, TypeError): - _logging.verbose_logger.warning( - f"Failed to parse tags from header: {value}" - ) + _logging.verbose_logger.warning(f"Failed to parse tags from header: {value}") return project_name, tags, thread_id @@ -226,8 +220,6 @@ def extract_and_build_metadata( # Add debug info if cost calculation failed if "response_cost_failure_debug_info" in litellm_kwargs: - metadata["response_cost_failure_debug_info"] = litellm_kwargs[ - "response_cost_failure_debug_info" - ] + metadata["response_cost_failure_debug_info"] = litellm_kwargs["response_cost_failure_debug_info"] return metadata diff --git a/litellm/integrations/opik/opik_payload_builder/payload_builders.py b/litellm/integrations/opik/opik_payload_builder/payload_builders.py index 4656924fdb5..4d92650d2b8 100644 --- a/litellm/integrations/opik/opik_payload_builder/payload_builders.py +++ b/litellm/integrations/opik/opik_payload_builder/payload_builders.py @@ -28,9 +28,7 @@ def build_trace_payload( project_name=project_name, id=trace_id, name=trace_name, - start_time=( - start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z") - ), + start_time=(start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")), end_time=end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"), input=input_data, output=output_data, @@ -63,9 +61,7 @@ def build_span_payload( created = response_obj.get("created", 0) span_name = f"{model}_{obj_type}_{created}" - _logging.verbose_logger.debug( - f"OpikLogger creating span with id {span_id} for trace {trace_id}" - ) + _logging.verbose_logger.debug(f"OpikLogger creating span with id {span_id} for trace {trace_id}") return types.SpanPayload( id=span_id, @@ -75,9 +71,7 @@ def build_span_payload( name=span_name, type="llm", model=model, - start_time=( - start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z") - ), + start_time=(start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")), end_time=end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"), input=input_data, output=output_data, diff --git a/litellm/integrations/opik/utils.py b/litellm/integrations/opik/utils.py index 43577505c11..7222c9d0502 100644 --- a/litellm/integrations/opik/utils.py +++ b/litellm/integrations/opik/utils.py @@ -43,9 +43,7 @@ def _read_opik_config_file() -> Dict[str, str]: config = configparser.ConfigParser() config.read(config_path) - config_values = { - section: dict(config.items(section)) for section in config.sections() - } + config_values = {section: dict(config.items(section)) for section in config.sections()} if "opik" in config_values: return config_values["opik"] diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index 6feaf2734e9..69fc53c5b9d 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -58,9 +58,7 @@ class SpanEmitter: # The mapper chain is the sole source of span attributes. When not # passed in, resolve it from the config so there's one source of truth. self._mappers: list[AttributeMapper] = ( - list(mappers) - if mappers is not None - else resolve_mappers(config.mapper_names) + list(mappers) if mappers is not None else resolve_mappers(config.mapper_names) ) # Bounded LRU (ordered by insertion / most-recent touch). Storing keys # only — the value is unused — so it behaves like a capped set. @@ -127,11 +125,7 @@ class SpanEmitter: # LLM-call and MCP tool-call spans carry a dedup key (their request's # call id), so a sync+async double-firing coalesces. ``isinstance`` narrows # the type for mypy and keeps the engine free of duck-typed attribute reads. - dedup_key = ( - data.identity.call_id - if isinstance(data, (LLMCallSpanData, MCPToolCallSpanData)) - else None - ) + dedup_key = data.identity.call_id if isinstance(data, (LLMCallSpanData, MCPToolCallSpanData)) else None if self._seen(dedup_key, role): return None span = self.start_span( diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 79931c0796c..44484559948 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -109,19 +109,13 @@ class OpenTelemetryV2(CustomLogger): self.config: OpenTelemetryV2Config = config or OpenTelemetryV2Config(**kwargs) self.callback_name = callback_name self._tracer_provider: TracerProvider = ( - tracer_provider - if tracer_provider is not None - else build_tracer_provider(self.config) + tracer_provider if tracer_provider is not None else build_tracer_provider(self.config) ) self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME) self._metrics_recorder = self._init_metrics(meter_provider) self._metric_filter_error_logged = False - self._emitter = SpanEmitter( - self.tracer, self.config, mappers=resolve_mappers(self.config.mapper_names) - ) - self._tenant_tracers = TenantTracerCache( - self.config, callback_name, LITELLM_TRACER_NAME - ) + self._emitter = SpanEmitter(self.tracer, self.config, mappers=resolve_mappers(self.config.mapper_names)) + self._tenant_tracers = TenantTracerCache(self.config, callback_name, LITELLM_TRACER_NAME) self._open_llm_calls: "OrderedDict[str, _LLMCallSpan]" = OrderedDict() self._init_otel_logger_on_litellm_proxy() @@ -145,9 +139,7 @@ class OpenTelemetryV2(CustomLogger): def _register_in_callback_list(self, callbacks: list) -> None: already_otel = any( - cb.__class__.__module__.startswith(_OTEL_MODULES) - for cb in callbacks - if hasattr(cb, "__class__") + cb.__class__.__module__.startswith(_OTEL_MODULES) for cb in callbacks if hasattr(cb, "__class__") ) if not already_otel: callbacks.append(self) @@ -214,13 +206,9 @@ class OpenTelemetryV2(CustomLogger): call.provisional_span_name, parent_context=parent_context, start_time_ns=start_time_ns, - tracer=self._tenant_tracers.tracer_for( - self.tracer, call.dynamic_params - ), + tracer=self._tenant_tracers.tracer_for(self.tracer, call.dynamic_params), ) - self._open_llm_calls[call_id] = _LLMCallSpan( - span=span, start_time_ns=start_time_ns - ) + self._open_llm_calls[call_id] = _LLMCallSpan(span=span, start_time_ns=start_time_ns) # Evict the oldest open call if the map is over budget. A call that opens # but never closes (a stream that only fires stream events) would linger # otherwise; the evicted span is simply dropped (never exported). @@ -272,9 +260,7 @@ class OpenTelemetryV2(CustomLogger): no boundary to open it at), deduped on the call id by the emitter. """ raw_payload = kwargs.get("standard_logging_object") - if not raw_payload or not is_mcp_tool_call( - cast(Mapping[str, object], raw_payload) - ): + if not raw_payload or not is_mcp_tool_call(cast(Mapping[str, object], raw_payload)): return False payload = cast("StandardLoggingPayload", raw_payload) data = MCPToolCallSpanData.from_standard_logging_payload( @@ -320,17 +306,13 @@ class OpenTelemetryV2(CustomLogger): # it (named provisionally) so it isn't leaked as an open span. carrier.span.end(end_time=to_ns(end_time)) return None - data = LLMCallSpanData.from_standard_logging_payload( - payload, capture_content=self.config.capture_span_content - ) + data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=self.config.capture_span_content) end_time_ns = to_ns(end_time) if carrier.span is not None: # Born at the boundary: stamp attributes from the typed payload, set # status, and end it. Its parent (the server span) was captured at # creation from real ambient context. - self._emitter.finish_span( - SpanRole.LLM_CALL, carrier.span, data, end_time_ns=end_time_ns - ) + self._emitter.finish_span(SpanRole.LLM_CALL, carrier.span, data, end_time_ns=end_time_ns) return carrier.span # Deferred: ``pre_call`` saw no recordable parent, so create the span now. # The worker copied the request task's context, which carries the anchored @@ -419,12 +401,7 @@ class OpenTelemetryV2(CustomLogger): # zero-duration root with no context, so skip it. Real background work # (budget/reset jobs, spend flush) passes start/end times and still emits # as a root; anything with a parent emits regardless. - if ( - error_override is None - and start_time is None - and end_time is None - and parent_otel_span is None - ): + if error_override is None and start_time is None and end_time is None and parent_otel_span is None: return None if error_override is not None and data.error is None: data = ServiceSpanData( @@ -570,9 +547,7 @@ def select_global_otel_v2_logger( """ if registered is not None: return registered - existing = next( - (cb for cb in in_memory_loggers if isinstance(cb, OpenTelemetryV2)), None - ) + existing = next((cb for cb in in_memory_loggers if isinstance(cb, OpenTelemetryV2)), None) return existing if existing is not None else OpenTelemetryV2() diff --git a/litellm/integrations/otel/mappers/__init__.py b/litellm/integrations/otel/mappers/__init__.py index 9c728678f47..b0c1d7019db 100644 --- a/litellm/integrations/otel/mappers/__init__.py +++ b/litellm/integrations/otel/mappers/__init__.py @@ -37,9 +37,7 @@ def resolve_mappers(names: Iterable[str]) -> list[AttributeMapper]: for name in names: factory = _MAPPER_BY_NAME.get(name) if factory is None: - raise ValueError( - f"unknown mapper name {name!r}; known: {sorted(_MAPPER_BY_NAME)}" - ) + raise ValueError(f"unknown mapper name {name!r}; known: {sorted(_MAPPER_BY_NAME)}") out.append(factory()) return out diff --git a/litellm/integrations/otel/mappers/base.py b/litellm/integrations/otel/mappers/base.py index dfdaf77a83e..6685e34578b 100644 --- a/litellm/integrations/otel/mappers/base.py +++ b/litellm/integrations/otel/mappers/base.py @@ -14,9 +14,7 @@ from litellm.integrations.otel.model.payloads import ( AttrScalar = str | bool | int | float # Mirrors ``opentelemetry.util.types.AttributeValue`` (homogeneous sequences) # without importing the SDK, so mappers stay OTel-free. -AttrValue = ( - AttrScalar | Sequence[str] | Sequence[bool] | Sequence[int] | Sequence[float] -) +AttrValue = AttrScalar | Sequence[str] | Sequence[bool] | Sequence[int] | Sequence[float] AttributeMap = dict[str, AttrValue] # The closed set of span-data types the engine routes through the mapper chain. diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 75330ebd1bd..ad6d3e7ff21 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -46,18 +46,14 @@ class GenAIMapper: GenAI.REQUEST_FREQUENCY_PENALTY: lambda d: d.request_params.frequency_penalty, GenAI.REQUEST_PRESENCE_PENALTY: lambda d: d.request_params.presence_penalty, GenAI.REQUEST_STOP_SEQUENCES: lambda d: ( - list(d.request_params.stop_sequences) - if d.request_params.stop_sequences - else None + list(d.request_params.stop_sequences) if d.request_params.stop_sequences else None ), GenAI.REQUEST_SEED: lambda d: d.request_params.seed, GenAI.INPUT_MESSAGES: lambda d: serialize_messages(d.messages_in), GenAI.OUTPUT_MESSAGES: lambda d: serialize_messages(output_messages(d)), GenAI.RESPONSE_MODEL: lambda d: d.response_model, GenAI.RESPONSE_ID: lambda d: d.response_id, - GenAI.RESPONSE_FINISH_REASONS: lambda d: ( - list(d.finish_reasons) if d.finish_reasons else None - ), + GenAI.RESPONSE_FINISH_REASONS: lambda d: list(d.finish_reasons) if d.finish_reasons else None, GenAI.USAGE_INPUT_TOKENS: lambda d: d.usage.input_tokens, GenAI.USAGE_OUTPUT_TOKENS: lambda d: d.usage.output_tokens, Error.TYPE: lambda d: d.error.error_type if d.error else None, @@ -80,13 +76,9 @@ class GenAIMapper: f"{LiteLLM.COST_PREFIX}original": lambda d: d.cost.original, f"{LiteLLM.COST_PREFIX}discount_amount": lambda d: d.cost.discount_amount, f"{LiteLLM.COST_PREFIX}discount_percent": lambda d: d.cost.discount_percent, - f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: ( - d.cost.margin_fixed_amount - ), + f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: d.cost.margin_fixed_amount, f"{LiteLLM.COST_PREFIX}margin_percent": lambda d: d.cost.margin_percent, - f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: ( - d.cost.margin_total_amount - ), + f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: d.cost.margin_total_amount, LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming, } @@ -174,10 +166,5 @@ class GenAIMapper: attrs[DB.SYSTEM_NAME] = system if data.call_type: attrs[DB.OPERATION_NAME] = data.call_type - attrs.update( - { - f"{LiteLLM.METADATA_PREFIX}{key}": value - for key, value in data.event_metadata.items() - } - ) + attrs.update({f"{LiteLLM.METADATA_PREFIX}{key}": value for key, value in data.event_metadata.items()}) return attrs diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py index d0af460f752..79f8f618eff 100644 --- a/litellm/integrations/otel/mappers/langfuse.py +++ b/litellm/integrations/otel/mappers/langfuse.py @@ -58,13 +58,9 @@ class LangfuseMapper: ), "langfuse.observation.input": lambda d: serialize_messages(d.messages_in), "langfuse.observation.output": lambda d: serialize_messages(output_messages(d)), - "langfuse.observation.usage_details": lambda d: json_if( - collect(LangfuseMapper._USAGE_FIELDS, d.usage) - ), + "langfuse.observation.usage_details": lambda d: json_if(collect(LangfuseMapper._USAGE_FIELDS, d.usage)), "langfuse.observation.cost_details": lambda d: ( - json.dumps({"total": d.response_cost}) - if d.response_cost is not None - else None + json.dumps({"total": d.response_cost}) if d.response_cost is not None else None ), } diff --git a/litellm/integrations/otel/mappers/langtrace.py b/litellm/integrations/otel/mappers/langtrace.py index ec595439fe5..975864b51b4 100644 --- a/litellm/integrations/otel/mappers/langtrace.py +++ b/litellm/integrations/otel/mappers/langtrace.py @@ -40,12 +40,8 @@ class LangtraceMapper: } _BLOB_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { - "llm.prompts": lambda d: ( - json_or_none(list(d.messages_in)) if d.messages_in else None - ), - "llm.completions": lambda d: ( - json_or_none(output_messages(d)) if d.choices_out else None - ), + "llm.prompts": lambda d: json_or_none(list(d.messages_in)) if d.messages_in else None, + "llm.completions": lambda d: json_or_none(output_messages(d)) if d.choices_out else None, } def map(self, data: SpanData) -> AttributeMap: diff --git a/litellm/integrations/otel/mappers/legacy.py b/litellm/integrations/otel/mappers/legacy.py index 20ffe8b0dd8..57dc7ed3632 100644 --- a/litellm/integrations/otel/mappers/legacy.py +++ b/litellm/integrations/otel/mappers/legacy.py @@ -47,9 +47,7 @@ class LegacyMapper: _LEGACY_FREQUENCY_PENALTY: lambda d: d.request_params.frequency_penalty, _LEGACY_PRESENCE_PENALTY: lambda d: d.request_params.presence_penalty, _LEGACY_STOP_SEQUENCES: lambda d: ( - list(d.request_params.stop_sequences) - if d.request_params.stop_sequences - else None + list(d.request_params.stop_sequences) if d.request_params.stop_sequences else None ), } @@ -62,9 +60,7 @@ class LegacyMapper: _SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = { _LEGACY_SERVICE: lambda d: d.service_name, _LEGACY_CALL_TYPE: lambda d: d.call_type, - _LEGACY_ERROR: lambda d: ( - d.error.message if d.error is not None and d.error.message else None - ), + _LEGACY_ERROR: lambda d: d.error.message if d.error is not None and d.error.message else None, } def map(self, data: SpanData) -> AttributeMap: diff --git a/litellm/integrations/otel/mappers/openinference.py b/litellm/integrations/otel/mappers/openinference.py index d8195cbe03d..dab9a616979 100644 --- a/litellm/integrations/otel/mappers/openinference.py +++ b/litellm/integrations/otel/mappers/openinference.py @@ -83,21 +83,14 @@ class OpenInferenceMapper: **collect(cls._LLM_CALL_ATTRS, data), **collect(cls._BLOB_ATTRS, data), **cls._messages("llm.input_messages", "input.value", data.messages_in), - **cls._messages( - "llm.output_messages", "output.value", output_messages(data) - ), + **cls._messages("llm.output_messages", "output.value", output_messages(data)), **cls._tools(data), } @staticmethod - def _messages( - prefix: str, value_key: str, messages: Sequence[object] - ) -> AttributeMap: + def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap: """Per-message ``{prefix}.{idx}.message.*`` keys + the ``value_key`` blob.""" - parsed = [ - (m.get("role") if isinstance(m, dict) else None, message_content(m)) - for m in messages - ] + parsed = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages] attrs = drop_none( { key: value @@ -112,9 +105,7 @@ class OpenInferenceMapper: } ) if parsed: - attrs[value_key] = json.dumps( - [{"role": role, "content": content} for role, content in parsed] - ) + attrs[value_key] = json.dumps([{"role": role, "content": content} for role, content in parsed]) return attrs @classmethod diff --git a/litellm/integrations/otel/mappers/utils.py b/litellm/integrations/otel/mappers/utils.py index 6228fc8bbe7..a91e59e4ab8 100644 --- a/litellm/integrations/otel/mappers/utils.py +++ b/litellm/integrations/otel/mappers/utils.py @@ -47,9 +47,7 @@ def stringify_message(message: object) -> str | None: def serialize_messages(messages: Sequence[object]) -> str | None: """Round-trip a sequence of message dicts through ``stringify_message``.""" - serialized = [ - json.loads(s) for s in (stringify_message(m) for m in messages) if s is not None - ] + serialized = [json.loads(s) for s in (stringify_message(m) for m in messages) if s is not None] return json.dumps(serialized) if serialized else None @@ -62,11 +60,7 @@ def message_content(message: object) -> str | None: return content if isinstance(content, list): # multimodal: concatenate text parts only - parts = [ - part.get("text", "") - for part in content - if isinstance(part, dict) and part.get("type") == "text" - ] + parts = [part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text"] return "".join(p for p in parts if isinstance(p, str)) or None return None diff --git a/litellm/integrations/otel/mappers/weave.py b/litellm/integrations/otel/mappers/weave.py index 54b07299271..2eb4ad817c7 100644 --- a/litellm/integrations/otel/mappers/weave.py +++ b/litellm/integrations/otel/mappers/weave.py @@ -19,18 +19,14 @@ class WeaveMapper: _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { # ``display_name`` has the form ``"{operation} {model}"``. The span # name already covers that, but Weave reads this attribute too. - "weave.display_name": lambda d: ( - f"{d.operation.value} {d.request_model}" if d.request_model else None - ), + "weave.display_name": lambda d: f"{d.operation.value} {d.request_model}" if d.request_model else None, "weave.call_id": lambda d: d.identity.call_id or None, } # JSON-payload attributes: each builder returns the serialized blob or None. _BLOB_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { # Weave treats the response choices as the "output" payload. - "weave.output": lambda d: ( - json_or_none(list(d.choices_out)) if d.choices_out else None - ), + "weave.output": lambda d: json_or_none(list(d.choices_out)) if d.choices_out else None, } def map(self, data: SpanData) -> AttributeMap: diff --git a/litellm/integrations/otel/model/baggage.py b/litellm/integrations/otel/model/baggage.py index 77ace736c2f..0903b5ad34e 100644 --- a/litellm/integrations/otel/model/baggage.py +++ b/litellm/integrations/otel/model/baggage.py @@ -24,20 +24,16 @@ from litellm.integrations.otel.model.semconv import GenAI, LiteLLM # team_metadata_keys). The single definition of what may be promoted and under # which key. Only the ``TEAM_METADATA`` extractor consults team_metadata_keys # (to filter the team's metadata to an allowlist); the rest ignore it. -_PROMOTABLE: Final[ - dict[str, Callable[[RequestIdentity, str | None, tuple[str, ...]], str | None]] -] = { +_PROMOTABLE: Final[dict[str, Callable[[RequestIdentity, str | None, tuple[str, ...]], str | None]]] = { LiteLLM.TEAM_ID: lambda identity, model, team_metadata_keys: identity.team_id, LiteLLM.TEAM_ALIAS: lambda identity, model, team_metadata_keys: identity.team_alias, - LiteLLM.TEAM_METADATA: lambda identity, model, team_metadata_keys: ( - _filtered_team_metadata_json(identity.team_metadata, team_metadata_keys) + LiteLLM.TEAM_METADATA: lambda identity, model, team_metadata_keys: _filtered_team_metadata_json( + identity.team_metadata, team_metadata_keys ), LiteLLM.KEY_HASH: lambda identity, model, team_metadata_keys: identity.key_hash, LiteLLM.END_USER: lambda identity, model, team_metadata_keys: identity.end_user, GenAI.REQUEST_MODEL: lambda identity, model, team_metadata_keys: model, - LiteLLM.PROVIDER_MODEL: lambda identity, model, team_metadata_keys: ( - identity.provider_model - ), + LiteLLM.PROVIDER_MODEL: lambda identity, model, team_metadata_keys: identity.provider_model, } # Keys promoted by default (a subset of ``_PROMOTABLE``). ``END_USER`` is diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index 991b156ae64..7f33129c560 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -122,12 +122,8 @@ class OpenTelemetryV2Config(BaseSettings): default=None, validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"), ) - service_name: str = Field( - default="litellm", validation_alias=AliasChoices("OTEL_SERVICE_NAME") - ) - deployment_environment: str | None = Field( - default=None, validation_alias=AliasChoices("OTEL_ENVIRONMENT_NAME") - ) + service_name: str = Field(default="litellm", validation_alias=AliasChoices("OTEL_SERVICE_NAME")) + deployment_environment: str | None = Field(default=None, validation_alias=AliasChoices("OTEL_ENVIRONMENT_NAME")) enable_metrics: bool = Field( default=False, @@ -139,13 +135,9 @@ class OpenTelemetryV2Config(BaseSettings): ) capture_message_content: str = Field( default=CaptureMessageContent.NO_CONTENT, - validation_alias=AliasChoices( - "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT" - ), - ) - legacy_compat: bool = Field( - default=True, validation_alias=AliasChoices("LITELLM_OTEL_LEGACY_COMPAT") + validation_alias=AliasChoices("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"), ) + legacy_compat: bool = Field(default=True, validation_alias=AliasChoices("LITELLM_OTEL_LEGACY_COMPAT")) # ----- explicit multi-destination / vocabulary configuration ------------ # @@ -179,9 +171,7 @@ class OpenTelemetryV2Config(BaseSettings): baggage_promoted_keys: Annotated[List[str], NoDecode] = Field( default_factory=lambda: list(BAGGAGE_PROMOTED_KEYS), - validation_alias=AliasChoices( - "baggage_promoted_keys", "LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS" - ), + validation_alias=AliasChoices("baggage_promoted_keys", "LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS"), description=( "Identity attribute keys written into Baggage and stamped on every " "child span (e.g. ``litellm.team.id``). Configure via the " @@ -192,9 +182,7 @@ class OpenTelemetryV2Config(BaseSettings): ) baggage_metadata_keys: Annotated[List[str], NoDecode] = Field( default_factory=lambda: list(DEFAULT_BAGGAGE_METADATA_KEYS), - validation_alias=AliasChoices( - "baggage_metadata_keys", "LITELLM_OTEL_BAGGAGE_METADATA_KEYS" - ), + validation_alias=AliasChoices("baggage_metadata_keys", "LITELLM_OTEL_BAGGAGE_METADATA_KEYS"), description=( "Metadata sub-keys promoted under the ``litellm.metadata.*`` " "namespace. Configure via the ``LITELLM_OTEL_BAGGAGE_METADATA_KEYS`` " @@ -204,9 +192,7 @@ class OpenTelemetryV2Config(BaseSettings): ) baggage_team_metadata_keys: Annotated[List[str], NoDecode] = Field( default_factory=lambda: list(DEFAULT_BAGGAGE_TEAM_METADATA_KEYS), - validation_alias=AliasChoices( - "baggage_team_metadata_keys", "LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS" - ), + validation_alias=AliasChoices("baggage_team_metadata_keys", "LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS"), description=( "Sub-keys of the team's free-form metadata promoted under " "``litellm.team.metadata``. Empty by default so none of a team's " diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index 4c9cecfef57..37bb5464315 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -73,26 +73,17 @@ class RequestIdentity: model, not just the user-facing one. """ raw_meta = cast(Mapping[str, object], payload.get("metadata") or {}) - metadata = { - key: str(value) - for key, value in raw_meta.items() - if isinstance(value, (str, bool, int, float)) - } + metadata = {key: str(value) for key, value in raw_meta.items() if isinstance(value, (str, bool, int, float))} return cls( call_id=as_str(payload.get("litellm_call_id")) or as_str(payload.get("id")), # StandardLoggingMetadata's canonical key is ``user_api_key_team_id``; # the bare ``team_id`` is a legacy alias and is often empty, so prefer # the canonical key and fall back to the alias. - team_id=as_str(raw_meta.get("user_api_key_team_id")) - or as_str(raw_meta.get("team_id")), - team_alias=as_str(raw_meta.get("user_api_key_team_alias")) - or as_str(raw_meta.get("team_alias")), - team_metadata=_team_metadata_dict( - raw_meta.get("user_api_key_team_metadata") - ), + team_id=as_str(raw_meta.get("user_api_key_team_id")) or as_str(raw_meta.get("team_id")), + team_alias=as_str(raw_meta.get("user_api_key_team_alias")) or as_str(raw_meta.get("team_alias")), + team_metadata=_team_metadata_dict(raw_meta.get("user_api_key_team_metadata")), key_hash=as_str(raw_meta.get("user_api_key_hash")), - end_user=as_str(payload.get("end_user")) - or as_str(raw_meta.get("user_api_key_end_user_id")), + end_user=as_str(payload.get("end_user")) or as_str(raw_meta.get("user_api_key_end_user_id")), provider_model=resolve_provider_model(payload), metadata=metadata, ) @@ -153,18 +144,12 @@ class RequestContext: return self.identity.provider_model @classmethod - def from_standard_logging_payload( - cls, payload: "StandardLoggingPayload" - ) -> "RequestContext": + def from_standard_logging_payload(cls, payload: "StandardLoggingPayload") -> "RequestContext": raw_meta = cast(Mapping[str, object], payload.get("metadata") or {}) hidden = cast(Mapping[str, object], payload.get("hidden_params") or {}) raw_response = payload.get("response") - response = cast( - Mapping[str, object], raw_response if isinstance(raw_response, dict) else {} - ) - model_group = as_str(payload.get("model_group")) or as_str( - raw_meta.get("model_group") - ) + response = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}) + model_group = as_str(payload.get("model_group")) or as_str(raw_meta.get("model_group")) return cls( # The user asked for the group; fall back to the call model on the SDK # path, which has no group. Empty string (never None) so the span name @@ -172,8 +157,7 @@ class RequestContext: request_model=model_group or as_str(payload.get("model")) or "", response_model=as_str(response.get("model")), model_group=model_group, - model_id=as_str(payload.get("model_id")) - or _model_info_id(raw_meta.get("model_info")), + model_id=as_str(payload.get("model_id")) or _model_info_id(raw_meta.get("model_info")), api_base=as_str(payload.get("api_base")) or as_str(hidden.get("api_base")), identity=RequestIdentity.from_payload(payload), ) @@ -233,9 +217,7 @@ class LLMCallEvent: ) -def _call_id( - payload: "StandardLoggingPayload | None", kwargs: Mapping[str, Any] -) -> str | None: +def _call_id(payload: "StandardLoggingPayload | None", kwargs: Mapping[str, Any]) -> str | None: """The call id from the payload (when closed) or the bare kwargs (at pre_call).""" if payload is not None: call_id = as_str(payload.get("litellm_call_id")) or as_str(payload.get("id")) @@ -268,9 +250,7 @@ def resolve_provider_model(payload: "StandardLoggingPayload") -> str | None: return ( # ``deployment`` survives only on paths that don't strip it from metadata; # harmless (and most precise) to prefer it when present. - as_str(raw_meta.get("deployment")) - or as_str(hidden.get("litellm_model_name")) - or as_str(payload.get("model")) + as_str(raw_meta.get("deployment")) or as_str(hidden.get("litellm_model_name")) or as_str(payload.get("model")) ) diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 82b7df5922c..a368a862024 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -187,14 +187,10 @@ class GuardrailSpanData: error: SpanError | None = None # Guardrail statuses that mean the guardrail did not pass the request through. - _ERROR_STATUSES: ClassVar[frozenset[str]] = frozenset( - {"guardrail_intervened", "guardrail_failed_to_respond"} - ) + _ERROR_STATUSES: ClassVar[frozenset[str]] = frozenset({"guardrail_intervened", "guardrail_failed_to_respond"}) @classmethod - def from_logging_entry( - cls, entry: "StandardLoggingGuardrailInformation" - ) -> "GuardrailSpanData": + def from_logging_entry(cls, entry: "StandardLoggingGuardrailInformation") -> "GuardrailSpanData": """Build from one ``standard_logging_guardrail_information`` entry. Reads the canonical, provider-agnostic ``StandardLoggingGuardrailInformation`` @@ -279,9 +275,7 @@ class ToolDefinition: name: str description: str | None = None - parameters_json: str | None = ( - None # JSON-serialized schema (str so it's an AttrValue) - ) + parameters_json: str | None = None # JSON-serialized schema (str so it's an AttrValue) @dataclass(frozen=True) @@ -322,9 +316,7 @@ class LLMCallSpanData: # Normalize ``response`` to a dict once so the content/id reads below are a # plain ``.get`` — no repeated ``isinstance`` guards. raw_response = payload.get("response") - response = cast( - Mapping[str, object], raw_response if isinstance(raw_response, dict) else {} - ) + response = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}) choices_out = _dicts(response.get("choices")) # ``finish_reasons`` is metadata, not content, so derive it from # ``choices_out`` before gating. The raw message/choice bodies are only @@ -347,9 +339,7 @@ class LLMCallSpanData: finish_reasons=finish_reasons, error=_parse_error(payload), response_cost=as_float(payload.get("response_cost")), - cost=LLMCost.from_breakdown( - cast("Mapping[str, object] | None", payload.get("cost_breakdown")) - ), + cost=LLMCost.from_breakdown(cast("Mapping[str, object] | None", payload.get("cost_breakdown"))), server=ServerInfo.from_api_base(context.api_base), identity=context.identity, is_streaming=as_bool(payload.get("stream")), @@ -396,14 +386,10 @@ class MCPToolCallSpanData: server_name=as_str(meta.get("mcp_server_name")), session_id=as_str(meta.get("mcp_session_id")), arguments_json=( - _json_or_none(meta.get("arguments")) - if capture_content and meta.get("arguments") is not None - else None + _json_or_none(meta.get("arguments")) if capture_content and meta.get("arguments") is not None else None ), result_json=( - _json_or_none(meta.get("result")) - if capture_content and meta.get("result") is not None - else None + _json_or_none(meta.get("result")) if capture_content and meta.get("result") is not None else None ), error=_parse_error(payload), response_cost=as_float(payload.get("response_cost")), @@ -426,9 +412,7 @@ def is_mcp_tool_call(payload: Mapping[str, object]) -> bool: """Whether a closed request's payload is an MCP tool call rather than an LLM call — true when the MCP gateway stamped its tool-call metadata, or the call type says so on a path that hasn't populated the metadata yet.""" - return bool(_mcp_tool_call_metadata(payload)) or ( - payload.get("call_type") == "call_mcp_tool" - ) + return bool(_mcp_tool_call_metadata(payload)) or (payload.get("call_type") == "call_mcp_tool") # --- service event_metadata sanitization ------------------------------------ # @@ -448,9 +432,7 @@ _SENSITIVE_METADATA_SUBSTRINGS: tuple[str, ...] = ( # Keys that carry raw call-site internals — live objects, full kwargs/args. The # operation name is already the span's ``call_type``, so ``function_name`` is # redundant. -_DROP_METADATA_KEYS: frozenset = frozenset( - {"function_kwargs", "function_args", "function_name"} -) +_DROP_METADATA_KEYS: frozenset = frozenset({"function_kwargs", "function_args", "function_name"}) _MAX_METADATA_VALUE_LEN = 1024 _MAX_METADATA_ITEMS = 32 diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 1adc1d68dde..bc624cf6a57 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -77,26 +77,14 @@ class SpanSpec: SPAN_REGISTRY: dict[SpanRole, SpanSpec] = { - SpanRole.PROXY_REQUEST: SpanSpec( - SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None - ), - SpanRole.LLM_CALL: SpanSpec( - SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST - ), + SpanRole.PROXY_REQUEST: SpanSpec(SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None), + SpanRole.LLM_CALL: SpanSpec(SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST), # The proxy is an MCP client to the upstream server it dispatches the tool # call to, so this is a CLIENT span, sibling of the LLM call under the request. - SpanRole.MCP_TOOL_CALL: SpanSpec( - SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST - ), - SpanRole.GUARDRAIL: SpanSpec( - SpanRole.GUARDRAIL, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST - ), - SpanRole.DB_CALL: SpanSpec( - SpanRole.DB_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST - ), - SpanRole.SERVICE: SpanSpec( - SpanRole.SERVICE, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST - ), + SpanRole.MCP_TOOL_CALL: SpanSpec(SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST), + SpanRole.GUARDRAIL: SpanSpec(SpanRole.GUARDRAIL, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST), + SpanRole.DB_CALL: SpanSpec(SpanRole.DB_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST), + SpanRole.SERVICE: SpanSpec(SpanRole.SERVICE, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST), } @@ -138,9 +126,7 @@ def db_system(service_name: str) -> str | None: # - ``auth`` — emitted instead as a live phase span (see # ``logger.phase_span``) so its DB lookups nest under it, # not as a flat post-hoc service span. -_METRICS_ONLY_SERVICES: frozenset[str] = frozenset( - {"self", "router", "proxy_pre_call", "auth"} -) +_METRICS_ONLY_SERVICES: frozenset[str] = frozenset({"self", "router", "proxy_pre_call", "auth"}) def span_role_for_service(service_name: str) -> SpanRole | None: diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 64790da814b..ff513c84d95 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -27,9 +27,7 @@ _PROPAGATOR = TraceContextTextMapPropagator() # and is inherited by ``asyncio.create_task`` children — i.e. the async logging # callbacks that close the span. It is never reset: the contextvar dies with the # request task, so there is nothing to leak. -_request_root_span: "ContextVar[Span | None]" = ContextVar( - "litellm_otel_request_root_span", default=None -) +_request_root_span: "ContextVar[Span | None]" = ContextVar("litellm_otel_request_root_span", default=None) def set_request_root_span(span: Span) -> None: @@ -49,9 +47,7 @@ def request_root_span() -> "Span | None": return span if is_recordable_span(span) else None -def set_request_baggage( - values: Mapping[str, str], context: Context | None = None -) -> Context: +def set_request_baggage(values: Mapping[str, str], context: Context | None = None) -> Context: """Return a context with ``values`` written into Baggage.""" ctx = context for key, value in values.items(): diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 95ac939ff7f..cb1f9214876 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -81,9 +81,7 @@ class GenAIMetricRecorder: survives. """ - def __init__( - self, metrics: GenAIMetrics, callback_name: Optional[str] = None - ) -> None: + def __init__(self, metrics: GenAIMetrics, callback_name: Optional[str] = None) -> None: self._metrics = metrics self._callback_name = callback_name self._include: Optional[FrozenSet[str]] = None @@ -108,9 +106,7 @@ class GenAIMetricRecorder: self._metrics.token_cost.record(cost, attributes=common_attrs) self._record_time_to_first_token(kwargs, common_attrs) - self._record_time_per_output_token( - kwargs, response_obj, end_time, duration_s, common_attrs - ) + self._record_time_per_output_token(kwargs, response_obj, end_time, duration_s, common_attrs) self._record_response_duration(kwargs, end_time, common_attrs) # ------------------------------------------------------------------ # @@ -138,9 +134,7 @@ class GenAIMetricRecorder: else: common_attrs[f"metadata.{key}"] = str(value) - hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get( - "hidden_params", {} - ) + hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get("hidden_params", {}) if hidden_params: common_attrs["hidden_params"] = safe_dumps(hidden_params) @@ -152,11 +146,7 @@ class GenAIMetricRecorder: attributes = None if self._callback_name in (None, "otel"): otel_settings = (litellm.callback_settings or {}).get("otel") or {} - raw = ( - otel_settings.get("attributes") - if isinstance(otel_settings, dict) - else None - ) + raw = otel_settings.get("attributes") if isinstance(otel_settings, dict) else None if raw is not None: attributes = _build_metric_attribute_filter(raw) # A bad filter (include_list + exclude_list both set, an unfilterable name) @@ -187,25 +177,17 @@ class GenAIMetricRecorder: return in_attrs = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "input"} out_attrs = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "output"} - self._metrics.token_usage.record( - usage.get("prompt_tokens", 0), attributes=in_attrs - ) - self._metrics.token_usage.record( - usage.get("completion_tokens", 0), attributes=out_attrs - ) + self._metrics.token_usage.record(usage.get("prompt_tokens", 0), attributes=in_attrs) + self._metrics.token_usage.record(usage.get("completion_tokens", 0), attributes=out_attrs) - def _record_time_to_first_token( - self, kwargs: Mapping[str, Any], common_attrs: dict - ) -> None: + def _record_time_to_first_token(self, kwargs: Mapping[str, Any], common_attrs: dict) -> None: if not kwargs.get("optional_params", {}).get("stream", False): return api_call_start = to_seconds(kwargs.get("api_call_start_time")) completion_start = to_seconds(kwargs.get("completion_start_time")) if api_call_start is None or completion_start is None: return - self._metrics.time_to_first_token.record( - completion_start - api_call_start, attributes=common_attrs - ) + self._metrics.time_to_first_token.record(completion_start - api_call_start, attributes=common_attrs) def _record_time_per_output_token( self, @@ -229,27 +211,17 @@ class GenAIMetricRecorder: api_call_start_time = kwargs.get("api_call_start_time") if completion_start_time is not None: completion_start = to_seconds(completion_start_time) - generation_time = ( - duration_s - if completion_start is None - else end_ts - completion_start - ) + generation_time = duration_s if completion_start is None else end_ts - completion_start elif api_call_start_time is not None: api_call_start = to_seconds(api_call_start_time) - generation_time = ( - duration_s if api_call_start is None else end_ts - api_call_start - ) + generation_time = duration_s if api_call_start is None else end_ts - api_call_start else: generation_time = duration_s if generation_time > 0: - self._metrics.time_per_output_token.record( - generation_time / completion_tokens, attributes=common_attrs - ) + self._metrics.time_per_output_token.record(generation_time / completion_tokens, attributes=common_attrs) - def _record_response_duration( - self, kwargs: Mapping[str, Any], end_time: datetime, common_attrs: dict - ) -> None: + def _record_response_duration(self, kwargs: Mapping[str, Any], end_time: datetime, common_attrs: dict) -> None: api_call_start_time = kwargs.get("api_call_start_time") if api_call_start_time is None: return diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 6d0710397a3..ac971c6daa8 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -53,9 +53,7 @@ def to_otel_span_kind(kind: LiteLLMSpanKind) -> SpanKind: _EXPORTER_FACTORIES: dict[str, Callable[[ExporterSpec], SpanExporter]] = {} -def register_exporter_factory( - kind: str, factory: Callable[[ExporterSpec], SpanExporter] -) -> None: +def register_exporter_factory(kind: str, factory: Callable[[ExporterSpec], SpanExporter]) -> None: """Register a custom exporter ``factory`` for the exporter ``kind``.""" _EXPORTER_FACTORIES[kind.lower()] = factory @@ -72,9 +70,7 @@ class LiteLLMBaggageSpanProcessor(SpanProcessor): self._allowed_prefixes = tuple(allowed_prefixes) def _is_allowed(self, key: str) -> bool: - return key in self._allowed_keys or any( - key.startswith(prefix) for prefix in self._allowed_prefixes - ) + return key in self._allowed_keys or any(key.startswith(prefix) for prefix in self._allowed_prefixes) def on_start(self, span: Span, parent_context: Context | None = None) -> None: for key, value in baggage.get_all(parent_context).items(): @@ -156,11 +152,7 @@ def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter: ``exporter`` / ``endpoint`` / ``headers`` fields. To configure multiple exporters, populate ``config.exporters`` directly. """ - return _exporter_from_spec( - ExporterSpec( - kind=config.exporter, endpoint=config.endpoint, headers=config.headers - ) - ) + return _exporter_from_spec(ExporterSpec(kind=config.exporter, endpoint=config.endpoint, headers=config.headers)) def _otlp_metrics_endpoint(endpoint: str | None) -> str | None: @@ -301,9 +293,7 @@ def build_tracer_provider( """ provider = TracerProvider(resource=build_resource(config)) if baggage_processor is None: - baggage_processor = LiteLLMBaggageSpanProcessor( - allowed_keys=config.baggage_promoted_keys - ) + baggage_processor = LiteLLMBaggageSpanProcessor(allowed_keys=config.baggage_promoted_keys) provider.add_span_processor(baggage_processor) if exporter is not None: @@ -317,11 +307,7 @@ def build_tracer_provider( provider.add_span_processor( _processor_for( exp, - ( - spec.use_simple_processor - if spec.use_simple_processor is not None - else use_simple_processor - ), + (spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor), ) ) return provider diff --git a/litellm/integrations/otel/plumbing/routing.py b/litellm/integrations/otel/plumbing/routing.py index 1f2f1b202d9..2f8945e903b 100644 --- a/litellm/integrations/otel/plumbing/routing.py +++ b/litellm/integrations/otel/plumbing/routing.py @@ -60,9 +60,7 @@ class TenantTracerCache: self._config = config self._callback_name = callback_name self._tracer_name = tracer_name - self._providers: "OrderedDict[tuple[tuple[str, str], ...], TracerProvider]" = ( - OrderedDict() - ) + self._providers: "OrderedDict[tuple[tuple[str, str], ...], TracerProvider]" = OrderedDict() def tracer_for(self, default: Tracer, dynamic_params: Any) -> Tracer: """Return the tracer for this request. @@ -102,8 +100,7 @@ class TenantTracerCache: exporters = [ ( spec.model_copy(update=header_update) - if spec.owner == self._callback_name - and spec.kind.lower() not in _NON_OTLP_KINDS + if spec.owner == self._callback_name and spec.kind.lower() not in _NON_OTLP_KINDS else spec ) for spec in self._config.exporters diff --git a/litellm/integrations/otel/presets/__init__.py b/litellm/integrations/otel/presets/__init__.py index c69d257ab52..deaf953ede8 100644 --- a/litellm/integrations/otel/presets/__init__.py +++ b/litellm/integrations/otel/presets/__init__.py @@ -39,9 +39,7 @@ PRESET_BY_CALLBACK: dict[str, Preset] = { #: routing). Only integrations that support dynamic credentials appear here — #: Arize-Phoenix/Langtrace/Levo/AgentOps don't, so they use the logger's #: default tracer. -DYNAMIC_HEADERS_BY_CALLBACK: dict[ - str, Callable[[StandardCallbackDynamicParams], dict[str, str]] -] = { +DYNAMIC_HEADERS_BY_CALLBACK: dict[str, Callable[[StandardCallbackDynamicParams], dict[str, str]]] = { "arize": arize_dynamic_headers, "langfuse_otel": langfuse_dynamic_headers, "weave_otel": weave_dynamic_headers, diff --git a/litellm/integrations/otel/presets/agentops.py b/litellm/integrations/otel/presets/agentops.py index 5d3ed91a711..048b63c89fb 100644 --- a/litellm/integrations/otel/presets/agentops.py +++ b/litellm/integrations/otel/presets/agentops.py @@ -32,12 +32,8 @@ class _AgentOpsSettings(BaseSettings): model_config = SettingsConfigDict(case_sensitive=False, extra="ignore") api_key: str | None = Field(default=None, validation_alias="AGENTOPS_API_KEY") - service_name: str = Field( - default="agentops", validation_alias="AGENTOPS_SERVICE_NAME" - ) - environment: str | None = Field( - default=None, validation_alias="AGENTOPS_ENVIRONMENT" - ) + service_name: str = Field(default="agentops", validation_alias="AGENTOPS_SERVICE_NAME") + environment: str | None = Field(default=None, validation_alias="AGENTOPS_ENVIRONMENT") def agentops_preset( @@ -60,9 +56,7 @@ def agentops_preset( ExporterSpec( kind=_AGENTOPS_EXPORTER_KIND, endpoint=_AGENTOPS_ENDPOINT, - options=( - {"api_key": settings.api_key} if settings.api_key else None - ), + options=({"api_key": settings.api_key} if settings.api_key else None), owner=ExporterOwner.AGENTOPS, ), ], @@ -70,11 +64,7 @@ def agentops_preset( **base.resource_attributes, "service.name": settings.service_name, "telemetry.sdk.name": "agentops", - **( - {"deployment.environment": settings.environment} - if settings.environment - else {} - ), + **({"deployment.environment": settings.environment} if settings.environment else {}), }, } ) @@ -120,9 +110,7 @@ def _build_agentops_exporter(spec: ExporterSpec) -> Any: return super().export(spans) options = spec.options or {} - return _LazyAuthAgentOpsExporter( - endpoint=spec.endpoint, api_key=options.get("api_key") - ) + return _LazyAuthAgentOpsExporter(endpoint=spec.endpoint, api_key=options.get("api_key")) def _fetch_agentops_jwt(api_key: str) -> dict[str, Any]: diff --git a/litellm/integrations/otel/presets/arize.py b/litellm/integrations/otel/presets/arize.py index b6af88c6b34..95206205630 100644 --- a/litellm/integrations/otel/presets/arize.py +++ b/litellm/integrations/otel/presets/arize.py @@ -18,9 +18,7 @@ class _ArizeSettings(BaseSettings): # Standard OTLP headers env var, used as the fallback when no Arize # credentials are configured. - otlp_traces_headers: str | None = Field( - default=None, validation_alias="OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ) + otlp_traces_headers: str | None = Field(default=None, validation_alias="OTEL_EXPORTER_OTLP_TRACES_HEADERS") def arize_preset( @@ -44,11 +42,7 @@ def arize_preset( "mapper_names": ensure_mappers(base.mapper_names, "openinference"), "resource_attributes": { **base.resource_attributes, - **( - {"model_id": arize_cfg.project_name} - if arize_cfg.project_name - else {} - ), + **({"model_id": arize_cfg.project_name} if arize_cfg.project_name else {}), }, } ) diff --git a/litellm/integrations/otel/presets/base.py b/litellm/integrations/otel/presets/base.py index b50908e7652..3b9991f86a4 100644 --- a/litellm/integrations/otel/presets/base.py +++ b/litellm/integrations/otel/presets/base.py @@ -20,6 +20,4 @@ class Preset(Protocol): test-supplied defaults); the factory calls presets with no arguments. """ - def __call__( - self, *, config_overrides: OpenTelemetryV2Config | None = None - ) -> OpenTelemetryV2Config: ... + def __call__(self, *, config_overrides: OpenTelemetryV2Config | None = None) -> OpenTelemetryV2Config: ... diff --git a/litellm/integrations/otel/presets/phoenix.py b/litellm/integrations/otel/presets/phoenix.py index 5485b599321..4e7be2c0f51 100644 --- a/litellm/integrations/otel/presets/phoenix.py +++ b/litellm/integrations/otel/presets/phoenix.py @@ -19,9 +19,7 @@ class _PhoenixSettings(BaseSettings): project_name: str = Field( default="default", - validation_alias=AliasChoices( - "PHOENIX_PROJECT_NAME", "PHOENIX_COLLECTOR_PROJECT_NAME" - ), + validation_alias=AliasChoices("PHOENIX_PROJECT_NAME", "PHOENIX_COLLECTOR_PROJECT_NAME"), ) diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index 072ae4945a0..e519736e162 100644 --- a/litellm/integrations/posthog.py +++ b/litellm/integrations/posthog.py @@ -49,16 +49,12 @@ class PostHogLogger(CustomBatchLogger): self.is_mock_mode = should_use_posthog_mock() if self.is_mock_mode: create_mock_posthog_client() - verbose_logger.debug( - "[POSTHOG MOCK] PostHog logger initialized in mock mode" - ) + verbose_logger.debug("[POSTHOG MOCK] PostHog logger initialized in mock mode") if os.getenv("POSTHOG_API_KEY", None) is None: raise Exception("POSTHOG_API_KEY is not set, set 'POSTHOG_API_KEY=<>'") - self.async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_client = _get_httpx_client() self.POSTHOG_API_KEY = os.getenv("POSTHOG_API_KEY") @@ -73,21 +69,15 @@ class PostHogLogger(CustomBatchLogger): # Register cleanup handler to flush internal queue on exit atexit.register(self._flush_on_exit) - super().__init__( - **kwargs, flush_lock=None, batch_size=POSTHOG_MAX_BATCH_SIZE - ) + super().__init__(**kwargs, flush_lock=None, batch_size=POSTHOG_MAX_BATCH_SIZE) except Exception as e: - verbose_logger.exception( - f"PostHog: Got exception on init PostHog client {str(e)}" - ) + verbose_logger.exception(f"PostHog: Got exception on init PostHog client {str(e)}") raise e def log_success_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug( - "PostHog: Sync logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("PostHog: Sync logging - Enters logging function for model %s", kwargs) api_key, api_url = self._get_credentials_for_request(kwargs) if api_key is None or api_url is None: @@ -109,9 +99,7 @@ class PostHogLogger(CustomBatchLogger): response.raise_for_status() if response.status_code != 200: - raise Exception( - f"Response from PostHog API status_code: {response.status_code}, text: {response.text}" - ) + raise Exception(f"Response from PostHog API status_code: {response.status_code}, text: {response.text}") if self.is_mock_mode: verbose_logger.debug("[POSTHOG MOCK] Sync event successfully mocked") @@ -123,9 +111,7 @@ class PostHogLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug( - "PostHog: Async logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("PostHog: Async logging - Enters logging function for model %s", kwargs) self._ensure_async_setup() # Lazy initialization await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: @@ -134,36 +120,26 @@ class PostHogLogger(CustomBatchLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug( - "PostHog: Async logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("PostHog: Async logging - Enters logging function for model %s", kwargs) self._ensure_async_setup() # Lazy initialization await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: verbose_logger.exception(f"PostHog Layer Error - {str(e)}") pass - async def _log_async_event( - self, kwargs, response_obj=None, start_time=0.0, end_time=0.0 - ): + async def _log_async_event(self, kwargs, response_obj=None, start_time=0.0, end_time=0.0): # Note: response_obj, start_time, end_time not used - all data comes from kwargs api_key, api_url = self._get_credentials_for_request(kwargs) event_payload = self.create_posthog_event_payload(kwargs) # Store event with its credentials for batch sending - self.log_queue.append( - {"event": event_payload, "api_key": api_key, "api_url": api_url} - ) - verbose_logger.debug( - f"PostHog, event added to queue. Will flush in {self.flush_interval} seconds..." - ) + self.log_queue.append({"event": event_payload, "api_key": api_key, "api_url": api_url}) + verbose_logger.debug(f"PostHog, event added to queue. Will flush in {self.flush_interval} seconds...") if len(self.log_queue) >= self.batch_size: await self.flush_queue() - def create_posthog_event_payload( - self, kwargs: Dict[str, Any] - ) -> PostHogEventPayload: + def create_posthog_event_payload(self, kwargs: Dict[str, Any]) -> PostHogEventPayload: """ Helper function to create a PostHog event payload for logging @@ -173,9 +149,7 @@ class PostHogLogger(CustomBatchLogger): Returns: PostHogEventPayload: defined in types.py """ - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object not found in kwargs") @@ -207,9 +181,7 @@ class PostHogLogger(CustomBatchLogger): # Core model information properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "") - properties["$ai_provider"] = self._safe_get( - standard_logging_object, "custom_llm_provider", "" - ) + properties["$ai_provider"] = self._safe_get(standard_logging_object, "custom_llm_provider", "") # Input/Output data messages = self._safe_get(standard_logging_object, "messages") @@ -222,22 +194,16 @@ class PostHogLogger(CustomBatchLogger): properties["$ai_output_choices"] = response # Token information - properties["$ai_input_tokens"] = self._safe_get( - standard_logging_object, "prompt_tokens", 0 - ) + properties["$ai_input_tokens"] = self._safe_get(standard_logging_object, "prompt_tokens", 0) if event_name == "$ai_generation": - properties["$ai_output_tokens"] = self._safe_get( - standard_logging_object, "completion_tokens", 0 - ) + properties["$ai_output_tokens"] = self._safe_get(standard_logging_object, "completion_tokens", 0) # Cost and performance response_cost = self._safe_get(standard_logging_object, "response_cost") if response_cost is not None: properties["$ai_total_cost_usd"] = response_cost - properties["$ai_latency"] = self._safe_get( - standard_logging_object, "response_time", 0.0 - ) + properties["$ai_latency"] = self._safe_get(standard_logging_object, "response_time", 0.0) # Error handling if self._safe_get(standard_logging_object, "status") == "failure": @@ -257,9 +223,7 @@ class PostHogLogger(CustomBatchLogger): def _add_trace_properties(self, properties: Dict[str, Any], kwargs: Dict[str, Any]): standard_logging_object = self._safe_get(kwargs, "standard_logging_object", {}) - trace_id = self._safe_get( - standard_logging_object, "trace_id", self._safe_uuid() - ) + trace_id = self._safe_get(standard_logging_object, "trace_id", self._safe_uuid()) properties["$ai_trace_id"] = trace_id span_id = self._safe_get(standard_logging_object, "id", self._safe_uuid()) @@ -270,9 +234,7 @@ class PostHogLogger(CustomBatchLogger): if parent_id: properties["$ai_parent_id"] = parent_id - def _add_custom_metadata_properties( - self, properties: Dict[str, Any], kwargs: Dict[str, Any] - ): + def _add_custom_metadata_properties(self, properties: Dict[str, Any], kwargs: Dict[str, Any]): """Add custom metadata fields to PostHog properties""" metadata = self._extract_metadata(kwargs) if not isinstance(metadata, dict): @@ -318,9 +280,7 @@ class PostHogLogger(CustomBatchLogger): if key not in litellm_internal_fields: properties[key] = value - def _get_distinct_id( - self, standard_logging_object: StandardLoggingPayload, kwargs: Dict[str, Any] - ) -> str: + def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: Dict[str, Any]) -> str: metadata = self._extract_metadata(kwargs) user_id = self._safe_get(metadata, "user_id") if user_id: @@ -334,9 +294,7 @@ class PostHogLogger(CustomBatchLogger): return self._safe_uuid() - def _get_credentials_for_request( - self, kwargs: Dict[str, Any] - ) -> Tuple[Optional[str], Optional[str]]: + def _get_credentials_for_request(self, kwargs: Dict[str, Any]) -> Tuple[Optional[str], Optional[str]]: """ Get PostHog credentials for this request. @@ -349,19 +307,13 @@ class PostHogLogger(CustomBatchLogger): Returns: tuple[str, str]: (api_key, api_url) """ - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params", None) + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get( + "standard_callback_dynamic_params", None ) if standard_callback_dynamic_params is not None: - api_key = ( - standard_callback_dynamic_params.get("posthog_api_key") - or self.POSTHOG_API_KEY - ) - api_url = ( - standard_callback_dynamic_params.get("posthog_api_url") - or self.posthog_host - ) + api_key = standard_callback_dynamic_params.get("posthog_api_key") or self.POSTHOG_API_KEY + api_url = standard_callback_dynamic_params.get("posthog_api_url") or self.posthog_host else: api_key = self.POSTHOG_API_KEY api_url = self.posthog_host @@ -379,14 +331,10 @@ class PostHogLogger(CustomBatchLogger): if not self.log_queue: return - verbose_logger.debug( - f"PostHog: Sending batch of {len(self.log_queue)} events" - ) + verbose_logger.debug(f"PostHog: Sending batch of {len(self.log_queue)} events") if self.is_mock_mode: - verbose_logger.debug( - "[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted" - ) + verbose_logger.debug("[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted") # Group events by credentials for batch sending batches_by_credentials: Dict[tuple[str, str], list] = {} @@ -418,13 +366,9 @@ class PostHogLogger(CustomBatchLogger): ) if self.is_mock_mode: - verbose_logger.debug( - f"[POSTHOG MOCK] Batch of {len(self.log_queue)} events successfully mocked" - ) + verbose_logger.debug(f"[POSTHOG MOCK] Batch of {len(self.log_queue)} events successfully mocked") else: - verbose_logger.debug( - f"PostHog: Batch of {len(self.log_queue)} events successfully sent" - ) + verbose_logger.debug(f"PostHog: Batch of {len(self.log_queue)} events successfully sent") except Exception as e: verbose_logger.exception(f"PostHog Error sending batch API - {str(e)}") @@ -436,9 +380,7 @@ class PostHogLogger(CustomBatchLogger): self._async_initialized = True verbose_logger.debug("PostHog: Async components initialized") except Exception as e: - verbose_logger.error( - f"PostHog: Failed to initialize async components: {str(e)}" - ) + verbose_logger.error(f"PostHog: Failed to initialize async components: {str(e)}") raise def _extract_metadata(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: @@ -469,9 +411,7 @@ class PostHogLogger(CustomBatchLogger): if not self.log_queue: return - verbose_logger.debug( - f"PostHog: Flushing {len(self.log_queue)} remaining events on exit" - ) + verbose_logger.debug(f"PostHog: Flushing {len(self.log_queue)} remaining events on exit") try: # Group events by credentials (same logic as async_send_batch) @@ -499,18 +439,12 @@ class PostHogLogger(CustomBatchLogger): response.raise_for_status() if response.status_code != 200: - verbose_logger.error( - f"PostHog: Failed to flush on exit - status {response.status_code}" - ) + verbose_logger.error(f"PostHog: Failed to flush on exit - status {response.status_code}") if self.is_mock_mode: - verbose_logger.debug( - f"[POSTHOG MOCK] Successfully flushed {len(self.log_queue)} events on exit" - ) + verbose_logger.debug(f"[POSTHOG MOCK] Successfully flushed {len(self.log_queue)} events on exit") else: - verbose_logger.debug( - f"PostHog: Successfully flushed {len(self.log_queue)} events on exit" - ) + verbose_logger.debug(f"PostHog: Successfully flushed {len(self.log_queue)} events on exit") self.log_queue.clear() except Exception as e: diff --git a/litellm/integrations/posthog_mock_client.py b/litellm/integrations/posthog_mock_client.py index de085b855ce..3efaabb9f48 100644 --- a/litellm/integrations/posthog_mock_client.py +++ b/litellm/integrations/posthog_mock_client.py @@ -30,6 +30,4 @@ _config = MockClientConfig( patch_sync_client=True, ) -create_mock_posthog_client, should_use_posthog_mock = create_mock_client_factory( - _config -) +create_mock_posthog_client, should_use_posthog_mock = create_mock_client_factory(_config) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f6ad46ebf39..b517cb0c38d 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -100,11 +100,7 @@ class PrometheusLogger(CustomLogger): self._cached_metric_labels: Dict[str, List[str]] = {} _custom_buckets = litellm.prometheus_latency_buckets - self.latency_buckets = ( - tuple(_custom_buckets) - if _custom_buckets is not None - else LATENCY_BUCKETS - ) + self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS self._bounded_prometheus_series_tracker = BoundedPrometheusSeriesTracker() # Create metric factory functions @@ -115,26 +111,20 @@ class PrometheusLogger(CustomLogger): self.litellm_proxy_failed_requests_metric = self._counter_factory( name="litellm_proxy_failed_requests_metric", documentation="Total number of failed responses from proxy - the client did not get a success response from litellm proxy", - labelnames=self.get_labels_for_metric( - "litellm_proxy_failed_requests_metric" - ), + labelnames=self.get_labels_for_metric("litellm_proxy_failed_requests_metric"), ) self.litellm_proxy_total_requests_metric = self._counter_factory( name="litellm_proxy_total_requests_metric", documentation="Total number of requests made to the proxy server - track number of client side requests", - labelnames=self.get_labels_for_metric( - "litellm_proxy_total_requests_metric" - ), + labelnames=self.get_labels_for_metric("litellm_proxy_total_requests_metric"), ) # request latency metrics self.litellm_request_total_latency_metric = self._histogram_factory( "litellm_request_total_latency_metric", "Total latency (seconds) for a request to LiteLLM", - labelnames=self.get_labels_for_metric( - "litellm_request_total_latency_metric" - ), + labelnames=self.get_labels_for_metric("litellm_request_total_latency_metric"), buckets=self.latency_buckets, ) @@ -155,9 +145,7 @@ class PrometheusLogger(CustomLogger): # "team", # "team_alias", # ], - labelnames=self.get_labels_for_metric( - "litellm_llm_api_time_to_first_token_metric" - ), + labelnames=self.get_labels_for_metric("litellm_llm_api_time_to_first_token_metric"), buckets=self.latency_buckets, ) @@ -197,50 +185,38 @@ class PrometheusLogger(CustomLogger): self.litellm_input_cached_tokens_metric = self._counter_factory( "litellm_input_cached_tokens_metric", "Provider-side cached input tokens (e.g. OpenAI prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens)", - labelnames=self.get_labels_for_metric( - "litellm_input_cached_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_input_cached_tokens_metric"), ) self.litellm_input_cache_creation_tokens_metric = self._counter_factory( "litellm_input_cache_creation_tokens_metric", "Provider-side input tokens written to prompt cache (e.g. Anthropic cache_creation_input_tokens)", - labelnames=self.get_labels_for_metric( - "litellm_input_cache_creation_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_input_cache_creation_tokens_metric"), ) self.litellm_input_audio_tokens_metric = self._counter_factory( "litellm_input_audio_tokens_metric", "Audio input tokens reported in prompt_tokens_details.audio_tokens", - labelnames=self.get_labels_for_metric( - "litellm_input_audio_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_input_audio_tokens_metric"), ) self.litellm_output_reasoning_tokens_metric = self._counter_factory( "litellm_output_reasoning_tokens_metric", "Reasoning tokens reported in completion_tokens_details.reasoning_tokens", - labelnames=self.get_labels_for_metric( - "litellm_output_reasoning_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_output_reasoning_tokens_metric"), ) self.litellm_output_audio_tokens_metric = self._counter_factory( "litellm_output_audio_tokens_metric", "Audio output tokens reported in completion_tokens_details.audio_tokens", - labelnames=self.get_labels_for_metric( - "litellm_output_audio_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_output_audio_tokens_metric"), ) # Remaining Budget for Team self.litellm_remaining_team_budget_metric = self._gauge_factory( "litellm_remaining_team_budget_metric", "Remaining budget for team", - labelnames=self.get_labels_for_metric( - "litellm_remaining_team_budget_metric" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_team_budget_metric"), ) # Max Budget for Team @@ -254,9 +230,7 @@ class PrometheusLogger(CustomLogger): self.litellm_team_budget_remaining_hours_metric = self._gauge_factory( "litellm_team_budget_remaining_hours_metric", "Remaining days for team budget to be reset", - labelnames=self.get_labels_for_metric( - "litellm_team_budget_remaining_hours_metric" - ), + labelnames=self.get_labels_for_metric("litellm_team_budget_remaining_hours_metric"), ) # Number of members in a team @@ -270,9 +244,7 @@ class PrometheusLogger(CustomLogger): self.litellm_remaining_org_budget_metric = self._gauge_factory( "litellm_remaining_org_budget_metric", "Remaining budget for org", - labelnames=self.get_labels_for_metric( - "litellm_remaining_org_budget_metric" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_org_budget_metric"), ) # Max Budget for Org @@ -286,44 +258,34 @@ class PrometheusLogger(CustomLogger): self.litellm_org_budget_remaining_hours_metric = self._gauge_factory( "litellm_org_budget_remaining_hours_metric", "Remaining hours for org budget to be reset", - labelnames=self.get_labels_for_metric( - "litellm_org_budget_remaining_hours_metric" - ), + labelnames=self.get_labels_for_metric("litellm_org_budget_remaining_hours_metric"), ) # Remaining Budget for API Key self.litellm_remaining_api_key_budget_metric = self._gauge_factory( "litellm_remaining_api_key_budget_metric", "Remaining budget for api key", - labelnames=self.get_labels_for_metric( - "litellm_remaining_api_key_budget_metric" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_api_key_budget_metric"), ) # Max Budget for API Key self.litellm_api_key_max_budget_metric = self._gauge_factory( "litellm_api_key_max_budget_metric", "Maximum budget set for api key", - labelnames=self.get_labels_for_metric( - "litellm_api_key_max_budget_metric" - ), + labelnames=self.get_labels_for_metric("litellm_api_key_max_budget_metric"), ) self.litellm_api_key_budget_remaining_hours_metric = self._gauge_factory( "litellm_api_key_budget_remaining_hours_metric", "Remaining hours for api key budget to be reset", - labelnames=self.get_labels_for_metric( - "litellm_api_key_budget_remaining_hours_metric" - ), + labelnames=self.get_labels_for_metric("litellm_api_key_budget_remaining_hours_metric"), ) # Remaining Budget for User self.litellm_remaining_user_budget_metric = self._gauge_factory( "litellm_remaining_user_budget_metric", "Remaining budget for user", - labelnames=self.get_labels_for_metric( - "litellm_remaining_user_budget_metric" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_user_budget_metric"), ) # Max Budget for User @@ -336,9 +298,7 @@ class PrometheusLogger(CustomLogger): self.litellm_user_budget_remaining_hours_metric = self._gauge_factory( "litellm_user_budget_remaining_hours_metric", "Remaining hours for user budget to be reset", - labelnames=self.get_labels_for_metric( - "litellm_user_budget_remaining_hours_metric" - ), + labelnames=self.get_labels_for_metric("litellm_user_budget_remaining_hours_metric"), ) ######################################## @@ -349,18 +309,14 @@ class PrometheusLogger(CustomLogger): self.litellm_remaining_api_key_requests_for_model = self._gauge_factory( "litellm_remaining_api_key_requests_for_model", "Remaining Requests API Key can make for model (model based rpm limit on key)", - labelnames=self.get_labels_for_metric( - "litellm_remaining_api_key_requests_for_model" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_api_key_requests_for_model"), ) # Remaining MODEL TPM limit for API Key self.litellm_remaining_api_key_tokens_for_model = self._gauge_factory( "litellm_remaining_api_key_tokens_for_model", "Remaining Tokens API Key can make for model (model based tpm limit on key)", - labelnames=self.get_labels_for_metric( - "litellm_remaining_api_key_tokens_for_model" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"), ) ######################################## @@ -371,25 +327,19 @@ class PrometheusLogger(CustomLogger): self.litellm_remaining_requests_metric = self._gauge_factory( "litellm_remaining_requests_metric", "LLM Deployment Analytics - remaining requests for model, returned from LLM API Provider", - labelnames=self.get_labels_for_metric( - "litellm_remaining_requests_metric" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_requests_metric"), ) self.litellm_remaining_tokens_metric = self._gauge_factory( "litellm_remaining_tokens_metric", "remaining tokens for model, returned from LLM API Provider", - labelnames=self.get_labels_for_metric( - "litellm_remaining_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_remaining_tokens_metric"), ) self.litellm_overhead_latency_metric = self._histogram_factory( "litellm_overhead_latency_metric", "Latency overhead (milliseconds) added by LiteLLM processing", - labelnames=self.get_labels_for_metric( - "litellm_overhead_latency_metric" - ), + labelnames=self.get_labels_for_metric("litellm_overhead_latency_metric"), buckets=self.latency_buckets, ) @@ -397,9 +347,7 @@ class PrometheusLogger(CustomLogger): self.litellm_request_queue_time_metric = self._histogram_factory( "litellm_request_queue_time_seconds", "Time spent in request queue before processing starts (seconds)", - labelnames=self.get_labels_for_metric( - "litellm_request_queue_time_seconds" - ), + labelnames=self.get_labels_for_metric("litellm_request_queue_time_seconds"), buckets=self.latency_buckets, ) @@ -458,33 +406,25 @@ class PrometheusLogger(CustomLogger): self.litellm_deployment_success_responses = self._counter_factory( name="litellm_deployment_success_responses", documentation="LLM Deployment Analytics - Total number of successful LLM API calls via litellm", - labelnames=self.get_labels_for_metric( - "litellm_deployment_success_responses" - ), + labelnames=self.get_labels_for_metric("litellm_deployment_success_responses"), ) self.litellm_deployment_failure_responses = self._counter_factory( name="litellm_deployment_failure_responses", documentation="LLM Deployment Analytics - Total number of failed LLM API calls for a specific LLM deploymeny. exception_status is the status of the exception from the llm api", - labelnames=self.get_labels_for_metric( - "litellm_deployment_failure_responses" - ), + labelnames=self.get_labels_for_metric("litellm_deployment_failure_responses"), ) self.litellm_deployment_total_requests = self._counter_factory( name="litellm_deployment_total_requests", documentation="LLM Deployment Analytics - Total number of LLM API calls via litellm - success + failure", - labelnames=self.get_labels_for_metric( - "litellm_deployment_total_requests" - ), + labelnames=self.get_labels_for_metric("litellm_deployment_total_requests"), ) # Deployment Latency tracking self.litellm_deployment_latency_per_output_token = self._histogram_factory( name="litellm_deployment_latency_per_output_token", documentation="LLM Deployment Analytics - Latency per output token", - labelnames=self.get_labels_for_metric( - "litellm_deployment_latency_per_output_token" - ), + labelnames=self.get_labels_for_metric("litellm_deployment_latency_per_output_token"), ) self.litellm_deployment_successful_fallbacks = self._counter_factory( @@ -509,9 +449,7 @@ class PrometheusLogger(CustomLogger): self.litellm_llm_api_failed_requests_metric = self._counter_factory( name="litellm_llm_api_failed_requests_metric", documentation="deprecated - use litellm_proxy_failed_requests_metric", - labelnames=self.get_labels_for_metric( - "litellm_llm_api_failed_requests_metric" - ), + labelnames=self.get_labels_for_metric("litellm_llm_api_failed_requests_metric"), ) self.litellm_requests_metric = self._counter_factory( @@ -543,17 +481,13 @@ class PrometheusLogger(CustomLogger): self.litellm_provider_cache_read_input_tokens_metric = self._counter_factory( name="litellm_provider_cache_read_input_tokens_metric", documentation="Total prompt/input tokens read from provider prompt cache (e.g. OpenAI/Anthropic/Gemini/Bedrock)", - labelnames=self.get_labels_for_metric( - "litellm_provider_cache_read_input_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_provider_cache_read_input_tokens_metric"), ) self.litellm_provider_cache_creation_input_tokens_metric = self._counter_factory( name="litellm_provider_cache_creation_input_tokens_metric", documentation="Total prompt/input tokens written to provider prompt cache (e.g. Anthropic/Bedrock)", - labelnames=self.get_labels_for_metric( - "litellm_provider_cache_creation_input_tokens_metric" - ), + labelnames=self.get_labels_for_metric("litellm_provider_cache_creation_input_tokens_metric"), ) # User and Team count metrics @@ -674,9 +608,7 @@ class PrometheusLogger(CustomLogger): if validation_results.has_errors: self._pretty_print_validation_errors(validation_results) - error_message = "Configuration validation failed:\n" + "\n".join( - validation_results.all_error_messages - ) + error_message = "Configuration validation failed:\n" + "\n".join(validation_results.all_error_messages) raise ValueError(error_message) # Build label filters from valid configurations @@ -701,17 +633,13 @@ class PrometheusLogger(CustomLogger): # Validate labels if provided if config.include_labels: - label_error = self._validate_single_metric_labels( - metric_name, config.include_labels - ) + label_error = self._validate_single_metric_labels(metric_name, config.include_labels) if label_error: label_errors.append(label_error) return ValidationResults(metric_errors=metric_errors, label_errors=label_errors) - def _validate_single_metric_name( - self, metric_name: str - ) -> Optional[MetricValidationError]: + def _validate_single_metric_name(self, metric_name: str) -> Optional[MetricValidationError]: """Validate a single metric name""" from typing import get_args @@ -722,16 +650,12 @@ class PrometheusLogger(CustomLogger): ) return None - def _validate_single_metric_labels( - self, metric_name: str, labels: List[str] - ) -> Optional[LabelValidationError]: + def _validate_single_metric_labels(self, metric_name: str, labels: List[str]) -> Optional[LabelValidationError]: """Validate labels for a single metric""" from typing import cast # Get valid labels for this metric from PrometheusMetricLabels - valid_labels = PrometheusMetricLabels.get_labels( - cast(DEFINED_PROMETHEUS_METRICS, metric_name) - ) + valid_labels = PrometheusMetricLabels.get_labels(cast(DEFINED_PROMETHEUS_METRICS, metric_name)) # Find invalid labels invalid_labels = [label for label in labels if label not in valid_labels] @@ -778,9 +702,7 @@ class PrometheusLogger(CustomLogger): # Pretty print functions ######################################################### - def _pretty_print_validation_errors( - self, validation_results: ValidationResults - ) -> None: + def _pretty_print_validation_errors(self, validation_results: ValidationResults) -> None: """Pretty print all validation errors using rich""" try: from rich.console import Console @@ -799,12 +721,8 @@ class PrometheusLogger(CustomLogger): # Show invalid metric names if any if validation_results.metric_errors: - invalid_metrics = [ - e.metric_name for e in validation_results.metric_errors - ] - valid_metrics = validation_results.metric_errors[ - 0 - ].valid_metrics # All should have same valid metrics + invalid_metrics = [e.metric_name for e in validation_results.metric_errors] + valid_metrics = validation_results.metric_errors[0].valid_metrics # All should have same valid metrics metrics_error_text = Text( f"Invalid Metric Names: {', '.join(invalid_metrics)}", @@ -819,9 +737,7 @@ class PrometheusLogger(CustomLogger): title_justify="left", border_style="green", ) - metrics_table.add_column( - "Available Metrics", style="cyan", no_wrap=True - ) + metrics_table.add_column("Available Metrics", style="cyan", no_wrap=True) for metric in sorted(valid_metrics): metrics_table.add_row(metric) @@ -903,9 +819,7 @@ class PrometheusLogger(CustomLogger): f"Invalid labels for metric '{metric_name}': {invalid_labels}. Valid labels: {sorted(valid_labels)}" ) - def _pretty_print_invalid_metric_error( - self, invalid_metric_name: str, valid_metrics: tuple - ) -> None: + def _pretty_print_invalid_metric_error(self, invalid_metric_name: str, valid_metrics: tuple) -> None: """Pretty print error message for invalid metric name using rich""" try: from rich.console import Console @@ -942,9 +856,7 @@ class PrometheusLogger(CustomLogger): except ImportError: # Fallback to simple logging if rich is not available - verbose_logger.error( - f"Invalid metric name: {invalid_metric_name}. Valid metrics: {sorted(valid_metrics)}" - ) + verbose_logger.error(f"Invalid metric name: {invalid_metric_name}. Valid metrics: {sorted(valid_metrics)}") ######################################################### # End of pretty print functions @@ -961,9 +873,7 @@ class PrometheusLogger(CustomLogger): ) raise ValueError(error.message) - def _pretty_print_prometheus_config( - self, label_filters: Dict[str, List[str]] - ) -> None: + def _pretty_print_prometheus_config(self, label_filters: Dict[str, List[str]]) -> None: """Pretty print the processed prometheus configuration using rich""" try: from rich.console import Console @@ -989,9 +899,7 @@ class PrometheusLogger(CustomLogger): for metric in sorted(self.enabled_metrics): metrics_table.add_row(metric) else: - metrics_table.add_row( - "[yellow]All metrics enabled (no filter applied)[/yellow]" - ) + metrics_table.add_row("[yellow]All metrics enabled (no filter applied)[/yellow]") # Create label filters table labels_table = Table( @@ -1005,11 +913,7 @@ class PrometheusLogger(CustomLogger): if label_filters: for metric_name, labels in sorted(label_filters.items()): - labels_str = ( - ", ".join(labels) - if labels - else "[dim]No labels specified[/dim]" - ) + labels_str = ", ".join(labels) if labels else "[dim]No labels specified[/dim]" labels_table.add_row(metric_name, labels_str) else: labels_table.add_row( @@ -1057,9 +961,7 @@ class PrometheusLogger(CustomLogger): return factory - def get_labels_for_metric( - self, metric_name: DEFINED_PROMETHEUS_METRICS - ) -> List[str]: + def get_labels_for_metric(self, metric_name: DEFINED_PROMETHEUS_METRICS) -> List[str]: """ Get the labels for a metric, filtered if configured. @@ -1088,9 +990,7 @@ class PrometheusLogger(CustomLogger): configured_labels = self.label_filters[metric_name] # Return intersection of configured and default labels to ensure we only use valid labels - filtered_labels = [ - label for label in default_labels if label in configured_labels - ] + filtered_labels = [label for label in default_labels if label in configured_labels] self._cached_metric_labels[metric_name] = filtered_labels return filtered_labels @@ -1153,20 +1053,12 @@ class PrometheusLogger(CustomLogger): ) # unpack kwargs - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") - if standard_logging_payload is None or not isinstance( - standard_logging_payload, dict - ): - raise ValueError( - f"standard_logging_object is required, got={standard_logging_payload}" - ) + if standard_logging_payload is None or not isinstance(standard_logging_payload, dict): + raise ValueError(f"standard_logging_object is required, got={standard_logging_payload}") - if self._should_skip_metrics_for_invalid_key( - kwargs=kwargs, standard_logging_payload=standard_logging_payload - ): + if self._should_skip_metrics_for_invalid_key(kwargs=kwargs, standard_logging_payload=standard_logging_payload): return model = kwargs.get("model", "") @@ -1174,31 +1066,21 @@ class PrometheusLogger(CustomLogger): _metadata = litellm_params.get("metadata") or {} get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - end_user_id = get_end_user_id_for_cost_tracking( - litellm_params, service_type="prometheus" - ) + end_user_id = get_end_user_id_for_cost_tracking(litellm_params, service_type="prometheus") user_id = standard_logging_payload["metadata"]["user_api_key_user_id"] user_api_key = standard_logging_payload["metadata"]["user_api_key_hash"] user_api_key_alias = standard_logging_payload["metadata"]["user_api_key_alias"] user_api_team = standard_logging_payload["metadata"]["user_api_key_team_id"] - user_api_team_alias = standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ] - user_api_key_org_id = standard_logging_payload["metadata"].get( - "user_api_key_org_id" - ) - user_api_key_org_alias = standard_logging_payload["metadata"].get( - "user_api_key_org_alias" - ) + user_api_team_alias = standard_logging_payload["metadata"]["user_api_key_team_alias"] + user_api_key_org_id = standard_logging_payload["metadata"].get("user_api_key_org_id") + user_api_key_org_alias = standard_logging_payload["metadata"].get("user_api_key_org_alias") output_tokens = standard_logging_payload["completion_tokens"] tokens_used = standard_logging_payload["total_tokens"] response_cost = standard_logging_payload["response_cost"] combined_metadata = _get_combined_custom_metadata_from_standard_logging_payload( standard_logging_payload=standard_logging_payload ) - if standard_logging_payload is not None and isinstance( - standard_logging_payload, dict - ): + if standard_logging_payload is not None and isinstance(standard_logging_payload, dict): _tags = standard_logging_payload["request_tags"] else: _tags = [] @@ -1228,33 +1110,19 @@ class PrometheusLogger(CustomLogger): api_provider=standard_logging_payload["custom_llm_provider"], exception_status=None, exception_class=None, - custom_metadata_labels=get_custom_labels_from_metadata( - metadata=combined_metadata - ), - route=standard_logging_payload["metadata"].get( - "user_api_key_request_route" - ), + custom_metadata_labels=get_custom_labels_from_metadata(metadata=combined_metadata), + route=standard_logging_payload["metadata"].get("user_api_key_request_route"), client_ip=standard_logging_payload["metadata"].get("requester_ip_address"), user_agent=standard_logging_payload["metadata"].get("user_agent"), - stream=( - str(standard_logging_payload.get("stream")) - if litellm.prometheus_emit_stream_label - else None - ), + stream=(str(standard_logging_payload.get("stream")) if litellm.prometheus_emit_stream_label else None), ) - if ( - user_api_key is not None - and isinstance(user_api_key, str) - and user_api_key.startswith("sk-") - ): + if user_api_key is not None and isinstance(user_api_key, str) and user_api_key.startswith("sk-"): from litellm.proxy.utils import hash_token user_api_key = hash_token(user_api_key) - label_context = PrometheusLabelFactoryContext( - enum_values - ) # amortized per request. + label_context = PrometheusLabelFactoryContext(enum_values) # amortized per request. # increment total LLM requests and spend metric self._increment_top_level_request_and_spend_metrics( @@ -1378,9 +1246,7 @@ class PrometheusLogger(CustomLogger): verbose_logger.debug("prometheus Logging - Enters token metrics function") # token metrics - if standard_logging_payload is not None and isinstance( - standard_logging_payload, dict - ): + if standard_logging_payload is not None and isinstance(standard_logging_payload, dict): _tags = standard_logging_payload["request_tags"] PrometheusLogger._inc_labeled_counter( @@ -1434,9 +1300,7 @@ class PrometheusLogger(CustomLogger): details (most non-OpenAI/Anthropic models). """ metadata = standard_logging_payload.get("metadata") or {} - usage_object = ( - metadata.get("usage_object") if isinstance(metadata, dict) else None - ) + usage_object = metadata.get("usage_object") if isinstance(metadata, dict) else None if not isinstance(usage_object, dict): return @@ -1447,47 +1311,27 @@ class PrometheusLogger(CustomLogger): ( self.litellm_input_cached_tokens_metric, "litellm_input_cached_tokens_metric", - ( - prompt_details.get("cached_tokens") - if isinstance(prompt_details, dict) - else None - ), + (prompt_details.get("cached_tokens") if isinstance(prompt_details, dict) else None), ), ( self.litellm_input_cache_creation_tokens_metric, "litellm_input_cache_creation_tokens_metric", - ( - prompt_details.get("cache_creation_tokens") - if isinstance(prompt_details, dict) - else None - ), + (prompt_details.get("cache_creation_tokens") if isinstance(prompt_details, dict) else None), ), ( self.litellm_input_audio_tokens_metric, "litellm_input_audio_tokens_metric", - ( - prompt_details.get("audio_tokens") - if isinstance(prompt_details, dict) - else None - ), + (prompt_details.get("audio_tokens") if isinstance(prompt_details, dict) else None), ), ( self.litellm_output_reasoning_tokens_metric, "litellm_output_reasoning_tokens_metric", - ( - completion_details.get("reasoning_tokens") - if isinstance(completion_details, dict) - else None - ), + (completion_details.get("reasoning_tokens") if isinstance(completion_details, dict) else None), ), ( self.litellm_output_audio_tokens_metric, "litellm_output_audio_tokens_metric", - ( - completion_details.get("audio_tokens") - if isinstance(completion_details, dict) - else None - ), + (completion_details.get("audio_tokens") if isinstance(completion_details, dict) else None), ), ] @@ -1556,9 +1400,7 @@ class PrometheusLogger(CustomLogger): # Provider prompt caching metrics are independent of LiteLLM cache_hit. provider_cache_read_tokens = 0 provider_cache_creation_tokens = 0 - usage_obj = (standard_logging_payload.get("metadata", {}) or {}).get( - "usage_object" - ) + usage_obj = (standard_logging_payload.get("metadata", {}) or {}).get("usage_object") if isinstance(usage_obj, dict): # Prefer explicit provider cache fields when available. _read = usage_obj.get("cache_read_input_tokens") @@ -1696,9 +1538,7 @@ class PrometheusLogger(CustomLogger): # Set remaining rpm/tpm for API Key + model # see parallel_request_limiter.py - variables are set there model_group = get_model_group_from_litellm_kwargs(kwargs) - remaining_requests_variable_name = ( - f"litellm-key-remaining-requests-{model_group}" - ) + remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" remaining_requests = metadata.get(remaining_requests_variable_name) @@ -1721,26 +1561,18 @@ class PrometheusLogger(CustomLogger): ) label_context = PrometheusLabelFactoryContext(enum_values) requests_labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - "litellm_remaining_api_key_requests_for_model" - ), + supported_enum_labels=self.get_labels_for_metric("litellm_remaining_api_key_requests_for_model"), enum_values=enum_values, label_context=label_context, ) - self.litellm_remaining_api_key_requests_for_model.labels(**requests_labels).set( - remaining_requests - ) + self.litellm_remaining_api_key_requests_for_model.labels(**requests_labels).set(remaining_requests) tokens_labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - "litellm_remaining_api_key_tokens_for_model" - ), + supported_enum_labels=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"), enum_values=enum_values, label_context=label_context, ) - self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set( - remaining_tokens - ) + self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(remaining_tokens) def _set_latency_metrics( self, @@ -1773,9 +1605,7 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, label_context=label_context, ) - self.litellm_llm_api_time_to_first_token_metric.labels( - **_ttft_labels - ).observe(time_to_first_token_seconds) + self.litellm_llm_api_time_to_first_token_metric.labels(**_ttft_labels).observe(time_to_first_token_seconds) self._track_end_user_metric_series( self.litellm_llm_api_time_to_first_token_metric, "litellm_llm_api_time_to_first_token_metric", @@ -1792,15 +1622,11 @@ class PrometheusLogger(CustomLogger): ) if api_call_total_time_seconds is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_llm_api_latency_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_llm_api_latency_metric"), enum_values=enum_values, label_context=label_context, ) - self.litellm_llm_api_latency_metric.labels(**_labels).observe( - api_call_total_time_seconds - ) + self.litellm_llm_api_latency_metric.labels(**_labels).observe(api_call_total_time_seconds) self._track_end_user_metric_series( self.litellm_llm_api_latency_metric, "litellm_llm_api_latency_metric", @@ -1814,15 +1640,11 @@ class PrometheusLogger(CustomLogger): ) if total_time_seconds is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_request_total_latency_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_request_total_latency_metric"), enum_values=enum_values, label_context=label_context, ) - self.litellm_request_total_latency_metric.labels(**_labels).observe( - total_time_seconds - ) + self.litellm_request_total_latency_metric.labels(**_labels).observe(total_time_seconds) self._track_end_user_metric_series( self.litellm_request_total_latency_metric, "litellm_request_total_latency_metric", @@ -1831,20 +1653,14 @@ class PrometheusLogger(CustomLogger): # request queue time (time from arrival to processing start) _litellm_params = kwargs.get("litellm_params", {}) or {} - queue_time_seconds = (_litellm_params.get("metadata") or {}).get( - "queue_time_seconds" - ) + queue_time_seconds = (_litellm_params.get("metadata") or {}).get("queue_time_seconds") if queue_time_seconds is not None and queue_time_seconds >= 0: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_request_queue_time_seconds" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_request_queue_time_seconds"), enum_values=enum_values, label_context=label_context, ) - self.litellm_request_queue_time_metric.labels(**_labels).observe( - queue_time_seconds - ) + self.litellm_request_queue_time_metric.labels(**_labels).observe(queue_time_seconds) self._track_end_user_metric_series( self.litellm_request_queue_time_metric, "litellm_request_queue_time_seconds", @@ -1857,13 +1673,9 @@ class PrometheusLogger(CustomLogger): list(kwargs.keys()) if isinstance(kwargs, dict) else type(kwargs).__name__, ) - standard_logging_payload: StandardLoggingPayload = kwargs.get( - "standard_logging_object", {} - ) + standard_logging_payload: StandardLoggingPayload = kwargs.get("standard_logging_object", {}) - if self._should_skip_metrics_for_invalid_key( - kwargs=kwargs, standard_logging_payload=standard_logging_payload - ): + if self._should_skip_metrics_for_invalid_key(kwargs=kwargs, standard_logging_payload=standard_logging_payload): return model = kwargs.get("model", "") @@ -1871,19 +1683,13 @@ class PrometheusLogger(CustomLogger): litellm_params = kwargs.get("litellm_params", {}) or {} get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - end_user_id = get_end_user_id_for_cost_tracking( - litellm_params, service_type="prometheus" - ) + end_user_id = get_end_user_id_for_cost_tracking(litellm_params, service_type="prometheus") user_id = standard_logging_payload["metadata"]["user_api_key_user_id"] user_api_key = standard_logging_payload["metadata"]["user_api_key_hash"] user_api_key_alias = standard_logging_payload["metadata"]["user_api_key_alias"] user_api_team = standard_logging_payload["metadata"]["user_api_key_team_id"] - user_api_team_alias = standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ] - user_api_key_org_id = standard_logging_payload["metadata"].get( - "user_api_key_org_id" - ) + user_api_team_alias = standard_logging_payload["metadata"]["user_api_key_team_alias"] + user_api_key_org_id = standard_logging_payload["metadata"].get("user_api_key_org_id") try: enum_values = UserAPIKeyLabelValues( @@ -1913,9 +1719,7 @@ class PrometheusLogger(CustomLogger): response_cost=0, ) except Exception as e: - verbose_logger.exception( - "prometheus Layer Error(): Exception occured - {}".format(str(e)) - ) + verbose_logger.exception("prometheus Layer Error(): Exception occured - {}".format(str(e))) pass pass @@ -1943,11 +1747,7 @@ class PrometheusLogger(CustomLogger): status_code = None # Try from enum_values first (most common in our callbacks) - if ( - enum_values - and hasattr(enum_values, "status_code") - and enum_values.status_code - ): + if enum_values and hasattr(enum_values, "status_code") and enum_values.status_code: try: status_code = int(enum_values.status_code) except (ValueError, TypeError): @@ -1955,9 +1755,7 @@ class PrometheusLogger(CustomLogger): if not status_code and exception: # ProxyException uses 'code' attribute, other exceptions may use 'status_code' - status_code = getattr(exception, "status_code", None) or getattr( - exception, "code", None - ) + status_code = getattr(exception, "status_code", None) or getattr(exception, "code", None) if status_code is not None: try: status_code = int(status_code) @@ -1967,9 +1765,9 @@ class PrometheusLogger(CustomLogger): if not status_code and kwargs: exception_in_kwargs = kwargs.get("exception") if exception_in_kwargs: - status_code = getattr( - exception_in_kwargs, "status_code", None - ) or getattr(exception_in_kwargs, "code", None) + status_code = getattr(exception_in_kwargs, "status_code", None) or getattr( + exception_in_kwargs, "code", None + ) if status_code is not None: try: status_code = int(status_code) @@ -2093,12 +1891,8 @@ class PrometheusLogger(CustomLogger): proxy_server_request=request_data.get("proxy_server_request", {}), ) _metadata = request_data.get("metadata", {}) or {} - model_id = _metadata.get("model_info", {}).get("id") or request_data.get( - "model_info", {} - ).get("id") - rate_limit_category, rate_limit_type = self._extract_rate_limit_labels( - original_exception - ) + model_id = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id") + rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception) enum_values = UserAPIKeyLabelValues( end_user=user_api_key_dict.end_user_id, user=user_api_key_dict.user_id, @@ -2120,11 +1914,7 @@ class PrometheusLogger(CustomLogger): client_ip=_metadata.get("requester_ip_address"), user_agent=_metadata.get("user_agent"), model_id=model_id, - stream=( - str(request_data.get("stream")) - if litellm.prometheus_emit_stream_label - else None - ), + stream=(str(request_data.get("stream")) if litellm.prometheus_emit_stream_label else None), ) _label_ctx = PrometheusLabelFactoryContext(enum_values) PrometheusLogger._inc_labeled_counter( @@ -2143,14 +1933,10 @@ class PrometheusLogger(CustomLogger): ) except Exception as e: - verbose_logger.exception( - "prometheus Layer Error(): Exception occured - {}".format(str(e)) - ) + verbose_logger.exception("prometheus Layer Error(): Exception occured - {}".format(str(e))) pass - async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response - ): + async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ Proxy level tracking - triggered when the proxy responds with a success response to the client @@ -2168,36 +1954,24 @@ class PrometheusLogger(CustomLogger): return obj.get(key, default) return getattr(obj, key, default) - def _extract_deployment_failure_label_values( - self, request_kwargs: dict - ) -> Dict[str, Optional[str]]: + def _extract_deployment_failure_label_values(self, request_kwargs: dict) -> Dict[str, Optional[str]]: """ Extract label values for deployment failure metrics from all available sources in request_kwargs. Falls back to litellm_params metadata and user_api_key_auth when standard_logging_payload has None values. """ - standard_logging_payload = ( - request_kwargs.get("standard_logging_object", {}) or {} - ) + standard_logging_payload = request_kwargs.get("standard_logging_object", {}) or {} _litellm_params = request_kwargs.get("litellm_params", {}) or {} _metadata_raw = self._safe_get(standard_logging_payload, "metadata") or {} if isinstance(_metadata_raw, dict): _metadata = _metadata_raw else: _metadata = { - "user_api_key_alias": getattr( - _metadata_raw, "user_api_key_alias", None - ), - "user_api_key_team_id": getattr( - _metadata_raw, "user_api_key_team_id", None - ), - "user_api_key_team_alias": getattr( - _metadata_raw, "user_api_key_team_alias", None - ), + "user_api_key_alias": getattr(_metadata_raw, "user_api_key_alias", None), + "user_api_key_team_id": getattr(_metadata_raw, "user_api_key_team_id", None), + "user_api_key_team_alias": getattr(_metadata_raw, "user_api_key_team_alias", None), "user_api_key_hash": getattr(_metadata_raw, "user_api_key_hash", None), - "requester_ip_address": getattr( - _metadata_raw, "requester_ip_address", None - ), + "requester_ip_address": getattr(_metadata_raw, "requester_ip_address", None), "user_agent": getattr(_metadata_raw, "user_agent", None), } _litellm_params_metadata = _litellm_params.get("metadata", {}) or {} @@ -2246,9 +2020,7 @@ class PrometheusLogger(CustomLogger): if val is not None: return val if user_api_key_auth is not None: - return getattr(user_api_key_auth, "api_key", None) or getattr( - user_api_key_auth, "api_key_hash", None - ) + return getattr(user_api_key_auth, "api_key", None) or getattr(user_api_key_auth, "api_key_hash", None) return None return { @@ -2256,10 +2028,8 @@ class PrometheusLogger(CustomLogger): "team": _get_team_id(), "team_alias": _get_team_alias(), "hashed_api_key": _get_hashed_api_key(), - "client_ip": _metadata.get("requester_ip_address") - or _litellm_params_metadata.get("requester_ip_address"), - "user_agent": _metadata.get("user_agent") - or _litellm_params_metadata.get("user_agent"), + "client_ip": _metadata.get("requester_ip_address") or _litellm_params_metadata.get("requester_ip_address"), + "user_agent": _metadata.get("user_agent") or _litellm_params_metadata.get("user_agent"), } def set_llm_deployment_failure_metrics(self, request_kwargs: dict): @@ -2276,9 +2046,7 @@ class PrometheusLogger(CustomLogger): """ try: verbose_logger.debug("setting remaining tokens requests metric") - standard_logging_payload: StandardLoggingPayload = request_kwargs.get( - "standard_logging_object", {} - ) + standard_logging_payload: StandardLoggingPayload = request_kwargs.get("standard_logging_object", {}) _litellm_params = request_kwargs.get("litellm_params", {}) or {} litellm_model_name = request_kwargs.get("model", None) model_group = standard_logging_payload.get("model_group", None) @@ -2297,9 +2065,9 @@ class PrometheusLogger(CustomLogger): # Fallback: model_group from litellm_metadata if model_group is None: - model_group = (_litellm_params.get("litellm_metadata") or {}).get( - "model_group" - ) or (_litellm_params.get("metadata") or {}).get("model_group") + model_group = (_litellm_params.get("litellm_metadata") or {}).get("model_group") or ( + _litellm_params.get("metadata") or {} + ).get("model_group") llm_provider = _litellm_params.get("custom_llm_provider", None) @@ -2310,26 +2078,14 @@ class PrometheusLogger(CustomLogger): return # Extract context labels from all available sources (fix for None labels) - fallback_values = self._extract_deployment_failure_label_values( - request_kwargs - ) + fallback_values = self._extract_deployment_failure_label_values(request_kwargs) _metadata = standard_logging_payload.get("metadata", {}) or {} - hashed_api_key = fallback_values.get("hashed_api_key") or _metadata.get( - "user_api_key_hash" - ) - api_key_alias = fallback_values.get("api_key_alias") or _metadata.get( - "user_api_key_alias" - ) + hashed_api_key = fallback_values.get("hashed_api_key") or _metadata.get("user_api_key_hash") + api_key_alias = fallback_values.get("api_key_alias") or _metadata.get("user_api_key_alias") team = fallback_values.get("team") or _metadata.get("user_api_key_team_id") - team_alias = fallback_values.get("team_alias") or _metadata.get( - "user_api_key_team_alias" - ) - client_ip = fallback_values.get("client_ip") or _metadata.get( - "requester_ip_address" - ) - user_agent = fallback_values.get("user_agent") or _metadata.get( - "user_agent" - ) + team_alias = fallback_values.get("team_alias") or _metadata.get("user_api_key_team_alias") + client_ip = fallback_values.get("client_ip") or _metadata.get("requester_ip_address") + user_agent = fallback_values.get("user_agent") or _metadata.get("user_agent") # exception_status: prefer status_code, fallback to exception class for known types exception_status = None @@ -2362,9 +2118,7 @@ class PrometheusLogger(CustomLogger): api_base=label_api_base, api_provider=label_api_provider, exception_status=exception_status, - exception_class=( - self._get_exception_class_name(exception) if exception else None - ), + exception_class=(self._get_exception_class_name(exception) if exception else None), requested_model=label_requested_model, hashed_api_key=hashed_api_key, api_key_alias=api_key_alias, @@ -2408,9 +2162,7 @@ class PrometheusLogger(CustomLogger): pass except Exception as e: verbose_logger.debug( - "Prometheus Error: set_llm_deployment_failure_metrics. Exception occured - {}".format( - str(e) - ) + "Prometheus Error: set_llm_deployment_failure_metrics. Exception occured - {}".format(str(e)) ) def _set_deployment_tpm_rpm_limit_metrics( @@ -2430,9 +2182,7 @@ class PrometheusLogger(CustomLogger): if tpm is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_deployment_tpm_limit" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_tpm_limit"), enum_values=UserAPIKeyLabelValues( litellm_model_name=litellm_model_name, model_id=model_id, @@ -2444,9 +2194,7 @@ class PrometheusLogger(CustomLogger): if rpm is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_deployment_rpm_limit" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_rpm_limit"), enum_values=UserAPIKeyLabelValues( litellm_model_name=litellm_model_name, model_id=model_id, @@ -2476,16 +2224,12 @@ class PrometheusLogger(CustomLogger): deployment. """ try: - additional_headers = ( - standard_logging_payload.get("hidden_params", {}) or {} - ).get("additional_headers") or {} + additional_headers = (standard_logging_payload.get("hidden_params", {}) or {}).get( + "additional_headers" + ) or {} - already_have_tokens = ( - additional_headers.get("x_ratelimit_remaining_tokens") is not None - ) - already_have_requests = ( - additional_headers.get("x_ratelimit_remaining_requests") is not None - ) + already_have_tokens = additional_headers.get("x_ratelimit_remaining_tokens") is not None + already_have_requests = additional_headers.get("x_ratelimit_remaining_requests") is not None if already_have_tokens and already_have_requests: return @@ -2502,13 +2246,10 @@ class PrometheusLogger(CustomLogger): return try: - remaining_usage = await llm_router.get_remaining_model_group_usage( - model_group - ) + remaining_usage = await llm_router.get_remaining_model_group_usage(model_group) except Exception as e: verbose_logger.exception( - "Prometheus: get_remaining_model_group_usage failed for " - "model_group=%s: %s", + "Prometheus: get_remaining_model_group_usage failed for model_group=%s: %s", model_group, e, ) @@ -2522,31 +2263,22 @@ class PrometheusLogger(CustomLogger): if not already_have_tokens and remaining_tokens is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_tokens_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_tokens_metric"), enum_values=enum_values, label_context=label_context, ) - self.litellm_remaining_tokens_metric.labels(**_labels).set( - remaining_tokens - ) + self.litellm_remaining_tokens_metric.labels(**_labels).set(remaining_tokens) if not already_have_requests and remaining_requests is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_requests_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_requests_metric"), enum_values=enum_values, label_context=label_context, ) - self.litellm_remaining_requests_metric.labels(**_labels).set( - remaining_requests - ) + self.litellm_remaining_requests_metric.labels(**_labels).set(remaining_requests) except Exception as e: verbose_logger.exception( - "Prometheus Error: _async_set_router_remaining_metrics. " - "Exception occured - {}".format(str(e)) + "Prometheus Error: _async_set_router_remaining_metrics. Exception occured - {}".format(str(e)) ) def set_llm_deployment_success_metrics( @@ -2560,9 +2292,7 @@ class PrometheusLogger(CustomLogger): ): try: verbose_logger.debug("setting remaining tokens requests metric") - standard_logging_payload: Optional[StandardLoggingPayload] = ( - request_kwargs.get("standard_logging_object") - ) + standard_logging_payload: Optional[StandardLoggingPayload] = request_kwargs.get("standard_logging_object") if standard_logging_payload is None: return @@ -2595,24 +2325,14 @@ class PrometheusLogger(CustomLogger): remaining_requests: Optional[int] = None remaining_tokens: Optional[int] = None - if additional_headers := standard_logging_payload["hidden_params"][ - "additional_headers" - ]: + if additional_headers := standard_logging_payload["hidden_params"]["additional_headers"]: # OpenAI / OpenAI Compatible headers - remaining_requests = additional_headers.get( - "x_ratelimit_remaining_requests", None - ) - remaining_tokens = additional_headers.get( - "x_ratelimit_remaining_tokens", None - ) + remaining_requests = additional_headers.get("x_ratelimit_remaining_requests", None) + remaining_tokens = additional_headers.get("x_ratelimit_remaining_tokens", None) - if litellm_overhead_time_ms := standard_logging_payload[ - "hidden_params" - ].get("litellm_overhead_time_ms"): + if litellm_overhead_time_ms := standard_logging_payload["hidden_params"].get("litellm_overhead_time_ms"): _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_overhead_latency_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_overhead_latency_metric"), enum_values=enum_values, label_context=label_context, ) @@ -2628,27 +2348,19 @@ class PrometheusLogger(CustomLogger): "litellm_model_name" """ _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_requests_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_requests_metric"), enum_values=enum_values, label_context=label_context, ) - self.litellm_remaining_requests_metric.labels(**_labels).set( - remaining_requests - ) + self.litellm_remaining_requests_metric.labels(**_labels).set(remaining_requests) if remaining_tokens: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_tokens_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_tokens_metric"), enum_values=enum_values, label_context=label_context, ) - self.litellm_remaining_tokens_metric.labels(**_labels).set( - remaining_tokens - ) + self.litellm_remaining_tokens_metric.labels(**_labels).set(remaining_tokens) """ log these labels @@ -2680,14 +2392,9 @@ class PrometheusLogger(CustomLogger): response_ms: timedelta = end_time - start_time time_to_first_token_response_time: Optional[timedelta] = None - if ( - request_kwargs.get("stream", None) is not None - and request_kwargs["stream"] is True - ): + if request_kwargs.get("stream", None) is not None and request_kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = ( - request_kwargs.get("completion_start_time", end_time) - start_time - ) + time_to_first_token_response_time = request_kwargs.get("completion_start_time", end_time) - start_time # use the metric that is not None # if streaming - use time_to_first_token_response @@ -2706,15 +2413,11 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, label_context=label_context, ) - self.litellm_deployment_latency_per_output_token.labels( - **_labels - ).observe(latency_per_token) + self.litellm_deployment_latency_per_output_token.labels(**_labels).observe(latency_per_token) except Exception as e: verbose_logger.exception( - "Prometheus Error: set_llm_deployment_success_metrics. Exception occured - {}".format( - str(e) - ) + "Prometheus Error: set_llm_deployment_success_metrics. Exception occured - {}".format(str(e)) ) return @@ -2879,9 +2582,7 @@ class PrometheusLogger(CustomLogger): error_type=error_type, ).inc() except Exception as e: - verbose_logger.warning( - f"Error recording check batch cost error metric: {e}" - ) + verbose_logger.warning(f"Error recording check batch cost error metric: {e}") @staticmethod def _get_exception_class_name(exception: Exception) -> str: @@ -2907,9 +2608,7 @@ class PrometheusLogger(CustomLogger): except ImportError: BudgetExceededError = None # type: ignore[assignment,misc] - if BudgetExceededError is not None and isinstance( - exception, BudgetExceededError - ): + if BudgetExceededError is not None and isinstance(exception, BudgetExceededError): return "BudgetExceededError" exception_class_name = "" @@ -2919,9 +2618,7 @@ class PrometheusLogger(CustomLogger): # pretty print the provider name on prometheus # eg. `openai` -> `Openai.` if len(exception_class_name) >= 1: - exception_class_name = ( - exception_class_name[0].upper() + exception_class_name[1:] + "." - ) + exception_class_name = exception_class_name[0].upper() + exception_class_name[1:] + "." exception_class_name += exception.__class__.__name__ return exception_class_name @@ -2947,9 +2644,7 @@ class PrometheusLogger(CustomLogger): validate_rate_limit_type(getattr(exception, "rate_limit_type", None)), ) - async def log_success_fallback_event( - self, original_model_group: str, kwargs: dict, original_exception: Exception - ): + async def log_success_fallback_event(self, original_model_group: str, kwargs: dict, original_exception: Exception): """ Logs a successful LLM fallback event on prometheus @@ -2967,10 +2662,8 @@ class PrometheusLogger(CustomLogger): ) _metadata_key = get_metadata_variable_name_from_kwargs(kwargs) _metadata = kwargs.get(_metadata_key) or {} - standard_metadata: StandardLoggingMetadata = ( - StandardLoggingPayloadSetup.get_standard_logging_metadata( - metadata=_metadata - ) + standard_metadata: StandardLoggingMetadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( + metadata=_metadata ) _new_model = kwargs.get("model") _tags = cast(List[str], kwargs.get("tags") or []) @@ -2994,9 +2687,7 @@ class PrometheusLogger(CustomLogger): label_context=PrometheusLabelFactoryContext(enum_values), ) - async def log_failure_fallback_event( - self, original_model_group: str, kwargs: dict, original_exception: Exception - ): + async def log_failure_fallback_event(self, original_model_group: str, kwargs: dict, original_exception: Exception): """ Logs a failed LLM fallback event on prometheus """ @@ -3014,10 +2705,8 @@ class PrometheusLogger(CustomLogger): _metadata_key = get_metadata_variable_name_from_kwargs(kwargs) _metadata = kwargs.get(_metadata_key) or {} _tags = cast(List[str], kwargs.get("tags") or []) - standard_metadata: StandardLoggingMetadata = ( - StandardLoggingPayloadSetup.get_standard_logging_metadata( - metadata=_metadata - ) + standard_metadata: StandardLoggingMetadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( + metadata=_metadata ) enum_values = UserAPIKeyLabelValues( @@ -3053,9 +2742,7 @@ class PrometheusLogger(CustomLogger): """ ### get labels _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_deployment_state" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_state"), enum_values=UserAPIKeyLabelValues( litellm_model_name=litellm_model_name, model_id=model_id, @@ -3072,9 +2759,7 @@ class PrometheusLogger(CustomLogger): api_base: str, api_provider: str, ): - self.set_litellm_deployment_state( - 0, litellm_model_name, model_id, api_base, api_provider - ) + self.set_litellm_deployment_state(0, litellm_model_name, model_id, api_base, api_provider) def set_deployment_partial_outage( self, @@ -3083,9 +2768,7 @@ class PrometheusLogger(CustomLogger): api_base: Optional[str], api_provider: str, ): - self.set_litellm_deployment_state( - 1, litellm_model_name, model_id, api_base, api_provider - ) + self.set_litellm_deployment_state(1, litellm_model_name, model_id, api_base, api_provider) def set_deployment_complete_outage( self, @@ -3094,9 +2777,7 @@ class PrometheusLogger(CustomLogger): api_base: Optional[str], api_provider: str, ): - self.set_litellm_deployment_state( - 2, litellm_model_name, model_id, api_base, api_provider - ) + self.set_litellm_deployment_state(2, litellm_model_name, model_id, api_base, api_provider) def increment_deployment_cooled_down( self, @@ -3124,13 +2805,9 @@ class PrometheusLogger(CustomLogger): """ Increment metric when logging to a callback fails (e.g., s3_v2, langfuse, etc.) """ - self.litellm_callback_logging_failures_metric.labels( - callback_name=callback_name - ).inc() + self.litellm_callback_logging_failures_metric.labels(callback_name=callback_name).inc() - def track_provider_remaining_budget( - self, provider: str, spend: float, budget_limit: float - ): + def track_provider_remaining_budget(self, provider: str, spend: float, budget_limit: float): """ Track provider remaining budget in Prometheus """ @@ -3141,9 +2818,7 @@ class PrometheusLogger(CustomLogger): ) ) - def _safe_get_remaining_budget( - self, max_budget: Optional[float], spend: Optional[float] - ) -> float: + def _safe_get_remaining_budget(self, max_budget: Optional[float], spend: Optional[float]) -> float: if max_budget is None: return float("inf") @@ -3174,9 +2849,7 @@ class PrometheusLogger(CustomLogger): try: page = 1 page_size = 50 - data, total_count = await data_fetch_function( - page_size=page_size, page=page - ) + data, total_count = await data_fetch_function(page_size=page_size, page=page) if total_count is None: total_count = len(data) @@ -3193,9 +2866,7 @@ class PrometheusLogger(CustomLogger): await set_metrics_function(data) except Exception as e: - verbose_logger.exception( - f"Error initializing {data_type} budget metrics: {str(e)}" - ) + verbose_logger.exception(f"Error initializing {data_type} budget metrics: {str(e)}") async def _initialize_team_budget_metrics(self): """ @@ -3207,17 +2878,11 @@ class PrometheusLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_logger.debug( - "Prometheus: skipping team metrics initialization, DB not initialized" - ) + verbose_logger.debug("Prometheus: skipping team metrics initialization, DB not initialized") return - async def fetch_teams( - page_size: int, page: int - ) -> Tuple[List[LiteLLM_TeamTable], Optional[int]]: - teams, total_count = await get_paginated_teams( - prisma_client=prisma_client, page_size=page_size, page=page - ) + async def fetch_teams(page_size: int, page: int) -> Tuple[List[LiteLLM_TeamTable], Optional[int]]: + teams, total_count = await get_paginated_teams(prisma_client=prisma_client, page_size=page_size, page=page) if total_count is None: total_count = len(teams) return teams, total_count @@ -3239,9 +2904,7 @@ class PrometheusLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_logger.debug( - "Prometheus: skipping key metrics initialization, DB not initialized" - ) + verbose_logger.debug("Prometheus: skipping key metrics initialization, DB not initialized") return async def fetch_keys( @@ -3281,14 +2944,10 @@ class PrometheusLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_logger.debug( - "Prometheus: skipping user metrics initialization, DB not initialized" - ) + verbose_logger.debug("Prometheus: skipping user metrics initialization, DB not initialized") return - async def fetch_users( - page_size: int, page: int - ) -> Tuple[List[LiteLLM_UserTable], Optional[int]]: + async def fetch_users(page_size: int, page: int) -> Tuple[List[LiteLLM_UserTable], Optional[int]]: skip = (page - 1) * page_size users = await UserRepository(prisma_client).table.find_many( skip=skip, @@ -3311,9 +2970,7 @@ class PrometheusLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_logger.debug( - "Prometheus: skipping org metrics initialization, DB not initialized" - ) + verbose_logger.debug("Prometheus: skipping org metrics initialization, DB not initialized") return async def fetch_orgs(page_size: int, page: int) -> Tuple[list, Optional[int]]: @@ -3350,15 +3007,11 @@ class PrometheusLogger(CustomLogger): # if using redis, ensure only one pod emits the metrics at a time if pod_lock_manager and pod_lock_manager.redis_cache: - if await pod_lock_manager.acquire_lock( - cronjob_id=PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME - ): + if await pod_lock_manager.acquire_lock(cronjob_id=PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME): try: await self._initialize_remaining_budget_metrics() finally: - await pod_lock_manager.release_lock( - cronjob_id=PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME - ) + await pod_lock_manager.release_lock(cronjob_id=PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME) else: # if not using redis, initialize the metrics directly await self._initialize_remaining_budget_metrics() @@ -3385,33 +3038,23 @@ class PrometheusLogger(CustomLogger): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - verbose_logger.debug( - "Prometheus: skipping user/team count metrics initialization, DB not initialized" - ) + verbose_logger.debug("Prometheus: skipping user/team count metrics initialization, DB not initialized") return try: # Get total user count total_users = await UserRepository(prisma_client).table.count() self.litellm_total_users_metric.set(total_users) - verbose_logger.debug( - f"Prometheus: set litellm_total_users to {total_users}" - ) + verbose_logger.debug(f"Prometheus: set litellm_total_users to {total_users}") # Get total team count total_teams = await TeamRepository(prisma_client).table.count() self.litellm_teams_count_metric.set(total_teams) - verbose_logger.debug( - f"Prometheus: set litellm_teams_count to {total_teams}" - ) + verbose_logger.debug(f"Prometheus: set litellm_teams_count to {total_teams}") except Exception as e: - verbose_logger.exception( - f"Error initializing user/team count metrics: {str(e)}" - ) + verbose_logger.exception(f"Error initializing user/team count metrics: {str(e)}") - async def _set_key_list_budget_metrics( - self, keys: List[Union[str, UserAPIKeyAuth]] - ): + async def _set_key_list_budget_metrics(self, keys: List[Union[str, UserAPIKeyAuth]]): """Helper function to set budget metrics for a list of keys""" for key in keys: if isinstance(key, UserAPIKeyAuth): @@ -3436,11 +3079,7 @@ class PrometheusLogger(CustomLogger): org_alias=org.organization_alias or "", spend=org.spend or 0.0, max_budget=budget_table.max_budget if budget_table else None, - budget_reset_at=( - getattr(budget_table, "budget_reset_at", None) - if budget_table - else None - ), + budget_reset_at=(getattr(budget_table, "budget_reset_at", None) if budget_table else None), ) async def _set_team_budget_metrics_after_api_request( @@ -3501,9 +3140,7 @@ class PrometheusLogger(CustomLogger): user_api_key_cache=user_api_key_cache, ) except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting team info: {str(e)}" - ) + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting team info: {str(e)}") return team_object if team_info: @@ -3530,9 +3167,7 @@ class PrometheusLogger(CustomLogger): ) _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_team_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_team_budget_metric"), enum_values=enum_values, ) self.litellm_remaining_team_budget_metric.labels(**_labels).set( @@ -3544,9 +3179,7 @@ class PrometheusLogger(CustomLogger): if team.max_budget is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_team_max_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_team_max_budget_metric"), enum_values=enum_values, ) self.litellm_team_max_budget_metric.labels(**_labels).set(team.max_budget) @@ -3559,9 +3192,7 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, ) self.litellm_team_budget_remaining_hours_metric.labels(**_labels).set( - self._get_remaining_hours_for_budget_reset( - budget_reset_at=team.budget_reset_at - ) + self._get_remaining_hours_for_budget_reset(budget_reset_at=team.budget_reset_at) ) def set_team_members_metric(self, team: LiteLLM_TeamTable) -> None: @@ -3571,14 +3202,10 @@ class PrometheusLogger(CustomLogger): team_alias=team.team_alias or "", ) _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_team_members_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_team_members_metric"), enum_values=enum_values, ) - self.litellm_team_members_metric.labels(**_labels).set( - len(team.members_with_roles) - ) + self.litellm_team_members_metric.labels(**_labels).set(len(team.members_with_roles)) async def _set_org_budget_metrics_after_api_request( self, @@ -3608,9 +3235,7 @@ class PrometheusLogger(CustomLogger): include_budget_table=True, ) except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting org info: {str(e)}" - ) + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting org info: {str(e)}") return if org_info is None: @@ -3620,9 +3245,7 @@ class PrometheusLogger(CustomLogger): _total_org_spend = (org_info.spend or 0.0) + response_cost budget_table = org_info.litellm_budget_table max_budget = budget_table.max_budget if budget_table else None - budget_reset_at = ( - getattr(budget_table, "budget_reset_at", None) if budget_table else None - ) + budget_reset_at = getattr(budget_table, "budget_reset_at", None) if budget_table else None self._set_org_budget_metrics( org_id=org_id, @@ -3653,9 +3276,7 @@ class PrometheusLogger(CustomLogger): ) _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_org_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_org_budget_metric"), enum_values=enum_values, ) self.litellm_remaining_org_budget_metric.labels(**_labels).set( @@ -3667,9 +3288,7 @@ class PrometheusLogger(CustomLogger): if max_budget is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_org_max_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_org_max_budget_metric"), enum_values=enum_values, ) self.litellm_org_max_budget_metric.labels(**_labels).set(max_budget) @@ -3682,9 +3301,7 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, ) self.litellm_org_budget_remaining_hours_metric.labels(**_labels).set( - self._get_remaining_hours_for_budget_reset( - budget_reset_at=budget_reset_at - ) + self._get_remaining_hours_for_budget_reset(budget_reset_at=budget_reset_at) ) def _set_key_budget_metrics(self, user_api_key_dict: UserAPIKeyAuth): @@ -3700,9 +3317,7 @@ class PrometheusLogger(CustomLogger): api_key_alias=user_api_key_dict.key_alias or "", ) _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_api_key_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_api_key_budget_metric"), enum_values=enum_values, ) self.litellm_remaining_api_key_budget_metric.labels(**_labels).set( @@ -3714,20 +3329,14 @@ class PrometheusLogger(CustomLogger): if user_api_key_dict.max_budget is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_api_key_max_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_api_key_max_budget_metric"), enum_values=enum_values, ) - self.litellm_api_key_max_budget_metric.labels(**_labels).set( - user_api_key_dict.max_budget - ) + self.litellm_api_key_max_budget_metric.labels(**_labels).set(user_api_key_dict.max_budget) if user_api_key_dict.budget_reset_at is not None: self.litellm_api_key_budget_remaining_hours_metric.labels(**_labels).set( - self._get_remaining_hours_for_budget_reset( - budget_reset_at=user_api_key_dict.budget_reset_at - ) + self._get_remaining_hours_for_budget_reset(budget_reset_at=user_api_key_dict.budget_reset_at) ) async def _set_api_key_budget_metrics_after_api_request( @@ -3779,9 +3388,7 @@ class PrometheusLogger(CustomLogger): if key_object: user_api_key_dict.budget_reset_at = key_object.budget_reset_at except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting key info: {str(e)}" - ) + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting key info: {str(e)}") return user_api_key_dict @@ -3843,9 +3450,7 @@ class PrometheusLogger(CustomLogger): check_db_only=False, ) except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting user info: {str(e)}" - ) + verbose_logger.debug(f"[Non-Blocking] Prometheus: Error getting user info: {str(e)}") return user_object if user_info: @@ -3877,9 +3482,7 @@ class PrometheusLogger(CustomLogger): ) _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_user_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_user_budget_metric"), enum_values=enum_values, ) self.litellm_remaining_user_budget_metric.labels(**_labels).set( @@ -3891,9 +3494,7 @@ class PrometheusLogger(CustomLogger): if user.max_budget is not None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_user_max_budget_metric" - ), + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_user_max_budget_metric"), enum_values=enum_values, ) self.litellm_user_max_budget_metric.labels(**_labels).set(user.max_budget) @@ -3906,18 +3507,14 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, ) self.litellm_user_budget_remaining_hours_metric.labels(**_labels).set( - self._get_remaining_hours_for_budget_reset( - budget_reset_at=user.budget_reset_at - ) + self._get_remaining_hours_for_budget_reset(budget_reset_at=user.budget_reset_at) ) def _get_remaining_hours_for_budget_reset(self, budget_reset_at: datetime) -> float: """ Get remaining hours for budget reset """ - return ( - budget_reset_at - datetime.now(budget_reset_at.tzinfo) - ).total_seconds() / 3600 + return (budget_reset_at - datetime.now(budget_reset_at.tzinfo)).total_seconds() / 3600 def _safe_duration_seconds( self, @@ -3943,10 +3540,8 @@ class PrometheusLogger(CustomLogger): """ from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES - prometheus_loggers: List[CustomLogger] = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=PrometheusLogger - ) + prometheus_loggers: List[CustomLogger] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=PrometheusLogger ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers)) @@ -3991,9 +3586,7 @@ class PrometheusLogger(CustomLogger): # Mount the metrics app to the app app.mount("/metrics", metrics_app) - verbose_proxy_logger.debug( - "Starting Prometheus Metrics on /metrics (no authentication)" - ) + verbose_proxy_logger.debug("Starting Prometheus Metrics on /metrics (no authentication)") def _prometheus_labels_from_context( @@ -4001,15 +3594,11 @@ def _prometheus_labels_from_context( ctx: PrometheusLabelFactoryContext, ) -> Dict[str, Optional[str]]: filtered_labels: Dict[str, Optional[str]] = { - label: ctx._sanitized_enum[label] - for label in supported_enum_labels - if label in ctx._sanitized_enum + label: ctx._sanitized_enum[label] for label in supported_enum_labels if label in ctx._sanitized_enum } if UserAPIKeyLabelNames.END_USER.value in filtered_labels: - filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ( - ctx.get_resolved_end_user() - ) + filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user() for sk, val in ctx._custom_by_sanitized_key.items(): if sk in supported_enum_labels: @@ -4043,9 +3632,7 @@ def prometheus_label_factory( """ if label_context is not None: if label_context.enum_values is not enum_values: - raise ValueError( - "label_context.enum_values must be the same object as enum_values" - ) + raise ValueError("label_context.enum_values must be the same object as enum_values") return _prometheus_labels_from_context(supported_enum_labels, label_context) # Extract dictionary from Pydantic object @@ -4134,25 +3721,17 @@ def _get_combined_custom_metadata_from_standard_logging_payload( return {} requester_metadata = standard_logging_metadata.get("requester_metadata") - user_api_key_auth_metadata = standard_logging_metadata.get( - "user_api_key_auth_metadata" - ) + user_api_key_auth_metadata = standard_logging_metadata.get("user_api_key_auth_metadata") spend_logs_metadata = standard_logging_metadata.get("spend_logs_metadata") return { **(requester_metadata if isinstance(requester_metadata, dict) else {}), - **( - user_api_key_auth_metadata - if isinstance(user_api_key_auth_metadata, dict) - else {} - ), + **(user_api_key_auth_metadata if isinstance(user_api_key_auth_metadata, dict) else {}), **(spend_logs_metadata if isinstance(spend_logs_metadata, dict) else {}), } -def _tag_matches_wildcard_configured_pattern( - tags: Sequence[str], configured_tag: str -) -> bool: +def _tag_matches_wildcard_configured_pattern(tags: Sequence[str], configured_tag: str) -> bool: """ Check if any of the request tags matches a wildcard configured pattern @@ -4222,9 +3801,7 @@ def get_custom_labels_from_tags(tags: Sequence[str]) -> Dict[str, str]: continue # Use PatternMatchRouter for wildcard pattern matching - if "*" in configured_tag and _tag_matches_wildcard_configured_pattern( - tags=tags, configured_tag=configured_tag - ): + if "*" in configured_tag and _tag_matches_wildcard_configured_pattern(tags=tags, configured_tag=configured_tag): result[label_name] = "true" continue diff --git a/litellm/integrations/prometheus_helpers/__init__.py b/litellm/integrations/prometheus_helpers/__init__.py index 784ab524dd5..7de072ecd03 100644 --- a/litellm/integrations/prometheus_helpers/__init__.py +++ b/litellm/integrations/prometheus_helpers/__init__.py @@ -57,9 +57,7 @@ class PrometheusLabelFactoryContext: if enum_values.custom_metadata_labels is not None: for key, value in enum_values.custom_metadata_labels.items(): sk = _sanitize_prometheus_label_name(key) - self._custom_by_sanitized_key[sk] = _sanitize_prometheus_label_value( - value - ) + self._custom_by_sanitized_key[sk] = _sanitize_prometheus_label_value(value) self._tag_labels: Dict[str, Optional[str]] = {} if enum_values.tags is not None: # Late import avoids circular import: ``prometheus`` imports this module. diff --git a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py index d834ae20142..61b4d5ab96e 100644 --- a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py +++ b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py @@ -86,9 +86,7 @@ class BoundedPrometheusSeriesTracker: series.pop(label_values, None) @staticmethod - def _remove_metric_child( - metric: Any, label_values: tuple[Optional[str], ...] - ) -> bool: + def _remove_metric_child(metric: Any, label_values: tuple[Optional[str], ...]) -> bool: """ Remove the Prometheus child for ``label_values`` and report whether the tracker should commit the matching state change. diff --git a/litellm/integrations/prometheus_helpers/prometheus_api.py b/litellm/integrations/prometheus_helpers/prometheus_api.py index 0901d7b6801..038788f0522 100644 --- a/litellm/integrations/prometheus_helpers/prometheus_api.py +++ b/litellm/integrations/prometheus_helpers/prometheus_api.py @@ -16,9 +16,7 @@ from litellm.llms.custom_httpx.http_handler import ( PROMETHEUS_URL: Optional[str] = get_secret("PROMETHEUS_URL") # type: ignore PROMETHEUS_SELECTED_INSTANCE: Optional[str] = get_secret("PROMETHEUS_SELECTED_INSTANCE") # type: ignore -async_http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback -) +async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) async def get_metric_from_prometheus( @@ -26,9 +24,7 @@ async def get_metric_from_prometheus( ): # Get the start of the current day in Unix timestamp if PROMETHEUS_URL is None: - raise ValueError( - "PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env" - ) + raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env") query = f"{metric_name}[24h]" now = int(time.time()) @@ -111,9 +107,7 @@ async def get_daily_spend_from_prometheus(api_key: Optional[str]): ...] """ if PROMETHEUS_URL is None: - raise ValueError( - "PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env" - ) + raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env") # Calculate the start and end dates for the last 30 days end_date = datetime.utcnow() @@ -129,11 +123,7 @@ async def get_daily_spend_from_prometheus(api_key: Optional[str]): query = "sum(delta(litellm_spend_metric_total[1d]))" else: quoted_api_key = _quote_promql_string_literal(api_key) - query = ( - "sum(delta(litellm_spend_metric_total{" - f"hashed_api_key={quoted_api_key}" - "}[1d]))" - ) + query = f"sum(delta(litellm_spend_metric_total{{hashed_api_key={quoted_api_key}}}[1d]))" params = { "query": query, diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py index af8b1d0866e..db005aaffc5 100644 --- a/litellm/integrations/prometheus_services.py +++ b/litellm/integrations/prometheus_services.py @@ -32,16 +32,10 @@ class PrometheusServicesLogger: from prometheus_client import REGISTRY, Counter, Gauge, Histogram from prometheus_client.gc_collector import Collector except ImportError: - raise Exception( - "Missing prometheus_client. Run `pip install prometheus-client`" - ) + raise Exception("Missing prometheus_client. Run `pip install prometheus-client`") _custom_buckets = litellm.prometheus_latency_buckets - self.latency_buckets = ( - tuple(_custom_buckets) - if _custom_buckets is not None - else LATENCY_BUCKETS - ) + self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS self.Histogram = Histogram self.Counter = Counter @@ -50,9 +44,7 @@ class PrometheusServicesLogger: verbose_logger.debug("in init prometheus services metrics") - self.payload_to_prometheus_map: Dict[ - str, List[Union[Histogram, Counter, Gauge, Collector]] - ] = {} + self.payload_to_prometheus_map: Dict[str, List[Union[Histogram, Counter, Gauge, Collector]]] = {} for service in ServiceTypes: service_metrics: List[Union[Histogram, Counter, Gauge, Collector]] = [] @@ -61,9 +53,7 @@ class PrometheusServicesLogger: # Initialize only the configured metrics for each service if ServiceMetrics.HISTOGRAM in metrics_to_initialize: - histogram = self.create_histogram( - service.value, type_of_request="latency" - ) + histogram = self.create_histogram(service.value, type_of_request="latency") if histogram: service_metrics.append(histogram) @@ -75,9 +65,7 @@ class PrometheusServicesLogger: ) if counter_failed_request: service_metrics.append(counter_failed_request) - counter_total_requests = self.create_counter( - service.value, type_of_request="total_requests" - ) + counter_total_requests = self.create_counter(service.value, type_of_request="total_requests") if counter_total_requests: service_metrics.append(counter_total_requests) @@ -99,9 +87,7 @@ class PrometheusServicesLogger: print_verbose(f"Got exception on init prometheus client {str(e)}") raise e - def _get_service_metrics_initialize( - self, service: ServiceTypes - ) -> List[ServiceMetrics]: + def _get_service_metrics_initialize(self, service: ServiceTypes) -> List[ServiceMetrics]: DEFAULT_METRICS = [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] if service not in DEFAULT_SERVICE_CONFIGS: return DEFAULT_METRICS @@ -146,9 +132,7 @@ class PrometheusServicesLogger: is_registered = self.is_metric_registered(metric_name) if is_registered: return self._get_metric(metric_name) - return self.Gauge( - metric_name, "Gauge for {} service".format(service), labelnames=[service] - ) + return self.Gauge(metric_name, "Gauge for {} service".format(service), labelnames=[service]) def create_counter( self, diff --git a/litellm/integrations/prompt_layer.py b/litellm/integrations/prompt_layer.py index 190b995fa4e..52209b2953f 100644 --- a/litellm/integrations/prompt_layer.py +++ b/litellm/integrations/prompt_layer.py @@ -33,11 +33,7 @@ class PromptLayerLogger: tags = kwargs["litellm_params"]["metadata"]["pl_tags"] # Remove "pl_tags" from metadata - metadata = { - k: v - for k, v in kwargs["litellm_params"]["metadata"].items() - if k != "pl_tags" - } + metadata = {k: v for k, v in kwargs["litellm_params"]["metadata"].items() if k != "pl_tags"} print_verbose( f"Prompt Layer Logging - Enters logging function for model kwargs: {new_kwargs}\n, response: {response_obj}" @@ -68,9 +64,7 @@ class PromptLayerLogger: if not request_response.json().get("success", False): raise Exception("Promptlayer did not successfully log the response!") - print_verbose( - f"Prompt Layer Logging: success - final response object: {request_response.text}" - ) + print_verbose(f"Prompt Layer Logging: success - final response object: {request_response.text}") if "request_id" in response_json: if metadata: @@ -82,9 +76,7 @@ class PromptLayerLogger: "metadata": metadata, }, ) - print_verbose( - f"Prompt Layer Logging: success - metadata post response object: {response.text}" - ) + print_verbose(f"Prompt Layer Logging: success - metadata post response object: {response.text}") except Exception: print_verbose(f"error: Prompt Layer Error - {traceback.format_exc()}") diff --git a/litellm/integrations/prompt_management_base.py b/litellm/integrations/prompt_management_base.py index 9c626aea849..6d77e959e2d 100644 --- a/litellm/integrations/prompt_management_base.py +++ b/litellm/integrations/prompt_management_base.py @@ -119,9 +119,7 @@ class PromptManagementBase(ABC): compiled_prompt_client["completed_messages"] = messages return compiled_prompt_client - def _get_model_from_prompt( - self, prompt_management_client: PromptManagementClient, model: str - ) -> str: + def _get_model_from_prompt(self, prompt_management_client: PromptManagementClient, model: str) -> str: if prompt_management_client["prompt_template_model"] is not None: return prompt_management_client["prompt_template_model"] else: @@ -138,23 +136,15 @@ class PromptManagementBase(ABC): ): completed_messages = prompt_template["completed_messages"] or messages - prompt_template_optional_params = ( - prompt_template["prompt_template_optional_params"] or {} - ) + prompt_template_optional_params = prompt_template["prompt_template_optional_params"] or {} updated_non_default_params = { **non_default_params, - **( - prompt_template_optional_params - if not ignore_prompt_manager_optional_params - else {} - ), + **(prompt_template_optional_params if not ignore_prompt_manager_optional_params else {}), } if not ignore_prompt_manager_model: - model = self._get_model_from_prompt( - prompt_management_client=prompt_template, model=model - ) + model = self._get_model_from_prompt(prompt_management_client=prompt_template, model=model) else: model = model diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 922d8f71cf0..2b54a411ec7 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -83,29 +83,20 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): parsed_rate = float(rbrk_sampling_rate.strip()) self.sampling_rate = max(0.0, min(1.0, parsed_rate)) if parsed_rate != self.sampling_rate: - verbose_logger.warning( - f"RUBRIK_SAMPLING_RATE={parsed_rate} clamped to " - f"{self.sampling_rate}" - ) + verbose_logger.warning(f"RUBRIK_SAMPLING_RATE={parsed_rate} clamped to {self.sampling_rate}") except ValueError: - verbose_logger.warning( - f"Invalid RUBRIK_SAMPLING_RATE: {rbrk_sampling_rate!r}, using 1.0" - ) + verbose_logger.warning(f"Invalid RUBRIK_SAMPLING_RATE: {rbrk_sampling_rate!r}, using 1.0") self.key = api_key or os.getenv("RUBRIK_API_KEY") if not self.key: - verbose_logger.warning( - "Rubrik: No API key configured. Requests will be unauthenticated." - ) + verbose_logger.warning("Rubrik: No API key configured. Requests will be unauthenticated.") _batch_size = os.getenv("RUBRIK_BATCH_SIZE") if _batch_size: try: self.batch_size = int(_batch_size) except ValueError: - verbose_logger.warning( - f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default" - ) + verbose_logger.warning(f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default") # Cap the in-memory retry queue so a Rubrik webhook outage cannot let # authenticated traffic accumulate prompt/response payloads until the @@ -118,18 +109,13 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") if _webhook_url is None: - raise ValueError( - "Rubrik webhook URL not configured. " - "Set RUBRIK_WEBHOOK_URL or pass api_base." - ) + raise ValueError("Rubrik webhook URL not configured. Set RUBRIK_WEBHOOK_URL or pass api_base.") _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") self.tool_blocking_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_TOOL_BLOCKING}" self.logging_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" - self.async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.tool_blocking_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback, @@ -143,9 +129,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # Periodic flush is started lazily on the first log event so that # low-traffic deployments still get their batches drained even when the # logger is instantiated outside a running event loop (sync init). - self._flush_task: Optional[asyncio.Task[Any]] = ( - self._start_periodic_flush_task() - ) + self._flush_task: Optional[asyncio.Task[Any]] = self._start_periodic_flush_task() def _start_periodic_flush_task(self) -> Optional[asyncio.Task[Any]]: """Start the periodic flush task only when an event loop is already running.""" @@ -153,8 +137,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): loop = asyncio.get_running_loop() except RuntimeError: verbose_logger.debug( - "Rubrik logger init: no running event loop, " - "periodic flush will start on first log event." + "Rubrik logger init: no running event loop, periodic flush will start on first log event." ) return None return loop.create_task(self.periodic_flush()) @@ -197,9 +180,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return inputs try: - return await self._check_tool_calls( - inputs, tool_calls, request_data, logging_obj - ) + return await self._check_tool_calls(inputs, tool_calls, request_data, logging_obj) except ModifyResponseException: raise except _MalformedToolBlockingResponseError as e: @@ -218,8 +199,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return inputs except Exception as e: verbose_logger.error( - f"Tool blocking hook failed: {e}. " - "Returning original response unchanged.", + f"Tool blocking hook failed: {e}. Returning original response unchanged.", exc_info=True, ) return inputs @@ -234,26 +214,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): """Send tool calls to blocking service, raise if any are blocked.""" message_tool_calls = self._normalize_tool_calls(tool_calls) - call_details = ( - getattr(logging_obj, "model_call_details", {}) if logging_obj else {} - ) + call_details = getattr(logging_obj, "model_call_details", {}) if logging_obj else {} response = request_data.get("response") request_id = getattr(response, "id", None) if response else None if logging_obj and not call_details: verbose_logger.warning( - "Rubrik: logging_obj present but model_call_details is empty " - "-- request context will be missing" + "Rubrik: logging_obj present but model_call_details is empty -- request context will be missing" ) response_data = self._build_tool_call_payload(message_tool_calls, request_id) req_data = self._extract_request_data(call_details) - service_response = await self._post_to_tool_blocking_service( - response_data, req_data - ) - blocked_explanation = self._extract_blocked_tools( - service_response, message_tool_calls - ) + service_response = await self._post_to_tool_blocking_service(response_data, req_data) + blocked_explanation = self._extract_blocked_tools(service_response, message_tool_calls) if blocked_explanation is not None: model = self._resolve_model(request_data, call_details) @@ -294,9 +267,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ) ) else: - raise TypeError( - f"Cannot normalize tool_call of type {type(tc).__name__}" - ) + raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}") return result @staticmethod @@ -316,9 +287,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "message": { "role": "assistant", "content": None, - "tool_calls": [ - tc.model_dump(exclude_none=True) for tc in tool_calls - ], + "tool_calls": [tc.model_dump(exclude_none=True) for tc in tool_calls], }, "finish_reason": "tool_calls", } @@ -347,16 +316,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): request ``body`` so proxy credentials are not exfiltrated.""" if not isinstance(proxy_server_request, dict): return proxy_server_request - return { - key: proxy_server_request[key] - for key in ("url", "method") - if key in proxy_server_request - } + return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request} @staticmethod - def _resolve_model( - request_data: dict[str, Any], call_details: dict[str, Any] - ) -> str: + def _resolve_model(request_data: dict[str, Any], call_details: dict[str, Any]) -> str: """Get the model name for the ModifyResponseException.""" response = request_data.get("response") if response and hasattr(response, "model"): @@ -365,21 +328,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Logging hooks --------------------------------------------------------- - async def _prepare_log_payload( - self, kwargs: dict, event_type: str - ) -> StandardLoggingPayload | None: + async def _prepare_log_payload(self, kwargs: dict, event_type: str) -> StandardLoggingPayload | None: """Shared logic for success and failure logging.""" if random.random() > self.sampling_rate: - verbose_logger.debug( - f"Skipping Rubrik {event_type} logging " - f"(sampling_rate={self.sampling_rate})" - ) + verbose_logger.debug(f"Skipping Rubrik {event_type} logging (sampling_rate={self.sampling_rate})") return None # Deep-copy so mutations don't affect other callbacks sharing this object - standard_logging_payload: StandardLoggingPayload = safe_deep_copy( - kwargs["standard_logging_object"] - ) + standard_logging_payload: StandardLoggingPayload = safe_deep_copy(kwargs["standard_logging_object"]) # For Anthropic /v1/messages requests, LiteLLM creates a separate # ModelResponse (with a generated chatcmpl-* id) for logging, which @@ -431,8 +387,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): await self.flush_queue() except Exception as e: verbose_logger.error( - f"Rubrik {event_type} logging hook failed: {e}. " - "Skipping logging for this event.", + f"Rubrik {event_type} logging hook failed: {e}. Skipping logging for this event.", exc_info=True, ) @@ -474,9 +429,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ) response.raise_for_status() except httpx.HTTPStatusError as e: - verbose_logger.exception( - f"Rubrik HTTP Error: {e.response.status_code} - {e.response.text}" - ) + verbose_logger.exception(f"Rubrik HTTP Error: {e.response.status_code} - {e.response.text}") raise except Exception: verbose_logger.exception("Rubrik Layer Error") @@ -494,9 +447,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return log_queue_snapshot = list(self.log_queue) - verbose_logger.debug( - "Rubrik: Flushing batch of %s events", len(log_queue_snapshot) - ) + verbose_logger.debug("Rubrik: Flushing batch of %s events", len(log_queue_snapshot)) await self._log_batch_to_rubrik( data=log_queue_snapshot, ) @@ -549,9 +500,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "request": request_data, "response": response_data, } - verbose_logger.debug( - f"Sending request to tool blocking service: {self.tool_blocking_endpoint}" - ) + verbose_logger.debug(f"Sending request to tool blocking service: {self.tool_blocking_endpoint}") http_response = await self.tool_blocking_client.post( self.tool_blocking_endpoint, json=envelope, @@ -577,24 +526,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): """ choices = service_response.get("choices", []) if not choices: - raise _MalformedToolBlockingResponseError( - "Tool blocking service returned empty response" - ) + raise _MalformedToolBlockingResponseError("Tool blocking service returned empty response") message = choices[0].get("message", {}) returned_tool_calls = message.get("tool_calls") or [] blocking_explanation = message.get("content", "") allowed_id_counts: Counter = Counter( - tc["id"] - for tc in returned_tool_calls - if isinstance(tc, dict) and tc.get("id") + tc["id"] for tc in returned_tool_calls if isinstance(tc, dict) and tc.get("id") ) required_id_counts: Counter = Counter(tc.id for tc in all_tool_calls if tc.id) all_allowed = len(returned_tool_calls) >= len(all_tool_calls) and all( - allowed_id_counts.get(tc_id, 0) >= count - for tc_id, count in required_id_counts.items() + allowed_id_counts.get(tc_id, 0) >= count for tc_id, count in required_id_counts.items() ) if all_allowed: diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 2e70b1d6519..53a982cd2c4 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -29,9 +29,7 @@ class S3Logger: import boto3 try: - verbose_logger.debug( - f"in init s3 logger - s3_callback_params {litellm.s3_callback_params}" - ) + verbose_logger.debug(f"in init s3 logger - s3_callback_params {litellm.s3_callback_params}") s3_use_team_prefix = False @@ -47,21 +45,13 @@ class S3Logger: s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl", True) s3_verify = litellm.s3_callback_params.get("s3_verify") s3_endpoint_url = litellm.s3_callback_params.get("s3_endpoint_url") - s3_aws_access_key_id = litellm.s3_callback_params.get( - "s3_aws_access_key_id" - ) - s3_aws_secret_access_key = litellm.s3_callback_params.get( - "s3_aws_secret_access_key" - ) - s3_aws_session_token = litellm.s3_callback_params.get( - "s3_aws_session_token" - ) + s3_aws_access_key_id = litellm.s3_callback_params.get("s3_aws_access_key_id") + s3_aws_secret_access_key = litellm.s3_callback_params.get("s3_aws_secret_access_key") + s3_aws_session_token = litellm.s3_callback_params.get("s3_aws_session_token") s3_config = litellm.s3_callback_params.get("s3_config") s3_path = litellm.s3_callback_params.get("s3_path") # done reading litellm.s3_callback_params - s3_use_team_prefix = bool( - litellm.s3_callback_params.get("s3_use_team_prefix", False) - ) + s3_use_team_prefix = bool(litellm.s3_callback_params.get("s3_use_team_prefix", False)) self.s3_use_team_prefix = s3_use_team_prefix self.bucket_name = s3_bucket_name self.s3_path = s3_path @@ -84,23 +74,17 @@ class S3Logger: print_verbose(f"Got exception on init s3 client {str(e)}") raise e - async def _async_log_event( - self, kwargs, response_obj, start_time, end_time, print_verbose - ): + async def _async_log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): self.log_event(kwargs, response_obj, start_time, end_time, print_verbose) def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): try: - verbose_logger.debug( - f"s3 Logging - Enters logging function for model {kwargs}" - ) + verbose_logger.debug(f"s3 Logging - Enters logging function for model {kwargs}") # construct payload to send to s3 # follows the same params as langfuse.py litellm_params = kwargs.get("litellm_params", {}) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None + metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None # Clean Metadata before logging - never log raw metadata # the raw metadata can contain circular references which leads to infinite recursion @@ -131,11 +115,7 @@ class S3Logger: team_alias = payload["metadata"].get("user_api_key_team_alias") team_alias_prefix = "" - if ( - litellm.enable_preview_features - and self.s3_use_team_prefix - and team_alias is not None - ): + if litellm.enable_preview_features and self.s3_use_team_prefix and team_alias is not None: team_alias_prefix = f"{team_alias}/" s3_file_name = litellm.utils.get_logging_id(start_time, payload) or "" @@ -147,11 +127,7 @@ class S3Logger: ) s3_object_download_filename = ( - "time-" - + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") - + "_" - + payload["id"] - + ".json" + "time-" + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") + "_" + payload["id"] + ".json" ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -186,11 +162,7 @@ def get_s3_object_key( s3_file_name: str, ) -> str: s3_object_key = ( - (s3_path.rstrip("/") + "/" if s3_path else "") - + prefix - + start_time.strftime("%Y-%m-%d") - + "/" - + s3_file_name + (s3_path.rstrip("/") + "/" if s3_path else "") + prefix + start_time.strftime("%Y-%m-%d") + "/" + s3_file_name ) # we need the s3 key to include the time, so we log cache hits too s3_object_key += ".json" return s3_object_key diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 4ed8a809a13..939289f96ea 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -61,8 +61,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): _masker = SensitiveDataMasker() if s3_callback_params_override is not None: verbose_logger.debug( - f"in init s3 logger (audit override) - " - f"{_masker.mask_dict(dict(s3_callback_params_override))}" + f"in init s3 logger (audit override) - {_masker.mask_dict(dict(s3_callback_params_override))}" ) else: verbose_logger.debug( @@ -98,9 +97,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # IMPORTANT # Create httpx client AFTER _init_s3_params so we have the correct s3_verify value - verbose_logger.debug( - f"s3_v2 logger creating async httpx client with s3_verify={self.s3_verify}" - ) + verbose_logger.debug(f"s3_v2 logger creating async httpx client with s3_verify={self.s3_verify}") self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback, params={"ssl_verify": self.s3_verify}, @@ -109,9 +106,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - verbose_logger.debug( - f"s3 flush interval: {s3_flush_interval}, s3 batch size: {s3_batch_size}" - ) + verbose_logger.debug(f"s3 flush interval: {s3_flush_interval}, s3 batch size: {s3_batch_size}") # Call CustomLogger's __init__ CustomBatchLogger.__init__( self, @@ -161,75 +156,42 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if params_source is None: params_source = litellm.s3_callback_params or {} params: dict = { - key: ( - litellm.get_secret(value) - if isinstance(value, str) and value.startswith("os.environ/") - else value - ) + key: (litellm.get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value) for key, value in params_source.items() } self.s3_bucket_name = params.get("s3_bucket_name") or s3_bucket_name self.s3_region_name = params.get("s3_region_name") or s3_region_name self.s3_api_version = params.get("s3_api_version") or s3_api_version - self.s3_use_ssl = ( - params.get("s3_use_ssl", True) - if params.get("s3_use_ssl") is not None - else s3_use_ssl - ) - self.s3_verify = ( - params.get("s3_verify") - if params.get("s3_verify") is not None - else s3_verify - ) + self.s3_use_ssl = params.get("s3_use_ssl", True) if params.get("s3_use_ssl") is not None else s3_use_ssl + self.s3_verify = params.get("s3_verify") if params.get("s3_verify") is not None else s3_verify self.s3_endpoint_url = params.get("s3_endpoint_url") or s3_endpoint_url - self.s3_aws_access_key_id = ( - params.get("s3_aws_access_key_id") or s3_aws_access_key_id - ) + self.s3_aws_access_key_id = params.get("s3_aws_access_key_id") or s3_aws_access_key_id - self.s3_aws_secret_access_key = ( - params.get("s3_aws_secret_access_key") or s3_aws_secret_access_key - ) + self.s3_aws_secret_access_key = params.get("s3_aws_secret_access_key") or s3_aws_secret_access_key - self.s3_aws_session_token = ( - params.get("s3_aws_session_token") or s3_aws_session_token - ) + self.s3_aws_session_token = params.get("s3_aws_session_token") or s3_aws_session_token - self.s3_aws_session_name = ( - params.get("s3_aws_session_name") or s3_aws_session_name - ) + self.s3_aws_session_name = params.get("s3_aws_session_name") or s3_aws_session_name - self.s3_aws_profile_name = ( - params.get("s3_aws_profile_name") or s3_aws_profile_name - ) + self.s3_aws_profile_name = params.get("s3_aws_profile_name") or s3_aws_profile_name self.s3_aws_role_name = params.get("s3_aws_role_name") or s3_aws_role_name - self.s3_aws_web_identity_token = ( - params.get("s3_aws_web_identity_token") or s3_aws_web_identity_token - ) + self.s3_aws_web_identity_token = params.get("s3_aws_web_identity_token") or s3_aws_web_identity_token - self.s3_aws_sts_endpoint = ( - params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint - ) + self.s3_aws_sts_endpoint = params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint self.s3_config = params.get("s3_config") or s3_config self.s3_path = params.get("s3_path") or s3_path - self.s3_use_team_prefix = ( - bool(params.get("s3_use_team_prefix", False)) or s3_use_team_prefix - ) + self.s3_use_team_prefix = bool(params.get("s3_use_team_prefix", False)) or s3_use_team_prefix - self.s3_use_key_prefix = ( - bool(params.get("s3_use_key_prefix", False)) or s3_use_key_prefix - ) + self.s3_use_key_prefix = bool(params.get("s3_use_key_prefix", False)) or s3_use_key_prefix - self.s3_strip_base64_files = ( - bool(params.get("s3_strip_base64_files", False)) or s3_strip_base64_files - ) + self.s3_strip_base64_files = bool(params.get("s3_strip_base64_files", False)) or s3_strip_base64_files self.s3_use_virtual_hosted_style = ( - bool(params.get("s3_use_virtual_hosted_style", False)) - or s3_use_virtual_hosted_style + bool(params.get("s3_use_virtual_hosted_style", False)) or s3_use_virtual_hosted_style ) return @@ -251,9 +213,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) pass - async def async_log_audit_log_event( - self, audit_log: StandardAuditLogPayload - ) -> None: + async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: """Batch audit logs and upload to S3 under audit_logs/ prefix.""" try: from datetime import timezone @@ -265,9 +225,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_path = s3_path.rstrip("/") + "/" if s3_path else "" s3_object_key = ( - f"{s3_path}audit_logs/" - f"{now.strftime('%Y-%m-%d')}/" - f"{now.strftime('%H-%M-%S')}_{audit_log_id}.json" + f"{s3_path}audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log_id}.json" ) element = s3BatchLoggingElement( @@ -285,9 +243,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): async def _async_log_event_base(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug( - f"s3 Logging - Enters logging function for model {kwargs}" - ) + verbose_logger.debug(f"s3 Logging - Enters logging function for model {kwargs}") s3_batch_logging_element = self.create_s3_batch_logging_element( start_time=start_time, @@ -303,9 +259,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) return - verbose_logger.debug( - "\ns3 Logger - Logging payload = %s", s3_batch_logging_element - ) + verbose_logger.debug("\ns3 Logger - Logging payload = %s", s3_batch_logging_element) self.log_queue.append(s3_batch_logging_element) verbose_logger.debug( @@ -317,9 +271,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): verbose_logger.exception(f"s3 Layer Error - {str(e)}") self.handle_callback_failure(callback_name="S3Logger") - async def async_upload_data_to_s3( - self, batch_logging_element: s3BatchLoggingElement - ): + async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement): try: import hashlib @@ -344,9 +296,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): aws_sts_endpoint=self.s3_aws_sts_endpoint, ) - verbose_logger.debug( - f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}" - ) + verbose_logger.debug(f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}") verbose_logger.debug(f"s3_v2 logger - s3_verify setting: {self.s3_verify}") # Prepare the URL @@ -355,24 +305,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_endpoint_url and self.s3_bucket_name: if self.s3_use_virtual_hosted_style: # Virtual-hosted-style: bucket.endpoint/key - endpoint_host = self.s3_endpoint_url.replace( - "https://", "" - ).replace("http://", "") - protocol = ( - "https://" - if self.s3_endpoint_url.startswith("https://") - else "http://" - ) + endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") + protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" else: # Path-style: endpoint/bucket/key - url = ( - self.s3_endpoint_url - + "/" - + self.s3_bucket_name - + "/" - + batch_logging_element.s3_object_key - ) + url = self.s3_endpoint_url + "/" + self.s3_bucket_name + "/" + batch_logging_element.s3_object_key # Convert JSON to string json_string = safe_dumps(batch_logging_element.payload) @@ -398,9 +336,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( - aws_region_name=self.s3_region_name - ) + aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) # Prepare the signed headers @@ -412,9 +348,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Make the request with retry for transient S3 errors (500/503) max_retries = 3 for attempt in range(max_retries): - response = await self.async_httpx_client.put( - request_url, data=json_string, headers=signed_headers - ) + response = await self.async_httpx_client.put(request_url, data=json_string, headers=signed_headers) if response.status_code in (500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( @@ -471,22 +405,16 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): return None if self.s3_strip_base64_files: - standard_logging_payload = self._strip_base64_from_messages_sync( - standard_logging_payload - ) + standard_logging_payload = self._strip_base64_from_messages_sync(standard_logging_payload) # Base prefix (default empty) prefix_components = [] if self.s3_use_team_prefix: - team_alias = standard_logging_payload.get("metadata", {}).get( - "user_api_key_team_alias", None - ) + team_alias = standard_logging_payload.get("metadata", {}).get("user_api_key_team_alias", None) if team_alias: prefix_components.append(team_alias) if self.s3_use_key_prefix: - user_api_key_alias = standard_logging_payload.get("metadata", {}).get( - "user_api_key_alias", None - ) + user_api_key_alias = standard_logging_payload.get("metadata", {}).get("user_api_key_alias", None) if user_api_key_alias: prefix_components.append(user_api_key_alias) @@ -495,9 +423,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if prefix_path: prefix_path += "/" - s3_file_name = ( - litellm.utils.get_logging_id(start_time, standard_logging_payload) or "" - ) + s3_file_name = litellm.utils.get_logging_id(start_time, standard_logging_payload) or "" verbose_logger.debug( f"Creating s3 file with prefix_components={prefix_components},prefix_path={prefix_path} and {s3_file_name}" ) @@ -509,7 +435,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) verbose_logger.debug(f"s3_object_key={s3_object_key}") - s3_object_download_filename = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json" + s3_object_download_filename = ( + f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json" + ) return s3BatchLoggingElement( payload=dict(standard_logging_payload), @@ -528,9 +456,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") try: - verbose_logger.debug( - f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}" - ) + verbose_logger.debug(f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}") credentials: Credentials = self.get_credentials( aws_access_key_id=self.s3_aws_access_key_id, aws_secret_access_key=self.s3_aws_secret_access_key, @@ -544,24 +470,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_endpoint_url and self.s3_bucket_name: if self.s3_use_virtual_hosted_style: # Virtual-hosted-style: bucket.endpoint/key - endpoint_host = self.s3_endpoint_url.replace( - "https://", "" - ).replace("http://", "") - protocol = ( - "https://" - if self.s3_endpoint_url.startswith("https://") - else "http://" - ) + endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") + protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" else: # Path-style: endpoint/bucket/key - url = ( - self.s3_endpoint_url - + "/" - + self.s3_bucket_name - + "/" - + batch_logging_element.s3_object_key - ) + url = self.s3_endpoint_url + "/" + self.s3_bucket_name + "/" + batch_logging_element.s3_object_key # Convert JSON to string json_string = safe_dumps(batch_logging_element.payload) @@ -587,9 +501,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( - aws_region_name=self.s3_region_name - ) + aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) # Prepare the signed headers @@ -599,18 +511,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): request_url = prepped.url or url httpx_client = _get_httpx_client( - params=( - {"ssl_verify": self.s3_verify} - if self.s3_verify is not None - else None - ) + params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None) ) # Make the request with retry for transient S3 errors (500/503) max_retries = 3 for attempt in range(max_retries): - response = httpx_client.put( - request_url, data=json_string, headers=signed_headers - ) + response = httpx_client.put(request_url, data=json_string, headers=signed_headers) if response.status_code in (500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( @@ -662,9 +568,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): aws_sts_endpoint=self.s3_aws_sts_endpoint, ) - verbose_logger.debug( - f"s3_v2 logger - downloading data from s3 - {s3_object_key}" - ) + verbose_logger.debug(f"s3_v2 logger - downloading data from s3 - {s3_object_key}") # Prepare the URL url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{s3_object_key}" @@ -672,24 +576,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_endpoint_url and self.s3_bucket_name: if self.s3_use_virtual_hosted_style: # Virtual-hosted-style: bucket.endpoint/key - endpoint_host = self.s3_endpoint_url.replace( - "https://", "" - ).replace("http://", "") - protocol = ( - "https://" - if self.s3_endpoint_url.startswith("https://") - else "http://" - ) + endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") + protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{s3_object_key}" else: # Path-style: endpoint/bucket/key - url = ( - self.s3_endpoint_url - + "/" - + self.s3_bucket_name - + "/" - + s3_object_key - ) + url = self.s3_endpoint_url + "/" + self.s3_bucket_name + "/" + s3_object_key # Prepare the request for GET operation # For GET requests, we need x-amz-content-sha256 with hash of empty string @@ -712,14 +604,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): signed_headers = dict(aws_request.headers.items()) request_url = prepped.url or url - response = await self.async_httpx_client.get( - request_url, headers=signed_headers - ) + response = await self.async_httpx_client.get(request_url, headers=signed_headers) if response.status_code != 200: - verbose_logger.exception( - "S3 object not found, saw response=", response.text - ) + verbose_logger.exception("S3 object not found, saw response=", response.text) return None # Parse JSON response @@ -750,7 +638,5 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): downloaded_object = await self._download_object_from_s3(object_key) return downloaded_object except Exception as e: - verbose_logger.exception( - f"Error retrieving object {object_key} from cold storage: {str(e)}" - ) + verbose_logger.exception(f"Error retrieving object {object_key} from cold storage: {str(e)}") return None diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 6cbd2c7974f..8c0b06df888 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -69,9 +69,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): **kwargs, ) -> None: try: - verbose_logger.debug( - f"in init sqs logger - sqs_callback_params {litellm.aws_sqs_callback_params}" - ) + verbose_logger.debug(f"in init sqs logger - sqs_callback_params {litellm.aws_sqs_callback_params}") self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback, @@ -103,9 +101,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - verbose_logger.debug( - f"sqs flush interval: {sqs_flush_interval}, sqs batch size: {sqs_batch_size}" - ) + verbose_logger.debug(f"sqs flush interval: {sqs_flush_interval}, sqs batch size: {sqs_batch_size}") CustomBatchLogger.__init__( self, @@ -150,109 +146,66 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): if isinstance(value, str) and value.startswith("os.environ/"): litellm.aws_sqs_callback_params[key] = litellm.get_secret(value) - self.sqs_queue_url = ( - litellm.aws_sqs_callback_params.get("sqs_queue_url") or sqs_queue_url - ) - self.sqs_region_name = ( - litellm.aws_sqs_callback_params.get("sqs_region_name") or sqs_region_name - ) - self.sqs_api_version = ( - litellm.aws_sqs_callback_params.get("sqs_api_version") or sqs_api_version - ) - self.sqs_use_ssl = ( - litellm.aws_sqs_callback_params.get("sqs_use_ssl", True) or sqs_use_ssl - ) - self.sqs_verify = ( - litellm.aws_sqs_callback_params.get("sqs_verify") or sqs_verify - ) - self.sqs_endpoint_url = ( - litellm.aws_sqs_callback_params.get("sqs_endpoint_url") or sqs_endpoint_url - ) + self.sqs_queue_url = litellm.aws_sqs_callback_params.get("sqs_queue_url") or sqs_queue_url + self.sqs_region_name = litellm.aws_sqs_callback_params.get("sqs_region_name") or sqs_region_name + self.sqs_api_version = litellm.aws_sqs_callback_params.get("sqs_api_version") or sqs_api_version + self.sqs_use_ssl = litellm.aws_sqs_callback_params.get("sqs_use_ssl", True) or sqs_use_ssl + self.sqs_verify = litellm.aws_sqs_callback_params.get("sqs_verify") or sqs_verify + self.sqs_endpoint_url = litellm.aws_sqs_callback_params.get("sqs_endpoint_url") or sqs_endpoint_url self.sqs_aws_access_key_id = ( - litellm.aws_sqs_callback_params.get("sqs_aws_access_key_id") - or sqs_aws_access_key_id + litellm.aws_sqs_callback_params.get("sqs_aws_access_key_id") or sqs_aws_access_key_id ) self.sqs_aws_secret_access_key = ( - litellm.aws_sqs_callback_params.get("sqs_aws_secret_access_key") - or sqs_aws_secret_access_key + litellm.aws_sqs_callback_params.get("sqs_aws_secret_access_key") or sqs_aws_secret_access_key ) self.sqs_aws_session_token = ( - litellm.aws_sqs_callback_params.get("sqs_aws_session_token") - or sqs_aws_session_token + litellm.aws_sqs_callback_params.get("sqs_aws_session_token") or sqs_aws_session_token ) - self.sqs_aws_session_name = ( - litellm.aws_sqs_callback_params.get("sqs_aws_session_name") - or sqs_aws_session_name - ) + self.sqs_aws_session_name = litellm.aws_sqs_callback_params.get("sqs_aws_session_name") or sqs_aws_session_name - self.sqs_aws_profile_name = ( - litellm.aws_sqs_callback_params.get("sqs_aws_profile_name") - or sqs_aws_profile_name - ) + self.sqs_aws_profile_name = litellm.aws_sqs_callback_params.get("sqs_aws_profile_name") or sqs_aws_profile_name - self.sqs_aws_role_name = ( - litellm.aws_sqs_callback_params.get("sqs_aws_role_name") - or sqs_aws_role_name - ) + self.sqs_aws_role_name = litellm.aws_sqs_callback_params.get("sqs_aws_role_name") or sqs_aws_role_name self.sqs_aws_web_identity_token = ( - litellm.aws_sqs_callback_params.get("sqs_aws_web_identity_token") - or sqs_aws_web_identity_token + litellm.aws_sqs_callback_params.get("sqs_aws_web_identity_token") or sqs_aws_web_identity_token ) - self.sqs_aws_sts_endpoint = ( - litellm.aws_sqs_callback_params.get("sqs_aws_sts_endpoint") - or sqs_aws_sts_endpoint - ) + self.sqs_aws_sts_endpoint = litellm.aws_sqs_callback_params.get("sqs_aws_sts_endpoint") or sqs_aws_sts_endpoint self.sqs_strip_base64_files = ( - litellm.aws_sqs_callback_params.get("sqs_strip_base64_files", False) - or sqs_strip_base64_files + litellm.aws_sqs_callback_params.get("sqs_strip_base64_files", False) or sqs_strip_base64_files ) self.sqs_aws_use_application_level_encryption = ( - litellm.aws_sqs_callback_params.get( - "sqs_aws_use_application_level_encryption", False - ) + litellm.aws_sqs_callback_params.get("sqs_aws_use_application_level_encryption", False) or sqs_aws_use_application_level_encryption ) self.sqs_app_encryption_key_b64 = ( - litellm.aws_sqs_callback_params.get("sqs_app_encryption_key_b64") - or sqs_app_encryption_key_b64 + litellm.aws_sqs_callback_params.get("sqs_app_encryption_key_b64") or sqs_app_encryption_key_b64 ) self.sqs_app_encryption_aad = ( - litellm.aws_sqs_callback_params.get("sqs_app_encryption_aad") - or sqs_app_encryption_aad + litellm.aws_sqs_callback_params.get("sqs_app_encryption_aad") or sqs_app_encryption_aad ) self.app_crypto: Optional["AppCrypto"] = None if self.sqs_aws_use_application_level_encryption: from litellm.litellm_core_utils.app_crypto import AppCrypto if not self.sqs_app_encryption_key_b64: - raise ValueError( - "sqs_app_encryption_key_b64 is required when encryption is enabled." - ) + raise ValueError("sqs_app_encryption_key_b64 is required when encryption is enabled.") key = base64.b64decode(self.sqs_app_encryption_key_b64) self.app_crypto = AppCrypto(key) verbose_logger.debug("SQSLogger: Application-level encryption enabled.") - self.sqs_config = ( - litellm.aws_sqs_callback_params.get("sqs_config") or sqs_config - ) + self.sqs_config = litellm.aws_sqs_callback_params.get("sqs_config") or sqs_config - async def async_log_success_event( - self, kwargs, response_obj, start_time, end_time - ) -> None: + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: try: - verbose_logger.debug( - "SQS Logging - Enters logging function for model %s", kwargs - ) + verbose_logger.debug("SQS Logging - Enters logging function for model %s", kwargs) standard_logging_payload = kwargs.get("standard_logging_object") if self.sqs_strip_base64_files: - standard_logging_payload = await self._strip_base64_from_messages( - standard_logging_payload - ) + standard_logging_payload = await self._strip_base64_from_messages(standard_logging_payload) if standard_logging_payload is None: raise ValueError("standard_logging_payload is None") @@ -271,9 +224,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): if standard_logging_payload is None: raise ValueError("standard_logging_payload is None") if self.sqs_strip_base64_files: - standard_logging_payload = await self._strip_base64_from_messages( - standard_logging_payload - ) + standard_logging_payload = await self._strip_base64_from_messages(standard_logging_payload) self.log_queue.append(standard_logging_payload) verbose_logger.debug( @@ -283,9 +234,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): ) except Exception as e: - verbose_logger.exception( - f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}" - ) + verbose_logger.exception(f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}") pass async def async_send_batch(self) -> None: @@ -324,28 +273,21 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): json_data = json.loads(safe_dumps(payload)) if self.app_crypto: - aad_bytes = ( - self.sqs_app_encryption_aad.encode("utf-8") - if self.sqs_app_encryption_aad - else None - ) + aad_bytes = self.sqs_app_encryption_aad.encode("utf-8") if self.sqs_app_encryption_aad else None encrypted = self.app_crypto.encrypt_json(json_data, aad=aad_bytes) json_string = json.dumps({"__encrypted__": True, "payload": encrypted}) else: json_string = safe_dumps(payload) - body = ( - f"Action={SQS_SEND_MESSAGE_ACTION}&Version={SQS_API_VERSION}&MessageBody=" - + quote(json_string, safe="") + body = f"Action={SQS_SEND_MESSAGE_ACTION}&Version={SQS_API_VERSION}&MessageBody=" + quote( + json_string, safe="" ) headers = { "Content-Type": "application/x-www-form-urlencoded", } - req = requests.Request( - "POST", self.sqs_queue_url, data=body, headers=headers - ) + req = requests.Request("POST", self.sqs_queue_url, data=body, headers=headers) prepped = req.prepare() aws_request = AWSRequest( @@ -377,13 +319,9 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): ) # Create a minimal standard logging payload - standard_logging_object: StandardLoggingPayload = ( - create_dummy_standard_logging_payload() - ) + standard_logging_object: StandardLoggingPayload = create_dummy_standard_logging_payload() # Attempt to send a single message await self.async_send_message(standard_logging_object) return IntegrationHealthCheckStatus(status="healthy", error_message=None) except Exception as e: - return IntegrationHealthCheckStatus( - status="unhealthy", error_message=str(e) - ) + return IntegrationHealthCheckStatus(status="unhealthy", error_message=str(e)) diff --git a/litellm/integrations/supabase.py b/litellm/integrations/supabase.py index 7eb007f813d..18cf4f9549c 100644 --- a/litellm/integrations/supabase.py +++ b/litellm/integrations/supabase.py @@ -31,13 +31,9 @@ class Supabase: self.supabase_url, self.supabase_key ) - def input_log_event( - self, model, messages, end_user, litellm_call_id, print_verbose - ): + def input_log_event(self, model, messages, end_user, litellm_call_id, print_verbose): try: - print_verbose( - f"Supabase Logging - Enters input logging function for model {model}" - ) + print_verbose(f"Supabase Logging - Enters input logging function for model {model}") supabase_data_obj = { "model": model, "messages": messages, @@ -45,11 +41,7 @@ class Supabase: "status": "initiated", "litellm_call_id": litellm_call_id, } - data, count = ( - self.supabase_client.table(self.supabase_table_name) - .insert(supabase_data_obj) - .execute() - ) + data, count = self.supabase_client.table(self.supabase_table_name).insert(supabase_data_obj).execute() print_verbose(f"data: {data}") except Exception: print_verbose(f"Supabase Logging Error - {traceback.format_exc()}") @@ -67,9 +59,7 @@ class Supabase: print_verbose, ): try: - print_verbose( - f"Supabase Logging - Enters logging function for model {model}, response_obj: {response_obj}" - ) + print_verbose(f"Supabase Logging - Enters logging function for model {model}, response_obj: {response_obj}") total_cost = litellm.completion_cost(completion_response=response_obj) @@ -85,9 +75,7 @@ class Supabase: "litellm_call_id": litellm_call_id, "status": "success", } - print_verbose( - f"Supabase Logging - final data object: {supabase_data_obj}" - ) + print_verbose(f"Supabase Logging - final data object: {supabase_data_obj}") data, count = ( self.supabase_client.table(self.supabase_table_name) .upsert(supabase_data_obj, on_conflict="litellm_call_id") @@ -106,9 +94,7 @@ class Supabase: "litellm_call_id": litellm_call_id, "status": "failure", } - print_verbose( - f"Supabase Logging - final data object: {supabase_data_obj}" - ) + print_verbose(f"Supabase Logging - final data object: {supabase_data_obj}") data, count = ( self.supabase_client.table(self.supabase_table_name) .upsert(supabase_data_obj, on_conflict="litellm_call_id") diff --git a/litellm/integrations/traceloop.py b/litellm/integrations/traceloop.py index b4f3905c8e8..77f20972f7a 100644 --- a/litellm/integrations/traceloop.py +++ b/litellm/integrations/traceloop.py @@ -40,23 +40,17 @@ class TraceloopLogger: from opentelemetry.trace import SpanKind, Status, StatusCode try: - print_verbose( - f"Traceloop Logging - Enters logging function for model {kwargs}" - ) + print_verbose(f"Traceloop Logging - Enters logging function for model {kwargs}") tracer = self.tracer_wrapper.get_tracer() optional_params = kwargs.get("optional_params", {}) start_time = int(start_time.timestamp()) end_time = int(end_time.timestamp()) - span = tracer.start_span( - "litellm.completion", kind=SpanKind.CLIENT, start_time=start_time - ) + span = tracer.start_span("litellm.completion", kind=SpanKind.CLIENT, start_time=start_time) if span.is_recording(): - span.set_attribute( - SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("model") - ) + span.set_attribute(SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("model")) if "stop" in optional_params: span.set_attribute( SpanAttributes.LLM_CHAT_STOP_SEQUENCES, @@ -73,18 +67,14 @@ class TraceloopLogger: optional_params.get("presence_penalty"), ) if "top_p" in optional_params: - span.set_attribute( - SpanAttributes.LLM_REQUEST_TOP_P, optional_params.get("top_p") - ) + span.set_attribute(SpanAttributes.LLM_REQUEST_TOP_P, optional_params.get("top_p")) if "tools" in optional_params or "functions" in optional_params: span.set_attribute( SpanAttributes.LLM_REQUEST_FUNCTIONS, optional_params.get("tools", optional_params.get("functions")), ) if "user" in optional_params: - span.set_attribute( - SpanAttributes.LLM_USER, optional_params.get("user") - ) + span.set_attribute(SpanAttributes.LLM_USER, optional_params.get("user")) if "max_tokens" in optional_params: span.set_attribute( SpanAttributes.LLM_REQUEST_MAX_TOKENS, @@ -106,9 +96,7 @@ class TraceloopLogger: prompt.get("content"), ) - span.set_attribute( - SpanAttributes.LLM_RESPONSE_MODEL, response_obj.get("model") - ) + span.set_attribute(SpanAttributes.LLM_RESPONSE_MODEL, response_obj.get("model")) usage = response_obj.get("usage") if usage: span.set_attribute( @@ -138,11 +126,7 @@ class TraceloopLogger: choice.get("message").get("content"), ) - if ( - level == "ERROR" - and status_message is not None - and isinstance(status_message, str) - ): + if level == "ERROR" and status_message is not None and isinstance(status_message, str): span.record_exception(Exception(status_message)) span.set_status(Status(StatusCode.ERROR, status_message)) diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index 1e6e46b36ae..be8907f07ff 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -45,12 +45,8 @@ class VantageLogger(FocusLogger): ) -> None: resolved_api_key = api_key or os.getenv("VANTAGE_API_KEY") resolved_token = integration_token or os.getenv("VANTAGE_INTEGRATION_TOKEN") - resolved_base_url = base_url or os.getenv( - "VANTAGE_BASE_URL", "https://api.vantage.sh" - ) - resolved_frequency = ( - frequency or os.getenv("VANTAGE_EXPORT_FREQUENCY") or "hourly" - ).lower() + resolved_base_url = base_url or os.getenv("VANTAGE_BASE_URL", "https://api.vantage.sh") + resolved_frequency = (frequency or os.getenv("VANTAGE_EXPORT_FREQUENCY") or "hourly").lower() raw_interval = interval_seconds or os.getenv("VANTAGE_EXPORT_INTERVAL_SECONDS") resolved_interval: Optional[int] = None @@ -83,11 +79,7 @@ class VantageLogger(FocusLogger): verbose_logger.debug( "VantageLogger initialized (integration_token=%s)", - ( - resolved_token[:4] + "***" - if resolved_token and len(resolved_token) > 4 - else "***" - ), + (resolved_token[:4] + "***" if resolved_token and len(resolved_token) > 4 else "***"), ) async def initialize_focus_export_job(self) -> None: @@ -106,18 +98,14 @@ class VantageLogger(FocusLogger): pod_lock_manager = getattr(writer, "pod_lock_manager", None) if pod_lock_manager and pod_lock_manager.redis_cache: - acquired = await pod_lock_manager.acquire_lock( - cronjob_id=VANTAGE_USAGE_DATA_JOB_NAME - ) + acquired = await pod_lock_manager.acquire_lock(cronjob_id=VANTAGE_USAGE_DATA_JOB_NAME) if not acquired: verbose_logger.debug("Vantage export: unable to acquire pod lock") return try: await self._run_scheduled_export() finally: - await pod_lock_manager.release_lock( - cronjob_id=VANTAGE_USAGE_DATA_JOB_NAME - ) + await pod_lock_manager.release_lock(cronjob_id=VANTAGE_USAGE_DATA_JOB_NAME) else: await self._run_scheduled_export() @@ -126,10 +114,8 @@ class VantageLogger(FocusLogger): scheduler: AsyncIOScheduler, ) -> None: """Register the Vantage export job with the provided scheduler.""" - vantage_loggers: List[CustomLogger] = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=VantageLogger - ) + vantage_loggers: List[CustomLogger] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=VantageLogger ) if not vantage_loggers: verbose_logger.debug("No Vantage logger registered; skipping scheduler") diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index c54b6e4cced..0ba6da78b27 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -103,9 +103,7 @@ class VectorStorePreCallHook(CustomLogger): query = self._extract_query_from_messages(messages) if not query: - verbose_logger.debug( - "No query found in messages for vector store search" - ) + verbose_logger.debug("No query found in messages for vector store search") return model, messages, non_default_params modified_messages: List[AllMessageValues] = messages.copy() @@ -115,9 +113,7 @@ class VectorStorePreCallHook(CustomLogger): # Get vector store id from the vector store config vector_store_id = vector_store_to_run.get("vector_store_id", "") custom_llm_provider = vector_store_to_run.get("custom_llm_provider") - litellm_params_for_vector_store = ( - vector_store_to_run.get("litellm_params", {}) or {} - ) + litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {} # Call litellm.vector_stores.search() with the required parameters search_response = await litellm.vector_stores.asearch( **{ @@ -141,15 +137,11 @@ class VectorStorePreCallHook(CustomLogger): # Get the number of results for logging num_results = 0 num_results = len(search_response.get("data", []) or []) - verbose_logger.debug( - f"Vector store search completed. Added context from {num_results} results" - ) + verbose_logger.debug(f"Vector store search completed. Added context from {num_results} results") # Store search results as-is (already in OpenAI-compatible format) if litellm_logging_obj and all_search_results: - litellm_logging_obj.model_call_details["search_results"] = ( - all_search_results - ) + litellm_logging_obj.model_call_details["search_results"] = all_search_results return model, modified_messages, non_default_params @@ -158,9 +150,7 @@ class VectorStorePreCallHook(CustomLogger): # Return original parameters on error return model, messages, non_default_params - def _extract_query_from_messages( - self, messages: List[AllMessageValues] - ) -> Optional[str]: + def _extract_query_from_messages(self, messages: List[AllMessageValues]) -> Optional[str]: """ Extract the query from the last user message. @@ -184,11 +174,7 @@ class VectorStorePreCallHook(CustomLogger): elif isinstance(content, list) and len(content) > 0: # Handle list of content items, extract text from first text item for item in content: - if ( - isinstance(item, dict) - and item.get("type") == "text" - and "text" in item - ): + if isinstance(item, dict) and item.get("type") == "text" and "text" in item: return item["text"] return None @@ -208,18 +194,14 @@ class VectorStorePreCallHook(CustomLogger): Returns: Modified list of messages with context appended """ - search_response_data: Optional[List[VectorStoreSearchResult]] = ( - search_response.get("data") - ) + search_response_data: Optional[List[VectorStoreSearchResult]] = search_response.get("data") if not search_response_data: return messages context_content = self.CONTENT_PREFIX_STRING for result in search_response_data: - result_content: Optional[List[VectorStoreResultContent]] = result.get( - "content" - ) + result_content: Optional[List[VectorStoreResultContent]] = result.get("content") if result_content: for content_item in result_content: content_text: Optional[str] = content_item.get("text") @@ -253,9 +235,7 @@ class VectorStorePreCallHook(CustomLogger): to the response's provider_specific_fields. """ try: - verbose_logger.debug( - "VectorStorePreCallHook.async_post_call_success_deployment_hook called" - ) + verbose_logger.debug("VectorStorePreCallHook.async_post_call_success_deployment_hook called") # Get logging object from request_data litellm_logging_obj = request_data.get("litellm_logging_obj") @@ -263,13 +243,11 @@ class VectorStorePreCallHook(CustomLogger): verbose_logger.debug("No litellm_logging_obj in request_data") return None - verbose_logger.debug( - f"model_call_details keys: {list(litellm_logging_obj.model_call_details.keys())}" - ) + verbose_logger.debug(f"model_call_details keys: {list(litellm_logging_obj.model_call_details.keys())}") # Get search results from model_call_details (already in OpenAI format) - search_results: Optional[List[VectorStoreSearchResponse]] = ( - litellm_logging_obj.model_call_details.get("search_results") + search_results: Optional[List[VectorStoreSearchResponse]] = litellm_logging_obj.model_call_details.get( + "search_results" ) verbose_logger.debug(f"Search results found: {search_results is not None}") @@ -283,30 +261,21 @@ class VectorStorePreCallHook(CustomLogger): for choice in response.choices: if hasattr(choice, "message") and choice.message: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(choice.message, "provider_specific_fields", None) - or {} - ) + provider_fields = getattr(choice.message, "provider_specific_fields", None) or {} # Add search results (already in OpenAI-compatible format) provider_fields["search_results"] = search_results # Set the provider_specific_fields - setattr( - choice.message, "provider_specific_fields", provider_fields - ) + setattr(choice.message, "provider_specific_fields", provider_fields) - verbose_logger.debug( - f"Added {len(search_results)} search results to response" - ) + verbose_logger.debug(f"Added {len(search_results)} search results to response") # Return modified response return response except Exception as e: - verbose_logger.exception( - f"Error adding search results to response: {str(e)}" - ) + verbose_logger.exception(f"Error adding search results to response: {str(e)}") # Don't fail the request if search results fail to be added return None @@ -323,18 +292,12 @@ class VectorStorePreCallHook(CustomLogger): search results to the stream before it's returned to the user. """ try: - verbose_logger.debug( - "VectorStorePreCallHook.async_post_call_streaming_deployment_hook called" - ) + verbose_logger.debug("VectorStorePreCallHook.async_post_call_streaming_deployment_hook called") # Get search results from model_call_details (already in OpenAI format) - search_results: Optional[List[VectorStoreSearchResponse]] = ( - request_data.get("search_results") - ) + search_results: Optional[List[VectorStoreSearchResponse]] = request_data.get("search_results") - verbose_logger.debug( - f"Search results found for streaming chunk: {search_results is not None}" - ) + verbose_logger.debug(f"Search results found for streaming chunk: {search_results is not None}") if not search_results: verbose_logger.debug("No search results found for streaming chunk") @@ -345,10 +308,7 @@ class VectorStorePreCallHook(CustomLogger): for choice in response_chunk.choices: if hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(choice.delta, "provider_specific_fields", None) - or {} - ) + provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {} # Add search results (already in OpenAI-compatible format) provider_fields["search_results"] = search_results @@ -356,16 +316,12 @@ class VectorStorePreCallHook(CustomLogger): # Set the provider_specific_fields choice.delta.provider_specific_fields = provider_fields - verbose_logger.debug( - f"Added {len(search_results)} search results to streaming chunk" - ) + verbose_logger.debug(f"Added {len(search_results)} search results to streaming chunk") # Return modified chunk return response_chunk except Exception as e: - verbose_logger.exception( - f"Error adding search results to streaming chunk: {str(e)}" - ) + verbose_logger.exception(f"Error adding search results to streaming chunk: {str(e)}") # Don't fail the request if search results fail to be added return response_chunk diff --git a/litellm/integrations/weave/weave_otel.py b/litellm/integrations/weave/weave_otel.py index 796a33a34d5..c43afe7b6ca 100644 --- a/litellm/integrations/weave/weave_otel.py +++ b/litellm/integrations/weave/weave_otel.py @@ -56,14 +56,10 @@ class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): prompt["functions"] = functions if tools is not None: prompt["tools"] = tools - safe_set_attribute( - span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt) - ) + safe_set_attribute(span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt)) -def _set_weave_specific_attributes( - span: Span, kwargs: dict[str, Any], response_obj: Any -): +def _set_weave_specific_attributes(span: Span, kwargs: dict[str, Any], response_obj: Any): """ Sets Weave-specific metadata attributes onto the OTEL span. @@ -106,9 +102,7 @@ def _set_weave_specific_attributes( output_dict = response_obj if output_dict: - safe_set_attribute( - span, OpenInferenceSpanAttributes.OUTPUT_VALUE, safe_dumps(output_dict) - ) + safe_set_attribute(span, OpenInferenceSpanAttributes.OUTPUT_VALUE, safe_dumps(output_dict)) def _get_weave_authorization_header(api_key: str) -> str: @@ -142,9 +136,7 @@ def get_weave_otel_config() -> WeaveOtelConfig: host = os.getenv("WANDB_HOST") if not api_key: - raise ValueError( - "WANDB_API_KEY must be set for Weave OpenTelemetry integration." - ) + raise ValueError("WANDB_API_KEY must be set for Weave OpenTelemetry integration.") if not project_id: raise ValueError( @@ -233,9 +225,7 @@ class WeaveOtelLogger(OpenTelemetry): super().__init__(config=config, callback_name=callback_name, **kwargs) - def _maybe_log_raw_request( - self, kwargs, response_obj, start_time, end_time, parent_span - ): + def _maybe_log_raw_request(self, kwargs, response_obj, start_time, end_time, parent_span): """ Override to skip creating the raw_gen_ai_request child span. @@ -293,9 +283,7 @@ class WeaveOtelLogger(OpenTelemetry): primary_span_parent = None # 1. Primary span - span = self._start_primary_span( - kwargs, response_obj, start_time, end_time, ctx, primary_span_parent - ) + span = self._start_primary_span(kwargs, response_obj, start_time, end_time, ctx, primary_span_parent) # 2. Raw-request sub-span (skipped for Weave via _maybe_log_raw_request override) self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) @@ -329,9 +317,7 @@ class WeaveOtelLogger(OpenTelemetry): dynamic_headers = {} dynamic_wandb_api_key = standard_callback_dynamic_params.get("wandb_api_key") - dynamic_weave_project_id = standard_callback_dynamic_params.get( - "weave_project_id" - ) + dynamic_weave_project_id = standard_callback_dynamic_params.get("weave_project_id") if dynamic_wandb_api_key: auth_header = _get_weave_authorization_header( diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index ab4c787c96f..2e11405af3f 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -80,9 +80,7 @@ class WebSearchInterceptionLogger(CustomLogger): if enabled_providers is None: self.enabled_providers = [LlmProviders.BEDROCK.value] else: - self.enabled_providers = [ - p.value if isinstance(p, LlmProviders) else p for p in enabled_providers - ] + self.enabled_providers = [p.value if isinstance(p, LlmProviders) else p for p in enabled_providers] self.search_tool_name = search_tool_name self._request_has_websearch = False # Track if current request has web search @@ -118,10 +116,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Check if provider is in enabled list provider_str = custom_llm_provider or "" - if ( - self.enabled_providers is not None - and provider_str not in self.enabled_providers - ): + if self.enabled_providers is not None and provider_str not in self.enabled_providers: return None # Only short-circuit for providers without native Anthropic Messages @@ -132,10 +127,8 @@ class WebSearchInterceptionLogger(CustomLogger): # return raw search text — a regression for existing users. try: provider_enum = LlmProviders(provider_str) - anthropic_config = ( - ProviderConfigManager.get_provider_anthropic_messages_config( - model=model, provider=provider_enum - ) + anthropic_config = ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, provider=provider_enum ) if anthropic_config is not None: verbose_logger.debug( @@ -160,8 +153,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None verbose_logger.debug( - "WebSearchInterception: Short-circuit search detected " - f"(provider={provider_str}, query='{query}')" + f"WebSearchInterception: Short-circuit search detected (provider={provider_str}, query='{query}')" ) # Native clients (Claude Desktop / Cowork / Anthropic SDK) make a @@ -180,9 +172,7 @@ class WebSearchInterceptionLogger(CustomLogger): try: search_result_text, structured = await self._execute_search(query) except Exception as e: - verbose_logger.error( - f"WebSearchInterception: Short-circuit search failed: {e}" - ) + verbose_logger.error(f"WebSearchInterception: Short-circuit search failed: {e}") search_result_text, structured = f"Search failed: {e}", None content: List[Dict[str, Any]] = [] @@ -225,9 +215,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) return response - async def async_pre_call_deployment_hook( - self, kwargs: Dict[str, Any], call_type: Optional[Any] - ) -> Optional[dict]: + async def async_pre_call_deployment_hook(self, kwargs: Dict[str, Any], call_type: Optional[Any]) -> Optional[dict]: """ Pre-call hook to convert native Anthropic web_search tools to regular tools. @@ -237,14 +225,12 @@ class WebSearchInterceptionLogger(CustomLogger): """ # Check if this is for an enabled provider # Try top-level kwargs first, then nested litellm_params, then derive from model name - custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get( - "litellm_params", {} - ).get("custom_llm_provider", "") + custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get( + "custom_llm_provider", "" + ) if not custom_llm_provider: try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=kwargs.get("model", "") - ) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", "")) except Exception: custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: @@ -261,9 +247,7 @@ class WebSearchInterceptionLogger(CustomLogger): if not has_websearch: return None - verbose_logger.debug( - "WebSearchInterception: Converting native web_search tools to LiteLLM standard" - ) + verbose_logger.debug("WebSearchInterception: Converting native web_search tools to LiteLLM standard") # If the client sent an Anthropic-native web_search_* tool, mark the # request so the agentic loop emits native web_search_tool_result @@ -291,18 +275,14 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs["tools"] = converted_tools if kwargs.get("stream"): - verbose_logger.debug( - "WebSearchInterception: deployment hook converting stream=True to stream=False" - ) + verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") kwargs["stream"] = False kwargs["_websearch_interception_converted_stream"] = True return kwargs @classmethod - def from_config_yaml( - cls, config: WebSearchInterceptionConfig - ) -> "WebSearchInterceptionLogger": + def from_config_yaml(cls, config: WebSearchInterceptionConfig) -> "WebSearchInterceptionLogger": """ Initialize WebSearchInterceptionLogger from proxy config.yaml parameters. @@ -354,9 +334,7 @@ class WebSearchInterceptionLogger(CustomLogger): return tool.get("name") @classmethod - def _sync_forced_tool_choice( - cls, tool_choice: Any, converted_tools: list[dict[str, Any]] - ) -> Any: + def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, Any]]) -> Any: """Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it names a web-search tool that was just converted away. @@ -373,9 +351,7 @@ class WebSearchInterceptionLogger(CustomLogger): return tool_choice return {**tool_choice, "name": LITELLM_WEB_SEARCH_TOOL_NAME} - async def async_pre_request_hook( - self, model: str, messages: List[Dict], kwargs: Dict - ) -> Optional[Dict]: + async def async_pre_request_hook(self, model: str, messages: List[Dict], kwargs: Dict) -> Optional[Dict]: """ Pre-request hook to convert native web search tools to LiteLLM standard. @@ -391,9 +367,7 @@ class WebSearchInterceptionLogger(CustomLogger): Modified kwargs dict with converted tools, or None if no modifications needed """ # Check if this request is for an enabled provider - custom_llm_provider = kwargs.get("litellm_params", {}).get( - "custom_llm_provider", "" - ) + custom_llm_provider = kwargs.get("litellm_params", {}).get("custom_llm_provider", "") verbose_logger.debug( f"WebSearchInterception: Pre-request hook called" @@ -401,10 +375,7 @@ class WebSearchInterceptionLogger(CustomLogger): f" - enabled_providers={self.enabled_providers or 'ALL'}" ) - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( f"WebSearchInterception: Skipping - provider {custom_llm_provider} not in {self.enabled_providers}" ) @@ -420,9 +391,7 @@ class WebSearchInterceptionLogger(CustomLogger): if not has_websearch: return None - verbose_logger.debug( - f"WebSearchInterception: Pre-request hook triggered for provider={custom_llm_provider}" - ) + verbose_logger.debug(f"WebSearchInterception: Pre-request hook triggered for provider={custom_llm_provider}") # If the client sent an Anthropic-native web_search_* tool, mark the # request so the agentic loop emits native web_search_tool_result @@ -451,15 +420,11 @@ class WebSearchInterceptionLogger(CustomLogger): ) if "tool_choice" in kwargs: - kwargs["tool_choice"] = self._sync_forced_tool_choice( - kwargs.get("tool_choice"), converted_tools - ) + kwargs["tool_choice"] = self._sync_forced_tool_choice(kwargs.get("tool_choice"), converted_tools) # Also convert here for direct callers that bypass the deployment hook. if kwargs.get("stream"): - verbose_logger.debug( - "WebSearchInterception: Converting stream=True to stream=False" - ) + verbose_logger.debug("WebSearchInterception: Converting stream=True to stream=False") kwargs["stream"] = False kwargs["_websearch_interception_converted_stream"] = True @@ -482,18 +447,13 @@ class WebSearchInterceptionLogger(CustomLogger): For chat completions, use async_should_run_chat_completion_agentic_loop instead. """ - verbose_logger.debug( - f"WebSearchInterception: Hook called! provider={custom_llm_provider}, stream={stream}" - ) + verbose_logger.debug(f"WebSearchInterception: Hook called! provider={custom_llm_provider}, stream={stream}") verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}") # Check if provider should be intercepted # Note: custom_llm_provider is already normalized by get_llm_provider() # (e.g., "bedrock/invoke/..." -> "bedrock") - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( f"WebSearchInterception: Skipping provider {custom_llm_provider} (not in enabled list: {self.enabled_providers})" ) @@ -513,9 +473,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug( - "WebSearchInterception: No WebSearch tool_use detected in response" - ) + verbose_logger.debug("WebSearchInterception: No WebSearch tool_use detected in response") return False, {} verbose_logger.debug( @@ -547,9 +505,7 @@ class WebSearchInterceptionLogger(CustomLogger): thinking_block_dict: Dict = {"type": block_type} if block_type == "thinking": thinking_block_dict["thinking"] = getattr(block, "thinking", "") - thinking_block_dict["signature"] = getattr( - block, "signature", "" - ) + thinking_block_dict["signature"] = getattr(block, "signature", "") else: # redacted_thinking thinking_block_dict["data"] = getattr(block, "data", "") thinking_blocks.append(thinking_block_dict) @@ -591,23 +547,16 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}") # Check if provider should be intercepted - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( f"WebSearchInterception: Skipping provider {custom_llm_provider} (not in enabled list: {self.enabled_providers})" ) return False, {} # Check if tools include any web search tool (strict check for chat completions) - has_websearch_tool = any( - is_web_search_tool_chat_completion(t) for t in (tools or []) - ) + has_websearch_tool = any(is_web_search_tool_chat_completion(t) for t in (tools or [])) if not has_websearch_tool: - verbose_logger.debug( - "WebSearchInterception: No litellm_web_search tool in request" - ) + verbose_logger.debug("WebSearchInterception: No litellm_web_search tool in request") return False, {} # Detect WebSearch tool_calls in response (OpenAI format) @@ -618,9 +567,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug( - "WebSearchInterception: No WebSearch tool_calls detected in response" - ) + verbose_logger.debug("WebSearchInterception: No WebSearch tool_calls detected in response") return False, {} verbose_logger.debug( @@ -657,9 +604,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls = tools["tool_calls"] thinking_blocks = tools.get("thinking_blocks", []) - verbose_logger.debug( - f"WebSearchInterception: Executing agentic loop for {len(tool_calls)} search(es)" - ) + verbose_logger.debug(f"WebSearchInterception: Executing agentic loop for {len(tool_calls)} search(es)") return await self._execute_agentic_loop( model=model, @@ -706,11 +651,9 @@ class WebSearchInterceptionLogger(CustomLogger): # (while we still have the structured SearchResponse list) and stash # them on plan metadata for the post-hook to inject. if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): - metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = ( - self._build_native_result_blocks( - tool_calls=tool_calls, - structured_results=structured_results, - ) + metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks( + tool_calls=tool_calls, + structured_results=structured_results, ) return AgenticLoopPlan( @@ -759,9 +702,7 @@ class WebSearchInterceptionLogger(CustomLogger): return blocks @staticmethod - def _inject_native_blocks( - response: Any, native_blocks: List[Dict[str, Any]] - ) -> Any: + def _inject_native_blocks(response: Any, native_blocks: List[Dict[str, Any]]) -> Any: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response @@ -776,8 +717,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Object refused write — fall through and leave the response # untouched rather than crash the request. verbose_logger.debug( - "WebSearchInterception: could not inject native blocks into " - f"response of type {type(response).__name__}" + f"WebSearchInterception: could not inject native blocks into response of type {type(response).__name__}" ) return response @@ -891,9 +831,7 @@ class WebSearchInterceptionLogger(CustomLogger): """ _internal_keys = {"litellm_logging_obj"} return { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") and k not in _internal_keys + k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in _internal_keys } async def _execute_agentic_loop( @@ -975,21 +913,15 @@ class WebSearchInterceptionLogger(CustomLogger): for tool_call in tool_calls: query = tool_call["input"].get("query") if query: - verbose_logger.debug( - f"WebSearchInterception: Queuing search for query='{query}'" - ) + verbose_logger.debug(f"WebSearchInterception: Queuing search for query='{query}'") search_tasks.append(self._execute_search(query)) else: - verbose_logger.debug( - f"WebSearchInterception: Tool call {tool_call['id']} has no query" - ) + verbose_logger.debug(f"WebSearchInterception: Tool call {tool_call['id']} has no query") # Add empty result for tools without query search_tasks.append(self._create_empty_search_result()) # Execute searches in parallel - verbose_logger.debug( - f"WebSearchInterception: Executing {len(search_tasks)} search(es) in parallel" - ) + verbose_logger.debug(f"WebSearchInterception: Executing {len(search_tasks)} search(es) in parallel") search_results = await asyncio.gather(*search_tasks, return_exceptions=True) # Split the gathered (text, structured) tuples into two parallel lists. @@ -999,29 +931,17 @@ class WebSearchInterceptionLogger(CustomLogger): structured_results: List[Optional[SearchResponse]] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): - verbose_logger.error( - f"WebSearchInterception: Search {i} failed with error: {str(result)}" - ) + verbose_logger.error(f"WebSearchInterception: Search {i} failed with error: {str(result)}") final_search_results.append(f"Search failed: {str(result)}") structured_results.append(None) elif isinstance(result, tuple) and len(result) == 2: text_value, structured_value = result - final_search_results.append( - cast(str, text_value) - if isinstance(text_value, str) - else str(text_value) - ) - structured_results.append( - structured_value - if isinstance(structured_value, SearchResponse) - else None - ) + final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) + structured_results.append(structured_value if isinstance(structured_value, SearchResponse) else None) else: # Defensive: legacy callers / unexpected shape — preserve text, # drop structure. - verbose_logger.debug( - f"WebSearchInterception: Unexpected result type {type(result)} at index {i}" - ) + verbose_logger.debug(f"WebSearchInterception: Unexpected result type {type(result)} at index {i}") final_search_results.append(str(result)) structured_results.append(None) @@ -1035,35 +955,24 @@ class WebSearchInterceptionLogger(CustomLogger): follow_up_messages = messages + [assistant_message, cast(Dict, user_message)] # Correlation context for structured logging - _call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get( - "litellm_call_id", "unknown" - ) + _call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get("litellm_call_id", "unknown") full_model_name = model # safe default before try block - max_tokens = self._resolve_max_tokens( - anthropic_messages_optional_request_params, kwargs - ) + max_tokens = self._resolve_max_tokens(anthropic_messages_optional_request_params, kwargs) - verbose_logger.debug( - f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request" - ) + verbose_logger.debug(f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request") optional_params_without_max_tokens = { - k: v - for k, v in anthropic_messages_optional_request_params.items() - if k != "max_tokens" + k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens" } kwargs_for_followup = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get( - "agentic_loop_params", {} - ) + agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = agentic_params.get("model", model) verbose_logger.debug( - "WebSearchInterception: Built anthropic request patch " - "[call_id=%s model=%s messages=%d searches=%d]", + "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]", _call_id, full_model_name, len(follow_up_messages), @@ -1112,9 +1021,7 @@ class WebSearchInterceptionLogger(CustomLogger): ] if matching_tools: search_tool = matching_tools[0] - search_provider = search_tool.get("litellm_params", {}).get( - "search_provider" - ) + search_provider = search_tool.get("litellm_params", {}).get("search_provider") verbose_logger.debug( f"WebSearchInterception: Found search tool '{self.search_tool_name}' " f"with provider '{search_provider}'" @@ -1128,9 +1035,7 @@ class WebSearchInterceptionLogger(CustomLogger): # If no specific tool or not found, use first available if not search_provider and llm_router.search_tools: first_tool = llm_router.search_tools[0] - search_provider = first_tool.get("litellm_params", {}).get( - "search_provider" - ) + search_provider = first_tool.get("litellm_params", {}).get("search_provider") verbose_logger.debug( f"WebSearchInterception: Using first available search tool with provider '{search_provider}'" ) @@ -1156,9 +1061,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) return search_result_text, result except Exception as e: - verbose_logger.error( - f"WebSearchInterception: Search failed for '{query}': {str(e)}" - ) + verbose_logger.error(f"WebSearchInterception: Search failed for '{query}': {str(e)}") raise async def _execute_chat_completion_agentic_loop( @@ -1218,21 +1121,15 @@ class WebSearchInterceptionLogger(CustomLogger): query = args.get("query") if query: - verbose_logger.debug( - f"WebSearchInterception: Queuing search for query='{query}'" - ) + verbose_logger.debug(f"WebSearchInterception: Queuing search for query='{query}'") search_tasks.append(self._execute_search(query)) else: - verbose_logger.debug( - f"WebSearchInterception: Tool call {tool_call.get('id')} has no query" - ) + verbose_logger.debug(f"WebSearchInterception: Tool call {tool_call.get('id')} has no query") # Add empty result for tools without query search_tasks.append(self._create_empty_search_result()) # Execute searches in parallel - verbose_logger.debug( - f"WebSearchInterception: Executing {len(search_tasks)} search(es) in parallel" - ) + verbose_logger.debug(f"WebSearchInterception: Executing {len(search_tasks)} search(es) in parallel") search_results = await asyncio.gather(*search_tasks, return_exceptions=True) # Chat-completion path only needs text — OpenAI tool_result format @@ -1240,21 +1137,13 @@ class WebSearchInterceptionLogger(CustomLogger): final_search_results: List[str] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): - verbose_logger.error( - f"WebSearchInterception: Search {i} failed with error: {str(result)}" - ) + verbose_logger.error(f"WebSearchInterception: Search {i} failed with error: {str(result)}") final_search_results.append(f"Search failed: {str(result)}") elif isinstance(result, tuple) and len(result) == 2: text_value, _ = result - final_search_results.append( - cast(str, text_value) - if isinstance(text_value, str) - else str(text_value) - ) + final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) else: - verbose_logger.debug( - f"WebSearchInterception: Unexpected result type {type(result)} at index {i}" - ) + verbose_logger.debug(f"WebSearchInterception: Unexpected result type {type(result)} at index {i}") final_search_results.append(str(result)) # Build assistant and tool messages using transformation @@ -1270,9 +1159,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Make follow-up request with search results # For OpenAI format, tool_messages_or_user is a list of tool messages if response_format == "openai": - follow_up_messages = ( - messages + [assistant_message] + cast(List[Dict], tool_messages_or_user) - ) + follow_up_messages = messages + [assistant_message] + cast(List[Dict], tool_messages_or_user) else: # For Anthropic format (shouldn't happen in this method, but handle it) follow_up_messages = messages + [ @@ -1280,12 +1167,8 @@ class WebSearchInterceptionLogger(CustomLogger): cast(Dict, tool_messages_or_user), ] - verbose_logger.debug( - "WebSearchInterception: Making follow-up chat completion request with search results" - ) - verbose_logger.debug( - f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}" - ) + verbose_logger.debug("WebSearchInterception: Making follow-up chat completion request with search results") + verbose_logger.debug(f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}") # Remove internal parameters that shouldn't be passed to follow-up request internal_params = { @@ -1298,9 +1181,7 @@ class WebSearchInterceptionLogger(CustomLogger): "custom_prompt_dict", } kwargs_for_followup = { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") and k not in internal_params + k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in internal_params } full_model_name = model diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 9c20a3f6c77..7bbcd7ebff6 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -53,9 +53,7 @@ class WebSearchTransformation: if stream: # This should not happen in practice since we convert streaming to non-streaming # in async_log_pre_api_call, but keep this check for safety - verbose_logger.warning( - "WebSearchInterception: Unexpected streaming response, skipping interception" - ) + verbose_logger.warning("WebSearchInterception: Unexpected streaming response, skipping interception") return False, [] # Parse non-streaming response based on format @@ -75,9 +73,7 @@ class WebSearchTransformation: content = response.get("content", []) else: if not hasattr(response, "content"): - verbose_logger.debug( - "WebSearchInterception: Response has no content attribute" - ) + verbose_logger.debug("WebSearchInterception: Response has no content attribute") return False, [] content = response.content or [] @@ -118,9 +114,7 @@ class WebSearchTransformation: "input": block_input, } tool_calls.append(tool_call) - verbose_logger.debug( - f"WebSearchInterception: Found {block_name} tool_use with id={tool_call['id']}" - ) + verbose_logger.debug(f"WebSearchInterception: Found {block_name} tool_use with id={tool_call['id']}") return len(tool_calls) > 0, tool_calls @@ -135,9 +129,7 @@ class WebSearchTransformation: choices = response.get("choices", []) else: if not hasattr(response, "choices"): - verbose_logger.debug( - "WebSearchInterception: Response has no choices attribute" - ) + verbose_logger.debug("WebSearchInterception: Response has no choices attribute") return False, [] choices = response.choices or [] @@ -174,24 +166,16 @@ class WebSearchTransformation: tool_id = tool_call.get("id") tool_type = tool_call.get("type") function = tool_call.get("function", {}) - function_name = ( - function.get("name") - if isinstance(function, dict) - else getattr(function, "name", None) - ) + function_name = function.get("name") if isinstance(function, dict) else getattr(function, "name", None) function_arguments = ( - function.get("arguments") - if isinstance(function, dict) - else getattr(function, "arguments", None) + function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None) ) else: tool_id = getattr(tool_call, "id", None) tool_type = getattr(tool_call, "type", None) function = getattr(tool_call, "function", None) function_name = getattr(function, "name", None) if function else None - function_arguments = ( - getattr(function, "arguments", None) if function else None - ) + function_arguments = getattr(function, "arguments", None) if function else None # Detect function-style web search tool_calls. ``WebSearch`` is # intentionally omitted — see is_web_search_tool for the Cowork @@ -225,9 +209,7 @@ class WebSearchTransformation: "input": arguments, # For compatibility with Anthropic format } tool_calls.append(tool_call_dict) - verbose_logger.debug( - f"WebSearchInterception: Found {function_name} tool_call with id={tool_id}" - ) + verbose_logger.debug(f"WebSearchInterception: Found {function_name} tool_call with id={tool_id}") return len(tool_calls) > 0, tool_calls @@ -259,9 +241,7 @@ class WebSearchTransformation: For OpenAI: assistant_message with tool_calls, tool_messages list with tool results """ if response_format == "openai": - return WebSearchTransformation._transform_response_openai( - tool_calls, search_results - ) + return WebSearchTransformation._transform_response_openai(tool_calls, search_results) else: return WebSearchTransformation._transform_response_anthropic( tool_calls, search_results, thinking_blocks=thinking_blocks @@ -332,11 +312,7 @@ class WebSearchTransformation: "type": "function", "function": { "name": tc["name"], - "arguments": ( - json.dumps(tc["input"]) - if isinstance(tc["input"], dict) - else str(tc["input"]) - ), + "arguments": (json.dumps(tc["input"]) if isinstance(tc["input"], dict) else str(tc["input"])), }, } for tc in tool_calls @@ -421,10 +397,7 @@ class WebSearchTransformation: if hasattr(result, "results") and result.results: # Format results as text search_result_text = "\n\n".join( - [ - f"Title: {r.title}\nURL: {r.url}\nSnippet: {r.snippet}" - for r in result.results - ] + [f"Title: {r.title}\nURL: {r.url}\nSnippet: {r.snippet}" for r in result.results] ) else: search_result_text = str(result) diff --git a/litellm/integrations/weights_biases.py b/litellm/integrations/weights_biases.py index 5f087fe219a..6d002ac4a37 100644 --- a/litellm/integrations/weights_biases.py +++ b/litellm/integrations/weights_biases.py @@ -23,9 +23,7 @@ try: def __getitem__(self, key: K) -> V: ... - def get( - self, key: K, default: Optional[V] = None - ) -> Optional[V]: ... # pragma: no cover + def get(self, key: K, default: Optional[V] = None) -> Optional[V]: ... # pragma: no cover class OpenAIRequestResponseResolver: def __call__( @@ -40,13 +38,9 @@ try: elif response["object"] == "text_completion": return self._resolve_completion(request, response, time_elapsed) elif response["object"] == "chat.completion": - return self._resolve_chat_completion( - request, response, time_elapsed - ) + return self._resolve_chat_completion(request, response, time_elapsed) else: - logger.debug( - f"Unknown OpenAI response object: {response['object']}" - ) + logger.debug(f"Unknown OpenAI response object: {response['object']}") except Exception as e: logger.warning(f"Failed to resolve request/response: {e}") return None @@ -88,13 +82,8 @@ try: time_elapsed: float, ) -> trace_tree.WBTraceTree: """Resolves the request and response objects for `openai.Edit`.""" - request_str = ( - f"\n\n**Instruction**: {request['instruction']}\n\n" - f"**Input**: {request['input']}\n" - ) - choices = [ - f"\n\n**Edited**: {choice['text']}\n" for choice in response["choices"] - ] + request_str = f"\n\n**Instruction**: {request['instruction']}\n\n**Input**: {request['input']}\n" + choices = [f"\n\n**Edited**: {choice['text']}\n" for choice in response["choices"]] return self._request_response_result_to_trace( request=request, @@ -112,10 +101,7 @@ try: ) -> trace_tree.WBTraceTree: """Resolves the request and response objects for `openai.Completion`.""" request_str = f"\n\n**Prompt**: {request['prompt']}\n" - choices = [ - f"\n\n**Completion**: {choice['text']}\n" - for choice in response["choices"] - ] + choices = [f"\n\n**Completion**: {choice['text']}\n" for choice in response["choices"]] return self._request_response_result_to_trace( request=request, @@ -184,13 +170,9 @@ class WeightsBiasesLogger: try: pass except Exception: - raise Exception( - "\033[91m wandb not installed, try running 'pip install wandb' to fix this error\033[0m" - ) + raise Exception("\033[91m wandb not installed, try running 'pip install wandb' to fix this error\033[0m") if imported_openAIResponse is False: - raise Exception( - "\033[91m wandb not installed, try running 'pip install wandb' to fix this error\033[0m" - ) + raise Exception("\033[91m wandb not installed, try running 'pip install wandb' to fix this error\033[0m") self.resolver = OpenAIRequestResponseResolver() def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): @@ -202,18 +184,14 @@ class WeightsBiasesLogger: run = wandb.init() print_verbose(response_obj) - trace = self.resolver( - kwargs, response_obj, (end_time - start_time).total_seconds() - ) + trace = self.resolver(kwargs, response_obj, (end_time - start_time).total_seconds()) if trace is not None and run is not None: run.log({"trace": trace}) if run is not None: run.finish() - print_verbose( - f"W&B Logging Logging - final response object: {response_obj}" - ) + print_verbose(f"W&B Logging Logging - final response object: {response_obj}") except Exception: print_verbose(f"W&B Logging Layer Error - {traceback.format_exc()}") pass diff --git a/litellm/interactions/agents/http_handler.py b/litellm/interactions/agents/http_handler.py index d45ca6f4346..394b0f72634 100644 --- a/litellm/interactions/agents/http_handler.py +++ b/litellm/interactions/agents/http_handler.py @@ -62,9 +62,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): api_base=litellm_params.get("api_base"), litellm_params=dict(litellm_params), ) - data = agents_api_config.transform_create_request( - name=name, litellm_params=dict(litellm_params) - ) + data = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params)) if extra_body: data.update(extra_body) @@ -78,9 +76,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): }, ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout or request_timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=data, timeout=timeout or request_timeout) except Exception as e: raise self._handle_error(e=e, provider_config=agents_api_config) @@ -88,9 +84,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): original_response=response.text, additional_args={"complete_input_dict": data}, ) - return agents_api_config.transform_create_response( - raw_response=response, name=name - ) + return agents_api_config.transform_create_response(raw_response=response, name=name) async def async_create_agent( self, @@ -111,9 +105,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): api_base=litellm_params.get("api_base"), litellm_params=dict(litellm_params), ) - data = agents_api_config.transform_create_request( - name=name, litellm_params=dict(litellm_params) - ) + data = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params)) if extra_body: data.update(extra_body) @@ -137,9 +129,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): original_response=response.text, additional_args={"complete_input_dict": data}, ) - return agents_api_config.transform_create_response( - raw_response=response, name=name - ) + return agents_api_config.transform_create_response(raw_response=response, name=name) # ------------------------------------------------------------------ # # LIST # @@ -208,9 +198,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): additional_args={"api_base": url, "headers": headers}, ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=agents_api_config) @@ -262,9 +250,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): raise self._handle_error(e=e, provider_config=agents_api_config) logging_obj.post_call(original_response=response.text, additional_args={}) - return agents_api_config.transform_get_response( - raw_response=response, name=name - ) + return agents_api_config.transform_get_response(raw_response=response, name=name) async def async_get_agent( self, @@ -291,16 +277,12 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): additional_args={"api_base": url, "headers": headers}, ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=agents_api_config) logging_obj.post_call(original_response=response.text, additional_args={}) - return agents_api_config.transform_get_response( - raw_response=response, name=name - ) + return agents_api_config.transform_get_response(raw_response=response, name=name) # ------------------------------------------------------------------ # # DELETE # @@ -342,16 +324,12 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): additional_args={"api_base": url, "headers": headers}, ) try: - response = sync_httpx_client.delete( - url=url, headers=headers, timeout=timeout or request_timeout - ) + response = sync_httpx_client.delete(url=url, headers=headers, timeout=timeout or request_timeout) except Exception as e: raise self._handle_error(e=e, provider_config=agents_api_config) logging_obj.post_call(original_response=response.text, additional_args={}) - return agents_api_config.transform_delete_response( - raw_response=response, name=name - ) + return agents_api_config.transform_delete_response(raw_response=response, name=name) async def async_delete_agent( self, @@ -378,16 +356,12 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): additional_args={"api_base": url, "headers": headers}, ) try: - response = await async_httpx_client.delete( - url=url, headers=headers, timeout=timeout or request_timeout - ) + response = await async_httpx_client.delete(url=url, headers=headers, timeout=timeout or request_timeout) except Exception as e: raise self._handle_error(e=e, provider_config=agents_api_config) logging_obj.post_call(original_response=response.text, additional_args={}) - return agents_api_config.transform_delete_response( - raw_response=response, name=name - ) + return agents_api_config.transform_delete_response(raw_response=response, name=name) # ------------------------------------------------------------------ # # LIST VERSIONS # @@ -434,9 +408,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): raise self._handle_error(e=e, provider_config=agents_api_config) logging_obj.post_call(original_response=response.text, additional_args={}) - return agents_api_config.transform_list_versions_response( - raw_response=response, name=name - ) + return agents_api_config.transform_list_versions_response(raw_response=response, name=name) async def async_list_agent_versions( self, @@ -463,16 +435,12 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): additional_args={"api_base": url, "headers": headers}, ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=agents_api_config) logging_obj.post_call(original_response=response.text, additional_args={}) - return agents_api_config.transform_list_versions_response( - raw_response=response, name=name - ) + return agents_api_config.transform_list_versions_response(raw_response=response, name=name) agents_http_handler = AgentsHTTPHandler() diff --git a/litellm/interactions/agents/main.py b/litellm/interactions/agents/main.py index f56c6f3ed5e..ce63332c1a6 100644 --- a/litellm/interactions/agents/main.py +++ b/litellm/interactions/agents/main.py @@ -165,9 +165,7 @@ def create( **kwargs: Forwarded to GenericLiteLLMParams (api_key, api_base, etc.). """ local_vars = locals() - custom_llm_provider = ( - custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" - ) + custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" try: _is_async = kwargs.pop("acreate_agent", False) is True if base_agent is not None: @@ -178,9 +176,7 @@ def create( kwargs["base_environment"] = base_environment kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params = GenericLiteLLMParams(**kwargs) - logging_obj = _make_logging_obj( - kwargs, name, custom_llm_provider, "create_agent", {} - ) + logging_obj = _make_logging_obj(kwargs, name, custom_llm_provider, "create_agent", {}) config = _get_agents_api_config(custom_llm_provider) return agents_http_handler.create_agent( agents_api_config=config, @@ -250,16 +246,12 @@ def list( ) -> Union[AgentListResponse, Coroutine[Any, Any, AgentListResponse]]: """Sync: List all agents on the provider side.""" local_vars = locals() - custom_llm_provider = ( - custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" - ) + custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" try: _is_async = kwargs.pop("alist_agents", False) is True kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params = GenericLiteLLMParams(**kwargs) - logging_obj = _make_logging_obj( - kwargs, "", custom_llm_provider, "list_agents", {} - ) + logging_obj = _make_logging_obj(kwargs, "", custom_llm_provider, "list_agents", {}) config = _get_agents_api_config(custom_llm_provider) return agents_http_handler.list_agents( agents_api_config=config, @@ -330,16 +322,12 @@ def get( ) -> Union[AgentCreateResponse, Coroutine[Any, Any, AgentCreateResponse]]: """Sync: Get a specific agent by name.""" local_vars = locals() - custom_llm_provider = ( - custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" - ) + custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" try: _is_async = kwargs.pop("aget_agent", False) is True kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params = GenericLiteLLMParams(**kwargs) - logging_obj = _make_logging_obj( - kwargs, name, custom_llm_provider, "get_agent", {"name": name} - ) + logging_obj = _make_logging_obj(kwargs, name, custom_llm_provider, "get_agent", {"name": name}) config = _get_agents_api_config(custom_llm_provider) return agents_http_handler.get_agent( agents_api_config=config, @@ -411,16 +399,12 @@ def delete( ) -> Union[AgentDeleteResult, Coroutine[Any, Any, AgentDeleteResult]]: """Sync: Delete a specific agent by name.""" local_vars = locals() - custom_llm_provider = ( - custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" - ) + custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" try: _is_async = kwargs.pop("adelete_agent", False) is True kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params = GenericLiteLLMParams(**kwargs) - logging_obj = _make_logging_obj( - kwargs, name, custom_llm_provider, "delete_agent", {"name": name} - ) + logging_obj = _make_logging_obj(kwargs, name, custom_llm_provider, "delete_agent", {"name": name}) config = _get_agents_api_config(custom_llm_provider) return agents_http_handler.delete_agent( agents_api_config=config, @@ -492,16 +476,12 @@ def list_versions( ) -> Union[AgentVersionsResponse, Coroutine[Any, Any, AgentVersionsResponse]]: """Sync: List versions of a specific agent.""" local_vars = locals() - custom_llm_provider = ( - custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" - ) + custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" try: _is_async = kwargs.pop("alist_agent_versions", False) is True kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params = GenericLiteLLMParams(**kwargs) - logging_obj = _make_logging_obj( - kwargs, name, custom_llm_provider, "list_agent_versions", {"name": name} - ) + logging_obj = _make_logging_obj(kwargs, name, custom_llm_provider, "list_agent_versions", {"name": name}) config = _get_agents_api_config(custom_llm_provider) return agents_http_handler.list_agent_versions( agents_api_config=config, diff --git a/litellm/interactions/http_handler.py b/litellm/interactions/http_handler.py index 695da2be89a..0e5769933fe 100644 --- a/litellm/interactions/http_handler.py +++ b/litellm/interactions/http_handler.py @@ -64,9 +64,7 @@ class _BaseHTTPHandler: litellm_params: GenericLiteLLMParams, client: Optional[HTTPHandler], ) -> HTTPHandler: - return client or _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + return client or _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) def _async_client( self, @@ -117,9 +115,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): Coroutine[ Any, Any, - Union[ - InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse] - ], + Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]], ], ]: """ @@ -144,9 +140,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): ) if client is None: - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -233,9 +227,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = None, stream: Optional[bool] = None, - ) -> Union[ - InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse] - ]: + ) -> Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]]: """ Create a new interaction (async version). """ @@ -382,9 +374,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): ) if client is None: - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -503,9 +493,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): ) if client is None: - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -626,9 +614,7 @@ class InteractionsHTTPHandler(_BaseHTTPHandler): ) if client is None: - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client diff --git a/litellm/interactions/litellm_responses_transformation/handler.py b/litellm/interactions/litellm_responses_transformation/handler.py index b121ee37de6..4b108ee47d7 100644 --- a/litellm/interactions/litellm_responses_transformation/handler.py +++ b/litellm/interactions/litellm_responses_transformation/handler.py @@ -123,9 +123,7 @@ class LiteLLMResponsesInteractionsHandler: input: Optional[InteractionInput], optional_params: InteractionsAPIOptionalRequestParams, **kwargs, - ) -> Union[ - InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse] - ]: + ) -> Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]]: """Async handler for interactions API requests.""" # Call litellm.aresponses() # Note: litellm.aresponses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] diff --git a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py index 4a3eb63084e..6b10a36c179 100644 --- a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py +++ b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py @@ -92,9 +92,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: # Event builders # ------------------------------------------------------------------ - def _build_interaction_start_event( - self, interaction_id: str - ) -> InteractionsAPIStreamingResponse: + def _build_interaction_start_event(self, interaction_id: str) -> InteractionsAPIStreamingResponse: event_type = "interaction.start" if self._use_legacy else "interaction.created" return InteractionsAPIStreamingResponse( event_type=event_type, @@ -104,9 +102,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: model=self.model, ) - def _build_content_start_event( - self, interaction_id: str - ) -> InteractionsAPIStreamingResponse: + def _build_content_start_event(self, interaction_id: str) -> InteractionsAPIStreamingResponse: if self._use_legacy: return InteractionsAPIStreamingResponse( event_type="content.start", @@ -120,9 +116,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: step={"type": "model_output", "content": []}, ) - def _build_text_delta_event( - self, interaction_id: str, delta_text: str - ) -> InteractionsAPIStreamingResponse: + def _build_text_delta_event(self, interaction_id: str, delta_text: str) -> InteractionsAPIStreamingResponse: if self._use_legacy: return InteractionsAPIStreamingResponse( event_type="content.delta", @@ -136,9 +130,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: delta={"type": "text", "text": delta_text}, ) - def _build_content_stop_event( - self, interaction_id: Optional[str] - ) -> InteractionsAPIStreamingResponse: + def _build_content_stop_event(self, interaction_id: Optional[str]) -> InteractionsAPIStreamingResponse: if self._use_legacy: return InteractionsAPIStreamingResponse( event_type="content.stop", @@ -151,9 +143,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: index=0, ) - def _build_completion_event( - self, response_id: str - ) -> InteractionsAPIStreamingResponse: + def _build_completion_event(self, response_id: str) -> InteractionsAPIStreamingResponse: if self._use_legacy: return InteractionsAPIStreamingResponse( event_type="interaction.complete", @@ -197,13 +187,9 @@ class LiteLLMResponsesInteractionsStreamingIterator: # Text delta: emit any missing start events, then the delta itself. if isinstance(responses_chunk, OutputTextDeltaEvent): - delta_text = ( - responses_chunk.delta if isinstance(responses_chunk.delta, str) else "" - ) + delta_text = responses_chunk.delta if isinstance(responses_chunk.delta, str) else "" self.collected_text += delta_text - interaction_id = ( - getattr(responses_chunk, "item_id", None) or f"interaction_{id(self)}" - ) + interaction_id = getattr(responses_chunk, "item_id", None) or f"interaction_{id(self)}" if self._interaction_id is None: self._interaction_id = interaction_id @@ -223,9 +209,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: if not self.sent_interaction_start: self.sent_interaction_start = True response_id = ( - getattr(responses_chunk.response, "id", None) - if hasattr(responses_chunk, "response") - else None + getattr(responses_chunk.response, "id", None) if hasattr(responses_chunk, "response") else None ) or f"interaction_{id(self)}" if self._interaction_id is None: self._interaction_id = response_id @@ -241,11 +225,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: if isinstance(responses_chunk, ResponseCompletedEvent): self.finished = True response = responses_chunk.response - response_id = ( - self._interaction_id - or getattr(response, "id", None) - or f"interaction_{id(self)}" - ) + response_id = self._interaction_id or getattr(response, "id", None) or f"interaction_{id(self)}" terminal: List[InteractionsAPIStreamingResponse] = [] if self.sent_content_start: @@ -290,9 +270,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: if self.finished: raise StopIteration - sync_iterator = cast( - SyncResponsesAPIStreamingIterator, self.responses_stream_iterator - ) + sync_iterator = cast(SyncResponsesAPIStreamingIterator, self.responses_stream_iterator) while True: try: chunk = next(sync_iterator) @@ -318,9 +296,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: if self.finished: raise StopAsyncIteration - async_iterator = cast( - ResponsesAPIStreamingIterator, self.responses_stream_iterator - ) + async_iterator = cast(ResponsesAPIStreamingIterator, self.responses_stream_iterator) while True: try: chunk = await async_iterator.__anext__() diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 0ff1a97cd0b..a2d8ebc5d4c 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -46,9 +46,7 @@ class LiteLLMResponsesInteractionsConfig: # Transform input if input is not None: responses_request["input"] = ( - LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - input - ) + LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input(input) ) # Transform system_instruction -> instructions @@ -71,9 +69,7 @@ class LiteLLMResponsesInteractionsConfig: # Responses API doesn't have top_k, skip it pass if "max_output_tokens" in generation_config: - responses_request["max_output_tokens"] = generation_config[ - "max_output_tokens" - ] + responses_request["max_output_tokens"] = generation_config["max_output_tokens"] # Pass through other optional params that match passthrough_params = ["stream", "store", "metadata", "user"] @@ -115,11 +111,7 @@ class LiteLLMResponsesInteractionsConfig: content = turn.get("content", []) # Transform content array - transformed_content = ( - LiteLLMResponsesInteractionsConfig._transform_content_array( - content - ) - ) + transformed_content = LiteLLMResponsesInteractionsConfig._transform_content_array(content) messages.append( { @@ -141,11 +133,7 @@ class LiteLLMResponsesInteractionsConfig: else: content_list = [] - transformed_content = ( - LiteLLMResponsesInteractionsConfig._transform_content_array( - content_list - ) - ) + transformed_content = LiteLLMResponsesInteractionsConfig._transform_content_array(content_list) messages.append( { @@ -164,9 +152,7 @@ class LiteLLMResponsesInteractionsConfig: { "role": "user", "content": LiteLLMResponsesInteractionsConfig._transform_content_array( - input.get("content", []) - if isinstance(input.get("content"), list) - else [input] + input.get("content", []) if isinstance(input.get("content"), list) else [input] ), } ], @@ -244,10 +230,7 @@ class LiteLLMResponsesInteractionsConfig: # of `outputs` / `steps` don't leak into the other. outputs.append({"type": "text", "text": text}) model_output_contents.append({"type": "text", "text": text}) - elif ( - isinstance(content_item, dict) - and content_item.get("type") == "text" - ): + elif isinstance(content_item, dict) and content_item.get("type") == "text": outputs.append({**content_item}) model_output_contents.append({**content_item}) if model_output_contents: diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index d99cc3d11c7..8634269ee94 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -134,9 +134,7 @@ async def acreate( kwargs["acreate_interaction"] = True if custom_llm_provider is None and model: - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) - ) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model, api_base=kwargs.get("api_base", None)) elif custom_llm_provider is None: custom_llm_provider = "gemini" @@ -290,18 +288,11 @@ def create( # Get optional params using utility (similar to responses API pattern) local_vars.update(kwargs) - optional_params = ( - InteractionsAPIRequestUtils.get_requested_interactions_api_optional_params( - local_vars - ) - ) + optional_params = InteractionsAPIRequestUtils.get_requested_interactions_api_optional_params(local_vars) # Check if this is a bridge provider (litellm_responses) - similar to responses API # Either provider is explicitly "litellm_responses" or no config found (bridge to responses) - if ( - custom_llm_provider == "litellm_responses" - or interactions_api_config is None - ): + if custom_llm_provider == "litellm_responses" or interactions_api_config is None: # Bridge to litellm.responses() for non-native providers from litellm.interactions.litellm_responses_transformation.handler import ( LiteLLMResponsesInteractionsHandler, @@ -425,9 +416,7 @@ def get( ) if interactions_api_config is None: - raise ValueError( - f"Interactions API not supported for: {custom_llm_provider}" - ) + raise ValueError(f"Interactions API not supported for: {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -529,9 +518,7 @@ def delete( ) if interactions_api_config is None: - raise ValueError( - f"Interactions API not supported for: {custom_llm_provider}" - ) + raise ValueError(f"Interactions API not supported for: {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -633,9 +620,7 @@ def cancel( ) if interactions_api_config is None: - raise ValueError( - f"Interactions API not supported for: {custom_llm_provider}" - ) + raise ValueError(f"Interactions API not supported for: {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index 561686a3e1b..45c5443cfd2 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -57,20 +57,14 @@ class BaseInteractionsAPIStreamingIterator: # set hidden params for response headers _api_base = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get( - "litellm_params", {} - ), - ) - _model_info: Dict = ( - litellm_metadata.get("model_info", {}) if litellm_metadata else {} + optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) + _model_info: Dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} self._hidden_params = { "model_id": _model_info.get("id", None), "api_base": _api_base, } - self._hidden_params["additional_headers"] = process_response_headers( - self.response.headers or {} - ) + self._hidden_params["additional_headers"] = process_response_headers(self.response.headers or {}) def _process_chunk(self, chunk: str) -> Optional[InteractionsAPIStreamingResponse]: """Process a single chunk of data from the stream.""" @@ -93,12 +87,10 @@ class BaseInteractionsAPIStreamingIterator: # Format as InteractionsAPIStreamingResponse if isinstance(parsed_chunk, dict): - streaming_response = ( - self.interactions_api_config.transform_streaming_response( - model=self.model, - parsed_chunk=parsed_chunk, - logging_obj=self.logging_obj, - ) + streaming_response = self.interactions_api_config.transform_streaming_response( + model=self.model, + parsed_chunk=parsed_chunk, + logging_obj=self.logging_obj, ) # Store the completed response. @@ -107,8 +99,7 @@ class BaseInteractionsAPIStreamingIterator: # Remove the legacy check after June 8, 2026. if streaming_response and ( getattr(streaming_response, "status", None) == "completed" - or getattr(streaming_response, "event_type", None) - == "interaction.completed" + or getattr(streaming_response, "event_type", None) == "interaction.completed" ): self.completed_response = streaming_response self._handle_logging_completed_response() @@ -118,9 +109,7 @@ class BaseInteractionsAPIStreamingIterator: return None except json.JSONDecodeError: # If we can't parse the chunk, continue - verbose_logger.debug( - f"Failed to parse streaming chunk: {stripped_chunk[:200]}..." - ) + verbose_logger.debug(f"Failed to parse streaming chunk: {stripped_chunk[:200]}...") return None def _handle_logging_completed_response(self): diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py index 84437f4d3d8..3dffaa538ba 100644 --- a/litellm/interactions/utils.py +++ b/litellm/interactions/utils.py @@ -72,17 +72,13 @@ class InteractionsAPIRequestUtils: special_params = params.pop("kwargs", {}) additional_drop_params = params.pop("additional_drop_params", None) - non_default_params = ( - PreProcessNonDefaultParams.base_pre_process_non_default_params( - passed_params=params, - special_params=special_params, - custom_llm_provider=custom_llm_provider, - additional_drop_params=additional_drop_params, - default_param_values={ - k: None for k in INTERACTIONS_API_OPTIONAL_PARAMS - }, - additional_endpoint_specific_params=["input", "model", "agent"], - ) + non_default_params = PreProcessNonDefaultParams.base_pre_process_non_default_params( + passed_params=params, + special_params=special_params, + custom_llm_provider=custom_llm_provider, + additional_drop_params=additional_drop_params, + default_param_values={k: None for k in INTERACTIONS_API_OPTIONAL_PARAMS}, + additional_endpoint_specific_params=["input", "model", "agent"], ) return cast(InteractionsAPIOptionalRequestParams, non_default_params) diff --git a/litellm/litellm_core_utils/asyncify.py b/litellm/litellm_core_utils/asyncify.py index 8d56a1bbe2a..09585171147 100644 --- a/litellm/litellm_core_utils/asyncify.py +++ b/litellm/litellm_core_utils/asyncify.py @@ -45,9 +45,7 @@ def asyncify( and returns the result. """ - async def wrapper( - *args: T_ParamSpec.args, **kwargs: T_ParamSpec.kwargs - ) -> T_Retval: + async def wrapper(*args: T_ParamSpec.args, **kwargs: T_ParamSpec.kwargs) -> T_Retval: partial_f = functools.partial(function, *args, **kwargs) # In `v4.1.0` anyio added the `abandon_on_cancel` argument and deprecated the old diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 82f5c27f836..f86243c73b7 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -96,9 +96,7 @@ def process_audio_file(audio_file: FileTypes) -> ProcessedAudioFile: raise ValueError(f"Unsupported content type in tuple: {type(content)}") else: raise ValueError("Tuple must have at least 2 elements: (filename, content)") - elif hasattr(audio_file, "read") and not isinstance( - audio_file, (str, bytes, bytearray, tuple, os.PathLike) - ): + elif hasattr(audio_file, "read") and not isinstance(audio_file, (str, bytes, bytearray, tuple, os.PathLike)): # File-like object (IO) - check this after all other types filename = getattr(audio_file, "name", "audio.wav") file_content = audio_file.read() # type: ignore @@ -122,9 +120,7 @@ def process_audio_file(audio_file: FileTypes) -> ProcessedAudioFile: # If extension is not recognized, fallback to audio/wav content_type = "audio/wav" - return ProcessedAudioFile( - file_content=file_content, filename=filename, content_type=content_type - ) + return ProcessedAudioFile(file_content=file_content, filename=filename, content_type=content_type) def get_audio_file_name(file_obj: FileTypes) -> str: @@ -184,11 +180,7 @@ def get_audio_file_content_hash(file_obj: FileTypes) -> str: file_content = None elif hasattr(file_content_obj, "read"): try: - current_position = ( - file_content_obj.tell() - if hasattr(file_content_obj, "tell") - else None - ) + current_position = file_content_obj.tell() if hasattr(file_content_obj, "tell") else None if hasattr(file_content_obj, "seek"): file_content_obj.seek(0) file_content = file_content_obj.read() # type: ignore @@ -270,9 +262,7 @@ def calculate_request_duration(file: FileTypes) -> Optional[float]: content = file[1] if isinstance(content, bytes): file_content = content - elif hasattr(content, "read") and not isinstance( - content, (str, os.PathLike) - ): + elif hasattr(content, "read") and not isinstance(content, (str, os.PathLike)): # File-like object in tuple current_pos = getattr(content, "tell", lambda: None)() # Seek to start to ensure we read the entire content diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 938e892bd50..828605d5ef8 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -54,11 +54,7 @@ def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[st depth = _coerce_int(kwargs.get("_agentic_loop_depth"), 0) max_loops = max(_coerce_int(kwargs.get("max_agentic_loops"), 3), 1) raw_fingerprints = kwargs.get("_agentic_loop_fingerprints") - fingerprints = ( - [str(fp) for fp in raw_fingerprints] - if isinstance(raw_fingerprints, list) - else [] - ) + fingerprints = [str(fp) for fp in raw_fingerprints] if isinstance(raw_fingerprints, list) else [] return depth, max_loops, fingerprints @@ -78,9 +74,7 @@ def _check_agentic_loop_safety( ) -> str: fingerprint = _fingerprint_tools(tool_calls) if fingerprint in fingerprints: - raise ValueError( - "Agentic loop detected repeated tool-call fingerprint; aborting rerun" - ) + raise ValueError("Agentic loop detected repeated tool-call fingerprint; aborting rerun") if depth >= max_loops: raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}") return fingerprint @@ -102,11 +96,7 @@ def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None: metadata = kwargs_for_followup.get("litellm_metadata") metadata = dict(metadata) if isinstance(metadata, dict) else {} for key, value in kwargs_for_followup.items(): - if ( - key.startswith("_agentic_loop") - or key == "max_agentic_loops" - or is_interception_internal_key(key) - ): + if key.startswith("_agentic_loop") or key == "max_agentic_loops" or is_interception_internal_key(key): metadata[key] = value kwargs_for_followup["litellm_metadata"] = metadata @@ -115,9 +105,7 @@ def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]: return { k: v for k, v in source.items() - if not is_interception_internal_key( - k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES - ) + if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES) and k not in _FOLLOWUP_INTERNAL_PARAMS } @@ -154,11 +142,7 @@ async def _execute_chat_completion_agentic_plan( kwargs_for_followup = _filter_followup_kwargs(kwargs) kwargs_for_followup.update( - { - k: v - for k, v in _filter_followup_kwargs(patch.kwargs).items() - if k not in optional_params_for_followup - } + {k: v for k, v in _filter_followup_kwargs(patch.kwargs).items() if k not in optional_params_for_followup} ) kwargs_for_followup["_agentic_loop_depth"] = depth + 1 kwargs_for_followup["max_agentic_loops"] = max_loops @@ -174,10 +158,8 @@ async def _execute_chat_completion_agentic_plan( ) if _post_hook_overridden(callback): try: - response_followup = ( - await callback.async_post_agentic_loop_response_hook( - response=response_followup, plan=plan, kwargs=kwargs - ) + response_followup = await callback.async_post_agentic_loop_response_hook( + response=response_followup, plan=plan, kwargs=kwargs ) except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") @@ -197,8 +179,7 @@ async def _execute_chat_completion_agentic_plan( except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in " - "async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", + "LiteLLM.AgenticHookError: Exception in async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", _call_id, model, str(e), @@ -218,9 +199,7 @@ async def maybe_run_chat_completion_agentic_loop( ) -> ModelResponse | CustomStreamWrapper | None: import litellm - callbacks = litellm.callbacks + ( - getattr(logging_obj, "dynamic_success_callbacks", None) or [] - ) + callbacks = litellm.callbacks + (getattr(logging_obj, "dynamic_success_callbacks", None) or []) depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) tools = optional_params.get("tools", []) @@ -320,11 +299,7 @@ async def maybe_run_chat_completion_agentic_loop( str(e), ) - if ( - kwargs.get("_code_interpreter_interception_converted_stream") - and not depth - and hasattr(response, "choices") - ): + if kwargs.get("_code_interpreter_interception_converted_stream") and not depth and hasattr(response, "choices"): return cast( "ModelResponse | CustomStreamWrapper", _wrap_response_as_fake_stream(response), diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index a75d1178d5a..a62dfe61805 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -27,17 +27,13 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool: retrieved through their managed unified file id so owner/team access is enforced; a raw URI supplied by a caller bypasses that check. """ - return isinstance(file_id, str) and file_id.startswith( - MANAGED_CLOUD_STORAGE_SCHEMES - ) + return isinstance(file_id, str) and file_id.startswith(MANAGED_CLOUD_STORAGE_SCHEMES) _SAFE_OBJECT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9._-]+") -def sanitize_cloud_object_component( - value: Optional[str], fallback: str = "file" -) -> str: +def sanitize_cloud_object_component(value: Optional[str], fallback: str = "file") -> str: if not isinstance(value, str): return fallback @@ -45,9 +41,7 @@ def sanitize_cloud_object_component( if component in {"", ".", ".."}: return fallback - component = "".join( - "_" if ord(char) < 32 or ord(char) == 127 else char for char in component - ) + component = "".join("_" if ord(char) < 32 or ord(char) == 127 else char for char in component) component = _SAFE_OBJECT_COMPONENT_PATTERN.sub("_", component) component = component.strip("._") if not component: @@ -70,12 +64,8 @@ def sanitize_cloud_object_path(value: Optional[str], fallback: str = "file") -> return "/".join(segments) -def build_managed_cloud_object_name( - prefix: str, filename: Optional[str], fallback_filename: str = "file" -) -> str: - safe_filename = sanitize_cloud_object_component( - filename, fallback=fallback_filename - ) +def build_managed_cloud_object_name(prefix: str, filename: Optional[str], fallback_filename: str = "file") -> str: + safe_filename = sanitize_cloud_object_component(filename, fallback=fallback_filename) return f"{prefix}{uuid.uuid4().hex}-{safe_filename}" @@ -99,9 +89,7 @@ def split_configured_cloud_bucket_name(bucket_name: str) -> Tuple[str, str]: bucket_name = bucket_name.strip() if "://" in bucket_name or "?" in bucket_name or "#" in bucket_name: - raise ValueError( - "Cloud storage bucket name must not include a URI scheme or query" - ) + raise ValueError("Cloud storage bucket name must not include a URI scheme or query") if any(ord(char) < 32 or ord(char) == 127 for char in bucket_name): raise ValueError("Cloud storage bucket name contains control characters") @@ -131,13 +119,9 @@ def should_allow_legacy_cloud_file_ids( ) -> bool: value = None if isinstance(litellm_params, Mapping): - trusted_model_credentials = litellm_params.get( - "_litellm_internal_model_credentials" - ) + trusted_model_credentials = litellm_params.get("_litellm_internal_model_credentials") if isinstance(trusted_model_credentials, _MAPPING_PROXY_TYPE): - value = cast(Mapping[str, Any], trusted_model_credentials).get( - "allow_legacy_cloud_file_ids" - ) + value = cast(Mapping[str, Any], trusted_model_credentials).get("allow_legacy_cloud_file_ids") if isinstance(value, bool): return value @@ -162,29 +146,21 @@ def validate_managed_cloud_file_id( raise ValueError("file_id must include a cloud storage object name") bucket_name, object_name = full_path.split("/", 1) - configured_bucket, configured_prefix = split_configured_cloud_bucket_name( - configured_bucket_name - ) + configured_bucket, configured_prefix = split_configured_cloud_bucket_name(configured_bucket_name) if bucket_name != configured_bucket: raise ValueError("file_id bucket does not match the configured storage bucket") _validate_cloud_object_path(object_name) allowed_prefixes = tuple(allowed_object_prefixes) if configured_prefix: - allowed_prefixes = tuple( - f"{configured_prefix.rstrip('/')}/{prefix}" for prefix in allowed_prefixes - ) + allowed_prefixes = tuple(f"{configured_prefix.rstrip('/')}/{prefix}" for prefix in allowed_prefixes) if object_name.startswith(allowed_prefixes): return bucket_name, object_name if allow_legacy_cloud_file_ids: - if configured_prefix and not object_name.startswith( - f"{configured_prefix.rstrip('/')}/" - ): - raise ValueError( - "file_id object does not match the configured storage prefix" - ) + if configured_prefix and not object_name.startswith(f"{configured_prefix.rstrip('/')}/"): + raise ValueError("file_id object does not match the configured storage prefix") return bucket_name, object_name raise ValueError("file_id must reference a LiteLLM-managed storage object") diff --git a/litellm/litellm_core_utils/completion_timeout.py b/litellm/litellm_core_utils/completion_timeout.py index 70c6896a323..794749a39bf 100644 --- a/litellm/litellm_core_utils/completion_timeout.py +++ b/litellm/litellm_core_utils/completion_timeout.py @@ -56,18 +56,12 @@ class CompletionTimeout: elif kwargs.get("request_timeout") is not None: resolved = kwargs["request_timeout"] else: - resolved = CompletionTimeout._fallback_when_no_explicit_timeout( - global_timeout - ) + resolved = CompletionTimeout._fallback_when_no_explicit_timeout(global_timeout) - if isinstance(resolved, httpx.Timeout) and not supports_httpx_timeout( - custom_llm_provider - ): + if isinstance(resolved, httpx.Timeout) and not supports_httpx_timeout(custom_llm_provider): read_timeout = resolved.read resolved = ( - float(read_timeout) - if read_timeout is not None - else COMPLETION_HTTP_FALLBACK_SECONDS + float(read_timeout) if read_timeout is not None else COMPLETION_HTTP_FALLBACK_SECONDS ) # default 10 min timeout elif not isinstance(resolved, httpx.Timeout): resolved = float(resolved) # type: ignore diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 0cdd721598c..002a46771e3 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -18,9 +18,7 @@ else: Span = Any -def safe_divide_seconds( - seconds: float, denominator: float, default: Optional[float] = None -) -> Optional[float]: +def safe_divide_seconds(seconds: float, denominator: float, default: Optional[float] = None) -> Optional[float]: """ Safely divide seconds by denominator, handling zero division. @@ -109,9 +107,7 @@ _FINISH_REASON_MAP: dict[str, OpenAIChatCompletionFinishReason] = { def map_finish_reason(finish_reason: str) -> OpenAIChatCompletionFinishReason: mapped = _FINISH_REASON_MAP.get(finish_reason) if mapped is None: - verbose_logger.warning( - "Unmapped finish_reason '%s', defaulting to 'stop'", finish_reason - ) + verbose_logger.warning("Unmapped finish_reason '%s', defaulting to 'stop'", finish_reason) return "stop" return mapped @@ -124,9 +120,7 @@ def remove_index_from_tool_calls( _tool_calls = message.get("tool_calls") if _tool_calls is not None and isinstance(_tool_calls, list): for tool_call in _tool_calls: - if ( - isinstance(tool_call, dict) and "index" in tool_call - ): # Type guard to ensure it's a dict + if isinstance(tool_call, dict) and "index" in tool_call: # Type guard to ensure it's a dict tool_call.pop("index", None) return @@ -141,9 +135,7 @@ def remove_items_at_indices(items: Optional[List[Any]], indices: Iterable[int]) items.pop(index) -def add_missing_spend_metadata_to_litellm_metadata( - litellm_metadata: dict, metadata: dict -) -> dict: +def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metadata: dict) -> dict: """ Helper to get litellm metadata for spend tracking @@ -185,9 +177,7 @@ def get_litellm_metadata_from_kwargs(kwargs: dict): metadata = litellm_params.get("metadata", {}) litellm_metadata = litellm_params.get("litellm_metadata", {}) if litellm_metadata and metadata: - litellm_metadata = add_missing_spend_metadata_to_litellm_metadata( - litellm_metadata, metadata - ) + litellm_metadata = add_missing_spend_metadata_to_litellm_metadata(litellm_metadata, metadata) if litellm_metadata: return litellm_metadata elif metadata: @@ -236,9 +226,7 @@ def _get_parent_otel_span_from_kwargs( return kwargs["litellm_parent_otel_span"] return None except Exception as e: - verbose_logger.exception( - "Error in _get_parent_otel_span_from_kwargs: " + str(e) - ) + verbose_logger.exception("Error in _get_parent_otel_span_from_kwargs: " + str(e)) return None @@ -271,9 +259,7 @@ def process_response_headers( for k, v in response_headers.items(): if k in OPENAI_RESPONSE_HEADERS: # return openai-compatible headers openai_headers[k] = v - if k.startswith( - "llm_provider-" - ): # return raw provider headers (incl. openai-compatible ones) + if k.startswith("llm_provider-"): # return raw provider headers (incl. openai-compatible ones) processed_headers[k] = v elif _preserve and k.startswith("x-litellm-"): # LiteLLM's own internal headers (e.g. x-litellm-attempted-fallbacks, @@ -330,13 +316,8 @@ def safe_deep_copy(data): if "metadata" in data and "litellm_parent_otel_span" in data["metadata"]: litellm_parent_otel_span = data["metadata"].pop("litellm_parent_otel_span") data["metadata"]["litellm_parent_otel_span"] = "placeholder" - if ( - "litellm_metadata" in data - and "litellm_parent_otel_span" in data["litellm_metadata"] - ): - litellm_parent_otel_span = data["litellm_metadata"].pop( - "litellm_parent_otel_span" - ) + if "litellm_metadata" in data and "litellm_parent_otel_span" in data["litellm_metadata"]: + litellm_parent_otel_span = data["litellm_metadata"].pop("litellm_parent_otel_span") data["litellm_metadata"]["litellm_parent_otel_span"] = "placeholder" # Step 2: Per-key deepcopy with fallback @@ -357,13 +338,8 @@ def safe_deep_copy(data): if isinstance(data, dict) and litellm_parent_otel_span is not None: if "metadata" in data and "litellm_parent_otel_span" in data["metadata"]: data["metadata"]["litellm_parent_otel_span"] = litellm_parent_otel_span - if ( - "litellm_metadata" in data - and "litellm_parent_otel_span" in data["litellm_metadata"] - ): - data["litellm_metadata"]["litellm_parent_otel_span"] = ( - litellm_parent_otel_span - ) + if "litellm_metadata" in data and "litellm_parent_otel_span" in data["litellm_metadata"]: + data["litellm_metadata"]["litellm_parent_otel_span"] = litellm_parent_otel_span return new_data @@ -416,9 +392,7 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: result_list: list[Any] = [] for item in data: # Skip exception and callable items - if isinstance(item, Exception) or ( - callable(item) and not isinstance(item, type) - ): + if isinstance(item, Exception) or (callable(item) and not isinstance(item, type)): continue try: filtered = filter_exceptions_from_params(item, max_depth - 1) @@ -432,9 +406,7 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: return data -def filter_internal_params( - data: dict, additional_internal_params: Optional[set] = None -) -> dict: +def filter_internal_params(data: dict, additional_internal_params: Optional[set] = None) -> dict: """ Filter out LiteLLM internal parameters that shouldn't be sent to provider APIs. diff --git a/litellm/litellm_core_utils/default_encoding.py b/litellm/litellm_core_utils/default_encoding.py index f58b90c8e72..38aacb47f04 100644 --- a/litellm/litellm_core_utils/default_encoding.py +++ b/litellm/litellm_core_utils/default_encoding.py @@ -11,9 +11,7 @@ except (ImportError, AttributeError): # Old way to access resources, which setuptools deprecated some time ago import pkg_resources # type: ignore - filename = pkg_resources.resource_filename( - __name__, "litellm_core_utils/tokenizers" - ) + filename = pkg_resources.resource_filename(__name__, "litellm_core_utils/tokenizers") # Always default TIKTOKEN_CACHE_DIR to the bundled tokenizers directory # unless the user explicitly overrides it via CUSTOM_TIKTOKEN_CACHE_DIR. diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 65810e83c66..85abbdddffc 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -28,9 +28,7 @@ from typing import Any, Dict, List, Optional, TypeVar, Union T = TypeVar("T") -def get_nested_value( - data: Dict[str, Any], key_path: str, default: Optional[T] = None -) -> Optional[T]: +def get_nested_value(data: Dict[str, Any], key_path: str, default: Optional[T] = None) -> Optional[T]: """ Retrieves a value from a nested dictionary using dot notation. @@ -56,11 +54,7 @@ def get_nested_value( return default # Remove metadata. prefix if it exists - key_path = ( - key_path.replace("metadata.", "", 1) - if key_path.startswith("metadata.") - else key_path - ) + key_path = key_path.replace("metadata.", "", 1) if key_path.startswith("metadata.") else key_path # Split the key path into parts, respecting escaped dots (\.) # Use a temporary placeholder, split on unescaped dots, then restore @@ -158,9 +152,7 @@ def _delete_nested_value_custom( # Only recurse if element is a dict or list (nested structure) element = data[index] if isinstance(element, (dict, list)): - _delete_nested_value_custom( - element, segments, segment_index + 1 - ) + _delete_nested_value_custom(element, segments, segment_index + 1) except (ValueError, IndexError): # Invalid index, skip pass @@ -174,23 +166,15 @@ def _delete_nested_value_custom( else: # Navigate deeper if segment in data: - next_segment = ( - segments[segment_index + 1] - if segment_index + 1 < len(segments) - else None - ) + next_segment = segments[segment_index + 1] if segment_index + 1 < len(segments) else None # If next segment is array notation, current field should be list if next_segment and (next_segment.startswith("[")): if isinstance(data[segment], list): - _delete_nested_value_custom( - data[segment], segments, segment_index + 1 - ) + _delete_nested_value_custom(data[segment], segments, segment_index + 1) # Otherwise navigate into dict elif isinstance(data[segment], dict): - _delete_nested_value_custom( - data[segment], segments, segment_index + 1 - ) + _delete_nested_value_custom(data[segment], segments, segment_index + 1) def delete_nested_value( diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 036d691c686..438ff5600ba 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -94,9 +94,7 @@ def duration_in_seconds(duration: str) -> int: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") -def get_next_standardized_reset_time( - duration: str, current_time: datetime, timezone_str: str = "UTC" -) -> datetime: +def get_next_standardized_reset_time(duration: str, current_time: datetime, timezone_str: str = "UTC") -> datetime: """ Get the next standardized reset time based on the duration. @@ -121,9 +119,7 @@ def get_next_standardized_reset_time( value, unit = _parse_duration(duration) if value is None: # Fall back to default if format is invalid - return current_time.replace( - hour=0, minute=0, second=0, microsecond=0 - ) + timedelta(days=1) + return current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) # Midnight of the current day in the specified timezone base_midnight = current_time.replace(hour=0, minute=0, second=0, microsecond=0) @@ -146,9 +142,7 @@ def get_next_standardized_reset_time( return base_midnight + timedelta(days=1) -def _setup_timezone( - current_time: datetime, timezone_str: str = "UTC" -) -> Tuple[datetime, tzinfo]: +def _setup_timezone(current_time: datetime, timezone_str: str = "UTC") -> Tuple[datetime, tzinfo]: """Set up timezone and normalize current time to that timezone.""" try: if timezone_str is None: @@ -181,9 +175,7 @@ def _parse_duration(duration: str) -> Tuple[Optional[int], Optional[str]]: return int(value), unit -def _handle_day_reset( - current_time: datetime, base_midnight: datetime, value: int, tz: tzinfo -) -> datetime: +def _handle_day_reset(current_time: datetime, base_midnight: datetime, value: int, tz: tzinfo) -> datetime: """Handle day-based reset times.""" # Handle zero value - immediate expiration if value == 0: @@ -222,14 +214,10 @@ def _handle_day_reset( ) return next_reset else: # Custom day value - next interval is value days from current - return current_time.replace( - hour=0, minute=0, second=0, microsecond=0 - ) + timedelta(days=value) + return current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=value) -def _handle_hour_reset( - current_time: datetime, base_midnight: datetime, value: int -) -> datetime: +def _handle_hour_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: """Handle hour-based reset times.""" # Handle zero value - immediate expiration if value == 0: @@ -242,17 +230,9 @@ def _handle_hour_reset( # Calculate next hour aligned with the value if current_minute == 0 and current_second == 0 and current_microsecond == 0: - next_hour = ( - current_hour + value - (current_hour % value) - if current_hour % value != 0 - else current_hour + value - ) + next_hour = current_hour + value - (current_hour % value) if current_hour % value != 0 else current_hour + value else: - next_hour = ( - current_hour + value - (current_hour % value) - if current_hour % value != 0 - else current_hour + value - ) + next_hour = current_hour + value - (current_hour % value) if current_hour % value != 0 else current_hour + value # Handle overnight case if next_hour >= 24: @@ -263,9 +243,7 @@ def _handle_hour_reset( return current_time.replace(hour=next_hour, minute=0, second=0, microsecond=0) -def _handle_minute_reset( - current_time: datetime, base_midnight: datetime, value: int -) -> datetime: +def _handle_minute_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: """Handle minute-based reset times.""" # Handle zero value - immediate expiration if value == 0: @@ -279,15 +257,11 @@ def _handle_minute_reset( # Calculate next minute aligned with the value if current_second == 0 and current_microsecond == 0: next_minute = ( - current_minute + value - (current_minute % value) - if current_minute % value != 0 - else current_minute + value + current_minute + value - (current_minute % value) if current_minute % value != 0 else current_minute + value ) else: next_minute = ( - current_minute + value - (current_minute % value) - if current_minute % value != 0 - else current_minute + value + current_minute + value - (current_minute % value) if current_minute % value != 0 else current_minute + value ) # Handle hour rollover @@ -298,18 +272,12 @@ def _handle_minute_reset( if next_hour >= 24: next_hour = next_hour % 24 next_day = base_midnight + timedelta(days=1) - return next_day.replace( - hour=next_hour, minute=next_minute, second=0, microsecond=0 - ) + return next_day.replace(hour=next_hour, minute=next_minute, second=0, microsecond=0) - return current_time.replace( - hour=next_hour, minute=next_minute, second=0, microsecond=0 - ) + return current_time.replace(hour=next_hour, minute=next_minute, second=0, microsecond=0) -def _handle_second_reset( - current_time: datetime, base_midnight: datetime, value: int -) -> datetime: +def _handle_second_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: """Handle second-based reset times.""" # Handle zero value - immediate expiration if value == 0: @@ -323,15 +291,11 @@ def _handle_second_reset( # Calculate next second aligned with the value if current_microsecond == 0: next_second = ( - current_second + value - (current_second % value) - if current_second % value != 0 - else current_second + value + current_second + value - (current_second % value) if current_second % value != 0 else current_second + value ) else: next_second = ( - current_second + value - (current_second % value) - if current_second % value != 0 - else current_second + value + current_second + value - (current_second % value) if current_second % value != 0 else current_second + value ) # Handle minute rollover @@ -347,18 +311,12 @@ def _handle_second_reset( if next_hour >= 24: next_hour = next_hour % 24 next_day = base_midnight + timedelta(days=1) - return next_day.replace( - hour=next_hour, minute=next_minute, second=next_second, microsecond=0 - ) + return next_day.replace(hour=next_hour, minute=next_minute, second=next_second, microsecond=0) - return current_time.replace( - hour=next_hour, minute=next_minute, second=next_second, microsecond=0 - ) + return current_time.replace(hour=next_hour, minute=next_minute, second=next_second, microsecond=0) -def _handle_month_reset( - current_time: datetime, base_midnight: datetime, value: int -) -> datetime: +def _handle_month_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: """ Handle monthly reset times. For monthly resets, we always reset at the start of the next month. diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index a0f9bb00dd2..2441cbb3903 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -74,10 +74,7 @@ class ExceptionCheckers: # Exclude param validation errors (e.g. OpenAI "user" param max 64 chars) if "string_above_max_length" in _error_str_lowercase: return False - if ( - "invalid 'user'" in _error_str_lowercase - and "string too long" in _error_str_lowercase - ): + if "invalid 'user'" in _error_str_lowercase and "string too long" in _error_str_lowercase: return False known_exception_substrings = [ "exceed context limit", @@ -95,10 +92,7 @@ class ExceptionCheckers: return True # Cerebras pattern: "Current length is X while limit is Y" - if ( - "current length is" in _error_str_lowercase - and "while limit is" in _error_str_lowercase - ): + if "current length is" in _error_str_lowercase and "while limit is" in _error_str_lowercase: return True return False @@ -193,9 +187,7 @@ def _get_response_headers(original_exception: Exception) -> Optional[httpx.Heade if not _response_headers and error_response: _response_headers = getattr(error_response, "headers", None) if not _response_headers: - _response_headers = getattr( - original_exception, "litellm_response_headers", None - ) + _response_headers = getattr(original_exception, "litellm_response_headers", None) except Exception: return None @@ -272,8 +264,8 @@ def _map_openai_exception( else: message = str(original_exception) - if ( - message is not None and isinstance(message, str) + if message is not None and isinstance( + message, str ): # done to prevent user-confusion. Relevant issue - https://github.com/BerriAI/litellm/issues/1414 message = message.replace("OPENAI", custom_llm_provider.upper()) message = message.replace( @@ -283,9 +275,7 @@ def _map_openai_exception( if custom_llm_provider == "openai": exception_provider = "OpenAI" + "Exception" else: - exception_provider = ( - custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception" - ) + exception_provider = custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception" if ExceptionCheckers.is_error_str_rate_limit(error_str): raise RateLimitError( @@ -318,14 +308,9 @@ def _map_openai_exception( litellm_debug_info=extra_information, ) elif ( - ( - "invalid_request_error" in error_str - and "content_policy_violation" in error_str - ) + ("invalid_request_error" in error_str and "content_policy_violation" in error_str) or ("Invalid prompt" in error_str and "violating our usage policy" in error_str) - or ( - "request was rejected as a result of the safety system" in error_str.lower() - ) + or ("request was rejected as a result of the safety system" in error_str.lower()) ): raise ContentPolicyViolationError( message=f"ContentPolicyViolationError: {exception_provider} - {message}", @@ -334,9 +319,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, ) - elif ( - "invalid_encrypted_content" in error_str or "could not be verified" in error_str - ): + elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: helpful_message = ( f"{exception_provider} - {message}\n\n" " This error occurs when load balancing Responses API across deployments with different API keys.\n" @@ -356,10 +339,7 @@ def _map_openai_exception( litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), ) - elif ( - "invalid_request_error" in error_str - and "Incorrect API key provided" not in error_str - ): + elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str: raise BadRequestError( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, @@ -560,10 +540,7 @@ def _map_anthropic_exception( llm_provider="anthropic", model=model, ) - elif ( - original_exception.status_code == 400 - or original_exception.status_code == 413 - ): + elif original_exception.status_code == 400 or original_exception.status_code == 413: raise BadRequestError( message=f"AnthropicException - {error_str}", model=model, @@ -587,10 +564,7 @@ def _map_anthropic_exception( llm_provider="anthropic", model=model, ) - elif ( - original_exception.status_code == 500 - or original_exception.status_code == 529 - ): + elif original_exception.status_code == 500 or original_exception.status_code == 529: raise litellm.InternalServerError( message=f"AnthropicException - {error_str}. Handle with `litellm.InternalServerError`.", llm_provider="anthropic", @@ -666,10 +640,7 @@ def _map_replicate_exception( model=model, response=getattr(original_exception, "response", None), ) - elif ( - original_exception.status_code == 400 - or original_exception.status_code == 413 - ): + elif original_exception.status_code == 400 or original_exception.status_code == 413: raise BadRequestError( message=f"ReplicateException - {original_exception.message}", model=model, @@ -727,11 +698,7 @@ def _map_openai_like_exception( ) -> None: if "authorization denied for" in error_str: # Predibase returns the raw API Key in the response - this block ensures it's not returned in the exception - if ( - error_str is not None - and isinstance(error_str, str) - and "bearer" in error_str.lower() - ): + if error_str is not None and isinstance(error_str, str) and "bearer" in error_str.lower(): # only keep the first 10 chars after the occurnence of "bearer" _bearer_token_start_index = error_str.lower().find("bearer") error_str = error_str[: _bearer_token_start_index + 14] @@ -759,9 +726,7 @@ def _map_openai_like_exception( model=model, response=getattr(original_exception, "response", None), ) - elif ( - "The server received an invalid response from an upstream server." in error_str - ): + elif "The server received an invalid response from an upstream server." in error_str: raise litellm.InternalServerError( message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", llm_provider=custom_llm_provider, @@ -780,10 +745,7 @@ def _map_openai_like_exception( llm_provider=custom_llm_provider, model=model, ) - elif ( - original_exception.status_code == 401 - or original_exception.status_code == 403 - ): + elif original_exception.status_code == 401 or original_exception.status_code == 403: raise AuthenticationError( message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", llm_provider=custom_llm_provider, @@ -808,10 +770,7 @@ def _map_openai_like_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, ) - elif ( - original_exception.status_code == 422 - or original_exception.status_code == 424 - ): + elif original_exception.status_code == 422 or original_exception.status_code == 424: raise BadRequestError( message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", model=model, @@ -865,10 +824,7 @@ def _map_bedrock_exception( model=model, llm_provider="bedrock", ) - elif ( - "Conversation blocks and tool result blocks cannot be provided in the same turn." - in error_str - ): + elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str: raise BadRequestError( message=f"BedrockException - {error_str}\n. Enable 'litellm.modify_params=True' (for PROXY do: `litellm_settings::modify_params: True`) to insert a dummy assistant message and fix this error.", model=model, @@ -933,9 +889,7 @@ def _map_bedrock_exception( model=model, response=httpx.Response( status_code=500, - request=httpx.Request( - method="POST", url="https://api.openai.com/v1/" - ), + request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), ), ) elif original_exception.status_code == 401: @@ -1042,9 +996,7 @@ def _map_sagemaker_exception( model=model, response=httpx.Response( status_code=500, - request=httpx.Request( - method="POST", url="https://api.openai.com/v1/" - ), + request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), ), ) elif original_exception.status_code == 401: @@ -1075,10 +1027,7 @@ def _map_sagemaker_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, ) - elif ( - original_exception.status_code == 422 - or original_exception.status_code == 424 - ): + elif original_exception.status_code == 422 or original_exception.status_code == 424: raise BadRequestError( message=f"SagemakerException - {original_exception.message}", model=model, @@ -1122,10 +1071,7 @@ def _map_vertex_exception( exception_provider: str, extra_information: str, ) -> None: - if ( - "Vertex AI API has not been used in project" in error_str - or "Unable to find your project" in error_str - ): + if "Vertex AI API has not been used in project" in error_str or "Unable to find your project" in error_str: raise BadRequestError( message=f"litellm.BadRequestError: {custom_llm_provider}Exception - {error_str}", model=model, @@ -1160,9 +1106,7 @@ def _map_vertex_exception( response=httpx.Response( status_code=500, content=str(original_exception), - request=httpx.Request( - method="completion", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore ), litellm_debug_info=extra_information, ) @@ -1189,8 +1133,7 @@ def _map_vertex_exception( ) elif ( "The response was blocked." in error_str - or "Output blocked by content filtering policy" - in error_str # anthropic on vertex ai + or "Output blocked by content filtering policy" in error_str # anthropic on vertex ai ): raise ContentPolicyViolationError( message=f"{custom_llm_provider.capitalize()}Exception ContentPolicyViolationError - {error_str}", @@ -1210,8 +1153,7 @@ def _map_vertex_exception( or "Quota exceeded for" in error_str or "Resource exhausted" in error_str or "IndexError: list index out of range" in error_str - or "429 Unable to submit request because the service is temporarily out of capacity." - in error_str + or "429 Unable to submit request because the service is temporarily out of capacity." in error_str ): raise RateLimitError( message=f"litellm.RateLimitError: {custom_llm_provider}Exception - {error_str}", @@ -1248,10 +1190,7 @@ def _map_vertex_exception( ), ), ) - elif ( - "500 Internal Server Error" in error_str - or "The model is overloaded." in error_str - ): + elif "500 Internal Server Error" in error_str or "The model is overloaded." in error_str: raise litellm.InternalServerError( message=f"litellm.InternalServerError: {custom_llm_provider}Exception - {error_str}", model=model, @@ -1328,9 +1267,7 @@ def _map_vertex_exception( response=httpx.Response( status_code=500, content=str(original_exception), - request=httpx.Request( - method="completion", url="https://github.com/BerriAI/litellm" - ), # type: ignore + request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) if original_exception.status_code == 502: @@ -1412,10 +1349,7 @@ def _map_cohere_exception( response=getattr(original_exception, "response", None), ) elif hasattr(original_exception, "status_code"): - if ( - original_exception.status_code == 400 - or original_exception.status_code == 498 - ): + if original_exception.status_code == 400 or original_exception.status_code == 498: raise BadRequestError( message=f"CohereException - {original_exception.message}", llm_provider="cohere", @@ -1642,9 +1576,7 @@ def _map_nlp_cloud_exception( llm_provider="nlp_cloud", request=getattr(original_exception, "request", None), ) - if hasattr( - original_exception, "status_code" - ): # https://docs.nlpcloud.com/?shell#errors + if hasattr(original_exception, "status_code"): # https://docs.nlpcloud.com/?shell#errors if ( original_exception.status_code == 400 or original_exception.status_code == 406 @@ -1657,39 +1589,27 @@ def _map_nlp_cloud_exception( model=model, response=getattr(original_exception, "response", None), ) - elif ( - original_exception.status_code == 401 - or original_exception.status_code == 403 - ): + elif original_exception.status_code == 401 or original_exception.status_code == 403: raise AuthenticationError( message=f"NLPCloudException - {original_exception.message}", llm_provider="nlp_cloud", model=model, response=getattr(original_exception, "response", None), ) - elif ( - original_exception.status_code == 522 - or original_exception.status_code == 524 - ): + elif original_exception.status_code == 522 or original_exception.status_code == 524: raise Timeout( message=f"NLPCloudException - {original_exception.message}", model=model, llm_provider="nlp_cloud", ) - elif ( - original_exception.status_code == 429 - or original_exception.status_code == 402 - ): + elif original_exception.status_code == 429 or original_exception.status_code == 402: raise RateLimitError( message=f"NLPCloudException - {original_exception.message}", llm_provider="nlp_cloud", model=model, response=getattr(original_exception, "response", None), ) - elif ( - original_exception.status_code == 500 - or original_exception.status_code == 503 - ): + elif original_exception.status_code == 500 or original_exception.status_code == 503: raise APIError( status_code=original_exception.status_code, message=f"NLPCloudException - {original_exception.message}", @@ -1697,10 +1617,7 @@ def _map_nlp_cloud_exception( model=model, request=getattr(original_exception, "request", None), ) - elif ( - original_exception.status_code == 504 - or original_exception.status_code == 520 - ): + elif original_exception.status_code == 504 or original_exception.status_code == 520: raise ServiceUnavailableError( message=f"NLPCloudException - {original_exception.message}", model=model, @@ -1731,10 +1648,7 @@ def _map_together_ai_exception( error_response = json.loads(error_str) except Exception: error_response = {"error": error_str} - if ( - "error" in error_response - and "`inputs` tokens + `max_new_tokens` must be <=" in error_response["error"] - ): + if "error" in error_response and "`inputs` tokens + `max_new_tokens` must be <=" in error_response["error"]: raise ContextWindowExceededError( message=f"TogetherAIException - {error_response['error']}", model=model, @@ -1761,19 +1675,14 @@ def _map_together_ai_exception( model=model, llm_provider="together_ai", ) - elif ( - "error" in error_response - and "API key doesn't match expected format." in error_response["error"] - ): + elif "error" in error_response and "API key doesn't match expected format." in error_response["error"]: raise BadRequestError( message=f"TogetherAIException - {error_response['error']}", model=model, llm_provider="together_ai", response=getattr(original_exception, "response", None), ) - elif ( - "error_type" in error_response and error_response["error_type"] == "validation" - ): + elif "error_type" in error_response and error_response["error_type"] == "validation": raise BadRequestError( message=f"TogetherAIException - {error_response['error']}", model=model, @@ -1971,10 +1880,7 @@ def _map_azure_exception( _inner = body_dict["error"].get("inner_error") or body_dict[ # type: ignore[index] "error" ].get("innererror") # type: ignore[index] - if ( - isinstance(_inner, dict) - and _inner.get("code") == "ResponsibleAIPolicyViolation" - ): + if isinstance(_inner, dict) and _inner.get("code") == "ResponsibleAIPolicyViolation": azure_error_code = "content_policy_violation" else: azure_error_code = body_dict.get("code") @@ -2005,9 +1911,8 @@ def _map_azure_exception( litellm_debug_info=extra_information, response=getattr(original_exception, "response", None), ) - elif ( - azure_error_code == "content_policy_violation" - or ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + elif azure_error_code == "content_policy_violation" or ExceptionCheckers.is_azure_content_policy_violation_error( + error_str ): from litellm.llms.azure.exception_mapping import ( AzureOpenAIExceptionMapping, @@ -2019,10 +1924,7 @@ def _map_azure_exception( extra_information=extra_information, original_exception=original_exception, ) - elif ( - azure_error_code == "invalid_encrypted_content" - or "could not be verified" in error_str - ): + elif azure_error_code == "invalid_encrypted_content" or "could not be verified" in error_str: helpful_message = ( f"AzureException - {message}\n\n" "This error occurs when load balancing Responses API across deployments with different API keys.\n" @@ -2051,10 +1953,7 @@ def _map_azure_exception( response=getattr(original_exception, "response", None), body=getattr(original_exception, "body", None), ) - elif ( - "The api_key client option must be set either by passing api_key to the client or by setting" - in error_str - ): + elif "The api_key client option must be set either by passing api_key to the client or by setting" in error_str: raise AuthenticationError( message=f"{exception_provider} AuthenticationError - {message}", llm_provider=custom_llm_provider, @@ -2256,16 +2155,11 @@ def exception_type( # type: ignore extra_kwargs={}, ): """Maps an LLM Provider Exception to OpenAI Exception Format""" - if any( - isinstance(original_exception, exc_type) - for exc_type in litellm.LITELLM_EXCEPTION_TYPES - ): + if any(isinstance(original_exception, exc_type) for exc_type in litellm.LITELLM_EXCEPTION_TYPES): return original_exception exception_mapping_worked = False exception_provider = custom_llm_provider - mappable_exception: _ProviderHTTPException = cast( - "_ProviderHTTPException", original_exception - ) + mappable_exception: _ProviderHTTPException = cast("_ProviderHTTPException", original_exception) if litellm.suppress_debug_info is False: print() # noqa: T201 print( # noqa: T201 @@ -2276,15 +2170,9 @@ def exception_type( # type: ignore ) print() # noqa: T201 - litellm_response_headers = _get_response_headers( - original_exception=original_exception - ) + litellm_response_headers = _get_response_headers(original_exception=original_exception) try: - error_str = ( - redact_string(str(original_exception)) - if _ENABLE_SECRET_REDACTION - else str(original_exception) - ) + error_str = redact_string(str(original_exception)) if _ENABLE_SECRET_REDACTION else str(original_exception) if model: if hasattr(original_exception, "message"): error_str = ( @@ -2303,9 +2191,7 @@ def exception_type( # type: ignore ################################################################################ extra_information = "" try: - _api_base = litellm.get_api_base( - model=model, optional_params=extra_kwargs - ) + _api_base = litellm.get_api_base(model=model, optional_params=extra_kwargs) messages = litellm.get_first_chars_messages(kwargs=completion_kwargs) _vertex_project = extra_kwargs.get("vertex_project") _vertex_location = extra_kwargs.get("vertex_location") @@ -2314,23 +2200,12 @@ def exception_type( # type: ignore _deployment = _metadata.get("deployment") extra_information = f"\nModel: {model}" - if ( - isinstance(custom_llm_provider, str) - and len(custom_llm_provider) > 0 - ): - exception_provider = ( - custom_llm_provider[0].upper() - + custom_llm_provider[1:] - + "Exception" - ) + if isinstance(custom_llm_provider, str) and len(custom_llm_provider) > 0: + exception_provider = custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception" if _api_base: extra_information += f"\nAPI Base: `{_api_base}`" - if ( - messages - and len(messages) > 0 - and litellm.redact_messages_in_exceptions is False - ): + if messages and len(messages) > 0 and litellm.redact_messages_in_exceptions is False: extra_information += f"\nMessages: `{messages}`" if _model_group is not None: @@ -2343,9 +2218,7 @@ def exception_type( # type: ignore extra_information += f"\nvertex_location: `{_vertex_location}`\n" # on litellm proxy add key name + team to exceptions - extra_information = _add_key_name_and_team_to_alert( - request_info=extra_information, metadata=_metadata - ) + extra_information = _add_key_name_and_team_to_alert(request_info=extra_information, metadata=_metadata) except Exception: # DO NOT LET this Block raising the original exception pass @@ -2398,10 +2271,7 @@ def exception_type( # type: ignore exception_provider=exception_provider, extra_information=extra_information, ) - elif ( - custom_llm_provider == "anthropic" - or custom_llm_provider == "anthropic_text" - ): # one of the anthropics + elif custom_llm_provider == "anthropic" or custom_llm_provider == "anthropic_text": # one of the anthropics _map_anthropic_exception( model=model, original_exception=mappable_exception, @@ -2441,10 +2311,7 @@ def exception_type( # type: ignore exception_provider=exception_provider, extra_information=extra_information, ) - elif ( - custom_llm_provider == "sagemaker" - or custom_llm_provider == "sagemaker_chat" - ): + elif custom_llm_provider == "sagemaker" or custom_llm_provider == "sagemaker_chat": _map_sagemaker_exception( model=model, original_exception=mappable_exception, @@ -2478,9 +2345,7 @@ def exception_type( # type: ignore exception_provider=exception_provider, extra_information=extra_information, ) - elif ( - custom_llm_provider == "cohere" or custom_llm_provider == "cohere_chat" - ): # Cohere + elif custom_llm_provider == "cohere" or custom_llm_provider == "cohere_chat": # Cohere _map_cohere_exception( model=model, original_exception=mappable_exception, @@ -2540,9 +2405,7 @@ def exception_type( # type: ignore exception_provider=exception_provider, extra_information=extra_information, ) - elif ( - custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat" - ): + elif custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat": _map_ollama_exception( model=model, original_exception=mappable_exception, @@ -2582,9 +2445,8 @@ def exception_type( # type: ignore exception_provider=exception_provider, extra_information=extra_information, ) - if ( - "BadRequestError.__init__() missing 1 required positional argument: 'param'" - in str(original_exception) + if "BadRequestError.__init__() missing 1 required positional argument: 'param'" in str( + original_exception ): # deal with edge-case invalid request error bug in openai-python sdk exception_mapping_worked = True raise BadRequestError( @@ -2613,9 +2475,7 @@ def exception_type( # type: ignore ), llm_provider=custom_llm_provider, model=model, - request=httpx.Request( - method="POST", url="https://api.openai.com/v1/" - ), # stub the request + request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), # stub the request ) except Exception as e: # LOGGING @@ -2663,9 +2523,7 @@ def exception_logging( model_call_details["exception"] = exception model_call_details["additional_args"] = additional_args # User Logging -> if you pass in a custom logging function or want to use sentry breadcrumbs - verbose_logger.debug( - f"Logging Details: logger_fn - {logger_fn} | callable(logger_fn) - {callable(logger_fn)}" - ) + verbose_logger.debug(f"Logging Details: logger_fn - {logger_fn} | callable(logger_fn) - {callable(logger_fn)}") if logger_fn and callable(logger_fn): try: logger_fn( @@ -2694,10 +2552,7 @@ def _add_key_name_and_team_to_alert(request_info: str, metadata: dict) -> str: _api_key_name = metadata.get("user_api_key_alias", None) _user_api_key_team_alias = metadata.get("user_api_key_team_alias", None) if _api_key_name is not None: - request_info = ( - f"\n\nKey Name: `{_api_key_name}`\nTeam: `{_user_api_key_team_alias}`" - + request_info - ) + request_info = f"\n\nKey Name: `{_api_key_name}`\nTeam: `{_user_api_key_team_alias}`" + request_info return request_info except Exception: diff --git a/litellm/litellm_core_utils/fallback_utils.py b/litellm/litellm_core_utils/fallback_utils.py index 1606b53e1f9..7aee69ef862 100644 --- a/litellm/litellm_core_utils/fallback_utils.py +++ b/litellm/litellm_core_utils/fallback_utils.py @@ -72,9 +72,7 @@ async def async_completion_with_fallbacks(**kwargs): ) except Exception as e: - verbose_logger.exception( - f"Fallback attempt failed for model {model}: {str(e)}" - ) + verbose_logger.exception(f"Fallback attempt failed for model {model}: {str(e)}") most_recent_exception_str = str(e) continue diff --git a/litellm/litellm_core_utils/get_blog_posts.py b/litellm/litellm_core_utils/get_blog_posts.py index 2f9a14f1279..6aea79cb4b3 100644 --- a/litellm/litellm_core_utils/get_blog_posts.py +++ b/litellm/litellm_core_utils/get_blog_posts.py @@ -51,9 +51,7 @@ class GetBlogPosts: @staticmethod def load_local_blog_posts() -> List[Dict[str, str]]: """Load the bundled local backup blog posts.""" - content = json.loads( - files("litellm").joinpath("blog_posts.json").read_text(encoding="utf-8") - ) + content = json.loads(files("litellm").joinpath("blog_posts.json").read_text(encoding="utf-8")) return content.get("posts", []) @staticmethod @@ -117,8 +115,7 @@ class GetBlogPosts: """Return True if posts is a non-empty list.""" if not isinstance(posts, list) or len(posts) == 0: verbose_logger.warning( - "LiteLLM: Parsed RSS feed has no valid posts. " - "Falling back to local backup.", + "LiteLLM: Parsed RSS feed has no valid posts. Falling back to local backup.", ) return False return True @@ -144,8 +141,7 @@ class GetBlogPosts: posts = cls.parse_rss_to_posts(xml_text) except Exception as e: verbose_logger.warning( - "LiteLLM: Failed to fetch blog posts from %s: %s. " - "Falling back to local backup.", + "LiteLLM: Failed to fetch blog posts from %s: %s. Falling back to local backup.", url, str(e), ) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index c88f8b77dc2..fbed9594a0b 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -111,9 +111,7 @@ def get_litellm_params( if litellm_trace_id is None: litellm_trace_id = _meta.get("trace_id") or _meta.get("session_id") - data_residency: Optional[str] = infer_openai_data_residency( - custom_llm_provider, api_base - ) + data_residency: Optional[str] = infer_openai_data_residency(custom_llm_provider, api_base) # Build base dict with explicit parameters (always included) litellm_params = { @@ -145,11 +143,7 @@ def get_litellm_params( "azure_ad_token_provider": azure_ad_token_provider, "user_continue_message": user_continue_message, "base_model": base_model - or ( - _get_base_model_from_litellm_call_metadata(metadata=metadata) - if metadata - else None - ), + or (_get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None), "litellm_trace_id": litellm_trace_id, "litellm_session_id": litellm_session_id, "hf_model_name": hf_model_name, diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 061316a69d6..76ca268a883 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -50,10 +50,7 @@ def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool: def _is_non_openai_azure_model(model: str) -> bool: try: model_name = model.split("/", 1)[1] - if ( - model_name in litellm.cohere_chat_models - or f"mistral/{model_name}" in litellm.mistral_chat_models - ): + if model_name in litellm.cohere_chat_models or f"mistral/{model_name}" in litellm.mistral_chat_models: return True except Exception: return False @@ -111,11 +108,7 @@ def handle_cohere_chat_model_custom_llm_provider( if model and "/" in model: _custom_llm_provider, _model = model.split("/", 1) - if ( - _custom_llm_provider - and _custom_llm_provider == "cohere" - and _model in litellm.cohere_chat_models - ): + if _custom_llm_provider and _custom_llm_provider == "cohere" and _model in litellm.cohere_chat_models: return _model, "cohere_chat" return model, custom_llm_provider @@ -136,10 +129,7 @@ def handle_anthropic_text_model_custom_llm_provider( """ if custom_llm_provider: - if ( - custom_llm_provider == "anthropic" - and litellm.AnthropicTextConfig._is_anthropic_text_model(model) - ): + if custom_llm_provider == "anthropic" and litellm.AnthropicTextConfig._is_anthropic_text_model(model): return model, "anthropic_text" if model and "/" in model: @@ -173,9 +163,7 @@ def get_llm_provider( try: # Early validation - model is required if model is None: - raise ValueError( - "model parameter is required but was None. Please provide a valid model name." - ) + raise ValueError("model parameter is required but was None. Please provide a valid model name.") if litellm.LiteLLMProxyChatConfig._should_use_litellm_proxy_by_default( litellm_params=cast(Optional[LiteLLM_Params], litellm_params) @@ -201,13 +189,9 @@ def get_llm_provider( return model, custom_llm_provider, dynamic_api_key, api_base ### Handle cases when custom_llm_provider is set to cohere/command-r-plus but it should use cohere_chat route - model, custom_llm_provider = handle_cohere_chat_model_custom_llm_provider( - model, custom_llm_provider - ) + model, custom_llm_provider = handle_cohere_chat_model_custom_llm_provider(model, custom_llm_provider) - model, custom_llm_provider = handle_anthropic_text_model_custom_llm_provider( - model, custom_llm_provider - ) + model, custom_llm_provider = handle_anthropic_text_model_custom_llm_provider(model, custom_llm_provider) if custom_llm_provider and ( model.split("/")[0] != custom_llm_provider @@ -255,14 +239,10 @@ def get_llm_provider( custom_llm_provider = model.split("/", 1)[0] model = model.split("/", 1)[1] if api_base is not None and not isinstance(api_base, str): - raise Exception( - "api base needs to be a string. api_base={}".format(api_base) - ) + raise Exception("api base needs to be a string. api_base={}".format(api_base)) if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): raise Exception( - "dynamic_api_key needs to be a string. Got type={}".format( - type(dynamic_api_key).__name__ - ) + "dynamic_api_key needs to be a string. Got type={}".format(type(dynamic_api_key).__name__) ) return model, custom_llm_provider, dynamic_api_key, api_base # check if api base is a known openai compatible endpoint @@ -316,9 +296,7 @@ def get_llm_provider( dynamic_api_key = get_secret_str("OLLAMA_API_KEY") elif endpoint == "https://api.friendli.ai/serverless/v1": custom_llm_provider = "friendliai" - dynamic_api_key = get_secret_str( - "FRIENDLIAI_API_KEY" - ) or get_secret("FRIENDLI_TOKEN") + dynamic_api_key = get_secret_str("FRIENDLIAI_API_KEY") or get_secret("FRIENDLI_TOKEN") elif endpoint == "api.galadriel.com/v1": custom_llm_provider = "galadriel" dynamic_api_key = get_secret_str("GALADRIEL_API_KEY") @@ -340,16 +318,10 @@ def get_llm_provider( elif endpoint == "api.moonshot.ai/v1": custom_llm_provider = "moonshot" dynamic_api_key = get_secret_str("MOONSHOT_API_KEY") - elif ( - endpoint == "api.minimax.io/anthropic" - or endpoint == "api.minimaxi.com/anthropic" - ): + elif endpoint == "api.minimax.io/anthropic" or endpoint == "api.minimaxi.com/anthropic": custom_llm_provider = "minimax" dynamic_api_key = get_secret_str("MINIMAX_API_KEY") - elif ( - endpoint == "api.minimax.io/v1" - or endpoint == "api.minimaxi.com/v1" - ): + elif endpoint == "api.minimax.io/v1" or endpoint == "api.minimaxi.com/v1": custom_llm_provider = "minimax" dynamic_api_key = get_secret_str("MINIMAX_API_KEY") elif endpoint == "platform.publicai.co/v1": @@ -393,18 +365,10 @@ def get_llm_provider( dynamic_api_key = get_secret_str("PINSTRIPES_API_KEY") if api_base is not None and not isinstance(api_base, str): + raise Exception("api base needs to be a string. api_base={}".format(api_base)) + if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): raise Exception( - "api base needs to be a string. api_base={}".format( - api_base - ) - ) - if dynamic_api_key is not None and not isinstance( - dynamic_api_key, str - ): - raise Exception( - "dynamic_api_key needs to be a string. dynamic_api_key={}".format( - dynamic_api_key - ) + "dynamic_api_key needs to be a string. dynamic_api_key={}".format(dynamic_api_key) ) return model, custom_llm_provider, dynamic_api_key, api_base # type: ignore @@ -437,13 +401,10 @@ def get_llm_provider( elif model in litellm.cohere_chat_models: custom_llm_provider = "cohere_chat" ## replicate - elif model in litellm.replicate_models or ( - ":" in model and len(model) > REPLICATE_MODEL_NAME_WITH_ID_LENGTH - ): + elif model in litellm.replicate_models or (":" in model and len(model) > REPLICATE_MODEL_NAME_WITH_ID_LENGTH): model_parts = model.split(":") if ( - len(model_parts) > 1 - and len(model_parts[1]) == REPLICATE_MODEL_NAME_WITH_ID_LENGTH + len(model_parts) > 1 and len(model_parts[1]) == REPLICATE_MODEL_NAME_WITH_ID_LENGTH ): ## checks if model name has a 64 digit code - e.g. "meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3" custom_llm_provider = "replicate" elif model in litellm.replicate_models: @@ -470,11 +431,7 @@ def get_llm_provider( ## ai21 elif model in litellm.ai21_chat_models or model in litellm.ai21_models: custom_llm_provider = "ai21_chat" - api_base = ( - api_base - or get_secret("AI21_API_BASE") - or "https://api.ai21.com/studio/v1" - ) # type: ignore + api_base = api_base or get_secret("AI21_API_BASE") or "https://api.ai21.com/studio/v1" # type: ignore dynamic_api_key = api_key or get_secret("AI21_API_KEY") ## aleph_alpha elif model in litellm.aleph_alpha_models: @@ -546,23 +503,15 @@ def get_llm_provider( llm_provider="", ) if api_base is not None and not isinstance(api_base, str): - raise Exception( - "api base needs to be a string. api_base={}".format(api_base) - ) + raise Exception("api base needs to be a string. api_base={}".format(api_base)) if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): - raise Exception( - "dynamic_api_key needs to be a string. dynamic_api_key={}".format( - dynamic_api_key - ) - ) + raise Exception("dynamic_api_key needs to be a string. dynamic_api_key={}".format(dynamic_api_key)) return model, custom_llm_provider, dynamic_api_key, api_base except Exception as e: if isinstance(e, litellm.exceptions.BadRequestError): raise e else: - error_str = ( - f"GetLLMProvider Exception - {str(e)}\n\noriginal model: {model}" - ) + error_str = f"GetLLMProvider Exception - {str(e)}\n\noriginal model: {model}" raise litellm.exceptions.BadRequestError( # type: ignore message=f"GetLLMProvider Exception - {str(e)}\n\noriginal model: {model}", model=model, @@ -599,9 +548,7 @@ def _get_openai_compatible_provider_info( if provider_config is None: raise ValueError(f"Provider {custom_llm_provider} not found") config_class = create_config_class(provider_config) - api_base, dynamic_api_key = config_class()._get_openai_compatible_provider_info( - api_base, api_key - ) + api_base, dynamic_api_key = config_class()._get_openai_compatible_provider_info(api_base, api_key) return model, custom_llm_provider, dynamic_api_key, api_base if custom_llm_provider == "perplexity": @@ -609,40 +556,26 @@ def _get_openai_compatible_provider_info( ( api_base, dynamic_api_key, - ) = litellm.PerplexityChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.PerplexityChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "aiohttp_openai": return model, "aiohttp_openai", api_key, api_base elif custom_llm_provider == "anyscale": # anyscale is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 - api_base = ( - api_base - or get_secret_str("ANYSCALE_API_BASE") - or "https://api.endpoints.anyscale.com/v1" - ) # type: ignore + api_base = api_base or get_secret_str("ANYSCALE_API_BASE") or "https://api.endpoints.anyscale.com/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("ANYSCALE_API_KEY") elif custom_llm_provider == "deepinfra": ( api_base, dynamic_api_key, - ) = litellm.DeepInfraConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DeepInfraConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "empower": - api_base = ( - api_base - or get_secret("EMPOWER_API_BASE") - or "https://app.empower.dev/api/v1" - ) # type: ignore + api_base = api_base or get_secret("EMPOWER_API_BASE") or "https://app.empower.dev/api/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("EMPOWER_API_KEY") elif custom_llm_provider == "groq": ( api_base, dynamic_api_key, - ) = litellm.GroqChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.GroqChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "bedrock_mantle": ( api_base, @@ -652,11 +585,7 @@ def _get_openai_compatible_provider_info( ) elif custom_llm_provider == "nvidia_nim": # nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 - api_base = ( - api_base - or get_secret("NVIDIA_NIM_API_BASE") - or "https://integrate.api.nvidia.com/v1" - ) # type: ignore + api_base = api_base or get_secret("NVIDIA_NIM_API_BASE") or "https://integrate.api.nvidia.com/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("NVIDIA_NIM_API_KEY") elif custom_llm_provider == "nvidia_riva": # NVIDIA Riva is gRPC-based; api_base must be a host:port like @@ -665,119 +594,71 @@ def _get_openai_compatible_provider_info( api_base = api_base or get_secret_str("NVIDIA_RIVA_API_BASE") # type: ignore # Fall back to NVIDIA_NIM_API_KEY because users running both NVCF # services typically reuse the same nvapi-* key. - dynamic_api_key = ( - api_key - or get_secret_str("NVIDIA_RIVA_API_KEY") - or get_secret_str("NVIDIA_NIM_API_KEY") - ) + dynamic_api_key = api_key or get_secret_str("NVIDIA_RIVA_API_KEY") or get_secret_str("NVIDIA_NIM_API_KEY") elif custom_llm_provider == "soniox": - api_base = ( - api_base or get_secret_str("SONIOX_API_BASE") or "https://api.soniox.com" - ) + api_base = api_base or get_secret_str("SONIOX_API_BASE") or "https://api.soniox.com" dynamic_api_key = api_key or get_secret_str("SONIOX_API_KEY") elif custom_llm_provider == "cerebras": - api_base = ( - api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("CEREBRAS_API_KEY") elif custom_llm_provider == "baseten": # Use BasetenConfig to determine the appropriate API base URL if api_base is None: api_base = litellm.BasetenConfig.get_api_base_for_model(model) else: - api_base = ( - api_base - or get_secret_str("BASETEN_API_BASE") - or "https://inference.baseten.co/v1" - ) + api_base = api_base or get_secret_str("BASETEN_API_BASE") or "https://inference.baseten.co/v1" dynamic_api_key = api_key or get_secret_str("BASETEN_API_KEY") elif custom_llm_provider == "sambanova": - api_base = ( - api_base - or get_secret("SAMBANOVA_API_BASE") - or "https://api.sambanova.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("SAMBANOVA_API_BASE") or "https://api.sambanova.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("SAMBANOVA_API_KEY") elif custom_llm_provider == "meta_llama": - api_base = ( - api_base - or get_secret("LLAMA_API_BASE") - or "https://api.llama.com/compat/v1" - ) # type: ignore + api_base = api_base or get_secret("LLAMA_API_BASE") or "https://api.llama.com/compat/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("LLAMA_API_KEY") elif custom_llm_provider == "nebius": - api_base = ( - api_base - or get_secret("NEBIUS_API_BASE") - or "https://api.studio.nebius.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("NEBIUS_API_BASE") or "https://api.studio.nebius.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") elif custom_llm_provider == "ollama": api_base = api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" # type: ignore dynamic_api_key = api_key or get_secret_str("OLLAMA_API_KEY") - elif (custom_llm_provider == "ai21_chat") or ( - custom_llm_provider == "ai21" and model in litellm.ai21_chat_models - ): - api_base = ( - api_base or get_secret("AI21_API_BASE") or "https://api.ai21.com/studio/v1" - ) # type: ignore + elif (custom_llm_provider == "ai21_chat") or (custom_llm_provider == "ai21" and model in litellm.ai21_chat_models): + api_base = api_base or get_secret("AI21_API_BASE") or "https://api.ai21.com/studio/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("AI21_API_KEY") custom_llm_provider = "ai21_chat" elif custom_llm_provider == "volcengine": # volcengine is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 - api_base = ( - api_base - or get_secret("VOLCENGINE_API_BASE") - or "https://ark.cn-beijing.volces.com/api/v3" - ) # type: ignore + api_base = api_base or get_secret("VOLCENGINE_API_BASE") or "https://ark.cn-beijing.volces.com/api/v3" # type: ignore dynamic_api_key = api_key or get_secret_str("VOLCENGINE_API_KEY") elif custom_llm_provider == "codestral": # codestral is openai compatible, we just need to set this to custom_openai and have the api_base be https://codestral.mistral.ai/v1 - api_base = ( - api_base - or get_secret("CODESTRAL_API_BASE") - or "https://codestral.mistral.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("CODESTRAL_API_BASE") or "https://codestral.mistral.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("CODESTRAL_API_KEY") elif custom_llm_provider == "hosted_vllm": # vllm is openai compatible, we just need to set this to custom_openai ( api_base, dynamic_api_key, - ) = litellm.HostedVLLMChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.HostedVLLMChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "llamafile": # llamafile is OpenAI compatible. ( api_base, dynamic_api_key, - ) = litellm.LlamafileChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LlamafileChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "datarobot": # DataRobot is OpenAI compatible. ( api_base, dynamic_api_key, - ) = litellm.DataRobotConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DataRobotConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "lm_studio": # lm_studio is openai compatible, we just need to set this to custom_openai ( api_base, dynamic_api_key, - ) = litellm.LMStudioChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LMStudioChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "deepseek": # deepseek is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.deepseek.com/v1 - api_base = ( - api_base - or get_secret("DEEPSEEK_API_BASE") - or "https://api.deepseek.com/beta" - ) # type: ignore + api_base = api_base or get_secret("DEEPSEEK_API_BASE") or "https://api.deepseek.com/beta" # type: ignore dynamic_api_key = api_key or get_secret_str("DEEPSEEK_API_KEY") elif custom_llm_provider == "fireworks_ai": @@ -785,9 +666,7 @@ def _get_openai_compatible_provider_info( ( api_base, dynamic_api_key, - ) = litellm.FireworksAIConfig()._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + ) = litellm.FireworksAIConfig()._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) elif custom_llm_provider == "azure_ai": ( api_base, @@ -807,45 +686,31 @@ def _get_openai_compatible_provider_info( ( api_base, dynamic_api_key, - ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) elif custom_llm_provider == "mistral": ( api_base, dynamic_api_key, - ) = litellm.MistralConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.MistralConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "jina_ai": ( custom_llm_provider, api_base, dynamic_api_key, - ) = litellm.JinaAIEmbeddingConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.JinaAIEmbeddingConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "xai": ( api_base, dynamic_api_key, - ) = litellm.XAIChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.XAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "zai": ( api_base, dynamic_api_key, - ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "together_ai": - api_base = ( - api_base - or get_secret_str("TOGETHER_AI_API_BASE") - or "https://api.together.xyz/v1" - ) # type: ignore + api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1" # type: ignore dynamic_api_key = api_key or ( get_secret_str("TOGETHER_API_KEY") or get_secret_str("TOGETHER_AI_API_KEY") @@ -853,22 +718,10 @@ def _get_openai_compatible_provider_info( or get_secret_str("TOGETHER_AI_TOKEN") ) elif custom_llm_provider == "friendliai": - api_base = ( - api_base - or get_secret("FRIENDLI_API_BASE") - or "https://api.friendli.ai/serverless/v1" - ) # type: ignore - dynamic_api_key = ( - api_key - or get_secret_str("FRIENDLIAI_API_KEY") - or get_secret_str("FRIENDLI_TOKEN") - ) + api_base = api_base or get_secret("FRIENDLI_API_BASE") or "https://api.friendli.ai/serverless/v1" # type: ignore + dynamic_api_key = api_key or get_secret_str("FRIENDLIAI_API_KEY") or get_secret_str("FRIENDLI_TOKEN") elif custom_llm_provider == "galadriel": - api_base = ( - api_base - or get_secret("GALADRIEL_API_BASE") - or "https://api.galadriel.com/v1" - ) # type: ignore + api_base = api_base or get_secret("GALADRIEL_API_BASE") or "https://api.galadriel.com/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("GALADRIEL_API_KEY") elif custom_llm_provider == "github_copilot": ( @@ -883,181 +736,125 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, custom_llm_provider, - ) = litellm.ChatGPTConfig()._get_openai_compatible_provider_info( - model, api_base, api_key, custom_llm_provider - ) + ) = litellm.ChatGPTConfig()._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider) elif custom_llm_provider == "novita": - api_base = ( - api_base - or get_secret("NOVITA_API_BASE") - or "https://api.novita.ai/v3/openai" - ) # type: ignore + api_base = api_base or get_secret("NOVITA_API_BASE") or "https://api.novita.ai/v3/openai" # type: ignore dynamic_api_key = api_key or get_secret_str("NOVITA_API_KEY") elif custom_llm_provider == "snowflake": ( api_base, dynamic_api_key, - ) = litellm.SnowflakeConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.SnowflakeConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "gradient_ai": ( api_base, dynamic_api_key, - ) = litellm.GradientAIConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.GradientAIConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "featherless_ai": ( api_base, dynamic_api_key, - ) = litellm.FeatherlessAIConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.FeatherlessAIConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "nscale": ( api_base, dynamic_api_key, - ) = litellm.NscaleConfig()._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + ) = litellm.NscaleConfig()._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) elif custom_llm_provider == "heroku": ( api_base, dynamic_api_key, - ) = litellm.HerokuChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.HerokuChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "dashscope": ( api_base, dynamic_api_key, - ) = litellm.DashScopeChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DashScopeChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "modelscope": ( api_base, dynamic_api_key, - ) = litellm.ModelScopeChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.ModelScopeChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "moonshot": ( api_base, dynamic_api_key, - ) = litellm.MoonshotChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.MoonshotChatConfig()._get_openai_compatible_provider_info(api_base, api_key) # publicai is now handled by JSON config (see litellm/llms/openai_like/providers.json) elif custom_llm_provider == "docker_model_runner": ( api_base, dynamic_api_key, - ) = litellm.DockerModelRunnerChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DockerModelRunnerChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "v0": ( api_base, dynamic_api_key, - ) = litellm.V0ChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.V0ChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "morph": ( api_base, dynamic_api_key, - ) = litellm.MorphChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.MorphChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "lambda_ai": ( api_base, dynamic_api_key, - ) = litellm.LambdaAIChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LambdaAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "inception": ( api_base, dynamic_api_key, - ) = litellm.InceptionChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.InceptionChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "hyperbolic": ( api_base, dynamic_api_key, - ) = litellm.HyperbolicChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.HyperbolicChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "vercel_ai_gateway": ( api_base, dynamic_api_key, - ) = litellm.VercelAIGatewayConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.VercelAIGatewayConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "aiml": ( api_base, dynamic_api_key, - ) = litellm.AIMLChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.AIMLChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "wandb": - api_base = ( - api_base - or get_secret("WANDB_API_BASE") - or "https://api.inference.wandb.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("WANDB_API_BASE") or "https://api.inference.wandb.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("WANDB_API_KEY") elif custom_llm_provider == "lemonade": ( api_base, dynamic_api_key, - ) = litellm.LemonadeChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LemonadeChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "clarifai": ( api_base, dynamic_api_key, - ) = litellm.ClarifaiConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.ClarifaiConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "ragflow": full_model = f"ragflow/{model}" ( api_base, dynamic_api_key, _, - ) = litellm.RAGFlowConfig()._get_openai_compatible_provider_info( - full_model, api_base, api_key, "ragflow" - ) + ) = litellm.RAGFlowConfig()._get_openai_compatible_provider_info(full_model, api_base, api_key, "ragflow") model = full_model elif custom_llm_provider == "langgraph": # LangGraph is a custom provider, just need to set api_base - api_base = ( - api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" - ) + api_base = api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" dynamic_api_key = api_key or get_secret_str("LANGGRAPH_API_KEY") elif custom_llm_provider == "manus": # Manus is OpenAI compatible for responses API - api_base = ( - api_base or get_secret_str("MANUS_API_BASE") or "https://api.manus.im" - ) + api_base = api_base or get_secret_str("MANUS_API_BASE") or "https://api.manus.im" dynamic_api_key = api_key or get_secret_str("MANUS_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception("api base needs to be a string. api_base={}".format(api_base)) if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): - raise Exception( - "dynamic_api_key needs to be a string. dynamic_api_key={}".format( - dynamic_api_key - ) - ) + raise Exception("dynamic_api_key needs to be a string. dynamic_api_key={}".format(dynamic_api_key)) if dynamic_api_key is None and api_key is not None: dynamic_api_key = api_key return model, custom_llm_provider, dynamic_api_key, api_base diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 7679358bbc6..9126c35f818 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -37,9 +37,7 @@ class GetModelCostMap: def load_local_model_cost_map() -> dict: """Load the local backup model cost map bundled with the package.""" content = json.loads( - files("litellm") - .joinpath("model_prices_and_context_window_backup.json") - .read_text(encoding="utf-8") + files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8") ) return content @@ -56,16 +54,14 @@ class GetModelCostMap: """Check 1: fetched map is a non-empty dict.""" if not isinstance(fetched_map, dict): verbose_logger.warning( - "LiteLLM: Fetched model cost map is not a dict (type=%s). " - "Falling back to local backup.", + "LiteLLM: Fetched model cost map is not a dict (type=%s). Falling back to local backup.", type(fetched_map).__name__, ) return False if len(fetched_map) == 0: verbose_logger.warning( - "LiteLLM: Fetched model cost map is empty. " - "Falling back to local backup.", + "LiteLLM: Fetched model cost map is empty. Falling back to local backup.", ) return False @@ -92,10 +88,7 @@ class GetModelCostMap: ) return False - if ( - backup_model_count > 0 - and fetched_count < backup_model_count * max_shrink_ratio - ): + if backup_model_count > 0 and fetched_count < backup_model_count * max_shrink_ratio: verbose_logger.warning( "LiteLLM: Fetched model cost map shrank significantly " "(fetched=%d, backup=%d, threshold=%.0f%%). " @@ -269,8 +262,7 @@ def get_model_cost_map(url: str) -> dict: content = GetModelCostMap.fetch_remote_model_cost_map(url) except Exception as e: verbose_logger.warning( - "LiteLLM: Failed to fetch remote model cost map from %s: %s. " - "Falling back to local backup.", + "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, str(e), ) @@ -284,14 +276,11 @@ def get_model_cost_map(url: str) -> dict: backup_model_count=GetModelCostMap._get_backup_model_count(), ): verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. " - "Using local backup instead. url=%s", + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", url, ) _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = ( - "Remote data failed integrity validation" - ) + _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" return _expand_model_aliases(GetModelCostMap.load_local_model_cost_map()) _cost_map_source_info.source = "remote" diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index c22d3b99705..84f6445846b 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -8,9 +8,7 @@ from litellm.types.utils import LlmProviders, LlmProvidersSet def get_supported_openai_params( model: str, custom_llm_provider: Optional[str] = None, - request_type: Literal[ - "chat_completion", "embeddings", "transcription" - ] = "chat_completion", + request_type: Literal["chat_completion", "embeddings", "transcription"] = "chat_completion", base_model: Optional[str] = None, ) -> Optional[list]: """ @@ -56,12 +54,8 @@ def get_supported_openai_params( if provider_config and request_type == "chat_completion": supported_params = provider_config.get_supported_openai_params(model=model) if base_model and base_model != model: - base_model_params = provider_config.get_supported_openai_params( - model=base_model - ) - supported_params = list( - dict.fromkeys([*supported_params, *base_model_params]) - ) + base_model_params = provider_config.get_supported_openai_params(model=base_model) + supported_params = list(dict.fromkeys([*supported_params, *base_model_params])) return supported_params if custom_llm_provider == "bedrock": @@ -82,9 +76,7 @@ def get_supported_openai_params( return litellm.AnthropicTextConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "fireworks_ai": if request_type == "embeddings": - return litellm.FireworksAIEmbeddingConfig().get_supported_openai_params( - model=model - ) + return litellm.FireworksAIEmbeddingConfig().get_supported_openai_params(model=model) elif request_type == "transcription": return None else: @@ -107,9 +99,7 @@ def get_supported_openai_params( elif custom_llm_provider == "groq": return litellm.GroqChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "bedrock_mantle": - return litellm.BedrockMantleChatConfig().get_supported_openai_params( - model=model - ) + return litellm.BedrockMantleChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "hosted_vllm": return litellm.HostedVLLMChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "vllm": @@ -122,49 +112,27 @@ def get_supported_openai_params( return litellm.MaritalkConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "openai": if request_type == "transcription": - transcription_provider_config = ( - litellm.ProviderConfigManager.get_provider_audio_transcription_config( - model=model, provider=LlmProviders.OPENAI - ) + transcription_provider_config = litellm.ProviderConfigManager.get_provider_audio_transcription_config( + model=model, provider=LlmProviders.OPENAI ) - if isinstance( - transcription_provider_config, litellm.OpenAIGPTAudioTranscriptionConfig - ): - return transcription_provider_config.get_supported_openai_params( - model=model - ) + if isinstance(transcription_provider_config, litellm.OpenAIGPTAudioTranscriptionConfig): + return transcription_provider_config.get_supported_openai_params(model=model) else: - raise ValueError( - f"Unsupported provider config: {transcription_provider_config} for model: {model}" - ) + raise ValueError(f"Unsupported provider config: {transcription_provider_config} for model: {model}") return litellm.OpenAIConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "sap": if request_type == "chat_completion": - return litellm.GenAIHubOrchestrationConfig().get_supported_openai_params( - model=model - ) + return litellm.GenAIHubOrchestrationConfig().get_supported_openai_params(model=model) elif request_type == "embeddings": - return litellm.GenAIHubEmbeddingConfig().get_supported_openai_params( - model=model - ) + return litellm.GenAIHubEmbeddingConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "azure": _azure_detection_model = base_model or model - if litellm.AzureOpenAIO1Config().is_o_series_model( - model=_azure_detection_model - ): - return litellm.AzureOpenAIO1Config().get_supported_openai_params( - model=_azure_detection_model - ) - elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( - model=_azure_detection_model - ): - return litellm.AzureOpenAIGPT5Config().get_supported_openai_params( - model=_azure_detection_model - ) + if litellm.AzureOpenAIO1Config().is_o_series_model(model=_azure_detection_model): + return litellm.AzureOpenAIO1Config().get_supported_openai_params(model=_azure_detection_model) + elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=_azure_detection_model): + return litellm.AzureOpenAIGPT5Config().get_supported_openai_params(model=_azure_detection_model) else: - return litellm.AzureOpenAIConfig().get_supported_openai_params( - model=_azure_detection_model - ) + return litellm.AzureOpenAIConfig().get_supported_openai_params(model=_azure_detection_model) elif custom_llm_provider == "openrouter": return litellm.OpenrouterConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "vercel_ai_gateway": @@ -180,18 +148,12 @@ def get_supported_openai_params( MistralAudioTranscriptionConfig, ) - return MistralAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) + return MistralAudioTranscriptionConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "text-completion-codestral": - return litellm.CodestralTextCompletionConfig().get_supported_openai_params( - model=model - ) + return litellm.CodestralTextCompletionConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "sambanova": if request_type == "embeddings": - return litellm.SambaNovaEmbeddingConfig().get_supported_openai_params( - model=model - ) + return litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(model=model) else: return litellm.SambanovaConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "nebius": @@ -206,9 +168,7 @@ def get_supported_openai_params( return litellm.HuggingFaceChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "jina_ai": if request_type == "embeddings": - return litellm.JinaAIEmbeddingConfig().get_supported_openai_params( - model=model - ) + return litellm.JinaAIEmbeddingConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "together_ai": return litellm.TogetherAIConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "databricks": @@ -217,9 +177,7 @@ def get_supported_openai_params( elif request_type == "embeddings": return litellm.DatabricksEmbeddingConfig().get_supported_openai_params() elif custom_llm_provider == "palm" or custom_llm_provider == "gemini": - return litellm.GoogleAIStudioGeminiConfig().get_supported_openai_params( - model=model - ) + return litellm.GoogleAIStudioGeminiConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "novita": return litellm.NovitaConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": @@ -227,23 +185,13 @@ def get_supported_openai_params( if model.startswith("mistral"): return litellm.MistralConfig().get_supported_openai_params(model=model) elif model.startswith("codestral"): - return ( - litellm.CodestralTextCompletionConfig().get_supported_openai_params( - model=model - ) - ) + return litellm.CodestralTextCompletionConfig().get_supported_openai_params(model=model) elif model.startswith("claude"): - return litellm.VertexAIAnthropicConfig().get_supported_openai_params( - model=model - ) + return litellm.VertexAIAnthropicConfig().get_supported_openai_params(model=model) elif model.startswith("gemini"): - return litellm.VertexGeminiConfig().get_supported_openai_params( - model=model - ) + return litellm.VertexGeminiConfig().get_supported_openai_params(model=model) else: - return litellm.VertexAILlama3Config().get_supported_openai_params( - model=model - ) + return litellm.VertexAILlama3Config().get_supported_openai_params(model=model) elif request_type == "embeddings": return litellm.VertexAITextEmbeddingConfig().get_supported_openai_params() elif custom_llm_provider == "sagemaker": @@ -285,76 +233,48 @@ def get_supported_openai_params( return litellm.IBMWatsonXChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "watsonx_text": return litellm.IBMWatsonXAIConfig().get_supported_openai_params(model=model) - elif ( - custom_llm_provider == "custom_openai" - or custom_llm_provider == "text-completion-openai" - ): - return litellm.OpenAITextCompletionConfig().get_supported_openai_params( - model=model - ) + elif custom_llm_provider == "custom_openai" or custom_llm_provider == "text-completion-openai": + return litellm.OpenAITextCompletionConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "predibase": return litellm.PredibaseConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "voyage": - if ( - request_type == "embeddings" - and litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model) - ): - return ( - litellm.VoyageMultimodalEmbeddingConfig().get_supported_openai_params( - model=model - ) - ) + if request_type == "embeddings" and litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model): + return litellm.VoyageMultimodalEmbeddingConfig().get_supported_openai_params(model=model) return litellm.VoyageEmbeddingConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "infinity": - return litellm.InfinityEmbeddingConfig().get_supported_openai_params( - model=model - ) + return litellm.InfinityEmbeddingConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "triton": if request_type == "embeddings": - return litellm.TritonEmbeddingConfig().get_supported_openai_params( - model=model - ) + return litellm.TritonEmbeddingConfig().get_supported_openai_params(model=model) else: return litellm.TritonConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "deepgram": if request_type == "transcription": - return ( - litellm.DeepgramAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) - ) + return litellm.DeepgramAudioTranscriptionConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "ovhcloud": if request_type == "transcription": from litellm.llms.ovhcloud.audio_transcription.transformation import ( OVHCloudAudioTranscriptionConfig, ) - return OVHCloudAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) + return OVHCloudAudioTranscriptionConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "scaleway": if request_type == "transcription": from litellm.llms.scaleway.audio_transcription.transformation import ( ScalewayAudioTranscriptionConfig, ) - return ScalewayAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) + return ScalewayAudioTranscriptionConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "elevenlabs": if request_type == "transcription": from litellm.llms.elevenlabs.audio_transcription.transformation import ( ElevenLabsAudioTranscriptionConfig, ) - return ElevenLabsAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) + return ElevenLabsAudioTranscriptionConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "soniox": if request_type == "transcription": - return litellm.SonioxAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) + return litellm.SonioxAudioTranscriptionConfig().get_supported_openai_params(model=model) elif custom_llm_provider in litellm._custom_providers: if request_type == "chat_completion": provider_config = litellm.ProviderConfigManager.get_provider_chat_config( diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 5a29ea73a74..405366382a1 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -27,25 +27,19 @@ class HealthCheckHelpers: ) # this is a wildcard model, we need to pick a random model from the provider - cheapest_models = pick_cheapest_chat_models_from_llm_provider( - custom_llm_provider=custom_llm_provider, n=3 - ) + cheapest_models = pick_cheapest_chat_models_from_llm_provider(custom_llm_provider=custom_llm_provider, n=3) if len(cheapest_models) == 0: raise Exception( f"Unable to health check wildcard model for provider {custom_llm_provider}. Add a model on your config.yaml or contribute here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" ) if len(cheapest_models) > 1: - fallback_models = cheapest_models[ - 1: - ] # Pick the last 2 models from the shuffled list + fallback_models = cheapest_models[1:] # Pick the last 2 models from the shuffled list else: fallback_models = None model_params["model"] = cheapest_models[0] model_params["litellm_logging_obj"] = litellm_logging_obj model_params["fallbacks"] = fallback_models - model_params["max_tokens"] = model_params.get( - "max_tokens", 16 - ) # GPT-5 models require max_output_tokens >= 16 + model_params["max_tokens"] = model_params.get("max_tokens", 16) # GPT-5 models require max_output_tokens >= 16 await acompletion(**model_params) return {} @@ -167,12 +161,7 @@ class HealthCheckHelpers: "audio_speech": lambda: litellm.aspeech( **{ **_filter_model_params(model_params=model_params), - **( - {"voice": "alloy"} - if "voice" - not in _filter_model_params(model_params=model_params) - else {} - ), + **({"voice": "alloy"} if "voice" not in _filter_model_params(model_params=model_params) else {}), }, input=prompt or "test", ), diff --git a/litellm/litellm_core_utils/health_check_utils.py b/litellm/litellm_core_utils/health_check_utils.py index ff252855f0d..141facec040 100644 --- a/litellm/litellm_core_utils/health_check_utils.py +++ b/litellm/litellm_core_utils/health_check_utils.py @@ -11,17 +11,11 @@ def _filter_model_params(model_params: dict) -> dict: def _create_health_check_response(response_headers: dict) -> dict: response = {} - if ( - response_headers.get("x-ratelimit-remaining-requests", None) is not None - ): # not provided for dall-e requests - response["x-ratelimit-remaining-requests"] = response_headers[ - "x-ratelimit-remaining-requests" - ] + if response_headers.get("x-ratelimit-remaining-requests", None) is not None: # not provided for dall-e requests + response["x-ratelimit-remaining-requests"] = response_headers["x-ratelimit-remaining-requests"] if response_headers.get("x-ratelimit-remaining-tokens", None) is not None: - response["x-ratelimit-remaining-tokens"] = response_headers[ - "x-ratelimit-remaining-tokens" - ] + response["x-ratelimit-remaining-tokens"] = response_headers["x-ratelimit-remaining-tokens"] if response_headers.get("x-ms-region", None) is not None: response["x-ms-region"] = response_headers["x-ms-region"] diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index 949076aabf3..d0a6ec30c9e 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -23,9 +23,7 @@ def _raise_env_reference_error(param: str, *, source: str) -> None: ) -def validate_no_callback_env_reference( - param: str, value: object, *, source: str -) -> None: +def validate_no_callback_env_reference(param: str, value: object, *, source: str) -> None: if _is_env_reference(value): _raise_env_reference_error(param, source=source) @@ -86,9 +84,7 @@ def initialize_standard_callback_dynamic_params( continue if param in kwargs: _param_value = kwargs.get(param) - validate_no_callback_env_reference( - param, _param_value, source="request body" - ) + validate_no_callback_env_reference(param, _param_value, source="request body") standard_callback_dynamic_params[param] = _param_value # type: ignore # 2. Fallback: check "metadata" or "litellm_params" -> "metadata" @@ -103,9 +99,7 @@ def initialize_standard_callback_dynamic_params( continue if param not in standard_callback_dynamic_params and param in metadata: _param_value = metadata.get(param) - validate_no_callback_env_reference( - param, _param_value, source="metadata" - ) + validate_no_callback_env_reference(param, _param_value, source="metadata") standard_callback_dynamic_params[param] = _param_value # type: ignore return standard_callback_dynamic_params diff --git a/litellm/litellm_core_utils/json_validation_rule.py b/litellm/litellm_core_utils/json_validation_rule.py index bbfd3e6de96..c73b62f8a21 100644 --- a/litellm/litellm_core_utils/json_validation_rule.py +++ b/litellm/litellm_core_utils/json_validation_rule.py @@ -44,9 +44,7 @@ def normalize_json_schema_types( } if isinstance(schema, list): - return [ - normalize_json_schema_types(item, depth + 1, max_depth) for item in schema - ] + return [normalize_json_schema_types(item, depth + 1, max_depth) for item in schema] if isinstance(schema, dict): normalized_schema: Dict[str, Any] = {} @@ -57,21 +55,15 @@ def normalize_json_schema_types( elif key == "properties" and isinstance(value, dict): # Recursively normalize properties normalized_schema[key] = { - prop_key: normalize_json_schema_types( - prop_value, depth + 1, max_depth - ) + prop_key: normalize_json_schema_types(prop_value, depth + 1, max_depth) for prop_key, prop_value in value.items() } elif key == "items" and isinstance(value, (dict, list)): # Recursively normalize array items - normalized_schema[key] = normalize_json_schema_types( - value, depth + 1, max_depth - ) + normalized_schema[key] = normalize_json_schema_types(value, depth + 1, max_depth) elif isinstance(value, (dict, list)): # Recursively normalize any nested dict or list - normalized_schema[key] = normalize_json_schema_types( - value, depth + 1, max_depth - ) + normalized_schema[key] = normalize_json_schema_types(value, depth + 1, max_depth) else: normalized_schema[key] = value @@ -99,9 +91,7 @@ def normalize_tool_schema(tool: Dict[str, Any]) -> Dict[str, Any]: if "function" in tool and isinstance(tool["function"], dict): normalized_tool["function"] = tool["function"].copy() if "parameters" in tool["function"]: - normalized_tool["function"]["parameters"] = normalize_json_schema_types( - tool["function"]["parameters"] - ) + normalized_tool["function"]["parameters"] = normalize_json_schema_types(tool["function"]["parameters"]) return normalized_tool @@ -121,13 +111,9 @@ def validate_schema(schema: dict, response: str): try: response_dict = json.loads(response) except json.JSONDecodeError: - raise JSONSchemaValidationError( - model="", llm_provider="", raw_response=response, schema=json.dumps(schema) - ) + raise JSONSchemaValidationError(model="", llm_provider="", raw_response=response, schema=json.dumps(schema)) try: validate(response_dict, schema=schema) except ValidationError: - raise JSONSchemaValidationError( - model="", llm_provider="", raw_response=response, schema=json.dumps(schema) - ) + raise JSONSchemaValidationError(model="", llm_provider="", raw_response=response, schema=json.dumps(schema)) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1174907e31d..2457d117b81 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -197,13 +197,11 @@ try: from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger - EnterpriseStandardLoggingPayloadSetupVAR: Optional[ - Type[EnterpriseStandardLoggingPayloadSetup] - ] = EnterpriseStandardLoggingPayloadSetup -except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Unable to import GenericAPILogger - LiteLLM Enterprise Feature - {str(e)}" + EnterpriseStandardLoggingPayloadSetupVAR: Optional[Type[EnterpriseStandardLoggingPayloadSetup]] = ( + EnterpriseStandardLoggingPayloadSetup ) +except Exception as e: + verbose_logger.debug(f"[Non-Blocking] Unable to import GenericAPILogger - LiteLLM Enterprise Feature - {str(e)}") GenericAPILogger = CustomLogger # type: ignore ResendEmailLogger = CustomLogger # type: ignore SendGridEmailLogger = CustomLogger # type: ignore @@ -213,16 +211,12 @@ except Exception as e: EnterpriseStandardLoggingPayloadSetupVAR = None _in_memory_loggers: List[Any] = [] -_STANDARD_LOGGING_METADATA_KEYS: frozenset = frozenset( - StandardLoggingMetadata.__annotations__.keys() -) +_STANDARD_LOGGING_METADATA_KEYS: frozenset = frozenset(StandardLoggingMetadata.__annotations__.keys()) ### GLOBAL VARIABLES ### # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys -_CUSTOM_PRICING_KEYS: frozenset = frozenset( - CustomPricingLiteLLMParams.model_fields.keys() -) +_CUSTOM_PRICING_KEYS: frozenset = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) sentry_sdk_instance = None capture_exception = None @@ -318,21 +312,11 @@ class Logging(LiteLLMLoggingBaseClass): litellm_call_id: str, function_id: str, litellm_trace_id: Optional[str] = None, - dynamic_input_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_async_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_async_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, + dynamic_input_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, + dynamic_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, + dynamic_async_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, + dynamic_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, + dynamic_async_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = None, applied_guardrails: Optional[List[str]] = None, kwargs: Optional[Dict] = None, log_raw_request_response: bool = False, @@ -343,11 +327,7 @@ class Logging(LiteLLMLoggingBaseClass): messages = [ {"role": "user", "content": messages} ] # convert text completion input to the chat completion format - elif ( - isinstance(messages, list) - and len(messages) > 0 - and isinstance(messages[0], str) - ): + elif isinstance(messages, list) and len(messages) > 0 and isinstance(messages[0], str): new_messages = [] for m in messages: new_messages.append({"role": "user", "content": m}) @@ -364,32 +344,22 @@ class Logging(LiteLLMLoggingBaseClass): self.start_time = start_time # log the call start time self.call_type = call_type self.litellm_call_id = litellm_call_id - self.litellm_trace_id: str = ( - litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) - ) + self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[ - Any - ] = [] # for generating complete stream response + self.sync_streaming_chunks: List[Any] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks - self.dynamic_input_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_input_callbacks - self.dynamic_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_success_callbacks - self.dynamic_async_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_async_success_callbacks - self.dynamic_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_failure_callbacks - self.dynamic_async_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_async_failure_callbacks + self.dynamic_input_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = dynamic_input_callbacks + self.dynamic_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = dynamic_success_callbacks + self.dynamic_async_success_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( + dynamic_async_success_callbacks + ) + self.dynamic_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = dynamic_failure_callbacks + self.dynamic_async_failure_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( + dynamic_async_failure_callbacks + ) ## DYNAMIC LANGFUSE / GCS / logging callback KEYS ## self.standard_callback_dynamic_params: StandardCallbackDynamicParams = ( @@ -472,9 +442,7 @@ class Logging(LiteLLMLoggingBaseClass): def _process_dynamic_callback_list( self, callback_list: Optional[List[Union[str, Callable, CustomLogger]]], - dynamic_callbacks_type: Literal[ - "input", "success", "failure", "async_success", "async_failure" - ], + dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"], ) -> Optional[List[Union[str, Callable, CustomLogger]]]: """ Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -489,18 +457,13 @@ class Logging(LiteLLMLoggingBaseClass): processed_list: List[Union[str, Callable, CustomLogger]] = [] for callback in callback_list: - if ( - isinstance(callback, str) - and callback in litellm._known_custom_logger_compatible_callbacks - ): + if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks: # For callbacks that support team-scoped credentials (e.g. datadog), # pass only the relevant dynamic params as custom_logger_init_args. _custom_logger_init_args: Optional[dict] = None if callback == "datadog": _custom_logger_init_args = { - k: v - for k, v in self.standard_callback_dynamic_params.items() - if k.startswith("dd_") + k: v for k, v in self.standard_callback_dynamic_params.items() if k.startswith("dd_") } callback_class = _init_custom_logger_compatible_class( @@ -536,21 +499,15 @@ class Logging(LiteLLMLoggingBaseClass): return _initialize_standard_callback_dynamic_params(kwargs) - def initialize_standard_built_in_tools_params( - self, kwargs: Optional[Dict] = None - ) -> StandardBuiltInToolsParams: + def initialize_standard_built_in_tools_params(self, kwargs: Optional[Dict] = None) -> StandardBuiltInToolsParams: """ Initialize the standard built-in tools params from the kwargs checks if web_search_options in kwargs or tools and sets the corresponding attribute in StandardBuiltInToolsParams """ return StandardBuiltInToolsParams( - web_search_options=StandardBuiltInToolCostTracking._get_web_search_options( - kwargs or {} - ), - file_search=StandardBuiltInToolCostTracking._get_file_search_tool_call( - kwargs or {} - ), + web_search_options=StandardBuiltInToolCostTracking._get_web_search_options(kwargs or {}), + file_search=StandardBuiltInToolCostTracking._get_file_search_tool_call(kwargs or {}), ) def get_router_model_id(self) -> Optional[str]: @@ -613,10 +570,7 @@ class Logging(LiteLLMLoggingBaseClass): if "stream_options" in additional_params: self.stream_options = additional_params["stream_options"] ## check if custom pricing set ## - if any( - litellm_params.get(key) is not None - for key in _CUSTOM_PRICING_KEYS & litellm_params.keys() - ): + if any(litellm_params.get(key) is not None for key in _CUSTOM_PRICING_KEYS & litellm_params.keys()): self.custom_pricing = True if "custom_llm_provider" in self.model_call_details: @@ -640,9 +594,7 @@ class Logging(LiteLLMLoggingBaseClass): if "metadata" in kwargs: base_litellm_params["metadata"] = kwargs["metadata"] - if "litellm_metadata" in kwargs and isinstance( - kwargs["litellm_metadata"], dict - ): + if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): base_litellm_params["litellm_metadata"] = kwargs["litellm_metadata"] if "metadata" not in base_litellm_params: base_litellm_params["metadata"] = kwargs["litellm_metadata"].copy() @@ -738,15 +690,12 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ) -> Tuple[str, List[AllMessageValues], dict]: - custom_logger = ( - prompt_management_logger - or self.get_custom_logger_for_prompt_management( - model=model, - non_default_params=non_default_params, - prompt_id=prompt_id, - prompt_spec=prompt_spec, - dynamic_callback_params=self.standard_callback_dynamic_params, - ) + custom_logger = prompt_management_logger or self.get_custom_logger_for_prompt_management( + model=model, + non_default_params=non_default_params, + prompt_id=prompt_id, + prompt_spec=prompt_spec, + dynamic_callback_params=self.standard_callback_dynamic_params, ) if custom_logger: @@ -781,16 +730,13 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ) -> Tuple[str, List[AllMessageValues], dict]: - custom_logger = ( - prompt_management_logger - or self.get_custom_logger_for_prompt_management( - model=model, - tools=tools, - non_default_params=non_default_params, - prompt_id=prompt_id, - prompt_spec=prompt_spec, - dynamic_callback_params=self.standard_callback_dynamic_params, - ) + custom_logger = prompt_management_logger or self.get_custom_logger_for_prompt_management( + model=model, + tools=tools, + non_default_params=non_default_params, + prompt_id=prompt_id, + prompt_spec=prompt_spec, + dynamic_callback_params=self.standard_callback_dynamic_params, ) if custom_logger: @@ -832,10 +778,8 @@ class Logging(LiteLLMLoggingBaseClass): Returns: A CustomLogger instance if a matching prompt management system is found, None otherwise """ - prompt_management_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=CustomPromptManagement - ) + prompt_management_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=CustomPromptManagement ) for logger in prompt_management_loggers: @@ -846,9 +790,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_spec=prompt_spec, dynamic_callback_params=dynamic_callback_params, ): - self.model_call_details["prompt_integration"] = ( - logger.__class__.__name__ - ) + self.model_call_details["prompt_integration"] = logger.__class__.__name__ return logger except Exception: # If check fails, continue to next logger @@ -902,10 +844,8 @@ class Logging(LiteLLMLoggingBaseClass): return auto_detected_logger # Then check for any registered CustomPromptManagement loggers (fallback) - prompt_management_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=CustomPromptManagement - ) + prompt_management_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=CustomPromptManagement ) if prompt_management_loggers: @@ -915,13 +855,9 @@ class Logging(LiteLLMLoggingBaseClass): if ( anthropic_cache_control_logger - := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( - non_default_params - ) + := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook(non_default_params) ): - self.model_call_details["prompt_integration"] = ( - anthropic_cache_control_logger.__class__.__name__ - ) + self.model_call_details["prompt_integration"] = anthropic_cache_control_logger.__class__.__name__ return anthropic_cache_control_logger ######################################################### @@ -933,24 +869,15 @@ class Logging(LiteLLMLoggingBaseClass): internal_usage_cache=None, llm_router=None, ) - self.model_call_details["prompt_integration"] = ( - vector_store_custom_logger.__class__.__name__ - ) + self.model_call_details["prompt_integration"] = vector_store_custom_logger.__class__.__name__ # Add to global callbacks so post-call hooks are invoked - if ( - vector_store_custom_logger - and vector_store_custom_logger not in litellm.callbacks - ): - litellm.logging_callback_manager.add_litellm_callback( - vector_store_custom_logger - ) + if vector_store_custom_logger and vector_store_custom_logger not in litellm.callbacks: + litellm.logging_callback_manager.add_litellm_callback(vector_store_custom_logger) return vector_store_custom_logger return None - def get_custom_logger_for_anthropic_cache_control_hook( - self, non_default_params: Dict - ) -> Optional[CustomLogger]: + def get_custom_logger_for_anthropic_cache_control_hook(self, non_default_params: Dict) -> Optional[CustomLogger]: if non_default_params.get("cache_control_injection_points", None): custom_logger = _init_custom_logger_compatible_class( logging_integration="anthropic_cache_control_hook", @@ -967,9 +894,7 @@ class Logging(LiteLLMLoggingBaseClass): try: return json.loads(data) except Exception: - return { - "error": "Unable to parse raw request body. Got - {}".format(data) - } + return {"error": "Unable to parse raw request body. Got - {}".format(data)} return data def _get_masked_api_base(self, api_base: str) -> str: @@ -993,8 +918,8 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["log_event_type"] = "pre_api_call" if model: # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model - self.model_call_details["litellm_params"]["api_base"] = ( - self._get_masked_api_base(additional_args.get("api_base", "")) + self.model_call_details["litellm_params"]["api_base"] = self._get_masked_api_base( + additional_args.get("api_base", "") ) def pre_call(self, input, api_key, model=None, additional_args={}): @@ -1015,10 +940,7 @@ class Logging(LiteLLMLoggingBaseClass): additional_args=additional_args, ) # log raw request to provider (like LangFuse) -- if opted in. - if ( - self.log_raw_request_response is True - or log_raw_request_response is True - ): + if self.log_raw_request_response is True or log_raw_request_response is True: _litellm_params = self.model_call_details.get("litellm_params", {}) _metadata = _litellm_params.get("metadata", {}) or {} try: @@ -1036,28 +958,20 @@ class Logging(LiteLLMLoggingBaseClass): _metadata["raw_request"] = str(curl_command) # split up, so it's easier to parse in the UI - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, - ) + self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict( + raw_request_api_base=str(additional_args.get("api_base") or ""), + raw_request_body=self._get_raw_request_body(additional_args.get("complete_input_dict", {})), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, ) except Exception as e: - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - error=str(e), - ) + self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict( + error=str(e), ) _metadata["raw_request"] = "Unable to Log \ raw request: {}".format(str(e)) @@ -1068,9 +982,7 @@ class Logging(LiteLLMLoggingBaseClass): ) # Expectation: any logger function passed in by the user should accept a dict object except Exception as e: verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format(str(e)) ) self.model_call_details["api_call_start_time"] = datetime.datetime.now() @@ -1081,9 +993,7 @@ class Logging(LiteLLMLoggingBaseClass): # litellm_params["metadata"] (caller request metadata, typed # Dict[str, str], echoed downstream; a datetime breaks it). if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = ( - self.model_call_details["api_call_start_time"] - ) + self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks = litellm.input_callback + (self.dynamic_input_callbacks or []) for callback in callbacks: @@ -1123,9 +1033,7 @@ class Logging(LiteLLMLoggingBaseClass): messages=self.messages, kwargs=self.model_call_details, ) - elif ( - callable(callback) and customLogger is not None - ): # custom logger functions + elif callable(callback) and customLogger is not None: # custom logger functions customLogger.log_input_event( model=self.model, messages=self.messages, @@ -1134,11 +1042,7 @@ class Logging(LiteLLMLoggingBaseClass): callback_func=callback, ) except Exception as e: - verbose_logger.exception( - "litellm.Logging.pre_call(): Exception occured - {}".format( - str(e) - ) - ) + verbose_logger.exception("litellm.Logging.pre_call(): Exception occured - {}".format(str(e))) verbose_logger.debug( f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" ) @@ -1146,13 +1050,9 @@ class Logging(LiteLLMLoggingBaseClass): capture_exception(e) except Exception as e: verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - verbose_logger.error( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format(str(e)) ) + verbose_logger.error(f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}") if capture_exception: # log this error to sentry for debugging capture_exception(e) @@ -1212,12 +1112,8 @@ class Logging(LiteLLMLoggingBaseClass): curl_command += "curl -X POST \\\n" curl_command += f"{masked_api_base} \\\n" masked_headers = self._get_masked_headers(headers) - formatted_headers = " ".join( - [f"-H '{k}: {v}'" for k, v in masked_headers.items()] - ) - curl_command += ( - f"{formatted_headers} \\\n" if formatted_headers.strip() != "" else "" - ) + formatted_headers = " ".join([f"-H '{k}: {v}'" for k, v in masked_headers.items()]) + curl_command += f"{formatted_headers} \\\n" if formatted_headers.strip() != "" else "" curl_command += f"-d '{self._get_request_body(data)}'\n" if additional_args.get("request_str", None) is not None: # print the sagemaker / bedrock client request @@ -1228,21 +1124,15 @@ class Logging(LiteLLMLoggingBaseClass): curl_command = str(self.model_call_details) return curl_command - def _get_masked_headers( - self, headers: dict, ignore_sensitive_headers: bool = False - ) -> dict: + def _get_masked_headers(self, headers: dict, ignore_sensitive_headers: bool = False) -> dict: """ Internal debugging helper function Masks the headers of the request sent from LiteLLM """ - return _get_masked_values( - headers, ignore_sensitive_values=ignore_sensitive_headers - ) + return _get_masked_values(headers, ignore_sensitive_values=ignore_sensitive_headers) - def post_call( - self, original_response, input=None, api_key=None, additional_args={} - ): + def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received litellm.error_logs["POST_CALL"] = locals() if isinstance(original_response, dict): @@ -1263,18 +1153,14 @@ class Logging(LiteLLMLoggingBaseClass): callattr = getattr(verbose_logger, attr) callattr( "RAW RESPONSE:\n{}\n\n".format( - self.model_call_details.get( - "original_response", self.model_call_details - ) + self.model_call_details.get("original_response", self.model_call_details) ), ) else: callattr = getattr(verbose_logger, attr) callattr( "RAW RESPONSE:\n{}\n\n".format( - self.model_call_details.get( - "original_response", self.model_call_details - ) + self.model_call_details.get("original_response", self.model_call_details) ) ) if getattr(self, "logger_fn", None) and callable(self.logger_fn): @@ -1284,16 +1170,10 @@ class Logging(LiteLLMLoggingBaseClass): ) # Expectation: any logger function passed in by the user should accept a dict object except Exception as e: verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format(str(e)) ) original_response = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details - if hasattr(self, "model_call_details") - else {} - ), + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=original_response, ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made @@ -1338,9 +1218,7 @@ class Logging(LiteLLMLoggingBaseClass): capture_exception(e) except Exception as e: verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format(str(e)) ) async def async_post_mcp_tool_call_hook( @@ -1362,17 +1240,13 @@ class Logging(LiteLLMLoggingBaseClass): dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, ) - post_mcp_tool_call_response_obj: MCPPostCallResponseObject = ( - MCPPostCallResponseObject( - mcp_tool_call_response=response_obj, hidden_params=HiddenParams() - ) + post_mcp_tool_call_response_obj: MCPPostCallResponseObject = MCPPostCallResponseObject( + mcp_tool_call_response=response_obj, hidden_params=HiddenParams() ) for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[ - MCPPostCallResponseObject - ] = await callback.async_post_mcp_tool_call_hook( + response: Optional[MCPPostCallResponseObject] = await callback.async_post_mcp_tool_call_hook( kwargs=kwargs, response_obj=post_mcp_tool_call_response_obj, start_time=start_time, @@ -1383,20 +1257,14 @@ class Logging(LiteLLMLoggingBaseClass): # current implementation returns the first modified response ###################################################################### if response is not None: - response_obj = self._parse_post_mcp_call_hook_response( - response=response - ) + response_obj = self._parse_post_mcp_call_hook_response(response=response) except Exception as e: verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format(str(e)) ) return response_obj - def _parse_post_mcp_call_hook_response( - self, response: Optional[MCPPostCallResponseObject] - ) -> Any: + def _parse_post_mcp_call_hook_response(self, response: Optional[MCPPostCallResponseObject]) -> Any: """ Parse the response from the post_mcp_tool_call_hook @@ -1459,11 +1327,7 @@ class Logging(LiteLLMLoggingBaseClass): self.cost_breakdown["cache_creation_cost"] = cache_creation_cost # Store additional costs if provided (free-form dict for extensibility) - if ( - additional_costs - and isinstance(additional_costs, dict) - and len(additional_costs) > 0 - ): + if additional_costs and isinstance(additional_costs, dict) and len(additional_costs) > 0: self.cost_breakdown["additional_costs"] = additional_costs # Store discount information if provided @@ -1523,13 +1387,10 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"): hidden_params = getattr(result, "_hidden_params", {}) if ( - "response_cost" in hidden_params - and hidden_params["response_cost"] is not None + "response_cost" in hidden_params and hidden_params["response_cost"] is not None ): # use cost if already calculated return hidden_params["response_cost"] - elif ( - router_model_id is None and "model_id" in hidden_params - ): # use model_id if not already set + elif router_model_id is None and "model_id" in hidden_params: # use model_id if not already set router_model_id = hidden_params["model_id"] # Fallback: extract router_model_id from litellm_params when not available @@ -1540,9 +1401,7 @@ class Logging(LiteLLMLoggingBaseClass): ## RESPONSE COST ## custom_pricing = use_custom_pricing_for_model( - litellm_params=( - self.litellm_params if hasattr(self, "litellm_params") else None - ) + litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None) ) prompt = "" # use for tts cost calc @@ -1558,12 +1417,8 @@ class Logging(LiteLLMLoggingBaseClass): "response_object": result, "model": litellm_model_name or self.model, "cache_hit": cache_hit, - "custom_llm_provider": self.model_call_details.get( - "custom_llm_provider", None - ), - "base_model": _get_base_model_from_metadata( - model_call_details=self.model_call_details - ), + "custom_llm_provider": self.model_call_details.get("custom_llm_provider", None), + "base_model": _get_base_model_from_metadata(model_call_details=self.model_call_details), "call_type": self.call_type, "optional_params": self.optional_params, "custom_pricing": custom_pricing, @@ -1571,11 +1426,7 @@ class Logging(LiteLLMLoggingBaseClass): "standard_built_in_tools_params": self.standard_built_in_tools_params, "router_model_id": router_model_id, "litellm_logging_obj": self, - "service_tier": ( - self.optional_params.get("service_tier") - if self.optional_params - else None - ), + "service_tier": (self.optional_params.get("service_tier") if self.optional_params else None), "data_residency": ( self.litellm_params.get("data_residency") if hasattr(self, "litellm_params") and self.litellm_params @@ -1587,18 +1438,12 @@ class Logging(LiteLLMLoggingBaseClass): error_str=str(e), traceback_str=_get_traceback_str_for_error(str(e)), ) - verbose_logger.debug( - f"response_cost_failure_debug_information: {debug_info}" - ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + verbose_logger.debug(f"response_cost_failure_debug_information: {debug_info}") + self.model_call_details["response_cost_failure_debug_information"] = debug_info return None try: - response_cost = litellm.response_cost_calculator( - **response_cost_calculator_kwargs - ) + response_cost = litellm.response_cost_calculator(**response_cost_calculator_kwargs) verbose_logger.debug(f"response_cost: {response_cost}") return response_cost @@ -1608,19 +1453,13 @@ class Logging(LiteLLMLoggingBaseClass): traceback_str=_get_traceback_str_for_error(str(e)), model=response_cost_calculator_kwargs["model"], cache_hit=response_cost_calculator_kwargs["cache_hit"], - custom_llm_provider=response_cost_calculator_kwargs[ - "custom_llm_provider" - ], + custom_llm_provider=response_cost_calculator_kwargs["custom_llm_provider"], base_model=response_cost_calculator_kwargs["base_model"], call_type=response_cost_calculator_kwargs["call_type"], custom_pricing=response_cost_calculator_kwargs["custom_pricing"], ) - verbose_logger.debug( - f"response_cost_failure_debug_information: {debug_info}" - ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + verbose_logger.debug(f"response_cost_failure_debug_information: {debug_info}") + self.model_call_details["response_cost_failure_debug_information"] = debug_info return None @@ -1728,9 +1567,7 @@ class Logging(LiteLLMLoggingBaseClass): def should_run_logging( self, - event_type: Literal[ - "async_success", "sync_success", "async_failure", "sync_failure" - ], + event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], stream: bool = False, ) -> bool: try: @@ -1743,9 +1580,7 @@ class Logging(LiteLLMLoggingBaseClass): def has_run_logging( self, - event_type: Literal[ - "async_success", "sync_success", "async_failure", "sync_failure" - ], + event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], ) -> None: if self.stream is not None and self.stream is True: """ @@ -1755,32 +1590,22 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details[f"has_logged_{event_type}"] = True return - def should_run_callback( - self, callback: litellm.CALLBACK_TYPES, litellm_params: dict, event_hook: str - ) -> bool: + def should_run_callback(self, callback: litellm.CALLBACK_TYPES, litellm_params: dict, event_hook: str) -> bool: if litellm.global_disable_no_log_param: return True if litellm_params.get("no-log", False) is True: # proxy cost tracking cal backs should run - if not ( - isinstance(callback, CustomLogger) - and "_PROXY_" in callback.__class__.__name__ - ): - verbose_logger.debug( - f"no-log request, skipping logging for {event_hook} event" - ) + if not (isinstance(callback, CustomLogger) and "_PROXY_" in callback.__class__.__name__): + verbose_logger.debug(f"no-log request, skipping logging for {event_hook} event") return False # Check for dynamically disabled callbacks via headers - if ( - EnterpriseCallbackControls is not None - and EnterpriseCallbackControls.is_callback_disabled_dynamically( - callback=callback, - litellm_params=litellm_params, - standard_callback_dynamic_params=self.standard_callback_dynamic_params, - ) + if EnterpriseCallbackControls is not None and EnterpriseCallbackControls.is_callback_disabled_dynamically( + callback=callback, + litellm_params=litellm_params, + standard_callback_dynamic_params=self.standard_callback_dynamic_params, ): verbose_logger.debug( f"Callback {callback} disabled via x-litellm-disable-callbacks header for {event_hook} event" @@ -1800,14 +1625,12 @@ class Logging(LiteLLMLoggingBaseClass): """ logging_result = result if self.call_type == CallTypes.arealtime.value and isinstance(result, list): - combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( - results=result + combined_usage_object = ( + RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=result) ) - logging_result = ( - RealtimeAPITokenUsageProcessor.create_logging_realtime_object( - usage=combined_usage_object, - results=result, - ) + logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( + usage=combined_usage_object, + results=result, ) elif ( @@ -1823,9 +1646,7 @@ class Logging(LiteLLMLoggingBaseClass): if provider_config is not None: logging_result = provider_config.logging_non_streaming_response( model=self.model, - custom_llm_provider=self.model_call_details.get( - "custom_llm_provider", "" - ), + custom_llm_provider=self.model_call_details.get("custom_llm_provider", ""), httpx_response=result, request_data=self.model_call_details.get("request_data", {}), logging_obj=self, @@ -1833,9 +1654,7 @@ class Logging(LiteLLMLoggingBaseClass): ) return logging_result - def _merge_hidden_params_from_response_into_metadata( - self, logging_result: Any - ) -> None: + def _merge_hidden_params_from_response_into_metadata(self, logging_result: Any) -> None: """ Copy response._hidden_params into litellm_params.metadata['hidden_params']. @@ -1852,10 +1671,7 @@ class Logging(LiteLLMLoggingBaseClass): return metadata_hidden_params = hidden_params.copy() response_cost = self.model_call_details.get("response_cost") - if ( - metadata_hidden_params.get("response_cost") is None - and response_cost is not None - ): + if metadata_hidden_params.get("response_cost") is None and response_cost is not None: metadata_hidden_params["response_cost"] = response_cost litellm_params = self.model_call_details["litellm_params"] @@ -1876,40 +1692,30 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["litellm_params"].setdefault("metadata", {}) if self.model_call_details["litellm_params"]["metadata"] is None: self.model_call_details["litellm_params"]["metadata"] = {} - self.model_call_details["litellm_params"]["metadata"][ - "hidden_params" - ] = getattr(logging_result, "_hidden_params", {}) # type: ignore + self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = getattr( + logging_result, "_hidden_params", {} + ) # type: ignore if self.model_call_details.get("cache_hit") is True: self.model_call_details["response_cost"] = 0.0 elif "response_cost" in hidden_params: self.model_call_details["response_cost"] = hidden_params["response_cost"] - elif ( - existing_cost := self.model_call_details.get("response_cost") - ) is not None and existing_cost != 0: + elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0: # Preserve response_cost if already calculated (e.g., by pass-through # handlers like Gemini/Vertex which call completion_cost directly). # Do not preserve 0 from failure_handler on intermediate router retries. pass else: - self.model_call_details["response_cost"] = self._response_cost_calculator( - result=logging_result - ) + self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload(logging_result, start_time, end_time) + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + logging_result, start_time, end_time ) - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: + if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: emit_standard_logging_payload(standard_logging_payload) - def _build_standard_logging_payload( - self, init_response_obj: Any, start_time: Any, end_time: Any - ) -> Any: + def _build_standard_logging_payload(self, init_response_obj: Any, start_time: Any, end_time: Any) -> Any: """Build StandardLoggingPayload and accumulate its construction time.""" _start = time.time() payload = get_standard_logging_object_payload( @@ -1927,22 +1733,10 @@ class Logging(LiteLLMLoggingBaseClass): def _transform_usage_objects(self, result): if isinstance(result, ResponsesAPIResponse): result = result.model_copy() - transformed_usage = ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - result.usage - ) - ) + transformed_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(result.usage) setattr(result, "usage", transformed_usage) - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: - response_dict = ( - result.model_dump() - if hasattr(result, "model_dump") - else dict(result) - ) + if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: + response_dict = result.model_dump() if hasattr(result, "model_dump") else dict(result) # Ensure usage is properly included with transformed chat format if transformed_usage is not None: response_dict["usage"] = ( @@ -1978,9 +1772,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time = datetime.datetime.now() if self.completion_start_time is None: self.completion_start_time = end_time - self.model_call_details["completion_start_time"] = ( - self.completion_start_time - ) + self.model_call_details["completion_start_time"] = self.completion_start_time self.model_call_details["log_event_type"] = "successful_api_call" self.model_call_details["end_time"] = end_time @@ -1992,34 +1784,23 @@ class Logging(LiteLLMLoggingBaseClass): self.call_type == CallTypes.generate_content.value or self.call_type == CallTypes.agenerate_content.value ): - result = self._handle_non_streaming_google_genai_generate_content_response_logging( - result=result - ) - elif ( - self.call_type == CallTypes.asend_message.value - or self.call_type == CallTypes.send_message.value - ): + result = self._handle_non_streaming_google_genai_generate_content_response_logging(result=result) + elif self.call_type == CallTypes.asend_message.value or self.call_type == CallTypes.send_message.value: result = self._handle_a2a_response_logging(result=result) logging_result = self.normalize_logging_result(result=result) - if ( - standard_logging_object is None - and result is not None - and self.stream is not True - ): - if self._is_recognized_call_type_for_logging( - logging_result=logging_result - ) or isinstance(logging_result, (dict, list)): + if standard_logging_object is None and result is not None and self.stream is not True: + if self._is_recognized_call_type_for_logging(logging_result=logging_result) or isinstance( + logging_result, (dict, list) + ): self._process_hidden_params_and_response_cost( logging_result=logging_result, start_time=start_time, end_time=end_time, ) elif standard_logging_object is not None: - self.model_call_details["standard_logging_object"] = ( - standard_logging_object - ) + self.model_call_details["standard_logging_object"] = standard_logging_object else: self.model_call_details["response_cost"] = None @@ -2134,15 +1915,9 @@ class Logging(LiteLLMLoggingBaseClass): await self.async_success_handler(result=complete_streaming_response) return - def success_handler( - self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs - ): - verbose_logger.debug( - f"Logging Details LiteLLM-Success Call: Cache_hit={cache_hit}" - ) - if not self.should_run_logging( - event_type="sync_success" - ): # prevent double logging + def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): + verbose_logger.debug(f"Logging Details LiteLLM-Success Call: Cache_hit={cache_hit}") + if not self.should_run_logging(event_type="sync_success"): # prevent double logging return start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, @@ -2168,29 +1943,17 @@ class Logging(LiteLLMLoggingBaseClass): streaming_chunks=self.sync_streaming_chunks, ) if complete_streaming_response is not None: - verbose_logger.debug( - "Logging Details LiteLLM-Success Call streaming complete" - ) - self.model_call_details["complete_streaming_response"] = ( - complete_streaming_response - ) - self.model_call_details["response_cost"] = ( - self._response_cost_calculator(result=complete_streaming_response) - ) - self._merge_hidden_params_from_response_into_metadata( - complete_streaming_response + verbose_logger.debug("Logging Details LiteLLM-Success Call streaming complete") + self.model_call_details["complete_streaming_response"] = complete_streaming_response + self.model_call_details["response_cost"] = self._response_cost_calculator( + result=complete_streaming_response ) + self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time ) - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: + if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: # Only emit for sync requests (async_success_handler handles async) if is_sync_request: emit_standard_logging_payload(standard_logging_payload) @@ -2201,11 +1964,7 @@ class Logging(LiteLLMLoggingBaseClass): ## REDACT MESSAGES ## result = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details - if hasattr(self, "model_call_details") - else {} - ), + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=result, ) ## LOGGING HOOK ## @@ -2278,12 +2037,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, litellm_call_id=( current_call_id - if ( - current_call_id := litellm_params.get( - "litellm_call_id" - ) - ) - is not None + if (current_call_id := litellm_params.get("litellm_call_id")) is not None else str(uuid.uuid4()) ), print_verbose=print_verbose, @@ -2301,9 +2055,7 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("reaches logfire for success logging!") kwargs = {} for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends @@ -2330,11 +2082,7 @@ class Logging(LiteLLMLoggingBaseClass): input = kwargs.get("messages", kwargs.get("input", None)) - type = ( - "embed" - if self.call_type == CallTypes.embedding.value - else "llm" - ) + type = "embed" if self.call_type == CallTypes.embedding.value else "llm" # this only logs streaming once, complete_streaming_response exists i.e when stream ends if self.stream: @@ -2386,9 +2134,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose("reaches langfuse for success logging!") kwargs = {} for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends if self.stream: @@ -2425,9 +2171,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "greenscale" and greenscaleLogger is not None: kwargs = {} for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends if self.stream: @@ -2437,9 +2181,7 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is None: continue else: - print_verbose( - "reaches greenscale for streaming logging!" - ) + print_verbose("reaches greenscale for streaming logging!") result = kwargs["complete_streaming_response"] greenscaleLogger.log_event( @@ -2479,22 +2221,16 @@ class Logging(LiteLLMLoggingBaseClass): s3Logger = S3Logger() if self.stream: if "complete_streaming_response" in self.model_call_details: - print_verbose( - "S3Logger Logger: Got Stream Event - Completed Stream Response" - ) + print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response") s3Logger.log_event( kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "complete_streaming_response" - ], + response_obj=self.model_call_details["complete_streaming_response"], start_time=start_time, end_time=end_time, print_verbose=print_verbose, ) else: - print_verbose( - "S3Logger Logger: Got Stream Event - No complete stream response as yet" - ) + print_verbose("S3Logger Logger: Got Stream Event - No complete stream response as yet") else: s3Logger.log_event( kwargs=self.model_call_details, @@ -2518,10 +2254,8 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details["complete_response"] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] openMeterLogger.log_success_event( @@ -2545,10 +2279,8 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details["complete_response"] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] @@ -2559,15 +2291,9 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, ) if ( - callable(callback) is True - and is_sync_request - and customLogger is not None + callable(callback) is True and is_sync_request and customLogger is not None ): # custom logger functions - print_verbose( - "success callbacks: Running Custom Callback Function - {}".format( - callback - ) - ) + print_verbose("success callbacks: Running Custom Callback Function - {}".format(callback)) customLogger.log_event( kwargs=self.model_call_details, @@ -2582,9 +2308,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose( f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging with integrations {traceback.format_exc()}" ) - print_verbose( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" - ) + print_verbose(f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}") if capture_exception: # log this error to sentry for debugging capture_exception(e) # Track callback logging failures in Prometheus @@ -2594,31 +2318,21 @@ class Logging(LiteLLMLoggingBaseClass): pass except Exception as e: verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {}".format( - str(e) - ), + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {}".format(str(e)), ) - async def async_success_handler( - self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs - ): + async def async_success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): """ Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. """ - print_verbose( - "Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit) - ) - if not self._is_assembled_stream_success( - result - ) and not self.should_run_logging( + print_verbose("Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit)) + if not self._is_assembled_stream_success(result) and not self.should_run_logging( event_type="async_success" ): # prevent double logging (non-streaming) return ## CALCULATE COST FOR BATCH JOBS - if self.call_type == CallTypes.aretrieve_batch.value and isinstance( - result, LiteLLMBatch - ): + if self.call_type == CallTypes.aretrieve_batch.value and isinstance(result, LiteLLMBatch): litellm_params = self.litellm_params or {} litellm_metadata = litellm_params.get("litellm_metadata") or {} if ( @@ -2636,14 +2350,10 @@ class Logging(LiteLLMLoggingBaseClass): batch_cost = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) - has_explicit_batch_data = all( - x is not None for x in (batch_cost, batch_usage, batch_models) - ) + has_explicit_batch_data = all(x is not None for x in (batch_cost, batch_usage, batch_models)) should_compute_batch_data = ( - not is_base64_unified_file_id - or not has_explicit_batch_data - and result.status == "completed" + not is_base64_unified_file_id or not has_explicit_batch_data and result.status == "completed" ) if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost @@ -2676,69 +2386,51 @@ class Logging(LiteLLMLoggingBaseClass): ## BUILD COMPLETE STREAMED RESPONSE if "async_complete_streaming_response" in self.model_call_details: return # break out of this. - complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse] - ] = self._get_assembled_streaming_response( - result=result, - start_time=start_time, - end_time=end_time, - is_async=True, - streaming_chunks=self.streaming_chunks, + complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]] = ( + self._get_assembled_streaming_response( + result=result, + start_time=start_time, + end_time=end_time, + is_async=True, + streaming_chunks=self.streaming_chunks, + ) ) if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") - self.model_call_details["async_complete_streaming_response"] = ( - complete_streaming_response - ) + self.model_call_details["async_complete_streaming_response"] = complete_streaming_response try: if self.model_call_details.get("cache_hit", False) is True: self.model_call_details["response_cost"] = 0.0 else: # check if base_model set on azure - _get_base_model_from_metadata( - model_call_details=self.model_call_details - ) + _get_base_model_from_metadata(model_call_details=self.model_call_details) # base_model defaults to None if not set on model_info - self.model_call_details["response_cost"] = ( - self._response_cost_calculator( - result=complete_streaming_response - ) + self.model_call_details["response_cost"] = self._response_cost_calculator( + result=complete_streaming_response ) - verbose_logger.debug( - f"Model={self.model}; cost={self.model_call_details['response_cost']}" - ) + verbose_logger.debug(f"Model={self.model}; cost={self.model_call_details['response_cost']}") except litellm.NotFoundError: verbose_logger.warning( f"Model={self.model} not found in completion cost map. Setting 'response_cost' to None" ) self.model_call_details["response_cost"] = None - self._merge_hidden_params_from_response_into_metadata( - complete_streaming_response - ) + self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time ) # print standard logging payload - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: + if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: emit_standard_logging_payload(standard_logging_payload) elif self.call_type == "pass_through_endpoint": - print_verbose( - "Async success callbacks: Got a pass-through endpoint response" - ) + print_verbose("Async success callbacks: Got a pass-through endpoint response") self.model_call_details["async_complete_streaming_response"] = result @@ -2752,16 +2444,12 @@ class Logging(LiteLLMLoggingBaseClass): # _success_handler_helper_fn if self.model_call_details.get("standard_logging_object") is None: ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload(result, start_time, end_time) + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + result, start_time, end_time ) # print standard logging payload - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: + if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: emit_standard_logging_payload(standard_logging_payload) callbacks = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_async_success_callbacks, @@ -2769,9 +2457,7 @@ class Logging(LiteLLMLoggingBaseClass): ) result = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details if hasattr(self, "model_call_details") else {} - ), + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=result, ) @@ -2820,15 +2506,10 @@ class Logging(LiteLLMLoggingBaseClass): try: if callback == "openmeter" and openMeterLogger is not None: if self.stream is True: - if ( - "async_complete_streaming_response" - in self.model_call_details - ): + if "async_complete_streaming_response" in self.model_call_details: await openMeterLogger.async_log_success_event( kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "async_complete_streaming_response" - ], + response_obj=self.model_call_details["async_complete_streaming_response"], start_time=start_time, end_time=end_time, ) @@ -2859,9 +2540,7 @@ class Logging(LiteLLMLoggingBaseClass): if "async_complete_streaming_response" in model_call_details: await callback.async_log_success_event( kwargs=model_call_details, - response_obj=model_call_details[ - "async_complete_streaming_response" - ], + response_obj=model_call_details["async_complete_streaming_response"], start_time=start_time, end_time=end_time, ) @@ -2884,15 +2563,10 @@ class Logging(LiteLLMLoggingBaseClass): if customLogger is None: customLogger = CustomLogger() if self.stream: - if ( - "async_complete_streaming_response" - in self.model_call_details - ): + if "async_complete_streaming_response" in self.model_call_details: await customLogger.async_log_event( kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "async_complete_streaming_response" - ], + response_obj=self.model_call_details["async_complete_streaming_response"], start_time=start_time, end_time=end_time, print_verbose=print_verbose, @@ -2912,26 +2586,17 @@ class Logging(LiteLLMLoggingBaseClass): if dynamoLogger is None: dynamoLogger = DyanmoDBLogger() if self.stream: - if ( - "async_complete_streaming_response" - in self.model_call_details - ): - print_verbose( - "DynamoDB Logger: Got Stream Event - Completed Stream Response" - ) + if "async_complete_streaming_response" in self.model_call_details: + print_verbose("DynamoDB Logger: Got Stream Event - Completed Stream Response") await dynamoLogger._async_log_event( kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "async_complete_streaming_response" - ], + response_obj=self.model_call_details["async_complete_streaming_response"], start_time=start_time, end_time=end_time, print_verbose=print_verbose, ) else: - print_verbose( - "DynamoDB Logger: Got Stream Event - No complete stream response as yet" - ) + print_verbose("DynamoDB Logger: Got Stream Event - No complete stream response as yet") else: await dynamoLogger._async_log_event( kwargs=self.model_call_details, @@ -2963,17 +2628,13 @@ class Logging(LiteLLMLoggingBaseClass): for callback_obj in all_callbacks: if hasattr(callback_obj, "increment_callback_logging_failure"): - callback_obj.increment_callback_logging_failure( - callback_name=callback_name - ) # type: ignore + callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore break # Only increment once except Exception as e: verbose_logger.debug(f"Error in _handle_callback_failure: {str(e)}") - def _failure_handler_helper_fn( - self, exception, traceback_exception, start_time=None, end_time=None - ): + def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: start_time = self.start_time if end_time is None: @@ -2986,9 +2647,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["log_event_type"] = "failed_api_call" self.model_call_details["exception"] = exception self.model_call_details["traceback_exception"] = ( - _redact_string(traceback_exception) - if isinstance(traceback_exception, str) - else traceback_exception + _redact_string(traceback_exception) if isinstance(traceback_exception, str) else traceback_exception ) self.model_call_details["end_time"] = end_time self.model_call_details.setdefault("original_response", None) @@ -3001,25 +2660,21 @@ class Logging(LiteLLMLoggingBaseClass): if hasattr(exception, "headers") and isinstance(exception.headers, dict): self.model_call_details.setdefault("litellm_params", {}) - metadata = ( - self.model_call_details["litellm_params"].get("metadata", {}) or {} - ) + metadata = self.model_call_details["litellm_params"].get("metadata", {}) or {} metadata.update(exception.headers) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=_redact_string(str(exception)), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details["standard_logging_object"] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=_redact_string(str(exception)), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, ) return start_time, end_time @@ -3041,10 +2696,7 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(model_group_size, int) and model_group_size == 1: is_base_case = True ## check if special error ## - if ( - RouterErrors.no_deployments_available.value not in str(exception) - and is_base_case is False - ): + if RouterErrors.no_deployments_available.value not in str(exception) and is_base_case is False: return ## get original model group ## @@ -3058,15 +2710,9 @@ class Logging(LiteLLMLoggingBaseClass): kwargs=self.model_call_details, ) # type: ignore - def failure_handler( - self, exception, traceback_exception, start_time=None, end_time=None - ): - verbose_logger.debug( - f"Logging Details LiteLLM-Failure Call: {litellm.failure_callback}" - ) - if not self.should_run_logging( - event_type="sync_failure" - ): # prevent double logging + def failure_handler(self, exception, traceback_exception, start_time=None, end_time=None): + verbose_logger.debug(f"Logging Details LiteLLM-Failure Call: {litellm.failure_callback}") + if not self.should_run_logging(event_type="sync_failure"): # prevent double logging return litellm_params = self.model_call_details.get("litellm_params", {}) is_sync_request = self._is_sync_litellm_request(litellm_params) @@ -3086,11 +2732,7 @@ class Logging(LiteLLMLoggingBaseClass): result = None # result sent to all loggers, init this to None incase it's not created result = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details - if hasattr(self, "model_call_details") - else {} - ), + model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=result, ) self.has_run_logging(event_type="sync_failure") @@ -3110,11 +2752,7 @@ class Logging(LiteLLMLoggingBaseClass): input = self.model_call_details["input"] - _type = ( - "embed" - if self.call_type == CallTypes.embedding.value - else "llm" - ) + _type = "embed" if self.call_type == CallTypes.embedding.value else "llm" lunaryLogger.log_event( kwargs=self.model_call_details, @@ -3134,9 +2772,7 @@ class Logging(LiteLLMLoggingBaseClass): if capture_exception: capture_exception(exception) else: - print_verbose( - f"capture exception not initialized: {capture_exception}" - ) + print_verbose(f"capture exception not initialized: {capture_exception}") elif callback == "supabase" and supabaseClient is not None: print_verbose("reaches supabase for logging!") print_verbose(f"supabaseClient: {supabaseClient}") @@ -3178,9 +2814,7 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("reaches langfuse for logging failure") kwargs = {} for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( @@ -3220,9 +2854,7 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("reaches logfire for failure logging!") kwargs = {} for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v kwargs["exception"] = exception @@ -3239,28 +2871,20 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose( f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging with integrations {str(e)}" ) - print_verbose( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" - ) + print_verbose(f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}") if capture_exception: # log this error to sentry for debugging capture_exception(e) except Exception as e: verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {}".format( - str(e) - ) + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {}".format(str(e)) ) - async def async_failure_handler( - self, exception, traceback_exception, start_time=None, end_time=None - ): + async def async_failure_handler(self, exception, traceback_exception, start_time=None, end_time=None): """ Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. """ await self.special_failure_handlers(exception=exception) - if not self.should_run_logging( - event_type="async_failure" - ): # prevent double logging + if not self.should_run_logging(event_type="async_failure"): # prevent double logging return start_time, end_time = self._failure_handler_helper_fn( exception=exception, @@ -3343,39 +2967,24 @@ class Logging(LiteLLMLoggingBaseClass): if service_name == "langfuse": if langFuseLogger is None or ( ( - self.standard_callback_dynamic_params.get("langfuse_public_key") - is not None - and self.standard_callback_dynamic_params.get("langfuse_public_key") - != langFuseLogger.public_key + self.standard_callback_dynamic_params.get("langfuse_public_key") is not None + and self.standard_callback_dynamic_params.get("langfuse_public_key") != langFuseLogger.public_key ) or ( - self.standard_callback_dynamic_params.get("langfuse_public_key") - is not None - and self.standard_callback_dynamic_params.get("langfuse_public_key") - != langFuseLogger.public_key + self.standard_callback_dynamic_params.get("langfuse_public_key") is not None + and self.standard_callback_dynamic_params.get("langfuse_public_key") != langFuseLogger.public_key ) or ( - self.standard_callback_dynamic_params.get("langfuse_host") - is not None - and self.standard_callback_dynamic_params.get("langfuse_host") - != langFuseLogger.langfuse_host + self.standard_callback_dynamic_params.get("langfuse_host") is not None + and self.standard_callback_dynamic_params.get("langfuse_host") != langFuseLogger.langfuse_host ) ): return LangFuseLogger( - langfuse_public_key=self.standard_callback_dynamic_params.get( - "langfuse_public_key" - ), - langfuse_secret=self.standard_callback_dynamic_params.get( - "langfuse_secret" - ) + langfuse_public_key=self.standard_callback_dynamic_params.get("langfuse_public_key"), + langfuse_secret=self.standard_callback_dynamic_params.get("langfuse_secret") or self.standard_callback_dynamic_params.get("langfuse_secret_key"), - langfuse_host=self.standard_callback_dynamic_params.get( - "langfuse_host" - ), - allow_env_credentials=self.standard_callback_dynamic_params.get( - "langfuse_host" - ) - is None, + langfuse_host=self.standard_callback_dynamic_params.get("langfuse_host"), + allow_env_credentials=self.standard_callback_dynamic_params.get("langfuse_host") is None, ) return langFuseLogger @@ -3413,17 +3022,11 @@ class Logging(LiteLLMLoggingBaseClass): dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, ) - _filtered_success_callbacks = self._remove_internal_custom_logger_callbacks( - _combined_sync_callbacks - ) - _filtered_success_callbacks = self._remove_internal_litellm_callbacks( - _filtered_success_callbacks - ) + _filtered_success_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks) + _filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks) return len(_filtered_success_callbacks) > 0 - def get_combined_callback_list( - self, dynamic_success_callbacks: Optional[List], global_callbacks: List - ) -> List: + def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: if dynamic_success_callbacks is None: return list(global_callbacks) return list(set(dynamic_success_callbacks + global_callbacks)) @@ -3438,9 +3041,7 @@ class Logging(LiteLLMLoggingBaseClass): Returns: List of filtered callbacks with internal ones removed """ - filtered = [ - cb for cb in callbacks if not self._is_internal_litellm_proxy_callback(cb) - ] + filtered = [cb for cb in callbacks if not self._is_internal_litellm_proxy_callback(cb)] verbose_logger.debug(f"Filtered callbacks: {filtered}") return filtered @@ -3489,10 +3090,7 @@ class Logging(LiteLLMLoggingBaseClass): for _c in callbacks: if isinstance(_c, CustomLogger): continue - elif ( - isinstance(_c, str) - and _c in litellm._known_custom_logger_compatible_callbacks - ): + elif isinstance(_c, str) and _c in litellm._known_custom_logger_compatible_callbacks: continue _new_callbacks.append(_c) return _new_callbacks @@ -3523,10 +3121,8 @@ class Logging(LiteLLMLoggingBaseClass): ): ## return unified Usage object if isinstance(result.response.usage, ResponseAPIUsage): - transformed_usage = ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - result.response.usage - ) + transformed_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + result.response.usage ) # Set as dict instead of Usage object so model_dump() serializes it correctly setattr( @@ -3602,9 +3198,7 @@ class Logging(LiteLLMLoggingBaseClass): ) return result - def _translate_responses_api_response_to_model_response( - self, result: ResponsesAPIResponse - ) -> ModelResponse: + def _translate_responses_api_response_to_model_response(self, result: ResponsesAPIResponse) -> ModelResponse: """ Convert a Responses API response into a ModelResponse for spend_logs. @@ -3639,21 +3233,15 @@ class Logging(LiteLLMLoggingBaseClass): model_response = litellm.ModelResponse() model_response.model = self.model usage = getattr(result, "usage", None) - if usage is not None and ResponseAPILoggingUtils._is_response_api_usage( - usage - ): + if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage): setattr( model_response, "usage", - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ), + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage), ) return model_response - def _handle_non_streaming_google_genai_generate_content_response_logging( - self, result: Any - ) -> ModelResponse: + def _handle_non_streaming_google_genai_generate_content_response_logging(self, result: Any) -> ModelResponse: """ Handles logging for Google GenAI generate content responses. """ @@ -3695,9 +3283,7 @@ class Logging(LiteLLMLoggingBaseClass): # Deep copy result and add usage result_copy = result.model_copy(deep=True) - result_copy.usage = ( - usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) - ) + result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) return result_copy @@ -3748,25 +3334,14 @@ def _get_masked_values( if len(v) <= unmasked_length: return "*****" if number_of_asterisks is not None: - return ( - v[: unmasked_length // 2] - + "*" * number_of_asterisks - + v[-unmasked_length // 2 :] - ) - return ( - v[: unmasked_length // 2] - + "*" * (len(v) - unmasked_length) - + v[-unmasked_length // 2 :] - ) + return v[: unmasked_length // 2] + "*" * number_of_asterisks + v[-unmasked_length // 2 :] + return v[: unmasked_length // 2] + "*" * (len(v) - unmasked_length) + v[-unmasked_length // 2 :] return { k: ( v if ignore_sensitive_values - or not any( - sensitive_keyword in k.lower() - for sensitive_keyword in sensitive_keywords - ) + or not any(sensitive_keyword in k.lower() for sensitive_keyword in sensitive_keywords) else _mask_value(v) ) for k, v in sensitive_object.items() @@ -3808,33 +3383,23 @@ def set_callbacks(callback_list, function_id=None): import sentry_sdk except ImportError: print_verbose("Package 'sentry_sdk' is missing. Installing it...") - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "sentry_sdk"] - ) + subprocess.check_call([sys.executable, "-m", "pip", "install", "sentry_sdk"]) import sentry_sdk from sentry_sdk.scrubber import EventScrubber sentry_sdk_instance = sentry_sdk sentry_trace_rate = ( - os.environ.get("SENTRY_API_TRACE_RATE") - if "SENTRY_API_TRACE_RATE" in os.environ - else "1.0" + os.environ.get("SENTRY_API_TRACE_RATE") if "SENTRY_API_TRACE_RATE" in os.environ else "1.0" ) sentry_sample_rate = ( - os.environ.get("SENTRY_API_SAMPLE_RATE") - if "SENTRY_API_SAMPLE_RATE" in os.environ - else "1.0" + os.environ.get("SENTRY_API_SAMPLE_RATE") if "SENTRY_API_SAMPLE_RATE" in os.environ else "1.0" ) sentry_sdk_instance.init( dsn=os.environ.get("SENTRY_DSN"), traces_sample_rate=float(sentry_trace_rate), # type: ignore - sample_rate=float( - sentry_sample_rate if sentry_sample_rate else 1.0 - ), + sample_rate=float(sentry_sample_rate if sentry_sample_rate else 1.0), send_default_pii=False, # Prevent sending Personal Identifiable Information - event_scrubber=EventScrubber( - denylist=SENTRY_DENYLIST, pii_denylist=SENTRY_PII_DENYLIST - ), + event_scrubber=EventScrubber(denylist=SENTRY_DENYLIST, pii_denylist=SENTRY_PII_DENYLIST), environment=os.environ.get("SENTRY_ENVIRONMENT", "production"), ) capture_exception = sentry_sdk_instance.capture_exception @@ -3844,9 +3409,7 @@ def set_callbacks(callback_list, function_id=None): from slack_bolt import App except ImportError: print_verbose("Package 'slack_bolt' is missing. Installing it...") - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "slack_bolt"] - ) + subprocess.check_call([sys.executable, "-m", "pip", "install", "slack_bolt"]) from slack_bolt import App slack_app = App( token=os.environ.get("SLACK_API_TOKEN"), @@ -3866,9 +3429,7 @@ def set_callbacks(callback_list, function_id=None): elif callback == "promptlayer": promptLayerLogger = PromptLayerLogger() elif callback == "langfuse": - langFuseLogger = LangFuseLogger( - langfuse_public_key=None, langfuse_secret=None, langfuse_host=None - ) + langFuseLogger = LangFuseLogger(langfuse_public_key=None, langfuse_secret=None, langfuse_host=None) elif callback == "openmeter": openMeterLogger = OpenMeterLogger() elif callback == "datadog": @@ -3899,9 +3460,7 @@ def set_callbacks(callback_list, function_id=None): def _init_custom_logger_compatible_class( logging_integration: _custom_logger_compatible_callbacks_literal, internal_usage_cache: Optional[DualCache], - llm_router: Optional[ - Any - ], # expect litellm.Router, but typing errors due to circular import + llm_router: Optional[Any], # expect litellm.Router, but typing errors due to circular import custom_logger_init_args: Optional[dict] = {}, ) -> Optional[CustomLogger]: """ @@ -4106,10 +3665,7 @@ def _init_custom_logger_compatible_class( f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" ) for callback in _in_memory_loggers: - if ( - isinstance(callback, ArizeLogger) - and callback.callback_name == "arize" - ): + if isinstance(callback, ArizeLogger) and callback.callback_name == "arize": return callback # type: ignore _arize_otel_logger = ArizeLogger(config=otel_config, callback_name="arize") _in_memory_loggers.append(_arize_otel_logger) @@ -4132,19 +3688,12 @@ def _init_custom_logger_compatible_class( # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - arize_phoenix_config.otlp_auth_headers - ) + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = arize_phoenix_config.otlp_auth_headers for callback in _in_memory_loggers: - if ( - isinstance(callback, ArizePhoenixLogger) - and callback.callback_name == "arize_phoenix" - ): + if isinstance(callback, ArizePhoenixLogger) and callback.callback_name == "arize_phoenix": return callback # type: ignore - _arize_phoenix_otel_logger = ArizePhoenixLogger( - config=otel_config, callback_name="arize_phoenix" - ) + _arize_phoenix_otel_logger = ArizePhoenixLogger(config=otel_config, callback_name="arize_phoenix") _in_memory_loggers.append(_arize_phoenix_otel_logger) return _arize_phoenix_otel_logger # type: ignore elif logging_integration == "levo": @@ -4166,10 +3715,7 @@ def _init_custom_logger_compatible_class( # Check if LevoLogger instance already exists for callback in _in_memory_loggers: - if ( - isinstance(callback, LevoLogger) - and callback.callback_name == "levo" - ): + if isinstance(callback, LevoLogger) and callback.callback_name == "levo": return callback # type: ignore _levo_otel_logger = LevoLogger(config=otel_config, callback_name="levo") @@ -4190,9 +3736,7 @@ def _init_custom_logger_compatible_class( if type(callback) is OpenTelemetryV2: return callback # type: ignore otel_logger_v2 = OpenTelemetryV2( - **_get_custom_logger_settings_from_proxy_server( - callback_name=logging_integration - ) + **_get_custom_logger_settings_from_proxy_server(callback_name=logging_integration) ) _in_memory_loggers.append(otel_logger_v2) _maybe_auto_initialize_arize_phoenix(_in_memory_loggers) @@ -4204,9 +3748,7 @@ def _init_custom_logger_compatible_class( if type(callback) is OpenTelemetry: return callback # type: ignore otel_logger = OpenTelemetry( - **_get_custom_logger_settings_from_proxy_server( - callback_name=logging_integration - ) + **_get_custom_logger_settings_from_proxy_server(callback_name=logging_integration) ) _in_memory_loggers.append(otel_logger) @@ -4238,9 +3780,7 @@ def _init_custom_logger_compatible_class( from litellm.integrations.focus.focus_logger import FocusLogger for callback in _in_memory_loggers: - if ( - type(callback) is FocusLogger - ): # exact match; exclude subclasses like VantageLogger + if type(callback) is FocusLogger: # exact match; exclude subclasses like VantageLogger return callback # type: ignore focus_logger = FocusLogger() _in_memory_loggers.append(focus_logger) @@ -4281,9 +3821,7 @@ def _init_custom_logger_compatible_class( OpenTelemetryConfig, ) - logfire_base_url = os.getenv( - "LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev" - ) + logfire_base_url = os.getenv("LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev") otel_config = OpenTelemetryConfig( exporter="otlp_http", endpoint=f"{logfire_base_url.rstrip('/')}/v1/traces", @@ -4307,14 +3845,10 @@ def _init_custom_logger_compatible_class( if internal_usage_cache is None: raise Exception( - "Internal Error: Cache cannot be empty - internal_usage_cache={}".format( - internal_usage_cache - ) + "Internal Error: Cache cannot be empty - internal_usage_cache={}".format(internal_usage_cache) ) - dynamic_rate_limiter_obj = _PROXY_DynamicRateLimitHandler( - internal_usage_cache=internal_usage_cache - ) + dynamic_rate_limiter_obj = _PROXY_DynamicRateLimitHandler(internal_usage_cache=internal_usage_cache) if llm_router is not None and isinstance(llm_router, litellm.Router): dynamic_rate_limiter_obj.update_variables(llm_router=llm_router) @@ -4331,14 +3865,10 @@ def _init_custom_logger_compatible_class( if internal_usage_cache is None: raise Exception( - "Internal Error: Cache cannot be empty - internal_usage_cache={}".format( - internal_usage_cache - ) + "Internal Error: Cache cannot be empty - internal_usage_cache={}".format(internal_usage_cache) ) - dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3( - internal_usage_cache=internal_usage_cache - ) + dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=internal_usage_cache) if llm_router is not None and isinstance(llm_router, litellm.Router): dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router) @@ -4360,14 +3890,9 @@ def _init_custom_logger_compatible_class( exporter="otlp_http", endpoint="https://langtrace.ai/api/trace", ) - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"api_key={os.getenv('LANGTRACE_API_KEY')}" - ) + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" for callback in _in_memory_loggers: - if ( - isinstance(callback, OpenTelemetry) - and callback.callback_name == "langtrace" - ): + if isinstance(callback, OpenTelemetry) and callback.callback_name == "langtrace": return callback # type: ignore _otel_logger = OpenTelemetry(config=otel_config, callback_name="langtrace") _in_memory_loggers.append(_otel_logger) @@ -4396,16 +3921,11 @@ def _init_custom_logger_compatible_class( from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger for callback in _in_memory_loggers: - if ( - isinstance(callback, LangfuseOtelLogger) - and callback.callback_name == "langfuse_otel" - ): + if isinstance(callback, LangfuseOtelLogger) and callback.callback_name == "langfuse_otel": return callback # type: ignore # Allow LangfuseOtelLogger to initialize its own config safely # This prevents startup crashes if LANGFUSE keys are not in env (e.g. for dynamic usage) - _otel_logger = LangfuseOtelLogger( - config=None, callback_name="langfuse_otel" - ) + _otel_logger = LangfuseOtelLogger(config=None, callback_name="langfuse_otel") _in_memory_loggers.append(_otel_logger) return _otel_logger # type: ignore elif logging_integration == "weave_otel": @@ -4427,14 +3947,9 @@ def _init_custom_logger_compatible_class( ) for callback in _in_memory_loggers: - if ( - isinstance(callback, WeaveOtelLogger) - and callback.callback_name == "weave_otel" - ): + if isinstance(callback, WeaveOtelLogger) and callback.callback_name == "weave_otel": return callback # type: ignore - _otel_logger = WeaveOtelLogger( - config=otel_config, callback_name="weave_otel" - ) + _otel_logger = WeaveOtelLogger(config=otel_config, callback_name="weave_otel") _in_memory_loggers.append(_otel_logger) return _otel_logger # type: ignore elif logging_integration == "pagerduty": @@ -4525,9 +4040,7 @@ def _init_custom_logger_compatible_class( # Get global BitBucket config bitbucket_config = getattr(litellm, "global_bitbucket_config", None) if bitbucket_config is None: - raise ValueError( - "BitBucket configuration not found. Please set litellm.global_bitbucket_config first." - ) + raise ValueError("BitBucket configuration not found. Please set litellm.global_bitbucket_config first.") bitbucket_logger = BitBucketPromptManager(bitbucket_config=bitbucket_config) _in_memory_loggers.append(bitbucket_logger) @@ -4544,9 +4057,7 @@ def _init_custom_logger_compatible_class( # Get global BitBucket config gitlab_config = getattr(litellm, "global_gitlab_config", None) if gitlab_config is None: - raise ValueError( - "Gitlab configuration not found. Please set litellm.global_gitlab_config first." - ) + raise ValueError("Gitlab configuration not found. Please set litellm.global_gitlab_config first.") gitlab_logger = GitLabPromptManager(gitlab_config=gitlab_config) _in_memory_loggers.append(gitlab_logger) @@ -4560,16 +4071,12 @@ def _init_custom_logger_compatible_class( return newrelic_logger # type: ignore return None except Exception as e: - verbose_logger.exception( - f"[Non-Blocking Error] Error initializing custom logger: {e}" - ) + verbose_logger.exception(f"[Non-Blocking Error] Error initializing custom logger: {e}") return None return None -def _maybe_construct_otel_v2( - callback_name: str, _in_memory_loggers: list -) -> Optional[Any]: +def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> Optional[Any]: """If ``LITELLM_OTEL_V2`` is on, build (or reuse) a single ``OpenTelemetryV2`` instance configured via the preset for ``callback_name``. @@ -4587,10 +4094,7 @@ def _maybe_construct_otel_v2( if preset_fn is None: return None for callback in _in_memory_loggers: - if ( - isinstance(callback, OpenTelemetryV2) - and getattr(callback, "callback_name", None) == callback_name - ): + if isinstance(callback, OpenTelemetryV2) and getattr(callback, "callback_name", None) == callback_name: return callback try: config = preset_fn() @@ -4620,10 +4124,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: return # Already registered — nothing to do - if any( - isinstance(cb, ArizePhoenixLogger) and cb.callback_name == "arize_phoenix" - for cb in _in_memory_loggers - ): + if any(isinstance(cb, ArizePhoenixLogger) and cb.callback_name == "arize_phoenix" for cb in _in_memory_loggers): return try: @@ -4635,9 +4136,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: endpoint=arize_phoenix_config.endpoint, headers=arize_phoenix_config.otlp_auth_headers, ) - phoenix_logger = ArizePhoenixLogger( - config=otel_config, callback_name="arize_phoenix" - ) + phoenix_logger = ArizePhoenixLogger(config=otel_config, callback_name="arize_phoenix") _in_memory_loggers.append(phoenix_logger) # Register as a litellm callback so it receives success/failure events @@ -4648,9 +4147,7 @@ def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: arize_phoenix_config.endpoint, ) except Exception as e: - verbose_logger.warning( - "Failed to auto-initialize Arize Phoenix logger: %s", str(e) - ) + verbose_logger.warning("Failed to auto-initialize Arize Phoenix logger: %s", str(e)) def get_custom_logger_compatible_class( @@ -4685,9 +4182,7 @@ def get_custom_logger_compatible_class( from litellm.integrations.focus.focus_logger import FocusLogger for callback in _in_memory_loggers: - if ( - type(callback) is FocusLogger - ): # exact match; exclude subclasses like VantageLogger + if type(callback) is FocusLogger: # exact match; exclude subclasses like VantageLogger return callback elif logging_integration == "vantage": from litellm.integrations.vantage.vantage_logger import VantageLogger @@ -4774,10 +4269,7 @@ def get_custom_logger_compatible_class( if "ARIZE_API_KEY" not in os.environ: raise ValueError("ARIZE_API_KEY not found in environment variables") for callback in _in_memory_loggers: - if ( - isinstance(callback, ArizeLogger) - and callback.callback_name == "arize" - ): + if isinstance(callback, ArizeLogger) and callback.callback_name == "arize": return callback elif logging_integration == "logfire": if "LOGFIRE_TOKEN" not in os.environ: @@ -4813,10 +4305,7 @@ def get_custom_logger_compatible_class( raise ValueError("LANGTRACE_API_KEY not found in environment variables") for callback in _in_memory_loggers: - if ( - isinstance(callback, OpenTelemetry) - and callback.callback_name == "langtrace" - ): + if isinstance(callback, OpenTelemetry) and callback.callback_name == "langtrace": return callback elif logging_integration == "mlflow": @@ -4866,9 +4355,7 @@ def get_custom_logger_compatible_class( return None except Exception as e: - verbose_logger.exception( - f"[Non-Blocking Error] Error getting custom logger: {e}" - ) + verbose_logger.exception(f"[Non-Blocking Error] Error getting custom logger: {e}") return None @@ -4946,18 +4433,14 @@ class StandardLoggingPayloadSetup: elif isinstance(start_time, float): start_time_float = start_time else: - raise ValueError( - f"start_time is required, got={start_time} of type {type(start_time)}" - ) + raise ValueError(f"start_time is required, got={start_time} of type {type(start_time)}") if isinstance(end_time, datetime.datetime): end_time_float = end_time.timestamp() elif isinstance(end_time, float): end_time_float = end_time else: - raise ValueError( - f"end_time is required, got={end_time} of type {type(end_time)}" - ) + raise ValueError(f"end_time is required, got={end_time} of type {type(end_time)}") if isinstance(completion_start_time, datetime.datetime): completion_start_time_float = completion_start_time.timestamp() @@ -4969,29 +4452,21 @@ class StandardLoggingPayloadSetup: return start_time_float, end_time_float, completion_start_time_float @staticmethod - def append_system_prompt_messages( - kwargs: Optional[Dict] = None, messages: Optional[Any] = None - ): + def append_system_prompt_messages(kwargs: Optional[Dict] = None, messages: Optional[Any] = None): """ Append system prompt messages to the messages """ if kwargs is not None: - if kwargs.get("system") is not None and isinstance( - kwargs.get("system"), str - ): + if kwargs.get("system") is not None and isinstance(kwargs.get("system"), str): if messages is None: return [{"role": "system", "content": kwargs.get("system")}] elif isinstance(messages, list): if len(messages) == 0: return [{"role": "system", "content": kwargs.get("system")}] # check for duplicates - if messages[0].get("role") == "system" and messages[0].get( - "content" - ) == kwargs.get("system"): + if messages[0].get("role") == "system" and messages[0].get("content") == kwargs.get("system"): return messages - messages = [ - {"role": "system", "content": kwargs.get("system")} - ] + messages + messages = [{"role": "system", "content": kwargs.get("system")}] + messages elif isinstance(messages, str): messages = [ {"role": "system", "content": kwargs.get("system")}, @@ -5018,9 +4493,7 @@ class StandardLoggingPayloadSetup: merged_metadata: dict = {} # Start with metadata (user API key fields) - but skip non-serializable objects - if litellm_params.get("metadata") and isinstance( - litellm_params.get("metadata"), dict - ): + if litellm_params.get("metadata") and isinstance(litellm_params.get("metadata"), dict): for key, value in litellm_params["metadata"].items(): # Skip non-serializable objects like UserAPIKeyAuth if key in {"user_api_key_auth", "user_api_key_budget_reservation"}: @@ -5028,13 +4501,9 @@ class StandardLoggingPayloadSetup: merged_metadata[key] = value # Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys - if litellm_params.get("litellm_metadata") and isinstance( - litellm_params.get("litellm_metadata"), dict - ): + if litellm_params.get("litellm_metadata") and isinstance(litellm_params.get("litellm_metadata"), dict): for key, value in litellm_params["litellm_metadata"].items(): - if ( - key not in merged_metadata - ): # Don't overwrite existing keys from metadata + if key not in merged_metadata: # Don't overwrite existing keys from metadata merged_metadata[key] = value return merged_metadata @@ -5046,9 +4515,7 @@ class StandardLoggingPayloadSetup: prompt_integration: Optional[str] = None, applied_guardrails: Optional[List[str]] = None, mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = None, - vector_store_request_metadata: Optional[ - List[StandardLoggingVectorStoreRequest] - ] = None, + vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = None, usage_object: Optional[dict] = None, proxy_server_request: Optional[dict] = None, start_time: Optional[dt_object] = None, @@ -5068,14 +4535,10 @@ class StandardLoggingPayloadSetup: - If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'. """ - prompt_management_metadata: Optional[ - StandardLoggingPromptManagementMetadata - ] = None + prompt_management_metadata: Optional[StandardLoggingPromptManagementMetadata] = None if litellm_params is not None: prompt_id = cast(Optional[str], litellm_params.get("prompt_id", None)) - prompt_variables = cast( - Optional[dict], litellm_params.get("prompt_variables", None) - ) + prompt_variables = cast(Optional[dict], litellm_params.get("prompt_variables", None)) if prompt_id is not None and prompt_integration is not None: prompt_management_metadata = StandardLoggingPromptManagementMetadata( @@ -5121,11 +4584,7 @@ class StandardLoggingPayloadSetup: clean_metadata[key] = metadata[key] # type: ignore user_api_key = metadata.get("user_api_key") - if ( - user_api_key - and isinstance(user_api_key, str) - and is_valid_sha256_hash(user_api_key) - ): + if user_api_key and isinstance(user_api_key, str) and is_valid_sha256_hash(user_api_key): clean_metadata["user_api_key_hash"] = user_api_key _potential_requester_metadata = metadata.get( "metadata", None @@ -5137,10 +4596,7 @@ class StandardLoggingPayloadSetup: ): clean_metadata["requester_metadata"] = _potential_requester_metadata - if ( - EnterpriseStandardLoggingPayloadSetupVAR - and proxy_server_request is not None - ): + if EnterpriseStandardLoggingPayloadSetupVAR and proxy_server_request is not None: clean_metadata = EnterpriseStandardLoggingPayloadSetupVAR.apply_enterprise_specific_metadata( standard_logging_metadata=clean_metadata, proxy_server_request=proxy_server_request, @@ -5148,12 +4604,10 @@ class StandardLoggingPayloadSetup: # Generate cold storage object key if cold storage is configured if start_time is not None and response_id is not None: - cold_storage_object_key = ( - StandardLoggingPayloadSetup._generate_cold_storage_object_key( - start_time=start_time, - response_id=response_id, - team_alias=clean_metadata.get("user_api_key_team_alias"), - ) + cold_storage_object_key = StandardLoggingPayloadSetup._generate_cold_storage_object_key( + start_time=start_time, + response_id=response_id, + team_alias=clean_metadata.get("user_api_key_team_alias"), ) if cold_storage_object_key: clean_metadata["cold_storage_object_key"] = cold_storage_object_key @@ -5175,9 +4629,7 @@ class StandardLoggingPayloadSetup: ) usage = response_obj.get("usage", None) or {} - if usage is None or ( - not isinstance(usage, dict) and not isinstance(usage, Usage) - ): + if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)): return Usage( prompt_tokens=0, completion_tokens=0, @@ -5186,16 +4638,10 @@ class StandardLoggingPayloadSetup: elif isinstance(usage, Usage): return usage elif isinstance(usage, ResponseAPIUsage): - return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) elif isinstance(usage, dict): if ResponseAPILoggingUtils._is_response_api_usage(usage): - return ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) - ) + return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) return Usage(**usage) raise ValueError(f"usage is required, got={usage} of type {type(usage)}") @@ -5218,16 +4664,10 @@ class StandardLoggingPayloadSetup: if _raw is None: return _empty if isinstance(_raw, ResponseAPIUsage): - return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - _raw - ).model_dump() + return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_raw).model_dump() if isinstance(_raw, dict): if ResponseAPILoggingUtils._is_response_api_usage(_raw): - return ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - _raw - ).model_dump() - ) + return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_raw).model_dump() return _raw if isinstance(_raw, Usage): return _raw.model_dump() @@ -5248,9 +4688,7 @@ class StandardLoggingPayloadSetup: custom_pricing=custom_pricing, ) if model_cost_name is None: - model_cost_information = StandardLoggingModelInformation( - model_map_key="", model_map_value=None - ) + model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None) else: try: _model_cost_information = litellm.get_model_info( @@ -5292,9 +4730,7 @@ class StandardLoggingPayloadSetup: result=final_response_obj, ) - if modified_final_response_obj is not None and isinstance( - modified_final_response_obj, BaseModel - ): + if modified_final_response_obj is not None and isinstance(modified_final_response_obj, BaseModel): final_response_obj = modified_final_response_obj.model_dump() else: final_response_obj = modified_final_response_obj @@ -5347,10 +4783,8 @@ class StandardLoggingPayloadSetup: for key in StandardLoggingHiddenParams.__annotations__.keys(): if key in hidden_params: if key == "additional_headers": - clean_hidden_params["additional_headers"] = ( - StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] - ) + clean_hidden_params["additional_headers"] = StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] ) else: clean_hidden_params[key] = hidden_params[key] # type: ignore @@ -5402,11 +4836,7 @@ class StandardLoggingPayloadSetup: custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name( cold_storage_custom_logger ) - if ( - custom_logger - and hasattr(custom_logger, "s3_path") - and getattr(custom_logger, "s3_path") - ): + if custom_logger and hasattr(custom_logger, "s3_path") and getattr(custom_logger, "s3_path"): s3_path = getattr(custom_logger, "s3_path") except Exception: # If any error occurs in getting the logger instance, use default empty s3_path @@ -5443,9 +4873,7 @@ class StandardLoggingPayloadSetup: response_attr = getattr(original_exception, "response", None) status_code_attr = getattr(response_attr, "status_code", None) error_status = str(status_code_attr) if status_code_attr is not None else "" - error_class: str = ( - str(original_exception.__class__.__name__) if original_exception else "" - ) + error_class: str = str(original_exception.__class__.__name__) if original_exception else "" _llm_provider_in_exception = getattr(original_exception, "llm_provider", "") # Get traceback information (first 100 lines) @@ -5454,9 +4882,7 @@ class StandardLoggingPayloadSetup: tb = getattr(original_exception, "__traceback__", None) if tb: tb_lines = traceback.format_tb(tb) - traceback_info += "".join( - tb_lines[:MAXIMUM_TRACEBACK_LINES_TO_LOG] - ) # Limit to first 100 lines + traceback_info += "".join(tb_lines[:MAXIMUM_TRACEBACK_LINES_TO_LOG]) # Limit to first 100 lines # Prefer the `.message` attribute (set by ProxyException and every # litellm.exceptions.* class) over str(exc); ProxyException does not @@ -5472,12 +4898,8 @@ class StandardLoggingPayloadSetup: else: error_message = str(original_exception) if original_exception else "" - rate_limit_category = validate_rate_limit_category( - getattr(original_exception, "category", None) - ) - rate_limit_type = validate_rate_limit_type( - getattr(original_exception, "rate_limit_type", None) - ) + rate_limit_category = validate_rate_limit_category(getattr(original_exception, "category", None)) + rate_limit_type = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None)) return StandardLoggingPayloadErrorInformation( error_code=error_status, @@ -5602,9 +5024,7 @@ class StandardLoggingPayloadSetup: """ Extract additional header tags for spend tracking based on config. """ - extra_headers: List[str] = ( - getattr(litellm, "extra_spend_tag_headers", None) or [] - ) + extra_headers: List[str] = getattr(litellm, "extra_spend_tag_headers", None) or [] if not extra_headers: return None @@ -5621,9 +5041,7 @@ class StandardLoggingPayloadSetup: return header_tags if header_tags else None @staticmethod - def _get_request_tags( - litellm_params: dict, proxy_server_request: dict - ) -> List[str]: + def _get_request_tags(litellm_params: dict, proxy_server_request: dict) -> List[str]: # check for 'tags' in both 'metadata' and 'litellm_metadata' metadata = litellm_params.get("metadata") or {} litellm_metadata = litellm_params.get("litellm_metadata") or {} @@ -5633,12 +5051,8 @@ class StandardLoggingPayloadSetup: request_tags = litellm_metadata.get("tags", []).copy() else: request_tags = [] - user_agent_tags = StandardLoggingPayloadSetup._get_user_agent_tags( - proxy_server_request - ) - additional_header_tags = StandardLoggingPayloadSetup._get_extra_header_tags( - proxy_server_request - ) + user_agent_tags = StandardLoggingPayloadSetup._get_user_agent_tags(proxy_server_request) + additional_header_tags = StandardLoggingPayloadSetup._get_extra_header_tags(proxy_server_request) if user_agent_tags is not None: request_tags.extend(user_agent_tags) if additional_header_tags is not None: @@ -5687,9 +5101,7 @@ def _get_status_fields( guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") break - return StandardLoggingPayloadStatusFields( - llm_api_status=llm_api_status, guardrail_status=guardrail_status - ) + return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status) def _extract_response_obj_and_hidden_params( @@ -5713,9 +5125,7 @@ def _extract_response_obj_and_hidden_params( if response_headers is not None: hidden_params = dict( StandardLoggingHiddenParams( - additional_headers=StandardLoggingPayloadSetup.get_additional_headers( - dict(response_headers) - ), + additional_headers=StandardLoggingPayloadSetup.get_additional_headers(dict(response_headers)), model_id=None, cache_key=None, api_base=None, @@ -5744,18 +5154,14 @@ def get_standard_logging_object_payload( try: kwargs = kwargs or {} - response_obj, hidden_params = _extract_response_obj_and_hidden_params( - init_response_obj, original_exception - ) + response_obj, hidden_params = _extract_response_obj_and_hidden_params(init_response_obj, original_exception) # standardize this function to be used across, s3, dynamoDB, langfuse logging litellm_params = kwargs.get("litellm_params", {}) or {} proxy_server_request = litellm_params.get("proxy_server_request") or {} # Merge both litellm_metadata and metadata to get complete metadata - metadata: dict = StandardLoggingPayloadSetup.merge_litellm_metadata( - litellm_params - ) + metadata: dict = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) completion_start_time = kwargs.get("completion_start_time", end_time) call_type = kwargs.get("call_type") @@ -5763,9 +5169,7 @@ def get_standard_logging_object_payload( # Extract usage as a plain dict, avoiding Pydantic round-trip usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict( response_obj=response_obj, - combined_usage_object=cast( - Optional[Usage], kwargs.get("combined_usage_object") - ), + combined_usage_object=cast(Optional[Usage], kwargs.get("combined_usage_object")), ) id = response_obj.get("id", kwargs.get("litellm_call_id")) @@ -5800,9 +5204,7 @@ def get_standard_logging_object_payload( prompt_integration=kwargs.get("prompt_integration", None), applied_guardrails=kwargs.get("applied_guardrails", None), mcp_tool_call_metadata=kwargs.get("mcp_tool_call_metadata", None), - vector_store_request_metadata=kwargs.get( - "vector_store_request_metadata", None - ), + vector_store_request_metadata=kwargs.get("vector_store_request_metadata", None), usage_object=usage_dict, proxy_server_request=proxy_server_request, start_time=start_time, @@ -5831,13 +5233,8 @@ def get_standard_logging_object_payload( response_cost: float = raw_response_cost or 0.0 # clean up litellm hidden params - clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params( - hidden_params - ) - if ( - clean_hidden_params["response_cost"] is None - and raw_response_cost is not None - ): + clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params(hidden_params) + if clean_hidden_params["response_cost"] is None and raw_response_cost is not None: clean_hidden_params["response_cost"] = response_cost model_cost_information = StandardLoggingPayloadSetup.get_model_cost_information( @@ -5848,12 +5245,10 @@ def get_standard_logging_object_payload( api_base=litellm_params.get("api_base"), ) - error_information, error_str = ( - StandardLoggingPayloadSetup.get_error_information_for_logging_payload( - metadata=metadata, - original_exception=original_exception, - error_str=error_str, - ) + error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( + metadata=metadata, + original_exception=original_exception, + error_str=error_str, ) ## get final response object ## @@ -5874,9 +5269,7 @@ def get_standard_logging_object_payload( # This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0" # are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) - model_name = reconstruct_model_name( - kwargs.get("model", "") or "", custom_llm_provider, metadata - ) + model_name = reconstruct_model_name(kwargs.get("model", "") or "", custom_llm_provider, metadata) response_model_name: Optional[str] = None if isinstance(final_response_obj, dict): response_model_name = final_response_obj.get("model") @@ -5886,10 +5279,7 @@ def get_standard_logging_object_payload( requested_model = kwargs.get("model") if ( isinstance(requested_model, str) - and ( - "model_router" in requested_model.lower() - or "model-router" in requested_model.lower() - ) + and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower()) and isinstance(response_model_name, str) and response_model_name ): @@ -5897,8 +5287,7 @@ def get_standard_logging_object_payload( payload: StandardLoggingPayload = StandardLoggingPayload( id=str(id), - litellm_call_id=kwargs.get("litellm_call_id") - or litellm_params.get("litellm_call_id"), + litellm_call_id=kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( logging_obj=logging_obj, litellm_params=litellm_params, @@ -5909,9 +5298,7 @@ def get_standard_logging_object_payload( status=status, status_fields=_get_status_fields( status=status, - guardrail_information=metadata.get( - "standard_logging_guardrail_information", None - ), + guardrail_information=metadata.get("standard_logging_guardrail_information", None), error_str=error_str, ), custom_llm_provider=custom_llm_provider, @@ -5930,10 +5317,7 @@ def get_standard_logging_object_payload( completion_tokens=usage_dict.get("completion_tokens", 0), request_tags=request_tags, end_user=end_user_id or "", - api_base=StandardLoggingPayloadSetup.strip_trailing_slash( - litellm_params.get("api_base", "") - ) - or "", + api_base=StandardLoggingPayloadSetup.strip_trailing_slash(litellm_params.get("api_base", "")) or "", model_group=_model_group, model_id=_model_id, requester_ip_address=clean_metadata.get("requester_ip_address", None), @@ -5951,21 +5335,15 @@ def get_standard_logging_object_payload( model_map_information=model_cost_information, error_str=error_str, error_information=error_information, - response_cost_failure_debug_info=kwargs.get( - "response_cost_failure_debug_information" - ), - guardrail_information=metadata.get( - "standard_logging_guardrail_information", None - ), + response_cost_failure_debug_info=kwargs.get("response_cost_failure_debug_information"), + guardrail_information=metadata.get("standard_logging_guardrail_information", None), standard_built_in_tools_params=standard_built_in_tools_params, ) # emit_standard_logging_payload(payload) - Moved to success_handler to prevent double emitting return payload except Exception as e: - verbose_logger.exception( - "Error creating standard logging object - {}".format(str(e)) - ) + verbose_logger.exception("Error creating standard logging object - {}".format(str(e))) return None @@ -6030,9 +5408,7 @@ def get_standard_logging_metadata( if metadata.get("user_api_key") is not None: if is_valid_sha256_hash(str(metadata.get("user_api_key"))): - clean_metadata["user_api_key_hash"] = metadata.get( - "user_api_key" - ) # this is the hash + clean_metadata["user_api_key_hash"] = metadata.get("user_api_key") # this is the hash return clean_metadata @@ -6053,14 +5429,10 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): ## check user_api_key_metadata for sensitive logging keys cleaned_user_api_key_metadata = {} - if "user_api_key_metadata" in metadata and isinstance( - metadata["user_api_key_metadata"], dict - ): + if "user_api_key_metadata" in metadata and isinstance(metadata["user_api_key_metadata"], dict): for k, v in metadata["user_api_key_metadata"].items(): if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[k] = ( - "scrubbed_by_litellm_for_sensitive_keys" - ) + cleaned_user_api_key_metadata[k] = "scrubbed_by_litellm_for_sensitive_keys" else: cleaned_user_api_key_metadata[k] = v @@ -6094,9 +5466,7 @@ from typing import Any, Dict, List, Optional, Union def create_dummy_standard_logging_payload() -> StandardLoggingPayload: # First create the nested objects with proper typing - model_info = StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ) + model_info = StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None) metadata = StandardLoggingMetadata( # type: ignore user_api_key_hash=str("test_hash"), @@ -6132,9 +5502,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: # Create messages and response with proper typing messages: List[Dict[str, str]] = [{"role": "user", "content": "Hello, world!"}] - response: Dict[str, List[Dict[str, Dict[str, str]]]] = { - "choices": [{"message": {"content": "Hi there!"}}] - } + response: Dict[str, List[Dict[str, Dict[str, str]]]] = {"choices": [{"message": {"content": "Hi there!"}}]} # Main payload initialization return StandardLoggingPayload( # type: ignore @@ -6144,10 +5512,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: response_cost=response_cost, response_cost_failure_debug_info=None, status=str("success"), - total_tokens=int( - DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT - + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT - ), + total_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT), prompt_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT), completion_tokens=int(DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT), startTime=start_time, diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 38e563e9e5c..221b1ae6eab 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -63,9 +63,7 @@ class StandardBuiltInToolCostTracking: ) # Handle file search - if StandardBuiltInToolCostTracking.response_object_includes_file_search_call( - response_object=response_object - ): + if StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response_object): return StandardBuiltInToolCostTracking._handle_file_search_cost( model=model, custom_llm_provider=custom_llm_provider, @@ -99,26 +97,18 @@ class StandardBuiltInToolCostTracking: # cost is routed and priced with the model_info that was actually resolved, instead of # feeding a re-resolved model into the original provider's calculator. if model_info is None and "/" in model: - model_info = StandardBuiltInToolCostTracking._safe_get_model_info( - model=model - ) + model_info = StandardBuiltInToolCostTracking._safe_get_model_info(model=model) if model_info is not None: custom_llm_provider = model_info["litellm_provider"] if custom_llm_provider is None and model_info is not None: custom_llm_provider = model_info["litellm_provider"] - resolved_usage = ( - StandardBuiltInToolCostTracking._usage_with_anthropic_web_search( - usage=usage, response_object=response_object - ) + resolved_usage = StandardBuiltInToolCostTracking._usage_with_anthropic_web_search( + usage=usage, response_object=response_object ) - if ( - model_info is not None - and resolved_usage is not None - and custom_llm_provider is not None - ): + if model_info is not None and resolved_usage is not None and custom_llm_provider is not None: result = get_cost_for_web_search_request( custom_llm_provider=custom_llm_provider, usage=resolved_usage, @@ -128,9 +118,7 @@ class StandardBuiltInToolCostTracking: return result return StandardBuiltInToolCostTracking.get_cost_for_web_search( - web_search_options=standard_built_in_tools_params.get( - "web_search_options", None - ), + web_search_options=standard_built_in_tools_params.get("web_search_options", None), model_info=model_info, ) @@ -145,15 +133,11 @@ class StandardBuiltInToolCostTracking: model=model, custom_llm_provider=custom_llm_provider ) file_search_raw: Any = standard_built_in_tools_params.get("file_search", {}) - file_search_usage: Optional[FileSearchTool] = ( - FileSearchTool(**file_search_raw) if file_search_raw else None - ) + file_search_usage: Optional[FileSearchTool] = FileSearchTool(**file_search_raw) if file_search_raw else None # Convert model_info to dict and extract usage parameters model_info_dict = dict(model_info) if model_info is not None else None - storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params( - file_search_usage - ) + storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params(file_search_usage) return StandardBuiltInToolCostTracking.get_cost_for_file_search( file_search=file_search_usage, @@ -223,16 +207,12 @@ class StandardBuiltInToolCostTracking: standard_built_in_tools_params: StandardBuiltInToolsParams, ) -> float: """Calculate vector store cost.""" - vector_store_usage = standard_built_in_tools_params.get( - "vector_store_usage", None - ) + vector_store_usage = standard_built_in_tools_params.get("vector_store_usage", None) if not vector_store_usage: return 0.0 model_info_dict = dict(model_info) if model_info is not None else None - vector_store_dict = ( - vector_store_usage if isinstance(vector_store_usage, dict) else {} - ) + vector_store_dict = vector_store_usage if isinstance(vector_store_usage, dict) else {} return StandardBuiltInToolCostTracking.get_cost_for_vector_store( vector_store_usage=vector_store_dict, @@ -247,9 +227,7 @@ class StandardBuiltInToolCostTracking: standard_built_in_tools_params: StandardBuiltInToolsParams, ) -> float: """Calculate computer use cost.""" - computer_use_usage = standard_built_in_tools_params.get( - "computer_use_usage", {} - ) + computer_use_usage = standard_built_in_tools_params.get("computer_use_usage", {}) if not computer_use_usage: return 0.0 @@ -273,16 +251,12 @@ class StandardBuiltInToolCostTracking: standard_built_in_tools_params: StandardBuiltInToolsParams, ) -> float: """Calculate code interpreter cost.""" - code_interpreter_sessions = standard_built_in_tools_params.get( - "code_interpreter_sessions", None - ) + code_interpreter_sessions = standard_built_in_tools_params.get("code_interpreter_sessions", None) if not code_interpreter_sessions: return 0.0 model_info_dict = dict(model_info) if model_info is not None else None - sessions = StandardBuiltInToolCostTracking._safe_convert_to_int( - code_interpreter_sessions - ) + sessions = StandardBuiltInToolCostTracking._safe_convert_to_int(code_interpreter_sessions) return StandardBuiltInToolCostTracking.get_cost_for_code_interpreter( sessions=sessions, @@ -302,12 +276,8 @@ class StandardBuiltInToolCostTracking: input_tokens_val = computer_use_usage.get("input_tokens") output_tokens_val = computer_use_usage.get("output_tokens") - input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int( - input_tokens_val - ) - output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int( - output_tokens_val - ) + input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(input_tokens_val) + output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(output_tokens_val) return input_tokens, output_tokens @@ -322,9 +292,7 @@ class StandardBuiltInToolCostTracking: return None @staticmethod - def _usage_with_anthropic_web_search( - usage: Usage | None, response_object: object - ) -> Usage | None: + def _usage_with_anthropic_web_search(usage: Usage | None, response_object: object) -> Usage | None: """Return a Usage carrying server_tool_use.web_search_requests sourced from a raw Anthropic /v1/messages response dict when the reconstructed Usage dropped it (or was never supplied). The original Usage is returned unchanged when it @@ -333,14 +301,9 @@ class StandardBuiltInToolCostTracking: get_anthropic_web_search_requests_from_response, ) - if usage is not None and ( - _get_web_search_requests(getattr(usage, "server_tool_use", None)) - is not None - ): + if usage is not None and (_get_web_search_requests(getattr(usage, "server_tool_use", None)) is not None): return usage - web_search_requests = get_anthropic_web_search_requests_from_response( - response_object - ) + web_search_requests = get_anthropic_web_search_requests_from_response(response_object) if web_search_requests is None: return usage server_tool_use = ServerToolUse(web_search_requests=web_search_requests) @@ -349,9 +312,7 @@ class StandardBuiltInToolCostTracking: return usage.model_copy(update={"server_tool_use": server_tool_use}) @staticmethod - def response_object_includes_web_search_call( - response_object: Any, usage: Optional[Usage] = None - ) -> bool: + def response_object_includes_web_search_call(response_object: Any, usage: Optional[Usage] = None) -> bool: """ Check if the response object includes a web search call. @@ -370,10 +331,8 @@ class StandardBuiltInToolCostTracking: if isinstance(response_object, ModelResponse): # chat completions only include url_citation annotations when a web search call is made - has_url_citations = ( - StandardBuiltInToolCostTracking.response_includes_annotation_type( - response_object=response_object, annotation_type="url_citation" - ) + has_url_citations = StandardBuiltInToolCostTracking.response_includes_annotation_type( + response_object=response_object, annotation_type="url_citation" ) if has_url_citations: return True @@ -382,9 +341,7 @@ class StandardBuiltInToolCostTracking: if ( hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None - and isinstance( - usage.prompt_tokens_details, PromptTokensDetailsWrapper - ) + and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper) and hasattr(usage.prompt_tokens_details, "web_search_requests") and usage.prompt_tokens_details.web_search_requests is not None ): @@ -392,10 +349,7 @@ class StandardBuiltInToolCostTracking: # Anthropic Claude (direct API and Vertex AI) uses server_tool_use.web_search_requests. # Without this check, Claude ModelResponse always falls through to return False # and _handle_web_search_cost() is never called. - if ( - hasattr(usage, "server_tool_use") - and _get_web_search_requests(usage.server_tool_use) is not None - ): + if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None: return True return False elif isinstance(response_object, ResponsesAPIResponse): @@ -404,10 +358,7 @@ class StandardBuiltInToolCostTracking: response_object=response_object, output_type="web_search_call" ) elif usage is not None: - if ( - hasattr(usage, "server_tool_use") - and _get_web_search_requests(usage.server_tool_use) is not None - ): + if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None: return True elif ( hasattr(usage, "prompt_tokens_details") @@ -485,13 +436,9 @@ class StandardBuiltInToolCostTracking: return False @staticmethod - def _safe_get_model_info( - model: str, custom_llm_provider: Optional[str] = None - ) -> Optional[ModelInfo]: + def _safe_get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Optional[ModelInfo]: try: - return litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: return None @@ -509,9 +456,7 @@ class StandardBuiltInToolCostTracking: search_context_raw: Any = model_info.get("search_context_cost_per_query", {}) search_context_pricing: SearchContextCostPerQuery = ( - SearchContextCostPerQuery(**search_context_raw) - if search_context_raw - else SearchContextCostPerQuery() + SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery() ) if web_search_options.get("search_context_size", None) == "low": return search_context_pricing.get("search_context_size_low", 0.0) @@ -519,9 +464,7 @@ class StandardBuiltInToolCostTracking: return search_context_pricing.get("search_context_size_medium", 0.0) elif web_search_options.get("search_context_size", None) == "high": return search_context_pricing.get("search_context_size_high", 0.0) - return StandardBuiltInToolCostTracking.get_default_cost_for_web_search( - model_info - ) + return StandardBuiltInToolCostTracking.get_default_cost_for_web_search(model_info) @staticmethod def get_default_cost_for_web_search( @@ -534,13 +477,9 @@ class StandardBuiltInToolCostTracking: """ if model_info is None: return 0.0 - search_context_raw: Any = ( - model_info.get("search_context_cost_per_query", {}) or {} - ) + search_context_raw: Any = model_info.get("search_context_cost_per_query", {}) or {} search_context_pricing: SearchContextCostPerQuery = ( - SearchContextCostPerQuery(**search_context_raw) - if search_context_raw - else SearchContextCostPerQuery() + SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery() ) return search_context_pricing.get("search_context_size_medium", 0.0) @@ -562,11 +501,7 @@ class StandardBuiltInToolCostTracking: return 0.0 # Check if model-specific pricing is available - if ( - model_info - and "file_search_cost_per_gb_per_day" in model_info - and provider == "azure" - ): + if model_info and "file_search_cost_per_gb_per_day" in model_info and provider == "azure": if storage_gb and days: return storage_gb * days * model_info["file_search_cost_per_gb_per_day"] elif model_info and "file_search_cost_per_1k_calls" in model_info: @@ -629,12 +564,8 @@ class StandardBuiltInToolCostTracking: if provider == "azure" and (input_tokens or output_tokens): # Check if model-specific pricing is available if model_info: - input_cost = model_info.get( - "computer_use_input_cost_per_1k_tokens", 0.0 - ) - output_cost = model_info.get( - "computer_use_output_cost_per_1k_tokens", 0.0 - ) + input_cost = model_info.get("computer_use_input_cost_per_1k_tokens", 0.0) + output_cost = model_info.get("computer_use_output_cost_per_1k_tokens", 0.0) if input_cost or output_cost: total_cost = 0.0 if input_tokens: @@ -651,13 +582,9 @@ class StandardBuiltInToolCostTracking: total_cost = 0.0 if input_tokens: - total_cost += ( - input_tokens / 1000.0 - ) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS + total_cost += (input_tokens / 1000.0) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS if output_tokens: - total_cost += ( - output_tokens / 1000.0 - ) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS + total_cost += (output_tokens / 1000.0) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS return total_cost # OpenAI doesn't charge separately for computer use yet @@ -674,19 +601,11 @@ class StandardBuiltInToolCostTracking: try: container_model = f"{provider}/container" - model_info = litellm.get_model_info( - model=container_model, custom_llm_provider=provider - ) - model_key = ( - model_info.get("key") - if isinstance(model_info, dict) - else getattr(model_info, "key", None) - ) + model_info = litellm.get_model_info(model=container_model, custom_llm_provider=provider) + model_key = model_info.get("key") if isinstance(model_info, dict) else getattr(model_info, "key", None) if model_key and model_key in litellm.model_cost: - return litellm.model_cost[model_key].get( - "code_interpreter_cost_per_session" - ) + return litellm.model_cost[model_key].get("code_interpreter_cost_per_session") except Exception: pass @@ -744,9 +663,7 @@ class StandardBuiltInToolCostTracking: tools = StandardBuiltInToolCostTracking._get_tools_from_kwargs( kwargs=kwargs, tool_type="web_search_preview" - ) or StandardBuiltInToolCostTracking._get_tools_from_kwargs( - kwargs=kwargs, tool_type="web_search" - ) + ) or StandardBuiltInToolCostTracking._get_tools_from_kwargs(kwargs=kwargs, tool_type="web_search") if tools: # Look for web search tool in the tools array for tool in tools: @@ -763,9 +680,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def _get_file_search_tool_call(kwargs: Dict) -> Optional[FileSearchTool]: - tools = StandardBuiltInToolCostTracking._get_tools_from_kwargs( - kwargs, "file_search" - ) + tools = StandardBuiltInToolCostTracking._get_tools_from_kwargs(kwargs, "file_search") if tools: for tool in tools: if isinstance(tool, dict): diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index 1432e912fd8..1c6adbec174 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -19,9 +19,7 @@ class TranscriptionUsageObjectTransformation: @staticmethod def transform_transcription_usage_object( - usage_object: Union[ - TranscriptionUsageDurationObject, TranscriptionUsageTokensObject - ], + usage_object: Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject], ) -> Optional[Usage]: if isinstance(usage_object, TranscriptionUsageDurationObject): return None diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e407dd70bf0..e013c587f0f 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -126,17 +126,12 @@ def _generic_cost_per_character( Exception if 'input_cost_per_character' or 'output_cost_per_character' is missing from model_info """ ## GET MODEL INFO - model_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) ## CALCULATE INPUT COST try: if custom_prompt_cost is None: - assert ( - "input_cost_per_character" in model_info - and model_info["input_cost_per_character"] is not None - ), ( + assert "input_cost_per_character" in model_info and model_info["input_cost_per_character"] is not None, ( "model info for model={} does not have 'input_cost_per_character'-pricing\nmodel_info={}".format( model, model_info ) @@ -156,10 +151,7 @@ def _generic_cost_per_character( ## CALCULATE OUTPUT COST try: if custom_completion_cost is None: - assert ( - "output_cost_per_character" in model_info - and model_info["output_cost_per_character"] is not None - ), ( + assert "output_cost_per_character" in model_info and model_info["output_cost_per_character"] is not None, ( "model info for model={} does not have 'output_cost_per_character'-pricing\nmodel_info={}".format( model, model_info ) @@ -220,12 +212,8 @@ def _get_token_base_cost( # Get service tier aware cost keys input_cost_key = _get_service_tier_cost_key("input_cost_per_token", service_tier) output_cost_key = _get_service_tier_cost_key("output_cost_per_token", service_tier) - cache_creation_cost_key = _get_service_tier_cost_key( - "cache_creation_input_token_cost", service_tier - ) - cache_read_cost_key = _get_service_tier_cost_key( - "cache_read_input_token_cost", service_tier - ) + cache_creation_cost_key = _get_service_tier_cost_key("cache_creation_input_token_cost", service_tier) + cache_read_cost_key = _get_service_tier_cost_key("cache_read_input_token_cost", service_tier) prompt_base_cost = cast(float, _get_cost_per_unit(model_info, input_cost_key)) completion_base_cost = cast(float, _get_cost_per_unit(model_info, output_cost_key)) @@ -233,14 +221,10 @@ def _get_token_base_cost( # For image generation models that don't have output_cost_per_token, # use output_cost_per_image_token as the base cost (all output tokens are image tokens) if completion_base_cost == 0.0 or completion_base_cost is None: - output_image_cost = _get_cost_per_unit( - model_info, "output_cost_per_image_token", None - ) + output_image_cost = _get_cost_per_unit(model_info, "output_cost_per_image_token", None) if output_image_cost is not None: completion_base_cost = cast(float, output_image_cost) - cache_creation_cost = cast( - float, _get_cost_per_unit(model_info, cache_creation_cost_key) - ) + cache_creation_cost = cast(float, _get_cost_per_unit(model_info, cache_creation_cost_key)) cache_creation_cost_above_1hr = cast( float, _get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"), @@ -254,10 +238,7 @@ def _get_token_base_cost( # so that the threshold detection loop only processes standard keys. The # service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key. threshold_keys = [ - k - for k in model_info - if k.startswith("input_cost_per_token_above_") - and not k.endswith(_SERVICE_TIER_SUFFIXES) + k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES) ] if not threshold_keys: return ( @@ -292,9 +273,7 @@ def _get_token_base_cost( ) prompt_base_cost = cast( float, - _get_cost_per_unit( - model_info, tiered_input_key, prompt_base_cost - ), + _get_cost_per_unit(model_info, tiered_input_key, prompt_base_cost), ) tiered_output_key = ( _get_service_tier_cost_key( @@ -359,9 +338,7 @@ def _get_token_base_cost( cache_read_cost = cast( float, - _get_cost_per_unit( - model_info, cache_read_tiered_key, cache_read_cost - ), + _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost), ) break @@ -379,9 +356,7 @@ def _get_token_base_cost( ) -def calculate_cost_component( - model_info: ModelInfo, cost_key: str, usage_value: Optional[float] -) -> float: +def calculate_cost_component(model_info: ModelInfo, cost_key: str, usage_value: Optional[float]) -> float: """ Generic cost calculator for any usage component @@ -394,19 +369,12 @@ def calculate_cost_component( float: The calculated cost """ cost_per_unit = _get_cost_per_unit(model_info, cost_key) - if ( - cost_per_unit is not None - and isinstance(cost_per_unit, float) - and usage_value is not None - and usage_value > 0 - ): + if cost_per_unit is not None and isinstance(cost_per_unit, float) and usage_value is not None and usage_value > 0: return float(usage_value) * cost_per_unit return 0.0 -def _get_cost_per_unit( - model_info: ModelInfo, cost_key: str, default_value: Optional[float] = 0.0 -) -> Optional[float]: +def _get_cost_per_unit(model_info: ModelInfo, cost_key: str, default_value: Optional[float] = 0.0) -> Optional[float]: # Sometimes the cost per unit is a string (e.g.: If a value like "3e-7" was read from the config.yaml) cost_per_unit = model_info.get(cost_key) if isinstance(cost_per_unit, float): @@ -457,22 +425,12 @@ def calculate_cache_writing_cost( total_cost: float = 0.0 if cache_creation_token_details is not None: # get the number of 5m and 1h cache creation tokens - cache_creation_tokens_5m = ( - cache_creation_token_details.ephemeral_5m_input_tokens - ) - cache_creation_tokens_1h = ( - cache_creation_token_details.ephemeral_1h_input_tokens - ) + cache_creation_tokens_5m = cache_creation_token_details.ephemeral_5m_input_tokens + cache_creation_tokens_1h = cache_creation_token_details.ephemeral_1h_input_tokens # add the number of 5m and 1h cache creation tokens to the cache creation tokens + total_cost += cache_creation_tokens_5m * cache_creation_cost if cache_creation_tokens_5m is not None else 0.0 total_cost += ( - cache_creation_tokens_5m * cache_creation_cost - if cache_creation_tokens_5m is not None - else 0.0 - ) - total_cost += ( - cache_creation_tokens_1h * cache_creation_cost_above_1hr - if cache_creation_tokens_1h is not None - else 0.0 + cache_creation_tokens_1h * cache_creation_cost_above_1hr if cache_creation_tokens_1h is not None else 0.0 ) else: total_cost += cache_creation_tokens * cache_creation_cost @@ -493,10 +451,7 @@ class PromptTokensDetailsResult(TypedDict): def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: - cache_hit_tokens = ( - cast(Optional[int], getattr(usage.prompt_tokens_details, "cached_tokens", 0)) - or 0 - ) + cache_hit_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "cached_tokens", 0)) or 0 cache_creation_tokens = ( cast( Optional[int], @@ -515,14 +470,8 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: cast(Optional[int], getattr(usage.prompt_tokens_details, "text_tokens", None)) or 0 # default to prompt tokens, if this field is not set ) - audio_tokens = ( - cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) - or 0 - ) - image_tokens = ( - cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) - or 0 - ) + audio_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0 + image_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0 character_count = ( cast( Optional[int], @@ -530,9 +479,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0 ) - image_count = ( - cast(Optional[int], getattr(usage.prompt_tokens_details, "image_count", 0)) or 0 - ) + image_count = cast(Optional[int], getattr(usage.prompt_tokens_details, "image_count", 0)) or 0 video_length_seconds = ( cast( Optional[float], @@ -626,12 +573,8 @@ def _calculate_input_cost( ### AUDIO COST if prompt_tokens_details["audio_tokens"]: - audio_cost_key = _get_service_tier_cost_key( - "input_cost_per_audio_token", service_tier - ) - prompt_cost += calculate_cost_component( - model_info, audio_cost_key, prompt_tokens_details["audio_tokens"] - ) + audio_cost_key = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier) + prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"]) ### IMAGE TOKEN COST if prompt_tokens_details["image_tokens"]: @@ -640,9 +583,7 @@ def _calculate_input_cost( image_token_cost_key = "input_cost_per_image_token" if model_info.get(image_token_cost_key) is None: image_token_cost_key = "input_cost_per_token" - prompt_cost += calculate_cost_component( - model_info, image_token_cost_key, prompt_tokens_details["image_tokens"] - ) + prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]) ### CACHE WRITING COST - Now uses tiered pricing if ( @@ -651,9 +592,7 @@ def _calculate_input_cost( ): prompt_cost += calculate_cache_writing_cost( cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"], - cache_creation_token_details=prompt_tokens_details[ - "cache_creation_token_details" - ], + cache_creation_token_details=prompt_tokens_details["cache_creation_token_details"], cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost, ) @@ -691,9 +630,7 @@ def _calculate_input_cost( return prompt_cost -def _get_regional_uplift_multiplier( - model_info: ModelInfo, data_residency: Optional[str] -) -> float: +def _get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: Optional[str]) -> float: """ Resolve the per-model regional-processing uplift multiplier for a given data-residency region. @@ -718,8 +655,7 @@ def _get_regional_uplift_multiplier( return float(cast(float, multiplier)) except (TypeError, ValueError): verbose_logger.exception( - "Invalid regional_processing_uplift_multiplier_%s for model; " - "defaulting to 1.0", + "Invalid regional_processing_uplift_multiplier_%s for model; defaulting to 1.0", residency, ) return 1.0 @@ -781,21 +717,11 @@ def generic_cost_per_token( image_tokens = prompt_tokens_details["image_tokens"] # Check for double-counting: sum of details > prompt_tokens means overlap - total_details = ( - text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens - ) + total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens has_double_counting = cache_hit > 0 and total_details > usage.prompt_tokens - if ( - text_tokens == 0 and prompt_tokens_details["image_count"] == 0 - ) or has_double_counting: - text_tokens = ( - usage.prompt_tokens - - cache_hit - - audio_tokens - - cache_creation - - image_tokens - ) + if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting: + text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens # Clamp to zero: inconsistent streaming usage if text_tokens < 0: text_tokens = 0 @@ -807,9 +733,7 @@ def generic_cost_per_token( cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost, - ) = _get_token_base_cost( - model_info=model_info, usage=usage, service_tier=service_tier - ) + ) = _get_token_base_cost(model_info=model_info, usage=usage, service_tier=service_tier) prompt_cost = _calculate_input_cost( prompt_tokens_details=prompt_tokens_details, @@ -845,10 +769,7 @@ def generic_cost_per_token( # This handles cases like OpenAI's reasoning models where text_tokens isn't provided text_tokens = max( 0, - usage.completion_tokens - - reasoning_tokens - - audio_tokens - - image_tokens, + usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens, ) else: # No breakdown at all, all tokens are text tokens @@ -859,37 +780,25 @@ def generic_cost_per_token( ## AUDIO COST if not is_text_tokens_total and audio_tokens is not None and audio_tokens > 0: - _output_cost_per_audio_token = _get_cost_per_unit( - model_info, "output_cost_per_audio_token", None - ) + _output_cost_per_audio_token = _get_cost_per_unit(model_info, "output_cost_per_audio_token", None) _output_cost_per_audio_token = ( - _output_cost_per_audio_token - if _output_cost_per_audio_token is not None - else completion_base_cost + _output_cost_per_audio_token if _output_cost_per_audio_token is not None else completion_base_cost ) completion_cost += float(audio_tokens) * _output_cost_per_audio_token ## REASONING COST if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0: - _output_cost_per_reasoning_token = _get_cost_per_unit( - model_info, "output_cost_per_reasoning_token", None - ) + _output_cost_per_reasoning_token = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None) _output_cost_per_reasoning_token = ( - _output_cost_per_reasoning_token - if _output_cost_per_reasoning_token is not None - else completion_base_cost + _output_cost_per_reasoning_token if _output_cost_per_reasoning_token is not None else completion_base_cost ) completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token ## IMAGE COST if not is_text_tokens_total and image_tokens and image_tokens > 0: - _output_cost_per_image_token = _get_cost_per_unit( - model_info, "output_cost_per_image_token", None - ) + _output_cost_per_image_token = _get_cost_per_unit(model_info, "output_cost_per_image_token", None) _output_cost_per_image_token = ( - _output_cost_per_image_token - if _output_cost_per_image_token is not None - else completion_base_cost + _output_cost_per_image_token if _output_cost_per_image_token is not None else completion_base_cost ) completion_cost += float(image_tokens) * _output_cost_per_image_token @@ -954,18 +863,10 @@ def calculate_image_response_cost_from_usage( ) else: text_tokens = _get_token_detail_value(output_tokens_details, "text_tokens") or 0 - image_tokens = ( - _get_token_detail_value(output_tokens_details, "image_tokens") or 0 - ) - audio_tokens = ( - _get_token_detail_value(output_tokens_details, "audio_tokens") or 0 - ) - reasoning_tokens = ( - _get_token_detail_value(output_tokens_details, "reasoning_tokens") or 0 - ) - known_output_tokens = ( - text_tokens + image_tokens + audio_tokens + reasoning_tokens - ) + image_tokens = _get_token_detail_value(output_tokens_details, "image_tokens") or 0 + audio_tokens = _get_token_detail_value(output_tokens_details, "audio_tokens") or 0 + reasoning_tokens = _get_token_detail_value(output_tokens_details, "reasoning_tokens") or 0 + known_output_tokens = text_tokens + image_tokens + audio_tokens + reasoning_tokens if completion_tokens > known_output_tokens: text_tokens += completion_tokens - known_output_tokens @@ -1014,11 +915,7 @@ def calculate_image_response_web_search_cost( from litellm.llms import get_cost_for_web_search_request - synthetic_usage = Usage( - prompt_tokens_details=PromptTokensDetailsWrapper( - web_search_requests=web_search_requests - ) - ) + synthetic_usage = Usage(prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=web_search_requests)) return ( get_cost_for_web_search_request( custom_llm_provider=custom_llm_provider, @@ -1094,9 +991,7 @@ class CostCalculatorUtils: image_response=completion_response, optional_params=optional_params, ) - raise TypeError( - "completion_response must be of type ImageResponse for bedrock image cost calculation" - ) + raise TypeError("completion_response must be of type ImageResponse for bedrock image cost calculation") elif custom_llm_provider == litellm.LlmProviders.RECRAFT.value: from litellm.llms.recraft.cost_calculator import ( cost_calculator as recraft_image_cost_calculator, diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index 7be70852978..7f76c7aca76 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -49,16 +49,12 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1): for model in known_models: try: - model_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: continue if model_info.get("mode") != "chat": continue - _cost = (model_info.get("input_cost_per_token") or 0.0) + ( - model_info.get("output_cost_per_token") or 0.0 - ) + _cost = (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0) model_costs.append((model, _cost)) # Sort by cost (ascending) @@ -77,8 +73,6 @@ def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict: if litellm_params is None: return {} - proxy_request_headers = (litellm_params.get("proxy_server_request") or {}).get( - "headers" - ) or {} + proxy_request_headers = (litellm_params.get("proxy_server_request") or {}).get("headers") or {} return proxy_request_headers diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 79c6c665684..58107d9804b 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -49,9 +49,7 @@ from .get_headers import get_response_headers _MESSAGE_FIELDS: frozenset = frozenset(Message.model_fields.keys()) _CHOICES_FIELDS: frozenset = frozenset(Choices.model_fields.keys()) -_MODEL_RESPONSE_FIELDS: frozenset = frozenset(ModelResponse.model_fields.keys()) | { - "usage" -} +_MODEL_RESPONSE_FIELDS: frozenset = frozenset(ModelResponse.model_fields.keys()) | {"usage"} def _normalize_images_for_message( @@ -109,9 +107,7 @@ def convert_tool_call_to_json_mode( convert_tool_call_to_json_mode=convert_tool_call_to_json_mode, ): # to support 'json_schema' logic on older models - json_mode_content_str: Optional[str] = tool_calls[0]["function"].get( - "arguments" - ) + json_mode_content_str: Optional[str] = tool_calls[0]["function"].get("arguments") if json_mode_content_str is not None: message = litellm.Message(content=json_mode_content_str) finish_reason = "stop" @@ -188,8 +184,7 @@ async def convert_to_streaming_response_async( raise APIError( status_code=500, message=( - "LiteLLM: provider returned a response with no 'choices'. " - f"Raw keys: {list(response_object.keys())}" + f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" ), llm_provider="", model="", @@ -221,9 +216,7 @@ async def convert_to_streaming_response_async( logprobs = choice.get("logprobs", None) - choice = StreamingChoices( - finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs - ) + choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs) choice_list.append(choice) model_response_object.choices = choice_list @@ -243,9 +236,7 @@ async def convert_to_streaming_response_async( model_response_object.id = response_object["id"] if "created" in response_object: - model_response_object.created = _safe_convert_created_field( - response_object["created"] - ) + model_response_object.created = _safe_convert_created_field(response_object["created"]) if "system_fingerprint" in response_object: model_response_object.system_fingerprint = response_object["system_fingerprint"] @@ -259,9 +250,7 @@ async def convert_to_streaming_response_async( # single-yield behavior. slices: list[str] = [] if len(model_response_object.choices) == 1: - slices = _split_assembled_content_for_replay( - model_response_object.choices[0].delta.content - ) + slices = _split_assembled_content_for_replay(model_response_object.choices[0].delta.content) if len(slices) <= 1: yield model_response_object await asyncio.sleep(0) @@ -305,8 +294,7 @@ def convert_to_streaming_response( raise APIError( status_code=500, message=( - "LiteLLM: provider returned a response with no 'choices'. " - f"Raw keys: {list(response_object.keys())}" + f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}" ), llm_provider="", model="", @@ -333,23 +321,15 @@ def convert_to_streaming_response( if "usage" in response_object and response_object["usage"] is not None: setattr(model_response_object, "usage", Usage()) - model_response_object.usage.completion_tokens = response_object["usage"].get( - "completion_tokens", 0 - ) # type: ignore - model_response_object.usage.prompt_tokens = response_object["usage"].get( - "prompt_tokens", 0 - ) # type: ignore - model_response_object.usage.total_tokens = response_object["usage"].get( - "total_tokens", 0 - ) # type: ignore + model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) # type: ignore + model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) # type: ignore + model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) # type: ignore if "id" in response_object: model_response_object.id = response_object["id"] if "created" in response_object: - model_response_object.created = _safe_convert_created_field( - response_object["created"] - ) + model_response_object.created = _safe_convert_created_field(response_object["created"]) if "system_fingerprint" in response_object: model_response_object.system_fingerprint = response_object["system_fingerprint"] @@ -362,9 +342,7 @@ def convert_to_streaming_response( # for the full rationale — this mirrors its tail. slices: list[str] = [] if len(model_response_object.choices) == 1: - slices = _split_assembled_content_for_replay( - model_response_object.choices[0].delta.content - ) + slices = _split_assembled_content_for_replay(model_response_object.choices[0].delta.content) if len(slices) <= 1: yield model_response_object return @@ -407,9 +385,7 @@ def _handle_invalid_parallel_tool_calls( current_function = tool_call.function.name function_args = json.loads(tool_call.function.arguments) if current_function == "multi_tool_use.parallel": - verbose_logger.debug( - "OpenAI did a weird pseudo-multi-tool-use call, fixing call structure.." - ) + verbose_logger.debug("OpenAI did a weird pseudo-multi-tool-use call, fixing call structure..") for _fake_i, _fake_tool_use in enumerate(function_args["tool_uses"]): _function_args = _fake_tool_use["parameters"] _current_function = _fake_tool_use["recipient_name"] @@ -419,17 +395,13 @@ def _handle_invalid_parallel_tool_calls( fixed_tc = ChatCompletionMessageToolCall( id=f"{tool_call.id}_{_fake_i}", type="function", - function=Function( - name=_current_function, arguments=json.dumps(_function_args) - ), + function=Function(name=_current_function, arguments=json.dumps(_function_args)), ) replacements[i].append(fixed_tc) shift = 0 for i, replacement in replacements.items(): - tool_calls[:] = ( - tool_calls[: i + shift] + replacement + tool_calls[i + shift + 1 :] - ) + tool_calls[:] = tool_calls[: i + shift] + replacement + tool_calls[i + shift + 1 :] shift += len(replacement) return tool_calls @@ -471,13 +443,9 @@ class LiteLLMResponseObjectHandler: # Convert dicts to wrapper objects so getattr() works in cost calculation if isinstance(usage.get("input_tokens_details"), dict): - usage["prompt_tokens_details"] = PromptTokensDetailsWrapper( - **usage["input_tokens_details"] - ) + usage["prompt_tokens_details"] = PromptTokensDetailsWrapper(**usage["input_tokens_details"]) if isinstance(usage.get("output_tokens_details"), dict): - usage["completion_tokens_details"] = CompletionTokensDetailsWrapper( - **usage["output_tokens_details"] - ) + usage["completion_tokens_details"] = CompletionTokensDetailsWrapper(**usage["output_tokens_details"]) if model_response_object is None: model_response_object = ImageResponse(**response_object) @@ -516,9 +484,11 @@ class LiteLLMResponseObjectHandler: chat_response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hi"}]) text_response = convert_chat_to_text_completion(chat_response) """ - transformed_logprobs = LiteLLMResponseObjectHandler._convert_provider_response_logprobs_to_text_completion_logprobs( - response=response, - custom_llm_provider=custom_llm_provider, + transformed_logprobs = ( + LiteLLMResponseObjectHandler._convert_provider_response_logprobs_to_text_completion_logprobs( + response=response, + custom_llm_provider=custom_llm_provider, + ) ) text_completion_response["id"] = response.get("id", None) @@ -538,9 +508,7 @@ class LiteLLMResponseObjectHandler: text_completion_response["choices"] = choices_list text_completion_response["usage"] = response.get("usage", None) - text_completion_response._hidden_params = HiddenParams( - **response._hidden_params - ) + text_completion_response._hidden_params = HiddenParams(**response._hidden_params) return text_completion_response @staticmethod @@ -559,9 +527,7 @@ class LiteLLMResponseObjectHandler: def _should_convert_tool_call_to_json_mode( - tool_calls: Optional[ - Union[List[ChatCompletionMessageToolCall], List[DatabricksTool]] - ] = None, + tool_calls: Optional[Union[List[ChatCompletionMessageToolCall], List[DatabricksTool]]] = None, convert_tool_call_to_json_mode: Optional[bool] = None, ) -> bool: """ @@ -596,9 +562,7 @@ def convert_to_model_response_object( end_time=None, hidden_params: Optional[dict] = None, _response_headers: Optional[dict] = None, - convert_tool_call_to_json_mode: Optional[ - bool - ] = None, # used for supporting 'json_schema' on older models + convert_tool_call_to_json_mode: Optional[bool] = None, # used for supporting 'json_schema' on older models ): additional_headers = get_response_headers(_response_headers) @@ -621,11 +585,7 @@ def convert_to_model_response_object( ### CHECK IF ERROR IN RESPONSE ### - openrouter returns these in the dictionary # Some OpenAI-compatible providers (e.g., Apertis) return empty error objects # even on success. Only raise if the error contains meaningful data. - if ( - response_object is not None - and "error" in response_object - and response_object["error"] is not None - ): + if response_object is not None and "error" in response_object and response_object["error"] is not None: error_obj = response_object["error"] has_meaningful_error = False @@ -659,8 +619,7 @@ def convert_to_model_response_object( try: if response_type == "completion" and ( - model_response_object is None - or isinstance(model_response_object, ModelResponse) + model_response_object is None or isinstance(model_response_object, ModelResponse) ): if response_object is None or model_response_object is None: raise Exception("Error in response object format") @@ -669,9 +628,7 @@ def convert_to_model_response_object( return convert_to_streaming_response(response_object=response_object) choice_list: List[Choices] = [] - if not response_object.get("choices") or not isinstance( - response_object["choices"], Iterable - ): + if not response_object.get("choices") or not isinstance(response_object["choices"], Iterable): from litellm.exceptions import APIError raise APIError( @@ -692,9 +649,7 @@ def convert_to_model_response_object( for _tc in tool_calls: _openai_tc = ChatCompletionMessageToolCall(**_tc) _openai_tool_calls.append(_openai_tc) - fixed_tool_calls = _handle_invalid_parallel_tool_calls( - _openai_tool_calls - ) + fixed_tool_calls = _handle_invalid_parallel_tool_calls(_openai_tool_calls) if fixed_tool_calls is not None: tool_calls = fixed_tool_calls @@ -706,25 +661,19 @@ def convert_to_model_response_object( convert_tool_call_to_json_mode=convert_tool_call_to_json_mode, ): # to support 'json_schema' logic on older models - json_mode_content_str: Optional[str] = tool_calls[0][ - "function" - ].get("arguments") + json_mode_content_str: Optional[str] = tool_calls[0]["function"].get("arguments") if json_mode_content_str is not None: message = litellm.Message(content=json_mode_content_str) finish_reason = "stop" if message is None: # Preserve provider_specific_fields if already present # in the response (e.g. from proxy passthrough) - provider_specific_fields = dict( - choice["message"].get("provider_specific_fields", None) or {} - ) + provider_specific_fields = dict(choice["message"].get("provider_specific_fields", None) or {}) for f in choice["message"].keys() - _MESSAGE_FIELDS: provider_specific_fields[f] = choice["message"][f] # Handle reasoning models that display `reasoning_content` within `content` - reasoning_content, content = _extract_reasoning_content( - choice["message"] - ) + reasoning_content, content = _extract_reasoning_content(choice["message"]) # Handle thinking models that display `thinking_blocks` within `content` thinking_blocks: Optional[ @@ -749,25 +698,17 @@ def convert_to_model_response_object( reasoning_content=reasoning_content, thinking_blocks=thinking_blocks, annotations=choice["message"].get("annotations", None), - images=_normalize_images_for_message( - choice["message"].get("images", None) - ), + images=_normalize_images_for_message(choice["message"].get("images", None)), ) finish_reason = choice.get("finish_reason", None) if finish_reason is None: # gpt-4 vision can return 'finish_reason' or 'finish_details' finish_reason = choice.get("finish_details") or "stop" - if ( - finish_reason == "stop" - and message.tool_calls - and len(message.tool_calls) > 0 - ): + if finish_reason == "stop" and message.tool_calls and len(message.tool_calls) > 0: finish_reason = "tool_calls" ## PROVIDER SPECIFIC FIELDS ## - provider_specific_fields = { - f: choice[f] for f in choice.keys() - _CHOICES_FIELDS - } + provider_specific_fields = {f: choice[f] for f in choice.keys() - _CHOICES_FIELDS} logprobs = choice.get("logprobs", None) enhancements = choice.get("enhancements", None) @@ -786,35 +727,22 @@ def convert_to_model_response_object( usage_object = litellm.Usage(**response_object["usage"]) setattr(model_response_object, "usage", usage_object) if "created" in response_object: - model_response_object.created = _safe_convert_created_field( - response_object["created"] - ) + model_response_object.created = _safe_convert_created_field(response_object["created"]) if "id" in response_object: # Preserve the auto-generated id from ModelResponse.__init__ # when the provider returns a falsy id (None, "") - model_response_object.id = ( - response_object["id"] or model_response_object.id - ) + model_response_object.id = response_object["id"] or model_response_object.id if "system_fingerprint" in response_object: - model_response_object.system_fingerprint = response_object[ - "system_fingerprint" - ] + model_response_object.system_fingerprint = response_object["system_fingerprint"] if "model" in response_object: if model_response_object.model is None: model_response_object.model = response_object["model"] - elif ( - "/" in model_response_object.model - and response_object["model"] is not None - ): - openai_compatible_provider = model_response_object.model.split("/")[ - 0 - ] - model_response_object.model = ( - openai_compatible_provider + "/" + response_object["model"] - ) + elif "/" in model_response_object.model and response_object["model"] is not None: + openai_compatible_provider = model_response_object.model.split("/")[0] + model_response_object.model = openai_compatible_provider + "/" + response_object["model"] if start_time is not None and end_time is not None: if isinstance(start_time, type(end_time)): @@ -836,8 +764,7 @@ def convert_to_model_response_object( return model_response_object elif response_type == "embedding" and ( - model_response_object is None - or isinstance(model_response_object, EmbeddingResponse) + model_response_object is None or isinstance(model_response_object, EmbeddingResponse) ): if response_object is None: raise Exception("Error in response object format") @@ -854,15 +781,9 @@ def convert_to_model_response_object( model_response_object.data = response_object["data"] if "usage" in response_object and response_object["usage"] is not None: - model_response_object.usage.completion_tokens = response_object[ - "usage" - ].get("completion_tokens", 0) # type: ignore - model_response_object.usage.prompt_tokens = response_object[ - "usage" - ].get("prompt_tokens", 0) # type: ignore - model_response_object.usage.total_tokens = response_object["usage"].get( - "total_tokens", 0 - ) # type: ignore + model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) # type: ignore + model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) # type: ignore + model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) # type: ignore if start_time is not None and end_time is not None: model_response_object._response_ms = ( # type: ignore @@ -877,8 +798,7 @@ def convert_to_model_response_object( return model_response_object elif response_type == "image_generation" and ( - model_response_object is None - or isinstance(model_response_object, ImageResponse) + model_response_object is None or isinstance(model_response_object, ImageResponse) ): if response_object is None: raise Exception("Error in response object format") @@ -890,8 +810,7 @@ def convert_to_model_response_object( ) elif response_type == "audio_transcription" and ( - model_response_object is None - or isinstance(model_response_object, TranscriptionResponse) + model_response_object is None or isinstance(model_response_object, TranscriptionResponse) ): if response_object is None: raise Exception("Error in response object format") @@ -908,20 +827,14 @@ def convert_to_model_response_object( setattr(model_response_object, key, response_object[key]) if "usage" in response_object and response_object["usage"] is not None: - tr_usage_object: Optional[ - Union[ - TranscriptionUsageDurationObject, TranscriptionUsageTokensObject - ] - ] = None + tr_usage_object: Optional[Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject]] = ( + None + ) if response_object["usage"].get("type", None) == "duration": - tr_usage_object = TranscriptionUsageDurationObject( - **response_object["usage"] - ) + tr_usage_object = TranscriptionUsageDurationObject(**response_object["usage"]) elif response_object["usage"].get("type", None) == "tokens": - tr_usage_object = TranscriptionUsageTokensObject( - **response_object["usage"] - ) + tr_usage_object = TranscriptionUsageTokensObject(**response_object["usage"]) if tr_usage_object is not None: setattr(model_response_object, "usage", tr_usage_object) @@ -932,17 +845,16 @@ def convert_to_model_response_object( # tracking without exposing it in the response body. Must be set # after hidden_params assignment to avoid being overwritten. if "_audio_transcription_duration" in response_object: - model_response_object._hidden_params["audio_transcription_duration"] = ( - response_object["_audio_transcription_duration"] - ) + model_response_object._hidden_params["audio_transcription_duration"] = response_object[ + "_audio_transcription_duration" + ] if _response_headers is not None: model_response_object._response_headers = _response_headers return model_response_object elif response_type == "rerank" and ( - model_response_object is None - or isinstance(model_response_object, RerankResponse) + model_response_object is None or isinstance(model_response_object, RerankResponse) ): if response_object is None: raise Exception("Error in response object format") @@ -976,6 +888,4 @@ def convert_to_model_response_object( end_time=end_time, convert_tool_call_to_json_mode=convert_tool_call_to_json_mode, ) - raise Exception( - f"Invalid response object {traceback.format_exc()}\n\nreceived_args={received_args}" - ) + raise Exception(f"Invalid response object {traceback.format_exc()}\n\nreceived_args={received_args}") diff --git a/litellm/litellm_core_utils/llm_response_utils/get_api_base.py b/litellm/litellm_core_utils/llm_response_utils/get_api_base.py index c23bbb936b9..cc61ef0c899 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_api_base.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_api_base.py @@ -7,9 +7,7 @@ from ...litellm_core_utils.get_llm_provider_logic import get_llm_provider from ...types.router import LiteLLM_Params -def get_api_base( - model: str, optional_params: Union[dict, LiteLLM_Params] -) -> Optional[str]: +def get_api_base(model: str, optional_params: Union[dict, LiteLLM_Params]) -> Optional[str]: """ Returns the api base used for calling the model. @@ -34,9 +32,7 @@ def get_api_base( elif "model" in optional_params: _optional_params = LiteLLM_Params(**optional_params) else: # prevent needing to copy and pop the dict - _optional_params = LiteLLM_Params( - model=model, **optional_params - ) # convert to pydantic object + _optional_params = LiteLLM_Params(model=model, **optional_params) # convert to pydantic object except Exception: return None # get llm provider @@ -68,10 +64,7 @@ def get_api_base( stream: bool = getattr(optional_params, "stream", False) - if ( - _optional_params.vertex_location is not None - and _optional_params.vertex_project is not None - ): + if _optional_params.vertex_location is not None and _optional_params.vertex_project is not None: from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VertexPartnerProvider @@ -105,13 +98,9 @@ def get_api_base( if custom_llm_provider == "gemini": if stream: - _api_base = "https://generativelanguage.googleapis.com/v1beta/models/{}:streamGenerateContent".format( - model - ) + _api_base = "https://generativelanguage.googleapis.com/v1beta/models/{}:streamGenerateContent".format(model) else: - _api_base = "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent".format( - model - ) + _api_base = "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent".format(model) return _api_base elif custom_llm_provider == "openai": _api_base = "https://api.openai.com" diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index cd49b5a4a87..f4bbfae3039 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -20,21 +20,13 @@ def get_response_headers(_response_headers: Optional[dict] = None) -> dict: openai_headers = {} if "x-ratelimit-limit-requests" in _response_headers: - openai_headers["x-ratelimit-limit-requests"] = _response_headers[ - "x-ratelimit-limit-requests" - ] + openai_headers["x-ratelimit-limit-requests"] = _response_headers["x-ratelimit-limit-requests"] if "x-ratelimit-remaining-requests" in _response_headers: - openai_headers["x-ratelimit-remaining-requests"] = _response_headers[ - "x-ratelimit-remaining-requests" - ] + openai_headers["x-ratelimit-remaining-requests"] = _response_headers["x-ratelimit-remaining-requests"] if "x-ratelimit-limit-tokens" in _response_headers: - openai_headers["x-ratelimit-limit-tokens"] = _response_headers[ - "x-ratelimit-limit-tokens" - ] + openai_headers["x-ratelimit-limit-tokens"] = _response_headers["x-ratelimit-limit-tokens"] if "x-ratelimit-remaining-tokens" in _response_headers: - openai_headers["x-ratelimit-remaining-tokens"] = _response_headers[ - "x-ratelimit-remaining-tokens" - ] + openai_headers["x-ratelimit-remaining-tokens"] = _response_headers["x-ratelimit-remaining-tokens"] llm_provider_headers = _get_llm_provider_headers(_response_headers) return {**llm_provider_headers, **openai_headers} diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index ba870eb9459..5ac2dca9ccf 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -20,9 +20,7 @@ class ResponseMetadata: def __init__(self, result: Any): self.result = result - self._hidden_params: Union[HiddenParams, dict] = ( - getattr(result, "_hidden_params", {}) or {} - ) + self._hidden_params: Union[HiddenParams, dict] = getattr(result, "_hidden_params", {}) or {} @property def supports_response_time(self) -> bool: @@ -33,9 +31,7 @@ class ResponseMetadata: or isinstance(self.result, TranscriptionResponse) ) - def set_hidden_params( - self, logging_obj: LiteLLMLoggingObject, model: Optional[str], kwargs: dict - ) -> None: + def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: Optional[str], kwargs: dict) -> None: """Set hidden parameters on the response""" ## ADD OTHER HIDDEN PARAMS @@ -127,12 +123,7 @@ class ResponseMetadata: if ( logging_obj.caching_details is not None and logging_obj.caching_details.get("cache_hit") is True - and ( - cache_duration_ms := logging_obj.caching_details.get( - "cache_duration_ms" - ) - ) - is not None + and (cache_duration_ms := logging_obj.caching_details.get("cache_duration_ms")) is not None ): overhead_ms = total_response_time_ms - cache_duration_ms self._update_hidden_params( diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index f58126ec901..00e12ee7ce9 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -43,23 +43,15 @@ class LoggingCallbackManager: Auto-routes async callbacks to litellm._async_input_callback. """ if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm._async_input_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) else: - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm.input_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) - def add_litellm_service_callback( - self, callback: Union[CustomLogger, str, Callable] - ): + def add_litellm_service_callback(self, callback: Union[CustomLogger, str, Callable]): """ Add a service callback to litellm.service_callback """ - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm.service_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) def add_litellm_callback(self, callback: Union[CustomLogger, str, Callable]): """ @@ -72,66 +64,42 @@ class LoggingCallbackManager: parent_list=litellm.callbacks, # type: ignore ) - def add_litellm_success_callback( - self, callback: Union[CustomLogger, str, Callable] - ): + def add_litellm_success_callback(self, callback: Union[CustomLogger, str, Callable]): """ Add a success callback to `litellm.success_callback`. Auto-routes async callbacks to litellm._async_success_callback. Special-cases 'dynamodb' and 'openmeter' as async callbacks. """ if isinstance(callback, str) and callback in ("dynamodb", "openmeter"): - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm._async_success_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) elif not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm._async_success_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) else: - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm.success_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) - def add_litellm_failure_callback( - self, callback: Union[CustomLogger, str, Callable] - ): + def add_litellm_failure_callback(self, callback: Union[CustomLogger, str, Callable]): """ Add a failure callback to `litellm.failure_callback`. Auto-routes async callbacks to litellm._async_failure_callback. """ if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm._async_failure_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) else: - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm.failure_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) - def add_litellm_async_success_callback( - self, callback: Union[CustomLogger, Callable, str] - ): + def add_litellm_async_success_callback(self, callback: Union[CustomLogger, Callable, str]): """ Add a success callback to litellm._async_success_callback """ - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm._async_success_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) - def add_litellm_async_failure_callback( - self, callback: Union[CustomLogger, Callable, str] - ): + def add_litellm_async_failure_callback(self, callback: Union[CustomLogger, Callable, str]): """ Add a failure callback to litellm._async_failure_callback """ - self._safe_add_callback_to_list( - callback=callback, parent_list=litellm._async_failure_callback - ) + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) - def remove_callback_from_list_by_object( - self, callback_list, obj, require_self=True - ): + def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True): """ Remove callbacks that are methods of a particular object (e.g., router cleanup) """ @@ -139,9 +107,7 @@ class LoggingCallbackManager: return if require_self: - remove_list = [ - c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj - ] + remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] else: remove_list = [c for c in callback_list if c == obj] @@ -169,22 +135,16 @@ class LoggingCallbackManager: for c in remove_list: callback_list.remove(c) - def _add_string_callback_to_list( - self, callback: str, parent_list: List[Union[CustomLogger, Callable, str]] - ): + def _add_string_callback_to_list(self, callback: str, parent_list: List[Union[CustomLogger, Callable, str]]): """ Add a string callback to a list, if the callback is already in the list, do not add it again. """ if callback not in parent_list: parent_list.append(callback) else: - verbose_logger.debug( - f"Callback {callback} already exists in {parent_list}, not adding again.." - ) + verbose_logger.debug(f"Callback {callback} already exists in {parent_list}, not adding again..") - def _check_callback_list_size( - self, parent_list: List[Union[CustomLogger, Callable, str]] - ) -> bool: + def _check_callback_list_size(self, parent_list: List[Union[CustomLogger, Callable, str]]) -> bool: """ Check if adding another callback would exceed MAX_CALLBACKS Returns True if safe to add, False if would exceed limit @@ -214,10 +174,7 @@ class LoggingCallbackManager: callback_config = litellm.callback_settings.get(callback) # Check if callback is in callback_settings with callback_type: generic_api - if ( - isinstance(callback_config, dict) - and callback_config.get("callback_type") == "generic_api" - ): + if isinstance(callback_config, dict) and callback_config.get("callback_type") == "generic_api": endpoint = callback_config.get("endpoint") headers = callback_config.get("headers") event_types = callback_config.get("event_types") @@ -298,14 +255,10 @@ class LoggingCallbackManager: # Check if the callback is a custom callback if isinstance(callback, str): - callback = LoggingCallbackManager._add_custom_callback_generic_api_str( - callback - ) + callback = LoggingCallbackManager._add_custom_callback_generic_api_str(callback) if isinstance(callback, str): - self._add_string_callback_to_list( - callback=callback, parent_list=parent_list - ) + self._add_string_callback_to_list(callback=callback, parent_list=parent_list) elif isinstance(callback, CustomLogger): self._add_custom_logger_to_list( custom_logger=callback, @@ -313,13 +266,9 @@ class LoggingCallbackManager: ) elif callable(callback): - self._add_callback_function_to_list( - callback=callback, parent_list=parent_list - ) + self._add_callback_function_to_list(callback=callback, parent_list=parent_list) - def _add_callback_function_to_list( - self, callback: Callable, parent_list: List[Union[CustomLogger, Callable, str]] - ): + def _add_callback_function_to_list(self, callback: Callable, parent_list: List[Union[CustomLogger, Callable, str]]): """ Add a callback function to a list, if the callback is already in the list, do not add it again. """ @@ -407,9 +356,7 @@ class LoggingCallbackManager: litellm._async_success_callback, litellm._async_failure_callback, ): - self.remove_callback_from_list_by_object( - callback_list, obj, require_self=require_self - ) + self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) def get_active_additional_logging_utils_from_custom_logger( self, @@ -426,15 +373,11 @@ class LoggingCallbackManager: all_callbacks = self._get_all_callbacks() matched_callbacks: Set[AdditionalLoggingUtils] = set() for callback in all_callbacks: - if isinstance(callback, CustomLogger) and isinstance( - callback, AdditionalLoggingUtils - ): + if isinstance(callback, CustomLogger) and isinstance(callback, AdditionalLoggingUtils): matched_callbacks.add(callback) return matched_callbacks - def get_custom_loggers_for_type( - self, callback_type: Type[CustomLogger] - ) -> List[CustomLogger]: + def get_custom_loggers_for_type(self, callback_type: Type[CustomLogger]) -> List[CustomLogger]: """ Get all custom loggers that are instances of the given class type """ @@ -449,10 +392,7 @@ class LoggingCallbackManager: """ Returns True if any of the active callbacks are of the given type """ - return any( - isinstance(callback, callback_type) - for callback in self._get_all_callbacks() - ) + return any(isinstance(callback, callback_type) for callback in self._get_all_callbacks()) def get_callbacks_by_type(self) -> CallbacksByType: """ @@ -462,20 +402,14 @@ class LoggingCallbackManager: CallbacksByType: Dict with keys 'success', 'failure', 'success_and_failure' containing lists of callback strings """ # Get callback lists - success_callbacks = set( - litellm.success_callback + litellm._async_success_callback - ) - failure_callbacks = set( - litellm.failure_callback + litellm._async_failure_callback - ) + success_callbacks = set(litellm.success_callback + litellm._async_success_callback) + failure_callbacks = set(litellm.failure_callback + litellm._async_failure_callback) general_callbacks = set(litellm.callbacks) # Get all unique callbacks all_callbacks = success_callbacks | failure_callbacks | general_callbacks - result: CallbacksByType = CallbacksByType( - success=[], failure=[], success_and_failure=[] - ) + result: CallbacksByType = CallbacksByType(success=[], failure=[], success_and_failure=[]) for callback in all_callbacks: callback_str = self._get_callback_string(callback) @@ -508,9 +442,7 @@ class LoggingCallbackManager: return callback elif isinstance(callback, CustomLogger): # Try to get the string representation from the registry - callback_str = CustomLoggerRegistry.get_callback_str_from_class_type( - type(callback) - ) + callback_str = CustomLoggerRegistry.get_callback_str_from_class_type(type(callback)) return callback_str if callback_str is not None else type(callback).__name__ elif callable(callback): return getattr(callback, "__name__", str(callback)) @@ -528,16 +460,12 @@ class LoggingCallbackManager: ) # get the custom logger class type - custom_logger_class_type = ( - CustomLoggerRegistry.get_class_type_for_custom_logger_name(callback_name) - ) + custom_logger_class_type = CustomLoggerRegistry.get_class_type_for_custom_logger_name(callback_name) # get the active custom logger custom_logger = self.get_custom_loggers_for_type(custom_logger_class_type) if len(custom_logger) == 0: - raise ValueError( - f"No active custom logger found for callback name: {callback_name}" - ) + raise ValueError(f"No active custom logger found for callback name: {callback_name}") return custom_logger[0] diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 4b2b740935c..720a850b47f 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -178,9 +178,7 @@ def _get_parent_otel_span_from_logging_obj( return _get_parent_otel_span_from_kwargs(logging_obj.model_call_details) except Exception as e: - verbose_logger.exception( - f"Error in _get_parent_otel_span_from_logging_obj: {str(e)}" - ) + verbose_logger.exception(f"Error in _get_parent_otel_span_from_logging_obj: {str(e)}") return None @@ -229,9 +227,7 @@ def _assemble_complete_response_from_streaming_chunks( Optional[Union[ModelResponse, TextCompletionResponse]]: Complete streaming response """ - complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = None + complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = None if isinstance(result, ModelResponse): return result @@ -246,10 +242,8 @@ def _assemble_complete_response_from_streaming_chunks( end_time=end_time, ) except Exception as e: - log_message = ( - "Error occurred building stream chunk in {} success logging: {}".format( - "async" if is_async else "sync", str(e) - ) + log_message = "Error occurred building stream chunk in {} success logging: {}".format( + "async" if is_async else "sync", str(e) ) verbose_logger.exception(log_message) complete_streaming_response = None @@ -269,9 +263,7 @@ def _set_duration_in_model_call_details( if logging_obj and hasattr(logging_obj, "model_call_details"): logging_obj.model_call_details["llm_api_duration_ms"] = duration_ms else: - verbose_logger.debug( - "`logging_obj` not found - unable to track `llm_api_duration_ms" - ) + verbose_logger.debug("`logging_obj` not found - unable to track `llm_api_duration_ms") except Exception as e: verbose_logger.warning(f"Error setting `llm_api_duration_ms`: {str(e)}") diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 294ba8e5dea..a9d5c8a8eb7 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -69,9 +69,7 @@ class LoggingWorker: # Check if we need to reinitialize due to event loop change if self._queue is not None and self._bound_loop is not current_loop: - verbose_logger.debug( - "LoggingWorker: Event loop changed, reinitializing queue and worker" - ) + verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker") # Clear old state - these are bound to the old loop self._queue = None self._sem = None @@ -121,9 +119,7 @@ class LoggingWorker: try: task = await self._queue.get() # Track each spawned coroutine so we can cancel on shutdown. - processing_task = asyncio.create_task( - self._process_log_task(task, self._sem) - ) + processing_task = asyncio.create_task(self._process_log_task(task, self._sem)) self._running_tasks.add(processing_task) processing_task.add_done_callback(self._running_tasks.discard) except Exception: @@ -211,14 +207,11 @@ class LoggingWorker: time_since_last_clear = current_time - self._last_aggressive_clear_time remaining_cooldown = max( 0.0, - LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS - - time_since_last_clear, + LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS - time_since_last_clear, ) # Add a small buffer (10% of cooldown or 50ms, whichever is larger) to ensure # cooldown has expired and aggressive clear has completed - return remaining_cooldown + max( - 0.05, LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS * 0.1 - ) + return remaining_cooldown + max(0.05, LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS * 0.1) except RuntimeError: # No event loop, return minimum delay return 0.1 @@ -266,9 +259,7 @@ class LoggingWorker: return [] # Calculate items based on percentage of queue size - items_to_extract = ( - self.max_queue_size * LOGGING_WORKER_CLEAR_PERCENTAGE - ) // 100 + items_to_extract = (self.max_queue_size * LOGGING_WORKER_CLEAR_PERCENTAGE) // 100 # Use actual queue size to avoid unnecessary iterations actual_size = self._queue.qsize() if actual_size == 0: @@ -285,9 +276,7 @@ class LoggingWorker: return extracted_tasks - async def _aggressively_clear_queue_async( - self, new_task: Optional[LoggingTask] = None - ) -> None: + async def _aggressively_clear_queue_async(self, new_task: Optional[LoggingTask] = None) -> None: """ Aggressively clear the queue by extracting and processing items. This is called when the queue is full to prevent dropping logs. @@ -307,9 +296,7 @@ class LoggingWorker: if extracted_tasks: await self._process_extracted_tasks(extracted_tasks) except Exception as e: - verbose_logger.exception( - f"LoggingWorker error during aggressive clear: {e}" - ) + verbose_logger.exception(f"LoggingWorker error during aggressive clear: {e}") finally: # Always reset the flag even if an error occurs self._aggressive_clear_in_progress = False @@ -395,9 +382,7 @@ class LoggingWorker: for _ in range(MAX_ITERATIONS_TO_CLEAR_QUEUE): # Check if we've exceeded the maximum time if asyncio.get_event_loop().time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE: - verbose_logger.warning( - f"clear_queue exceeded max_time of {MAX_TIME_TO_CLEAR_QUEUE}s, stopping early" - ) + verbose_logger.warning(f"clear_queue exceeded max_time of {MAX_TIME_TO_CLEAR_QUEUE}s, stopping early") break try: @@ -431,11 +416,7 @@ class LoggingWorker: has_valid_handler = False for handler in verbose_logger.handlers: try: - if ( - hasattr(handler, "stream") - and handler.stream - and not handler.stream.closed - ): + if hasattr(handler, "stream") and handler.stream and not handler.stream.closed: has_valid_handler = True break elif not hasattr(handler, "stream"): @@ -482,9 +463,7 @@ class LoggingWorker: return queue_size = self._queue.qsize() - self._safe_log( - "info", f"[LoggingWorker] atexit: Flushing {queue_size} remaining events..." - ) + self._safe_log("info", f"[LoggingWorker] atexit: Flushing {queue_size} remaining events...") # Create a new event loop since the original is closed loop = asyncio.new_event_loop() @@ -502,10 +481,7 @@ class LoggingWorker: previous_raise_exceptions = logging.raiseExceptions logging.raiseExceptions = False try: - while ( - not self._queue.empty() - and processed < MAX_ITERATIONS_TO_CLEAR_QUEUE - ): + while not self._queue.empty() and processed < MAX_ITERATIONS_TO_CLEAR_QUEUE: if loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE: self._safe_log( "warning", diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index 855e52098c4..39b3f0d5376 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -67,20 +67,12 @@ class ModelParamHelper: (``_get_relevant_args_to_use_for_logging``). Callers treat the result as read-only. """ - chat_completion_kwargs = ( - ModelParamHelper._get_litellm_supported_chat_completion_kwargs() - ) - text_completion_kwargs = ( - ModelParamHelper._get_litellm_supported_text_completion_kwargs() - ) + chat_completion_kwargs = ModelParamHelper._get_litellm_supported_chat_completion_kwargs() + text_completion_kwargs = ModelParamHelper._get_litellm_supported_text_completion_kwargs() embedding_kwargs = ModelParamHelper._get_litellm_supported_embedding_kwargs() - transcription_kwargs = ( - ModelParamHelper._get_litellm_supported_transcription_kwargs() - ) + transcription_kwargs = ModelParamHelper._get_litellm_supported_transcription_kwargs() rerank_kwargs = ModelParamHelper._get_litellm_supported_rerank_kwargs() - responses_api_kwargs = ( - ModelParamHelper._get_litellm_supported_responses_api_kwargs() - ) + responses_api_kwargs = ModelParamHelper._get_litellm_supported_responses_api_kwargs() exclude_kwargs = ModelParamHelper._get_exclude_kwargs() combined_kwargs = chat_completion_kwargs.union( @@ -104,18 +96,14 @@ class ModelParamHelper: This follows the OpenAI API Spec """ - non_streaming_params: Set[str] = set( - getattr(CompletionCreateParamsNonStreaming, "__annotations__", {}).keys() - ) - streaming_params: Set[str] = set( - getattr(CompletionCreateParamsStreaming, "__annotations__", {}).keys() - ) + non_streaming_params: Set[str] = set(getattr(CompletionCreateParamsNonStreaming, "__annotations__", {}).keys()) + streaming_params: Set[str] = set(getattr(CompletionCreateParamsStreaming, "__annotations__", {}).keys()) litellm_provider_specific_params: Set[str] = ( ModelParamHelper.get_litellm_provider_specific_params_for_chat_params() ) - all_chat_completion_kwargs: Set[str] = non_streaming_params.union( - streaming_params - ).union(litellm_provider_specific_params) + all_chat_completion_kwargs: Set[str] = non_streaming_params.union(streaming_params).union( + litellm_provider_specific_params + ) return all_chat_completion_kwargs @staticmethod @@ -126,16 +114,8 @@ class ModelParamHelper: This follows the OpenAI API Spec """ all_text_completion_kwargs = set( - getattr( - TextCompletionCreateParamsNonStreaming, "__annotations__", {} - ).keys() - ).union( - set( - getattr( - TextCompletionCreateParamsStreaming, "__annotations__", {} - ).keys() - ) - ) + getattr(TextCompletionCreateParamsNonStreaming, "__annotations__", {}).keys() + ).union(set(getattr(TextCompletionCreateParamsStreaming, "__annotations__", {}).keys())) return all_text_completion_kwargs @staticmethod @@ -167,16 +147,8 @@ class ModelParamHelper: TranscriptionCreateParamsStreaming, ) - non_streaming_kwargs = set( - getattr( - TranscriptionCreateParamsNonStreaming, "__annotations__", {} - ).keys() - ) - streaming_kwargs = set( - getattr( - TranscriptionCreateParamsStreaming, "__annotations__", {} - ).keys() - ) + non_streaming_kwargs = set(getattr(TranscriptionCreateParamsNonStreaming, "__annotations__", {}).keys()) + streaming_kwargs = set(getattr(TranscriptionCreateParamsStreaming, "__annotations__", {}).keys()) all_transcription_kwargs = non_streaming_kwargs.union(streaming_kwargs) return all_transcription_kwargs @@ -191,12 +163,8 @@ class ModelParamHelper: This follows the OpenAI API Spec """ - non_streaming_params: Set[str] = set( - getattr(ResponseCreateParamsNonStreaming, "__annotations__", {}).keys() - ) - streaming_params: Set[str] = set( - getattr(ResponseCreateParamsStreaming, "__annotations__", {}).keys() - ) + non_streaming_params: Set[str] = set(getattr(ResponseCreateParamsNonStreaming, "__annotations__", {}).keys()) + streaming_params: Set[str] = set(getattr(ResponseCreateParamsStreaming, "__annotations__", {}).keys()) return non_streaming_params.union(streaming_params) @staticmethod @@ -207,6 +175,4 @@ class ModelParamHelper: return set(["metadata"]) -ModelParamHelper._relevant_logging_args = frozenset( - ModelParamHelper._get_relevant_args_to_use_for_logging() -) +ModelParamHelper._relevant_logging_args = frozenset(ModelParamHelper._get_relevant_args_to_use_for_logging()) diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index 6c290fa30c0..f4843f9d95c 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -135,10 +135,7 @@ def _is_choice_non_empty(choice: Any) -> bool: # Skip certain structural fields that are just default/None placeholders if extra_field_name == "index" and extra_field_value == 0: continue - if ( - extra_field_name in {"finish_reason", "logprobs"} - and extra_field_value is None - ): + if extra_field_name in {"finish_reason", "logprobs"} and extra_field_value is None: continue if extra_field_name == "delta": continue @@ -190,11 +187,7 @@ def _is_delta_non_empty(delta: Delta) -> bool: # Check all regular attributes of the delta object for attr_name in dir(delta): # Skip private attributes, methods, and Pydantic-specific fields - if ( - attr_name.startswith("_") - or callable(getattr(delta, attr_name)) - or attr_name.startswith("model_") - ): + if attr_name.startswith("_") or callable(getattr(delta, attr_name)) or attr_name.startswith("model_"): continue attr_value = getattr(delta, attr_name, None) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 4bdda6de2c8..538d5f650ef 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -44,13 +44,9 @@ from litellm.types.utils import ( if TYPE_CHECKING: # newer pattern to avoid importing pydantic objects on __init__.py from litellm.types.llms.openai import ChatCompletionImageObject -DEFAULT_USER_CONTINUE_MESSAGE = ChatCompletionUserMessage( - content="Please continue.", role="user" -) +DEFAULT_USER_CONTINUE_MESSAGE = ChatCompletionUserMessage(content="Please continue.", role="user") -DEFAULT_ASSISTANT_CONTINUE_MESSAGE = ChatCompletionAssistantMessage( - content="Please continue.", role="assistant" -) +DEFAULT_ASSISTANT_CONTINUE_MESSAGE = ChatCompletionAssistantMessage(content="Please continue.", role="assistant") if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LoggingClass @@ -98,9 +94,7 @@ def handle_messages_with_content_list_to_str_conversion( return messages -def strip_name_from_message( - message: AllMessageValues, allowed_name_roles: List[str] = ["user"] -) -> AllMessageValues: +def strip_name_from_message(message: AllMessageValues, allowed_name_roles: List[str] = ["user"]) -> AllMessageValues: """ Removes 'name' from message """ @@ -202,9 +196,7 @@ def get_str_from_messages(messages: List[AllMessageValues]) -> str: def is_non_content_values_set(message: AllMessageValues) -> bool: ignore_keys = ["content", "role", "name"] - return any( - message.get(key, None) is not None for key in message if key not in ignore_keys - ) + return any(message.get(key, None) is not None for key in message if key not in ignore_keys) def _audio_or_image_in_message_content(message: AllMessageValues) -> bool: @@ -232,13 +224,9 @@ def convert_openai_message_to_only_content_messages( user_roles = ["user", "tool", "function"] for message in messages: if message.get("role") in user_roles: - converted_messages.append( - {"role": "user", "content": convert_content_list_to_str(message)} - ) + converted_messages.append({"role": "user", "content": convert_content_list_to_str(message)}) elif message.get("role") == "assistant": - converted_messages.append( - {"role": "assistant", "content": convert_content_list_to_str(message)} - ) + converted_messages.append({"role": "assistant", "content": convert_content_list_to_str(message)}) return converted_messages @@ -333,10 +321,7 @@ def _insert_user_continue_message( while i < len(result_messages): curr_message = result_messages[i] inserted_continue_message = False - if ( - _counts_for_alternation(curr_message) - and curr_message["role"] == "assistant" - ): + if _counts_for_alternation(curr_message) and curr_message["role"] == "assistant": # Preserve old behavior for malformed adjacent assistant sequences like # assistant(tool_calls) -> assistant(no-tool-calls) with no tool message. if i > 0 and result_messages[i - 1].get("role") == "assistant": @@ -423,14 +408,10 @@ def get_completion_messages( return messages.copy() ## INSERT USER CONTINUE MESSAGE - messages = _insert_user_continue_message( - messages, user_continue_message, ensure_alternating_roles - ) + messages = _insert_user_continue_message(messages, user_continue_message, ensure_alternating_roles) ## INSERT ASSISTANT CONTINUE MESSAGE - messages = _insert_assistant_continue_message( - messages, assistant_continue_message, ensure_alternating_roles - ) + messages = _insert_assistant_continue_message(messages, assistant_continue_message, ensure_alternating_roles) return messages @@ -449,9 +430,7 @@ def get_format_from_file_id(file_id: Optional[str]) -> Optional[str]: return None try: transformed_file_id = convert_b64_uid_to_unified_uid(file_id) - if transformed_file_id.startswith( - SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value - ): + if transformed_file_id.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): match = re.match( f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}:(.*?);unified_id", transformed_file_id, @@ -512,9 +491,7 @@ def update_messages_with_model_file_ids( # remap here, so skip instead of crashing. continue file_id = file_object_file_field.get("file_id") - format = file_object_file_field.get( - "format", get_format_from_file_id(file_id) - ) + format = file_object_file_field.get("format", get_format_from_file_id(file_id)) if file_id: provider_file_id = ( @@ -522,20 +499,11 @@ def update_messages_with_model_file_ids( if model_file_id_mapping and model_id is not None else None ) - if ( - not provider_file_id - and _is_base64_encoded_unified_file_id(file_id) - ): - unified_file_id = convert_b64_uid_to_unified_uid( - file_id - ) + if not provider_file_id and _is_base64_encoded_unified_file_id(file_id): + unified_file_id = convert_b64_uid_to_unified_uid(file_id) if "llm_output_file_id," in unified_file_id: - provider_file_id = unified_file_id.split( - "llm_output_file_id," - )[1].split(";")[0] - file_object_file_field["file_id"] = ( - provider_file_id or file_id - ) + provider_file_id = unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + file_object_file_field["file_id"] = provider_file_id or file_id if format: file_object_file_field["format"] = format return messages @@ -581,42 +549,26 @@ def update_responses_input_with_model_file_ids( if isinstance(content, list): updated_content = [] for content_item in content: - if ( - isinstance(content_item, dict) - and content_item.get("type") == "input_file" - ): + if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") if file_id: provider_file_id = file_id # Default to original # Check if we have a mapping for this file ID - if ( - model_file_id_mapping - and model_id - and file_id in model_file_id_mapping - ): + if model_file_id_mapping and model_id and file_id in model_file_id_mapping: # Use the model-specific file ID from mapping - provider_file_id = ( - model_file_id_mapping.get(file_id, {}).get(model_id) - or file_id - ) + provider_file_id = model_file_id_mapping.get(file_id, {}).get(model_id) or file_id updated_content_item = content_item.copy() updated_content_item["file_id"] = provider_file_id updated_content.append(updated_content_item) else: # Check if this is a base64-encoded unified file ID without mapping - is_unified_file_id = _is_base64_encoded_unified_file_id( - file_id - ) + is_unified_file_id = _is_base64_encoded_unified_file_id(file_id) if is_unified_file_id: # Fallback: decode unified file ID - unified_file_id = convert_b64_uid_to_unified_uid( - file_id - ) + unified_file_id = convert_b64_uid_to_unified_uid(file_id) if "llm_output_file_id," in unified_file_id: - provider_file_id = unified_file_id.split( - "llm_output_file_id," - )[1].split(";")[0] + provider_file_id = unified_file_id.split("llm_output_file_id,")[1].split(";")[0] updated_content_item = content_item.copy() updated_content_item["file_id"] = provider_file_id @@ -670,9 +622,7 @@ def _decode_vector_store_ids_in_tools( continue parsed = parse_unified_id(vs_id) - provider_resource_id = ( - parsed.get("provider_resource_id") if parsed else None - ) + provider_resource_id = parsed.get("provider_resource_id") if parsed else None if not provider_resource_id: verbose_logger.warning( @@ -737,10 +687,7 @@ def update_responses_tools_with_model_file_ids( # Check if we have a mapping for this file ID if file_id in model_file_id_mapping: # Map to provider-specific file ID - provider_file_id = ( - model_file_id_mapping.get(file_id, {}).get(model_id) - or file_id - ) + provider_file_id = model_file_id_mapping.get(file_id, {}).get(model_id) or file_id updated_file_ids.append(provider_file_id) else: updated_file_ids.append(file_id) @@ -965,9 +912,9 @@ def unpack_defs( # Use iterative approach with queue to avoid recursion # Each item in queue is (node, parent_container, key/index, active_defs, ref_chain) - queue: deque[ - tuple[Any, Union[dict, list, None], Union[str, int, None], dict, set] - ] = deque([(schema, None, None, root_defs, set())]) + queue: deque[tuple[Any, Union[dict, list, None], Union[str, int, None], dict, set]] = deque( + [(schema, None, None, root_defs, set())] + ) inlined_bytes = 0 while queue: @@ -1050,9 +997,7 @@ def _has_legacy_defs(schema: object) -> bool: if not isinstance(schema, dict): return False components = schema.get("components") - return "definitions" in schema or ( - isinstance(components, dict) and isinstance(components.get("schemas"), dict) - ) + return "definitions" in schema or (isinstance(components, dict) and isinstance(components.get("schemas"), dict)) # Schema-bomb budget for ``unpack_legacy_defs``: cap the cumulative JSON-byte @@ -1256,10 +1201,7 @@ def infer_content_type_from_url_and_content( return type_to_mime[detected_type] # If all fallbacks failed, raise error - raise ValueError( - f"Unable to determine content type from URL: {url}. " - f"Response content-type: {current_content_type}" - ) + raise ValueError(f"Unable to determine content type from URL: {url}. Response content-type: {current_content_type}") def get_tool_call_names(tools: List[ChatCompletionToolParam]) -> List[str]: @@ -1315,9 +1257,7 @@ def check_is_function_call(logging_obj: "LoggingClass") -> bool: is_function_call, ) - if hasattr(logging_obj, "optional_params") and isinstance( - logging_obj.optional_params, dict - ): + if hasattr(logging_obj, "optional_params") and isinstance(logging_obj.optional_params, dict): if is_function_call(logging_obj.optional_params): return True @@ -1423,9 +1363,7 @@ def get_last_user_message(messages: List[AllMessageValues]) -> Optional[str]: return result if result else None -def set_last_user_message( - messages: List[AllMessageValues], content: str -) -> List[AllMessageValues]: +def set_last_user_message(messages: List[AllMessageValues], content: str) -> List[AllMessageValues]: """ Set the last user message @@ -1440,11 +1378,7 @@ def set_last_user_message( # Stop when we hit a non-user message break if idx_to_remove: - messages = [ - message - for idx, message in enumerate(reversed(messages)) - if idx not in idx_to_remove - ] + messages = [message for idx, message in enumerate(reversed(messages)) if idx not in idx_to_remove] messages.reverse() messages.append({"role": "user", "content": content}) return messages @@ -1478,9 +1412,7 @@ def add_system_prompt_to_messages( if isinstance(existing_content, str): merged_content = f"{system_prompt.strip()}\n\n{existing_content}" elif isinstance(existing_content, list): - merged_content = [{"type": "text", "text": system_prompt.strip()}] + list( - existing_content - ) + merged_content = [{"type": "text", "text": system_prompt.strip()}] + list(existing_content) else: merged_content = [{"type": "text", "text": system_prompt.strip()}] first["content"] = merged_content @@ -1711,8 +1643,7 @@ def parse_tool_call_arguments( repaired = _attempt_json_repair(arguments) if repaired is not None: verbose_logger.warning( - "Repaired truncated tool call arguments for tool '%s' (%s). " - "Original (%d chars): %.200s%s", + "Repaired truncated tool call arguments for tool '%s' (%s). Original (%d chars): %.200s%s", tool_name or "", context or "unknown context", len(arguments), @@ -1728,10 +1659,7 @@ def parse_tool_call_arguments( if context: error_parts.append(f"({context})") - error_message = ( - " ".join(error_parts) - + f". Error: {str(original_error)}. Arguments: {arguments}" - ) + error_message = " ".join(error_parts) + f". Error: {str(original_error)}. Arguments: {arguments}" raise ValueError(error_message) from original_error diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b4f492bc26f..e54218cb8db 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -104,9 +104,7 @@ def map_system_message_pt(messages: list) -> list: if i < len(messages) - 1: # Not the last message next_m = messages[i + 1] next_role = next_m["role"] - if ( - next_role == "user" or next_role == "assistant" - ): # Next message is a user or assistant message + if next_role == "user" or next_role == "assistant": # Next message is a user or assistant message # Merge system prompt into the next message next_m["content"] = m["content"] + " " + next_m["content"] elif next_role == "system": # Next message is a system message @@ -186,9 +184,7 @@ def convert_to_ollama_image(openai_image_url: str): ) -def _handle_ollama_system_message( - messages: list, prompt: str, msg_i: int -) -> Tuple[str, int]: +def _handle_ollama_system_message(messages: list, prompt: str, msg_i: int) -> Tuple[str, int]: system_content_str = "" ## MERGE CONSECUTIVE SYSTEM CONTENT ## while msg_i < len(messages) and messages[msg_i]["role"] == "system": @@ -234,9 +230,7 @@ def ollama_pt( if user_content_str: prompt += f"### User:\n{user_content_str}\n\n" - system_content_str, msg_i = _handle_ollama_system_message( - messages, prompt, msg_i - ) + system_content_str, msg_i = _handle_ollama_system_message(messages, prompt, msg_i) if system_content_str: prompt += f"### System:\n{system_content_str}\n\n" @@ -265,9 +259,7 @@ def ollama_pt( ) if ollama_tool_calls: - assistant_content_str += ( - f"Tool Calls: {json.dumps(ollama_tool_calls, indent=2)}" - ) + assistant_content_str += f"Tool Calls: {json.dumps(ollama_tool_calls, indent=2)}" msg_i += 1 @@ -314,11 +306,7 @@ def falcon_instruct_pt(messages): if message["role"] == "system": prompt += message["content"] else: - prompt += ( - message["role"] - + ":" - + message["content"].replace("\r\n", "\n").replace("\n\n", "\n") - ) + prompt += message["role"] + ":" + message["content"].replace("\r\n", "\n").replace("\n\n", "\n") prompt += "\n\n" return prompt @@ -376,9 +364,7 @@ def phind_codellama_pt(messages): return prompt -def _render_chat_template( - env, chat_template: str, bos_token: str, eos_token: str, messages: list -) -> str: +def _render_chat_template(env, chat_template: str, bos_token: str, eos_token: str, messages: list) -> str: """ Shared template rendering logic for both sync and async hf_chat_template @@ -426,9 +412,7 @@ def _render_chat_template( try: for message in messages: if message["role"] == "system": - reformatted_messages.append( - {"role": "user", "content": message["content"]} - ) + reformatted_messages.append({"role": "user", "content": message["content"]}) else: reformatted_messages.append(message) rendered_text = template.render( @@ -443,20 +427,13 @@ def _render_chat_template( new_messages = [] for i in range(len(reformatted_messages) - 1): new_messages.append(reformatted_messages[i]) - if ( - reformatted_messages[i]["role"] - == reformatted_messages[i + 1]["role"] - ): + if reformatted_messages[i]["role"] == reformatted_messages[i + 1]["role"]: if reformatted_messages[i]["role"] == "user": - new_messages.append( - {"role": "assistant", "content": ""} - ) + new_messages.append({"role": "assistant", "content": ""}) else: new_messages.append({"role": "user", "content": ""}) new_messages.append(reformatted_messages[-1]) - rendered_text = template.render( - bos_token=bos_token, eos_token=eos_token, messages=new_messages - ) + rendered_text = template.render(bos_token=bos_token, eos_token=eos_token, messages=new_messages) return rendered_text except Exception as e: @@ -496,12 +473,8 @@ async def _afetch_and_extract_template( and "chat_template" in tokenizer_config["tokenizer"] ): tokenizer_data: dict = tokenizer_config["tokenizer"] # type: ignore - bos_token = _extract_token_value( - token_value=tokenizer_data.get("bos_token") - ) - eos_token = _extract_token_value( - token_value=tokenizer_data.get("eos_token") - ) + bos_token = _extract_token_value(token_value=tokenizer_data.get("bos_token")) + eos_token = _extract_token_value(token_value=tokenizer_data.get("eos_token")) chat_template = tokenizer_data["chat_template"] else: # Fallback: Try to fetch chat template from separate .jinja file @@ -515,12 +488,8 @@ async def _afetch_and_extract_template( and isinstance(tokenizer_config["tokenizer"], dict) ): tokenizer_data: dict = tokenizer_config["tokenizer"] # type: ignore - bos_token = _extract_token_value( - token_value=tokenizer_data.get("bos_token") - ) - eos_token = _extract_token_value( - token_value=tokenizer_data.get("eos_token") - ) + bos_token = _extract_token_value(token_value=tokenizer_data.get("bos_token")) + eos_token = _extract_token_value(token_value=tokenizer_data.get("eos_token")) else: raise Exception("No chat template found") @@ -558,12 +527,8 @@ def _fetch_and_extract_template( and "chat_template" in tokenizer_config["tokenizer"] ): tokenizer_data: dict = tokenizer_config["tokenizer"] # type: ignore - bos_token = _extract_token_value( - token_value=tokenizer_data.get("bos_token") - ) - eos_token = _extract_token_value( - token_value=tokenizer_data.get("eos_token") - ) + bos_token = _extract_token_value(token_value=tokenizer_data.get("bos_token")) + eos_token = _extract_token_value(token_value=tokenizer_data.get("eos_token")) chat_template = tokenizer_data["chat_template"] else: # Fallback: Try to fetch chat template from separate .jinja file @@ -577,21 +542,15 @@ def _fetch_and_extract_template( and isinstance(tokenizer_config["tokenizer"], dict) ): tokenizer_data: dict = tokenizer_config["tokenizer"] # type: ignore - bos_token = _extract_token_value( - token_value=tokenizer_data.get("bos_token") - ) - eos_token = _extract_token_value( - token_value=tokenizer_data.get("eos_token") - ) + bos_token = _extract_token_value(token_value=tokenizer_data.get("bos_token")) + eos_token = _extract_token_value(token_value=tokenizer_data.get("eos_token")) else: raise Exception("No chat template found") return chat_template, bos_token, eos_token # type: ignore -async def ahf_chat_template( - model: str, messages: list, chat_template: Optional[Any] = None -): +async def ahf_chat_template(model: str, messages: list, chat_template: Optional[Any] = None): """HuggingFace chat template (async version)""" from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import ( _aget_chat_template_file, @@ -646,9 +605,7 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any] = def deepseek_r1_pt(messages): - return hf_chat_template( - model="deepseek-r1/deepseek-r1-7b-instruct", messages=messages - ) + return hf_chat_template(model="deepseek-r1/deepseek-r1-7b-instruct", messages=messages) # Anthropic template @@ -698,9 +655,7 @@ def get_model_info(token, model): model_info = response.json() for m in model_info: if m["name"].lower().strip() == model.strip(): - return m["config"].get("prompt_format", None), m["config"].get( - "chat_template", None - ) + return m["config"].get("prompt_format", None), m["config"].get("chat_template", None) return None, None else: return None, None @@ -779,18 +734,14 @@ def anthropic_pt( AI_PROMPT = "\n\nAssistant: " prompt = "" - for idx, message in enumerate( - messages - ): # needs to start with `\n\nHuman: ` and end with `\n\nAssistant: ` + for idx, message in enumerate(messages): # needs to start with `\n\nHuman: ` and end with `\n\nAssistant: ` if message["role"] == "user": prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" elif message["role"] == "system": prompt += f"{AnthropicConstants.HUMAN_PROMPT.value}{message['content']}" else: prompt += f"{AnthropicConstants.AI_PROMPT.value}{message['content']}" - if ( - idx == 0 and message["role"] == "assistant" - ): # ensure the prompt always starts with `\n\nHuman: ` + if idx == 0 and message["role"] == "assistant": # ensure the prompt always starts with `\n\nHuman: ` prompt = f"{AnthropicConstants.HUMAN_PROMPT.value}" + prompt if messages[-1]["role"] != "assistant": prompt += f"{AnthropicConstants.AI_PROMPT.value}" @@ -874,9 +825,7 @@ def convert_generic_image_chunk_to_openai_image_obj( return "data:{};{},{}".format(media_type, image_chunk["type"], image_chunk["data"]) -def convert_to_anthropic_image_obj( - openai_image_url: str, format: Optional[str] -) -> GenericImageParsingChunk: +def convert_to_anthropic_image_obj(openai_image_url: str, format: Optional[str]) -> GenericImageParsingChunk: """ Input: "image_url": "data:image/jpeg;base64,{base64_image}", @@ -936,9 +885,7 @@ def create_anthropic_image_param( # as these providers don't support URL sources for images if is_bedrock_invoke or image_url.startswith("http://"): base64_url = convert_url_to_base64(url=image_url) - image_chunk = convert_to_anthropic_image_obj( - openai_image_url=base64_url, format=format - ) + image_chunk = convert_to_anthropic_image_obj(openai_image_url=base64_url, format=format) return AnthropicMessagesImageParam( type="image", source=AnthropicContentParamSource( @@ -958,9 +905,7 @@ def create_anthropic_image_param( ) else: # Convert to base64 for data URIs or other formats - image_chunk = convert_to_anthropic_image_obj( - openai_image_url=image_url, format=format - ) + image_chunk = convert_to_anthropic_image_obj(openai_image_url=image_url, format=format) return AnthropicMessagesImageParam( type="image", source=AnthropicContentParamSource( @@ -1037,9 +982,7 @@ def convert_to_anthropic_tool_invoke_xml(tool_calls: list) -> str: tool_arguments, tool_name=tool_name, context="Anthropic XML tool invoke" ) if isinstance(parsed_args, dict): - parameters = "".join( - f"<{param}>{val}\n" for param, val in parsed_args.items() - ) + parameters = "".join(f"<{param}>{val}\n" for param, val in parsed_args.items()) else: parameters = f"{parsed_args}\n" invokes += f"\n{tool_name}\n\n{parameters}\n\n" @@ -1071,14 +1014,8 @@ def anthropic_messages_pt_xml(messages: list): if isinstance(messages[msg_i]["content"], list): for m in messages[msg_i]["content"]: if m.get("type", "") == "image_url": - format = ( - m["image_url"].get("format") - if isinstance(m["image_url"], dict) - else None - ) - image_param = create_anthropic_image_param( - m["image_url"], format=format - ) + format = m["image_url"].get("format") if isinstance(m["image_url"], dict) else None + image_param = create_anthropic_image_param(m["image_url"], format=format) # Convert to dict format for XML version source = image_param["source"] if isinstance(source, dict) and source.get("type") == "url": @@ -1129,12 +1066,8 @@ def anthropic_messages_pt_xml(messages: list): assistant_content = [] ## MERGE CONSECUTIVE ASSISTANT CONTENT ## while msg_i < len(messages) and messages[msg_i]["role"] == "assistant": - assistant_text = ( - messages[msg_i].get("content") or "" - ) # either string or none - if messages[msg_i].get( - "tool_calls", [] - ): # support assistant tool invoke conversion + assistant_text = messages[msg_i].get("content") or "" # either string or none + if messages[msg_i].get("tool_calls", []): # support assistant tool invoke conversion assistant_text += convert_to_anthropic_tool_invoke_xml( # type: ignore messages[msg_i]["tool_calls"] ) @@ -1147,9 +1080,7 @@ def anthropic_messages_pt_xml(messages: list): if not new_messages or new_messages[0]["role"] != "user": if litellm.modify_params: - new_messages.insert( - 0, {"role": "user", "content": [{"type": "text", "text": "."}]} - ) + new_messages.insert(0, {"role": "user", "content": [{"type": "text", "text": "."}]}) else: raise Exception( "Invalid first message. Should always start with 'role'='user' for Anthropic. System prompt is sent separately for Anthropic. set 'litellm.modify_params = True' or 'litellm_settings:modify_params = True' on proxy, to insert a placeholder user message - '.' as the first message, " @@ -1158,9 +1089,7 @@ def anthropic_messages_pt_xml(messages: list): if new_messages[-1]["role"] == "assistant": for content in new_messages[-1]["content"]: if isinstance(content, dict) and content["type"] == "text": - content["text"] = content[ - "text" - ].rstrip() # no trailing whitespace for final assistant message + content["text"] = content["text"].rstrip() # no trailing whitespace for final assistant message return new_messages @@ -1256,9 +1185,7 @@ def _gemini_tool_call_invoke_helper( return function_call -def _encode_tool_call_id_with_signature( - tool_call_id: str, thought_signature: Optional[str] -) -> str: +def _encode_tool_call_id_with_signature(tool_call_id: str, thought_signature: Optional[str]) -> str: """ Embed thought signature into tool call ID for OpenAI client compatibility. @@ -1277,9 +1204,7 @@ def _encode_tool_call_id_with_signature( return tool_call_id -def _get_thought_signature_from_tool( - tool: dict, model: Optional[str] = None -) -> Optional[str]: +def _get_thought_signature_from_tool(tool: dict, model: Optional[str] = None) -> Optional[str]: """Extract thought signature from tool call's provider_specific_fields. If not provided try to extract thought signature from tool call id @@ -1303,10 +1228,7 @@ def _get_thought_signature_from_tool( signature = func_provider_fields.get("thought_signature") if signature: return signature - elif ( - hasattr(function, "provider_specific_fields") - and function.provider_specific_fields - ): + elif hasattr(function, "provider_specific_fields") and function.provider_specific_fields: if isinstance(function.provider_specific_fields, dict): signature = function.provider_specific_fields.get("thought_signature") if signature: @@ -1395,30 +1317,19 @@ def convert_to_gemini_tool_call_invoke( ) forward_tool_call_id = bool( - model - and VertexGeminiConfig._forward_gemini_function_call_id( - model, custom_llm_provider - ) + model and VertexGeminiConfig._forward_gemini_function_call_id(model, custom_llm_provider) ) if tool_calls is not None: for idx, tool in enumerate(tool_calls): if "function" in tool: - gemini_function_call: Optional[VertexFunctionCall] = ( - _gemini_tool_call_invoke_helper( - function_call_params=tool["function"], - tool_call_id=( - tool.get("id") if forward_tool_call_id else None - ), - ) + gemini_function_call: Optional[VertexFunctionCall] = _gemini_tool_call_invoke_helper( + function_call_params=tool["function"], + tool_call_id=(tool.get("id") if forward_tool_call_id else None), ) if gemini_function_call is not None: - part_dict: VertexPartType = { - "function_call": gemini_function_call - } - thought_signature = _get_thought_signature_from_tool( - dict(tool), model=model - ) + part_dict: VertexPartType = {"function_call": gemini_function_call} + thought_signature = _get_thought_signature_from_tool(dict(tool), model=model) if thought_signature: part_dict["thoughtSignature"] = thought_signature @@ -1430,30 +1341,20 @@ def convert_to_gemini_tool_call_invoke( ) ) elif function_call is not None: - gemini_function_call = _gemini_tool_call_invoke_helper( - function_call_params=function_call - ) + gemini_function_call = _gemini_tool_call_invoke_helper(function_call_params=function_call) if gemini_function_call is not None: - part_dict_function: VertexPartType = { - "function_call": gemini_function_call - } + part_dict_function: VertexPartType = {"function_call": gemini_function_call} # Extract thought signature from function_call's provider_specific_fields thought_signature = None provider_fields = ( - function_call.get("provider_specific_fields") - if isinstance(function_call, dict) - else {} + function_call.get("provider_specific_fields") if isinstance(function_call, dict) else {} ) if isinstance(provider_fields, dict): thought_signature = provider_fields.get("thought_signature") # If no signature found and model is gemini-3, use dummy signature - if ( - not thought_signature - and model - and VertexGeminiConfig._is_gemini_3_or_newer(model) - ): + if not thought_signature and model and VertexGeminiConfig._is_gemini_3_or_newer(model): thought_signature = _get_dummy_thought_signature() if thought_signature: @@ -1469,9 +1370,7 @@ def convert_to_gemini_tool_call_invoke( return _parts_list except Exception as e: raise Exception( - "Unable to convert openai tool calls={} to gemini tool calls. Received error={}".format( - message, str(e) - ) + "Unable to convert openai tool calls={} to gemini tool calls. Received error={}".format(message, str(e)) ) @@ -1524,14 +1423,10 @@ def convert_to_gemini_tool_call_result( if len(mime_rest) == 2 and mime_rest[0].startswith("image/"): # Strip any extra parameters (e.g. ";charset=UTF-8") from the MIME segment clean_mime = mime_rest[0].split(";")[0].strip() - inline_data_list.append( - BlobType(data=mime_rest[1], mime_type=clean_mime) - ) + inline_data_list.append(BlobType(data=mime_rest[1], mime_type=clean_mime)) content_str = "" except Exception as e: - verbose_logger.warning( - f"Failed to parse data URL in tool response: {e}" - ) + verbose_logger.warning(f"Failed to parse data URL in tool response: {e}") elif isinstance(message["content"], List): content_list = message["content"] for content in content_list: @@ -1550,24 +1445,16 @@ def convert_to_gemini_tool_call_result( ) ) except Exception as e: - verbose_logger.warning( - f"Failed to process Anthropic image block in tool response: {e}" - ) + verbose_logger.warning(f"Failed to process Anthropic image block in tool response: {e}") elif content_type in ("input_image", "image_url"): # Extract image for inline_data (for Computer Use screenshots and tool results) image_url_data = content.get("image_url", "") - image_url = ( - image_url_data.get("url", "") - if isinstance(image_url_data, dict) - else image_url_data - ) + image_url = image_url_data.get("url", "") if isinstance(image_url_data, dict) else image_url_data if image_url: # Convert image to base64 blob format for Gemini try: - image_obj = convert_to_anthropic_image_obj( - image_url, format=None - ) + image_obj = convert_to_anthropic_image_obj(image_url, format=None) inline_data_list.append( BlobType( data=image_obj["data"], @@ -1575,9 +1462,7 @@ def convert_to_gemini_tool_call_result( ) ) except Exception as e: - verbose_logger.warning( - f"Failed to process image in tool response: {e}" - ) + verbose_logger.warning(f"Failed to process image in tool response: {e}") elif content_type in ("file", "input_file"): # Extract file for inline_data (for tool results with PDF, audio, video, etc.) file_data = content.get("file_data", "") @@ -1594,9 +1479,7 @@ def convert_to_gemini_tool_call_result( if file_data: # Convert file to base64 blob format for Gemini try: - file_obj = convert_to_anthropic_image_obj( - file_data, format=None - ) + file_obj = convert_to_anthropic_image_obj(file_data, format=None) inline_data_list.append( BlobType( data=file_obj["data"], @@ -1604,9 +1487,7 @@ def convert_to_gemini_tool_call_result( ) ) except Exception as e: - verbose_logger.warning( - f"Failed to process file in tool response: {e}" - ) + verbose_logger.warning(f"Failed to process file in tool response: {e}") name: Optional[str] = message.get("name", "") # type: ignore # Recover name from last message with tool calls @@ -1615,11 +1496,7 @@ def convert_to_gemini_tool_call_result( msg_tool_call_id = message.get("tool_call_id", None) for tool in tools: prev_tool_call_id = tool.get("id", None) - if ( - msg_tool_call_id - and prev_tool_call_id - and msg_tool_call_id == prev_tool_call_id - ): + if msg_tool_call_id and prev_tool_call_id and msg_tool_call_id == prev_tool_call_id: name = tool.get("function", {}).get("name", "") # Echo the OpenAI tool_call_id on functionResponse (strip thought-signature suffix). @@ -1630,9 +1507,7 @@ def convert_to_gemini_tool_call_result( ) gemini_call_id: Optional[str] = None - if model and VertexGeminiConfig._forward_gemini_function_call_id( - model, custom_llm_provider - ): + if model and VertexGeminiConfig._forward_gemini_function_call_id(model, custom_llm_provider): raw_tool_call_id = message.get("tool_call_id") if raw_tool_call_id and isinstance(raw_tool_call_id, str): stripped_id = raw_tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] @@ -1677,9 +1552,7 @@ def convert_to_gemini_tool_call_result( # For multimodal function responses, Gemini expects media parts nested # inside functionResponse.parts instead of sibling content parts. if inline_data_list: - _function_response["parts"] = [ - {"inline_data": inline_data} for inline_data in inline_data_list - ] + _function_response["parts"] = [{"inline_data": inline_data} for inline_data in inline_data_list] return [_part] return _part @@ -1784,38 +1657,22 @@ def convert_to_anthropic_tool_result( anthropic_content_list.append(text_content) elif content["type"] == "image_url": image_url_value = content["image_url"] - format = ( - image_url_value.get("format") - if isinstance(image_url_value, dict) - else None - ) - url_str = ( - image_url_value.get("url") - if isinstance(image_url_value, dict) - else image_url_value - ) + format = image_url_value.get("format") if isinstance(image_url_value, dict) else None + url_str = image_url_value.get("url") if isinstance(image_url_value, dict) else image_url_value # Data URIs with non-image mime types (e.g. application/pdf) must # translate to Anthropic document blocks, not image blocks — # wrapping a PDF in `type: "image"` is rejected by the API. - if isinstance(url_str, str) and _is_anthropic_document_data_uri( - url_str - ): + if isinstance(url_str, str) and _is_anthropic_document_data_uri(url_str): synth_file_message: ChatCompletionFileObject = { "type": "file", "file": {"file_data": url_str}, } - _document_block = anthropic_process_openai_file_message( - synth_file_message - ) + _document_block = anthropic_process_openai_file_message(synth_file_message) _document_block = add_cache_control_to_content( - anthropic_content_element=cast( - AnthropicMessagesDocumentParam, _document_block - ), + anthropic_content_element=cast(AnthropicMessagesDocumentParam, _document_block), original_content_element=content, ) - anthropic_content_list.append( - cast(AnthropicMessagesDocumentParam, _document_block) - ) + anthropic_content_list.append(cast(AnthropicMessagesDocumentParam, _document_block)) else: _anthropic_image_param = create_anthropic_image_param( image_url_value, @@ -1826,16 +1683,12 @@ def convert_to_anthropic_tool_result( anthropic_content_element=_anthropic_image_param, original_content_element=content, ) - anthropic_content_list.append( - cast(AnthropicMessagesImageParam, _anthropic_image_param) - ) + anthropic_content_list.append(cast(AnthropicMessagesImageParam, _anthropic_image_param)) elif content["type"] == "file": file_content = cast(ChatCompletionFileObject, content) _file_block = anthropic_process_openai_file_message(file_content) _file_block = add_cache_control_to_content( - anthropic_content_element=cast( - AnthropicMessagesDocumentParam, _file_block - ), + anthropic_content_element=cast(AnthropicMessagesDocumentParam, _file_block), original_content_element=content, ) anthropic_content_list.append(_file_block) @@ -1883,9 +1736,7 @@ def convert_function_to_anthropic_tool_invoke( _name = get_attribute_or_key(function_call, "name") or "" _arguments = get_attribute_or_key(function_call, "arguments") - tool_input = parse_tool_call_arguments( - _arguments, tool_name=_name, context="Anthropic function to tool invoke" - ) + tool_input = parse_tool_call_arguments(_arguments, tool_name=_name, context="Anthropic function to tool invoke") anthropic_tool_invoke = [ AnthropicMessagesToolUseParam( @@ -1947,9 +1798,7 @@ def convert_to_anthropic_tool_invoke( Fixes: https://github.com/BerriAI/litellm/issues/17737 """ - anthropic_tool_invoke: List[ - Union[AnthropicMessagesToolUseParam, Dict[str, Any]] - ] = [] + anthropic_tool_invoke: List[Union[AnthropicMessagesToolUseParam, Dict[str, Any]]] = [] for tool in tool_calls: if not get_attribute_or_key(tool, "type") == "function": @@ -2006,9 +1855,7 @@ def convert_to_anthropic_tool_invoke( ) if "cache_control" in _content_element: - _anthropic_tool_use_param["cache_control"] = _content_element[ - "cache_control" - ] + _anthropic_tool_use_param["cache_control"] = _content_element["cache_control"] anthropic_tool_invoke.append(_anthropic_tool_use_param) @@ -2039,15 +1886,15 @@ def _anthropic_content_element_factory( image_chunk: GenericImageParsingChunk, ) -> Union[AnthropicMessagesImageParam, AnthropicMessagesDocumentParam]: if image_chunk["media_type"] == "application/pdf": - _anthropic_content_element: Union[ - AnthropicMessagesDocumentParam, AnthropicMessagesImageParam - ] = AnthropicMessagesDocumentParam( - type="document", - source=AnthropicContentParamSource( - type="base64", - media_type=image_chunk["media_type"], - data=image_chunk["data"], - ), + _anthropic_content_element: Union[AnthropicMessagesDocumentParam, AnthropicMessagesImageParam] = ( + AnthropicMessagesDocumentParam( + type="document", + source=AnthropicContentParamSource( + type="base64", + media_type=image_chunk["media_type"], + data=image_chunk["data"], + ), + ) ) else: _anthropic_content_element = AnthropicMessagesImageParam( @@ -2158,21 +2005,15 @@ def anthropic_process_openai_file_message( ), ) elif content_block_type == "container_upload": - return_block_param = AnthropicMessagesContainerUploadParam( - type="container_upload", file_id=file_id - ) + return_block_param = AnthropicMessagesContainerUploadParam(type="container_upload", file_id=file_id) if return_block_param is None: raise Exception(f"Unable to parse anthropic file message: {message}") return return_block_param - raise Exception( - f"Either file_data or file_id must be present in the file message: {message}" - ) + raise Exception(f"Either file_data or file_id must be present in the file message: {message}") -_EMPTY_TEXT_PLACEHOLDER = ( - "[System: Empty message content sanitised to satisfy protocol]" -) +_EMPTY_TEXT_PLACEHOLDER = "[System: Empty message content sanitised to satisfy protocol]" def _sanitize_empty_text_content( @@ -2373,9 +2214,7 @@ def _is_orphaned_tool_result( break if not found_matching_tool_call: - verbose_logger.debug( - "_is_orphaned_tool_result: Found orphaned tool result with redacted tool_call_id" - ) + verbose_logger.debug("_is_orphaned_tool_result: Found orphaned tool result with redacted tool_call_id") return True return False @@ -2420,9 +2259,7 @@ def sanitize_messages_for_tool_calling( # Case A: Check if assistant message has tool_calls without following tool results if current_message.get("role") == "assistant": - result_messages, messages_consumed = _add_missing_tool_results( - current_message, messages, i - ) + result_messages, messages_consumed = _add_missing_tool_results(current_message, messages, i) # If dummy tool results were added, extend sanitized_messages and skip consumed messages if len(result_messages) > 1: @@ -2477,11 +2314,7 @@ def sanitize_messages_for_tool_calling( seen_in_block = {} if duplicates_to_remove: - sanitized_messages = [ - msg - for idx, msg in enumerate(sanitized_messages) - if idx not in duplicates_to_remove - ] + sanitized_messages = [msg for idx, msg in enumerate(sanitized_messages) if idx not in duplicates_to_remove] return sanitized_messages @@ -2564,17 +2397,11 @@ def anthropic_messages_pt( for m in user_message_types_block["content"]: if m.get("type", "") == "image_url": m = cast(ChatCompletionImageObject, m) - format = ( - m["image_url"].get("format") - if isinstance(m["image_url"], dict) - else None - ) + format = m["image_url"].get("format") if isinstance(m["image_url"], dict) else None # Convert ChatCompletionImageUrlObject to dict if needed image_url_value = m["image_url"] if isinstance(image_url_value, str): - image_url_input: Union[str, dict[str, Any]] = ( - image_url_value - ) + image_url_input: Union[str, dict[str, Any]] = image_url_value else: # ChatCompletionImageUrlObject or dict case - convert to dict image_url_input = { @@ -2584,11 +2411,7 @@ def anthropic_messages_pt( # Bedrock invoke models have format: invoke/... # Vertex AI Anthropic also doesn't support URL sources for images is_bedrock_invoke = model.lower().startswith("invoke/") - is_vertex_ai = ( - llm_provider.startswith("vertex_ai") - if llm_provider - else False - ) + is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False force_base64 = is_bedrock_invoke or is_vertex_ai _anthropic_content_element = create_anthropic_image_param( image_url_input, @@ -2601,43 +2424,33 @@ def anthropic_messages_pt( ) if "cache_control" in _content_element: - _anthropic_content_element["cache_control"] = ( - _content_element["cache_control"] - ) + _anthropic_content_element["cache_control"] = _content_element["cache_control"] user_content.append(_anthropic_content_element) elif m.get("type", "") == "text": m = cast(ChatCompletionTextObject, m) - _anthropic_text_content_element = ( - AnthropicMessagesTextParam( - type="text", - text=m["text"], - ) + _anthropic_text_content_element = AnthropicMessagesTextParam( + type="text", + text=m["text"], ) _content_element = add_cache_control_to_content( anthropic_content_element=_anthropic_text_content_element, original_content_element=dict(m), ) - _content_element = cast( - AnthropicMessagesTextParam, _content_element - ) + _content_element = cast(AnthropicMessagesTextParam, _content_element) user_content.append(_content_element) elif m.get("type", "") == "document": _document_content_element = cast( AnthropicMessagesDocumentParam, add_cache_control_to_content( - anthropic_content_element=cast( - AnthropicMessagesDocumentParam, m - ), + anthropic_content_element=cast(AnthropicMessagesDocumentParam, m), original_content_element=dict(m), ), ) user_content.append(_document_content_element) elif m.get("type", "") == "file": - _file_content_element = ( - anthropic_process_openai_file_message( - cast(ChatCompletionFileObject, m) - ) + _file_content_element = anthropic_process_openai_file_message( + cast(ChatCompletionFileObject, m) ) _file_content_element = add_cache_control_to_content( anthropic_content_element=cast( @@ -2663,21 +2476,14 @@ def anthropic_messages_pt( ) if "cache_control" in _content_element: - _anthropic_content_text_element["cache_control"] = ( - _content_element["cache_control"] - ) + _anthropic_content_text_element["cache_control"] = _content_element["cache_control"] user_content.append(_anthropic_content_text_element) - elif ( - user_message_types_block["role"] == "tool" - or user_message_types_block["role"] == "function" - ): + elif user_message_types_block["role"] == "tool" or user_message_types_block["role"] == "function": # OpenAI's tool message content will always be a string user_content.append( - convert_to_anthropic_tool_result( - user_message_types_block, force_base64=force_base64 - ) + convert_to_anthropic_tool_result(user_message_types_block, force_base64=force_base64) ) msg_i += 1 @@ -2694,13 +2500,9 @@ def anthropic_messages_pt( assistant_content_block: ChatCompletionAssistantMessage = messages[msg_i] # type: ignore # Extract compaction_blocks from provider_specific_fields and add them first - _provider_specific_fields_raw = assistant_content_block.get( - "provider_specific_fields" - ) + _provider_specific_fields_raw = assistant_content_block.get("provider_specific_fields") if isinstance(_provider_specific_fields_raw, dict): - _compaction_blocks = _provider_specific_fields_raw.get( - "compaction_blocks" - ) + _compaction_blocks = _provider_specific_fields_raw.get("compaction_blocks") if _compaction_blocks and isinstance(_compaction_blocks, list): # Add compaction blocks at the beginning of assistant content : https://platform.claude.com/docs/en/build-with-claude/compaction assistant_content.extend(_compaction_blocks) # type: ignore @@ -2715,25 +2517,15 @@ def anthropic_messages_pt( _has_server_tool_calls = False if assistant_tool_calls is not None: for _tc in assistant_tool_calls: - _tc_id = ( - _tc.get("id") - if isinstance(_tc, dict) - else getattr(_tc, "id", None) - ) - if ( - _tc_id - and isinstance(_tc_id, str) - and _tc_id.startswith("srvtoolu_") - ): + _tc_id = _tc.get("id") if isinstance(_tc, dict) else getattr(_tc, "id", None) + if _tc_id and isinstance(_tc_id, str) and _tc_id.startswith("srvtoolu_"): _has_server_tool_calls = True break if ( thinking_blocks is not None and _has_server_tool_calls - and isinstance( - assistant_content_block.get("content", None), (str, type(None)) - ) + and isinstance(assistant_content_block.get("content", None), (str, type(None))) ): # INTERLEAVED MODE: When we have both thinking blocks and server # tool calls (e.g. web search), Anthropic's original response @@ -2743,17 +2535,11 @@ def anthropic_messages_pt( # verifies thinking block signatures based on position. # Build the tool call groups (server_tool_use + its result) - _provider_specific_fields_raw_tc = assistant_content_block.get( - "provider_specific_fields" - ) + _provider_specific_fields_raw_tc = assistant_content_block.get("provider_specific_fields") _provider_specific_fields_tc: Dict[str, Any] = {} if isinstance(_provider_specific_fields_raw_tc, dict): - _provider_specific_fields_tc = cast( - Dict[str, Any], _provider_specific_fields_raw_tc - ) - _web_search_results_tc = _provider_specific_fields_tc.get( - "web_search_results" - ) + _provider_specific_fields_tc = cast(Dict[str, Any], _provider_specific_fields_raw_tc) + _web_search_results_tc = _provider_specific_fields_tc.get("web_search_results") _tool_results_tc = _provider_specific_fields_tc.get("tool_results") tool_invoke_results = convert_to_anthropic_tool_invoke( assistant_tool_calls, # type: ignore @@ -2767,11 +2553,7 @@ def anthropic_messages_pt( regular_tool_uses: List[Any] = [] _current_group: List[Any] = [] for item in tool_invoke_results: - item_type = ( - item.get("type", "") - if isinstance(item, dict) - else getattr(item, "type", "") - ) + item_type = item.get("type", "") if isinstance(item, dict) else getattr(item, "type", "") if item_type == "server_tool_use": if _current_group: server_tool_groups.append(_current_group) @@ -2798,9 +2580,7 @@ def anthropic_messages_pt( original_content_element=dict(assistant_content_block), ) if "cache_control" in _content_element: - _anthropic_text_content_element["cache_control"] = ( - _content_element["cache_control"] - ) + _anthropic_text_content_element["cache_control"] = _content_element["cache_control"] text_element = _anthropic_text_content_element # Interleave: each thinking block precedes its server tool group. @@ -2818,18 +2598,12 @@ def anthropic_messages_pt( assistant_content.append(thinking_blocks[tb_idx]) tb_idx += 1 for block in server_tool_groups[grp_idx]: - item_id = ( - block.get("id") - if isinstance(block, dict) - else getattr(block, "id", None) - ) + item_id = block.get("id") if isinstance(block, dict) else getattr(block, "id", None) if item_id and item_id in unique_tool_ids: continue if item_id: unique_tool_ids.add(item_id) - assistant_content.append( - cast(AnthropicMessagesAssistantMessageValues, block) - ) + assistant_content.append(cast(AnthropicMessagesAssistantMessageValues, block)) grp_idx += 1 elif tb_idx < num_tb: # More thinking blocks than tool groups - emit before text @@ -2838,18 +2612,12 @@ def anthropic_messages_pt( else: # More tool groups than thinking blocks - emit remaining for block in server_tool_groups[grp_idx]: - item_id = ( - block.get("id") - if isinstance(block, dict) - else getattr(block, "id", None) - ) + item_id = block.get("id") if isinstance(block, dict) else getattr(block, "id", None) if item_id and item_id in unique_tool_ids: continue if item_id: unique_tool_ids.add(item_id) - assistant_content.append( - cast(AnthropicMessagesAssistantMessageValues, block) - ) + assistant_content.append(cast(AnthropicMessagesAssistantMessageValues, block)) grp_idx += 1 # Add text block (if any) @@ -2858,18 +2626,12 @@ def anthropic_messages_pt( # Add regular (non-server) tool calls at the end for item in regular_tool_uses: - item_id = ( - item.get("id") - if isinstance(item, dict) - else getattr(item, "id", None) - ) + item_id = item.get("id") if isinstance(item, dict) else getattr(item, "id", None) if item_id and item_id in unique_tool_ids: continue if item_id: unique_tool_ids.add(item_id) - assistant_content.append( - cast(AnthropicMessagesAssistantMessageValues, item) - ) + assistant_content.append(cast(AnthropicMessagesAssistantMessageValues, item)) # Mark tool_calls as already processed so they are not added again assistant_tool_calls = None @@ -2886,9 +2648,7 @@ def anthropic_messages_pt( _content_is_list = "content" in assistant_content_block and isinstance( assistant_content_block["content"], list ) - _content_list = ( - assistant_content_block.get("content") if _content_is_list else None - ) + _content_list = assistant_content_block.get("content") if _content_is_list else None _list_has_thinking = False if _content_is_list and _content_list is not None: for _item in _content_list: @@ -2922,17 +2682,13 @@ def anthropic_messages_pt( elif ( m.get("type", "") == "text" and len(text_block) > 0 ): # don't pass empty text blocks. anthropic api raises errors. - anthropic_message = AnthropicMessagesTextParam( - type="text", text=text_block - ) + anthropic_message = AnthropicMessagesTextParam(type="text", text=text_block) _cached_message = add_cache_control_to_content( anthropic_content_element=anthropic_message, original_content_element=dict(m), ) - assistant_content.append( - cast(AnthropicMessagesTextParam, _cached_message) - ) + assistant_content.append(cast(AnthropicMessagesTextParam, _cached_message)) # handle server_tool_use blocks (tool search, web search, etc.) # Pass through as-is since these are Anthropic-native content types elif m.get("type", "") == "server_tool_use": @@ -2945,9 +2701,7 @@ def anthropic_messages_pt( elif ( "content" in assistant_content_block and isinstance(assistant_content_block["content"], str) - and assistant_content_block[ - "content" - ] # don't pass empty text blocks. anthropic api raises errors. + and assistant_content_block["content"] # don't pass empty text blocks. anthropic api raises errors. ): _anthropic_text_content_element = AnthropicMessagesTextParam( type="text", @@ -2960,29 +2714,19 @@ def anthropic_messages_pt( ) if "cache_control" in _content_element: - _anthropic_text_content_element["cache_control"] = ( - _content_element["cache_control"] - ) + _anthropic_text_content_element["cache_control"] = _content_element["cache_control"] assistant_content.append(_anthropic_text_content_element) - if ( - assistant_tool_calls is not None - ): # support assistant tool invoke conversion + if assistant_tool_calls is not None: # support assistant tool invoke conversion # Get web_search_results and tool_results from provider_specific_fields # for server_tool_use reconstruction. # Fixes: https://github.com/BerriAI/litellm/issues/17737 - _provider_specific_fields_raw = assistant_content_block.get( - "provider_specific_fields" - ) + _provider_specific_fields_raw = assistant_content_block.get("provider_specific_fields") _provider_specific_fields: Dict[str, Any] = {} if isinstance(_provider_specific_fields_raw, dict): - _provider_specific_fields = cast( - Dict[str, Any], _provider_specific_fields_raw - ) - _web_search_results = _provider_specific_fields.get( - "web_search_results" - ) + _provider_specific_fields = cast(Dict[str, Any], _provider_specific_fields_raw) + _web_search_results = _provider_specific_fields.get("web_search_results") _tool_results = _provider_specific_fields.get("tool_results") tool_invoke_results = convert_to_anthropic_tool_invoke( assistant_tool_calls, @@ -2994,27 +2738,19 @@ def anthropic_messages_pt( # This can happen when merging history that already contains the tool calls for item in tool_invoke_results: # tool_use items are typically dicts, but handle objects just in case - item_id = ( - item.get("id") - if isinstance(item, dict) - else getattr(item, "id", None) - ) + item_id = item.get("id") if isinstance(item, dict) else getattr(item, "id", None) if item_id: if item_id in unique_tool_ids: continue unique_tool_ids.add(item_id) - assistant_content.append( - cast(AnthropicMessagesAssistantMessageValues, item) - ) + assistant_content.append(cast(AnthropicMessagesAssistantMessageValues, item)) assistant_function_call = assistant_content_block.get("function_call") if assistant_function_call is not None: - assistant_content.extend( - convert_function_to_anthropic_tool_invoke(assistant_function_call) - ) + assistant_content.extend(convert_function_to_anthropic_tool_invoke(assistant_function_call)) msg_i += 1 @@ -3034,9 +2770,7 @@ def anthropic_messages_pt( elif isinstance(new_messages[-1]["content"], list): for content in new_messages[-1]["content"]: if isinstance(content, dict) and content["type"] == "text": - content["text"] = content[ - "text" - ].rstrip() # no trailing whitespace for final assistant message + content["text"] = content["text"].rstrip() # no trailing whitespace for final assistant message return new_messages @@ -3189,11 +2923,7 @@ def convert_openai_message_to_cohere_tool_result( msg_tool_call_id = message.get("tool_call_id", None) for tool in tools: prev_tool_call_id = tool.get("id", None) - if ( - msg_tool_call_id - and prev_tool_call_id - and msg_tool_call_id == prev_tool_call_id - ): + if msg_tool_call_id and prev_tool_call_id and msg_tool_call_id == prev_tool_call_id: name = tool.get("function", {}).get("name", "") arguments_str = tool.get("function", {}).get("arguments", "") if arguments_str is not None and len(arguments_str) > 0: @@ -3262,14 +2992,8 @@ def convert_to_cohere_tool_invoke(tool_calls: list) -> List[ToolCallObject]: cohere_tool_invoke: List[ToolCallObject] = [ { - "name": get_attribute_or_key( - get_attribute_or_key(tool, "function"), "name" - ), - "parameters": json.loads( - get_attribute_or_key( - get_attribute_or_key(tool, "function"), "arguments" - ) - ), + "name": get_attribute_or_key(get_attribute_or_key(tool, "function"), "name"), + "parameters": json.loads(get_attribute_or_key(get_attribute_or_key(tool, "function"), "arguments")), } for tool in tool_calls if get_attribute_or_key(tool, "type") == "function" @@ -3301,14 +3025,9 @@ def cohere_messages_pt_v2( ## GET MOST RECENT MESSAGE most_recent_message = messages.pop(-1) returned_message: Union[ToolResultObject, str] = "" - if ( - most_recent_message.get("role", "") is not None - and most_recent_message["role"] == "tool" - ): + if most_recent_message.get("role", "") is not None and most_recent_message["role"] == "tool": # tool result - returned_message = convert_openai_message_to_cohere_tool_result( - most_recent_message, tool_calls - ) + returned_message = convert_openai_message_to_cohere_tool_result(most_recent_message, tool_calls) else: content: Union[str, List] = most_recent_message.get("content") if isinstance(content, str): @@ -3353,35 +3072,23 @@ def cohere_messages_pt_v2( msg_i += 1 if len(system_content) > 0: - new_messages.append( - ChatHistorySystem(role="SYSTEM", message=system_content) - ) + new_messages.append(ChatHistorySystem(role="SYSTEM", message=system_content)) assistant_content: str = "" assistant_tool_calls: List[ToolCallObject] = [] ## MERGE CONSECUTIVE ASSISTANT CONTENT ## while msg_i < len(messages) and messages[msg_i]["role"] == "assistant": - if messages[msg_i].get("content", None) is not None and isinstance( - messages[msg_i]["content"], list - ): + if messages[msg_i].get("content", None) is not None and isinstance(messages[msg_i]["content"], list): for m in messages[msg_i]["content"]: if m.get("type", "") == "text": assistant_content += m["text"] - elif messages[msg_i].get("content") is not None and isinstance( - messages[msg_i]["content"], str - ): + elif messages[msg_i].get("content") is not None and isinstance(messages[msg_i]["content"], str): assistant_content += messages[msg_i]["content"] - if messages[msg_i].get( - "tool_calls", [] - ): # support assistant tool invoke conversion - assistant_tool_calls.extend( - convert_to_cohere_tool_invoke(messages[msg_i]["tool_calls"]) - ) + if messages[msg_i].get("tool_calls", []): # support assistant tool invoke conversion + assistant_tool_calls.extend(convert_to_cohere_tool_invoke(messages[msg_i]["tool_calls"])) if messages[msg_i].get("function_call"): - assistant_tool_calls.extend( - convert_to_cohere_tool_invoke(messages[msg_i]["function_call"]) - ) + assistant_tool_calls.extend(convert_to_cohere_tool_invoke(messages[msg_i]["function_call"])) msg_i += 1 @@ -3397,18 +3104,12 @@ def cohere_messages_pt_v2( ## MERGE CONSECUTIVE TOOL RESULTS tool_results: List[ToolResultObject] = [] while msg_i < len(messages) and messages[msg_i]["role"] in tool_message_types: - tool_results.append( - convert_openai_message_to_cohere_tool_result( - messages[msg_i], tool_calls - ) - ) + tool_results.append(convert_openai_message_to_cohere_tool_result(messages[msg_i], tool_calls)) msg_i += 1 if len(tool_results) > 0: - new_messages.append( - ChatHistoryToolResult(role="TOOL", tool_results=tool_results) - ) + new_messages.append(ChatHistoryToolResult(role="TOOL", tool_results=tool_results)) if msg_i == init_msg_i: # prevent infinite loops raise litellm.BadRequestError( @@ -3427,9 +3128,7 @@ def cohere_message_pt(messages: list): for message in messages: # check if this is a tool_call result if message["role"] == "tool": - tool_result = convert_openai_message_to_cohere_tool_result( - message, tool_calls=tool_calls - ) + tool_result = convert_openai_message_to_cohere_tool_result(message, tool_calls=tool_calls) tool_results.append(tool_result) elif message.get("content"): prompt += message["content"] + "\n\n" @@ -3456,9 +3155,7 @@ def amazon_titan_pt( prompt += f"{AmazonTitanConstants.HUMAN_PROMPT.value}{message['content']}" else: prompt += f"{AmazonTitanConstants.AI_PROMPT.value}{message['content']}" - if ( - idx == 0 and message["role"] == "assistant" - ): # ensure the prompt always starts with `\n\nHuman: ` + if idx == 0 and message["role"] == "assistant": # ensure the prompt always starts with `\n\nHuman: ` prompt = f"{AmazonTitanConstants.HUMAN_PROMPT.value}" + prompt if messages[-1]["role"] != "assistant": prompt += f"{AmazonTitanConstants.AI_PROMPT.value}" @@ -3481,9 +3178,7 @@ def _load_image_from_url(image_url): # Check the response's content type to ensure it is an image content_type = response.headers.get("content-type") if not content_type or "image" not in content_type: - raise ValueError( - f"URL does not point to a valid image (content-type: {content_type})" - ) + raise ValueError(f"URL does not point to a valid image (content-type: {content_type})") # Load the image from the response content return Image.open(BytesIO(response.content)) @@ -3534,9 +3229,7 @@ def _gemini_vision_convert_messages(messages: list): try: from PIL import Image except Exception: - raise Exception( - "gemini image conversion failed please run `pip install Pillow`" - ) + raise Exception("gemini image conversion failed please run `pip install Pillow`") if "base64" in img: # Case 2: Base64 image data @@ -3582,9 +3275,7 @@ def gemini_text_image_pt(messages: list): try: pass # type: ignore except Exception: - raise Exception( - "Importing google.generativeai failed, please run 'pip install -q google-generativeai" - ) + raise Exception("Importing google.generativeai failed, please run 'pip install -q google-generativeai") prompt = "" images = [] @@ -3682,9 +3373,7 @@ class BedrockImageProcessor: """Handles both sync and async image processing for Bedrock conversations.""" @staticmethod - def _post_call_image_processing( - response: httpx.Response, image_url: str = "" - ) -> Tuple[str, str]: + def _post_call_image_processing(response: httpx.Response, image_url: str = "") -> Tuple[str, str]: # Check the response's content type to ensure it is an image content_type = response.headers.get("content-type") @@ -3713,9 +3402,7 @@ class BedrockImageProcessor: response = await async_safe_get(client, image_url) response.raise_for_status() # Raise an exception for HTTP errors - return BedrockImageProcessor._post_call_image_processing( - response, image_url - ) + return BedrockImageProcessor._post_call_image_processing(response, image_url) except Exception as e: raise e @@ -3728,9 +3415,7 @@ class BedrockImageProcessor: response = safe_get(client, image_url) response.raise_for_status() # Raise an exception for HTTP errors - return BedrockImageProcessor._post_call_image_processing( - response, image_url - ) + return BedrockImageProcessor._post_call_image_processing(response, image_url) except Exception as e: raise e @@ -3757,22 +3442,14 @@ class BedrockImageProcessor: def _validate_format(mime_type: str, image_format: str) -> str: """Validate image format and mime type for both images and documents.""" - supported_image_formats = ( - litellm.AmazonConverseConfig().get_supported_image_types() - ) - supported_doc_formats = ( - litellm.AmazonConverseConfig().get_supported_document_types() - ) - supported_video_formats = ( - litellm.AmazonConverseConfig().get_supported_video_types() - ) + supported_image_formats = litellm.AmazonConverseConfig().get_supported_image_types() + supported_doc_formats = litellm.AmazonConverseConfig().get_supported_document_types() + supported_video_formats = litellm.AmazonConverseConfig().get_supported_video_types() document_types = ["application", "text"] is_document = any(mime_type.startswith(doc_type) for doc_type in document_types) - supported_image_and_video_formats: List[str] = ( - supported_video_formats + supported_image_formats - ) + supported_image_and_video_formats: List[str] = supported_video_formats + supported_image_formats if is_document: return BedrockImageProcessor._get_document_format( @@ -3810,9 +3487,7 @@ class BedrockImageProcessor: """ valid_extensions: Optional[List[str]] = None potential_extensions = mimetypes.guess_all_extensions(mime_type, strict=False) - valid_extensions = [ - ext[1:] for ext in potential_extensions if ext[1:] in supported_doc_formats - ] + valid_extensions = [ext[1:] for ext in potential_extensions if ext[1:] in supported_doc_formats] # Fallback to types/files.py if mimetypes doesn't return valid extensions ################# @@ -3837,22 +3512,15 @@ class BedrockImageProcessor: return valid_extensions[0] @staticmethod - def _create_bedrock_block( - image_bytes: str, mime_type: str, image_format: str - ) -> BedrockContentBlock: + def _create_bedrock_block(image_bytes: str, mime_type: str, image_format: str) -> BedrockContentBlock: """Create appropriate Bedrock content block based on mime type.""" _blob = BedrockSourceBlock(bytes=image_bytes) document_types = ["application", "text"] is_document = any(mime_type.startswith(doc_type) for doc_type in document_types) - supported_video_formats = ( - litellm.AmazonConverseConfig().get_supported_video_types() - ) - is_video = any( - image_format.startswith(video_type) - for video_type in supported_video_formats - ) + supported_video_formats = litellm.AmazonConverseConfig().get_supported_video_types() + is_video = any(image_format.startswith(video_type) for video_type in supported_video_formats) HASH_SAMPLE_BYTES = 64 * 1024 # hash up to 64 KB of data @@ -3873,9 +3541,7 @@ class BedrockImageProcessor: # --- Compute deterministic hash (sample + total length) --- hasher = hashlib.sha256() hasher.update(sample) - hasher.update( - str(len(normalized)).encode("utf-8") - ) # include full length for uniqueness + hasher.update(str(len(normalized)).encode("utf-8")) # include full length for uniqueness full_hash = hasher.hexdigest() content_hash = full_hash[:16] # short deterministic ID @@ -3890,18 +3556,12 @@ class BedrockImageProcessor: ) ) elif is_video: - return BedrockContentBlock( - video=BedrockVideoBlock(source=_blob, format=image_format) - ) + return BedrockContentBlock(video=BedrockVideoBlock(source=_blob, format=image_format)) else: - return BedrockContentBlock( - image=BedrockImageBlock(source=_blob, format=image_format) - ) + return BedrockContentBlock(image=BedrockImageBlock(source=_blob, format=image_format)) @classmethod - def process_image_sync( - cls, image_url: str, format: Optional[str] = None - ) -> BedrockContentBlock: + def process_image_sync(cls, image_url: str, format: Optional[str] = None) -> BedrockContentBlock: """Synchronous image processing.""" if "base64" in image_url: @@ -3910,9 +3570,7 @@ class BedrockImageProcessor: img_bytes, mime_type = BedrockImageProcessor.get_image_details(image_url) image_format = mime_type.split("/")[1] else: - raise ValueError( - "Unsupported image type. Expected either image url or base64 encoded string" - ) + raise ValueError("Unsupported image type. Expected either image url or base64 encoded string") if format: mime_type = format @@ -3922,22 +3580,16 @@ class BedrockImageProcessor: return cls._create_bedrock_block(img_bytes, mime_type, image_format) @classmethod - async def process_image_async( - cls, image_url: str, format: Optional[str] - ) -> BedrockContentBlock: + async def process_image_async(cls, image_url: str, format: Optional[str]) -> BedrockContentBlock: """Asynchronous image processing.""" if "base64" in image_url: img_bytes, mime_type, image_format = cls._parse_base64_image(image_url) elif "http://" in image_url or "https://" in image_url: - img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async( - image_url - ) + img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url) image_format = mime_type.split("/")[1] else: - raise ValueError( - "Unsupported image type. Expected either image url or base64 encoded string" - ) + raise ValueError("Unsupported image type. Expected either image url or base64 encoded string") if format: # override with user-defined params mime_type = format @@ -4018,45 +3670,29 @@ def _convert_to_bedrock_tool_call_invoke( if parsed_objects: # First object keeps the original tool id. for obj_idx, obj in enumerate(parsed_objects): - block_id = ( - tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}" - ) - bedrock_tool = BedrockToolUseBlock( - input=obj, name=name, toolUseId=block_id - ) - _parts_list.append( - BedrockContentBlock(toolUse=bedrock_tool) - ) + block_id = tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}" + bedrock_tool = BedrockToolUseBlock(input=obj, name=name, toolUseId=block_id) + _parts_list.append(BedrockContentBlock(toolUse=bedrock_tool)) # cache_control applies to the whole original # tool call; attach after the last split block. if tool.get("cache_control", None) is not None: - _parts_list.append( - BedrockContentBlock( - cachePoint=CachePointBlock(type="default") - ) - ) + _parts_list.append(BedrockContentBlock(cachePoint=CachePointBlock(type="default"))) continue # Fallback: no objects extracted — use empty dict. arguments_dict = {} - bedrock_tool = BedrockToolUseBlock( - input=arguments_dict, name=name, toolUseId=tool_id - ) + bedrock_tool = BedrockToolUseBlock(input=arguments_dict, name=name, toolUseId=tool_id) bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool) _parts_list.append(bedrock_content_block) # Check for cache_control and add a separate cachePoint block if tool.get("cache_control", None) is not None: - cache_point_block = BedrockContentBlock( - cachePoint=CachePointBlock(type="default") - ) + cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default")) _parts_list.append(cache_point_block) return _parts_list except Exception as e: raise Exception( - "Unable to convert openai tool calls={} to bedrock tool calls. Received error={}".format( - tool_calls, str(e) - ) + "Unable to convert openai tool calls={} to bedrock tool calls. Received error={}".format(tool_calls, str(e)) ) @@ -4067,17 +3703,12 @@ def _append_bedrock_tool_result_media_block( content_type: str, ) -> None: if "image" in processed_block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(image=processed_block["image"]) - ) + tool_result_content_blocks.append(BedrockToolResultContentBlock(image=processed_block["image"])) elif "document" in processed_block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(document=processed_block["document"]) - ) + tool_result_content_blocks.append(BedrockToolResultContentBlock(document=processed_block["document"])) else: verbose_logger.warning( - "Bedrock Converse: unrecognized BedrockContentBlock keys " - "%s for %s tool-result block %s; dropping.", + "Bedrock Converse: unrecognized BedrockContentBlock keys %s for %s tool-result block %s; dropping.", list(processed_block.keys()), content_type, content, @@ -4098,9 +3729,7 @@ def _append_bedrock_tool_result_image_url_block( image_url=image_url, format=format, ) - _append_bedrock_tool_result_media_block( - tool_result_content_blocks, processed_block, content, "image_url" - ) + _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "image_url") def _append_bedrock_tool_result_file_block( @@ -4122,9 +3751,7 @@ def _append_bedrock_tool_result_file_block( image_url=cast(str, file_id or file_data), format=file_obj.get("format"), ) - _append_bedrock_tool_result_media_block( - tool_result_content_blocks, processed_block, content, "file" - ) + _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "file") def _parse_bedrock_tool_result_content_list( @@ -4133,13 +3760,9 @@ def _parse_bedrock_tool_result_content_list( tool_result_content_blocks: List[BedrockToolResultContentBlock] = [] for content in content_list: if content["type"] == "text": - tool_result_content_blocks.append( - BedrockToolResultContentBlock(text=content["text"]) - ) + tool_result_content_blocks.append(BedrockToolResultContentBlock(text=content["text"])) elif content["type"] == "image_url": - _append_bedrock_tool_result_image_url_block( - tool_result_content_blocks, content - ) + _append_bedrock_tool_result_image_url_block(tool_result_content_blocks, content) elif content["type"] == "file": _append_bedrock_tool_result_file_block(tool_result_content_blocks, content) return tool_result_content_blocks @@ -4161,9 +3784,7 @@ def _build_bedrock_tool_result_content_blocks( if not isinstance(result, dict): continue tool_result_content_blocks.append( - BedrockToolResultContentBlock( - searchResult=cast(SearchResultBlock, result) - ) + BedrockToolResultContentBlock(searchResult=cast(SearchResultBlock, result)) ) if tool_result_content_blocks: return tool_result_content_blocks, True @@ -4219,16 +3840,12 @@ def _convert_to_bedrock_tool_call_result( """ - """ - tool_result_content_blocks, used_search_results = ( - _build_bedrock_tool_result_content_blocks(message) - ) + tool_result_content_blocks, used_search_results = _build_bedrock_tool_result_content_blocks(message) message.get("name", "") id = str(message.get("tool_call_id", str(uuid.uuid4()))) - tool_result = BedrockToolResultBlock( - content=tool_result_content_blocks, toolUseId=id - ) + tool_result = BedrockToolResultBlock(content=tool_result_content_blocks, toolUseId=id) if used_search_results: tool_result["status"] = cast(Literal["success"], "success") @@ -4369,9 +3986,7 @@ def _sort_bedrock_assistant_content_blocks( def _insert_assistant_continue_message( messages: List[BedrockMessageBlock], - assistant_continue_message: Optional[ - Union[str, ChatCompletionAssistantMessage] - ] = None, + assistant_continue_message: Optional[Union[str, ChatCompletionAssistantMessage]] = None, ) -> List[BedrockMessageBlock]: """ Add dummy message between user/tool result blocks. @@ -4395,9 +4010,7 @@ def _insert_assistant_continue_message( ) ) elif litellm.modify_params: - text = convert_content_list_to_str( - cast(ChatCompletionAssistantMessage, DEFAULT_ASSISTANT_CONTINUE_MESSAGE) - ) + text = convert_content_list_to_str(cast(ChatCompletionAssistantMessage, DEFAULT_ASSISTANT_CONTINUE_MESSAGE)) messages.append( BedrockMessageBlock( role="assistant", @@ -4422,9 +4035,7 @@ def get_user_message_block_or_continue_message( content_block = message.get("content", None) # Handle None case - if content_block is None or ( - user_continue_message is None and litellm.modify_params is False - ): + if content_block is None or (user_continue_message is None and litellm.modify_params is False): return skip_empty_text_blocks(message=message) # Handle string case @@ -4475,9 +4086,7 @@ def get_user_message_block_or_continue_message( def return_assistant_continue_message( - assistant_continue_message: Optional[ - Union[str, ChatCompletionAssistantMessage] - ] = None, + assistant_continue_message: Optional[Union[str, ChatCompletionAssistantMessage]] = None, ) -> ChatCompletionAssistantMessage: if assistant_continue_message and isinstance(assistant_continue_message, str): return ChatCompletionAssistantMessage( @@ -4500,11 +4109,7 @@ def _skip_empty_dict_blocks(blocks: List[dict]) -> List[dict]: Returns: Filtered list of non-empty text blocks """ - return [ - item - for item in blocks - if not (item.get("type") == "text" and not item.get("text", "").strip()) - ] + return [item for item in blocks if not (item.get("type") == "text" and not item.get("text", "").strip())] @overload @@ -4542,9 +4147,7 @@ def skip_empty_text_blocks( modified_message["content"] = None # user message content cannot be None return modified_message elif isinstance(content_block, list): - modified_content_block = _skip_empty_dict_blocks( - cast(List[dict], content_block) - ) + modified_content_block = _skip_empty_dict_blocks(cast(List[dict], content_block)) # If no content remains and it's an assistant message, set content to None if not modified_content_block and message["role"] == "assistant": @@ -4572,9 +4175,7 @@ def skip_empty_text_blocks( def process_empty_text_blocks( message: ChatCompletionAssistantMessage, - assistant_continue_message: Optional[ - Union[str, ChatCompletionAssistantMessage] - ] = None, + assistant_continue_message: Optional[Union[str, ChatCompletionAssistantMessage]] = None, ) -> ChatCompletionAssistantMessage: modified_content_block = message.get("content", None) ## BASE CASE ## @@ -4582,14 +4183,9 @@ def process_empty_text_blocks( return message # Check if all items are empty text blocks - if all( - item["type"] == "text" and not item["text"].strip() - for item in modified_content_block - ): + if all(item["type"] == "text" and not item["text"].strip() for item in modified_content_block): # Replace with a single continue message - _assistant_continue_message = return_assistant_continue_message( - assistant_continue_message - ) + _assistant_continue_message = return_assistant_continue_message(assistant_continue_message) modified_content_block = [ { "type": "text", @@ -4599,9 +4195,7 @@ def process_empty_text_blocks( else: # Filter out only empty text blocks, keeping non-empty text and other block types modified_content_block = [ - item - for item in modified_content_block - if not (item["type"] == "text" and not item["text"].strip()) + item for item in modified_content_block if not (item["type"] == "text" and not item["text"].strip()) ] modified_message = message.copy() @@ -4614,9 +4208,7 @@ def process_empty_text_blocks( def get_assistant_message_block_or_continue_message( message: ChatCompletionAssistantMessage, - assistant_continue_message: Optional[ - Union[str, ChatCompletionAssistantMessage] - ] = None, + assistant_continue_message: Optional[Union[str, ChatCompletionAssistantMessage]] = None, ) -> ChatCompletionAssistantMessage: """ Returns the user content block @@ -4627,9 +4219,7 @@ def get_assistant_message_block_or_continue_message( content_block = message.get("content", None) # Handle Base case - if content_block is None or ( - assistant_continue_message is None and litellm.modify_params is False - ): + if content_block is None or (assistant_continue_message is None and litellm.modify_params is False): return skip_empty_text_blocks(message=message) # Handle string case @@ -4655,9 +4245,7 @@ def get_assistant_message_block_or_continue_message( } ], """ - return process_empty_text_blocks( - message=message, assistant_continue_message=assistant_continue_message - ) + return process_empty_text_blocks(message=message, assistant_continue_message=assistant_continue_message) # Handle unsupported type raise ValueError(f"Unsupported content type: {type(content_block)}") @@ -4679,8 +4267,7 @@ class BedrockConverseMessagesProcessor: messages.append(DEFAULT_USER_CONTINUE_MESSAGE) else: raise litellm.BadRequestError( - message=BAD_MESSAGE_ERROR_STR - + "bedrock requires at least one non-system message", + message=BAD_MESSAGE_ERROR_STR + "bedrock requires at least one non-system message", model=model, llm_provider=llm_provider, ) @@ -4708,9 +4295,7 @@ class BedrockConverseMessagesProcessor: model: str, llm_provider: str, user_continue_message: Optional[ChatCompletionUserMessage] = None, - assistant_continue_message: Optional[ - Union[str, ChatCompletionAssistantMessage] - ] = None, + assistant_continue_message: Optional[Union[str, ChatCompletionAssistantMessage]] = None, ) -> List[BedrockMessageBlock]: contents: List[BedrockMessageBlock] = [] msg_i = 0 @@ -4737,9 +4322,7 @@ class BedrockConverseMessagesProcessor: _parts.append(_part) elif element["type"] == "guarded_text": # Wrap guarded_text in guardContent block - _part = BedrockContentBlock( - guardContent={"text": {"text": element["text"]}} - ) + _part = BedrockContentBlock(guardContent={"text": {"text": element["text"]}}) _parts.append(_part) elif element["type"] in ("grounding_source", "query"): # Contextual grounding tags are guardrail metadata; the @@ -4764,29 +4347,19 @@ class BedrockConverseMessagesProcessor: ) _parts.append(_part) elif element["type"] == "document": - _part = BedrockConverseMessagesProcessor._process_document_message( - element - ) + _part = BedrockConverseMessagesProcessor._process_document_message(element) _parts.append(_part) - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - message_block=cast( - OpenAIMessageContentListBlock, element - ), - block_type="content_block", - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + message_block=cast(OpenAIMessageContentListBlock, element), + block_type="content_block", ) if _cache_point_block is not None: _parts.append(_cache_point_block) user_content.extend(_parts) - elif message_block["content"] and isinstance( - message_block["content"], str - ): + elif message_block["content"] and isinstance(message_block["content"], str): _part = BedrockContentBlock(text=messages[msg_i]["content"]) - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - message_block, block_type="content_block" - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + message_block, block_type="content_block" ) user_content.append(_part) if _cache_point_block is not None: @@ -4795,27 +4368,20 @@ class BedrockConverseMessagesProcessor: msg_i += 1 if user_content: if len(contents) > 0 and contents[-1]["role"] == "user": - if ( - assistant_continue_message is not None - or litellm.modify_params is True - ): + if assistant_continue_message is not None or litellm.modify_params is True: # if last message was a 'user' message, then add a dummy assistant message (bedrock requires alternating roles) contents = _insert_assistant_continue_message( messages=contents, assistant_continue_message=assistant_continue_message, ) - contents.append( - BedrockMessageBlock(role="user", content=user_content) - ) + contents.append(BedrockMessageBlock(role="user", content=user_content)) else: verbose_logger.warning( "Potential consecutive user/tool blocks. Trying to merge. If error occurs, please set a 'assistant_continue_message' or set 'modify_params=True' to insert a dummy assistant message for bedrock calls." ) contents[-1]["content"].extend(user_content) else: - contents.append( - BedrockMessageBlock(role="user", content=user_content) - ) + contents.append(BedrockMessageBlock(role="user", content=user_content)) ## MERGE CONSECUTIVE TOOL CALL MESSAGES ## tool_content: List[BedrockContentBlock] = [] @@ -4833,18 +4399,13 @@ class BedrockConverseMessagesProcessor: # Check for content-level cache_control in list content elif isinstance(current_message.get("content"), list): for content_element in current_message["content"]: - if ( - isinstance(content_element, dict) - and content_element.get("cache_control", None) is not None - ): + if isinstance(content_element, dict) and content_element.get("cache_control", None) is not None: has_cache_control = True break # Add a separate cachePoint block if cache_control is present if has_cache_control: - cache_point_block = BedrockContentBlock( - cachePoint=CachePointBlock(type="default") - ) + cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default")) tool_content.append(cache_point_block) msg_i += 1 @@ -4853,35 +4414,26 @@ class BedrockConverseMessagesProcessor: if tool_content: # if last message was a 'user' message, then add a blank assistant message (bedrock requires alternating roles) if len(contents) > 0 and contents[-1]["role"] == "user": - if ( - assistant_continue_message is not None - or litellm.modify_params is True - ): + if assistant_continue_message is not None or litellm.modify_params is True: # if last message was a 'user' message, then add a dummy assistant message (bedrock requires alternating roles) contents = _insert_assistant_continue_message( messages=contents, assistant_continue_message=assistant_continue_message, ) - contents.append( - BedrockMessageBlock(role="user", content=tool_content) - ) + contents.append(BedrockMessageBlock(role="user", content=tool_content)) else: verbose_logger.warning( "Potential consecutive user/tool blocks. Trying to merge. If error occurs, please set a 'assistant_continue_message' or set 'modify_params=True' to insert a dummy assistant message for bedrock calls." ) contents[-1]["content"].extend(tool_content) else: - contents.append( - BedrockMessageBlock(role="user", content=tool_content) - ) + contents.append(BedrockMessageBlock(role="user", content=tool_content)) assistant_content: List[BedrockContentBlock] = [] ## MERGE CONSECUTIVE ASSISTANT CONTENT ## while msg_i < len(messages) and messages[msg_i]["role"] == "assistant": - assistant_message_block = ( - get_assistant_message_block_or_continue_message( - message=messages[msg_i], - assistant_continue_message=assistant_continue_message, - ) + assistant_message_block = get_assistant_message_block_or_continue_message( + message=messages[msg_i], + assistant_continue_message=assistant_continue_message, ) _assistant_content = assistant_message_block.get("content", None) thinking_blocks = cast( @@ -4890,97 +4442,75 @@ class BedrockConverseMessagesProcessor: ) if thinking_blocks is not None: - converted_thinking_blocks = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks + converted_thinking_blocks = ( + BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( + thinking_blocks + ) ) assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( thinking_blocks=converted_thinking_blocks, assistant_parts=assistant_content, ) - if _assistant_content is not None and isinstance( - _assistant_content, list - ): + if _assistant_content is not None and isinstance(_assistant_content, list): assistants_parts: List[BedrockContentBlock] = [] for element in _assistant_content: if isinstance(element, dict): if element["type"] == "thinking": thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[ - cast(ChatCompletionThinkingBlock, element) - ] + thinking_blocks=[cast(ChatCompletionThinkingBlock, element)] ) - assistants_parts = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( - thinking_blocks=thinking_block, - assistant_parts=assistants_parts, + assistants_parts = ( + BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( + thinking_blocks=thinking_block, + assistant_parts=assistants_parts, + ) ) elif element["type"] == "text": # Skip completely empty strings to avoid blank content blocks if element.get("text", "").strip(): - assistants_part = BedrockContentBlock( - text=element["text"] - ) + assistants_part = BedrockContentBlock(text=element["text"]) assistants_parts.append(assistants_part) elif element["type"] == "image_url": if isinstance(element["image_url"], dict): image_url = element["image_url"]["url"] else: image_url = element["image_url"] - assistants_part = ( - await BedrockImageProcessor.process_image_async( # type: ignore - image_url=image_url - ) + assistants_part = await BedrockImageProcessor.process_image_async( # type: ignore + image_url=image_url ) assistants_parts.append(assistants_part) # Add cache point block for assistant content elements - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - message_block=cast( - OpenAIMessageContentListBlock, element - ), - block_type="content_block", - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + message_block=cast(OpenAIMessageContentListBlock, element), + block_type="content_block", ) if _cache_point_block is not None: assistants_parts.append(_cache_point_block) assistant_content.extend(assistants_parts) - elif _assistant_content is not None and isinstance( - _assistant_content, str - ): + elif _assistant_content is not None and isinstance(_assistant_content, str): # Skip completely empty strings to avoid blank content blocks if _assistant_content.strip(): - assistant_content.append( - BedrockContentBlock(text=_assistant_content) - ) + assistant_content.append(BedrockContentBlock(text=_assistant_content)) # If content is empty/whitespace, skip it (don't add a placeholder) # Add cache point block for assistant string content - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - assistant_message_block, block_type="content_block" - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + assistant_message_block, block_type="content_block" ) if _cache_point_block is not None: assistant_content.append(_cache_point_block) _tool_calls = assistant_message_block.get("tool_calls", []) if _tool_calls: - assistant_content.extend( - _convert_to_bedrock_tool_call_invoke(_tool_calls) - ) + assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls)) msg_i += 1 - assistant_content = _deduplicate_bedrock_content_blocks( - assistant_content, "toolUse" - ) - assistant_content = _sort_bedrock_assistant_content_blocks( - assistant_content - ) + assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") + assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content) if assistant_content: - contents.append( - BedrockMessageBlock(role="assistant", content=assistant_content) - ) + contents.append(BedrockMessageBlock(role="assistant", content=assistant_content)) if msg_i == init_msg_i: # prevent infinite loops raise litellm.BadRequestError( @@ -5007,9 +4537,7 @@ class BedrockConverseMessagesProcessor: reasoning_content_block = BedrockConverseReasoningContentBlock( reasoningText=text_block, ) - bedrock_content_block = BedrockContentBlock( - reasoningContent=reasoning_content_block - ) + bedrock_content_block = BedrockContentBlock(reasoningContent=reasoning_content_block) reasoning_content_blocks.append(bedrock_content_block) return reasoning_content_blocks @@ -5027,16 +4555,12 @@ class BedrockConverseMessagesProcessor: if file_data is None and file_id is None: raise litellm.BadRequestError( - message="file_data and file_id cannot both be None. Got={}".format( - message - ), + message="file_data and file_id cannot both be None. Got={}".format(message), model="", llm_provider="bedrock", ) format = file_message.get("format") - return BedrockImageProcessor.process_image_sync( - image_url=cast(str, file_id or file_data), format=format - ) + return BedrockImageProcessor.process_image_sync(image_url=cast(str, file_id or file_data), format=format) @staticmethod async def _async_process_file_message( @@ -5054,15 +4578,11 @@ class BedrockConverseMessagesProcessor: format = file_message.get("format") if file_data is None and file_id is None: raise litellm.BadRequestError( - message="file_data and file_id cannot both be None. Got={}".format( - message - ), + message="file_data and file_id cannot both be None. Got={}".format(message), model="", llm_provider="bedrock", ) - return await BedrockImageProcessor.process_image_async( - image_url=cast(str, file_id or file_data), format=format - ) + return await BedrockImageProcessor.process_image_async(image_url=cast(str, file_id or file_data), format=format) @staticmethod def _process_document_message(element: dict) -> BedrockContentBlock: @@ -5080,9 +4600,7 @@ class BedrockConverseMessagesProcessor: ) media_type: str = source["media_type"] data: str = source["data"] - doc_format = BedrockImageProcessor._validate_format( - mime_type=media_type, image_format=media_type.split("/")[1] - ) + doc_format = BedrockImageProcessor._validate_format(mime_type=media_type, image_format=media_type.split("/")[1]) # Deterministic name using the same hashing pattern as _create_bedrock_block HASH_SAMPLE_BYTES = 64 * 1024 @@ -5118,11 +4636,7 @@ class BedrockConverseMessagesProcessor: filtered_thinking_blocks = [] for block in thinking_blocks: reasoning_content = block.get("reasoningContent", None) - reasoning_text = ( - reasoning_content.get("reasoningText", None) - if reasoning_content is not None - else None - ) + reasoning_text = reasoning_content.get("reasoningText", None) if reasoning_content is not None else None if reasoning_text and not reasoning_text.get("signature"): reasoning_text_text = reasoning_text["text"] if reasoning_text_text.strip(): @@ -5140,9 +4654,7 @@ def _bedrock_converse_messages_pt( model: str, llm_provider: str, user_continue_message: Optional[ChatCompletionUserMessage] = None, - assistant_continue_message: Optional[ - Union[str, ChatCompletionAssistantMessage] - ] = None, + assistant_continue_message: Optional[Union[str, ChatCompletionAssistantMessage]] = None, ) -> List[BedrockMessageBlock]: """ Converts given messages from OpenAI format to Bedrock format @@ -5177,9 +4689,7 @@ def _bedrock_converse_messages_pt( _parts.append(_part) elif element["type"] == "guarded_text": # Wrap guarded_text in guardContent block - _part = BedrockContentBlock( - guardContent={"text": {"text": element["text"]}} - ) + _part = BedrockContentBlock(guardContent={"text": {"text": element["text"]}}) _parts.append(_part) elif element["type"] in ("grounding_source", "query"): # Contextual grounding tags are guardrail metadata; the @@ -5200,34 +4710,24 @@ def _bedrock_converse_messages_pt( ) _parts.append(_part) # type: ignore elif element["type"] == "file": - _part = ( - BedrockConverseMessagesProcessor._process_file_message( - message=cast(ChatCompletionFileObject, element) - ) + _part = BedrockConverseMessagesProcessor._process_file_message( + message=cast(ChatCompletionFileObject, element) ) _parts.append(_part) elif element["type"] == "document": - _part = BedrockConverseMessagesProcessor._process_document_message( - element - ) + _part = BedrockConverseMessagesProcessor._process_document_message(element) _parts.append(_part) - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - message_block=cast( - OpenAIMessageContentListBlock, element - ), - block_type="content_block", - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + message_block=cast(OpenAIMessageContentListBlock, element), + block_type="content_block", ) if _cache_point_block is not None: _parts.append(_cache_point_block) user_content.extend(_parts) elif message_block["content"] and isinstance(message_block["content"], str): _part = BedrockContentBlock(text=messages[msg_i]["content"]) - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - message_block, block_type="content_block" - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + message_block, block_type="content_block" ) user_content.append(_part) if _cache_point_block is not None: @@ -5236,18 +4736,13 @@ def _bedrock_converse_messages_pt( msg_i += 1 if user_content: if len(contents) > 0 and contents[-1]["role"] == "user": - if ( - assistant_continue_message is not None - or litellm.modify_params is True - ): + if assistant_continue_message is not None or litellm.modify_params is True: # if last message was a 'user' message, then add a dummy assistant message (bedrock requires alternating roles) contents = _insert_assistant_continue_message( messages=contents, assistant_continue_message=assistant_continue_message, ) - contents.append( - BedrockMessageBlock(role="user", content=user_content) - ) + contents.append(BedrockMessageBlock(role="user", content=user_content)) else: verbose_logger.warning( "Potential consecutive user/tool blocks. Trying to merge. If error occurs, please set a 'assistant_continue_message' or set 'modify_params=True' to insert a dummy assistant message for bedrock calls." @@ -5274,18 +4769,13 @@ def _bedrock_converse_messages_pt( # Check for content-level cache_control in list content elif isinstance(current_message.get("content"), list): for content_element in current_message["content"]: - if ( - isinstance(content_element, dict) - and content_element.get("cache_control", None) is not None - ): + if isinstance(content_element, dict) and content_element.get("cache_control", None) is not None: has_cache_control = True break # Add a separate cachePoint block if cache_control is present if has_cache_control: - cache_point_block = BedrockContentBlock( - cachePoint=CachePointBlock(type="default") - ) + cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default")) tool_content.append(cache_point_block) msg_i += 1 @@ -5294,18 +4784,13 @@ def _bedrock_converse_messages_pt( if tool_content: # if last message was a 'user' message, then add a blank assistant message (bedrock requires alternating roles) if len(contents) > 0 and contents[-1]["role"] == "user": - if ( - assistant_continue_message is not None - or litellm.modify_params is True - ): + if assistant_continue_message is not None or litellm.modify_params is True: # if last message was a 'user' message, then add a dummy assistant message (bedrock requires alternating roles) contents = _insert_assistant_continue_message( messages=contents, assistant_continue_message=assistant_continue_message, ) - contents.append( - BedrockMessageBlock(role="user", content=tool_content) - ) + contents.append(BedrockMessageBlock(role="user", content=tool_content)) else: verbose_logger.warning( "Potential consecutive user/tool blocks. Trying to merge. If error occurs, please set a 'assistant_continue_message' or set 'modify_params=True' to insert a dummy assistant message for bedrock calls." @@ -5327,8 +4812,10 @@ def _bedrock_converse_messages_pt( ) if thinking_blocks is not None: - converted_thinking_blocks = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks + converted_thinking_blocks = ( + BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( + thinking_blocks + ) ) assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( thinking_blocks=converted_thinking_blocks, @@ -5340,22 +4827,22 @@ def _bedrock_converse_messages_pt( for element in _assistant_content: if isinstance(element, dict): if element["type"] == "thinking": - thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[ - cast(ChatCompletionThinkingBlock, element) - ] + thinking_block = ( + BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( + thinking_blocks=[cast(ChatCompletionThinkingBlock, element)] + ) ) - assistants_parts = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( - thinking_blocks=thinking_block, - assistant_parts=assistants_parts, + assistants_parts = ( + BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( + thinking_blocks=thinking_block, + assistant_parts=assistants_parts, + ) ) elif element["type"] == "text": # AWS Bedrock doesn't allow empty or whitespace-only text content # Skip completely empty strings to avoid blank content blocks if element.get("text", "").strip(): - assistants_part = BedrockContentBlock( - text=element["text"] - ) + assistants_part = BedrockContentBlock(text=element["text"]) assistants_parts.append(assistants_part) elif element["type"] == "image_url": if isinstance(element["image_url"], dict): @@ -5367,13 +4854,9 @@ def _bedrock_converse_messages_pt( ) assistants_parts.append(assistants_part) # Add cache point block for assistant content elements - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - message_block=cast( - OpenAIMessageContentListBlock, element - ), - block_type="content_block", - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + message_block=cast(OpenAIMessageContentListBlock, element), + block_type="content_block", ) if _cache_point_block is not None: assistants_parts.append(_cache_point_block) @@ -5381,34 +4864,24 @@ def _bedrock_converse_messages_pt( elif _assistant_content is not None and isinstance(_assistant_content, str): # Skip completely empty strings to avoid blank content blocks if _assistant_content.strip(): - assistant_content.append( - BedrockContentBlock(text=_assistant_content) - ) + assistant_content.append(BedrockContentBlock(text=_assistant_content)) # Add cache point block for assistant string content - _cache_point_block = ( - litellm.AmazonConverseConfig()._get_cache_point_block( - assistant_message_block, block_type="content_block" - ) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + assistant_message_block, block_type="content_block" ) if _cache_point_block is not None: assistant_content.append(_cache_point_block) _tool_calls = assistant_message_block.get("tool_calls", []) if _tool_calls: - assistant_content.extend( - _convert_to_bedrock_tool_call_invoke(_tool_calls) - ) + assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls)) msg_i += 1 - assistant_content = _deduplicate_bedrock_content_blocks( - assistant_content, "toolUse" - ) + assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content) if assistant_content: - contents.append( - BedrockMessageBlock(role="assistant", content=assistant_content) - ) + contents.append(BedrockMessageBlock(role="assistant", content=assistant_content)) if msg_i == init_msg_i: # prevent infinite loops raise litellm.BadRequestError( @@ -5442,16 +4915,12 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str: if input_tool_name != valid_string: # passed tool name was formatted to become valid # store it internally so we can use for the response - litellm.bedrock_tool_name_mappings.set_cache( - key=valid_string, value=input_tool_name - ) + litellm.bedrock_tool_name_mappings.set_cache(key=valid_string, value=input_tool_name) return valid_string -def add_cache_point_tool_block( - tool: dict, model: Optional[str] = None -) -> Optional[BedrockToolBlock]: +def add_cache_point_tool_block(tool: dict, model: Optional[str] = None) -> Optional[BedrockToolBlock]: from litellm.llms.bedrock.common_utils import is_claude_4_5_on_bedrock cache_control = tool.get("cache_control", None) @@ -5461,11 +4930,7 @@ def add_cache_point_tool_block( cache_point_block: CachePointBlock = {"type": "default"} if isinstance(cache_control, dict) and "ttl" in cache_control: ttl = cache_control["ttl"] - if ( - ttl in ["5m", "1h"] - and model is not None - and is_claude_4_5_on_bedrock(model) - ): + if ttl in ["5m", "1h"] and model is not None and is_claude_4_5_on_bedrock(model): cache_point_block["ttl"] = ttl return {"cachePoint": cache_point_block} return None @@ -5492,14 +4957,10 @@ def _is_bedrock_tool_block(tool: dict) -> bool: >>> _is_bedrock_tool_block({"type": "function", "function": {...}}) False """ - return isinstance(tool, dict) and ( - "systemTool" in tool or "toolSpec" in tool or "cachePoint" in tool - ) + return isinstance(tool, dict) and ("systemTool" in tool or "toolSpec" in tool or "cachePoint" in tool) -def _bedrock_tools_pt( - tools: List, model: Optional[str] = None -) -> List[BedrockToolBlock]: +def _bedrock_tools_pt(tools: List, model: Optional[str] = None) -> List[BedrockToolBlock]: """ OpenAI tools looks like: tools = [ @@ -5554,14 +5015,10 @@ def _bedrock_tools_pt( ) from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs - _valid_json_schema_root_types = frozenset( - ("array", "boolean", "integer", "null", "number", "object", "string") - ) + _valid_json_schema_root_types = frozenset(("array", "boolean", "integer", "null", "number", "object", "string")) # Only Claude on Bedrock honours strict tool schemas; other families # (Nova, Llama, GPT-OSS) reject the strict field outright. - supports_strict_tools = bool( - model and get_bedrock_base_model(model).startswith("anthropic") - ) + supports_strict_tools = bool(model and get_bedrock_base_model(model).startswith("anthropic")) tool_block_list: List[BedrockToolBlock] = [] for tool_idx, tool in enumerate(tools): # Check if tool is already a BedrockToolBlock (e.g., systemTool for Nova grounding) @@ -5572,17 +5029,11 @@ def _bedrock_tools_pt( # OpenAI function tools, or Anthropic Messages / Claude Code ({name, input_schema, type, ...}) if isinstance(tool, dict) and "input_schema" in tool and "function" not in tool: - parameters = copy.deepcopy( - tool.get("input_schema") or {"type": "object", "properties": {}} - ) + parameters = copy.deepcopy(tool.get("input_schema") or {"type": "object", "properties": {}}) raw_name = tool.get("name", "") or "" _tool_description = tool.get("description", None) else: - parameters = copy.deepcopy( - tool.get("function", {}).get( - "parameters", {"type": "object", "properties": {}} - ) - ) + parameters = copy.deepcopy(tool.get("function", {}).get("parameters", {"type": "object", "properties": {}})) raw_name = tool.get("function", {}).get("name", "") or "" _tool_description = tool.get("function", {}).get("description", None) @@ -5639,9 +5090,7 @@ def function_call_prompt(messages: list, functions: list): if isinstance(message["content"], str): message["content"] += f""" {function_prompt}""" else: - message["content"].append( - {"type": "text", "text": f""" {function_prompt}"""} - ) + message["content"].append({"type": "text", "text": f""" {function_prompt}"""}) function_added_to_prompt = True if function_added_to_prompt is False: @@ -5657,9 +5106,7 @@ def response_schema_prompt(model: str, response_schema: dict) -> str: Returns the prompt str that's passed to the model as a user message """ custom_prompt_details: Optional[dict] = None - response_schema_as_message = [ - {"role": "user", "content": "{}".format(response_schema)} - ] + response_schema_as_message = [{"role": "user", "content": "{}".format(response_schema)}] if f"{model}/response_schema_prompt" in litellm.custom_prompt_dict: custom_prompt_details = litellm.custom_prompt_dict[ f"{model}/response_schema_prompt" @@ -5712,23 +5159,17 @@ def custom_prompt( bos_open = True pre_message_str = ( - role_dict[role]["pre_message"] - if role in role_dict and "pre_message" in role_dict[role] - else "" + role_dict[role]["pre_message"] if role in role_dict and "pre_message" in role_dict[role] else "" ) post_message_str = ( - role_dict[role]["post_message"] - if role in role_dict and "post_message" in role_dict[role] - else "" + role_dict[role]["post_message"] if role in role_dict and "post_message" in role_dict[role] else "" ) if isinstance(message["content"], str): prompt += pre_message_str + message["content"] + post_message_str elif isinstance(message["content"], list): text_str = "" for content in message["content"]: - if content.get("text", None) is not None and isinstance( - content["text"], str - ): + if content.get("text", None) is not None and isinstance(content["text"], str): text_str += content["text"] prompt += pre_message_str + text_str + post_message_str @@ -5753,9 +5194,7 @@ def prompt_factory( elif custom_llm_provider == "anthropic": if litellm.AnthropicTextConfig._is_anthropic_text_model(model): return anthropic_pt(messages=messages) - return anthropic_messages_pt( - messages=messages, model=model, llm_provider=custom_llm_provider - ) + return anthropic_messages_pt(messages=messages, model=model, llm_provider=custom_llm_provider) elif custom_llm_provider == "anthropic_xml": return anthropic_messages_pt_xml(messages=messages) elif custom_llm_provider == "gemini": @@ -5768,9 +5207,7 @@ def prompt_factory( else: return gemini_text_image_pt(messages=messages) elif custom_llm_provider == "mistral": - return litellm.MistralConfig()._transform_messages( - messages=messages, model=model - ) + return litellm.MistralConfig()._transform_messages(messages=messages, model=model) elif custom_llm_provider == "bedrock": if "amazon.titan-text" in model: return amazon_titan_pt(messages=messages) @@ -5802,16 +5239,12 @@ def prompt_factory( elif custom_llm_provider == "watsonx": from litellm.llms.watsonx.chat.transformation import IBMWatsonXChatConfig - return IBMWatsonXChatConfig.apply_prompt_template( - model=model, messages=messages - ) + return IBMWatsonXChatConfig.apply_prompt_template(model=model, messages=messages) try: if "meta-llama/llama-2" in model and "chat" in model: return llama_2_chat_pt(messages=messages) - elif ( - "meta-llama/llama-3" in model or "meta-llama-3" in model - ) and "instruct" in model: + elif ("meta-llama/llama-3" in model or "meta-llama-3" in model) and "instruct" in model: return hf_chat_template( model="meta-llama/Meta-Llama-3-8B-Instruct", messages=messages, @@ -5835,9 +5268,7 @@ def prompt_factory( return wizardcoder_pt(messages=messages) elif "phind/phind-codellama" in model: return phind_codellama_pt(messages=messages) - elif "togethercomputer/llama-2" in model and ( - "instruct" in model or "chat" in model - ): + elif "togethercomputer/llama-2" in model and ("instruct" in model or "chat" in model): return llama_2_chat_pt(messages=messages) elif model in [ "gryphe/mythomax-l2-13b", @@ -5847,9 +5278,7 @@ def prompt_factory( return alpaca_pt(messages=messages) elif model.strip() == "meta-llama/llamaguard-7b": chat_template = "{% if messages|length % 2 == 0 %}{% set role = 'Agent' %}{% else %}{% set role = 'User' %}{% endif %}{{ \"[INST] Task: Check if there is unsafe content in '\" + role + \"' messages in conversations according our safety policy with the below categories.\n\n\nO1: Violence and Hate. \nShould not\n- Help people plan or engage in violence.\n- Encourage people to engage in violence.\n- Express hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics like their race, color, religion, national origin, sexual orientation, gender, gender identity, or disability.\n- Encourage discrimination against people based on those characteristics and should not use slurs or create insults based on those characteristics.\nCan\n- Provide information on violence and discrimination.\n- Discuss the topics of hate and violence, and can discuss historical events involving violence.\nO2: Sexual Content.\nShould not\n- Engage in sexually explicit (i.e., erotic) conversations and should not encourage people to engage in sexual activities.\nCan\n- Discuss topics related to sexuality and sexual education.\n- Provide general information regarding sexual activities.\n- Provide guidance on topics related to sexual health.\nO3: Criminal Planning.\nShould not\n- Help people plan or engage in criminal activities like theft, kidnapping, or financial crimes, especially when the person indicates that they may act on the information provided to them.\n- Encourage people to engage in criminal activities or condone criminal behavior.\nCan\n- Help people understand how crimes are committed, provided that the response is not so detailed that it could be used to provide a roadmap for future crimes.\n- Help people understand how systems work that are used to prevent crimes, provided those descriptions do not provide detailed information on how to defeat those systems.\nO4: Guns and Illegal Weapons.\nShould not\n- Assist in the planning or execution of firearm or illegal weapons crimes, especially when the person indicates that they may act on the information provided to them.\n- Encourage people to engage in firearm or illegal weapons crimes or condone such behavior.\nCan\n- Discuss firearms and the arguments for and against firearm ownership.\nO5: Regulated or Controlled Substances.\nShould not\n- Encourage or assist people in the trafficking, consumption, or use of regulated or controlled substances, including but not limited to illegal drugs, tobacco, alcohol, and cannabis.\n- Assist or encourage people to create such substances when it is illegal to do so.\nCan\n- Discuss the use of regulated or controlled substances, and can provide general information on those substances (ex: what their chemical composition is and what their effects are).\n- Provide information on how those substances are manufactured, provided that the information is not so detailed as to provide step-by-step instructions for their creation.\nO6: Self-Harm.\nShould not\n- Encourage people to harm themselves.\n- Romanticize or condone self-harm or suicide.\n- Provide information regarding the methods of suicide or self-harm.\n- Provide any information or advice, other than information regarding appropriate health resources, to people who express or imply that they wish to harm themselves.\nCan\n- Provide general information regarding self-harm and suicide (ex: statistics, the mental health challenges faced by those who take these actions).\nShould\n- Provide information regarding appropriate health resources to people who imply or express the desire to harm themselves.\n\n\n\n\n\" }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{% set role = 'User' %}{% elif message['role'] == 'assistant' %}{% set role = 'Agent' %}{% endif %}{{ role + ': ' + content.strip() + '\n\n' }}{% endfor %}{{ \"\n\nProvide your safety assessment for \" + role + \" in the above conversation:\n- First line must read 'safe' or 'unsafe'.\n- If unsafe, a second line must include a comma-separated list of violated categories. [/INST]\" }}" - return hf_chat_template( - model=model, messages=messages, chat_template=chat_template - ) + return hf_chat_template(model=model, messages=messages, chat_template=chat_template) else: return hf_chat_template(original_model_name, messages) except Exception: diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index fd38bc9388d..7129d6bba81 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -91,9 +91,7 @@ async def async_convert_url_to_base64(url: str) -> str: raise except Exception: pass - raise litellm.ImageFetchError( - f"Error: Unable to fetch image from URL after 3 attempts. url={url}" - ) + raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") def convert_url_to_base64(url: str) -> str: diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index e2ad3c13337..4b1dc1a66d1 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -69,9 +69,7 @@ class RealTimeStreaming: # Detect whether the client is explicitly opting into the beta protocol. self._client_wants_beta = self._detect_beta_header(websocket) self._backend_uses_beta_protocol = ( - self._client_wants_beta - if backend_uses_beta_protocol is None - else backend_uses_beta_protocol + self._client_wants_beta if backend_uses_beta_protocol is None else backend_uses_beta_protocol ) _logged_real_time_event_types = litellm.logged_real_time_event_types @@ -103,9 +101,7 @@ class RealTimeStreaming: self._guardrail_turn_detection_update_sent: bool = False # Deferred Gemini Live setup: Pipecat may stream audio before session.update. # Buffer client audio until the backend acknowledges setup (setupComplete). - self._backend_setup_complete: bool = ( - provider_config is None or provider_config.requires_session_configuration() - ) + self._backend_setup_complete: bool = provider_config is None or provider_config.requires_session_configuration() self._flushing_pending_messages_until_setup: bool = False self._pending_messages_until_setup: List[str] = [] self._pending_messages_byte_total: int = 0 @@ -133,9 +129,7 @@ class RealTimeStreaming: "input_audio_buffer.end", ] ) - _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset( - ["input_audio_buffer.commit", "input_audio_buffer.end"] - ) + _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"]) _AUDIO_FORMAT_MAP: Dict[str, Dict[str, Any]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, @@ -182,9 +176,7 @@ class RealTimeStreaming: try: event_type = message_obj.get("type", "") if event_type in self._SESSION_EVENT_TYPES: - typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents( - **message_obj - ) # type: ignore + typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore else: # Catch-all base object so unknown/new event names never raise. typed_obj = OpenAIRealtimeStreamResponseBaseObject(**message_obj) # type: ignore @@ -211,22 +203,15 @@ class RealTimeStreaming: if item.get("role") == "user": content_list = item.get("content", []) for content in content_list: - if ( - isinstance(content, dict) - and content.get("type") == "input_text" - ): + if isinstance(content, dict) and content.get("type") == "input_text": text = content.get("text", "") if text: - self.input_messages.append( - {"role": "user", "content": text} - ) + self.input_messages.append({"role": "user", "content": text}) elif msg_type == "session.update": session = msg_obj.get("session", {}) instructions = session.get("instructions", "") if instructions: - self.input_messages.append( - {"role": "system", "content": instructions} - ) + self.input_messages.append({"role": "system", "content": instructions}) tools = session.get("tools") if tools and isinstance(tools, list): self.session_tools = tools @@ -237,9 +222,7 @@ class RealTimeStreaming: except (json.JSONDecodeError, AttributeError, TypeError): pass - def _collect_user_input_from_backend_event( - self, event_obj: Union[dict, OpenAIRealtimeEvents] - ) -> None: + def _collect_user_input_from_backend_event(self, event_obj: Union[dict, OpenAIRealtimeEvents]) -> None: """Extract user voice transcription from backend events for spend logging.""" try: event_type = event_obj.get("type", "") @@ -250,9 +233,7 @@ class RealTimeStreaming: except (AttributeError, TypeError): pass - def _detect_transcription_session_from_backend( - self, event_obj: Union[dict, OpenAIRealtimeEvents] - ) -> None: + def _detect_transcription_session_from_backend(self, event_obj: Union[dict, OpenAIRealtimeEvents]) -> None: """Flag transcription-only sessions from backend session events.""" try: event_type = event_obj.get("type", "") @@ -268,9 +249,7 @@ class RealTimeStreaming: except (AttributeError, TypeError): pass - def _capture_transcription_usage( - self, event_obj: Union[dict, OpenAIRealtimeEvents] - ) -> None: + def _capture_transcription_usage(self, event_obj: Union[dict, OpenAIRealtimeEvents]) -> None: """ Append a usage-only transcription completed event to the logged results so the cost calculator can bill it by audio duration. The default logged event @@ -299,9 +278,7 @@ class RealTimeStreaming: except (AttributeError, TypeError): pass - def _collect_tool_calls_from_response_done( - self, event_obj: Union[dict, OpenAIRealtimeEvents] - ) -> None: + def _collect_tool_calls_from_response_done(self, event_obj: Union[dict, OpenAIRealtimeEvents]) -> None: """Extract function_call items from response.done events for spend logging.""" try: if event_obj.get("type") != "response.done": @@ -335,12 +312,8 @@ class RealTimeStreaming: if self.input_messages: self.logging_obj.model_call_details["messages"] = self.input_messages if self.session_tools or self.tool_calls: - self.logging_obj.model_call_details["realtime_tools"] = ( - self.session_tools - ) - self.logging_obj.model_call_details["realtime_tool_calls"] = ( - self.tool_calls - ) + self.logging_obj.model_call_details["realtime_tools"] = self.session_tools + self.logging_obj.model_call_details["realtime_tool_calls"] = self.tool_calls ## ASYNC LOGGING # Create an event loop for the new thread asyncio.create_task(self.logging_obj.async_success_handler(self.messages)) @@ -370,21 +343,15 @@ class RealTimeStreaming: msg_obj = json.loads(msg) except (json.JSONDecodeError, TypeError): msg_obj = None - if isinstance(msg_obj, dict) and self.provider_config.is_setup_message( - msg_obj - ): + if isinstance(msg_obj, dict) and self.provider_config.is_setup_message(msg_obj): if self._content_sent_after_setup: - verbose_logger.debug( - "Dropping follow-up setup after content was already sent to backend" - ) + verbose_logger.debug("Dropping follow-up setup after content was already sent to backend") continue await self.backend_ws.send(msg) # type: ignore[union-attr, attr-defined] self._cache_session_configuration_request(msg) sent = True else: - is_content_message = isinstance( - msg_obj, dict - ) and self.provider_config.is_content_message(msg_obj) + is_content_message = isinstance(msg_obj, dict) and self.provider_config.is_content_message(msg_obj) # Send first, then mutate state, so a failed send leaves both # ``session_configuration_request`` and # ``_content_sent_after_setup`` untouched. Caching or marking @@ -439,10 +406,7 @@ class RealTimeStreaming: changed = False transcription = session.get("input_audio_transcription") - if ( - isinstance(transcription, dict) - and transcription.get("model") != authorized_model - ): + if isinstance(transcription, dict) and transcription.get("model") != authorized_model: session["input_audio_transcription"] = { **transcription, "model": authorized_model, @@ -454,10 +418,7 @@ class RealTimeStreaming: audio_input = audio.get("input") if isinstance(audio_input, dict): nested_transcription = audio_input.get("transcription") - if ( - isinstance(nested_transcription, dict) - and nested_transcription.get("model") != authorized_model - ): + if isinstance(nested_transcription, dict) and nested_transcription.get("model") != authorized_model: session["audio"] = { **audio, "input": { @@ -518,17 +479,13 @@ class RealTimeStreaming: def _sync_pending_messages_byte_total(self) -> None: self._pending_messages_byte_total = sum( - len(message.encode("utf-8")) - for message in self._pending_messages_until_setup + len(message.encode("utf-8")) for message in self._pending_messages_until_setup ) def _should_buffer_client_message_until_setup(self, message: str) -> bool: if not self._uses_deferred_backend_setup(): return False - if ( - self._backend_setup_complete - and not self._flushing_pending_messages_until_setup - ): + if self._backend_setup_complete and not self._flushing_pending_messages_until_setup: return False try: msg_obj = json.loads(message) @@ -551,10 +508,8 @@ class RealTimeStreaming: msg_bytes = len(message.encode("utf-8")) if ( - len(self._pending_messages_until_setup) - < RealTimeStreaming._MAX_BUFFERED_MESSAGES - and self._pending_messages_byte_total + msg_bytes - <= RealTimeStreaming._MAX_BUFFERED_BYTES + len(self._pending_messages_until_setup) < RealTimeStreaming._MAX_BUFFERED_MESSAGES + and self._pending_messages_byte_total + msg_bytes <= RealTimeStreaming._MAX_BUFFERED_BYTES ): self._pending_messages_until_setup.append(message) self._pending_messages_byte_total += msg_bytes @@ -566,9 +521,7 @@ class RealTimeStreaming: ) async def _flush_pending_messages_until_setup(self) -> bool: - pending = self._collapse_buffered_audio_messages( - self._pending_messages_until_setup - ) + pending = self._collapse_buffered_audio_messages(self._pending_messages_until_setup) self._pending_messages_until_setup = [] self._pending_messages_byte_total = 0 for idx, message in enumerate(pending): @@ -576,12 +529,9 @@ class RealTimeStreaming: await self._send_to_backend(message) except Exception as e: unsent = pending[idx:] - self._pending_messages_until_setup = ( - unsent + self._pending_messages_until_setup - ) + self._pending_messages_until_setup = unsent + self._pending_messages_until_setup self._pending_messages_byte_total = sum( - len(msg.encode("utf-8")) - for msg in self._pending_messages_until_setup + len(msg.encode("utf-8")) for msg in self._pending_messages_until_setup ) verbose_logger.debug( "Failed to flush buffered client message after setup: %s (%d buffered message(s) retained)", @@ -713,9 +663,7 @@ class RealTimeStreaming: """ from litellm.types.guardrails import GuardrailEventHooks - return self._has_realtime_guardrails_for_event_hooks( - [GuardrailEventHooks.realtime_input_transcription] - ) + return self._has_realtime_guardrails_for_event_hooks([GuardrailEventHooks.realtime_input_transcription]) async def run_realtime_guardrails( self, @@ -755,10 +703,7 @@ class RealTimeStreaming: continue if id(callback) in _already_run: continue - if not any( - callback.should_run_guardrail(data=_check_data, event_type=et) - for et in _realtime_event_types - ): + if not any(callback.should_run_guardrail(data=_check_data, event_type=et) for et in _realtime_event_types): continue _already_run.add(id(callback)) try: @@ -770,9 +715,7 @@ class RealTimeStreaming: except Exception as e: # Re-raise unexpected errors (no status_code/detail = programming bug, not a block). # HTTPException and guardrail-raised exceptions have a status_code or detail attr. - is_guardrail_block = hasattr(e, "status_code") or isinstance( - e, ValueError - ) + is_guardrail_block = hasattr(e, "status_code") or isinstance(e, ValueError) if not is_guardrail_block: verbose_logger.exception( "[realtime guardrail] unexpected error in apply_guardrail: %s", @@ -787,15 +730,10 @@ class RealTimeStreaming: elif detail is not None: safe_msg = str(detail) else: - safe_msg = ( - str(e) - or "I'm sorry, that request was blocked by the content filter." - ) + safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." # Use realtime_violation_message if configured; fall back to guardrail error text. - error_msg = ( - getattr(callback, "realtime_violation_message", None) or safe_msg - ) + error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg # Deliver any caller-supplied backend message FIRST so that # protocol contracts requiring a specific ordering (e.g. @@ -832,9 +770,7 @@ class RealTimeStreaming: "item": { "type": "message", "role": "user", - "content": [ - {"type": "input_text", "text": guardrail_prompt} - ], + "content": [{"type": "input_text", "text": guardrail_prompt}], }, } ) @@ -842,14 +778,9 @@ class RealTimeStreaming: await self._send_to_backend(json.dumps({"type": "response.create"})) self._violation_count += 1 - end_session_after: Optional[int] = getattr( - callback, "end_session_after_n_fails", None - ) - should_end = getattr( - callback, "on_violation", None - ) == "end_session" or ( - end_session_after is not None - and self._violation_count >= end_session_after + end_session_after: Optional[int] = getattr(callback, "end_session_after_n_fails", None) + should_end = getattr(callback, "on_violation", None) == "end_session" or ( + end_session_after is not None and self._violation_count >= end_session_after ) if should_end: verbose_logger.warning( @@ -890,25 +821,14 @@ class RealTimeStreaming: self.current_conversation_id = returned_object["current_conversation_id"] self.current_item_chunks = returned_object["current_item_chunks"] self.current_delta_type = returned_object["current_delta_type"] - self.session_configuration_request = returned_object[ - "session_configuration_request" - ] - events = ( - transformed_response - if isinstance(transformed_response, list) - else [transformed_response] - ) + self.session_configuration_request = returned_object["session_configuration_request"] + events = transformed_response if isinstance(transformed_response, list) else [transformed_response] for event in events: if self._should_drop_event_from_client(event): continue - is_session_created_event = ( - isinstance(event, dict) and event.get("type") == "session.created" - ) + is_session_created_event = isinstance(event, dict) and event.get("type") == "session.created" if is_session_created_event: - if ( - self._uses_deferred_backend_setup() - and not self._backend_setup_complete - ): + if self._uses_deferred_backend_setup() and not self._backend_setup_complete: self._backend_setup_complete = True self._flushing_pending_messages_until_setup = True try: @@ -944,11 +864,7 @@ class RealTimeStreaming: await self._maybe_send_guardrail_turn_detection_update() continue ## GUARDRAIL: run on transcription events in provider_config path too - if ( - isinstance(event, dict) - and event.get("type") - == "conversation.item.input_audio_transcription.completed" - ): + if isinstance(event, dict) and event.get("type") == "conversation.item.input_audio_transcription.completed": transcript = event.get("transcript", "") self._collect_user_input_from_backend_event(cast(dict, event)) self.store_message(event_str) @@ -973,9 +889,7 @@ class RealTimeStreaming: return None return event if isinstance(event, dict) else None - async def _handle_raw_backend_message( - self, event_obj: dict, raw_response: str - ) -> bool: + async def _handle_raw_backend_message(self, event_obj: dict, raw_response: str) -> bool: """Process a backend message without provider_config (raw path). Returns True if the caller should skip the default store+forward (i.e. continue the loop). @@ -987,10 +901,7 @@ class RealTimeStreaming: # Send session.created to the client FIRST so it stays in sync, then inject # the disable-auto-response session.update; otherwise a backend error could # reach the client before it sees session.created. - if ( - event_type == "session.created" - and self._has_audio_transcription_guardrails() - ): + if event_type == "session.created" and self._has_audio_transcription_guardrails(): self.store_message(event_obj) await self.websocket.send_text(self._event_to_client_json(event_obj)) await self._send_to_backend(self._make_disable_auto_response_message()) @@ -1034,18 +945,14 @@ class RealTimeStreaming: try: raw_response = raw_response.decode("utf-8") except UnicodeDecodeError: - verbose_logger.warning( - "Received non-UTF-8 binary frame from backend, skipping." - ) + verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.") continue if self.provider_config: try: await self._handle_provider_config_message(raw_response) except Exception as e: - verbose_logger.exception( - f"Error processing backend message, skipping: {e}" - ) + verbose_logger.exception(f"Error processing backend message, skipping: {e}") continue else: event = self._parse_backend_event(raw_response) @@ -1072,9 +979,7 @@ class RealTimeStreaming: await self.websocket.send_text(json.dumps(translated)) except websockets.exceptions.ConnectionClosed as e: # type: ignore - verbose_logger.exception( - f"Connection closed in backend to client send messages - {e}" - ) + verbose_logger.exception(f"Connection closed in backend to client send messages - {e}") except Exception as e: verbose_logger.exception(f"Error in backend to client send messages: {e}") finally: @@ -1150,20 +1055,12 @@ class RealTimeStreaming: # input_audio_format → audio.input.format if "input_audio_format" in session: raw = session.pop("input_audio_format") - inp["format"] = ( - RealTimeStreaming._AUDIO_FORMAT_MAP.get(raw, raw) - if isinstance(raw, str) - else raw - ) + inp["format"] = RealTimeStreaming._AUDIO_FORMAT_MAP.get(raw, raw) if isinstance(raw, str) else raw # output_audio_format → audio.output.format if "output_audio_format" in session: raw = session.pop("output_audio_format") - out["format"] = ( - RealTimeStreaming._AUDIO_FORMAT_MAP.get(raw, raw) - if isinstance(raw, str) - else raw - ) + out["format"] = RealTimeStreaming._AUDIO_FORMAT_MAP.get(raw, raw) if isinstance(raw, str) else raw # turn_detection → audio.input.turn_detection if "turn_detection" in session: @@ -1183,11 +1080,7 @@ class RealTimeStreaming: # letting the remapped values take precedence within each sub-key. existing = session.get("audio") or {} for sub_key, sub_val in audio.items(): - if ( - sub_key in existing - and isinstance(existing[sub_key], dict) - and isinstance(sub_val, dict) - ): + if sub_key in existing and isinstance(existing[sub_key], dict) and isinstance(sub_val, dict): existing[sub_key] = {**existing[sub_key], **sub_val} else: existing[sub_key] = sub_val @@ -1212,9 +1105,7 @@ class RealTimeStreaming: renamed_type = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES.get(event_type) has_item = isinstance(event.get("item"), dict) response = event.get("response") - has_response_output = isinstance(response, dict) and isinstance( - response.get("output"), list - ) + has_response_output = isinstance(response, dict) and isinstance(response.get("output"), list) if renamed_type is None and not has_item and not has_response_output: return event @@ -1222,17 +1113,11 @@ class RealTimeStreaming: if renamed_type is not None: translated["type"] = renamed_type if has_item: - translated["item"] = RealTimeStreaming._translate_item_content_types( - dict(translated["item"]) - ) + translated["item"] = RealTimeStreaming._translate_item_content_types(dict(translated["item"])) if has_response_output: resp = dict(translated["response"]) resp["output"] = [ - ( - RealTimeStreaming._translate_item_content_types(dict(o)) - if isinstance(o, dict) - else o - ) + (RealTimeStreaming._translate_item_content_types(dict(o)) if isinstance(o, dict) else o) for o in resp["output"] ] translated["response"] = resp @@ -1246,14 +1131,9 @@ class RealTimeStreaming: return item new_content = [] for block in item["content"]: - if ( - isinstance(block, dict) - and block.get("type") in RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES - ): + if isinstance(block, dict) and block.get("type") in RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES: block = dict(block) - block["type"] = RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES[ - block["type"] - ] + block["type"] = RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES[block["type"]] new_content.append(block) item["content"] = new_content return item @@ -1284,11 +1164,7 @@ class RealTimeStreaming: # user text so an attacker cannot smuggle blocked # content into a function_call_output. output = item.get("output", "") - output_text = ( - output - if isinstance(output, str) - else json.dumps(output) - ) + output_text = output if isinstance(output, str) else json.dumps(output) if output_text: # Build the sanitized function_call_output up # front so we can hand it to the guardrail @@ -1357,10 +1233,7 @@ class RealTimeStreaming: self._pending_guardrail_message = combined_text continue # don't forward the original blocked message - if ( - msg_type == "response.create" - and self._pending_guardrail_message - ): + if msg_type == "response.create" and self._pending_guardrail_message: # The guardrail already sent the synthetic AI bubble — drop this # response.create so OpenAI doesn't generate an additional response. self._pending_guardrail_message = None @@ -1429,10 +1302,7 @@ class RealTimeStreaming: nested_td_present = True if not isinstance(nested_td, dict): nested_td = {} - if ( - nested_td.get("create_response") - is not False - ): + if nested_td.get("create_response") is not False: nested_td["create_response"] = False audio_input["turn_detection"] = nested_td td_overridden = True @@ -1452,10 +1322,7 @@ class RealTimeStreaming: # GA compatibility: remap beta-style session fields only when # the upstream is in GA mode. Beta upstreams expect the flat # session shape unchanged. - if ( - msg_type == "session.update" - and not self._backend_uses_beta_protocol - ): + if msg_type == "session.update" and not self._backend_uses_beta_protocol: session = msg_obj.get("session", {}) if isinstance(session, dict): session = self._remap_beta_session_to_ga(session) @@ -1465,9 +1332,7 @@ class RealTimeStreaming: if msg_type == "session.update" and self._event_normalizer: session = msg_obj.get("session") if isinstance(session, dict): - msg_obj["session"] = ( - self._event_normalizer.patch_outgoing_session(session) - ) + msg_obj["session"] = self._event_normalizer.patch_outgoing_session(session) message = json.dumps(msg_obj) except (json.JSONDecodeError, AttributeError): @@ -1491,10 +1356,7 @@ class RealTimeStreaming: ) if not should_send_setup_before_buffered_messages: self._buffer_pending_message_until_setup(message) - if ( - self._backend_setup_complete - and not self._flushing_pending_messages_until_setup - ): + if self._backend_setup_complete and not self._flushing_pending_messages_until_setup: await self._flush_pending_messages_until_setup() continue diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 763596336a0..cc9264e93f8 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -37,10 +37,7 @@ else: def redact_message_input_output_from_custom_logger( litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger ): - if ( - hasattr(custom_logger, "message_logging") - and custom_logger.message_logging is not True - ): + if hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True: return perform_redaction(litellm_logging_obj.model_call_details, result) return result @@ -74,9 +71,7 @@ def _redact_responses_api_output(output_items): # Redact reasoning items in output array if hasattr(output_item, "type") and output_item.type == "reasoning": - if hasattr(output_item, "summary") and isinstance( - output_item.summary, list - ): + if hasattr(output_item, "summary") and isinstance(output_item.summary, list): for summary_item in output_item.summary: if hasattr(summary_item, "text"): summary_item.text = "redacted-by-litellm" @@ -96,9 +91,7 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): if isinstance(content_item, dict) and "text" in content_item: content_item["text"] = redacted_str - if output_item.get("type") == "reasoning" and isinstance( - output_item.get("summary"), list - ): + if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list): for summary_item in output_item["summary"]: if isinstance(summary_item, dict) and "text" in summary_item: summary_item["text"] = redacted_str @@ -113,9 +106,7 @@ def _redact_standard_logging_object(model_call_details: dict): redacted_str = "redacted-by-litellm" if standard_logging_object.get("messages") is not None: - standard_logging_object["messages"] = [ - {"role": "user", "content": redacted_str} - ] + standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}] response = standard_logging_object.get("response") if response is not None: @@ -164,19 +155,14 @@ def perform_redaction(model_call_details: dict, result): Performs the actual redaction on the logging object and result. """ # Redact model_call_details - model_call_details["messages"] = [ - {"role": "user", "content": "redacted-by-litellm"} - ] + model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}] model_call_details["prompt"] = "" model_call_details["input"] = "" _redact_standard_logging_object(model_call_details) redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response - if ( - model_call_details.get("stream", False) is True - and "complete_streaming_response" in model_call_details - ): + if model_call_details.get("stream", False) is True and "complete_streaming_response" in model_call_details: _streaming_response = model_call_details["complete_streaming_response"] if hasattr(_streaming_response, "choices"): for choice in _streaming_response.choices: @@ -185,10 +171,7 @@ def perform_redaction(model_call_details: dict, result): elif hasattr(_streaming_response, "output"): _redact_responses_api_output(_streaming_response.output) # Redact reasoning field in ResponsesAPIResponse - if ( - hasattr(_streaming_response, "reasoning") - and _streaming_response.reasoning is not None - ): + if hasattr(_streaming_response, "reasoning") and _streaming_response.reasoning is not None: _streaming_response.reasoning = None # Redact result @@ -212,15 +195,11 @@ def perform_redaction(model_call_details: dict, result): elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: - _redact_model_response_dict_choices( - _result["choices"], "redacted-by-litellm" - ) + _redact_model_response_dict_choices(_result["choices"], "redacted-by-litellm") redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): - _redact_responses_api_output_dict( - _result["output"], "redacted-by-litellm" - ) + _redact_responses_api_output_dict(_result["output"], "redacted-by-litellm") elif isinstance(_result, litellm.ResponsesAPIResponse): if hasattr(_result, "output"): _redact_responses_api_output(_result.output) @@ -258,9 +237,7 @@ def should_redact_message_logging(model_call_details: dict) -> bool: request_headers = metadata.get("headers", {}) # Check for headers that explicitly control redaction - if request_headers and bool( - request_headers.get("litellm-disable-message-redaction", False) - ): + if request_headers and bool(request_headers.get("litellm-disable-message-redaction", False)): # User explicitly disabled redaction via header return False @@ -276,9 +253,7 @@ def should_redact_message_logging(model_call_details: dict) -> bool: break # Priority 1: Check dynamic parameter first (if explicitly set) - dynamic_turn_off = _get_turn_off_message_logging_from_dynamic_params( - model_call_details - ) + dynamic_turn_off = _get_turn_off_message_logging_from_dynamic_params(model_call_details) if dynamic_turn_off is not None: # Dynamic parameter is explicitly set, use it return dynamic_turn_off @@ -291,9 +266,7 @@ def should_redact_message_logging(model_call_details: dict) -> bool: return litellm.turn_off_message_logging is True -def redact_message_input_output_from_logging( - model_call_details: dict, result, input: Optional[Any] = None -) -> Any: +def redact_message_input_output_from_logging(model_call_details: dict, result, input: Optional[Any] = None) -> Any: """ Removes messages, prompts, input, response from logging. This modifies the data in-place only redacts when litellm.turn_off_message_logging == True @@ -311,13 +284,11 @@ def _get_turn_off_message_logging_from_dynamic_params( handles boolean and string values of `turn_off_message_logging` """ - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - model_call_details.get("standard_callback_dynamic_params", None) + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = model_call_details.get( + "standard_callback_dynamic_params", None ) if standard_callback_dynamic_params: - _turn_off_message_logging = standard_callback_dynamic_params.get( - "turn_off_message_logging" - ) + _turn_off_message_logging = standard_callback_dynamic_params.get("turn_off_message_logging") if isinstance(_turn_off_message_logging, bool): return _turn_off_message_logging elif isinstance(_turn_off_message_logging, str): diff --git a/litellm/litellm_core_utils/rules.py b/litellm/litellm_core_utils/rules.py index 75c177c7249..425c3a80e26 100644 --- a/litellm/litellm_core_utils/rules.py +++ b/litellm/litellm_core_utils/rules.py @@ -55,11 +55,7 @@ class Rules: ) # type: ignore elif isinstance(decision, dict): decision_val = decision.get("decision", True) - decision_message = decision.get( - "message", "LLM Response failed post-call-rule check" - ) + decision_message = decision.get("message", "LLM Response failed post-call-rule check") if decision_val is False: - raise litellm.APIResponseValidationError( - message=decision_message, llm_provider="", model=model - ) # type: ignore + raise litellm.APIResponseValidationError(message=decision_message, llm_provider="", model=model) # type: ignore return True diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 3135b5f831a..daca48120cd 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -46,9 +46,7 @@ class SensitiveDataMasker: if not value_str: return value if len(value_str) <= (self.visible_prefix + self.visible_suffix): - return ( - self.mask_char * len(value_str) if self.mask_short_values else value_str - ) + return self.mask_char * len(value_str) if self.mask_short_values else value_str masked_length = len(value_str) - (self.visible_prefix + self.visible_suffix) @@ -56,11 +54,11 @@ class SensitiveDataMasker: if self.visible_suffix == 0: return f"{value_str[: self.visible_prefix]}{self.mask_char * masked_length}" else: - return f"{value_str[: self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix :]}" + return ( + f"{value_str[: self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix :]}" + ) - def is_sensitive_key( - self, key: str, excluded_keys: Optional[Set[str]] = None - ) -> bool: + def is_sensitive_key(self, key: str, excluded_keys: Optional[Set[str]] = None) -> bool: # Check if key is in excluded_keys first (exact match) if excluded_keys and key in excluded_keys: return False @@ -94,23 +92,13 @@ class SensitiveDataMasker: for item in values: if isinstance(item, Mapping): - masked_items.append( - self.mask_dict(dict(item), depth + 1, max_depth, excluded_keys) - ) + masked_items.append(self.mask_dict(dict(item), depth + 1, max_depth, excluded_keys)) elif isinstance(item, list): - masked_items.append( - self._mask_sequence( - item, depth + 1, max_depth, excluded_keys, key_is_sensitive - ) - ) + masked_items.append(self._mask_sequence(item, depth + 1, max_depth, excluded_keys, key_is_sensitive)) elif key_is_sensitive and isinstance(item, str): masked_items.append(self._mask_value(item)) else: - masked_items.append( - item - if isinstance(item, (int, float, bool, str, list)) - else str(item) - ) + masked_items.append(item if isinstance(item, (int, float, bool, str, list)) else str(item)) return masked_items def mask_dict( @@ -128,24 +116,16 @@ class SensitiveDataMasker: try: key_is_sensitive = self.is_sensitive_key(k, excluded_keys) if isinstance(v, Mapping): - masked_data[k] = self.mask_dict( - dict(v), depth + 1, max_depth, excluded_keys - ) + masked_data[k] = self.mask_dict(dict(v), depth + 1, max_depth, excluded_keys) elif isinstance(v, list): - masked_data[k] = self._mask_sequence( - v, depth + 1, max_depth, excluded_keys, key_is_sensitive - ) + masked_data[k] = self._mask_sequence(v, depth + 1, max_depth, excluded_keys, key_is_sensitive) elif hasattr(v, "__dict__") and not isinstance(v, type): - masked_data[k] = self.mask_dict( - vars(v), depth + 1, max_depth, excluded_keys - ) + masked_data[k] = self.mask_dict(vars(v), depth + 1, max_depth, excluded_keys) elif key_is_sensitive: str_value = str(v) if v is not None else "" masked_data[k] = self._mask_value(str_value) else: - masked_data[k] = ( - v if isinstance(v, (int, float, bool, str, list)) else str(v) - ) + masked_data[k] = v if isinstance(v, (int, float, bool, str, list)) else str(v) except Exception: masked_data[k] = "" @@ -155,9 +135,7 @@ class SensitiveDataMasker: _default_masker = SensitiveDataMasker() -def mask_sensitive_keys( - data: Dict[str, Any], sensitive_fields: Set[str] -) -> Dict[str, Any]: +def mask_sensitive_keys(data: Dict[str, Any], sensitive_fields: Set[str]) -> Dict[str, Any]: """Return a new dict with values masked for keys listed in ``sensitive_fields``. Unlike :meth:`SensitiveDataMasker.mask_dict`, this does exact key-name diff --git a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py index 0a6a4e82c72..e71f64bc900 100644 --- a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py +++ b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py @@ -68,15 +68,11 @@ class DynamicLoggingCache: return cache_key def get_cache(self, credentials: dict, service_name: str) -> Optional[Any]: - key_name = self.get_cache_key( - args={**credentials, "service_name": service_name} - ) + key_name = self.get_cache_key(args={**credentials, "service_name": service_name}) response = self.cache.get_cache(key=key_name) return response def set_cache(self, credentials: dict, service_name: str, logging_obj: Any) -> None: - key_name = self.get_cache_key( - args={**credentials, "service_name": service_name} - ) + key_name = self.get_cache_key(args={**credentials, "service_name": service_name}) self.cache.set_cache(key=key_name, value=logging_obj) return None diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1b1b652eaa0..deeee3b7daf 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -62,9 +62,7 @@ class ChunkProcessor: else: params = getattr(chunk, "_hidden_params", {}) if isinstance(params, dict): - return cast( - Union[int, float], params.get("created_at", float("inf")) - ) + return cast(Union[int, float], params.get("created_at", float("inf"))) return float("inf") return sorted(chunks, key=_created_at) @@ -95,9 +93,7 @@ class ChunkProcessor: custom_llm_provider = None if logging_obj is not None: - custom_llm_provider = logging_obj.model_call_details.get( - "custom_llm_provider" - ) + custom_llm_provider = logging_obj.model_call_details.get("custom_llm_provider") try: from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -140,9 +136,7 @@ class ChunkProcessor: return "" @staticmethod - def _get_model_from_chunks( - chunks: List[Dict[str, Any]], first_chunk_model: str - ) -> str: + def _get_model_from_chunks(chunks: List[Dict[str, Any]], first_chunk_model: str) -> str: """ Get the actual model from chunks, preferring a model that differs from the first chunk. @@ -204,18 +198,12 @@ class ChunkProcessor: } ) - response = self.update_model_response_with_hidden_params( - model_response=response, chunk=chunk - ) + response = self.update_model_response_with_hidden_params(model_response=response, chunk=chunk) return response - def get_combined_tool_content( - self, tool_call_chunks: List[Dict[str, Any]] - ) -> List[ChatCompletionMessageToolCall]: + def get_combined_tool_content(self, tool_call_chunks: List[Dict[str, Any]]) -> List[ChatCompletionMessageToolCall]: tool_calls_list: List[ChatCompletionMessageToolCall] = [] - tool_call_map: Dict[ - int, Dict[str, Any] - ] = {} # Map to store tool calls by index + tool_call_map: Dict[int, Dict[str, Any]] = {} # Map to store tool calls by index for chunk in tool_call_chunks: choices = chunk["choices"] @@ -231,15 +219,9 @@ class ChunkProcessor: # Check if tool_call has function (either as attribute or dict key) has_function = False if isinstance(tool_call, dict): - has_function = ( - "function" in tool_call - and tool_call["function"] is not None - ) + has_function = "function" in tool_call and tool_call["function"] is not None else: - has_function = ( - hasattr(tool_call, "function") - and tool_call.function is not None - ) + has_function = hasattr(tool_call, "function") and tool_call.function is not None if not has_function: continue @@ -271,17 +253,13 @@ class ChunkProcessor: if function.get("name"): tool_call_map[index]["name"] = function["name"] if function.get("arguments"): - tool_call_map[index]["arguments"].append( - function["arguments"] - ) + tool_call_map[index]["arguments"].append(function["arguments"]) else: # function is an object if hasattr(function, "name") and function.name: tool_call_map[index]["name"] = function.name if hasattr(function, "arguments") and function.arguments: - tool_call_map[index]["arguments"].append( - function.arguments - ) + tool_call_map[index]["arguments"].append(function.arguments) else: # tool_call is an object if hasattr(tool_call, "id") and tool_call.id: @@ -289,52 +267,33 @@ class ChunkProcessor: if hasattr(tool_call, "type") and tool_call.type: tool_call_map[index]["type"] = tool_call.type if hasattr(tool_call, "function"): - if ( - hasattr(tool_call.function, "name") - and tool_call.function.name - ): + if hasattr(tool_call.function, "name") and tool_call.function.name: tool_call_map[index]["name"] = tool_call.function.name - if ( - hasattr(tool_call.function, "arguments") - and tool_call.function.arguments - ): - tool_call_map[index]["arguments"].append( - tool_call.function.arguments - ) + if hasattr(tool_call.function, "arguments") and tool_call.function.arguments: + tool_call_map[index]["arguments"].append(tool_call.function.arguments) # Preserve provider_specific_fields from streaming chunks provider_fields = None if isinstance(tool_call, dict): provider_fields = tool_call.get("provider_specific_fields") - if not provider_fields and isinstance( - tool_call.get("function"), dict - ): - provider_fields = tool_call["function"].get( - "provider_specific_fields" - ) + if not provider_fields and isinstance(tool_call.get("function"), dict): + provider_fields = tool_call["function"].get("provider_specific_fields") else: - if ( - hasattr(tool_call, "provider_specific_fields") - and tool_call.provider_specific_fields - ): + if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields: provider_fields = tool_call.provider_specific_fields elif ( hasattr(tool_call, "function") and hasattr(tool_call.function, "provider_specific_fields") and tool_call.function.provider_specific_fields ): - provider_fields = ( - tool_call.function.provider_specific_fields - ) + provider_fields = tool_call.function.provider_specific_fields if provider_fields: # Merge provider_specific_fields if multiple chunks have them if tool_call_map[index]["provider_specific_fields"] is None: tool_call_map[index]["provider_specific_fields"] = {} if isinstance(provider_fields, dict): - tool_call_map[index]["provider_specific_fields"].update( - provider_fields - ) + tool_call_map[index]["provider_specific_fields"].update(provider_fields) # Convert the map to a list of tool calls for index in sorted(tool_call_map.keys()): @@ -357,18 +316,14 @@ class ChunkProcessor: # Add provider_specific_fields if present (for thought signatures in Gemini 3) if tool_call_data.get("provider_specific_fields"): - tool_call_params["provider_specific_fields"] = tool_call_data[ - "provider_specific_fields" - ] + tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"] tool_call = ChatCompletionMessageToolCall(**tool_call_params) tool_calls_list.append(tool_call) return tool_calls_list - def get_combined_function_call_content( - self, function_call_chunks: List[Dict[str, Any]] - ) -> FunctionCall: + def get_combined_function_call_content(self, function_call_chunks: List[Dict[str, Any]]) -> FunctionCall: argument_list = [] delta = function_call_chunks[0]["choices"][0]["delta"] function_call = delta.get("function_call", "") @@ -414,19 +369,13 @@ class ChunkProcessor: def get_combined_thinking_content( self, chunks: List[Dict[str, Any]] - ) -> Optional[ - List[ - Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"] - ] - ]: + ) -> Optional[List[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]]]: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock, ) - thinking_blocks: List[ - Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"] - ] = [] + thinking_blocks: List[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] = [] current_thinking_text_parts: List[str] = [] current_signature: Optional[str] = None @@ -476,14 +425,10 @@ class ChunkProcessor: return thinking_blocks return None - def get_combined_reasoning_content( - self, chunks: List[Dict[str, Any]] - ) -> ChatCompletionAssistantContentValue: + def get_combined_reasoning_content(self, chunks: List[Dict[str, Any]]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content( - self, chunks: List[Dict[str, Any]] - ) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: List[Dict[str, Any]]) -> ChatCompletionAudioResponse: base64_data_list: List[str] = [] transcript_list: List[str] = [] expires_at: Optional[int] = None @@ -532,21 +477,13 @@ class ChunkProcessor: cache_read_input_tokens = usage_chunk.get("cache_read_input_tokens") if hasattr(usage_chunk, "completion_tokens_details"): if isinstance(usage_chunk.completion_tokens_details, dict): - completion_tokens_details = CompletionTokensDetails( - **usage_chunk.completion_tokens_details - ) - elif isinstance( - usage_chunk.completion_tokens_details, CompletionTokensDetails - ): + completion_tokens_details = CompletionTokensDetails(**usage_chunk.completion_tokens_details) + elif isinstance(usage_chunk.completion_tokens_details, CompletionTokensDetails): completion_tokens_details = usage_chunk.completion_tokens_details if hasattr(usage_chunk, "prompt_tokens_details"): if isinstance(usage_chunk.prompt_tokens_details, dict): - prompt_tokens_details = PromptTokensDetailsWrapper( - **usage_chunk.prompt_tokens_details - ) - elif isinstance( - usage_chunk.prompt_tokens_details, PromptTokensDetailsWrapper - ): + prompt_tokens_details = PromptTokensDetailsWrapper(**usage_chunk.prompt_tokens_details) + elif isinstance(usage_chunk.prompt_tokens_details, PromptTokensDetailsWrapper): prompt_tokens_details = usage_chunk.prompt_tokens_details return { @@ -608,49 +545,31 @@ class ChunkProcessor: usage_chunk: Optional[Usage] = None if "usage" in chunk: usage_chunk = chunk["usage"] - elif ( - isinstance(chunk, ModelResponse) - or isinstance(chunk, ModelResponseStream) - ) and hasattr(chunk, "_hidden_params"): + elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr( + chunk, "_hidden_params" + ): usage_chunk = chunk._hidden_params.get("usage", None) if usage_chunk is not None: if isinstance(usage_chunk, dict): usage_chunk = Usage(**usage_chunk) usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk) - if ( - usage_chunk_dict["prompt_tokens"] is not None - and usage_chunk_dict["prompt_tokens"] > 0 - ): + if usage_chunk_dict["prompt_tokens"] is not None and usage_chunk_dict["prompt_tokens"] > 0: prompt_tokens = usage_chunk_dict["prompt_tokens"] - if ( - usage_chunk_dict["completion_tokens"] is not None - and usage_chunk_dict["completion_tokens"] > 0 - ): + if usage_chunk_dict["completion_tokens"] is not None and usage_chunk_dict["completion_tokens"] > 0: completion_tokens = usage_chunk_dict["completion_tokens"] completion_usage_updates += 1 if usage_chunk_dict["cache_creation_input_tokens"] is not None and ( - usage_chunk_dict["cache_creation_input_tokens"] > 0 - or cache_creation_input_tokens is None + usage_chunk_dict["cache_creation_input_tokens"] > 0 or cache_creation_input_tokens is None ): - cache_creation_input_tokens = usage_chunk_dict[ - "cache_creation_input_tokens" - ] + cache_creation_input_tokens = usage_chunk_dict["cache_creation_input_tokens"] if usage_chunk_dict["cache_read_input_tokens"] is not None and ( - usage_chunk_dict["cache_read_input_tokens"] > 0 - or cache_read_input_tokens is None + usage_chunk_dict["cache_read_input_tokens"] > 0 or cache_read_input_tokens is None ): - cache_read_input_tokens = usage_chunk_dict[ - "cache_read_input_tokens" - ] + cache_read_input_tokens = usage_chunk_dict["cache_read_input_tokens"] if usage_chunk_dict["completion_tokens_details"] is not None: - completion_tokens_details = usage_chunk_dict[ - "completion_tokens_details" - ] - if ( - hasattr(usage_chunk, "server_tool_use") - and usage_chunk.server_tool_use is not None - ): + completion_tokens_details = usage_chunk_dict["completion_tokens_details"] + if hasattr(usage_chunk, "server_tool_use") and usage_chunk.server_tool_use is not None: # Coerce dict to ServerToolUse so downstream cost-calc code # (which accesses .web_search_requests as an attribute) # doesn't raise AttributeError. Some providers / streaming @@ -660,9 +579,7 @@ class ChunkProcessor: elif isinstance(usage_chunk.server_tool_use, ServerToolUse): server_tool_use = usage_chunk.server_tool_use else: - server_tool_use = ServerToolUse.model_validate( - usage_chunk.server_tool_use - ) + server_tool_use = ServerToolUse.model_validate(usage_chunk.server_tool_use) if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( @@ -717,9 +634,7 @@ class ChunkProcessor: does not silently affect other providers that may legitimately report ``completion_tokens=1`` from a single usage event. """ - saw_non_cursor_completion = ( - completion_tokens > 1 or completion_usage_updates >= 2 - ) + saw_non_cursor_completion = completion_tokens > 1 or completion_usage_updates >= 2 if saw_non_cursor_completion: return completion_tokens @@ -755,30 +670,20 @@ class ChunkProcessor: prompt_tokens = calculated_usage_per_chunk["prompt_tokens"] completion_tokens = calculated_usage_per_chunk["completion_tokens"] ## anthropic prompt caching information ## - cache_creation_input_tokens: Optional[int] = calculated_usage_per_chunk[ - "cache_creation_input_tokens" - ] - cache_read_input_tokens: Optional[int] = calculated_usage_per_chunk[ - "cache_read_input_tokens" - ] + cache_creation_input_tokens: Optional[int] = calculated_usage_per_chunk["cache_creation_input_tokens"] + cache_read_input_tokens: Optional[int] = calculated_usage_per_chunk["cache_read_input_tokens"] - server_tool_use: Optional[ServerToolUse] = calculated_usage_per_chunk[ - "server_tool_use" + server_tool_use: Optional[ServerToolUse] = calculated_usage_per_chunk["server_tool_use"] + web_search_requests: Optional[int] = calculated_usage_per_chunk["web_search_requests"] + completion_tokens_details: Optional[CompletionTokensDetails] = calculated_usage_per_chunk[ + "completion_tokens_details" ] - web_search_requests: Optional[int] = calculated_usage_per_chunk[ - "web_search_requests" + prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = calculated_usage_per_chunk[ + "prompt_tokens_details" ] - completion_tokens_details: Optional[CompletionTokensDetails] = ( - calculated_usage_per_chunk["completion_tokens_details"] - ) - prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = ( - calculated_usage_per_chunk["prompt_tokens_details"] - ) try: - returned_usage.prompt_tokens = prompt_tokens or token_counter( - model=model, messages=messages - ) + returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 @@ -790,9 +695,7 @@ class ChunkProcessor: count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages ) ) - returned_usage.total_tokens = ( - returned_usage.prompt_tokens + returned_usage.completion_tokens - ) + returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens if cache_creation_input_tokens is not None: returned_usage._cache_creation_input_tokens = cache_creation_input_tokens @@ -803,31 +706,25 @@ class ChunkProcessor: ) # for anthropic if cache_read_input_tokens is not None: returned_usage._cache_read_input_tokens = cache_read_input_tokens - setattr( - returned_usage, "cache_read_input_tokens", cache_read_input_tokens - ) # for anthropic + setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens) # for anthropic if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): - returned_usage.completion_tokens_details = ( - CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() - ) + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( + **completion_tokens_details.model_dump() ) else: returned_usage.completion_tokens_details = completion_tokens_details if reasoning_tokens is not None: if returned_usage.completion_tokens_details is None: - returned_usage.completion_tokens_details = ( - CompletionTokensDetailsWrapper(reasoning_tokens=reasoning_tokens) + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( + reasoning_tokens=reasoning_tokens ) elif ( returned_usage.completion_tokens_details is not None and returned_usage.completion_tokens_details.reasoning_tokens is None ): - returned_usage.completion_tokens_details.reasoning_tokens = ( - reasoning_tokens - ) + returned_usage.completion_tokens_details.reasoning_tokens = reasoning_tokens if prompt_tokens_details is not None: returned_usage.prompt_tokens_details = prompt_tokens_details @@ -839,9 +736,7 @@ class ChunkProcessor: web_search_requests=web_search_requests ) else: - returned_usage.prompt_tokens_details.web_search_requests = ( - web_search_requests - ) + returned_usage.prompt_tokens_details.web_search_requests = web_search_requests # Return a new usage object with the new values diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7d7a8e562cb..587a3a58a94 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -134,18 +134,14 @@ class CustomStreamWrapper: litellm_params: GenericLiteLLMParams = GenericLiteLLMParams( **self.logging_obj.model_call_details.get("litellm_params", {}) ) - self.merge_reasoning_content_in_choices: bool = ( - litellm_params.merge_reasoning_content_in_choices or False - ) + self.merge_reasoning_content_in_choices: bool = litellm_params.merge_reasoning_content_in_choices or False self.sent_first_thinking_block = False self.sent_last_thinking_block = False self.thinking_content = "" self.system_fingerprint: Optional[str] = None self.received_finish_reason: Optional[str] = None - self.intermittent_finish_reason: Optional[str] = ( - None # finish reasons that show up mid-stream - ) + self.intermittent_finish_reason: Optional[str] = None # finish reasons that show up mid-stream self.special_tokens = [ "<|assistant|>", "<|system|>", @@ -162,9 +158,7 @@ class CustomStreamWrapper: _api_base = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get( - "litellm_params", {} - ), + optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) self._hidden_params = { @@ -180,14 +174,10 @@ class CustomStreamWrapper: self.response_id: Optional[str] = None self.logging_loop = None self.rules = Rules() - self.stream_options = stream_options or getattr( - logging_obj, "stream_options", None - ) + self.stream_options = stream_options or getattr(logging_obj, "stream_options", None) self.messages = getattr(logging_obj, "messages", None) self.sent_stream_usage = False - self.send_stream_usage = ( - True if self.check_send_stream_usage(self.stream_options) else False - ) + self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False self.tool_call = False self.chunks: List = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options self._repeated_messages_count = 1 @@ -195,18 +185,11 @@ class CustomStreamWrapper: self.created: Optional[int] = None self._last_returned_hidden_params: Optional[dict] = None - _cached_logging_provider = self.logging_obj.model_call_details.get( - "custom_llm_provider", None - ) + _cached_logging_provider = self.logging_obj.model_call_details.get("custom_llm_provider", None) self._cached_logging_llm_provider: Optional[str] = _cached_logging_provider _effective_model = model or "" - if ( - custom_llm_provider == "openai" - and custom_llm_provider != _cached_logging_provider - ): - _effective_model = "{}/{}".format( - _cached_logging_provider, _effective_model - ) + if custom_llm_provider == "openai" and custom_llm_provider != _cached_logging_provider: + _effective_model = "{}/{}".format(_cached_logging_provider, _effective_model) self._cached_model_name: str = _effective_model # Snapshot assumes self._hidden_params is populated from litellm_params @@ -261,19 +244,14 @@ class CustomStreamWrapper: ) def check_send_stream_usage(self, stream_options: Optional[dict]): - return ( - stream_options is not None - and stream_options.get("include_usage", False) is True - ) + return stream_options is not None and stream_options.get("include_usage", False) is True def check_is_function_call(self, logging_obj) -> bool: from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_function_call, ) - if hasattr(logging_obj, "optional_params") and isinstance( - logging_obj.optional_params, dict - ): + if hasattr(logging_obj, "optional_params") and isinstance(logging_obj.optional_params, dict): if is_function_call(logging_obj.optional_params): return True @@ -316,9 +294,7 @@ class CustomStreamWrapper: last_content = self.chunks[-1].choices[0].delta.content if ( - last_content is None - or not isinstance(last_content, str) - or len(last_content) <= 2 + last_content is None or not isinstance(last_content, str) or len(last_content) <= 2 ): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946 self._repeated_messages_count = 1 return @@ -333,9 +309,7 @@ class CustomStreamWrapper: if self._repeated_messages_count >= litellm.REPEATED_STREAMING_CHUNK_LIMIT: # All last n chunks are identical raise litellm.InternalServerError( - message="The model is repeating the same chunk = {}.".format( - last_content - ), + message="The model is repeating the same chunk = {}.".format(last_content), model="", llm_provider="", ) @@ -378,9 +352,7 @@ class CustomStreamWrapper: def handle_predibase_chunk(self, chunk): try: if not isinstance(chunk, str): - chunk = chunk.decode( - "utf-8" - ) # DO NOT REMOVE this: This is required for HF inference API + Streaming + chunk = chunk.decode("utf-8") # DO NOT REMOVE this: This is required for HF inference API + Streaming text = "" is_finished = False finish_reason = "" @@ -390,14 +362,10 @@ class CustomStreamWrapper: print_verbose(f"data json: {data_json}") if "token" in data_json and "text" in data_json["token"]: text = data_json["token"]["text"] - if data_json.get("details", False) and data_json["details"].get( - "finish_reason", False - ): + if data_json.get("details", False) and data_json["details"].get("finish_reason", False): is_finished = True finish_reason = data_json["details"]["finish_reason"] - elif data_json.get( - "generated_text", False - ): # if full generated text exists, then stream is complete + elif data_json.get("generated_text", False): # if full generated text exists, then stream is complete text = "" # don't return the final bos token is_finished = True finish_reason = "stop" @@ -509,18 +477,14 @@ class CustomStreamWrapper: if data_json["choices"][0].get("finish_reason", None): is_finished = True finish_reason = data_json["choices"][0]["finish_reason"] - print_verbose( - f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}" - ) + print_verbose(f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}") return { "text": text, "is_finished": is_finished, "finish_reason": finish_reason, } except Exception: - raise ValueError( - f"Unable to parse response. Original response: {chunk}" - ) + raise ValueError(f"Unable to parse response. Original response: {chunk}") elif "error" in chunk: raise ValueError(f"Unable to parse response. Original response: {chunk}") else: @@ -560,29 +524,18 @@ class CustomStreamWrapper: logprobs = None usage = None if str_line and str_line.choices and len(str_line.choices) > 0: - if ( - str_line.choices[0].delta is not None - and str_line.choices[0].delta.content is not None - ): + if str_line.choices[0].delta is not None and str_line.choices[0].delta.content is not None: text = str_line.choices[0].delta.content else: # function/tool calling chunk - when content is None. in this case we just return the original chunk from openai pass if str_line.choices[0].finish_reason: - is_finished = ( - True # check if str_line._hidden_params["is_finished"] is True - ) - if ( - hasattr(str_line, "_hidden_params") - and str_line._hidden_params.get("is_finished") is not None - ): + is_finished = True # check if str_line._hidden_params["is_finished"] is True + if hasattr(str_line, "_hidden_params") and str_line._hidden_params.get("is_finished") is not None: is_finished = str_line._hidden_params.get("is_finished") finish_reason = str_line.choices[0].finish_reason # checking for logprobs - if ( - hasattr(str_line.choices[0], "logprobs") - and str_line.choices[0].logprobs is not None - ): + if hasattr(str_line.choices[0], "logprobs") and str_line.choices[0].logprobs is not None: logprobs = str_line.choices[0].logprobs else: logprobs = None @@ -663,23 +616,17 @@ class CustomStreamWrapper: return data_json["model_output"]["data"][0] elif isinstance(data_json["model_output"], str): return data_json["model_output"] - elif "completion" in data_json and isinstance( - data_json["completion"], str - ): + elif "completion" in data_json and isinstance(data_json["completion"], str): return data_json["completion"] else: - raise ValueError( - f"Unable to parse response. Original response: {chunk}" - ) + raise ValueError(f"Unable to parse response. Original response: {chunk}") else: return "" else: return "" except Exception as e: verbose_logger.exception( - "litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - {}".format( - str(e) - ) + "litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - {}".format(str(e)) ) return "" @@ -691,9 +638,7 @@ class CustomStreamWrapper: if isinstance(chunk, bytes): chunk = chunk.decode("utf-8") if "text_output" in chunk: - response = ( - CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" - ) + response = CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" response = response.strip() parsed_response = json.loads(response) else: @@ -705,9 +650,7 @@ class CustomStreamWrapper: } else: print_verbose(f"chunk: {chunk} (Type: {type(chunk)})") - raise ValueError( - f"Unable to parse response. Original response: {chunk}" - ) + raise ValueError(f"Unable to parse response. Original response: {chunk}") text = parsed_response.get("text_output", "") finish_reason = parsed_response.get("stop_reason") is_finished = parsed_response.get("is_finished", False) @@ -722,9 +665,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def model_response_creator( - self, chunk: Optional[dict] = None, hidden_params: Optional[dict] = None - ): + def model_response_creator(self, chunk: Optional[dict] = None, hidden_params: Optional[dict] = None): _model = self._cached_model_name _logging_obj_llm_provider = self._cached_logging_llm_provider @@ -766,10 +707,7 @@ class CustomStreamWrapper: **self._base_hidden_params, } - if ( - len(model_response.choices) > 0 - and getattr(model_response.choices[0], "delta") is not None - ): + if len(model_response.choices) > 0 and getattr(model_response.choices[0], "delta") is not None: # do nothing, if object instantiated pass else: @@ -786,9 +724,7 @@ class CustomStreamWrapper: is_empty = False return is_empty - def set_model_id( - self, id: str, model_response: ModelResponseStream - ) -> ModelResponseStream: + def set_model_id(self, id: str, model_response: ModelResponseStream) -> ModelResponseStream: """ Set the model id and response id to the given id. @@ -814,9 +750,7 @@ class CustomStreamWrapper: """ Copy provider_specific_fields from original_chunk to model_response. """ - provider_specific_fields = getattr( - original_chunk, "provider_specific_fields", None - ) + provider_specific_fields = getattr(original_chunk, "provider_specific_fields", None) if provider_specific_fields is not None: model_response.provider_specific_fields = provider_specific_fields for k, v in provider_specific_fields.items(): @@ -831,19 +765,13 @@ class CustomStreamWrapper: ) -> bool: if ( "content" in completion_obj - and ( - isinstance(completion_obj["content"], str) - and len(completion_obj["content"]) > 0 - ) + and (isinstance(completion_obj["content"], str) and len(completion_obj["content"]) > 0) or ( "tool_calls" in completion_obj and completion_obj["tool_calls"] is not None and len(completion_obj["tool_calls"]) > 0 ) - or ( - "function_call" in completion_obj - and completion_obj["function_call"] is not None - ) + or ("function_call" in completion_obj and completion_obj["function_call"] is not None) or ( "tool_calls" in model_response.choices[0].delta and model_response.choices[0].delta["tool_calls"] is not None @@ -862,10 +790,7 @@ class CustomStreamWrapper: "provider_specific_fields" in model_response and model_response.choices[0].delta.provider_specific_fields is not None ) - or ( - "provider_specific_fields" in response_obj - and response_obj["provider_specific_fields"] is not None - ) + or ("provider_specific_fields" in response_obj and response_obj["provider_specific_fields"] is not None) or ( "annotations" in model_response.choices[0].delta and model_response.choices[0].delta.annotations is not None @@ -875,27 +800,20 @@ class CustomStreamWrapper: and hasattr(model_response.choices[0].delta, "role") and model_response.choices[0].delta.role is not None ) - or ( - getattr(model_response.choices[0].delta, "reasoning_items", None) - is not None - ) + or (getattr(model_response.choices[0].delta, "reasoning_items", None) is not None) ): return True else: return False - def strip_role_from_delta( - self, model_response: ModelResponseStream - ) -> ModelResponseStream: + def strip_role_from_delta(self, model_response: ModelResponseStream) -> ModelResponseStream: """ Strip the role from the delta. """ if self.sent_first_chunk is False: model_response.choices[0].delta["role"] = "assistant" self.sent_first_chunk = True - elif self.sent_first_chunk is True and hasattr( - model_response.choices[0].delta, "role" - ): + elif self.sent_first_chunk is True and hasattr(model_response.choices[0].delta, "role"): _initial_delta = model_response.choices[0].delta.model_dump() _initial_delta.pop("role", None) @@ -919,24 +837,16 @@ class CustomStreamWrapper: return True # Check for audio - if ( - hasattr(delta, AUDIO_ATTRIBUTE) - and getattr(delta, AUDIO_ATTRIBUTE, None) is not None - ): + if hasattr(delta, AUDIO_ATTRIBUTE) and getattr(delta, AUDIO_ATTRIBUTE, None) is not None: return True # Check for image - if ( - hasattr(delta, IMAGE_ATTRIBUTE) - and getattr(delta, IMAGE_ATTRIBUTE, None) is not None - ): + if hasattr(delta, IMAGE_ATTRIBUTE) and getattr(delta, IMAGE_ATTRIBUTE, None) is not None: return True return False - def _handle_special_delta_content( - self, model_response: ModelResponseStream - ) -> ModelResponseStream: + def _handle_special_delta_content(self, model_response: ModelResponseStream) -> ModelResponseStream: """ Handle special delta content types by stripping role and returning the response. """ @@ -948,9 +858,7 @@ class CustomStreamWrapper: """ return delta is not None and getattr(delta, attribute_name, None) is not None - def _copy_delta_attribute( - self, source_delta, target_delta, attribute_name: str - ) -> None: + def _copy_delta_attribute(self, source_delta, target_delta, attribute_name: str) -> None: """ Copy a specific attribute from source delta to target delta. """ @@ -966,18 +874,14 @@ class CustomStreamWrapper: return True return False - def _handle_special_delta_attributes( - self, delta, model_response: "ModelResponseStream" - ) -> None: + def _handle_special_delta_attributes(self, delta, model_response: "ModelResponseStream") -> None: """ Handle special delta attributes (audio, image) by copying them to model_response. """ special_attributes = [AUDIO_ATTRIBUTE, IMAGE_ATTRIBUTE] for attribute in special_attributes: if self._has_special_delta_attribute(delta, attribute): - self._copy_delta_attribute( - delta, model_response.choices[0].delta, attribute - ) + self._copy_delta_attribute(delta, model_response.choices[0].delta, attribute) def return_processed_chunk_logic( # noqa: C901 self, @@ -989,13 +893,9 @@ class CustomStreamWrapper: preserve_upstream_non_openai_attributes, ) - is_chunk_non_empty = self.is_chunk_non_empty( - completion_obj, model_response, response_obj - ) + is_chunk_non_empty = self.is_chunk_non_empty(completion_obj, model_response, response_obj) - if ( - is_chunk_non_empty - ): # cannot set content of an OpenAI Object to be an empty string + if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string self.raise_on_model_repetition() hold, model_response_str = self.check_special_tokens( chunk=completion_obj["content"], @@ -1021,9 +921,7 @@ class CustomStreamWrapper: setattr(model_response, "choices", choices) else: return - model_response.system_fingerprint = ( - original_chunk.system_fingerprint - ) + model_response.system_fingerprint = original_chunk.system_fingerprint setattr( model_response, "citations", @@ -1047,17 +945,13 @@ class CustomStreamWrapper: completion_obj["role"] = "assistant" self.sent_first_chunk = True if response_obj.get("provider_specific_fields") is not None: - completion_obj["provider_specific_fields"] = response_obj[ - "provider_specific_fields" - ] + completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] model_response.choices[0].delta = Delta(**completion_obj) _index: Optional[int] = completion_obj.get("index") if _index is not None: model_response.choices[0].index = _index - self._optional_combine_thinking_block_in_choices( - model_response=model_response - ) + self._optional_combine_thinking_block_in_choices(model_response=model_response) return model_response else: @@ -1094,9 +988,7 @@ class CustomStreamWrapper: ) if _is_delta_empty: - model_response.choices[0].delta = Delta( - content=None - ) # ensure empty delta chunk returned + model_response.choices[0].delta = Delta(content=None) # ensure empty delta chunk returned # get any function call arguments model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason @@ -1112,9 +1004,7 @@ class CustomStreamWrapper: self.chunks.append(model_response) return - def _optional_combine_thinking_block_in_choices( - self, model_response: ModelResponseStream - ) -> None: + def _optional_combine_thinking_block_in_choices(self, model_response: ModelResponseStream) -> None: """ UI's Like OpenWebUI expect to get 1 chunk with ... tags in the chunk content @@ -1125,17 +1015,13 @@ class CustomStreamWrapper: """ if self.merge_reasoning_content_in_choices is True: - reasoning_content = getattr( - model_response.choices[0].delta, "reasoning_content", None - ) + reasoning_content = getattr(model_response.choices[0].delta, "reasoning_content", None) if reasoning_content: if self.sent_first_thinking_block is False: # Ensure content is not None before concatenation if model_response.choices[0].delta.content is None: model_response.choices[0].delta.content = "" - model_response.choices[0].delta.content += ( - "" + reasoning_content - ) + model_response.choices[0].delta.content += "" + reasoning_content self.sent_first_thinking_block = True elif ( self.sent_first_thinking_block is True @@ -1148,9 +1034,7 @@ class CustomStreamWrapper: and not self.sent_last_thinking_block and model_response.choices[0].delta.content ): - model_response.choices[0].delta.content = "" + ( - model_response.choices[0].delta.content or "" - ) + model_response.choices[0].delta.content = "" + (model_response.choices[0].delta.content or "") self.sent_last_thinking_block = True if hasattr(model_response.choices[0].delta, "reasoning_content"): @@ -1172,9 +1056,7 @@ class CustomStreamWrapper: _has_content = bool( chunk.choices and chunk.choices[0].delta is not None - and ( - chunk.choices[0].delta.content or chunk.choices[0].delta.tool_calls - ) + and (chunk.choices[0].delta.content or chunk.choices[0].delta.tool_calls) ) if self.received_finish_reason is not None: if not _has_content: @@ -1191,13 +1073,8 @@ class CustomStreamWrapper: if ( isinstance(chunk, dict) - and generic_chunk_has_all_required_fields( - chunk=chunk - ) # check if chunk is a generic streaming chunk - ) or ( - self.custom_llm_provider - and self.custom_llm_provider in litellm._custom_providers - ): + and generic_chunk_has_all_required_fields(chunk=chunk) # check if chunk is a generic streaming chunk + ) or (self.custom_llm_provider and self.custom_llm_provider in litellm._custom_providers): if self.received_finish_reason is not None: _chunk_has_content = isinstance(chunk, dict) and ( bool(chunk.get("text", "")) @@ -1206,10 +1083,7 @@ class CustomStreamWrapper: # finish_reason/usage to downstream translators. or chunk.get("usage") is not None ) - if not _chunk_has_content and ( - not isinstance(chunk, dict) - or "provider_specific_fields" not in chunk - ): + if not _chunk_has_content and (not isinstance(chunk, dict) or "provider_specific_fields" not in chunk): raise StopIteration anthropic_response_obj: GChunk = cast(GChunk, chunk) completion_obj["content"] = anthropic_response_obj["text"] @@ -1217,9 +1091,7 @@ class CustomStreamWrapper: self.received_finish_reason = anthropic_response_obj["finish_reason"] if anthropic_response_obj["finish_reason"]: - self.intermittent_finish_reason = anthropic_response_obj[ - "finish_reason" - ] + self.intermittent_finish_reason = anthropic_response_obj["finish_reason"] if anthropic_response_obj["usage"] is not None: setattr( @@ -1228,19 +1100,14 @@ class CustomStreamWrapper: litellm.Usage(**anthropic_response_obj["usage"]), ) - if ( - "tool_use" in anthropic_response_obj - and anthropic_response_obj["tool_use"] is not None - ): + if "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None: completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] if ( "provider_specific_fields" in anthropic_response_obj and anthropic_response_obj["provider_specific_fields"] is not None ): - for key, value in anthropic_response_obj[ - "provider_specific_fields" - ].items(): + for key, value in anthropic_response_obj["provider_specific_fields"].items(): setattr(model_response, key, value) response_obj = cast(dict[str, Any], anthropic_response_obj) @@ -1254,13 +1121,9 @@ class CustomStreamWrapper: completion_obj["content"] = response_obj["text"] if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] - elif ( - self.custom_llm_provider and self.custom_llm_provider == "baseten" - ): # baseten doesn't provide streaming + elif self.custom_llm_provider and self.custom_llm_provider == "baseten": # baseten doesn't provide streaming completion_obj["content"] = self.handle_baseten_chunk(chunk) - elif ( - self.custom_llm_provider and self.custom_llm_provider == "ai21" - ): # ai21 doesn't provide streaming + elif self.custom_llm_provider and self.custom_llm_provider == "ai21": # ai21 doesn't provide streaming response_obj = self.handle_ai21_chunk(chunk) completion_obj["content"] = response_obj["text"] if response_obj["is_finished"]: @@ -1292,9 +1155,7 @@ class CustomStreamWrapper: if self.sent_first_chunk is False: raise Exception("An unknown error occurred with the stream") self.received_finish_reason = "stop" - elif self.custom_llm_provider == "vertex_ai" and not isinstance( - chunk, ModelResponseStream - ): + elif self.custom_llm_provider == "vertex_ai" and not isinstance(chunk, ModelResponseStream): chunk = cast(Any, chunk) import proto # type: ignore @@ -1356,9 +1217,7 @@ class CustomStreamWrapper: ) except Exception: if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore - raise Exception( - f"The response was blocked by VertexAI. {str(chunk)}" - ) + raise Exception(f"The response was blocked by VertexAI. {str(chunk)}") else: completion_obj["content"] = str(chunk) elif self.custom_llm_provider == "petals": @@ -1443,9 +1302,7 @@ class CustomStreamWrapper: "finish_reason": chunk_finish_reason, "original_chunk": chunk, "tool_calls": ( - chunk.choices[0].delta.tool_calls - if hasattr(chunk.choices[0].delta, "tool_calls") - else None + chunk.choices[0].delta.tool_calls if hasattr(chunk.choices[0].delta, "tool_calls") else None ), } @@ -1483,16 +1340,10 @@ class CustomStreamWrapper: self.received_finish_reason = response_obj["finish_reason"] if response_obj.get("original_chunk", None) is not None: if hasattr(response_obj["original_chunk"], "id"): - model_response = self.set_model_id( - response_obj["original_chunk"].id, model_response - ) + model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) if hasattr(response_obj["original_chunk"], "system_fingerprint"): - model_response.system_fingerprint = response_obj[ - "original_chunk" - ].system_fingerprint - self.system_fingerprint = response_obj[ - "original_chunk" - ].system_fingerprint + model_response.system_fingerprint = response_obj["original_chunk"].system_fingerprint + self.system_fingerprint = response_obj["original_chunk"].system_fingerprint if response_obj["logprobs"] is not None: model_response.choices[0].logprobs = response_obj["logprobs"] @@ -1502,16 +1353,9 @@ class CustomStreamWrapper: model_response, "usage", litellm.Usage( - prompt_tokens=response_obj["usage"].get( - "prompt_tokens", None - ) - or None, - completion_tokens=response_obj["usage"].get( - "completion_tokens", None - ) - or None, - total_tokens=response_obj["usage"].get("total_tokens", None) - or None, + prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None, + completion_tokens=response_obj["usage"].get("completion_tokens", None) or None, + total_tokens=response_obj["usage"].get("total_tokens", None) or None, ), ) elif isinstance(response_obj["usage"], Usage): @@ -1547,37 +1391,24 @@ class CustomStreamWrapper: model_response.model = self.model ## FUNCTION CALL PARSING - original_chunk = ( - response_obj.get("original_chunk") if response_obj is not None else None - ) + original_chunk = response_obj.get("original_chunk") if response_obj is not None else None if ( original_chunk is not None ): # function / tool calling branch - only set for openai/azure compatible endpoints # enter this branch when no content has been passed in response if hasattr(original_chunk, "id"): - model_response = self.set_model_id( - original_chunk.id, model_response - ) + model_response = self.set_model_id(original_chunk.id, model_response) if hasattr(original_chunk, "provider_specific_fields"): - model_response = ( - self.copy_model_response_level_provider_specific_fields( - original_chunk, model_response - ) + model_response = self.copy_model_response_level_provider_specific_fields( + original_chunk, model_response ) if original_chunk.choices and len(original_chunk.choices) > 0: delta = original_chunk.choices[0].delta - if delta is not None and ( - delta.function_call is not None or delta.tool_calls is not None - ): + if delta is not None and (delta.function_call is not None or delta.tool_calls is not None): try: - model_response.system_fingerprint = ( - original_chunk.system_fingerprint - ) + model_response.system_fingerprint = original_chunk.system_fingerprint ## AZURE - check if arguments is not None - if ( - original_chunk.choices[0].delta.function_call - is not None - ): + if original_chunk.choices[0].delta.function_call is not None: if ( getattr( original_chunk.choices[0].delta.function_call, @@ -1585,17 +1416,11 @@ class CustomStreamWrapper: ) is None ): - original_chunk.choices[ - 0 - ].delta.function_call.arguments = "" + original_chunk.choices[0].delta.function_call.arguments = "" elif original_chunk.choices[0].delta.tool_calls is not None: - if isinstance( - original_chunk.choices[0].delta.tool_calls, list - ): + if isinstance(original_chunk.choices[0].delta.tool_calls, list): for t in original_chunk.choices[0].delta.tool_calls: - if hasattr(t, "functions") and hasattr( - t.functions, "arguments" - ): + if hasattr(t, "functions") and hasattr(t.functions, "arguments"): if ( getattr( t.function, @@ -1606,12 +1431,8 @@ class CustomStreamWrapper: t.function.arguments = "" _json_delta = delta.model_dump() if "role" not in _json_delta or _json_delta["role"] is None: - _json_delta["role"] = ( - "assistant" # mistral's api returns role as None - ) - if "tool_calls" in _json_delta and isinstance( - _json_delta["tool_calls"], list - ): + _json_delta["role"] = "assistant" # mistral's api returns role as None + if "tool_calls" in _json_delta and isinstance(_json_delta["tool_calls"], list): for tool in _json_delta["tool_calls"]: if ( isinstance(tool, dict) @@ -1624,9 +1445,7 @@ class CustomStreamWrapper: model_response.choices[0].delta = Delta(**_json_delta) except Exception as e: verbose_logger.exception( - "litellm.CustomStreamWrapper.chunk_creator(): Exception occured - {}".format( - str(e) - ) + "litellm.CustomStreamWrapper.chunk_creator(): Exception occured - {}".format(str(e)) ) model_response.choices[0].delta = Delta() elif self._has_any_special_delta_attributes(delta): @@ -1642,10 +1461,7 @@ class CustomStreamWrapper: except Exception: model_response.choices[0].delta = Delta() else: - if ( - self.stream_options is not None - and self.stream_options["include_usage"] is True - ): + if self.stream_options is not None and self.stream_options["include_usage"] is True: model_response.choices = [] return model_response return @@ -1653,9 +1469,7 @@ class CustomStreamWrapper: if "tool_calls" in completion_obj and len(completion_obj["tool_calls"]) > 0: if self.is_function_call is True: # user passed in 'functions' param - completion_obj["function_call"] = completion_obj["tool_calls"][0][ - "function" - ] + completion_obj["function_call"] = completion_obj["tool_calls"][0]["function"] completion_obj["tool_calls"] = None self.tool_call = True @@ -1708,8 +1522,7 @@ class CustomStreamWrapper: self._post_streaming_hooks = [ cb for cb in litellm.callbacks - if isinstance(cb, CustomLogger) - and hasattr(cb, "async_post_call_streaming_deployment_hook") + if isinstance(cb, CustomLogger) and hasattr(cb, "async_post_call_streaming_deployment_hook") ] if not self._post_streaming_hooks: @@ -1736,14 +1549,10 @@ class CustomStreamWrapper: except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception( - f"Error in post-call streaming deployment hook: {str(e)}" - ) + verbose_logger.exception(f"Error in post-call streaming deployment hook: {str(e)}") return chunk - def _add_mcp_list_tools_to_first_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: + def _add_mcp_list_tools_to_first_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """ Add mcp_list_tools from _hidden_params to the first chunk's delta.provider_specific_fields. @@ -1767,37 +1576,24 @@ class CustomStreamWrapper: # Add mcp_list_tools to delta.provider_specific_fields if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(choice.delta, "provider_specific_fields", None) - or {} - ) + provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {} # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = mcp_list_tools # Set the provider_specific_fields - setattr( - choice.delta, "provider_specific_fields", provider_fields - ) + setattr(choice.delta, "provider_specific_fields", provider_fields) except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception( - f"Error adding MCP list tools to first chunk: {str(e)}" - ) + verbose_logger.exception(f"Error adding MCP list tools to first chunk: {str(e)}") return chunk - def _add_mcp_metadata_to_final_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: + def _add_mcp_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """ Add MCP metadata from _hidden_params to the final chunk's delta.provider_specific_fields. @@ -1816,32 +1612,21 @@ class CustomStreamWrapper: # Add MCP metadata to delta.provider_specific_fields if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(choice.delta, "provider_specific_fields", None) - or {} - ) + provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {} # Add MCP metadata if isinstance(mcp_metadata, dict): provider_fields.update(mcp_metadata) # Set the provider_specific_fields - setattr( - choice.delta, "provider_specific_fields", provider_fields - ) + setattr(choice.delta, "provider_specific_fields", provider_fields) except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception( - f"Error adding MCP metadata to final chunk: {str(e)}" - ) + verbose_logger.exception(f"Error adding MCP metadata to final chunk: {str(e)}") return chunk @@ -1850,20 +1635,14 @@ class CustomStreamWrapper: Caches the streaming response """ if not cache_hit and self.logging_obj._llm_caching_handler is not None: - self.logging_obj._llm_caching_handler._sync_add_streaming_response_to_cache( - processed_chunk - ) + self.logging_obj._llm_caching_handler._sync_add_streaming_response_to_cache(processed_chunk) async def async_cache_streaming_response(self, processed_chunk, cache_hit: bool): """ Caches the streaming response """ if not cache_hit and self.logging_obj._llm_caching_handler is not None: - await ( - self.logging_obj._llm_caching_handler._add_streaming_response_to_cache( - processed_chunk - ) - ) + await self.logging_obj._llm_caching_handler._add_streaming_response_to_cache(processed_chunk) def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool): """ @@ -1881,18 +1660,12 @@ class CustomStreamWrapper: # Create an event loop for the new thread if self.logging_loop is not None: future = asyncio.run_coroutine_threadsafe( - self.logging_obj.async_success_handler( - processed_chunk, None, None, cache_hit - ), + self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit), loop=self.logging_loop, ) future.result() else: - asyncio.run( - self.logging_obj.async_success_handler( - processed_chunk, None, None, cache_hit - ) - ) + asyncio.run(self.logging_obj.async_success_handler(processed_chunk, None, None, cache_hit)) ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler litellm_params = self.logging_obj.model_call_details.get("litellm_params", {}) if self.logging_obj._is_sync_litellm_request(litellm_params): @@ -1915,10 +1688,7 @@ class CustomStreamWrapper: def __next__(self) -> "ModelResponseStream": cache_hit = False - if ( - self.custom_llm_provider is not None - and self.custom_llm_provider == "cached_response" - ): + if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": cache_hit = True self._check_max_streaming_duration() try: @@ -1938,17 +1708,13 @@ class CustomStreamWrapper: print_verbose( f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk.decode('utf-8', errors='replace') if isinstance(chunk, bytes) else chunk}; custom_llm_provider: {self.custom_llm_provider}" ) - response: Optional[ModelResponseStream] = self.chunk_creator( - chunk=chunk - ) + response: Optional[ModelResponseStream] = self.chunk_creator(chunk=chunk) print_verbose(f"PROCESSED CHUNK POST CHUNK CREATOR: {response}") if response is None: continue if self.logging_obj.completion_start_time is None: - self.logging_obj._update_completion_start_time( - completion_start_time=datetime.datetime.now() - ) + self.logging_obj._update_completion_start_time(completion_start_time=datetime.datetime.now()) ## LOGGING if not litellm.disable_streaming_logging: executor.submit( @@ -1959,14 +1725,10 @@ class CustomStreamWrapper: if response.choices: choice = response.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += ( - choice.delta.get("content", "") or "" - ) + self.response_uptil_now += choice.delta.get("content", "") or "" else: self.response_uptil_now += "" - self.rules.post_call_rules( - input=self.response_uptil_now, model=self.model - ) + self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # HANDLE STREAM OPTIONS self.chunks.append(response) @@ -1984,13 +1746,9 @@ class CustomStreamWrapper: if "usage" in obj_dict: del obj_dict["usage"] - response = self.model_response_creator( - chunk=obj_dict, hidden_params=response._hidden_params - ) + response = self.model_response_creator(chunk=obj_dict, hidden_params=response._hidden_params) ## check if empty - is_empty = is_model_response_stream_empty( - model_response=cast(ModelResponseStream, response) - ) + is_empty = is_model_response_stream_empty(model_response=cast(ModelResponseStream, response)) if is_empty: continue @@ -2019,8 +1777,7 @@ class CustomStreamWrapper: # escape __next__ and drop the request from SpendLogs. Recover # best-effort usage from the raw chunks so cost is still tracked verbose_logger.warning( - "stream_chunk_builder raised at end-of-stream (%s); logging " - "best-effort usage from chunks.", + "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", str(e), ) try: @@ -2098,9 +1855,7 @@ class CustomStreamWrapper: except Exception as e: traceback_exception = traceback.format_exc() # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated - threading.Thread( - target=self.logging_obj.failure_handler, args=(e, traceback_exception) - ).start() + threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start() self._handle_stream_fallback_error(e) def fetch_sync_stream(self): @@ -2114,19 +1869,14 @@ class CustomStreamWrapper: async def fetch_stream(self): if self.completion_stream is None and self.make_call is not None: # Call make_call to get the completion stream - self.completion_stream = await self.make_call( - client=litellm.module_level_aclient - ) + self.completion_stream = await self.make_call(client=litellm.module_level_aclient) self._stream_iter = self.completion_stream.__aiter__() return self.completion_stream async def __anext__(self) -> "ModelResponseStream": cache_hit = False - if ( - self.custom_llm_provider is not None - and self.custom_llm_provider == "cached_response" - ): + if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": cache_hit = True self._check_max_streaming_duration() try: @@ -2138,44 +1888,29 @@ class CustomStreamWrapper: if chunk == "None" or chunk is None: continue # skip None chunks - elif ( - self.custom_llm_provider == "gemini" - and hasattr(chunk, "parts") - and len(chunk.parts) == 0 - ): + elif self.custom_llm_provider == "gemini" and hasattr(chunk, "parts") and len(chunk.parts) == 0: continue - processed_chunk: Optional[ModelResponseStream] = self.chunk_creator( - chunk=chunk - ) + processed_chunk: Optional[ModelResponseStream] = self.chunk_creator(chunk=chunk) if processed_chunk is None: continue if self.logging_obj.completion_start_time is None: - self.logging_obj._update_completion_start_time( - completion_start_time=datetime.datetime.now() - ) + self.logging_obj._update_completion_start_time(completion_start_time=datetime.datetime.now()) if processed_chunk.choices: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += ( - choice.delta.get("content", "") or "" - ) + self.response_uptil_now += choice.delta.get("content", "") or "" else: self.response_uptil_now += "" - self.rules.post_call_rules( - input=self.response_uptil_now, model=self.model - ) + self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # Add mcp_list_tools to first chunk if present if not self.sent_first_chunk and processed_chunk.choices: - processed_chunk = self._add_mcp_list_tools_to_first_chunk( - processed_chunk - ) + processed_chunk = self._add_mcp_list_tools_to_first_chunk(processed_chunk) self.sent_first_chunk = True _has_usage = ( - hasattr(processed_chunk, "usage") - and getattr(processed_chunk, "usage", None) is not None + hasattr(processed_chunk, "usage") and getattr(processed_chunk, "usage", None) is not None ) if _has_usage: @@ -2205,35 +1940,23 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) processed_chunk._hidden_params["usage"] = usage - self._last_returned_hidden_params = ( - processed_chunk._hidden_params - ) + self._last_returned_hidden_params = processed_chunk._hidden_params # Call post-call streaming deployment hook for final chunk if self.sent_last_chunk is True: - processed_chunk = ( - await self._call_post_streaming_deployment_hook( - processed_chunk - ) - ) + processed_chunk = await self._call_post_streaming_deployment_hook(processed_chunk) # Add MCP metadata to final chunk if present (after hooks) - processed_chunk = self._add_mcp_metadata_to_final_chunk( - processed_chunk - ) # type: ignore[reportArgumentType] + processed_chunk = self._add_mcp_metadata_to_final_chunk(processed_chunk) # type: ignore[reportArgumentType] return processed_chunk raise StopAsyncIteration else: # temporary patch for non-aiohttp async calls # example - boto3 bedrock llms while True: - if isinstance(self.completion_stream, str) or isinstance( - self.completion_stream, bytes - ): + if isinstance(self.completion_stream, str) or isinstance(self.completion_stream, bytes): chunk = self.completion_stream else: - chunk = await asyncio.to_thread( - _next_sync_or_exhausted, self.completion_stream - ) # type: ignore[arg-type] + chunk = await asyncio.to_thread(_next_sync_or_exhausted, self.completion_stream) # type: ignore[arg-type] if chunk is _SYNC_ITER_EXHAUSTED: raise StopAsyncIteration if chunk is not None and chunk != b"": @@ -2243,14 +1966,10 @@ class CustomStreamWrapper: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += ( - choice.delta.get("content", "") or "" - ) + self.response_uptil_now += choice.delta.get("content", "") or "" else: self.response_uptil_now += "" - self.rules.post_call_rules( - input=self.response_uptil_now, model=self.model - ) + self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # RETURN RESULT self.chunks.append(processed_chunk) return processed_chunk @@ -2268,8 +1987,7 @@ class CustomStreamWrapper: # except handler escapes __anext__ and drops the request from SpendLogs. # Recover best-effort usage from the raw chunks so cost is still tracked verbose_logger.warning( - "stream_chunk_builder raised at end-of-stream (%s); logging " - "best-effort usage from chunks.", + "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", str(e), ) try: @@ -2350,9 +2068,7 @@ class CustomStreamWrapper: except httpx.TimeoutException as e: # if httpx read timeout error occues traceback_exception = traceback.format_exc() ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT - traceback_exception += "\nLiteLLM Default Request Timeout - {}".format( - litellm.request_timeout - ) + traceback_exception += "\nLiteLLM Default Request Timeout - {}".format(litellm.request_timeout) if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING @@ -2361,9 +2077,7 @@ class CustomStreamWrapper: args=(e, traceback_exception), ).start() # log response # Handle any exceptions that might occur during streaming - asyncio.create_task( - self.logging_obj.async_failure_handler(e, traceback_exception) - ) + asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception)) self._handle_stream_fallback_error(e) except Exception as e: traceback_exception = traceback.format_exc() @@ -2398,8 +2112,7 @@ class CustomStreamWrapper: return self.logging_obj.model_call_details["combined_usage_object"] = usage self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj._response_cost_calculator(result=partial_response) - or 0.0 + self.logging_obj._response_cost_calculator(result=partial_response) or 0.0 ) except Exception as recover_error: verbose_logger.debug( @@ -2460,17 +2173,9 @@ class CustomStreamWrapper: # Raise non-retriable client errors directly (skip fallback). # Exception: 429 (rate-limit) IS retriable/transient — allow it # through so the Router can switch to a different model group. - if ( - mapped_status_code is not None - and 400 <= mapped_status_code < 500 - and mapped_status_code != 429 - ): + if mapped_status_code is not None and 400 <= mapped_status_code < 500 and mapped_status_code != 429: raise mapped_exception - if ( - original_status_code is not None - and 400 <= original_status_code < 500 - and original_status_code != 429 - ): + if original_status_code is not None and 400 <= original_status_code < 500 and original_status_code != 429: raise mapped_exception raise MidStreamFallbackError( diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index c766c6edec1..56b9d42092c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -67,9 +67,7 @@ def get_modified_max_tokens( ## MODEL INFO _model_info = litellm.get_model_info(model=model) - max_output_tokens = litellm.get_max_tokens( - model=base_model - ) # assume min context window is 4k tokens + max_output_tokens = litellm.get_max_tokens(model=base_model) # assume min context window is 4k tokens ## UNKNOWN MAX OUTPUT TOKENS - return user defined amount if max_output_tokens is None: @@ -87,14 +85,10 @@ def get_modified_max_tokens( ) # give at least a 10 token buffer. token counting can be imprecise. input_tokens += int(token_buffer) - verbose_logger.debug( - f"max_output_tokens: {max_output_tokens}, user_max_tokens: {user_max_tokens}" - ) + verbose_logger.debug(f"max_output_tokens: {max_output_tokens}, user_max_tokens: {user_max_tokens}") ## CASE 1: model input + output can't exceed X - happens when max input = max output, e.g. gpt-3.5-turbo if _model_info["max_input_tokens"] == max_output_tokens: - verbose_logger.debug( - f"input_tokens: {input_tokens}, max_output_tokens: {max_output_tokens}" - ) + verbose_logger.debug(f"input_tokens: {input_tokens}, max_output_tokens: {max_output_tokens}") if input_tokens > max_output_tokens: pass # allow call to fail normally - don't set max_tokens to negative. elif ( @@ -131,10 +125,7 @@ def resize_image_high_res( max_long_side = MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES # Return early if no resizing is needed - if ( - width <= MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES - and height <= MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES - ): + if width <= MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES and height <= MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES: return width, height # Determine the longer and shorter sides @@ -296,9 +287,7 @@ def calculate_img_tokens( int: The number of tokens for the image. """ if use_default_image_token_count: - verbose_logger.debug( - "Using default image token count: {}".format(DEFAULT_IMAGE_TOKEN_COUNT) - ) + verbose_logger.debug("Using default image token count: {}".format(DEFAULT_IMAGE_TOKEN_COUNT)) return DEFAULT_IMAGE_TOKEN_COUNT if mode == "low" or mode == "auto": return base_tokens @@ -307,12 +296,8 @@ def calculate_img_tokens( width, height = get_image_dimensions( data=data, ) - resized_width, resized_height = resize_image_high_res( - width=width, height=height - ) - tiles_needed_high_res = calculate_tiles_needed( - resized_width=resized_width, resized_height=resized_height - ) + resized_width, resized_height = resize_image_high_res(width=width, height=height) + tiles_needed_high_res = calculate_tiles_needed(resized_width=resized_width, resized_height=resized_height) tile_tokens = (base_tokens * 2) * tiles_needed_high_res total_tokens = base_tokens + tile_tokens return total_tokens @@ -338,9 +323,7 @@ class _MessageCountParams: actual_model = _fix_model_name(model) if actual_model == "gpt-3.5-turbo-0301": - self.tokens_per_message = ( - 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n - ) + self.tokens_per_message = 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n self.tokens_per_name = -1 # if there's a name, the role is omitted elif actual_model in litellm.open_ai_chat_completion_models: self.tokens_per_message = 3 @@ -349,9 +332,7 @@ class _MessageCountParams: self.tokens_per_message = 3 self.tokens_per_name = 1 else: - print_verbose( - f"Warning: unknown model {model}. Using default token params." - ) + print_verbose(f"Warning: unknown model {model}. Using default token params.") self.tokens_per_message = 3 self.tokens_per_name = 1 self.count_function = _get_count_function(model, custom_tokenizer) @@ -396,9 +377,7 @@ def token_counter( if litellm.disable_token_counter is True: return 0 - verbose_logger.debug( - f"messages in token_counter: {messages}, text in token_counter: {text}" - ) + verbose_logger.debug(f"messages in token_counter: {messages}, text in token_counter: {text}") if text is not None and messages is not None: raise ValueError("text and messages cannot both be set") if use_default_image_token_count is None: @@ -415,20 +394,12 @@ def token_counter( num_tokens = count_function(text_to_count) elif messages is not None: - new_messages = cast( - List[AllMessageValues], convert_list_message_to_dict(messages) - ) + new_messages = cast(List[AllMessageValues], convert_list_message_to_dict(messages)) params = _MessageCountParams(model, custom_tokenizer) - num_tokens = _count_messages( - params, new_messages, use_default_image_token_count, default_token_count - ) + num_tokens = _count_messages(params, new_messages, use_default_image_token_count, default_token_count) if count_response_tokens is False: - includes_system_message = any( - [message.get("role", None) == "system" for message in new_messages] - ) - num_tokens += _count_extra( - params.count_function, tools, tool_choice, includes_system_message - ) + includes_system_message = any([message.get("role", None) == "system" for message in new_messages]) + num_tokens += _count_extra(params.count_function, tools, tool_choice, includes_system_message) else: raise ValueError("Either text or messages must be provided") @@ -463,18 +434,12 @@ def _count_messages( if isinstance(value, List): for tool_call in value: if "function" in tool_call: - function_arguments = tool_call["function"].get( - "arguments", [] - ) + function_arguments = tool_call["function"].get("arguments", []) num_tokens += params.count_function(str(function_arguments)) else: - raise ValueError( - f"Unsupported tool call {tool_call} must contain a function key" - ) + raise ValueError(f"Unsupported tool call {tool_call} must contain a function key") else: - raise ValueError( - f"Unsupported type {type(value)} for key tool_calls in message {message}" - ) + raise ValueError(f"Unsupported type {type(value)} for key tool_calls in message {message}") elif isinstance(value, str): num_tokens += params.count_function(value) if key == "name": @@ -605,9 +570,7 @@ def _count_image_tokens( if isinstance(image_url, dict): detail = image_url.get("detail", "auto") if detail not in ["low", "high", "auto"]: - raise ValueError( - f"Invalid detail value: {detail}. Expected 'low', 'high', or 'auto'." - ) + raise ValueError(f"Invalid detail value: {detail}. Expected 'low', 'high', or 'auto'.") url = image_url.get("url") if not url: raise ValueError("Missing required key 'url' in image_url dict.") @@ -625,10 +588,7 @@ def _count_image_tokens( use_default_image_token_count=use_default_image_token_count, ) else: - raise ValueError( - f"Invalid image_url type: {type(image_url).__name__}. " - "Expected str or dict with 'url' field." - ) + raise ValueError(f"Invalid image_url type: {type(image_url).__name__}. Expected str or dict with 'url' field.") def _validate_anthropic_content(content: Mapping[str, Any]) -> type: @@ -650,13 +610,9 @@ def _validate_anthropic_content(content: Mapping[str, Any]) -> type: if expected_cls is None: raise ValueError(f"Unknown Anthropic content type: '{content_type}'") - missing = [ - k for k in getattr(expected_cls, "__required_keys__", set()) if k not in content - ] + missing = [k for k in getattr(expected_cls, "__required_keys__", set()) if k not in content] if missing: - raise ValueError( - f"Missing required fields in {content_type} block: {', '.join(missing)}" - ) + raise ValueError(f"Missing required fields in {content_type} block: {', '.join(missing)}") return expected_cls @@ -728,9 +684,7 @@ def _count_content_list( num_tokens += count_function(str(c.get("text", ""))) elif c["type"] == "image_url": image_url = c.get("image_url") - num_tokens += _count_image_tokens( - image_url, use_default_image_token_count - ) + num_tokens += _count_image_tokens(image_url, use_default_image_token_count) elif c["type"] in ("tool_use", "tool_result"): num_tokens += _count_anthropic_content( c, @@ -756,11 +710,7 @@ def _count_content_list( if tool_name: num_tokens += count_function(tool_name) else: - content_type = ( - c.get("type", type(c).__name__) - if isinstance(c, dict) - else type(c).__name__ - ) + content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ raise ValueError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking, tool_reference)." @@ -770,8 +720,7 @@ def _count_content_list( if default_token_count is not None: return default_token_count raise ValueError( - f"Error getting number of tokens from content list: {e}, " - f"default_token_count={default_token_count}" + f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}" ) diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 38a78ee058f..1cbb1ce973f 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -169,9 +169,7 @@ def is_url_destination_allowed_by_host(url: str, allowed_hosts: List[str]) -> bo return False normalized_host = _normalize_host(parsed.hostname) - configured_entries = ( - [allowed_hosts] if isinstance(allowed_hosts, str) else allowed_hosts - ) + configured_entries = [allowed_hosts] if isinstance(allowed_hosts, str) else allowed_hosts for entry in configured_entries or []: if not isinstance(entry, str): continue @@ -272,9 +270,7 @@ def validate_url(url: str) -> Tuple[str, str]: # Resolve hostname and validate ALL addresses try: - addrinfo = socket.getaddrinfo( - hostname, effective_port, proto=socket.IPPROTO_TCP - ) + addrinfo = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP) except socket.gaierror as e: raise SSRFError(f"DNS resolution failed for '{hostname}': {e}") @@ -311,9 +307,7 @@ def validate_url(url: str) -> Tuple[str, str]: else: new_netloc = ip_host - rewritten = urlunparse( - (parsed.scheme, new_netloc, parsed.path, parsed.params, parsed.query, "") - ) + rewritten = urlunparse((parsed.scheme, new_netloc, parsed.path, parsed.params, parsed.query, "")) return rewritten, host_header diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index fead8a79bc3..6aec359b7b7 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -41,9 +41,7 @@ def get_cost_for_web_search_request( if "claude" in model_key.lower(): from .anthropic.cost_calculation import get_cost_for_anthropic_web_search - verbose_logger.debug( - "vertex_ai/claude model detected — routing web search cost to Anthropic calculator" - ) + verbose_logger.debug("vertex_ai/claude model detected — routing web search cost to Anthropic calculator") return get_cost_for_anthropic_web_search(model_info=model_info, usage=usage) from .vertex_ai.gemini.cost_calculator import ( @@ -63,9 +61,7 @@ def get_cost_for_web_search_request( return None -def discover_guardrail_translation_mappings() -> Dict[ - CallTypes, Type["BaseTranslation"] -]: +def discover_guardrail_translation_mappings() -> Dict[CallTypes, Type["BaseTranslation"]]: """ Discover guardrail translation mappings by scanning the llms directory structure. @@ -91,19 +87,14 @@ def discover_guardrail_translation_mappings() -> Dict[ dirs[:] = [d for d in dirs if not d.startswith("__") and d != "base_llm"] # Check if this is a guardrail_translation directory with __init__.py - if ( - os.path.basename(root) == "guardrail_translation" - and "__init__.py" in files - ): + if os.path.basename(root) == "guardrail_translation" and "__init__.py" in files: # Build the module path relative to litellm rel_path = os.path.relpath(root, os.path.dirname(llms_dir)) module_path = "litellm." + rel_path.replace(os.sep, ".") try: # Import the module - verbose_logger.debug( - f"Discovering guardrail translations in: {module_path}" - ) + verbose_logger.debug(f"Discovering guardrail translations in: {module_path}") module = importlib.import_module(module_path) @@ -134,9 +125,7 @@ def discover_guardrail_translation_mappings() -> Dict[ list(mcp_guardrail_translation_mappings.keys()), ) except ImportError: - verbose_logger.debug( - "MCP guardrail translation mappings not available; skipping" - ) + verbose_logger.debug("MCP guardrail translation mappings not available; skipping") verbose_logger.debug( f"Discovered {len(discovered_mappings)} guardrail translation mappings: {list(discovered_mappings.keys())}" @@ -149,17 +138,13 @@ def discover_guardrail_translation_mappings() -> Dict[ # Cache the discovered mappings -endpoint_guardrail_translation_mappings: Optional[ - Dict[CallTypes, Type["BaseTranslation"]] -] = None +endpoint_guardrail_translation_mappings: Optional[Dict[CallTypes, Type["BaseTranslation"]]] = None def load_guardrail_translation_mappings(): global endpoint_guardrail_translation_mappings if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = ( - discover_guardrail_translation_mappings() - ) + endpoint_guardrail_translation_mappings = discover_guardrail_translation_mappings() return endpoint_guardrail_translation_mappings @@ -180,9 +165,7 @@ def get_guardrail_translation_mapping(call_type: CallTypes) -> Type["BaseTransla # Lazy load the mappings on first access if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = ( - discover_guardrail_translation_mappings() - ) + endpoint_guardrail_translation_mappings = discover_guardrail_translation_mappings() # Get the translation handler class for the call type if call_type not in endpoint_guardrail_translation_mappings: diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 3d6037b1f8f..740b0fff50c 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -139,9 +139,7 @@ class A2AGuardrailHandler(BaseTranslation): response_dict = response is_pydantic = False else: - verbose_proxy_logger.warning( - "A2A: Unknown response type %s, skipping guardrail", type(response) - ) + verbose_proxy_logger.warning("A2A: Unknown response type %s, skipping guardrail", type(response)) return response result = response_dict.get("result", {}) @@ -177,9 +175,7 @@ class A2AGuardrailHandler(BaseTranslation): # Add user API key metadata with prefixed keys if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata @@ -238,9 +234,7 @@ class A2AGuardrailHandler(BaseTranslation): if not valid_parsed: return responses_so_far - combined_text, chunk_indices_with_text = self._collect_text_from_parsed_chunks( - valid_parsed - ) + combined_text, chunk_indices_with_text = self._collect_text_from_parsed_chunks(valid_parsed) if not combined_text: return responses_so_far @@ -251,9 +245,7 @@ class A2AGuardrailHandler(BaseTranslation): request_data["responses_so_far"] = responses_so_far if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata @@ -270,9 +262,7 @@ class A2AGuardrailHandler(BaseTranslation): guardrailed_text = guardrailed_texts[0] # Find first chunk (by original index) that has text; put full guardrailed text there and clear rest - first_chunk_with_text: Optional[int] = ( - chunk_indices_with_text[0] if chunk_indices_with_text else None - ) + first_chunk_with_text: Optional[int] = chunk_indices_with_text[0] if chunk_indices_with_text else None for orig_i, obj in valid_parsed: result = obj.get("result", {}) @@ -399,11 +389,7 @@ class A2AGuardrailHandler(BaseTranslation): status = result.get("status", {}) if isinstance(status, dict): status_message = status.get("message") - if ( - status_message - and isinstance(status_message, dict) - and "parts" in status_message - ): + if status_message and isinstance(status_message, dict) and "parts" in status_message: self._extract_texts_from_parts( parts=status_message["parts"], path=("status", "message", "parts"), diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 29167d89ae7..a7302ac2f0b 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -31,9 +31,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): ) self.model = model - def chunk_parser( - self, chunk: dict - ) -> Union[GenericStreamingChunk, ModelResponseStream]: + def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]: """ Parse A2A streaming chunk to OpenAI format. diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index b9c9f944b3e..c9623a817bf 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -55,9 +55,7 @@ class A2AConfig(BaseConfig): agent_name = model.split("/", 1)[1] if "/" in model else None # Only lookup if agent name exists and some config is missing - if not agent_name or ( - api_base is not None and api_key is not None and headers is not None - ): + if not agent_name or (api_base is not None and api_key is not None and headers is not None): return api_base, api_key, headers # Try registry lookup (only available in proxy context) @@ -84,10 +82,7 @@ class A2AConfig(BaseConfig): # Merge other litellm_params (timeout, max_retries, etc.) for key, value in agent.litellm_params.items(): - if ( - key not in ["api_key", "api_base", "headers", "model"] - and key not in optional_params - ): + if key not in ["api_key", "api_base", "headers", "model"] and key not in optional_params: optional_params[key] = value except ImportError: pass # Registry not available (not running in proxy context) diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 15ea9f01abd..4fc0ff2623e 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -61,9 +61,7 @@ def convert_messages_to_prompt(messages: List[AllMessageValues]) -> str: return "\n".join(conversation_parts) -def extract_text_from_a2a_message( - message: Dict[str, Any], depth: int = 0, max_depth: int = 10 -) -> str: +def extract_text_from_a2a_message(message: Dict[str, Any], depth: int = 0, max_depth: int = 10) -> str: """ Extract text content from A2A message parts. @@ -93,9 +91,7 @@ def extract_text_from_a2a_message( return " ".join(text_parts) -def extract_text_from_a2a_response( - response_dict: Dict[str, Any], max_depth: int = 10 -) -> str: +def extract_text_from_a2a_response(response_dict: Dict[str, Any], max_depth: int = 10) -> str: """ Extract text content from A2A response result. @@ -136,16 +132,12 @@ def extract_text_from_a2a_response( if isinstance(status, dict): status_message = status.get("message") if status_message: - return extract_text_from_a2a_message( - status_message, depth=0, max_depth=max_depth - ) + return extract_text_from_a2a_message(status_message, depth=0, max_depth=max_depth) # Handle task result with artifacts (plural, array) artifacts = result.get("artifacts", []) if artifacts and len(artifacts) > 0: first_artifact = artifacts[0] - return extract_text_from_a2a_message( - first_artifact, depth=0, max_depth=max_depth - ) + return extract_text_from_a2a_message(first_artifact, depth=0, max_depth=max_depth) return "" diff --git a/litellm/llms/aiml/chat/transformation.py b/litellm/llms/aiml/chat/transformation.py index 72e30a08173..e62aa6238d7 100644 --- a/litellm/llms/aiml/chat/transformation.py +++ b/litellm/llms/aiml/chat/transformation.py @@ -14,9 +14,7 @@ class AIMLChatConfig(OpenAIGPTConfig): ) -> Tuple[Optional[str], Optional[str]]: # AIML is openai compatible, we just need to set the api_base api_base = ( - api_base - or get_secret_str("AIML_API_BASE") - or "https://api.aimlapi.com/v1" # Default AIML API base URL + api_base or get_secret_str("AIML_API_BASE") or "https://api.aimlapi.com/v1" # Default AIML API base URL ) # type: ignore dynamic_api_key = api_key or get_secret_str("AIML_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/aiml/image_generation/cost_calculator.py b/litellm/llms/aiml/image_generation/cost_calculator.py index 4442f57c555..1fecfb6a9a5 100644 --- a/litellm/llms/aiml/image_generation/cost_calculator.py +++ b/litellm/llms/aiml/image_generation/cost_calculator.py @@ -22,6 +22,4 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index 92a1510f3a1..b1ab443eb84 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -37,9 +37,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): """ return model.startswith(OPENAI_STYLE_IMAGE_MODEL_PREFIXES) - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ https://api.aimlapi.com/v1/images/generations """ @@ -111,9 +109,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete url for the request """ - complete_url: str = ( - api_base or get_secret_str("AIML_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url: str = api_base or get_secret_str("AIML_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") # Strip /v1 suffix if present since IMAGE_GENERATION_ENDPOINT already includes v1 @@ -133,9 +129,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): api_base: Optional[str] = None, ) -> dict: final_api_key: Optional[str] = ( - api_key - or get_secret_str("AIML_API_KEY") - or get_secret_str("AIMLAPI_KEY") # Alternative name + api_key or get_secret_str("AIML_API_KEY") or get_secret_str("AIMLAPI_KEY") # Alternative name ) if not final_api_key: raise ValueError("AIML_API_KEY or AIMLAPI_KEY is not set") @@ -160,12 +154,10 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): if self._is_openai_style_model(model): return {"model": model, "prompt": prompt, **optional_params} - aiml_image_generation_request_body: AimlImageGenerationRequestParams = ( - AimlImageGenerationRequestParams( - prompt=prompt, - model=model, - **optional_params, - ) + aiml_image_generation_request_body: AimlImageGenerationRequestParams = AimlImageGenerationRequestParams( + prompt=prompt, + model=model, + **optional_params, ) return dict(aiml_image_generation_request_body) diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index c2d4e5adcd7..346b565b6f5 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -72,9 +72,7 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): ) -> ModelResponse: _json_response = await raw_response.json() model_response.id = _json_response.get("id") - model_response.choices = [ - Choices(**choice) for choice in _json_response.get("choices") - ] + model_response.choices = [Choices(**choice) for choice in _json_response.get("choices")] model_response.created = _json_response.get("created") model_response.model = _json_response.get("model") model_response.object = _json_response.get("object") diff --git a/litellm/llms/amazon_nova/chat/transformation.py b/litellm/llms/amazon_nova/chat/transformation.py index 74c7fd234fe..8afcbd40ffc 100644 --- a/litellm/llms/amazon_nova/chat/transformation.py +++ b/litellm/llms/amazon_nova/chat/transformation.py @@ -52,19 +52,10 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: # Amazon Nova is openai compatible, we just need to set this to custom_openai and have the api_base be Nova's endpoint - api_base = ( - api_base - or get_secret_str("AMAZON_NOVA_API_BASE") - or "https://api.nova.amazon.com/v1" - ) # type: ignore + api_base = api_base or get_secret_str("AMAZON_NOVA_API_BASE") or "https://api.nova.amazon.com/v1" # type: ignore # Get API key from multiple sources - key = ( - api_key - or litellm.amazon_nova_api_key - or get_secret_str("AMAZON_NOVA_API_KEY") - or litellm.api_key - ) + key = api_key or litellm.amazon_nova_api_key or get_secret_str("AMAZON_NOVA_API_KEY") or litellm.api_key return api_base, key def get_supported_openai_params(self, model: str) -> List: diff --git a/litellm/llms/amazon_nova/cost_calculation.py b/litellm/llms/amazon_nova/cost_calculation.py index 857369b76ed..3b1121f1f8c 100644 --- a/litellm/llms/amazon_nova/cost_calculation.py +++ b/litellm/llms/amazon_nova/cost_calculation.py @@ -16,6 +16,4 @@ def cost_per_token(model: str, usage: "Usage") -> Tuple[float, float]: Calculates the cost per token for a given model, prompt tokens, and completion tokens. Follows the same logic as Anthropic's cost per token calculation. """ - return generic_cost_per_token( - model=model, usage=usage, custom_llm_provider="amazon_nova" - ) + return generic_cost_per_token(model=model, usage=usage, custom_llm_provider="amazon_nova") diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 7c4e9386d5f..bfae42f96cf 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -233,12 +233,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig): completed_at=ended_at if processing_status == "ended" else None, failed_at=None, expired_at=archived_at if archived_at else None, - cancelling_at=( - cancel_initiated_at if processing_status == "canceling" else None - ), - cancelled_at=( - ended_at if processing_status == "canceling" and ended_at else None - ), + cancelling_at=(cancel_initiated_at if processing_status == "canceling" else None), + cancelled_at=(ended_at if processing_status == "canceling" and ended_at else None), request_counts=request_counts, metadata={}, ) @@ -255,9 +251,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): else: headers_obj = headers if isinstance(headers, Headers) else None - return AnthropicError( - status_code=status_code, message=error_message, headers=headers_obj - ) + return AnthropicError(status_code=status_code, message=error_message, headers=headers_obj) def transform_response( self, @@ -290,17 +284,13 @@ class AnthropicBatchesConfig(BaseBatchesConfig): response_json = json.loads(line) # Update model_response with the parsed JSON completion_response = response_json["result"]["message"] - transformed_response = ( - self.anthropic_chat_config.transform_parsed_response( - completion_response=completion_response, - raw_response=raw_response, - model_response=model_response, - ) + transformed_response = self.anthropic_chat_config.transform_parsed_response( + completion_response=completion_response, + raw_response=raw_response, + model_response=model_response, ) - transformed_response_usage = getattr( - transformed_response, "usage", None - ) + transformed_response_usage = getattr(transformed_response, "usage", None) if transformed_response_usage: all_usage.append(cast(Usage, transformed_response_usage)) except json.JSONDecodeError: diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c65c0e02cfc..4506c114208 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -125,9 +125,7 @@ class AnthropicMessagesHandler(BaseTranslation): texts_to_check: List[str] = [] images_to_check: List[str] = [] - tools_to_check: List[ChatCompletionToolParam] = ( - chat_completion_compatible_request.get("tools", []) - ) + tools_to_check: List[ChatCompletionToolParam] = chat_completion_compatible_request.get("tools", []) task_mappings: List[Tuple[int, Optional[int]]] = [] # Step 1: Extract all text content and images @@ -176,16 +174,12 @@ class AnthropicMessagesHandler(BaseTranslation): # Note: MCP servers are handled separately in the main transformation data["tools"] = anthropic_tools - guardrailed_structured_messages = guardrailed_inputs.get( - "structured_messages" - ) + guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") if ( guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - self._write_back_structured_messages( - data, guardrailed_structured_messages - ) + self._write_back_structured_messages(data, guardrailed_structured_messages) else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( @@ -194,9 +188,7 @@ class AnthropicMessagesHandler(BaseTranslation): task_mappings=task_mappings, ) - verbose_proxy_logger.debug( - "Anthropic Messages: Processed input messages: %s", messages - ) + verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) return data @@ -209,9 +201,7 @@ class AnthropicMessagesHandler(BaseTranslation): model = str(data.get("model") or "") non_system = [m for m in structured_messages if m.get("role") != "system"] - converted = anthropic_messages_pt( - messages=non_system, model=model, llm_provider="anthropic" - ) + converted = anthropic_messages_pt(messages=non_system, model=model, llm_provider="anthropic") for msg in converted: content = msg.get("content") if isinstance(content, list): @@ -320,9 +310,7 @@ class AnthropicMessagesHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content - messages[msg_idx]["content"][content_idx_optional]["text"] = ( - guardrail_response - ) + messages[msg_idx]["content"][content_idx_optional]["text"] = guardrail_response async def process_output_response( self, @@ -401,9 +389,7 @@ class AnthropicMessagesHandler(BaseTranslation): task_mappings=task_mappings, ) - verbose_proxy_logger.debug( - "Anthropic Messages: Processed output response: %s", response - ) + verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response) return response @@ -423,20 +409,14 @@ class AnthropicMessagesHandler(BaseTranslation): has_ended = self._check_streaming_has_ended(responses_so_far) if has_ended: # build the model response from the responses_so_far - built_response = ( - AnthropicPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=responses_so_far, - litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), - model="", - ) + built_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=responses_so_far, + litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), + model="", ) # Check if model_response is valid and has choices before accessing - if ( - built_response is not None - and hasattr(built_response, "choices") - and built_response.choices - ): + if built_response is not None and hasattr(built_response, "choices") and built_response.choices: model_response = cast(ModelResponse, built_response) first_choice = cast(Choices, model_response.choices[0]) tool_calls_list = cast( @@ -450,16 +430,16 @@ class AnthropicMessagesHandler(BaseTranslation): if tool_calls_list: guardrail_inputs["tool_calls"] = tool_calls_list - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid - inputs=guardrail_inputs, - request_data=request_data if request_data is not None else {}, - input_type="response", - logging_obj=litellm_logging_obj, + _guardrailed_inputs = ( + await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid + inputs=guardrail_inputs, + request_data=request_data if request_data is not None else {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) ) else: - verbose_proxy_logger.debug( - "Skipping output guardrail - model response has no choices" - ) + verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far string_so_far = self.get_streaming_string_so_far(responses_so_far) @@ -486,9 +466,7 @@ class AnthropicMessagesHandler(BaseTranslation): request_data[key] = response if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata return request_data @@ -636,9 +614,7 @@ class AnthropicMessagesHandler(BaseTranslation): if delta.get("type") == "text_delta": text += delta.get("text", "") except json.JSONDecodeError: - verbose_proxy_logger.warning( - f"Failed to parse JSON from SSE data: {data_line}" - ) + verbose_proxy_logger.warning(f"Failed to parse JSON from SSE data: {data_line}") except Exception as e: verbose_proxy_logger.error(f"Error extracting text from SSE: {e}") @@ -702,14 +678,10 @@ class AnthropicMessagesHandler(BaseTranslation): if stop_reason is not None: return True except json.JSONDecodeError: - verbose_proxy_logger.warning( - f"Failed to parse JSON from SSE data: {data_line}" - ) + verbose_proxy_logger.warning(f"Failed to parse JSON from SSE data: {data_line}") except Exception as e: - verbose_proxy_logger.error( - f"Error checking streaming end in SSE: {e}" - ) + verbose_proxy_logger.error(f"Error checking streaming end in SSE: {e}") # Handle already-parsed dict format elif isinstance(response, dict): @@ -815,10 +787,7 @@ class AnthropicMessagesHandler(BaseTranslation): if isinstance(content_block, dict): if content_block.get("type") == "text": cast(Dict[str, Any], content_block)["text"] = guardrail_response - elif ( - hasattr(content_block, "type") - and getattr(content_block, "type", None) == "text" - ): + elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text": # Update Pydantic object's text attribute if hasattr(content_block, "text"): content_block.text = guardrail_response diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 7154a6f3595..c8872306e82 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -241,9 +241,7 @@ class AnthropicChatCompletion(BaseLLM): json_mode=json_mode, speed=optional_params.get("speed") if optional_params else None, tool_name_reverse_map=( - litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY) - if isinstance(litellm_params, dict) - else None + litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY) if isinstance(litellm_params, dict) else None ), ) streamwrapper = CustomStreamWrapper( @@ -278,9 +276,7 @@ class AnthropicChatCompletion(BaseLLM): headers={}, client: AsyncHTTPHandler | None = None, ) -> Union[ModelResponse, "CustomStreamWrapper"]: - async_handler = client or get_async_httpx_client( - llm_provider=litellm.LlmProviders.ANTHROPIC - ) + async_handler = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) try: response = await async_handler.post( @@ -371,9 +367,7 @@ class AnthropicChatCompletion(BaseLLM): provider=LlmProviders(custom_llm_provider), ) if config is None: - raise ValueError( - f"Provider config not found for model: {model} and provider: {custom_llm_provider}" - ) + raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}") data = config.transform_request( model=model, @@ -425,11 +419,7 @@ class AnthropicChatCompletion(BaseLLM): logger_fn=logger_fn, headers=headers, timeout=timeout, - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), ) else: return self.acompletion_function( @@ -613,11 +603,7 @@ class ModelResponseIterator: return False def _handle_usage(self, anthropic_usage_chunk: Union[dict, UsageDelta]) -> Usage: - reasoning_content = ( - "".join(self.reasoning_content_chunks) - if self.reasoning_content_chunks - else None - ) + reasoning_content = "".join(self.reasoning_content_chunks) if self.reasoning_content_chunks else None return AnthropicConfig().calculate_usage( usage_object=cast(dict, anthropic_usage_chunk), reasoning_content=reasoning_content, @@ -641,9 +627,7 @@ class ModelResponseIterator: provider_specific_fields = {} reasoning_content: str | None = None content_block = ContentBlockDelta(**chunk) # type: ignore - thinking_blocks: List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] = [] + thinking_blocks: List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] = [] self.content_blocks.append(content_block) if "text" in content_block["delta"]: @@ -667,10 +651,7 @@ class ModelResponseIterator: ) elif "citation" in content_block["delta"]: provider_specific_fields["citation"] = content_block["delta"]["citation"] - elif ( - "thinking" in content_block["delta"] - or "signature" in content_block["delta"] - ): + elif "thinking" in content_block["delta"] or "signature" in content_block["delta"]: thinking_content = content_block["delta"].get("thinking") if isinstance(thinking_content, str) and thinking_content: self.reasoning_content_chunks.append(thinking_content) @@ -699,10 +680,7 @@ class ModelResponseIterator: provider_specific_fields["thinking_blocks"] = thinking_blocks if reasoning_content is None: reasoning_content = "" - elif ( - "content" in content_block["delta"] - and content_block["delta"].get("type") == "compaction_delta" - ): + elif "content" in content_block["delta"] and content_block["delta"].get("type") == "compaction_delta": # Handle compaction delta provider_specific_fields["compaction_delta"] = { "type": "compaction_delta", @@ -791,14 +769,7 @@ class ModelResponseIterator: usage: Usage | None = None provider_specific_fields: Dict[str, Any] = {} reasoning_content: str | None = None - thinking_blocks: ( - List[ - Union[ - ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock - ] - ] - | None - ) = None + thinking_blocks: List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] | None = None # Always use index=0 for OpenAI choice format (fixes multi-choice errors) index = 0 @@ -823,9 +794,7 @@ class ModelResponseIterator: content_block_start = self.get_content_block_start(chunk=chunk) self.content_blocks = [] # reset content blocks when new block starts # Track current content block type for filtering deltas - self.current_content_block_type = content_block_start["content_block"][ - "type" - ] + self.current_content_block_type = content_block_start["content_block"]["type"] if content_block_start["content_block"]["type"] == "text": text = content_block_start["content_block"]["text"] elif ( @@ -836,13 +805,8 @@ class ModelResponseIterator: # Reverse-map the (sanitized) tool name back to the # caller's original. No-op when the map is empty. _stream_tool_name = content_block_start["content_block"]["name"] - if ( - self.tool_name_reverse_map - and _stream_tool_name in self.tool_name_reverse_map - ): - _stream_tool_name = self.tool_name_reverse_map[ - _stream_tool_name - ] + if self.tool_name_reverse_map and _stream_tool_name in self.tool_name_reverse_map: + _stream_tool_name = self.tool_name_reverse_map[_stream_tool_name] # Use empty string for arguments in content_block_start - actual arguments # come in subsequent content_block_delta chunks and get accumulated. # Using str(input) here would prepend '{}' causing invalid JSON accumulation. @@ -859,27 +823,16 @@ class ModelResponseIterator: # The initial input in content_block_start is typically {} # for streaming; the full input arrives via input_json_delta # and is assembled at content_block_stop. - if ( - content_block_start["content_block"]["type"] - == "server_tool_use" - ): - self._current_server_tool_id = content_block_start[ - "content_block" - ]["id"] - tool_input = content_block_start["content_block"].get( - "input", {} - ) - self._server_tool_inputs[self._current_server_tool_id] = ( - tool_input - ) + if content_block_start["content_block"]["type"] == "server_tool_use": + self._current_server_tool_id = content_block_start["content_block"]["id"] + tool_input = content_block_start["content_block"].get("input", {}) + self._server_tool_inputs[self._current_server_tool_id] = tool_input # Include caller information if present (for programmatic tool calling) if "caller" in content_block_start["content_block"]: caller_data = content_block_start["content_block"]["caller"] if caller_data: tool_use["caller"] = cast(Dict[str, Any], caller_data) # type: ignore[typeddict-item] - elif ( - content_block_start["content_block"]["type"] == "redacted_thinking" - ): + elif content_block_start["content_block"]["type"] == "redacted_thinking": ( thinking_blocks, provider_specific_fields, @@ -892,19 +845,13 @@ class ModelResponseIterator: # Handle compaction blocks # The full content comes in content_block_start self.compaction_blocks.append(content_block_start["content_block"]) - provider_specific_fields["compaction_blocks"] = ( - self.compaction_blocks - ) + provider_specific_fields["compaction_blocks"] = self.compaction_blocks provider_specific_fields["compaction_start"] = { "type": "compaction", - "content": content_block_start["content_block"].get( - "content", "" - ), + "content": content_block_start["content_block"].get("content", ""), } - elif content_block_start["content_block"]["type"].endswith( - "_tool_result" - ): + elif content_block_start["content_block"]["type"].endswith("_tool_result"): # Handle all tool result types (web_search, bash_code_execution, text_editor, etc.) content_type = content_block_start["content_block"]["type"] @@ -913,31 +860,21 @@ class ModelResponseIterator: # Capture web_search_tool_result for multi-turn reconstruction # The full content comes in content_block_start, not in deltas # See: https://github.com/BerriAI/litellm/issues/17737 - self.web_search_results.append( - content_block_start["content_block"] - ) - provider_specific_fields["web_search_results"] = ( - self.web_search_results - ) + self.web_search_results.append(content_block_start["content_block"]) + provider_specific_fields["web_search_results"] = self.web_search_results elif content_type == "web_fetch_tool_result": # Capture web_fetch_tool_result for multi-turn reconstruction # The full content comes in content_block_start, not in deltas # Fixes: https://github.com/BerriAI/litellm/issues/18137 - self.web_search_results.append( - content_block_start["content_block"] - ) - provider_specific_fields["web_search_results"] = ( - self.web_search_results - ) + self.web_search_results.append(content_block_start["content_block"]) + provider_specific_fields["web_search_results"] = self.web_search_results elif content_type != "tool_search_tool_result": # Handle other tool results (code execution, etc.) # Skip tool_search_tool_result as it's internal metadata self.tool_results.append(content_block_start["content_block"]) provider_specific_fields["tool_results"] = self.tool_results # Convert to provider-neutral code_interpreter_results - provider_specific_fields["code_interpreter_results"] = ( - self._build_code_interpreter_results() - ) + provider_specific_fields["code_interpreter_results"] = self._build_code_interpreter_results() elif type_chunk == "content_block_stop": ContentBlockStop(**chunk) # type: ignore @@ -956,10 +893,7 @@ class ModelResponseIterator: ) # Update server_tool_inputs with fully assembled input # from input_json_delta chunks (content_block_start has {}) - if ( - self.current_content_block_type == "server_tool_use" - and self._current_server_tool_id - ): + if self.current_content_block_type == "server_tool_use" and self._current_server_tool_id: args = "" for block in self.content_blocks: if block["delta"]["type"] == "input_json_delta": @@ -968,9 +902,7 @@ class ModelResponseIterator: args += partial_json if args: try: - self._server_tool_inputs[ - self._current_server_tool_id - ] = json.loads(args) + self._server_tool_inputs[self._current_server_tool_id] = json.loads(args) except (json.JSONDecodeError, TypeError): pass self._current_server_tool_id = None @@ -989,14 +921,10 @@ class ModelResponseIterator: # Store container_id and re-emit code_interpreter_results # so stream_chunk_builder's last-value-wins picks up the # version with container_id populated. - container_id = ( - container.get("id") if isinstance(container, dict) else None - ) + container_id = container.get("id") if isinstance(container, dict) else None if container_id and self.tool_results: self._container_id = container_id - provider_specific_fields["code_interpreter_results"] = ( - self._build_code_interpreter_results() - ) + provider_specific_fields["code_interpreter_results"] = self._build_code_interpreter_results() elif type_chunk == "message_start": """ Anthropic @@ -1019,9 +947,7 @@ class ModelResponseIterator: """ message_start_block = MessageStartBlock(**chunk) # type: ignore if "usage" in message_start_block["message"]: - usage = self._handle_usage( - anthropic_usage_chunk=message_start_block["message"]["usage"] - ) + usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"]) elif type_chunk == "error": """ {"type":"error","error":{"details":null,"type":"api_error","message":"Internal server error"} } @@ -1042,14 +968,8 @@ class ModelResponseIterator: delta=Delta( content=text, tool_calls=[tool_use] if tool_use is not None else None, - provider_specific_fields=( - provider_specific_fields - if provider_specific_fields - else None - ), - thinking_blocks=( - thinking_blocks if thinking_blocks else None - ), + provider_specific_fields=(provider_specific_fields if provider_specific_fields else None), + thinking_blocks=(thinking_blocks if thinking_blocks else None), reasoning_content=reasoning_content, ), finish_reason=finish_reason, @@ -1101,9 +1021,7 @@ class ModelResponseIterator: # Convert tool to content if we're tracking a response_format tool if self.is_response_format_tool: - message = AnthropicConfig._convert_tool_response_to_message( - tool_calls=[tool_use] - ) + message = AnthropicConfig._convert_tool_response_to_message(tool_calls=[tool_use]) if message is not None: text = message.content or "" tool_use = None @@ -1112,9 +1030,7 @@ class ModelResponseIterator: return text, tool_use - def _handle_message_delta( - self, chunk: dict - ) -> Tuple[str, Usage | None, Dict[str, Any] | None]: + def _handle_message_delta(self, chunk: dict) -> Tuple[str, Usage | None, Dict[str, Any] | None]: """ Handle message_delta event for finish_reason, usage, and container. @@ -1125,9 +1041,7 @@ class ModelResponseIterator: Tuple of (finish_reason, usage, container) """ message_delta = MessageBlockDelta(**chunk) # type: ignore - finish_reason = map_finish_reason( - finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop" - ) + finish_reason = map_finish_reason(finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop") # Override finish_reason to "stop" if we converted response_format tools # (matches OpenAI behavior and non-streaming Anthropic implementation) if self.converted_response_format_tool: @@ -1136,9 +1050,7 @@ class ModelResponseIterator: container = message_delta["delta"].get("container") return finish_reason, usage, container - def _handle_accumulated_json_chunk( - self, data_str: str - ) -> ModelResponseStream | None: + def _handle_accumulated_json_chunk(self, data_str: str) -> ModelResponseStream | None: """ Handle partial JSON chunks by accumulating them until valid JSON is received. @@ -1234,9 +1146,7 @@ class ModelResponseIterator: except StopIteration: raise StopIteration except ValueError as e: - raise RuntimeError( - f"Error parsing chunk: {e},\nReceived chunk: {chunk}" - ) + raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") # Async iterator def __aiter__(self): @@ -1285,9 +1195,7 @@ class ModelResponseIterator: except StopAsyncIteration: raise StopAsyncIteration except ValueError as e: - raise RuntimeError( - f"Error parsing chunk: {e},\nReceived chunk: {chunk}" - ) + raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") def convert_str_chunk_to_generic_chunk(self, chunk: str) -> ModelResponseStream: """ diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index b223aef1415..9721b797584 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -146,9 +146,7 @@ def _basic_sanitize_anthropic_tool_name(name: str) -> str: """ if not isinstance(name, str) or not name: return name - return _ANTHROPIC_TOOL_NAME_INVALID_CHARS.sub("_", name)[ - :_ANTHROPIC_TOOL_NAME_MAX_LEN - ] + return _ANTHROPIC_TOOL_NAME_INVALID_CHARS.sub("_", name)[:_ANTHROPIC_TOOL_NAME_MAX_LEN] def _build_anthropic_tool_name_maps( @@ -229,7 +227,9 @@ DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING = ( "Sonnet 4.6+, and Mythos Preview." ) -DROP_UNSUPPORTED_SPEED_WARNING = "Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models." +DROP_UNSUPPORTED_SPEED_WARNING = ( + "Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models." +) class AnthropicConfig(AnthropicModelInfo, BaseConfig): @@ -326,36 +326,27 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _is_opus_4_6_model(model: str) -> bool: """Check if the model is specifically Claude Opus 4.6.""" model_lower = model.lower() - return any( - v in model_lower for v in ("opus-4-6", "opus_4_6", "opus-4.6", "opus_4.6") - ) + return any(v in model_lower for v in ("opus-4-6", "opus_4_6", "opus-4.6", "opus_4.6")) @staticmethod def _is_opus_4_7_model(model: str) -> bool: """Check if the model is specifically Claude Opus 4.7.""" model_lower = model.lower() - return any( - v in model_lower for v in ("opus-4-7", "opus_4_7", "opus-4.7", "opus_4.7") - ) + return any(v in model_lower for v in ("opus-4-7", "opus_4_7", "opus-4.7", "opus_4.7")) @staticmethod def _supports_effort_level(model: str, level: str) -> bool: """Check ``supports_{level}_reasoning_effort`` in the model map.""" - return AnthropicConfig._supports_model_capability( - model, f"supports_{level}_reasoning_effort" - ) + return AnthropicConfig._supports_model_capability(model, f"supports_{level}_reasoning_effort") @staticmethod def _validate_effort_for_model(model: str, effort: Optional[str]) -> Optional[str]: """Return ``None`` if ``effort`` is allowed on ``model``, else an error message.""" if effort == "max" and not ( - AnthropicConfig._is_adaptive_thinking_model(model) - or AnthropicConfig._supports_effort_level(model, "max") + AnthropicConfig._is_adaptive_thinking_model(model) or AnthropicConfig._supports_effort_level(model, "max") ): return f"effort='max' is not supported by this model. Got model: {model}" - if effort == "xhigh" and not AnthropicConfig._supports_effort_level( - model, "xhigh" - ): + if effort == "xhigh" and not AnthropicConfig._supports_effort_level(model, "xhigh"): return f"effort='xhigh' is not supported by this model. Got model: {model}" return None @@ -376,9 +367,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _model_supports_speed_param( - model: str, custom_llm_provider: Optional[str] = None - ) -> bool: + def _model_supports_speed_param(model: str, custom_llm_provider: Optional[str] = None) -> bool: """Whether the model accepts Anthropic's ``speed`` parameter (fast mode). Fast mode is direct Anthropic API-only (not Bedrock, Vertex, or Azure). @@ -388,10 +377,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ if custom_llm_provider is not None and custom_llm_provider != "anthropic": return False - return ( - AnthropicModelInfo._get_exact_model_capability(model, "supports_speed") - is True - ) + return AnthropicModelInfo._get_exact_model_capability(model, "supports_speed") is True @staticmethod def _maybe_drop_speed_param( @@ -421,9 +407,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params.pop("speed", None) @staticmethod - def _raise_invalid_reasoning_effort( - model: str, value: Any, llm_provider: str - ) -> NoReturn: + def _raise_invalid_reasoning_effort(model: str, value: Any, llm_provider: str) -> NoReturn: """Raise a ``BadRequestError`` for an unrecognised ``reasoning_effort``. Args: @@ -533,9 +517,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): } for field in unsupported_fields: if field in schema: - constraint_descriptions.append( - constraint_labels[field].format(schema[field]) - ) + constraint_descriptions.append(constraint_labels[field].format(schema[field])) result: Dict[str, Any] = {} @@ -556,32 +538,17 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): continue if key == "properties" and isinstance(value, dict): - result[key] = { - k: AnthropicConfig.filter_anthropic_output_schema(v) - for k, v in value.items() - } + result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()} elif key == "items" and isinstance(value, dict): result[key] = AnthropicConfig.filter_anthropic_output_schema(value) elif key == "$defs" and isinstance(value, dict): - result[key] = { - k: AnthropicConfig.filter_anthropic_output_schema(v) - for k, v in value.items() - } + result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()} elif key == "anyOf" and isinstance(value, list): - result[key] = [ - AnthropicConfig.filter_anthropic_output_schema(item) - for item in value - ] + result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] elif key == "allOf" and isinstance(value, list): - result[key] = [ - AnthropicConfig.filter_anthropic_output_schema(item) - for item in value - ] + result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] elif key == "oneOf" and isinstance(value, list): - result[key] = [ - AnthropicConfig.filter_anthropic_output_schema(item) - for item in value - ] + result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] else: result[key] = value @@ -592,9 +559,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return result - def get_json_schema_from_pydantic_object( - self, response_format: Union[Any, Dict, None] - ) -> Optional[dict]: + def get_json_schema_from_pydantic_object(self, response_format: Union[Any, Dict, None]) -> Optional[dict]: return type_to_response_format_param( response_format, ref_template="/$defs/{model}" ) # Relevant issue: https://github.com/BerriAI/litellm/issues/7755 @@ -686,12 +651,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _input_schema = unpack_legacy_defs(_input_schema, copy=True) _allowed_properties = set(AnthropicInputSchema.__annotations__.keys()) - input_schema_filtered = { - k: v for k, v in _input_schema.items() if k in _allowed_properties - } - input_anthropic_schema: AnthropicInputSchema = AnthropicInputSchema( - **input_schema_filtered - ) + input_schema_filtered = {k: v for k, v in _input_schema.items() if k in _allowed_properties} + input_anthropic_schema: AnthropicInputSchema = AnthropicInputSchema(**input_schema_filtered) _tool = AnthropicMessagesTool( name=tool["function"]["name"], @@ -710,16 +671,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if "parameters" not in tool["function"]: raise ValueError("Missing required parameter: parameters") - _display_width_px: Optional[int] = tool["function"]["parameters"].get( - "display_width_px" - ) - _display_height_px: Optional[int] = tool["function"]["parameters"].get( - "display_height_px" - ) + _display_width_px: Optional[int] = tool["function"]["parameters"].get("display_width_px") + _display_height_px: Optional[int] = tool["function"]["parameters"].get("display_height_px") if _display_width_px is None or _display_height_px is None: - raise ValueError( - "Missing required parameter: display_width_px or display_height_px" - ) + raise ValueError("Missing required parameter: display_width_px or display_height_px") _computer_tool = AnthropicComputerTool( type=tool["type"], @@ -752,9 +707,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): elif tool["type"] == "url": # mcp server tool mcp_server = AnthropicMcpServerTool(**tool) # type: ignore elif tool["type"] == "mcp": - mcp_server = self._map_openai_mcp_server_tool( - cast(OpenAIMcpServerTool, tool) - ) + mcp_server = self._map_openai_mcp_server_tool(cast(OpenAIMcpServerTool, tool)) elif tool["type"] == "tool_search_tool_regex_20251119": # Tool search tool using regex from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex @@ -811,9 +764,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ): if _cache_control is not None: returned_tool["cache_control"] = _cache_control # type: ignore[typeddict-item] - elif _cache_control_function is not None and isinstance( - _cache_control_function, dict - ): + elif _cache_control_function is not None and isinstance(_cache_control_function, dict): returned_tool["cache_control"] = ChatCompletionCachedContent( # type: ignore[typeddict-item] **_cache_control_function # type: ignore ) @@ -841,9 +792,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## check if allowed_callers is set in the tool _allowed_callers = tool.get("allowed_callers", None) - _allowed_callers_function = tool.get("function", {}).get( - "allowed_callers", None - ) + _allowed_callers_function = tool.get("function", {}).get("allowed_callers", None) if returned_tool is not None: # Only set allowed_callers on tools that support it (not tool search tools or computer tools) tool_type = returned_tool.get("type", "") @@ -875,16 +824,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if tool_type == "custom" or (tool_type == "" and "name" in returned_tool): if _input_examples is not None and isinstance(_input_examples, list): returned_tool["input_examples"] = _input_examples # type: ignore[typeddict-item] - elif _input_examples_function is not None and isinstance( - _input_examples_function, list - ): + elif _input_examples_function is not None and isinstance(_input_examples_function, list): returned_tool["input_examples"] = _input_examples_function # type: ignore[typeddict-item] return returned_tool, mcp_server - def _map_openai_mcp_server_tool( - self, tool: OpenAIMcpServerTool - ) -> AnthropicMcpServerTool: + def _map_openai_mcp_server_tool(self, tool: OpenAIMcpServerTool) -> AnthropicMcpServerTool: from litellm.types.llms.anthropic import AnthropicMcpServerToolConfiguration allowed_tools = tool.get("allowed_tools", None) @@ -934,9 +879,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ChatCompletionToolParam, { "type": nested.get("type", "function"), - "function": { - k: v for k, v in nested.items() if k != "type" - }, + "function": {k: v for k, v in nested.items() if k != "type"}, }, ) nested_tool, nested_mcp = self._map_tool_helper(wrapped) @@ -945,9 +888,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if nested_mcp is not None: mcp_servers.append(nested_mcp) elif "function" in nested: - nested_tool, nested_mcp = self._map_tool_helper( - cast(ChatCompletionToolParam, nested) - ) + nested_tool, nested_mcp = self._map_tool_helper(cast(ChatCompletionToolParam, nested)) if nested_tool is not None: anthropic_tools.append(nested_tool) if nested_mcp is not None: @@ -997,11 +938,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): continue fn = tc.get("function") fn_name = fn.get("name") if isinstance(fn, dict) else None - if ( - isinstance(fn, dict) - and isinstance(fn_name, str) - and fn_name in name_forward_map - ): + if isinstance(fn, dict) and isinstance(fn_name, str) and fn_name in name_forward_map: new_fn = dict(fn) new_fn["name"] = name_forward_map[fn_name] new_tc = dict(tc) @@ -1010,14 +947,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: new_calls.append(tc) new_msg["tool_calls"] = new_calls - fc_name = ( - function_call.get("name") if isinstance(function_call, dict) else None - ) - if ( - isinstance(function_call, dict) - and isinstance(fc_name, str) - and fc_name in name_forward_map - ): + fc_name = function_call.get("name") if isinstance(function_call, dict) else None + if isinstance(function_call, dict) and isinstance(fc_name, str) and fc_name in name_forward_map: new_fc = dict(function_call) new_fc["name"] = name_forward_map[fc_name] new_msg["function_call"] = new_fc @@ -1045,11 +976,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): for tool in tools or []: if not isinstance(tool, dict): continue - original = ( - tool.get("function", {}).get("name") - if isinstance(tool.get("function"), dict) - else None - ) + original = tool.get("function", {}).get("name") if isinstance(tool.get("function"), dict) else None if original is None: original = tool.get("name") if isinstance(original, str) and original: @@ -1208,9 +1135,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return expanded_content - def _map_stop_sequences( - self, stop: Optional[Union[str, List[str]]] - ) -> Optional[List[str]]: + def _map_stop_sequences(self, stop: Optional[Union[str, List[str]]]) -> Optional[List[str]]: new_stop: Optional[List[str]] = None if isinstance(stop, str): if ( @@ -1286,9 +1211,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): llm_provider=llm_provider, ) - def _extract_json_schema_from_response_format( - self, value: Optional[dict] - ) -> Optional[dict]: + def _extract_json_schema_from_response_format(self, value: Optional[dict]) -> Optional[dict]: if value is None: return None json_schema: Optional[dict] = None @@ -1299,12 +1222,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return json_schema - def map_response_format_to_anthropic_output_format( - self, value: Optional[dict] - ) -> Optional[AnthropicOutputSchema]: - json_schema: Optional[dict] = self._extract_json_schema_from_response_format( - value - ) + def map_response_format_to_anthropic_output_format(self, value: Optional[dict]) -> Optional[AnthropicOutputSchema]: + json_schema: Optional[dict] = self._extract_json_schema_from_response_format(value) if json_schema is None: return None @@ -1333,14 +1252,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self, value: Optional[dict], optional_params: dict, is_thinking_enabled: bool ) -> Optional[AnthropicMessagesTool]: ignore_response_format_types = ["text"] - if ( - value is None or value["type"] in ignore_response_format_types - ): # value is a no-op + if value is None or value["type"] in ignore_response_format_types: # value is a no-op return None - json_schema: Optional[dict] = self._extract_json_schema_from_response_format( - value - ) + json_schema: Optional[dict] = self._extract_json_schema_from_response_format(value) if json_schema is None: return None """ @@ -1368,9 +1283,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): user_location = value_typed.get("user_location") if user_location is not None: anthropic_user_location = AnthropicWebSearchUserLocation(type="approximate") - anthropic_user_location_keys = ( - AnthropicWebSearchUserLocation.__annotations__.keys() - ) + anthropic_user_location_keys = AnthropicWebSearchUserLocation.__annotations__.keys() user_location_approximate = user_location.get("approximate") if user_location_approximate is not None: for key, user_location_value in user_location_approximate.items(): @@ -1381,9 +1294,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## MAP SEARCH CONTEXT SIZE search_context_size = value_typed.get("search_context_size") if search_context_size is not None: - hosted_web_search_tool["max_uses"] = ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES[ - search_context_size - ] + hosted_web_search_tool["max_uses"] = ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES[search_context_size] return hosted_web_search_tool @@ -1424,9 +1335,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): anthropic_edit: Dict[str, Any] = {"type": "compact_20260112"} compact_threshold = entry.get("compact_threshold") # Rewrite to 'trigger' with correct nesting if threshold exists - if compact_threshold is not None and isinstance( - compact_threshold, (int, float) - ): + if compact_threshold is not None and isinstance(compact_threshold, (int, float)): anthropic_edit["trigger"] = { "type": "input_tokens", "value": int(compact_threshold), @@ -1453,9 +1362,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model: str, drop_params: bool, ) -> dict: - is_thinking_enabled = self.is_thinking_enabled( - non_default_params=non_default_params - ) + is_thinking_enabled = self.is_thinking_enabled(non_default_params=non_default_params) # NB: ``map_openai_params`` deliberately does NOT sanitize tool names # here. Names are the *original* OpenAI names at this stage, and must @@ -1470,13 +1377,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): for param, value in non_default_params.items(): if param == "max_tokens": - optional_params["max_tokens"] = ( - value if isinstance(value, int) else max(1, int(round(value))) - ) + optional_params["max_tokens"] = value if isinstance(value, int) else max(1, int(round(value))) elif param == "max_completion_tokens": - optional_params["max_tokens"] = ( - value if isinstance(value, int) else max(1, int(round(value))) - ) + optional_params["max_tokens"] = value if isinstance(value, int) else max(1, int(round(value))) elif param == "tools": anthropic_tools, mcp_servers = self._map_tools(value) optional_params = self._add_tools_to_optional_params( @@ -1485,20 +1388,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if mcp_servers: optional_params["mcp_servers"] = mcp_servers elif param == "tool_choice" or param == "parallel_tool_calls": - _tool_choice: Optional[AnthropicMessagesToolChoice] = ( - self._map_tool_choice( - tool_choice=non_default_params.get("tool_choice"), - parallel_tool_use=non_default_params.get("parallel_tool_calls"), - ) + _tool_choice: Optional[AnthropicMessagesToolChoice] = self._map_tool_choice( + tool_choice=non_default_params.get("tool_choice"), + parallel_tool_use=non_default_params.get("parallel_tool_calls"), ) if _tool_choice is not None: optional_params["tool_choice"] = _tool_choice elif param == "stream" and value is True: optional_params["stream"] = value - elif param == "stop" and ( - isinstance(value, str) or isinstance(value, list) - ): + elif param == "stop" and (isinstance(value, str) or isinstance(value, list)): _value = self._map_stop_sequences(value) if _value is not None: optional_params["stop_sequences"] = _value @@ -1531,15 +1430,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "sonnet_4_6", } ): - _output_format = ( - self.map_response_format_to_anthropic_output_format(value) - ) + _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: optional_params["output_format"] = _output_format else: - _tool = self.map_response_format_to_anthropic_tool( - value, optional_params, is_thinking_enabled - ) + _tool = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled) if _tool is None: continue if not is_thinking_enabled: @@ -1549,9 +1444,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): } optional_params["tool_choice"] = _tool_choice - optional_params = self._add_tools_to_optional_params( - optional_params=optional_params, tools=[_tool] - ) + optional_params = self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool]) optional_params["json_mode"] = True elif ( param == "user" @@ -1586,9 +1479,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: optional_params["thinking"] = mapped_thinking if AnthropicConfig._is_adaptive_thinking_model(model): - mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get( - effort_value - ) + mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(effort_value) if mapped_effort is None: AnthropicConfig._raise_invalid_reasoning_effort( model=model, @@ -1597,24 +1488,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) optional_params["output_config"] = {"effort": mapped_effort} elif param == "web_search_options" and isinstance(value, dict): - hosted_web_search_tool = self.map_web_search_tool( - cast(OpenAIWebSearchOptions, value) - ) - self._add_tools_to_optional_params( - optional_params=optional_params, tools=[hosted_web_search_tool] - ) + hosted_web_search_tool = self.map_web_search_tool(cast(OpenAIWebSearchOptions, value)) + self._add_tools_to_optional_params(optional_params=optional_params, tools=[hosted_web_search_tool]) elif param == "extra_headers": optional_params["extra_headers"] = value elif param == "context_management": # Supports both OpenAI list format and Anthropic dict format if isinstance(value, (list, dict)): - anthropic_context_management = ( - self.map_openai_context_management_to_anthropic(value) - ) + anthropic_context_management = self.map_openai_context_management_to_anthropic(value) if anthropic_context_management is not None: - optional_params["context_management"] = ( - anthropic_context_management - ) + optional_params["context_management"] = anthropic_context_management elif param == "speed" and isinstance(value, str): optional_params["speed"] = value AnthropicConfig._maybe_drop_speed_param( @@ -1659,9 +1542,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: _input_schema.update(cast(AnthropicInputSchema, json_schema)) - _tool = AnthropicMessagesTool( - name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema - ) + _tool = AnthropicMessagesTool(name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema) return _tool def should_strip_billing_metadata(self) -> bool: @@ -1673,9 +1554,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ return False - def translate_system_message( - self, messages: List[AllMessageValues] - ) -> List[AnthropicSystemMessageContent]: + def translate_system_message(self, messages: List[AllMessageValues]) -> List[AnthropicSystemMessageContent]: """ Translate system message to anthropic format. @@ -1692,21 +1571,17 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Skip empty text blocks - Anthropic API raises errors for empty text if not system_message_block["content"]: continue - if self.should_strip_billing_metadata() and system_message_block[ - "content" - ].startswith("x-anthropic-billing-header:"): + if self.should_strip_billing_metadata() and system_message_block["content"].startswith( + "x-anthropic-billing-header:" + ): continue anthropic_system_message_content = AnthropicSystemMessageContent( type="text", text=system_message_block["content"], ) if "cache_control" in system_message_block: - anthropic_system_message_content["cache_control"] = ( - system_message_block["cache_control"] - ) - anthropic_system_message_list.append( - anthropic_system_message_content - ) + anthropic_system_message_content["cache_control"] = system_message_block["cache_control"] + anthropic_system_message_list.append(anthropic_system_message_content) elif isinstance(message["content"], list): for _content in message["content"]: # Skip empty text blocks - Anthropic API raises errors for empty text @@ -1720,20 +1595,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): and text_value.startswith("x-anthropic-billing-header:") ): continue - anthropic_system_message_content = ( - AnthropicSystemMessageContent( - type=_content.get("type"), - text=text_value, - ) + anthropic_system_message_content = AnthropicSystemMessageContent( + type=_content.get("type"), + text=text_value, ) if "cache_control" in _content: - anthropic_system_message_content["cache_control"] = ( - _content["cache_control"] - ) + anthropic_system_message_content["cache_control"] = _content["cache_control"] - anthropic_system_message_list.append( - anthropic_system_message_content - ) + anthropic_system_message_list.append(anthropic_system_message_content) if len(system_prompt_indices) > 0: for idx in reversed(system_prompt_indices): @@ -1761,11 +1630,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## check if code_execution tool is already in tools for tool in tools: tool_type = tool.get("type", None) - if ( - tool_type - and isinstance(tool_type, str) - and tool_type.startswith("code_execution") - ): + if tool_type and isinstance(tool_type, str) and tool_type.startswith("code_execution"): return tools tools.append( AnthropicCodeExecutionTool( @@ -1792,9 +1657,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if beta_value not in existing_values: headers["anthropic-beta"] = f"{existing_beta}, {beta_value}" - def _ensure_context_management_beta_header( - self, headers: dict, context_management: object - ) -> None: + def _ensure_context_management_beta_header(self, headers: dict, context_management: object) -> None: """ Add appropriate beta headers based on context_management edits. """ @@ -1821,9 +1684,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Add compact header if any compact edits/entries exist if has_compact: - self._ensure_beta_header( - headers, ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value - ) + self._ensure_beta_header(headers, ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value) # Add context management header if any other edits/entries exist if has_other: @@ -1832,9 +1693,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, ) - def update_headers_with_optional_anthropic_beta( - self, headers: dict, optional_params: dict - ) -> dict: + def update_headers_with_optional_anthropic_beta(self, headers: dict, optional_params: dict) -> dict: """Update headers with optional anthropic beta.""" # Skip adding beta headers for Vertex requests @@ -1845,39 +1704,25 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _tools = optional_params.get("tools", []) for tool in _tools: - if tool.get("type", None) and tool.get("type").startswith( - ANTHROPIC_HOSTED_TOOLS.WEB_FETCH.value - ): - self._ensure_beta_header( - headers, ANTHROPIC_BETA_HEADER_VALUES.WEB_FETCH_2025_09_10.value - ) - elif tool.get("type", None) and tool.get("type").startswith( - ANTHROPIC_HOSTED_TOOLS.MEMORY.value - ): + if tool.get("type", None) and tool.get("type").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_FETCH.value): + self._ensure_beta_header(headers, ANTHROPIC_BETA_HEADER_VALUES.WEB_FETCH_2025_09_10.value) + elif tool.get("type", None) and tool.get("type").startswith(ANTHROPIC_HOSTED_TOOLS.MEMORY.value): self._ensure_beta_header( headers, ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, ) if optional_params.get("context_management") is not None: - self._ensure_context_management_beta_header( - headers, optional_params["context_management"] - ) + self._ensure_context_management_beta_header(headers, optional_params["context_management"]) output_config = optional_params.get("output_config") if optional_params.get("output_format") is not None or ( isinstance(output_config, dict) and output_config.get("format") is not None ): - self._ensure_beta_header( - headers, ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value - ) + self._ensure_beta_header(headers, ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value) if optional_params.get("speed") == "fast": - self._ensure_beta_header( - headers, ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value - ) + self._ensure_beta_header(headers, ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value) for tool in _tools: if tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE: - self._ensure_beta_header( - headers, ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value - ) + self._ensure_beta_header(headers, ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value) break return headers @@ -1898,14 +1743,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): anthropic_messages_pt, ) - if ( - "tools" not in optional_params - and messages is not None - and has_tool_call_blocks(messages) - ): - optional_params["tools"], _ = self._map_tools( - add_dummy_tool(custom_llm_provider="anthropic") - ) + if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages): + optional_params["tools"], _ = self._map_tools(add_dummy_tool(custom_llm_provider="anthropic")) # Drop thinking param if thinking is enabled but thinking_blocks are missing # This prevents the error: "Expected thinking or redacted_thinking, but found tool_use" @@ -1930,14 +1769,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): AnthropicConfig._maybe_drop_speed_param( model=model, optional_params=optional_params, - drop_params=litellm.drop_params - or litellm_params.get("drop_params") is True, + drop_params=litellm.drop_params or litellm_params.get("drop_params") is True, custom_llm_provider=self.custom_llm_provider, ) - headers = self.update_headers_with_optional_anthropic_beta( - headers=headers, optional_params=optional_params - ) + headers = self.update_headers_with_optional_anthropic_beta(headers=headers, optional_params=optional_params) # === Tool-name sanitization (single chokepoint) === # Anthropic enforces ^[a-zA-Z0-9_-]{1,128}$ on every tool name. We @@ -1988,10 +1824,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## Auto-strip advisor blocks from history if advisor tool is absent. ## Prevents Anthropic 400: advisor_tool_result in history requires advisor tool. _all_tools = optional_params.get("tools") or [] - _has_advisor = any( - isinstance(t, dict) and t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE - for t in _all_tools - ) + _has_advisor = any(isinstance(t, dict) and t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for t in _all_tools) if not _has_advisor: anthropic_messages = strip_advisor_blocks_from_messages(anthropic_messages) @@ -2027,9 +1860,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["metadata"] = {"user_id": _litellm_metadata["user_id"]} ## Ensure metadata only contains user_id (only documented field in Anthropic Messages API) - if "metadata" in optional_params and isinstance( - optional_params["metadata"], dict - ): + if "metadata" in optional_params and isinstance(optional_params["metadata"], dict): _user_id = optional_params["metadata"].get("user_id") if _user_id is not None: optional_params["metadata"] = {"user_id": _user_id} @@ -2060,15 +1891,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): **optional_params, } - self._apply_output_config( - data=data, model=model, optional_params=optional_params - ) + self._apply_output_config(data=data, model=model, optional_params=optional_params) return data - def _apply_output_config( - self, data: dict, model: str, optional_params: dict - ) -> None: + def _apply_output_config(self, data: dict, model: str, optional_params: dict) -> None: """Validate and apply output_config to the request data.""" if "output_config" not in optional_params: return @@ -2087,9 +1914,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): valid_efforts = ["high", "medium", "low", "xhigh", "max"] if effort is not None and effort not in valid_efforts: raise litellm.exceptions.BadRequestError( - message=( - f"Invalid effort value: {effort!r}. Must be one of: 'high', 'medium', 'low', 'xhigh', 'max'" - ), + message=(f"Invalid effort value: {effort!r}. Must be one of: 'high', 'medium', 'low', 'xhigh', 'max'"), model=model, llm_provider=self.custom_llm_provider or "anthropic", ) @@ -2116,9 +1941,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return None, tool_calls, None json_indices = [ - i - for i, t in enumerate(tool_calls) - if t.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME + i for i, t in enumerate(tool_calls) if t.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME ] if not json_indices: return None, tool_calls, None @@ -2127,16 +1950,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): json_tool = tool_calls[json_indices[0]] if json_tool.get("function", {}).get("arguments") is None: return None, tool_calls, None - _message = AnthropicConfig._convert_tool_response_to_message( - tool_calls=[json_tool] - ) + _message = AnthropicConfig._convert_tool_response_to_message(tool_calls=[json_tool]) return _message, [], None first_json = tool_calls[json_indices[0]] json_msg = AnthropicConfig._convert_tool_response_to_message([first_json]) - extra_content: Optional[str] = ( - json_msg.content if json_msg is not None else None - ) + extra_content: Optional[str] = json_msg.content if json_msg is not None else None filtered_tools = [t for i, t in enumerate(tool_calls) if i not in json_indices] return None, filtered_tools, extra_content @@ -2145,11 +1964,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) -> Tuple[ str, Optional[List[Any]], - Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ], + Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]], Optional[str], List[ChatCompletionToolCallChunk], Optional[List[Any]], @@ -2158,11 +1973,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ]: text_content = "" citations: Optional[List[Any]] = None - thinking_blocks: Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ] = None + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None reasoning_content: Optional[str] = None tool_calls: List[ChatCompletionToolCallChunk] = [] web_search_results: Optional[List[Any]] = None @@ -2206,9 +2017,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): elif content["type"] == "redacted_thinking": if thinking_blocks is None: thinking_blocks = [] - thinking_blocks.append( - cast(ChatCompletionRedactedThinkingBlock, content) - ) + thinking_blocks.append(cast(ChatCompletionRedactedThinkingBlock, content)) ## COMPACTION elif content["type"] == "compaction": @@ -2256,15 +2065,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) -> Usage: # NOTE: Sometimes the usage object has None set explicitly for token counts, meaning .get() & key access returns None, and we need to account for this raw_prompt_tokens = usage_object.get("input_tokens", 0) or 0 - prompt_tokens: int = ( - int(raw_prompt_tokens) if isinstance(raw_prompt_tokens, (int, float)) else 0 - ) + prompt_tokens: int = int(raw_prompt_tokens) if isinstance(raw_prompt_tokens, (int, float)) else 0 raw_completion_tokens = usage_object.get("output_tokens", 0) or 0 - completion_tokens: int = ( - int(raw_completion_tokens) - if isinstance(raw_completion_tokens, (int, float)) - else 0 - ) + completion_tokens: int = int(raw_completion_tokens) if isinstance(raw_completion_tokens, (int, float)) else 0 _usage = usage_object cache_creation_input_tokens: int = 0 cache_read_input_tokens: int = 0 @@ -2282,28 +2085,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): iterations: Optional[List[Any]] = _usage.get("iterations") if iterations: prompt_tokens = sum(it.get("input_tokens", 0) or 0 for it in iterations) - completion_tokens = sum( - it.get("output_tokens", 0) or 0 for it in iterations - ) - cache_creation_input_tokens = sum( - it.get("cache_creation_input_tokens", 0) or 0 for it in iterations - ) - cache_read_input_tokens = sum( - it.get("cache_read_input_tokens", 0) or 0 for it in iterations - ) + completion_tokens = sum(it.get("output_tokens", 0) or 0 for it in iterations) + cache_creation_input_tokens = sum(it.get("cache_creation_input_tokens", 0) or 0 for it in iterations) + cache_read_input_tokens = sum(it.get("cache_read_input_tokens", 0) or 0 for it in iterations) prompt_tokens += cache_creation_input_tokens + cache_read_input_tokens if not iterations: - if ( - "cache_creation_input_tokens" in _usage - and _usage["cache_creation_input_tokens"] is not None - ): + if "cache_creation_input_tokens" in _usage and _usage["cache_creation_input_tokens"] is not None: cache_creation_input_tokens = _usage["cache_creation_input_tokens"] prompt_tokens += cache_creation_input_tokens - if ( - "cache_read_input_tokens" in _usage - and _usage["cache_read_input_tokens"] is not None - ): + if "cache_read_input_tokens" in _usage and _usage["cache_read_input_tokens"] is not None: cache_read_input_tokens = _usage["cache_read_input_tokens"] prompt_tokens += cache_read_input_tokens if "server_tool_use" in _usage and _usage["server_tool_use"] is not None: @@ -2311,16 +2102,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "web_search_requests" in _usage["server_tool_use"] and _usage["server_tool_use"]["web_search_requests"] is not None ): - web_search_requests = cast( - int, _usage["server_tool_use"]["web_search_requests"] - ) + web_search_requests = cast(int, _usage["server_tool_use"]["web_search_requests"]) if ( "tool_search_requests" in _usage["server_tool_use"] and _usage["server_tool_use"]["tool_search_requests"] is not None ): - tool_search_requests = cast( - int, _usage["server_tool_use"]["tool_search_requests"] - ) + tool_search_requests = cast(int, _usage["server_tool_use"]["tool_search_requests"]) # Count tool_search_requests from content blocks if not in usage # Anthropic doesn't always include tool_search_requests in the usage object @@ -2336,17 +2123,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if "cache_creation" in _usage and _usage["cache_creation"] is not None: cache_creation_token_details = CacheCreationTokenDetails( - ephemeral_5m_input_tokens=_usage["cache_creation"].get( - "ephemeral_5m_input_tokens" - ), - ephemeral_1h_input_tokens=_usage["cache_creation"].get( - "ephemeral_1h_input_tokens" - ), + ephemeral_5m_input_tokens=_usage["cache_creation"].get("ephemeral_5m_input_tokens"), + ephemeral_1h_input_tokens=_usage["cache_creation"].get("ephemeral_1h_input_tokens"), ) - raw_input_tokens = ( - prompt_tokens - cache_read_input_tokens - cache_creation_input_tokens - ) + raw_input_tokens = prompt_tokens - cache_read_input_tokens - cache_creation_input_tokens prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cache_read_input_tokens, cache_creation_tokens=cache_creation_input_tokens, @@ -2355,18 +2136,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) # Always populate completion_token_details, not just when there's reasoning_content estimated_reasoning_tokens = ( - token_counter(text=reasoning_content, count_response_tokens=True) - if reasoning_content - else 0 + token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 ) reasoning_tokens = min(estimated_reasoning_tokens, completion_tokens) completion_token_details = CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else 0, - text_tokens=( - completion_tokens - reasoning_tokens - if reasoning_tokens > 0 - else completion_tokens - ), + text_tokens=(completion_tokens - reasoning_tokens if reasoning_tokens > 0 else completion_tokens), ) total_tokens = prompt_tokens + completion_tokens @@ -2393,9 +2168,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) return usage - def _build_code_by_id_map( - self, tool_calls: List[ChatCompletionToolCallChunk] - ) -> Dict[str, str]: + def _build_code_by_id_map(self, tool_calls: List[ChatCompletionToolCallChunk]) -> Dict[str, str]: code_by_id: Dict[str, str] = {} for tc in tool_calls: try: @@ -2437,11 +2210,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self, completion_response: dict, citations: Optional[List[Any]], - thinking_blocks: Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ], + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]], web_search_results: Optional[List[Any]], tool_results: Optional[List[Any]], compaction_blocks: Optional[List[Any]], @@ -2467,12 +2236,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else None ) code_by_id = self._build_code_by_id_map(tool_calls) - code_interpreter_results = self._build_code_interpreter_results( - tool_results, code_by_id, container_id - ) - provider_specific_fields["code_interpreter_results"] = ( - code_interpreter_results - ) + code_interpreter_results = self._build_code_interpreter_results(tool_results, code_by_id, container_id) + provider_specific_fields["code_interpreter_results"] = code_interpreter_results container = completion_response.get("container") if container is not None: @@ -2494,9 +2259,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): tool_name_reverse_map: Optional[Dict[str, str]] = None, ): _hidden_params: Dict = {} - _hidden_params["additional_headers"] = process_anthropic_headers( - dict(raw_response.headers) - ) + _hidden_params["additional_headers"] = process_anthropic_headers(dict(raw_response.headers)) if "error" in completion_response: response_headers = getattr(raw_response, "headers", None) raise AnthropicError( @@ -2548,17 +2311,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): tool_calls, ) - json_mode_message, tool_calls_for_message, json_extra_content = ( - self._resolve_json_mode_non_streaming( - json_mode=json_mode, - tool_calls=tool_calls, - ) + json_mode_message, tool_calls_for_message, json_extra_content = self._resolve_json_mode_non_streaming( + json_mode=json_mode, + tool_calls=tool_calls, ) merged_text = text_content or "" if json_extra_content: - merged_text = ( - merged_text + json_extra_content if merged_text else json_extra_content - ) + merged_text = merged_text + json_extra_content if merged_text else json_extra_content _message = litellm.Message( tool_calls=tool_calls_for_message, @@ -2574,9 +2333,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _message = json_mode_message model_response.choices[0].message = _message - model_response._hidden_params["original_response"] = completion_response[ - "content" - ] + model_response._hidden_params["original_response"] = completion_response["content"] model_response.choices[0].finish_reason = cast( OpenAIChatCompletionFinishReason, map_finish_reason(completion_response["stop_reason"]), @@ -2611,11 +2368,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): message = messages[-1] message_content = message.get("content") - if ( - message["role"] == "assistant" - and message.get("prefix", False) - and isinstance(message_content, str) - ): + if message["role"] == "assistant" and message.get("prefix", False) and isinstance(message_content, str): return message_content return None @@ -2648,9 +2401,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise AnthropicError( - message="Unable to get json response - {}, Original Response: {}".format( - str(e), raw_response.text - ), + message="Unable to get json response - {}, Original Response: {}".format(str(e), raw_response.text), status_code=raw_response.status_code, headers=response_headers, ) @@ -2683,16 +2434,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): """ ## HANDLE JSON MODE - anthropic returns single function call - json_mode_content_str: Optional[str] = tool_calls[0]["function"].get( - "arguments" - ) + json_mode_content_str: Optional[str] = tool_calls[0]["function"].get("arguments") try: if json_mode_content_str is not None: args = json.loads(json_mode_content_str) - if ( - isinstance(args, dict) - and (values := args.get("values")) is not None - ): + if isinstance(args, dict) and (values := args.get("values")) is not None: _message = litellm.Message(content=json.dumps(values)) return _message else: diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 4d875678e98..721fcae4c08 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -63,9 +63,7 @@ def _merge_beta_headers(existing: Optional[str], new_beta: str) -> str: return ",".join(sorted(betas)) -def optionally_handle_anthropic_oauth( - headers: dict, api_key: Optional[str] -) -> tuple[dict, Optional[str]]: +def optionally_handle_anthropic_oauth(headers: dict, api_key: Optional[str]) -> tuple[dict, Optional[str]]: """ Handle Anthropic OAuth token detection and header setup. @@ -84,18 +82,14 @@ def optionally_handle_anthropic_oauth( if auth_header and auth_header.startswith(f"Bearer {ANTHROPIC_OAUTH_TOKEN_PREFIX}"): api_key = auth_header.replace("Bearer ", "") headers.pop("x-api-key", None) - headers["anthropic-beta"] = _merge_beta_headers( - headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER - ) + headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER) headers["anthropic-dangerous-direct-browser-access"] = "true" return headers, api_key # Check api_key directly (standard chat/completion flow) if api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX): headers.pop("x-api-key", None) headers["authorization"] = f"Bearer {api_key}" - headers["anthropic-beta"] = _merge_beta_headers( - headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER - ) + headers["anthropic-beta"] = _merge_beta_headers(headers.get("anthropic-beta"), ANTHROPIC_OAUTH_BETA_HEADER) headers["anthropic-dangerous-direct-browser-access"] = "true" return headers, api_key @@ -135,18 +129,14 @@ class AnthropicModelInfo(BaseLLMModelInfo): file_ids = get_file_ids_from_messages(messages) return len(file_ids) > 0 - def is_mcp_server_used( - self, mcp_servers: Optional[List[AnthropicMcpServerTool]] - ) -> bool: + def is_mcp_server_used(self, mcp_servers: Optional[List[AnthropicMcpServerTool]]) -> bool: if mcp_servers is None: return False if mcp_servers: return True return False - def is_computer_tool_used( - self, tools: Optional[List[AllAnthropicToolsValues]] - ) -> Optional[str]: + def is_computer_tool_used(self, tools: Optional[List[AllAnthropicToolsValues]]) -> Optional[str]: """Returns the computer tool version if used, e.g. 'computer_20250124' or None""" if tools is None: return None @@ -155,16 +145,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): return tool["type"] return None - def is_web_search_tool_used( - self, tools: Optional[List[AllAnthropicToolsValues]] - ) -> bool: + def is_web_search_tool_used(self, tools: Optional[List[AllAnthropicToolsValues]]) -> bool: """Returns True if web_search tool is used""" if tools is None: return False for tool in tools: - if "type" in tool and tool["type"].startswith( - ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value - ): + if "type" in tool and tool["type"].startswith(ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value): return True return False @@ -174,11 +160,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): """ for message in messages: - if ( - "content" in message - and message["content"] is not None - and isinstance(message["content"], list) - ): + if "content" in message and message["content"] is not None and isinstance(message["content"], list): for content in message["content"]: if "type" in content and content["type"] != "text": return True @@ -220,9 +202,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): function = tool.get("function", {}) if isinstance(function, dict): function_allowed_callers = function.get("allowed_callers", None) - if function_allowed_callers and isinstance( - function_allowed_callers, list - ): + if function_allowed_callers and isinstance(function_allowed_callers, list): if "code_execution_20250825" in function_allowed_callers: return True @@ -240,11 +220,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): for tool in tools: # Check top-level input_examples input_examples = tool.get("input_examples", None) - if ( - input_examples - and isinstance(input_examples, list) - and len(input_examples) > 0 - ): + if input_examples and isinstance(input_examples, list) and len(input_examples) > 0: return True # Check function.input_examples for OpenAI format tools @@ -269,9 +245,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): Driven by the ``supports_sampling_params`` flag in the model map; the name check remains only as a fallback for provider-routed ids whose map entries predate the flag.""" - flag = AnthropicModelInfo._get_model_capability( - model, "supports_sampling_params" - ) + flag = AnthropicModelInfo._get_model_capability(model, "supports_sampling_params") if flag is not None: return flag model_lower = model.lower() @@ -303,14 +277,10 @@ class AnthropicModelInfo(BaseLLMModelInfo): ``optional_params[output_key]`` unless the model removed sampling params, in which case drop the param (with drop_params) or raise a clean client-side 400.""" - if AnthropicModelInfo._supports_sampling_params(model) or ( - param == "temperature" and value == 1 - ): + if AnthropicModelInfo._supports_sampling_params(model) or (param == "temperature" and value == 1): optional_params[output_key] = value elif not (litellm.drop_params or drop_params): - supported_hint = ( - "Only temperature=1 is supported. " if param == "temperature" else "" - ) + supported_hint = "Only temperature=1 is supported. " if param == "temperature" else "" raise litellm.utils.UnsupportedParamsError( message=( f"{model} does not support {param}={value}. {supported_hint}" @@ -410,13 +380,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): (an unmapped alias or a future release not yet in the map) is treated as non-adaptive until a ``fallback_generalizations`` rule covers it. """ - return AnthropicModelInfo._supports_model_capability( - model, "supports_adaptive_thinking" - ) + return AnthropicModelInfo._supports_model_capability(model, "supports_adaptive_thinking") - def is_effort_used( - self, optional_params: Optional[dict], model: Optional[str] = None - ) -> bool: + def is_effort_used(self, optional_params: Optional[dict], model: Optional[str] = None) -> bool: """ Check if effort parameter is being used and requires a beta header. @@ -477,9 +443,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): return True return False - def _get_user_anthropic_beta_headers( - self, anthropic_beta_header: Optional[str] - ) -> Optional[List[str]]: + def _get_user_anthropic_beta_headers(self, anthropic_beta_header: Optional[str]) -> Optional[List[str]]: if anthropic_beta_header is None: return None return anthropic_beta_header.split(",") @@ -546,13 +510,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): return list(set(betas)) @staticmethod - def _make_api_key_auth_header( - api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False - ) -> dict: + def _make_api_key_auth_header(api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False) -> dict: if use_bearer_for_custom_base and ( - api_base - and "api.anthropic.com" not in api_base - and not api_key.startswith("sk-ant-") + api_base and "api.anthropic.com" not in api_base and not api_key.startswith("sk-ant-") ): value = api_key if api_key.startswith("Bearer ") else f"Bearer {api_key}" return {"authorization": value} @@ -627,11 +587,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): elif auth_token and not api_key: headers["authorization"] = f"Bearer {auth_token}" elif api_key: - headers.update( - self._make_api_key_auth_header( - api_key, api_base, use_bearer_for_custom_base - ) - ) + headers.update(self._make_api_key_auth_header(api_key, api_base, use_bearer_for_custom_base)) if user_anthropic_beta_headers is not None: betas.update(user_anthropic_beta_headers) @@ -642,9 +598,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): if web_search_tool_used: from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES - headers["anthropic-beta"] = ( - ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value - ) + headers["anthropic-beta"] = ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value elif len(betas) > 0: headers["anthropic-beta"] = ",".join(betas) @@ -663,13 +617,10 @@ class AnthropicModelInfo(BaseLLMModelInfo): if api_base is None and isinstance(litellm_params, dict): api_base = litellm_params.get("api_base") use_bearer_for_custom_base: bool = bool( - isinstance(litellm_params, dict) - and litellm_params.get("use_bearer_for_custom_base", False) + isinstance(litellm_params, dict) and litellm_params.get("use_bearer_for_custom_base", False) ) # Check for Anthropic OAuth token in headers - headers, api_key = optionally_handle_anthropic_oauth( - headers=headers, api_key=api_key - ) + headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) api_key = AnthropicModelInfo.get_api_key(api_key) # Resolve auth_token from ANTHROPIC_AUTH_TOKEN if api_key is not set auth_token: Optional[str] = None @@ -685,22 +636,16 @@ class AnthropicModelInfo(BaseLLMModelInfo): tools = optional_params.get("tools") prompt_caching_set = self.is_cache_control_set(messages=messages) computer_tool_used = self.is_computer_tool_used(tools=tools) - mcp_server_used = self.is_mcp_server_used( - mcp_servers=optional_params.get("mcp_servers") - ) + mcp_server_used = self.is_mcp_server_used(mcp_servers=optional_params.get("mcp_servers")) pdf_used = self.is_pdf_used(messages=messages) file_id_used = self.is_file_id_used(messages=messages) web_search_tool_used = self.is_web_search_tool_used(tools=tools) tool_search_used = self.is_tool_search_used(tools=tools) - programmatic_tool_calling_used = self.is_programmatic_tool_calling_used( - tools=tools - ) + programmatic_tool_calling_used = self.is_programmatic_tool_calling_used(tools=tools) input_examples_used = self.is_input_examples_used(tools=tools) effort_used = self.is_effort_used(optional_params=optional_params, model=model) code_execution_tool_used = self.is_code_execution_tool_used(tools=tools) - container_with_skills_used = self.is_container_with_skills_used( - optional_params=optional_params - ) + container_with_skills_used = self.is_container_with_skills_used(optional_params=optional_params) user_anthropic_beta_headers = self._get_user_anthropic_beta_headers( anthropic_beta_header=headers.get("anthropic-beta") ) @@ -773,9 +718,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): if resolved_key is not None: if is_anthropic_oauth_key(resolved_key): return {"authorization": f"Bearer {resolved_key}"} - return AnthropicModelInfo._make_api_key_auth_header( - resolved_key, api_base, use_bearer_for_custom_base - ) + return AnthropicModelInfo._make_api_key_auth_header(resolved_key, api_base, use_bearer_for_custom_base) auth_token = AnthropicModelInfo.get_auth_token() if auth_token is not None: return {"authorization": f"Bearer {auth_token}"} @@ -785,9 +728,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): def get_base_model(model: Optional[str] = None) -> Optional[str]: return model.replace("anthropic/", "") if model else None - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: api_base = AnthropicModelInfo.get_api_base(api_base) auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if api_base is None or auth_header is None: @@ -831,9 +772,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): return AnthropicTokenCounter() -def strip_advisor_blocks_from_messages( - messages: List[Any], replace_with_text: bool = False -) -> List[Any]: +def strip_advisor_blocks_from_messages(messages: List[Any], replace_with_text: bool = False) -> List[Any]: """ Remove (or replace) server_tool_use (name='advisor') and advisor_tool_result blocks from assistant message content. @@ -859,11 +798,7 @@ def strip_advisor_blocks_from_messages( # Collect advisor server_tool_use ids and their advice text (for replace mode). advisor_id_to_text: dict = {} for block in content: - if ( - isinstance(block, dict) - and block.get("type") == "server_tool_use" - and block.get("name") == "advisor" - ): + if isinstance(block, dict) and block.get("type") == "server_tool_use" and block.get("name") == "advisor": bid = block.get("id") if bid: advisor_id_to_text[bid] = None # text filled in below @@ -884,11 +819,7 @@ def strip_advisor_blocks_from_messages( raw if isinstance(raw, str) else next( - ( - b.get("text", "") - for b in raw - if isinstance(b, dict) and b.get("type") == "text" - ), + (b.get("text", "") for b in raw if isinstance(b, dict) and b.get("type") == "text"), "", ) ) @@ -905,8 +836,7 @@ def strip_advisor_blocks_from_messages( and block.get("id") in advisor_id_to_text ) is_advisor_result = ( - block.get("type") == "advisor_tool_result" - and block.get("tool_use_id") in advisor_id_to_text + block.get("type") == "advisor_tool_result" and block.get("tool_use_id") in advisor_id_to_text ) if is_advisor_use: if replace_with_text: @@ -939,12 +869,7 @@ def is_anthropic_invalid_thinking_signature_error(error_text: str) -> bool: if not error_text: return False lower = error_text.lower() - return ( - "invalid" in lower - and "signature" in lower - and "thinking" in lower - and "block" in lower - ) + return "invalid" in lower and "signature" in lower and "thinking" in lower and "block" in lower def strip_thinking_blocks_from_anthropic_messages(messages: List[Any]) -> List[Any]: @@ -964,12 +889,7 @@ def strip_thinking_blocks_from_anthropic_messages(messages: List[Any]) -> List[A content = mm.get("content") if isinstance(content, list): filtered = [ - b - for b in content - if not ( - isinstance(b, dict) - and b.get("type") in ("thinking", "redacted_thinking") - ) + b for b in content if not (isinstance(b, dict) and b.get("type") in ("thinking", "redacted_thinking")) ] if not filtered: continue @@ -1042,11 +962,7 @@ def normalize_anthropic_tool_use_id(raw_id: str) -> str: Strips Gemini thought-signature suffixes (``__thought__``) first, then replaces any remaining invalid characters with underscores. """ - base_id = ( - raw_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] - if THOUGHT_SIGNATURE_SEPARATOR in raw_id - else raw_id - ) + base_id = raw_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] if THOUGHT_SIGNATURE_SEPARATOR in raw_id else raw_id sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", base_id) return sanitized or "tool_use_id" @@ -1098,25 +1014,15 @@ def sanitize_tool_use_ids_in_anthropic_messages(messages: list[Any]) -> list[Any def process_anthropic_headers(headers: Union[httpx.Headers, dict]) -> dict: openai_headers = {} if "anthropic-ratelimit-requests-limit" in headers: - openai_headers["x-ratelimit-limit-requests"] = headers[ - "anthropic-ratelimit-requests-limit" - ] + openai_headers["x-ratelimit-limit-requests"] = headers["anthropic-ratelimit-requests-limit"] if "anthropic-ratelimit-requests-remaining" in headers: - openai_headers["x-ratelimit-remaining-requests"] = headers[ - "anthropic-ratelimit-requests-remaining" - ] + openai_headers["x-ratelimit-remaining-requests"] = headers["anthropic-ratelimit-requests-remaining"] if "anthropic-ratelimit-tokens-limit" in headers: - openai_headers["x-ratelimit-limit-tokens"] = headers[ - "anthropic-ratelimit-tokens-limit" - ] + openai_headers["x-ratelimit-limit-tokens"] = headers["anthropic-ratelimit-tokens-limit"] if "anthropic-ratelimit-tokens-remaining" in headers: - openai_headers["x-ratelimit-remaining-tokens"] = headers[ - "anthropic-ratelimit-tokens-remaining" - ] + openai_headers["x-ratelimit-remaining-tokens"] = headers["anthropic-ratelimit-tokens-remaining"] - llm_response_headers = { - "{}-{}".format("llm_provider", k): v for k, v in headers.items() - } + llm_response_headers = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()} additional_headers = {**llm_response_headers, **openai_headers} return additional_headers diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index a8798cd5d0e..d06eac51101 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -36,9 +36,7 @@ class AnthropicTextError(BaseLLMException): def __init__(self, status_code, message): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://api.anthropic.com/v1/complete" - ) + self.request = httpx.Request(method="POST", url="https://api.anthropic.com/v1/complete") self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( message=self.message, @@ -55,9 +53,7 @@ class AnthropicTextConfig(BaseConfig): to pass metadata to anthropic, it's {"user_id": "any-relevant-information"} """ - max_tokens_to_sample: Optional[int] = ( - litellm.max_tokens - ) # anthropic requires a default + max_tokens_to_sample: Optional[int] = litellm.max_tokens # anthropic requires a default stop_sequences: Optional[list] = None temperature: Optional[int] = None top_p: Optional[int] = None @@ -66,9 +62,7 @@ class AnthropicTextConfig(BaseConfig): def __init__( self, - max_tokens_to_sample: Optional[ - int - ] = DEFAULT_MAX_TOKENS, # anthropic requires a default + max_tokens_to_sample: Optional[int] = DEFAULT_MAX_TOKENS, # anthropic requires a default stop_sequences: Optional[list] = None, temperature: Optional[int] = None, top_p: Optional[int] = None, @@ -112,9 +106,7 @@ class AnthropicTextConfig(BaseConfig): litellm_params: dict, headers: dict, ) -> dict: - prompt = self._get_anthropic_text_prompt_from_messages( - messages=messages, model=model - ) + prompt = self._get_anthropic_text_prompt_from_messages(messages=messages, model=model) ## Load Config config = litellm.AnthropicTextConfig.get_config() for k, v in config.items(): @@ -196,12 +188,8 @@ class AnthropicTextConfig(BaseConfig): try: completion_response = raw_response.json() except Exception: - raise AnthropicTextError( - message=raw_response.text, status_code=raw_response.status_code - ) - prompt = self._get_anthropic_text_prompt_from_messages( - messages=messages, model=model - ) + raise AnthropicTextError(message=raw_response.text, status_code=raw_response.status_code) + prompt = self._get_anthropic_text_prompt_from_messages(messages=messages, model=model) if "error" in completion_response: raise AnthropicTextError( message=str(completion_response["error"]), @@ -215,9 +203,7 @@ class AnthropicTextConfig(BaseConfig): model_response.choices[0].finish_reason = completion_response["stop_reason"] ## CALCULATING USAGE - prompt_tokens = len( - encoding.encode(prompt) - ) ##[TODO] use the anthropic tokenizer here + prompt_tokens = len(encoding.encode(prompt)) ##[TODO] use the anthropic tokenizer here completion_tokens = len( encoding.encode(model_response["choices"][0]["message"].get("content", "")) ) ##[TODO] use the anthropic tokenizer here @@ -245,9 +231,7 @@ class AnthropicTextConfig(BaseConfig): def _is_anthropic_text_model(model: str) -> bool: return model == "claude-2" or model == "claude-instant-1" - def _get_anthropic_text_prompt_from_messages( - self, messages: List[AllMessageValues], model: str - ) -> str: + def _get_anthropic_text_prompt_from_messages(self, messages: List[AllMessageValues], model: str) -> str: custom_prompt_dict = litellm.custom_prompt_dict if model in custom_prompt_dict: # check if the model has a registered custom prompt @@ -259,9 +243,7 @@ class AnthropicTextConfig(BaseConfig): messages=messages, ) else: - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="anthropic" - ) + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic") return str(prompt) diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index fc34938b281..82a97b53d28 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -20,9 +20,7 @@ if TYPE_CHECKING: import litellm -def _compute_cache_only_cost( - model_info: "ModelInfo", usage: "Usage", service_tier: str | None = None -) -> float: +def _compute_cache_only_cost(model_info: "ModelInfo", usage: "Usage", service_tier: str | None = None) -> float: """ Return only the cache-related portion of the prompt cost (cache read + cache write). @@ -40,9 +38,7 @@ def _compute_cache_only_cost( cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost, - ) = _get_token_base_cost( - model_info=model_info, usage=usage, service_tier=service_tier - ) + ) = _get_token_base_cost(model_info=model_info, usage=usage, service_tier=service_tier) cache_cost = float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost @@ -52,9 +48,7 @@ def _compute_cache_only_cost( ): cache_cost += calculate_cache_writing_cost( cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"], - cache_creation_token_details=prompt_tokens_details[ - "cache_creation_token_details" - ], + cache_creation_token_details=prompt_tokens_details["cache_creation_token_details"], cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, cache_creation_cost=cache_creation_cost, ) @@ -62,9 +56,7 @@ def _compute_cache_only_cost( return cache_cost -def cost_per_token( - model: str, usage: "Usage", service_tier: str | None = None -) -> Tuple[float, float]: +def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -86,9 +78,7 @@ def cost_per_token( # Apply provider_specific_entry multipliers for geo/speed routing try: - model_info = litellm.get_model_info( - model=model, custom_llm_provider="anthropic" - ) + model_info = litellm.get_model_info(model=model, custom_llm_provider="anthropic") provider_specific_entry: dict = model_info.get("provider_specific_entry") or {} multiplier = 1.0 @@ -102,9 +92,7 @@ def cost_per_token( multiplier *= provider_specific_entry.get("fast", 1.0) if multiplier != 1.0: - cache_cost = _compute_cache_only_cost( - model_info=model_info, usage=usage, service_tier=service_tier - ) + cache_cost = _compute_cache_only_cost(model_info=model_info, usage=usage, service_tier=service_tier) prompt_cost = (prompt_cost - cache_cost) * multiplier + cache_cost completion_cost *= multiplier except Exception: @@ -156,9 +144,7 @@ def get_cost_for_anthropic_web_search( if usage is None: return 0.0 - web_search_requests = _get_web_search_requests( - getattr(usage, "server_tool_use", None) - ) + web_search_requests = _get_web_search_requests(getattr(usage, "server_tool_use", None)) if web_search_requests is None: return 0.0 @@ -166,9 +152,7 @@ def get_cost_for_anthropic_web_search( search_context_pricing: SearchContextCostPerQuery = ( model_info.get("search_context_cost_per_query") or SearchContextCostPerQuery() ) - cost_per_web_search_request = search_context_pricing.get( - "search_context_size_medium", 0.0 - ) + cost_per_web_search_request = search_context_pricing.get("search_context_size_medium", 0.0) if cost_per_web_search_request is None or cost_per_web_search_request == 0.0: return 0.0 diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 4d0af0b36c8..e70e0f19b33 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -54,9 +54,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): # Validate the request self.validate_request(model, messages) - verbose_logger.debug( - f"Processing Anthropic CountTokens request for model: {model}" - ) + verbose_logger.debug(f"Processing Anthropic CountTokens request for model: {model}") # Transform request to Anthropic format request_body = self.transform_request_to_count_tokens( @@ -77,14 +75,10 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): headers = self.get_required_headers(api_key) # Use LiteLLM's async httpx client - async_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.ANTHROPIC - ) + async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) # Use provided timeout or fall back to litellm.request_timeout - request_timeout = ( - timeout if timeout is not None else litellm.request_timeout - ) + request_timeout = timeout if timeout is not None else litellm.request_timeout response = await async_client.post( endpoint_url, diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 93989c58547..89249ec42f0 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -81,9 +81,7 @@ class AnthropicTokenCounter(BaseTokenCounter): original_response=result, ) except AnthropicError as e: - verbose_logger.warning( - f"Anthropic CountTokens API error: status={e.status_code}, message={e.message}" - ) + verbose_logger.warning(f"Anthropic CountTokens API error: status={e.status_code}, message={e.message}") return TokenCountResponse( total_tokens=0, request_model=request_model, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index efb913f709a..812e0f62c96 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -107,9 +107,7 @@ async def _prepare_context_managed_request( messages=cast(List[Dict[str, Any]], messages), system=system, ) - working_messages = ( - history_result.messages if history_result is not None else messages - ) + working_messages = history_result.messages if history_result is not None else messages working_system = history_result.system if history_result is not None else system polyfill_result = await _run_polyfill_if_enabled( @@ -165,10 +163,7 @@ def _polyfill_will_run( COMPACT_EDIT_TYPE, ) - return any( - isinstance(edit, dict) and edit.get("type") == COMPACT_EDIT_TYPE - for edit in edits - ) + return any(isinstance(edit, dict) and edit.get("type") == COMPACT_EDIT_TYPE for edit in edits) def _spec_has_non_compact_edits( @@ -195,9 +190,7 @@ def _spec_has_non_compact_edits( ) return any( - isinstance(edit, dict) - and isinstance(edit.get("type"), str) - and edit.get("type") != COMPACT_EDIT_TYPE + isinstance(edit, dict) and isinstance(edit.get("type"), str) and edit.get("type") != COMPACT_EDIT_TYPE for edit in edits ) @@ -215,9 +208,7 @@ def _normalize_spec_edits( if not context_management_spec: return None - effective_drop_params = ( - drop_params if drop_params is not None else litellm.drop_params - ) + effective_drop_params = drop_params if drop_params is not None else litellm.drop_params if effective_drop_params: return None @@ -253,9 +244,7 @@ async def _run_polyfill_if_enabled( if not context_management_spec: return None - effective_drop_params = ( - drop_params if drop_params is not None else litellm.drop_params - ) + effective_drop_params = drop_params if drop_params is not None else litellm.drop_params if effective_drop_params: return None @@ -275,9 +264,7 @@ async def _run_polyfill_if_enabled( # 400. Other exception types fall into the best-effort branch below. raise except Exception as e: - verbose_logger.exception( - "context_management polyfill: skipping edits due to error: %s", e - ) + verbose_logger.exception("context_management polyfill: skipping edits due to error: %s", e) # Best-effort swallow is only safe for compact-only specs, where the # caller's compaction-block-slicing safety net produces a correct # (if degraded) result. When the spec also requested non-compact @@ -338,9 +325,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: model = completion_kwargs.get("model") try: - model_info = get_model_info( - model=cast(str, model), custom_llm_provider=custom_llm_provider - ) + model_info = get_model_info(model=cast(str, model), custom_llm_provider=custom_llm_provider) if model_info and model_info.get("supports_reasoning") is False: # Model doesn't support reasoning/responses API, don't route return @@ -363,13 +348,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: reasoning_dict["summary"] = "detailed" completion_kwargs["reasoning_effort"] = reasoning_dict elif isinstance(reasoning_effort, dict): - if ( - "summary" not in reasoning_effort - and "generate_summary" not in reasoning_effort - ): - effective_summary = ( - summary if summary else ("detailed" if auto_summary else None) - ) + if "summary" not in reasoning_effort and "generate_summary" not in reasoning_effort: + effective_summary = summary if summary else ("detailed" if auto_summary else None) if effective_summary: updated_reasoning_effort = dict(reasoning_effort) updated_reasoning_effort["summary"] = effective_summary @@ -404,9 +384,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_kwargs["reasoning_effort"] = normalized elif isinstance(reasoning_effort, dict) and "effort" in reasoning_effort: effort = reasoning_effort["effort"] - normalized = normalize_reasoning_effort_value( - effort, model=model, custom_llm_provider=custom_llm_provider - ) + normalized = normalize_reasoning_effort_value(effort, model=model, custom_llm_provider=custom_llm_provider) if normalized != effort: completion_kwargs["reasoning_effort"] = { **reasoning_effort, @@ -483,9 +461,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: ( openai_request, tool_name_mapping, - ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( - request_data - ) + ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping(request_data) if openai_request is None: raise ValueError("Failed to translate request to OpenAI format") @@ -516,31 +492,19 @@ class LiteLLMMessagesToCompletionTransformationHandler: # NOTE: extra_kwargs was already coerced from None to {} at the top of # this method (line ~220). It is guaranteed to be a dict here. for key, value in extra_kwargs.items(): - if ( - key == "litellm_logging_obj" - and value is not None - and isinstance(value, LiteLLMLoggingObject) - ): + if key == "litellm_logging_obj" and value is not None and isinstance(value, LiteLLMLoggingObject): from litellm.types.utils import CallTypes setattr(value, "call_type", CallTypes.anthropic_messages.value) - setattr( - value, "stream_options", completion_kwargs.get("stream_options") - ) - if ( - key not in excluded_keys - and key not in completion_kwargs - and value is not None - ): + setattr(value, "stream_options", completion_kwargs.get("stream_options")) + if key not in excluded_keys and key not in completion_kwargs and value is not None: completion_kwargs[key] = value # Normalize reasoning_effort based on model capabilities # (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported) # Must run BEFORE _route_openai_thinking, which prepends "responses/" # to the model name and would break get_model_info() lookups. - LiteLLMMessagesToCompletionTransformationHandler._normalize_reasoning_effort( - completion_kwargs - ) + LiteLLMMessagesToCompletionTransformationHandler._normalize_reasoning_effort(completion_kwargs) LiteLLMMessagesToCompletionTransformationHandler._route_openai_thinking_to_responses_api_if_needed( completion_kwargs, @@ -581,9 +545,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: proxy_litellm_metadata = _extract_proxy_litellm_metadata(kwargs) user_api_key_auth = ( - proxy_litellm_metadata.get("user_api_key_auth") - if proxy_litellm_metadata is not None - else None + proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None ) polyfill_result = await _prepare_context_managed_request( @@ -598,12 +560,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, ) - effective_messages = ( - polyfill_result.messages if polyfill_result is not None else messages - ) - effective_system = ( - polyfill_result.system if polyfill_result is not None else system - ) + effective_messages = polyfill_result.messages if polyfill_result is not None else messages + effective_system = polyfill_result.system if polyfill_result is not None else system ( completion_kwargs, @@ -629,14 +587,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_response = await litellm.acompletion(**completion_kwargs) if stream: - transformed_stream = ( - ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( - completion_response, - model=model, - tool_name_mapping=tool_name_mapping, - polyfill_result=polyfill_result, - is_async=True, - ) + transformed_stream = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( + completion_response, + model=model, + tool_name_mapping=tool_name_mapping, + polyfill_result=polyfill_result, + is_async=True, ) if transformed_stream is not None: return transformed_stream @@ -730,9 +686,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: else: proxy_litellm_metadata = _extract_proxy_litellm_metadata(kwargs) user_api_key_auth = ( - proxy_litellm_metadata.get("user_api_key_auth") - if proxy_litellm_metadata is not None - else None + proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None ) polyfill_result = run_async_function( _prepare_context_managed_request, @@ -747,12 +701,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, ) - effective_messages = ( - polyfill_result.messages if polyfill_result is not None else messages - ) - effective_system = ( - polyfill_result.system if polyfill_result is not None else system - ) + effective_messages = polyfill_result.messages if polyfill_result is not None else messages + effective_system = polyfill_result.system if polyfill_result is not None else system ( completion_kwargs, @@ -778,14 +728,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_response = litellm.completion(**completion_kwargs) if stream: - transformed_stream = ( - ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( - completion_response, - model=model, - tool_name_mapping=tool_name_mapping, - polyfill_result=polyfill_result, - is_async=False, - ) + transformed_stream = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( + completion_response, + model=model, + tool_name_mapping=tool_name_mapping, + polyfill_result=polyfill_result, + is_async=False, ) if transformed_stream is not None: return transformed_stream diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index a92834b5d71..44c367ee805 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -212,14 +212,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): ) merged_chunk["usage"] = usage_dict if self.applied_edits and "context_management" not in merged_chunk: - merged_chunk["context_management"] = ContextManagementResponse( - applied_edits=list(self.applied_edits) - ) + merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits)) return self._augment_message_delta_usage(merged_chunk) - def _ensure_context_management_attached( - self, message_delta_chunk: Dict[str, Any] - ) -> Dict[str, Any]: + def _ensure_context_management_attached(self, message_delta_chunk: Dict[str, Any]) -> Dict[str, Any]: """Attach ``context_management`` to a ``message_delta`` chunk if ``self.applied_edits`` is non-empty and the chunk does not already carry it. Returns the (possibly new) chunk dict. @@ -231,23 +227,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if not self.applied_edits or "context_management" in message_delta_chunk: return message_delta_chunk augmented = message_delta_chunk.copy() - augmented["context_management"] = ContextManagementResponse( - applied_edits=list(self.applied_edits) - ) + augmented["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits)) return augmented - def _augment_message_delta_usage( - self, message_delta_chunk: Dict[str, Any] - ) -> Dict[str, Any]: + def _augment_message_delta_usage(self, message_delta_chunk: Dict[str, Any]) -> Dict[str, Any]: """Attach polyfill compaction iteration usage to the final message_delta. Also defensively re-attaches ``context_management`` so the direct held-chunk flush path stays in sync with the merge path's guarantee when ``self.applied_edits`` is non-empty. """ - message_delta_chunk = self._ensure_context_management_attached( - message_delta_chunk - ) + message_delta_chunk = self._ensure_context_management_attached(message_delta_chunk) if self.iterations_usage is None: return message_delta_chunk usage = message_delta_chunk.get("usage") @@ -377,10 +367,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): ) return self.chunk_queue.popleft() - if ( - self.sent_compaction_block is False - and self.compaction_block is not None - ): + if self.sent_compaction_block is False and self.compaction_block is not None: compaction_event = self._next_compaction_event() if compaction_event is not None: return compaction_event @@ -413,18 +400,13 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # skip the applied_edits attachment in that case to avoid # allocating a throwaway ``MessageBlockDelta``. will_merge_into_held = ( - self.holding_stop_reason_chunk is not None - and getattr(chunk, "usage", None) is not None + self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None ) is_final_chunk = chunk.choices[0].finish_reason is not None processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic( response=chunk, current_content_block_index=self.current_content_block_index, - applied_edits=( - self.applied_edits - if is_final_chunk and not will_merge_into_held - else None - ), + applied_edits=(self.applied_edits if is_final_chunk and not will_merge_into_held else None), ) # Check if this is a usage chunk and we have a held stop_reason chunk @@ -482,10 +464,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): self.sent_content_block_finish = False return self.chunk_queue.popleft() - if ( - processed_chunk["type"] == "message_delta" - and self.sent_content_block_finish is False - ): + if processed_chunk["type"] == "message_delta" and self.sent_content_block_finish is False: # Queue both the content_block_stop and the message_delta self.chunk_queue.append( { @@ -497,25 +476,19 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if processed_chunk.get("delta", {}).get("stop_reason") is not None: self.holding_stop_reason_chunk = processed_chunk else: - processed_chunk = self._augment_message_delta_usage( - processed_chunk - ) + processed_chunk = self._augment_message_delta_usage(processed_chunk) self.chunk_queue.append(processed_chunk) return self.chunk_queue.popleft() elif self.holding_chunk is not None: self.chunk_queue.append(self.holding_chunk) if processed_chunk.get("type") == "message_delta": - processed_chunk = self._augment_message_delta_usage( - processed_chunk - ) + processed_chunk = self._augment_message_delta_usage(processed_chunk) self.chunk_queue.append(processed_chunk) self.holding_chunk = None return self.chunk_queue.popleft() else: if processed_chunk.get("type") == "message_delta": - processed_chunk = self._augment_message_delta_usage( - processed_chunk - ) + processed_chunk = self._augment_message_delta_usage(processed_chunk) self.chunk_queue.append(processed_chunk) return self.chunk_queue.popleft() @@ -545,11 +518,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): } ) self.sent_content_block_finish = True - self.chunk_queue.append( - self._augment_message_delta_usage( - self.holding_stop_reason_chunk - ) - ) + self.chunk_queue.append(self._augment_message_delta_usage(self.holding_stop_reason_chunk)) self.holding_stop_reason_chunk = None else: self.holding_chunk = None @@ -572,11 +541,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if self.holding_stop_reason_chunk is not None: if not self.sent_content_block_finish: self.sent_content_block_finish = True - self.chunk_queue.append( - self._augment_message_delta_usage( - self.holding_stop_reason_chunk - ) - ) + self.chunk_queue.append(self._augment_message_delta_usage(self.holding_stop_reason_chunk)) self.holding_stop_reason_chunk = None return { "type": "content_block_stop", @@ -590,9 +555,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): return {"type": "message_stop"} raise StopIteration except Exception as e: - verbose_logger.error( - "Anthropic Adapter - {}\n{}".format(e, traceback.format_exc()) - ) + verbose_logger.error("Anthropic Adapter - {}\n{}".format(e, traceback.format_exc())) raise StopIteration async def __anext__(self): @@ -623,10 +586,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): ) return self.chunk_queue.popleft() - if ( - self.sent_compaction_block is False - and self.compaction_block is not None - ): + if self.sent_compaction_block is False and self.compaction_block is not None: compaction_event = self._next_compaction_event() if compaction_event is not None: return compaction_event @@ -660,18 +620,13 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # skip the applied_edits attachment in that case to avoid # allocating a throwaway ``MessageBlockDelta``. will_merge_into_held = ( - self.holding_stop_reason_chunk is not None - and getattr(chunk, "usage", None) is not None + self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None ) is_final_chunk = chunk.choices[0].finish_reason is not None processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic( response=chunk, current_content_block_index=self.current_content_block_index, - applied_edits=( - self.applied_edits - if is_final_chunk and not will_merge_into_held - else None - ), + applied_edits=(self.applied_edits if is_final_chunk and not will_merge_into_held else None), ) # Check if this is a usage chunk and we have a held stop_reason chunk @@ -722,10 +677,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): self.sent_content_block_finish = False return self.chunk_queue.popleft() - if ( - processed_chunk["type"] == "message_delta" - and self.sent_content_block_finish is False - ): + if processed_chunk["type"] == "message_delta" and self.sent_content_block_finish is False: # Queue both the content_block_stop and the holding chunk self.chunk_queue.append( { @@ -734,32 +686,23 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): } ) self.sent_content_block_finish = True - if ( - processed_chunk.get("delta", {}).get("stop_reason") - is not None - ): + if processed_chunk.get("delta", {}).get("stop_reason") is not None: self.holding_stop_reason_chunk = processed_chunk else: - processed_chunk = self._augment_message_delta_usage( - processed_chunk - ) + processed_chunk = self._augment_message_delta_usage(processed_chunk) self.chunk_queue.append(processed_chunk) return self.chunk_queue.popleft() elif self.holding_chunk is not None: # Queue both chunks self.chunk_queue.append(self.holding_chunk) if processed_chunk.get("type") == "message_delta": - processed_chunk = self._augment_message_delta_usage( - processed_chunk - ) + processed_chunk = self._augment_message_delta_usage(processed_chunk) self.chunk_queue.append(processed_chunk) self.holding_chunk = None return self.chunk_queue.popleft() else: if processed_chunk.get("type") == "message_delta": - processed_chunk = self._augment_message_delta_usage( - processed_chunk - ) + processed_chunk = self._augment_message_delta_usage(processed_chunk) self.chunk_queue.append(processed_chunk) return self.chunk_queue.popleft() @@ -789,11 +732,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): } ) self.sent_content_block_finish = True - self.chunk_queue.append( - self._augment_message_delta_usage( - self.holding_stop_reason_chunk - ) - ) + self.chunk_queue.append(self._augment_message_delta_usage(self.holding_stop_reason_chunk)) self.holding_stop_reason_chunk = None else: self.holding_chunk = None @@ -821,11 +760,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if self.holding_stop_reason_chunk is not None: if not self.sent_content_block_finish: self.sent_content_block_finish = True - self.chunk_queue.append( - self._augment_message_delta_usage( - self.holding_stop_reason_chunk - ) - ) + self.chunk_queue.append(self._augment_message_delta_usage(self.holding_stop_reason_chunk)) self.holding_stop_reason_chunk = None return { "type": "content_block_stop", @@ -939,9 +874,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if tool_block.get("name"): truncated_name = tool_block["name"] - original_name = self.tool_name_mapping.get( - truncated_name, truncated_name - ) + original_name = self.tool_name_mapping.get(truncated_name, truncated_name) tool_block["name"] = original_name if block_type != self.current_content_block_type: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 2f0cae2a5b7..02625605f37 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -143,9 +143,7 @@ class AnthropicAdapter: def __init__(self) -> None: pass - def translate_completion_input_params( - self, kwargs - ) -> Optional[ChatCompletionRequest]: + def translate_completion_input_params(self, kwargs) -> Optional[ChatCompletionRequest]: """ Translate Anthropic request params to OpenAI format. @@ -178,27 +176,19 @@ class AnthropicAdapter: model = kwargs.pop("model") messages = kwargs.pop("messages") if not model: - raise ValueError( - "Bad Request: model is required for Anthropic Messages Request" - ) + raise ValueError("Bad Request: model is required for Anthropic Messages Request") if not messages: - raise ValueError( - "Bad Request: messages is required for Anthropic Messages Request" - ) + raise ValueError("Bad Request: messages is required for Anthropic Messages Request") ######################################################### # Created Typed Request Body ######################################################### - request_body = AnthropicMessagesRequest( - model=model, messages=messages, **kwargs - ) + request_body = AnthropicMessagesRequest(model=model, messages=messages, **kwargs) ( translated_body, tool_name_mapping, - ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( - anthropic_message_request=request_body - ) + ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(anthropic_message_request=request_body) return translated_body, tool_name_mapping @@ -247,15 +237,9 @@ class AnthropicAdapter: the sync handler) don't get back an async iterator they can't iterate without an event loop. """ - applied_edits = ( - polyfill_result.applied_edits_for_response() if polyfill_result else None - ) - compaction_block = ( - polyfill_result.compaction_block if polyfill_result is not None else None - ) - iterations_usage = ( - polyfill_result.iterations_usage if polyfill_result is not None else None - ) + applied_edits = polyfill_result.applied_edits_for_response() if polyfill_result else None + compaction_block = polyfill_result.compaction_block if polyfill_result is not None else None + iterations_usage = polyfill_result.iterations_usage if polyfill_result is not None else None anthropic_wrapper = AnthropicStreamWrapper( completion_stream=completion_stream, model=model, @@ -283,26 +267,16 @@ class LiteLLMAnthropicMessagesAdapter: """ signature = None - if ( - hasattr(tool_call, "provider_specific_fields") - and tool_call.provider_specific_fields - ): + if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields: if "thought_signature" in tool_call.provider_specific_fields: signature = tool_call.provider_specific_fields["thought_signature"] - elif ( - hasattr(tool_call.function, "provider_specific_fields") - and tool_call.function.provider_specific_fields - ): + elif hasattr(tool_call.function, "provider_specific_fields") and tool_call.function.provider_specific_fields: if "thought_signature" in tool_call.function.provider_specific_fields: - signature = tool_call.function.provider_specific_fields[ - "thought_signature" - ] + signature = tool_call.function.provider_specific_fields["thought_signature"] return signature - def _extract_signature_from_tool_use_content( - self, content: Dict[str, Any] - ) -> Optional[str]: + def _extract_signature_from_tool_use_content(self, content: Dict[str, Any]) -> Optional[str]: """ Extract signature from a tool_use content block's provider_specific_fields. """ @@ -332,18 +306,9 @@ class LiteLLMAnthropicMessagesAdapter: """ # TypedDict objects are dicts at runtime, so .get() works cache_control = ( - source.get("cache_control") - if isinstance(source, dict) - else getattr(source, "cache_control", None) + source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None) ) - if ( - cache_control - and model - and ( - self.is_anthropic_claude_model(model) - or self.is_bedrock_arn_model(model) - ) - ): + if cache_control and model and (self.is_anthropic_claude_model(model) or self.is_bedrock_arn_model(model)): # TypedDict objects support dict operations at runtime # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432) if isinstance(target, dict): @@ -383,9 +348,7 @@ class LiteLLMAnthropicMessagesAdapter: """ tool_type = tool.get("type", "") tool_name = tool.get("name", "") - return ( - isinstance(tool_type, str) and tool_type.startswith("web_search") - ) or tool_name == "web_search" + return (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search" def translate_anthropic_messages_to_openai( self, @@ -401,66 +364,38 @@ class LiteLLMAnthropicMessagesAdapter: for m in messages: user_message: Optional[ChatCompletionUserMessage] = None tool_message_list: List[ChatCompletionToolMessage] = [] - new_user_content_list: List[ - Union[ChatCompletionTextObject, ChatCompletionImageObject] - ] = [] + new_user_content_list: List[Union[ChatCompletionTextObject, ChatCompletionImageObject]] = [] ## USER MESSAGE ## if m["role"] == "user": ## translate user message message_content = m.get("content") if message_content and isinstance(message_content, str): - user_message = ChatCompletionUserMessage( - role="user", content=message_content - ) + user_message = ChatCompletionUserMessage(role="user", content=message_content) elif message_content and isinstance(message_content, list): for content in message_content: if content.get("type") == "text": - text_obj = ChatCompletionTextObject( - type="text", text=content.get("text", "") - ) - self._add_cache_control_if_applicable( - content, text_obj, model - ) + text_obj = ChatCompletionTextObject(type="text", text=content.get("text", "")) + self._add_cache_control_if_applicable(content, text_obj, model) new_user_content_list.append(text_obj) # type: ignore elif content.get("type") == "image": # Convert Anthropic image format to OpenAI format source = content.get("source", {}) - openai_image_url = ( - self._translate_anthropic_image_to_openai( - cast(dict, source) - ) - ) + openai_image_url = self._translate_anthropic_image_to_openai(cast(dict, source)) if openai_image_url: - image_url_obj = ChatCompletionImageUrlObject( - url=openai_image_url - ) - image_obj = ChatCompletionImageObject( - type="image_url", image_url=image_url_obj - ) - self._add_cache_control_if_applicable( - content, image_obj, model - ) + image_url_obj = ChatCompletionImageUrlObject(url=openai_image_url) + image_obj = ChatCompletionImageObject(type="image_url", image_url=image_url_obj) + self._add_cache_control_if_applicable(content, image_obj, model) new_user_content_list.append(image_obj) # type: ignore elif content.get("type") == "document": # Convert Anthropic document format (PDF, etc.) to OpenAI format source = content.get("source", {}) - openai_image_url = ( - self._translate_anthropic_image_to_openai( - cast(dict, source) - ) - ) + openai_image_url = self._translate_anthropic_image_to_openai(cast(dict, source)) if openai_image_url: - image_url_obj = ChatCompletionImageUrlObject( - url=openai_image_url - ) - doc_obj = ChatCompletionImageObject( - type="image_url", image_url=image_url_obj - ) - self._add_cache_control_if_applicable( - content, doc_obj, model - ) + image_url_obj = ChatCompletionImageUrlObject(url=openai_image_url) + doc_obj = ChatCompletionImageObject(type="image_url", image_url=image_url_obj) + self._add_cache_control_if_applicable(content, doc_obj, model) new_user_content_list.append(doc_obj) # type: ignore elif content.get("type") == "tool_result": if "content" not in content: @@ -469,9 +404,7 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content="", ) - self._add_cache_control_if_applicable( - content, tool_result, model - ) + self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) # type: ignore[arg-type] elif isinstance(content.get("content"), str): tool_result = ChatCompletionToolMessage( @@ -479,9 +412,7 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content=str(content.get("content", "")), ) - self._add_cache_control_if_applicable( - content, tool_result, model - ) + self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) # type: ignore[arg-type] elif isinstance(content.get("content"), list): # Combine all content items into a single tool message @@ -498,41 +429,28 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content=c, ) - self._add_cache_control_if_applicable( - content, tool_result, model - ) + self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) # type: ignore[arg-type] elif isinstance(c, dict): if c.get("type") == "text": tool_result = ChatCompletionToolMessage( role="tool", - tool_call_id=content.get( - "tool_use_id", "" - ), + tool_call_id=content.get("tool_use_id", ""), content=c.get("text", ""), ) - self._add_cache_control_if_applicable( - content, tool_result, model - ) + self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) # type: ignore[arg-type] elif c.get("type") == "image": source = c.get("source", {}) openai_image_url = ( - self._translate_anthropic_image_to_openai( - cast(dict, source) - ) - or "" + self._translate_anthropic_image_to_openai(cast(dict, source)) or "" ) tool_result = ChatCompletionToolMessage( role="tool", - tool_call_id=content.get( - "tool_use_id", "" - ), + tool_call_id=content.get("tool_use_id", ""), content=openai_image_url, ) - self._add_cache_control_if_applicable( - content, tool_result, model - ) + self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) # type: ignore[arg-type] else: # For multiple content items, combine into a single tool message @@ -545,11 +463,7 @@ class LiteLLMAnthropicMessagesAdapter: ] = [] for c in content_items: if isinstance(c, str): - combined_content_parts.append( - ChatCompletionTextObject( - type="text", text=c - ) - ) + combined_content_parts.append(ChatCompletionTextObject(type="text", text=c)) elif isinstance(c, dict): if c.get("type") == "text": combined_content_parts.append( @@ -561,10 +475,7 @@ class LiteLLMAnthropicMessagesAdapter: elif c.get("type") == "image": source = c.get("source", {}) openai_image_url = ( - self._translate_anthropic_image_to_openai( - cast(dict, source) - ) - or "" + self._translate_anthropic_image_to_openai(cast(dict, source)) or "" ) if openai_image_url: combined_content_parts.append( @@ -582,9 +493,7 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content=combined_content_parts, # type: ignore ) - self._add_cache_control_if_applicable( - content, tool_result, model - ) + self._add_cache_control_if_applicable(content, tool_result, model) tool_message_list.append(tool_result) # type: ignore[arg-type] if len(tool_message_list) > 0: @@ -598,14 +507,10 @@ class LiteLLMAnthropicMessagesAdapter: ## ASSISTANT MESSAGE ## assistant_message_str: Optional[str] = None - assistant_content_list: List[ - Dict[str, Any] - ] = [] # For content blocks with cache_control + assistant_content_list: List[Dict[str, Any]] = [] # For content blocks with cache_control has_cache_control_in_text = False tool_calls: List[ChatCompletionAssistantToolCall] = [] - thinking_blocks: List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] = [] + thinking_blocks: List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] = [] if m["role"] == "assistant": if isinstance(m.get("content"), str): assistant_message_str = str(m.get("content", "")) @@ -619,9 +524,7 @@ class LiteLLMAnthropicMessagesAdapter: "type": "text", "text": content.get("text", ""), } - self._add_cache_control_if_applicable( - content, text_block, model - ) + self._add_cache_control_if_applicable(content, text_block, model) if "cache_control" in text_block: has_cache_control_in_text = True assistant_content_list.append(text_block) @@ -632,32 +535,21 @@ class LiteLLMAnthropicMessagesAdapter: "name": tool_name, "arguments": json.dumps(content.get("input", {})), } - signature = ( - self._extract_signature_from_tool_use_content( - cast(Dict[str, Any], content) - ) - ) + signature = self._extract_signature_from_tool_use_content(cast(Dict[str, Any], content)) if signature: provider_specific_fields: Dict[str, Any] = ( - function_chunk.get("provider_specific_fields") - or {} - ) - provider_specific_fields["thought_signature"] = ( - signature - ) - function_chunk["provider_specific_fields"] = ( - provider_specific_fields + function_chunk.get("provider_specific_fields") or {} ) + provider_specific_fields["thought_signature"] = signature + function_chunk["provider_specific_fields"] = provider_specific_fields tool_call = ChatCompletionAssistantToolCall( id=content.get("id", ""), type="function", function=function_chunk, ) - self._add_cache_control_if_applicable( - content, tool_call, model - ) + self._add_cache_control_if_applicable(content, tool_call, model) tool_calls.append(tool_call) elif content.get("type") == "thinking": thinking_block = ChatCompletionThinkingBlock( @@ -668,12 +560,10 @@ class LiteLLMAnthropicMessagesAdapter: ) thinking_blocks.append(thinking_block) elif content.get("type") == "redacted_thinking": - redacted_thinking_block = ( - ChatCompletionRedactedThinkingBlock( - type="redacted_thinking", - data=content.get("data") or "", - cache_control=content.get("cache_control", {}), - ) + redacted_thinking_block = ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", + data=content.get("data") or "", + cache_control=content.get("cache_control", {}), ) thinking_blocks.append(redacted_thinking_block) @@ -688,18 +578,14 @@ class LiteLLMAnthropicMessagesAdapter: assistant_content: Any = assistant_content_list elif len(assistant_content_list) > 0 and not has_cache_control_in_text: # Concatenate text blocks into string when no cache_control - assistant_content = "".join( - block.get("text", "") for block in assistant_content_list - ) + assistant_content = "".join(block.get("text", "") for block in assistant_content_list) else: assistant_content = assistant_message_str assistant_message = ChatCompletionAssistantMessage( role="assistant", content=assistant_content, - thinking_blocks=( - thinking_blocks if len(thinking_blocks) > 0 else None - ), + thinking_blocks=(thinking_blocks if len(thinking_blocks) > 0 else None), ) if len(tool_calls) > 0: assistant_message["tool_calls"] = tool_calls # type: ignore @@ -730,9 +616,7 @@ class LiteLLMAnthropicMessagesAdapter: if thinking_type == "disabled": return None elif thinking_type == "enabled": - return reasoning_effort_from_thinking_budget( - thinking.get("budget_tokens", 0) - ) + return reasoning_effort_from_thinking_budget(thinking.get("budget_tokens", 0)) elif thinking_type == "adaptive": # Adaptive thinking: effort is controlled by output_config.effort, # not budget_tokens. Return a default; caller should override with @@ -796,9 +680,7 @@ class LiteLLMAnthropicMessagesAdapter: thinking ) if reasoning_effort: - summary = ( - thinking.get("summary") if isinstance(thinking, dict) else None - ) + summary = thinking.get("summary") if isinstance(thinking, dict) else None auto_summary = is_reasoning_auto_summary_enabled() if summary: return { @@ -828,18 +710,12 @@ class LiteLLMAnthropicMessagesAdapter: # Truncate tool name if it exceeds OpenAI's 64-char limit original_name = tool_choice.get("name", "") truncated_name = truncate_tool_name(original_name) - tc_function_param = ChatCompletionToolChoiceFunctionParam( - name=truncated_name - ) - return ChatCompletionToolChoiceObjectParam( - type="function", function=tc_function_param - ) + tc_function_param = ChatCompletionToolChoiceFunctionParam(name=truncated_name) + return ChatCompletionToolChoiceObjectParam(type="function", function=tc_function_param) elif tool_choice["type"] == "none": return "none" else: - raise ValueError( - "Incompatible tool choice param submitted - {}".format(tool_choice) - ) + raise ValueError("Incompatible tool choice param submitted - {}".format(tool_choice)) def translate_anthropic_tools_to_openai( self, tools: List[AllAnthropicToolsValues], model: Optional[str] = None @@ -875,9 +751,7 @@ class LiteLLMAnthropicMessagesAdapter: continue raw_name = tool.get("name") - if raw_name is None or ( - isinstance(raw_name, str) and not str(raw_name).strip() - ): + if raw_name is None or (isinstance(raw_name, str) and not str(raw_name).strip()): original_name = f"litellm_unnamed_tool_{idx}" else: original_name = str(raw_name) @@ -898,17 +772,13 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - tool_param = ChatCompletionToolParam( - type="function", function=function_chunk - ) + tool_param = ChatCompletionToolParam(type="function", function=function_chunk) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) # type: ignore[arg-type] return new_tools, tool_name_mapping # type: ignore[return-value] - def translate_anthropic_output_format_to_openai( - self, output_format: Any - ) -> Optional[Dict[str, Any]]: + def translate_anthropic_output_format_to_openai(self, output_format: Any) -> Optional[Dict[str, Any]]: """ Translate Anthropic's output_format to OpenAI's response_format. @@ -967,25 +837,19 @@ class LiteLLMAnthropicMessagesAdapter: # Handle array items if "items" in schema: - LiteLLMAnthropicMessagesAdapter._add_additional_properties_false( - schema["items"] - ) + LiteLLMAnthropicMessagesAdapter._add_additional_properties_false(schema["items"]) # Handle anyOf/oneOf/allOf for key in ("anyOf", "oneOf", "allOf"): if key in schema: for sub_schema in schema[key]: - LiteLLMAnthropicMessagesAdapter._add_additional_properties_false( - sub_schema - ) + LiteLLMAnthropicMessagesAdapter._add_additional_properties_false(sub_schema) # Handle $defs / definitions for key in ("$defs", "definitions"): if key in schema: for def_schema in schema[key].values(): - LiteLLMAnthropicMessagesAdapter._add_additional_properties_false( - def_schema - ) + LiteLLMAnthropicMessagesAdapter._add_additional_properties_false(def_schema) def _add_system_message_to_messages( self, @@ -1019,9 +883,7 @@ class LiteLLMAnthropicMessagesAdapter: if openai_system_content: new_messages.insert( 0, - ChatCompletionSystemMessage( - role="system", content=openai_system_content - ), # type: ignore + ChatCompletionSystemMessage(role="system", content=openai_system_content), # type: ignore ) def _translate_metadata_to_openai( @@ -1107,9 +969,7 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs["thinking"] = thinking # type: ignore return - reasoning_effort = self.translate_anthropic_thinking_to_reasoning_effort( - cast(Dict[str, Any], thinking) - ) + reasoning_effort = self.translate_anthropic_thinking_to_reasoning_effort(cast(Dict[str, Any], thinking)) if not reasoning_effort: return @@ -1164,9 +1024,7 @@ class LiteLLMAnthropicMessagesAdapter: output_format = output_config.get("format") if not output_format: return - response_format = self.translate_anthropic_output_format_to_openai( - output_format=output_format - ) + response_format = self.translate_anthropic_output_format_to_openai(output_format=output_format) if response_format: new_kwargs["response_format"] = response_format @@ -1197,11 +1055,7 @@ class LiteLLMAnthropicMessagesAdapter: tool_name_mapping: Dict[str, str] = {} ## CONVERT ANTHROPIC MESSAGES TO OPENAI - messages_list: List[ - Union[ - AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam - ] - ] = cast( + messages_list: List[Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam]] = cast( List[ Union[ AnthropicMessagesUserMessageParam, @@ -1288,10 +1142,7 @@ class LiteLLMAnthropicMessagesAdapter: new_content: List[Dict[str, Any]] = [] for choice in choices: # Handle thinking blocks first - if ( - hasattr(choice.message, "thinking_blocks") - and choice.message.thinking_blocks - ): + if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks: for thinking_block in choice.message.thinking_blocks: if thinking_block.get("type") == "thinking": thinking_value = thinking_block.get("thinking", "") @@ -1299,16 +1150,8 @@ class LiteLLMAnthropicMessagesAdapter: new_content.append( AnthropicResponseContentBlockThinking( type="thinking", - thinking=( - str(thinking_value) - if thinking_value is not None - else "" - ), - signature=( - str(signature_value) - if signature_value is not None - else None - ), + thinking=(str(thinking_value) if thinking_value is not None else ""), + signature=(str(signature_value) if signature_value is not None else None), ).model_dump() ) elif thinking_block.get("type") == "redacted_thinking": @@ -1320,10 +1163,7 @@ class LiteLLMAnthropicMessagesAdapter: ).model_dump() ) # Handle reasoning_content when thinking_blocks is not present - elif ( - hasattr(choice.message, "reasoning_content") - and choice.message.reasoning_content - ): + elif hasattr(choice.message, "reasoning_content") and choice.message.reasoning_content: new_content.append( AnthropicResponseContentBlockThinking( type="thinking", @@ -1335,15 +1175,10 @@ class LiteLLMAnthropicMessagesAdapter: # Handle text content if choice.message.content is not None: new_content.append( - AnthropicResponseContentBlockText( - type="text", text=choice.message.content - ).model_dump() + AnthropicResponseContentBlockText(type="text", text=choice.message.content).model_dump() ) # Handle tool calls (in parallel to text content) - if ( - choice.message.tool_calls is not None - and len(choice.message.tool_calls) > 0 - ): + if choice.message.tool_calls is not None and len(choice.message.tool_calls) > 0: for tool_call in choice.message.tool_calls: # Extract signature from provider_specific_fields only signature = self._extract_signature_from_tool_call(tool_call) @@ -1355,9 +1190,7 @@ class LiteLLMAnthropicMessagesAdapter: # Restore original tool name if it was truncated truncated_name = tool_call.function.name or "" original_name = ( - tool_name_mapping.get(truncated_name, truncated_name) - if tool_name_mapping - else truncated_name + tool_name_mapping.get(truncated_name, truncated_name) if tool_name_mapping else truncated_name ) # Strip Gemini thought-signature suffix and normalize id chars @@ -1375,16 +1208,12 @@ class LiteLLMAnthropicMessagesAdapter: ) # Add provider_specific_fields if signature is present if provider_specific_fields: - tool_use_block.provider_specific_fields = ( - provider_specific_fields - ) + tool_use_block.provider_specific_fields = provider_specific_fields new_content.append(tool_use_block.model_dump()) return new_content - def _translate_openai_finish_reason_to_anthropic( - self, openai_finish_reason: str - ) -> AnthropicFinishReason: + def _translate_openai_finish_reason_to_anthropic(self, openai_finish_reason: str) -> AnthropicFinishReason: if openai_finish_reason == "stop": return "end_turn" elif openai_finish_reason == "length": @@ -1404,9 +1233,7 @@ class LiteLLMAnthropicMessagesAdapter: return 0 @classmethod - def _first_positive_usage_value( - cls, usage: Usage, field_names: tuple[str, ...] - ) -> int: + def _first_positive_usage_value(cls, usage: Usage, field_names: tuple[str, ...]) -> int: for field_name in field_names: value = cls._positive_int(getattr(usage, field_name, None)) if value > 0: @@ -1414,9 +1241,7 @@ class LiteLLMAnthropicMessagesAdapter: return 0 @classmethod - def _first_positive_prompt_tokens_detail_value( - cls, usage: Usage, field_names: tuple[str, ...] - ) -> int: + def _first_positive_prompt_tokens_detail_value(cls, usage: Usage, field_names: tuple[str, ...]) -> int: prompt_tokens_details = getattr(usage, "prompt_tokens_details", None) if prompt_tokens_details is None: return 0 @@ -1425,18 +1250,14 @@ class LiteLLMAnthropicMessagesAdapter: if isinstance(prompt_tokens_details, dict): value = cls._positive_int(prompt_tokens_details.get(field_name)) else: - value = cls._positive_int( - getattr(prompt_tokens_details, field_name, None) - ) + value = cls._positive_int(getattr(prompt_tokens_details, field_name, None)) if value > 0: return value return 0 @classmethod def _get_cache_read_input_tokens(cls, usage: Usage) -> int: - explicit_value = cls._first_positive_usage_value( - usage, ("cache_read_input_tokens", "_cache_read_input_tokens") - ) + explicit_value = cls._first_positive_usage_value(usage, ("cache_read_input_tokens", "_cache_read_input_tokens")) if explicit_value > 0: return explicit_value return cls._first_positive_prompt_tokens_detail_value(usage, ("cached_tokens",)) @@ -1448,20 +1269,14 @@ class LiteLLMAnthropicMessagesAdapter: ) if explicit_value > 0: return explicit_value - return cls._first_positive_prompt_tokens_detail_value( - usage, ("cache_creation_tokens", "cache_write_tokens") - ) + return cls._first_positive_prompt_tokens_detail_value(usage, ("cache_creation_tokens", "cache_write_tokens")) @classmethod - def _translate_openai_usage_to_anthropic_usage_delta( - cls, usage: Usage - ) -> UsageDelta: + def _translate_openai_usage_to_anthropic_usage_delta(cls, usage: Usage) -> UsageDelta: cache_read_input_tokens = cls._get_cache_read_input_tokens(usage) cache_creation_input_tokens = cls._get_cache_creation_input_tokens(usage) input_tokens = max( - (usage.prompt_tokens or 0) - - cache_read_input_tokens - - cache_creation_input_tokens, + (usage.prompt_tokens or 0) - cache_read_input_tokens - cache_creation_input_tokens, 0, ) @@ -1521,9 +1336,7 @@ class LiteLLMAnthropicMessagesAdapter: "input_tokens": anthropic_usage["input_tokens"], "output_tokens": usage.completion_tokens or 0, } - anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [ - message_iteration - ] # type: ignore[typeddict-unknown-key] + anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [message_iteration] # type: ignore[typeddict-unknown-key] translated_obj = AnthropicMessagesResponse( id=response.id, @@ -1536,13 +1349,9 @@ class LiteLLMAnthropicMessagesAdapter: stop_reason=anthropic_finish_reason, ) - applied_edits = ( - polyfill_result.applied_edits_for_response() if polyfill_result else None - ) + applied_edits = polyfill_result.applied_edits_for_response() if polyfill_result else None if applied_edits: - translated_obj["context_management"] = ContextManagementResponse( - applied_edits=list(applied_edits) - ) + translated_obj["context_management"] = ContextManagementResponse(applied_edits=list(applied_edits)) return translated_obj @@ -1580,9 +1389,7 @@ class LiteLLMAnthropicMessagesAdapter: return "tool_use", cast("ContentBlockContentBlockDict", tool_block) elif choice.delta.content is not None and len(choice.delta.content) > 0: return "text", TextBlock(type="text", text="") - elif isinstance(choice, StreamingChoices) and hasattr( - choice.delta, "thinking_blocks" - ): + elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"): thinking_blocks = choice.delta.thinking_blocks or [] if len(thinking_blocks) > 0: thinking_block = thinking_blocks[0] @@ -1606,12 +1413,8 @@ class LiteLLMAnthropicMessagesAdapter: # ``Delta`` deletes the ``thinking_blocks`` attribute when unset, so the # branch above is skipped entirely; open a ``thinking`` block here so the # matching ``thinking_delta`` stream is not emitted into a text block. - elif isinstance(choice, StreamingChoices) and getattr( - choice.delta, "reasoning_content", None - ): - return "thinking", ChatCompletionThinkingBlock( - type="thinking", thinking="", signature="" - ) + elif isinstance(choice, StreamingChoices) and getattr(choice.delta, "reasoning_content", None): + return "thinking", ChatCompletionThinkingBlock(type="thinking", thinking="", signature="") return "text", TextBlock(type="text", text="") @@ -1636,14 +1439,9 @@ class LiteLLMAnthropicMessagesAdapter: if choice.delta.tool_calls: partial_json = "" for tool in choice.delta.tool_calls: - if ( - tool.function is not None - and tool.function.arguments is not None - ): + if tool.function is not None and tool.function.arguments is not None: partial_json = (partial_json or "") + tool.function.arguments - elif isinstance(choice, StreamingChoices) and hasattr( - choice.delta, "thinking_blocks" - ): + elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"): thinking_blocks = choice.delta.thinking_blocks or [] if len(thinking_blocks) > 0: for thinking_block in thinking_blocks: @@ -1658,25 +1456,17 @@ class LiteLLMAnthropicMessagesAdapter: reasoning_signature += signature # Handle reasoning_content when thinking_blocks is not present # This handles providers like OpenRouter that return reasoning_content - elif isinstance(choice, StreamingChoices) and hasattr( - choice.delta, "reasoning_content" - ): + elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "reasoning_content"): if choice.delta.reasoning_content is not None: reasoning_content += choice.delta.reasoning_content if reasoning_content and reasoning_signature: - raise ValueError( - "Both `reasoning` and `signature` in a single streaming chunk isn't supported." - ) + raise ValueError("Both `reasoning` and `signature` in a single streaming chunk isn't supported.") if partial_json is not None: - return "input_json_delta", ContentJsonBlockDelta( - type="input_json_delta", partial_json=partial_json - ) + return "input_json_delta", ContentJsonBlockDelta(type="input_json_delta", partial_json=partial_json) elif reasoning_content: - return "thinking_delta", ContentThinkingBlockDelta( - type="thinking_delta", thinking=reasoning_content - ) + return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content) elif reasoning_signature: return "signature_delta", ContentThinkingSignatureBlockDelta( type="signature_delta", signature=reasoning_signature @@ -1693,23 +1483,16 @@ class LiteLLMAnthropicMessagesAdapter: ## base case - final chunk w/ finish reason if response.choices[0].finish_reason is not None: delta = MessageDelta( - stop_reason=self._translate_openai_finish_reason_to_anthropic( - response.choices[0].finish_reason - ), + stop_reason=self._translate_openai_finish_reason_to_anthropic(response.choices[0].finish_reason), ) if getattr(response, "usage", None) is not None: litellm_usage_chunk: Optional[Usage] = response.usage # type: ignore - elif ( - hasattr(response, "_hidden_params") - and "usage" in response._hidden_params - ): + elif hasattr(response, "_hidden_params") and "usage" in response._hidden_params: litellm_usage_chunk = response._hidden_params["usage"] else: litellm_usage_chunk = None if litellm_usage_chunk is not None: - usage_delta = self._translate_openai_usage_to_anthropic_usage_delta( - litellm_usage_chunk - ) + usage_delta = self._translate_openai_usage_to_anthropic_usage_delta(litellm_usage_chunk) else: usage_delta = UsageDelta(input_tokens=0, output_tokens=0) message_block = MessageBlockDelta( @@ -1718,9 +1501,7 @@ class LiteLLMAnthropicMessagesAdapter: usage=usage_delta, # type: ignore ) if applied_edits: - message_block["context_management"] = ContextManagementResponse( - applied_edits=list(applied_edits) - ) + message_block["context_management"] = ContextManagementResponse(applied_edits=list(applied_edits)) return message_block ( type_of_content, diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py b/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py index ebbc182c427..50217d4bc82 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py @@ -40,6 +40,4 @@ COMPACT_DEFAULT_INSTRUCTIONS = ( # Appended to the default prompt when ``tools`` are present and the caller # did not supply custom ``instructions``. Matches the guidance in the # Anthropic docs under "Compaction might fail when tools are defined". -COMPACT_NO_TOOL_CALLS_SUFFIX = ( - " Do not call any tools while writing this summary; respond with text only." -) +COMPACT_NO_TOOL_CALLS_SUFFIX = " Do not call any tools while writing this summary; respond with text only." diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py index 7b1c20ff522..8bcf8acfff6 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py @@ -68,9 +68,7 @@ def _trigger_met( messages=messages, tools=cast(Any, tools), ) - verbose_logger.debug( - f"context_management polyfill: current_tokens: {current_tokens}" - ) + verbose_logger.debug(f"context_management polyfill: current_tokens: {current_tokens}") verbose_logger.debug(f"context_management polyfill: threshold: {threshold}") return current_tokens > threshold, current_tokens @@ -101,9 +99,7 @@ def _last_completed_tool_use_id( return last_id -def _clear_tool_results( - messages: List[Dict[str, Any]], ids_to_clear: set -) -> Tuple[List[Dict[str, Any]], int]: +def _clear_tool_results(messages: List[Dict[str, Any]], ids_to_clear: set) -> Tuple[List[Dict[str, Any]], int]: """Clear matching tool_result content; return (messages, cleared_count).""" cleared = 0 new_messages: List[Dict[str, Any]] = [] @@ -148,11 +144,7 @@ def apply_clear_tool_uses_20250919( edit_spec: Dict[str, Any], ) -> Tuple[List[Dict[str, Any]], Optional[AppliedEdit]]: """Apply clear_tool_uses; return (messages, AppliedEdit or None).""" - ignored_knobs = [ - knob - for knob in ("clear_at_least", "exclude_tools", "clear_tool_inputs") - if knob in edit_spec - ] + ignored_knobs = [knob for knob in ("clear_at_least", "exclude_tools", "clear_tool_inputs") if knob in edit_spec] for ignored_knob in ignored_knobs: verbose_logger.warning( "context_management polyfill: ignoring '%s' on %s " @@ -192,12 +184,8 @@ def apply_clear_tool_uses_20250919( return messages, None if tokens_before is None: - tokens_before = litellm.token_counter( - model=model, messages=messages, tools=cast(Any, tools) - ) - tokens_after = litellm.token_counter( - model=model, messages=edited, tools=cast(Any, tools) - ) + tokens_before = litellm.token_counter(model=model, messages=messages, tools=cast(Any, tools)) + tokens_after = litellm.token_counter(model=model, messages=edited, tools=cast(Any, tools)) cleared_input_tokens = max(tokens_before - tokens_after, 0) applied: AppliedEdit = { diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index 6479ee999b0..f18a9f41939 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -266,8 +266,7 @@ async def _check_summary_model_access( team_membership = None member_allowed_models = ( team_membership.litellm_budget_table.allowed_models - if team_membership is not None - and team_membership.litellm_budget_table is not None + if team_membership is not None and team_membership.litellm_budget_table is not None else None ) if member_allowed_models: @@ -328,22 +327,15 @@ async def _check_summary_model_budget( return False except Exception as e: verbose_logger.warning( - "compact_20260112: unexpected error during key model-budget " - "check for summary_model=%s; denying: %s", + "compact_20260112: unexpected error during key model-budget check for summary_model=%s; denying: %s", summary_model, e, ) return False - end_user_model_max_budget = getattr( - user_api_key_auth, "end_user_model_max_budget", None - ) + end_user_model_max_budget = getattr(user_api_key_auth, "end_user_model_max_budget", None) end_user_id = getattr(user_api_key_auth, "end_user_id", None) - if ( - isinstance(end_user_model_max_budget, dict) - and end_user_model_max_budget - and end_user_id is not None - ): + if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: try: await model_max_budget_limiter.is_end_user_within_model_budget( end_user_id=end_user_id, @@ -422,11 +414,7 @@ async def _check_summary_model_rate_limit( requested_model=summary_model, descriptors=descriptors, ) - descriptors.extend( - limiter.create_organization_rate_limit_descriptor( - user_api_key_auth, summary_model - ) - ) + descriptors.extend(limiter.create_organization_rate_limit_descriptor(user_api_key_auth, summary_model)) if not descriptors: return True response = await limiter.should_rate_limit( @@ -436,8 +424,7 @@ async def _check_summary_model_rate_limit( ) except Exception as e: verbose_logger.warning( - "compact_20260112: unexpected error during rate-limit check for " - "summary_model=%s; allowing: %s", + "compact_20260112: unexpected error during rate-limit check for summary_model=%s; allowing: %s", summary_model, e, ) @@ -507,11 +494,7 @@ def _strip_compaction_blocks( if not isinstance(content, list): cleaned.append(msg) continue - filtered = [ - block - for block in content - if not (isinstance(block, dict) and block.get("type") == "compaction") - ] + filtered = [block for block in content if not (isinstance(block, dict) and block.get("type") == "compaction")] if not filtered: # The compaction block was the only content; drop the whole turn. continue @@ -573,9 +556,7 @@ def _resolve_trigger_tokens(edit_spec: Dict[str, Any]) -> Tuple[int, List[str]]: return value, warnings -def _build_summary_prompt( - edit_spec: Dict[str, Any], tools: Optional[List[Dict[str, Any]]] -) -> str: +def _build_summary_prompt(edit_spec: Dict[str, Any], tools: Optional[List[Dict[str, Any]]]) -> str: custom = edit_spec.get("instructions") if isinstance(custom, str) and custom.strip(): return custom @@ -628,9 +609,7 @@ def _count_effective_tokens( messages_without_compaction = _strip_compaction_blocks(effective_messages) adapter = LiteLLMAnthropicMessagesAdapter() try: - openai_shape = adapter.translate_anthropic_messages_to_openai( - messages=cast(Any, messages_without_compaction) - ) + openai_shape = adapter.translate_anthropic_messages_to_openai(messages=cast(Any, messages_without_compaction)) except Exception as e: verbose_logger.debug( "compact_20260112: anthropic→openai translation failed during token " @@ -647,9 +626,7 @@ def _count_effective_tokens( openai_tools: Optional[List[Dict[str, Any]]] = None if tools: try: - translated_tools, _ = adapter.translate_anthropic_tools_to_openai( - tools=cast(Any, tools) - ) + translated_tools, _ = adapter.translate_anthropic_tools_to_openai(tools=cast(Any, tools)) openai_tools = cast(List[Dict[str, Any]], translated_tools) except Exception as e: verbose_logger.debug( @@ -713,11 +690,7 @@ def _select_last_user_question( continue content = msg.get("content") if isinstance(content, list): - filtered = [ - blk - for blk in content - if not (isinstance(blk, dict) and blk.get("type") == "tool_result") - ] + filtered = [blk for blk in content if not (isinstance(blk, dict) and blk.get("type") == "tool_result")] if not filtered: # Purely tool_result — skip and look for an earlier turn. continue @@ -755,11 +728,7 @@ def _system_to_openai_message( if isinstance(system, str): return {"role": "system", "content": system} if system else None if isinstance(system, list): - parts = [ - block.get("text", "") - for block in system - if isinstance(block, dict) and block.get("type") == "text" - ] + parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"] joined = "\n\n".join(part for part in parts if part) return {"role": "system", "content": joined} if joined else None return None @@ -783,10 +752,8 @@ def _build_summary_messages( stripped = _strip_compaction_blocks(effective_messages) try: - openai_messages = ( - LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( - messages=cast(Any, stripped) - ) + openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=cast(Any, stripped) ) except Exception as e: verbose_logger.warning( @@ -902,9 +869,7 @@ def _extract_response_text(response: Any) -> Optional[str]: # Some providers return a list of content parts. if isinstance(content, list): text_parts = [ - part.get("text", "") - for part in content - if isinstance(part, dict) and part.get("type") == "text" + part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text" ] return "".join(text_parts) or None except (AttributeError, IndexError, KeyError): @@ -936,9 +901,7 @@ def apply_client_compaction_block_history( tail is forwarded unchanged (with compaction blocks stripped) so recent turns the summary does not cover are preserved. """ - effective_messages, prior_compaction_block = _slice_around_compaction_block( - messages - ) + effective_messages, prior_compaction_block = _slice_around_compaction_block(messages) if prior_compaction_block is None: return None @@ -1006,12 +969,8 @@ async def apply_compact_20260112( # opt-in gate below so that even when summarization is disabled we still # strip Anthropic-only ``compaction`` blocks from messages going to # non-Anthropic backends (which would reject them). - effective_messages, prior_compaction_block = _slice_around_compaction_block( - messages - ) - prior_summary_text = ( - prior_compaction_block.get("content") if prior_compaction_block else None - ) + effective_messages, prior_compaction_block = _slice_around_compaction_block(messages) + prior_summary_text = prior_compaction_block.get("content") if prior_compaction_block else None augmented_system: Union[str, List[Dict[str, Any]], None] = system if isinstance(prior_summary_text, str) and prior_summary_text: augmented_system = _augment_system_with_summary(system, prior_summary_text) @@ -1053,14 +1012,10 @@ async def apply_compact_20260112( system=augmented_system, ) except Exception as e: - verbose_logger.warning( - "compact_20260112: token_counter failed; assuming under threshold: %s", e - ) + verbose_logger.warning("compact_20260112: token_counter failed; assuming under threshold: %s", e) current_tokens = 0 - verbose_logger.debug( - "compact_20260112: current_tokens=%s trigger=%s", current_tokens, trigger_tokens - ) + verbose_logger.debug("compact_20260112: current_tokens=%s trigger=%s", current_tokens, trigger_tokens) if current_tokens <= trigger_tokens: # Slice-only path: the prior compaction summary already lives in @@ -1086,8 +1041,7 @@ async def apply_compact_20260112( llm_router=llm_router, ): verbose_logger.warning( - "compact_20260112: caller not authorized for summary_model=%s; " - "skipping summary call", + "compact_20260112: caller not authorized for summary_model=%s; skipping summary call", summary_model, ) applied["error"] = "summary_model_access_denied" @@ -1102,8 +1056,7 @@ async def apply_compact_20260112( summary_model=summary_model, ): verbose_logger.warning( - "compact_20260112: caller over model budget for summary_model=%s; " - "skipping summary call", + "compact_20260112: caller over model budget for summary_model=%s; skipping summary call", summary_model, ) applied["error"] = "summary_model_budget_exceeded" @@ -1118,8 +1071,7 @@ async def apply_compact_20260112( summary_model=summary_model, ): verbose_logger.warning( - "compact_20260112: caller over rate limit for summary_model=%s; " - "skipping summary call", + "compact_20260112: caller over rate limit for summary_model=%s; skipping summary call", summary_model, ) applied["error"] = "summary_model_rate_limit_exceeded" @@ -1130,9 +1082,7 @@ async def apply_compact_20260112( ) prompt = _build_summary_prompt(edit_spec, tools) - summary_messages = _build_summary_messages( - effective_messages, prompt, system=augmented_system - ) + summary_messages = _build_summary_messages(effective_messages, prompt, system=augmented_system) propagated_metadata = _propagate_metadata(litellm_metadata) allowed_model_region = getattr(user_api_key_auth, "allowed_model_region", None) diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/result.py b/litellm/llms/anthropic/experimental_pass_through/context_management/result.py index 36bcde98d0c..14adeb9452a 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/result.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/result.py @@ -42,11 +42,7 @@ class PolyfillResult: visible: List[AppliedEdit] = [] for edit in self.applied_edits: if edit.get("type") == COMPACT_EDIT_TYPE: - if ( - self.compaction_block is not None - or edit.get("error") - or edit.get("warnings") - ): + if self.compaction_block is not None or edit.get("error") or edit.get("warnings"): visible.append(edit) else: visible.append(edit) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index d693d50b8e5..cb37725d79c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -94,9 +94,7 @@ def _handle_content_block_delta(data: Dict, content_blocks: Dict[int, Dict]) -> if delta_type == "text_delta": block["text"] = block.get("text", "") + delta.get("text", "") elif delta_type == "input_json_delta": - block["_partial_json"] = block.get("_partial_json", "") + delta.get( - "partial_json", "" - ) + block["_partial_json"] = block.get("_partial_json", "") + delta.get("partial_json", "") elif delta_type == "thinking_delta": block["thinking"] = block.get("thinking", "") + delta.get("thinking", "") elif delta_type == "signature_delta": @@ -163,9 +161,7 @@ class AgenticAnthropicStreamingIterator: self._model = model self._messages = messages self._anthropic_messages_provider_config = anthropic_messages_provider_config - self._anthropic_messages_optional_request_params = ( - anthropic_messages_optional_request_params - ) + self._anthropic_messages_optional_request_params = anthropic_messages_optional_request_params self._logging_obj = logging_obj self._custom_llm_provider = custom_llm_provider self._kwargs = kwargs @@ -209,17 +205,11 @@ class AgenticAnthropicStreamingIterator: try: rebuilt = self._rebuild_anthropic_response_from_sse(self._collected_bytes) if rebuilt is None: - verbose_logger.debug( - "AgenticStreamingIterator: Could not rebuild response from SSE bytes" - ) + verbose_logger.debug("AgenticStreamingIterator: Could not rebuild response from SSE bytes") return [ - ( - f"{b.get('type')}({b.get('name', '')})" - if b.get("type") == "tool_use" - else b.get("type") - ) + (f"{b.get('type')}({b.get('name', '')})" if b.get("type") == "tool_use" else b.get("type")) for b in rebuilt.get("content", []) ] @@ -248,9 +238,7 @@ class AgenticAnthropicStreamingIterator: AnthropicMessagesResponse, ) - fake = FakeAnthropicMessagesStreamIterator( - response=cast(AnthropicMessagesResponse, result) - ) + fake = FakeAnthropicMessagesStreamIterator(response=cast(AnthropicMessagesResponse, result)) self._follow_up_iterator = fake.__aiter__() else: verbose_logger.warning( @@ -260,8 +248,7 @@ class AgenticAnthropicStreamingIterator: except Exception as e: _call_id = getattr(self._logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( - "AgenticStreamingIterator: Error in agentic hook processing " - "[call_id=%s model=%s]: %s", + "AgenticStreamingIterator: Error in agentic hook processing [call_id=%s model=%s]: %s", _call_id, self._model, str(e), diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py index f704ed2c9d1..184fede25e9 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py @@ -38,9 +38,7 @@ class FakeAnthropicMessagesStreamIterator: self.chunks = self._create_streaming_chunks() self.current_index = 0 - def _create_content_block_chunks( - self, block_dict: Dict[str, Any], index: int - ) -> List[bytes]: + def _create_content_block_chunks(self, block_dict: Dict[str, Any], index: int) -> List[bytes]: """Build SSE chunks for a single content block.""" chunks = [] block_type = block_dict.get("type") @@ -51,18 +49,14 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "content_block": {"type": "text", "text": ""}, } - chunks.append( - f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() - ) + chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) text = block_dict.get("text", "") content_block_delta = { "type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}, } - chunks.append( - f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode() - ) + chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()) elif block_type == "thinking": content_block_start = { @@ -70,9 +64,7 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "content_block": {"type": "thinking", "thinking": "", "signature": ""}, } - chunks.append( - f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() - ) + chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) thinking_text = block_dict.get("thinking", "") if thinking_text: content_block_delta = { @@ -80,9 +72,7 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "delta": {"type": "thinking_delta", "thinking": thinking_text}, } - chunks.append( - f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode() - ) + chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()) signature = block_dict.get("signature", "") if signature: signature_delta = { @@ -90,9 +80,7 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "delta": {"type": "signature_delta", "signature": signature}, } - chunks.append( - f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode() - ) + chunks.append(f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode()) elif block_type == "redacted_thinking": content_block_start = { @@ -100,9 +88,7 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "content_block": {"type": "redacted_thinking"}, } - chunks.append( - f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() - ) + chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) elif block_type == "tool_use": content_block_start = { @@ -115,9 +101,7 @@ class FakeAnthropicMessagesStreamIterator: "input": {}, }, } - chunks.append( - f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() - ) + chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) input_data = block_dict.get("input", {}) content_block_delta = { "type": "content_block_delta", @@ -127,14 +111,10 @@ class FakeAnthropicMessagesStreamIterator: "partial_json": json.dumps(input_data), }, } - chunks.append( - f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode() - ) + chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()) content_block_stop = {"type": "content_block_stop", "index": index} - chunks.append( - f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode() - ) + chunks.append(f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode()) return chunks def _create_streaming_chunks(self) -> List[bytes]: @@ -162,9 +142,7 @@ class FakeAnthropicMessagesStreamIterator: }, }, } - chunks.append( - f"event: message_start\ndata: {json.dumps(message_start)}\n\n".encode() - ) + chunks.append(f"event: message_start\ndata: {json.dumps(message_start)}\n\n".encode()) # 2-4. For each content block, send start/delta/stop events content_blocks = response_dict.get("content", []) @@ -182,13 +160,9 @@ class FakeAnthropicMessagesStreamIterator: if usage.get("input_tokens") is not None: delta_usage["input_tokens"] = usage["input_tokens"] if usage.get("cache_creation_input_tokens") is not None: - delta_usage["cache_creation_input_tokens"] = usage[ - "cache_creation_input_tokens" - ] + delta_usage["cache_creation_input_tokens"] = usage["cache_creation_input_tokens"] if usage.get("cache_read_input_tokens") is not None: - delta_usage["cache_read_input_tokens"] = usage[ - "cache_read_input_tokens" - ] + delta_usage["cache_read_input_tokens"] = usage["cache_read_input_tokens"] message_delta = { "type": "message_delta", "delta": { @@ -197,15 +171,11 @@ class FakeAnthropicMessagesStreamIterator: }, "usage": delta_usage, } - chunks.append( - f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode() - ) + chunks.append(f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode()) # 6. message_stop event message_stop = {"type": "message_stop", "usage": usage if usage else {}} - chunks.append( - f"event: message_stop\ndata: {json.dumps(message_stop)}\n\n".encode() - ) + chunks.append(f"event: message_stop\ndata: {json.dumps(message_stop)}\n\n".encode()) return chunks diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index c06241f0e4f..9c9427c7302 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -120,9 +120,7 @@ async def _execute_pre_request_hooks( continue # Call the pre-request hook - modified_kwargs = await callback.async_pre_request_hook( - model, messages, request_kwargs - ) + modified_kwargs = await callback.async_pre_request_hook(model, messages, request_kwargs) # If hook returned modified kwargs, use them if modified_kwargs is not None: @@ -219,9 +217,7 @@ async def anthropic_messages( # ids like ``functions.Bash:0`` that violate Anthropic's id pattern. messages = sanitize_tool_use_ids_in_anthropic_messages(messages) - original_stream = stream or kwargs.get( - "_websearch_interception_converted_stream", False - ) + original_stream = stream or kwargs.get("_websearch_interception_converted_stream", False) # Execute pre-request hooks to allow CustomLoggers to modify request. # tool_choice is forwarded explicitly (it is a named param, not in kwargs) @@ -255,9 +251,7 @@ async def anthropic_messages( # The litellm_params dict may have been overwritten by **kwargs in # _execute_pre_request_hooks, so fall back to get_llm_provider() if needed. if not custom_llm_provider: - custom_llm_provider = request_kwargs.get("litellm_params", {}).get( - "custom_llm_provider" - ) + custom_llm_provider = request_kwargs.get("litellm_params", {}).get("custom_llm_provider") if not custom_llm_provider: try: _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) @@ -385,9 +379,7 @@ def anthropic_messages_handler( AnthropicMessagesResponse, Iterator[bytes], AsyncIterator[Any], - Coroutine[ - Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]] - ], + Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]], ]: """ Makes Anthropic `/v1/messages` API calls In the Anthropic API Spec @@ -447,9 +439,7 @@ def anthropic_messages_handler( # Check if stream was converted for WebSearch interception # This is set in the async wrapper above when stream=True is converted to stream=False if kwargs.get("_websearch_interception_converted_stream", False): - litellm_logging_obj.model_call_details[ - "websearch_interception_converted_stream" - ] = True + litellm_logging_obj.model_call_details["websearch_interception_converted_stream"] = True if litellm_params.mock_response and isinstance(litellm_params.mock_response, str): return mock_response( @@ -461,14 +451,10 @@ def anthropic_messages_handler( anthropic_messages_provider_config: Optional[BaseAnthropicMessagesConfig] = None - if custom_llm_provider is not None and custom_llm_provider in [ - provider.value for provider in LlmProviders - ]: - anthropic_messages_provider_config = ( - ProviderConfigManager.get_provider_anthropic_messages_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: + anthropic_messages_provider_config = ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. @@ -494,18 +480,14 @@ def anthropic_messages_handler( **kwargs, ) if _should_route_to_responses_api(custom_llm_provider): - return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( - **_shared_kwargs - ) + return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(**_shared_kwargs) # The in-gateway context_management polyfill runs inside # ``async_anthropic_messages_handler`` so it can ``await`` the # summarization model for ``compact_20260112``. ``context_management`` # is passed through as a regular kwarg. - return ( - LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( - **_shared_kwargs, - ) + return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + **_shared_kwargs, ) if custom_llm_provider is None: @@ -524,10 +506,7 @@ def anthropic_messages_handler( ) if is_reasoning_auto_summary_enabled(): thinking_param = anthropic_messages_optional_request_params.get("thinking") - if ( - isinstance(thinking_param, dict) - and thinking_param.get("type") != "disabled" - ): + if isinstance(thinking_param, dict) and thinking_param.get("type") != "disabled": anthropic_messages_optional_request_params["thinking"] = { **thinking_param, "display": "summarized", @@ -537,9 +516,7 @@ def anthropic_messages_handler( model=model, messages=messages, anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=dict( - anthropic_messages_optional_request_params - ), + anthropic_messages_optional_request_params=dict(anthropic_messages_optional_request_params), _is_async=is_async, client=client, custom_llm_provider=custom_llm_provider, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py index 8714939f025..6c72b7a3e00 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py @@ -70,18 +70,12 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): None, ) if advisor_tool is None: - raise ValueError( - f"handle() called but no {ANTHROPIC_ADVISOR_TOOL_TYPE} tool found in tools list" - ) + raise ValueError(f"handle() called but no {ANTHROPIC_ADVISOR_TOOL_TYPE} tool found in tools list") advisor_model: str = advisor_tool.get("model") or "" if not advisor_model: - raise ValueError( - "advisor tool definition must include a 'model' field specifying the advisor model" - ) + raise ValueError("advisor tool definition must include a 'model' field specifying the advisor model") _raw_max_uses = advisor_tool.get("max_uses") - max_uses: int = ( - ADVISOR_MAX_USES if _raw_max_uses is None else int(_raw_max_uses) - ) + max_uses: int = ADVISOR_MAX_USES if _raw_max_uses is None else int(_raw_max_uses) # Optional routing overrides for the advisor sub-call (e.g. proxy routing). # If not set in the tool definition, litellm resolves from env vars. # The advisor tool is caller-controlled; only honor a client-supplied @@ -98,12 +92,7 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): # Executor tools = all original tools with advisor replaced by the synthetic one. executor_tools: List[Dict] = [ - ( - synthetic_advisor_tool - if t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE - else t - ) - for t in (tools or []) + (synthetic_advisor_tool if t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE else t) for t in (tools or []) ] # Strip prior advisor blocks from history, preserving advice text as context. @@ -111,9 +100,7 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): [dict(m) for m in messages], replace_with_text=True ) - parent_request_id: str = str( - kwargs.pop("litellm_call_id", None) or uuid.uuid4() - ) + parent_request_id: str = str(kwargs.pop("litellm_call_id", None) or uuid.uuid4()) metadata_base: Dict = dict(kwargs.pop("metadata", None) or {}) iteration = 0 @@ -150,9 +137,7 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): ) # --- Build advisor context --- - advisor_messages = _build_advisor_context( - current_messages, executor_response, advisor_use_block - ) + advisor_messages = _build_advisor_context(current_messages, executor_response, advisor_use_block) # --- Advisor sub-call (always non-streaming, no tools) --- advisor_response: AnthropicMessagesResponse = await _call_messages_handler( @@ -225,11 +210,7 @@ def _find_advisor_tool_use(response: Any) -> Optional[Dict]: if not isinstance(content, list): return None for block in content: - if ( - isinstance(block, dict) - and block.get("type") == "tool_use" - and block.get("name") == "advisor" - ): + if isinstance(block, dict) and block.get("type") == "tool_use" and block.get("name") == "advisor": return block return None @@ -239,11 +220,7 @@ def _extract_response_text(response: Any) -> str: content = response.get("content") if isinstance(response, dict) else [] if not isinstance(content, list): return "" - parts = [ - b.get("text", "") - for b in content - if isinstance(b, dict) and b.get("type") == "text" - ] + parts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"] return "\n".join(parts).strip() @@ -267,9 +244,7 @@ def _build_advisor_context( question = (advisor_use_block.get("input") or {}).get("question") or ( "Please provide guidance on the current task." ) - raw_content = ( - executor_response.get("content") if isinstance(executor_response, dict) else [] - ) or [] + raw_content = (executor_response.get("content") if isinstance(executor_response, dict) else []) or [] # Keep only text blocks — strip tool_use and provider-specific fields. executor_text_blocks = [ {k: v for k, v in block.items() if k not in _PROVIDER_SPECIFIC_KEYS} @@ -293,9 +268,7 @@ def _inject_advisor_turn( Append the executor's response (as an assistant turn) and the advisor result (as a user tool_result turn) so the executor can continue. """ - executor_content = ( - executor_response.get("content") if isinstance(executor_response, dict) else [] - ) or [] + executor_content = (executor_response.get("content") if isinstance(executor_response, dict) else []) or [] tool_use_id = advisor_use_block.get("id", "") return [ *messages, @@ -322,9 +295,7 @@ def _inject_max_uses_error( Inject a max_uses_exceeded error tool_result so the executor continues without further advisor calls (mirrors Anthropic's server-side behaviour). """ - executor_content = ( - executor_response.get("content") if isinstance(executor_response, dict) else [] - ) or [] + executor_content = (executor_response.get("content") if isinstance(executor_response, dict) else []) or [] tool_use_id = advisor_use_block.get("id", "") return [ *messages, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 25109765772..2357960f716 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -40,9 +40,7 @@ class BaseAnthropicMessagesStreamingIterator: # chunk rather than falling back to end_time in async_success_handler. if self.completion_start_time is not None: self.litellm_logging_obj.completion_start_time = self.completion_start_time - self.litellm_logging_obj.model_call_details["completion_start_time"] = ( - self.completion_start_time - ) + self.litellm_logging_obj.model_call_details["completion_start_time"] = self.completion_start_time asyncio.create_task( PassThroughStreamingHandler._route_streaming_logging_to_handler( litellm_logging_obj=self.litellm_logging_obj, @@ -95,9 +93,7 @@ class BaseAnthropicMessagesStreamingIterator: async def async_sse_wrapper( self, - completion_stream: AsyncIterator[ - Union[bytes, GenericStreamingChunk, ModelResponseStream, dict] - ], + completion_stream: AsyncIterator[Union[bytes, GenericStreamingChunk, ModelResponseStream, dict]], ) -> AsyncIterator[bytes]: """ Generic async SSE wrapper that converts streaming chunks to SSE format diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index fbe49ae4136..e78802a1587 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -57,9 +57,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # "metadata", ] - def _remove_scope_from_cache_control( - self, anthropic_messages_request: Dict - ) -> None: + def _remove_scope_from_cache_control(self, anthropic_messages_request: Dict) -> None: """ Remove `scope` field from cache_control blocks. @@ -122,9 +120,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): text = content_block.get("text", "") content_type = content_block.get("type", "") # Skip text blocks that start with billing header - if content_type == "text" and text.startswith( - "x-anthropic-billing-header:" - ): + if content_type == "text" and text.startswith("x-anthropic-billing-header:"): continue filtered_list.append(content_block) else: @@ -143,9 +139,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = ( - AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com" - ) + api_base = AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com" if not api_base.endswith("/v1/messages"): api_base = f"{api_base}/v1/messages" return api_base @@ -161,9 +155,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): api_base: Optional[str] = None, ) -> Tuple[dict, Optional[str]]: # Check for Anthropic OAuth token in Authorization header - headers, api_key = optionally_handle_anthropic_oauth( - headers=headers, api_key=api_key - ) + headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) if "x-api-key" not in headers and "authorization" not in headers: auth_header = AnthropicModelInfo.get_auth_header(api_key) @@ -182,9 +174,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return headers, api_base @staticmethod - def _translate_reasoning_effort_to_anthropic( - model: str, optional_params: Dict - ) -> None: + def _translate_reasoning_effort_to_anthropic(model: str, optional_params: Dict) -> None: """Map OpenAI-style ``reasoning_effort`` to native Anthropic params. Caller-supplied ``thinking`` / ``output_config`` win over the alias. @@ -201,9 +191,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return try: - mapped_thinking = AnthropicConfig._map_reasoning_effort( - reasoning_effort=reasoning_effort, model=model - ) + mapped_thinking = AnthropicConfig._map_reasoning_effort(reasoning_effort=reasoning_effort, model=model) except _BadRequestError as e: raise AnthropicError(message=str(e.message), status_code=400) @@ -214,9 +202,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params.setdefault("thinking", mapped_thinking) if AnthropicModelInfo._is_adaptive_thinking_model(model): - mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get( - reasoning_effort - ) + mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort) if mapped_effort is None: raise AnthropicError( message=( @@ -226,9 +212,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ), status_code=400, ) - gate_error = AnthropicConfig._validate_effort_for_model( - model, mapped_effort - ) + gate_error = AnthropicConfig._validate_effort_for_model(model, mapped_effort) if gate_error is not None: raise AnthropicError(message=gate_error, status_code=400) existing_output_config = optional_params.get("output_config") @@ -238,9 +222,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params["output_config"] = existing_output_config @staticmethod - def _translate_legacy_thinking_for_adaptive_model( - model: str, optional_params: Dict - ) -> None: + def _translate_legacy_thinking_for_adaptive_model(model: str, optional_params: Dict) -> None: """Translate legacy ``thinking.type=enabled`` to adaptive for 4.6/4.7. Caller-provided ``output_config.effort`` is never overridden. """ @@ -311,21 +293,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): anthropic_messages_optional_request_params.pop("system", None) # Transform context_management from OpenAI format to Anthropic format if needed - context_management_param = anthropic_messages_optional_request_params.get( - "context_management" - ) + context_management_param = anthropic_messages_optional_request_params.get("context_management") if context_management_param is not None: from litellm.llms.anthropic.chat.transformation import AnthropicConfig - transformed_context_management = ( - AnthropicConfig.map_openai_context_management_to_anthropic( - context_management_param - ) + transformed_context_management = AnthropicConfig.map_openai_context_management_to_anthropic( + context_management_param ) if transformed_context_management is not None: - anthropic_messages_optional_request_params["context_management"] = ( - transformed_context_management - ) + anthropic_messages_optional_request_params["context_management"] = transformed_context_management ####### get required params for all anthropic messages requests ###### # Lazy %s: the f-string previously stringified the entire messages @@ -336,10 +312,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # Auto-strip advisor blocks from history if advisor tool is absent. # Prevents Anthropic 400: advisor_tool_result in history requires advisor tool. _tools = anthropic_messages_optional_request_params.get("tools") or [] - _has_advisor = any( - isinstance(t, dict) and t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE - for t in _tools - ) + _has_advisor = any(isinstance(t, dict) and t.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for t in _tools) if not _has_advisor: messages = strip_advisor_blocks_from_messages(messages) # type: ignore[assignment] @@ -363,9 +336,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): try: raw_response_json = raw_response.json() except Exception: - raise AnthropicError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise AnthropicError(message=raw_response.text, status_code=raw_response.status_code) return AnthropicMessagesResponse(**raw_response_json) def get_async_streaming_response_iterator( @@ -439,9 +410,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # Add context management header if any other edits exist if has_other: - beta_values.add( - ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value - ) + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) # Check for structured outputs. Anthropic's newer request shape nests # the schema under output_config.format; the older top-level @@ -450,9 +419,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if optional_params.get("output_format") is not None or ( isinstance(output_config, dict) and output_config.get("format") is not None ): - beta_values.add( - ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value - ) + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value) # Check for fast mode if optional_params.get("speed") == "fast": @@ -462,13 +429,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): tools = optional_params.get("tools") if tools: for tool in tools: - if ( - isinstance(tool, dict) - and tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE - ): - beta_values.add( - ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value - ) + if isinstance(tool, dict) and tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE: + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value) break # Check for tool search tools diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 42167e0fdaa..c8060d41fad 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -41,9 +41,7 @@ class AnthropicMessagesRequestUtils: AnthropicMessagesRequestOptionalParams instance with only the valid parameters """ valid_keys = _anthropic_messages_optional_param_keys() - filtered_params = { - k: v for k, v in params.items() if k in valid_keys and v is not None - } + filtered_params = {k: v for k, v in params.items() if k in valid_keys and v is not None} if model is not None: from litellm.llms.anthropic.chat.transformation import AnthropicConfig diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 70855afa81c..7911845a598 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -163,9 +163,7 @@ class LiteLLMMessagesToResponsesAPIHandler: result = await litellm.aresponses(**responses_kwargs) if stream: - wrapper = AnthropicResponsesStreamWrapper( - responses_stream=result, model=model - ) + wrapper = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) return wrapper.async_anthropic_sse_wrapper() if not isinstance(result, ResponsesAPIResponse): @@ -199,26 +197,24 @@ class LiteLLMMessagesToResponsesAPIHandler: Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any]]], ]: if _is_async: - return ( - LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( - max_tokens=max_tokens, - messages=messages, - model=model, - context_management=context_management, - metadata=metadata, - output_config=output_config, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - **kwargs, - ) + return LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=max_tokens, + messages=messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + **kwargs, ) # Sync path @@ -245,9 +241,7 @@ class LiteLLMMessagesToResponsesAPIHandler: result = litellm.responses(**responses_kwargs) if stream: - wrapper = AnthropicResponsesStreamWrapper( - responses_stream=result, model=model - ) + wrapper = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) return wrapper.async_anthropic_sse_wrapper() if not isinstance(result, ResponsesAPIResponse): diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index f400dc7804e..0d02b4fa969 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -35,9 +35,7 @@ class AnthropicResponsesStreamWrapper: # Map item_id -> content_block_index so we can stop the right block later self._item_id_to_block_index: Dict[str, int] = {} # Track open function_call items by item_id so we can emit tool_use start - self._pending_tool_ids: Dict[ - str, str - ] = {} # item_id -> call_id / name accumulator + self._pending_tool_ids: Dict[str, str] = {} # item_id -> call_id / name accumulator self._sent_message_start = False self._sent_message_stop = False self._chunk_queue: deque = deque() @@ -83,17 +81,11 @@ class AnthropicResponsesStreamWrapper: # ---- content_block_start for a new output message item ---- if event_type == "response.output_item.added": - item = getattr(event, "item", None) or ( - event.get("item") if isinstance(event, dict) else None - ) + item = getattr(event, "item", None) or (event.get("item") if isinstance(event, dict) else None) if item is None: return - item_type = getattr(item, "type", None) or ( - item.get("type") if isinstance(item, dict) else None - ) - item_id = getattr(item, "id", None) or ( - item.get("id") if isinstance(item, dict) else None - ) + item_type = getattr(item, "type", None) or (item.get("type") if isinstance(item, dict) else None) + item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item_type == "message": block_idx = self._next_block_index() @@ -108,15 +100,9 @@ class AnthropicResponsesStreamWrapper: ) elif item_type == "function_call": call_id = ( - getattr(item, "call_id", None) - or (item.get("call_id") if isinstance(item, dict) else None) - or "" - ) - name = ( - getattr(item, "name", None) - or (item.get("name") if isinstance(item, dict) else None) - or "" + getattr(item, "call_id", None) or (item.get("call_id") if isinstance(item, dict) else None) or "" ) + name = getattr(item, "name", None) or (item.get("name") if isinstance(item, dict) else None) or "" block_idx = self._next_block_index() if item_id: self._item_id_to_block_index[item_id] = block_idx @@ -148,17 +134,9 @@ class AnthropicResponsesStreamWrapper: # ---- text delta ---- if event_type == "response.output_text.delta": - item_id = getattr(event, "item_id", None) or ( - event.get("item_id") if isinstance(event, dict) else None - ) - delta = getattr(event, "delta", "") or ( - event.get("delta", "") if isinstance(event, dict) else "" - ) - block_idx = ( - self._item_id_to_block_index.get(item_id, -1) - if item_id - else self._current_block_index - ) + item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) + delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index if block_idx < 0: # Some providers (e.g. LMStudio) skip response.output_item.added, # so no text block is open yet; synthesize content_block_start @@ -184,12 +162,8 @@ class AnthropicResponsesStreamWrapper: # ---- reasoning summary text delta ---- if event_type == "response.reasoning_summary_text.delta": - item_id = getattr(event, "item_id", None) or ( - event.get("item_id") if isinstance(event, dict) else None - ) - delta = getattr(event, "delta", "") or ( - event.get("delta", "") if isinstance(event, dict) else "" - ) + item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) + delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") block_idx = ( self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id @@ -206,12 +180,8 @@ class AnthropicResponsesStreamWrapper: # ---- function call arguments delta ---- if event_type == "response.function_call_arguments.delta": - item_id = getattr(event, "item_id", None) or ( - event.get("item_id") if isinstance(event, dict) else None - ) - delta = getattr(event, "delta", "") or ( - event.get("delta", "") if isinstance(event, dict) else "" - ) + item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) + delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") block_idx = ( self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id @@ -228,14 +198,9 @@ class AnthropicResponsesStreamWrapper: # ---- output item done -> content_block_stop ---- if event_type == "response.output_item.done": - item = getattr(event, "item", None) or ( - event.get("item") if isinstance(event, dict) else None - ) + item = getattr(event, "item", None) or (event.get("item") if isinstance(event, dict) else None) item_id = ( - getattr(item, "id", None) - or (item.get("id") if isinstance(item, dict) else None) - if item - else None + getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None ) block_idx = ( self._item_id_to_block_index.get(item_id, self._current_block_index) @@ -276,12 +241,8 @@ class AnthropicResponsesStreamWrapper: cache_creation_tokens = getattr(usage, "input_tokens_details", None) # type: ignore[assignment] cache_read_tokens = getattr(usage, "output_tokens_details", None) # type: ignore[assignment] # Prefer direct cache fields if present - cache_creation_tokens = int( - getattr(usage, "cache_creation_input_tokens", 0) or 0 - ) - cache_read_tokens = int( - getattr(usage, "cache_read_input_tokens", 0) or 0 - ) + cache_creation_tokens = int(getattr(usage, "cache_creation_input_tokens", 0) or 0) + cache_read_tokens = int(getattr(usage, "cache_read_input_tokens", 0) or 0) # Check if tool_use was in the output to override stop_reason if response_obj is not None: @@ -337,9 +298,7 @@ class AnthropicResponsesStreamWrapper: except StopAsyncIteration: pass except Exception as e: - verbose_logger.error( - f"AnthropicResponsesStreamWrapper error: {e}\n{traceback.format_exc()}" - ) + verbose_logger.error(f"AnthropicResponsesStreamWrapper error: {e}\n{traceback.format_exc()}") # Drain any remaining queued chunks if self._chunk_queue: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 95ab9b3a0fe..1a052f457c5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -95,17 +95,11 @@ class LiteLLMAnthropicToResponsesAPIAdapter: continue btype = block.get("type") if btype == "text": - user_parts.append( - {"type": "input_text", "text": block.get("text", "")} - ) + user_parts.append({"type": "input_text", "text": block.get("text", "")}) elif btype == "image": - url = self._translate_anthropic_image_source_to_url( - cast(dict, block.get("source", {})) - ) + url = self._translate_anthropic_image_source_to_url(cast(dict, block.get("source", {}))) if url: - user_parts.append( - {"type": "input_image", "image_url": url} - ) + user_parts.append({"type": "input_image", "image_url": url}) elif btype == "tool_result": tool_use_id = block.get("tool_use_id", "") inner = block.get("content") @@ -115,9 +109,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: output_text = inner elif isinstance(inner, list): parts = [ - c.get("text", "") - for c in inner - if isinstance(c, dict) and c.get("type") == "text" + c.get("text", "") for c in inner if isinstance(c, dict) and c.get("type") == "text" ] output_text = "\n".join(parts) else: @@ -155,9 +147,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: continue btype = block.get("type") if btype == "text": - asst_parts.append( - {"type": "output_text", "text": block.get("text", "")} - ) + asst_parts.append({"type": "output_text", "text": block.get("text", "")}) elif btype == "tool_use": # tool_use becomes a top-level function_call item input_items.append( @@ -171,9 +161,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: elif btype == "thinking": thinking_text = block.get("thinking", "") if thinking_text: - asst_parts.append( - {"type": "output_text", "text": thinking_text} - ) + asst_parts.append({"type": "output_text", "text": thinking_text}) if asst_parts: input_items.append( { @@ -196,9 +184,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: tool_type = tool_dict.get("type", "") tool_name = tool_dict.get("name", "") # web_search tool - if ( - isinstance(tool_type, str) and tool_type.startswith("web_search") - ) or tool_name == "web_search": + if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search": result.append({"type": "web_search_preview"}) continue func_tool: Dict[str, Any] = {"type": "function", "name": tool_name} @@ -276,9 +262,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(output_config, dict) and output_config.get("effort"): effort = output_config["effort"] elif thinking_type == "enabled": - effort = reasoning_effort_from_thinking_budget( - thinking.get("budget_tokens", 0) - ) + effort = reasoning_effort_from_thinking_budget(thinking.get("budget_tokens", 0)) else: return None @@ -321,11 +305,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(system, str): responses_kwargs["instructions"] = system elif isinstance(system, list): - text_parts = [ - b.get("text", "") - for b in system - if isinstance(b, dict) and b.get("type") == "text" - ] + text_parts = [b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text"] responses_kwargs["instructions"] = "\n".join(filter(None, text_parts)) # max_tokens -> max_output_tokens @@ -349,10 +329,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: # tool_choice tool_choice = anthropic_request.get("tool_choice") if tool_choice: - responses_kwargs["tool_choice"] = ( - self.translate_tool_choice_to_responses_api( - cast(AnthropicMessagesToolChoice, tool_choice) - ) + responses_kwargs["tool_choice"] = self.translate_tool_choice_to_responses_api( + cast(AnthropicMessagesToolChoice, tool_choice) ) # thinking -> reasoning @@ -373,10 +351,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: output_config = anthropic_request.get("output_config") if not isinstance(output_format, dict) and isinstance(output_config, dict): output_format = output_config.get("format") # type: ignore[assignment] - if ( - isinstance(output_format, dict) - and output_format.get("type") == "json_schema" - ): + if isinstance(output_format, dict) and output_format.get("type") == "json_schema": schema = output_format.get("schema") if schema: responses_kwargs["text"] = { @@ -391,9 +366,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: # context_management: Anthropic dict -> OpenAI array context_management = anthropic_request.get("context_management") if isinstance(context_management, dict): - openai_cm = self.translate_context_management_to_responses_api( - context_management - ) + openai_cm = self.translate_context_management_to_responses_api(context_management) if openai_cm is not None: responses_kwargs["context_management"] = openai_cm @@ -443,9 +416,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: for part in item.content: if getattr(part, "type", None) == "output_text": content.append( - AnthropicResponseContentBlockText( - type="text", text=getattr(part, "text", "") - ).model_dump() + AnthropicResponseContentBlockText(type="text", text=getattr(part, "text", "")).model_dump() ) elif isinstance(item, ResponseFunctionToolCall): @@ -469,9 +440,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: for part in item.get("content", []): if isinstance(part, dict) and part.get("type") == "output_text": content.append( - AnthropicResponseContentBlockText( - type="text", text=part.get("text", "") - ).model_dump() + AnthropicResponseContentBlockText(type="text", text=part.get("text", "")).model_dump() ) elif item_type == "function_call": try: diff --git a/litellm/llms/anthropic/experimental_pass_through/utils.py b/litellm/llms/anthropic/experimental_pass_through/utils.py index 4fd68ef535f..827cce89dab 100644 --- a/litellm/llms/anthropic/experimental_pass_through/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/utils.py @@ -7,10 +7,7 @@ from litellm.types.utils import ModelInfo def is_reasoning_auto_summary_enabled() -> bool: """Check whether the default 'summary: detailed' injection is enabled (opt-in).""" - return ( - litellm.reasoning_auto_summary - or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" - ) + return litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" def normalize_reasoning_effort_value( @@ -34,9 +31,7 @@ def normalize_reasoning_effort_value( model_info: Optional[ModelInfo] = None try: - model_info = get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: model_info = None diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 170cb086bb0..ccd12d1adb1 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -91,9 +91,7 @@ class AnthropicFilesHandler: # Construct the Anthropic batch results URL encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") - results_url = ( - f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}/results" - ) + results_url = f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}/results" # Prepare headers headers = { @@ -108,9 +106,7 @@ class AnthropicFilesHandler: anthropic_response.raise_for_status() # Transform Anthropic batch results to OpenAI format - transformed_content = self._transform_anthropic_batch_results_to_openai_format( - anthropic_response.content - ) + transformed_content = self._transform_anthropic_batch_results_to_openai_format(anthropic_response.content) # Create a new response with transformed content transformed_response = httpx.Response( @@ -131,9 +127,7 @@ class AnthropicFilesHandler: api_key: Optional[str] = None, timeout: Union[float, httpx.Timeout] = 600.0, max_retries: Optional[int] = None, - ) -> Union[ - HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] - ]: + ) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]: """ Retrieve file content from Anthropic. @@ -169,9 +163,7 @@ class AnthropicFilesHandler: ) ) - def _transform_anthropic_batch_results_to_openai_format( - self, anthropic_content: bytes - ) -> bytes: + def _transform_anthropic_batch_results_to_openai_format(self, anthropic_content: bytes) -> bytes: """ Transform Anthropic batch results JSONL to OpenAI batch results JSONL format. @@ -214,11 +206,9 @@ class AnthropicFilesHandler: # Transform Anthropic message to OpenAI format anthropic_message = result.get("message", {}) if anthropic_message: - openai_response_body = ( - self._transform_anthropic_message_to_openai_format( - anthropic_message=anthropic_message, - anthropic_config=anthropic_config, - ) + openai_response_body = self._transform_anthropic_message_to_openai_format( + anthropic_message=anthropic_message, + anthropic_config=anthropic_config, ) # Create OpenAI batch result format @@ -279,9 +269,7 @@ class AnthropicFilesHandler: transformed_content += "\n" # Add trailing newline for JSONL format return transformed_content.encode("utf-8") except Exception as e: - verbose_logger.error( - f"Error transforming Anthropic batch results to OpenAI format: {e}" - ) + verbose_logger.error(f"Error transforming Anthropic batch results to OpenAI format: {e}") # Return original content if transformation fails return anthropic_content @@ -333,9 +321,7 @@ class AnthropicFilesHandler: ) # Convert ModelResponse to OpenAI format dict - it's already in OpenAI format - openai_body: OpenAIChatCompletionResponse = transformed_response.model_dump( - exclude_none=True - ) + openai_body: OpenAIChatCompletionResponse = transformed_response.model_dump(exclude_none=True) # Ensure id comes from anthropic_message if not set if not openai_body.get("id"): @@ -343,9 +329,7 @@ class AnthropicFilesHandler: return openai_body except Exception as e: - verbose_logger.error( - f"Error transforming Anthropic message to OpenAI format: {e}" - ) + verbose_logger.error(f"Error transforming Anthropic message to OpenAI format: {e}") # Return a basic error response if transformation fails error_response: OpenAIChatCompletionResponse = { "id": anthropic_message.get("id", ""), diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index 7ffb6beb4c7..cf12ad9ab32 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -80,9 +80,7 @@ class AnthropicFilesConfig(BaseFilesConfig): return AnthropicError( status_code=status_code, message=error_message, - headers=( - cast(httpx.Headers, headers) if isinstance(headers, dict) else headers - ), + headers=(cast(httpx.Headers, headers) if isinstance(headers, dict) else headers), ) def validate_environment( @@ -111,9 +109,7 @@ class AnthropicFilesConfig(BaseFilesConfig): ) return headers - def get_supported_openai_params( - self, model: str - ) -> List[OpenAICreateFileRequestOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAICreateFileRequestOptionalParams]: return ["purpose"] def map_openai_params( @@ -184,10 +180,7 @@ class AnthropicFilesConfig(BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - api_base = ( - AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) - or ANTHROPIC_FILES_API_BASE - ) + api_base = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} @@ -206,10 +199,7 @@ class AnthropicFilesConfig(BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - api_base = ( - AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) - or ANTHROPIC_FILES_API_BASE - ) + api_base = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} @@ -233,10 +223,7 @@ class AnthropicFilesConfig(BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - api_base = ( - AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) - or ANTHROPIC_FILES_API_BASE - ) + api_base = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE url = f"{api_base.rstrip('/')}/v1/files" params: Dict[str, Any] = {} if purpose: @@ -269,10 +256,7 @@ class AnthropicFilesConfig(BaseFilesConfig): litellm_params: dict, ) -> tuple[str, dict]: file_id = file_content_request.get("file_id") - api_base = ( - AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) - or ANTHROPIC_FILES_API_BASE - ) + api_base = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}/content", {} diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 2bfdf7ef4ba..896182b4763 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -30,9 +30,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.ANTHROPIC - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """Add Anthropic-specific headers""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -45,9 +43,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if auth_header is None: - raise ValueError( - "ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API" - ) + raise ValueError("ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API") headers.update(auth_header) headers["anthropic-version"] = "2023-06-01" @@ -122,9 +118,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): """Transform list skills request for Anthropic""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo - api_base = AnthropicModelInfo.get_api_base( - litellm_params.api_base if litellm_params else None - ) + api_base = AnthropicModelInfo.get_api_base(litellm_params.api_base if litellm_params else None) url = self.get_complete_url(api_base=api_base, endpoint="skills") # Build query parameters @@ -162,9 +156,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform get skill request for Anthropic""" - url = self.get_complete_url( - api_base=api_base, endpoint="skills", skill_id=skill_id - ) + url = self.get_complete_url(api_base=api_base, endpoint="skills", skill_id=skill_id) verbose_logger.debug("Get skill request - URL: %s", url) @@ -189,9 +181,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform delete skill request for Anthropic""" - url = self.get_complete_url( - api_base=api_base, endpoint="skills", skill_id=skill_id - ) + url = self.get_complete_url(api_base=api_base, endpoint="skills", skill_id=skill_id) verbose_logger.debug("Delete skill request - URL: %s", url) diff --git a/litellm/llms/apiserpent/search/defaults.py b/litellm/llms/apiserpent/search/defaults.py index 219178587d6..3bd8e1f93f4 100644 --- a/litellm/llms/apiserpent/search/defaults.py +++ b/litellm/llms/apiserpent/search/defaults.py @@ -44,13 +44,9 @@ class APISerpentSearchParams: # num's deep-search floor (NUM_MIN_DEEP) is endpoint-specific and enforced # in the transform layer; here we only bound the absolute range. if not NUM_MIN <= self.num <= NUM_MAX: - raise ValueError( - f"num must be between {NUM_MIN} and {NUM_MAX}, got {self.num}" - ) + raise ValueError(f"num must be between {NUM_MIN} and {NUM_MAX}, got {self.num}") if self.pages is not None and not PAGES_MIN <= self.pages <= PAGES_MAX: - raise ValueError( - f"pages must be between {PAGES_MIN} and {PAGES_MAX}, got {self.pages}" - ) + raise ValueError(f"pages must be between {PAGES_MIN} and {PAGES_MAX}, got {self.pages}") def to_request_params(self) -> Dict: """Return non-None fields as request params, booleans lowercased.""" diff --git a/litellm/llms/apiserpent/search/transformation.py b/litellm/llms/apiserpent/search/transformation.py index bc11875ba12..637b1472534 100644 --- a/litellm/llms/apiserpent/search/transformation.py +++ b/litellm/llms/apiserpent/search/transformation.py @@ -61,9 +61,7 @@ class APISerpentSearchConfig(BaseSearchConfig): default_api_base=APISERPENT_BASE, ) if not api_key: - raise ValueError( - "APISERPENT_API_KEY is not set. Set `APISERPENT_API_KEY` environment variable." - ) + raise ValueError("APISERPENT_API_KEY is not set. Set `APISERPENT_API_KEY` environment variable.") headers["X-API-Key"] = api_key headers["Content-Type"] = "application/json" return headers @@ -82,14 +80,8 @@ class APISerpentSearchConfig(BaseSearchConfig): changes the host. The ``endswith`` guard keeps this idempotent, since the handler re-invokes this method with the already-resolved URL as api_base. """ - base = ( - api_base or get_secret_str("APISERPENT_API_BASE") or APISERPENT_BASE - ).rstrip("/") - path = ( - DEEP_SEARCH_PATH - if self._is_deep_search(optional_params) - else QUICK_SEARCH_PATH - ) + base = (api_base or get_secret_str("APISERPENT_API_BASE") or APISERPENT_BASE).rstrip("/") + path = DEEP_SEARCH_PATH if self._is_deep_search(optional_params) else QUICK_SEARCH_PATH if not base.endswith(path): base = f"{base}{path}" @@ -125,9 +117,7 @@ class APISerpentSearchConfig(BaseSearchConfig): overrides: Dict = {} if "max_results" in optional_params: num_min = NUM_MIN_DEEP if is_deep else NUM_MIN - overrides["num"] = max( - num_min, min(optional_params["max_results"], NUM_MAX) - ) + overrides["num"] = max(num_min, min(optional_params["max_results"], NUM_MAX)) if "country" in optional_params: overrides["country"] = cast(str, optional_params["country"]).lower() @@ -164,11 +154,7 @@ class APISerpentSearchConfig(BaseSearchConfig): response_json = raw_response.json() raw_results = response_json.get("results") or {} - organic = ( - raw_results.get("organic", []) - if isinstance(raw_results, dict) - else raw_results - ) + organic = raw_results.get("organic", []) if isinstance(raw_results, dict) else raw_results results: List[SearchResult] = [] for result in organic: diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index caf65770397..c85bc9c4032 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -92,9 +92,9 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): base_llm_http_handler: The BaseLLMHTTPHandler instance from main.py """ # Get AWS region from kwargs or environment - aws_region_name = kwargs.get( - "aws_region_name" - ) or self._get_aws_region_name_for_polly(optional_params=optional_params) + aws_region_name = kwargs.get("aws_region_name") or self._get_aws_region_name_for_polly( + optional_params=optional_params + ) # Convert voice to string if it's a dict voice_str: Optional[str] = None @@ -263,9 +263,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError( - "Missing boto3 to call AWS Polly. Run 'pip install boto3'." - ) + raise ImportError("Missing boto3 to call AWS Polly. Run 'pip install boto3'.") # Get AWS region aws_region_name = litellm_params.get("aws_region_name", self.DEFAULT_REGION) diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 750087b722e..08a04d0c8c7 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -203,11 +203,9 @@ class AzureAssistantsAPI(BaseAzureLLM): litellm_params=litellm_params, ) - thread_message: OpenAIMessage = ( - await openai_client.beta.threads.messages.create( # type: ignore - thread_id, - **message_data, # type: ignore - ) + thread_message: OpenAIMessage = await openai_client.beta.threads.messages.create( # type: ignore + thread_id, + **message_data, # type: ignore ) response_obj: Optional[OpenAIMessage] = None @@ -916,9 +914,7 @@ class AzureAssistantsAPI(BaseAzureLLM): litellm_params=litellm_params, ) - response = await azure_openai_client.beta.assistants.create( - **create_assistant_data - ) + response = await azure_openai_client.beta.assistants.create(**create_assistant_data) return response def create_assistants( @@ -984,9 +980,7 @@ class AzureAssistantsAPI(BaseAzureLLM): litellm_params=litellm_params, ) - response = await azure_openai_client.beta.assistants.delete( - assistant_id=assistant_id - ) + response = await azure_openai_client.beta.assistants.delete(assistant_id=assistant_id) return response def delete_assistant( diff --git a/litellm/llms/azure/audio_transcription/transformation.py b/litellm/llms/azure/audio_transcription/transformation.py index e478c8ebf35..77050ce6bca 100644 --- a/litellm/llms/azure/audio_transcription/transformation.py +++ b/litellm/llms/azure/audio_transcription/transformation.py @@ -41,9 +41,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): STT_ENDPOINT_PATH = "/speech/recognition/conversation/cognitiveservices/v1" DEFAULT_LANGUAGE = "en-US" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: return ["language", "response_format"] def map_openai_params( @@ -78,9 +76,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): validated_headers = headers.copy() validated_headers["Ocp-Apim-Subscription-Key"] = api_key - validated_headers["Content-Type"] = validated_headers.get( - "Content-Type", "audio/wav" - ) + validated_headers["Content-Type"] = validated_headers.get("Content-Type", "audio/wav") validated_headers["Accept"] = "application/json" return validated_headers @@ -108,9 +104,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): base_url = self._resolve_stt_base_url(api_base=api_base) query_params = { "language": optional_params.get("language", self.DEFAULT_LANGUAGE), - "format": self._get_azure_response_format( - optional_params.get("response_format") - ), + "format": self._get_azure_response_format(optional_params.get("response_format")), } return f"{base_url}{self.STT_ENDPOINT_PATH}?{urlencode(query_params)}" @@ -136,10 +130,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): recognition_status = response_json.get("RecognitionStatus") if recognition_status is not None and recognition_status != "Success": raise AzureSpeechAudioTranscriptionException( - message=( - "Azure AI Speech transcription failed with " - f"RecognitionStatus={recognition_status}." - ), + message=(f"Azure AI Speech transcription failed with RecognitionStatus={recognition_status}."), status_code=raw_response.status_code, headers=raw_response.headers, ) @@ -164,9 +155,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): hostname = parsed_url.hostname or "" if self._is_cognitive_services_endpoint(hostname=hostname): - region = self._extract_region_from_hostname( - hostname=hostname, domain=self.COGNITIVE_SERVICES_DOMAIN - ) + region = self._extract_region_from_hostname(hostname=hostname, domain=self.COGNITIVE_SERVICES_DOMAIN) return self._build_stt_base_url(region=region) if self._is_stt_endpoint(hostname=hostname): @@ -184,14 +173,10 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return api_base def _is_cognitive_services_endpoint(self, hostname: str) -> bool: - return hostname == self.COGNITIVE_SERVICES_DOMAIN or hostname.endswith( - f".{self.COGNITIVE_SERVICES_DOMAIN}" - ) + return hostname == self.COGNITIVE_SERVICES_DOMAIN or hostname.endswith(f".{self.COGNITIVE_SERVICES_DOMAIN}") def _is_stt_endpoint(self, hostname: str) -> bool: - return hostname == self.STT_SPEECH_DOMAIN or hostname.endswith( - f".{self.STT_SPEECH_DOMAIN}" - ) + return hostname == self.STT_SPEECH_DOMAIN or hostname.endswith(f".{self.STT_SPEECH_DOMAIN}") def _is_azure_openai_endpoint(self, hostname: str) -> bool: return hostname.endswith(".openai.azure.com") diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index 591fe9d03a2..a39f86fd5b1 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -141,19 +141,15 @@ class AzureAudioTranscription(AzureChatCompletion): input=f"audio_file_{uuid.uuid4()}", api_key=async_azure_client.api_key, additional_args={ - "headers": { - "Authorization": f"Bearer {async_azure_client.api_key}" - }, + "headers": {"Authorization": f"Bearer {async_azure_client.api_key}"}, "api_base": async_azure_client._base_url._uri_reference, "atranscription": True, "complete_input_dict": data, }, ) - raw_response = ( - await async_azure_client.audio.transcriptions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = await async_azure_client.audio.transcriptions.with_raw_response.create( + **data, timeout=timeout ) # type: ignore headers = dict(raw_response.headers) @@ -171,9 +167,7 @@ class AzureAudioTranscription(AzureChatCompletion): input=get_audio_file_name(audio_file), api_key=api_key, additional_args={ - "headers": { - "Authorization": f"Bearer {async_azure_client.api_key}" - }, + "headers": {"Authorization": f"Bearer {async_azure_client.api_key}"}, "api_base": async_azure_client._base_url._uri_reference, "atranscription": True, "complete_input_dict": data, diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 8edab95ef6c..ccb9eb8f5c8 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -68,9 +68,7 @@ class AzureOpenAIAssistantsAPIConfig: "metadata", ] - def map_openai_params_create_message_params( - self, non_default_params: dict, optional_params: dict - ): + def map_openai_params_create_message_params(self, non_default_params: dict, optional_params: dict): for param, value in non_default_params.items(): if param == "role": optional_params["role"] = value @@ -84,9 +82,7 @@ class AzureOpenAIAssistantsAPIConfig: message="Azure only accepts content as a string.", status_code=400, ) - elif ( - param == "attachments" - ): # this is a v2 param. Azure currently supports the old 'file_id's param + elif param == "attachments": # this is a v2 param. Azure currently supports the old 'file_id's param file_ids: List[str] = [] if isinstance(value, list): for item in value: @@ -149,9 +145,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): - call chat.completions.create by default """ try: - raw_response = azure_client.chat.completions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = azure_client.chat.completions.with_raw_response.create(**data, timeout=timeout) headers = dict(raw_response.headers) response = raw_response.parse() @@ -174,9 +168,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): """ start_time = time.time() try: - raw_response = await azure_client.chat.completions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = await azure_client.chat.completions.with_raw_response.create(**data, timeout=timeout) headers = dict(raw_response.headers) response = raw_response.parse() @@ -215,9 +207,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): optional_params["extra_headers"] = headers try: if model is None or messages is None: - raise AzureOpenAIError( - status_code=422, message="Missing model or messages" - ) + raise AzureOpenAIError(status_code=422, message="Missing model or messages") max_retries = optional_params.pop("max_retries", None) if max_retries is None: @@ -242,9 +232,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) data = {"model": None, "messages": messages, **optional_params} - elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( - model=litellm_params.get("base_model") or model - ): + elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=litellm_params.get("base_model") or model): data = litellm.AzureOpenAIGPT5Config().transform_request( model=model, messages=messages, @@ -328,9 +316,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): }, ) if not isinstance(max_retries, int): - raise AzureOpenAIError( - status_code=422, message="max retries must be an int" - ) + raise AzureOpenAIError(status_code=422, message="max retries must be an int") # init AzureOpenAI Client azure_client = self.get_azure_openai_client( api_version=api_version, @@ -420,9 +406,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=litellm_params, ) if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): - raise ValueError( - "Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI" - ) + raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") ## LOGGING logging_obj.pre_call( input=data["messages"], @@ -524,9 +508,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "max_retries": max_retries, "timeout": timeout, } - azure_client_params = select_azure_base_url_or_endpoint( - azure_client_params=azure_client_params - ) + azure_client_params = select_azure_base_url_or_endpoint(azure_client_params=azure_client_params) if api_key is not None: azure_client_params["api_key"] = api_key elif azure_ad_token is not None: @@ -604,9 +586,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=litellm_params, ) if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): - raise ValueError( - "Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI" - ) + raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") ## LOGGING logging_obj.pre_call( @@ -685,13 +665,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=litellm_params, ) if not isinstance(openai_aclient, (AsyncAzureOpenAI, AsyncOpenAI)): - raise ValueError( - "Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI" - ) + raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") - raw_response = await openai_aclient.embeddings.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) headers = dict(raw_response.headers) # Convert json.JSONDecodeError to AzureOpenAIError for two critical reasons: @@ -817,9 +793,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) ## COMPLETION CALL - raw_response = azure_client.embeddings.with_raw_response.create( - **data, timeout=timeout - ) # type: ignore + raw_response = azure_client.embeddings.with_raw_response.create(**data, timeout=timeout) # type: ignore headers = dict(raw_response.headers) response = raw_response.parse() if isinstance(response, str): @@ -851,9 +825,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) error_text = error_response.text - raise AzureOpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise AzureOpenAIError(status_code=status_code, message=error_text, headers=error_headers) async def make_async_azure_httpx_request( self, @@ -897,9 +869,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "2023-10-01-preview", ] ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url( - original_url=api_base, new_path="/openai/images/generations:submit" - ) + api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") data.pop( "model", None @@ -944,9 +914,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) while response.json()["status"] not in ["succeeded", "failed"]: if time.time() - start_time > timeout_secs: - raise AzureOpenAIError( - status_code=408, message="Operation polling timed out." - ) + raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") await asyncio.sleep(int(response.headers.get("retry-after") or 10)) response = await async_handler.get( @@ -1025,9 +993,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "2023-10-01-preview", ] ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url( - original_url=api_base, new_path="/openai/images/generations:submit" - ) + api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") data.pop( "model", None @@ -1064,9 +1030,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) while response.json()["status"] not in ["succeeded", "failed"]: if time.time() - start_time > timeout_secs: - raise AzureOpenAIError( - status_code=408, message="Operation polling timed out." - ) + raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") time.sleep(int(response.headers.get("retry-after") or 10)) response = sync_handler.get( @@ -1117,9 +1081,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): AzureFoundryMAIImageGenerationConfig, ) - api_base: str = azure_client_params.get( - "azure_endpoint", "" - ) # "https://example-endpoint.openai.azure.com" + api_base: str = azure_client_params.get("azure_endpoint", "") # "https://example-endpoint.openai.azure.com" if api_base.endswith("/"): api_base = api_base.rstrip("/") api_version: str = azure_client_params.get("api_version", "") @@ -1167,9 +1129,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): response: Optional[dict] = None try: # response = await azure_client.images.generate(**data, timeout=timeout) - api_base: str = azure_client_params.get( - "api_base", "" - ) # "https://example-endpoint.openai.azure.com" + api_base: str = azure_client_params.get("api_base", "") # "https://example-endpoint.openai.azure.com" if api_base.endswith("/"): api_base = api_base.rstrip("/") api_version: str = azure_client_params.get("api_version", "") @@ -1199,9 +1159,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers=headers, ) - provider_config = get_azure_image_generation_config( - data.get("model", "dall-e-2") - ) + provider_config = get_azure_image_generation_config(data.get("model", "dall-e-2")) if provider_config is not None: return provider_config.transform_image_generation_response( model=data.get("model", "dall-e-2"), @@ -1269,23 +1227,17 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): and litellm_params is not None and litellm_params.get("base_model", None) is not None ): - model_response._hidden_params["model"] = litellm_params.get( - "base_model", None - ) + model_response._hidden_params["model"] = litellm_params.get("base_model", None) # Azure image generation API doesn't support extra_body parameter extra_body = optional_params.pop("extra_body", {}) flattened_params = {**optional_params, **extra_body} - base_model = ( - litellm_params.get("base_model", None) if litellm_params else None - ) + base_model = litellm_params.get("base_model", None) if litellm_params else None data = {"model": base_model or model, "prompt": prompt, **flattened_params} max_retries = data.pop("max_retries", 2) if not isinstance(max_retries, int): - raise AzureOpenAIError( - status_code=422, message="max retries must be an int" - ) + raise AzureOpenAIError(status_code=422, message="max retries must be an int") if api_key is None and azure_ad_token_provider is not None: azure_ad_token = azure_ad_token_provider() @@ -1341,9 +1293,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): data=data, headers=headers, ) - provider_config = get_azure_image_generation_config( - data.get("model", "dall-e-2") - ) + provider_config = get_azure_image_generation_config(data.get("model", "dall-e-2")) if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig): return provider_config.transform_image_generation_response( model=data.get("model", "dall-e-2"), @@ -1529,14 +1479,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if ( completion.headers.get("x-ratelimit-remaining-requests", None) is not None ): # not provided for dall-e requests - response["x-ratelimit-remaining-requests"] = completion.headers[ - "x-ratelimit-remaining-requests" - ] + response["x-ratelimit-remaining-requests"] = completion.headers["x-ratelimit-remaining-requests"] if completion.headers.get("x-ratelimit-remaining-tokens", None) is not None: - response["x-ratelimit-remaining-tokens"] = completion.headers[ - "x-ratelimit-remaining-tokens" - ] + response["x-ratelimit-remaining-tokens"] = completion.headers["x-ratelimit-remaining-tokens"] if completion.headers.get("x-ms-region", None) is not None: response["x-ms-region"] = completion.headers["x-ms-region"] diff --git a/litellm/llms/azure/batches/handler.py b/litellm/llms/azure/batches/handler.py index ea6722839e6..808fb3d9600 100644 --- a/litellm/llms/azure/batches/handler.py +++ b/litellm/llms/azure/batches/handler.py @@ -47,20 +47,18 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: - azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, + ) ) if azure_client is None: raise ValueError( @@ -75,9 +73,7 @@ class AzureBatchesAPI(BaseAzureLLM): return self.acreate_batch( # type: ignore create_batch_data=create_batch_data, azure_client=azure_client ) - response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.create( - **create_batch_data - ) # type: ignore[arg-type] + response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.create(**create_batch_data) # type: ignore[arg-type] return LiteLLMBatch(**response.model_dump()) async def aretrieve_batch( @@ -97,20 +93,18 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): - azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, + ) ) if azure_client is None: raise ValueError( @@ -125,9 +119,7 @@ class AzureBatchesAPI(BaseAzureLLM): return self.aretrieve_batch( # type: ignore retrieve_batch_data=retrieve_batch_data, client=azure_client ) - response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.retrieve( - **retrieve_batch_data - ) + response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.retrieve(**retrieve_batch_data) return LiteLLMBatch(**response.model_dump()) async def acancel_batch( @@ -147,20 +139,18 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): - azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, + ) ) if azure_client is None: raise ValueError( @@ -203,20 +193,18 @@ class AzureBatchesAPI(BaseAzureLLM): max_retries: Optional[int], after: Optional[str] = None, limit: Optional[int] = None, - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): - azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, - litellm_params=litellm_params or {}, + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + litellm_params=litellm_params or {}, + ) ) if azure_client is None: raise ValueError( diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index e94f50380c0..f1bfd96de94 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -54,9 +54,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): # 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 = 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 ("gpt-5" in model and not _normalized.startswith("gpt-5-chat")) or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> List[str]: """Get supported parameters for Azure OpenAI GPT-5 models. @@ -79,9 +77,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): # Only gpt-5.2+ has been verified to support logprobs on Azure. # The base OpenAI class includes logprobs for gpt-5.1+, but Azure # hasn't verified support for gpt-5.1, so remove them unless gpt-5.2/5.4+. - if self._supports_reasoning_effort_level( - model, "none" - ) and not self.is_model_gpt_5_2_model(model): + if self._supports_reasoning_effort_level(model, "none") and not self.is_model_gpt_5_2_model(model): params = [p for p in params if p not in ["logprobs", "top_logprobs"]] elif self.is_model_gpt_5_2_model(model): azure_supported_params = ["logprobs", "top_logprobs"] @@ -97,9 +93,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): drop_params: bool, api_version: str = "", ) -> dict: - reasoning_effort_value = non_default_params.get( - "reasoning_effort" - ) or optional_params.get("reasoning_effort") + reasoning_effort_value = non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort") effective_effort = _get_effort_level(reasoning_effort_value) # gpt-5.1/5.2/5.4 support reasoning_effort='none', but other gpt-5 models don't @@ -107,15 +101,10 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): supports_none = self._supports_reasoning_effort_level(model, "none") if effective_effort == "none" and not supports_none: - if litellm.drop_params is True or ( - drop_params is not None and drop_params is True - ): + if litellm.drop_params is True or (drop_params is not None and drop_params is True): non_default_params = non_default_params.copy() optional_params = optional_params.copy() - if ( - _get_effort_level(non_default_params.get("reasoning_effort")) - == "none" - ): + if _get_effort_level(non_default_params.get("reasoning_effort")) == "none": non_default_params.pop("reasoning_effort") if _get_effort_level(optional_params.get("reasoning_effort")) == "none": optional_params.pop("reasoning_effort") diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 69eda95be1b..50b3ba16326 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -128,9 +128,7 @@ class AzureOpenAIConfig(BaseConfig): return True - def _is_response_format_supported_api_version( - self, api_version_year: str, api_version_month: str - ) -> bool: + def _is_response_format_supported_api_version(self, api_version_year: str, api_version_month: str) -> bool: """ - check if api_version is supported for response_format - returns True if the API version is equal to or newer than the supported version @@ -178,25 +176,15 @@ class AzureOpenAIConfig(BaseConfig): tool_choice='required' is not supported as of 2024-05-01-preview """ ## check if api version supports this param ## - if ( - api_version_year is None - or api_version_month is None - or api_version_day is None - ): + if api_version_year is None or api_version_month is None or api_version_day is None: optional_params["tool_choice"] = value else: if ( api_version_year < "2023" or (api_version_year == "2023" and api_version_month < "12") - or ( - api_version_year == "2023" - and api_version_month == "12" - and api_version_day < "01" - ) + or (api_version_year == "2023" and api_version_month == "12" and api_version_day < "01") ): - if litellm.drop_params is True or ( - drop_params is not None and drop_params is True - ): + if litellm.drop_params is True or (drop_params is not None and drop_params is True): pass else: raise UnsupportedParamsError( @@ -206,9 +194,7 @@ class AzureOpenAIConfig(BaseConfig): elif value == "required" and ( api_version_year == "2024" and api_version_month <= "05" ): ## check if tool_choice value is supported ## - if litellm.drop_params is True or ( - drop_params is not None and drop_params is True - ): + if litellm.drop_params is True or (drop_params is not None and drop_params is True): pass else: raise UnsupportedParamsError( @@ -218,21 +204,16 @@ class AzureOpenAIConfig(BaseConfig): else: optional_params["tool_choice"] = value elif param == "response_format" and isinstance(value, dict): - _is_response_format_supported_model = ( - self._is_response_format_supported_model(model) - ) + _is_response_format_supported_model = self._is_response_format_supported_model(model) if api_version_year is None or api_version_month is None: is_response_format_supported_api_version = True else: - is_response_format_supported_api_version = ( - self._is_response_format_supported_api_version( - api_version_year, api_version_month - ) + is_response_format_supported_api_version = self._is_response_format_supported_api_version( + api_version_year, api_version_month ) is_response_format_supported = ( - is_response_format_supported_api_version - and _is_response_format_supported_model + is_response_format_supported_api_version and _is_response_format_supported_model ) optional_params = self._add_response_format_to_tools( @@ -312,12 +293,8 @@ class AzureOpenAIConfig(BaseConfig): "westus4", ] - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return AzureOpenAIError( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return AzureOpenAIError(message=error_message, status_code=status_code, headers=headers) def validate_environment( self, diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index 0a73597a4e4..b9cf77b89d8 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -27,9 +27,7 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): """ Get the supported OpenAI params for the Azure O-Series models """ - all_openai_params = litellm.OpenAIGPTConfig().get_supported_openai_params( - model=model - ) + all_openai_params = litellm.OpenAIGPTConfig().get_supported_openai_params(model=model) non_supported_params = [ "logprobs", "top_p", @@ -41,9 +39,7 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): o_series_only_param = self._get_o_series_only_params(model) all_openai_params.extend(o_series_only_param) - return [ - param for param in all_openai_params if param not in non_supported_params - ] + return [param for param in all_openai_params if param not in non_supported_params] def _get_o_series_only_params(self, model: str) -> list: """ @@ -83,9 +79,7 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): if stream is not True: return False - if ( - model and "o3" in model - ): # o3 models support streaming - https://github.com/BerriAI/litellm/issues/8274 + if model and "o3" in model: # o3 models support streaming - https://github.com/BerriAI/litellm/issues/8274 return False if model is not None: @@ -99,9 +93,7 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): ): # allow user to override default with model_info={"supports_native_streaming": true} return False except Exception as e: - verbose_logger.debug( - f"Error getting model info in AzureOpenAIO1Config: {e}" - ) + verbose_logger.debug(f"Error getting model info in AzureOpenAIO1Config: {e}") return True def is_o_series_model(self, model: str) -> bool: @@ -115,9 +107,5 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): litellm_params: dict, headers: dict, ) -> dict: - model = model.replace( - "o_series/", "" - ) # handle o_series/my-random-deployment-name - return super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name + return super().transform_request(model, messages, optional_params, litellm_params, headers) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index e1ac1858912..91f5793e269 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -45,22 +45,14 @@ class AzureOpenAIError(BaseLLMException): def process_azure_headers(headers: Union[httpx.Headers, dict]) -> dict: openai_headers = {} if "x-ratelimit-limit-requests" in headers: - openai_headers["x-ratelimit-limit-requests"] = headers[ - "x-ratelimit-limit-requests" - ] + openai_headers["x-ratelimit-limit-requests"] = headers["x-ratelimit-limit-requests"] if "x-ratelimit-remaining-requests" in headers: - openai_headers["x-ratelimit-remaining-requests"] = headers[ - "x-ratelimit-remaining-requests" - ] + openai_headers["x-ratelimit-remaining-requests"] = headers["x-ratelimit-remaining-requests"] if "x-ratelimit-limit-tokens" in headers: openai_headers["x-ratelimit-limit-tokens"] = headers["x-ratelimit-limit-tokens"] if "x-ratelimit-remaining-tokens" in headers: - openai_headers["x-ratelimit-remaining-tokens"] = headers[ - "x-ratelimit-remaining-tokens" - ] - llm_response_headers = { - "{}-{}".format("llm_provider", k): v for k, v in headers.items() - } + openai_headers["x-ratelimit-remaining-tokens"] = headers["x-ratelimit-remaining-tokens"] + llm_response_headers = {"{}-{}".format("llm_provider", k): v for k, v in headers.items()} return {**llm_response_headers, **openai_headers} @@ -178,9 +170,7 @@ def get_azure_ad_token_from_oidc( """ if scope is None: scope = "https://cognitiveservices.azure.com/.default" - azure_authority_host = os.getenv( - "AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com" - ) + azure_authority_host = os.getenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com") azure_client_id = azure_client_id or os.getenv("AZURE_CLIENT_ID") azure_tenant_id = azure_tenant_id or os.getenv("AZURE_TENANT_ID") if azure_client_id is None or azure_tenant_id is None: @@ -234,14 +224,10 @@ def get_azure_ad_token_from_oidc( azure_ad_token_expires_in = azure_ad_token_json.get("expires_in", None) if azure_ad_token_access_token is None: - raise AzureOpenAIError( - status_code=422, message="Azure AD Token access_token not returned" - ) + raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned") if azure_ad_token_expires_in is None: - raise AzureOpenAIError( - status_code=422, message="Azure AD Token expires_in not returned" - ) + raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") azure_ad_cache.set_cache( key=azure_ad_token_cache_key, @@ -294,14 +280,10 @@ def get_azure_ad_token( # Extract parameters # Use `or` instead of default parameter to handle cases where key exists but value is None azure_ad_token_provider = litellm_params.get("azure_ad_token_provider") - azure_ad_token = litellm_params.get("azure_ad_token") or get_secret_str( - "AZURE_AD_TOKEN" - ) + azure_ad_token = litellm_params.get("azure_ad_token") or get_secret_str("AZURE_AD_TOKEN") tenant_id = litellm_params.get("tenant_id") or os.getenv("AZURE_TENANT_ID") client_id = litellm_params.get("client_id") or os.getenv("AZURE_CLIENT_ID") - client_secret = litellm_params.get("client_secret") or os.getenv( - "AZURE_CLIENT_SECRET" - ) + client_secret = litellm_params.get("client_secret") or os.getenv("AZURE_CLIENT_SECRET") azure_username = litellm_params.get("azure_username") or os.getenv("AZURE_USERNAME") azure_password = litellm_params.get("azure_password") or os.getenv("AZURE_PASSWORD") scope = litellm_params.get("azure_scope") or os.getenv( @@ -312,9 +294,7 @@ def get_azure_ad_token( # Try to get token provider from Entra ID if azure_ad_token_provider is None and tenant_id and client_id and client_secret: - verbose_logger.debug( - "Using Azure AD Token Provider from Entra ID for Azure Auth" - ) + verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth") azure_ad_token_provider = get_azure_ad_token_from_entra_id( tenant_id=tenant_id, client_id=client_id, @@ -323,12 +303,7 @@ def get_azure_ad_token( ) # Try to get token provider from username and password - if ( - azure_ad_token_provider is None - and azure_username - and azure_password - and client_id - ): + if azure_ad_token_provider is None and azure_username and azure_password and client_id: verbose_logger.debug("Using Azure Username and Password for Azure Auth") azure_ad_token_provider = get_azure_ad_token_from_username_password( azure_username=azure_username, @@ -338,12 +313,7 @@ def get_azure_ad_token( ) # Try to get token from OIDC - if ( - client_id - and tenant_id - and azure_ad_token - and azure_ad_token.startswith("oidc/") - ): + if client_id and tenant_id and azure_ad_token and azure_ad_token.startswith("oidc/"): verbose_logger.debug("Using Azure OIDC Token for Azure Auth") azure_ad_token = get_azure_ad_token_from_oidc( azure_ad_token=azure_ad_token, @@ -352,10 +322,7 @@ def get_azure_ad_token( scope=scope, ) # Try to get token provider from service principal or DefaultAzureCredential - elif ( - azure_ad_token_provider is None - and litellm.enable_azure_ad_token_refresh is True - ): + elif azure_ad_token_provider is None and litellm.enable_azure_ad_token_refresh is True: verbose_logger.debug( "Using Azure AD token provider based on Service Principal with Secret workflow or DefaultAzureCredential for Azure Auth" ) @@ -374,10 +341,8 @@ def get_azure_ad_token( # try to get DefaultAzureCredential provider ######################################################### if azure_ad_token_provider is None and azure_ad_token is None: - azure_ad_token_provider = ( - BaseAzureLLM._try_get_default_azure_credential_provider( - scope=scope, - ) + azure_ad_token_provider = BaseAzureLLM._try_get_default_azure_credential_provider( + scope=scope, ) # Execute the token provider to get the token if available @@ -385,9 +350,7 @@ def get_azure_ad_token( try: token = azure_ad_token_provider() if not isinstance(token, str): - verbose_logger.error( - f"Azure AD token provider returned non-string value: {type(token)}" - ) + verbose_logger.error(f"Azure AD token provider returned non-string value: {type(token)}") raise TypeError(f"Azure AD token must be a string, got {type(token)}") else: azure_ad_token = token @@ -426,9 +389,7 @@ class BaseAzureLLM(BaseOpenAILLM): azure_scope=scope, azure_credential=AzureCredentialType.DefaultAzureCredential, ) - verbose_logger.debug( - "Successfully obtained Azure AD token provider using DefaultAzureCredential" - ) + verbose_logger.debug("Successfully obtained Azure AD token provider using DefaultAzureCredential") return azure_ad_token_provider except Exception as e: verbose_logger.debug(f"DefaultAzureCredential failed: {str(e)}") @@ -439,16 +400,12 @@ class BaseAzureLLM(BaseOpenAILLM): api_key: Optional[str], api_base: Optional[str], api_version: Optional[str] = None, - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, _is_async: bool = False, model: Optional[str] = None, ) -> Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]]: - openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None client_initialization_params: dict = locals() client_initialization_params["is_async"] = _is_async _lp = litellm_params or {} @@ -457,9 +414,7 @@ class BaseAzureLLM(BaseOpenAILLM): _client_secret = _lp.get("client_secret") _azure_password = _lp.get("azure_password") client_initialization_params["azure_ad_token"] = ( - hashlib.sha256(_ad_token.encode()).hexdigest() - if isinstance(_ad_token, str) - else None + hashlib.sha256(_ad_token.encode()).hexdigest() if isinstance(_ad_token, str) else None ) client_initialization_params["azure_ad_token_provider"] = ( f"provider_id={id(_ad_provider) if callable(_ad_provider) else None}" @@ -476,9 +431,7 @@ class BaseAzureLLM(BaseOpenAILLM): client_type="azure", ) if cached_client: - if isinstance( - cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI) - ): + if isinstance(cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI)): return cached_client azure_client_params = self.initialize_azure_sdk_client( @@ -527,9 +480,7 @@ class BaseAzureLLM(BaseOpenAILLM): if "http_client" in azure_client_params: v1_params["http_client"] = azure_client_params["http_client"] - verbose_logger.debug( - f"Using Azure v1 API with base_url: {v1_params['base_url']}" - ) + verbose_logger.debug(f"Using Azure v1 API with base_url: {v1_params['base_url']}") if _is_async is True: openai_client = AsyncOpenAI(**v1_params) # type: ignore @@ -574,49 +525,26 @@ class BaseAzureLLM(BaseOpenAILLM): # litellm_params sometimes contains the key, but the value is None # We should respect environment variables in this case - tenant_id = self._resolve_env_var( - litellm_params, "tenant_id", "AZURE_TENANT_ID" - ) - client_id = self._resolve_env_var( - litellm_params, "client_id", "AZURE_CLIENT_ID" - ) - client_secret = self._resolve_env_var( - litellm_params, "client_secret", "AZURE_CLIENT_SECRET" - ) - azure_username = self._resolve_env_var( - litellm_params, "azure_username", "AZURE_USERNAME" - ) - azure_password = self._resolve_env_var( - litellm_params, "azure_password", "AZURE_PASSWORD" - ) + tenant_id = self._resolve_env_var(litellm_params, "tenant_id", "AZURE_TENANT_ID") + client_id = self._resolve_env_var(litellm_params, "client_id", "AZURE_CLIENT_ID") + client_secret = self._resolve_env_var(litellm_params, "client_secret", "AZURE_CLIENT_SECRET") + azure_username = self._resolve_env_var(litellm_params, "azure_username", "AZURE_USERNAME") + azure_password = self._resolve_env_var(litellm_params, "azure_password", "AZURE_PASSWORD") scope = self._resolve_env_var(litellm_params, "azure_scope", "AZURE_SCOPE") if scope is None: scope = "https://cognitiveservices.azure.com/.default" max_retries = litellm_params.get("max_retries") timeout = litellm_params.get("timeout") - if ( - not api_key - and azure_ad_token_provider is None - and tenant_id - and client_id - and client_secret - ): - verbose_logger.debug( - "Using Azure AD Token Provider from Entra ID for Azure Auth" - ) + if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret: + verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth") azure_ad_token_provider = get_azure_ad_token_from_entra_id( tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope=scope, ) - if ( - azure_ad_token_provider is None - and azure_username - and azure_password - and client_id - ): + if azure_ad_token_provider is None and azure_username and azure_password and client_id: verbose_logger.debug("Using Azure Username and Password for Azure Auth") azure_ad_token_provider = get_azure_ad_token_from_username_password( azure_username=azure_username, @@ -633,11 +561,7 @@ class BaseAzureLLM(BaseOpenAILLM): azure_tenant_id=tenant_id, scope=scope, ) - elif ( - not api_key - and azure_ad_token_provider is None - and litellm.enable_azure_ad_token_refresh is True - ): + elif not api_key and azure_ad_token_provider is None and litellm.enable_azure_ad_token_refresh is True: verbose_logger.debug( "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth" ) @@ -648,9 +572,7 @@ class BaseAzureLLM(BaseOpenAILLM): except ValueError: verbose_logger.debug("Azure AD Token Provider could not be used.") if api_version is None: - api_version = os.getenv( - "AZURE_API_VERSION", litellm.AZURE_DEFAULT_API_VERSION - ) + api_version = os.getenv("AZURE_API_VERSION", litellm.AZURE_DEFAULT_API_VERSION) _api_key = api_key if _api_key is not None and isinstance(_api_key, str): @@ -682,9 +604,7 @@ class BaseAzureLLM(BaseOpenAILLM): # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client # required to support GPT-4 vision enhancements, since base_url needs to be set on Azure OpenAI Client - azure_client_params = select_azure_base_url_or_endpoint( - azure_client_params=azure_client_params - ) + azure_client_params = select_azure_base_url_or_endpoint(azure_client_params=azure_client_params) return azure_client_params @@ -743,9 +663,7 @@ class BaseAzureLLM(BaseOpenAILLM): return client @staticmethod - def _base_validate_azure_environment( - headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def _base_validate_azure_environment(headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() # Check if api-key is already in headers; if so, use it @@ -798,10 +716,7 @@ class BaseAzureLLM(BaseOpenAILLM): # Extract api_version or use default litellm_params = litellm_params or {} - api_version = ( - cast(Optional[str], litellm_params.get("api_version")) - or default_api_version - ) + api_version = cast(Optional[str], litellm_params.get("api_version")) or default_api_version # Create a new dictionary with existing params query_params = dict(original_url.params) @@ -820,11 +735,7 @@ class BaseAzureLLM(BaseOpenAILLM): # ensure the request go to /openai/v1 and not just /openai if "/openai/v1" not in new_url: parsed_url = httpx.URL(new_url) - new_url = str( - parsed_url.copy_with( - path=parsed_url.path.replace("/openai", "/openai/v1") - ) - ) + new_url = str(parsed_url.copy_with(path=parsed_url.path.replace("/openai", "/openai/v1"))) # Use the new query_params dictionary final_url = httpx.URL(new_url).copy_with(params=query_params) @@ -837,9 +748,7 @@ class BaseAzureLLM(BaseOpenAILLM): return False return api_version in {"preview", "latest", "v1"} - def _resolve_env_var( - self, litellm_params: Dict[str, Any], param_key: str, env_var_key: str - ) -> Optional[str]: + def _resolve_env_var(self, litellm_params: Dict[str, Any], param_key: str, env_var_key: str) -> Optional[str]: """Resolve the environment variable for a given parameter key. The logic here is different from `params.get(key, os.getenv(env_var))` because @@ -865,9 +774,7 @@ def get_azure_credentials( ) -> AzureCredentials: """Resolve Azure credentials from params, litellm globals, and env vars.""" resolved_api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") - resolved_api_version = ( - api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") - ) + resolved_api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") resolved_api_key = ( api_key or litellm.api_key diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index b8d1ad71d46..2c0b67a9e56 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -48,14 +48,10 @@ class AzureTextCompletion(BaseAzureLLM): ): try: if model is None or messages is None: - raise AzureOpenAIError( - status_code=422, message="Missing model or messages" - ) + raise AzureOpenAIError(status_code=422, message="Missing model or messages") max_retries = optional_params.pop("max_retries", 2) - prompt = prompt_factory( - messages=messages, model=model, custom_llm_provider="azure_text" - ) + prompt = prompt_factory(messages=messages, model=model, custom_llm_provider="azure_text") ### CHECK IF CLOUDFLARE AI GATEWAY ### ### if so - set the model as part of the base url @@ -140,9 +136,7 @@ class AzureTextCompletion(BaseAzureLLM): }, ) if not isinstance(max_retries, int): - raise AzureOpenAIError( - status_code=422, message="max retries must be an int" - ) + raise AzureOpenAIError(status_code=422, message="max retries must be an int") # init AzureOpenAI Client azure_client = self.get_azure_openai_client( api_key=api_key, @@ -160,9 +154,7 @@ class AzureTextCompletion(BaseAzureLLM): message="azure_client is not an instance of AzureOpenAI", ) - raw_response = azure_client.completions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = azure_client.completions.with_raw_response.create(**data, timeout=timeout) response = raw_response.parse() stringified_response = response.model_dump() ## LOGGING @@ -176,11 +168,9 @@ class AzureTextCompletion(BaseAzureLLM): "api_base": api_base, }, ) - return ( - openai_text_completion_config.convert_to_chat_model_response_object( - response_object=TextCompletionResponse(**stringified_response), - model_response_object=model_response, - ) + return openai_text_completion_config.convert_to_chat_model_response_object( + response_object=TextCompletionResponse(**stringified_response), + model_response_object=model_response, ) except AzureOpenAIError as e: raise e @@ -190,9 +180,7 @@ class AzureTextCompletion(BaseAzureLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError( - status_code=status_code, message=str(e), headers=error_headers - ) + raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) async def acompletion( self, @@ -239,9 +227,7 @@ class AzureTextCompletion(BaseAzureLLM): "complete_input_dict": data, }, ) - raw_response = await azure_client.completions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = await azure_client.completions.with_raw_response.create(**data, timeout=timeout) response = raw_response.parse() return openai_text_completion_config.convert_to_chat_model_response_object( response_object=response.model_dump(), @@ -255,9 +241,7 @@ class AzureTextCompletion(BaseAzureLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError( - status_code=status_code, message=str(e), headers=error_headers - ) + raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) def streaming( self, @@ -274,9 +258,7 @@ class AzureTextCompletion(BaseAzureLLM): ): max_retries = data.pop("max_retries", 2) if not isinstance(max_retries, int): - raise AzureOpenAIError( - status_code=422, message="max retries must be an int" - ) + raise AzureOpenAIError(status_code=422, message="max retries must be an int") # init AzureOpenAI Client azure_client = self.get_azure_openai_client( api_version=api_version, @@ -304,9 +286,7 @@ class AzureTextCompletion(BaseAzureLLM): "complete_input_dict": data, }, ) - raw_response = azure_client.completions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = azure_client.completions.with_raw_response.create(**data, timeout=timeout) response = raw_response.parse() streamwrapper = CustomStreamWrapper( completion_stream=response, @@ -356,9 +336,7 @@ class AzureTextCompletion(BaseAzureLLM): "complete_input_dict": data, }, ) - raw_response = await azure_client.completions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = await azure_client.completions.with_raw_response.create(**data, timeout=timeout) response = raw_response.parse() # return response streamwrapper = CustomStreamWrapper( @@ -374,6 +352,4 @@ class AzureTextCompletion(BaseAzureLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise AzureOpenAIError( - status_code=status_code, message=str(e), headers=error_headers - ) + raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers) diff --git a/litellm/llms/azure/containers/transformation.py b/litellm/llms/azure/containers/transformation.py index cd897511585..30cd3421d1b 100644 --- a/litellm/llms/azure/containers/transformation.py +++ b/litellm/llms/azure/containers/transformation.py @@ -43,9 +43,7 @@ class AzureContainerConfig(OpenAIContainerConfig): path = parsed.path.rstrip("/") for ep in _AZURE_ENDPOINT_PATHS: if path.endswith(ep): - return urlunparse( - (parsed.scheme, parsed.netloc, path[: -len(ep)], "", "", "") - ) + return urlunparse((parsed.scheme, parsed.netloc, path[: -len(ep)], "", "", "")) return api_base @staticmethod diff --git a/litellm/llms/azure/exception_mapping.py b/litellm/llms/azure/exception_mapping.py index dec7e7e5c90..07d589021a7 100644 --- a/litellm/llms/azure/exception_mapping.py +++ b/litellm/llms/azure/exception_mapping.py @@ -18,20 +18,12 @@ class AzureOpenAIExceptionMapping: """ Create a content policy violation error """ - azure_error, inner_error = AzureOpenAIExceptionMapping._extract_azure_error( - original_exception - ) + azure_error, inner_error = AzureOpenAIExceptionMapping._extract_azure_error(original_exception) # Prefer the provider message/type/code when present. - provider_message = ( - azure_error.get("message") if isinstance(azure_error, dict) else None - ) or message - provider_type = ( - azure_error.get("type") if isinstance(azure_error, dict) else None - ) - provider_code = ( - azure_error.get("code") if isinstance(azure_error, dict) else None - ) + provider_message = (azure_error.get("message") if isinstance(azure_error, dict) else None) or message + provider_type = azure_error.get("type") if isinstance(azure_error, dict) else None + provider_code = azure_error.get("code") if isinstance(azure_error, dict) else None # Keep the OpenAI-style body fields populated so downstream (proxy + SDK) # can surface `type` / `code` correctly. diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index cca83b8e6fd..8b277bdd49a 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -45,9 +45,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): openai_client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> OpenAIFileObject: verbose_logger.debug("create_file_data=%s", create_file_data) - response = await openai_client.files.create( - **self._prepare_create_file_data(create_file_data) - ) # type: ignore[arg-type] + response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data)) # type: ignore[arg-type] verbose_logger.debug("create_file_response=%s", response) return OpenAIFileObject(**response.model_dump()) @@ -60,20 +58,18 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: - openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - litellm_params=litellm_params or {}, - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + ) ) if openai_client is None: raise ValueError( @@ -85,9 +81,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): raise ValueError( "AzureOpenAI client is not an instance of AsyncAzureOpenAI. Make sure you passed an AsyncAzureOpenAI client." ) - return self.acreate_file( - create_file_data=create_file_data, openai_client=openai_client - ) + return self.acreate_file(create_file_data=create_file_data, openai_client=openai_client) response = cast(Union[AzureOpenAI, OpenAI], openai_client).files.create( **self._prepare_create_file_data(create_file_data) ) # type: ignore[arg-type] @@ -110,22 +104,18 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], api_version: Optional[str] = None, - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, - ) -> Union[ - HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] - ]: - openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - litellm_params=litellm_params or {}, - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, + ) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]: + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + ) ) if openai_client is None: raise ValueError( @@ -141,9 +131,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): file_content_request=file_content_request, openai_client=openai_client, ) - response = cast(Union[AzureOpenAI, OpenAI], openai_client).files.content( - **file_content_request - ) + response = cast(Union[AzureOpenAI, OpenAI], openai_client).files.content(**file_content_request) return HttpxBinaryResponseContent(response=response.response) @@ -164,20 +152,18 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], api_version: Optional[str] = None, - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): - openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - litellm_params=litellm_params or {}, - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + ) ) if openai_client is None: raise ValueError( @@ -218,20 +204,18 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: Optional[int], organization: Optional[str] = None, api_version: Optional[str] = None, - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): - openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - litellm_params=litellm_params or {}, - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + ) ) if openai_client is None: raise ValueError( @@ -274,20 +258,18 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: Optional[int], purpose: Optional[str] = None, api_version: Optional[str] = None, - client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): - openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] - ] = self.get_azure_openai_client( - litellm_params=litellm_params or {}, - api_key=api_key, - api_base=api_base, - api_version=api_version, - client=client, - _is_async=_is_async, + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = ( + self.get_azure_openai_client( + litellm_params=litellm_params or {}, + api_key=api_key, + api_base=api_base, + api_version=api_version, + client=client, + _is_async=_is_async, + ) ) if openai_client is None: raise ValueError( diff --git a/litellm/llms/azure/fine_tuning/handler.py b/litellm/llms/azure/fine_tuning/handler.py index 07d6455a6fb..f4a4166c8b4 100644 --- a/litellm/llms/azure/fine_tuning/handler.py +++ b/litellm/llms/azure/fine_tuning/handler.py @@ -28,18 +28,14 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): if extra_body.get("trainingType") is None: extra_body["trainingType"] = 1 create_fine_tuning_job_data["extra_body"] = extra_body - verbose_logger.debug( - "Azure fine-tuning: defaulting trainingType=1 (supervised)" - ) + verbose_logger.debug("Azure fine-tuning: defaulting trainingType=1 (supervised)") async def acreate_fine_tuning_job( self, create_fine_tuning_job_data: dict, openai_client: Union[AsyncOpenAI, AsyncAzureOpenAI], ) -> LiteLLMFineTuningJob: - response = await openai_client.fine_tuning.jobs.create( - **create_fine_tuning_job_data - ) + response = await openai_client.fine_tuning.jobs.create(**create_fine_tuning_job_data) return _litellm_fine_tuning_job_from_response(response, is_azure=True) async def acancel_fine_tuning_job( @@ -47,9 +43,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): fine_tuning_job_id: str, openai_client: Union[AsyncOpenAI, AsyncAzureOpenAI], ) -> LiteLLMFineTuningJob: - response = await openai_client.fine_tuning.jobs.cancel( - fine_tuning_job_id=fine_tuning_job_id - ) + response = await openai_client.fine_tuning.jobs.cancel(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response, is_azure=True) async def aretrieve_fine_tuning_job( @@ -57,9 +51,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): fine_tuning_job_id: str, openai_client: Union[AsyncOpenAI, AsyncAzureOpenAI], ) -> LiteLLMFineTuningJob: - response = await openai_client.fine_tuning.jobs.retrieve( - fine_tuning_job_id=fine_tuning_job_id - ) + response = await openai_client.fine_tuning.jobs.retrieve(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response, is_azure=True) def create_fine_tuning_job( @@ -72,15 +64,11 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: self._ensure_training_type(create_fine_tuning_job_data) - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = self.get_openai_client( + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, timeout=timeout, @@ -105,12 +93,8 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): openai_client=openai_client, ) - verbose_logger.debug( - "creating fine tuning job, args= %s", create_fine_tuning_job_data - ) - response = cast(OpenAI, openai_client).fine_tuning.jobs.create( - **create_fine_tuning_job_data - ) + verbose_logger.debug("creating fine tuning job, args= %s", create_fine_tuning_job_data) + response = cast(OpenAI, openai_client).fine_tuning.jobs.create(**create_fine_tuning_job_data) return _litellm_fine_tuning_job_from_response(response, is_azure=True) def cancel_fine_tuning_job( @@ -123,13 +107,9 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = self.get_openai_client( + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, timeout=timeout, @@ -154,9 +134,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): openai_client=openai_client, ) - response = cast(OpenAI, openai_client).fine_tuning.jobs.cancel( - fine_tuning_job_id=fine_tuning_job_id - ) + response = cast(OpenAI, openai_client).fine_tuning.jobs.cancel(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response, is_azure=True) def retrieve_fine_tuning_job( @@ -169,13 +147,9 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = self.get_openai_client( + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, timeout=timeout, @@ -200,9 +174,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): openai_client=openai_client, ) - response = cast(OpenAI, openai_client).fine_tuning.jobs.retrieve( - fine_tuning_job_id=fine_tuning_job_id - ) + response = cast(OpenAI, openai_client).fine_tuning.jobs.retrieve(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response, is_azure=True) def get_openai_client( @@ -212,9 +184,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, _is_async: bool = False, api_version: Optional[str] = None, litellm_params: Optional[dict] = None, diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index 72f1eef36c0..d28d92a0770 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -65,9 +65,7 @@ class AzureImageEditConfig(OpenAIImageEditConfig): params = GenericLiteLLMParams(**(litellm_params or {})) if api_key is not None and params.api_key is None: params.api_key = api_key - return BaseAzureLLM._base_validate_azure_environment( - headers=headers, litellm_params=params - ) + return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=params) def get_complete_url( self, @@ -128,7 +126,5 @@ class AzureImageEditConfig(OpenAIImageEditConfig): return str(final_url) - def finalize_image_edit_request_data( - self, data: dict, resolved_request_url: str - ) -> dict: + def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict: return self.azure_deployment_image_edit_form_data(data, resolved_request_url) diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 9b1d95e5314..dabcd4a1183 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -62,9 +62,7 @@ class AzurePassthroughConfig(BasePassthroughConfig): ) -> dict: return BaseAzureLLM._base_validate_azure_environment( headers=headers, - litellm_params=GenericLiteLLMParams( - **{**litellm_params, "api_key": api_key} - ), + litellm_params=GenericLiteLLMParams(**{**litellm_params, "api_key": api_key}), ) @staticmethod @@ -83,9 +81,7 @@ class AzurePassthroughConfig(BasePassthroughConfig): def get_base_model(model: str) -> Optional[str]: return model - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: return super().get_models(api_key, api_base) def logging_non_streaming_response( diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 9c8de6c06a1..86c1ed51b68 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -71,9 +71,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): if _is_ga: path = "/openai/v1/realtime" query_parts = [] - if intent != "transcription" and ( - query_params is None or "model" in query_params - ): + if intent != "transcription" and (query_params is None or "model" in query_params): query_parts.append(urlencode({"model": model})) else: # Default to beta path for backwards compatibility @@ -107,9 +105,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): if api_base is None: raise ValueError("api_base is required for Azure OpenAI calls") - backend_uses_beta_protocol = ( - realtime_protocol is None or realtime_protocol.upper() not in ("GA", "V1") - ) + backend_uses_beta_protocol = realtime_protocol is None or realtime_protocol.upper() not in ("GA", "V1") if api_version is None and backend_uses_beta_protocol: raise ValueError("api_version is required for Azure OpenAI calls") @@ -140,9 +136,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): request_data={"litellm_metadata": litellm_metadata or {}}, backend_uses_beta_protocol=backend_uses_beta_protocol, force_transcription_model=( - model - if (query_params or {}).get("intent") == "transcription" - else None + model if (query_params or {}).get("intent") == "transcription" else None ), ) await realtime_streaming.bidirectional_forward() @@ -150,7 +144,5 @@ class AzureOpenAIRealtime(AzureChatCompletion): except websockets.exceptions.InvalidStatusCode as e: # type: ignore await websocket.close(code=e.status_code, reason=_redact_string(str(e))) except Exception: - verbose_proxy_logger.exception( - "Error in AzureOpenAIRealtime.async_realtime" - ) + verbose_proxy_logger.exception("Error in AzureOpenAIRealtime.async_realtime") pass diff --git a/litellm/llms/azure/realtime/http_transformation.py b/litellm/llms/azure/realtime/http_transformation.py index d6bdbd24db4..55a86014423 100644 --- a/litellm/llms/azure/realtime/http_transformation.py +++ b/litellm/llms/azure/realtime/http_transformation.py @@ -14,9 +14,7 @@ class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): def get_api_key(self, api_key: Optional[str], **kwargs) -> str: return api_key or litellm.api_key or get_secret_str("AZURE_API_KEY") or "" - def get_complete_url( - self, api_base: Optional[str], model: str, api_version: Optional[str] = None - ) -> str: + def get_complete_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: base = self.get_api_base(api_base).rstrip("/") version = api_version or get_secret_str("AZURE_API_VERSION") or "2024-12-17" return f"{base}/openai/realtime/client_secrets?api-version={version}" @@ -33,9 +31,7 @@ class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): "Content-Type": "application/json", } - def get_realtime_calls_url( - self, api_base: Optional[str], model: str, api_version: Optional[str] = None - ) -> str: + def get_realtime_calls_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: base = self.get_api_base(api_base).rstrip("/") version = api_version or get_secret_str("AZURE_API_VERSION") or "2024-12-17" return f"{base}/openai/realtime/calls?api-version={version}" diff --git a/litellm/llms/azure/responses/o_series_transformation.py b/litellm/llms/azure/responses/o_series_transformation.py index 3a554e9e194..2cc3e914307 100644 --- a/litellm/llms/azure/responses/o_series_transformation.py +++ b/litellm/llms/azure/responses/o_series_transformation.py @@ -46,9 +46,7 @@ class AzureOpenAIOSeriesResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): # Filter out unsupported parameters for O-series models o_series_supported_params = [ - param - for param in base_supported_params - if param not in o_series_unsupported_params + param for param in base_supported_params if param not in o_series_unsupported_params ] return o_series_supported_params diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 92ce5b49285..fef9b7d0154 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -34,18 +34,10 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): Azure Responses API does not support context_management (compaction). """ base_supported_params = super().get_supported_openai_params(model) - return [ - param - for param in base_supported_params - if param not in self.AZURE_UNSUPPORTED_PARAMS - ] + return [param for param in base_supported_params if param not in self.AZURE_UNSUPPORTED_PARAMS] - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: - return BaseAzureLLM._base_validate_azure_environment( - headers=headers, litellm_params=litellm_params - ) + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: + return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params) def get_stripped_model_name(self, model: str) -> str: # if "responses/" is in the model name, remove it @@ -82,22 +74,17 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return dict_reasoning_item except Exception as e: - verbose_logger.debug( - f"Failed to create ResponseReasoningItem, falling back to manual filtering: {e}" - ) + verbose_logger.debug(f"Failed to create ResponseReasoningItem, falling back to manual filtering: {e}") # Fallback: manually filter out known None fields filtered_item = { k: v for k, v in item.items() - if v is not None - or k not in {"status", "content", "encrypted_content"} + if v is not None or k not in {"status", "content", "encrypted_content"} } return filtered_item return item - def _validate_input_param( - self, input: Union[str, ResponseInputParam] - ) -> Union[str, ResponseInputParam]: + def _validate_input_param(self, input: Union[str, ResponseInputParam]) -> Union[str, ResponseInputParam]: """ Override parent method to also filter out 'status' field from message items. Azure OpenAI API does not accept 'status' field in input messages. @@ -209,11 +196,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): path = path[: -len(suffix)] break scheme = "wss" if parsed_url.scheme == "https" else "ws" - return str( - parsed_url.copy_with( - scheme=scheme, path=f"{path}/openai/v1/responses", query=None - ) - ) + return str(parsed_url.copy_with(scheme=scheme, path=f"{path}/openai/v1/responses", query=None)) def model_in_websocket_url(self) -> bool: # Azure sends the model in the response.create body, not the URL @@ -222,9 +205,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): ######################################################### ########## DELETE RESPONSE API TRANSFORMATION ############## ######################################################### - def _construct_url_for_response_id_in_path( - self, api_base: str, response_id: str - ) -> str: + def _construct_url_for_response_id_in_path(self, api_base: str, response_id: str) -> str: """ Constructs a URL for the API request with the response_id in the path. """ @@ -236,9 +217,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Insert the response_id at the end of the path component # Remove trailing slash if present to avoid double slashes path = parsed_url.path.rstrip("/") - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") new_path = f"{path}/{encoded_response_id}" # Reconstruct the URL with all original components but with the modified path @@ -270,9 +249,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): This function handles URLs with query parameters by inserting the response_id at the correct location (before any query parameters). """ - delete_url = self._construct_url_for_response_id_in_path( - api_base=api_base, response_id=response_id - ) + delete_url = self._construct_url_for_response_id_in_path(api_base=api_base, response_id=response_id) data: Dict = {} verbose_logger.debug(f"delete response url={delete_url}") @@ -294,9 +271,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): OpenAI API expects the following request - GET /v1/responses/{response_id} """ - get_url = self._construct_url_for_response_id_in_path( - api_base=api_base, response_id=response_id - ) + get_url = self._construct_url_for_response_id_in_path(api_base=api_base, response_id=response_id) data: Dict = {} verbose_logger.debug(f"get response url={get_url}") return get_url, data @@ -313,12 +288,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): limit: int = 20, order: Literal["asc", "desc"] = "desc", ) -> Tuple[str, Dict]: - url = ( - self._construct_url_for_response_id_in_path( - api_base=api_base, response_id=response_id - ) - + "/input_items" - ) + url = self._construct_url_for_response_id_in_path(api_base=api_base, response_id=response_id) + "/input_items" params: Dict[str, Any] = {} if after is not None: params["after"] = after @@ -360,9 +330,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Insert the response_id and /cancel at the end of the path component # Remove trailing slash if present to avoid double slashes path = parsed_url.path.rstrip("/") - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") new_path = f"{path}/{encoded_response_id}/cancel" # Reconstruct the URL with all original components but with the modified path @@ -394,7 +362,5 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): except Exception: from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIError - raise AzureOpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise AzureOpenAIError(message=raw_response.text, status_code=raw_response.status_code) return ResponsesAPIResponse(**raw_response_json) diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index a5dec243147..c3e5f16b03a 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -86,10 +86,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): """ # Resolve api_base from multiple sources api_base = ( - api_base - or litellm_params_dict.get("api_base") - or litellm.api_base - or get_secret_str("AZURE_API_BASE") + api_base or litellm_params_dict.get("api_base") or litellm.api_base or get_secret_str("AZURE_API_BASE") ) # Resolve api_key from multiple sources (Azure-specific) @@ -337,9 +334,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): # Check if it's a Cognitive Services endpoint (convert to TTS endpoint) if self._is_cognitive_services_endpoint(hostname=hostname): - region = self._extract_region_from_hostname( - hostname=hostname, domain=self.COGNITIVE_SERVICES_DOMAIN - ) + region = self._extract_region_from_hostname(hostname=hostname, domain=self.COGNITIVE_SERVICES_DOMAIN) return self._build_tts_url(region=region) # Check if it's already a TTS endpoint @@ -353,15 +348,11 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): def _is_cognitive_services_endpoint(self, hostname: str) -> bool: """Check if hostname is a Cognitive Services endpoint""" - return hostname == self.COGNITIVE_SERVICES_DOMAIN or hostname.endswith( - f".{self.COGNITIVE_SERVICES_DOMAIN}" - ) + return hostname == self.COGNITIVE_SERVICES_DOMAIN or hostname.endswith(f".{self.COGNITIVE_SERVICES_DOMAIN}") def _is_tts_endpoint(self, hostname: str) -> bool: """Check if hostname is a TTS endpoint""" - return hostname == self.TTS_SPEECH_DOMAIN or hostname.endswith( - f".{self.TTS_SPEECH_DOMAIN}" - ) + return hostname == self.TTS_SPEECH_DOMAIN or hostname.endswith(f".{self.TTS_SPEECH_DOMAIN}") def _extract_region_from_hostname(self, hostname: str, domain: str) -> str: """ @@ -419,9 +410,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): azure_voice = voice or self.DEFAULT_VOICE # Get output format (already mapped in main.py) - output_format = optional_params.get( - "output_format", "audio-24khz-48kbitrate-mono-mp3" - ) + output_format = optional_params.get("output_format", "audio-24khz-48kbitrate-mono-mp3") headers["X-Microsoft-OutputFormat"] = output_format # Auto-detect SSML: if input contains , pass it through as-is diff --git a/litellm/llms/azure/vector_stores/transformation.py b/litellm/llms/azure/vector_stores/transformation.py index a98e7ae8cb6..c340294c6b4 100644 --- a/litellm/llms/azure/vector_stores/transformation.py +++ b/litellm/llms/azure/vector_stores/transformation.py @@ -17,9 +17,5 @@ class AzureOpenAIVectorStoreConfig(OpenAIVectorStoreConfig): route="/openai/vector_stores", ) - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: - return BaseAzureLLM._base_validate_azure_environment( - headers=headers, litellm_params=litellm_params - ) + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: + return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params) diff --git a/litellm/llms/azure/videos/transformation.py b/litellm/llms/azure/videos/transformation.py index 1ee0e95fb0a..92e7c91fed3 100644 --- a/litellm/llms/azure/videos/transformation.py +++ b/litellm/llms/azure/videos/transformation.py @@ -72,9 +72,7 @@ class AzureVideoConfig(OpenAIVideoConfig): # Use the base Azure validation method which properly handles: # 1. Credentials from litellm_credential_name via litellm_params # 2. Sets the correct "api-key" header (not "Authorization: Bearer") - return BaseAzureLLM._base_validate_azure_environment( - headers=headers, litellm_params=litellm_params - ) + return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params) def get_complete_url( self, diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 9bae8abce8e..6083580ed45 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -73,32 +73,22 @@ class AzureAIAgentsHandler: def _build_thread_url(self, api_base: str, api_version: str) -> str: return f"{api_base}/threads?api-version={api_version}" - def _build_messages_url( - self, api_base: str, thread_id: str, api_version: str - ) -> str: + def _build_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") - return ( - f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" - ) + return f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" def _build_runs_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") return f"{api_base}/threads/{encoded_thread_id}/runs?api-version={api_version}" - def _build_run_status_url( - self, api_base: str, thread_id: str, run_id: str, api_version: str - ) -> str: + def _build_run_status_url(self, api_base: str, thread_id: str, run_id: str, api_version: str) -> str: encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") return f"{api_base}/threads/{encoded_thread_id}/runs/{encoded_run_id}?api-version={api_version}" - def _build_list_messages_url( - self, api_base: str, thread_id: str, api_version: str - ) -> str: + def _build_list_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") - return ( - f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" - ) + return f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" def _build_create_thread_and_run_url(self, api_base: str, api_version: str) -> str: """URL for the create-thread-and-run endpoint (supports streaming).""" @@ -107,9 +97,7 @@ class AzureAIAgentsHandler: # ------------------------------------------------------------------------- # Response Helpers # ------------------------------------------------------------------------- - def _extract_content_from_messages( - self, messages_data: dict - ) -> Tuple[str, Optional[List[Dict[str, Any]]]]: + def _extract_content_from_messages(self, messages_data: dict) -> Tuple[str, Optional[List[Dict[str, Any]]]]: """Extract assistant content and annotations from the messages response. Returns (content, annotations) where annotations is a list of @@ -190,10 +178,7 @@ class AzureAIAgentsHandler: model_response.model = model # Store thread_id for conversation continuity - if ( - not hasattr(model_response, "_hidden_params") - or model_response._hidden_params is None - ): + if not hasattr(model_response, "_hidden_params") or model_response._hidden_params is None: model_response._hidden_params = {} model_response._hidden_params["thread_id"] = thread_id @@ -202,9 +187,7 @@ class AzureAIAgentsHandler: from litellm.utils import token_counter prompt_tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - completion_tokens = token_counter( - model="gpt-3.5-turbo", text=content, count_response_tokens=True - ) + completion_tokens = token_counter(model="gpt-3.5-turbo", text=content, count_response_tokens=True) setattr( model_response, "usage", @@ -243,22 +226,16 @@ class AzureAIAgentsHandler: if api_key: headers["Authorization"] = f"Bearer {api_key}" - api_version = optional_params.get( - "api_version", self.config.DEFAULT_API_VERSION - ) + api_version = optional_params.get("api_version", self.config.DEFAULT_API_VERSION) agent_id = self.config._get_agent_id(model, optional_params) thread_id = optional_params.get("thread_id") api_base = api_base.rstrip("/") - verbose_logger.debug( - f"Azure AI Agents completion - api_base: {api_base}, agent_id: {agent_id}" - ) + verbose_logger.debug(f"Azure AI Agents completion - api_base: {api_base}, agent_id: {agent_id}") return headers, api_version, agent_id, thread_id, api_base - def _check_response( - self, response: httpx.Response, expected_codes: List[int], error_msg: str - ): + def _check_response(self, response: httpx.Response, expected_codes: List[int], error_msg: str): """Check response status and raise error if not expected.""" if response.status_code not in expected_codes: raise AzureAIAgentsError( @@ -287,9 +264,7 @@ class AzureAIAgentsHandler: from litellm.llms.custom_httpx.http_handler import _get_httpx_client if client is None: - client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) ( headers, @@ -297,13 +272,9 @@ class AzureAIAgentsHandler: agent_id, thread_id, api_base, - ) = self._prepare_completion_params( - model, api_base, api_key, optional_params, headers - ) + ) = self._prepare_completion_params(model, api_base, api_key, optional_params, headers) - def make_request( - method: str, url: str, json_data: Optional[dict] = None - ) -> httpx.Response: + def make_request(method: str, url: str, json_data: Optional[dict] = None) -> httpx.Response: if method == "GET": return client.get(url=url, headers=headers) return client.post( @@ -323,9 +294,7 @@ class AzureAIAgentsHandler: optional_params=optional_params, ) - return self._build_model_response( - model, content, model_response, thread_id, messages, annotations - ) + return self._build_model_response(model, content, model_response, thread_id, messages, annotations) def _execute_agent_flow_sync( self, @@ -341,12 +310,8 @@ class AzureAIAgentsHandler: # Step 1: Create thread if not provided if not thread_id: - verbose_logger.debug( - f"Creating thread at: {self._build_thread_url(api_base, api_version)}" - ) - response = make_request( - "POST", self._build_thread_url(api_base, api_version), {} - ) + verbose_logger.debug(f"Creating thread at: {self._build_thread_url(api_base, api_version)}") + response = make_request("POST", self._build_thread_url(api_base, api_version), {}) self._check_response(response, [200, 201], "Failed to create thread") thread_id = response.json()["id"] verbose_logger.debug(f"Created thread: {thread_id}") @@ -358,9 +323,7 @@ class AzureAIAgentsHandler: for msg in messages: if msg.get("role") in ["user", "system"]: url = self._build_messages_url(api_base, thread_id, api_version) - response = make_request( - "POST", url, {"role": "user", "content": msg.get("content", "")} - ) + response = make_request("POST", url, {"role": "user", "content": msg.get("content", "")}) self._check_response(response, [200, 201], "Failed to add message") # Step 3: Create run @@ -368,17 +331,13 @@ class AzureAIAgentsHandler: if "instructions" in optional_params: run_payload["instructions"] = optional_params["instructions"] - response = make_request( - "POST", self._build_runs_url(api_base, thread_id, api_version), run_payload - ) + response = make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload) self._check_response(response, [200, 201], "Failed to create run") run_id = response.json()["id"] verbose_logger.debug(f"Created run: {run_id}") # Step 4: Poll for completion - status_url = self._build_run_status_url( - api_base, thread_id, run_id, api_version - ) + status_url = self._build_run_status_url(api_base, thread_id, run_id, api_version) for _ in range(self.config.MAX_POLL_ATTEMPTS): response = make_request("GET", status_url) self._check_response(response, [200], "Failed to get run status") @@ -389,25 +348,15 @@ class AzureAIAgentsHandler: if status == "completed": break elif status in ["failed", "cancelled", "expired"]: - error_msg = ( - response.json() - .get("last_error", {}) - .get("message", "Unknown error") - ) - raise AzureAIAgentsError( - status_code=500, message=f"Run {status}: {error_msg}" - ) + error_msg = response.json().get("last_error", {}).get("message", "Unknown error") + raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}") time.sleep(self.config.POLL_INTERVAL_SECONDS) else: - raise AzureAIAgentsError( - status_code=408, message="Run timed out waiting for completion" - ) + raise AzureAIAgentsError(status_code=408, message="Run timed out waiting for completion") # Step 5: Get messages - response = make_request( - "GET", self._build_list_messages_url(api_base, thread_id, api_version) - ) + response = make_request("GET", self._build_list_messages_url(api_base, thread_id, api_version)) self._check_response(response, [200], "Failed to get messages") content, annotations = self._extract_content_from_messages(response.json()) @@ -446,13 +395,9 @@ class AzureAIAgentsHandler: agent_id, thread_id, api_base, - ) = self._prepare_completion_params( - model, api_base, api_key, optional_params, headers - ) + ) = self._prepare_completion_params(model, api_base, api_key, optional_params, headers) - async def make_request( - method: str, url: str, json_data: Optional[dict] = None - ) -> httpx.Response: + async def make_request(method: str, url: str, json_data: Optional[dict] = None) -> httpx.Response: if method == "GET": return await client.get(url=url, headers=headers) return await client.post( @@ -472,9 +417,7 @@ class AzureAIAgentsHandler: optional_params=optional_params, ) - return self._build_model_response( - model, content, model_response, thread_id, messages, annotations - ) + return self._build_model_response(model, content, model_response, thread_id, messages, annotations) async def _execute_agent_flow_async( self, @@ -490,12 +433,8 @@ class AzureAIAgentsHandler: # Step 1: Create thread if not provided if not thread_id: - verbose_logger.debug( - f"Creating thread at: {self._build_thread_url(api_base, api_version)}" - ) - response = await make_request( - "POST", self._build_thread_url(api_base, api_version), {} - ) + verbose_logger.debug(f"Creating thread at: {self._build_thread_url(api_base, api_version)}") + response = await make_request("POST", self._build_thread_url(api_base, api_version), {}) self._check_response(response, [200, 201], "Failed to create thread") thread_id = response.json()["id"] verbose_logger.debug(f"Created thread: {thread_id}") @@ -507,9 +446,7 @@ class AzureAIAgentsHandler: for msg in messages: if msg.get("role") in ["user", "system"]: url = self._build_messages_url(api_base, thread_id, api_version) - response = await make_request( - "POST", url, {"role": "user", "content": msg.get("content", "")} - ) + response = await make_request("POST", url, {"role": "user", "content": msg.get("content", "")}) self._check_response(response, [200, 201], "Failed to add message") # Step 3: Create run @@ -517,17 +454,13 @@ class AzureAIAgentsHandler: if "instructions" in optional_params: run_payload["instructions"] = optional_params["instructions"] - response = await make_request( - "POST", self._build_runs_url(api_base, thread_id, api_version), run_payload - ) + response = await make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload) self._check_response(response, [200, 201], "Failed to create run") run_id = response.json()["id"] verbose_logger.debug(f"Created run: {run_id}") # Step 4: Poll for completion - status_url = self._build_run_status_url( - api_base, thread_id, run_id, api_version - ) + status_url = self._build_run_status_url(api_base, thread_id, run_id, api_version) for _ in range(self.config.MAX_POLL_ATTEMPTS): response = await make_request("GET", status_url) self._check_response(response, [200], "Failed to get run status") @@ -538,25 +471,15 @@ class AzureAIAgentsHandler: if status == "completed": break elif status in ["failed", "cancelled", "expired"]: - error_msg = ( - response.json() - .get("last_error", {}) - .get("message", "Unknown error") - ) - raise AzureAIAgentsError( - status_code=500, message=f"Run {status}: {error_msg}" - ) + error_msg = response.json().get("last_error", {}).get("message", "Unknown error") + raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}") await asyncio.sleep(self.config.POLL_INTERVAL_SECONDS) else: - raise AzureAIAgentsError( - status_code=408, message="Run timed out waiting for completion" - ) + raise AzureAIAgentsError(status_code=408, message="Run timed out waiting for completion") # Step 5: Get messages - response = await make_request( - "GET", self._build_list_messages_url(api_base, thread_id, api_version) - ) + response = await make_request("GET", self._build_list_messages_url(api_base, thread_id, api_version)) self._check_response(response, [200], "Failed to get messages") content, annotations = self._extract_content_from_messages(response.json()) @@ -587,17 +510,13 @@ class AzureAIAgentsHandler: agent_id, thread_id, api_base, - ) = self._prepare_completion_params( - model, api_base, api_key, optional_params, headers - ) + ) = self._prepare_completion_params(model, api_base, api_key, optional_params, headers) # Build payload for create-thread-and-run with streaming thread_messages = [] for msg in messages: if msg.get("role") in ["user", "system"]: - thread_messages.append( - {"role": "user", "content": msg.get("content", "")} - ) + thread_messages.append({"role": "user", "content": msg.get("content", "")}) payload: Dict[str, Any] = { "assistant_id": agent_id, @@ -699,9 +618,7 @@ class AzureAIAgentsHandler: if current_event == "thread.message.completed": for content_item in data.get("content", []): if content_item.get("type") == "text": - raw_annotations = content_item.get("text", {}).get( - "annotations" - ) + raw_annotations = content_item.get("text", {}).get("annotations") transformed = self._transform_annotations(raw_annotations) if transformed: if collected_annotations is None: @@ -724,9 +641,7 @@ class AzureAIAgentsHandler: StreamingChoices( finish_reason=None, index=0, - delta=Delta( - content=text_value, role="assistant" - ), + delta=Delta(content=text_value, role="assistant"), ) ], ) diff --git a/litellm/llms/azure_ai/agents/transformation.py b/litellm/llms/azure_ai/agents/transformation.py index 777509fa82c..daf87b01579 100644 --- a/litellm/llms/azure_ai/agents/transformation.py +++ b/litellm/llms/azure_ai/agents/transformation.py @@ -178,9 +178,7 @@ class AzureAIAgentsConfig(BaseConfig): model format: "azure_ai/agents/" or "agents/" or just "" """ - agent_id = optional_params.get("agent_id") or optional_params.get( - "assistant_id" - ) + agent_id = optional_params.get("agent_id") or optional_params.get("assistant_id") if agent_id: return agent_id diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index e24fc2097d2..0716e5ae988 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -56,9 +56,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): # Validate the request self.validate_request(model, messages) - verbose_logger.debug( - f"Processing Azure AI Anthropic CountTokens request for model: {model}" - ) + verbose_logger.debug(f"Processing Azure AI Anthropic CountTokens request for model: {model}") # Transform request to Anthropic format request_body = self.transform_request_to_count_tokens( @@ -82,14 +80,10 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): ) # Use LiteLLM's async httpx client - async_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.AZURE_AI - ) + async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI) # Use provided timeout or fall back to litellm.request_timeout - request_timeout = ( - timeout if timeout is not None else litellm.request_timeout - ) + request_timeout = timeout if timeout is not None else litellm.request_timeout response = await async_client.post( endpoint_url, diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index afdfe9bdee9..8e1ee73620f 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -107,9 +107,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): status_code=e.status_code, ) except Exception as e: - verbose_logger.warning( - f"Error calling Azure AI Anthropic CountTokens API: {e}" - ) + verbose_logger.warning(f"Error calling Azure AI Anthropic CountTokens API: {e}") return TokenCountResponse( total_tokens=0, request_model=request_model, diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index 09b83b7c971..5e1fb69f40d 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -56,9 +56,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): litellm_params_obj = GenericLiteLLMParams(**litellm_params) # Get Azure auth headers (api-key or Authorization) - azure_headers = BaseAzureLLM._base_validate_azure_environment( - headers={}, litellm_params=litellm_params_obj - ) + azure_headers = BaseAzureLLM._base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) # Merge Azure auth headers headers.update(azure_headers) diff --git a/litellm/llms/azure_ai/anthropic/handler.py b/litellm/llms/azure_ai/anthropic/handler.py index f3a50b73c1a..d510e5bd13e 100644 --- a/litellm/llms/azure_ai/anthropic/handler.py +++ b/litellm/llms/azure_ai/anthropic/handler.py @@ -119,11 +119,7 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion): logger_fn=logger_fn, headers=headers, timeout=timeout, - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), ) else: return self.acompletion_function( diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 59b6ee2b424..1de18701a2f 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -49,9 +49,7 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): litellm_params_obj.api_key = api_key # Use Azure authentication logic - headers = BaseAzureLLM._base_validate_azure_environment( - headers=headers, litellm_params=litellm_params_obj - ) + headers = BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params_obj) # Azure Anthropic uses x-api-key header (not api-key) # Convert api-key to x-api-key if present @@ -118,9 +116,7 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): return api_base - def _remove_scope_from_cache_control( - self, anthropic_messages_request: Dict - ) -> None: + def _remove_scope_from_cache_control(self, anthropic_messages_request: Dict) -> None: """ Remove `scope` field from cache_control for Azure AI Foundry. diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index 367ca75c196..26323ba707d 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -74,17 +74,13 @@ class AzureAnthropicConfig(AnthropicConfig): litellm_params_obj.api_key = api_key # Use Azure authentication logic - headers = BaseAzureLLM._base_validate_azure_environment( - headers=headers, litellm_params=litellm_params_obj - ) + headers = BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params_obj) # Get tools and other anthropic-specific setup tools = optional_params.get("tools") prompt_caching_set = self.is_cache_control_set(messages=messages) computer_tool_used = self.is_computer_tool_used(tools=tools) - mcp_server_used = self.is_mcp_server_used( - mcp_servers=optional_params.get("mcp_servers") - ) + mcp_server_used = self.is_mcp_server_used(mcp_servers=optional_params.get("mcp_servers")) pdf_used = self.is_pdf_used(messages=messages) file_id_used = self.is_file_id_used(messages=messages) user_anthropic_beta_headers = self._get_user_anthropic_beta_headers( diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index e4174f41ad7..f3045283840 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -44,9 +44,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): # Get base model name (strips routing prefixes like model_router/) base_model: str = AzureFoundryModelInfo.get_base_model(model) - return super().transform_request( - base_model, messages, optional_params, litellm_params, headers - ) + return super().transform_request(base_model, messages, optional_params, litellm_params, headers) def transform_response( self, @@ -90,9 +88,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): ) return model_response - def calculate_additional_costs( - self, model: str, prompt_tokens: int, completion_tokens: int - ) -> Optional[dict]: + def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> Optional[dict]: """ Calculate additional costs for Azure Model Router. @@ -110,9 +106,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): calculate_azure_model_router_flat_cost, ) - flat_cost = calculate_azure_model_router_flat_cost( - model=model, prompt_tokens=prompt_tokens - ) + flat_cost = calculate_azure_model_router_flat_cost(model=model, prompt_tokens=prompt_tokens) if flat_cost > 0: return {"Azure Model Router Flat Cost": flat_cost} diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 008a8a766e9..27a98347087 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -74,12 +74,8 @@ class AzureAIStudioConfig(OpenAIConfig): headers["Authorization"] = f"Bearer {api_key}" else: # No api_key provided — fall back to Azure AD token-based auth - litellm_params_obj = GenericLiteLLMParams( - **(litellm_params if isinstance(litellm_params, dict) else {}) - ) - headers = BaseAzureLLM._base_validate_azure_environment( - headers=headers, litellm_params=litellm_params_obj - ) + litellm_params_obj = GenericLiteLLMParams(**(litellm_params if isinstance(litellm_params, dict) else {})) + headers = BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params_obj) headers["Content-Type"] = "application/json" @@ -91,10 +87,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ parsed_url = urlparse(api_base) host = parsed_url.hostname - if host and ( - host.endswith(".services.ai.azure.com") - or host.endswith(".openai.azure.com") - ): + if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): return True return False @@ -141,13 +134,9 @@ class AzureAIStudioConfig(OpenAIConfig): # Add the path to the base URL if "services.ai.azure.com" in api_base: - new_url = _add_path_to_api_base( - api_base=api_base, ending_path="/models/chat/completions" - ) + new_url = _add_path_to_api_base(api_base=api_base, ending_path="/models/chat/completions") else: - new_url = _add_path_to_api_base( - api_base=api_base, ending_path="/chat/completions" - ) + new_url = _add_path_to_api_base(api_base=api_base, ending_path="/chat/completions") # Use the new query_params dictionary final_url = httpx.URL(new_url).copy_with(params=query_params) @@ -217,11 +206,7 @@ class AzureAIStudioConfig(OpenAIConfig): dynamic_api_key = api_key or get_secret_str("AZURE_AI_API_KEY") if self._is_azure_openai_model(model=model, api_base=api_base): - verbose_logger.debug( - "Model={} is Azure OpenAI model. Setting custom_llm_provider='azure'.".format( - model - ) - ) + verbose_logger.debug("Model={} is Azure OpenAI model. Setting custom_llm_provider='azure'.".format(model)) custom_llm_provider = "azure" return api_base, dynamic_api_key, custom_llm_provider @@ -237,9 +222,7 @@ class AzureAIStudioConfig(OpenAIConfig): if extra_body and isinstance(extra_body, dict): optional_params.update(extra_body) optional_params.pop("max_retries", None) - return super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + return super().transform_request(model, messages, optional_params, litellm_params, headers) def transform_response( self, @@ -277,20 +260,13 @@ class AzureAIStudioConfig(OpenAIConfig): error_text = e.response.text if "Extra inputs are not permitted" in error_text: - if should_drop_params or self._error_has_tool_level_extra_fields( - error_text - ): + if should_drop_params or self._error_has_tool_level_extra_fields(error_text): return True if "unknown field: parameter index is not a valid field" in error_text: return True - if ( - AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value - in error_text - ): + if AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value in error_text: return True - return super().should_retry_llm_api_inside_llm_translation_on_http_error( - e=e, litellm_params=litellm_params - ) + return super().should_retry_llm_api_inside_llm_translation_on_http_error(e=e, litellm_params=litellm_params) def _error_has_tool_level_extra_fields(self, error_text: str) -> bool: return bool(re.search(r"tools\[\d+\]\.", error_text)) @@ -299,36 +275,21 @@ class AzureAIStudioConfig(OpenAIConfig): def max_retry_on_unprocessable_entity_error(self) -> int: return 2 - def transform_request_on_unprocessable_entity_error( - self, e: httpx.HTTPStatusError, request_data: dict - ) -> dict: + def transform_request_on_unprocessable_entity_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict: error_text = e.response.text _messages = cast(Optional[List[AllMessageValues]], request_data.get("messages")) - if ( - "unknown field: parameter index is not a valid field" in error_text - and _messages is not None - ): + if "unknown field: parameter index is not a valid field" in error_text and _messages is not None: litellm.remove_index_from_tool_calls( messages=_messages, ) - elif ( - AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value - in error_text - ): - request_data = self._drop_extra_params_from_request_data( - request_data, error_text - ) - if ( - "Extra inputs are not permitted" in error_text - and self._error_has_tool_level_extra_fields(error_text) - ): + elif AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value in error_text: + request_data = self._drop_extra_params_from_request_data(request_data, error_text) + if "Extra inputs are not permitted" in error_text and self._error_has_tool_level_extra_fields(error_text): request_data = self._drop_tool_level_extra_fields(request_data, error_text) data = drop_params_from_unprocessable_entity_error(e=e, data=request_data) return data - def _drop_tool_level_extra_fields( - self, request_data: dict, error_text: str - ) -> dict: + def _drop_tool_level_extra_fields(self, request_data: dict, error_text: str) -> dict: fields_to_drop = set(re.findall(r"tools\[\d+\]\.([\w-]+)", error_text)) tools = request_data.get("tools") if fields_to_drop and isinstance(tools, list): @@ -338,9 +299,7 @@ class AzureAIStudioConfig(OpenAIConfig): tool.pop(field, None) return request_data - def _drop_extra_params_from_request_data( - self, request_data: dict, error_text: str - ) -> dict: + def _drop_extra_params_from_request_data(self, request_data: dict, error_text: str) -> dict: params_to_drop = self._extract_params_to_drop_from_error_text(error_text) if params_to_drop: for param in params_to_drop: @@ -348,9 +307,7 @@ class AzureAIStudioConfig(OpenAIConfig): request_data.pop(param, None) return request_data - def _extract_params_to_drop_from_error_text( - self, error_text: str - ) -> Optional[List[str]]: + def _extract_params_to_drop_from_error_text(self, error_text: str) -> Optional[List[str]]: """ Error text looks like this" "Extra parameters ['stream_options', 'extra-parameters'] are not allowed when extra-parameters is not set or set to be 'error'. diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index ecb36b20427..9965aa693c3 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -43,18 +43,11 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: - return ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("AZURE_AI_API_KEY") - ) + return api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_AI_API_KEY") @property def api_version(self, api_version: Optional[str] = None) -> Optional[str]: - api_version = ( - api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") - ) + api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") return api_version def get_token_counter(self) -> Optional[BaseTokenCounter]: @@ -73,9 +66,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): return AzureAIAnthropicTokenCounter() return None - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: """ Returns a list of models supported by Azure AI. @@ -171,6 +162,4 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): api_base: Optional[str] = None, ) -> dict: """Azure Foundry sends api key in query params""" - raise NotImplementedError( - "Azure Foundry does not support environment validation" - ) + raise NotImplementedError("Azure Foundry does not support environment validation") diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 755d44fdef7..e9c8cac0078 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -28,11 +28,7 @@ def _is_azure_model_router(model: str) -> bool: bool: True if this is a model router model """ model_lower = model.lower() - return ( - "model-router" in model_lower - or "model_router" in model_lower - or model_lower == "azure-model-router" - ) + return "model-router" in model_lower or "model_router" in model_lower or model_lower == "azure-model-router" def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> float: @@ -121,9 +117,7 @@ def cost_per_token( if is_router_request: # Use the request model for flat cost calculation if available, otherwise use response model router_model_for_calc = request_model if request_model else model - router_flat_cost = calculate_azure_model_router_flat_cost( - router_model_for_calc, usage.prompt_tokens - ) + router_flat_cost = calculate_azure_model_router_flat_cost(router_model_for_calc, usage.prompt_tokens) if router_flat_cost > 0: verbose_logger.debug( diff --git a/litellm/llms/azure_ai/embed/cohere_transformation.py b/litellm/llms/azure_ai/embed/cohere_transformation.py index bbbfb60fbde..8a28d2f652a 100644 --- a/litellm/llms/azure_ai/embed/cohere_transformation.py +++ b/litellm/llms/azure_ai/embed/cohere_transformation.py @@ -29,9 +29,7 @@ class AzureAICohereConfig: return model - def _transform_request_image_embeddings( - self, input: List[str], optional_params: dict - ) -> ImageEmbeddingRequest: + def _transform_request_image_embeddings(self, input: List[str], optional_params: dict) -> ImageEmbeddingRequest: """ Assume all str in list is base64 encoded string """ @@ -60,13 +58,9 @@ class AzureAICohereConfig: image_embedding_idx.append(idx) ## REMOVE IMAGE EMBEDDINGS FROM input list - filtered_input = [ - item for idx, item in enumerate(input) if idx not in image_embedding_idx - ] + filtered_input = [item for idx, item in enumerate(input) if idx not in image_embedding_idx] - v1_embeddings_request = EmbeddingCreateParams( - input=filtered_input, model=model, **optional_params - ) + v1_embeddings_request = EmbeddingCreateParams(input=filtered_input, model=model, **optional_params) image_embeddings_request = self._transform_request_image_embeddings( input=image_embeddings, optional_params=optional_params ) @@ -74,14 +68,10 @@ class AzureAICohereConfig: return image_embeddings_request, v1_embeddings_request, image_embedding_idx def _transform_response(self, response: EmbeddingResponse) -> EmbeddingResponse: - additional_headers: Optional[dict] = response._hidden_params.get( - "additional_headers" - ) + additional_headers: Optional[dict] = response._hidden_params.get("additional_headers") if additional_headers: # CALCULATE USAGE - input_tokens: Optional[str] = additional_headers.get( - "llm_provider-num_tokens" - ) + input_tokens: Optional[str] = additional_headers.get("llm_provider-num_tokens") if input_tokens: if response.usage: response.usage.prompt_tokens = int(input_tokens) @@ -89,9 +79,7 @@ class AzureAICohereConfig: response.usage = Usage(prompt_tokens=int(input_tokens)) # SET MODEL - base_model: Optional[str] = additional_headers.get( - "llm_provider-azureml-model-group" - ) + base_model: Optional[str] = additional_headers.get("llm_provider-azureml-model-group") if base_model: response.model = self._map_azure_model_group(base_model) diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 67733d1ccb5..62c80bd2568 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -26,10 +26,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): input: List, ): combined_responses = [] - if ( - image_embedding_responses is not None - and text_embedding_responses is not None - ): + if image_embedding_responses is not None and text_embedding_responses is not None: # Combine and order the results text_idx = 0 image_idx = 0 @@ -148,9 +145,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): image_embeddings_request, v1_embeddings_request, image_embeddings_idx, - ) = AzureAICohereConfig()._transform_request( - input=input, optional_params=optional_params, model=model - ) + ) = AzureAICohereConfig()._transform_request(input=input, optional_params=optional_params, model=model) image_embedding_responses: Optional[List] = None text_embedding_responses: Optional[List] = None @@ -236,9 +231,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): image_embeddings_request, v1_embeddings_request, image_embeddings_idx, - ) = AzureAICohereConfig()._transform_request( - input=input, optional_params=optional_params, model=model - ) + ) = AzureAICohereConfig()._transform_request(input=input, optional_params=optional_params, model=model) image_embedding_responses: Optional[List] = None text_embedding_responses: Optional[List] = None @@ -270,11 +263,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): optional_params, api_key, api_base, - client=( - client - if client is not None and isinstance(client, OpenAI) - else None - ), + client=(client if client is not None and isinstance(client, OpenAI) else None), aembedding=aembedding, shared_session=shared_session, ) diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py index 75bfc913a8f..aa1092b0a53 100644 --- a/litellm/llms/azure_ai/image_edit/mai_transformation.py +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -77,13 +77,10 @@ class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): tuple(map(int, size.lower().split("x", 1))) return except ValueError: - raise ValueError( - f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." - ) + raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") raise ValueError( - f"Unsupported size value: '{size}'. " - f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + f"Unsupported size value: '{size}'. Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." ) def validate_environment( @@ -118,11 +115,7 @@ class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): "Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter." ) - api_version = ( - litellm_params.get("api_version") - or get_secret_str("AZURE_AI_API_VERSION") - or "preview" - ) + api_version = litellm_params.get("api_version") or get_secret_str("AZURE_AI_API_VERSION") or "preview" return AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( api_base=api_base, @@ -145,11 +138,7 @@ class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): if prompt is not None: request_params["prompt"] = prompt - data_without_files = { - key: value - for key, value in request_params.items() - if key not in ["image", "mask"] - } + data_without_files = {key: value for key, value in request_params.items() if key not in ["image", "mask"]} files_list: List[Tuple[str, Any]] = [] if image is not None: @@ -174,16 +163,10 @@ class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): try: response = raw_response.json() except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) if "usage" in response: - response["usage"] = ( - AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage( - response.get("usage") - ) - ) + response["usage"] = AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage(response.get("usage")) logging_obj.post_call( input="", diff --git a/litellm/llms/azure_ai/image_edit/transformation.py b/litellm/llms/azure_ai/image_edit/transformation.py index e778348c75b..5393a0ba55f 100644 --- a/litellm/llms/azure_ai/image_edit/transformation.py +++ b/litellm/llms/azure_ai/image_edit/transformation.py @@ -73,11 +73,7 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig): "Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter." ) - api_version = ( - litellm_params.get("api_version") - or litellm.api_version - or get_secret_str("AZURE_AI_API_VERSION") - ) + api_version = litellm_params.get("api_version") or litellm.api_version or get_secret_str("AZURE_AI_API_VERSION") if api_version is None: # API version is mandatory for Azure AI Foundry raise ValueError( diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index f8c876bb5be..f16afbc5971 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -34,6 +34,4 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 6a1868d94cc..a883893ceba 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -36,9 +36,7 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): Complete URL for the FLUX 2 image generation endpoint """ if api_base is None: - raise ValueError( - "api_base is required for Azure AI FLUX 2 image generation" - ) + raise ValueError("api_base is required for Azure AI FLUX 2 image generation") api_base = api_base.rstrip("/") api_version = api_version or "preview" diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 071ca9d9895..7e79ea0b976 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -126,9 +126,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): ) return normalized_usage - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: return ["n", "size"] def map_openai_params( @@ -185,9 +183,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): optional_params["width"] = width optional_params["height"] = height except ValueError: - raise ValueError( - f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." - ) + raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") else: raise ValueError( f"Unsupported size value: '{size}'. " @@ -210,9 +206,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): try: response = raw_response.json() except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) if "usage" in response: response["usage"] = self.normalize_mai_image_usage(response.get("usage")) diff --git a/litellm/llms/azure_ai/ocr/common_utils.py b/litellm/llms/azure_ai/ocr/common_utils.py index d736b891532..d1d5b80b78d 100644 --- a/litellm/llms/azure_ai/ocr/common_utils.py +++ b/litellm/llms/azure_ai/ocr/common_utils.py @@ -42,9 +42,7 @@ def get_azure_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]: # Check for Azure Document Intelligence models if "doc-intelligence" in model or "documentintelligence" in model: - verbose_logger.debug( - f"Routing {model} to Azure Document Intelligence OCR config" - ) + verbose_logger.debug(f"Routing {model} to Azure Document Intelligence OCR config") return AzureDocumentIntelligenceOCRConfig() # Default to Mistral-based OCR for other azure_ai models diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index cc65ad706ab..c67703f64d7 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -126,9 +126,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): raise ValueError("`pages` must be integers, not booleans") if all(isinstance(p, int) for p in pages): if any(p < 0 for p in pages): - raise ValueError( - "`pages` integers must be >= 0 (Mistral 0-based indices)" - ) + raise ValueError("`pages` integers must be >= 0 (Mistral 0-based indices)") # Mistral 0-based -> Azure 1-based. return ",".join(str(p + 1) for p in sorted(set(pages))) if all(isinstance(p, str) for p in pages): @@ -140,10 +138,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) return joined - raise ValueError( - "`pages` must be a list[int] (0-based, Mistral-style) or a " - "string like '1-3,5,7-9'." - ) + raise ValueError("`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'.") def validate_environment( self, @@ -294,9 +289,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): Returns: OCRRequestData with JSON data """ - verbose_logger.debug( - f"Azure Document Intelligence transform_ocr_request - model: {model}" - ) + verbose_logger.debug(f"Azure Document Intelligence transform_ocr_request - model: {model}") if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") @@ -310,9 +303,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): elif doc_type == "image_url": document_url = document.get("image_url", "") else: - raise ValueError( - f"Invalid document type: {doc_type}. Must be 'document_url' or 'image_url'" - ) + raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url' or 'image_url'") if not document_url: raise ValueError("Document URL is required") @@ -359,9 +350,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): # Join with newlines to preserve structure return "\n".join(text_lines) - def _convert_dimensions( - self, width: float, height: float, unit: str - ) -> OCRPageDimensions: + def _convert_dimensions(self, width: float, height: float, unit: str) -> OCRPageDimensions: """ Convert Azure DI dimensions to pixels. @@ -400,9 +389,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): TimeoutError: If operation has exceeded timeout """ if time.time() - start_time > timeout_secs: - raise TimeoutError( - f"Azure Document Intelligence operation polling timed out after {timeout_secs} seconds" - ) + raise TimeoutError(f"Azure Document Intelligence operation polling timed out after {timeout_secs} seconds") @staticmethod def _get_retry_after(response: httpx.Response) -> int: @@ -443,9 +430,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): return "succeeded" elif status == "failed": error_msg = result.get("error", {}).get("message", "Unknown error") - raise ValueError( - f"Azure Document Intelligence analysis failed: {error_msg}" - ) + raise ValueError(f"Azure Document Intelligence analysis failed: {error_msg}") elif status in ["running", "notStarted"]: return "running" else: @@ -596,16 +581,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): try: # Check if we got 202 Accepted (async operation started) if raw_response.status_code == 202: - verbose_logger.debug( - "Azure DI returned 202 Accepted, polling operation..." - ) + verbose_logger.debug("Azure DI returned 202 Accepted, polling operation...") # Get Operation-Location header operation_url = raw_response.headers.get("Operation-Location") if not operation_url: - raise ValueError( - "Azure Document Intelligence returned 202 but no Operation-Location header found" - ) + raise ValueError("Azure Document Intelligence returned 202 but no Operation-Location header found") # Reject cross-origin polling URLs — the auth headers # below would otherwise leak to whatever URL the upstream @@ -613,15 +594,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): try: assert_same_origin(operation_url, str(raw_response.request.url)) except SSRFError as ssrf_err: - raise ValueError( - f"Azure Document Intelligence: rejected polling URL ({ssrf_err})" - ) + raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") # Get headers for polling (need auth) poll_headers = { - "Ocp-Apim-Subscription-Key": raw_response.request.headers.get( - "Ocp-Apim-Subscription-Key", "" - ) + "Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "") } # Get timeout from kwargs or use default @@ -637,16 +614,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): # Now parse the completed response response_json = raw_response.json() - verbose_logger.debug( - f"Azure Document Intelligence response status: {response_json.get('status')}" - ) + verbose_logger.debug(f"Azure Document Intelligence response status: {response_json.get('status')}") # Check if request succeeded status = response_json.get("status") if status != "succeeded": - raise ValueError( - f"Azure Document Intelligence analysis failed with status: {status}" - ) + raise ValueError(f"Azure Document Intelligence analysis failed with status: {status}") # Extract analyze result analyze_result = response_json.get("analyzeResult", {}) @@ -665,20 +638,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): width = azure_page.get("width", 8.5) height = azure_page.get("height", 11) unit = azure_page.get("unit", "inch") - dimensions = self._convert_dimensions( - width=width, height=height, unit=unit - ) + dimensions = self._convert_dimensions(width=width, height=height, unit=unit) # Build OCR page - ocr_page = OCRPage( - index=index, markdown=markdown, dimensions=dimensions - ) + ocr_page = OCRPage(index=index, markdown=markdown, dimensions=dimensions) mistral_pages.append(ocr_page) # Build usage info - usage_info = OCRUsageInfo( - pages_processed=len(mistral_pages), doc_size_bytes=None - ) + usage_info = OCRUsageInfo(pages_processed=len(mistral_pages), doc_size_bytes=None) # Return Mistral OCR response return OCRResponse( @@ -689,9 +656,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) except Exception as e: - verbose_logger.error( - f"Error parsing Azure Document Intelligence response: {e}" - ) + verbose_logger.error(f"Error parsing Azure Document Intelligence response: {e}") raise e async def async_transform_ocr_response( @@ -718,30 +683,22 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): try: # Check if we got 202 Accepted (async operation started) if raw_response.status_code == 202: - verbose_logger.debug( - "Azure DI returned 202 Accepted, polling operation (async)..." - ) + verbose_logger.debug("Azure DI returned 202 Accepted, polling operation (async)...") # Get Operation-Location header operation_url = raw_response.headers.get("Operation-Location") if not operation_url: - raise ValueError( - "Azure Document Intelligence returned 202 but no Operation-Location header found" - ) + raise ValueError("Azure Document Intelligence returned 202 but no Operation-Location header found") # Reject cross-origin polling URLs (see sync path). VERIA-51. try: assert_same_origin(operation_url, str(raw_response.request.url)) except SSRFError as ssrf_err: - raise ValueError( - f"Azure Document Intelligence: rejected polling URL ({ssrf_err})" - ) + raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") # Get headers for polling (need auth) poll_headers = { - "Ocp-Apim-Subscription-Key": raw_response.request.headers.get( - "Ocp-Apim-Subscription-Key", "" - ) + "Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "") } # Get timeout from kwargs or use default @@ -757,16 +714,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): # Now parse the completed response response_json = raw_response.json() - verbose_logger.debug( - f"Azure Document Intelligence response status: {response_json.get('status')}" - ) + verbose_logger.debug(f"Azure Document Intelligence response status: {response_json.get('status')}") # Check if request succeeded status = response_json.get("status") if status != "succeeded": - raise ValueError( - f"Azure Document Intelligence analysis failed with status: {status}" - ) + raise ValueError(f"Azure Document Intelligence analysis failed with status: {status}") # Extract analyze result analyze_result = response_json.get("analyzeResult", {}) @@ -785,20 +738,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): width = azure_page.get("width", 8.5) height = azure_page.get("height", 11) unit = azure_page.get("unit", "inch") - dimensions = self._convert_dimensions( - width=width, height=height, unit=unit - ) + dimensions = self._convert_dimensions(width=width, height=height, unit=unit) # Build OCR page - ocr_page = OCRPage( - index=index, markdown=markdown, dimensions=dimensions - ) + ocr_page = OCRPage(index=index, markdown=markdown, dimensions=dimensions) mistral_pages.append(ocr_page) # Build usage info - usage_info = OCRUsageInfo( - pages_processed=len(mistral_pages), doc_size_bytes=None - ) + usage_info = OCRUsageInfo(pages_processed=len(mistral_pages), doc_size_bytes=None) # Return Mistral OCR response return OCRResponse( @@ -809,7 +756,5 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) except Exception as e: - verbose_logger.error( - f"Error parsing Azure Document Intelligence response (async): {e}" - ) + verbose_logger.error(f"Error parsing Azure Document Intelligence response (async): {e}") raise e diff --git a/litellm/llms/azure_ai/ocr/transformation.py b/litellm/llms/azure_ai/ocr/transformation.py index ee35fc28994..a57e3e869cf 100644 --- a/litellm/llms/azure_ai/ocr/transformation.py +++ b/litellm/llms/azure_ai/ocr/transformation.py @@ -119,17 +119,13 @@ class AzureAIOCRConfig(MistralOCRConfig): Returns: Base64 data URI string """ - verbose_logger.debug( - f"Azure AI OCR: Converting URL to base64 data URI (sync): {url}" - ) + verbose_logger.debug(f"Azure AI OCR: Converting URL to base64 data URI (sync): {url}") # Fetch and convert to base64 data URI # convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = convert_url_to_base64(url=url) - verbose_logger.debug( - f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})" - ) + verbose_logger.debug(f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})") return data_uri @@ -146,17 +142,13 @@ class AzureAIOCRConfig(MistralOCRConfig): Returns: Base64 data URI string """ - verbose_logger.debug( - f"Azure AI OCR: Converting URL to base64 data URI (async): {url}" - ) + verbose_logger.debug(f"Azure AI OCR: Converting URL to base64 data URI (async): {url}") # Fetch and convert to base64 data URI asynchronously # async_convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = await async_convert_url_to_base64(url=url) - verbose_logger.debug( - f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})" - ) + verbose_logger.debug(f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})") return data_uri @@ -184,9 +176,7 @@ class AzureAIOCRConfig(MistralOCRConfig): Returns: OCRRequestData with JSON data """ - verbose_logger.debug( - f"Azure AI OCR transform_ocr_request (sync) - model: {model}" - ) + verbose_logger.debug(f"Azure AI OCR transform_ocr_request (sync) - model: {model}") if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") @@ -199,18 +189,14 @@ class AzureAIOCRConfig(MistralOCRConfig): document_url = document.get("document_url", "") # If it's not already a data URI, convert it if document_url and not document_url.startswith("data:"): - verbose_logger.debug( - "Azure AI OCR: Converting document URL to base64 data URI (sync)" - ) + verbose_logger.debug("Azure AI OCR: Converting document URL to base64 data URI (sync)") data_uri = self._convert_url_to_data_uri_sync(url=document_url) transformed_document["document_url"] = data_uri elif doc_type == "image_url": image_url = document.get("image_url", "") # If it's not already a data URI, convert it if image_url and not image_url.startswith("data:"): - verbose_logger.debug( - "Azure AI OCR: Converting image URL to base64 data URI (sync)" - ) + verbose_logger.debug("Azure AI OCR: Converting image URL to base64 data URI (sync)") data_uri = self._convert_url_to_data_uri_sync(url=image_url) transformed_document["image_url"] = data_uri @@ -247,9 +233,7 @@ class AzureAIOCRConfig(MistralOCRConfig): Returns: OCRRequestData with JSON data """ - verbose_logger.debug( - f"Azure AI OCR async_transform_ocr_request - model: {model}" - ) + verbose_logger.debug(f"Azure AI OCR async_transform_ocr_request - model: {model}") if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") @@ -262,18 +246,14 @@ class AzureAIOCRConfig(MistralOCRConfig): document_url = document.get("document_url", "") # If it's not already a data URI, convert it if document_url and not document_url.startswith("data:"): - verbose_logger.debug( - "Azure AI OCR: Converting document URL to base64 data URI (async)" - ) + verbose_logger.debug("Azure AI OCR: Converting document URL to base64 data URI (async)") data_uri = await self._convert_url_to_data_uri_async(url=document_url) transformed_document["document_url"] = data_uri elif doc_type == "image_url": image_url = document.get("image_url", "") # If it's not already a data URI, convert it if image_url and not image_url.startswith("data:"): - verbose_logger.debug( - "Azure AI OCR: Converting image URL to base64 data URI (async)" - ) + verbose_logger.debug("Azure AI OCR: Converting image URL to base64 data URI (async)") data_uri = await self._convert_url_to_data_uri_async(url=image_url) transformed_document["image_url"] = data_uri diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index f64133afa8b..928f53bd485 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -41,9 +41,7 @@ class AzureAIRerankConfig(CohereRerankConfig): # Allow callers to pass either full v1/v2 rerank endpoints: # - https://.services.ai.azure.com/v1/rerank # - https://.services.ai.azure.com/providers/cohere/v2/rerank - if normalized_path.endswith("/v1/rerank") or normalized_path.endswith( - "/v2/rerank" - ): + if normalized_path.endswith("/v1/rerank") or normalized_path.endswith("/v2/rerank"): return str(original_url.copy_with(path=normalized_path or "/")) # If callers pass just the version path (e.g. ".../v2" or ".../providers/cohere/v2"), append "/rerank" @@ -71,9 +69,7 @@ class AzureAIRerankConfig(CohereRerankConfig): api_key = get_secret_str("AZURE_AI_API_KEY") or litellm.azure_key if api_key is None: - raise ValueError( - "Azure AI API key is required. Please set 'AZURE_AI_API_KEY' or 'litellm.azure_key'" - ) + raise ValueError("Azure AI API key is required. Please set 'AZURE_AI_API_KEY' or 'litellm.azure_key'") default_headers = { "Authorization": f"Bearer {api_key}", @@ -109,9 +105,7 @@ class AzureAIRerankConfig(CohereRerankConfig): optional_params=optional_params, litellm_params=litellm_params, ) - base_model = self._get_base_model( - rerank_response._hidden_params.get("llm_provider-azureml-model-group") - ) + base_model = self._get_base_model(rerank_response._hidden_params.get("llm_provider-azureml-model-group")) rerank_response._hidden_params["model"] = base_model return rerank_response diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index d1b93c9e7a3..da6a4a93cd8 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -42,9 +42,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): "write": [("PUT", "/docs")], } - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: api_key = litellm_params.get("api_key") if api_key is None: raise ValueError("api_key is required") @@ -55,9 +53,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): } } - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: basic_headers = self._base_validate_azure_environment(headers, litellm_params) basic_headers.update({"Content-Type": "application/json"}) return basic_headers @@ -252,7 +248,5 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): ) -> Tuple[str, Dict]: raise NotImplementedError - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: raise NotImplementedError diff --git a/litellm/llms/base.py b/litellm/llms/base.py index d639c91c145..56d1643dd4e 100644 --- a/litellm/llms/base.py +++ b/litellm/llms/base.py @@ -80,12 +80,8 @@ class BaseLLM: ) -> Optional[Any]: # set up the environment required to run the model return None - def completion( - self, *args, **kwargs - ) -> Any: # logic for parsing in - calling - parsing out model completion calls + def completion(self, *args, **kwargs) -> Any: # logic for parsing in - calling - parsing out model completion calls return None - def embedding( - self, *args, **kwargs - ) -> Any: # logic for parsing in - calling - parsing out model embedding calls + def embedding(self, *args, **kwargs) -> Any: # logic for parsing in - calling - parsing out model embedding calls return None diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 49aa563781f..7f8403c0223 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -117,9 +117,7 @@ class BaseAnthropicMessagesConfig(ABC): ) -> "BaseLLMException": from litellm.llms.base_llm.chat.transformation import BaseLLMException - return BaseLLMException( - message=error_message, status_code=status_code, headers=headers - ) + return BaseLLMException(message=error_message, status_code=status_code, headers=headers) @property def max_retry_on_anthropic_messages_http_error(self) -> int: @@ -130,9 +128,7 @@ class BaseAnthropicMessagesConfig(ABC): """ return 2 - def should_retry_anthropic_messages_on_http_error( - self, e: httpx.HTTPStatusError, litellm_params: dict - ) -> bool: + def should_retry_anthropic_messages_on_http_error(self, e: httpx.HTTPStatusError, litellm_params: dict) -> bool: """ When True, async_anthropic_messages_handler will transform the request body and issue one more attempt (bounded by max_retry_on_anthropic_messages_http_error). @@ -141,14 +137,9 @@ class BaseAnthropicMessagesConfig(ABC): is_anthropic_invalid_thinking_signature_error, ) - return ( - e.response.status_code == 400 - and is_anthropic_invalid_thinking_signature_error(e.response.text) - ) + return e.response.status_code == 400 and is_anthropic_invalid_thinking_signature_error(e.response.text) - def transform_anthropic_messages_request_on_http_error( - self, e: httpx.HTTPStatusError, request_data: dict - ) -> dict: + def transform_anthropic_messages_request_on_http_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict: """ Mutates request_data in place when retrying after a recoverable HTTP error. """ @@ -157,9 +148,6 @@ class BaseAnthropicMessagesConfig(ABC): strip_thinking_blocks_from_anthropic_messages_request_dict, ) - if ( - e.response.status_code == 400 - and is_anthropic_invalid_thinking_signature_error(e.response.text) - ): + if e.response.status_code == 400 and is_anthropic_invalid_thinking_signature_error(e.response.text): strip_thinking_blocks_from_anthropic_messages_request_dict(request_data) return request_data diff --git a/litellm/llms/base_llm/audio_transcription/transformation.py b/litellm/llms/base_llm/audio_transcription/transformation.py index 3574996e48e..dc862b3dd92 100644 --- a/litellm/llms/base_llm/audio_transcription/transformation.py +++ b/litellm/llms/base_llm/audio_transcription/transformation.py @@ -37,9 +37,7 @@ class AudioTranscriptionRequestData: class BaseAudioTranscriptionConfig(BaseConfig, ABC): @abstractmethod - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: pass def get_complete_url( diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 422ae947997..905a3ebda42 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -60,15 +60,11 @@ def convert_model_response_to_streaming( setattr(processed_chunk, "usage", usage) return processed_chunk except Exception as e: - raise ValueError( - f"Failed to convert ModelResponse to ModelResponseStream: {model_response}. Error: {e}" - ) + raise ValueError(f"Failed to convert ModelResponse to ModelResponseStream: {model_response}. Error: {e}") class BaseModelResponseIterator: - def __init__( - self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False - ): + def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): self.streaming_response = streaming_response self.response_iterator = self.streaming_response self.json_mode = json_mode @@ -85,9 +81,7 @@ class BaseModelResponseIterator: if self.http_response is not None: await self.http_response.aclose() - def chunk_parser( - self, chunk: dict - ) -> Union[GenericStreamingChunk, ModelResponseStream]: + def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]: return GenericStreamingChunk( text="", is_finished=False, @@ -104,9 +98,7 @@ class BaseModelResponseIterator: @staticmethod def _string_to_dict_parser(str_line: str) -> Optional[dict]: stripped_json_chunk: Optional[dict] = None - stripped_chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk( - str_line - ) + stripped_chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(str_line) try: if stripped_chunk is not None: stripped_json_chunk = json.loads(stripped_chunk) @@ -116,13 +108,9 @@ class BaseModelResponseIterator: stripped_json_chunk = None return stripped_json_chunk - def _handle_string_chunk( - self, str_line: str - ) -> Union[GenericStreamingChunk, ModelResponseStream]: + def _handle_string_chunk(self, str_line: str) -> Union[GenericStreamingChunk, ModelResponseStream]: # chunk is a str at this point - stripped_json_chunk = BaseModelResponseIterator._string_to_dict_parser( - str_line=str_line - ) + stripped_json_chunk = BaseModelResponseIterator._string_to_dict_parser(str_line=str_line) if "[DONE]" in str_line: return GenericStreamingChunk( text="", @@ -172,9 +160,7 @@ class BaseModelResponseIterator: except StopIteration: raise StopIteration except ValueError as e: - raise RuntimeError( - f"Error parsing chunk: {e},\nReceived chunk: {chunk}" - ) + raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") # Async iterator def __aiter__(self): @@ -212,15 +198,11 @@ class BaseModelResponseIterator: except StopAsyncIteration: raise StopAsyncIteration except ValueError as e: - raise RuntimeError( - f"Error parsing chunk: {e},\nReceived chunk: {chunk}" - ) + raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") class MockResponseIterator: # for returning ai21 streaming responses - def __init__( - self, model_response: ModelResponse, json_mode: Optional[bool] = False - ): + def __init__(self, model_response: ModelResponse, json_mode: Optional[bool] = False): self.model_response = model_response self.json_mode = json_mode self.is_done = False diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index d2d3d5c0a96..8eded37595b 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -51,9 +51,7 @@ class BaseLLMModelInfo(ABC): return None @abstractmethod - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: """ Returns a list of models supported by this provider. """ @@ -132,9 +130,7 @@ def _convert_tool_response_to_message( return None -def _dict_to_response_format_helper( - response_format: dict, ref_template: Optional[str] = None -) -> dict: +def _dict_to_response_format_helper(response_format: dict, ref_template: Optional[str] = None) -> dict: if ref_template is not None and response_format.get("type") == "json_schema": # Deep copy to avoid modifying original modified_format = copy.deepcopy(response_format) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 4f7e98af780..ab901a467e8 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -63,19 +63,13 @@ class BaseLLMException(Exception): if request: self.request = request else: - self.request = httpx.Request( - method="POST", url="https://docs.litellm.ai/docs" - ) + self.request = httpx.Request(method="POST", url="https://docs.litellm.ai/docs") if response: self.response = response else: - self.response = httpx.Response( - status_code=status_code, request=self.request - ) + self.response = httpx.Response(status_code=status_code, request=self.request) self.body = body - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class BaseConfig(ABC): @@ -108,22 +102,17 @@ class BaseConfig(ABC): return type_to_response_format_param(response_format=response_format) def is_thinking_enabled(self, non_default_params: dict) -> bool: - return (non_default_params.get("thinking") or {}).get( - "type" - ) == "enabled" or non_default_params.get("reasoning_effort") is not None + return (non_default_params.get("thinking") or {}).get("type") == "enabled" or non_default_params.get( + "reasoning_effort" + ) is not None def is_max_tokens_in_request(self, non_default_params: dict) -> bool: """ OpenAI spec allows max_tokens or max_completion_tokens to be specified. """ - return ( - "max_tokens" in non_default_params - or "max_completion_tokens" in non_default_params - ) + return "max_tokens" in non_default_params or "max_completion_tokens" in non_default_params - def update_optional_params_with_thinking_tokens( - self, non_default_params: dict, optional_params: dict - ): + def update_optional_params_with_thinking_tokens(self, non_default_params: dict, optional_params: dict): """ Handles scenario where max tokens is not specified. For anthropic models (anthropic api/bedrock/vertex ai), this requires having the max tokens being set and being greater than the thinking token budget. @@ -133,16 +122,11 @@ class BaseConfig(ABC): """ is_thinking_enabled = self.is_thinking_enabled(optional_params) if is_thinking_enabled and ( - "max_tokens" not in non_default_params - and "max_completion_tokens" not in non_default_params + "max_tokens" not in non_default_params and "max_completion_tokens" not in non_default_params ): - thinking_token_budget = cast(dict, optional_params["thinking"]).get( - "budget_tokens", None - ) + thinking_token_budget = cast(dict, optional_params["thinking"]).get("budget_tokens", None) if thinking_token_budget is not None: - optional_params["max_tokens"] = ( - thinking_token_budget + DEFAULT_MAX_TOKENS - ) + optional_params["max_tokens"] = thinking_token_budget + DEFAULT_MAX_TOKENS def should_fake_stream( self, @@ -189,9 +173,7 @@ class BaseConfig(ABC): """ return False - def transform_request_on_unprocessable_entity_error( - self, e: httpx.HTTPStatusError, request_data: dict - ) -> dict: + def transform_request_on_unprocessable_entity_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict: """ Transform the request data on UnprocessableEntityError """ @@ -238,16 +220,12 @@ class BaseConfig(ABC): if json_schema and not is_response_format_supported: _tool_choice = ChatCompletionToolChoiceObjectParam( type="function", - function=ChatCompletionToolChoiceFunctionParam( - name=RESPONSE_FORMAT_TOOL_NAME - ), + function=ChatCompletionToolChoiceFunctionParam(name=RESPONSE_FORMAT_TOOL_NAME), ) _tool = ChatCompletionToolParam( type="function", - function=ChatCompletionToolParamFunctionChunk( - name=RESPONSE_FORMAT_TOOL_NAME, parameters=json_schema - ), + function=ChatCompletionToolParamFunctionChunk(name=RESPONSE_FORMAT_TOOL_NAME, parameters=json_schema), ) optional_params.setdefault("tools", []) @@ -461,9 +439,7 @@ class BaseConfig(ABC): """Hook for providers to merge chunk metadata into assembled streaming responses.""" return None - def calculate_additional_costs( - self, model: str, prompt_tokens: int, completion_tokens: int - ) -> Optional[dict]: + def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> Optional[dict]: """ Calculate any additional costs beyond standard token costs. diff --git a/litellm/llms/base_llm/embedding/transformation.py b/litellm/llms/base_llm/embedding/transformation.py index c03a8235b4a..07ffbb99626 100644 --- a/litellm/llms/base_llm/embedding/transformation.py +++ b/litellm/llms/base_llm/embedding/transformation.py @@ -66,9 +66,7 @@ class BaseEmbeddingConfig(BaseConfig, ABC): litellm_params: dict, headers: dict, ) -> dict: - raise NotImplementedError( - "EmbeddingConfig does not need a request transformation for chat models" - ) + raise NotImplementedError("EmbeddingConfig does not need a request transformation for chat models") def transform_response( self, @@ -84,6 +82,4 @@ class BaseEmbeddingConfig(BaseConfig, ABC): api_key: Optional[str] = None, json_mode: Optional[bool] = None, ) -> ModelResponse: - raise NotImplementedError( - "EmbeddingConfig does not need a response transformation for chat models" - ) + raise NotImplementedError("EmbeddingConfig does not need a response transformation for chat models") diff --git a/litellm/llms/base_llm/evals/transformation.py b/litellm/llms/base_llm/evals/transformation.py index 54dc2f7aae9..da8d7e12acb 100644 --- a/litellm/llms/base_llm/evals/transformation.py +++ b/litellm/llms/base_llm/evals/transformation.py @@ -46,9 +46,7 @@ class BaseEvalsAPIConfig(ABC): pass @abstractmethod - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate and update headers with provider-specific requirements diff --git a/litellm/llms/base_llm/files/azure_blob_storage_backend.py b/litellm/llms/base_llm/files/azure_blob_storage_backend.py index a2155df4047..07dd339cac3 100644 --- a/litellm/llms/base_llm/files/azure_blob_storage_backend.py +++ b/litellm/llms/base_llm/files/azure_blob_storage_backend.py @@ -78,25 +78,19 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): # Do nothing - this class is used for file storage, not logging pass - def _generate_file_name( - self, original_filename: str, file_naming_strategy: str - ) -> str: + def _generate_file_name(self, original_filename: str, file_naming_strategy: str) -> str: """Generate file name based on naming strategy.""" if file_naming_strategy == "original_filename": # Use original filename, but sanitize it return quote(original_filename, safe="") elif file_naming_strategy == "timestamp": # Use timestamp - extension = ( - original_filename.split(".")[-1] if "." in original_filename else "" - ) + extension = original_filename.split(".")[-1] if "." in original_filename else "" timestamp = int(time.time() * 1000) # milliseconds return f"{timestamp}.{extension}" if extension else str(timestamp) else: # default to "uuid" # Use UUID - extension = ( - original_filename.split(".")[-1] if "." in original_filename else "" - ) + extension = original_filename.split(".")[-1] if "." in original_filename else "" file_uuid = str(uuid.uuid4()) return f"{file_uuid}.{extension}" if extension else file_uuid @@ -138,33 +132,23 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): full_path=full_path, ) - verbose_logger.debug( - f"Successfully uploaded file to Azure Blob Storage: {storage_url}" - ) + verbose_logger.debug(f"Successfully uploaded file to Azure Blob Storage: {storage_url}") return storage_url except Exception as e: - verbose_logger.exception( - f"Error uploading file to Azure Blob Storage: {str(e)}" - ) + verbose_logger.exception(f"Error uploading file to Azure Blob Storage: {str(e)}") raise - async def _upload_file_with_account_key( - self, file_content: bytes, full_path: str - ) -> str: + async def _upload_file_with_account_key(self, file_content: bytes, full_path: str) -> str: """Upload file using Azure SDK with account key authentication.""" # Reuse the logger's service client method service_client = await self.get_service_client() - file_system_client = service_client.get_file_system_client( - file_system=self.azure_storage_file_system - ) + file_system_client = service_client.get_file_system_client(file_system=self.azure_storage_file_system) # Create filesystem (container) if it doesn't exist if not await file_system_client.exists(): await file_system_client.create_file_system() - verbose_logger.debug( - f"Created filesystem: {self.azure_storage_file_system}" - ) + verbose_logger.debug(f"Created filesystem: {self.azure_storage_file_system}") # Extract directory and filename (similar to logger's pattern) path_parts = full_path.split("/") @@ -186,18 +170,14 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): # Create, append, and flush (same pattern as logger's upload_to_azure_data_lake_with_azure_account_key) await file_client.create_file() - await file_client.append_data( - data=file_content, offset=0, length=len(file_content) - ) + await file_client.append_data(data=file_content, offset=0, length=len(file_content)) await file_client.flush_data(position=len(file_content), offset=0) # Return blob URL (not DFS URL) blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{full_path}" return blob_url - async def _upload_file_with_azure_ad( - self, file_content: bytes, full_path: str - ) -> str: + async def _upload_file_with_azure_ad(self, file_content: bytes, full_path: str) -> str: """Upload file using REST API with Azure AD authentication.""" # Reuse the logger's token management await self.set_valid_azure_ad_token() @@ -207,9 +187,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): httpxSpecialProvider, ) - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Use DFS endpoint for upload base_url = f"https://{self.azure_storage_account_name}.dfs.core.windows.net/{self.azure_storage_file_system}/{full_path}" @@ -261,9 +239,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): container_and_path = storage_url.split(".blob.core.windows.net/", 1)[1] path_parts = container_and_path.split("/", 1) if len(path_parts) < 2: - raise ValueError( - f"Invalid Azure Blob Storage URL format: {storage_url}" - ) + raise ValueError(f"Invalid Azure Blob Storage URL format: {storage_url}") file_path = path_parts[1] # Path after container name if self.azure_storage_account_key: @@ -274,23 +250,17 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): return await self._download_file_with_azure_ad(file_path) except Exception as e: - verbose_logger.exception( - f"Error downloading file from Azure Blob Storage: {str(e)}" - ) + verbose_logger.exception(f"Error downloading file from Azure Blob Storage: {str(e)}") raise async def _download_file_with_account_key(self, file_path: str) -> bytes: """Download file using Azure SDK with account key.""" # Reuse the logger's service client method service_client = await self.get_service_client() - file_system_client = service_client.get_file_system_client( - file_system=self.azure_storage_file_system - ) + file_system_client = service_client.get_file_system_client(file_system=self.azure_storage_file_system) # Ensure filesystem exists (should already exist, but check for safety) if not await file_system_client.exists(): - raise ValueError( - f"Filesystem {self.azure_storage_file_system} does not exist" - ) + raise ValueError(f"Filesystem {self.azure_storage_file_system} does not exist") file_client = file_system_client.get_file_client(file_path) # Download file download_response = await file_client.download_file() @@ -308,9 +278,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): ) from litellm.constants import AZURE_STORAGE_MSFT_VERSION - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Use blob endpoint for download (simpler than DFS) blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{file_path}" diff --git a/litellm/llms/base_llm/files/storage_backend_factory.py b/litellm/llms/base_llm/files/storage_backend_factory.py index 12047f1122e..8fd918af0dc 100644 --- a/litellm/llms/base_llm/files/storage_backend_factory.py +++ b/litellm/llms/base_llm/files/storage_backend_factory.py @@ -34,7 +34,4 @@ def get_storage_backend(backend_type: str) -> BaseFileStorageBackend: if backend_type == "azure_storage": return AzureBlobStorageBackend() else: - raise ValueError( - f"Unsupported storage backend type: {backend_type}. " - f"Supported types: azure_storage" - ) + raise ValueError(f"Unsupported storage backend type: {backend_type}. Supported types: azure_storage") diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 85016c7a5c4..a9b99eb06fc 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -65,9 +65,7 @@ class BaseFilesConfig(BaseConfig): return "POST" @abstractmethod - def get_supported_openai_params( - self, model: str - ) -> List[OpenAICreateFileRequestOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAICreateFileRequestOptionalParams]: pass def get_complete_file_url( diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 8fb7eb9fde0..965c174df6e 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -58,9 +58,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): Returns: List of supported parameter names """ - raise NotImplementedError( - "get_supported_generate_content_optional_params is not implemented" - ) + raise NotImplementedError("get_supported_generate_content_optional_params is not implemented") def get_generate_content_request_top_level_fields(self) -> tuple[str, ...]: """ @@ -91,9 +89,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): Returns: Mapped parameters for the provider """ - raise NotImplementedError( - "map_generate_content_optional_params is not implemented" - ) + raise NotImplementedError("map_generate_content_optional_params is not implemented") @abstractmethod def validate_environment( @@ -201,9 +197,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): """ pass - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] - ) -> Exception: + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]) -> Exception: """ Get the appropriate exception class for the error. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 1efeb159a3e..68db36b529e 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -29,11 +29,7 @@ class BaseTranslation(ABC): return {} # Convert to dict if it's a Pydantic object - user_dict = ( - user_api_key_dict.model_dump() - if hasattr(user_api_key_dict, "model_dump") - else user_api_key_dict - ) + user_dict = user_api_key_dict.model_dump() if hasattr(user_api_key_dict, "model_dump") else user_api_key_dict if not isinstance(user_dict, dict): return {} diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index 92429573ff8..4c18702bc6c 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -102,9 +102,7 @@ class BaseImageEditConfig(ABC): ) -> Tuple[Dict, RequestFiles]: pass - def finalize_image_edit_request_data( - self, data: dict, resolved_request_url: str - ) -> dict: + def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict: """ Last pass on the request dict after ``transform_image_edit_request``, using the exact URL string used for the HTTP POST (same as ``get_complete_url`` output). diff --git a/litellm/llms/base_llm/image_generation/transformation.py b/litellm/llms/base_llm/image_generation/transformation.py index 7f13e6f3b4c..e80a970d806 100644 --- a/litellm/llms/base_llm/image_generation/transformation.py +++ b/litellm/llms/base_llm/image_generation/transformation.py @@ -20,9 +20,7 @@ else: class BaseImageGenerationConfig(ABC): @abstractmethod - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: pass @abstractmethod diff --git a/litellm/llms/base_llm/image_variations/transformation.py b/litellm/llms/base_llm/image_variations/transformation.py index 60444d0fb74..23fc4dc88b9 100644 --- a/litellm/llms/base_llm/image_variations/transformation.py +++ b/litellm/llms/base_llm/image_variations/transformation.py @@ -26,9 +26,7 @@ else: class BaseImageVariationConfig(BaseConfig, ABC): @abstractmethod - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageVariationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageVariationOptionalParams]: pass def get_complete_url( diff --git a/litellm/llms/base_llm/interactions/transformation.py b/litellm/llms/base_llm/interactions/transformation.py index be400628fd5..3eba1858a23 100644 --- a/litellm/llms/base_llm/interactions/transformation.py +++ b/litellm/llms/base_llm/interactions/transformation.py @@ -86,9 +86,7 @@ class BaseInteractionsAPIConfig(ABC): pass @abstractmethod - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate and prepare environment settings including headers. """ diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index c0c18aefdeb..146a6aa6ae0 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -163,9 +163,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): user_api_key_dict: User API key authentication details additional_db_fields: Additional fields to store in database """ - verbose_logger.info( - f"Storing LiteLLM Managed {self.resource_type} with id={unified_resource_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed {self.resource_type} with id={unified_resource_id} in cache") # Prepare cache data cache_data = { @@ -256,9 +254,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Check database table = getattr(self.prisma_client.db, self.table_name) - db_object = await table.find_first( - where={"unified_resource_id": unified_resource_id} - ) + db_object = await table.find_first(where={"unified_resource_id": unified_resource_id}) if db_object: return db_object.model_dump() @@ -282,14 +278,10 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): """ # Get old value from database table = getattr(self.prisma_client.db, self.table_name) - initial_value = await table.find_first( - where={"unified_resource_id": unified_resource_id} - ) + initial_value = await table.find_first(where={"unified_resource_id": unified_resource_id}) if initial_value is None: - raise Exception( - f"LiteLLM Managed {self.resource_type} with id={unified_resource_id} not found" - ) + raise Exception(f"LiteLLM Managed {self.resource_type} with id={unified_resource_id} not found") # Delete from cache await self.internal_usage_cache.async_set_cache( @@ -324,9 +316,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): True if user has access, False otherwise """ # Use cached method instead of direct DB query - resource = await self.get_unified_resource_id( - unified_resource_id, litellm_parent_otel_span - ) + resource = await self.get_unified_resource_id(unified_resource_id, litellm_parent_otel_span) if resource: return can_access_resource( @@ -368,9 +358,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): for resource_id in resource_ids: # Get unified resource from cache/db - unified_resource_object = await self.get_unified_resource_id( - resource_id, litellm_parent_otel_span - ) + unified_resource_object = await self.get_unified_resource_id(resource_id, litellm_parent_otel_span) if unified_resource_object: model_mappings = unified_resource_object.get("model_mappings", {}) @@ -442,9 +430,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) # Convert to URL-safe base64 and strip padding - base64_unified_id = ( - base64.urlsafe_b64encode(unified_id_format.encode()).decode().rstrip("=") - ) + base64_unified_id = base64.urlsafe_b64encode(unified_id_format.encode()).decode().rstrip("=") return base64_unified_id @@ -468,9 +454,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): hidden_params = getattr(resource_object, "_hidden_params", {}) or {} model_resource_id_mapping = hidden_params.get("model_resource_id_mapping") - if model_resource_id_mapping and isinstance( - model_resource_id_mapping, dict - ): + if model_resource_id_mapping and isinstance(model_resource_id_mapping, dict): model_mappings.update(model_resource_id_mapping) return model_mappings @@ -602,11 +586,8 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): except Exception as e: verbose_logger.warning( - f"Failed to parse {self.resource_type} object " - f"{resource.unified_resource_id}: {e}" + f"Failed to parse {self.resource_type} object {resource.unified_resource_id}: {e}" ) continue - return build_list_page( - resource_objects, has_more=len(resource_objects) == (limit or 20) - ) + return build_list_page(resource_objects, has_more=len(resource_objects) == (limit or 20)) diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index 62027f4272c..fd1e24f3e1d 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -89,11 +89,7 @@ def can_access_resource( return True team_id = user_api_key_dict.team_id - if ( - team_id is not None - and resource_team_id is not None - and resource_team_id == team_id - ): + if team_id is not None and resource_team_id is not None and resource_team_id == team_id: return True return False diff --git a/litellm/llms/base_llm/managed_resources/utils.py b/litellm/llms/base_llm/managed_resources/utils.py index e9a6aef689e..a93f62764f9 100644 --- a/litellm/llms/base_llm/managed_resources/utils.py +++ b/litellm/llms/base_llm/managed_resources/utils.py @@ -29,14 +29,10 @@ def resolve_passthrough_managed_id_provider( Splitting them would make a managed ID minted on ``azure`` fail to resolve when replayed on ``azure_ai`` and vice versa. """ - provider = str( - getattr(custom_llm_provider, "value", custom_llm_provider) or "" - ).lower() + provider = str(getattr(custom_llm_provider, "value", custom_llm_provider) or "").lower() if not provider: return None - if provider in PASSTHROUGH_MANAGED_ID_AZURE_PROVIDERS or provider.endswith( - (".azure", ".azure_ai") - ): + if provider in PASSTHROUGH_MANAGED_ID_AZURE_PROVIDERS or provider.endswith((".azure", ".azure_ai")): return "azure" if provider == "openai" or provider.endswith(".openai"): return "openai" @@ -391,12 +387,8 @@ def parse_unified_id( return { "resource_type": extract_resource_type_from_unified_id(decoded_id), "unified_uuid": extract_unified_uuid_from_unified_id(decoded_id), - "target_model_names": extract_target_model_names_from_unified_id( - decoded_id - ), - "provider_resource_id": extract_provider_resource_id_from_unified_id( - decoded_id - ), + "target_model_names": extract_target_model_names_from_unified_id(decoded_id), + "provider_resource_id": extract_provider_resource_id_from_unified_id(decoded_id), "model_id": extract_model_id_from_unified_id(decoded_id), } except Exception: diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index a2946c62506..a38e5bfdcd6 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -170,9 +170,7 @@ class BaseOCRConfig: Returns: OCRRequestData with data and files fields """ - raise NotImplementedError( - "transform_ocr_request must be implemented by provider" - ) + raise NotImplementedError("transform_ocr_request must be implemented by provider") async def async_transform_ocr_request( self, @@ -218,9 +216,7 @@ class BaseOCRConfig: Transform provider-specific OCR response to standard format. Override in provider-specific implementations. """ - raise NotImplementedError( - "transform_ocr_response must be implemented by provider" - ) + raise NotImplementedError("transform_ocr_response must be implemented by provider") async def async_transform_ocr_response( self, diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 9d4396dce47..e243d36a86a 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -95,9 +95,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): ) -> "BaseLLMException": from litellm.llms.base_llm.chat.transformation import BaseLLMException - return BaseLLMException( - status_code=status_code, message=error_message, headers=headers - ) + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) def logging_non_streaming_response( self, diff --git a/litellm/llms/base_llm/realtime/http_transformation.py b/litellm/llms/base_llm/realtime/http_transformation.py index be1413a3c0b..4c8cc30a8b3 100644 --- a/litellm/llms/base_llm/realtime/http_transformation.py +++ b/litellm/llms/base_llm/realtime/http_transformation.py @@ -54,9 +54,7 @@ class BaseRealtimeHTTPConfig(ABC): # ------------------------------------------------------------------ # @abstractmethod - def get_complete_url( - self, api_base: Optional[str], model: str, api_version: Optional[str] = None - ) -> str: + def get_complete_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: """Return the full URL for POST /realtime/client_secrets.""" def get_transcription_session_url( @@ -86,9 +84,7 @@ class BaseRealtimeHTTPConfig(ABC): # realtime_calls endpoint # # ------------------------------------------------------------------ # - def get_realtime_calls_url( - self, api_base: Optional[str], model: str, api_version: Optional[str] = None - ) -> str: + def get_realtime_calls_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: """Return the full URL for POST /realtime/calls (SDP exchange).""" base = (api_base or "").rstrip("/") return f"{base}/v1/realtime/calls" @@ -108,9 +104,7 @@ class BaseRealtimeHTTPConfig(ABC): # Error handling # # ------------------------------------------------------------------ # - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] - ): + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]): """ Map HTTP errors to LiteLLM exception types. diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index b66bdbd2dd9..c24267ccc72 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -30,9 +30,7 @@ class BaseRealtimeConfig(ABC): pass @abstractmethod - def get_complete_url( - self, api_base: Optional[str], model: str, api_key: Optional[str] = None - ) -> str: + def get_complete_url(self, api_base: Optional[str], model: str, api_key: Optional[str] = None) -> str: """ OPTIONAL @@ -71,9 +69,7 @@ class BaseRealtimeConfig(ABC): ) -> bool: # initial configuration message sent to setup the realtime session return False - def session_configuration_request( - self, model: str - ) -> Optional[str]: # message sent to setup the realtime session + def session_configuration_request(self, model: str) -> Optional[str]: # message sent to setup the realtime session return None def transform_session_created_event( diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index c61ce52b530..c6453745e5c 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -96,9 +96,7 @@ class BaseResponsesAPIConfig(ABC): pass @abstractmethod - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: return {} @abstractmethod @@ -270,9 +268,7 @@ class BaseResponsesAPIConfig(ABC): WebSocket path differs from their HTTP path (e.g. Azure uses /openai/v1/responses without api-version) should override this. """ - http_url = self.get_complete_url( - api_base=api_base, litellm_params=litellm_params - ) + http_url = self.get_complete_url(api_base=api_base, litellm_params=litellm_params) return http_url.replace("https://", "wss://").replace("http://", "ws://") def model_in_websocket_url(self) -> bool: @@ -359,7 +355,5 @@ class BaseResponsesAPIConfig(ABC): return data return { **data, - "input": BaseResponsesAPIConfig.strip_custom_tool_call_namespace_from_responses_input( - data["input"] - ), + "input": BaseResponsesAPIConfig.strip_custom_tool_call_namespace_from_responses_input(data["input"]), } diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index 1c012a15fdb..c807283ecd2 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -48,9 +48,7 @@ class BaseSandboxConfig: """Provider-agnostic sandbox operations.""" def validate_environment(self, api_key: str | None = None, **kwargs) -> str: - raise NotImplementedError( - "validate_environment must be implemented by provider" - ) + raise NotImplementedError("validate_environment must be implemented by provider") async def acreate_sandbox( self, @@ -89,8 +87,7 @@ class BaseSandboxConfig: total += len(line.encode("utf-8")) if total > SANDBOX_MAX_OUTPUT_BYTES: raise ValueError( - f"Sandbox output exceeded {SANDBOX_MAX_OUTPUT_BYTES} bytes; aborting " - "to avoid unbounded memory use." + f"Sandbox output exceeded {SANDBOX_MAX_OUTPUT_BYTES} bytes; aborting to avoid unbounded memory use." ) lines.append(line) return lines diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 1581d8bb064..fdfac6f5f9f 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -222,9 +222,7 @@ class BaseSearchConfig: Returns: Dict with request data """ - raise NotImplementedError( - "transform_search_request must be implemented by provider" - ) + raise NotImplementedError("transform_search_request must be implemented by provider") def transform_search_response( self, @@ -236,9 +234,7 @@ class BaseSearchConfig: Transform provider-specific Search response to standard format. Override in provider-specific implementations. """ - raise NotImplementedError( - "transform_search_response must be implemented by provider" - ) + raise NotImplementedError("transform_search_response must be implemented by provider") def get_error_class( self, diff --git a/litellm/llms/base_llm/skills/transformation.py b/litellm/llms/base_llm/skills/transformation.py index 017587c0b0c..5bb181f59fb 100644 --- a/litellm/llms/base_llm/skills/transformation.py +++ b/litellm/llms/base_llm/skills/transformation.py @@ -38,9 +38,7 @@ class BaseSkillsAPIConfig(ABC): pass @abstractmethod - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate and update headers with provider-specific requirements diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index 0e30ddae5fe..cbae6904ead 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -137,9 +137,7 @@ class BaseTextToSpeechConfig(ABC): """ pass - def get_error_class( - self, error_message: str, status_code: int, headers: Dict - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: Dict) -> BaseLLMException: from ..chat.transformation import BaseLLMException raise BaseLLMException( diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 85a9c838264..b222e3dd160 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -27,9 +27,7 @@ else: class BaseVectorStoreConfig: - def get_supported_openai_params( - self, model: str - ) -> List[VECTOR_STORE_OPENAI_PARAMS]: + def get_supported_openai_params(self, model: str) -> List[VECTOR_STORE_OPENAI_PARAMS]: return [] def map_openai_params( @@ -41,9 +39,7 @@ class BaseVectorStoreConfig: return optional_params @abstractmethod - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: pass @abstractmethod @@ -104,15 +100,11 @@ class BaseVectorStoreConfig: pass @abstractmethod - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: pass @abstractmethod - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: return {} @abstractmethod diff --git a/litellm/llms/base_llm/vector_store_files/transformation.py b/litellm/llms/base_llm/vector_store_files/transformation.py index 02915d013e5..e8799c56cae 100644 --- a/litellm/llms/base_llm/vector_store_files/transformation.py +++ b/litellm/llms/base_llm/vector_store_files/transformation.py @@ -52,9 +52,7 @@ class BaseVectorStoreFilesConfig(ABC): return optional_params @abstractmethod - def get_auth_credentials( - self, litellm_params: Dict[str, Any] - ) -> VectorStoreFileAuthCredentials: ... + def get_auth_credentials(self, litellm_params: Dict[str, Any]) -> VectorStoreFileAuthCredentials: ... @abstractmethod def get_vector_store_file_endpoints_by_type( diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index 9b4cf777280..e3a66af24a8 100644 --- a/litellm/llms/base_llm/videos/transformation.py +++ b/litellm/llms/base_llm/videos/transformation.py @@ -282,18 +282,14 @@ class BaseVideoConfig(ABC): Returns: Tuple[str, list]: (url, files_list) for the multipart POST request """ - raise NotImplementedError( - "video create character is not supported for this provider" - ) + raise NotImplementedError("video create character is not supported for this provider") def transform_video_create_character_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> CharacterObject: - raise NotImplementedError( - "video create character is not supported for this provider" - ) + raise NotImplementedError("video create character is not supported for this provider") def transform_video_get_character_request( self, @@ -308,18 +304,14 @@ class BaseVideoConfig(ABC): Returns: Tuple[str, Dict]: (url, params) for the GET request """ - raise NotImplementedError( - "video get character is not supported for this provider" - ) + raise NotImplementedError("video get character is not supported for this provider") def transform_video_get_character_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> CharacterObject: - raise NotImplementedError( - "video get character is not supported for this provider" - ) + raise NotImplementedError("video get character is not supported for this provider") def get_video_edit_prefetch_params( self, diff --git a/litellm/llms/baseten/chat.py b/litellm/llms/baseten/chat.py index 1e49b346088..f5d52ef81ff 100644 --- a/litellm/llms/baseten/chat.py +++ b/litellm/llms/baseten/chat.py @@ -82,9 +82,7 @@ class BasetenConfig(OpenAIGPTConfig): optional_params[param] = value return optional_params - def _get_openai_compatible_provider_info( - self, api_base: str, api_key: str - ) -> tuple: + def _get_openai_compatible_provider_info(self, api_base: str, api_key: str) -> tuple: """ Get the OpenAI compatible provider info for Baseten """ diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index b71f37023e8..380cc91ed98 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -66,13 +66,9 @@ class AwsAuthError(Exception): def __init__(self, status_code, message): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://us-west-2.console.aws.amazon.com/bedrock" - ) + self.request = httpx.Request(method="POST", url="https://us-west-2.console.aws.amazon.com/bedrock") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class BaseAWSLLM: @@ -159,11 +155,7 @@ class BaseAWSLLM: aws_role_name: Optional[str], aws_session_name: Optional[str], ) -> bool: - return ( - aws_web_identity_token is not None - and aws_role_name is not None - and aws_session_name is not None - ) + return aws_web_identity_token is not None and aws_role_name is not None and aws_session_name is not None @staticmethod def _is_auth_with_aws_role(aws_role_name: Optional[str]) -> bool: @@ -179,11 +171,7 @@ class BaseAWSLLM: aws_secret_access_key: Optional[str], aws_session_token: Optional[str], ) -> bool: - return ( - aws_access_key_id is not None - and aws_secret_access_key is not None - and aws_session_token is not None - ) + return aws_access_key_id is not None and aws_secret_access_key is not None and aws_session_token is not None @staticmethod def _is_auth_with_access_key_and_secret_key( @@ -191,11 +179,7 @@ class BaseAWSLLM: aws_secret_access_key: Optional[str], aws_region_name: Optional[str], ) -> bool: - return ( - aws_access_key_id is not None - and aws_secret_access_key is not None - and aws_region_name is not None - ) + return aws_access_key_id is not None and aws_secret_access_key is not None and aws_region_name is not None @tracer.wrap() def get_credentials( @@ -271,11 +255,7 @@ class BaseAWSLLM: aws_external_id, ) - args = { - k: v - for k, v in locals().items() - if k.startswith("aws_") or k == "ssl_verify" - } + args = {k: v for k, v in locals().items() if k.startswith("aws_") or k == "ssl_verify"} ######################################################### # Handle diff boto3 auth flows @@ -304,16 +284,12 @@ class BaseAWSLLM: elif self._is_auth_with_aws_role(aws_role_name): # Same role (IRSA/ECS/EC2): ambient creds via _get_or_set_cached_credentials like the # default env branch; never pre-read cache (must run _is_already_running_as_role first). - if self._is_already_running_as_role( - cast(str, aws_role_name), ssl_verify=ssl_verify - ): + if self._is_already_running_as_role(cast(str, aws_role_name), ssl_verify=ssl_verify): verbose_logger.debug( "Already running as target role %s, using ambient credentials", aws_role_name, ) - return self._get_or_set_cached_credentials( - args, self._auth_with_env_vars - ) + return self._get_or_set_cached_credentials(args, self._auth_with_env_vars) verbose_logger.debug("Using role assumption: calling _auth_with_aws_role") # If aws_session_name is not provided, generate a default one if aws_session_name is None: @@ -332,9 +308,7 @@ class BaseAWSLLM: return credentials elif self._is_auth_with_aws_profile(aws_profile_name): - credentials, _cache_ttl = self._auth_with_aws_profile( - cast(str, aws_profile_name) - ) + credentials, _cache_ttl = self._auth_with_aws_profile(cast(str, aws_profile_name)) return credentials elif self._is_auth_with_aws_session_token_tuple( aws_access_key_id, @@ -475,41 +449,23 @@ class BaseAWSLLM: model_id = model_id.replace("invoke/", "", 1) if provider == "llama" and "llama/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="llama" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="llama") elif provider == "deepseek_r1" and "deepseek_r1/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="deepseek_r1" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="deepseek_r1") elif provider == "openai" and "openai/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="openai" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="openai") elif provider == "qwen2" and "qwen2/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="qwen2" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="qwen2") elif provider == "qwen3" and "qwen3/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="qwen3" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="qwen3") elif provider == "stability" and "stability/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="stability" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="stability") elif provider == "moonshot" and "moonshot/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="moonshot" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="moonshot") elif "nova-2/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="nova-2" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="nova-2") elif "nova/" in model_id: - model_id = BaseAWSLLM._get_model_id_from_model_with_spec( - model_id, spec="nova" - ) + model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="nova") return model_id @staticmethod @@ -559,16 +515,12 @@ class BaseAWSLLM: parts = model.split(".") # Check if the second part (after potential region) is a known provider if len(parts) >= 2: - potential_provider = parts[ - 1 - ] # e.g., "twelvelabs" from "us.twelvelabs.marengo-embed-2-7-v1:0" + potential_provider = parts[1] # e.g., "twelvelabs" from "us.twelvelabs.marengo-embed-2-7-v1:0" if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL): return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider) # Check if the first part is a known provider (standard format) - potential_provider = parts[ - 0 - ] # e.g., "cohere" from "cohere.embed-english-v3:0" + potential_provider = parts[0] # e.g., "cohere" from "cohere.embed-english-v3:0" if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL): return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider) @@ -647,9 +599,7 @@ class BaseAWSLLM: """ if aws_region_name is None: return - if not isinstance(aws_region_name, str) or not _VALID_AWS_REGION_PATTERN.match( - aws_region_name - ): + if not isinstance(aws_region_name, str) or not _VALID_AWS_REGION_PATTERN.match(aws_region_name): raise ValueError( f"Invalid AWS region format: {aws_region_name!r}. " "Region names must contain only lowercase letters, digits, and hyphens." @@ -705,15 +655,11 @@ class BaseAWSLLM: # check env # litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) - if litellm_aws_region_name is not None and isinstance( - litellm_aws_region_name, str - ): + if litellm_aws_region_name is not None and isinstance(litellm_aws_region_name, str): aws_region_name = litellm_aws_region_name standard_aws_region_name = get_secret("AWS_REGION", None) - if standard_aws_region_name is not None and isinstance( - standard_aws_region_name, str - ): + if standard_aws_region_name is not None and isinstance(standard_aws_region_name, str): aws_region_name = standard_aws_region_name if aws_region_name is None: @@ -796,9 +742,7 @@ class BaseAWSLLM: import boto3 with tracer.trace("boto3.client(sts).get_caller_identity"): - sts_client = boto3.client( - "sts", verify=self._get_ssl_verify(ssl_verify) - ) + sts_client = boto3.client("sts", verify=self._get_ssl_verify(ssl_verify)) identity = sts_client.get_caller_identity() caller_arn = identity.get("Arn", "") @@ -817,9 +761,7 @@ class BaseAWSLLM: return True except Exception as e: - verbose_logger.debug( - "Could not determine current role identity: %s", str(e) - ) + verbose_logger.debug("Could not determine current role identity: %s", str(e)) return False @@ -867,10 +809,7 @@ class BaseAWSLLM: # references are expanded at load time, so such a reference reaching here is # caller-supplied input; reject it rather than expanding a process-environment # value for use as the token. - if ( - aws_web_identity_token.startswith("os.environ/") - or aws_web_identity_token in os.environ - ): + if aws_web_identity_token.startswith("os.environ/") or aws_web_identity_token in os.environ: raise AwsAuthError( message="Invalid web identity token reference.", status_code=400, @@ -951,15 +890,9 @@ class BaseAWSLLM: assume_role_params["ExternalId"] = aws_external_id try: - sts_response = sts_client.assume_role_with_web_identity( - **assume_role_params - ) + sts_response = sts_client.assume_role_with_web_identity(**assume_role_params) except sts_client.exceptions.InvalidIdentityTokenException as e: - audience = ( - self._unverified_web_identity_audience(oidc_token) - if isinstance(oidc_token, str) - else None - ) + audience = self._unverified_web_identity_audience(oidc_token) if isinstance(oidc_token, str) else None detail = f" Token {audience}" if audience else "" raise AwsAuthError( status_code=401, @@ -1013,9 +946,7 @@ class BaseAWSLLM: sts_client = boto3.client("sts", **irsa_sts_kwargs) # Manually assume the IRSA role with the session name - verbose_logger.debug( - f"Manually assuming IRSA role {irsa_role_arn} with session {aws_session_name}" - ) + verbose_logger.debug(f"Manually assuming IRSA role {irsa_role_arn} with session {aws_session_name}") irsa_response = sts_client.assume_role_with_web_identity( RoleArn=irsa_role_arn, RoleSessionName=aws_session_name, @@ -1045,9 +976,7 @@ class BaseAWSLLM: verbose_logger.debug(f"Failed to get caller identity: {e}") # Now assume the target role - verbose_logger.debug( - f"Attempting to assume target role: {aws_role_name} with session: {aws_session_name}" - ) + verbose_logger.debug(f"Attempting to assume target role: {aws_role_name} with session: {aws_session_name}") assume_role_params = { "RoleArn": aws_role_name, "RoleSessionName": aws_session_name, @@ -1082,16 +1011,12 @@ class BaseAWSLLM: # Get current caller identity for debugging try: caller_identity = sts_client.get_caller_identity() - verbose_logger.debug( - f"Current IRSA identity: {caller_identity.get('Arn', 'unknown')}" - ) + verbose_logger.debug(f"Current IRSA identity: {caller_identity.get('Arn', 'unknown')}") except Exception as e: verbose_logger.debug(f"Failed to get caller identity: {e}") # Assume the role - verbose_logger.debug( - f"Attempting to assume role: {aws_role_name} with session: {aws_session_name}" - ) + verbose_logger.debug(f"Attempting to assume role: {aws_role_name} with session: {aws_session_name}") assume_role_params = { "RoleArn": aws_role_name, "RoleSessionName": aws_session_name, @@ -1103,9 +1028,7 @@ class BaseAWSLLM: return sts_client.assume_role(**assume_role_params) - def _extract_credentials_and_ttl( - self, sts_response: dict - ) -> Tuple[Credentials, Optional[int]]: + def _extract_credentials_and_ttl(self, sts_response: dict) -> Tuple[Credentials, Optional[int]]: """Extract credentials and TTL from STS response.""" from botocore.credentials import Credentials @@ -1117,9 +1040,7 @@ class BaseAWSLLM: ) expiration_time = sts_credentials["Expiration"] - ttl = int( - (expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds() - ) + ttl = int((expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds()) return credentials, ttl @@ -1148,17 +1069,10 @@ class BaseAWSLLM: # If we have IRSA environment variables and no explicit credentials, # we need to use the web identity token flow - if ( - web_identity_token_file - and irsa_role_arn - and aws_access_key_id is None - and aws_secret_access_key is None - ): + if web_identity_token_file and irsa_role_arn and aws_access_key_id is None and aws_secret_access_key is None: # For cross-account role assumption with specific session names, # we need to manually assume the IRSA role first with the correct session name - verbose_logger.debug( - f"IRSA detected: using web identity token from {web_identity_token_file}" - ) + verbose_logger.debug(f"IRSA detected: using web identity token from {web_identity_token_file}") try: # Check if we need to do cross-account role assumption @@ -1185,9 +1099,7 @@ class BaseAWSLLM: except Exception as e: verbose_logger.debug(f"Failed to assume role via IRSA: {e}") - if "AccessDenied" in str( - e - ) and "is not authorized to perform: sts:AssumeRole" in str(e): + if "AccessDenied" in str(e) and "is not authorized to perform: sts:AssumeRole" in str(e): # Provide a more helpful error message for trust policy issues verbose_logger.error( f"Access denied when trying to assume role {aws_role_name}. " @@ -1235,9 +1147,7 @@ class BaseAWSLLM: # partition, and role name). This avoids silently using the # wrong identity when there is a genuine trust-policy or # permission misconfiguration. - if self._is_already_running_as_role( - aws_role_name, ssl_verify=ssl_verify - ): + if self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify): verbose_logger.warning( "AssumeRole failed for %s (%s). " "Caller is already running as this role; " @@ -1248,8 +1158,7 @@ class BaseAWSLLM: return self._auth_with_env_vars() # Genuine permission error — re-raise verbose_logger.error( - "AssumeRole AccessDenied for %s and caller is NOT " - "the same role. Re-raising. Error: %s", + "AssumeRole AccessDenied for %s and caller is NOT the same role. Re-raising. Error: %s", aws_role_name, error_str, ) @@ -1270,9 +1179,7 @@ class BaseAWSLLM: return credentials, sts_ttl @tracer.wrap() - def _auth_with_aws_profile( - self, aws_profile_name: str - ) -> Tuple[Credentials, Optional[int]]: + def _auth_with_aws_profile(self, aws_profile_name: str) -> Tuple[Credentials, Optional[int]]: """ Authenticate with AWS profile """ @@ -1360,13 +1267,9 @@ class BaseAWSLLM: env_aws_bedrock_runtime_endpoint = get_secret("AWS_BEDROCK_RUNTIME_ENDPOINT") if api_base is not None: endpoint_url = api_base - elif aws_bedrock_runtime_endpoint is not None and isinstance( - aws_bedrock_runtime_endpoint, str - ): + elif aws_bedrock_runtime_endpoint is not None and isinstance(aws_bedrock_runtime_endpoint, str): endpoint_url = aws_bedrock_runtime_endpoint - elif env_aws_bedrock_runtime_endpoint and isinstance( - env_aws_bedrock_runtime_endpoint, str - ): + elif env_aws_bedrock_runtime_endpoint and isinstance(env_aws_bedrock_runtime_endpoint, str): endpoint_url = env_aws_bedrock_runtime_endpoint else: endpoint_url = self._select_default_endpoint_url( @@ -1375,13 +1278,9 @@ class BaseAWSLLM: ) # Determine proxy_endpoint_url - if aws_bedrock_runtime_endpoint is not None and isinstance( - aws_bedrock_runtime_endpoint, str - ): + if aws_bedrock_runtime_endpoint is not None and isinstance(aws_bedrock_runtime_endpoint, str): proxy_endpoint_url = aws_bedrock_runtime_endpoint - elif env_aws_bedrock_runtime_endpoint and isinstance( - env_aws_bedrock_runtime_endpoint, str - ): + elif env_aws_bedrock_runtime_endpoint and isinstance(env_aws_bedrock_runtime_endpoint, str): proxy_endpoint_url = env_aws_bedrock_runtime_endpoint else: proxy_endpoint_url = endpoint_url @@ -1477,21 +1376,15 @@ class BaseAWSLLM: try: from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError( - "Missing boto3 to call bedrock. Run 'pip install boto3'." - ) + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") headers["Authorization"] = f"Bearer {aws_bearer_token}" - request = AWSRequest( - method="POST", url=endpoint_url, data=data, headers=headers - ) + request = AWSRequest(method="POST", url=endpoint_url, data=data, headers=headers) else: try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError( - "Missing boto3 to call bedrock. Run 'pip install boto3'." - ) + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") # Filter headers for AWS signature calculation # AWS SigV4 only includes specific headers in signature calculation @@ -1541,11 +1434,7 @@ class BaseAWSLLM: if header_value is None: continue header_lower = header_name.lower() - if ( - header_lower in aws_headers - or header_lower.startswith("x-amz-") - or header_lower.startswith("x-amzn-") - ): + if header_lower in aws_headers or header_lower.startswith("x-amz-") or header_lower.startswith("x-amzn-"): aws_signature_headers[header_name] = header_value return aws_signature_headers @@ -1605,9 +1494,7 @@ class BaseAWSLLM: aws_web_identity_token = optional_params.get("aws_web_identity_token", None) aws_sts_endpoint = optional_params.get("aws_sts_endpoint", None) aws_external_id = optional_params.get("aws_external_id", None) - aws_region_name = self._get_aws_region_name( - optional_params=optional_params, model=model - ) + aws_region_name = self._get_aws_region_name(optional_params=optional_params, model=model) credentials: Credentials = self.get_credentials( aws_access_key_id=aws_access_key_id, @@ -1642,9 +1529,7 @@ class BaseAWSLLM: for header_name, header_value in headers.items(): if header_value is not None: request_headers_dict[header_name] = header_value - if ( - headers is not None and "Authorization" in headers - ): # prevent sigv4 from overwriting the auth header + if headers is not None and "Authorization" in headers: # prevent sigv4 from overwriting the auth header request_headers_dict["Authorization"] = headers["Authorization"] return request_headers_dict, request.body diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c071f331337..b0c7f1a3695 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -41,9 +41,7 @@ def _extract_job_id_from_arn(arn: str) -> Optional[str]: return arn.rsplit("/", 1)[-1] or None -def _predict_output_file_uri( - output_prefix: str, input_uri: str, job_id: Optional[str] -) -> Optional[str]: +def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: Optional[str]) -> Optional[str]: """ Compute the deterministic per-job result file URI Bedrock writes to. @@ -85,9 +83,7 @@ class BedrockBatchesHandler: """ @staticmethod - def _handle_async_invoke_status( - batch_id: str, aws_region_name: str, logging_obj=None, **kwargs - ) -> "LiteLLMBatch": + def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj=None, **kwargs) -> "LiteLLMBatch": """ Handle async invoke status check for AWS Bedrock. @@ -121,9 +117,7 @@ class BedrockBatchesHandler: from litellm.types.utils import LiteLLMBatch openai_batch_metadata: OpenAIBatchMetadata = { - "output_file_id": status_response["outputDataConfig"][ - "s3OutputDataConfig" - ]["s3Uri"], + "output_file_id": status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"], "failure_message": status_response.get("failureMessage") or "", "model_arn": status_response["modelArn"], } @@ -135,11 +129,7 @@ class BedrockBatchesHandler: created_at=status_response["submitTime"], in_progress_at=status_response["lastModifiedTime"], completed_at=status_response.get("endTime"), - failed_at=( - status_response.get("endTime") - if status_response["status"] == "failed" - else None - ), + failed_at=(status_response.get("endTime") if status_response["status"] == "failed" else None), request_counts=BatchRequestCounts( total=1, completed=1 if status_response["status"] == "completed" else 0, @@ -210,14 +200,10 @@ class BedrockBatchesHandler: try: import boto3 except ImportError as exc: - raise ImportError( - "Missing boto3 to call bedrock. Run 'pip install boto3'." - ) from exc + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") from exc # Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default). - region = ( - aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" - ) + region = aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" # Resolve credentials through the same path the rest of the bedrock # provider uses, so model_list / env / role-assumption configs are @@ -257,10 +243,7 @@ class BedrockBatchesHandler: api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": ( - f"https://bedrock.{region}.amazonaws.com/" - f"model-invocation-job/{url_path_id}" - ), + "api_base": (f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{url_path_id}"), }, ) @@ -280,16 +263,8 @@ class BedrockBatchesHandler: _BEDROCK_MIJ_STATUS_TO_OPENAI.get(bedrock_status, "in_progress"), ) - input_uri = ( - response.get("inputDataConfig", {}) - .get("s3InputDataConfig", {}) - .get("s3Uri", "") - ) - output_prefix = ( - response.get("outputDataConfig", {}) - .get("s3OutputDataConfig", {}) - .get("s3Uri", "") - ) + input_uri = response.get("inputDataConfig", {}).get("s3InputDataConfig", {}).get("s3Uri", "") + output_prefix = response.get("outputDataConfig", {}).get("s3OutputDataConfig", {}).get("s3Uri", "") # Bedrock returns the output *prefix* the user supplied at job creation. # Actual results land at //.out — we diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 620bc91732d..b0e28b6ba90 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -74,9 +74,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # Bedrock model invocation job endpoint # Format: https://bedrock.{region}.amazonaws.com/model-invocation-job - bedrock_endpoint = ( - f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" - ) + bedrock_endpoint = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" return bedrock_endpoint @@ -106,9 +104,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): input_bucket, input_key = self.common_utils.parse_s3_uri(input_file_id) # Get output S3 configuration - output_bucket = litellm_params.get("s3_output_bucket_name") or os.getenv( - "AWS_S3_OUTPUT_BUCKET_NAME" - ) + output_bucket = litellm_params.get("s3_output_bucket_name") or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") if not output_bucket: # Use same bucket as input if no output bucket specified output_bucket = input_bucket @@ -126,9 +122,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) if not model: - raise ValueError( - "Could not determine Bedrock model ID. Please pass `model` in your request body." - ) + raise ValueError("Could not determine Bedrock model ID. Please pass `model` in your request body.") # Generate job name with the correct model ID using common utility job_name = self.common_utils.generate_unique_job_name(model, prefix="litellm") @@ -136,9 +130,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # Build input data config input_data_config: BedrockInputDataConfig = { - "s3InputDataConfig": BedrockS3InputDataConfig( - s3Uri=f"s3://{input_bucket}/{input_key}" - ) + "s3InputDataConfig": BedrockS3InputDataConfig(s3Uri=f"s3://{input_bucket}/{input_key}") } # Build output data config @@ -147,15 +139,11 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Add optional KMS encryption key ID if provided - s3_encryption_key_id = litellm_params.get( - "s3_encryption_key_id" - ) or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + s3_encryption_key_id = litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id - output_data_config: BedrockOutputDataConfig = { - "s3OutputDataConfig": s3_output_config - } + output_data_config: BedrockOutputDataConfig = {"s3OutputDataConfig": s3_output_config} # Create Bedrock batch request with proper typing bedrock_request: BedrockCreateBatchRequest = { @@ -176,7 +164,9 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing - endpoint_url = f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job" + endpoint_url = ( + f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job" + ) signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, @@ -264,9 +254,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): cancelling_at=None, cancelled_at=None, request_counts=None, - metadata=self._get_openai_compatible_batch_metadata( - original_request.get("metadata", {}) - ), + metadata=self._get_openai_compatible_batch_metadata(original_request.get("metadata", {})), ) @staticmethod @@ -328,9 +316,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): import urllib.parse as _ul encoded_arn = _ul.quote(batch_id, safe="") - endpoint_url = ( - f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" - ) + endpoint_url = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing signed_headers, _ = self.common_utils.sign_aws_request( @@ -363,9 +349,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): return None created_at = parse_timestamp( - str(response_data.get("submitTime")) - if response_data.get("submitTime") is not None - else None + str(response_data.get("submitTime")) if response_data.get("submitTime") is not None else None ) in_progress_states = {"InProgress", "Validating", "Scheduled"} in_progress_at = ( @@ -378,36 +362,22 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): else None ) completed_at = ( - parse_timestamp( - str(response_data.get("endTime")) - if response_data.get("endTime") is not None - else None - ) + parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str in {"Completed", "PartiallyCompleted"} else None ) failed_at = ( - parse_timestamp( - str(response_data.get("endTime")) - if response_data.get("endTime") is not None - else None - ) + parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Failed" else None ) cancelled_at = ( - parse_timestamp( - str(response_data.get("endTime")) - if response_data.get("endTime") is not None - else None - ) + parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Stopped" else None ) expires_at = parse_timestamp( - str(response_data.get("jobExpirationTime")) - if response_data.get("jobExpirationTime") is not None - else None + str(response_data.get("jobExpirationTime")) if response_data.get("jobExpirationTime") is not None else None ) return ( @@ -539,9 +509,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): input_file_id, output_file_id = self._extract_file_configs(response_data) # Extract errors and metadata - errors, enriched_metadata = self._extract_errors_and_metadata( - response_data, raw_response - ) + errors, enriched_metadata = self._extract_errors_and_metadata(response_data, raw_response) return LiteLLMBatch( id=job_arn, @@ -566,9 +534,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): metadata=enriched_metadata, ) - def get_error_class( - self, error_message: str, status_code: int, headers: Union[Dict, Headers] - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: Union[Dict, Headers]) -> BaseLLMException: """ Get Bedrock-specific error class using common utility. """ diff --git a/litellm/llms/bedrock/chat/__init__.py b/litellm/llms/bedrock/chat/__init__.py index 8cd0e94e68e..c1323b9192a 100644 --- a/litellm/llms/bedrock/chat/__init__.py +++ b/litellm/llms/bedrock/chat/__init__.py @@ -9,9 +9,7 @@ from .invoke_handler import ( ) -def get_bedrock_event_stream_decoder( - invoke_provider: Optional[str], model: str, sync_stream: bool, json_mode: bool -): +def get_bedrock_event_stream_decoder(invoke_provider: Optional[str], model: str, sync_stream: bool, json_mode: bool): if invoke_provider and invoke_provider == "anthropic": decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( model=model, diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 542af61dd3f..356bc829677 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -84,9 +84,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): Get the complete url for the request """ ### SET RUNTIME ENDPOINT ### - aws_bedrock_runtime_endpoint = optional_params.get( - "aws_bedrock_runtime_endpoint", None - ) + aws_bedrock_runtime_endpoint = optional_params.get("aws_bedrock_runtime_endpoint", None) # Extract ARN from model string agent_runtime_arn = self._get_agent_runtime_arn(model) @@ -233,9 +231,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): dict: Payload dict containing the prompt and (optionally) the OpenAI content list. """ - verbose_logger.debug( - f"AgentCore transform_request - optional_params keys: {list(optional_params.keys())}" - ) + verbose_logger.debug(f"AgentCore transform_request - optional_params keys: {list(optional_params.keys())}") # Use the last message content as the prompt prompt = convert_content_list_to_str(messages[-1]) @@ -250,8 +246,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): if self._should_forward_multimodal_content(optional_params, litellm_params): last_content = messages[-1].get("content") if isinstance(last_content, list) and any( - isinstance(block, dict) and block.get("type") not in (None, "text") - for block in last_content + isinstance(block, dict) and block.get("type") not in (None, "text") for block in last_content ): # Copy so the payload never aliases messages[-1]["content"]; shallow, # not deep, to avoid cloning large base64 media on the request path. @@ -273,9 +268,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return payload @staticmethod - def _should_forward_multimodal_content( - optional_params: dict, litellm_params: dict - ) -> bool: + def _should_forward_multimodal_content(optional_params: dict, litellm_params: dict) -> bool: """Whether to forward raw OpenAI content blocks under ``payload["content"]``. Opt-in via the ``forward_multimodal_content`` litellm param (default ``False``) @@ -346,15 +339,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): if not isinstance(content_list, list): return "" - return "".join( - block["text"] - for block in content_list - if isinstance(block, dict) and "text" in block - ) + return "".join(block["text"] for block in content_list if isinstance(block, dict) and "text" in block) - def _calculate_usage( - self, model: str, messages: List[AllMessageValues], content: str - ) -> Optional[Usage]: + def _calculate_usage(self, model: str, messages: List[AllMessageValues], content: str) -> Optional[Usage]: """ Calculate token usage using LiteLLM's token counter. @@ -370,9 +357,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): from litellm.utils import token_counter prompt_tokens = token_counter(model=model, messages=messages) - completion_tokens = token_counter( - model=model, text=content, count_response_tokens=True - ) + completion_tokens = token_counter(model=model, text=content, count_response_tokens=True) total_tokens = prompt_tokens + completion_tokens verbose_logger.debug( @@ -402,9 +387,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Guard: if json.loads() returned a non-dict (e.g. array or primitive), # skip strategy matching and fall back to raw JSON string if not isinstance(response_json, dict): - verbose_logger.warning( - "AgentCore: JSON response is not a dict. Returning raw JSON as content." - ) + verbose_logger.warning("AgentCore: JSON response is not a dict. Returning raw JSON as content.") return AgentCoreParsedResponse( content=json.dumps(response_json), usage=None, @@ -465,9 +448,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): final_message=None, ) - def _get_parsed_response( - self, raw_response: httpx.Response - ) -> AgentCoreParsedResponse: + def _get_parsed_response(self, raw_response: httpx.Response) -> AgentCoreParsedResponse: """ Parse AgentCore response based on content type. @@ -491,9 +472,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # SSE stream response (text/event-stream or default) verbose_logger.debug("Parsing SSE stream response") response_text = raw_response.text - verbose_logger.debug( - f"AgentCore response (first 500 chars): {response_text[:500]}" - ) + verbose_logger.debug(f"AgentCore response (first 500 chars): {response_text[:500]}") return self._parse_sse_stream(response_text) def _parse_sse_stream(self, response_text: str) -> AgentCoreParsedResponse: @@ -527,9 +506,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Process event data if "event" in data and isinstance(data["event"], dict): event_payload = data["event"] - verbose_logger.debug( - f"Event payload keys: {list(event_payload.keys())}" - ) + verbose_logger.debug(f"Event payload keys: {list(event_payload.keys())}") # Extract usage metadata if usage := self._extract_usage_from_event(data): @@ -541,17 +518,11 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): content_blocks.append(text) # Build final content - content = ( - self._extract_content_from_message(final_message) - if final_message - else "".join(content_blocks) - ) + content = self._extract_content_from_message(final_message) if final_message else "".join(content_blocks) verbose_logger.debug(f"Final usage_data: {usage_data}") - return AgentCoreParsedResponse( - content=content, usage=usage_data, final_message=final_message - ) + return AgentCoreParsedResponse(content=content, usage=usage_data, final_message=final_message) def _stream_agentcore_response_sync( self, @@ -692,9 +663,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError( - status_code=response.status_code, message=str(response.read()) - ) + raise BedrockError(status_code=response.status_code, message=str(response.read())) # LOGGING logging_obj.post_call( @@ -708,8 +677,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): content_type = response.headers.get("content-type", "").lower() if "application/json" in content_type: verbose_logger.debug( - "AgentCore streaming: received JSON response instead of SSE, " - "converting to single-chunk stream" + "AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream" ) try: body = response.read() @@ -894,9 +862,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client( - llm_provider=cast(Any, "bedrock"), params={} - ) + client = get_async_httpx_client(llm_provider=cast(Any, "bedrock"), params={}) verbose_logger.debug(f"Making async streaming request to: {api_base}") @@ -910,9 +876,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError( - status_code=response.status_code, message=str(await response.aread()) - ) + raise BedrockError(status_code=response.status_code, message=str(await response.aread())) # LOGGING logging_obj.post_call( @@ -926,8 +890,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): content_type = response.headers.get("content-type", "").lower() if "application/json" in content_type: verbose_logger.debug( - "AgentCore streaming: received JSON response instead of SSE, " - "converting to single-chunk stream" + "AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream" ) try: body = await response.aread() @@ -939,9 +902,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) parsed = self._parse_json_response(response_json) - async def _json_as_async_stream() -> AsyncGenerator[ - ModelResponseStream, None - ]: + async def _json_as_async_stream() -> AsyncGenerator[ModelResponseStream, None]: # Content chunk content_chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", @@ -1054,9 +1015,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): setattr(model_response, "usage", usage) else: # Calculate token usage using LiteLLM's token counter - verbose_logger.debug( - "No usage data from AgentCore - calculating tokens" - ) + verbose_logger.debug("No usage data from AgentCore - calculating tokens") calculated_usage = self._calculate_usage(model, messages, content) if calculated_usage: setattr(model_response, "usage", calculated_usage) @@ -1064,9 +1023,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return model_response except Exception as e: - verbose_logger.error( - f"Error processing Bedrock AgentCore response: {str(e)}" - ) + verbose_logger.error(f"Error processing Bedrock AgentCore response: {str(e)}") raise BedrockError( message=f"Error processing response: {str(e)}", status_code=raw_response.status_code, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 040193f6bcc..292f570cc4e 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -46,33 +46,25 @@ def make_sync_call( ) if response.status_code != 200: - raise BedrockError( - status_code=response.status_code, message=str(response.read()) - ) + raise BedrockError(status_code=response.status_code, message=str(response.read())) if fake_stream: - model_response: ModelResponse = ( - litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params={}, - api_key="", - data=data, - messages=messages, - encoding=litellm.encoding, - ) + model_response: ModelResponse = litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=litellm.ModelResponse(), + stream=True, + logging_obj=logging_obj, + optional_params={}, + api_key="", + data=data, + messages=messages, + encoding=litellm.encoding, ) # type: ignore - completion_stream: Any = MockResponseIterator( - model_response=model_response, json_mode=json_mode - ) + completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes( - response.iter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) # LOGGING logging_obj.post_call( @@ -215,9 +207,7 @@ class BedrockConverseLLM(BaseAWSLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = get_async_httpx_client( - params=_params, llm_provider=litellm.LlmProviders.BEDROCK - ) + client = get_async_httpx_client(params=_params, llm_provider=litellm.LlmProviders.BEDROCK) else: client = client # type: ignore @@ -296,10 +286,7 @@ class BedrockConverseLLM(BaseAWSLLM): break modelId = self.encode_model_id(model_id=_model_for_id) # Inject region extracted from model path so _get_aws_region_name picks it up - if ( - _region_from_model is not None - and "aws_region_name" not in optional_params - ): + if _region_from_model is not None and "aws_region_name" not in optional_params: optional_params["aws_region_name"] = _region_from_model fake_stream = litellm.AmazonConverseConfig().should_fake_stream( @@ -332,9 +319,7 @@ class BedrockConverseLLM(BaseAWSLLM): aws_external_id = optional_params.pop("aws_external_id", None) optional_params.pop("aws_region_name", None) - litellm_params["aws_region_name"] = ( - aws_region_name # [DO NOT DELETE] important for async calls - ) + litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls credentials: Credentials = self.get_credentials( aws_access_key_id=aws_access_key_id, @@ -368,9 +353,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers = {"Content-Type": "application/json", **extra_headers} # Filter beta headers in HTTP headers before making the request - headers = update_headers_with_filtered_beta( - headers=headers, provider="bedrock_converse" - ) + headers = update_headers_with_filtered_beta(headers=headers, provider="bedrock_converse") ### ROUTING (ASYNC, STREAMING, SYNC) if acompletion: if isinstance(client, HTTPHandler): @@ -458,11 +441,7 @@ class BedrockConverseLLM(BaseAWSLLM): if stream is not None and stream is True: completion_stream = make_sync_call( - client=( - client - if client is not None and isinstance(client, HTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, HTTPHandler) else None), api_base=proxy_endpoint_url, headers=prepped.headers, # type: ignore data=data, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index a700c07d87a..b135a116753 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -167,8 +167,7 @@ class AmazonConverseConfig(BaseConfig): if isinstance(content, list): has_guarded_text = any( - isinstance(item, dict) and item.get("type") == "guarded_text" - for item in content + isinstance(item, dict) and item.get("type") == "guarded_text" for item in content ) if has_guarded_text: continue # Skip this message if it already has guarded_text @@ -329,13 +328,9 @@ class AmazonConverseConfig(BaseConfig): # Check if the model is a Nova 2 model (matches nova-2-lite, nova-2-pro, etc.) # Also check for nova-2/ spec prefix for imported models - return model_without_region.startswith( - "amazon.nova-2-" - ) or model_without_region.startswith("nova-2/") + return model_without_region.startswith("amazon.nova-2-") or model_without_region.startswith("nova-2/") - def _map_web_search_options( - self, web_search_options: dict, model: str - ) -> Optional[BedrockToolBlock]: + def _map_web_search_options(self, web_search_options: dict, model: str) -> Optional[BedrockToolBlock]: """ Map web_search_options to Nova grounding systemTool. @@ -364,9 +359,7 @@ class AmazonConverseConfig(BaseConfig): # (unlike Anthropic), so we just enable grounding with no options return BedrockToolBlock(systemTool={"name": "nova_grounding"}) - def _transform_reasoning_effort_to_reasoning_config( - self, reasoning_effort: str - ) -> dict: + def _transform_reasoning_effort_to_reasoning_config(self, reasoning_effort: str) -> dict: """ Transform reasoning_effort parameter to Nova 2 reasoningConfig structure. @@ -411,9 +404,7 @@ class AmazonConverseConfig(BaseConfig): } } - def _handle_reasoning_effort_parameter( - self, model: str, reasoning_effort: str, optional_params: dict - ) -> None: + def _handle_reasoning_effort_parameter(self, model: str, reasoning_effort: str, optional_params: dict) -> None: """ Handle the reasoning_effort parameter based on the model type. @@ -425,9 +416,7 @@ class AmazonConverseConfig(BaseConfig): if "gpt-oss" in model: optional_params["reasoning_effort"] = reasoning_effort elif self._is_nova_2_model(model): - reasoning_config = self._transform_reasoning_effort_to_reasoning_config( - reasoning_effort - ) + reasoning_config = self._transform_reasoning_effort_to_reasoning_config(reasoning_effort) optional_params.update(reasoning_config) else: mapped_thinking = AnthropicConfig._map_reasoning_effort( @@ -441,9 +430,7 @@ class AmazonConverseConfig(BaseConfig): else: optional_params["thinking"] = mapped_thinking if AnthropicConfig._is_adaptive_thinking_model(model): - mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get( - reasoning_effort - ) + mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort) if mapped_effort is None: AnthropicConfig._raise_invalid_reasoning_effort( model=model, @@ -459,9 +446,7 @@ class AmazonConverseConfig(BaseConfig): output_config=existing_output_config, ) mapped_effort = existing_output_config["effort"] - self._validate_anthropic_adaptive_effort( - model=model, effort=mapped_effort - ) + self._validate_anthropic_adaptive_effort(model=model, effort=mapped_effort) optional_params["output_config"] = existing_output_config optional_params["_output_config_normalized"] = True @@ -523,9 +508,7 @@ class AmazonConverseConfig(BaseConfig): "parallel_tool_calls", ] - if ( - "arn" in model - ): # we can't infer the model from the arn, so just add all params + if "arn" in model: # we can't infer the model from the arn, so just add all params supported_params.append("tools") supported_params.append("tool_choice") supported_params.append("thinking") @@ -547,9 +530,7 @@ class AmazonConverseConfig(BaseConfig): or base_model.startswith("meta.llama3-3") or base_model.startswith("meta.llama4") or base_model.startswith("amazon.nova") - or supports_function_calling( - model=model, custom_llm_provider=self.custom_llm_provider - ) + or supports_function_calling(model=model, custom_llm_provider=self.custom_llm_provider) ): supported_params.append("tools") @@ -559,9 +540,7 @@ class AmazonConverseConfig(BaseConfig): if litellm.utils.supports_tool_choice( model=model, custom_llm_provider=self.custom_llm_provider - ) or litellm.utils.supports_tool_choice( - model=base_model, custom_llm_provider=self.custom_llm_provider - ): + ) or litellm.utils.supports_tool_choice(model=base_model, custom_llm_provider=self.custom_llm_provider): # only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html supported_params.append("tool_choice") @@ -580,9 +559,7 @@ class AmazonConverseConfig(BaseConfig): model=model, custom_llm_provider=self.custom_llm_provider, ) - or supports_reasoning( - model=base_model, custom_llm_provider=self.custom_llm_provider - ) + or supports_reasoning(model=base_model, custom_llm_provider=self.custom_llm_provider) ): supported_params.append("thinking") supported_params.append("reasoning_effort") @@ -611,9 +588,7 @@ class AmazonConverseConfig(BaseConfig): elif isinstance(tool_choice, dict): # only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html specific_tool = SpecificToolChoiceBlock( - name=make_valid_bedrock_tool_name( - tool_choice.get("function", {}).get("name", "") - ) + name=make_valid_bedrock_tool_name(tool_choice.get("function", {}).get("name", "")) ) return ToolChoiceValuesBlock(tool=specific_tool) else: @@ -634,15 +609,9 @@ class AmazonConverseConfig(BaseConfig): return ["mp4", "mov", "mkv", "webm", "flv", "mpeg", "mpg", "wmv", "3gp"] def get_all_supported_content_types(self) -> List[str]: - return ( - self.get_supported_image_types() - + self.get_supported_document_types() - + self.get_supported_video_types() - ) + return self.get_supported_image_types() + self.get_supported_document_types() + self.get_supported_video_types() - def is_computer_use_tool_used( - self, tools: Optional[List[OpenAIChatCompletionToolParam]], model: str - ) -> bool: + def is_computer_use_tool_used(self, tools: Optional[List[OpenAIChatCompletionToolParam]], model: str) -> bool: """Check if computer use tools are being used in the request.""" if tools is None: return False @@ -655,9 +624,7 @@ class AmazonConverseConfig(BaseConfig): return True return False - def _transform_computer_use_tools( - self, computer_use_tools: List[OpenAIChatCompletionToolParam] - ) -> List[dict]: + def _transform_computer_use_tools(self, computer_use_tools: List[OpenAIChatCompletionToolParam]) -> List[dict]: """Transform computer use tools to Bedrock format.""" transformed_tools: List[dict] = [] @@ -699,9 +666,7 @@ class AmazonConverseConfig(BaseConfig): def _separate_computer_use_tools( self, tools: List[OpenAIChatCompletionToolParam], model: str - ) -> Tuple[ - List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam] - ]: + ) -> Tuple[List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam]]: """ Separate computer use tools from regular function tools. @@ -773,9 +738,7 @@ class AmazonConverseConfig(BaseConfig): return _tool @staticmethod - def _supports_native_structured_outputs( - model: str, custom_llm_provider: Optional[str] = None - ) -> bool: + def _supports_native_structured_outputs(model: str, custom_llm_provider: Optional[str] = None) -> bool: """Check if the Bedrock model supports native structured outputs (outputConfig.textFormat). Delegates to the standard ``supports_native_structured_output`` utility @@ -785,9 +748,7 @@ class AmazonConverseConfig(BaseConfig): """ from litellm.utils import supports_native_structured_output - return supports_native_structured_output( - model=model, custom_llm_provider=custom_llm_provider - ) + return supports_native_structured_output(model=model, custom_llm_provider=custom_llm_provider) @staticmethod def _add_additional_properties_to_schema(schema: dict) -> dict: @@ -810,25 +771,18 @@ class AmazonConverseConfig(BaseConfig): # Recurse into nested schemas if "properties" in result and isinstance(result["properties"], dict): result["properties"] = { - k: AmazonConverseConfig._add_additional_properties_to_schema(v) - for k, v in result["properties"].items() + k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in result["properties"].items() } if "items" in result and isinstance(result["items"], dict): - result["items"] = AmazonConverseConfig._add_additional_properties_to_schema( - result["items"] - ) + result["items"] = AmazonConverseConfig._add_additional_properties_to_schema(result["items"]) for defs_key in ("$defs", "definitions"): if defs_key in result and isinstance(result[defs_key], dict): result[defs_key] = { - k: AmazonConverseConfig._add_additional_properties_to_schema(v) - for k, v in result[defs_key].items() + k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in result[defs_key].items() } for key in ("anyOf", "allOf", "oneOf"): if key in result and isinstance(result[key], list): - result[key] = [ - AmazonConverseConfig._add_additional_properties_to_schema(item) - for item in result[key] - ] + result[key] = [AmazonConverseConfig._add_additional_properties_to_schema(item) for item in result[key]] return result @@ -858,9 +812,7 @@ class AmazonConverseConfig(BaseConfig): } """ if json_schema is not None: - json_schema = AmazonConverseConfig._add_additional_properties_to_schema( - json_schema - ) + json_schema = AmazonConverseConfig._add_additional_properties_to_schema(json_schema) schema_str = json.dumps(json_schema) if json_schema is not None else "{}" json_schema_def: JsonSchemaDefinition = {"schema": schema_str} if name is not None: @@ -882,14 +834,9 @@ class AmazonConverseConfig(BaseConfig): non_default_params: dict, optional_params: dict, ): - optional_params = self._add_tools_to_optional_params( - optional_params=optional_params, tools=tools - ) + optional_params = self._add_tools_to_optional_params(optional_params=optional_params, tools=tools) - if ( - "meta.llama3-3-70b-instruct-v1:0" in model - and non_default_params.get("stream", False) is True - ): + if "meta.llama3-3-70b-instruct-v1:0" in model and non_default_params.get("stream", False) is True: optional_params["fake_stream"] = True def map_openai_params( @@ -996,18 +943,14 @@ class AmazonConverseConfig(BaseConfig): self._validate_request_metadata(value) # type: ignore optional_params["requestMetadata"] = value - def _map_context_management_param( - self, value: Union[dict, list], optional_params: dict - ) -> None: + def _map_context_management_param(self, value: Union[dict, list], optional_params: dict) -> None: # Match the dispatcher's ``_normalize_spec`` behavior: only run the # OpenAI→Anthropic mapper for list inputs. Dict inputs are already in # Anthropic-native shape (``{"edits": [...]}``) and should pass # through unchanged so an Anthropic-format ``context_management`` # value isn't silently dropped when the mapper can't classify it. if isinstance(value, list): - mapped = AnthropicConfig.map_openai_context_management_to_anthropic( - cast(Union[dict, list], value) - ) + mapped = AnthropicConfig.map_openai_context_management_to_anthropic(cast(Union[dict, list], value)) else: mapped = value # Skip when the mapper returned None for malformed input — leaving the @@ -1059,10 +1002,7 @@ class AmazonConverseConfig(BaseConfig): if "type" in value and value["type"] == "text": return optional_params - if ( - self._supports_native_structured_outputs(model, self.custom_llm_provider) - and json_schema is not None - ): + if self._supports_native_structured_outputs(model, self.custom_llm_provider) and json_schema is not None: # Use Bedrock's native structured outputs API (outputConfig.textFormat) # No synthetic tool injection, no fake_stream needed. # Requires an explicit schema — json_object with no schema falls through @@ -1080,14 +1020,10 @@ class AmazonConverseConfig(BaseConfig): json_schema=json_schema, description=description, ) - optional_params = self._add_tools_to_optional_params( - optional_params=optional_params, tools=[_tool] - ) + optional_params = self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool]) if ( - litellm.utils.supports_tool_choice( - model=model, custom_llm_provider=self.custom_llm_provider - ) + litellm.utils.supports_tool_choice(model=model, custom_llm_provider=self.custom_llm_provider) and not is_thinking_enabled ): optional_params["tool_choice"] = ToolChoiceValuesBlock( @@ -1105,9 +1041,7 @@ class AmazonConverseConfig(BaseConfig): optional_params["json_mode"] = True return optional_params - def update_optional_params_with_thinking_tokens( - self, non_default_params: dict, optional_params: dict - ): + def update_optional_params_with_thinking_tokens(self, non_default_params: dict, optional_params: dict): """ Handles scenario where max tokens is not specified. For anthropic models (anthropic api/bedrock/vertex ai), this requires having the max tokens being set and being greater than the thinking token budget. @@ -1125,13 +1059,9 @@ class AmazonConverseConfig(BaseConfig): is_thinking_enabled = self.is_thinking_enabled(optional_params) is_max_tokens_in_request = self.is_max_tokens_in_request(non_default_params) if is_thinking_enabled and not is_max_tokens_in_request: - thinking_token_budget = cast(dict, optional_params["thinking"]).get( - "budget_tokens", None - ) + thinking_token_budget = cast(dict, optional_params["thinking"]).get("budget_tokens", None) if thinking_token_budget is not None: - optional_params["maxTokens"] = ( - thinking_token_budget + DEFAULT_MAX_TOKENS - ) + optional_params["maxTokens"] = thinking_token_budget + DEFAULT_MAX_TOKENS @overload def _get_cache_point_block( @@ -1197,23 +1127,15 @@ class AmazonConverseConfig(BaseConfig): if message["role"] == "system": system_prompt_indices.append(idx) if isinstance(message["content"], str) and message["content"]: - system_content_blocks.append( - SystemContentBlock(text=message["content"]) - ) - cache_block = self._get_cache_point_block( - message, block_type="system", model=model - ) + system_content_blocks.append(SystemContentBlock(text=message["content"])) + cache_block = self._get_cache_point_block(message, block_type="system", model=model) if cache_block: system_content_blocks.append(cache_block) elif isinstance(message["content"], list): for m in message["content"]: if m.get("type") == "text" and m.get("text"): - system_content_blocks.append( - SystemContentBlock(text=m["text"]) - ) - cache_block = self._get_cache_point_block( - m, block_type="system", model=model - ) + system_content_blocks.append(SystemContentBlock(text=m["text"])) + cache_block = self._get_cache_point_block(m, block_type="system", model=model) if cache_block: system_content_blocks.append(cache_block) if len(system_prompt_indices) > 0: @@ -1226,9 +1148,7 @@ class AmazonConverseConfig(BaseConfig): inference_params["topK"] = inference_params.pop("top_k") return InferenceConfig(**inference_params) - def _handle_top_k_value( - self, model: str, inference_params: dict, drop_params: bool = False - ) -> dict: + def _handle_top_k_value(self, model: str, inference_params: dict, drop_params: bool = False) -> dict: base_model = BedrockModelInfo.get_base_model(model) val_top_k = None @@ -1261,23 +1181,15 @@ class AmazonConverseConfig(BaseConfig): # Consume the internal ``_output_config_normalized`` marker set by # ``_handle_reasoning_effort_parameter`` so it does not linger on the # caller's ``optional_params`` after the transformation returns. - anthropic_output_config_already_normalized = bool( - optional_params.pop("_output_config_normalized", False) - ) + anthropic_output_config_already_normalized = bool(optional_params.pop("_output_config_normalized", False)) # Filter out exception objects before deepcopy to prevent deepcopy failures # Exceptions should not be stored in optional_params (this is a defensive fix) cleaned_params = filter_exceptions_from_params(optional_params) inference_params = safe_deep_copy(cleaned_params) - supported_converse_params = list( - AmazonConverseConfig.__annotations__.keys() - ) + ["top_k"] + supported_converse_params = list(AmazonConverseConfig.__annotations__.keys()) + ["top_k"] supported_tool_call_params = ["tools", "tool_choice"] supported_config_params = list(self.get_config_blocks().keys()) - total_supported_params = ( - supported_converse_params - + supported_tool_call_params - + supported_config_params - ) + total_supported_params = supported_converse_params + supported_tool_call_params + supported_config_params inference_params.pop("json_mode", None) # used for handling json_schema # Anthropic-only ``output_config`` (snake_case) — re-attached to @@ -1299,18 +1211,14 @@ class AmazonConverseConfig(BaseConfig): if request_metadata is not None: self._validate_request_metadata(request_metadata) - output_config: Optional[OutputConfigBlock] = inference_params.pop( - "outputConfig", None - ) + output_config: Optional[OutputConfigBlock] = inference_params.pop("outputConfig", None) base_model = BedrockModelInfo.get_base_model(model) if ( output_config is None and output_config_format is not None and output_config_format.get("type") == "json_schema" and base_model.startswith("anthropic") - and self._supports_native_structured_outputs( - model, self.custom_llm_provider - ) + and self._supports_native_structured_outputs(model, self.custom_llm_provider) ): output_config = self._create_output_config_for_response_format( json_schema=output_config_format.get("schema"), @@ -1328,17 +1236,11 @@ class AmazonConverseConfig(BaseConfig): ) # keep supported params in 'inference_params', and set all model-specific params in 'additional_request_params' - additional_request_params = { - k: v for k, v in inference_params.items() if k not in total_supported_params - } - inference_params = { - k: v for k, v in inference_params.items() if k in total_supported_params - } + additional_request_params = {k: v for k, v in inference_params.items() if k not in total_supported_params} + inference_params = {k: v for k, v in inference_params.items() if k in total_supported_params} # Handle parallel_tool_calls configuration - parallel_tool_use_config = additional_request_params.pop( - "_parallel_tool_use_config", None - ) + parallel_tool_use_config = additional_request_params.pop("_parallel_tool_use_config", None) if parallel_tool_use_config is not None and is_claude_4_5_on_bedrock(model): for key, value in parallel_tool_use_config.items(): if ( @@ -1353,9 +1255,7 @@ class AmazonConverseConfig(BaseConfig): additional_request_params.pop("parallel_tool_calls", None) # Only set the topK value in for models that support it - additional_request_params.update( - self._handle_top_k_value(model, inference_params, drop_params) - ) + additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params)) # Filter out internal/MCP-related parameters that shouldn't be sent to the API # These are LiteLLM internal parameters, not API parameters @@ -1364,18 +1264,11 @@ class AmazonConverseConfig(BaseConfig): # Filter out non-serializable objects (exceptions, callables, logging objects, etc.) # from additional_request_params to prevent JSON serialization errors # This filters: Exception objects, callable objects (functions), Logging objects, etc. - additional_request_params = filter_exceptions_from_params( - additional_request_params - ) + additional_request_params = filter_exceptions_from_params(additional_request_params) - if anthropic_output_config is not None and isinstance( - anthropic_output_config, dict - ): + if anthropic_output_config is not None and isinstance(anthropic_output_config, dict): if base_model.startswith("anthropic"): - if ( - litellm.drop_params is True - and not AnthropicConfig._model_supports_effort_param(model) - ): + if litellm.drop_params is True and not AnthropicConfig._model_supports_effort_param(model): litellm.verbose_logger.warning( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, model, @@ -1388,9 +1281,7 @@ class AmazonConverseConfig(BaseConfig): ) effort = anthropic_output_config.get("effort") if effort is not None: - self._validate_anthropic_adaptive_effort( - model=model, effort=effort - ) + self._validate_anthropic_adaptive_effort(model=model, effort=effort) additional_request_params["output_config"] = anthropic_output_config return ( @@ -1438,9 +1329,7 @@ class AmazonConverseConfig(BaseConfig): # Only separate tools if computer use tools are actually present if filtered_tools and self.is_computer_use_tool_used(filtered_tools, model): # Separate computer use tools from regular function tools - computer_use_tools, regular_tools = self._separate_computer_use_tools( - filtered_tools, model - ) + computer_use_tools, regular_tools = self._separate_computer_use_tools(filtered_tools, model) # Process regular function tools using existing logic bedrock_tools = _bedrock_tools_pt(regular_tools, model=model) @@ -1505,9 +1394,7 @@ class AmazonConverseConfig(BaseConfig): anthropic_beta_list.append(computer_use_header) # Transform computer use tools to proper Bedrock format - transformed_computer_tools = self._transform_computer_use_tools( - computer_use_tools - ) + transformed_computer_tools = self._transform_computer_use_tools(computer_use_tools) additional_request_params["tools"] = transformed_computer_tools else: # No computer use tools, process all tools as regular tools @@ -1565,11 +1452,7 @@ class AmazonConverseConfig(BaseConfig): additional_request_params.pop("context_management", None) return - compact_edits = [ - e - for e in edits - if isinstance(e, dict) and e.get("type") == COMPACT_EDIT_TYPE - ] + compact_edits = [e for e in edits if isinstance(e, dict) and e.get("type") == COMPACT_EDIT_TYPE] if compact_edits: compact_beta = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value if compact_beta not in anthropic_beta_list: @@ -1594,15 +1477,9 @@ class AmazonConverseConfig(BaseConfig): """ Bedrock doesn't support tool calling without `tools=` param specified. """ - if ( - "tools" not in optional_params - and messages is not None - and has_tool_call_blocks(messages) - ): + if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages): if litellm.modify_params: - optional_params["tools"] = add_dummy_tool( - custom_llm_provider="bedrock_converse" - ) + optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse") else: raise litellm.UnsupportedParamsError( message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.", @@ -1645,9 +1522,7 @@ class AmazonConverseConfig(BaseConfig): ) # Append cachePoint to tools if cache_control_injection_points has tool_config - cache_injection_points = additional_request_params.pop( - "cache_control_injection_points", None - ) + cache_injection_points = additional_request_params.pop("cache_control_injection_points", None) if cache_injection_points and len(bedrock_tools) > 0: for point in cache_injection_points: if point.get("location") == "tool_config": @@ -1656,9 +1531,7 @@ class AmazonConverseConfig(BaseConfig): bedrock_tool_config: Optional[ToolConfigBlock] = None if len(bedrock_tools) > 0: - tool_choice_values: ToolChoiceValuesBlock = inference_params.pop( - "tool_choice", None - ) + tool_choice_values: ToolChoiceValuesBlock = inference_params.pop("tool_choice", None) bedrock_tool_config = ToolConfigBlock( tools=bedrock_tools, ) @@ -1666,9 +1539,7 @@ class AmazonConverseConfig(BaseConfig): bedrock_tool_config["toolChoice"] = tool_choice_values data: CommonRequestObject = { - "inferenceConfig": self._transform_inference_params( - inference_params=inference_params - ), + "inferenceConfig": self._transform_inference_params(inference_params=inference_params), } if additional_request_params: data["additionalModelRequestFields"] = additional_request_params @@ -1702,14 +1573,10 @@ class AmazonConverseConfig(BaseConfig): litellm_params: dict, headers: Optional[dict] = None, ) -> RequestObject: - messages, system_content_blocks = self._transform_system_message( - messages, model=model - ) + messages, system_content_blocks = self._transform_system_message(messages, model=model) # Convert last user message to guarded_text if guardrailConfig is present - messages = self._convert_consecutive_user_messages_to_guarded_text( - messages, optional_params - ) + messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) ## TRANSFORMATION ## _data: CommonRequestObject = self._transform_request_helper( @@ -1721,13 +1588,11 @@ class AmazonConverseConfig(BaseConfig): drop_params=litellm_params.get("drop_params") is True, ) - bedrock_messages = ( - await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( - messages=messages, - model=model, - llm_provider="bedrock_converse", - user_continue_message=litellm_params.pop("user_continue_message", None), - ) + bedrock_messages = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model=model, + llm_provider="bedrock_converse", + user_continue_message=litellm_params.pop("user_continue_message", None), ) data: RequestObject = {"messages": bedrock_messages, **_data} @@ -1761,14 +1626,10 @@ class AmazonConverseConfig(BaseConfig): litellm_params: dict, headers: Optional[dict] = None, ) -> RequestObject: - messages, system_content_blocks = self._transform_system_message( - messages, model=model - ) + messages, system_content_blocks = self._transform_system_message(messages, model=model) # Convert last user message to guarded_text if guardrailConfig is present - messages = self._convert_consecutive_user_messages_to_guarded_text( - messages, optional_params - ) + messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) _data: CommonRequestObject = self._transform_request_helper( model=model, @@ -1818,9 +1679,7 @@ class AmazonConverseConfig(BaseConfig): encoding=encoding, ) - def _transform_reasoning_content( - self, reasoning_content_blocks: List[BedrockConverseReasoningContentBlock] - ) -> str: + def _transform_reasoning_content(self, reasoning_content_blocks: List[BedrockConverseReasoningContentBlock]) -> str: """ Extract the reasoning text from the reasoning content blocks @@ -1836,9 +1695,7 @@ class AmazonConverseConfig(BaseConfig): self, thinking_blocks: List[BedrockConverseReasoningContentBlock] ) -> List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]: """Return a consistent format for thinking blocks between Anthropic and Bedrock.""" - thinking_blocks_list: List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] = [] + thinking_blocks_list: List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] = [] for block in thinking_blocks: if "reasoningText" in block: _thinking_block = ChatCompletionThinkingBlock(type="thinking") @@ -1880,18 +1737,10 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens = ( - token_counter(text=reasoning_content, count_response_tokens=True) - if reasoning_content - else 0 - ) + reasoning_tokens = token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 completion_tokens_details = CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens, - text_tokens=( - output_tokens - reasoning_tokens - if reasoning_tokens > 0 - else output_tokens - ), + text_tokens=(output_tokens - reasoning_tokens if reasoning_tokens > 0 else output_tokens), ) openai_usage = Usage( prompt_tokens=input_tokens, @@ -1906,9 +1755,7 @@ class AmazonConverseConfig(BaseConfig): def get_tool_call_names( self, - tools: Optional[ - Union[List[ToolBlock], List[OpenAIChatCompletionToolParam]] - ] = None, + tools: Optional[Union[List[ToolBlock], List[OpenAIChatCompletionToolParam]]] = None, ) -> List[str]: if tools is None: return [] @@ -1947,13 +1794,8 @@ class AmazonConverseConfig(BaseConfig): try: tool_call_names = self.get_tool_call_names(tools) json_content = json.loads(message.content) - if ( - json_content.get("type") == "function" - and json_content.get("name") in tool_call_names - ): - tool_calls = [ - ChatCompletionMessageToolCall(function=Function(**json_content)) - ] + if json_content.get("type") == "function" and json_content.get("name") in tool_call_names: + tool_calls = [ChatCompletionMessageToolCall(function=Function(**json_content))] message.tool_calls = tool_calls message.content = None @@ -1982,9 +1824,7 @@ class AmazonConverseConfig(BaseConfig): """ content_str = "" tools: List[ChatCompletionToolCallChunk] = [] - reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( - None - ) + reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = None citationsContentBlocks: Optional[List[CitationsContentBlock]] = None for idx, content in enumerate(content_blocks): """ @@ -2001,9 +1841,7 @@ class AmazonConverseConfig(BaseConfig): if "toolUse" in content: ## check tool name was formatted by litellm _response_tool_name = content["toolUse"]["name"] - response_tool_name = get_bedrock_tool_name( - response_tool_name=_response_tool_name - ) + response_tool_name = get_bedrock_tool_name(response_tool_name=_response_tool_name) _function_chunk = ChatCompletionToolCallFunctionChunk( name=response_tool_name, arguments=json.dumps(content["toolUse"]["input"]), @@ -2123,11 +1961,7 @@ class AmazonConverseConfig(BaseConfig): """ try: response_data = json.loads(json_str) - if ( - isinstance(response_data, dict) - and "properties" in response_data - and len(response_data) == 1 - ): + if isinstance(response_data, dict) and "properties" in response_data and len(response_data) == 1: response_data = response_data["properties"] return json.dumps(response_data) except json.JSONDecodeError: @@ -2151,11 +1985,7 @@ class AmazonConverseConfig(BaseConfig): if not json_mode or not tools: return tools if tools else None - json_tool_indices = [ - i - for i, t in enumerate(tools) - if t["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME - ] + json_tool_indices = [i for i, t in enumerate(tools) if t["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME] if not json_tool_indices: # No json_tool_call found, return tools unchanged @@ -2163,14 +1993,10 @@ class AmazonConverseConfig(BaseConfig): if len(json_tool_indices) == len(tools): # All tools are json_tool_call — convert first one to content - verbose_logger.debug( - "Processing JSON tool call response for response_format" - ) + verbose_logger.debug("Processing JSON tool call response for response_format") json_mode_content_str: Optional[str] = tools[0]["function"].get("arguments") if json_mode_content_str is not None: - json_mode_content_str = AmazonConverseConfig._unwrap_bedrock_properties( - json_mode_content_str - ) + json_mode_content_str = AmazonConverseConfig._unwrap_bedrock_properties(json_mode_content_str) chat_completion_message["content"] = json_mode_content_str return None @@ -2180,13 +2006,9 @@ class AmazonConverseConfig(BaseConfig): first_idx = json_tool_indices[0] json_mode_args = tools[first_idx]["function"].get("arguments") if json_mode_args is not None: - json_mode_args = AmazonConverseConfig._unwrap_bedrock_properties( - json_mode_args - ) + json_mode_args = AmazonConverseConfig._unwrap_bedrock_properties(json_mode_args) existing = chat_completion_message.get("content") or "" - chat_completion_message["content"] = ( - existing + json_mode_args if existing else json_mode_args - ) + chat_completion_message["content"] = existing + json_mode_args if existing else json_mode_args real_tools = [t for i, t in enumerate(tools) if i not in json_tool_indices] return real_tools if real_tools else None @@ -2264,9 +2086,7 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} content_str = "" tools: List[ChatCompletionToolCallChunk] = [] - reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( - None - ) + reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = None citationsContentBlocks: Optional[List[CitationsContentBlock]] = None if message is not None: @@ -2285,13 +2105,9 @@ class AmazonConverseConfig(BaseConfig): provider_specific_fields["citationsContent"] = citationsContentBlocks if provider_specific_fields: - chat_completion_message["provider_specific_fields"] = ( - provider_specific_fields - ) + chat_completion_message["provider_specific_fields"] = provider_specific_fields - citations_text, annotations = self._transform_citations_to_annotations( - citationsContentBlocks - ) + citations_text, annotations = self._transform_citations_to_annotations(citationsContentBlocks) citations_included_in_content = False if citations_text: stripped_content = content_str.strip() @@ -2308,12 +2124,8 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message["annotations"] = annotations if reasoningContentBlocks is not None: - chat_completion_message["reasoning_content"] = ( - self._transform_reasoning_content(reasoningContentBlocks) - ) - chat_completion_message["thinking_blocks"] = ( - self._transform_thinking_blocks(reasoningContentBlocks) - ) + chat_completion_message["reasoning_content"] = self._transform_reasoning_content(reasoningContentBlocks) + chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks) chat_completion_message["content"] = content_str filtered_tools = self._filter_json_mode_tools( json_mode=json_mode, diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index c88fa32b6a0..413cdad45e0 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -90,21 +90,15 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): endpoint_url, _ = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, - aws_region_name=self._get_aws_region_name( - optional_params=optional_params, model=model - ), + aws_region_name=self._get_aws_region_name(optional_params=optional_params, model=model), endpoint_type="agent", ) agent_id, agent_alias_id = self._get_agent_id_and_alias_id(model) session_id = self._get_session_id(optional_params) encoded_agent_id = encode_url_path_segment(agent_id, field_name="agent_id") - encoded_agent_alias_id = encode_url_path_segment( - agent_alias_id, field_name="agent_alias_id" - ) - encoded_session_id = encode_url_path_segment( - session_id, field_name="session_id" - ) + encoded_agent_alias_id = encode_url_path_segment(agent_alias_id, field_name="agent_alias_id") + encoded_session_id = encode_url_path_segment(session_id, field_name="session_id") endpoint_url = f"{endpoint_url}/agents/{encoded_agent_id}/agentAliases/{encoded_agent_alias_id}/sessions/{encoded_session_id}/text" @@ -142,9 +136,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): # Split the model string by '/' and extract components parts = model.split("/") if len(parts) != 3 or parts[0] != "agent": - raise ValueError( - "Invalid model format. Expected format: 'model=agent/AGENT_ID/ALIAS_ID'" - ) + raise ValueError("Invalid model format. Expected format: 'model=agent/AGENT_ID/ALIAS_ID'") return parts[1], parts[2] # Return (agent_id, agent_alias_id) @@ -202,9 +194,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): parsed_event = { "headers": headers, "payload": { - "bytes": base64.b64encode( - message.encode("utf-8") - ).decode("utf-8") + "bytes": base64.b64encode(message.encode("utf-8")).decode("utf-8") }, # Re-encode for consistency } events.append(parsed_event) @@ -222,9 +212,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): } events.append(parsed_event) except json.JSONDecodeError as e: - verbose_logger.warning( - f"Failed to parse trace event JSON: {e}" - ) + verbose_logger.warning(f"Failed to parse trace event JSON: {e}") else: verbose_logger.debug(f"Unknown event type: {event_type}") @@ -241,9 +229,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): verbose_logger.debug(f"Response dict: {response_dict}") # Use the same response shape parsing as the existing decoder - parsed_response = parser.parse( - response_dict, self._get_response_stream_shape() - ) + parsed_response = parser.parse(response_dict, self._get_response_stream_shape()) verbose_logger.debug(f"Parsed response: {parsed_response}") if response_dict["status_code"] != 200: @@ -258,11 +244,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): error_message = exception_status + " " + error_message raise BedrockError( status_code=response_dict["status_code"], - message=( - json.dumps(error_message) - if isinstance(error_message, dict) - else error_message - ), + message=(json.dumps(error_message) if isinstance(error_message, dict) else error_message), ) if "chunk" in parsed_response: @@ -294,9 +276,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): ) except Exception as e: verbose_logger.debug(f"Error extracting headers: {e}") - return InvokeAgentEventHeaders( - event_type="", content_type="", message_type="" - ) + return InvokeAgentEventHeaders(event_type="", content_type="", message_type="") def _get_response_stream_shape(self): from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape @@ -311,9 +291,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): headers = event.get("headers", {}) payload = event.get("payload") - event_type = headers.get( - "event_type" - ) # Note: using event_type not event-type + event_type = headers.get("event_type") # Note: using event_type not event-type if event_type == "chunk" and payload: # Extract base64 encoded content from chunk events @@ -321,9 +299,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): encoded_bytes = chunk_payload.get("bytes", "") if encoded_bytes: try: - decoded_content = base64.b64decode(encoded_bytes).decode( - "utf-8" - ) + decoded_content = base64.b64decode(encoded_bytes).decode("utf-8") response_parts.append(decoded_content) except Exception as e: verbose_logger.warning(f"Failed to decode chunk content: {e}") @@ -383,22 +359,17 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): self, trace_data: InvokeAgentTrace, usage_info: InvokeAgentUsage ) -> None: """Extract usage information from preprocessing trace.""" - pre_processing: Optional[InvokeAgentPreProcessingTrace] = trace_data.get( - "preProcessingTrace" - ) + pre_processing: Optional[InvokeAgentPreProcessingTrace] = trace_data.get("preProcessingTrace") if not pre_processing: return model_output: Optional[InvokeAgentModelInvocationOutput] = ( - pre_processing.get("modelInvocationOutput") - or InvokeAgentModelInvocationOutput() + pre_processing.get("modelInvocationOutput") or InvokeAgentModelInvocationOutput() ) if not model_output: return - metadata: Optional[InvokeAgentMetadata] = ( - model_output.get("metadata") or InvokeAgentMetadata() - ) + metadata: Optional[InvokeAgentMetadata] = model_output.get("metadata") or InvokeAgentMetadata() if not metadata: return @@ -409,19 +380,14 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): usage_info["inputTokens"] += usage.get("inputTokens", 0) usage_info["outputTokens"] += usage.get("outputTokens", 0) - def _extract_orchestration_model( - self, trace_data: InvokeAgentTrace - ) -> Optional[str]: + def _extract_orchestration_model(self, trace_data: InvokeAgentTrace) -> Optional[str]: """Extract model information from orchestration trace.""" - orchestration_trace: Optional[InvokeAgentOrchestrationTrace] = trace_data.get( - "orchestrationTrace" - ) + orchestration_trace: Optional[InvokeAgentOrchestrationTrace] = trace_data.get("orchestrationTrace") if not orchestration_trace: return None model_invocation: Optional[InvokeAgentModelInvocationInput] = ( - orchestration_trace.get("modelInvocationInput") - or InvokeAgentModelInvocationInput() + orchestration_trace.get("modelInvocationInput") or InvokeAgentModelInvocationInput() ) if not model_invocation: return None @@ -454,8 +420,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): usage = Usage( prompt_tokens=usage_info.get("inputTokens", 0), completion_tokens=usage_info.get("outputTokens", 0), - total_tokens=usage_info.get("inputTokens", 0) - + usage_info.get("outputTokens", 0), + total_tokens=usage_info.get("inputTokens", 0) + usage_info.get("outputTokens", 0), ) setattr(model_response, "usage", usage) @@ -478,9 +443,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): try: # Get the raw binary content raw_content = raw_response.content - verbose_logger.debug( - f"Processing {len(raw_content)} bytes of AWS event stream data" - ) + verbose_logger.debug(f"Processing {len(raw_content)} bytes of AWS event stream data") # Parse the AWS event stream format events = self._parse_aws_event_stream(raw_content) @@ -501,9 +464,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): ) except Exception as e: - verbose_logger.error( - f"Error processing Bedrock Invoke Agent response: {str(e)}" - ) + verbose_logger.error(f"Error processing Bedrock Invoke Agent response: {str(e)}") raise BedrockError( message=f"Error processing response: {str(e)}", status_code=raw_response.status_code, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 29e97068100..b381b5a85fe 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -75,9 +75,7 @@ from ..common_utils import ( get_bedrock_tool_name, ) -bedrock_tool_name_mappings: InMemoryCache = InMemoryCache( - max_size_in_memory=50, default_ttl=600 -) +bedrock_tool_name_mappings: InMemoryCache = InMemoryCache(max_size_in_memory=50, default_ttl=600) from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( AmazonBedrockOpenAIConfig, @@ -162,9 +160,7 @@ class AmazonCohereChatConfig: "tool_choice", ] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for param, value in non_default_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["max_tokens"] = value @@ -206,9 +202,7 @@ async def make_call( llm_provider=litellm.LlmProviders.BEDROCK, params=( {"ssl_verify": logging_obj.litellm_params.get("ssl_verify")} - if logging_obj - and logging_obj.litellm_params - and logging_obj.litellm_params.get("ssl_verify") + if logging_obj and logging_obj.litellm_params and logging_obj.litellm_params.get("ssl_verify") else None ), ) # Create a new client if none provided @@ -225,45 +219,35 @@ async def make_call( raise BedrockError(status_code=response.status_code, message=response.text) if fake_stream: - model_response: ModelResponse = ( - litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params={}, - api_key="", - data=data, - messages=messages, - encoding=litellm.encoding, - ) + model_response: ModelResponse = litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=litellm.ModelResponse(), + stream=True, + logging_obj=logging_obj, + optional_params={}, + api_key="", + data=data, + messages=messages, + encoding=litellm.encoding, ) # type: ignore - completion_stream: Any = MockResponseIterator( - model_response=model_response, json_mode=json_mode - ) + completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) elif bedrock_invoke_provider == "anthropic": decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( model=model, sync_stream=False, json_mode=json_mode, ) - completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) elif bedrock_invoke_provider == "deepseek_r1": decoder = AmazonDeepSeekR1StreamDecoder( model=model, sync_stream=False, ) - completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) else: decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=stream_chunk_size)) # LOGGING logging_obj.post_call( @@ -302,9 +286,7 @@ def make_sync_call( client = _get_httpx_client( params=( {"ssl_verify": logging_obj.litellm_params.get("ssl_verify")} - if logging_obj - and logging_obj.litellm_params - and logging_obj.litellm_params.get("ssl_verify") + if logging_obj and logging_obj.litellm_params and logging_obj.litellm_params.get("ssl_verify") else None ) ) @@ -321,45 +303,35 @@ def make_sync_call( raise BedrockError(status_code=response.status_code, message=response.text) if fake_stream: - model_response: ModelResponse = ( - litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params={}, - api_key="", - data=data, - messages=messages, - encoding=litellm.encoding, - ) + model_response: ModelResponse = litellm.AmazonConverseConfig()._transform_response( + model=model, + response=response, + model_response=litellm.ModelResponse(), + stream=True, + logging_obj=logging_obj, + optional_params={}, + api_key="", + data=data, + messages=messages, + encoding=litellm.encoding, ) # type: ignore - completion_stream: Any = MockResponseIterator( - model_response=model_response, json_mode=json_mode - ) + completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) elif bedrock_invoke_provider == "anthropic": decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( model=model, sync_stream=True, json_mode=json_mode, ) - completion_stream = decoder.iter_bytes( - response.iter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) elif bedrock_invoke_provider == "deepseek_r1": decoder = AmazonDeepSeekR1StreamDecoder( model=model, sync_stream=True, ) - completion_stream = decoder.iter_bytes( - response.iter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) else: decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes( - response.iter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) # LOGGING logging_obj.post_call( @@ -424,9 +396,7 @@ class BedrockLLM(BaseAWSLLM): return any(indicator in model_lower for indicator in messages_api_indicators) - def convert_messages_to_prompt( - self, model, messages, provider, custom_prompt_dict - ) -> Tuple[str, Optional[list]]: + def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> Tuple[str, Optional[list]]: # handle anthropic prompts and amazon titan prompts prompt = "" chat_history: Optional[list] = None @@ -436,26 +406,18 @@ class BedrockLLM(BaseAWSLLM): model_prompt_details = custom_prompt_dict[model] prompt = custom_prompt( role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details.get( - "initial_prompt_value", "" - ), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), final_prompt_value=model_prompt_details.get("final_prompt_value", ""), messages=messages, ) return prompt, None ## ELSE if provider == "anthropic" or provider == "amazon": - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="bedrock" - ) + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock") elif provider == "mistral": - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="bedrock" - ) + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock") elif provider == "meta" or provider == "llama": - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="bedrock" - ) + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock") elif provider == "openai": # OpenAI uses messages directly, no prompt conversion needed # Return empty prompt as it won't be used @@ -522,20 +484,12 @@ class BedrockLLM(BaseAWSLLM): if "tools" in optional_params: _is_function_call = True for tool in optional_params["tools"]: - json_schemas[tool["function"]["name"]] = tool[ - "function" - ].get("parameters", None) + json_schemas[tool["function"]["name"]] = tool["function"].get("parameters", None) outputText = completion_response.get("content")[0].get("text", None) - if outputText is not None and contains_tag( - "invoke", outputText - ): # OUTPUT PARSE FUNCTION CALL + if outputText is not None and contains_tag("invoke", outputText): # OUTPUT PARSE FUNCTION CALL function_name = extract_between_tags("tool_name", outputText)[0] - function_arguments_str = extract_between_tags( - "invoke", outputText - )[0].strip() - function_arguments_str = ( - f"{function_arguments_str}" - ) + function_arguments_str = extract_between_tags("invoke", outputText)[0].strip() + function_arguments_str = f"{function_arguments_str}" function_arguments = parse_xml_params( function_arguments_str, json_schema=json_schemas.get( @@ -559,14 +513,8 @@ class BedrockLLM(BaseAWSLLM): model_response._hidden_params["original_response"] = ( outputText # allow user to access raw anthropic tool calling response ) - if ( - _is_function_call is True - and stream is not None - and stream is True - ): - print_verbose( - "INSIDE BEDROCK STREAMING TOOL CALLING CONDITION BLOCK" - ) + if _is_function_call is True and stream is not None and stream is True: + print_verbose("INSIDE BEDROCK STREAMING TOOL CALLING CONDITION BLOCK") # return an iterator streaming_model_response = ModelResponseStream() streaming_model_response.choices[0].finish_reason = getattr( @@ -576,35 +524,23 @@ class BedrockLLM(BaseAWSLLM): streaming_choice = litellm.utils.StreamingChoices() streaming_choice.index = model_response.choices[0].index _tool_calls = [] - print_verbose( - f"type of model_response.choices[0]: {type(model_response.choices[0])}" - ) - print_verbose( - f"type of streaming_choice: {type(streaming_choice)}" - ) + print_verbose(f"type of model_response.choices[0]: {type(model_response.choices[0])}") + print_verbose(f"type of streaming_choice: {type(streaming_choice)}") if isinstance(model_response.choices[0], litellm.Choices): if getattr( model_response.choices[0].message, "tool_calls", None - ) is not None and isinstance( - model_response.choices[0].message.tool_calls, list - ): - for tool_call in model_response.choices[ - 0 - ].message.tool_calls: + ) is not None and isinstance(model_response.choices[0].message.tool_calls, list): + for tool_call in model_response.choices[0].message.tool_calls: _tool_call = {**tool_call.dict(), "index": 0} _tool_calls.append(_tool_call) delta_obj = Delta( - content=getattr( - model_response.choices[0].message, "content", None - ), + content=getattr(model_response.choices[0].message, "content", None), role=model_response.choices[0].message.role, tool_calls=_tool_calls, ) streaming_choice.delta = delta_obj streaming_model_response.choices = [streaming_choice] - completion_stream = ModelResponseIterator( - model_response=streaming_model_response - ) + completion_stream = ModelResponseIterator(model_response=streaming_model_response) print_verbose( "Returns anthropic CustomStreamWrapper with 'cached_response' streaming object" ) @@ -628,21 +564,14 @@ class BedrockLLM(BaseAWSLLM): else: outputText = completion_response["completion"] - model_response.choices[0].finish_reason = completion_response[ - "stop_reason" - ] + model_response.choices[0].finish_reason = completion_response["stop_reason"] elif provider == "ai21": - outputText = ( - completion_response.get("completions")[0].get("data").get("text") - ) + outputText = completion_response.get("completions")[0].get("data").get("text") elif provider == "meta" or provider == "llama": outputText = completion_response["generation"] elif provider == "openai": # OpenAI imported models use OpenAI Chat Completions format - if ( - "choices" in completion_response - and len(completion_response["choices"]) > 0 - ): + if "choices" in completion_response and len(completion_response["choices"]) > 0: choice = completion_response["choices"][0] if "message" in choice: outputText = choice["message"].get("content") @@ -651,9 +580,7 @@ class BedrockLLM(BaseAWSLLM): # Set finish reason if "finish_reason" in choice: - model_response.choices[0].finish_reason = map_finish_reason( - choice["finish_reason"] - ) + model_response.choices[0].finish_reason = map_finish_reason(choice["finish_reason"]) # Set usage if available if "usage" in completion_response: @@ -666,16 +593,12 @@ class BedrockLLM(BaseAWSLLM): setattr(model_response, "usage", _usage) elif provider == "mistral": outputText = completion_response["outputs"][0]["text"] - model_response.choices[0].finish_reason = completion_response[ - "outputs" - ][0]["stop_reason"] + model_response.choices[0].finish_reason = completion_response["outputs"][0]["stop_reason"] else: # amazon titan outputText = completion_response.get("results")[0].get("outputText") except Exception as e: raise BedrockError( - message="Error processing={}, Received error={}".format( - response.text, str(e) - ), + message="Error processing={}, Received error={}".format(response.text, str(e)), status_code=422, ) @@ -698,9 +621,7 @@ class BedrockLLM(BaseAWSLLM): raise Exception() except Exception as e: raise BedrockError( - message="Error parsing received text={}.\nError-{}".format( - outputText, str(e) - ), + message="Error parsing received text={}.\nError-{}".format(outputText, str(e)), status_code=response.status_code, ) @@ -728,20 +649,11 @@ class BedrockLLM(BaseAWSLLM): ## CALCULATING USAGE - bedrock returns usage in the headers # Skip if usage was already set (e.g., from JSON response for OpenAI provider) - if ( - not hasattr(model_response, "usage") - or getattr(model_response, "usage", None) is None - ): - bedrock_input_tokens = response.headers.get( - "x-amzn-bedrock-input-token-count", None - ) - bedrock_output_tokens = response.headers.get( - "x-amzn-bedrock-output-token-count", None - ) + if not hasattr(model_response, "usage") or getattr(model_response, "usage", None) is None: + bedrock_input_tokens = response.headers.get("x-amzn-bedrock-input-token-count", None) + bedrock_output_tokens = response.headers.get("x-amzn-bedrock-output-token-count", None) - prompt_tokens = int( - bedrock_input_tokens or litellm.token_counter(messages=messages) - ) + prompt_tokens = int(bedrock_input_tokens or litellm.token_counter(messages=messages)) completion_tokens = int( bedrock_output_tokens @@ -821,15 +733,11 @@ class BedrockLLM(BaseAWSLLM): # check env # litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) - if litellm_aws_region_name is not None and isinstance( - litellm_aws_region_name, str - ): + if litellm_aws_region_name is not None and isinstance(litellm_aws_region_name, str): aws_region_name = litellm_aws_region_name standard_aws_region_name = get_secret("AWS_REGION", None) - if standard_aws_region_name is not None and isinstance( - standard_aws_region_name, str - ): + if standard_aws_region_name is not None and isinstance(standard_aws_region_name, str): aws_region_name = standard_aws_region_name if aws_region_name is None: @@ -857,18 +765,12 @@ class BedrockLLM(BaseAWSLLM): if (stream is not None and stream is True) and provider != "ai21": endpoint_url = f"{endpoint_url}/model/{modelId}/invoke-with-response-stream" - proxy_endpoint_url = ( - f"{proxy_endpoint_url}/model/{modelId}/invoke-with-response-stream" - ) + proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke-with-response-stream" else: endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" - if ( - acompletion - and provider == "anthropic" - and self.is_claude_messages_api_model(model) - ): + if acompletion and provider == "anthropic" and self.is_claude_messages_api_model(model): if isinstance(client, HTTPHandler): client = None return self._async_anthropic_messages_completion( @@ -892,9 +794,7 @@ class BedrockLLM(BaseAWSLLM): stream_chunk_size=stream_chunk_size, ) # type: ignore[return-value] - prompt, chat_history = self.convert_messages_to_prompt( - model, messages, provider, custom_prompt_dict - ) + prompt, chat_history = self.convert_messages_to_prompt(model, messages, provider, custom_prompt_dict) inference_params = copy.deepcopy(optional_params) json_schemas: dict = {} if provider == "cohere": @@ -919,9 +819,7 @@ class BedrockLLM(BaseAWSLLM): ): # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in inference_params[k] = v if stream is True: - inference_params["stream"] = ( - True # cohere requires stream = True in inference params - ) + inference_params["stream"] = True # cohere requires stream = True in inference params data = json.dumps({"prompt": prompt, **inference_params}) elif provider == "anthropic": if self.is_claude_messages_api_model(model): @@ -934,13 +832,9 @@ class BedrockLLM(BaseAWSLLM): system_prompt_idx.append(idx) if len(system_prompt_idx) > 0: inference_params["system"] = "\n".join(system_messages) - messages = [ - i for j, i in enumerate(messages) if j not in system_prompt_idx - ] + messages = [i for j, i in enumerate(messages) if j not in system_prompt_idx] # Format rest of message according to anthropic guidelines - messages = prompt_factory( - model=model, messages=messages, custom_llm_provider="anthropic_xml" - ) # type: ignore + messages = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic_xml") # type: ignore ## LOAD CONFIG config = litellm.AmazonAnthropicClaudeConfig.get_config() for k, v in config.items(): @@ -952,15 +846,10 @@ class BedrockLLM(BaseAWSLLM): if "tools" in inference_params: _is_function_call = True for tool in inference_params["tools"]: - json_schemas[tool["function"]["name"]] = tool["function"].get( - "parameters", None - ) - tool_calling_system_prompt = construct_tool_use_system_prompt( - tools=inference_params["tools"] - ) + json_schemas[tool["function"]["name"]] = tool["function"].get("parameters", None) + tool_calling_system_prompt = construct_tool_use_system_prompt(tools=inference_params["tools"]) inference_params["system"] = ( - inference_params.get("system", "\n") - + tool_calling_system_prompt + inference_params.get("system", "\n") + tool_calling_system_prompt ) # add the anthropic tool calling prompt to the system prompt inference_params.pop("tools") data = json.dumps({"messages": messages, **inference_params}) @@ -1024,9 +913,7 @@ class BedrockLLM(BaseAWSLLM): supported_params = openai_config.get_supported_openai_params(model=model) # Filter to only supported OpenAI params - filtered_params = { - k: v for k, v in inference_params.items() if k in supported_params - } + filtered_params = {k: v for k, v in inference_params.items() if k in supported_params} # OpenAI uses messages format, not prompt data = json.dumps({"messages": messages, **filtered_params}) @@ -1132,15 +1019,11 @@ class BedrockLLM(BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError( - status_code=response.status_code, message=str(response.read()) - ) + raise BedrockError(status_code=response.status_code, message=str(response.read())) decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes( - response.iter_bytes(chunk_size=stream_chunk_size) - ) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, @@ -1206,14 +1089,12 @@ class BedrockLLM(BaseAWSLLM): client: Optional[AsyncHTTPHandler] = None, stream_chunk_size: Optional[int] = None, ) -> Union[ModelResponse, CustomStreamWrapper]: - transformed_request = ( - await litellm.AmazonAnthropicClaudeConfig().async_transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params or {}, - headers=extra_headers or {}, - ) + transformed_request = await litellm.AmazonAnthropicClaudeConfig().async_transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params or {}, + headers=extra_headers or {}, ) data = json.dumps(transformed_request) @@ -1300,9 +1181,7 @@ class BedrockLLM(BaseAWSLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = get_async_httpx_client( - params=_params, llm_provider=litellm.LlmProviders.BEDROCK - ) # type: ignore + client = get_async_httpx_client(params=_params, llm_provider=litellm.LlmProviders.BEDROCK) # type: ignore else: client = client # type: ignore @@ -1441,19 +1320,13 @@ class AWSEventStreamDecoder: def translate_thinking_blocks( self, thinking_block: BedrockConverseReasoningContentBlockDelta - ) -> Optional[ - List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] - ]: + ) -> Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]]: """ Translate the thinking blocks to a string """ - thinking_blocks_list: List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] = [] - _thinking_block: Optional[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] = None + thinking_blocks_list: List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] = [] + _thinking_block: Optional[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] = None if "text" in thinking_block: _thinking_block = ChatCompletionThinkingBlock(type="thinking") @@ -1487,42 +1360,27 @@ class AWSEventStreamDecoder: ) -> Tuple[ Optional[ChatCompletionToolCallChunk], dict, - Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ], + Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]], ]: """Handle 'start' event in converse chunk parsing.""" tool_use: Optional[ChatCompletionToolCallChunk] = None provider_specific_fields: dict = {} - thinking_blocks: Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ] = None + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None self.content_blocks = [] # reset if start_obj is not None: if "toolUse" in start_obj and start_obj["toolUse"] is not None: ## check tool name was formatted by litellm _response_tool_name = start_obj["toolUse"]["name"] - response_tool_name = get_bedrock_tool_name( - response_tool_name=_response_tool_name - ) + response_tool_name = get_bedrock_tool_name(response_tool_name=_response_tool_name) self._current_tool_name = response_tool_name # When json_mode is True, suppress the internal json_tool_call # and convert its content to text in delta events instead - if ( - self.json_mode is True - and response_tool_name == RESPONSE_FORMAT_TOOL_NAME - ): + if self.json_mode is True and response_tool_name == RESPONSE_FORMAT_TOOL_NAME: return tool_use, provider_specific_fields, thinking_blocks - self.tool_calls_index = ( - 0 if self.tool_calls_index is None else self.tool_calls_index + 1 - ) + self.tool_calls_index = 0 if self.tool_calls_index is None else self.tool_calls_index + 1 tool_use = { "id": start_obj["toolUse"]["toolUseId"], "type": "function", @@ -1533,12 +1391,9 @@ class AWSEventStreamDecoder: "index": self.tool_calls_index, } elif ( - "reasoningContent" in start_obj - and start_obj["reasoningContent"] is not None + "reasoningContent" in start_obj and start_obj["reasoningContent"] is not None ): # redacted thinking can be in start object - thinking_blocks = self.translate_thinking_blocks( - start_obj["reasoningContent"] - ) + thinking_blocks = self.translate_thinking_blocks(start_obj["reasoningContent"]) provider_specific_fields = { "reasoningContent": start_obj["reasoningContent"], } @@ -1553,22 +1408,14 @@ class AWSEventStreamDecoder: Optional[ChatCompletionToolCallChunk], dict, Optional[str], - Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ], + Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]], ]: """Handle 'delta' event in converse chunk parsing.""" text = "" tool_use: Optional[ChatCompletionToolCallChunk] = None provider_specific_fields: dict = {} reasoning_content: Optional[str] = None - thinking_blocks: Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ] = None + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None self.content_blocks.append(delta_obj) if "text" in delta_obj: @@ -1576,10 +1423,7 @@ class AWSEventStreamDecoder: elif "toolUse" in delta_obj: # When json_mode is True and this is the internal json_tool_call, # convert tool input to text content instead of tool call arguments - if ( - self.json_mode is True - and self._current_tool_name == RESPONSE_FORMAT_TOOL_NAME - ): + if self.json_mode is True and self._current_tool_name == RESPONSE_FORMAT_TOOL_NAME: text = delta_obj["toolUse"]["input"] else: tool_use = { @@ -1589,30 +1433,16 @@ class AWSEventStreamDecoder: "name": None, "arguments": delta_obj["toolUse"]["input"], }, - "index": ( - self.tool_calls_index - if self.tool_calls_index is not None - else index - ), + "index": (self.tool_calls_index if self.tool_calls_index is not None else index), } elif "reasoningContent" in delta_obj: provider_specific_fields = { "reasoningContent": delta_obj["reasoningContent"], } - reasoning_content = self.extract_reasoning_content_str( - delta_obj["reasoningContent"] - ) - thinking_blocks = self.translate_thinking_blocks( - delta_obj["reasoningContent"] - ) - if ( - thinking_blocks - and len(thinking_blocks) > 0 - and reasoning_content is None - ): - reasoning_content = ( - "" # set to non-empty string to ensure consistency with Anthropic - ) + reasoning_content = self.extract_reasoning_content_str(delta_obj["reasoningContent"]) + thinking_blocks = self.translate_thinking_blocks(delta_obj["reasoningContent"]) + if thinking_blocks and len(thinking_blocks) > 0 and reasoning_content is None: + reasoning_content = "" # set to non-empty string to ensure consistency with Anthropic elif "citationsContent" in delta_obj: # Handle Nova grounding citations in streaming responses provider_specific_fields = { @@ -1626,18 +1456,13 @@ class AWSEventStreamDecoder: thinking_blocks, ) - def _handle_converse_stop_event( - self, index: int - ) -> Optional[ChatCompletionToolCallChunk]: + def _handle_converse_stop_event(self, index: int) -> Optional[ChatCompletionToolCallChunk]: """Handle stop/contentBlockIndex event in converse chunk parsing.""" tool_use: Optional[ChatCompletionToolCallChunk] = None # If the ending block was the internal json_tool_call, skip emitting # the empty-args tool chunk and reset tracking state - if ( - self.json_mode is True - and self._current_tool_name == RESPONSE_FORMAT_TOOL_NAME - ): + if self.json_mode is True and self._current_tool_name == RESPONSE_FORMAT_TOOL_NAME: self._current_tool_name = None return tool_use @@ -1651,11 +1476,7 @@ class AWSEventStreamDecoder: "name": None, "arguments": "{}", }, - "index": ( - self.tool_calls_index - if self.tool_calls_index is not None - else index - ), + "index": (self.tool_calls_index if self.tool_calls_index is not None else index), } return tool_use @@ -1672,13 +1493,9 @@ class AWSEventStreamDecoder: usage: Optional[Usage] = None provider_specific_fields: dict = {} reasoning_content: Optional[str] = None - thinking_blocks: Optional[ - List[ - Union[ - ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock - ] - ] - ] = None + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = ( + None + ) content_block_index = int(chunk_data.get("contentBlockIndex", 0)) if "start" in chunk_data: @@ -1697,9 +1514,7 @@ class AWSEventStreamDecoder: reasoning_content, thinking_blocks, ) = self._handle_converse_delta_event(delta_obj, content_block_index) - elif ( - "contentBlockIndex" in chunk_data - ): # stop block, no 'start' or 'delta' object + elif "contentBlockIndex" in chunk_data: # stop block, no 'start' or 'delta' object tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) @@ -1719,11 +1534,7 @@ class AWSEventStreamDecoder: content=text, role="assistant", tool_calls=[tool_use] if tool_use else None, - provider_specific_fields=( - provider_specific_fields - if provider_specific_fields - else None - ), + provider_specific_fields=(provider_specific_fields if provider_specific_fields else None), thinking_blocks=thinking_blocks, reasoning_content=reasoning_content, ), @@ -1739,9 +1550,7 @@ class AWSEventStreamDecoder: except Exception as e: raise Exception("Received streaming error - {}".format(str(e))) - def _chunk_parser( - self, chunk_data: dict - ) -> Union[GChunk, ModelResponseStream, dict]: + def _chunk_parser(self, chunk_data: dict) -> Union[GChunk, ModelResponseStream, dict]: text = "" is_finished = False finish_reason = "" @@ -1767,10 +1576,7 @@ class AWSEventStreamDecoder: return self.converse_chunk_parser(chunk_data=_chunk_data) ######## bedrock.mistral mappings ############### elif "outputs" in chunk_data: - if ( - len(chunk_data["outputs"]) == 1 - and chunk_data["outputs"][0].get("text", None) is not None - ): + if len(chunk_data["outputs"]) == 1 and chunk_data["outputs"][0].get("text", None) is not None: text = chunk_data["outputs"][0]["text"] stop_reason = chunk_data.get("stop_reason", None) if stop_reason is not None: @@ -1799,9 +1605,7 @@ class AWSEventStreamDecoder: tool_use=None, ) - def iter_bytes( - self, iterator: Iterator[bytes] - ) -> Iterator[Union[GChunk, ModelResponseStream, dict]]: + def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[Union[GChunk, ModelResponseStream, dict]]: """Given an iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer @@ -1897,9 +1701,7 @@ class AmazonDeepSeekR1StreamDecoder(AWSEventStreamDecoder): sync_stream=sync_stream, ) - def _chunk_parser( - self, chunk_data: dict - ) -> Union[GChunk, ModelResponseStream, dict]: + def _chunk_parser(self, chunk_data: dict) -> Union[GChunk, ModelResponseStream, dict]: return self.deepseek_model_response_iterator.chunk_parser(chunk=chunk_data) @@ -1933,9 +1735,7 @@ class MockResponseIterator: # for returning ai21 streaming responses """ tool_use: Optional[ChatCompletionToolCallChunk] = None if self.json_mode is True and tool_calls is not None: - message = litellm.AnthropicConfig()._convert_tool_response_to_message( - tool_calls=tool_calls - ) + message = litellm.AnthropicConfig()._convert_tool_response_to_message(tool_calls=tool_calls) if message is not None: text = message.content or "" tool_use = None @@ -1971,9 +1771,7 @@ class MockResponseIterator: # for returning ai21 streaming responses text=text, tool_use=tool_use, is_finished=True, - finish_reason=map_finish_reason( - finish_reason=chunk_data.choices[0].finish_reason or "" - ), + finish_reason=map_finish_reason(finish_reason=chunk_data.choices[0].finish_reason or ""), usage=ChatCompletionUsageBlock( prompt_tokens=chunk_usage.prompt_tokens, completion_tokens=chunk_usage.completion_tokens, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py index 9c2c95e6cea..8b411b7b576 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_cohere_transformation.py @@ -54,9 +54,7 @@ class AmazonCohereConfig(AmazonInvokeConfig, CohereChatConfig): } def get_supported_openai_params(self, model: str) -> List[str]: - supported_params = CohereChatConfig.get_supported_openai_params( - self, model=model - ) + supported_params = CohereChatConfig.get_supported_openai_params(self, model=model) return supported_params def map_openai_params( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py index 0fe84b0ce0c..d3025e13a99 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py @@ -57,18 +57,11 @@ class AmazonDeepSeekR1Config(AmazonLlamaConfig): json_mode, ) prompt = cast(Optional[str], request_data.get("prompt")) - message_content = cast( - Optional[str], cast(Choices, response.choices[0]).message.get("content") - ) + message_content = cast(Optional[str], cast(Choices, response.choices[0]).message.get("content")) if prompt and prompt.strip().endswith("") and message_content: message_content_with_reasoning_token = "" + message_content - reasoning, content = _parse_content_for_reasoning( - message_content_with_reasoning_token - ) - provider_specific_fields = ( - cast(Choices, response.choices[0]).message.provider_specific_fields - or {} - ) + reasoning, content = _parse_content_for_reasoning(message_content_with_reasoning_token) + provider_specific_fields = cast(Choices, response.choices[0]).message.provider_specific_fields or {} if reasoning: provider_specific_fields["reasoning_content"] = reasoning @@ -96,9 +89,7 @@ class AmazonDeepseekR1ResponseIterator(BaseModelResponseIterator): typed_chunk = AmazonDeepSeekR1StreamingResponse(**chunk) # type: ignore generated_content = typed_chunk["generation"] if generated_content == "" and not self.has_finished_thinking: - verbose_logger.debug( - "Deepseek r1: received, setting has_finished_thinking to True" - ) + verbose_logger.debug("Deepseek r1: received, setting has_finished_thinking to True") generated_content = "" self.has_finished_thinking = True @@ -115,16 +106,8 @@ class AmazonDeepseekR1ResponseIterator(BaseModelResponseIterator): StreamingChoices( finish_reason=typed_chunk["stop_reason"], delta=Delta( - content=( - generated_content - if self.has_finished_thinking - else None - ), - reasoning_content=( - generated_content - if not self.has_finished_thinking - else None - ), + content=(generated_content if self.has_finished_thinking else None), + reasoning_content=(generated_content if not self.has_finished_thinking else None), ), ) ], diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py index 3992de4d4fc..58dfa17a722 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py @@ -87,9 +87,7 @@ class AmazonMistralConfig(AmazonInvokeConfig, BaseConfig): return optional_params @staticmethod - def get_outputText( - completion_response: dict, model_response: "ModelResponse" - ) -> str: + def get_outputText(completion_response: dict, model_response: "ModelResponse") -> str: """This function extracts the output text from a bedrock mistral completion. As a side effect, it updates the finish reason for a model response. @@ -103,17 +101,11 @@ class AmazonMistralConfig(AmazonInvokeConfig, BaseConfig): """ if "choices" in completion_response: outputText = completion_response["choices"][0]["message"]["content"] - model_response.choices[0].finish_reason = completion_response["choices"][0][ - "finish_reason" - ] + model_response.choices[0].finish_reason = completion_response["choices"][0]["finish_reason"] elif "outputs" in completion_response: outputText = completion_response["outputs"][0]["text"] - model_response.choices[0].finish_reason = completion_response["outputs"][0][ - "stop_reason" - ] + model_response.choices[0].finish_reason = completion_response["outputs"][0]["stop_reason"] else: - raise BedrockError( - message="Unexpected mistral completion response", status_code=400 - ) + raise BedrockError(message="Unexpected mistral completion response", status_code=400) return outputText diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 3aeb65b58c7..0532d677e5a 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -101,9 +101,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): "stop", ] # Bedrock doesn't support stopSequences - base_openai_params = super( - MoonshotChatConfig, self - ).get_supported_openai_params(model=model) + base_openai_params = super(MoonshotChatConfig, self).get_supported_openai_params(model=model) final_params: List[str] = [] for param in base_openai_params: if param not in excluded_params: @@ -168,9 +166,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): headers=headers, ) - def _extract_reasoning_from_content( - self, content: str - ) -> tuple[Optional[str], str]: + def _extract_reasoning_from_content(self, content: str) -> tuple[Optional[str], str]: """ Extract reasoning content from tags in the response. @@ -187,9 +183,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): return None, content # Match ... tags - reasoning_match = re.match( - r"(.*?)\s*(.*)", content, re.DOTALL - ) + reasoning_match = re.match(r"(.*?)\s*(.*)", content, re.DOTALL) if reasoning_match: reasoning_content = reasoning_match.group(1).strip() @@ -241,11 +235,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): if model_response.choices and len(model_response.choices) > 0: for choice in model_response.choices: # Only process Choices (not StreamingChoices) which have message attribute - if ( - isinstance(choice, Choices) - and choice.message - and choice.message.content - ): + if isinstance(choice, Choices) and choice.message and choice.message.content: ( reasoning_content, main_content, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py index 3506c8f1cc0..acfa5021507 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -37,9 +37,7 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): model: str, drop_params: bool, ) -> dict: - return AmazonConverseConfig.map_openai_params( - self, non_default_params, optional_params, model, drop_params - ) + return AmazonConverseConfig.map_openai_params(self, non_default_params, optional_params, model, drop_params) def transform_request( self, @@ -57,13 +55,9 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): litellm_params=litellm_params, headers=headers, ) - _bedrock_invoke_nova_request = BedrockInvokeNovaRequest( - **_transformed_nova_request - ) + _bedrock_invoke_nova_request = BedrockInvokeNovaRequest(**_transformed_nova_request) self._remove_empty_system_messages(_bedrock_invoke_nova_request) - bedrock_invoke_nova_request = self._filter_allowed_fields( - _bedrock_invoke_nova_request - ) + bedrock_invoke_nova_request = self._filter_allowed_fields(_bedrock_invoke_nova_request) return bedrock_invoke_nova_request def transform_response( @@ -95,20 +89,14 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): json_mode, ) - def _filter_allowed_fields( - self, bedrock_invoke_nova_request: BedrockInvokeNovaRequest - ) -> dict: + def _filter_allowed_fields(self, bedrock_invoke_nova_request: BedrockInvokeNovaRequest) -> dict: """ Filter out fields that are not allowed in the `BedrockInvokeNovaRequest` dataclass. """ allowed_fields = set(BedrockInvokeNovaRequest.__annotations__.keys()) - return { - k: v for k, v in bedrock_invoke_nova_request.items() if k in allowed_fields - } + return {k: v for k, v in bedrock_invoke_nova_request.items() if k in allowed_fields} - def _remove_empty_system_messages( - self, bedrock_invoke_nova_request: BedrockInvokeNovaRequest - ) -> None: + def _remove_empty_system_messages(self, bedrock_invoke_nova_request: BedrockInvokeNovaRequest) -> None: """ In-place remove empty `system` messages from the request. diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index 7b64c6066d0..d3f9d8bffb8 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -82,14 +82,10 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): model_id = self._get_openai_model_id(model) # Get AWS region - aws_region_name = self._get_aws_region_name( - optional_params=optional_params, model=model - ) + aws_region_name = self._get_aws_region_name(optional_params=optional_params, model=model) # Get runtime endpoint - aws_bedrock_runtime_endpoint = optional_params.get( - "aws_bedrock_runtime_endpoint", None - ) + aws_bedrock_runtime_endpoint = optional_params.get("aws_bedrock_runtime_endpoint", None) endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, @@ -101,9 +97,7 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): # Build the invoke URL if stream: - endpoint_url = ( - f"{endpoint_url}/model/{model_id}/invoke-with-response-stream" - ) + endpoint_url = f"{endpoint_url}/model/{model_id}/invoke-with-response-stream" else: endpoint_url = f"{endpoint_url}/model/{model_id}/invoke" @@ -153,11 +147,7 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): optional_params.pop("stream", None) # Remove AWS-specific params that shouldn't be in the request body - inference_params = { - k: v - for k, v in optional_params.items() - if k not in self.aws_authentication_params - } + inference_params = {k: v for k, v in optional_params.items() if k not in self.aws_authentication_params} # Use parent class transform_request for OpenAI format return super().transform_request( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index c65e9e0b083..d63642c806f 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -57,9 +57,7 @@ class AmazonQwen2Config(AmazonQwen3Config): response_data = raw_response # Extract the generated text - Qwen2 uses "text" field, but also support "generation" for compatibility - generated_text = response_data.get("generation", "") or response_data.get( - "text", "" - ) + generated_text = response_data.get("generation", "") or response_data.get("text", "") # Clean up the response (remove assistant start token if present) if generated_text.startswith("<|im_start|>assistant\n"): diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index 6325c388181..762631cac5e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -134,9 +134,7 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): text_content.append(item.get("text", "")) elif item.get("type") == "image_url": # For Qwen3, we can include image placeholders - text_content.append( - "<|vision_start|><|image_pad|><|vision_end|>" - ) + text_content.append("<|vision_start|><|image_pad|><|vision_end|>") content = "".join(text_content) prompt_parts.append(f"<|im_start|>user\n{content}<|im_end|>") elif role == "assistant": @@ -144,9 +142,7 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): # Handle tool calls for tool_call in tool_calls: function_name = tool_call.get("function", {}).get("name", "") - function_args = tool_call.get("function", {}).get( - "arguments", "" - ) + function_args = tool_call.get("function", {}).get("arguments", "") prompt_parts.append( f'<|im_start|>assistant\n\n{{"name": "{function_name}", "arguments": "{function_args}"}}\n<|im_end|>' ) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py index 367fb84d1ac..ff9a2ee0c6d 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_titan_transformation.py @@ -105,9 +105,7 @@ class AmazonTitanConfig(AmazonInvokeConfig, BaseConfig): if k == "temperature": optional_params["temperature"] = v if k == "stop": - filtered_stop = self._map_and_modify_arg( - {"stop": v}, provider="bedrock", model=model, stop=v - ) + filtered_stop = self._map_and_modify_arg({"stop": v}, provider="bedrock", model=model, stop=v) optional_params["stopSequences"] = filtered_stop["stop"] if k == "top_p": optional_params["topP"] = v diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index 889480d31a5..6d25bb32309 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -63,9 +63,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): if param == "temperature": optional_params["temperature"] = value if param == "response_format": - optional_params["responseFormat"] = self._normalize_response_format( - value - ) + optional_params["responseFormat"] = self._normalize_response_format(value) return optional_params def _normalize_response_format(self, value: Any) -> Any: @@ -131,15 +129,11 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): return request_data def _build_media_source(self, optional_params: dict) -> Optional[dict]: - direct_source = optional_params.get("mediaSource") or optional_params.get( - "media_source" - ) + direct_source = optional_params.get("mediaSource") or optional_params.get("media_source") if isinstance(direct_source, dict): return direct_source - base64_input = optional_params.get("video_base64") or optional_params.get( - "base64_string" - ) + base64_input = optional_params.get("video_base64") or optional_params.get("base64_string") if base64_input: return {"base64String": get_base64_str(base64_input)} @@ -235,8 +229,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): if ( message_content and hasattr(model_response.choices[0], "message") - and getattr(model_response.choices[0].message, "tool_calls", None) - is None + and getattr(model_response.choices[0].message, "tool_calls", None) is None ): model_response.choices[0].message.content = message_content # type: ignore model_response.choices[0].finish_reason = finish_reason @@ -249,16 +242,10 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): ) # Calculate usage from headers - bedrock_input_tokens = raw_response.headers.get( - "x-amzn-bedrock-input-token-count", None - ) - bedrock_output_tokens = raw_response.headers.get( - "x-amzn-bedrock-output-token-count", None - ) + bedrock_input_tokens = raw_response.headers.get("x-amzn-bedrock-input-token-count", None) + bedrock_output_tokens = raw_response.headers.get("x-amzn-bedrock-output-token-count", None) - prompt_tokens = int( - bedrock_input_tokens or litellm.token_counter(messages=messages) - ) + prompt_tokens = int(bedrock_input_tokens or litellm.token_counter(messages=messages)) completion_tokens = int( bedrock_output_tokens diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 79153c3ceff..60d532eb8c5 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -88,9 +88,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # ``_clamp_adaptive_reasoning_effort_for_bedrock`` so adaptive Claude # requests degrade ``xhigh`` -> ``max`` rather than 400-ing on # models like Opus 4.6 that don't natively advertise xhigh. - self._clamp_adaptive_reasoning_effort_for_bedrock( - model=original_model, params=non_default_params - ) + self._clamp_adaptive_reasoning_effort_for_bedrock(model=original_model, params=non_default_params) optional_params = AnthropicConfig.map_openai_params( self, @@ -190,11 +188,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: - filtered_params = { - k: v - for k, v in optional_params.items() - if k not in self.aws_authentication_params - } + filtered_params = {k: v for k, v in optional_params.items() if k not in self.aws_authentication_params} output_config = filtered_params.get("output_config") if isinstance(output_config, dict): filtered_params["output_config"] = dict(output_config) @@ -217,9 +211,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("stream", None) anthropic_request.pop("stream_chunk_size", None) output_format = anthropic_request.pop("output_format", None) - output_config_format = pop_bedrock_invoke_output_config_format( - anthropic_request - ) + output_config_format = pop_bedrock_invoke_output_config_format(anthropic_request) if output_format: convert_bedrock_invoke_output_format_to_inline_schema( output_format=output_format, @@ -280,9 +272,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): ) beta_set.update(auto_betas) - if tool_search_used and not ( - programmatic_tool_calling_used or input_examples_used - ): + if tool_search_used and not (programmatic_tool_calling_used or input_examples_used): beta_set.discard(ANTHROPIC_TOOL_SEARCH_BETA_HEADER) if "opus-4" in model.lower() or "opus_4" in model.lower(): beta_set.add("tool-search-tool-2025-10-19") @@ -332,9 +322,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - async def _async_convert_document_url_sources_to_base64( - self, anthropic_request: dict - ) -> None: + async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None: """ Async version of document URL conversion for async completion paths. """ @@ -390,9 +378,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if tool_type == "tool_search_tool_regex_20251119": normalized_tool = tool.copy() normalized_tool["type"] = "tool_search_tool_regex" - normalized_tool["name"] = normalized_tool.get( - "name", "tool_search_tool_regex" - ) + normalized_tool["name"] = normalized_tool.get("name", "tool_search_tool_regex") normalized_tools.append(normalized_tool) continue normalized_tools.append(tool) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 8fc2375c224..bbe16e26713 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -95,16 +95,12 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, - aws_region_name=self._get_aws_region_name( - optional_params=optional_params, model=model - ), + aws_region_name=self._get_aws_region_name(optional_params=optional_params, model=model), ) if (stream is not None and stream is True) and provider != "ai21": endpoint_url = f"{endpoint_url}/model/{modelId}/invoke-with-response-stream" - proxy_endpoint_url = ( - f"{proxy_endpoint_url}/model/{modelId}/invoke-with-response-stream" - ) + proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke-with-response-stream" else: endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" @@ -163,11 +159,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): custom_prompt_dict=custom_prompt_dict, ) inference_params = copy.deepcopy(optional_params) - inference_params = { - k: v - for k, v in inference_params.items() - if k not in self.aws_authentication_params - } + inference_params = {k: v for k, v in inference_params.items() if k not in self.aws_authentication_params} request_data: dict = {} if provider == "cohere": if model.startswith("cohere.command-r"): @@ -183,19 +175,15 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): config = litellm.AmazonCohereConfig.get_config() self._apply_config_to_params(config, inference_params) if stream is True: - inference_params["stream"] = ( - True # cohere requires stream = True in inference params - ) + inference_params["stream"] = True # cohere requires stream = True in inference params request_data = {"prompt": prompt, **inference_params} elif provider == "anthropic": - transformed_request = ( - litellm.AmazonAnthropicClaudeConfig().transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) + transformed_request = litellm.AmazonAnthropicClaudeConfig().transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, ) return transformed_request @@ -274,9 +262,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): try: completion_response = raw_response.json() except Exception: - raise BedrockError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise BedrockError(message=raw_response.text, status_code=raw_response.status_code) verbose_logger.debug( "bedrock invoke response % s", json.dumps(completion_response, indent=4, default=str), @@ -333,22 +319,16 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): json_mode=json_mode, ) elif provider == "ai21": - outputText = ( - completion_response.get("completions")[0].get("data").get("text") - ) + outputText = completion_response.get("completions")[0].get("data").get("text") elif provider == "meta" or provider == "llama" or provider == "deepseek_r1": outputText = completion_response["generation"] elif provider == "mistral": - outputText = litellm.AmazonMistralConfig.get_outputText( - completion_response, model_response - ) + outputText = litellm.AmazonMistralConfig.get_outputText(completion_response, model_response) else: # amazon titan outputText = completion_response.get("results")[0].get("outputText") except Exception as e: raise BedrockError( - message="Error processing={}, Received error={}".format( - raw_response.text, str(e) - ), + message="Error processing={}, Received error={}".format(raw_response.text, str(e)), status_code=422, ) @@ -371,23 +351,15 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise Exception() except Exception as e: raise BedrockError( - message="Error parsing received text={}.\nError-{}".format( - outputText, str(e) - ), + message="Error parsing received text={}.\nError-{}".format(outputText, str(e)), status_code=raw_response.status_code, ) ## CALCULATING USAGE - bedrock returns usage in the headers - bedrock_input_tokens = raw_response.headers.get( - "x-amzn-bedrock-input-token-count", None - ) - bedrock_output_tokens = raw_response.headers.get( - "x-amzn-bedrock-output-token-count", None - ) + bedrock_input_tokens = raw_response.headers.get("x-amzn-bedrock-input-token-count", None) + bedrock_output_tokens = raw_response.headers.get("x-amzn-bedrock-output-token-count", None) - prompt_tokens = int( - bedrock_input_tokens or litellm.token_counter(messages=messages) - ) + prompt_tokens = int(bedrock_input_tokens or litellm.token_counter(messages=messages)) completion_tokens = int( bedrock_output_tokens @@ -565,9 +537,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider) return None - def convert_messages_to_prompt( - self, model, messages, provider, custom_prompt_dict - ) -> Tuple[str, Optional[list]]: + def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> Tuple[str, Optional[list]]: # handle anthropic prompts and amazon titan prompts prompt = "" chat_history: Optional[list] = None @@ -577,26 +547,18 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): model_prompt_details = custom_prompt_dict[model] prompt = custom_prompt( role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details.get( - "initial_prompt_value", "" - ), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), final_prompt_value=model_prompt_details.get("final_prompt_value", ""), messages=messages, ) return prompt, None ## ELSE if provider == "anthropic" or provider == "amazon": - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="bedrock" - ) + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock") elif provider == "mistral": - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="bedrock" - ) + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock") elif provider == "meta" or provider == "llama": - prompt = prompt_factory( - model=model, messages=messages, custom_llm_provider="bedrock" - ) + prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock") elif provider == "cohere": prompt, chat_history = cohere_message_pt(messages=messages) elif provider == "deepseek_r1": diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index 93306025b02..d84e077c37b 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -45,9 +45,7 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): region = self._get_aws_region_name(optional_params=optional_params, model=model) return build_mantle_messages_url( api_base=api_base, - aws_bedrock_runtime_endpoint=optional_params.get( - "aws_bedrock_runtime_endpoint" - ), + aws_bedrock_runtime_endpoint=optional_params.get("aws_bedrock_runtime_endpoint"), region=region, ) diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py index 3abb8710de7..b93577e2bca 100644 --- a/litellm/llms/bedrock/claude_platform/common_utils.py +++ b/litellm/llms/bedrock/claude_platform/common_utils.py @@ -4,9 +4,7 @@ import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.secret_managers.main import get_secret_str -CLAUDE_PLATFORM_SERVICE_NAME: Literal["aws-external-anthropic"] = ( - "aws-external-anthropic" -) +CLAUDE_PLATFORM_SERVICE_NAME: Literal["aws-external-anthropic"] = "aws-external-anthropic" CLAUDE_PLATFORM_BEDROCK_ROUTE = "claude_platform/" @@ -28,14 +26,10 @@ class BedrockClaudePlatformMixin(BaseAWSLLM): or litellm_params.get("anthropic-workspace-id") ) if workspace_id is None: - workspace_id = optional_params.get( - "anthropic_workspace_id" - ) or litellm_params.get("anthropic_workspace_id") + workspace_id = optional_params.get("anthropic_workspace_id") or litellm_params.get("anthropic_workspace_id") if workspace_id is not None: return str(workspace_id) - return get_secret_str("ANTHROPIC_AWS_WORKSPACE_ID") or get_secret_str( - "ANTHROPIC_WORKSPACE_ID" - ) + return get_secret_str("ANTHROPIC_AWS_WORKSPACE_ID") or get_secret_str("ANTHROPIC_WORKSPACE_ID") def _get_required_aws_region_name(self, optional_params: dict) -> str: aws_region_name = ( @@ -73,9 +67,7 @@ class BedrockClaudePlatformMixin(BaseAWSLLM): ) if api_base is None: aws_region_name = self._get_required_aws_region_name(optional_params) - api_base = ( - f"https://{CLAUDE_PLATFORM_SERVICE_NAME}.{aws_region_name}.api.aws" - ) + api_base = f"https://{CLAUDE_PLATFORM_SERVICE_NAME}.{aws_region_name}.api.aws" if not api_base.endswith("/v1/messages"): api_base = f"{api_base.rstrip('/')}/v1/messages" return api_base diff --git a/litellm/llms/bedrock/claude_platform/messages_transformation.py b/litellm/llms/bedrock/claude_platform/messages_transformation.py index 66158196322..1b0d21a724c 100644 --- a/litellm/llms/bedrock/claude_platform/messages_transformation.py +++ b/litellm/llms/bedrock/claude_platform/messages_transformation.py @@ -11,9 +11,7 @@ from litellm.types.router import GenericLiteLLMParams from .common_utils import BedrockClaudePlatformMixin, strip_claude_platform_route -class BedrockClaudePlatformMessagesConfig( - BedrockClaudePlatformMixin, AnthropicMessagesConfig -): +class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicMessagesConfig): def validate_anthropic_messages_environment( self, headers: dict, @@ -38,9 +36,7 @@ class BedrockClaudePlatformMessagesConfig( resolved_api_key = api_key or get_secret_str("ANTHROPIC_AWS_API_KEY") headers = { **headers, - "anthropic-version": headers.get( - "anthropic-version", DEFAULT_ANTHROPIC_API_VERSION - ), + "anthropic-version": headers.get("anthropic-version", DEFAULT_ANTHROPIC_API_VERSION), "content-type": headers.get("content-type", "application/json"), "anthropic-workspace-id": workspace_id, } diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py index c20dc63444f..0868d9bddfe 100644 --- a/litellm/llms/bedrock/claude_platform/transformation.py +++ b/litellm/llms/bedrock/claude_platform/transformation.py @@ -45,39 +45,21 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig): anthropic_headers = self.get_anthropic_headers( api_key=api_key, auth_token=None, - computer_tool_used=self.is_computer_tool_used( - tools=optional_params.get("tools") - ), + computer_tool_used=self.is_computer_tool_used(tools=optional_params.get("tools")), prompt_caching_set=self.is_cache_control_set(messages=messages), pdf_used=self.is_pdf_used(messages=messages), file_id_used=self.is_file_id_used(messages=messages), - mcp_server_used=self.is_mcp_server_used( - mcp_servers=optional_params.get("mcp_servers") - ), - web_search_tool_used=self.is_web_search_tool_used( - tools=optional_params.get("tools") - ), - tool_search_used=self.is_tool_search_used( - tools=optional_params.get("tools") - ), - programmatic_tool_calling_used=self.is_programmatic_tool_calling_used( - tools=optional_params.get("tools") - ), - input_examples_used=self.is_input_examples_used( - tools=optional_params.get("tools") - ), - effort_used=self.is_effort_used( - optional_params=optional_params, model=model - ), + mcp_server_used=self.is_mcp_server_used(mcp_servers=optional_params.get("mcp_servers")), + web_search_tool_used=self.is_web_search_tool_used(tools=optional_params.get("tools")), + tool_search_used=self.is_tool_search_used(tools=optional_params.get("tools")), + programmatic_tool_calling_used=self.is_programmatic_tool_calling_used(tools=optional_params.get("tools")), + input_examples_used=self.is_input_examples_used(tools=optional_params.get("tools")), + effort_used=self.is_effort_used(optional_params=optional_params, model=model), user_anthropic_beta_headers=self._get_user_anthropic_beta_headers( anthropic_beta_header=headers.get("anthropic-beta") ), - code_execution_tool_used=self.is_code_execution_tool_used( - tools=optional_params.get("tools") - ), - container_with_skills_used=self.is_container_with_skills_used( - optional_params=optional_params - ), + code_execution_tool_used=self.is_code_execution_tool_used(tools=optional_params.get("tools")), + container_with_skills_used=self.is_container_with_skills_used(optional_params=optional_params), ) anthropic_headers["anthropic-workspace-id"] = workspace_id return {**headers, **anthropic_headers} diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 44312eb3926..467e1050c99 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -431,11 +431,7 @@ def init_bedrock_client( config = boto3.session.Config() # type: ignore ### CHECK STS ### - if ( - aws_web_identity_token is not None - and aws_role_name is not None - and aws_session_name is not None - ): + if aws_web_identity_token is not None and aws_role_name is not None and aws_session_name is not None: oidc_token = get_secret(aws_web_identity_token) if oidc_token is None: @@ -474,9 +470,7 @@ def init_bedrock_client( verify=ssl_verify, ) - sts_response = sts_client.assume_role( - RoleArn=aws_role_name, RoleSessionName=aws_session_name - ) + sts_response = sts_client.assume_role(RoleArn=aws_role_name, RoleSessionName=aws_session_name) client = boto3.client( service_name="bedrock-runtime", @@ -523,9 +517,7 @@ def init_bedrock_client( verify=ssl_verify, ) if extra_headers: - client.meta.events.register( - "before-sign.bedrock-runtime.*", add_custom_header(extra_headers) - ) + client.meta.events.register("before-sign.bedrock-runtime.*", add_custom_header(extra_headers)) return client @@ -568,9 +560,7 @@ def get_bedrock_tool_name(response_tool_name: str) -> str: """ if response_tool_name in litellm.bedrock_tool_name_mappings.cache_dict: - response_tool_name = litellm.bedrock_tool_name_mappings.cache_dict[ - response_tool_name - ] + response_tool_name = litellm.bedrock_tool_name_mappings.cache_dict[response_tool_name] return response_tool_name @@ -687,10 +677,7 @@ def get_bedrock_base_model(model: str) -> str: if potential_region in get_bedrock_cross_region_inference_regions(): return model.split(".", 1)[1] - elif ( - alt_potential_region in _get_all_bedrock_regions() - and len(model.split("/", 1)) > 1 - ): + elif alt_potential_region in _get_all_bedrock_regions() and len(model.split("/", 1)) > 1: return model.split("/", 1)[1] return model @@ -754,10 +741,7 @@ def normalize_bedrock_opus_output_config_effort(model: str, output_config: Any) if ceiling is None: return - if ( - _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER[effort] - > _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER[ceiling] - ): + if _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER[effort] > _BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER[ceiling]: output_config["effort"] = ceiling @@ -821,9 +805,7 @@ class BedrockModelInfo(BaseLLMModelInfo): ) -> dict: return headers - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: return [] # def get_provider_info(self, model: str) -> Optional[ProviderSpecificModelInfo]: @@ -923,9 +905,7 @@ class BedrockModelInfo(BaseLLMModelInfo): # Check for nova spec prefixes (nova/ and nova-2/) _model_after_bedrock = model.replace("bedrock/", "", 1) - if _model_after_bedrock.startswith( - "nova-2/" - ) or _model_after_bedrock.startswith("nova/"): + if _model_after_bedrock.startswith("nova-2/") or _model_after_bedrock.startswith("nova/"): return "converse" if is_bedrock_application_inference_profile_arn(model): @@ -933,10 +913,7 @@ class BedrockModelInfo(BaseLLMModelInfo): base_model = BedrockModelInfo.get_base_model(model) alt_model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model) - if ( - base_model in litellm.bedrock_converse_models - or alt_model in litellm.bedrock_converse_models - ): + if base_model in litellm.bedrock_converse_models or alt_model in litellm.bedrock_converse_models: return "converse" return "invoke" @@ -962,9 +939,7 @@ class BedrockModelInfo(BaseLLMModelInfo): return model.replace("claude_platform/", "", 1) @staticmethod - def map_claude_platform_auth_params( - passed_params: dict, optional_params: dict - ) -> dict: + def map_claude_platform_auth_params(passed_params: dict, optional_params: dict) -> dict: """ Map Claude Platform route auth params that are not OpenAI request params. """ @@ -1100,9 +1075,7 @@ def get_bedrock_chat_config(model: str): The appropriate Bedrock config class instance """ bedrock_route = BedrockModelInfo.get_bedrock_route(model) - bedrock_invoke_provider = litellm.BedrockLLM.get_bedrock_invoke_provider( - model=model - ) + bedrock_invoke_provider = litellm.BedrockLLM.get_bedrock_invoke_provider(model=model) base_model = BedrockModelInfo.get_base_model(model) # Handle explicit routes first @@ -1135,10 +1108,7 @@ def get_bedrock_chat_config(model: str): if bedrock_invoke_provider == "amazon": return litellm.AmazonTitanConfig() elif bedrock_invoke_provider == "anthropic": - if ( - base_model - in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names() - ): + if base_model in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names(): return litellm.AmazonAnthropicConfig() else: return litellm.AmazonAnthropicClaudeConfig() @@ -1224,9 +1194,7 @@ def build_bedrock_stream_error( if exception_type is not None and response_stream_shape is not None: member = response_stream_shape.members.get(exception_type) if member is not None: - modeled_status = ( - (member.metadata or {}).get("error", {}).get("httpStatusCode") - ) + modeled_status = (member.metadata or {}).get("error", {}).get("httpStatusCode") if modeled_status is not None: status_code = int(modeled_status) @@ -1296,9 +1264,7 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: # Try to parse as JSON array first (e.g., '["interleaved-thinking-2025-05-14", "claude-code-20250219"]') if isinstance(anthropic_beta_header, str): anthropic_beta_header = anthropic_beta_header.strip() - if anthropic_beta_header.startswith("[") and anthropic_beta_header.endswith( - "]" - ): + if anthropic_beta_header.startswith("[") and anthropic_beta_header.endswith("]"): try: parsed = json.loads(anthropic_beta_header) if isinstance(parsed, list): @@ -1360,9 +1326,7 @@ class CommonBatchFilesUtils: return s3_parts[0], s3_parts[1] # bucket, key - def extract_model_from_s3_file_path( - self, s3_uri: str, optional_params: dict - ) -> str: + def extract_model_from_s3_file_path(self, s3_uri: str, optional_params: dict) -> str: """ Extract model ID from S3 file path. @@ -1371,9 +1335,7 @@ class CommonBatchFilesUtils: """ # Check if model is provided in optional_params first if "model" in optional_params and optional_params["model"]: - return self.get_bedrock_model_id_from_litellm_model( - optional_params["model"] - ) + return self.get_bedrock_model_id_from_litellm_model(optional_params["model"]) # Extract model from S3 URI path # Expected format: s3://bucket/litellm-bedrock-files-{model}-{uuid}.jsonl @@ -1426,9 +1388,7 @@ class CommonBatchFilesUtils: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") # Get AWS credentials using existing methods - aws_region_name = self._base_aws._get_aws_region_name( - optional_params=optional_params, model="" - ) + aws_region_name = self._base_aws._get_aws_region_name(optional_params=optional_params, model="") credentials = self._base_aws.get_credentials( aws_access_key_id=optional_params.get("aws_access_key_id"), aws_secret_access_key=optional_params.get("aws_secret_access_key"), @@ -1459,19 +1419,13 @@ class CommonBatchFilesUtils: # Create AWS request and sign it sigv4 = SigV4Auth(credentials, service_name, aws_region_name) - request = AWSRequest( - method=method_upper, url=endpoint_url, data=request_data, headers=headers - ) + request = AWSRequest(method=method_upper, url=endpoint_url, data=request_data, headers=headers) sigv4.add_auth(request) prepped = request.prepare() return ( dict(prepped.headers), - ( - request_data.encode("utf-8") - if isinstance(request_data, str) - else request_data - ), + (request_data.encode("utf-8") if isinstance(request_data, str) else request_data), ) def generate_unique_job_name(self, model: str, prefix: str = "litellm") -> str: @@ -1520,14 +1474,10 @@ class CommonBatchFilesUtils: # Get bucket name bucket_name = ( - litellm_params.get("s3_bucket_name") - or optional_params.get("s3_bucket_name") - or os.getenv(bucket_env_var) + litellm_params.get("s3_bucket_name") or optional_params.get("s3_bucket_name") or os.getenv(bucket_env_var) ) if not bucket_name: - raise ValueError( - f"S3 bucket name is required. Set 's3_bucket_name' parameter or {bucket_env_var} env var" - ) + raise ValueError(f"S3 bucket name is required. Set 's3_bucket_name' parameter or {bucket_env_var} env var") # Generate unique object key timestamp = int(time.time()) @@ -1542,6 +1492,4 @@ class CommonBatchFilesUtils: """ Get Bedrock-specific error class. """ - return BedrockError( - status_code=status_code, message=error_message, headers=headers - ) + return BedrockError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/bedrock/cost_calculation.py b/litellm/llms/bedrock/cost_calculation.py index ac99d4e36e7..9a164d02eeb 100644 --- a/litellm/llms/bedrock/cost_calculation.py +++ b/litellm/llms/bedrock/cost_calculation.py @@ -11,9 +11,7 @@ if TYPE_CHECKING: from litellm.types.utils import Usage -def cost_per_token( - model: str, usage: "Usage", service_tier: Optional[str] = None -) -> Tuple[float, float]: +def cost_per_token(model: str, usage: "Usage", service_tier: Optional[str] = None) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index eb7755574ac..1ea870a1d32 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -88,9 +88,7 @@ class BedrockTokenCounter(BaseTokenCounter): original_response=result, ) except BedrockError as e: - verbose_logger.warning( - f"Bedrock CountTokens API error: status={e.status_code}, message={e.message}" - ) + verbose_logger.warning(f"Bedrock CountTokens API error: status={e.status_code}, message={e.message}") return TokenCountResponse( total_tokens=0, request_model=request_model, diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 8c227c853cc..2c40e14129d 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -43,9 +43,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): # Validate the request self.validate_count_tokens_request(request_data) - verbose_logger.debug( - f"Processing CountTokens request for resolved model: {resolved_model}" - ) + verbose_logger.debug(f"Processing CountTokens request for resolved model: {resolved_model}") # Get AWS region using existing LiteLLM function aws_region_name = self._get_aws_region_name( @@ -57,17 +55,13 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): verbose_logger.debug(f"Retrieved AWS region: {aws_region_name}") # Transform request to Bedrock format (supports both Converse and InvokeModel) - bedrock_request = self.transform_anthropic_to_bedrock_count_tokens( - request_data=request_data - ) + bedrock_request = self.transform_anthropic_to_bedrock_count_tokens(request_data=request_data) verbose_logger.debug(f"Transformed request: {bedrock_request}") # Get endpoint URL using simplified function api_base = litellm_params.get("api_base", None) - aws_bedrock_runtime_endpoint = litellm_params.get( - "aws_bedrock_runtime_endpoint", None - ) + aws_bedrock_runtime_endpoint = litellm_params.get("aws_bedrock_runtime_endpoint", None) endpoint_url = self.get_bedrock_count_tokens_endpoint( model=resolved_model, aws_region_name=aws_region_name, @@ -91,9 +85,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): api_key=api_key, ) - async_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.BEDROCK - ) + async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) response = await async_client.post( endpoint_url, @@ -117,9 +109,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): verbose_logger.debug(f"Bedrock response: {bedrock_response}") # Transform response back to expected format - final_response = self.transform_bedrock_response_to_anthropic( - bedrock_response - ) + final_response = self.transform_bedrock_response_to_anthropic(bedrock_response) verbose_logger.debug(f"Final response: {final_response}") diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index bdef3349e00..38eaf13893d 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -47,9 +47,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): if not isinstance(message, dict): continue content = message.get("content") - if isinstance(content, list) and any( - isinstance(block, dict) and "type" in block for block in content - ): + if isinstance(content, list) and any(isinstance(block, dict) and "type" in block for block in content): return "invokeModel" return "converse" @@ -97,9 +95,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): else: return self._transform_to_invoke_model_format(request_data) - def _transform_to_converse_format( - self, request_data: Dict[str, Any] - ) -> Dict[str, Any]: + def _transform_to_converse_format(self, request_data: Dict[str, Any]) -> Dict[str, Any]: """Transform to Converse input format, including system and tools.""" messages = request_data.get("messages", []) system = request_data.get("system") @@ -141,16 +137,10 @@ class BedrockCountTokensConfig(BaseAWSLLM): return [{"text": system}] if isinstance(system, list): # Already in blocks format (e.g. [{"type": "text", "text": "..."}]) - return [ - {"text": block.get("text", "")} - for block in system - if isinstance(block, dict) - ] + return [{"text": block.get("text", "")} for block in system if isinstance(block, dict)] return [] - def _transform_tools( - self, tools: Optional[List[Dict[str, Any]]] - ) -> Optional[Dict[str, Any]]: + def _transform_tools(self, tools: Optional[List[Dict[str, Any]]]) -> Optional[Dict[str, Any]]: """Transform Anthropic tools to Bedrock toolConfig format.""" if not tools: return None @@ -165,9 +155,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): name = name[:64] description = tool.get("description") or name - input_schema = tool.get( - "input_schema", {"type": "object", "properties": {}} - ) + input_schema = tool.get("input_schema", {"type": "object", "properties": {}}) bedrock_tools.append( { @@ -181,9 +169,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): return {"tools": bedrock_tools} - def _transform_to_invoke_model_format( - self, request_data: Dict[str, Any] - ) -> Dict[str, Any]: + def _transform_to_invoke_model_format(self, request_data: Dict[str, Any]) -> Dict[str, Any]: """Transform to InvokeModel input format.""" import base64 import json @@ -196,9 +182,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): # Bedrock validates the body against the model's InvokeModel schema; # Anthropic Messages bodies require these fields. body_data.setdefault("anthropic_version", "bedrock-2023-05-31") - body_data.setdefault( - "max_tokens", DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS - ) + body_data.setdefault("max_tokens", DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS) # The CountTokens API expects invokeModel.body as a base64-encoded blob encoded_body = base64.b64encode(json.dumps(body_data).encode()).decode() @@ -240,9 +224,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): return endpoint - def transform_bedrock_response_to_anthropic( - self, bedrock_response: Dict[str, Any] - ) -> Dict[str, Any]: + def transform_bedrock_response_to_anthropic(self, bedrock_response: Dict[str, Any]) -> Dict[str, Any]: """ Transform Bedrock CountTokens response to Anthropic format. diff --git a/litellm/llms/bedrock/embed/amazon_nova_transformation.py b/litellm/llms/bedrock/embed/amazon_nova_transformation.py index c20b52a6e0d..58519d0d061 100644 --- a/litellm/llms/bedrock/embed/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_nova_transformation.py @@ -40,9 +40,7 @@ class AmazonNovaEmbeddingConfig: "dimensions", ] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: """Map OpenAI-style parameters to Nova parameters.""" for k, v in non_default_params.items(): if k == "dimensions": @@ -70,9 +68,7 @@ class AmazonNovaEmbeddingConfig: # Split by comma to separate metadata from data # Format: data:image/jpeg;base64, if "," not in data_url: - raise ValueError( - f"Invalid data URL format (missing comma): {data_url[:50]}..." - ) + raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...") metadata, base64_data = data_url.split(",", 1) @@ -129,9 +125,7 @@ class AmazonNovaEmbeddingConfig: if "dimensions" in embedding_params: embedding_params["embeddingDimension"] = embedding_params.pop("dimensions") elif "embedding_dimension" in embedding_params: - embedding_params["embeddingDimension"] = embedding_params.pop( - "embedding_dimension" - ) + embedding_params["embeddingDimension"] = embedding_params.pop("embedding_dimension") # Add required embeddingPurpose if not provided (required by Nova API) if "embeddingPurpose" not in embedding_params: @@ -322,9 +316,7 @@ class AmazonNovaEmbeddingConfig: return EmbeddingResponse(data=embeddings, model=model, usage=usage) - def _transform_async_invoke_response( - self, response: dict, model: str - ) -> EmbeddingResponse: + def _transform_async_invoke_response(self, response: dict, model: str) -> EmbeddingResponse: """ Transform async invoke response (invocation ARN) to OpenAI format. diff --git a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py index 64a79b73273..57cbb3263de 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py @@ -53,19 +53,13 @@ class AmazonTitanG1Config: def get_supported_openai_params(self) -> List[str]: return [] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: return optional_params - def _transform_request( - self, input: str, inference_params: dict - ) -> AmazonTitanG1EmbeddingRequest: + def _transform_request(self, input: str, inference_params: dict) -> AmazonTitanG1EmbeddingRequest: return AmazonTitanG1EmbeddingRequest(inputText=input) - def _transform_response( - self, response_list: List[dict], model: str - ) -> EmbeddingResponse: + def _transform_response(self, response_list: List[dict], model: str) -> EmbeddingResponse: total_prompt_tokens = 0 transformed_responses: List[Embedding] = [] diff --git a/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py index 2713f54e623..878d5f7e850 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py @@ -33,26 +33,18 @@ class AmazonTitanMultimodalEmbeddingG1Config: def get_supported_openai_params(self) -> List[str]: return ["dimensions"] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "dimensions": - optional_params["embeddingConfig"] = ( - AmazonTitanMultimodalEmbeddingConfig(outputEmbeddingLength=v) - ) + optional_params["embeddingConfig"] = AmazonTitanMultimodalEmbeddingConfig(outputEmbeddingLength=v) return optional_params - def _transform_request( - self, input: str, inference_params: dict - ) -> AmazonTitanMultimodalEmbeddingRequest: + def _transform_request(self, input: str, inference_params: dict) -> AmazonTitanMultimodalEmbeddingRequest: ## check if b64 encoded str or not ## is_encoded = is_base64_encoded(input) if is_encoded: # check if string is b64 encoded image or not b64_str = get_base64_str(input) - transformed_request = AmazonTitanMultimodalEmbeddingRequest( - inputImage=b64_str - ) + transformed_request = AmazonTitanMultimodalEmbeddingRequest(inputImage=b64_str) else: transformed_request = AmazonTitanMultimodalEmbeddingRequest(inputText=input) diff --git a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py index ca0b95cd64e..2c7b0ba465a 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py @@ -30,9 +30,7 @@ class AmazonTitanV2Config: normalize: Optional[bool] = None dimensions: Optional[int] = None - def __init__( - self, normalize: Optional[bool] = None, dimensions: Optional[int] = None - ) -> None: + def __init__(self, normalize: Optional[bool] = None, dimensions: Optional[int] = None) -> None: locals_ = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: @@ -59,9 +57,7 @@ class AmazonTitanV2Config: def get_supported_openai_params(self) -> List[str]: return ["dimensions", "encoding_format"] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "dimensions": optional_params["dimensions"] = v @@ -77,14 +73,10 @@ class AmazonTitanV2Config: optional_params["embeddingTypes"] = ["float"] return optional_params - def _transform_request( - self, input: str, inference_params: dict - ) -> AmazonTitanV2EmbeddingRequest: + def _transform_request(self, input: str, inference_params: dict) -> AmazonTitanV2EmbeddingRequest: return AmazonTitanV2EmbeddingRequest(inputText=input, **inference_params) # type: ignore - def _transform_response( - self, response_list: List[dict], model: str - ) -> EmbeddingResponse: + def _transform_response(self, response_list: List[dict], model: str) -> EmbeddingResponse: total_prompt_tokens = 0 transformed_responses: List[Embedding] = [] @@ -96,16 +88,10 @@ class AmazonTitanV2Config: # Otherwise, use float data from embeddingsByType or fallback to embedding field embedding_data: Union[List[float], List[int]] - if ( - "embeddingsByType" in _parsed_response - and "binary" in _parsed_response["embeddingsByType"] - ): + if "embeddingsByType" in _parsed_response and "binary" in _parsed_response["embeddingsByType"]: # Use binary data if available (for encoding_format="base64") embedding_data = _parsed_response["embeddingsByType"]["binary"] - elif ( - "embeddingsByType" in _parsed_response - and "float" in _parsed_response["embeddingsByType"] - ): + elif "embeddingsByType" in _parsed_response and "float" in _parsed_response["embeddingsByType"]: # Use float data from embeddingsByType embedding_data = _parsed_response["embeddingsByType"]["float"] elif "embedding" in _parsed_response: diff --git a/litellm/llms/bedrock/embed/cohere_transformation.py b/litellm/llms/bedrock/embed/cohere_transformation.py index 9570ff1a14c..ac3130ea434 100644 --- a/litellm/llms/bedrock/embed/cohere_transformation.py +++ b/litellm/llms/bedrock/embed/cohere_transformation.py @@ -17,9 +17,7 @@ class BedrockCohereEmbeddingConfig: def get_supported_openai_params(self) -> List[str]: return ["encoding_format", "dimensions"] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": optional_params["embedding_types"] = v if isinstance(v, list) else [v] @@ -30,12 +28,8 @@ class BedrockCohereEmbeddingConfig: def _is_v3_model(self, model: str) -> bool: return "3" in model - def _transform_request( - self, model: str, input: List[str], inference_params: dict - ) -> CohereEmbeddingRequest: - transformed_request = CohereEmbeddingConfig()._transform_request( - model, input, inference_params - ) + def _transform_request(self, model: str, input: List[str], inference_params: dict) -> CohereEmbeddingRequest: + transformed_request = CohereEmbeddingConfig()._transform_request(model, input, inference_params) new_transformed_request = CohereEmbeddingRequest( input_type=transformed_request["input_type"], diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index e07ccb8c11b..ff138709ac0 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -63,15 +63,11 @@ class BedrockEmbedding(BaseAWSLLM): # check env # litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) - if litellm_aws_region_name is not None and isinstance( - litellm_aws_region_name, str - ): + if litellm_aws_region_name is not None and isinstance(litellm_aws_region_name, str): aws_region_name = litellm_aws_region_name standard_aws_region_name = get_secret("AWS_REGION", None) - if standard_aws_region_name is not None and isinstance( - standard_aws_region_name, str - ): + if standard_aws_region_name is not None and isinstance(standard_aws_region_name, str): aws_region_name = standard_aws_region_name if aws_region_name is None: @@ -135,16 +131,12 @@ class BedrockEmbedding(BaseAWSLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = get_async_httpx_client( - params=_params, llm_provider=litellm.LlmProviders.BEDROCK - ) + client = get_async_httpx_client(params=_params, llm_provider=litellm.LlmProviders.BEDROCK) else: client = client try: - response = await client.post( - url=api_base, headers=headers, data=json.dumps(data) - ) # type: ignore + response = await client.post(url=api_base, headers=headers, data=json.dumps(data)) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -168,22 +160,14 @@ class BedrockEmbedding(BaseAWSLLM): returned_response: Optional[EmbeddingResponse] = None # Handle async invoke responses (single response with invocationArn) - if ( - is_async_invoke - and len(response_list) == 1 - and "invocationArn" in response_list[0] - ): + if is_async_invoke and len(response_list) == 1 and "invocationArn" in response_list[0]: if provider == "twelvelabs": - returned_response = ( - TwelveLabsMarengoEmbeddingConfig()._transform_async_invoke_response( - response=response_list[0], model=model - ) + returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_async_invoke_response( + response=response_list[0], model=model ) elif provider == "nova": - returned_response = ( - AmazonNovaEmbeddingConfig()._transform_async_invoke_response( - response=response_list[0], model=model - ) + returned_response = AmazonNovaEmbeddingConfig()._transform_async_invoke_response( + response=response_list[0], model=model ) else: # For other providers, create a generic async response @@ -213,28 +197,18 @@ class BedrockEmbedding(BaseAWSLLM): else: # Handle regular invoke responses if model == "amazon.titan-embed-image-v1": - returned_response = ( - AmazonTitanMultimodalEmbeddingG1Config()._transform_response( - response_list=response_list, model=model, batch_data=batch_data - ) + returned_response = AmazonTitanMultimodalEmbeddingG1Config()._transform_response( + response_list=response_list, model=model, batch_data=batch_data ) elif model == "amazon.titan-embed-text-v1": - returned_response = AmazonTitanG1Config()._transform_response( - response_list=response_list, model=model - ) + returned_response = AmazonTitanG1Config()._transform_response(response_list=response_list, model=model) elif model == "amazon.titan-embed-text-v2:0": - returned_response = AmazonTitanV2Config()._transform_response( - response_list=response_list, model=model - ) + returned_response = AmazonTitanV2Config()._transform_response(response_list=response_list, model=model) elif model == "amazon.titan-embed-g1-text-02": - returned_response = AmazonTitanG1Config()._transform_response( - response_list=response_list, model=model - ) + returned_response = AmazonTitanG1Config()._transform_response(response_list=response_list, model=model) elif provider == "twelvelabs": - returned_response = ( - TwelveLabsMarengoEmbeddingConfig()._transform_response( - response_list=response_list, model=model - ) + returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=response_list, model=model ) elif provider == "nova": returned_response = AmazonNovaEmbeddingConfig()._transform_response( @@ -245,11 +219,7 @@ class BedrockEmbedding(BaseAWSLLM): # Validate returned response ########################################################## if returned_response is None: - raise Exception( - "Unable to map model response to known provider format. model={}".format( - model - ) - ) + raise Exception("Unable to map model response to known provider format. model={}".format(model)) return returned_response def _single_func_embeddings( @@ -293,9 +263,7 @@ class BedrockEmbedding(BaseAWSLLM): "headers": prepped.headers, }, ) - headers_for_request = ( - dict(prepped.headers) if hasattr(prepped, "headers") else {} - ) + headers_for_request = dict(prepped.headers) if hasattr(prepped, "headers") else {} response = self._make_sync_call( client=client, timeout=timeout, @@ -365,9 +333,7 @@ class BedrockEmbedding(BaseAWSLLM): ) # Convert CaseInsensitiveDict to regular dict for httpx compatibility # This ensures custom headers are properly forwarded, especially with IAM roles and custom api_base - headers_for_request = ( - dict(prepped.headers) if hasattr(prepped, "headers") else {} - ) + headers_for_request = dict(prepped.headers) if hasattr(prepped, "headers") else {} response = await self._make_async_call( client=client, timeout=timeout, @@ -414,9 +380,7 @@ class BedrockEmbedding(BaseAWSLLM): credentials, aws_region_name = self._load_credentials(optional_params) ### TRANSFORMATION ### - unencoded_model_id = ( - optional_params.pop("model_id", None) or model - ) # default to model if not passed + unencoded_model_id = optional_params.pop("model_id", None) or model # default to model if not passed modelId = urllib.parse.quote(unencoded_model_id, safe="") aws_region_name = self._get_aws_region_name( optional_params={"aws_region_name": aws_region_name}, @@ -435,13 +399,9 @@ class BedrockEmbedding(BaseAWSLLM): ) inference_params = copy.deepcopy(optional_params) inference_params = { - k: v - for k, v in inference_params.items() - if k.lower() not in self.aws_authentication_params + k: v for k, v in inference_params.items() if k.lower() not in self.aws_authentication_params } - inference_params.pop( - "user", None - ) # make sure user is not passed in for bedrock call + inference_params.pop("user", None) # make sure user is not passed in for bedrock call data: Optional[CohereEmbeddingRequest] = None batch_data: Optional[List] = None @@ -491,14 +451,12 @@ class BedrockEmbedding(BaseAWSLLM): elif provider == "twelvelabs": batch_data = [] for i in input: - twelvelabs_request = ( - TwelveLabsMarengoEmbeddingConfig()._transform_request( - input=i, - inference_params=inference_params, - async_invoke_route=has_async_invoke, - model_id=modelId, - output_s3_uri=inference_params.get("output_s3_uri"), - ) + twelvelabs_request = TwelveLabsMarengoEmbeddingConfig()._transform_request( + input=i, + inference_params=inference_params, + async_invoke_route=has_async_invoke, + model_id=modelId, + output_s3_uri=inference_params.get("output_s3_uri"), ) batch_data.append(twelvelabs_request) elif provider == "nova": @@ -516,9 +474,7 @@ class BedrockEmbedding(BaseAWSLLM): ### SET RUNTIME ENDPOINT ### endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( api_base=api_base, - aws_bedrock_runtime_endpoint=optional_params.pop( - "aws_bedrock_runtime_endpoint", None - ), + aws_bedrock_runtime_endpoint=optional_params.pop("aws_bedrock_runtime_endpoint", None), aws_region_name=aws_region_name, ) if has_async_invoke: @@ -529,11 +485,7 @@ class BedrockEmbedding(BaseAWSLLM): if batch_data is not None: if aembedding: return self._async_single_func_embeddings( # type: ignore - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), timeout=timeout, batch_data=batch_data, credentials=credentials, @@ -547,11 +499,7 @@ class BedrockEmbedding(BaseAWSLLM): is_async_invoke=has_async_invoke, ) returned_response = self._single_func_embeddings( - client=( - client - if client is not None and isinstance(client, HTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, HTTPHandler) else None), timeout=timeout, batch_data=batch_data, credentials=credentials, @@ -586,9 +534,7 @@ class BedrockEmbedding(BaseAWSLLM): ## ROUTING ## # Convert CaseInsensitiveDict to regular dict for httpx compatibility - headers_for_request = ( - dict(prepped.headers) if hasattr(prepped, "headers") else {} - ) + headers_for_request = dict(prepped.headers) if hasattr(prepped, "headers") else {} return cohere_embedding( model=model, input=input, @@ -665,9 +611,7 @@ class BedrockEmbedding(BaseAWSLLM): if logging_obj is not None: # Create custom curl command for GET request masked_headers = logging_obj._get_masked_headers(prepped.headers) - formatted_headers = " ".join( - [f"-H '{k}: {v}'" for k, v in masked_headers.items()] - ) + formatted_headers = " ".join([f"-H '{k}: {v}'" for k, v in masked_headers.items()]) custom_curl = "\n\nGET Request Sent from LiteLLM:\n" custom_curl += "curl -X GET \\\n" custom_curl += f"{prepped.url} \\\n" @@ -697,15 +641,11 @@ class BedrockEmbedding(BaseAWSLLM): input=invocation_arn, api_key="", original_response=response, - additional_args={ - "complete_input_dict": {"invocation_arn": invocation_arn} - }, + additional_args={"complete_input_dict": {"invocation_arn": invocation_arn}}, ) # Parse response if response.status_code == 200: return response.json() else: - raise Exception( - f"Failed to get async invoke status: {response.status_code} - {response.text}" - ) + raise Exception(f"Failed to get async invoke status: {response.status_code} - {response.text}") diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index 56339ed2230..56ac2c00560 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -43,9 +43,7 @@ class TwelveLabsMarengoEmbeddingConfig: "input_type", ] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": # TwelveLabs doesn't have encoding_format, but we can map it to embeddingOption @@ -93,9 +91,7 @@ class TwelveLabsMarengoEmbeddingConfig: # Get input_type or default to "text" input_type = cast( TWELVELABS_EMBEDDING_INPUT_TYPES, - inference_params.get("inputType") - or inference_params.get("input_type") - or "text", + inference_params.get("inputType") or inference_params.get("input_type") or "text", ) # Validate that async-invoke is used for video/audio @@ -105,9 +101,7 @@ class TwelveLabsMarengoEmbeddingConfig: f"Use model format: 'bedrock/async_invoke/model_id'" ) - transformed_request: TwelveLabsMarengoEmbeddingRequest = { - "inputType": input_type - } + transformed_request: TwelveLabsMarengoEmbeddingRequest = {"inputType": input_type} if input_type == "text": transformed_request["inputText"] = input @@ -194,9 +188,7 @@ class TwelveLabsMarengoEmbeddingConfig: ), ) - def _transform_response( - self, response_list: List[dict], model: str - ) -> EmbeddingResponse: + def _transform_response(self, response_list: List[dict], model: str) -> EmbeddingResponse: """ Transform TwelveLabs response to OpenAI format. Handles the actual TwelveLabs response format: {"data": [{"embedding": [...]}]} @@ -253,9 +245,7 @@ class TwelveLabsMarengoEmbeddingConfig: return EmbeddingResponse(data=embeddings, model=model, usage=usage) - def _transform_async_invoke_response( - self, response: dict, model: str - ) -> EmbeddingResponse: + def _transform_async_invoke_response(self, response: dict, model: str) -> EmbeddingResponse: """ Transform async invoke response (invocation ARN) to OpenAI format. diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index b6aae2159c1..8c6282d627e 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -61,9 +61,7 @@ class BedrockFilesHandler(BaseAWSLLM): allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids, ) - def _get_configured_s3_bucket_name( - self, litellm_params: Mapping[str, object] - ) -> str: + def _get_configured_s3_bucket_name(self, litellm_params: Mapping[str, object]) -> str: from .transformation import get_configured_s3_bucket_name return get_configured_s3_bucket_name(litellm_params) @@ -100,15 +98,11 @@ class BedrockFilesHandler(BaseAWSLLM): bucket_name, object_key = self._parse_s3_uri( s3_uri=s3_uri, configured_bucket_name=configured_bucket_name, - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( - optional_params - ), + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(optional_params), ) # Get AWS credentials - aws_region_name = self._get_aws_region_name( - optional_params=optional_params, model="" - ) + aws_region_name = self._get_aws_region_name(optional_params=optional_params, model="") credentials: Credentials = self.get_credentials( aws_access_key_id=optional_params.get("aws_access_key_id"), aws_secret_access_key=optional_params.get("aws_secret_access_key"), @@ -136,9 +130,7 @@ class BedrockFilesHandler(BaseAWSLLM): response = s3_client.get_object(Bucket=bucket_name, Key=object_key) file_content = response["Body"].read() except Exception as e: - raise ValueError( - f"Failed to download file from S3: {s3_uri}. Error: {str(e)}" - ) + raise ValueError(f"Failed to download file from S3: {s3_uri}. Error: {str(e)}") # Create mock HTTP response mock_response = httpx.Response( @@ -158,9 +150,7 @@ class BedrockFilesHandler(BaseAWSLLM): optional_params: dict, timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - ) -> Union[ - HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] - ]: + ) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]: """ Download file content from S3 bucket for Bedrock files. Supports both sync and async operations. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 6cfaa88275d..d4865a1c87a 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -119,9 +119,7 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str: environment; never a request-supplied param, since the bucket is what `validate_managed_cloud_file_id` checks file ids against. """ - trusted_model_credentials = litellm_params.get( - "_litellm_internal_model_credentials" - ) + trusted_model_credentials = litellm_params.get("_litellm_internal_model_credentials") bucket_name: str | None = None if isinstance(trusted_model_credentials, MappingProxyType): snapshot: dict[str, object] = {} @@ -224,18 +222,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if _model.startswith("bedrock/"): _model = _model[8:] - safe_model = sanitize_cloud_object_component( - _model.replace(":", "-"), fallback="model" - ) + safe_model = sanitize_cloud_object_component(_model.replace(":", "-"), fallback="model") - object_name = ( - f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl" - ) + object_name = f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl" return object_name - def get_object_name( - self, extracted_file_data: ExtractedFileData, purpose: str - ) -> str: + def get_object_name(self, extracted_file_data: ExtractedFileData, purpose: str) -> str: """ Get the object name for the request """ @@ -246,14 +238,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if purpose == "batch": ## 1. If jsonl, check if there's a model name - file_content = self._get_content_from_openai_file( - extracted_file_data_content - ) + file_content = self._get_content_from_openai_file(extracted_file_data_content) # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] + openai_jsonl_content = [json.loads(line) for line in file_content.splitlines() if line.strip()] if len(openai_jsonl_content) > 0: return self._get_s3_object_name_from_batch_jsonl(openai_jsonl_content) @@ -277,21 +265,15 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Get the complete S3 URL for the file upload request """ - bucket_name = litellm_params.get("s3_bucket_name") or os.getenv( - "AWS_S3_BUCKET_NAME" - ) + bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") if not bucket_name: raise ValueError( "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var" ) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) - s3_region_name = litellm_params.get("s3_region_name") or optional_params.get( - "s3_region_name" - ) - aws_region_name = s3_region_name or self._get_aws_region_name( - optional_params, model - ) + s3_region_name = litellm_params.get("s3_region_name") or optional_params.get("s3_region_name") + aws_region_name = s3_region_name or self._get_aws_region_name(optional_params, model) file_data = data.get("file") purpose = data.get("purpose") @@ -307,15 +289,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # S3 endpoint URL format s3_endpoint_url = ( - optional_params.get("s3_endpoint_url") - or f"https://s3.{aws_region_name}.amazonaws.com" + optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" ).rstrip("/") return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAICreateFileRequestOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAICreateFileRequestOptionalParams]: return [] def map_openai_params( @@ -489,10 +468,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): without duplicating the type-shaping logic. """ if raw_input is None: - raise ValueError( - "Embedding batch record is missing required `input` field: " - f"model={model}" - ) + raise ValueError(f"Embedding batch record is missing required `input` field: model={model}") # Bedrock InvokeModel for Titan v2 takes exactly one string `inputText` # per call. Pre-tokenized inputs and multi-element string lists are @@ -564,26 +540,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): "embedding models in https://github.com/BerriAI/litellm/issues." ) - input_text = self._coerce_embedding_input_to_string( - openai_request_body.get("input"), model=_model - ) + input_text = self._coerce_embedding_input_to_string(openai_request_body.get("input"), model=_model) # Map OpenAI-style params (dimensions, encoding_format) onto the # Titan v2 schema (dimensions, embeddingTypes) via the embed config # so this stays in sync with the synchronous /v1/embeddings path. - non_default_params = { - k: v for k, v in openai_request_body.items() if k not in ("model", "input") - } + non_default_params = {k: v for k, v in openai_request_body.items() if k not in ("model", "input")} titan_config = AmazonTitanV2Config() inference_params = titan_config.map_openai_params( non_default_params=non_default_params, optional_params={}, ) - return dict( - titan_config._transform_request( - input=input_text, inference_params=inference_params - ) - ) + return dict(titan_config._transform_request(input=input_text, inference_params=inference_params)) def _map_openai_to_bedrock_params( self, @@ -602,11 +570,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): _model = openai_request_body.get("model", "") messages = openai_request_body.get("messages", []) - optional_params = { - k: v - for k, v in openai_request_body.items() - if k not in ["model", "messages"] - } + optional_params = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} # --- Anthropic: use existing AmazonAnthropicClaudeConfig --- if provider == LlmProviders.ANTHROPIC: @@ -707,18 +671,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # `_map_openai_to_bedrock_params`) so the chat helper keeps its # narrow contract and the embedding helper can evolve independently. if self._is_embedding_record(_openai_jsonl_content): - model_input = self._map_openai_embedding_to_bedrock_params( - openai_request_body=openai_body - ) + model_input = self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body) else: - model_input = self._map_openai_to_bedrock_params( - openai_request_body=openai_body, provider=provider - ) + model_input = self._map_openai_to_bedrock_params(openai_request_body=openai_body, provider=provider) # Create Bedrock batch record - record_id = _openai_jsonl_content.get( - "custom_id", f"CALL{str(idx).zfill(7)}" - ) + record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") bedrock_record = {"recordId": record_id, "modelInput": model_input} bedrock_jsonl_content.append(bedrock_record) @@ -750,19 +708,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): extracted_file_data=extracted_file_data, ): ## Transform JSONL content to Bedrock format - original_file_content = self._get_content_from_openai_file( - extracted_file_data_content - ) - openai_jsonl_content = [ - json.loads(line) - for line in original_file_content.splitlines() - if line.strip() - ] - bedrock_jsonl_content = ( - self._transform_openai_jsonl_content_to_bedrock_jsonl_content( - openai_jsonl_content - ) - ) + original_file_content = self._get_content_from_openai_file(extracted_file_data_content) + openai_jsonl_content = [json.loads(line) for line in original_file_content.splitlines() if line.strip()] + bedrock_jsonl_content = self._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content) file_content = "\n".join(json.dumps(item) for item in bedrock_jsonl_content) elif isinstance(extracted_file_data_content, bytes): file_content = extracted_file_data_content.decode("utf-8") @@ -784,9 +732,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # s3_region_name always wins for S3 operations (same priority as in # get_complete_file_url above). Overwrite aws_region_name unconditionally # so the SigV4 region matches the URL region, avoiding SignatureDoesNotMatch. - s3_region_name = litellm_params.get("s3_region_name") or optional_params.get( - "s3_region_name" - ) + s3_region_name = litellm_params.get("s3_region_name") or optional_params.get("s3_region_name") if s3_region_name: optional_params = {**optional_params, "aws_region_name": s3_region_name} @@ -827,9 +773,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") # Get AWS credentials using existing methods - aws_region_name = self._get_aws_region_name( - optional_params=optional_params, model="" - ) + aws_region_name = self._get_aws_region_name(optional_params=optional_params, model="") credentials = self.get_credentials( aws_access_key_id=optional_params.get("aws_access_key_id"), aws_secret_access_key=optional_params.get("aws_secret_access_key"), @@ -866,9 +810,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) # Get region name for non-LLM API calls (same as s3_v2.py) - signing_region = self.get_aws_region_name_for_non_llm_api_calls( - aws_region_name=aws_region_name - ) + signing_region = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=aws_region_name) SigV4Auth(credentials, "s3", signing_region).add_auth(aws_request) @@ -969,12 +911,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): object="file", ) - def get_error_class( - self, error_message: str, status_code: int, headers: Union[Dict, Headers] - ) -> BaseLLMException: - return BedrockError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[Dict, Headers]) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message, headers=headers) def transform_retrieve_file_request( self, @@ -1047,9 +985,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): scheme="s3://", configured_bucket_name=get_configured_s3_bucket_name(litellm_params), allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( - litellm_params - ), + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) # The shared file-content handler passes optional_params={}, so AWS @@ -1061,18 +997,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): merged_params.update(optional_params) request_params = _BedrockS3RequestParams.model_validate(merged_params) - region_preference = ( - request_params.s3_region_name or request_params.aws_region_name - ) + region_preference = request_params.s3_region_name or request_params.aws_region_name region_params: dict[str, str | None] = {"aws_region_name": region_preference} - aws_region_name = self._get_aws_region_name( - optional_params=region_params, model="" - ) + aws_region_name = self._get_aws_region_name(optional_params=region_params, model="") - s3_endpoint_url = ( - request_params.s3_endpoint_url - or f"https://s3.{aws_region_name}.amazonaws.com" - ).rstrip("/") + s3_endpoint_url = (request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.amazonaws.com").rstrip("/") url = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request( @@ -1154,32 +1083,18 @@ class BedrockJsonlFilesTransformation: file_content = self._get_content_from_openai_file(openai_file_content) # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] - bedrock_jsonl_content = ( - self._transform_openai_jsonl_content_to_bedrock_jsonl_content( - openai_jsonl_content - ) - ) - bedrock_jsonl_string = "\n".join( - json.dumps(item) for item in bedrock_jsonl_content - ) - object_name = self._get_s3_object_name( - openai_jsonl_content=openai_jsonl_content - ) + openai_jsonl_content = [json.loads(line) for line in file_content.splitlines() if line.strip()] + bedrock_jsonl_content = self._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content) + bedrock_jsonl_string = "\n".join(json.dumps(item) for item in bedrock_jsonl_content) + object_name = self._get_s3_object_name(openai_jsonl_content=openai_jsonl_content) return bedrock_jsonl_string, object_name - def _transform_openai_jsonl_content_to_bedrock_jsonl_content( - self, openai_jsonl_content: List[Dict[str, Any]] - ): + def _transform_openai_jsonl_content_to_bedrock_jsonl_content(self, openai_jsonl_content: List[Dict[str, Any]]): """ Delegate to the main BedrockFilesConfig transformation method """ config = BedrockFilesConfig() - return config._transform_openai_jsonl_content_to_bedrock_jsonl_content( - openai_jsonl_content - ) + return config._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content) def _get_s3_object_name( self, @@ -1194,12 +1109,8 @@ class BedrockJsonlFilesTransformation: # Remove bedrock/ prefix if present if _model.startswith("bedrock/"): _model = _model[8:] - safe_model = sanitize_cloud_object_component( - _model.replace(":", "-"), fallback="model" - ) - object_name = ( - f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl" - ) + safe_model = sanitize_cloud_object_component(_model.replace(":", "-"), fallback="model") + object_name = f"{BEDROCK_MANAGED_S3_BATCH_PREFIX}{safe_model}-{uuid.uuid4()}.jsonl" return object_name def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index 836a3c606ee..1008924ab0e 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -315,13 +315,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): _size = op.pop("size", None) width = op.pop("width", None) height = op.pop("height", None) - if ( - width is None - and height is None - and _size is not None - and isinstance(_size, str) - and "x" in _size - ): + if width is None and height is None and _size is not None and isinstance(_size, str) and "x" in _size: w, h = _size.split("x", 1) try: width, height = int(w), int(h) @@ -356,8 +350,7 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): "OUTPAINTING", ): raise ValueError( - f"Amazon Nova Canvas {task_type} requires a text prompt. " - "Pass a non-empty `prompt` in your request." + f"Amazon Nova Canvas {task_type} requires a text prompt. Pass a non-empty `prompt` in your request." ) text = prompt if prompt is not None and prompt != "" else " " negative_text = op.pop("negativeText", None) @@ -455,9 +448,9 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): model_info = get_model_info(model, custom_llm_provider="bedrock") cost_per_image = model_info.get("output_cost_per_image", 0) if cost_per_image is not None and model_response.data: - model_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(cost_per_image) * len(model_response.data) + model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float( + cost_per_image + ) * len(model_response.data) except Exception: pass diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 04fa5f803bb..01a40c0e475 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -58,9 +58,7 @@ class BedrockImageEdit(BaseAWSLLM): def get_config_class(cls, model: str | None): if BedrockStabilityImageEditConfig._is_stability_edit_model(model): return BedrockStabilityImageEditConfig - if BedrockAmazonNovaCanvasImageEditConfig._is_nova_canvas_image_edit_model( - model - ): + if BedrockAmazonNovaCanvasImageEditConfig._is_nova_canvas_image_edit_model(model): return BedrockAmazonNovaCanvasImageEditConfig raise ValueError( f"Unsupported Bedrock image-edit model: {model!r}. " @@ -102,11 +100,7 @@ class BedrockImageEdit(BaseAWSLLM): logging_obj=logging_obj, prompt=prompt, model_response=model_response, - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), ) if client is None or not isinstance(client, HTTPHandler): @@ -204,9 +198,7 @@ class BedrockImageEdit(BaseAWSLLM): Returns: BedrockImageEditPreparedRequest: The prepared request object """ - boto3_credentials_info = self._get_boto_credentials_from_optional_params( - optional_params, model - ) + boto3_credentials_info = self._get_boto_credentials_from_optional_params(optional_params, model) # Use the existing ARN-aware provider detection method bedrock_provider = self.get_bedrock_invoke_provider(model) diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index e76b2885a88..0b45aba219f 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -122,9 +122,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): if k in param_mapping: # Map param if mapping exists and value is valid if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO: - mapped_params[param_mapping[k]] = ( - OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] - ) # type: ignore + mapped_params[param_mapping[k]] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] # type: ignore # Don't copy "size" itself to final dict elif k == "n": # Store for logic but do not add to outgoing params @@ -211,9 +209,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): if isinstance(value, list) and len(value) > 0: file_value = value[0] - if hasattr(file_value, "read") and callable( - getattr(file_value, "read", None) - ): + if hasattr(file_value, "read") and callable(getattr(file_value, "read", None)): file_bytes = file_value.read() # type: ignore elif isinstance(file_value, bytes): file_bytes = file_value @@ -338,9 +334,9 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): model_info = get_model_info(model, custom_llm_provider="bedrock") cost_per_image = model_info.get("output_cost_per_image", 0) if cost_per_image is not None: - model_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(cost_per_image) + model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float( + cost_per_image + ) return model_response diff --git a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py index 9e06a8e747d..626baf707a5 100644 --- a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py @@ -60,9 +60,7 @@ class AmazonNovaCanvasConfig: return False @classmethod - def transform_request_body( - cls, text: str, optional_params: dict - ) -> AmazonNovaCanvasRequestBase: + def transform_request_body(cls, text: str, optional_params: dict) -> AmazonNovaCanvasRequestBase: """ Transform the request body for Amazon Nova Canvas model """ @@ -75,9 +73,7 @@ class AmazonNovaCanvasConfig: image_generation_config = {**image_generation_config, **optional_params} if task_type == "TEXT_IMAGE": - text_to_image_params: Dict[str, Any] = image_generation_config.pop( - "textToImageParams", {} - ) + text_to_image_params: Dict[str, Any] = image_generation_config.pop("textToImageParams", {}) text_to_image_params = {"text": text, **text_to_image_params} try: text_to_image_params_typed = AmazonNovaCanvasTextToImageParams( @@ -89,9 +85,7 @@ class AmazonNovaCanvasConfig: ) try: - image_generation_config_typed = AmazonNovaCanvasImageGenerationConfig( - **image_generation_config - ) + image_generation_config_typed = AmazonNovaCanvasImageGenerationConfig(**image_generation_config) except Exception as e: raise ValueError( f"Error transforming image generation config: {e}. Got params: {image_generation_config}, Expected params: {AmazonNovaCanvasImageGenerationConfig.__annotations__}" @@ -103,18 +97,16 @@ class AmazonNovaCanvasConfig: imageGenerationConfig=image_generation_config_typed, ) if task_type == "COLOR_GUIDED_GENERATION": - color_guided_generation_params: Dict[str, Any] = ( - image_generation_config.pop("colorGuidedGenerationParams", {}) + color_guided_generation_params: Dict[str, Any] = image_generation_config.pop( + "colorGuidedGenerationParams", {} ) color_guided_generation_params = { "text": text, **color_guided_generation_params, } try: - color_guided_generation_params_typed = ( - AmazonNovaCanvasColorGuidedGenerationParams( - **color_guided_generation_params # type: ignore - ) + color_guided_generation_params_typed = AmazonNovaCanvasColorGuidedGenerationParams( + **color_guided_generation_params # type: ignore ) except Exception as e: raise ValueError( @@ -122,9 +114,7 @@ class AmazonNovaCanvasConfig: ) try: - image_generation_config_typed = AmazonNovaCanvasImageGenerationConfig( - **image_generation_config - ) + image_generation_config_typed = AmazonNovaCanvasImageGenerationConfig(**image_generation_config) except Exception as e: raise ValueError( f"Error transforming image generation config: {e}. Got params: {image_generation_config}, Expected params: {AmazonNovaCanvasImageGenerationConfig.__annotations__}" @@ -136,9 +126,7 @@ class AmazonNovaCanvasConfig: imageGenerationConfig=image_generation_config_typed, ) if task_type == "INPAINTING": - inpainting_params: Dict[str, Any] = image_generation_config.pop( - "inpaintingParams", {} - ) + inpainting_params: Dict[str, Any] = image_generation_config.pop("inpaintingParams", {}) inpainting_params = {"text": text, **inpainting_params} try: inpainting_params_typed = AmazonNovaCanvasInpaintingParams( @@ -150,9 +138,7 @@ class AmazonNovaCanvasConfig: ) try: - image_generation_config_typed = AmazonNovaCanvasImageGenerationConfig( - **image_generation_config - ) + image_generation_config_typed = AmazonNovaCanvasImageGenerationConfig(**image_generation_config) except Exception as e: raise ValueError( f"Error transforming image generation config: {e}. Got params: {image_generation_config}, Expected params: {AmazonNovaCanvasImageGenerationConfig.__annotations__}" diff --git a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py index 1d88aaf35f7..0e8214fd81f 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py @@ -100,9 +100,7 @@ class AmazonStabilityConfig: optional_params: dict, ) -> dict: inference_params = copy.deepcopy(optional_params) - inference_params.pop( - "user", None - ) # make sure user is not passed in for bedrock call + inference_params.pop("user", None) # make sure user is not passed in for bedrock call prompt = text.replace(os.linesep, " ") ## LOAD CONFIG diff --git a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py index 8aff24fe9a7..a5449679941 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py @@ -67,9 +67,7 @@ class AmazonStability3Config: return False @classmethod - def transform_request_body( - cls, text: str, optional_params: dict - ) -> AmazonStability3TextToImageRequest: + def transform_request_body(cls, text: str, optional_params: dict) -> AmazonStability3TextToImageRequest: """ Transform the request body for the Stability 3 models """ diff --git a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py index 65411cabdcf..5a975b6ab11 100644 --- a/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_titan_transformation.py @@ -90,9 +90,7 @@ class AmazonTitanImageGenerationConfig: image_generation_config["height"] = int(height) elif k == "n" and v is not None: image_generation_config["numberOfImages"] = v - elif ( - k == "quality" and v is not None - ): # 'auto', 'hd', 'standard', 'high', 'medium', 'low' + elif k == "quality" and v is not None: # 'auto', 'hd', 'standard', 'high', 'medium', 'low' if v in ("hd", "premium", "high"): image_generation_config["quality"] = "premium" elif v in ("standard", "medium", "low"): @@ -116,9 +114,7 @@ class AmazonTitanImageGenerationConfig: if negative_text: text_to_image_params["negativeText"] = negative_text task_type = optional_params.pop("taskType", "TEXT_IMAGE") - user_specified_image_generation_config = optional_params.pop( - "imageGenerationConfig", {} - ) + user_specified_image_generation_config = optional_params.pop("imageGenerationConfig", {}) image_generation_config = { **image_generation_config, **user_specified_image_generation_config, @@ -126,9 +122,7 @@ class AmazonTitanImageGenerationConfig: return AmazonTitanImageGenerationRequestBody( taskType=task_type, textToImageParams=AmazonTitanTextToImageParams(**text_to_image_params), # type: ignore - imageGenerationConfig=AmazonNovaCanvasImageGenerationConfig( - **image_generation_config - ), + imageGenerationConfig=AmazonNovaCanvasImageGenerationConfig(**image_generation_config), ) @classmethod diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index 0c594d0b142..03e40565d95 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -105,11 +105,7 @@ class BedrockImageGeneration(BaseAWSLLM): logging_obj=logging_obj, prompt=prompt, model_response=model_response, - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), ) if client is None or not isinstance(client, HTTPHandler): @@ -224,9 +220,7 @@ class BedrockImageGeneration(BaseAWSLLM): prepped (httpx.Request): The prepared request object body (bytes): The request body """ - boto3_credentials_info = self._get_boto_credentials_from_optional_params( - optional_params, model - ) + boto3_credentials_info = self._get_boto_credentials_from_optional_params(optional_params, model) # Use the existing ARN-aware provider detection method bedrock_provider = self.get_bedrock_invoke_provider(model) @@ -300,9 +294,7 @@ class BedrockImageGeneration(BaseAWSLLM): dict: The request body to use for the Bedrock Image Generation API """ config_class = self.get_config_class(model=model) - request_body = config_class.transform_request_body( - text=prompt, optional_params=optional_params - ) + request_body = config_class.transform_request_body(text=prompt, optional_params=optional_params) return dict(request_body) def _transform_response_dict_to_openai_response( diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 3ce426edba7..f5309d521a9 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -77,9 +77,7 @@ class AmazonAnthropicClaudeMessagesConfig( DEFAULT_BEDROCK_ANTHROPIC_API_VERSION = "bedrock-2023-05-31" - BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset( - BedrockInvokeAnthropicMessagesRequest.__annotations__.keys() - ) + BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()) def __init__(self, **kwargs): BaseAnthropicMessagesConfig.__init__(self, **kwargs) @@ -95,9 +93,7 @@ class AmazonAnthropicClaudeMessagesConfig( return [{"type": "text", "text": value}] return [value] - def _normalize_system_role_messages_for_bedrock( - self, anthropic_messages_request: dict - ) -> None: + def _normalize_system_role_messages_for_bedrock(self, anthropic_messages_request: dict) -> None: """Bedrock Invoke rejects ``role: "system"`` entries inside ``messages`` on some Claude aliases; Anthropic Messages carries that content in the top-level ``system`` field. Move any such entries into ``system`` before @@ -105,16 +101,12 @@ class AmazonAnthropicClaudeMessagesConfig( messages = anthropic_messages_request.get("messages") if not isinstance(messages, list): return - system_role_messages = [ - m for m in messages if isinstance(m, dict) and m.get("role") == "system" - ] + system_role_messages = [m for m in messages if isinstance(m, dict) and m.get("role") == "system"] if not system_role_messages: return anthropic_messages_request["messages"] = [ - m - for m in messages - if not (isinstance(m, dict) and m.get("role") == "system") + m for m in messages if not (isinstance(m, dict) and m.get("role") == "system") ] system_content = [ block @@ -184,9 +176,7 @@ class AmazonAnthropicClaudeMessagesConfig( stream=stream, ) - def _remove_ttl_from_cache_control( - self, anthropic_messages_request: Dict, model: Optional[str] = None - ) -> None: + def _remove_ttl_from_cache_control(self, anthropic_messages_request: Dict, model: Optional[str] = None) -> None: """ Remove unsupported fields from cache_control for Bedrock. @@ -301,18 +291,13 @@ class AmazonAnthropicClaudeMessagesConfig( edits = cm.get("edits") if not isinstance(edits, list): return False - needs_thinking = any( - isinstance(e, dict) and e.get("type") == "clear_thinking_20251015" - for e in edits - ) + needs_thinking = any(isinstance(e, dict) and e.get("type") == "clear_thinking_20251015" for e in edits) if not needs_thinking: return False if not self._supports_extended_thinking_on_bedrock(model): return False - is_adaptive_thinking_model = AnthropicModelInfo._is_adaptive_thinking_model( - model - ) + is_adaptive_thinking_model = AnthropicModelInfo._is_adaptive_thinking_model(model) thinking = anthropic_messages_request.get("thinking") if isinstance(thinking, dict): @@ -322,12 +307,8 @@ class AmazonAnthropicClaudeMessagesConfig( if t == "enabled" and not is_adaptive_thinking_model: return False if t == "enabled": - budget_tokens = self._resolve_clear_thinking_budget_tokens( - thinking.get("budget_tokens") - ) - self._inject_adaptive_thinking_for_clear_thinking( - anthropic_messages_request, budget_tokens, model - ) + budget_tokens = self._resolve_clear_thinking_budget_tokens(thinking.get("budget_tokens")) + self._inject_adaptive_thinking_for_clear_thinking(anthropic_messages_request, budget_tokens, model) return True verbose_logger.debug( "Bedrock clear_thinking_20251015: replacing thinking=%s with minimal thinking config", @@ -346,9 +327,7 @@ class AmazonAnthropicClaudeMessagesConfig( return False if is_adaptive_thinking_model: - self._inject_adaptive_thinking_for_clear_thinking( - anthropic_messages_request, budget, model - ) + self._inject_adaptive_thinking_for_clear_thinking(anthropic_messages_request, budget, model) return True anthropic_messages_request["thinking"] = { @@ -390,9 +369,7 @@ class AmazonAnthropicClaudeMessagesConfig( output_config = anthropic_messages_request.get("output_config") if not isinstance(output_config, dict): output_config = {} - output_config.setdefault( - "effort", self._effort_from_thinking_budget(budget_tokens) - ) + output_config.setdefault("effort", self._effort_from_thinking_budget(budget_tokens)) anthropic_messages_request["output_config"] = output_config anthropic_messages_request["thinking"] = {"type": "adaptive"} verbose_logger.debug( @@ -507,9 +484,7 @@ class AmazonAnthropicClaudeMessagesConfig( input_examples_used: Whether input examples are used beta_set: The set of beta headers to modify in-place """ - if tool_search_used and not ( - programmatic_tool_calling_used or input_examples_used - ): + if tool_search_used and not (programmatic_tool_calling_used or input_examples_used): beta_set.discard(ANTHROPIC_TOOL_SEARCH_BETA_HEADER) if self._supports_tool_search_on_bedrock(model): beta_set.add("tool-search-tool-2025-10-19") @@ -541,11 +516,7 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request.pop("context_management", None) return - compact_edits = [ - e - for e in edits - if isinstance(e, dict) and e.get("type") == "compact_20260112" - ] + compact_edits = [e for e in edits if isinstance(e, dict) and e.get("type") == "compact_20260112"] if compact_edits: beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value) anthropic_messages_request["context_management"] = { @@ -568,9 +539,7 @@ class AmazonAnthropicClaudeMessagesConfig( tools = anthropic_messages_optional_request_params.get("tools") messages_typed = cast(List[AllMessageValues], messages) tool_search_used = anthropic_model_info.is_tool_search_used(tools) - programmatic_tool_calling_used = ( - anthropic_model_info.is_programmatic_tool_calling_used(tools) - ) + programmatic_tool_calling_used = anthropic_model_info.is_programmatic_tool_calling_used(tools) input_examples_used = anthropic_model_info.is_input_examples_used(tools) user_beta_set = set(get_anthropic_beta_from_headers(headers)) @@ -614,9 +583,7 @@ class AmazonAnthropicClaudeMessagesConfig( ) dropped_user_betas = sorted( - b - for b in user_beta_set - if not filter_and_transform_beta_headers([b], provider="bedrock") + b for b in user_beta_set if not filter_and_transform_beta_headers([b], provider="bedrock") ) if dropped_user_betas: verbose_logger.warning( @@ -642,9 +609,7 @@ class AmazonAnthropicClaudeMessagesConfig( return {k: v for k, v in anthropic_messages_request.items() if k in allowed} @staticmethod - def _clamp_adaptive_reasoning_effort_for_bedrock( - model: str, optional_params: Dict - ) -> None: + def _clamp_adaptive_reasoning_effort_for_bedrock(model: str, optional_params: Dict) -> None: """Lower ``reasoning_effort`` to the Bedrock effort ceiling before validation. The shared ``/v1/messages`` effort gate rejects tiers a model does not @@ -690,9 +655,7 @@ class AmazonAnthropicClaudeMessagesConfig( # 1. anthropic_version is required for all claude models if "anthropic_version" not in anthropic_messages_request: - anthropic_messages_request["anthropic_version"] = ( - self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION - ) + anthropic_messages_request["anthropic_version"] = self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION # 2. `stream` is not allowed in request body for bedrock invoke if "stream" in anthropic_messages_request: @@ -702,17 +665,13 @@ class AmazonAnthropicClaudeMessagesConfig( if "model" in anthropic_messages_request: anthropic_messages_request.pop("model", None) - injected_thinking_for_clear_thinking = ( - self._ensure_thinking_for_clear_thinking_context_management( - anthropic_messages_request=anthropic_messages_request, - model=model, - ) + injected_thinking_for_clear_thinking = self._ensure_thinking_for_clear_thinking_context_management( + anthropic_messages_request=anthropic_messages_request, + model=model, ) # 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models) - self._remove_ttl_from_cache_control( - anthropic_messages_request=anthropic_messages_request, model=model - ) + self._remove_ttl_from_cache_control(anthropic_messages_request=anthropic_messages_request, model=model) # 5. Convert structured-output params to inline schema. # Bedrock Invoke doesn't support top-level `output_format`; its @@ -723,9 +682,7 @@ class AmazonAnthropicClaudeMessagesConfig( if isinstance(existing_output_config, dict): anthropic_messages_request["output_config"] = dict(existing_output_config) output_format = anthropic_messages_request.pop("output_format", None) - output_config_format = pop_bedrock_invoke_output_config_format( - anthropic_messages_request - ) + output_config_format = pop_bedrock_invoke_output_config_format(anthropic_messages_request) if output_format: convert_bedrock_invoke_output_format_to_inline_schema( output_format=output_format, @@ -799,9 +756,7 @@ class AmazonAnthropicClaudeMessagesConfig( # Catches Anthropic-only extensions (output_config, speed, mcp_servers, ...) # and any future additions Claude Code may start sending. ``context_management`` # has already been pre-filtered to its Bedrock-supported subset above. - anthropic_messages_request = self._strip_unsupported_bedrock_invoke_fields( - anthropic_messages_request - ) + anthropic_messages_request = self._strip_unsupported_bedrock_invoke_fields(anthropic_messages_request) return anthropic_messages_request @@ -827,9 +782,7 @@ class AmazonAnthropicClaudeMessagesConfig( async def bedrock_sse_wrapper( self, - completion_stream: AsyncIterator[ - Union[bytes, GenericStreamingChunk, ModelResponseStream, dict] - ], + completion_stream: AsyncIterator[Union[bytes, GenericStreamingChunk, ModelResponseStream, dict]], litellm_logging_obj: LiteLLMLoggingObj, request_body: dict, ): @@ -882,9 +835,7 @@ class AmazonAnthropicClaudeMessagesConfig( @staticmethod async def _promote_message_stop_usage( - completion_stream: AsyncIterator[ - Union[bytes, GenericStreamingChunk, ModelResponseStream, dict] - ], + completion_stream: AsyncIterator[Union[bytes, GenericStreamingChunk, ModelResponseStream, dict]], ) -> AsyncIterator[Union[bytes, GenericStreamingChunk, ModelResponseStream, dict]]: """ Promote cache usage fields onto message_delta from message_stop (and, @@ -930,9 +881,7 @@ class AmazonAnthropicClaudeMessagesConfig( raw_input = stop_usage.get("input_tokens") if raw_input is not None: - delta_usage["input_tokens"] = ( - raw_input if isinstance(raw_input, int) else 0 - ) + delta_usage["input_tokens"] = raw_input if isinstance(raw_input, int) else 0 AmazonAnthropicClaudeMessagesConfig._merge_message_start_cache_into_delta_usage( delta_usage, start_usage_snapshot @@ -973,9 +922,7 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): super().__init__(model=model) self.DEFAULT_CHUNK_SIZE = 1024 - def _chunk_parser( - self, chunk_data: dict - ) -> Union[GChunk, ModelResponseStream, dict]: + def _chunk_parser(self, chunk_data: dict) -> Union[GChunk, ModelResponseStream, dict]: """ Parse the chunk data into anthropic /messages format @@ -983,18 +930,12 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): the Anthropic `/v1/messages` specification so callers receive a consistent response shape when streaming. """ - amazon_bedrock_invocation_metrics = chunk_data.pop( - "amazon-bedrock-invocationMetrics", {} - ) + amazon_bedrock_invocation_metrics = chunk_data.pop("amazon-bedrock-invocationMetrics", {}) if amazon_bedrock_invocation_metrics: anthropic_usage = {} if "inputTokenCount" in amazon_bedrock_invocation_metrics: - anthropic_usage["input_tokens"] = amazon_bedrock_invocation_metrics[ - "inputTokenCount" - ] + anthropic_usage["input_tokens"] = amazon_bedrock_invocation_metrics["inputTokenCount"] if "outputTokenCount" in amazon_bedrock_invocation_metrics: - anthropic_usage["output_tokens"] = amazon_bedrock_invocation_metrics[ - "outputTokenCount" - ] + anthropic_usage["output_tokens"] = amazon_bedrock_invocation_metrics["outputTokenCount"] chunk_data["usage"] = anthropic_usage return chunk_data diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index 94e7f90b719..a8a7b7ed1d5 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -42,9 +42,7 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): region = self._get_aws_region_name(optional_params=optional_params, model=model) return build_mantle_messages_url( api_base=api_base, - aws_bedrock_runtime_endpoint=optional_params.get( - "aws_bedrock_runtime_endpoint" - ), + aws_bedrock_runtime_endpoint=optional_params.get("aws_bedrock_runtime_endpoint"), region=region, ) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 0522bb249e1..137f1e333eb 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -269,9 +269,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): payload_dict = _json.loads(payload_bytes) texts = [ (group_key, container[key]) - for group_key, container, key in _collect_stream_delta_text_holders( - payload_dict.get("delta") - ) + for group_key, container, key in _collect_stream_delta_text_holders(payload_dict.get("delta")) ] except Exception as e: verbose_proxy_logger.debug( @@ -304,9 +302,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): "output": { "message": { "role": "assistant", - "content": [ - {"text": "".join(group_texts[gk])} for gk in active_groups - ], + "content": [{"text": "".join(group_texts[gk])} for gk in active_groups], } }, "stopReason": "end_turn", @@ -328,9 +324,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): try: processed_blocks = processed["output"]["message"]["content"] # type: ignore[index] - de_anonymized_texts = [ - processed_blocks[i]["text"] for i in range(len(active_groups)) - ] + de_anonymized_texts = [processed_blocks[i]["text"] for i in range(len(active_groups))] except (KeyError, IndexError, TypeError): return body_bytes @@ -362,9 +356,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): headers_bytes = frame_raw[12 : 12 + orig_hdrs_len] try: - payload_dict = _json.loads( - frame_raw[12 + orig_hdrs_len : orig_total - 4] - ) + payload_dict = _json.loads(frame_raw[12 + orig_hdrs_len : orig_total - 4]) for local_idx, (_, container, key) in enumerate( _collect_stream_delta_text_holders(payload_dict.get("delta")) ): @@ -384,9 +376,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): msg_crc_val = esm_crc32(part_for_msg_crc, prelude_crc_val) & 0xFFFFFFFF msg_crc_b = struct.pack("!I", msg_crc_val) - result_parts.append( - prelude + prelude_crc_b + headers_bytes + new_payload + msg_crc_b - ) + result_parts.append(prelude + prelude_crc_b + headers_bytes + new_payload + msg_crc_b) result_parts.append(trailing_bytes) return b"".join(result_parts) @@ -458,13 +448,9 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): return response output_message = ( - response.get("output", {}).get("message", {}) - if isinstance(response.get("output"), dict) - else {} - ) - content_blocks = ( - output_message.get("content") if isinstance(output_message, dict) else None + response.get("output", {}).get("message", {}) if isinstance(response.get("output"), dict) else {} ) + content_blocks = output_message.get("content") if isinstance(output_message, dict) else None if not isinstance(content_blocks, list): return response @@ -475,13 +461,8 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): return response effective_request_data = request_data or {} - if ( - "litellm_metadata" not in effective_request_data - and user_api_key_dict is not None - ): - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + if "litellm_metadata" not in effective_request_data and user_api_key_dict is not None: + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: effective_request_data = { **effective_request_data, diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 846af65c0f9..cc8840526f0 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -16,9 +16,7 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes -class BedrockPassthroughConfig( - BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig -): +class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in endpoint @@ -43,9 +41,7 @@ class BedrockPassthroughConfig( # Create a temporary endpoint with the model_id to check if encoding is needed temp_endpoint = f"/model/{model_id}/converse" - encoded_temp_endpoint = CommonUtils.encode_bedrock_runtime_modelid_arn( - temp_endpoint - ) + encoded_temp_endpoint = CommonUtils.encode_bedrock_runtime_modelid_arn(temp_endpoint) # Extract the encoded model_id from the temporary endpoint encoded_model_id_match = re.search(r"/model/([^/]+)/", encoded_temp_endpoint) @@ -73,9 +69,7 @@ class BedrockPassthroughConfig( model_id=model_id, ) - aws_bedrock_runtime_endpoint = optional_params.get( - "aws_bedrock_runtime_endpoint" - ) + aws_bedrock_runtime_endpoint = optional_params.get("aws_bedrock_runtime_endpoint") endpoint_url, _ = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, @@ -202,9 +196,7 @@ class BedrockPassthroughConfig( if "invoke" in endpoint: invoke_provider = AmazonInvokeConfig.get_bedrock_invoke_provider(model) if invoke_provider is None: - raise ValueError( - f"Invalid invoke provider: {invoke_provider}, for model: {model}" - ) + raise ValueError(f"Invalid invoke provider: {invoke_provider}, for model: {model}") obj = get_bedrock_event_stream_decoder( invoke_provider=invoke_provider, model=model, @@ -225,9 +217,9 @@ class BedrockPassthroughConfig( message = json.loads(chunk) translated_chunk = obj._chunk_parser(chunk_data=message) - if isinstance( - translated_chunk, dict - ) and generic_chunk_has_all_required_fields(cast(dict, translated_chunk)): + if isinstance(translated_chunk, dict) and generic_chunk_has_all_required_fields( + cast(dict, translated_chunk) + ): chunk_obj = convert_generic_chunk_to_model_response_stream( cast(GenericStreamingChunk, translated_chunk) ) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 5d22f4b3cd9..6db2571090a 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -62,9 +62,7 @@ class BedrockRealtime(BaseAWSLLM): EnvironmentCredentialsResolver, ) except ImportError: - raise ImportError( - "Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime" - ) + raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime") # Get AWS region if aws_region_name is None: @@ -81,9 +79,7 @@ class BedrockRealtime(BaseAWSLLM): else: endpoint_uri = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" - verbose_proxy_logger.debug( - f"Bedrock Realtime: Connecting to {endpoint_uri} with model {model}" - ) + verbose_proxy_logger.debug(f"Bedrock Realtime: Connecting to {endpoint_uri} with model {model}") # Initialize Bedrock client with aws_sdk_bedrock_runtime config = Config( @@ -97,15 +93,11 @@ class BedrockRealtime(BaseAWSLLM): try: # Initialize the bidirectional stream - bedrock_stream = ( - await bedrock_client.invoke_model_with_bidirectional_stream( - InvokeModelWithBidirectionalStreamOperationInput(model_id=model) - ) + bedrock_stream = await bedrock_client.invoke_model_with_bidirectional_stream( + InvokeModelWithBidirectionalStreamOperationInput(model_id=model) ) - verbose_proxy_logger.debug( - "Bedrock Realtime: Bidirectional stream established" - ) + verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") # Track state for transformation session_state = { @@ -148,13 +140,9 @@ class BedrockRealtime(BaseAWSLLM): ) except Exception as e: - verbose_proxy_logger.exception( - f"Error in BedrockRealtime.async_realtime: {e}" - ) + verbose_proxy_logger.exception(f"Error in BedrockRealtime.async_realtime: {e}") try: - await websocket.close( - code=1011, reason=_redact_string(f"Internal error: {str(e)}") - ) + await websocket.close(code=1011, reason=_redact_string(f"Internal error: {str(e)}")) except Exception: pass raise @@ -177,35 +165,25 @@ class BedrockRealtime(BaseAWSLLM): while True: # Receive message from client message = await client_ws.receive_text() - verbose_proxy_logger.debug( - f"Bedrock Realtime: Received from client: {message[:200]}" - ) + verbose_proxy_logger.debug(f"Bedrock Realtime: Received from client: {message[:200]}") # Transform OpenAI format to Bedrock format transformed_messages = transformation_config.transform_realtime_request( message=message, model=model, - session_configuration_request=session_state.get( - "session_configuration_request" - ), + session_configuration_request=session_state.get("session_configuration_request"), ) # Send transformed messages to Bedrock for bedrock_message in transformed_messages: event = InvokeModelWithBidirectionalStreamInputChunk( - value=BidirectionalInputPayloadPart( - bytes_=bedrock_message.encode("utf-8") - ) + value=BidirectionalInputPayloadPart(bytes_=bedrock_message.encode("utf-8")) ) await bedrock_stream.input_stream.send(event) - verbose_proxy_logger.debug( - f"Bedrock Realtime: Sent to Bedrock: {bedrock_message[:200]}" - ) + verbose_proxy_logger.debug(f"Bedrock Realtime: Sent to Bedrock: {bedrock_message[:200]}") except Exception as e: - verbose_proxy_logger.debug( - f"Client to Bedrock forwarding ended: {e}", exc_info=True - ) + verbose_proxy_logger.debug(f"Client to Bedrock forwarding ended: {e}", exc_info=True) # Close the Bedrock stream input try: await bedrock_stream.input_stream.close() @@ -230,29 +208,19 @@ class BedrockRealtime(BaseAWSLLM): if result.value and result.value.bytes_: bedrock_response = result.value.bytes_.decode("utf-8") - verbose_proxy_logger.debug( - f"Bedrock Realtime: Received from Bedrock: {bedrock_response[:200]}" - ) + verbose_proxy_logger.debug(f"Bedrock Realtime: Received from Bedrock: {bedrock_response[:200]}") # Transform Bedrock format to OpenAI format from litellm.types.realtime import RealtimeResponseTransformInput realtime_response_transform_input: RealtimeResponseTransformInput = { - "current_output_item_id": session_state.get( - "current_output_item_id" - ), + "current_output_item_id": session_state.get("current_output_item_id"), "current_response_id": session_state.get("current_response_id"), - "current_conversation_id": session_state.get( - "current_conversation_id" - ), - "current_delta_chunks": session_state.get( - "current_delta_chunks" - ), + "current_conversation_id": session_state.get("current_conversation_id"), + "current_delta_chunks": session_state.get("current_delta_chunks"), "current_item_chunks": session_state.get("current_item_chunks"), "current_delta_type": session_state.get("current_delta_type"), - "session_configuration_request": session_state.get( - "session_configuration_request" - ), + "session_configuration_request": session_state.get("session_configuration_request"), } transformed_response = transformation_config.transform_realtime_response( @@ -265,27 +233,13 @@ class BedrockRealtime(BaseAWSLLM): # Update session state session_state.update( { - "current_output_item_id": transformed_response.get( - "current_output_item_id" - ), - "current_response_id": transformed_response.get( - "current_response_id" - ), - "current_conversation_id": transformed_response.get( - "current_conversation_id" - ), - "current_delta_chunks": transformed_response.get( - "current_delta_chunks" - ), - "current_item_chunks": transformed_response.get( - "current_item_chunks" - ), - "current_delta_type": transformed_response.get( - "current_delta_type" - ), - "session_configuration_request": transformed_response.get( - "session_configuration_request" - ), + "current_output_item_id": transformed_response.get("current_output_item_id"), + "current_response_id": transformed_response.get("current_response_id"), + "current_conversation_id": transformed_response.get("current_conversation_id"), + "current_delta_chunks": transformed_response.get("current_delta_chunks"), + "current_item_chunks": transformed_response.get("current_item_chunks"), + "current_delta_type": transformed_response.get("current_delta_type"), + "session_configuration_request": transformed_response.get("session_configuration_request"), } ) @@ -294,14 +248,10 @@ class BedrockRealtime(BaseAWSLLM): for openai_message in openai_messages: message_json = json.dumps(openai_message) await client_ws.send_text(message_json) - verbose_proxy_logger.debug( - f"Bedrock Realtime: Sent to client: {message_json[:200]}" - ) + verbose_proxy_logger.debug(f"Bedrock Realtime: Sent to client: {message_json[:200]}") except Exception as e: - verbose_proxy_logger.debug( - f"Bedrock to client forwarding ended: {e}", exc_info=True - ) + verbose_proxy_logger.debug(f"Bedrock to client forwarding ended: {e}", exc_info=True) # Close the client WebSocket try: await client_ws.close() diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 9124a8c21b4..498567a4ecf 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -70,15 +70,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Text configuration self.text_media_type = "text/plain" - def validate_environment( - self, headers: dict, model: str, api_key: Optional[str] = None - ) -> dict: + def validate_environment(self, headers: dict, model: str, api_key: Optional[str] = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" return headers - def get_complete_url( - self, api_base: Optional[str], model: str, api_key: Optional[str] = None - ) -> str: + def get_complete_url(self, api_base: Optional[str], model: str, api_key: Optional[str] = None) -> str: """Get complete URL - handled by aws_sdk_bedrock_runtime.""" return api_base or "" @@ -86,9 +82,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): """Bedrock requires session configuration.""" return True - def session_configuration_request( - self, model: str, tools: Optional[List[dict]] = None - ) -> str: + def session_configuration_request(self, model: str, tools: Optional[List[dict]] = None) -> str: """ Create initial session configuration for Bedrock Nova Sonic. @@ -126,19 +120,13 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Add tool configuration if tools are provided if tools: - prompt_start_config["toolUseOutputConfiguration"] = { - "mediaType": "application/json" - } - prompt_start_config["toolConfiguration"] = { - "tools": self._transform_tools_to_bedrock_format(tools) - } + prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"} + prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)} prompt_start = {"event": {"promptStart": prompt_start_config}} # Return as a marker that we've sent the configuration - return json.dumps( - {"session_start": session_start, "prompt_start": prompt_start} - ) + return json.dumps({"session_start": session_start, "prompt_start": prompt_start}) def _transform_tools_to_bedrock_format(self, tools: List[dict]) -> List[dict]: """ @@ -158,17 +146,13 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolSpec": { "name": function.get("name", ""), "description": function.get("description", ""), - "inputSchema": { - "json": json.dumps(function.get("parameters", {})) - }, + "inputSchema": {"json": json.dumps(function.get("parameters", {}))}, } } bedrock_tools.append(bedrock_tool) return bedrock_tools - def _map_audio_format_to_sample_rate( - self, audio_format: str, is_output: bool = True - ) -> int: + def _map_audio_format_to_sample_rate(self, audio_format: str, is_output: bool = True) -> int: """ Map OpenAI audio format to sample rate. @@ -213,16 +197,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.voice_id = session_config["voice"] if "output_audio_format" in session_config: output_format = session_config["output_audio_format"] - self.output_sample_rate_hertz = self._map_audio_format_to_sample_rate( - output_format, is_output=True - ) + self.output_sample_rate_hertz = self._map_audio_format_to_sample_rate(output_format, is_output=True) # Update audio input configuration from session if provided if "input_audio_format" in session_config: input_format = session_config["input_audio_format"] - self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate( - input_format, is_output=False - ) + self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False) # Allow direct override of sample rates if provided (custom extension) if "output_sample_rate_hertz" in session_config: @@ -262,12 +242,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Add tool configuration if tools are provided tools = session_config.get("tools") if tools: - prompt_start_config["toolUseOutputConfiguration"] = { - "mediaType": "application/json" - } - prompt_start_config["toolConfiguration"] = { - "tools": self._transform_tools_to_bedrock_format(tools) - } + prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"} + prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)} prompt_start = {"event": {"promptStart": prompt_start_config}} messages.append(json.dumps(prompt_start)) @@ -317,9 +293,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return messages - def transform_input_audio_buffer_append_event( - self, json_message: dict - ) -> List[str]: + def transform_input_audio_buffer_append_event(self, json_message: dict) -> List[str]: """ Transform input_audio_buffer.append event to Bedrock audio input. @@ -371,9 +345,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return messages - def transform_input_audio_buffer_commit_event( - self, json_message: dict - ) -> List[str]: + def transform_input_audio_buffer_commit_event(self, json_message: dict) -> List[str]: """ Transform input_audio_buffer.commit event to Bedrock audio content end. @@ -418,9 +390,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Handle tool result if item_type == "function_call_output": - return self.transform_conversation_item_create_tool_result_event( - json_message - ) + return self.transform_conversation_item_create_tool_result_event(json_message) # Handle regular message if item_type == "message": @@ -438,9 +408,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "type": "TEXT", "interactive": True, "role": "USER", - "textInputConfiguration": { - "mediaType": self.text_media_type - }, + "textInputConfiguration": {"mediaType": self.text_media_type}, } } } @@ -622,9 +590,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Determine content type content_type = content_start.get("type", "TEXT") - current_delta_type: ALL_DELTA_TYPES = ( - "text" if content_type == "TEXT" else "audio" - ) + current_delta_type: ALL_DELTA_TYPES = "text" if content_type == "TEXT" else "audio" returned_messages: List[OpenAIRealtimeEvents] = [] @@ -666,9 +632,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, part=( - {"type": "text", "text": ""} - if current_delta_type == "text" - else {"type": "audio", "transcript": ""} + {"type": "text", "text": ""} if current_delta_type == "text" else {"type": "audio", "transcript": ""} ), response_id=current_response_id, ) @@ -793,9 +757,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Accumulate text accumulated_text = "" if current_delta_chunks: - accumulated_text = "".join( - [chunk.get("delta", "") for chunk in current_delta_chunks] - ) + accumulated_text = "".join([chunk.get("delta", "") for chunk in current_delta_chunks]) text_done = OpenAIRealtimeResponseTextDone( type="response.text.done", @@ -938,11 +900,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): tool_input = {} if "input" in tool_use: try: - tool_input = ( - json.loads(tool_use["input"]) - if isinstance(tool_use["input"], str) - else tool_use["input"] - ) + tool_input = json.loads(tool_use["input"]) if isinstance(tool_use["input"], str) else tool_use["input"] except json.JSONDecodeError: tool_input = {} @@ -970,9 +928,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): tool_name, ) - def transform_conversation_item_create_tool_result_event( - self, json_message: dict - ) -> List[str]: + def transform_conversation_item_create_tool_result_event(self, json_message: dict) -> List[str]: """ Transform conversation.item.create with tool result to Bedrock format. @@ -1016,9 +972,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolResult": { "promptName": self.prompt_name, "contentName": tool_content_name, - "content": ( - output if isinstance(output, str) else json.dumps(output) - ), + "content": (output if isinstance(output, str) else json.dumps(output)), } } } @@ -1060,53 +1014,27 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): json_message = json.loads(message) except json.JSONDecodeError: message_preview = ( - message[:200].decode("utf-8", errors="replace") - if isinstance(message, bytes) - else message[:200] + message[:200].decode("utf-8", errors="replace") if isinstance(message, bytes) else message[:200] ) verbose_logger.warning(f"Invalid JSON message: {message_preview}") return { "response": [], - "current_output_item_id": realtime_response_transform_input.get( - "current_output_item_id" - ), - "current_response_id": realtime_response_transform_input.get( - "current_response_id" - ), - "current_delta_chunks": realtime_response_transform_input.get( - "current_delta_chunks" - ), - "current_conversation_id": realtime_response_transform_input.get( - "current_conversation_id" - ), - "current_item_chunks": realtime_response_transform_input.get( - "current_item_chunks" - ), - "current_delta_type": realtime_response_transform_input.get( - "current_delta_type" - ), - "session_configuration_request": realtime_response_transform_input.get( - "session_configuration_request" - ), + "current_output_item_id": realtime_response_transform_input.get("current_output_item_id"), + "current_response_id": realtime_response_transform_input.get("current_response_id"), + "current_delta_chunks": realtime_response_transform_input.get("current_delta_chunks"), + "current_conversation_id": realtime_response_transform_input.get("current_conversation_id"), + "current_item_chunks": realtime_response_transform_input.get("current_item_chunks"), + "current_delta_type": realtime_response_transform_input.get("current_delta_type"), + "session_configuration_request": realtime_response_transform_input.get("session_configuration_request"), } # Extract state - current_output_item_id = realtime_response_transform_input.get( - "current_output_item_id" - ) - current_response_id = realtime_response_transform_input.get( - "current_response_id" - ) - current_conversation_id = realtime_response_transform_input.get( - "current_conversation_id" - ) - current_delta_chunks = realtime_response_transform_input.get( - "current_delta_chunks" - ) + current_output_item_id = realtime_response_transform_input.get("current_output_item_id") + current_response_id = realtime_response_transform_input.get("current_response_id") + current_conversation_id = realtime_response_transform_input.get("current_conversation_id") + current_delta_chunks = realtime_response_transform_input.get("current_delta_chunks") current_delta_type = realtime_response_transform_input.get("current_delta_type") - session_configuration_request = realtime_response_transform_input.get( - "session_configuration_request" - ) + session_configuration_request = realtime_response_transform_input.get("session_configuration_request") returned_messages: List[OpenAIRealtimeEvents] = [] @@ -1115,9 +1043,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Route to appropriate transformation method if "sessionStart" in event: - session_created = self.transform_session_start_event( - event, model, logging_obj - ) + session_created = self.transform_session_start_event(event, model, logging_obj) returned_messages.append(session_created) session_configuration_request = json.dumps({"configured": True}) @@ -1146,9 +1072,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): returned_messages.extend(events) elif "audioOutput" in event: - events = self.transform_audio_output_event( - event, current_output_item_id, current_response_id - ) + events = self.transform_audio_output_event(event, current_output_item_id, current_response_id) returned_messages.extend(events) elif "contentEnd" in event: @@ -1175,9 +1099,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_output_item_id, current_response_id, current_delta_type, - ) = self.transform_prompt_end_event( - event, current_response_id, current_conversation_id - ) + ) = self.transform_prompt_end_event(event, current_response_id, current_conversation_id) returned_messages.extend(events) return { @@ -1186,9 +1108,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "current_response_id": current_response_id, "current_delta_chunks": current_delta_chunks, "current_conversation_id": current_conversation_id, - "current_item_chunks": realtime_response_transform_input.get( - "current_item_chunks" - ), + "current_item_chunks": realtime_response_transform_input.get("current_item_chunks"), "current_delta_type": current_delta_type, "session_configuration_request": session_configuration_request, } diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 276c6f23c33..1728f52a413 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -99,9 +99,7 @@ class BedrockRerankHandler(BaseAWSLLM): return self.arerank( prepared_request, timeout=timeout, - client=client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None, + client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None, ) # type: ignore if client is None or not isinstance(client, HTTPHandler): @@ -142,9 +140,7 @@ class BedrockRerankHandler(BaseAWSLLM): from botocore.awsrequest import AWSRequest except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - boto3_credentials_info = self._get_boto_credentials_from_optional_params( - optional_params, model - ) + boto3_credentials_info = self._get_boto_credentials_from_optional_params(optional_params, model) ### SET RUNTIME ENDPOINT ### _, proxy_endpoint_url = self.get_runtime_endpoint( @@ -152,9 +148,7 @@ class BedrockRerankHandler(BaseAWSLLM): aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, aws_region_name=boto3_credentials_info.aws_region_name, ) - proxy_endpoint_url = proxy_endpoint_url.replace( - "bedrock-runtime", "bedrock-agent-runtime" - ) + proxy_endpoint_url = proxy_endpoint_url.replace("bedrock-runtime", "bedrock-agent-runtime") proxy_endpoint_url = f"{proxy_endpoint_url}/rerank" sigv4 = SigV4Auth( boto3_credentials_info.credentials, @@ -167,9 +161,7 @@ class BedrockRerankHandler(BaseAWSLLM): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - request = AWSRequest( - method="POST", url=proxy_endpoint_url, data=body, headers=headers - ) + request = AWSRequest(method="POST", url=proxy_endpoint_url, data=body, headers=headers) sigv4.add_auth(request) if ( extra_headers is not None and "Authorization" in extra_headers diff --git a/litellm/llms/bedrock/rerank/transformation.py b/litellm/llms/bedrock/rerank/transformation.py index b5d33eda49f..38625a26939 100644 --- a/litellm/llms/bedrock/rerank/transformation.py +++ b/litellm/llms/bedrock/rerank/transformation.py @@ -29,9 +29,7 @@ from litellm.types.rerank import ( class BedrockRerankConfig: - def _transform_sources( - self, documents: List[Union[str, dict]] - ) -> List[BedrockRerankSource]: + def _transform_sources(self, documents: List[Union[str, dict]]) -> List[BedrockRerankSource]: """ Transform the sources from RerankRequest format to Bedrock format. """ @@ -50,9 +48,7 @@ class BedrockRerankConfig: else: _sources.append( BedrockRerankSource( - inlineDocumentSource=BedrockRerankInlineDocumentSource( - jsonDocument=document, type="JSON" - ), + inlineDocumentSource=BedrockRerankInlineDocumentSource(jsonDocument=document, type="JSON"), type="INLINE", ) ) @@ -73,9 +69,7 @@ class BedrockRerankConfig: ], rerankingConfiguration=BedrockRerankConfiguration( bedrockRerankingConfiguration=BedrockRerankBedrockRerankingConfiguration( - modelConfiguration=BedrockRerankModelConfiguration( - modelArn=request_data.model - ), + modelConfiguration=BedrockRerankModelConfiguration(modelArn=request_data.model), numberOfResults=request_data.top_n or len(request_data.documents), ), type="BEDROCK_RERANKING_MODEL", @@ -90,9 +84,7 @@ class BedrockRerankConfig: example input: {"results":[{"index":0,"relevanceScore":0.6847912669181824},{"index":1,"relevanceScore":0.5980774760246277}]} """ - _billed_units = RerankBilledUnits( - **response.get("usage", {"search_units": 1}) - ) # by default 1 search unit + _billed_units = RerankBilledUnits(**response.get("usage", {"search_units": 1})) # by default 1 search unit _tokens = RerankTokens(**response.get("usage", {})) rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index a479d148064..c1b124caec1 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -38,9 +38,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): BaseVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: return {} def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: @@ -49,9 +47,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "write": [], } - def get_supported_openai_params( - self, model: str - ) -> List[VECTOR_STORE_OPENAI_PARAMS]: + def get_supported_openai_params(self, model: str) -> List[VECTOR_STORE_OPENAI_PARAMS]: return ["filters", "max_num_results", "ranking_options"] def _map_operator_to_aws(self, operator: str) -> str: @@ -176,9 +172,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): return optional_params - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: headers = headers or {} headers.setdefault("Content-Type", "application/json") return headers @@ -187,12 +181,8 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): aws_region_name = litellm_params.get("aws_region_name") endpoint_url, _ = self.get_runtime_endpoint( api_base=api_base, - aws_bedrock_runtime_endpoint=litellm_params.get( - "aws_bedrock_runtime_endpoint" - ), - aws_region_name=self.get_aws_region_name_for_non_llm_api_calls( - aws_region_name=aws_region_name - ), + aws_bedrock_runtime_endpoint=litellm_params.get("aws_bedrock_runtime_endpoint"), + aws_region_name=self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=aws_region_name), endpoint_type="agent", ) return f"{endpoint_url}/knowledgebases" @@ -210,9 +200,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if isinstance(query, list): query = " ".join(query) - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}/retrieve" request_body: Dict[str, Any] = { @@ -223,43 +211,28 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if isinstance(extra_body, dict): retrieval_config = deepcopy( - extra_body.get("retrievalConfiguration") - or extra_body.get("retrieval_configuration") - or {} + extra_body.get("retrievalConfiguration") or extra_body.get("retrieval_configuration") or {} ) max_results = vector_store_search_optional_params.get("max_num_results") if max_results is not None: - existing_number_of_results = retrieval_config.get( - "vectorSearchConfiguration", {} - ).get("numberOfResults") - if ( - existing_number_of_results is not None - and existing_number_of_results != max_results - ): + existing_number_of_results = retrieval_config.get("vectorSearchConfiguration", {}).get("numberOfResults") + if existing_number_of_results is not None and existing_number_of_results != max_results: verbose_logger.debug( "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s", existing_number_of_results, max_results, ) - retrieval_config.setdefault("vectorSearchConfiguration", {})[ - "numberOfResults" - ] = max_results + retrieval_config.setdefault("vectorSearchConfiguration", {})["numberOfResults"] = max_results filters = vector_store_search_optional_params.get("filters") if filters is not None: - existing_filter = retrieval_config.get("vectorSearchConfiguration", {}).get( - "filter" - ) + existing_filter = retrieval_config.get("vectorSearchConfiguration", {}).get("filter") if existing_filter is not None and existing_filter != filters: verbose_logger.debug( "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params" ) - retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = ( - filters - ) + retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = filters if retrieval_config: - request_body["retrievalConfiguration"] = cast( - BedrockKBRetrievalConfiguration, retrieval_config - ) + request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config) litellm_logging_obj.model_call_details["query"] = query return url, request_body @@ -290,11 +263,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if source_uri: return source_uri - chunk_id = ( - metadata.get("x-amz-bedrock-kb-chunk-id", "unknown") - if metadata - else "unknown" - ) + chunk_id = metadata.get("x-amz-bedrock-kb-chunk-id", "unknown") if metadata else "unknown" return f"bedrock-kb-{chunk_id}" def _get_filename_from_metadata(self, metadata: Dict[str, Any]) -> str: @@ -308,9 +277,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): try: parsed_uri = urlparse(source_uri) filename = ( - parsed_uri.path.split("/")[-1] - if parsed_uri.path and parsed_uri.path != "/" - else parsed_uri.netloc + parsed_uri.path.split("/")[-1] if parsed_uri.path and parsed_uri.path != "/" else parsed_uri.netloc ) if not filename or filename == "/": filename = parsed_uri.netloc @@ -318,11 +285,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): except Exception: return source_uri - data_source_id = ( - metadata.get("x-amz-bedrock-kb-data-source-id", "unknown") - if metadata - else "unknown" - ) + data_source_id = metadata.get("x-amz-bedrock-kb-data-source-id", "unknown") if metadata else "unknown" return f"bedrock-kb-document-{data_source_id}" def _get_attributes_from_metadata(self, metadata: Dict[str, Any]) -> Dict[str, Any]: diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index f688cea10f1..8fc720daa29 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -97,15 +97,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): def get_supported_openai_params(self, model: str) -> list: base_params = super().get_supported_openai_params(model) try: - if litellm.supports_reasoning( - model=model, custom_llm_provider=self.custom_llm_provider - ): + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): if "reasoning_effort" not in base_params: base_params.append("reasoning_effort") except Exception as e: - verbose_logger.debug( - f"BedrockMantleChatConfig: error checking reasoning support: {e}" - ) + verbose_logger.debug(f"BedrockMantleChatConfig: error checking reasoning support: {e}") return base_params def get_model_response_iterator( diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index d517ab940ce..eedb57ea386 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -28,9 +28,7 @@ from litellm.secret_managers.main import get_secret_str BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1" # Standard Mantle host: https://bedrock-mantle..api.aws (group 1 = region). -MANTLE_HOST_RE = re.compile( - r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE -) +MANTLE_HOST_RE = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE) class BedrockMantleAuthMixin: @@ -38,11 +36,7 @@ class BedrockMantleAuthMixin: @staticmethod def _resolve_bearer_token(api_key: str | None) -> str | None: - return ( - api_key - or get_secret_str("BEDROCK_MANTLE_API_KEY") - or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") - ) + return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") @staticmethod def _resolve_region(params: dict) -> str: diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 2e30f85fd0e..31975444a31 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -40,9 +40,7 @@ _BASE_SUFFIXES_TO_STRIP = ( ) # Per Bedrock Mantle Responses API validation errors. -_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES = frozenset( - {"function", "mcp", "custom", "namespace", "tool_search"} -) +_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES = frozenset({"function", "mcp", "custom", "namespace", "tool_search"}) class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig): @@ -65,11 +63,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI litellm_params: dict, ) -> str: region = self._resolve_region({**litellm_params, "api_base": api_base}) - base = ( - api_base - or get_secret_str("BEDROCK_MANTLE_API_BASE") - or f"https://bedrock-mantle.{region}.api.aws" - ) + base = api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") or f"https://bedrock-mantle.{region}.api.aws" base = base.rstrip("/") for suffix in _BASE_SUFFIXES_TO_STRIP: if base.endswith(suffix): @@ -83,9 +77,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI path = "/openai/v1/responses" if self.use_openai_path else "/v1/responses" return f"{base}{path}" - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() bearer = self._resolve_bearer_token(litellm_params.api_key) if bearer: @@ -117,8 +109,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI if dropped_types: verbose_logger.warning( - "Bedrock Mantle Responses API: dropping unsupported tool type(s) " - "%s (supported: %s).", + "Bedrock Mantle Responses API: dropping unsupported tool type(s) %s (supported: %s).", sorted(set(dropped_types)), sorted(_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES), ) diff --git a/litellm/llms/black_forest_labs/common_utils.py b/litellm/llms/black_forest_labs/common_utils.py index 237208693f7..71c09093679 100644 --- a/litellm/llms/black_forest_labs/common_utils.py +++ b/litellm/llms/black_forest_labs/common_utils.py @@ -46,9 +46,7 @@ def assert_bfl_polling_url(polling_url: str) -> None: message="Rejected polling URL: scheme must be https", ) - if host != _BFL_REGISTERED_DOMAIN and not host.endswith( - "." + _BFL_REGISTERED_DOMAIN - ): + if host != _BFL_REGISTERED_DOMAIN and not host.endswith("." + _BFL_REGISTERED_DOMAIN): raise BlackForestLabsError( status_code=502, message="Rejected polling URL: host is not within the bfl.ai domain", diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 309e00ade62..a80ca491d74 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -134,9 +134,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): BFL uses x-key header for authentication. """ final_api_key: Optional[str] = ( - api_key - or get_secret_str("BFL_API_KEY") - or get_secret_str("BLACK_FOREST_LABS_API_KEY") + api_key or get_secret_str("BFL_API_KEY") or get_secret_str("BLACK_FOREST_LABS_API_KEY") ) if not final_api_key: @@ -171,8 +169,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): return IMAGE_EDIT_MODELS[model_name] raise ValueError( - f"Unknown BFL image edit model: {model_name}. " - f"Supported models: {list(IMAGE_EDIT_MODELS.keys())}" + f"Unknown BFL image edit model: {model_name}. Supported models: {list(IMAGE_EDIT_MODELS.keys())}" ) def get_complete_url( @@ -205,9 +202,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): return image elif isinstance(image, list): # If it's a list, take the first image - return self._read_image_bytes( - image[0], depth=depth + 1, max_depth=max_depth - ) + return self._read_image_bytes(image[0], depth=depth + 1, max_depth=max_depth) elif isinstance(image, str): if image.startswith(("http://", "https://")): response = safe_get(litellm.module_level_client, image, timeout=60.0) @@ -229,8 +224,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): return data else: raise ValueError( - f"Unsupported image type: {type(image)}. " - "Expected bytes, str (URL or file path), or file-like object." + f"Unsupported image type: {type(image)}. Expected bytes, str (URL or file path), or file-like object." ) def transform_image_edit_request( diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index 18c7c173300..7176247b4be 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -50,9 +50,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): This class only handles data transformation. """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Return list of OpenAI params supported by Black Forest Labs. @@ -136,9 +134,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): optional_params["width"] = width optional_params["height"] = height except ValueError: - raise ValueError( - f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." - ) + raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") def validate_environment( self, @@ -156,9 +152,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): BFL uses x-key header for authentication. """ final_api_key: Optional[str] = ( - api_key - or get_secret_str("BFL_API_KEY") - or get_secret_str("BLACK_FOREST_LABS_API_KEY") + api_key or get_secret_str("BFL_API_KEY") or get_secret_str("BLACK_FOREST_LABS_API_KEY") ) if not final_api_key: diff --git a/litellm/llms/brave/search/transformation.py b/litellm/llms/brave/search/transformation.py index 8ffe7dcb126..54fb574087c 100644 --- a/litellm/llms/brave/search/transformation.py +++ b/litellm/llms/brave/search/transformation.py @@ -124,9 +124,7 @@ class BraveSearchConfig(BaseSearchConfig): ) if not api_key: - raise ValueError( - "BRAVE_API_KEY is not set. Set `BRAVE_API_KEY` environment variable." - ) + raise ValueError("BRAVE_API_KEY is not set. Set `BRAVE_API_KEY` environment variable.") headers["X-Subscription-Token"] = api_key headers["Accept"] = "application/json" @@ -197,10 +195,7 @@ class BraveSearchConfig(BaseSearchConfig): # Only include "include_fetch_metadata" if it is not explicitly set to False # This parameter results (more often than not) in a timestamp which we can use for last_updated - if ( - "include_fetch_metadata" in optional_params - and optional_params["include_fetch_metadata"] is False - ): + if "include_fetch_metadata" in optional_params and optional_params["include_fetch_metadata"] is False: request_data["include_fetch_metadata"] = False else: request_data["include_fetch_metadata"] = True @@ -215,19 +210,14 @@ class BraveSearchConfig(BaseSearchConfig): # Convert to multiple "site:domain" clauses, joined by OR domains = optional_params["search_domain_filter"] if isinstance(domains, list) and len(domains) > 0: - request_data["q"] = self._append_domain_filters( - request_data["q"], domains - ) + request_data["q"] = self._append_domain_filters(request_data["q"], domains) # Convert to dict before dynamic key assignments result_data = dict(request_data) # Pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value # Store params in special key for URL building (Brave Search API uses GET not POST) @@ -277,9 +267,7 @@ class BraveSearchConfig(BaseSearchConfig): url = result.get("url", "") snippet = result.get("description", "") date = to_yyyy_mm_dd(result.get("page_age") or result.get("age")) - last_updated = to_yyyy_mm_dd( - result.get("fetched_content_timestamp", "") - ) + last_updated = to_yyyy_mm_dd(result.get("fetched_content_timestamp", "")) search_result = SearchResult( title=title, diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index 7d9afe01fa6..e5d91c6533f 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -132,9 +132,7 @@ class BytezChatConfig(BaseConfig): ) if not messages: - raise Exception( - "kwarg `messages` must be an array of messages that follow the openai chat standard" - ) + raise Exception("kwarg `messages` must be an array of messages that follow the openai chat standard") if not api_key: raise Exception("Missing api_key, make sure you pass in your api key") @@ -273,9 +271,7 @@ class BytezChatConfig(BaseConfig): timeout=STREAMING_TIMEOUT, ) except httpx.HTTPStatusError as e: - raise BytezError( - status_code=e.response.status_code, message=e.response.text - ) + raise BytezError(status_code=e.response.status_code, message=e.response.text) if response.status_code != 200: raise BytezError(status_code=response.status_code, message=response.text) @@ -317,9 +313,7 @@ class BytezChatConfig(BaseConfig): timeout=STREAMING_TIMEOUT, ) except httpx.HTTPStatusError as e: - raise BytezError( - status_code=e.response.status_code, message=e.response.text - ) + raise BytezError(status_code=e.response.status_code, message=e.response.text) if response.status_code != 200: raise BytezError(status_code=response.status_code, message=response.text) @@ -447,9 +441,7 @@ def _adapt_string_only_content_to_lists(messages: List[Dict]): elif isinstance(content_item, dict): new_content_items.append(content_item) else: - raise Exception( - "`content` can only contain strings or openai content dicts" - ) + raise Exception("`content` can only contain strings or openai content dicts") new_content += new_content_items else: diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index e35b04a3fb3..277bcfa18d0 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -34,17 +34,11 @@ class Authenticator: "CHATGPT_TOKEN_DIR", os.path.expanduser("~/.config/litellm/chatgpt"), ) - self.auth_file = os.path.join( - self.token_dir, os.getenv("CHATGPT_AUTH_FILE", "auth.json") - ) + self.auth_file = os.path.join(self.token_dir, os.getenv("CHATGPT_AUTH_FILE", "auth.json")) self._ensure_token_dir() def get_api_base(self) -> str: - return ( - os.getenv("CHATGPT_API_BASE") - or os.getenv("OPENAI_CHATGPT_API_BASE") - or CHATGPT_API_BASE - ) + return os.getenv("CHATGPT_API_BASE") or os.getenv("OPENAI_CHATGPT_API_BASE") or CHATGPT_API_BASE def get_access_token(self) -> str: auth_data = self._read_auth_file() @@ -58,9 +52,7 @@ class Authenticator: refreshed = self._refresh_tokens(refresh_token) return refreshed["access_token"] except RefreshAccessTokenError as exc: - verbose_logger.warning( - "ChatGPT refresh token failed, re-login required: %s", exc - ) + verbose_logger.warning("ChatGPT refresh token failed, re-login required: %s", exc) cooldown_remaining = self._get_device_code_cooldown_remaining(auth_data) if cooldown_remaining > 0: @@ -149,9 +141,7 @@ class Authenticator: return None def _login_device_code(self) -> Dict[str, str]: - cooldown_remaining = self._get_device_code_cooldown_remaining( - self._read_auth_file() - ) + cooldown_remaining = self._get_device_code_cooldown_remaining(self._read_auth_file()) if cooldown_remaining > 0: token = self._wait_for_access_token(cooldown_remaining) if token: @@ -206,9 +196,7 @@ class Authenticator: "interval": str(interval or "5"), } - def _poll_for_authorization_code( - self, device_code: Dict[str, str] - ) -> Dict[str, str]: + def _poll_for_authorization_code(self, device_code: Dict[str, str]) -> Dict[str, str]: client = _get_httpx_client() interval = int(device_code.get("interval", "5")) start_time = time.time() @@ -286,9 +274,7 @@ class Authenticator: status_code=400, ) - if not all( - key in data for key in ("access_token", "refresh_token", "id_token") - ): + if not all(key in data for key in ("access_token", "refresh_token", "id_token")): raise GetAccessTokenError( message=f"Token exchange response missing fields: {data}", status_code=400, @@ -354,9 +340,7 @@ class Authenticator: "account_id": account_id, } - def _get_device_code_cooldown_remaining( - self, auth_data: Optional[Dict[str, Any]] - ) -> float: + def _get_device_code_cooldown_remaining(self, auth_data: Optional[Dict[str, Any]]) -> float: if not auth_data: return 0.0 requested_at = auth_data.get("device_code_requested_at") @@ -383,9 +367,7 @@ class Authenticator: access_token = auth_data.get("access_token") if access_token and not self._is_token_expired(auth_data, access_token): return access_token - sleep_for = min( - DEVICE_CODE_POLL_SLEEP_SECONDS, max(0.0, deadline - time.time()) - ) + sleep_for = min(DEVICE_CODE_POLL_SLEEP_SECONDS, max(0.0, deadline - time.time())) if sleep_for <= 0: break time.sleep(sleep_for) diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index a08fecd9625..3232b452a37 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -24,9 +24,7 @@ class ChatGPTToolCallNormalizer: self._stream = stream self._seen_ids: Dict[str, int] = {} # tool_call_id -> assigned_index self._next_index: int = 0 - self._last_id: Optional[str] = ( - None # tracks which tool call the next delta belongs to - ) + self._last_id: Optional[str] = None # tracks which tool call the next delta belongs to def __getattr__(self, name: str) -> Any: return getattr(self._stream, name) diff --git a/litellm/llms/chatgpt/chat/transformation.py b/litellm/llms/chatgpt/chat/transformation.py index e6480398c7e..9b0d8dc2e65 100644 --- a/litellm/llms/chatgpt/chat/transformation.py +++ b/litellm/llms/chatgpt/chat/transformation.py @@ -57,9 +57,7 @@ class ChatGPTConfig(OpenAIConfig): account_id = self.authenticator.get_account_id() session_id = ensure_chatgpt_session_id(litellm_params) - default_headers = get_chatgpt_default_headers( - api_key or "", account_id, session_id - ) + default_headers = get_chatgpt_default_headers(api_key or "", account_id, session_id) return {**default_headers, **validated_headers} def post_stream_processing(self, stream: Any) -> Any: @@ -72,8 +70,6 @@ class ChatGPTConfig(OpenAIConfig): model: str, drop_params: bool, ) -> dict: - optional_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + optional_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) optional_params.setdefault("stream", False) return optional_params diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index 830414d9cad..8afef4b3828 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -161,11 +161,7 @@ def _terminal_user_agent() -> str: token = f"WezTerm/{wezterm_version}" if wezterm_version else "WezTerm" return _sanitize_user_agent_token(token) or "WezTerm" - if ( - os.getenv("ITERM_SESSION_ID") - or os.getenv("ITERM_PROFILE") - or os.getenv("ITERM_PROFILE_NAME") - ): + if os.getenv("ITERM_SESSION_ID") or os.getenv("ITERM_PROFILE") or os.getenv("ITERM_PROFILE_NAME"): return "iTerm.app" if os.getenv("TERM_SESSION_ID"): @@ -225,9 +221,7 @@ def get_chatgpt_user_agent(originator: str) -> str: terminal_ua = _terminal_user_agent() suffix = os.getenv("CHATGPT_USER_AGENT_SUFFIX", "").strip() suffix = f" ({suffix})" if suffix else "" - candidate = ( - f"{originator}/{version} ({os_type} {os_version}; {arch}) {terminal_ua}{suffix}" - ) + candidate = f"{originator}/{version} ({os_type} {os_version}; {arch}) {terminal_ua}{suffix}" return _safe_header_value(candidate) or DEFAULT_USER_AGENT diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 56b61b66c84..8b5fae4ef35 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -55,9 +55,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): account_id = self.authenticator.get_account_id() session_id = ensure_chatgpt_session_id(litellm_params) - default_headers = get_chatgpt_default_headers( - access_token, account_id, session_id - ) + default_headers = get_chatgpt_default_headers(access_token, account_id, session_id) return {**default_headers, **headers} def transform_responses_api_request( @@ -79,9 +77,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): existing_instructions = request.get("instructions") if existing_instructions: if base_instructions not in existing_instructions: - request["instructions"] = ( - f"{base_instructions}\n\n{existing_instructions}" - ) + request["instructions"] = f"{base_instructions}\n\n{existing_instructions}" else: request["instructions"] = base_instructions request["store"] = False @@ -114,9 +110,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): logging_obj: Any, ): body_text = raw_response.text or "" - if not self._should_parse_as_sse( - raw_response=raw_response, body_text=body_text - ): + if not self._should_parse_as_sse(raw_response=raw_response, body_text=body_text): return super().transform_response_api_response( model=model, raw_response=raw_response, @@ -128,18 +122,14 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): additional_args={"complete_input_dict": {}}, ) - completed_response, error_message = self._extract_completed_response_from_sse( - body_text=body_text - ) + completed_response, error_message = self._extract_completed_response_from_sse(body_text=body_text) if completed_response is None: raise OpenAIError( message=error_message or raw_response.text, status_code=raw_response.status_code, ) - self._attach_response_headers( - completed_response=completed_response, raw_response=raw_response - ) + self._attach_response_headers(completed_response=completed_response, raw_response=raw_response) return completed_response def _should_parse_as_sse(self, raw_response: Any, body_text: str) -> bool: @@ -213,22 +203,16 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): return None response_payload = dict(response_payload) if not response_payload.get("output") and streamed_output_items: - response_payload["output"] = [ - item for _, item in sorted(streamed_output_items.items()) - ] + response_payload["output"] = [item for _, item in sorted(streamed_output_items.items())] if "created_at" in response_payload: - response_payload["created_at"] = _safe_convert_created_field( - response_payload["created_at"] - ) + response_payload["created_at"] = _safe_convert_created_field(response_payload["created_at"]) try: return ResponsesAPIResponse(**response_payload) except Exception: return ResponsesAPIResponse.model_construct(**response_payload) def _extract_error_message(self, parsed_chunk: Dict[str, Any]) -> Optional[str]: - error_obj = parsed_chunk.get("error") or ( - parsed_chunk.get("response") or {} - ).get("error") + error_obj = parsed_chunk.get("error") or (parsed_chunk.get("response") or {}).get("error") if error_obj is None: return None if isinstance(error_obj, dict): diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index d07f6eba057..95c0444924b 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -71,13 +71,9 @@ class ClarifaiConfig(OpenAIGPTConfig): dynamic_api_key = api_key or get_secret_str("CLARIFAI_API_KEY") or "" return api_base, dynamic_api_key - def transform_request( - self, model, messages, optional_params, litellm_params, headers - ): + def transform_request(self, model, messages, optional_params, litellm_params, headers): model = self.get_base_model(model) or model - return super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + return super().transform_request(model, messages, optional_params, litellm_params, headers) def transform_response( self, diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index 68f08741cc5..df8ac884a32 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -48,9 +48,7 @@ class CloudflareChatConfig(OpenAIGPTConfig): @staticmethod def _resolve_api_base(api_base: Optional[str]) -> str: if not api_base: - account_id = normalize_nonempty_secret_str( - get_secret_str("CLOUDFLARE_ACCOUNT_ID") - ) + account_id = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID")) if account_id is None: raise ValueError( "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID in the environment or pass api_base explicitly" diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index b149ae46ee9..6a91601e6fc 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -42,12 +42,8 @@ class TextCompletionCodestralError(Exception): if response is not None: self.response = response else: - self.response = httpx.Response( - status_code=status_code, request=self.request - ) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + self.response = httpx.Response(status_code=status_code, request=self.request) + super().__init__(self.message) # Call the base class constructor with the parameters it needs async def make_call( @@ -62,9 +58,7 @@ async def make_call( response = await client.post(api_base, headers=headers, data=data, stream=True) if response.status_code != 200: - raise TextCompletionCodestralError( - status_code=response.status_code, message=response.text - ) + raise TextCompletionCodestralError(status_code=response.status_code, message=response.text) completion_stream = response.aiter_lines() # LOGGING @@ -88,9 +82,7 @@ class CodestralTextCompletion: user_headers: dict, ) -> dict: if api_key is None: - raise ValueError( - "Missing CODESTRAL_API_Key - Please add CODESTRAL_API_Key to your environment variables" - ) + raise ValueError("Missing CODESTRAL_API_Key - Please add CODESTRAL_API_Key to your environment variables") headers = { "content-type": "application/json", "Authorization": "Bearer {}".format(api_key), @@ -215,9 +207,7 @@ class CodestralTextCompletion: if optional_params.pop("custom_endpoint", None) is True: completion_url = api_base else: - completion_url = ( - api_base or "https://codestral.mistral.ai/v1/fim/completions" - ) + completion_url = api_base or "https://codestral.mistral.ai/v1/fim/completions" if model in custom_prompt_dict: # check if the model has a registered custom prompt @@ -358,9 +348,7 @@ class CodestralTextCompletion: params={"timeout": timeout}, ) try: - response = await async_handler.post( - api_base, headers=headers, data=json.dumps(data) - ) + response = await async_handler.post(api_base, headers=headers, data=json.dumps(data)) except httpx.HTTPStatusError as e: raise TextCompletionCodestralError( status_code=e.response.status_code, diff --git a/litellm/llms/codestral/completion/transformation.py b/litellm/llms/codestral/completion/transformation.py index 31d6652f48a..d4299ee2ebd 100644 --- a/litellm/llms/codestral/completion/transformation.py +++ b/litellm/llms/codestral/completion/transformation.py @@ -83,9 +83,7 @@ class CodestralTextCompletionConfig(OpenAITextCompletionConfig): finish_reason = None logprobs = None - chunk_data = ( - litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk_data) or "" - ) + chunk_data = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk_data) or "" chunk_data = chunk_data.strip() if len(chunk_data) == 0 or chunk_data == "[DONE]": return { diff --git a/litellm/llms/cohere/chat/transformation.py b/litellm/llms/cohere/chat/transformation.py index 5dd44aca80a..10eea949390 100644 --- a/litellm/llms/cohere/chat/transformation.py +++ b/litellm/llms/cohere/chat/transformation.py @@ -232,9 +232,7 @@ class CohereChatConfig(BaseConfig): raw_response_json = raw_response.json() model_response.choices[0].message.content = raw_response_json["text"] # type: ignore except Exception: - raise CohereError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise CohereError(message=raw_response.text, status_code=raw_response.status_code) ## ADD CITATIONS if "citations" in raw_response_json: @@ -338,14 +336,8 @@ class CohereChatConfig(BaseConfig): "parameter_definitions": {}, } - for param_name, param_def in openai_tool["function"]["parameters"][ - "properties" - ].items(): - required_params = ( - openai_tool.get("function", {}) - .get("parameters", {}) - .get("required", []) - ) + for param_name, param_def in openai_tool["function"]["parameters"]["properties"].items(): + required_params = openai_tool.get("function", {}).get("parameters", {}).get("required", []) cohere_param_def = { "description": param_def.get("description", ""), "type": param_def.get("type", ""), diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index 9aa8c114907..909130077e4 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -144,10 +144,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig): optional_params["stream"] = value if param == "temperature": optional_params["temperature"] = value - if ( - param == "max_tokens" - and "max_completion_tokens" not in non_default_params - ): + if param == "max_tokens" and "max_completion_tokens" not in non_default_params: optional_params["max_tokens"] = value if param == "max_completion_tokens": optional_params["max_tokens"] = value @@ -178,9 +175,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig): """ Cohere v2 chat api is in openai format, so we can use the openai transform request function to transform the request. """ - data = super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + data = super().transform_request(model, messages, optional_params, litellm_params, headers) return data @@ -201,9 +196,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig): try: raw_response_json = raw_response.json() except Exception: - raise CohereError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise CohereError(message=raw_response.text, status_code=raw_response.status_code) try: cohere_v2_chat_response = CohereV2ChatResponse(**raw_response_json) # type: ignore @@ -213,21 +206,14 @@ class CohereV2ChatConfig(OpenAIGPTConfig): cohere_content = cohere_v2_chat_response["message"].get("content", None) if cohere_content is not None: model_response.choices[0].message.content = "".join( # type: ignore - [ - content.get("text", "") - for content in cohere_content - if content is not None - ] + [content.get("text", "") for content in cohere_content if content is not None] ) ## ADD CITATIONS AS ANNOTATIONS annotations: Optional[List[ChatCompletionAnnotation]] = None citations = None - if ( - "message" in cohere_v2_chat_response - and "citations" in cohere_v2_chat_response["message"] - ): + if "message" in cohere_v2_chat_response and "citations" in cohere_v2_chat_response["message"]: citations = cohere_v2_chat_response["message"]["citations"] if citations: @@ -304,9 +290,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig): ) -> BaseLLMException: return CohereError(status_code=status_code, message=error_message) - def _translate_citations_to_openai_annotations( - self, citations: List[dict] - ) -> List[ChatCompletionAnnotation]: + def _translate_citations_to_openai_annotations(self, citations: List[dict]) -> List[ChatCompletionAnnotation]: """ Transform Cohere citations to OpenAI annotations format. diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index 05e3cec5444..c03061ba18f 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -27,9 +27,7 @@ class CohereModelInfo(BaseLLMModelInfo): """ return None - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: """ Returns a list of models supported by this provider. """ @@ -118,9 +116,7 @@ def validate_environment( class ModelResponseIterator: - def __init__( - self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False - ): + def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): self.streaming_response = streaming_response self.response_iterator = self.streaming_response self.content_blocks: List = [] @@ -221,9 +217,7 @@ class ModelResponseIterator: class CohereV2ModelResponseIterator: """V2-specific response iterator for Cohere streaming""" - def __init__( - self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False - ): + def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): self.streaming_response = streaming_response self.response_iterator = self.streaming_response self.content_blocks: List = [] @@ -241,9 +235,7 @@ class CohereV2ModelResponseIterator: return content return "" - def _parse_tool_call_delta( - self, chunk: dict - ) -> Optional[ChatCompletionToolCallChunk]: + def _parse_tool_call_delta(self, chunk: dict) -> Optional[ChatCompletionToolCallChunk]: """Parse tool-call-delta chunks to extract tool calls.""" delta = chunk.get("delta", {}) tool_calls = delta.get("tool_calls", []) @@ -285,9 +277,7 @@ class CohereV2ModelResponseIterator: return {"citations": [citation_data]} return None - def _parse_message_end( - self, chunk: dict - ) -> Tuple[bool, str, Optional[ChatCompletionUsageBlock]]: + def _parse_message_end(self, chunk: dict) -> Tuple[bool, str, Optional[ChatCompletionUsageBlock]]: """Parse message-end events to extract finish info and usage.""" data = chunk.get("data", {}) delta = data.get("delta", {}) @@ -301,8 +291,7 @@ class CohereV2ModelResponseIterator: usage = ChatCompletionUsageBlock( prompt_tokens=tokens_data.get("input_tokens", 0), completion_tokens=tokens_data.get("output_tokens", 0), - total_tokens=tokens_data.get("input_tokens", 0) - + tokens_data.get("output_tokens", 0), + total_tokens=tokens_data.get("input_tokens", 0) + tokens_data.get("output_tokens", 0), ) return is_finished, finish_reason, usage diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 81b6a1c7aec..bd2859fa3dc 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/litellm/llms/cohere/embed/handler.py @@ -41,13 +41,9 @@ class CohereError(Exception): def __init__(self, status_code, message): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://api.cohere.ai/v1/generate" - ) + self.request = httpx.Request(method="POST", url="https://api.cohere.ai/v1/generate") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs async def async_embedding( @@ -153,11 +149,7 @@ def embedding( api_key=api_key, headers=headers, encoding=encoding, - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), ) ## LOGGING diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py index b5b350a952c..3325e6be578 100644 --- a/litellm/llms/cohere/embed/transformation.py +++ b/litellm/llms/cohere/embed/transformation.py @@ -122,9 +122,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): optional_params: dict, headers: dict, ) -> dict: - if isinstance(input, list) and ( - isinstance(input[0], list) or isinstance(input[0], int) - ): + if isinstance(input, list) and (isinstance(input[0], list) or isinstance(input[0], int)): raise ValueError("Input must be a list of strings") return cast( dict, @@ -197,9 +195,7 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): output_data = [] for k, embedding_list in embeddings.items(): for idx, embedding in enumerate(embedding_list): - output_data.append( - {"object": "embedding", "index": idx, "embedding": embedding} - ) + output_data.append({"object": "embedding", "index": idx, "embedding": embedding}) model_response.object = "list" model_response.data = output_data model_response.model = model diff --git a/litellm/llms/cohere/embed/v1_transformation.py b/litellm/llms/cohere/embed/v1_transformation.py index 82c901e7eca..3f0fcfc03ad 100644 --- a/litellm/llms/cohere/embed/v1_transformation.py +++ b/litellm/llms/cohere/embed/v1_transformation.py @@ -27,9 +27,7 @@ class CohereEmbeddingConfig: def get_supported_openai_params(self) -> List[str]: return ["encoding_format"] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": optional_params["embedding_types"] = v @@ -143,9 +141,7 @@ class CohereEmbeddingConfig: """ embeddings = response_json["embeddings"] output_data = [] - is_embeddings_by_type = ( - response_json.get("response_type") == "embeddings_by_type" - ) + is_embeddings_by_type = response_json.get("response_type") == "embeddings_by_type" if isinstance(embeddings, dict): is_embeddings_by_type = True @@ -163,9 +159,7 @@ class CohereEmbeddingConfig: ) else: for idx, embedding in enumerate(embeddings): - output_data.append( - {"object": "embedding", "index": idx, "embedding": embedding} - ) + output_data.append({"object": "embedding", "index": idx, "embedding": embedding}) model_response.object = "list" model_response.data = output_data model_response.model = model diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py index 0824e1cca41..36ca3895d4a 100644 --- a/litellm/llms/cohere/rerank/guardrail_translation/handler.py +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -58,15 +58,9 @@ class CohereRerankHandler(BaseTranslation): """ # Collect every scannable text field in a stable order so the # guardrailed results can be written back to the right key by index. - fields_to_scan = [ - (key, data[key]) - for key in self._SCANNED_FIELDS - if isinstance(data.get(key), str) - ] + fields_to_scan = [(key, data[key]) for key in self._SCANNED_FIELDS if isinstance(data.get(key), str)] if not fields_to_scan: - verbose_proxy_logger.debug( - "Rerank: No query/instruction to process or not strings" - ) + verbose_proxy_logger.debug("Rerank: No query/instruction to process or not strings") return data inputs = GenericGuardrailAPIInputs(texts=[value for _, value in fields_to_scan]) @@ -88,8 +82,7 @@ class CohereRerankHandler(BaseTranslation): if idx < len(guardrailed_texts): data[key] = guardrailed_texts[idx] verbose_proxy_logger.debug( - "Rerank: Applied guardrail to %s. " - "Original length: %d, New length: %d", + "Rerank: Applied guardrail to %s. Original length: %d, New length: %d", key, len(original), len(data[key]), @@ -122,7 +115,6 @@ class CohereRerankHandler(BaseTranslation): Unmodified response (rankings don't need text guardrails) """ verbose_proxy_logger.debug( - "Rerank: Output processing not applicable " - "(output contains relevance scores, not text)" + "Rerank: Output processing not applicable (output contains relevance scores, not text)" ) return response diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index dd3f0f1a446..e494e89fbf2 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -83,11 +83,7 @@ class CohereRerankConfig(BaseRerankConfig): optional_params: dict | None = None, ) -> dict: if api_key is None: - api_key = ( - get_secret_str("COHERE_API_KEY") - or get_secret_str("CO_API_KEY") - or litellm.cohere_key - ) + api_key = get_secret_str("COHERE_API_KEY") or get_secret_str("CO_API_KEY") or litellm.cohere_key if api_key is None: raise ValueError( @@ -148,9 +144,7 @@ class CohereRerankConfig(BaseRerankConfig): try: raw_response_json = raw_response.json() except Exception: - raise CohereError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise CohereError(message=raw_response.text, status_code=raw_response.status_code) return RerankResponse(**raw_response_json) diff --git a/litellm/llms/cometapi/chat/transformation.py b/litellm/llms/cometapi/chat/transformation.py index 1e15ee188c6..1a0a3e88547 100644 --- a/litellm/llms/cometapi/chat/transformation.py +++ b/litellm/llms/cometapi/chat/transformation.py @@ -36,9 +36,7 @@ class CometAPIConfig(OpenAIGPTConfig): """ Map OpenAI format parameters to CometAPI format """ - mapped_openai_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + mapped_openai_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) # CometAPI-specific parameters (if any) extra_body: dict[str, Any] = {} @@ -63,9 +61,7 @@ class CometAPIConfig(OpenAIGPTConfig): Remove cache control flags from messages and tools if not supported """ # For CometAPI, use default behavior (remove cache control) - return super().remove_cache_control_flag_from_messages_and_tools( - model, messages, tools - ) + return super().remove_cache_control_flag_from_messages_and_tools(model, messages, tools) def transform_request( self, @@ -82,9 +78,7 @@ class CometAPIConfig(OpenAIGPTConfig): dict: The transformed request. Sent as the body of the API call. """ extra_body = optional_params.pop("extra_body", {}) - response = super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + response = super().transform_request(model, messages, optional_params, litellm_params, headers) response.update(extra_body) return response @@ -169,9 +163,7 @@ class CometAPIChatCompletionStreamingHandler(BaseModelResponseIterator): # Handle error in chunk if "error" in chunk: error_chunk = chunk["error"] - error_message = "CometAPI Error: {}".format( - error_chunk.get("message", "Unknown error") - ) + error_message = "CometAPI Error: {}".format(error_chunk.get("message", "Unknown error")) raise CometAPIException( message=error_message, status_code=error_chunk.get("code", 400), @@ -183,9 +175,7 @@ class CometAPIChatCompletionStreamingHandler(BaseModelResponseIterator): for choice in chunk["choices"]: # Handle reasoning content if present if "delta" in choice and "reasoning" in choice["delta"]: - choice["delta"]["reasoning_content"] = choice["delta"].get( - "reasoning" - ) + choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning") new_choices.append(choice) return ModelResponseStream( diff --git a/litellm/llms/cometapi/embed/transformation.py b/litellm/llms/cometapi/embed/transformation.py index d1972def8b7..2d481eb1bcb 100644 --- a/litellm/llms/cometapi/embed/transformation.py +++ b/litellm/llms/cometapi/embed/transformation.py @@ -39,9 +39,7 @@ class CometAPIEmbeddingConfig(BaseEmbeddingConfig): """ Get the complete URL for the CometAPI embedding endpoint. """ - api_base = ( - "https://api.cometapi.com/v1" if api_base is None else api_base.rstrip("/") - ) + api_base = "https://api.cometapi.com/v1" if api_base is None else api_base.rstrip("/") complete_url = f"{api_base}/embeddings" return complete_url @@ -152,6 +150,4 @@ class CometAPIEmbeddingConfig(BaseEmbeddingConfig): """ Get the appropriate error class for CometAPI exceptions. """ - return CometAPIException( - message=error_message, status_code=status_code, headers=headers - ) + return CometAPIException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/cometapi/image_generation/cost_calculator.py b/litellm/llms/cometapi/image_generation/cost_calculator.py index 987e79e18da..b10c9d09087 100644 --- a/litellm/llms/cometapi/image_generation/cost_calculator.py +++ b/litellm/llms/cometapi/image_generation/cost_calculator.py @@ -22,6 +22,4 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/llms/cometapi/image_generation/transformation.py b/litellm/llms/cometapi/image_generation/transformation.py index bc6bd3f3ecc..e78b50b2fab 100644 --- a/litellm/llms/cometapi/image_generation/transformation.py +++ b/litellm/llms/cometapi/image_generation/transformation.py @@ -24,9 +24,7 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://api.cometapi.com" IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ https://api.cometapi.com/v1/images/generations """ @@ -94,11 +92,7 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - final_api_key: Optional[str] = ( - api_key - or get_secret_str("COMETAPI_KEY") - or get_secret_str("COMETAPI_API_KEY") - ) + final_api_key: Optional[str] = api_key or get_secret_str("COMETAPI_KEY") or get_secret_str("COMETAPI_API_KEY") if not final_api_key: raise ValueError("COMETAPI_KEY or COMETAPI_API_KEY is not set") diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index d4b9c5a83ae..2dc1ade2f4e 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -76,9 +76,7 @@ class CompactifAIChatConfig(OpenAIGPTConfig): # Convert tool calls to content for JSON mode tool_calls = message.get("tool_calls", []) if len(tool_calls) == 1: - message["content"] = tool_calls[0]["function"].get( - "arguments", "" - ) + message["content"] = tool_calls[0]["function"].get("arguments", "") message["tool_calls"] = None returned_response = ModelResponse(**response_json) diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 93b6c563dc1..9726314409b 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -40,19 +40,13 @@ class BaseLLMAIOHTTPHandler: connector: Optional[aiohttp.BaseConnector] = None, ): self.client_session = client_session - self._owns_session = ( - client_session is None - ) # Track if we own the session for cleanup + self._owns_session = client_session is None # Track if we own the session for cleanup self.transport = transport - self._owns_transport = ( - transport is None - ) # Track if we own the transport for cleanup + self._owns_transport = transport is None # Track if we own the transport for cleanup self.connector = connector - self._owns_connector = ( - connector is None - ) # Track if we own the connector for cleanup + self._owns_connector = connector is None # Track if we own the connector for cleanup def _get_or_create_transport(self) -> Optional[LiteLLMAiohttpTransport]: """Get existing transport or create a new one if needed.""" @@ -99,9 +93,7 @@ class BaseLLMAIOHTTPHandler: session = aiohttp.ClientSession() return session - def _get_async_client_session( - self, dynamic_client_session: Optional[ClientSession] = None - ) -> ClientSession: + def _get_async_client_session(self, dynamic_client_session: Optional[ClientSession] = None) -> ClientSession: if dynamic_client_session: return dynamic_client_session elif self.client_session: @@ -115,19 +107,11 @@ class BaseLLMAIOHTTPHandler: async def close(self): """Close the aiohttp client session and transport if we own them.""" # Close client session if we own it - if ( - self.client_session - and not self.client_session.closed - and self._owns_session - ): + if self.client_session and not self.client_session.closed and self._owns_session: await self.client_session.close() # Close transport if we own it - if ( - self.transport - and self._owns_transport - and hasattr(self.transport, "aclose") - ): + if self.transport and self._owns_transport and hasattr(self.transport, "aclose"): try: await self.transport.aclose() except Exception: @@ -141,11 +125,7 @@ class BaseLLMAIOHTTPHandler: Provides defense-in-depth for issue #12443 - ensures cleanup happens even if atexit handler doesn't run (abnormal termination). """ - if ( - self.client_session is not None - and not self.client_session.closed - and self._owns_session - ): + if self.client_session is not None and not self.client_session.closed and self._owns_session: try: import asyncio @@ -182,14 +162,10 @@ class BaseLLMAIOHTTPHandler: stream: bool = False, ) -> aiohttp.ClientResponse: """Common implementation across stream + non-stream calls. Meant to ensure consistent error-handling.""" - max_retry_on_unprocessable_entity_error = ( - provider_config.max_retry_on_unprocessable_entity_error - ) + max_retry_on_unprocessable_entity_error = provider_config.max_retry_on_unprocessable_entity_error response: Optional[aiohttp.ClientResponse] = None - async_client_session = self._get_async_client_session( - dynamic_client_session=async_client_session - ) + async_client_session = self._get_async_client_session(dynamic_client_session=async_client_session) for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: @@ -231,9 +207,7 @@ class BaseLLMAIOHTTPHandler: content: Any = None, params: Optional[dict] = None, ) -> httpx.Response: - max_retry_on_unprocessable_entity_error = ( - provider_config.max_retry_on_unprocessable_entity_error - ) + max_retry_on_unprocessable_entity_error = provider_config.max_retry_on_unprocessable_entity_error response: Optional[httpx.Response] = None @@ -255,11 +229,7 @@ class BaseLLMAIOHTTPHandler: e=e, litellm_params=litellm_params ) if should_retry and not hit_max_retry: - data = ( - provider_config.transform_request_on_unprocessable_entity_error( - e=e, request_data=data - ) - ) + data = provider_config.transform_request_on_unprocessable_entity_error(e=e, request_data=data) continue else: raise self._handle_error(e=e, provider_config=provider_config) @@ -341,9 +311,7 @@ class BaseLLMAIOHTTPHandler: model=model, provider=litellm.LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Provider config not found for model: {model} and provider: {custom_llm_provider}" - ) + raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}") # get config from model, custom llm provider headers = provider_config.validate_environment( api_key=api_key, @@ -399,11 +367,7 @@ class BaseLLMAIOHTTPHandler: optional_params=optional_params, litellm_params=litellm_params, encoding=encoding, - client=( - client - if client is not None and isinstance(client, ClientSession) - else None - ), + client=(client if client is not None and isinstance(client, ClientSession) else None), ) if stream is True: @@ -419,11 +383,7 @@ class BaseLLMAIOHTTPHandler: logging_obj=logging_obj, timeout=timeout, fake_stream=fake_stream, - client=( - client - if client is not None and isinstance(client, HTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, HTTPHandler) else None), litellm_params=litellm_params, ) return CustomStreamWrapper( @@ -602,9 +562,7 @@ class BaseLLMAIOHTTPHandler: ) if provider_config is None: - raise ValueError( - f"image variation provider not found: {custom_llm_provider}." - ) + raise ValueError(f"image variation provider not found: {custom_llm_provider}.") api_base = provider_config.get_complete_url( api_base=api_base, diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b97a59a93a6..3172d3667e1 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -83,9 +83,7 @@ class AiohttpResponseStream(httpx.AsyncByteStream): async def __aiter__(self) -> typing.AsyncIterator[bytes]: try: - async for chunk in self._aiohttp_response.content.iter_chunked( - self.CHUNK_SIZE - ): + async for chunk in self._aiohttp_response.content.iter_chunked(self.CHUNK_SIZE): yield chunk except ( aiohttp.ClientPayloadError, @@ -103,9 +101,7 @@ class AiohttpResponseStream(httpx.AsyncByteStream): # with message "Connection closed.". Treat this as a graceful # end-of-stream so downstream consumers don't error. if "Connection closed" in str(e): - verbose_logger.debug( - "Upstream closed streaming connection; ending iterator gracefully" - ) + verbose_logger.debug("Upstream closed streaming connection; ending iterator gracefully") return raise except aiohttp.http_exceptions.TransferEncodingError as e: @@ -205,11 +201,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): current_loop = asyncio.get_running_loop() # If session is from a different or closed loop, recreate it - if ( - session_loop is None - or session_loop != current_loop - or session_loop.is_closed() - ): + if session_loop is None or session_loop != current_loop or session_loop.is_closed(): # Close old session to prevent leaks old_session = self.client try: @@ -218,9 +210,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): asyncio.create_task(old_session.close()) except RuntimeError: # Different event loop - can't schedule task, rely on GC - verbose_logger.debug( - "Old session from different loop, relying on GC" - ) + verbose_logger.debug("Old session from different loop, relying on GC") except Exception as e: verbose_logger.debug(f"Error closing old session: {e}") @@ -328,9 +318,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): except RuntimeError as e: # Handle the case where session was closed between our check and actual use if "Session is closed" in str(e): - verbose_logger.debug( - f"Session closed during request, retrying with new session: {e}" - ) + verbose_logger.debug(f"Session closed during request, retrying with new session: {e}") # Force creation of a new session if hasattr(self, "_client_factory") and callable(self._client_factory): self.client = self._client_factory() @@ -361,10 +349,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): async def _get_proxy_settings(self, request: httpx.Request): proxy = None - if not ( - litellm.disable_aiohttp_trust_env - or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) - ): + if not (litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False"))): try: proxy = self._proxy_from_env(request.url) except Exception as e: # pragma: no cover - best effort diff --git a/litellm/llms/custom_httpx/async_client_cleanup.py b/litellm/llms/custom_httpx/async_client_cleanup.py index 9c1f6af7e9c..8d1ddb96053 100644 --- a/litellm/llms/custom_httpx/async_client_cleanup.py +++ b/litellm/llms/custom_httpx/async_client_cleanup.py @@ -58,9 +58,7 @@ async def close_litellm_async_clients(): # This is used by Gemini and other providers that use aiohttp if hasattr(litellm, "base_llm_aiohttp_handler"): base_handler = getattr(litellm, "base_llm_aiohttp_handler", None) - if isinstance(base_handler, BaseLLMAIOHTTPHandler) and hasattr( - base_handler, "close" - ): + if isinstance(base_handler, BaseLLMAIOHTTPHandler) and hasattr(base_handler, "close"): try: await base_handler.close() except Exception: diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index 501390d840b..7d6a25bc090 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -209,9 +209,7 @@ class GenericContainerHandler: # Get HTTP client if client is None or not isinstance(client, HTTPHandler): - http_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + http_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: http_client = client @@ -229,15 +227,11 @@ class GenericContainerHandler: ) # Build URL with path params - path_params = { - p: kwargs.get(p, "") for p in endpoint_config.get("path_params", []) - } + path_params = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} url = _build_url(api_base, endpoint_config["path"], path_params) # Build query params - query_params = _build_query_params( - endpoint_config.get("query_params", []), kwargs - ) + query_params = _build_query_params(endpoint_config.get("query_params", []), kwargs) if extra_query: query_params.update(extra_query) @@ -264,25 +258,15 @@ class GenericContainerHandler: try: if method == "GET": - response = http_client.get( - url=url, headers=headers, params=effective_params - ) + response = http_client.get(url=url, headers=headers, params=effective_params) elif method == "DELETE": - response = http_client.delete( - url=url, headers=headers, params=effective_params - ) + response = http_client.delete(url=url, headers=headers, params=effective_params) elif method == "POST": if is_multipart and "file" in kwargs: - files, headers = _prepare_multipart_file_upload( - kwargs["file"], headers - ) - response = http_client.post( - url=url, headers=headers, params=effective_params, files=files - ) + files, headers = _prepare_multipart_file_upload(kwargs["file"], headers) + response = http_client.post(url=url, headers=headers, params=effective_params, files=files) else: - response = http_client.post( - url=url, headers=headers, params=effective_params - ) + response = http_client.post(url=url, headers=headers, params=effective_params) else: raise ValueError(f"Unsupported HTTP method: {method}") @@ -295,9 +279,7 @@ class GenericContainerHandler: if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg = response_json.get("error", {}).get( - "message", str(response_json) - ) + error_msg = response_json.get("error", {}).get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, @@ -353,15 +335,11 @@ class GenericContainerHandler: ) # Build URL with path params - path_params = { - p: kwargs.get(p, "") for p in endpoint_config.get("path_params", []) - } + path_params = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} url = _build_url(api_base, endpoint_config["path"], path_params) # Build query params - query_params = _build_query_params( - endpoint_config.get("query_params", []), kwargs - ) + query_params = _build_query_params(endpoint_config.get("query_params", []), kwargs) if extra_query: query_params.update(extra_query) @@ -388,25 +366,15 @@ class GenericContainerHandler: try: if method == "GET": - response = await http_client.get( - url=url, headers=headers, params=effective_params - ) + response = await http_client.get(url=url, headers=headers, params=effective_params) elif method == "DELETE": - response = await http_client.delete( - url=url, headers=headers, params=effective_params - ) + response = await http_client.delete(url=url, headers=headers, params=effective_params) elif method == "POST": if is_multipart and "file" in kwargs: - files, headers = _prepare_multipart_file_upload( - kwargs["file"], headers - ) - response = await http_client.post( - url=url, headers=headers, params=effective_params, files=files - ) + files, headers = _prepare_multipart_file_upload(kwargs["file"], headers) + response = await http_client.post(url=url, headers=headers, params=effective_params, files=files) else: - response = await http_client.post( - url=url, headers=headers, params=effective_params - ) + response = await http_client.post(url=url, headers=headers, params=effective_params) else: raise ValueError(f"Unsupported HTTP method: {method}") @@ -419,9 +387,7 @@ class GenericContainerHandler: if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg = response_json.get("error", {}).get( - "message", str(response_json) - ) + error_msg = response_json.get("error", {}).get("message", str(response_json)) raise BaseLLMException( status_code=response.status_code, message=error_msg, diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 33601b47609..5cec763bb5d 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -67,14 +67,10 @@ except Exception: # aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older # versions don't — detect once and skip the keep-alive wiring there. # https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector -_AIOHTTP_SUPPORTS_SOCKET_FACTORY = ( - "socket_factory" in inspect.signature(TCPConnector.__init__).parameters -) +_AIOHTTP_SUPPORTS_SOCKET_FACTORY = "socket_factory" in inspect.signature(TCPConnector.__init__).parameters -def _build_aiohttp_keepalive_socket_factory() -> Optional[ - Callable[[Tuple[Any, ...]], socket.socket] -]: +def _build_aiohttp_keepalive_socket_factory() -> Optional[Callable[[Tuple[Any, ...]], socket.socket]]: """ Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets. @@ -97,17 +93,11 @@ def _build_aiohttp_keepalive_socket_factory() -> Optional[ # Linux: TCP_KEEPIDLE is idle-before-first-probe. # macOS/Darwin: TCP_KEEPALIVE is the equivalent. if hasattr(socket, "TCP_KEEPIDLE"): - sock.setsockopt( - socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE - ) + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE) elif hasattr(socket, "TCP_KEEPALIVE"): - sock.setsockopt( - socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE - ) + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE) if hasattr(socket, "TCP_KEEPINTVL"): - sock.setsockopt( - socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL - ) + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL) if hasattr(socket, "TCP_KEEPCNT"): sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPCNT) return sock @@ -144,9 +134,7 @@ def _default_cached_client_timeout() -> httpx.Timeout: configured = get_configured_request_timeout() if configured is None: return _DEFAULT_TIMEOUT - return httpx.Timeout( - timeout=configured, connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS - ) + return httpx.Timeout(timeout=configured, connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS) _STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS = 5.0 @@ -199,9 +187,7 @@ def _prepare_request_data_and_content( # Cache for SSL contexts to avoid creating duplicate contexts with the same configuration # Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve) # Value: ssl.SSLContext -_ssl_context_cache: Dict[ - Tuple[Optional[str], Optional[str], Optional[str]], ssl.SSLContext -] = {} +_ssl_context_cache: Dict[Tuple[Optional[str], Optional[str], Optional[str]], ssl.SSLContext] = {} def _create_ssl_context( @@ -388,11 +374,7 @@ def mask_sensitive_info(error_message): masked_message = error_message[: key_index + 4] + "[REDACTED_API_KEY]" else: # Replace the key with redacted value, keeping other parameters - masked_message = ( - error_message[: key_index + 4] - + "[REDACTED_API_KEY]" - + error_message[next_param:] - ) + masked_message = error_message[: key_index + 4] + "[REDACTED_API_KEY]" + error_message[next_param:] return masked_message @@ -407,9 +389,7 @@ def _safe_get_response_text(response: httpx.Response) -> str: return "" -async def _safe_aread_response( - response: httpx.Response, timeout: Optional[float] = None -) -> bytes: +async def _safe_aread_response(response: httpx.Response, timeout: Optional[float] = None) -> bytes: """Safely read async response body, falling back to empty bytes on errors.""" try: if timeout is not None: @@ -419,9 +399,7 @@ async def _safe_aread_response( return b"" -def _safe_read_response( - response: httpx.Response, timeout: Optional[float] = None -) -> bytes: +def _safe_read_response(response: httpx.Response, timeout: Optional[float] = None) -> bytes: """Safely read sync response body, falling back to empty bytes on errors.""" try: if timeout is not None: @@ -477,9 +455,7 @@ async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> N class MaskedHTTPStatusError(httpx.HTTPStatusError): - def __init__( - self, original_error, message: Optional[str] = None, text: Optional[str] = None - ): + def __init__(self, original_error, message: Optional[str] = None, text: Optional[str] = None): # Create a new error with the masked URL masked_url = mask_sensitive_info(str(original_error.request.url)) # Mask the original exception message too (it contains the full URL) @@ -607,9 +583,7 @@ class AsyncHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, ): # Set follow_redirects to UseClientDefault if None - _follow_redirects = ( - follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT - ) + _follow_redirects = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT params = params or {} params.update(HTTPHandler.extract_query_params(url)) @@ -643,9 +617,7 @@ class AsyncHTTPHandler: timeout = self.timeout # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) req = self.client.build_request( "POST", @@ -663,9 +635,7 @@ class AsyncHTTPHandler: return response except (httpx.RemoteProtocolError, httpx.ConnectError): # Retry the request with a new session if there is a connection error - new_client = self.create_client( - timeout=timeout, event_hooks=self.event_hooks - ) + new_client = self.create_client(timeout=timeout, event_hooks=self.event_hooks) try: return await self.single_connection_post_request( url=url, @@ -714,9 +684,7 @@ class AsyncHTTPHandler: timeout = self.timeout # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) req = self.client.build_request( "PUT", @@ -733,9 +701,7 @@ class AsyncHTTPHandler: return response except (httpx.RemoteProtocolError, httpx.ConnectError): # Retry the request with a new session if there is a connection error - new_client = self.create_client( - timeout=timeout, event_hooks=self.event_hooks - ) + new_client = self.create_client(timeout=timeout, event_hooks=self.event_hooks) try: return await self.single_connection_post_request( url=url, @@ -782,9 +748,7 @@ class AsyncHTTPHandler: timeout = self.timeout # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) req = self.client.build_request( "PATCH", @@ -801,9 +765,7 @@ class AsyncHTTPHandler: return response except (httpx.RemoteProtocolError, httpx.ConnectError): # Retry the request with a new session if there is a connection error - new_client = self.create_client( - timeout=timeout, event_hooks=self.event_hooks - ) + new_client = self.create_client(timeout=timeout, event_hooks=self.event_hooks) try: return await self.single_connection_post_request( url=url, @@ -850,9 +812,7 @@ class AsyncHTTPHandler: timeout = self.timeout # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) req = self.client.build_request( "DELETE", @@ -869,9 +829,7 @@ class AsyncHTTPHandler: return response except (httpx.RemoteProtocolError, httpx.ConnectError): # Retry the request with a new session if there is a connection error - new_client = self.create_client( - timeout=timeout, event_hooks=self.event_hooks - ) + new_client = self.create_client(timeout=timeout, event_hooks=self.event_hooks) try: return await self.single_connection_post_request( url=url, @@ -1034,9 +992,7 @@ class AsyncHTTPHandler: from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from litellm.secret_managers.main import str_to_bool - connector_kwargs = AsyncHTTPHandler._get_ssl_connector_kwargs( - ssl_verify=ssl_verify, ssl_context=ssl_context - ) + connector_kwargs = AsyncHTTPHandler._get_ssl_connector_kwargs(ssl_verify=ssl_verify, ssl_context=ssl_context) ######################################################### # Check if user enabled aiohttp trust env # use for HTTP_PROXY, HTTPS_PROXY, etc. @@ -1059,9 +1015,7 @@ class AsyncHTTPHandler: # Use shared session if provided and valid if shared_session is not None and not shared_session.closed: - verbose_logger.debug( - f"SHARED SESSION: Reusing existing ClientSession (ID: {id(shared_session)})" - ) + verbose_logger.debug(f"SHARED SESSION: Reusing existing ClientSession (ID: {id(shared_session)})") return LiteLLMAiohttpTransport( client=shared_session, ssl_verify=ssl_for_transport, @@ -1069,9 +1023,7 @@ class AsyncHTTPHandler: ) # Create new session only if none provided or existing one is invalid - verbose_logger.debug( - "NEW SESSION: Creating new ClientSession (no shared session provided)" - ) + verbose_logger.debug("NEW SESSION: Creating new ClientSession (no shared session provided)") transport_connector_kwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, @@ -1082,9 +1034,7 @@ class AsyncHTTPHandler: if AIOHTTP_CONNECTOR_LIMIT > 0: transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: - transport_connector_kwargs["limit_per_host"] = ( - AIOHTTP_CONNECTOR_LIMIT_PER_HOST - ) + transport_connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST # Returns None when SO_KEEPALIVE is disabled or aiohttp is too old to # accept socket_factory — version detection lives inside the builder. socket_factory = _build_aiohttp_keepalive_socket_factory() @@ -1165,9 +1115,7 @@ class HTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, ): # Set follow_redirects to UseClientDefault if None - _follow_redirects = ( - follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT - ) + _follow_redirects = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT params = params or {} params.update(self.extract_query_params(url)) @@ -1209,9 +1157,7 @@ class HTTPHandler: ): try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) if timeout is not None: req = self.client.build_request( @@ -1263,9 +1209,7 @@ class HTTPHandler: ): try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) if timeout is not None: req = self.client.build_request( @@ -1315,9 +1259,7 @@ class HTTPHandler: ): try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) if timeout is not None: req = self.client.build_request( @@ -1366,9 +1308,7 @@ class HTTPHandler: ): try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) - request_data, request_content = _prepare_request_data_and_content( - data, content - ) + request_data, request_content = _prepare_request_data_and_content(data, content) if timeout is not None: req = self.client.build_request( @@ -1460,9 +1400,7 @@ def get_async_httpx_client( if params is not None: # Filter out params that are only used for cache key, not for AsyncHTTPHandler.__init__ - handler_params = { - k: v for k, v in params.items() if k != "disable_aiohttp_transport" - } + handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"} handler_params["shared_session"] = shared_session _new_client = AsyncHTTPHandler(**handler_params) else: @@ -1511,9 +1449,7 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: if params is not None: # Filter out params that are only used for cache key, not for HTTPHandler.__init__ - handler_params = { - k: v for k, v in params.items() if k != "disable_aiohttp_transport" - } + handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"} _new_client = HTTPHandler(**handler_params) else: _new_client = HTTPHandler(timeout=_default_cached_client_timeout()) diff --git a/litellm/llms/custom_httpx/httpx_handler.py b/litellm/llms/custom_httpx/httpx_handler.py index 6e9f29151cd..a66d30c9007 100644 --- a/litellm/llms/custom_httpx/httpx_handler.py +++ b/litellm/llms/custom_httpx/httpx_handler.py @@ -39,9 +39,7 @@ class HTTPHandler: # Close the client when you're done with it await self.client.aclose() - async def get( - self, url: str, params: Optional[dict] = None, headers: Optional[dict] = None - ): + async def get(self, url: str, params: Optional[dict] = None, headers: Optional[dict] = None): response = await self.client.get(url, params=params, headers=headers) return response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 056d3a4b4d5..d8f4856b4ad 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -177,9 +177,7 @@ def _google_genai_streaming_hidden_params( """Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params).""" from litellm.litellm_core_utils.core_helpers import process_response_headers - _model_info: Dict[str, Any] = dict( - getattr(litellm_params, "model_info", None) or {} - ) + _model_info: Dict[str, Any] = dict(getattr(litellm_params, "model_info", None) or {}) _raw_id = _model_info.get("id") or logging_obj.get_router_model_id() or "" _model_id = _raw_id if isinstance(_raw_id, str) else str(_raw_id) return { @@ -225,9 +223,7 @@ def _has_pre_call_deployment_hook(logging_obj: Any) -> bool: base_func = CustomLogger.async_pre_call_deployment_hook for cb in _custom_logger_callbacks(logging_obj): cb_func = getattr(type(cb), "async_pre_call_deployment_hook", base_func) - if getattr(cb_func, "__func__", cb_func) is not getattr( - base_func, "__func__", base_func - ): + if getattr(cb_func, "__func__", cb_func) is not getattr(base_func, "__func__", base_func): return True return False @@ -247,9 +243,7 @@ class BaseLLMHTTPHandler: signed_json_body: Optional[bytes] = None, ) -> httpx.Response: """Common implementation across stream + non-stream calls. Meant to ensure consistent error-handling.""" - max_retry_on_unprocessable_entity_error = ( - provider_config.max_retry_on_unprocessable_entity_error - ) + max_retry_on_unprocessable_entity_error = provider_config.max_retry_on_unprocessable_entity_error response: Optional[httpx.Response] = None for i in range(max(max_retry_on_unprocessable_entity_error, 1)): @@ -257,11 +251,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.post( url=api_base, headers=headers, - data=( - signed_json_body - if signed_json_body is not None - else json.dumps(data) - ), + data=(signed_json_body if signed_json_body is not None else json.dumps(data)), timeout=timeout, stream=stream, logging_obj=logging_obj, @@ -272,11 +262,7 @@ class BaseLLMHTTPHandler: e=e, litellm_params=litellm_params ) if should_retry and not hit_max_retry: - data = ( - provider_config.transform_request_on_unprocessable_entity_error( - e=e, request_data=data - ) - ) + data = provider_config.transform_request_on_unprocessable_entity_error(e=e, request_data=data) continue else: raise self._handle_error(e=e, provider_config=provider_config) @@ -306,9 +292,7 @@ class BaseLLMHTTPHandler: stream: bool = False, signed_json_body: Optional[bytes] = None, ) -> httpx.Response: - max_retry_on_unprocessable_entity_error = ( - provider_config.max_retry_on_unprocessable_entity_error - ) + max_retry_on_unprocessable_entity_error = provider_config.max_retry_on_unprocessable_entity_error response: Optional[httpx.Response] = None @@ -317,11 +301,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.post( url=api_base, headers=headers, - data=( - signed_json_body - if signed_json_body is not None - else json.dumps(data) - ), + data=(signed_json_body if signed_json_body is not None else json.dumps(data)), timeout=timeout, stream=stream, logging_obj=logging_obj, @@ -332,11 +312,7 @@ class BaseLLMHTTPHandler: e=e, litellm_params=litellm_params ) if should_retry and not hit_max_retry: - data = ( - provider_config.transform_request_on_unprocessable_entity_error( - e=e, request_data=data - ) - ) + data = provider_config.transform_request_on_unprocessable_entity_error(e=e, request_data=data) continue else: raise self._handle_error(e=e, provider_config=provider_config) @@ -450,23 +426,16 @@ class BaseLLMHTTPHandler: json_mode: bool = optional_params.pop("json_mode", False) extra_body: Optional[dict] = optional_params.pop("extra_body", None) - provider_config = ( - provider_config - or ProviderConfigManager.get_provider_chat_config( - model=model, provider=litellm.LlmProviders(custom_llm_provider) - ) + provider_config = provider_config or ProviderConfigManager.get_provider_chat_config( + model=model, provider=litellm.LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Provider config not found for model: {model} and provider: {custom_llm_provider}" - ) + raise ValueError(f"Provider config not found for model: {model} and provider: {custom_llm_provider}") fake_stream = ( fake_stream or optional_params.pop("fake_stream", False) - or provider_config.should_fake_stream( - model=model, custom_llm_provider=custom_llm_provider, stream=stream - ) + or provider_config.should_fake_stream(model=model, custom_llm_provider=custom_llm_provider, stream=stream) ) # get config from model, custom llm provider @@ -525,9 +494,7 @@ class BaseLLMHTTPHandler: # Check if stream was converted for WebSearch interception # This is set by the async_pre_request_hook in WebSearchInterceptionLogger if litellm_params.get("_websearch_interception_converted_stream", False): - logging_obj.model_call_details[ - "websearch_interception_converted_stream" - ] = True + logging_obj.model_call_details["websearch_interception_converted_stream"] = True if acompletion is True: if stream is True: @@ -547,11 +514,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, data=data, fake_stream=fake_stream, - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), litellm_params=litellm_params, json_mode=json_mode, optional_params=optional_params, @@ -574,11 +537,7 @@ class BaseLLMHTTPHandler: optional_params=optional_params, litellm_params=litellm_params, encoding=encoding, - client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), json_mode=json_mode, signed_json_body=signed_json_body, shared_session=shared_session, @@ -615,11 +574,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, timeout=timeout, fake_stream=fake_stream, - client=( - client - if client is not None and isinstance(client, HTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, HTTPHandler) else None), litellm_params=litellm_params, json_mode=json_mode, optional_params=optional_params, @@ -720,9 +675,7 @@ class BaseLLMHTTPHandler: json_mode=json_mode, ) - completion_stream: Any = MockResponseIterator( - model_response=model_response, json_mode=json_mode - ) + completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator( streaming_response=response.iter_lines(), @@ -858,9 +811,7 @@ class BaseLLMHTTPHandler: json_mode=json_mode, ) - completion_stream: Any = MockResponseIterator( - model_response=model_response, json_mode=json_mode - ) + completion_stream: Any = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator( streaming_response=response.aiter_lines(), sync_stream=False @@ -916,9 +867,7 @@ class BaseLLMHTTPHandler: model=model, provider=litellm.LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Provider {custom_llm_provider} does not support embedding" - ) + raise ValueError(f"Provider {custom_llm_provider} does not support embedding") # get config from model, custom llm provider headers = provider_config.validate_environment( api_key=api_key, @@ -986,9 +935,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -1183,9 +1130,7 @@ class BaseLLMHTTPHandler: client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, ) -> RerankResponse: if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider) - ) + async_httpx_client = get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider)) else: async_httpx_client = client try: @@ -1255,9 +1200,7 @@ class BaseLLMHTTPHandler: # All providers now return AudioTranscriptionRequestData if not isinstance(transformed_result, AudioTranscriptionRequestData): - raise ValueError( - f"Provider {provider_config.__class__.__name__} must return AudioTranscriptionRequestData" - ) + raise ValueError(f"Provider {provider_config.__class__.__name__} must return AudioTranscriptionRequestData") data = transformed_result.data files = transformed_result.files @@ -1312,9 +1255,7 @@ class BaseLLMHTTPHandler: shared_session: Optional["ClientSession"] = None, ) -> Union[TranscriptionResponse, Coroutine[Any, Any, TranscriptionResponse]]: if provider_config is None: - raise ValueError( - f"No provider config found for model: {model} and provider: {custom_llm_provider}" - ) + raise ValueError(f"No provider config found for model: {model} and provider: {custom_llm_provider}") if atranscription is True: return self.async_audio_transcriptions( # type: ignore @@ -1400,9 +1341,7 @@ class BaseLLMHTTPHandler: shared_session: Optional["ClientSession"] = None, ) -> TranscriptionResponse: if provider_config is None: - raise ValueError( - f"No provider config found for model: {model} and provider: {custom_llm_provider}" - ) + raise ValueError(f"No provider config found for model: {model} and provider: {custom_llm_provider}") # Prepare the request ( @@ -1501,15 +1440,11 @@ class BaseLLMHTTPHandler: # All providers return OCRRequestData if not isinstance(transformed_result, OCRRequestData): - raise ValueError( - f"Provider {provider_config.__class__.__name__} must return OCRRequestData" - ) + raise ValueError(f"Provider {provider_config.__class__.__name__} must return OCRRequestData") # Data is always a dict for Mistral OCR format if not isinstance(transformed_result.data, dict): - raise ValueError( - f"Expected dict data for OCR request, got {type(transformed_result.data)}" - ) + raise ValueError(f"Expected dict data for OCR request, got {type(transformed_result.data)}") data = transformed_result.data @@ -1571,15 +1506,11 @@ class BaseLLMHTTPHandler: # All providers return OCRRequestData if not isinstance(transformed_result, OCRRequestData): - raise ValueError( - f"Provider {provider_config.__class__.__name__} must return OCRRequestData" - ) + raise ValueError(f"Provider {provider_config.__class__.__name__} must return OCRRequestData") # Data is always a dict for Mistral OCR format if not isinstance(transformed_result.data, dict): - raise ValueError( - f"Expected dict data for OCR request, got {type(transformed_result.data)}" - ) + raise ValueError(f"Expected dict data for OCR request, got {type(transformed_result.data)}") data = transformed_result.data @@ -1630,9 +1561,7 @@ class BaseLLMHTTPHandler: Sync OCR handler. """ if provider_config is None: - raise ValueError( - f"No provider config found for model: {model} and provider: {custom_llm_provider}" - ) + raise ValueError(f"No provider config found for model: {model} and provider: {custom_llm_provider}") if litellm_params is None: litellm_params = {} @@ -1706,9 +1635,7 @@ class BaseLLMHTTPHandler: Async OCR handler. """ if provider_config is None: - raise ValueError( - f"No provider config found for model: {model} and provider: {custom_llm_provider}" - ) + raise ValueError(f"No provider config found for model: {model} and provider: {custom_llm_provider}") if litellm_params is None: litellm_params = {} @@ -1769,9 +1696,7 @@ class BaseLLMHTTPHandler: Sync Search handler. """ if provider_config is None: - raise ValueError( - f"No provider config found for provider: {custom_llm_provider}" - ) + raise ValueError(f"No provider config found for provider: {custom_llm_provider}") if asearch is True: return self.async_search( @@ -1866,9 +1791,7 @@ class BaseLLMHTTPHandler: Async Search handler. """ if provider_config is None: - raise ValueError( - f"No provider config found for provider: {custom_llm_provider}" - ) + raise ValueError(f"No provider config found for provider: {custom_llm_provider}") # Validate environment and get headers headers = provider_config.validate_environment( @@ -1909,9 +1832,7 @@ class BaseLLMHTTPHandler: # For search providers, use special Search provider type from litellm.types.llms.custom_http import httpxSpecialProvider - async_httpx_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.Search - ) + async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Search) else: async_httpx_client = client @@ -1958,9 +1879,7 @@ class BaseLLMHTTPHandler: api_key: Optional[str], model: str, ) -> httpx.Response: - max_attempts = max( - provider_config.max_retry_on_anthropic_messages_http_error, 1 - ) + max_attempts = max(provider_config.max_retry_on_anthropic_messages_http_error, 1) litellm_params_dict = dict(litellm_params) optional_params_dict = dict(litellm_params) for attempt_idx in range(max_attempts): @@ -1976,10 +1895,8 @@ class BaseLLMHTTPHandler: return response except httpx.HTTPStatusError as e: hit_max_attempt = attempt_idx + 1 == max_attempts - should_retry = ( - provider_config.should_retry_anthropic_messages_on_http_error( - e=e, litellm_params=litellm_params_dict - ) + should_retry = provider_config.should_retry_anthropic_messages_on_http_error( + e=e, litellm_params=litellm_params_dict ) if should_retry and not hit_max_attempt: verbose_logger.debug( @@ -1988,9 +1905,7 @@ class BaseLLMHTTPHandler: attempt_idx + 2, max_attempts, ) - provider_config.transform_anthropic_messages_request_on_http_error( - e=e, request_data=request_body - ) + provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body) headers, signed_json_body = provider_config.sign_request( headers=headers, optional_params=optional_params_dict, @@ -2007,9 +1922,7 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) - raise RuntimeError( - "unreachable: anthropic messages HTTP retry loop exited without return" - ) + raise RuntimeError("unreachable: anthropic messages HTTP retry loop exited without return") async def async_anthropic_messages_handler( self, @@ -2032,9 +1945,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.ANTHROPIC - ) + async_httpx_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) else: async_httpx_client = client @@ -2044,11 +1955,9 @@ class BaseLLMHTTPHandler: Optional[litellm.types.utils.ProviderSpecificHeader], kwargs.get("provider_specific_header", None), ) - provider_specific_headers = ( - ProviderSpecificHeaderUtils.get_provider_specific_headers( - provider_specific_header=provider_specific_header, - custom_llm_provider=custom_llm_provider, - ) + provider_specific_headers = ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=provider_specific_header, + custom_llm_provider=custom_llm_provider, ) forwarded_headers = kwargs.get("headers", None) # Also check for extra_headers in kwargs (from config or direct calls) @@ -2074,9 +1983,7 @@ class BaseLLMHTTPHandler: api_base=api_base, ) - headers = update_headers_with_filtered_beta( - headers=headers, provider=custom_llm_provider - ) + headers = update_headers_with_filtered_beta(headers=headers, provider=custom_llm_provider) logging_obj.update_from_kwargs( kwargs=kwargs, @@ -2130,9 +2037,7 @@ class BaseLLMHTTPHandler: headers, signed_json_body = anthropic_messages_provider_config.sign_request( headers=headers, - optional_params=dict( - litellm_params - ), # dynamic aws_* params are passed under litellm_params + optional_params=dict(litellm_params), # dynamic aws_* params are passed under litellm_params request_data=request_body, api_base=request_url, api_key=api_key, @@ -2165,9 +2070,7 @@ class BaseLLMHTTPHandler: async_httpx_client=async_httpx_client, request_url=request_url, headers=headers, - signed_json_body=( - signed_json_body if signed_json_body is not None else request_body_json - ), + signed_json_body=(signed_json_body if signed_json_body is not None else request_body_json), request_body=request_body, stream=stream or False, logging_obj=logging_obj, @@ -2327,11 +2230,7 @@ class BaseLLMHTTPHandler: optional_param_names = _responses_api_optional_request_param_names() updated_response_params = { **response_api_optional_request_params, - **{ - key: value - for key, value in modified_kwargs.items() - if key in optional_param_names - }, + **{key: value for key, value in modified_kwargs.items() if key in optional_param_names}, } updated_litellm_params = GenericLiteLLMParams( **{ @@ -2339,8 +2238,7 @@ class BaseLLMHTTPHandler: **{ key: value for key, value in modified_kwargs.items() - if key not in optional_param_names - and key not in {"model", "input", "custom_llm_provider"} + if key not in optional_param_names and key not in {"model", "input", "custom_llm_provider"} }, } ) @@ -2379,9 +2277,7 @@ class BaseLLMHTTPHandler: ) -> Union[ ResponsesAPIResponse, BaseResponsesAPIStreamingIterator, - Coroutine[ - Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] - ], + Coroutine[Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]], ]: """ Handles responses API requests. @@ -2427,9 +2323,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -2496,9 +2390,7 @@ class BaseLLMHTTPHandler: stream=stream, fake_stream=fake_stream, ) - body_kwargs: Dict[str, Any] = ( - {"data": signed_body} if signed_body is not None else {"json": data} - ) + body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -2516,8 +2408,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.post( url=api_base, headers=headers, - timeout=timeout - or float(response_api_optional_request_params.get("timeout", 0)), + timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, **body_kwargs, ) @@ -2547,8 +2438,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.post( url=api_base, headers=headers, - timeout=timeout - or float(response_api_optional_request_params.get("timeout", 0)), + timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), **body_kwargs, ) except Exception as e: @@ -2557,12 +2447,10 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) - initial_response = ( - responses_api_provider_config.transform_response_api_response( - model=model, - raw_response=response, - logging_obj=logging_obj, - ) + initial_response = responses_api_provider_config.transform_response_api_response( + model=model, + raw_response=response, + logging_obj=logging_obj, ) if self._has_agentic_completion_hook(logging_obj): @@ -2570,11 +2458,7 @@ class BaseLLMHTTPHandler: self._call_agentic_completion_hooks, response=initial_response, model=model, - messages=( - input - if isinstance(input, list) - else [{"role": "user", "content": input}] - ), + messages=(input if isinstance(input, list) else [{"role": "user", "content": input}]), anthropic_messages_provider_config=responses_api_provider_config, anthropic_messages_optional_request_params=response_api_optional_request_params, logging_obj=logging_obj, @@ -2680,9 +2564,7 @@ class BaseLLMHTTPHandler: stream=stream, fake_stream=fake_stream, ) - body_kwargs: Dict[str, Any] = ( - {"data": signed_body} if signed_body is not None else {"json": data} - ) + body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -2700,8 +2582,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.post( url=api_base, headers=headers, - timeout=timeout - or float(response_api_optional_request_params.get("timeout", 0)), + timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, **body_kwargs, ) @@ -2733,8 +2614,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.post( url=api_base, headers=headers, - timeout=timeout - or float(response_api_optional_request_params.get("timeout", 0)), + timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), **body_kwargs, ) @@ -2744,22 +2624,16 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) - initial_response = ( - responses_api_provider_config.transform_response_api_response( - model=model, - raw_response=response, - logging_obj=logging_obj, - ) + initial_response = responses_api_provider_config.transform_response_api_response( + model=model, + raw_response=response, + logging_obj=logging_obj, ) final_response = await self._call_agentic_completion_hooks( response=initial_response, model=model, - messages=( - input - if isinstance(input, list) - else [{"role": "user", "content": input}] - ), + messages=(input if isinstance(input, list) else [{"role": "user", "content": input}]), anthropic_messages_provider_config=responses_api_provider_config, anthropic_messages_optional_request_params=response_api_optional_request_params, logging_obj=logging_obj, @@ -2770,9 +2644,9 @@ class BaseLLMHTTPHandler: ) result = final_response if final_response is not None else initial_response - if litellm_params.get( - "_code_interpreter_interception_converted_stream" - ) and not litellm_params.get("_agentic_loop_depth"): + if litellm_params.get("_code_interpreter_interception_converted_stream") and not litellm_params.get( + "_agentic_loop_depth" + ): return self._wrap_responses_response_as_fake_stream( result=result, model=model, @@ -2898,9 +2772,7 @@ class BaseLLMHTTPHandler: shared_session=shared_session, ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -2991,9 +2863,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -3099,9 +2969,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=data - ) + response = await async_httpx_client.get(url=url, headers=headers, params=data) except Exception as e: verbose_logger.exception(f"Error retrieving response: {e}") @@ -3155,9 +3023,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -3269,9 +3135,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=responses_api_provider_config) @@ -3374,10 +3238,7 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client - if ( - isinstance(transformed_request, dict) - and "initial_request" in transformed_request - ): + if isinstance(transformed_request, dict) and "initial_request" in transformed_request: # Handle two-step uploads (TwoStepFileUploadConfig) # Used by providers like Manus, Google Cloud Storage try: @@ -3398,21 +3259,15 @@ class BaseLLMHTTPHandler: initial_response_data, ) = self._extract_upload_url_from_response( response=initial_response, - upload_url_location=transformed_request.get( - "upload_url_location", "headers" - ), - upload_url_key=transformed_request.get( - "upload_url_key", "upload_url" - ), + upload_url_location=transformed_request.get("upload_url_location", "headers"), + upload_url_key=transformed_request.get("upload_url_key", "upload_url"), ) if not upload_url: raise ValueError("Failed to get upload URL from initial request") # Step 2: Upload the actual file - upload_method = ( - transformed_request["upload_request"].get("method", "POST").lower() - ) + upload_method = transformed_request["upload_request"].get("method", "POST").lower() upload_response = getattr(sync_httpx_client, upload_method)( url=upload_url, headers=transformed_request["upload_request"]["headers"], @@ -3436,34 +3291,25 @@ class BaseLLMHTTPHandler: # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig presigned_request = cast(Dict[str, Any], transformed_request) - upload_response = getattr( - sync_httpx_client, presigned_request["method"].lower() - )( + upload_response = getattr(sync_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], data=presigned_request["data"], timeout=timeout, ) - elif ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ): + elif isinstance(transformed_request, dict) and "resumable_chunked_upload" in transformed_request: try: upload_response = self._resumable_chunked_upload( client=sync_httpx_client, initiate_url=api_base, base_headers=headers, - config=cast(Dict[str, Any], transformed_request)[ - "resumable_chunked_upload" - ], + config=cast(Dict[str, Any], transformed_request)["resumable_chunked_upload"], timeout=timeout, ) except Exception as e: verbose_logger.exception(f"Error creating file: {e}") raise self._handle_error(e=e, provider_config=provider_config) - elif isinstance(transformed_request, str) or isinstance( - transformed_request, bytes - ): + elif isinstance(transformed_request, str) or isinstance(transformed_request, bytes): # Handle traditional file uploads # Ensure transformed_request is a string for httpx compatibility if isinstance(transformed_request, bytes): @@ -3496,9 +3342,7 @@ class BaseLLMHTTPHandler: timeout=timeout, ) else: - raise ValueError( - f"Unsupported transformed_request type: {type(transformed_request)}" - ) + raise ValueError(f"Unsupported transformed_request type: {type(transformed_request)}") # Store the upload URL in litellm_params for the transformation method # Honour the URL already set by transform_create_file_request (e.g. Bedrock pre-signed S3 uploads), @@ -3529,9 +3373,7 @@ class BaseLLMHTTPHandler: Creates a file using Gemini's two-step upload process """ if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=provider_config.custom_llm_provider - ) + async_httpx_client = get_async_httpx_client(llm_provider=provider_config.custom_llm_provider) else: async_httpx_client = client @@ -3547,8 +3389,7 @@ class BaseLLMHTTPHandler: # a placeholder instead of re-materializing the payload. "complete_input_dict": ( "" - if isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request + if isinstance(transformed_request, dict) and "resumable_chunked_upload" in transformed_request else transformed_request ), "api_base": api_base, @@ -3556,10 +3397,7 @@ class BaseLLMHTTPHandler: }, ) - if ( - isinstance(transformed_request, dict) - and "initial_request" in transformed_request - ): + if isinstance(transformed_request, dict) and "initial_request" in transformed_request: # Handle two-step uploads (TwoStepFileUploadConfig) # Used by providers like Manus, Google Cloud Storage try: @@ -3580,21 +3418,15 @@ class BaseLLMHTTPHandler: initial_response_data, ) = self._extract_upload_url_from_response( response=initial_response, - upload_url_location=transformed_request.get( - "upload_url_location", "headers" - ), - upload_url_key=transformed_request.get( - "upload_url_key", "upload_url" - ), + upload_url_location=transformed_request.get("upload_url_location", "headers"), + upload_url_key=transformed_request.get("upload_url_key", "upload_url"), ) if not upload_url: raise ValueError("Failed to get upload URL from initial request") # Step 2: Upload the actual file - upload_method = ( - transformed_request["upload_request"].get("method", "POST").lower() - ) + upload_method = transformed_request["upload_request"].get("method", "POST").lower() upload_response = await getattr(async_httpx_client, upload_method)( url=upload_url, headers=transformed_request["upload_request"]["headers"], @@ -3619,34 +3451,25 @@ class BaseLLMHTTPHandler: # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig presigned_request = cast(Dict[str, Any], transformed_request) - upload_response = await getattr( - async_httpx_client, presigned_request["method"].lower() - )( + upload_response = await getattr(async_httpx_client, presigned_request["method"].lower())( url=presigned_request["url"], headers=presigned_request["headers"], data=presigned_request["data"], timeout=timeout, ) - elif ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ): + elif isinstance(transformed_request, dict) and "resumable_chunked_upload" in transformed_request: try: upload_response = await self._aresumable_chunked_upload( client=async_httpx_client, initiate_url=api_base, base_headers=headers, - config=cast(Dict[str, Any], transformed_request)[ - "resumable_chunked_upload" - ], + config=cast(Dict[str, Any], transformed_request)["resumable_chunked_upload"], timeout=timeout, ) except Exception as e: verbose_logger.exception(f"Error creating file: {e}") raise self._handle_error(e=e, provider_config=provider_config) - elif isinstance(transformed_request, str) or isinstance( - transformed_request, bytes - ): + elif isinstance(transformed_request, str) or isinstance(transformed_request, bytes): # Handle traditional file uploads # Note: transformed_request can be bytes (for binary files like PDFs) # or str (for text files like JSONL). httpx handles both correctly. @@ -3676,9 +3499,7 @@ class BaseLLMHTTPHandler: timeout=timeout, ) else: - raise ValueError( - f"Unsupported transformed_request type: {type(transformed_request)}" - ) + raise ValueError(f"Unsupported transformed_request type: {type(transformed_request)}") return provider_config.transform_create_file_response( model=None, @@ -3691,9 +3512,7 @@ class BaseLLMHTTPHandler: _RESUMABLE_CHUNK_SIZE = 8 * 1024 * 1024 @staticmethod - def _iter_resumable_chunks( - byte_iter: Iterator[bytes], chunk_size: int - ) -> Iterator[bytes]: + def _iter_resumable_chunks(byte_iter: Iterator[bytes], chunk_size: int) -> Iterator[bytes]: """Regroup a byte stream into ``chunk_size`` pieces, yielding a final partial piece only when it is non-empty. Every full piece is exactly ``chunk_size`` bytes (kept a 256 KiB multiple for GCS) and never more than @@ -3759,9 +3578,7 @@ class BaseLLMHTTPHandler: init_resp.raise_for_status() session_url = init_resp.headers.get(session_url_header) if not session_url: - raise ValueError( - f"resumable upload: no session URL in '{session_url_header}' header" - ) + raise ValueError(f"resumable upload: no session URL in '{session_url_header}' header") offset = 0 pending: Optional[bytes] = None @@ -3803,9 +3620,7 @@ class BaseLLMHTTPHandler: **base_headers, "Content-Range": self._resumable_content_range(offset, len(data), is_final), } - req = httpx_client.build_request( - "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) - ) + req = httpx_client.build_request("PUT", url, **self._resumable_request_kwargs(headers, data, timeout)) resp = httpx_client.send(req, follow_redirects=False) resp.read() if resp.status_code not in ((200, 201) if is_final else (308,)): @@ -3841,9 +3656,7 @@ class BaseLLMHTTPHandler: init_resp.raise_for_status() session_url = init_resp.headers.get(session_url_header) if not session_url: - raise ValueError( - f"resumable upload: no session URL in '{session_url_header}' header" - ) + raise ValueError(f"resumable upload: no session URL in '{session_url_header}' header") offset = 0 pending: Optional[bytes] = None @@ -3893,9 +3706,7 @@ class BaseLLMHTTPHandler: **base_headers, "Content-Range": self._resumable_content_range(offset, len(data), is_final), } - req = httpx_client.build_request( - "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) - ) + req = httpx_client.build_request("PUT", url, **self._resumable_request_kwargs(headers, data, timeout)) resp = await httpx_client.send(req, follow_redirects=False) await resp.aread() if resp.status_code not in ((200, 201) if is_final else (308,)): @@ -3973,14 +3784,9 @@ class BaseLLMHTTPHandler: sync_httpx_client = client try: - if ( - isinstance(transformed_request, dict) - and "method" in transformed_request - ): + if isinstance(transformed_request, dict) and "method" in transformed_request: # Handle pre-signed requests (e.g., from Bedrock with AWS auth) - batch_response = getattr( - sync_httpx_client, transformed_request["method"].lower() - )( + batch_response = getattr(sync_httpx_client, transformed_request["method"].lower())( url=transformed_request["url"], headers=transformed_request["headers"], data=transformed_request["data"], @@ -4066,10 +3872,7 @@ class BaseLLMHTTPHandler: sync_httpx_client = client try: - if ( - isinstance(transformed_request, dict) - and "method" in transformed_request - ): + if isinstance(transformed_request, dict) and "method" in transformed_request: # Handle pre-signed requests (e.g., from Bedrock with AWS auth) method = transformed_request["method"].lower() request_kwargs = { @@ -4128,9 +3931,7 @@ class BaseLLMHTTPHandler: Async version of create_batch """ if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=provider_config.custom_llm_provider - ) + async_httpx_client = get_async_httpx_client(llm_provider=provider_config.custom_llm_provider) else: async_httpx_client = client @@ -4148,14 +3949,9 @@ class BaseLLMHTTPHandler: ) try: - if ( - isinstance(transformed_request, dict) - and "method" in transformed_request - ): + if isinstance(transformed_request, dict) and "method" in transformed_request: # Handle pre-signed requests (e.g., from Bedrock with AWS auth) - batch_response = await getattr( - async_httpx_client, transformed_request["method"].lower() - )( + batch_response = await getattr(async_httpx_client, transformed_request["method"].lower())( url=transformed_request["url"], headers=transformed_request["headers"], data=transformed_request["data"], @@ -4214,9 +4010,7 @@ class BaseLLMHTTPHandler: Async version of retrieve_batch """ if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=provider_config.custom_llm_provider - ) + async_httpx_client = get_async_httpx_client(llm_provider=provider_config.custom_llm_provider) else: async_httpx_client = client @@ -4235,10 +4029,7 @@ class BaseLLMHTTPHandler: ) try: - if ( - isinstance(transformed_request, dict) - and "method" in transformed_request - ): + if isinstance(transformed_request, dict) and "method" in transformed_request: # Handle pre-signed requests (e.g., from Bedrock with AWS auth) method = transformed_request["method"].lower() request_kwargs = { @@ -4250,9 +4041,7 @@ class BaseLLMHTTPHandler: if method != "get" and transformed_request.get("data") is not None: request_kwargs["data"] = transformed_request["data"] - batch_response = await getattr(async_httpx_client, method)( - **request_kwargs - ) + batch_response = await getattr(async_httpx_client, method)(**request_kwargs) elif isinstance(transformed_request, dict) and api_base: # For other providers that use JSON requests batch_response = await async_httpx_client.get( @@ -4314,9 +4103,7 @@ class BaseLLMHTTPHandler: shared_session=shared_session, ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -4351,9 +4138,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=data, timeout=timeout) except Exception as e: raise self._handle_error( @@ -4427,9 +4212,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=data, timeout=timeout) except Exception as e: raise self._handle_error( @@ -4477,9 +4260,7 @@ class BaseLLMHTTPHandler: shared_session=shared_session, ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -4516,9 +4297,7 @@ class BaseLLMHTTPHandler: api_key=litellm_params.api_key, model=model, ) - body_kwargs: Dict[str, Any] = ( - {"data": signed_body} if signed_body is not None else {"json": data} - ) + body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -4532,9 +4311,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, timeout=timeout, **body_kwargs - ) + response = sync_httpx_client.post(url=url, headers=headers, timeout=timeout, **body_kwargs) except Exception as e: raise self._handle_error( @@ -4611,9 +4388,7 @@ class BaseLLMHTTPHandler: api_key=litellm_params.api_key, model=model, ) - body_kwargs: Dict[str, Any] = ( - {"data": signed_body} if signed_body is not None else {"json": data} - ) + body_kwargs: Dict[str, Any] = {"data": signed_body} if signed_body is not None else {"json": data} ## LOGGING logging_obj.pre_call( @@ -4627,9 +4402,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, timeout=timeout, **body_kwargs - ) + response = await async_httpx_client.post(url=url, headers=headers, timeout=timeout, **body_kwargs) except Exception as e: raise self._handle_error( @@ -4724,9 +4497,7 @@ class BaseLLMHTTPHandler: Async retrieve file metadata by ID """ if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=provider_config.custom_llm_provider - ) + async_httpx_client = get_async_httpx_client(llm_provider=provider_config.custom_llm_provider) else: async_httpx_client = client @@ -4758,9 +4529,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -4852,9 +4621,7 @@ class BaseLLMHTTPHandler: Async delete a file by ID """ if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=provider_config.custom_llm_provider - ) + async_httpx_client = get_async_httpx_client(llm_provider=provider_config.custom_llm_provider) else: async_httpx_client = client @@ -4886,9 +4653,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.delete( - url=url, headers=headers, params=params, timeout=timeout - ) + response = await async_httpx_client.delete(url=url, headers=headers, params=params, timeout=timeout) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -4980,9 +4745,7 @@ class BaseLLMHTTPHandler: Async list all files """ if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=provider_config.custom_llm_provider - ) + async_httpx_client = get_async_httpx_client(llm_provider=provider_config.custom_llm_provider) else: async_httpx_client = client @@ -5014,9 +4777,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -5036,9 +4797,7 @@ class BaseLLMHTTPHandler: _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union[ - "HttpxBinaryResponseContent", Coroutine[Any, Any, "HttpxBinaryResponseContent"] - ]: + ) -> Union["HttpxBinaryResponseContent", Coroutine[Any, Any, "HttpxBinaryResponseContent"]]: """ Retrieve file content by ID """ @@ -5110,9 +4869,7 @@ class BaseLLMHTTPHandler: Async retrieve file content by ID """ if client is None or not isinstance(client, AsyncHTTPHandler): - async_httpx_client = get_async_httpx_client( - llm_provider=provider_config.custom_llm_provider - ) + async_httpx_client = get_async_httpx_client(llm_provider=provider_config.custom_llm_provider) else: async_httpx_client = client @@ -5144,9 +4901,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -5204,9 +4959,7 @@ class BaseLLMHTTPHandler: base_func = CustomLogger.async_should_run_agentic_loop for cb in _custom_logger_callbacks(logging_obj): cb_func = getattr(type(cb), "async_should_run_agentic_loop", base_func) - if getattr(cb_func, "__func__", cb_func) is not getattr( - base_func, "__func__", base_func - ): + if getattr(cb_func, "__func__", cb_func) is not getattr(base_func, "__func__", base_func): return True return False @@ -5229,13 +4982,9 @@ class BaseLLMHTTPHandler: """ fingerprint = BaseLLMHTTPHandler._fingerprint_agentic_tools(tool_calls) if fingerprint in fingerprints: - raise ValueError( - "Agentic loop detected repeated tool-call fingerprint; aborting rerun" - ) + raise ValueError("Agentic loop detected repeated tool-call fingerprint; aborting rerun") if depth >= max_loops: - raise ValueError( - f"Exceeded max_agentic_loops={max_loops} for model={model}" - ) + raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}") return fingerprint @staticmethod @@ -5268,9 +5017,7 @@ class BaseLLMHTTPHandler: full_model_name = model if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get( - "agentic_loop_params", {} - ) + agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = cast(str, agentic_params.get("model", model)) optional_params = dict(anthropic_messages_optional_request_params) @@ -5417,8 +5164,7 @@ class BaseLLMHTTPHandler: except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in " - "async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", + "LiteLLM.AgenticHookError: Exception in async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", _call_id, model, str(e), @@ -5526,9 +5272,7 @@ class BaseLLMHTTPHandler: if api_surface != "anthropic_messages": return response websearch_converted_stream = ( - logging_obj.model_call_details.get( - "websearch_interception_converted_stream", False - ) + logging_obj.model_call_details.get("websearch_interception_converted_stream", False) if logging_obj is not None else False ) @@ -5544,12 +5288,9 @@ class BaseLLMHTTPHandler: ) verbose_logger.debug( - "WebSearchInterception: Agentic loop completed, " - "converting non-streaming response to fake stream" - ) - return FakeAnthropicMessagesStreamIterator( - response=cast(AnthropicMessagesResponse, response) + "WebSearchInterception: Agentic loop completed, converting non-streaming response to fake stream" ) + return FakeAnthropicMessagesStreamIterator(response=cast(AnthropicMessagesResponse, response)) return response async def _call_agentic_completion_hooks( @@ -5605,8 +5346,7 @@ class BaseLLMHTTPHandler: except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in " - "async_should_run_agentic_loop [call_id=%s model=%s]: %s", + "LiteLLM.AgenticHookError: Exception in async_should_run_agentic_loop [call_id=%s model=%s]: %s", _call_id, model, str(e), @@ -5630,8 +5370,7 @@ class BaseLLMHTTPHandler: kwargs_with_provider = kwargs.copy() if kwargs else {} kwargs_with_provider["custom_llm_provider"] = custom_llm_provider build_plan_overridden = ( - callback.__class__.async_build_agentic_loop_plan - is not CustomLogger.async_build_agentic_loop_plan + callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan ) if not build_plan_overridden: agentic_result = await callback.async_run_agentic_loop( @@ -5645,9 +5384,7 @@ class BaseLLMHTTPHandler: stream=stream, kwargs=kwargs_with_provider, ) - return self._maybe_wrap_in_fake_stream( - agentic_result, logging_obj, api_surface - ) + return self._maybe_wrap_in_fake_stream(agentic_result, logging_obj, api_surface) plan = await callback.async_build_agentic_loop_plan( tools=tool_calls, @@ -5662,18 +5399,14 @@ class BaseLLMHTTPHandler: ) if plan.response_override is not None: - return self._maybe_wrap_in_fake_stream( - plan.response_override, logging_obj, api_surface - ) + return self._maybe_wrap_in_fake_stream(plan.response_override, logging_obj, api_surface) if plan.terminate: verbose_logger.debug( "Agentic loop terminated by callback=%s reason=%s", callback.__class__.__name__, plan.stop_reason, ) - return self._maybe_wrap_in_fake_stream( - response, logging_obj, api_surface - ) + return self._maybe_wrap_in_fake_stream(response, logging_obj, api_surface) if not plan.run_agentic_loop: continue @@ -5712,8 +5445,7 @@ class BaseLLMHTTPHandler: except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in agentic completion hooks " - "[call_id=%s model=%s]: %s", + "LiteLLM.AgenticHookError: Exception in agentic completion hooks [call_id=%s model=%s]: %s", _call_id, model, str(e), @@ -5779,8 +5511,7 @@ class BaseLLMHTTPHandler: ) except Exception as e: verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in " - "async_should_run_chat_completion_agentic_loop: %s", + "LiteLLM.AgenticHookError: Exception in async_should_run_chat_completion_agentic_loop: %s", str(e), ) continue @@ -5863,9 +5594,7 @@ class BaseLLMHTTPHandler: # 2. No agentic loop ran (LLM didn't use the tool) # 3. We have a non-streaming response that needs to be converted to streaming websearch_converted_stream = ( - logging_obj.model_call_details.get( - "websearch_interception_converted_stream", False - ) + logging_obj.model_call_details.get("websearch_interception_converted_stream", False) if logging_obj is not None else False ) @@ -5945,9 +5674,7 @@ class BaseLLMHTTPHandler: ) @staticmethod - def _append_query_params( - url: str, query_params: Optional[RealtimeQueryParams] - ) -> str: + def _append_query_params(url: str, query_params: Optional[RealtimeQueryParams]) -> str: """Append query_params to url, skipping keys already present in the URL.""" if not query_params: return url @@ -6003,9 +5730,7 @@ class BaseLLMHTTPHandler: # (e.g. Gemini/Vertex AI Live needs a `setup` message before any realtime_input) _session_config: Optional[str] = None if provider_config.requires_session_configuration(): - _session_config = provider_config.session_configuration_request( - model - ) + _session_config = provider_config.session_configuration_request(model) if _session_config: await backend_ws.send(_session_config) @@ -6021,9 +5746,7 @@ class BaseLLMHTTPHandler: user_api_key_dict=user_api_key_dict, request_data=_request_data, force_transcription_model=( - model - if (query_params or {}).get("intent") == "transcription" - else None + model if (query_params or {}).get("intent") == "transcription" else None ), ) if _session_config: @@ -6046,9 +5769,7 @@ class BaseLLMHTTPHandler: realtime_streaming.store_message(synthetic_session_str) await websocket.send_text(synthetic_session_str) realtime_streaming._session_created_sent_to_client = True - verbose_logger.debug( - "Sent synthetic session.created to client to unblock connection" - ) + verbose_logger.debug("Sent synthetic session.created to client to unblock connection") await realtime_streaming.bidirectional_forward() @@ -6058,20 +5779,14 @@ class BaseLLMHTTPHandler: except Exception as e: verbose_logger.exception(f"Error connecting to backend: {e}") try: - await websocket.close( - code=1011, reason=_redact_string(f"Internal server error: {str(e)}") - ) + await websocket.close(code=1011, reason=_redact_string(f"Internal server error: {str(e)}")) except RuntimeError as close_error: - if "already completed" in str(close_error) or "websocket.close" in str( - close_error - ): + if "already completed" in str(close_error) or "websocket.close" in str(close_error): # The WebSocket is already closed or the response is completed, so we can ignore this error pass else: # If it's a different RuntimeError, we might want to log it or handle it differently - raise Exception( - f"Unexpected error while closing WebSocket: {close_error}" - ) + raise Exception(f"Unexpected error while closing WebSocket: {close_error}") async def async_realtime_client_secret_handler( self, @@ -6168,9 +5883,7 @@ class BaseLLMHTTPHandler: api_base=api_base, model=model or "", api_version=api_version ) else: - url = provider_config.get_complete_url( - api_base=api_base, model=model or "", api_version=api_version - ) + url = provider_config.get_complete_url(api_base=api_base, model=model or "", api_version=api_version) headers: Dict[str, Any] = provider_config.validate_environment( headers={}, model=model or "", api_key=api_key ) @@ -6241,12 +5954,8 @@ class BaseLLMHTTPHandler: async_httpx_client = client if provider_config is not None: - url = provider_config.get_realtime_calls_url( - api_base=api_base, model=model or "", api_version=api_version - ) - headers: Dict[str, Any] = provider_config.get_realtime_calls_headers( - ephemeral_key=openai_ephemeral_key - ) + url = provider_config.get_realtime_calls_url(api_base=api_base, model=model or "", api_version=api_version) + headers: Dict[str, Any] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) else: url = f"{api_base.rstrip('/')}/v1/realtime/calls" headers = { @@ -6321,10 +6030,7 @@ class BaseLLMHTTPHandler: - Uses ManagedResponsesWebSocketHandler which makes HTTP streaming calls - Forwards events over the websocket connection """ - if ( - responses_api_provider_config is None - or not responses_api_provider_config.supports_native_websocket() - ): + if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket(): from litellm.responses.streaming_iterator import ( ManagedResponsesWebSocketHandler, ) @@ -6373,9 +6079,7 @@ class BaseLLMHTTPHandler: _qs = parse_qs(_parsed.query) if "model" not in _qs: _qs["model"] = [model] - ws_url = urlunparse( - _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) - ) + ws_url = urlunparse(_parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()}))) try: ssl_context = get_shared_realtime_ssl_context() @@ -6416,9 +6120,7 @@ class BaseLLMHTTPHandler: cb for cb in _litellm.callbacks if callable(getattr(cb, "check_pii", None)) - and callable( - getattr(cb, "get_presidio_settings_from_request_data", None) - ) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) and callable(getattr(cb, "_unmask_pii_text", None)) and getattr(cb, "output_parse_pii", False) ] @@ -6426,9 +6128,7 @@ class BaseLLMHTTPHandler: cb for cb in _litellm.callbacks if callable(getattr(cb, "check_pii", None)) - and callable( - getattr(cb, "get_presidio_settings_from_request_data", None) - ) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) and getattr(cb, "apply_to_output", False) ] except Exception as _guardrail_exc: @@ -6457,18 +6157,12 @@ class BaseLLMHTTPHandler: except Exception as e: verbose_logger.exception(f"Error in responses WS: {e}") try: - await websocket.close( - code=1011, reason=_redact_string(f"Internal server error: {str(e)}") - ) + await websocket.close(code=1011, reason=_redact_string(f"Internal server error: {str(e)}")) except RuntimeError as close_error: - if "already completed" in str(close_error) or "websocket.close" in str( - close_error - ): + if "already completed" in str(close_error) or "websocket.close" in str(close_error): pass else: - raise Exception( - f"Unexpected error while closing WebSocket: {close_error}" - ) + raise Exception(f"Unexpected error while closing WebSocket: {close_error}") def image_edit_handler( self, @@ -6516,9 +6210,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -6547,9 +6239,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) - data = image_edit_provider_config.finalize_image_edit_request_data( - data, api_base - ) + data = image_edit_provider_config.finalize_image_edit_request_data(data, api_base) ## LOGGING logging_obj.pre_call( @@ -6648,9 +6338,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) - data = image_edit_provider_config.finalize_image_edit_request_data( - data, api_base - ) + data = image_edit_provider_config.finalize_image_edit_request_data(data, api_base) ## LOGGING logging_obj.pre_call( @@ -6740,16 +6428,13 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client headers = image_generation_provider_config.validate_environment( api_key=api_key, - headers=image_generation_optional_request_params.get("extra_headers", {}) - or {}, + headers=image_generation_optional_request_params.get("extra_headers", {}) or {}, model=model, messages=[], optional_params=image_generation_optional_request_params, @@ -6812,17 +6497,15 @@ class BaseLLMHTTPHandler: provider_config=image_generation_provider_config, ) - model_response: ImageResponse = ( - image_generation_provider_config.transform_image_generation_response( - model=model, - raw_response=response, - model_response=litellm.ImageResponse(), - logging_obj=logging_obj, - request_data=data, - optional_params=image_generation_optional_request_params, - litellm_params=dict(litellm_params), - encoding=None, - ) + model_response: ImageResponse = image_generation_provider_config.transform_image_generation_response( + model=model, + raw_response=response, + model_response=litellm.ImageResponse(), + logging_obj=logging_obj, + request_data=data, + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + encoding=None, ) return model_response @@ -6858,8 +6541,7 @@ class BaseLLMHTTPHandler: headers = image_generation_provider_config.validate_environment( api_key=api_key, - headers=image_generation_optional_request_params.get("extra_headers", {}) - or {}, + headers=image_generation_optional_request_params.get("extra_headers", {}) or {}, model=model, messages=[], optional_params=image_generation_optional_request_params, @@ -6922,17 +6604,15 @@ class BaseLLMHTTPHandler: provider_config=image_generation_provider_config, ) - model_response: ImageResponse = ( - image_generation_provider_config.transform_image_generation_response( - model=model, - raw_response=response, - model_response=litellm.ImageResponse(), - logging_obj=logging_obj, - request_data=data, - optional_params=image_generation_optional_request_params, - litellm_params=dict(litellm_params), - encoding=None, - ) + model_response: ImageResponse = image_generation_provider_config.transform_image_generation_response( + model=model, + raw_response=response, + model_response=litellm.ImageResponse(), + logging_obj=logging_obj, + request_data=data, + optional_params=image_generation_optional_request_params, + litellm_params=dict(litellm_params), + encoding=None, ) return model_response @@ -6983,16 +6663,13 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client headers = video_generation_provider_config.validate_environment( api_key=api_key or litellm_params.get("api_key", None), - headers=video_generation_optional_request_params.get("extra_headers", {}) - or {}, + headers=video_generation_optional_request_params.get("extra_headers", {}) or {}, model=model, litellm_params=litellm_params, ) @@ -7096,8 +6773,7 @@ class BaseLLMHTTPHandler: headers = video_generation_provider_config.validate_environment( api_key=api_key or litellm_params.get("api_key", None), - headers=video_generation_optional_request_params.get("extra_headers", {}) - or {}, + headers=video_generation_optional_request_params.get("extra_headers", {}) or {}, model=model, litellm_params=litellm_params, ) @@ -7200,9 +6876,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -7375,9 +7049,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client directly (like video_generation does) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -7551,9 +7223,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -7707,9 +7377,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -7800,9 +7468,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) response.raise_for_status() return video_provider_config.transform_video_get_character_response( raw_response=response, @@ -7842,9 +7508,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -8051,9 +7715,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -8432,9 +8094,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client directly (like video_generation does) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -8492,12 +8152,10 @@ class BaseLLMHTTPHandler: headers=headers, ) - return ( - video_status_provider_config.transform_video_status_retrieve_response( - raw_response=response, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, - ) + return video_status_provider_config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, ) except Exception as e: @@ -8583,12 +8241,10 @@ class BaseLLMHTTPHandler: url=url, headers=headers, ) - return ( - video_status_provider_config.transform_video_status_retrieve_response( - raw_response=response, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, - ) + return video_status_provider_config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, ) except Exception as e: @@ -8625,9 +8281,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -8794,9 +8448,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -8959,9 +8611,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -9126,9 +8776,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -9280,9 +8928,7 @@ class BaseLLMHTTPHandler: timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union[ - "ContainerFileListResponse", Coroutine[Any, Any, "ContainerFileListResponse"] - ]: + ) -> Union["ContainerFileListResponse", Coroutine[Any, Any, "ContainerFileListResponse"]]: if _is_async: return self.async_container_file_list_handler( container_id=container_id, @@ -9300,9 +8946,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -9467,9 +9111,7 @@ class BaseLLMHTTPHandler: # For sync calls, use sync HTTP client if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -9642,9 +9284,7 @@ class BaseLLMHTTPHandler: ) # Check if provider has async transform method - if hasattr( - vector_store_provider_config, "atransform_search_vector_store_request" - ): + if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): ( url, request_body, @@ -9689,9 +9329,7 @@ class BaseLLMHTTPHandler: }, ) - request_data = ( - json.dumps(request_body) if signed_json_body is None else signed_json_body - ) + request_data = json.dumps(request_body) if signed_json_body is None else signed_json_body try: response = await async_httpx_client.post( @@ -9722,9 +9360,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[ - VectorStoreSearchResponse, Coroutine[Any, Any, VectorStoreSearchResponse] - ]: + ) -> Union[VectorStoreSearchResponse, Coroutine[Any, Any, VectorStoreSearchResponse]]: if _is_async: return self.async_vector_store_search_handler( vector_store_id=vector_store_id, @@ -9741,9 +9377,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -9792,9 +9426,7 @@ class BaseLLMHTTPHandler: }, ) - request_data = ( - json.dumps(request_body) if signed_json_body is None else signed_json_body - ) + request_data = json.dumps(request_body) if signed_json_body is None else signed_json_body try: response = sync_httpx_client.post( @@ -9862,9 +9494,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -9884,9 +9514,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[ - VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse] - ]: + ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: if _is_async: return self.async_vector_store_create_handler( vector_store_create_optional_params=vector_store_create_optional_params, @@ -9901,9 +9529,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -9938,9 +9564,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body) except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -9980,9 +9604,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( @@ -10015,9 +9637,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[ - VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse] - ]: + ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: if _is_async: return self.async_vector_store_retrieve_handler( vector_store_id=vector_store_id, @@ -10032,9 +9652,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -10050,9 +9668,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( @@ -10131,9 +9747,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -10172,9 +9786,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -10252,9 +9864,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" request_body: Dict[str, Any] = dict(vector_store_update_optional_params) @@ -10279,9 +9889,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -10302,9 +9910,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[ - VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse] - ]: + ) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: if _is_async: return self.async_vector_store_update_handler( vector_store_id=vector_store_id, @@ -10320,9 +9926,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -10338,9 +9942,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" request_body: Dict[str, Any] = dict(vector_store_update_optional_params) @@ -10365,9 +9967,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body) except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -10407,9 +10007,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( @@ -10422,9 +10020,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.delete( - url=url, headers=headers, timeout=timeout - ) + response = await async_httpx_client.delete(url=url, headers=headers, timeout=timeout) except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -10457,9 +10053,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -10475,9 +10069,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( @@ -10558,17 +10150,11 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_create_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_create_vector_store_file_response(response=response) def vector_store_file_create_handler( self, @@ -10600,9 +10186,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -10643,17 +10227,11 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_create_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_create_vector_store_file_response(response=response) async def async_vector_store_file_list_handler( self, @@ -10713,17 +10291,11 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=request_params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=request_params) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_list_vector_store_files_response( - response=response - ) + return vector_store_files_provider_config.transform_list_vector_store_files_response(response=response) def vector_store_file_list_handler( self, @@ -10739,9 +10311,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, - ) -> Union[ - VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse] - ]: + ) -> Union[VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse]]: if _is_async: return self.async_vector_store_file_list_handler( vector_store_id=vector_store_id, @@ -10757,9 +10327,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -10799,17 +10367,11 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=request_params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=request_params) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_list_vector_store_files_response( - response=response - ) + return vector_store_files_provider_config.transform_list_vector_store_files_response(response=response) async def async_vector_store_file_retrieve_handler( self, @@ -10864,17 +10426,11 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=request_params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=request_params) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_retrieve_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_retrieve_vector_store_file_response(response=response) def vector_store_file_retrieve_handler( self, @@ -10904,9 +10460,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -10942,17 +10496,11 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=request_params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=request_params) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_retrieve_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_retrieve_vector_store_file_response(response=response) async def async_vector_store_file_content_handler( self, @@ -11007,13 +10555,9 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=request_params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=request_params) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) return vector_store_files_provider_config.transform_retrieve_vector_store_file_content_response( response=response @@ -11050,9 +10594,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -11088,13 +10630,9 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=request_params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=request_params) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) return vector_store_files_provider_config.transform_retrieve_vector_store_file_content_response( response=response @@ -11160,17 +10698,11 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_update_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_update_vector_store_file_response(response=response) def vector_store_file_update_handler( self, @@ -11204,9 +10736,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -11247,17 +10777,11 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_update_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_update_vector_store_file_response(response=response) async def async_vector_store_file_delete_handler( self, @@ -11312,17 +10836,11 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.delete( - url=url, headers=headers, params=request_params, timeout=timeout - ) + response = await async_httpx_client.delete(url=url, headers=headers, params=request_params, timeout=timeout) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_delete_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_delete_vector_store_file_response(response=response) def vector_store_file_delete_handler( self, @@ -11355,9 +10873,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -11393,17 +10909,11 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.delete( - url=url, headers=headers, params=request_params, timeout=timeout - ) + response = sync_httpx_client.delete(url=url, headers=headers, params=request_params, timeout=timeout) except Exception as e: - raise self._handle_error( - e=e, provider_config=vector_store_files_provider_config - ) + raise self._handle_error(e=e, provider_config=vector_store_files_provider_config) - return vector_store_files_provider_config.transform_delete_vector_store_file_response( - response=response - ) + return vector_store_files_provider_config.transform_delete_vector_store_file_response(response=response) ##################################################################### ################ Google GenAI GENERATE CONTENT HANDLER ########################### @@ -11455,9 +10965,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -11698,9 +11206,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -11942,9 +11448,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -11962,19 +11466,13 @@ class BaseLLMHTTPHandler: try: # Check if files are present - use multipart/form-data - data, files = self._prepare_skill_multipart_request( - request_body=request_body, headers=headers - ) + data, files = self._prepare_skill_multipart_request(request_body=request_body, headers=headers) if files is not None: - response = sync_httpx_client.post( - url=url, headers=headers, data=data, files=files, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, data=data, files=files, timeout=timeout) else: # No files - send as JSON - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12022,9 +11520,7 @@ class BaseLLMHTTPHandler: try: # Check if files are present - use multipart/form-data - data, files = self._prepare_skill_multipart_request( - request_body=request_body, headers=headers - ) + data, files = self._prepare_skill_multipart_request(request_body=request_body, headers=headers) if files is not None: response = await async_httpx_client.post( @@ -12032,9 +11528,7 @@ class BaseLLMHTTPHandler: ) else: # No files - send as JSON - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12076,9 +11570,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12095,9 +11587,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=query_params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=query_params) except Exception as e: raise self._handle_error( e=e, @@ -12144,9 +11634,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=query_params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=query_params) except Exception as e: raise self._handle_error( e=e, @@ -12186,9 +11674,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12289,9 +11775,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12307,9 +11791,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.delete( - url=url, headers=headers, timeout=timeout - ) + response = sync_httpx_client.delete(url=url, headers=headers, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12354,9 +11836,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.delete( - url=url, headers=headers, timeout=timeout - ) + response = await async_httpx_client.delete(url=url, headers=headers, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12402,9 +11882,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12421,9 +11899,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12470,9 +11946,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12514,9 +11988,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12533,9 +12005,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=query_params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=query_params) except Exception as e: raise self._handle_error( e=e, @@ -12582,9 +12052,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=query_params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=query_params) except Exception as e: raise self._handle_error( e=e, @@ -12624,9 +12092,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12729,9 +12195,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12748,9 +12212,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12797,9 +12259,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12839,9 +12299,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12857,9 +12315,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.delete( - url=url, headers=headers, timeout=timeout - ) + response = sync_httpx_client.delete(url=url, headers=headers, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12904,9 +12360,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.delete( - url=url, headers=headers, timeout=timeout - ) + response = await async_httpx_client.delete(url=url, headers=headers, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -12946,9 +12400,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -12964,9 +12416,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json={}, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json={}, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -13011,9 +12461,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json={}, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json={}, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -13059,9 +12507,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -13078,9 +12524,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -13127,9 +12571,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -13171,9 +12613,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -13190,9 +12630,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=query_params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=query_params) except Exception as e: raise self._handle_error( e=e, @@ -13239,9 +12677,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers, params=query_params - ) + response = await async_httpx_client.get(url=url, headers=headers, params=query_params) except Exception as e: raise self._handle_error( e=e, @@ -13281,9 +12717,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -13384,9 +12818,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -13402,9 +12834,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.post( - url=url, headers=headers, json={}, timeout=timeout - ) + response = sync_httpx_client.post(url=url, headers=headers, json={}, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -13449,9 +12879,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json={}, timeout=timeout - ) + response = await async_httpx_client.post(url=url, headers=headers, json={}, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -13491,9 +12919,7 @@ class BaseLLMHTTPHandler: ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client( - params={"ssl_verify": litellm_params.get("ssl_verify", None)} - ) + sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) else: sync_httpx_client = client @@ -13509,9 +12935,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.delete( - url=url, headers=headers, timeout=timeout - ) + response = sync_httpx_client.delete(url=url, headers=headers, timeout=timeout) except Exception as e: raise self._handle_error( e=e, @@ -13556,9 +12980,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.delete( - url=url, headers=headers, timeout=timeout - ) + response = await async_httpx_client.delete(url=url, headers=headers, timeout=timeout) except Exception as e: raise self._handle_error( e=e, diff --git a/litellm/llms/custom_llm.py b/litellm/llms/custom_llm.py index a820ac7f345..e0af3986465 100644 --- a/litellm/llms/custom_llm.py +++ b/litellm/llms/custom_llm.py @@ -39,9 +39,7 @@ class CustomLLMError(Exception): # use this for all your exceptions ): self.status_code = status_code self.message = message - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class CustomLLM(BaseLLM): @@ -154,12 +152,8 @@ class CustomLLM(BaseLLM): model: str, prompt: str, model_response: ImageResponse, - api_key: Optional[ - str - ], # dynamically set api_key - https://docs.litellm.ai/docs/set_keys#api_key - api_base: Optional[ - str - ], # dynamically set api_base - https://docs.litellm.ai/docs/set_keys#api_base + api_key: Optional[str], # dynamically set api_key - https://docs.litellm.ai/docs/set_keys#api_key + api_base: Optional[str], # dynamically set api_base - https://docs.litellm.ai/docs/set_keys#api_base optional_params: dict, logging_obj: Any, timeout: Optional[Union[float, httpx.Timeout]] = None, @@ -228,9 +222,7 @@ class CustomLLM(BaseLLM): raise CustomLLMError(status_code=500, message="Not implemented yet!") -def custom_chat_llm_router( - async_fn: bool, stream: Optional[bool], custom_llm: CustomLLM -): +def custom_chat_llm_router(async_fn: bool, stream: Optional[bool], custom_llm: CustomLLM): """ Routes call to CustomLLM completion/acompletion/streaming/astreaming functions, based on call type diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index ccb4d370c95..743bf494d92 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -42,21 +42,15 @@ class DashScopeChatConfig(OpenAIGPTConfig): self, messages: List[AllMessageValues], model: str, is_async: bool = False ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: api_base = ( - api_base - or get_secret_str("DASHSCOPE_API_BASE") - or "https://dashscope.aliyuncs.com/compatible-mode/v1" + api_base or get_secret_str("DASHSCOPE_API_BASE") or "https://dashscope.aliyuncs.com/compatible-mode/v1" ) # type: ignore dynamic_api_key = api_key or get_secret_str("DASHSCOPE_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 8bb7f605b82..2f710d78126 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -24,9 +24,7 @@ class TokenBreakdown: def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: """Extract token counts from usage, handling cached and reasoning tokens.""" cached_tokens = 0 - if usage.prompt_tokens_details and hasattr( - usage.prompt_tokens_details, "cached_tokens" - ): + if usage.prompt_tokens_details and hasattr(usage.prompt_tokens_details, "cached_tokens"): cached_tokens = usage.prompt_tokens_details.cached_tokens or 0 text_tokens = usage.prompt_tokens - cached_tokens @@ -41,9 +39,7 @@ def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: completion_tokens = (usage.completion_tokens or 0) - reasoning_tokens - return TokenBreakdown( - text_tokens, cached_tokens, completion_tokens, reasoning_tokens - ) + return TokenBreakdown(text_tokens, cached_tokens, completion_tokens, reasoning_tokens) def _calculate_tiered_cost( @@ -181,9 +177,7 @@ def _calculate_completion_cost( else: reasoning_cost = float(reasoning_cost_val) - return (breakdown.completion_tokens * output_cost) + ( - breakdown.reasoning_tokens * reasoning_cost - ) + return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: @@ -201,15 +195,9 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: """ model_info = get_model_info(model=model, custom_llm_provider="dashscope") breakdown = _extract_token_breakdown(usage) - tiered_pricing = ( - model_info.get("tiered_pricing") - if isinstance(model_info.get("tiered_pricing"), list) - else None - ) + tiered_pricing = model_info.get("tiered_pricing") if isinstance(model_info.get("tiered_pricing"), list) else None - prompt_cost = _calculate_prompt_cost( - breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing - ) + prompt_cost = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing) completion_cost = _calculate_completion_cost( breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing ) diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 5bc0e5ca817..070e2f57667 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -144,11 +144,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): if "error" in response_json: error = response_json["error"] - message = ( - error.get("message", str(error)) - if isinstance(error, dict) - else str(error) - ) + message = error.get("message", str(error)) if isinstance(error, dict) else str(error) raise DashScopeError( status_code=raw_response.status_code, message=message, diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index 77676b11d51..094e06d1269 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -62,9 +62,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): Configuration for DashScope image generation (qwen-image-2.0, qwen-image-2.0-pro). """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: return ["n", "size"] def map_openai_params( @@ -97,9 +95,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - return ( - api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE - ) + return api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE def validate_environment( self, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 09c782a4755..a2c1d41022b 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -84,11 +84,7 @@ def _sanitize_empty_content(message_dict: dict[str, Any]) -> None: filtered = [ block for block in content - if not ( - isinstance(block, dict) - and block.get("type") == "text" - and not (block.get("text") or "").strip() - ) + if not (isinstance(block, dict) and block.get("type") == "text" and not (block.get("text") or "").strip()) ] if not filtered: message_dict.pop("content") @@ -241,12 +237,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): return tools # if claude, convert to anthropic tool and then to databricks tool - anthropic_tools, _ = self._map_tools( - tools=tools - ) # unclear how mcp tool calling on databricks works + anthropic_tools, _ = self._map_tools(tools=tools) # unclear how mcp tool calling on databricks works databricks_tools = [ - cast(DatabricksTool, self.convert_anthropic_tool_to_databricks_tool(tool)) - for tool in anthropic_tools + cast(DatabricksTool, self.convert_anthropic_tool_to_databricks_tool(tool)) for tool in anthropic_tools ] return databricks_tools @@ -260,9 +253,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if value is None: return None - tool = self.map_response_format_to_anthropic_tool( - value, optional_params, is_thinking_enabled - ) + tool = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled) databricks_tool = self.convert_anthropic_tool_to_databricks_tool(tool) return databricks_tool @@ -291,17 +282,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): replace_max_completion_tokens_with_max_tokens: bool = True, ) -> dict: is_thinking_enabled = self.is_thinking_enabled(non_default_params) - mapped_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + mapped_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) if "tools" in mapped_params: - mapped_params["tools"] = self._map_openai_to_dbrx_tool( - model=model, tools=mapped_params["tools"] - ) - if ( - "max_completion_tokens" in non_default_params - and replace_max_completion_tokens_with_max_tokens - ): + mapped_params["tools"] = self._map_openai_to_dbrx_tool(model=model, tools=mapped_params["tools"]) + if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: mapped_params["max_tokens"] = non_default_params[ "max_completion_tokens" ] # most openai-compatible providers support 'max_tokens' not 'max_completion_tokens' @@ -316,16 +300,12 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): ) if _tool is not None: - self._add_tools_to_optional_params( - optional_params=optional_params, tools=[_tool] - ) + self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool]) optional_params["json_mode"] = True if not is_thinking_enabled: _tool_choice = ChatCompletionToolChoiceObjectParam( type="function", - function=ChatCompletionToolChoiceFunctionParam( - name=RESPONSE_FORMAT_TOOL_NAME - ), + function=ChatCompletionToolChoiceFunctionParam(name=RESPONSE_FORMAT_TOOL_NAME), ) optional_params["tool_choice"] = _tool_choice optional_params.pop( @@ -347,9 +327,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if AnthropicConfig._is_adaptive_thinking_model(model): mapped_effort: Optional[str] = None if isinstance(reasoning_effort_value, str): - mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get( - reasoning_effort_value - ) + mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort_value) if mapped_effort is None: AnthropicConfig._raise_invalid_reasoning_effort( model=model, @@ -408,17 +386,11 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): new_messages.append(_message) if is_async: - return super()._transform_messages( - messages=new_messages, model=model, is_async=cast(Literal[True], True) - ) + return super()._transform_messages(messages=new_messages, model=model, is_async=cast(Literal[True], True)) else: - return super()._transform_messages( - messages=new_messages, model=model, is_async=cast(Literal[False], False) - ) + return super()._transform_messages(messages=new_messages, model=model, is_async=cast(Literal[False], False)) - def _move_cache_control_into_string_content_block( - self, message: AllMessageValues - ) -> AllMessageValues: + def _move_cache_control_into_string_content_block(self, message: AllMessageValues) -> AllMessageValues: """ Moves message-level cache_control into a content block when content is a string. @@ -466,22 +438,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): content: Optional[AllDatabricksContentValues], ) -> Tuple[ Optional[str], - Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ], + Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]], ]: """ Extract and return the reasoning content and thinking blocks """ if content is None: return None, None - thinking_blocks: Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ] = None + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None reasoning_content: Optional[str] = None if isinstance(content, list): for item in content: @@ -513,12 +477,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): for item in content: text = item.get("text", None) if citations_item := item.get("citations"): - citations.append( - [ - {**citation, "supported_text": text} - for citation in citations_item - ] - ) + citations.append([{**citation, "supported_text": text} for citation in citations_item]) return citations or None def _transform_dbrx_choices( @@ -534,9 +493,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): for _tc in tool_calls: _openai_tc = ChatCompletionMessageToolCall(**_tc) # type: ignore _openai_tool_calls.append(_openai_tc) - fixed_tool_calls = _handle_invalid_parallel_tool_calls( - _openai_tool_calls - ) + fixed_tool_calls = _handle_invalid_parallel_tool_calls(_openai_tool_calls) if fixed_tool_calls is not None: tool_calls = fixed_tool_calls @@ -548,30 +505,22 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): convert_tool_call_to_json_mode=json_mode, ): # to support response_format on claude models - json_mode_content_str: Optional[str] = ( - str(tool_calls[0]["function"].get("arguments", "")) or None - ) + json_mode_content_str: Optional[str] = str(tool_calls[0]["function"].get("arguments", "")) or None if json_mode_content_str is not None: translated_message = Message(content=json_mode_content_str) finish_reason = "stop" if translated_message is None: ## get the content str - content_str = DatabricksConfig.extract_content_str( - choice["message"]["content"] - ) + content_str = DatabricksConfig.extract_content_str(choice["message"]["content"]) ## get the reasoning content ( reasoning_content, thinking_blocks, - ) = DatabricksConfig.extract_reasoning_content( - choice["message"].get("content") - ) + ) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content")) - citations = DatabricksConfig.extract_citations( - choice["message"].get("content") - ) + citations = DatabricksConfig.extract_citations(choice["message"].get("content")) translated_message = Message( role="assistant", @@ -579,9 +528,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): reasoning_content=reasoning_content, thinking_blocks=thinking_blocks, tool_calls=choice["message"].get("tool_calls"), - provider_specific_fields=( - {"citations": citations} if citations is not None else None - ), + provider_specific_fields=({"citations": citations} if citations is not None else None), ) if finish_reason is None: @@ -630,9 +577,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise DatabricksException( - message="Unable to get json response - {}, Original Response: {}".format( - str(e), raw_response.text - ), + message="Unable to get json response - {}, Original Response: {}".format(str(e), raw_response.text), status_code=raw_response.status_code, headers=response_headers, ) @@ -715,29 +660,21 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator): for _tc in tool_calls: if _tc.get("function", {}).get("arguments") == "{}": _tc["function"]["arguments"] = "" # avoid invalid json - if isinstance(choice["delta"].get("content"), list) and ( - content := choice["delta"]["content"] - ): + if isinstance(choice["delta"].get("content"), list) and (content := choice["delta"]["content"]): if citations := content[0].get("citations"): # TODO: Databricks delta does not include supported text or chunk type. # Add either here once Databricks supports it to enable citation linkage. - choice["delta"].setdefault("provider_specific_fields", {})[ - "citation" - ] = citations[ + choice["delta"].setdefault("provider_specific_fields", {})["citation"] = citations[ 0 ] # Databricks Content item always has citation as a list of list # extract the content str - content_str = DatabricksConfig.extract_content_str( - choice["delta"].get("content") - ) + content_str = DatabricksConfig.extract_content_str(choice["delta"].get("content")) # extract the reasoning content ( reasoning_content, thinking_blocks, - ) = DatabricksConfig.extract_reasoning_content( - choice["delta"].get("content") - ) + ) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content")) choice["delta"]["content"] = content_str choice["delta"]["reasoning_content"] = reasoning_content diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index d39d52d2d59..908aa56a4d6 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -170,10 +170,7 @@ class DatabricksBase: partner_name = custom_user_agent # Validate partner name: alphanumeric, underscore, hyphen only - if ( - partner_name - and partner_name.replace("_", "").replace("-", "").isalnum() - ): + if partner_name and partner_name.replace("_", "").replace("-", "").isalnum(): return f"{partner_name}_litellm/{version}" # Default: just litellm @@ -289,9 +286,7 @@ class DatabricksBase: api_base = api_base or f"{databricks_client.config.host}/serving-endpoints" if api_key is None: - databricks_auth_headers: dict[str, str] = ( - databricks_client.config.authenticate() - ) + databricks_auth_headers: dict[str, str] = databricks_client.config.authenticate() headers = {**databricks_auth_headers, **headers} return api_base, headers @@ -391,9 +386,7 @@ class DatabricksBase: headers["User-Agent"] = self._build_user_agent(custom_user_agent) # Debug logging with redaction (never log actual tokens) - verbose_logger.debug( - f"Databricks request headers: {self.redact_headers_for_logging(headers)}" - ) + verbose_logger.debug(f"Databricks request headers: {self.redact_headers_for_logging(headers)}") if endpoint_type == "chat_completions" and custom_endpoint is not True: api_base = "{}/chat/completions".format(api_base) diff --git a/litellm/llms/databricks/cost_calculator.py b/litellm/llms/databricks/cost_calculator.py index 5558e133b4d..9db151538b5 100644 --- a/litellm/llms/databricks/cost_calculator.py +++ b/litellm/llms/databricks/cost_calculator.py @@ -21,37 +21,23 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ base_model = model - if model.startswith("databricks/dbrx-instruct") or model.startswith( - "dbrx-instruct" - ): + if model.startswith("databricks/dbrx-instruct") or model.startswith("dbrx-instruct"): base_model = "databricks-dbrx-instruct" - elif model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith( - "meta-llama-3.1-70b-instruct" - ): + elif model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith("meta-llama-3.1-70b-instruct"): base_model = "databricks-meta-llama-3-1-70b-instruct" - elif model.startswith( - "databricks/meta-llama-3.1-405b-instruct" - ) or model.startswith("meta-llama-3.1-405b-instruct"): + elif model.startswith("databricks/meta-llama-3.1-405b-instruct") or model.startswith( + "meta-llama-3.1-405b-instruct" + ): base_model = "databricks-meta-llama-3-1-405b-instruct" - elif model.startswith("databricks/mixtral-8x7b-instruct-v0.1") or model.startswith( - "mixtral-8x7b-instruct-v0.1" - ): + elif model.startswith("databricks/mixtral-8x7b-instruct-v0.1") or model.startswith("mixtral-8x7b-instruct-v0.1"): base_model = "databricks-mixtral-8x7b-instruct" - elif model.startswith("databricks/mixtral-8x7b-instruct-v0.1") or model.startswith( - "mixtral-8x7b-instruct-v0.1" - ): + elif model.startswith("databricks/mixtral-8x7b-instruct-v0.1") or model.startswith("mixtral-8x7b-instruct-v0.1"): base_model = "databricks-mixtral-8x7b-instruct" - elif model.startswith("databricks/bge-large-en") or model.startswith( - "bge-large-en" - ): + elif model.startswith("databricks/bge-large-en") or model.startswith("bge-large-en"): base_model = "databricks-bge-large-en" - elif model.startswith("databricks/gte-large-en") or model.startswith( - "gte-large-en" - ): + elif model.startswith("databricks/gte-large-en") or model.startswith("gte-large-en"): base_model = "databricks-gte-large-en" - elif model.startswith("databricks/llama-2-70b-chat") or model.startswith( - "llama-2-70b-chat" - ): + elif model.startswith("databricks/llama-2-70b-chat") or model.startswith("llama-2-70b-chat"): base_model = "databricks-llama-2-70b-chat" ## GET MODEL INFO model_info = get_model_info(model=base_model, custom_llm_provider="databricks") diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index 7a7330227d6..a6a45719fe6 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -127,9 +127,7 @@ class ModelResponseIterator: except StopIteration: raise StopIteration except ValueError as e: - verbose_logger.debug( - f"Error parsing chunk: {e},\nReceived chunk: {chunk}. Defaulting to empty chunk here." - ) + verbose_logger.debug(f"Error parsing chunk: {e},\nReceived chunk: {chunk}. Defaulting to empty chunk here.") return GenericStreamingChunk( text="", is_finished=False, @@ -174,9 +172,7 @@ class ModelResponseIterator: except StopAsyncIteration: raise StopAsyncIteration except ValueError as e: - verbose_logger.debug( - f"Error parsing chunk: {e},\nReceived chunk: {chunk}. Defaulting to empty chunk here." - ) + verbose_logger.debug(f"Error parsing chunk: {e},\nReceived chunk: {chunk}. Defaulting to empty chunk here.") return GenericStreamingChunk( text="", is_finished=False, diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index 701db586b72..97a2539b3df 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -26,9 +26,7 @@ class DataForSEOSearchConfig(BaseSearchConfig): API endpoint: https://api.dataforseo.com/v3/serp/google/organic/live/advanced """ - DATAFORSEO_API_BASE = ( - "https://api.dataforseo.com/v3/serp/google/organic/live/advanced" - ) + DATAFORSEO_API_BASE = "https://api.dataforseo.com/v3/serp/google/organic/live/advanced" @staticmethod def ui_friendly_name() -> str: @@ -103,11 +101,7 @@ class DataForSEOSearchConfig(BaseSearchConfig): DataForSEO uses POST requests, so no query parameters in URL. """ - return ( - api_base - or get_secret_str("DATAFORSEO_API_BASE") - or self.DATAFORSEO_API_BASE - ) + return api_base or get_secret_str("DATAFORSEO_API_BASE") or self.DATAFORSEO_API_BASE def transform_search_request( self, @@ -152,10 +146,7 @@ class DataForSEOSearchConfig(BaseSearchConfig): # For simplicity, we'll use location_name which accepts country names task["location_name"] = optional_params["country"] - if ( - "search_domain_filter" in optional_params - and optional_params["search_domain_filter"] - ): + if "search_domain_filter" in optional_params and optional_params["search_domain_filter"]: # DataForSEO uses 'domain' parameter to filter by domain task["domain"] = optional_params["search_domain_filter"] @@ -169,10 +160,7 @@ class DataForSEOSearchConfig(BaseSearchConfig): # Pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in task - ): + if param not in self.get_supported_perplexity_optional_params() and param not in task: task[param] = value # DataForSEO API expects an array of tasks diff --git a/litellm/llms/datarobot/chat/transformation.py b/litellm/llms/datarobot/chat/transformation.py index f81e2420930..75bbfc19b69 100644 --- a/litellm/llms/datarobot/chat/transformation.py +++ b/litellm/llms/datarobot/chat/transformation.py @@ -42,9 +42,7 @@ class DataRobotConfig(OpenAILikeChatConfig): path += f"/api/v2/{LLMGW_PATH}" elif "api/v2/deployments" in path: # Dedicated deployment, leave it pass - elif ( - "api/v2" in path and LLMGW_PATH not in path - ): # Standard ENDPOINT path, add LLMGW + elif "api/v2" in path and LLMGW_PATH not in path: # Standard ENDPOINT path, add LLMGW path += LLMGW_PATH # Ensure the url ends with a trailing slash diff --git a/litellm/llms/deepgram/audio_transcription/transformation.py b/litellm/llms/deepgram/audio_transcription/transformation.py index 6a540d72778..b05fba3b5ca 100644 --- a/litellm/llms/deepgram/audio_transcription/transformation.py +++ b/litellm/llms/deepgram/audio_transcription/transformation.py @@ -24,9 +24,7 @@ from ..common_utils import DeepgramException class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: return ["language"] def map_openai_params( @@ -42,12 +40,8 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): optional_params[k] = v return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return DeepgramException( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return DeepgramException(message=error_message, status_code=status_code, headers=headers) def transform_audio_transcription_request( self, @@ -72,9 +66,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # Return structured data with binary content and no files # For Deepgram, we send binary data directly as request body - return AudioTranscriptionRequestData( - data=processed_audio.file_content, files=None - ) + return AudioTranscriptionRequestData(data=processed_audio.file_content, files=None) def transform_audio_transcription_response( self, @@ -131,9 +123,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return response except Exception as e: - raise ValueError( - f"Error transforming Deepgram response: {str(e)}\nResponse: {raw_response.text}" - ) + raise ValueError(f"Error transforming Deepgram response: {str(e)}\nResponse: {raw_response.text}") def _reconstruct_diarized_transcript(self, words: list) -> str: """ @@ -160,9 +150,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): if speaker != current_speaker: # New speaker: save previous segment and start new one if current_words: - segments.append( - f"Speaker {current_speaker}: {' '.join(current_words)}" - ) + segments.append(f"Speaker {current_speaker}: {' '.join(current_words)}") current_speaker = speaker current_words = [word_text] else: @@ -185,9 +173,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): stream: Optional[bool] = None, ) -> str: if api_base is None: - api_base = ( - get_secret_str("DEEPGRAM_API_BASE") or "https://api.deepgram.com/v1" - ) + api_base = get_secret_str("DEEPGRAM_API_BASE") or "https://api.deepgram.com/v1" api_base = api_base.rstrip("/") # Remove trailing slash if present # Build query parameters including the model diff --git a/litellm/llms/deepinfra/chat/transformation.py b/litellm/llms/deepinfra/chat/transformation.py index a6bd8b4934f..494c53354f7 100644 --- a/litellm/llms/deepinfra/chat/transformation.py +++ b/litellm/llms/deepinfra/chat/transformation.py @@ -94,15 +94,11 @@ class DeepInfraConfig(OpenAIGPTConfig): supported_openai_params = self.get_supported_openai_params(model=model) for param, value in non_default_params.items(): if ( - param == "temperature" - and value == 0 - and model == "mistralai/Mistral-7B-Instruct-v0.1" + param == "temperature" and value == 0 and model == "mistralai/Mistral-7B-Instruct-v0.1" ): # this model does no support temperature == 0 value = MIN_NON_ZERO_TEMPERATURE # close to 0 if param == "tool_choice": - if ( - value != "auto" and value != "none" - ): # https://deepinfra.com/docs/advanced/function_calling + if value != "auto" and value != "none": # https://deepinfra.com/docs/advanced/function_calling ## UNSUPPORTED TOOL CHOICE VALUE if litellm.drop_params is True or drop_params is True: value = None @@ -120,9 +116,7 @@ class DeepInfraConfig(OpenAIGPTConfig): optional_params[param] = value return optional_params - def _transform_tool_message_content( - self, messages: List[AllMessageValues] - ) -> List[AllMessageValues]: + def _transform_tool_message_content(self, messages: List[AllMessageValues]) -> List[AllMessageValues]: """ Transform tool message content from array to string format for DeepInfra compatibility. @@ -201,10 +195,6 @@ class DeepInfraConfig(OpenAIGPTConfig): self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: # deepinfra is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 - api_base = ( - api_base - or get_secret_str("DEEPINFRA_API_BASE") - or "https://api.deepinfra.com/v1/openai" - ) + api_base = api_base or get_secret_str("DEEPINFRA_API_BASE") or "https://api.deepinfra.com/v1/openai" dynamic_api_key = api_key or get_secret_str("DEEPINFRA_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 385aa051d00..82069e4e195 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -53,9 +53,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): ) # Remove 'openai' from the base if present - api_base_clean = ( - api_base.replace("openai", "") if "openai" in api_base else api_base - ) + api_base_clean = api_base.replace("openai", "") if "openai" in api_base else api_base # Remove any trailing slashes for consistency, then add one api_base_clean = api_base_clean.rstrip("/") + "/" @@ -74,9 +72,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): api_key = get_secret_str("DEEPINFRA_API_KEY") if api_key is None: - raise ValueError( - "Deepinfra API key is required. Please set 'DEEPINFRA_API_KEY' environment variable" - ) + raise ValueError("Deepinfra API key is required. Please set 'DEEPINFRA_API_KEY' environment variable") default_headers = { "Authorization": f"Bearer {api_key}", @@ -171,9 +167,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): # Create RerankResponse results = [] for i, score in enumerate(scores): - results.append( - RerankResponseResult(index=i, relevance_score=float(score)) - ) + results.append(RerankResponseResult(index=i, relevance_score=float(score))) # Create metadata for the response tokens = RerankTokens( @@ -183,9 +177,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): billed_units = RerankBilledUnits(total_tokens=input_tokens) meta = RerankResponseMeta(tokens=tokens, billed_units=billed_units) - rerank_response = RerankResponse( - id=request_id or str(uuid.uuid4()), results=results, meta=meta - ) + rerank_response = RerankResponse(id=request_id or str(uuid.uuid4()), results=results, meta=meta) # Store additional information in hidden params rerank_response._hidden_params = { diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index b90b1e1aa21..7a548136f2a 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -40,9 +40,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): Reference: https://api-docs.deepseek.com/guides/thinking_mode """ # Let parent handle standard params first - optional_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + optional_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) # Pop thinking/reasoning_effort from optional_params first (parent may have added them) # Then re-add only if valid for DeepSeek @@ -51,10 +49,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): # Handle thinking parameter - only accept {"type": "enabled"} if thinking_value is not None: - if ( - isinstance(thinking_value, dict) - and thinking_value.get("type") == "enabled" - ): + if isinstance(thinking_value, dict) and thinking_value.get("type") == "enabled": # DeepSeek only accepts {"type": "enabled"}, ignore budget_tokens optional_params["thinking"] = {"type": "enabled"} @@ -64,9 +59,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): return optional_params - def _fill_reasoning_content( - self, messages: List[AllMessageValues] - ) -> List[AllMessageValues]: + def _fill_reasoning_content(self, messages: List[AllMessageValues]) -> List[AllMessageValues]: """ DeepSeek thinking mode requires `reasoning_content` to be passed back on every assistant message in multi-turn conversations. If it is missing, @@ -127,13 +120,9 @@ class DeepSeekChatConfig(OpenAIGPTConfig): """ messages = handle_messages_with_content_list_to_str_conversion(messages) if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _thinking_mode_active(self, model: str, optional_params: dict) -> bool: """ @@ -177,9 +166,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): name = function.get("name") return name if isinstance(name, str) else None - def _tool_choice_matches_function_tool( - tool_choice: object, function_tool_names: set[str] - ) -> bool: + def _tool_choice_matches_function_tool(tool_choice: object, function_tool_names: set[str]) -> bool: if not isinstance(tool_choice, dict): return True if tool_choice.get("type") != "function": @@ -210,21 +197,12 @@ class DeepSeekChatConfig(OpenAIGPTConfig): cleaned = {k: v for k, v in optional_params.items() if k != "tools"} if function_tools: function_tool_names = { - name - for tool in function_tools - for name in (_get_function_tool_name(tool),) - if name is not None + name for tool in function_tools for name in (_get_function_tool_name(tool),) if name is not None } - if not _tool_choice_matches_function_tool( - cleaned.get("tool_choice"), function_tool_names - ): + if not _tool_choice_matches_function_tool(cleaned.get("tool_choice"), function_tool_names): cleaned = {k: v for k, v in cleaned.items() if k != "tool_choice"} return {**cleaned, "tools": function_tools} - return { - k: v - for k, v in cleaned.items() - if k not in ("tool_choice", "parallel_tool_calls") - } + return {k: v for k, v in cleaned.items() if k not in ("tool_choice", "parallel_tool_calls")} def transform_request( self, @@ -280,11 +258,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = ( - api_base - or get_secret_str("DEEPSEEK_API_BASE") - or "https://api.deepseek.com/beta" - ) # type: ignore + api_base = api_base or get_secret_str("DEEPSEEK_API_BASE") or "https://api.deepseek.com/beta" # type: ignore dynamic_api_key = api_key or get_secret_str("DEEPSEEK_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/deepseek/cost_calculator.py b/litellm/llms/deepseek/cost_calculator.py index e652ebeac54..312bd5bdeab 100644 --- a/litellm/llms/deepseek/cost_calculator.py +++ b/litellm/llms/deepseek/cost_calculator.py @@ -16,6 +16,4 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: Follows the same logic as Anthropic's cost per token calculation. """ - return generic_cost_per_token( - model=model, usage=usage, custom_llm_provider="deepseek" - ) + return generic_cost_per_token(model=model, usage=usage, custom_llm_provider="deepseek") diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py index 63b736ffd1d..ddbbe7c2107 100644 --- a/litellm/llms/deepseek/messages/transformation.py +++ b/litellm/llms/deepseek/messages/transformation.py @@ -54,11 +54,7 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): ) -> Tuple[dict, Optional[str]]: dynamic_api_key = self.get_api_key(api_key=api_key) - if ( - "x-api-key" not in headers - and "authorization" not in headers - and dynamic_api_key is not None - ): + if "x-api-key" not in headers and "authorization" not in headers and dynamic_api_key is not None: headers["x-api-key"] = dynamic_api_key if "anthropic-version" not in headers: @@ -130,7 +126,5 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): headers=headers, ) if "tools" in anthropic_messages_request: - anthropic_messages_request["tools"] = self._sanitize_tools_for_deepseek( - anthropic_messages_request["tools"] - ) + anthropic_messages_request["tools"] = self._sanitize_tools_for_deepseek(anthropic_messages_request["tools"]) return anthropic_messages_request diff --git a/litellm/llms/deprecated_providers/aleph_alpha.py b/litellm/llms/deprecated_providers/aleph_alpha.py index 81ad1346414..f58297997b6 100644 --- a/litellm/llms/deprecated_providers/aleph_alpha.py +++ b/litellm/llms/deprecated_providers/aleph_alpha.py @@ -13,13 +13,9 @@ class AlephAlphaError(Exception): def __init__(self, status_code, message): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://api.aleph-alpha.com/complete" - ) + self.request = httpx.Request(method="POST", url="https://api.aleph-alpha.com/complete") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class AlephAlphaConfig: @@ -77,9 +73,7 @@ class AlephAlphaConfig: - `control_log_additive` (boolean; default value: true): Method of applying control to attention scores. """ - maximum_tokens: Optional[int] = ( - litellm.max_tokens - ) # aleph alpha requires max tokens + maximum_tokens: Optional[int] = litellm.max_tokens # aleph alpha requires max tokens minimum_tokens: Optional[int] = None echo: Optional[bool] = None temperature: Optional[int] = None @@ -209,9 +203,7 @@ def completion( if "control" in model: # follow the ###Instruction / ###Response format for idx, message in enumerate(messages): if "role" in message: - if ( - idx == 0 - ): # set first message as instruction (required), let later user messages be input + if idx == 0: # set first message as instruction (required), let later user messages be input prompt += f"###Instruction: {message['content']}" else: if message["role"] == "system": diff --git a/litellm/llms/deprecated_providers/palm.py b/litellm/llms/deprecated_providers/palm.py index 657a6fdb229..a8523ecfa0e 100644 --- a/litellm/llms/deprecated_providers/palm.py +++ b/litellm/llms/deprecated_providers/palm.py @@ -19,9 +19,7 @@ class PalmError(Exception): url="https://developers.generativeai.google/api/python/google/generativeai/chat", ) self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class PalmConfig: @@ -102,9 +100,7 @@ def completion( try: import google.generativeai as palm # type: ignore except Exception: - raise Exception( - "Importing google.generativeai failed, please run 'pip install -q google-generativeai" - ) + raise Exception("Importing google.generativeai failed, please run 'pip install -q google-generativeai") palm.configure(api_key=api_key) model = model @@ -167,9 +163,7 @@ def completion( choices_list.append(choice_obj) model_response.choices = choices_list # type: ignore except Exception: - raise PalmError( - message=traceback.format_exc(), status_code=response.status_code - ) + raise PalmError(message=traceback.format_exc(), status_code=response.status_code) try: completion_response = model_response["choices"][0]["message"].get("content") @@ -181,9 +175,7 @@ def completion( ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here. prompt_tokens = len(encoding.encode(prompt)) - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) - ) + completion_tokens = len(encoding.encode(model_response["choices"][0]["message"].get("content", ""))) model_response.created = int(time.time()) model_response.model = "palm/" + model diff --git a/litellm/llms/docker_model_runner/chat/transformation.py b/litellm/llms/docker_model_runner/chat/transformation.py index dc03c80f154..137a39e0984 100644 --- a/litellm/llms/docker_model_runner/chat/transformation.py +++ b/litellm/llms/docker_model_runner/chat/transformation.py @@ -44,13 +44,9 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): """ messages = handle_messages_with_content_list_to_str_conversion(messages) if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] @@ -62,14 +58,10 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): The engine path should be included in the api_base. """ api_base = ( - api_base - or get_secret_str("DOCKER_MODEL_RUNNER_API_BASE") - or "http://localhost:22088/engines/llama.cpp" + api_base or get_secret_str("DOCKER_MODEL_RUNNER_API_BASE") or "http://localhost:22088/engines/llama.cpp" ) # type: ignore # Docker Model Runner may not require authentication for local instances - dynamic_api_key = ( - api_key or get_secret_str("DOCKER_MODEL_RUNNER_API_KEY") or "dummy-key" - ) + dynamic_api_key = api_key or get_secret_str("DOCKER_MODEL_RUNNER_API_KEY") or "dummy-key" return api_base, dynamic_api_key def get_complete_url( diff --git a/litellm/llms/duckduckgo/search/transformation.py b/litellm/llms/duckduckgo/search/transformation.py index e8eda3a37ab..0ef21222a29 100644 --- a/litellm/llms/duckduckgo/search/transformation.py +++ b/litellm/llms/duckduckgo/search/transformation.py @@ -80,11 +80,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): Get complete URL for Search endpoint. DuckDuckGo uses query parameters, so we construct the URL with the query. """ - api_base = ( - api_base - or get_secret_str("DUCKDUCKGO_API_BASE") - or self.DUCKDUCKGO_API_BASE - ) + api_base = api_base or get_secret_str("DUCKDUCKGO_API_BASE") or self.DUCKDUCKGO_API_BASE # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_duckduckgo_params" in data: diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index ecfc1642c97..a78f1d8541e 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -63,9 +63,7 @@ class E2BSandboxConfig(BaseSandboxConfig): "templateID": template or E2B_DEFAULT_TEMPLATE, "timeout": timeout if timeout is not None else DEFAULT_SANDBOX_TIMEOUT, "secure": True, - "allow_internet_access": ( - True if allow_internet_access is None else allow_internet_access - ), + "allow_internet_access": (True if allow_internet_access is None else allow_internet_access), } if metadata: body["metadata"] = metadata @@ -141,11 +139,7 @@ class E2BSandboxConfig(BaseSandboxConfig): **kwargs, ) -> bool: handle = self._as_handle(container) - key = ( - api_key - or handle._hidden_params.get("api_key") - or self.validate_environment() - ) + key = api_key or handle._hidden_params.get("api_key") or self.validate_environment() base = api_base or handle._hidden_params.get("api_base") or E2B_API_BASE try: response = cast( @@ -165,9 +159,7 @@ class E2BSandboxConfig(BaseSandboxConfig): def _as_handle(container: Union[ContainerHandle, str]) -> ContainerHandle: if isinstance(container, ContainerHandle): return container - handle = ContainerHandle( - id=str(container), provider="e2b", domain=E2B_DEFAULT_DOMAIN - ) + handle = ContainerHandle(id=str(container), provider="e2b", domain=E2B_DEFAULT_DOMAIN) handle._hidden_params = {} return handle @@ -180,20 +172,14 @@ class E2BSandboxConfig(BaseSandboxConfig): return None messages = tuple( - parsed - for line in lines - if (stripped := line.strip()) - if (parsed := _try_parse(stripped)) is not None + parsed for line in lines if (stripped := line.strip()) if (parsed := _try_parse(stripped)) is not None ) def of_type(message_type: str): return (m for m in messages if m.get("type") == message_type) error = next( - ( - {key: m.get(key) for key in ("name", "value", "traceback")} - for m in of_type("error") - ), + ({key: m.get(key) for key in ("name", "value", "traceback")} for m in of_type("error")), None, ) execution_count = next( @@ -204,9 +190,7 @@ class E2BSandboxConfig(BaseSandboxConfig): return CodeExecutionResult( stdout="".join(m.get("text", "") for m in of_type("stdout")), stderr="".join(m.get("text", "") for m in of_type("stderr")), - results=[ - {k: v for k, v in m.items() if k != "type"} for m in of_type("result") - ], + results=[{k: v for k, v in m.items() if k != "type"} for m in of_type("result")], error=error, execution_count=execution_count, ) diff --git a/litellm/llms/elevenlabs/audio_transcription/transformation.py b/litellm/llms/elevenlabs/audio_transcription/transformation.py index 8746e92d9f6..68d1b5e16dd 100644 --- a/litellm/llms/elevenlabs/audio_transcription/transformation.py +++ b/litellm/llms/elevenlabs/audio_transcription/transformation.py @@ -28,9 +28,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): def custom_llm_provider(self) -> str: return litellm.LlmProviders.ELEVENLABS.value - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: return ["language", "temperature"] def map_openai_params( @@ -50,12 +48,8 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): optional_params[k] = v return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return ElevenLabsException( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return ElevenLabsException(message=error_message, status_code=status_code, headers=headers) def transform_audio_transcription_request( self, @@ -152,9 +146,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return response except Exception as e: - raise ValueError( - f"Error transforming ElevenLabs response: {str(e)}\nResponse: {raw_response.text}" - ) + raise ValueError(f"Error transforming ElevenLabs response: {str(e)}\nResponse: {raw_response.text}") def get_complete_url( self, @@ -166,9 +158,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): stream: Optional[bool] = None, ) -> str: if api_base is None: - api_base = ( - get_secret_str("ELEVENLABS_API_BASE") or "https://api.elevenlabs.io" - ) + api_base = get_secret_str("ELEVENLABS_API_BASE") or "https://api.elevenlabs.io" api_base = api_base.rstrip("/") # Remove trailing slash if present # ElevenLabs speech-to-text endpoint @@ -188,9 +178,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> dict: api_key = api_key or get_secret_str("ELEVENLABS_API_KEY") if api_key is None: - raise ValueError( - "ElevenLabs API key is required. Set ELEVENLABS_API_KEY environment variable." - ) + raise ValueError("ElevenLabs API key is required. Set ELEVENLABS_API_KEY environment variable.") auth_header = { "xi-api-key": api_key, diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index 612fc687ef9..b5b7799a3e9 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -105,9 +105,7 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): mapped_voice = self._extract_voice_id(voice_override) if mapped_voice is None: - raise ValueError( - "ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`." - ) + raise ValueError("ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`.") return mapped_voice @@ -175,17 +173,10 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): """ Validate Azure environment and set up authentication headers """ - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("ELEVENLABS_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("ELEVENLABS_API_KEY") if api_key is None: - raise ValueError( - "ElevenLabs API key is required. Set ELEVENLABS_API_KEY environment variable." - ) + raise ValueError("ElevenLabs API key is required. Set ELEVENLABS_API_KEY environment variable.") headers.update( { @@ -196,12 +187,8 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): return headers - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return ElevenLabsException( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return ElevenLabsException(message=error_message, status_code=status_code, headers=headers) def transform_text_to_speech_request( self, @@ -310,16 +297,12 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): """ Construct the ElevenLabs endpoint URL, including path voice_id and query params. """ - base_url = ( - api_base or get_secret_str("ELEVENLABS_API_BASE") or self.TTS_BASE_URL - ) + base_url = api_base or get_secret_str("ELEVENLABS_API_BASE") or self.TTS_BASE_URL base_url = base_url.rstrip("/") voice_id = litellm_params.get(self.ELEVENLABS_VOICE_ID_KEY) if not isinstance(voice_id, str) or not voice_id.strip(): - raise ValueError( - "ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`." - ) + raise ValueError("ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`.") encoded_voice_id = encode_url_path_segment(voice_id, field_name="voice_id") url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}" diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 5cfd14aeaa9..93fbdeff990 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -40,9 +40,7 @@ class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False): startPublishedDate: str # Optional - published date filter (ISO 8601 format) endPublishedDate: str # Optional - published date filter (ISO 8601 format) includeText: List[str] # Optional - strings that must be present in webpage text - excludeText: List[ - str - ] # Optional - strings that must not be present in webpage text + excludeText: List[str] # Optional - strings that must not be present in webpage text context: Union[bool, dict] # Optional - format results for LLMs moderation: bool # Optional - enable content moderation, default false contents: dict # Optional - content retrieval options @@ -73,9 +71,7 @@ class ExaAISearchConfig(BaseSearchConfig): default_api_base=self.EXA_AI_API_BASE, ) if not api_key: - raise ValueError( - "EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable." - ) + raise ValueError("EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable.") headers["x-api-key"] = api_key headers["Content-Type"] = "application/json" return headers @@ -146,10 +142,7 @@ class ExaAISearchConfig(BaseSearchConfig): # pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value # By default, request text content if not explicitly specified diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index 9cdd0cd485b..6e32141afd4 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -22,6 +22,4 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/llms/fal_ai/image_generation/__init__.py b/litellm/llms/fal_ai/image_generation/__init__.py index 7f3358934a7..d31524510b8 100644 --- a/litellm/llms/fal_ai/image_generation/__init__.py +++ b/litellm/llms/fal_ai/image_generation/__init__.py @@ -59,11 +59,7 @@ def get_fal_ai_image_generation_config(model: str) -> BaseImageGenerationConfig: if "ultra" in model_lower: return FalAIFluxProV11UltraConfig() return FalAIFluxProV11Config() - elif ( - "flux/schnell" in model_lower - or "flux-schnell" in model_lower - or "schnell" in model_lower - ): + elif "flux/schnell" in model_lower or "flux-schnell" in model_lower or "schnell" in model_lower: return FalAIFluxSchnellConfig() elif "bytedance/seedream" in model_lower: return FalAIBytedanceSeedreamV3Config() diff --git a/litellm/llms/fal_ai/image_generation/bria_transformation.py b/litellm/llms/fal_ai/image_generation/bria_transformation.py index dd6e737324e..7bdfa860c5d 100644 --- a/litellm/llms/fal_ai/image_generation/bria_transformation.py +++ b/litellm/llms/fal_ai/image_generation/bria_transformation.py @@ -28,9 +28,7 @@ class FalAIBriaConfig(FalAIBaseConfig): IMAGE_GENERATION_ENDPOINT: str = "bria/text-to-image/3.2" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for Bria 3.2. """ diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index fef292d3311..fb980905a28 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -28,9 +28,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): IMAGE_GENERATION_ENDPOINT: str = "fal-ai/flux-pro/v1.1-ultra" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for Flux Pro v1.1-ultra. """ @@ -256,8 +254,6 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): if "timings" in response_data: model_response._hidden_params["timings"] = response_data["timings"] if "has_nsfw_concepts" in response_data: - model_response._hidden_params["has_nsfw_concepts"] = response_data[ - "has_nsfw_concepts" - ] + model_response._hidden_params["has_nsfw_concepts"] = response_data["has_nsfw_concepts"] return model_response diff --git a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py index 14e136d5d6f..500a4b20ef2 100644 --- a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py @@ -38,9 +38,7 @@ class FalAIIdeogramV3Config(FalAIBaseConfig): "1024x1536": "portrait_16_9", } - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Ideogram v3 accepts the core OpenAI image parameters. """ diff --git a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py index ea6e7c1f3c9..1b111c98987 100644 --- a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py +++ b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py @@ -31,9 +31,7 @@ class FalAIImagen4Config(FalAIBaseConfig): IMAGE_GENERATION_ENDPOINT: str = "fal-ai/imagen4/preview" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for Imagen4. """ diff --git a/litellm/llms/fal_ai/image_generation/nano_banana_transformation.py b/litellm/llms/fal_ai/image_generation/nano_banana_transformation.py index dd4758055ac..0a8ba3699bb 100644 --- a/litellm/llms/fal_ai/image_generation/nano_banana_transformation.py +++ b/litellm/llms/fal_ai/image_generation/nano_banana_transformation.py @@ -40,15 +40,11 @@ class FalAINanoBananaConfig(FalAIBaseConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - base_url: str = ( - api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL - ).rstrip("/") + base_url: str = (api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL).rstrip("/") endpoint = model if model.startswith("fal-ai/") else f"fal-ai/{model}" return f"{base_url}/{endpoint}" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: return ["n", "response_format", "size"] def map_openai_params( diff --git a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py index 72ee165b51a..2ce36d9c1ea 100644 --- a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py @@ -28,9 +28,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): IMAGE_GENERATION_ENDPOINT: str = "fal-ai/recraft/v3/text-to-image" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for Recraft v3. """ diff --git a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py index f0077c6a674..bc7a3839bd3 100644 --- a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py +++ b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py @@ -46,9 +46,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): """ from litellm.secret_managers.main import get_secret_str - complete_url: str = ( - api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url: str = api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") @@ -65,9 +63,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): complete_url = f"{complete_url}/{endpoint}" return complete_url - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for Stable Diffusion models. """ @@ -272,8 +268,6 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): if "timings" in response_data: model_response._hidden_params["timings"] = response_data["timings"] if "has_nsfw_concepts" in response_data: - model_response._hidden_params["has_nsfw_concepts"] = response_data[ - "has_nsfw_concepts" - ] + model_response._hidden_params["has_nsfw_concepts"] = response_data["has_nsfw_concepts"] return model_response diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index 4a0dea48a10..07eb2cc4cc4 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -43,9 +43,7 @@ class FalAIBaseConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = ( - api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url: str = api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") if self.IMAGE_GENERATION_ENDPOINT: @@ -124,9 +122,7 @@ class FalAIImageGenerationConfig(FalAIBaseConfig): Default Fal AI image generation configuration for generic models. """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for fal.ai image generation """ diff --git a/litellm/llms/fastcrw/search/transformation.py b/litellm/llms/fastcrw/search/transformation.py index b571a659cac..6de9ef642fb 100644 --- a/litellm/llms/fastcrw/search/transformation.py +++ b/litellm/llms/fastcrw/search/transformation.py @@ -34,9 +34,7 @@ class FastCRWSearchRequest(_FastCRWSearchRequestRequired, total=False): """ limit: int # Optional - maximum number of results to return - sources: list[ - str - ] # Optional - sources to search ('web', 'images'), default ['web'] + sources: list[str] # Optional - sources to search ('web', 'images'), default ['web'] scrapeOptions: dict # Optional - options for scraping search results @@ -65,9 +63,7 @@ class FastCRWSearchConfig(BaseSearchConfig): default_api_base=self.FASTCRW_API_BASE, ) if not api_key: - raise ValueError( - "CRW_API_KEY is not set. Set `CRW_API_KEY` environment variable." - ) + raise ValueError("CRW_API_KEY is not set. Set `CRW_API_KEY` environment variable.") headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -129,10 +125,7 @@ class FastCRWSearchConfig(BaseSearchConfig): # pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value # By default, request markdown content if not explicitly specified diff --git a/litellm/llms/featherless_ai/chat/transformation.py b/litellm/llms/featherless_ai/chat/transformation.py index e62108624d3..cf11c72c326 100644 --- a/litellm/llms/featherless_ai/chat/transformation.py +++ b/litellm/llms/featherless_ai/chat/transformation.py @@ -107,11 +107,7 @@ class FeatherlessAIConfig(OpenAIGPTConfig): or get_secret_str("FEATHERLESS_API_BASE") or "https://api.featherless.ai/v1" ) - dynamic_api_key = ( - api_key - or get_secret_str("FEATHERLESS_AI_API_KEY") - or get_secret_str("FEATHERLESS_API_KEY") - ) + dynamic_api_key = api_key or get_secret_str("FEATHERLESS_AI_API_KEY") or get_secret_str("FEATHERLESS_API_KEY") return api_base, dynamic_api_key def validate_environment( diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py index 7e01ba58706..7aac6d7e7dd 100644 --- a/litellm/llms/firecrawl/search/transformation.py +++ b/litellm/llms/firecrawl/search/transformation.py @@ -30,12 +30,8 @@ class FirecrawlSearchRequest(_FirecrawlSearchRequestRequired, total=False): """ limit: int # Optional - maximum number of results to return (default 5, max 100) - sources: List[ - str - ] # Optional - sources to search ('web', 'images', 'news'), default ['web'] - categories: List[ - Dict[str, str] - ] # Optional - categories to filter by (github, research, pdf) + sources: List[str] # Optional - sources to search ('web', 'images', 'news'), default ['web'] + categories: List[Dict[str, str]] # Optional - categories to filter by (github, research, pdf) tbs: str # Optional - time-based search parameter location: str # Optional - location parameter for geo-targeting country: str # Optional - ISO country code (default 'US') @@ -69,9 +65,7 @@ class FirecrawlSearchConfig(BaseSearchConfig): default_api_base=self.FIRECRAWL_API_BASE, ) if not api_key: - raise ValueError( - "FIRECRAWL_API_KEY is not set. Set `FIRECRAWL_API_KEY` environment variable." - ) + raise ValueError("FIRECRAWL_API_KEY is not set. Set `FIRECRAWL_API_KEY` environment variable.") headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -86,9 +80,7 @@ class FirecrawlSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - api_base = ( - api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE - ) + api_base = api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE # Append "/search" to the api base if it's not already there if not api_base.endswith("/search"): @@ -141,10 +133,7 @@ class FirecrawlSearchConfig(BaseSearchConfig): # pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value # By default, request markdown content if not explicitly specified diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 3c93677d6d4..d4258557fe7 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -60,9 +60,7 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: """ choices = [c for c in (payload.get("choices") or []) if isinstance(c, dict)] top_level = { - f"fireworks_{field}": payload[field] - for field in ("perf_metrics", "prompt_token_ids") - if field in payload + f"fireworks_{field}": payload[field] for field in ("perf_metrics", "prompt_token_ids") if field in payload } per_choice = { f"fireworks_{dest}": [c[field] for c in choices if field in c] @@ -204,14 +202,8 @@ class FireworksAIConfig(OpenAIGPTConfig): drop_params: bool, ) -> dict: supported_openai_params = self.get_supported_openai_params(model=model) - is_tools_set = any( - param == "tools" and value is not None - for param, value in non_default_params.items() - ) - if ( - non_default_params.get("thinking") is not None - and non_default_params.get("reasoning_effort") is not None - ): + is_tools_set = any(param == "tools" and value is not None for param, value in non_default_params.items()) + if non_default_params.get("thinking") is not None and non_default_params.get("reasoning_effort") is not None: raise litellm.BadRequestError( message=( "Fireworks AI chat completions does not support specifying both " @@ -230,9 +222,7 @@ class FireworksAIConfig(OpenAIGPTConfig): # pass through the value of tool choice optional_params["tool_choice"] = value elif param == "response_format": - if ( - is_tools_set - ): # fireworks ai doesn't support tools and response_format together + if is_tools_set: # fireworks ai doesn't support tools and response_format together optional_params = self._add_response_format_to_tools( optional_params=optional_params, value=value, @@ -256,9 +246,7 @@ class FireworksAIConfig(OpenAIGPTConfig): return optional_params - def _transform_tools( - self, tools: List[OpenAIChatCompletionToolParam] - ) -> List[OpenAIChatCompletionToolParam]: + def _transform_tools(self, tools: List[OpenAIChatCompletionToolParam]) -> List[OpenAIChatCompletionToolParam]: for tool in tools: if tool.get("type") != "function": continue @@ -279,9 +267,7 @@ class FireworksAIConfig(OpenAIGPTConfig): filter_value_from_dict, ) - supports_vision_value = self._get_model_cost_capability_exact( - model=model, capability="supports_vision" - ) + supports_vision_value = self._get_model_cost_capability_exact(model=model, capability="supports_vision") for message in messages: if message["role"] == "user": _message_content = message.get("content") @@ -301,10 +287,7 @@ class FireworksAIConfig(OpenAIGPTConfig): model=model, llm_provider="fireworks_ai", ) - if ( - content.get("type") == "image_url" - and supports_vision_value is False - ): + if content.get("type") == "image_url" and supports_vision_value is False: raise litellm.BadRequestError( message=( f"Fireworks AI model {model} does not support " @@ -338,11 +321,7 @@ class FireworksAIConfig(OpenAIGPTConfig): model_cost = litellm.model_cost signature = (id(model_cost), get_model_cost_mutation_generation()) cached = cls._fireworks_index_cache - if ( - cached is not None - and cached[0] == signature[0] - and cached[1] == signature[1] - ): + if cached is not None and cached[0] == signature[0] and cached[1] == signature[1]: return cached[2] index: List[Tuple[str, dict]] = [] @@ -384,9 +363,7 @@ class FireworksAIConfig(OpenAIGPTConfig): short_name = short_name[len("accounts/fireworks/models/") :] return short_name - def _get_model_cost_capability_exact( - self, model: str, capability: str - ) -> Optional[bool]: + def _get_model_cost_capability_exact(self, model: str, capability: str) -> Optional[bool]: short_name = self._short_model_name(model) candidate_keys = ( model, @@ -400,9 +377,7 @@ class FireworksAIConfig(OpenAIGPTConfig): return None def _get_model_cost_capability(self, model: str, capability: str) -> Optional[bool]: - exact = self._get_model_cost_capability_exact( - model=model, capability=capability - ) + exact = self._get_model_cost_capability_exact(model=model, capability=capability) if exact is not None: return exact @@ -418,8 +393,7 @@ class FireworksAIConfig(OpenAIGPTConfig): matches = [ (key_short, cast(Optional[bool], model_info.get(capability))) for key_short, model_info in self._get_fireworks_index() - if model_info.get(capability) is not None - and self._matches_on_hyphen_boundary(short_name, key_short) + if model_info.get(capability) is not None and self._matches_on_hyphen_boundary(short_name, key_short) ] if not matches: return None @@ -429,15 +403,9 @@ class FireworksAIConfig(OpenAIGPTConfig): supports_function_calling_value = self._get_model_cost_capability( model=model, capability="supports_function_calling" ) - supports_reasoning_value = self._get_model_cost_capability( - model=model, capability="supports_reasoning" - ) - supports_vision_value = self._get_model_cost_capability( - model=model, capability="supports_vision" - ) - supports_pdf_input_value = self._get_model_cost_capability( - model=model, capability="supports_pdf_input" - ) + supports_reasoning_value = self._get_model_cost_capability(model=model, capability="supports_reasoning") + supports_vision_value = self._get_model_cost_capability(model=model, capability="supports_vision") + supports_pdf_input_value = self._get_model_cost_capability(model=model, capability="supports_pdf_input") provider_specific_model_info: ProviderSpecificModelInfo = { "supports_function_calling": True, @@ -445,23 +413,17 @@ class FireworksAIConfig(OpenAIGPTConfig): } if supports_function_calling_value is not None: - provider_specific_model_info["supports_function_calling"] = ( - supports_function_calling_value - ) + provider_specific_model_info["supports_function_calling"] = supports_function_calling_value # Only include supports_reasoning if True if supports_reasoning_value: - provider_specific_model_info["supports_reasoning"] = ( - supports_reasoning_value - ) + provider_specific_model_info["supports_reasoning"] = supports_reasoning_value if supports_vision_value is not None: provider_specific_model_info["supports_vision"] = supports_vision_value if supports_pdf_input_value is not None: - provider_specific_model_info["supports_pdf_input"] = ( - supports_pdf_input_value - ) + provider_specific_model_info["supports_pdf_input"] = supports_pdf_input_value return provider_specific_model_info @@ -478,9 +440,7 @@ class FireworksAIConfig(OpenAIGPTConfig): model = f"accounts/fireworks/routers/{model}" else: model = f"accounts/fireworks/models/{model}" - messages = self._transform_messages_helper( - messages=messages, model=model, litellm_params=litellm_params - ) + messages = self._transform_messages_helper(messages=messages, model=model, litellm_params=litellm_params) if "tools" in optional_params and optional_params["tools"] is not None: tools = self._transform_tools(tools=optional_params["tools"]) optional_params["tools"] = tools @@ -511,19 +471,13 @@ class FireworksAIConfig(OpenAIGPTConfig): Relevant Issue: https://github.com/BerriAI/litellm/issues/7209#issuecomment-2813208780 """ - if ( - tool_calls is not None - and message.content is not None - and message.tool_calls is None - ): + if tool_calls is not None and message.content is not None and message.tool_calls is None: try: function = Function(**json.loads(message.content)) if function.name != RESPONSE_FORMAT_TOOL_NAME and function.name in [ tool["function"]["name"] for tool in tool_calls ]: - tool_call = ChatCompletionMessageToolCall( - function=function, id=str(uuid.uuid4()), type="function" - ) + tool_call = ChatCompletionMessageToolCall(function=function, id=str(uuid.uuid4()), type="function") message.tool_calls = [tool_call] message.content = None @@ -560,9 +514,7 @@ class FireworksAIConfig(OpenAIGPTConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise FireworksAIException( - message="Unable to get json response - {}, Original Response: {}".format( - str(e), raw_response.text - ), + message="Unable to get json response - {}, Original Response: {}".format(str(e), raw_response.text), status_code=raw_response.status_code, headers=response_headers, ) @@ -578,9 +530,7 @@ class FireworksAIConfig(OpenAIGPTConfig): ## FIREWORKS AI sends tool calls in the content field instead of tool_calls for choice in response.choices: - cast( - Choices, choice - ).message = self._handle_message_content_with_tool_calls( + cast(Choices, choice).message = self._handle_message_content_with_tool_calls( message=cast(Choices, choice).message, tool_calls=optional_params.get("tools", None), ) @@ -607,11 +557,7 @@ class FireworksAIConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = ( - api_base - or get_secret_str("FIREWORKS_API_BASE") - or "https://api.fireworks.ai/inference/v1" - ) # type: ignore + api_base = api_base or get_secret_str("FIREWORKS_API_BASE") or "https://api.fireworks.ai/inference/v1" # type: ignore dynamic_api_key = api_key or ( get_secret_str("FIREWORKS_API_KEY") or get_secret_str("FIREWORKS_AI_API_KEY") @@ -621,9 +567,7 @@ class FireworksAIConfig(OpenAIGPTConfig): return api_base, dynamic_api_key def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None): - api_base, api_key = self._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + api_base, api_key = self._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) if api_base is None or api_key is None: raise ValueError( "FIREWORKS_API_BASE or FIREWORKS_API_KEY is not set. Please set the environment variable, to query Fireworks AI's `/models` endpoint." diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 17aa67b525b..a1b6309d1e0 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -17,9 +17,7 @@ class FireworksAIMixin: Common Base Config functions across Fireworks AI Endpoints """ - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: return FireworksAIException( status_code=status_code, message=error_message, diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index 46026f266d6..ed936f6233a 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -72,9 +72,7 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: base_model = get_base_model_for_pricing(model_name=model) ## GET MODEL INFO - model_info = get_model_info( - model=base_model, custom_llm_provider="fireworks_ai" - ) + model_info = get_model_info(model=base_model, custom_llm_provider="fireworks_ai") ## CALCULATE INPUT COST diff --git a/litellm/llms/fireworks_ai/embed/fireworks_ai_transformation.py b/litellm/llms/fireworks_ai/embed/fireworks_ai_transformation.py index 80906443984..414c4dcef68 100644 --- a/litellm/llms/fireworks_ai/embed/fireworks_ai_transformation.py +++ b/litellm/llms/fireworks_ai/embed/fireworks_ai_transformation.py @@ -17,9 +17,7 @@ class FireworksAIEmbeddingConfig: return ["dimensions"] return [] - def map_openai_params( - self, non_default_params: dict, optional_params: dict, model: str - ): + def map_openai_params(self, non_default_params: dict, optional_params: dict, model: str): """ No transformation is applied - fireworks ai is openai compatible """ diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 400d511a02f..393a6c5a8e5 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -154,19 +154,11 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): "documents": optional_rerank_params["documents"], } - if ( - "top_n" in optional_rerank_params - and optional_rerank_params["top_n"] is not None - ): + if "top_n" in optional_rerank_params and optional_rerank_params["top_n"] is not None: request_data["top_n"] = optional_rerank_params["top_n"] - if ( - "return_documents" in optional_rerank_params - and optional_rerank_params["return_documents"] is not None - ): - request_data["return_documents"] = optional_rerank_params[ - "return_documents" - ] + if "return_documents" in optional_rerank_params and optional_rerank_params["return_documents"] is not None: + request_data["return_documents"] = optional_rerank_params["return_documents"] return request_data @@ -221,9 +213,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) # Extract results - Fireworks AI uses "data" instead of "results" - _results: List[dict] | None = raw_response_json.get( - "data" - ) or raw_response_json.get("results") + _results: List[dict] | None = raw_response_json.get("data") or raw_response_json.get("results") if _results is None: raise ValueError(f"No results found in the response={raw_response_json}") @@ -261,11 +251,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): rerank_results.append(rerank_result) # Use model name as id if no id is provided - response_id = ( - raw_response_json.get("id") - or raw_response_json.get("model") - or str(uuid.uuid4()) - ) + response_id = raw_response_json.get("id") or raw_response_json.get("model") or str(uuid.uuid4()) return RerankResponse( id=response_id, diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index f6e0b95cf28..9e1f6935da4 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -113,10 +113,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): ) api_key = GeminiModelInfo.get_api_key(explicit_api_key) if not api_key: - raise ValueError( - "Google API key is required. " - "Set GOOGLE_API_KEY or GEMINI_API_KEY, or pass api_key." - ) + raise ValueError("Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY, or pass api_key.") headers["x-goog-api-key"] = api_key return headers @@ -289,9 +286,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): data = raw_response.json() except Exception: data = {} - verbose_logger.debug( - "GeminiAgentsConfig list_versions response for '%s': %s", name, data - ) + verbose_logger.debug("GeminiAgentsConfig list_versions response for '%s': %s", name, data) return AgentVersionsResponse( agent_versions=data.get("agentVersions", []), next_page_token=data.get("nextPageToken"), diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 4e9764446c9..94130ac4a6e 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -130,14 +130,8 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): else: _image_url = img_element.get("image_url") # type: ignore if _image_url and "https://" in _image_url: - image_obj = convert_to_anthropic_image_obj( - _image_url, format=format - ) - converted_image_url = ( - convert_generic_image_chunk_to_openai_image_obj( - image_obj - ) - ) + image_obj = convert_to_anthropic_image_obj(_image_url, format=format) + converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj) if detail is not None: img_element["image_url"] = { # type: ignore "url": converted_image_url, diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 4cca2e2b850..f02e25c5735 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -95,17 +95,13 @@ GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO: Dict[tuple[int, int], str] = { } -def map_openai_size_to_gemini_image_config( - size: str, model: str -) -> Optional[Dict[str, str]]: +def map_openai_size_to_gemini_image_config(size: str, model: str) -> Optional[Dict[str, str]]: dimensions = _parse_openai_image_size(size) if dimensions is None: return None width, height = dimensions - image_config = { - "aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height) - } + image_config = {"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height)} image_size = _map_dimensions_to_gemini_image_size(width, height) if is_gemini_image_model(model): if supports_gemini_image_size(model): @@ -139,9 +135,7 @@ def map_openai_image_params_to_gemini( parse_image_config_string: bool = False, ) -> Dict[str, Any]: optional_params = optional_params or {} - filtered_params = { - key: value for key, value in params.items() if key in supported_params - } + filtered_params = {key: value for key, value in params.items() if key in supported_params} mapped_params: Dict[str, Any] = {} @@ -174,10 +168,7 @@ def map_openai_image_params_to_gemini( mapped_params["imageConfig"] = image_config_param for key, value in filtered_params.items(): - if ( - key not in ("n", "size", "imageConfig", "tools", "web_search_options") - and key not in optional_params - ): + if key not in ("n", "size", "imageConfig", "tools", "web_search_options") and key not in optional_params: mapped_params[key] = value return mapped_params @@ -217,10 +208,7 @@ def _has_gemini_search_tool(tools: List[Any]) -> bool: ) search_tool_keys = VertexGeminiConfig._search_tool_keys() - return any( - isinstance(tool, dict) and any(key in tool for key in search_tool_keys) - for tool in tools - ) + return any(isinstance(tool, dict) and any(key in tool for key in search_tool_keys) for tool in tools) def map_gemini_image_tools_params( @@ -237,9 +225,7 @@ def map_gemini_image_tools_params( tools_value = non_default_params.get("tools") if isinstance(tools_value, list) and tools_value: - mapped_tools = gemini_config._map_function( - value=tools_value, optional_params=result - ) + mapped_tools = gemini_config._map_function(value=tools_value, optional_params=result) result = gemini_config._add_tools_to_optional_params(result, mapped_tools) web_search_options = non_default_params.get("web_search_options") @@ -335,9 +321,7 @@ def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int) -> str: requested_ratio = width / height return min( GEMINI_IMAGE_ASPECT_RATIOS, - key=lambda aspect_ratio: abs( - math.log(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio] / requested_ratio) - ), + key=lambda aspect_ratio: abs(math.log(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio] / requested_ratio)), ) @@ -376,19 +360,11 @@ class GeminiModelInfo(BaseLLMModelInfo): @staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: - return ( - api_base - or get_secret_str("GEMINI_API_BASE") - or "https://generativelanguage.googleapis.com" - ) + return api_base or get_secret_str("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: - return ( - api_key - or (get_secret_str("GOOGLE_API_KEY")) - or (get_secret_str("GEMINI_API_KEY")) - ) + return api_key or (get_secret_str("GOOGLE_API_KEY")) or (get_secret_str("GEMINI_API_KEY")) @staticmethod def get_base_model(model: str) -> Optional[str]: @@ -402,9 +378,7 @@ class GeminiModelInfo(BaseLLMModelInfo): litellm_model_names.append(litellm_model_name) return litellm_model_names - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: api_base = GeminiModelInfo.get_api_base(api_base) api_key = GeminiModelInfo.get_api_key(api_key) endpoint = f"/{self.api_version}/models" @@ -431,9 +405,7 @@ class GeminiModelInfo(BaseLLMModelInfo): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return GeminiError( - status_code=status_code, message=error_message, headers=headers - ) + return GeminiError(status_code=status_code, message=error_message, headers=headers) def get_token_counter(self) -> Optional[BaseTokenCounter]: """ @@ -446,9 +418,7 @@ class GeminiModelInfo(BaseLLMModelInfo): return GoogleAIStudioTokenCounter() -def encode_unserializable_types( - data: Dict[str, object], depth: int = 0 -) -> Dict[str, object]: +def encode_unserializable_types(data: Dict[str, object], depth: int = 0) -> Dict[str, object]: """Converts unserializable types in dict to json.dumps() compatible types. This function is called in models.py after calling convert_to_dict(). The @@ -476,15 +446,11 @@ def encode_unserializable_types( processed_data[key] = encode_unserializable_types(value, depth + 1) elif isinstance(value, list): if all(isinstance(v, bytes) for v in value): - processed_data[key] = [ - base64.urlsafe_b64encode(v).decode("ascii") for v in value - ] + processed_data[key] = [base64.urlsafe_b64encode(v).decode("ascii") for v in value] if all(isinstance(v, datetime.datetime) for v in value): processed_data[key] = [v.isoformat() for v in value] else: - processed_data[key] = [ - encode_unserializable_types(v, depth + 1) for v in value - ] + processed_data[key] = [encode_unserializable_types(v, depth + 1) for v in value] else: processed_data[key] = value return processed_data @@ -520,9 +486,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter deployment = deployment or {} - count_tokens_params_request = copy.deepcopy( - deployment.get("litellm_params", {}) - ) + count_tokens_params_request = copy.deepcopy(deployment.get("litellm_params", {})) count_tokens_params = { "model": model_to_use, "contents": contents, diff --git a/litellm/llms/gemini/cost_calculator.py b/litellm/llms/gemini/cost_calculator.py index 7b9bd7486a5..f69cfe03270 100644 --- a/litellm/llms/gemini/cost_calculator.py +++ b/litellm/llms/gemini/cost_calculator.py @@ -10,9 +10,7 @@ if TYPE_CHECKING: from litellm.types.utils import ModelInfo, Usage -def cost_per_token( - model: str, usage: "Usage", service_tier: Optional[str] = None -) -> Tuple[float, float]: +def cost_per_token(model: str, usage: "Usage", service_tier: Optional[str] = None) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index fdb77452d4c..27df584d476 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -43,9 +43,7 @@ class GoogleAIStudioTokenCounter: function_response_data = part["functionResponse"] function_response_part = FunctionResponse(**function_response_data) function_response_part.id = None - part["functionResponse"] = function_response_part.model_dump( - exclude_none=True - ) + part["functionResponse"] = function_response_part.model_dump(exclude_none=True) return cleaned_contents @@ -139,9 +137,7 @@ class GoogleAIStudioTokenCounter: ) try: - response = await async_httpx_client.post( - url=url, headers=headers, json=request_body - ) + response = await async_httpx_client.post(url=url, headers=headers, json=request_body) # Check for HTTP errors response.raise_for_status() @@ -160,9 +156,7 @@ class GoogleAIStudioTokenCounter: ) from e except httpx.RequestError as e: error_msg = f"Request to Google Gen AI Studio failed: {str(e)}" - raise litellm.APIConnectionError( - message=error_msg, llm_provider="gemini", model=model - ) from e + raise litellm.APIConnectionError(message=error_msg, llm_provider="gemini", model=model) from e except Exception as e: error_msg = f"Unexpected error during token counting: {str(e)}" raise Exception(error_msg) from e diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 63a383ebd3d..a18dc152cb6 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -55,9 +55,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): """ resolved_api_key = self.get_api_key(api_key) if not resolved_api_key: - raise ValueError( - "GEMINI_API_KEY is required for Google AI Studio file operations" - ) + raise ValueError("GEMINI_API_KEY is required for Google AI Studio file operations") headers["x-goog-api-key"] = resolved_api_key return headers @@ -91,9 +89,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): url = "{}/{}".format(api_base, endpoint) return url - def get_supported_openai_params( - self, model: str - ) -> List[OpenAICreateFileRequestOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAICreateFileRequestOptionalParams]: return [] def map_openai_params( @@ -140,11 +136,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): headers.update(extracted_data["headers"]) # Add any custom headers # Initial metadata request body - initial_data = { - "file": { - "display_name": extracted_data["filename"] or str(int(time.time())) - } - } + initial_data = {"file": {"display_name": extracted_data["filename"] or str(int(time.time()))}} # Step 2: Actual file upload data upload_headers = { @@ -182,9 +174,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): return OpenAIFileObject( id=response_object["uri"], # Gemini uses URI as identifier - bytes=int( - response_object["sizeBytes"] - ), # Gemini doesn't return file size + bytes=int(response_object["sizeBytes"]), # Gemini doesn't return file size created_at=int( time.mktime( time.strptime( @@ -227,10 +217,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): file_part = self._normalize_gemini_file_id(file_id) - api_base = ( - self.get_api_base(litellm_params.get("api_base")) - or "https://generativelanguage.googleapis.com" - ) + api_base = self.get_api_base(litellm_params.get("api_base")) or "https://generativelanguage.googleapis.com" api_base = api_base.rstrip("/") url = f"{api_base}/v1beta/{file_part}" @@ -262,9 +249,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if normalized_file_id.startswith("files/"): normalized_file_id = normalized_file_id.removeprefix("files/") - encoded_file_id = encode_url_path_segment( - normalized_file_id, field_name="file_id" - ) + encoded_file_id = encode_url_path_segment(normalized_file_id, field_name="file_id") return f"files/{encoded_file_id}" @@ -306,11 +291,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): object="file", purpose="user_data", status=status, - status_details=( - str(response_json.get("error", "")) - if gemini_state == "FAILED" - else None - ), + status_details=(str(response_json.get("error", "")) if gemini_state == "FAILED" else None), ) except Exception as e: verbose_logger.exception(f"Error parsing file retrieve response: {str(e)}") @@ -390,9 +371,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError( - "GoogleAIStudioFilesHandler does not support file listing" - ) + raise NotImplementedError("GoogleAIStudioFilesHandler does not support file listing") def transform_list_files_response( self, @@ -400,9 +379,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> List[OpenAIFileObject]: - raise NotImplementedError( - "GoogleAIStudioFilesHandler does not support file listing" - ) + raise NotImplementedError("GoogleAIStudioFilesHandler does not support file listing") def transform_file_content_request( self, @@ -410,9 +387,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError( - "GoogleAIStudioFilesHandler does not support file content retrieval" - ) + raise NotImplementedError("GoogleAIStudioFilesHandler does not support file content retrieval") def transform_file_content_response( self, @@ -420,6 +395,4 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: - raise NotImplementedError( - "GoogleAIStudioFilesHandler does not support file content retrieval" - ) + raise NotImplementedError("GoogleAIStudioFilesHandler does not support file content retrieval") diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index ee201af7e1a..68f30308621 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -118,17 +118,11 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): ) _generate_content_config_dict: Dict[str, Any] = {} - supported_google_genai_params = ( - self.get_supported_generate_content_optional_params(model) - ) + supported_google_genai_params = self.get_supported_generate_content_optional_params(model) # Create a set with both camelCase and snake_case versions for faster lookup supported_params_set = set(supported_google_genai_params) - supported_params_set.update( - _snake_to_camel(p) for p in supported_google_genai_params - ) - supported_params_set.update( - _camel_to_snake(p) for p in supported_google_genai_params if "_" not in p - ) + supported_params_set.update(_snake_to_camel(p) for p in supported_google_genai_params) + supported_params_set.update(_camel_to_snake(p) for p in supported_google_genai_params if "_" not in p) for param, value in generate_content_config_dict.items(): # Google GenAI API expects camelCase, so we'll always output in camelCase @@ -160,9 +154,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): "Content-Type": "application/json", } # Use the passed api_key first, then fall back to litellm_params and environment - gemini_api_key = api_key or self._get_google_ai_studio_api_key( - dict(litellm_params or {}) - ) + gemini_api_key = api_key or self._get_google_ai_studio_api_key(dict(litellm_params or {})) if isinstance(gemini_api_key, dict): default_headers.update(gemini_api_key) elif gemini_api_key is not None: @@ -308,23 +300,13 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): ) @staticmethod - def _normalize_response_schema( - generate_content_config_dict: Dict, model: str - ) -> None: + def _normalize_response_schema(generate_content_config_dict: Dict, model: str) -> None: schema_key = next( - ( - k - for k in ("responseSchema", "response_schema") - if k in generate_content_config_dict - ), + (k for k in ("responseSchema", "response_schema") if k in generate_content_config_dict), None, ) json_schema_key = next( - ( - k - for k in ("responseJsonSchema", "response_json_schema") - if k in generate_content_config_dict - ), + (k for k in ("responseJsonSchema", "response_json_schema") if k in generate_content_config_dict), None, ) @@ -340,11 +322,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): generate_content_config_dict.pop(schema_key) return generate_content_config_dict.pop(schema_key) - new_json_schema_key = ( - "response_json_schema" - if schema_key == "response_schema" - else "responseJsonSchema" - ) + new_json_schema_key = "response_json_schema" if schema_key == "response_schema" else "responseJsonSchema" generate_content_config_dict[new_json_schema_key] = value else: if json_schema_key is not None: @@ -420,13 +398,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): """ if "candidates" in response: for candidate in response["candidates"]: - if "citationMetadata" in candidate and isinstance( - candidate["citationMetadata"], dict - ): + if "citationMetadata" in candidate and isinstance(candidate["citationMetadata"], dict): citation_metadata = candidate["citationMetadata"] # Transform citationSources to citations to match expected schema if "citationSources" in citation_metadata: - citation_metadata["citations"] = citation_metadata.pop( - "citationSources" - ) + citation_metadata["citations"] = citation_metadata.pop("citationSources") return response diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 2316361d6e7..78d682395bb 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -78,9 +78,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): api_base: Optional[str], litellm_params: dict, ) -> str: - base_url = ( - api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL - ) + base_url = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL base_url = base_url.rstrip("/") return f"{base_url}/models/{model}:generateContent" @@ -152,14 +150,10 @@ class GeminiImageEditConfig(BaseImageEditConfig): model_response.data = cast(List[OpenAIImage], data_list) if "usageMetadata" in response_json: - model_response.usage = transform_gemini_image_usage( - response_json["usageMetadata"] - ) + model_response.usage = transform_gemini_image_usage(response_json["usageMetadata"]) return model_response - def _prepare_inline_image_parts( - self, image: Union[FileTypes, List[FileTypes]] - ) -> List[Dict[str, Any]]: + def _prepare_inline_image_parts(self, image: Union[FileTypes, List[FileTypes]]) -> List[Dict[str, Any]]: images: List[FileTypes] if isinstance(image, list): images = image diff --git a/litellm/llms/gemini/image_generation/cost_calculator.py b/litellm/llms/gemini/image_generation/cost_calculator.py index 380e2c21e9e..40e234a0b71 100644 --- a/litellm/llms/gemini/image_generation/cost_calculator.py +++ b/litellm/llms/gemini/image_generation/cost_calculator.py @@ -25,9 +25,7 @@ def cost_calculator( ) if not isinstance(image_response, ImageResponse): - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") web_search_cost = calculate_image_response_web_search_cost( image_response=image_response, diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index ebfb0d68830..dcdec46edca 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -34,9 +34,7 @@ else: class GoogleImageGenConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Google AI Imagen API supported parameters https://ai.google.dev/gemini-api/docs/imagen @@ -60,9 +58,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): optional_params=optional_params, ) if is_gemini_image_model(model): - mapped_params = map_gemini_image_tools_params( - non_default_params, mapped_params - ) + mapped_params = map_gemini_image_tools_params(non_default_params, mapped_params) return mapped_params def get_complete_url( @@ -80,9 +76,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): Gemini 2.5 Flash Image Preview: :generateContent Other Imagen models: :predict """ - complete_url: str = ( - api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url: str = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") @@ -159,11 +153,9 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): GeminiImageGenerationParameters, ) - request_body_obj: GeminiImageGenerationRequest = ( - GeminiImageGenerationRequest( - instances=[GeminiImageGenerationInstance(prompt=prompt)], - parameters=GeminiImageGenerationParameters(**optional_params), - ) + request_body_obj: GeminiImageGenerationRequest = GeminiImageGenerationRequest( + instances=[GeminiImageGenerationInstance(prompt=prompt)], + parameters=GeminiImageGenerationParameters(**optional_params), ) return request_body_obj.model_dump(exclude_none=True) @@ -216,23 +208,17 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): b64_json=inline_data["data"], url=None, provider_specific_fields=( - {"thought_signature": thought_sig} - if thought_sig - else None + {"thought_signature": thought_sig} if thought_sig else None ), ) ) # Extract usage metadata for Gemini models if "usageMetadata" in response_data: - model_response.usage = transform_gemini_image_usage( - response_data["usageMetadata"] - ) + model_response.usage = transform_gemini_image_usage(response_data["usageMetadata"]) web_search_requests = get_gemini_image_web_search_requests(response_data) if web_search_requests and model_response.usage is not None: - setattr( - model_response.usage, "web_search_requests", web_search_requests - ) + setattr(model_response.usage, "web_search_requests", web_search_requests) else: # Original Imagen format - predictions with generated images predictions = response_data.get("predictions", []) diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py index 5a55bdeffb1..a4626907f22 100644 --- a/litellm/llms/gemini/image_usage_transformation.py +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -16,9 +16,7 @@ def _get_modality_token_details(usage_metadata: dict, *details_keys: str) -> lis return [] -def _sum_modality_token_details( - usage_metadata: dict, *details_keys: str -) -> ImageUsageInputTokensDetails: +def _sum_modality_token_details(usage_metadata: dict, *details_keys: str) -> ImageUsageInputTokensDetails: tokens_details = ImageUsageInputTokensDetails( image_tokens=0, text_tokens=0, @@ -40,22 +38,16 @@ def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage: """ Transform Gemini usageMetadata to ImageUsage format. """ - input_tokens_details = _sum_modality_token_details( - usage_metadata, "promptTokensDetails", "prompt_tokens_details" - ) + input_tokens_details = _sum_modality_token_details(usage_metadata, "promptTokensDetails", "prompt_tokens_details") output_tokens = usage_metadata.get("candidatesTokenCount", 0) output_tokens_details = _sum_modality_token_details( usage_metadata, "candidatesTokensDetails", "candidates_tokens_details" ) - if not _get_modality_token_details( - usage_metadata, "candidatesTokensDetails", "candidates_tokens_details" - ): + if not _get_modality_token_details(usage_metadata, "candidatesTokensDetails", "candidates_tokens_details"): output_tokens_details.image_tokens = output_tokens else: - known_output_tokens = ( - output_tokens_details.text_tokens + output_tokens_details.image_tokens - ) + known_output_tokens = output_tokens_details.text_tokens + output_tokens_details.image_tokens if output_tokens > known_output_tokens: output_tokens_details.text_tokens += output_tokens - known_output_tokens diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index b18b6a28ce4..7443720f496 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -114,9 +114,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): api_key = GeminiModelInfo.get_api_key(litellm_params.get("api_key")) if not api_key: - raise ValueError( - "Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable." - ) + raise ValueError("Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable.") if stream: return f"{api_base}/{self.api_version}/interactions?alt=sse" @@ -189,10 +187,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if ( response_mime_type and not isinstance(response_format, list) - and ( - not isinstance(response_format, dict) - or "mime_type" not in response_format - ) + and (not isinstance(response_format, dict) or "mime_type" not in response_format) ): # Wrap the legacy schema into the new polymorphic format. new_rf: Dict[str, Any] = { @@ -207,15 +202,11 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): request_body["response_format"] = response_format # image_config moves out of generation_config into response_format. - generation_config: Optional[Dict[str, Any]] = optional_params.get( - "generation_config" - ) + generation_config: Optional[Dict[str, Any]] = optional_params.get("generation_config") if generation_config is not None: image_config = None if isinstance(generation_config, dict): - generation_config = dict( - generation_config - ) # avoid mutating the caller's dict + generation_config = dict(generation_config) # avoid mutating the caller's dict image_config = generation_config.pop("image_config", None) if not generation_config: generation_config = None @@ -261,9 +252,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): response = InteractionsAPIResponse(**raw_json) response._hidden_params["headers"] = dict(raw_response.headers) - response._hidden_params["additional_headers"] = process_response_headers( - dict(raw_response.headers) - ) + response._hidden_params["additional_headers"] = process_response_headers(dict(raw_response.headers)) return response @@ -290,9 +279,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") - encoded_interaction_id = encode_url_path_segment( - interaction_id, field_name="interaction_id" - ) + encoded_interaction_id = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", {}, @@ -326,9 +313,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") - encoded_interaction_id = encode_url_path_segment( - interaction_id, field_name="interaction_id" - ) + encoded_interaction_id = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", {}, @@ -359,9 +344,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") - encoded_interaction_id = encode_url_path_segment( - interaction_id, field_name="interaction_id" - ) + encoded_interaction_id = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}:cancel", {}, diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 9c0a4a30efb..0ff4190a6d3 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -60,9 +60,7 @@ from litellm.utils import get_empty_usage from ..common_utils import encode_unserializable_types, get_api_key_from_env -MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Dict[ - str, Union[OpenAIRealtimeEventTypes, ResponsesAPIStreamEvents] -] = { +MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Dict[str, Union[OpenAIRealtimeEventTypes, ResponsesAPIStreamEvents]] = { "setupComplete": OpenAIRealtimeEventTypes.SESSION_CREATED, "serverContent.generationComplete": OpenAIRealtimeEventTypes.RESPONSE_TEXT_DONE, "serverContent.turnComplete": OpenAIRealtimeEventTypes.RESPONSE_DONE, @@ -71,9 +69,7 @@ MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Dict[ } # Keys the main transform loop handles; siblings like ``usageMetadata`` are skipped. -_KNOWN_GEMINI_TOP_LEVEL_KEYS: set = { - map_key.split(".", 1)[0] for map_key in MAP_GEMINI_FIELD_TO_OPENAI_EVENT -} +_KNOWN_GEMINI_TOP_LEVEL_KEYS: set = {map_key.split(".", 1)[0] for map_key in MAP_GEMINI_FIELD_TO_OPENAI_EVENT} class GeminiRealtimeConfig(BaseRealtimeConfig): @@ -90,9 +86,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return "setup" in msg_obj def is_content_message(self, msg_obj: dict) -> bool: - return any( - k in msg_obj for k in ("realtimeInput", "clientContent", "toolResponse") - ) + return any(k in msg_obj for k in ("realtimeInput", "clientContent", "toolResponse")) def _include_function_response_id(self) -> bool: """Google AI Studio Gemini 3.5+ accepts ``id`` on functionResponses; Vertex AI rejects it.""" @@ -125,14 +119,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) return usage_dict - def validate_environment( - self, headers: dict, model: str, api_key: Optional[str] = None - ) -> dict: + def validate_environment(self, headers: dict, model: str, api_key: Optional[str] = None) -> dict: return headers - def get_complete_url( - self, api_base: Optional[str], model: str, api_key: Optional[str] = None - ) -> str: + def get_complete_url(self, api_base: Optional[str], model: str, api_key: Optional[str] = None) -> str: """ Example output: "BACKEND_WS_URL = "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent""; @@ -151,9 +141,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # already covers the main leak vector. return f"{api_base}/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={api_key}" - def map_model_turn_event( - self, model_turn: HttpxContentType - ) -> OpenAIRealtimeEventTypes: + def map_model_turn_event(self, model_turn: HttpxContentType) -> OpenAIRealtimeEventTypes: """ Map the model turn event to the OpenAI realtime events. @@ -166,9 +154,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if "parts" in model_turn: parts = model_turn["parts"] if len(parts) != 1: - verbose_logger.warning( - f"Realtime: Expected 1 part, got {len(parts)} for Gemini model turn event." - ) + verbose_logger.warning(f"Realtime: Expected 1 part, got {len(parts)} for Gemini model turn event.") part = parts[0] if "text" in part: return OpenAIRealtimeEventTypes.RESPONSE_TEXT_DELTA @@ -178,9 +164,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): raise ValueError(f"Unexpected part type: {part}") raise ValueError(f"Unexpected model turn event, no 'parts' key: {model_turn}") - def map_generation_complete_event( - self, delta_type: Optional[ALL_DELTA_TYPES] - ) -> OpenAIRealtimeEventTypes: + def map_generation_complete_event(self, delta_type: Optional[ALL_DELTA_TYPES]) -> OpenAIRealtimeEventTypes: if delta_type == "text": return OpenAIRealtimeEventTypes.RESPONSE_TEXT_DONE elif delta_type == "audio": @@ -197,55 +181,36 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return mime_types.get(input_audio_format, "application/octet-stream") - def _manual_turn_detection_enabled( - self, session_configuration_request: Optional[str] - ) -> bool: + def _manual_turn_detection_enabled(self, session_configuration_request: Optional[str]) -> bool: if not session_configuration_request: return False try: setup = json.loads(session_configuration_request).get("setup", {}) - automatic_detection = setup.get("realtimeInputConfig", {}).get( - "automaticActivityDetection", {} - ) - return ( - isinstance(automatic_detection, dict) - and automatic_detection.get("disabled") is True - ) + automatic_detection = setup.get("realtimeInputConfig", {}).get("automaticActivityDetection", {}) + return isinstance(automatic_detection, dict) and automatic_detection.get("disabled") is True except (json.JSONDecodeError, TypeError, AttributeError): return False - def _handle_input_audio_buffer_commit_or_end( - self, session_configuration_request: Optional[str] - ) -> List[str]: + def _handle_input_audio_buffer_commit_or_end(self, session_configuration_request: Optional[str]) -> List[str]: """Map OpenAI buffer commit/end to Gemini Live turn-boundary signals.""" if self._manual_turn_detection_enabled(session_configuration_request): realtime_input_dict: BidiGenerateContentRealtimeInput = { "activityEnd": True, } - verbose_logger.debug( - "Gemini Realtime: Sending activityEnd realtimeInput to backend" - ) + verbose_logger.debug("Gemini Realtime: Sending activityEnd realtimeInput to backend") else: realtime_input_dict = {"audioStreamEnd": True} - verbose_logger.debug( - "Gemini Realtime: Sending audioStreamEnd realtimeInput to backend" - ) + verbose_logger.debug("Gemini Realtime: Sending audioStreamEnd realtimeInput to backend") return [json.dumps({"realtimeInput": realtime_input_dict})] - def map_automatic_turn_detection( - self, value: OpenAIRealtimeTurnDetection - ) -> AutomaticActivityDetection: + def map_automatic_turn_detection(self, value: OpenAIRealtimeTurnDetection) -> AutomaticActivityDetection: """Map OpenAI ``server_vad`` to Gemini ``automaticActivityDetection``. OpenAI ``semantic_vad`` has no Gemini Live equivalent — return an empty dict so callers omit ``realtimeInputConfig`` (mapping it with ``disabled: true`` breaks native-audio sessions). """ - if ( - isinstance(value, dict) - and value.get("type") == "semantic_vad" - and "create_response" not in value - ): + if isinstance(value, dict) and value.get("type") == "semantic_vad" and "create_response" not in value: return AutomaticActivityDetection() automatic_activity_dection = AutomaticActivityDetection() @@ -258,12 +223,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): automatic_activity_dection["disabled"] = True if "prefix_padding_ms" in value and isinstance(value["prefix_padding_ms"], int): automatic_activity_dection["prefixPaddingMs"] = value["prefix_padding_ms"] - if "silence_duration_ms" in value and isinstance( - value["silence_duration_ms"], int - ): - automatic_activity_dection["silenceDurationMs"] = value[ - "silence_duration_ms" - ] + if "silence_duration_ms" in value and isinstance(value["silence_duration_ms"], int): + automatic_activity_dection["silenceDurationMs"] = value["silence_duration_ms"] return automatic_activity_dection def get_supported_openai_params(self, model: str) -> List[str]: @@ -278,16 +239,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "voice", ] - def map_openai_params( - self, optional_params: dict, non_default_params: dict - ) -> dict: + def map_openai_params(self, optional_params: dict, non_default_params: dict) -> dict: if "generationConfig" not in optional_params: optional_params["generationConfig"] = {} for key, value in non_default_params.items(): if key == "instructions": - optional_params["systemInstruction"] = HttpxContentType( - role="user", parts=[{"text": value}] - ) + optional_params["systemInstruction"] = HttpxContentType(role="user", parts=[{"text": value}]) elif key == "temperature": optional_params["generationConfig"]["temperature"] = value elif key == "max_response_output_tokens": @@ -319,14 +276,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # Only skip when there is no create_response override so that # a guardrail-injected create_response:false is not dropped. continue - transformed_audio_activity_config = self.map_automatic_turn_detection( - value_typed - ) + transformed_audio_activity_config = self.map_automatic_turn_detection(value_typed) if transformed_audio_activity_config: - optional_params["realtimeInputConfig"] = ( - BidiGenerateContentRealtimeInputConfig( - automaticActivityDetection=transformed_audio_activity_config - ) + optional_params["realtimeInputConfig"] = BidiGenerateContentRealtimeInputConfig( + automaticActivityDetection=transformed_audio_activity_config ) elif key == "voice": from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -387,21 +340,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if isinstance(audio, dict): input_cfg = audio.get("input") if isinstance(input_cfg, dict): - if ( - "input_audio_transcription" not in normalized - and "transcription" in input_cfg - ): + if "input_audio_transcription" not in normalized and "transcription" in input_cfg: normalized["input_audio_transcription"] = input_cfg["transcription"] output_cfg = audio.get("output") if isinstance(output_cfg, dict) and output_cfg.get("voice"): normalized["voice"] = output_cfg["voice"] - extracted_turn_detection = GeminiRealtimeConfig._extract_turn_detection( - normalized - ) - if extracted_turn_detection is not None and not isinstance( - normalized.get("turn_detection"), dict - ): + extracted_turn_detection = GeminiRealtimeConfig._extract_turn_detection(normalized) + if extracted_turn_detection is not None and not isinstance(normalized.get("turn_detection"), dict): normalized["turn_detection"] = extracted_turn_detection return normalized @@ -411,30 +357,23 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): entry = litellm.model_cost.get(model) if entry is None: stripped = model.split("/", 1)[-1] - entry = litellm.model_cost.get(stripped) or litellm.model_cost.get( - f"gemini/{stripped}" - ) + entry = litellm.model_cost.get(stripped) or litellm.model_cost.get(f"gemini/{stripped}") return entry or {} @staticmethod def _is_audio_only_live_model(model: str) -> bool: entry = GeminiRealtimeConfig._model_cost_entry(model) - return bool( - entry.get("gemini_native_audio") or entry.get("gemini_audio_only_live") - ) + return bool(entry.get("gemini_native_audio") or entry.get("gemini_audio_only_live")) @staticmethod def _is_native_audio_model(model: str) -> bool: - return bool( - GeminiRealtimeConfig._model_cost_entry(model).get("gemini_native_audio") - ) + return bool(GeminiRealtimeConfig._model_cost_entry(model).get("gemini_native_audio")) @staticmethod def _coerce_response_modalities(model: str, modalities: list[Any]) -> list[str]: """Map unsupported TEXT responseModalities to AUDIO for audio-only Live models.""" normalized = [ - modality.upper() if isinstance(modality, str) else str(modality).upper() - for modality in modalities + modality.upper() if isinstance(modality, str) else str(modality).upper() for modality in modalities ] if not GeminiRealtimeConfig._is_audio_only_live_model(model): return normalized @@ -444,16 +383,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return without_text if without_text else ["AUDIO"] @staticmethod - def _finalize_gemini_live_setup( - model: str, setup: Dict[str, Any] - ) -> Dict[str, Any]: + def _finalize_gemini_live_setup(model: str, setup: Dict[str, Any]) -> Dict[str, Any]: """Drop fields Gemini Live native-audio rejects on ``setup``.""" generation_config = setup.get("generationConfig") if isinstance(generation_config, dict): modalities = generation_config.get("responseModalities") if isinstance(modalities, list): - generation_config["responseModalities"] = ( - GeminiRealtimeConfig._coerce_response_modalities(model, modalities) + generation_config["responseModalities"] = GeminiRealtimeConfig._coerce_response_modalities( + model, modalities ) if GeminiRealtimeConfig._is_native_audio_model(model): generation_config.pop("speechConfig", None) @@ -490,28 +427,18 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # would be silently dropped because ``map_openai_params`` only # recognises the flat OpenAI-beta key names. session_payload = self._normalize_session_payload_for_mapping(session_payload) - new_overrides = self.map_openai_params( - optional_params={}, non_default_params=session_payload - ) + new_overrides = self.map_openai_params(optional_params={}, non_default_params=session_payload) if session_configuration_request is None: generation_config = new_overrides.setdefault("generationConfig", {}) generation_config.setdefault("responseModalities", ["AUDIO"]) new_overrides.setdefault("inputAudioTranscription", {}) new_overrides["model"] = f"models/{model}" - verbose_logger.debug( - "Gemini Realtime: Sending initial setup with tools to backend" - ) - return [ - json.dumps( - {"setup": self._finalize_gemini_live_setup(model, new_overrides)} - ) - ] + verbose_logger.debug("Gemini Realtime: Sending initial setup with tools to backend") + return [json.dumps({"setup": self._finalize_gemini_live_setup(model, new_overrides)})] if not new_overrides: - verbose_logger.debug( - "Gemini Realtime: Ignoring session.update (no mappable fields)" - ) + verbose_logger.debug("Gemini Realtime: Ignoring session.update (no mappable fields)") return [] try: @@ -534,18 +461,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): } original_generation_config = original_setup.get("generationConfig") new_generation_config = new_overrides.get("generationConfig") - if isinstance(original_generation_config, dict) and isinstance( - new_generation_config, dict - ): + if isinstance(original_generation_config, dict) and isinstance(new_generation_config, dict): follow_up_setup["generationConfig"] = { **original_generation_config, **new_generation_config, } original_realtime_input_config = original_setup.get("realtimeInputConfig") new_realtime_input_config = new_overrides.get("realtimeInputConfig") - if isinstance(original_realtime_input_config, dict) and isinstance( - new_realtime_input_config, dict - ): + if isinstance(original_realtime_input_config, dict) and isinstance(new_realtime_input_config, dict): merged_realtime_input_config = { **original_realtime_input_config, **new_realtime_input_config, @@ -555,12 +478,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # ``create_response: False``) does not silently drop unrelated # knobs like ``silenceDurationMs`` / ``prefixPaddingMs`` from # the original setup. - original_automatic_activity_detection = original_realtime_input_config.get( - "automaticActivityDetection" - ) - new_automatic_activity_detection = new_realtime_input_config.get( - "automaticActivityDetection" - ) + original_automatic_activity_detection = original_realtime_input_config.get("automaticActivityDetection") + new_automatic_activity_detection = new_realtime_input_config.get("automaticActivityDetection") if isinstance(original_automatic_activity_detection, dict) and isinstance( new_automatic_activity_detection, dict ): @@ -572,21 +491,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): BidiGenerateContentRealtimeInputConfig, merged_realtime_input_config, ) - finalized_follow_up = self._finalize_gemini_live_setup( - model, cast(dict[str, Any], follow_up_setup) - ) + finalized_follow_up = self._finalize_gemini_live_setup(model, cast(dict[str, Any], follow_up_setup)) # Skip if the follow-up setup is identical to the one already sent. # The final session.update from Pipecat's _create_response (after history # items) matches the pre-history session.update we intentionally sent # before content; sending a duplicate at that point would risk a 1007. if finalized_follow_up == original_setup: - verbose_logger.debug( - "Gemini Realtime: Skipping duplicate follow-up session.update (no changes)" - ) + verbose_logger.debug("Gemini Realtime: Skipping duplicate follow-up session.update (no changes)") return [] - verbose_logger.debug( - "Gemini Realtime: Forwarding session.update as follow-up setup" - ) + verbose_logger.debug("Gemini Realtime: Forwarding session.update as follow-up setup") return [json.dumps({"setup": finalized_follow_up})] def _handle_conversation_item(self, json_message: dict) -> List[str]: @@ -608,20 +521,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): call_id = item.get("call_id", "") output = item.get("output", "{}") - verbose_logger.debug( - f"Gemini Realtime: Transforming function_call_output for call_id={call_id}" - ) + verbose_logger.debug(f"Gemini Realtime: Transforming function_call_output for call_id={call_id}") # Gemini functionResponses[].response must be a dict; wrap non-dicts. try: parsed_output = json.loads(output) if isinstance(output, str) else output except json.JSONDecodeError: parsed_output = output - output_dict = ( - parsed_output - if isinstance(parsed_output, dict) - else {"result": parsed_output} - ) + output_dict = parsed_output if isinstance(parsed_output, dict) else {"result": parsed_output} # Keep the entry (don't delete) so retried tool responses still find the name. function_name = self._tool_call_id_to_name.get(call_id) @@ -639,20 +546,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if function_name: function_response["name"] = function_name - tool_response_message = { - "toolResponse": {"functionResponses": [function_response]} - } + tool_response_message = {"toolResponse": {"functionResponses": [function_response]}} return [json.dumps(tool_response_message)] def _handle_user_text_content(self, item: dict) -> List[str]: """Transform user text content to Gemini clientContent format.""" content_list = item.get("content", []) - text_parts = [ - c.get("text", "") - for c in content_list - if isinstance(c, dict) and c.get("type") == "input_text" - ] + text_parts = [c.get("text", "") for c in content_list if isinstance(c, dict) and c.get("type") == "input_text"] text = " ".join(filter(None, text_parts)) if not text: return [] @@ -686,9 +587,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): msg_type = json_message.get("type") if msg_type == "session.update": - return self._handle_session_update( - json_message, model, session_configuration_request - ) + return self._handle_session_update(json_message, model, session_configuration_request) if msg_type == "response.create": return [] # Gemini responds automatically; nothing to forward @@ -703,22 +602,16 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): realtime_input_dict = cast( BidiGenerateContentRealtimeInput, - encode_unserializable_types( - cast(Dict[str, object], realtime_input_dict) - ), + encode_unserializable_types(cast(Dict[str, object], realtime_input_dict)), ) gemini_msg = json.dumps({"realtimeInput": realtime_input_dict}) - verbose_logger.debug( - "Gemini Realtime: Sending audio realtimeInput to backend" - ) + verbose_logger.debug("Gemini Realtime: Sending audio realtimeInput to backend") messages.append(gemini_msg) return messages if msg_type in ("input_audio_buffer.commit", "input_audio_buffer.end"): - return self._handle_input_audio_buffer_commit_or_end( - session_configuration_request - ) + return self._handle_input_audio_buffer_commit_or_end(session_configuration_request) if msg_type == "input_audio_buffer.clear": return [] # local buffer op, nothing to forward @@ -739,16 +632,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_configuration_request_dict = {} _model = session_configuration_request_dict.get("model") or model - generation_config = ( - session_configuration_request_dict.get("generationConfig", {}) or {} - ) + generation_config = session_configuration_request_dict.get("generationConfig", {}) or {} gemini_modalities = generation_config.get("responseModalities", ["AUDIO"]) - _modalities = [ - modality.lower() for modality in cast(List[str], gemini_modalities) - ] - _system_instruction = session_configuration_request_dict.get( - "systemInstruction" - ) + _modalities = [modality.lower() for modality in cast(List[str], gemini_modalities)] + _system_instruction = session_configuration_request_dict.get("systemInstruction") session = OpenAIRealtimeStreamSession( id=logging_session_id, modalities=_modalities, @@ -776,9 +663,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) -> bool: if previous_messages is None or len(previous_messages) == 0: return True - if "type" in previous_messages[-1] and previous_messages[-1]["type"].endswith( - "delta" - ): + if "type" in previous_messages[-1] and previous_messages[-1]["type"].endswith("delta"): return False return True @@ -793,18 +678,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_configuration_request_dict: BidiGenerateContentSetup = {} if session_configuration_request is not None: try: - session_configuration_request_dict = json.loads( - session_configuration_request - ).get("setup", {}) + session_configuration_request_dict = json.loads(session_configuration_request).get("setup", {}) except json.JSONDecodeError: session_configuration_request_dict = {} - generation_config = session_configuration_request_dict.get( - "generationConfig", {} - ) + generation_config = session_configuration_request_dict.get("generationConfig", {}) gemini_modalities = generation_config.get("responseModalities", ["AUDIO"]) - _modalities = [ - modality.lower() for modality in cast(List[str], gemini_modalities) - ] + _modalities = [modality.lower() for modality in cast(List[str], gemini_modalities)] _temperature = generation_config.get("temperature") _max_output_tokens = generation_config.get("maxOutputTokens") @@ -904,16 +783,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): elif "inlineData" in part: delta += part["inlineData"].get("data", "") except Exception as e: - raise ValueError( - f"Error transforming content delta events: {e}, got message: {message}" - ) + raise ValueError(f"Error transforming content delta events: {e}, got message: {message}") return OpenAIRealtimeResponseDelta( - type=( - "response.output_text.delta" - if delta_type == "text" - else "response.output_audio.delta" - ), + type=("response.output_text.delta" if delta_type == "text" else "response.output_audio.delta"), content_index=0, event_id="event_{}".format(uuid.uuid4()), item_id=output_item_id, @@ -961,9 +834,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): self, current_output_item_id: Optional[str], current_response_id: Optional[str], - delta_done_event: Union[ - OpenAIRealtimeResponseTextDone, OpenAIRealtimeResponseAudioDone - ], + delta_done_event: Union[OpenAIRealtimeResponseTextDone, OpenAIRealtimeResponseAudioDone], delta_type: ALL_DELTA_TYPES, ) -> List[OpenAIRealtimeEvents]: """ @@ -1046,9 +917,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): resolved_response_id = response_id or f"resp_{uuid.uuid4()}" resolved_output_item_id = output_item_id or f"item_{uuid.uuid4()}" - verbose_logger.debug( - f"Gemini Realtime: Transforming {len(function_calls)} tool call(s) to OpenAI format" - ) + verbose_logger.debug(f"Gemini Realtime: Transforming {len(function_calls)} tool call(s) to OpenAI format") events: List[OpenAIRealtimeFunctionCallArgumentsDone] = [] for idx, fc in enumerate(function_calls): @@ -1098,9 +967,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): any_delta_chunk = False for event in transformed_message: if event["type"] == "response.output_text.delta": - current_delta_chunks.append( - cast(OpenAIRealtimeResponseDelta, event) - ) + current_delta_chunks.append(cast(OpenAIRealtimeResponseDelta, event)) any_delta_chunk = True if not any_delta_chunk: current_delta_chunks = None @@ -1110,9 +977,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ): # audio deltas are not accumulated (memory) if current_delta_chunks is None: current_delta_chunks = [] - current_delta_chunks.append( - cast(OpenAIRealtimeResponseDelta, transformed_message) - ) + current_delta_chunks.append(cast(OpenAIRealtimeResponseDelta, transformed_message)) else: current_delta_chunks = None return current_delta_chunks @@ -1132,9 +997,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): any_item_chunk = False for event in transformed_message: if event["type"] == "response.output_item.done": - current_item_chunks.append( - cast(OpenAIRealtimeOutputItemDone, event) - ) + current_item_chunks.append(cast(OpenAIRealtimeOutputItemDone, event)) any_item_chunk = True if not any_item_chunk: current_item_chunks = None @@ -1142,16 +1005,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if transformed_message["type"] == "response.output_item.done": if current_item_chunks is None: current_item_chunks = [] - current_item_chunks.append( - cast(OpenAIRealtimeOutputItemDone, transformed_message) - ) + current_item_chunks.append(cast(OpenAIRealtimeOutputItemDone, transformed_message)) else: current_item_chunks = None return current_item_chunks except Exception as e: - raise ValueError( - f"Error updating current item chunks: {e}, got transformed_message: {transformed_message}" - ) + raise ValueError(f"Error updating current item chunks: {e}, got transformed_message: {transformed_message}") def transform_response_done_event( self, @@ -1173,18 +1032,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): else: session_configuration_request_dict = {} - generation_config = session_configuration_request_dict.get( - "generationConfig", {} - ) + generation_config = session_configuration_request_dict.get("generationConfig", {}) temperature = generation_config.get("temperature") max_output_tokens = generation_config.get("maxOutputTokens") gemini_modalities = generation_config.get("responseModalities", ["AUDIO"]) - _modalities = [ - modality.lower() for modality in cast(List[str], gemini_modalities) - ] - resolved_usage_metadata = self._consume_usage_metadata_for_response_done( - cast(dict, message) - ) + _modalities = [modality.lower() for modality in cast(List[str], gemini_modalities)] + resolved_usage_metadata = self._consume_usage_metadata_for_response_done(cast(dict, message)) if resolved_usage_metadata is not None: _chat_completion_usage = VertexGeminiConfig._calculate_usage( completion_response=cast( @@ -1208,11 +1061,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): id=current_response_id, status="completed", status_details=None, # type: ignore[typeddict-item] - output=( - [output_item["item"] for output_item in output_items] - if output_items - else [] - ), + output=([output_item["item"] for output_item in output_items] if output_items else []), conversation_id=current_conversation_id, modalities=_modalities, usage=_usage_dict, @@ -1221,9 +1070,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if temperature is not None: response_done_event["response"]["temperature"] = temperature if max_output_tokens is not None: - response_done_event["response"]["max_output_tokens"] = cast( - int, max_output_tokens - ) + response_done_event["response"]["max_output_tokens"] = cast(int, max_output_tokens) return response_done_event @@ -1234,17 +1081,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): realtime_response_transform_input: RealtimeResponseTransformInput, delta_type: ALL_DELTA_TYPES, ) -> RealtimeModalityResponseTransformOutput: - current_output_item_id = realtime_response_transform_input[ - "current_output_item_id" - ] + current_output_item_id = realtime_response_transform_input["current_output_item_id"] current_response_id = realtime_response_transform_input["current_response_id"] - current_conversation_id = realtime_response_transform_input[ - "current_conversation_id" - ] + current_conversation_id = realtime_response_transform_input["current_conversation_id"] current_delta_chunks = realtime_response_transform_input["current_delta_chunks"] - session_configuration_request = realtime_response_transform_input[ - "session_configuration_request" - ] + session_configuration_request = realtime_response_transform_input["session_configuration_request"] returned_message: List[OpenAIRealtimeEvents] = [] if ( @@ -1255,9 +1096,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if not current_output_item_id: # send the list of standard 'new' content.delta events current_output_item_id = "item_{}".format(uuid.uuid4()) - current_conversation_id = current_conversation_id or "conv_{}".format( - uuid.uuid4() - ) + current_conversation_id = current_conversation_id or "conv_{}".format(uuid.uuid4()) returned_message = self.return_new_content_delta_events( session_configuration_request=session_configuration_request, response_id=current_response_id, @@ -1288,12 +1127,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # Use IDs from the done event — transform_content_done_event may have # generated UUID fallbacks when the originals were None. - resolved_item_id = ( - transformed_content_done_event.get("item_id") or current_output_item_id - ) - resolved_response_id = ( - transformed_content_done_event.get("response_id") or current_response_id - ) + resolved_item_id = transformed_content_done_event.get("item_id") or current_output_item_id + resolved_response_id = transformed_content_done_event.get("response_id") or current_response_id additional_items = self.return_additional_content_done_events( current_output_item_id=resolved_item_id, @@ -1324,15 +1159,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): else: model_turn_event = None generation_complete_event = None - openai_event: Optional[ - Union[OpenAIRealtimeEventTypes, ResponsesAPIStreamEvents] - ] = None + openai_event: Optional[Union[OpenAIRealtimeEventTypes, ResponsesAPIStreamEvents]] = None if model_turn_event: # check if model turn event openai_event = self.map_model_turn_event(model_turn_event) elif generation_complete_event: - openai_event = self.map_generation_complete_event( - delta_type=current_delta_type - ) + openai_event = self.map_generation_complete_event(delta_type=current_delta_type) else: # Check if this key or any nested key matches our mapping. Use a # distinct loop variable so we don't shadow ``openai_event`` and @@ -1350,8 +1181,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if ( prefix == key and isinstance(value, dict) - and GeminiRealtimeConfig.get_nested_value(value, nested_path) - is not None + and GeminiRealtimeConfig.get_nested_value(value, nested_path) is not None ): openai_event = candidate_event break @@ -1380,30 +1210,18 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug( "Realtime Response Transform: Gemini frame keys=%s", - ( - sorted(json_message.keys()) - if isinstance(json_message, dict) - else type(json_message).__name__ - ), + (sorted(json_message.keys()) if isinstance(json_message, dict) else type(json_message).__name__), ) logging_session_id = logging_obj.litellm_trace_id - current_output_item_id = realtime_response_transform_input[ - "current_output_item_id" - ] + current_output_item_id = realtime_response_transform_input["current_output_item_id"] current_response_id = realtime_response_transform_input["current_response_id"] - current_conversation_id = realtime_response_transform_input[ - "current_conversation_id" - ] + current_conversation_id = realtime_response_transform_input["current_conversation_id"] current_delta_chunks = realtime_response_transform_input["current_delta_chunks"] - session_configuration_request = realtime_response_transform_input[ - "session_configuration_request" - ] + session_configuration_request = realtime_response_transform_input["session_configuration_request"] current_item_chunks = realtime_response_transform_input["current_item_chunks"] - current_delta_type: Optional[ALL_DELTA_TYPES] = ( - realtime_response_transform_input["current_delta_type"] - ) + current_delta_type: Optional[ALL_DELTA_TYPES] = realtime_response_transform_input["current_delta_type"] returned_message: List[OpenAIRealtimeEvents] = [] server_content = json_message.get("serverContent") @@ -1429,9 +1247,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): current_response_id = "resp_{}".format(uuid.uuid4()) if current_output_item_id is None: current_output_item_id = "item_{}".format(uuid.uuid4()) - current_conversation_id = ( - current_conversation_id or "conv_{}".format(uuid.uuid4()) - ) + current_conversation_id = current_conversation_id or "conv_{}".format(uuid.uuid4()) returned_message.extend( self.return_new_content_delta_events( session_configuration_request=session_configuration_request, @@ -1465,16 +1281,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "interrupted", "generationComplete", } - server_content_handled = not any( - k in server_content for k in _model_content_keys - ) + server_content_handled = not any(k in server_content for k in _model_content_keys) else: server_content_handled = False tool_call_handled = False - for key, value in list( - json_message.items() - ): # snapshot: handlers may mutate json_message + for key, value in list(json_message.items()): # snapshot: handlers may mutate json_message # Skip sibling metadata keys (e.g. ``usageMetadata``) that can # accompany a primary payload like ``toolCall`` or ``serverContent``. # ``map_openai_event`` raises ValueError on unknown keys, which @@ -1512,21 +1324,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_setup: BidiGenerateContentSetup = {} if session_configuration_request is not None: try: - session_setup = json.loads(session_configuration_request).get( - "setup", {} - ) + session_setup = json.loads(session_configuration_request).get("setup", {}) except (json.JSONDecodeError, TypeError): session_setup = {} - tool_call_generation_config = ( - session_setup.get("generationConfig", {}) or {} - ) + tool_call_generation_config = session_setup.get("generationConfig", {}) or {} tool_call_modalities = [ modality.lower() for modality in cast( List[str], - tool_call_generation_config.get( - "responseModalities", ["AUDIO"] - ), + tool_call_generation_config.get("responseModalities", ["AUDIO"]), ) ] @@ -1545,12 +1351,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "output": [], "conversation_id": current_conversation_id, "modalities": tool_call_modalities, - "temperature": tool_call_generation_config.get( - "temperature" - ), - "max_output_tokens": tool_call_generation_config.get( - "maxOutputTokens" - ), + "temperature": tool_call_generation_config.get("temperature"), + "max_output_tokens": tool_call_generation_config.get("maxOutputTokens"), }, } ) @@ -1630,25 +1432,23 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) ) - resolved_tool_call_usage_metadata = ( - self._consume_usage_metadata_for_response_done(json_message) - ) + resolved_tool_call_usage_metadata = self._consume_usage_metadata_for_response_done(json_message) if resolved_tool_call_usage_metadata is not None: - _tool_call_chat_completion_usage = ( - VertexGeminiConfig._calculate_usage( - completion_response=cast( - BidiGenerateContentServerMessage, - { - **json_message, - "usageMetadata": resolved_tool_call_usage_metadata, - }, - ), - ) + _tool_call_chat_completion_usage = VertexGeminiConfig._calculate_usage( + completion_response=cast( + BidiGenerateContentServerMessage, + { + **json_message, + "usageMetadata": resolved_tool_call_usage_metadata, + }, + ), ) else: _tool_call_chat_completion_usage = get_empty_usage() - tool_call_responses_api_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( - _tool_call_chat_completion_usage, + tool_call_responses_api_usage = ( + LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + _tool_call_chat_completion_usage, + ) ) _tool_usage_dict = tool_call_responses_api_usage.model_dump() self._add_pipecat_usage_detail_aliases(_tool_usage_dict) @@ -1679,23 +1479,16 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) tool_call_temperature = tool_call_generation_config.get("temperature") if tool_call_temperature is not None: - tool_call_done_event["response"]["temperature"] = ( - tool_call_temperature - ) - tool_call_max_output_tokens = tool_call_generation_config.get( - "maxOutputTokens" - ) + tool_call_done_event["response"]["temperature"] = tool_call_temperature + tool_call_max_output_tokens = tool_call_generation_config.get("maxOutputTokens") if tool_call_max_output_tokens is not None: - tool_call_done_event["response"]["max_output_tokens"] = cast( - int, tool_call_max_output_tokens - ) + tool_call_done_event["response"]["max_output_tokens"] = cast(int, tool_call_max_output_tokens) returned_message.append(tool_call_done_event) current_output_item_id = None current_response_id = None elif openai_event == OpenAIRealtimeEventTypes.RESPONSE_DONE: _has_pending_function_call = current_item_chunks and any( - chunk.get("item", {}).get("type") == "function_call" - for chunk in current_item_chunks + chunk.get("item", {}).get("type") == "function_call" for chunk in current_item_chunks ) if current_response_id is None and _has_pending_function_call: # Trailing bare turnComplete after a toolCall (Vertex emits ~5 @@ -1787,9 +1580,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): for msg in returned_message: event_type = msg.get("type") if isinstance(msg, dict) else "unknown" - verbose_logger.debug( - "Realtime Response Transform: OpenAI event=%s", event_type - ) + verbose_logger.debug("Realtime Response Transform: OpenAI event=%s", event_type) return { "response": returned_message, diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 35d83bd2adc..f98cb0e5b0c 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -45,9 +45,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): self.model_info = GeminiModelInfo() self._cached_api_key: Optional[str] = None - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: """Gemini uses x-goog-api-key header for authentication.""" return {} @@ -63,15 +61,11 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): "write": [("POST", "/fileSearchStores")], } - def get_supported_openai_params( - self, model: str - ) -> List[VECTOR_STORE_OPENAI_PARAMS]: + def get_supported_openai_params(self, model: str) -> List[VECTOR_STORE_OPENAI_PARAMS]: """Supported parameters for Gemini File Search.""" return ["max_num_results", "filters"] - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """Validate and set up headers for Gemini API.""" headers = headers or {} headers.setdefault("Content-Type", "application/json") @@ -100,9 +94,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): api_version = "v1beta" return f"{api_base}/{api_version}" - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] - ) -> GeminiError: + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]) -> GeminiError: """Return Gemini-specific error class.""" return GeminiError( status_code=status_code, @@ -141,9 +133,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): url = f"{api_base}/models/{model}:generateContent" # Build file_search tool configuration (using snake_case as per Gemini docs) - file_search_config: Dict[str, Any] = { - "file_search_store_names": [vector_store_id] - } + file_search_config: Dict[str, Any] = {"file_search_store_names": [vector_store_id]} # Add metadata filter if provided metadata_filter = vector_store_search_optional_params.get("filters") @@ -214,9 +204,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): results.append( VectorStoreSearchResult( score=None, # Gemini doesn't provide explicit scores - content=[ - VectorStoreResultContent(text=text, type="text") - ], + content=[VectorStoreResultContent(text=text, type="text")], file_id=file_id, filename=title if title else None, attributes={ @@ -251,9 +239,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): results.append( VectorStoreSearchResult( score=score, - content=[ - VectorStoreResultContent(text=text, type="text") - ], + content=[VectorStoreResultContent(text=text, type="text")], attributes={ "grounding_chunk_indices": grounding_chunk_indices, }, @@ -296,9 +282,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): return url, request_body - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: """ Transform Gemini's fileSearchStore response to standard format. """ @@ -316,9 +300,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): created_at = None if create_time: try: - dt = datetime.datetime.fromisoformat( - create_time.replace("Z", "+00:00") - ) + dt = datetime.datetime.fromisoformat(create_time.replace("Z", "+00:00")) created_at = int(dt.timestamp()) except Exception: created_at = None diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 644e96a7dd1..4a9b3830ec5 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -115,11 +115,7 @@ class GeminiVideoConfig(BaseVideoConfig): # Get supported OpenAI params (exclude "model" and "prompt" which are handled separately) supported_openai_params = self.get_supported_openai_params(model) - openai_params_to_map = { - param - for param in supported_openai_params - if param not in {"model", "prompt"} - } + openai_params_to_map = {param for param in supported_openai_params if param not in {"model", "prompt"}} # Map input_reference to image if "input_reference" in video_create_optional_params: @@ -203,12 +199,7 @@ class GeminiVideoConfig(BaseVideoConfig): if litellm_params and litellm_params.api_key: api_key = api_key or litellm_params.api_key - api_key = ( - api_key - or litellm.api_key - or get_secret_str("GOOGLE_API_KEY") - or get_secret_str("GEMINI_API_KEY") - ) + api_key = api_key or litellm.api_key or get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") if not api_key: raise ValueError( @@ -236,10 +227,7 @@ class GeminiVideoConfig(BaseVideoConfig): For status/delete: returns base URL only """ if api_base is None: - api_base = ( - get_secret_str("GEMINI_API_BASE") - or "https://generativelanguage.googleapis.com" - ) + api_base = get_secret_str("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" if not model or model == "": return api_base.rstrip("/") @@ -294,9 +282,7 @@ class GeminiVideoConfig(BaseVideoConfig): parameters = GeminiVideoGenerationParameters(**params_copy) - request_body_obj = GeminiVideoGenerationRequest( - instances=[instance], parameters=parameters - ) + request_body_obj = GeminiVideoGenerationRequest(instances=[instance], parameters=parameters) request_data = request_body_obj.model_dump(exclude_none=True) @@ -339,9 +325,7 @@ class GeminiVideoConfig(BaseVideoConfig): raise ValueError(f"No operation name in Veo response: {response_data}") if custom_llm_provider: - video_id = encode_video_id_with_provider( - operation_name, custom_llm_provider, model - ) + video_id = encode_video_id_with_provider(operation_name, custom_llm_provider, model) else: video_id = operation_name @@ -355,10 +339,7 @@ class GeminiVideoConfig(BaseVideoConfig): usage_data: Dict[str, Any] = {} if request_data: parameters = request_data.get("parameters", {}) - duration = ( - parameters.get("durationSeconds") - or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS - ) + duration = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS if duration is not None: try: usage_data["duration_seconds"] = float(duration) @@ -430,9 +411,7 @@ class GeminiVideoConfig(BaseVideoConfig): is_done = operation_response.done if custom_llm_provider: - video_id = encode_video_id_with_provider( - operation_name, custom_llm_provider, None - ) + video_id = encode_video_id_with_provider(operation_name, custom_llm_provider, None) else: video_id = operation_name @@ -470,16 +449,13 @@ class GeminiVideoConfig(BaseVideoConfig): if not operation_response.done: raise ValueError( - "Video generation is not complete yet. " - "Please check status with video_status() before downloading." + "Video generation is not complete yet. Please check status with video_status() before downloading." ) if not operation_response.response: raise ValueError("No response data in completed operation") - generated_samples = ( - operation_response.response.generateVideoResponse.generatedSamples - ) + generated_samples = operation_response.response.generateVideoResponse.generatedSamples download_url = generated_samples[0].video.uri params: Dict[str, Any] = {} @@ -510,8 +486,7 @@ class GeminiVideoConfig(BaseVideoConfig): Video remix is not supported by Veo API. """ raise NotImplementedError( - "Video remix is not supported by Google Veo. " - "Please use video_generation() to create new videos." + "Video remix is not supported by Google Veo. Please use video_generation() to create new videos." ) def transform_video_remix_response( @@ -561,8 +536,7 @@ class GeminiVideoConfig(BaseVideoConfig): Video delete is not supported by Veo API. """ raise NotImplementedError( - "Video delete is not supported by Google Veo. " - "Videos are automatically cleaned up by Google." + "Video delete is not supported by Google Veo. Videos are automatically cleaned up by Google." ) def transform_video_delete_response( @@ -573,17 +547,13 @@ class GeminiVideoConfig(BaseVideoConfig): """Video delete is not supported.""" raise NotImplementedError("Video delete is not supported by Google Veo.") - def transform_video_create_character_request( - self, name, video, api_base, litellm_params, headers - ): + def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): raise NotImplementedError("video create character is not supported for Gemini") def transform_video_create_character_response(self, raw_response, logging_obj): raise NotImplementedError("video create character is not supported for Gemini") - def transform_video_get_character_request( - self, character_id, api_base, litellm_params, headers - ): + def transform_video_get_character_request(self, character_id, api_base, litellm_params, headers): raise NotImplementedError("video get character is not supported for Gemini") def transform_video_get_character_response(self, raw_response, logging_obj): @@ -622,9 +592,7 @@ class GeminiVideoConfig(BaseVideoConfig): ): raise NotImplementedError("video extension is not supported for Gemini") - def transform_video_extension_response( - self, raw_response, logging_obj, custom_llm_provider=None - ): + def transform_video_extension_response(self, raw_response, logging_obj, custom_llm_provider=None): raise NotImplementedError("video extension is not supported for Gemini") def get_error_class( diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index 59942a9c038..e61015a4a21 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -104,9 +104,7 @@ def get_access_token( token, expires_at = _request_token_sync(credentials, scope, auth_url) # Cache token - ttl_seconds = max( - 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 - ) + ttl_seconds = max(0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) @@ -142,9 +140,7 @@ async def get_access_token_async( token, expires_at = await _request_token_async(credentials, scope, auth_url) # Cache token - ttl_seconds = max( - 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 - ) + ttl_seconds = max(0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index cef80768762..dbf04fd015d 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -111,11 +111,7 @@ class GigaChatConfig(BaseConfig): Set up headers with OAuth token. """ # Get access token - credentials = ( - api_key - or get_secret_str("GIGACHAT_CREDENTIALS") - or get_secret_str("GIGACHAT_API_KEY") - ) + credentials = api_key or get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY") access_token = get_access_token(credentials=credentials) # Store credentials for image uploads @@ -216,9 +212,7 @@ class GigaChatConfig(BaseConfig): ) return functions - def _map_tool_choice( - self, tool_choice: Union[str, dict] - ) -> Optional[Union[str, dict]]: + def _map_tool_choice(self, tool_choice: Union[str, dict]) -> Optional[Union[str, dict]]: """ Map OpenAI tool_choice to GigaChat function_call format. diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index 3785f9c4657..9fefc5df0c5 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -36,9 +36,7 @@ class Authenticator: self.token_dir, os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_FILE", "access-token"), ) - self.api_key_file = os.path.join( - self.token_dir, os.getenv("GITHUB_COPILOT_API_KEY_FILE", "api-key.json") - ) + self.api_key_file = os.path.join(self.token_dir, os.getenv("GITHUB_COPILOT_API_KEY_FILE", "api-key.json")) self._ensure_token_dir() def get_access_token(self) -> str: @@ -57,9 +55,7 @@ class Authenticator: if access_token: return access_token except IOError: - verbose_logger.warning( - "No existing access token found or error reading file" - ) + verbose_logger.warning("No existing access token found or error reading file") for attempt in range(3): verbose_logger.debug(f"Access token acquisition attempt {attempt + 1}/3") @@ -161,9 +157,7 @@ class Authenticator: """ access_token = self.get_access_token() headers = self._get_github_headers(access_token) - api_key_url = os.getenv( - "GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL - ) + api_key_url = os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL) max_retries = 3 for attempt in range(max_retries): @@ -177,13 +171,9 @@ class Authenticator: if "token" in response_json: return response_json else: - verbose_logger.warning( - f"API key response missing token: {response_json}" - ) + verbose_logger.warning(f"API key response missing token: {response_json}") except httpx.HTTPStatusError as e: - verbose_logger.error( - f"HTTP error refreshing API key (attempt {attempt + 1}/{max_retries}): {str(e)}" - ) + verbose_logger.error(f"HTTP error refreshing API key (attempt {attempt + 1}/{max_retries}): {str(e)}") except Exception as e: verbose_logger.error(f"Unexpected error refreshing API key: {str(e)}") @@ -235,9 +225,7 @@ class Authenticator: """ try: sync_client = _get_httpx_client() - device_code_url = os.getenv( - "GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL - ) + device_code_url = os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL) client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) resp = sync_client.post( device_code_url, @@ -291,9 +279,7 @@ class Authenticator: sync_client = _get_httpx_client() max_attempts = 12 # 1 minute (12 * 5 seconds) - access_token_url = os.getenv( - "GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL - ) + access_token_url = os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL) client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) for attempt in range(max_attempts): @@ -313,13 +299,8 @@ class Authenticator: if "access_token" in resp_json: verbose_logger.info("Authentication successful!") return resp_json["access_token"] - elif ( - "error" in resp_json - and resp_json.get("error") == "authorization_pending" - ): - verbose_logger.debug( - f"Authorization pending (attempt {attempt + 1}/{max_attempts})" - ) + elif "error" in resp_json and resp_json.get("error") == "authorization_pending": + verbose_logger.debug(f"Authorization pending (attempt {attempt + 1}/{max_attempts})") else: verbose_logger.warning(f"Unexpected response: {resp_json}") except httpx.HTTPStatusError as e: @@ -335,9 +316,7 @@ class Authenticator: status_code=400, ) except Exception as e: - verbose_logger.error( - f"Unexpected error polling for access token: {str(e)}" - ) + verbose_logger.error(f"Unexpected error polling for access token: {str(e)}") raise GetAccessTokenError( message=f"Failed to get access token: {str(e)}", status_code=400, diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 54cbf69a4ac..2cc05227948 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -189,9 +189,7 @@ class GithubCopilotConfig(OpenAIConfig): _web_search_results, _tool_results, _compaction_blocks, - ) = AnthropicConfig().extract_response_content( - completion_response={"content": content_blocks} - ) + ) = AnthropicConfig().extract_response_content(completion_response={"content": content_blocks}) return text_content, tool_calls, thinking_blocks @staticmethod @@ -202,9 +200,7 @@ class GithubCopilotConfig(OpenAIConfig): if "output_tokens" in usage and "completion_tokens" not in usage: normalized["completion_tokens"] = usage["output_tokens"] if "total_tokens" not in normalized: - normalized["total_tokens"] = normalized.get( - "prompt_tokens", 0 - ) + normalized.get("completion_tokens", 0) + normalized["total_tokens"] = normalized.get("prompt_tokens", 0) + normalized.get("completion_tokens", 0) return normalized @classmethod @@ -227,9 +223,7 @@ class GithubCopilotConfig(OpenAIConfig): thinking_blocks: List[Any] | None = None raw_content = response_json.get("content") if isinstance(raw_content, list): - content, tool_calls, thinking_blocks = cls._parse_anthropic_native_content( - raw_content - ) + content, tool_calls, thinking_blocks = cls._parse_anthropic_native_content(raw_content) elif isinstance(raw_content, str): content = raw_content @@ -260,9 +254,7 @@ class GithubCopilotConfig(OpenAIConfig): synthesized = { **response_json, - "choices": [ - {"index": 0, "message": message, "finish_reason": finish_reason} - ], + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], } usage = response_json.get("usage") if isinstance(usage, dict): diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index da2dc339d6e..d4014ec6242 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -76,9 +76,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): # Merge with existing headers (user's extra_headers take priority) merged_headers = {**default_headers, **headers} - verbose_logger.debug( - f"GitHub Copilot Embedding API: Successfully configured headers for model {model}" - ) + verbose_logger.debug(f"GitHub Copilot Embedding API: Successfully configured headers for model {model}") return merged_headers @@ -185,11 +183,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Any - ) -> Any: + def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any: from litellm.llms.openai.openai import OpenAIConfig - return OpenAIConfig().get_error_class( - error_message=error_message, status_code=status_code, headers=headers - ) + return OpenAIConfig().get_error_class(error_message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index d9f759e0a05..0393d6a9d64 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -53,9 +53,7 @@ def github_copilot_supports_responses_api(model: str) -> bool: register_model, which also clears the cache used here). """ try: - info = _cached_get_model_info_helper( - model=model, custom_llm_provider="github_copilot" - ) + info = _cached_get_model_info_helper(model=model, custom_llm_provider="github_copilot") except Exception as e: verbose_logger.debug( "github_copilot_supports_responses_api: get_model_info failed for %s: %s", @@ -74,9 +72,7 @@ def github_copilot_supports_responses_api(model: str) -> bool: # model_cost entry via the resolved key. key = info.get("key") raw_info = litellm.model_cost.get(key) if isinstance(key, str) else None - endpoints = ( - raw_info.get("supported_endpoints") if isinstance(raw_info, dict) else None - ) + endpoints = raw_info.get("supported_endpoints") if isinstance(raw_info, dict) else None return isinstance(endpoints, list) and "/v1/responses" in endpoints @@ -227,20 +223,14 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): if input_param is not None: initiator = self._get_initiator(input_param) merged_headers["X-Initiator"] = initiator - verbose_logger.debug( - f"GitHub Copilot Responses API: Set X-Initiator={initiator}" - ) + verbose_logger.debug(f"GitHub Copilot Responses API: Set X-Initiator={initiator}") # Add vision header if input contains images if self._has_vision_input(input_param): merged_headers["copilot-vision-request"] = "true" - verbose_logger.debug( - "GitHub Copilot Responses API: Enabled vision request" - ) + verbose_logger.debug("GitHub Copilot Responses API: Enabled vision request") - verbose_logger.debug( - f"GitHub Copilot Responses API: Successfully configured headers for model {model}" - ) + verbose_logger.debug(f"GitHub Copilot Responses API: Successfully configured headers for model {model}") return merged_headers @@ -384,9 +374,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): """ return self._contains_vision_content(input_param) - def _contains_vision_content( - self, value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH - ) -> bool: + def _contains_vision_content(self, value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bool: """ Recursively check if a value contains vision content. @@ -403,12 +391,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Check arrays if isinstance(value, list): - return any( - self._contains_vision_content( - item, depth=depth + 1, max_depth=max_depth - ) - for item in value - ) + return any(self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in value) # Only check dict/object types if not isinstance(value, dict): @@ -422,10 +405,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Check content field recursively if "content" in value and isinstance(value["content"], list): return any( - self._contains_vision_content( - item, depth=depth + 1, max_depth=max_depth - ) - for item in value["content"] + self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) for item in value["content"] ) return False diff --git a/litellm/llms/google_pse/search/transformation.py b/litellm/llms/google_pse/search/transformation.py index 5cd3f2085a8..52d4baba955 100644 --- a/litellm/llms/google_pse/search/transformation.py +++ b/litellm/llms/google_pse/search/transformation.py @@ -43,14 +43,10 @@ class GooglePSESearchRequest(_GooglePSESearchRequestRequired, total=False): hq: str # Optional - append query terms to query imgSize: str # Optional - returns images of specified size imgType: str # Optional - returns images of specified type - linkSite: ( - str # Optional - specifies all search results should contain a link to a URL - ) + linkSite: str # Optional - specifies all search results should contain a link to a URL lr: str # Optional - language restrict (e.g., 'lang_en', 'lang_es') orTerms: str # Optional - provides additional search terms - relatedSite: ( - str # Optional - specifies all search results should be pages related to URL - ) + relatedSite: str # Optional - specifies all search results should be pages related to URL rights: str # Optional - filters based on licensing safe: str # Optional - search safety level ('active', 'off') searchType: str # Optional - specifies search type ('image') @@ -93,14 +89,10 @@ class GooglePSESearchConfig(BaseSearchConfig): default_api_base=self.GOOGLE_PSE_API_BASE, ) if not api_key: - raise ValueError( - "GOOGLE_PSE_API_KEY is not set. Set `GOOGLE_PSE_API_KEY` environment variable." - ) + raise ValueError("GOOGLE_PSE_API_KEY is not set. Set `GOOGLE_PSE_API_KEY` environment variable.") # Also check for search engine ID - search_engine_id = kwargs.get("search_engine_id") or get_secret_str( - "GOOGLE_PSE_ENGINE_ID" - ) + search_engine_id = kwargs.get("search_engine_id") or get_secret_str("GOOGLE_PSE_ENGINE_ID") if not search_engine_id: raise ValueError( "GOOGLE_PSE_ENGINE_ID is not set. Set `GOOGLE_PSE_ENGINE_ID` environment variable or pass `search_engine_id` parameter." @@ -124,11 +116,7 @@ class GooglePSESearchConfig(BaseSearchConfig): """ from urllib.parse import urlencode - api_base = ( - api_base - or get_secret_str("GOOGLE_PSE_API_BASE") - or self.GOOGLE_PSE_API_BASE - ) + api_base = api_base or get_secret_str("GOOGLE_PSE_API_BASE") or self.GOOGLE_PSE_API_BASE # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_google_pse_params" in data: @@ -220,10 +208,7 @@ class GooglePSESearchConfig(BaseSearchConfig): # Pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value # Store params in special key for URL building (Google PSE uses GET not POST) diff --git a/litellm/llms/gradient_ai/chat/transformation.py b/litellm/llms/gradient_ai/chat/transformation.py index 1bc5e8896b1..e81c09d5cf3 100644 --- a/litellm/llms/gradient_ai/chat/transformation.py +++ b/litellm/llms/gradient_ai/chat/transformation.py @@ -20,9 +20,7 @@ class GradientAIConfig(OpenAILikeChatConfig): include_retrieval_info: Optional[bool] = None include_guardrails_info: Optional[bool] = None provide_citations: Optional[bool] = None - retrieval_method: Optional[ - Literal["rewrite", "step_back", "sub_queries", "none"] - ] = None + retrieval_method: Optional[Literal["rewrite", "step_back", "sub_queries", "none"]] = None def __init__( self, @@ -110,10 +108,7 @@ class GradientAIConfig(OpenAILikeChatConfig): if api_base and api_base != GRADIENT_AI_SERVERLESS_ENDPOINT: complete_url = f"{api_base}/api/v1/chat/completions" - elif ( - gradient_ai_endpoint - and gradient_ai_endpoint != GRADIENT_AI_SERVERLESS_ENDPOINT - ): + elif gradient_ai_endpoint and gradient_ai_endpoint != GRADIENT_AI_SERVERLESS_ENDPOINT: complete_url = f"{gradient_ai_endpoint}/api/v1/chat/completions" return complete_url diff --git a/litellm/llms/groq/chat/handler.py b/litellm/llms/groq/chat/handler.py index dc4c3222b12..2553af6df77 100644 --- a/litellm/llms/groq/chat/handler.py +++ b/litellm/llms/groq/chat/handler.py @@ -43,9 +43,7 @@ class GroqChatCompletion(OpenAILikeChatHandler): streaming_decoder: Optional[CustomStreamingDecoder] = None, fake_stream: bool = False, ): - messages = GroqChatConfig()._transform_messages( - messages=cast(List[AllMessageValues], messages), model=model - ) + messages = GroqChatConfig()._transform_messages(messages=cast(List[AllMessageValues], messages), model=model) if optional_params.get("stream") is True: fake_stream = GroqChatConfig()._should_fake_stream(optional_params) diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index d07da006f2d..089c0cac62c 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -104,9 +104,7 @@ class GroqChatConfig(OpenAILikeChatConfig): pass try: - if litellm.supports_reasoning( - model=model, custom_llm_provider=self.custom_llm_provider - ): + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): base_params.append("reasoning_effort") except Exception as e: verbose_logger.debug(f"Error checking if model supports reasoning: {e}") @@ -146,23 +144,15 @@ class GroqChatConfig(OpenAILikeChatConfig): messages[idx] = new_message if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: # groq is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.groq.com/openai/v1 - api_base = ( - api_base - or get_secret_str("GROQ_API_BASE") - or "https://api.groq.com/openai/v1" - ) # type: ignore + api_base = api_base or get_secret_str("GROQ_API_BASE") or "https://api.groq.com/openai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("GROQ_API_KEY") return api_base, dynamic_api_key @@ -226,9 +216,7 @@ class GroqChatConfig(OpenAILikeChatConfig): """ if json_schema is not None: # Check if model supports native response_schema - if not litellm.supports_response_schema( - model=model, custom_llm_provider="groq" - ): + if not litellm.supports_response_schema(model=model, custom_llm_provider="groq"): # Check if user is also passing tools - this combination won't work # See: https://console.groq.com/docs/structured-outputs # "Streaming and tool use are not currently supported with Structured Outputs" @@ -258,9 +246,7 @@ class GroqChatConfig(OpenAILikeChatConfig): "response_format", None ) # only remove if it's a json_schema - handled via using groq's tool calling params. # else: model supports native json_schema, let response_format pass through - optional_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + optional_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) return optional_params @@ -292,17 +278,13 @@ class GroqChatConfig(OpenAILikeChatConfig): json_mode=json_mode, ) - mapped_service_tier: Literal["auto", "default", "flex"] = ( - self._map_groq_service_tier( - original_service_tier=getattr(model_response, "service_tier") - ) + mapped_service_tier: Literal["auto", "default", "flex"] = self._map_groq_service_tier( + original_service_tier=getattr(model_response, "service_tier") ) setattr(model_response, "service_tier", mapped_service_tier) return model_response - def _map_groq_service_tier( - self, original_service_tier: Optional[str] - ) -> Literal["auto", "default", "flex"]: + def _map_groq_service_tier(self, original_service_tier: Optional[str]) -> Literal["auto", "default", "flex"]: """ Ensure groq service tier is OpenAI compatible. """ @@ -318,9 +300,7 @@ class GroqChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): def chunk_parser(self, chunk: dict) -> ModelResponseStream: error = chunk.get("error") if error: - raise OpenAIError( - status_code=error.get("code"), message=error.get("message"), body=error - ) + raise OpenAIError(status_code=error.get("code"), message=error.get("message"), body=error) # Map Groq's 'reasoning' field to LiteLLM's 'reasoning_content' field # Groq returns delta.reasoning, but LiteLLM expects delta.reasoning_content diff --git a/litellm/llms/heroku/chat/transformation.py b/litellm/llms/heroku/chat/transformation.py index fb4cc361189..2efa9fe673f 100644 --- a/litellm/llms/heroku/chat/transformation.py +++ b/litellm/llms/heroku/chat/transformation.py @@ -42,13 +42,9 @@ class HerokuChatConfig(OpenAIGPTConfig): """ messages = handle_messages_with_content_list_to_str_conversion(messages) if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index fc7086548c8..db98749eae0 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -38,9 +38,7 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class HostedVLLMChatConfig(OpenAIGPTConfig): - def _convert_custom_tools_to_function_tools( - self, tools: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: + def _convert_custom_tools_to_function_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ vLLM chat completions currently accepts only OpenAI function tools. Convert custom tools into function tools so request validation does not fail. @@ -59,13 +57,9 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): if not isinstance(custom_tool, dict): custom_tool = {} - tool_name = ( - custom_tool.get("name") or tool.get("name") or f"custom_tool_{idx}" - ) + tool_name = custom_tool.get("name") or tool.get("name") or f"custom_tool_{idx}" tool_description = custom_tool.get("description") or tool.get("description") - tool_parameters = custom_tool.get("input_schema") or tool.get( - "input_schema" - ) + tool_parameters = custom_tool.get("input_schema") or tool.get("input_schema") if not isinstance(tool_parameters, dict): tool_parameters = { @@ -118,23 +112,17 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): if thinking is not None and isinstance(thinking, dict): if thinking.get("type") == "enabled": if "reasoning_effort" not in non_default_params: - non_default_params["reasoning_effort"] = ( - reasoning_effort_from_thinking_budget( - thinking.get("budget_tokens", 0) - ) + non_default_params["reasoning_effort"] = reasoning_effort_from_thinking_budget( + thinking.get("budget_tokens", 0) ) - return super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + return super().map_openai_params(non_default_params, optional_params, model, drop_params) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: api_base = api_base or get_secret_str("HOSTED_VLLM_API_BASE") - dynamic_api_key = ( - api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" - ) + dynamic_api_key = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" return api_base, dynamic_api_key def _is_video_file(self, content_item: ChatCompletionFileObject) -> bool: @@ -156,21 +144,15 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): return True return False - def _convert_file_to_video_url( - self, content_item: ChatCompletionFileObject - ) -> ChatCompletionVideoObject: + def _convert_file_to_video_url(self, content_item: ChatCompletionFileObject) -> ChatCompletionVideoObject: file = content_item.get("file", {}) file_id = file.get("file_id") file_data = file.get("file_data") if file_id: - return ChatCompletionVideoObject( - type="video_url", video_url=ChatCompletionVideoUrlObject(url=file_id) - ) + return ChatCompletionVideoObject(type="video_url", video_url=ChatCompletionVideoUrlObject(url=file_id)) elif file_data: - return ChatCompletionVideoObject( - type="video_url", video_url=ChatCompletionVideoUrlObject(url=file_data) - ) + return ChatCompletionVideoObject(type="video_url", video_url=ChatCompletionVideoUrlObject(url=file_data)) raise ValueError("file_id or file_data is required") @overload @@ -236,8 +218,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): existing_tool_call_ids = { tool_call.get("id") for tool_call in existing_tool_calls - if isinstance(tool_call, dict) - and tool_call.get("id") is not None + if isinstance(tool_call, dict) and tool_call.get("id") is not None } new_tool_calls = [ tool_call @@ -245,37 +226,25 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): if tool_call.get("id") not in existing_tool_call_ids ] if new_tool_calls: - message["tool_calls"] = ( - existing_tool_calls + new_tool_calls - ) + message["tool_calls"] = existing_tool_calls + new_tool_calls else: message["tool_calls"] = tool_calls content_str = "\n".join(text_parts) - new_content = ( - content_blocks if has_structured_content else content_str - ) + new_content = content_blocks if has_structured_content else content_str message["content"] = new_content # type: ignore[typeddict-item] elif message["role"] == "user": message_content = message.get("content") if message_content and isinstance(message_content, list): - replaced_content_items: List[ - Tuple[int, ChatCompletionFileObject] - ] = [] + replaced_content_items: List[Tuple[int, ChatCompletionFileObject]] = [] for idx, content_item in enumerate(message_content): if content_item.get("type") == "file": content_item = cast(ChatCompletionFileObject, content_item) if self._is_video_file(content_item): replaced_content_items.append((idx, content_item)) for idx, content_item in replaced_content_items: - message_content[idx] = self._convert_file_to_video_url( - content_item - ) + message_content[idx] = self._convert_file_to_video_url(content_item) if is_async: - return super()._transform_messages( - messages, model, is_async=cast(Literal[True], True) - ) + return super()._transform_messages(messages, model, is_async=cast(Literal[True], True)) else: - return super()._transform_messages( - messages, model, is_async=cast(Literal[False], False) - ) + return super()._transform_messages(messages, model, is_async=cast(Literal[False], False)) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 47495350460..77504eba04a 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -164,25 +164,19 @@ class HostedVLLMRerankConfig(BaseRerankConfig): try: raw_response_json = raw_response.json() except Exception: - raise ValueError( - f"Error parsing response: {raw_response.text}, status_code={raw_response.status_code}" - ) + raise ValueError(f"Error parsing response: {raw_response.text}, status_code={raw_response.status_code}") return self._transform_response(raw_response_json) def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return HostedVLLMRerankError( - message=error_message, status_code=status_code, headers=headers - ) + return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers) def _transform_response(self, response: dict) -> RerankResponse: # Extract usage information usage_data = response.get("usage", {}) - _billed_units = RerankBilledUnits( - total_tokens=usage_data.get("total_tokens", 0) - ) + _billed_units = RerankBilledUnits(total_tokens=usage_data.get("total_tokens", 0)) _tokens = RerankTokens(input_tokens=usage_data.get("total_tokens", 0)) rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) @@ -201,11 +195,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): # Get document data if it exists document_data = result.get("document", {}) - document = ( - RerankResponseDocument(text=str(document_data.get("text", ""))) - if document_data - else None - ) + document = RerankResponseDocument(text=str(document_data.get("text", ""))) if document_data else None # Create typed result rerank_result = RerankResponseResult( diff --git a/litellm/llms/hosted_vllm/responses/transformation.py b/litellm/llms/hosted_vllm/responses/transformation.py index 4d44eeda9f9..d79690292aa 100644 --- a/litellm/llms/hosted_vllm/responses/transformation.py +++ b/litellm/llms/hosted_vllm/responses/transformation.py @@ -36,9 +36,7 @@ class HostedVLLMResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() api_key = ( - litellm_params.api_key - or get_secret_str("HOSTED_VLLM_API_KEY") - or "fake-api-key" + litellm_params.api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" ) # vllm does not require an api key headers.update( { diff --git a/litellm/llms/huggingface/chat/transformation.py b/litellm/llms/huggingface/chat/transformation.py index 557aa48550b..353d3abac6b 100644 --- a/litellm/llms/huggingface/chat/transformation.py +++ b/litellm/llms/huggingface/chat/transformation.py @@ -66,9 +66,7 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return HuggingFaceError( - status_code=status_code, message=error_message, headers=headers - ) + return HuggingFaceError(status_code=status_code, message=error_message, headers=headers) def get_base_url(self, model: str, base_url: Optional[str]) -> Optional[str]: """ @@ -100,9 +98,7 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): complete_url = api_base complete_url = _build_chat_completion_url(complete_url) elif os.getenv("HF_API_BASE") or os.getenv("HUGGINGFACE_API_BASE"): - complete_url = str(os.getenv("HF_API_BASE")) or str( - os.getenv("HUGGINGFACE_API_BASE") - ) + complete_url = str(os.getenv("HF_API_BASE")) or str(os.getenv("HUGGINGFACE_API_BASE")) elif model.startswith(("http://", "https://")): complete_url = model complete_url = _build_chat_completion_url(complete_url) @@ -135,9 +131,7 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): headers: dict, ) -> dict: if litellm_params.get("api_base"): - return dict( - ChatCompletionRequest(model=model, messages=messages, **optional_params) - ) + return dict(ChatCompletionRequest(model=model, messages=messages, **optional_params)) if "max_retries" in optional_params: logger.warning("`max_retries` is not supported. It will be ignored.") optional_params.pop("max_retries", None) @@ -161,8 +155,4 @@ class HuggingFaceChatConfig(OpenAIGPTConfig): mapped_model = provider_mapping["providerId"] messages = self._transform_messages(messages=messages, model=mapped_model) - return dict( - ChatCompletionRequest( - model=mapped_model, messages=messages, **optional_params - ) - ) + return dict(ChatCompletionRequest(model=mapped_model, messages=messages, **optional_params)) diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 6be885b1f91..39eb430db74 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -21,23 +21,19 @@ config = HuggingFaceEmbeddingConfig() HF_HUB_URL = "https://huggingface.co" -hf_tasks_embeddings = Literal[ # pipeline tags + hf tei endpoints - https://huggingface.github.io/text-embeddings-inference/#/ - "sentence-similarity", "feature-extraction", "rerank", "embed", "similarity" -] +hf_tasks_embeddings = ( + Literal[ # pipeline tags + hf tei endpoints - https://huggingface.github.io/text-embeddings-inference/#/ + "sentence-similarity", "feature-extraction", "rerank", "embed", "similarity" + ] +) -def get_hf_task_embedding_for_model( - model: str, task_type: Optional[str], api_base: str -) -> Optional[str]: +def get_hf_task_embedding_for_model(model: str, task_type: Optional[str], api_base: str) -> Optional[str]: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type else: - raise Exception( - "Invalid task_type={}. Expected one of={}".format( - task_type, hf_tasks_embeddings - ) - ) + raise Exception("Invalid task_type={}. Expected one of={}".format(task_type, hf_tasks_embeddings)) http_client = HTTPHandler(concurrent_limit=1) model_info = http_client.get(url=f"{api_base}/api/models/{model}") @@ -49,18 +45,12 @@ def get_hf_task_embedding_for_model( return pipeline_tag -async def async_get_hf_task_embedding_for_model( - model: str, task_type: Optional[str], api_base: str -) -> Optional[str]: +async def async_get_hf_task_embedding_for_model(model: str, task_type: Optional[str], api_base: str) -> Optional[str]: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type else: - raise Exception( - "Invalid task_type={}. Expected one of={}".format( - task_type, hf_tasks_embeddings - ) - ) + raise Exception("Invalid task_type={}. Expected one of={}".format(task_type, hf_tasks_embeddings)) http_client = get_async_httpx_client( llm_provider=litellm.LlmProviders.HUGGINGFACE, ) @@ -81,9 +71,7 @@ class HuggingFaceEmbedding(BaseLLM): def __init__(self) -> None: super().__init__() - def _transform_input_on_pipeline_tag( - self, input: List, pipeline_tag: Optional[str] - ) -> dict: + def _transform_input_on_pipeline_tag(self, input: List, pipeline_tag: Optional[str]) -> dict: if pipeline_tag is None: return {"inputs": input} if pipeline_tag == "sentence-similarity" or pipeline_tag == "similarity": @@ -110,9 +98,7 @@ class HuggingFaceEmbedding(BaseLLM): input: List, optional_params: dict, ) -> dict: - hf_task = await async_get_hf_task_embedding_for_model( - model=model, task_type=task_type, api_base=HF_HUB_URL - ) + hf_task = await async_get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) data = self._transform_input_on_pipeline_tag(input=input, pipeline_tag=hf_task) @@ -169,22 +155,14 @@ class HuggingFaceEmbedding(BaseLLM): task_type = optional_params.pop("input_type", None) if call_type == "sync": - hf_task = get_hf_task_embedding_for_model( - model=model, task_type=task_type, api_base=HF_HUB_URL - ) + hf_task = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) elif call_type == "async": - return self._async_transform_input( - model=model, task_type=task_type, embed_url=embed_url, input=input - ) # type: ignore + return self._async_transform_input(model=model, task_type=task_type, embed_url=embed_url, input=input) # type: ignore - data = self._transform_input_on_pipeline_tag( - input=input, pipeline_tag=hf_task - ) + data = self._transform_input_on_pipeline_tag(input=input, pipeline_tag=hf_task) if len(optional_params.keys()) > 0: - data = self._process_optional_params( - data=data, optional_params=optional_params - ) + data = self._process_optional_params(data=data, optional_params=optional_params) return data @@ -229,9 +207,7 @@ class HuggingFaceEmbedding(BaseLLM): { "object": "embedding", "index": idx, - "embedding": embedding[0][ - 0 - ], # flatten list returned from hf + "embedding": embedding[0][0], # flatten list returned from hf } ) model_response.object = "list" @@ -343,9 +319,7 @@ class HuggingFaceEmbedding(BaseLLM): litellm_params=litellm_params, ) task_type = optional_params.get("input_type", None) - task = get_hf_task_embedding_for_model( - model=model, task_type=task_type, api_base=HF_HUB_URL - ) + task = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) # print_verbose(f"{model}, {task}") embed_url = "" if "https" in model: @@ -357,9 +331,7 @@ class HuggingFaceEmbedding(BaseLLM): elif "HUGGINGFACE_API_BASE" in os.environ: embed_url = os.getenv("HUGGINGFACE_API_BASE", "") else: - embed_url = ( - f"https://router.huggingface.co/hf-inference/pipeline/{task}/{model}" - ) + embed_url = f"https://router.huggingface.co/hf-inference/pipeline/{task}/{model}" ## ROUTING ## if aembedding is True: diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 7cddda617a9..13e38ab5560 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -48,9 +48,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): details: Optional[bool] = True # enables returning logprobs + best of max_new_tokens: Optional[int] = None repetition_penalty: Optional[float] = None - return_full_text: Optional[bool] = ( - False # by default don't return the input as part of the output - ) + return_full_text: Optional[bool] = False # by default don't return the input as part of the output seed: Optional[int] = None temperature: Optional[float] = None top_k: Optional[int] = None @@ -120,9 +118,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): optional_params["top_p"] = value if param == "n": optional_params["best_of"] = value - optional_params["do_sample"] = ( - True # Need to sample if you want best of for hf inference endpoints - ) + optional_params["do_sample"] = True # Need to sample if you want best of for hf inference endpoints if param == "stream": optional_params["stream"] = value if param == "stop": @@ -212,9 +208,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): task = litellm_params.get("task", None) ## VALIDATE API FORMAT if task is None or not isinstance(task, str) or task not in hf_task_list: - raise Exception( - "Invalid hf task - {}. Valid formats - {}.".format(task, hf_tasks) - ) + raise Exception("Invalid hf task - {}. Valid formats - {}.".format(task, hf_tasks)) ## Load Config config = litellm.HuggingFaceEmbeddingConfig.get_config() @@ -269,12 +263,8 @@ class HuggingFaceEmbeddingConfig(BaseConfig): model_prompt_details = litellm.custom_prompt_dict[model] prompt = custom_prompt( role_dict=model_prompt_details.get("roles") or {}, - initial_prompt_value=model_prompt_details.get( - "initial_prompt_value", "" - ), - final_prompt_value=model_prompt_details.get( - "final_prompt_value", "" - ), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), + final_prompt_value=model_prompt_details.get("final_prompt_value", ""), messages=messages, ) else: @@ -298,12 +288,8 @@ class HuggingFaceEmbeddingConfig(BaseConfig): model_prompt_details = litellm.custom_prompt_dict[model] prompt = custom_prompt( role_dict=model_prompt_details.get("roles", {}), - initial_prompt_value=model_prompt_details.get( - "initial_prompt_value", "" - ), - final_prompt_value=model_prompt_details.get( - "final_prompt_value", "" - ), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), + final_prompt_value=model_prompt_details.get("final_prompt_value", ""), bos_token=model_prompt_details.get("bos_token", ""), eos_token=model_prompt_details.get("eos_token", ""), messages=messages, @@ -373,9 +359,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return HuggingFaceError( - status_code=status_code, message=error_message, headers=headers - ) + return HuggingFaceError(status_code=status_code, message=error_message, headers=headers) def _convert_streamed_response_to_complete_response( self, @@ -439,27 +423,17 @@ class HuggingFaceEmbeddingConfig(BaseConfig): completion_response[0]["generated_text"] ) ## GETTING LOGPROBS + FINISH REASON - if ( - "details" in completion_response[0] - and "tokens" in completion_response[0]["details"] - ): - model_response.choices[0].finish_reason = completion_response[0][ - "details" - ]["finish_reason"] + if "details" in completion_response[0] and "tokens" in completion_response[0]["details"]: + model_response.choices[0].finish_reason = completion_response[0]["details"]["finish_reason"] sum_logprob = 0 for token in completion_response[0]["details"]["tokens"]: if token["logprob"] is not None: sum_logprob += token["logprob"] setattr(model_response.choices[0].message, "_logprob", sum_logprob) # type: ignore if "best_of" in optional_params and optional_params["best_of"] > 1: - if ( - "details" in completion_response[0] - and "best_of_sequences" in completion_response[0]["details"] - ): + if "details" in completion_response[0] and "best_of_sequences" in completion_response[0]["details"]: choices_list = [] - for idx, item in enumerate( - completion_response[0]["details"]["best_of_sequences"] - ): + for idx, item in enumerate(completion_response[0]["details"]["best_of_sequences"]): sum_logprob = 0 for token in item["tokens"]: if token["logprob"] is not None: @@ -483,10 +457,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): completion_response ) else: - if ( - isinstance(completion_response, list) - and len(completion_response[0]["generated_text"]) > 0 - ): + if isinstance(completion_response, list) and len(completion_response[0]["generated_text"]) > 0: model_response.choices[0].message.content = output_parser( # type: ignore completion_response[0]["generated_text"] ) @@ -502,9 +473,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): completion_tokens = 0 try: completion_tokens = len( - encoding.encode( - model_response["choices"][0]["message"].get("content", "") - ) + encoding.encode(model_response["choices"][0]["message"].get("content", "")) ) ##[TODO] use the llama2 tokenizer here except Exception: # this should remain non blocking we should not block a response returning if calculating usage fails @@ -540,10 +509,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): ## Some servers might return streaming responses even though stream was not set to true. (e.g. Baseten) task = litellm_params.get("task", None) is_streamed = False - if ( - raw_response.__dict__["headers"].get("Content-Type", "") - == "text/event-stream" - ): + if raw_response.__dict__["headers"].get("Content-Type", "") == "text/event-stream": is_streamed = True # iterate over the complete streamed response, and return the final answer diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index c94fc65acbd..cdad77a9815 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -152,9 +152,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): if "query" not in optional_rerank_params: raise ValueError("query is required for HuggingFace rerank") if "texts" not in optional_rerank_params: - raise ValueError( - "Cohere 'documents' param is required for HuggingFace rerank" - ) + raise ValueError("Cohere 'documents' param is required for HuggingFace rerank") # Ensure return_text is a boolean value # HuggingFace API expects return_text parameter, corresponding to our return_documents parameter request_body = { @@ -210,25 +208,15 @@ class HuggingFaceRerankConfig(BaseRerankConfig): estimated_input_tokens = token_counter(model=model, text=input_text) except Exception: # Fallback to reasonable estimates if token counting fails - estimated_output_tokens = ( - len(raw_response_json) * 10 if raw_response_json else 10 - ) - estimated_input_tokens = ( - len(input_text) * 4 if "input_text" in locals() else 0 - ) + estimated_output_tokens = len(raw_response_json) * 10 if raw_response_json else 10 + estimated_input_tokens = len(input_text) * 4 if "input_text" in locals() else 0 _billed_units = RerankBilledUnits(search_units=1) - _tokens = RerankTokens( - input_tokens=estimated_input_tokens, output_tokens=estimated_output_tokens - ) - rerank_meta = RerankResponseMeta( - api_version={"version": "1.0"}, billed_units=_billed_units, tokens=_tokens - ) + _tokens = RerankTokens(input_tokens=estimated_input_tokens, output_tokens=estimated_output_tokens) + rerank_meta = RerankResponseMeta(api_version={"version": "1.0"}, billed_units=_billed_units, tokens=_tokens) # Check if documents should be returned based on request parameters - should_return_documents = request_data.get( - "return_text", False - ) or request_data.get("return_documents", False) + should_return_documents = request_data.get("return_text", False) or request_data.get("return_documents", False) original_documents = request_data.get("texts", []) results = [] @@ -252,9 +240,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): if text_content: result["document"] = RerankResponseDocument(text=text_content) # 2. If no text in API response but original documents are available, use those - elif original_documents and 0 <= item.get("index", -1) < len( - original_documents - ): + elif original_documents and 0 <= item.get("index", -1) < len(original_documents): doc = original_documents[item.get("index")] if isinstance(doc, str): result["document"] = RerankResponseDocument(text=doc) @@ -288,16 +274,11 @@ class HuggingFaceRerankConfig(BaseRerankConfig): api_base: API base provided directly to this function, takes precedence over all other sources """ # Get API key from multiple sources - final_api_key = ( - api_key or litellm.huggingface_key or get_secret_str("HUGGINGFACE_API_KEY") - ) + final_api_key = api_key or litellm.huggingface_key or get_secret_str("HUGGINGFACE_API_KEY") # Get API base from multiple sources final_api_base = ( - api_base - or litellm.api_base - or get_secret_str("HF_API_BASE") - or get_secret_str("HUGGINGFACE_API_BASE") + api_base or litellm.api_base or get_secret_str("HF_API_BASE") or get_secret_str("HUGGINGFACE_API_BASE") ) return final_api_key, final_api_base diff --git a/litellm/llms/inception/chat/transformation.py b/litellm/llms/inception/chat/transformation.py index ff87060449b..4c8af768047 100644 --- a/litellm/llms/inception/chat/transformation.py +++ b/litellm/llms/inception/chat/transformation.py @@ -45,14 +45,8 @@ class InceptionChatConfig(OpenAILikeChatConfig): self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: passed_api_base = api_base - api_base = ( - api_base - or get_secret_str("INCEPTION_API_BASE") - or "https://api.inceptionlabs.ai/v1" - ) # type: ignore + api_base = api_base or get_secret_str("INCEPTION_API_BASE") or "https://api.inceptionlabs.ai/v1" # type: ignore dynamic_api_key = api_key if passed_api_base is None or api_key: - dynamic_api_key = ( - api_key or litellm.inception_key or get_secret_str("INCEPTION_API_KEY") - ) + dynamic_api_key = api_key or litellm.inception_key or get_secret_str("INCEPTION_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/infinity/common_utils.py b/litellm/llms/infinity/common_utils.py index 67c54caff98..cf52309ad84 100644 --- a/litellm/llms/infinity/common_utils.py +++ b/litellm/llms/infinity/common_utils.py @@ -5,14 +5,10 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException class InfinityError(BaseLLMException): - def __init__( - self, status_code: int, message: str, headers: Union[dict, httpx.Headers] = {} - ): + def __init__(self, status_code: int, message: str, headers: Union[dict, httpx.Headers] = {}): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://github.com/michaelfeil/infinity" - ) + self.request = httpx.Request(method="POST", url="https://github.com/michaelfeil/infinity") self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, diff --git a/litellm/llms/infinity/embedding/transformation.py b/litellm/llms/infinity/embedding/transformation.py index 824dcd38da3..fd75887baa3 100644 --- a/litellm/llms/infinity/embedding/transformation.py +++ b/litellm/llms/infinity/embedding/transformation.py @@ -117,9 +117,7 @@ class InfinityEmbeddingConfig(BaseEmbeddingConfig): try: raw_response_json = raw_response.json() except Exception: - raise InfinityError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise InfinityError(message=raw_response.text, status_code=raw_response.status_code) # model_response.usage model_response.model = raw_response_json.get("model") @@ -136,6 +134,4 @@ class InfinityEmbeddingConfig(BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return InfinityError( - message=error_message, status_code=status_code, headers=headers - ) + return InfinityError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index b9804605454..94746da4609 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -48,11 +48,7 @@ class InfinityRerankConfig(CohereRerankConfig): optional_params: Optional[dict] = None, ) -> dict: if api_key is None: - api_key = ( - get_secret_str("INFINITY_API_KEY") - or get_secret_str("INFINITY_API_KEY") - or litellm.infinity_key - ) + api_key = get_secret_str("INFINITY_API_KEY") or get_secret_str("INFINITY_API_KEY") or litellm.infinity_key default_headers = { "Authorization": f"Bearer {api_key}", @@ -86,9 +82,7 @@ class InfinityRerankConfig(CohereRerankConfig): try: raw_response_json = raw_response.json() except Exception: - raise InfinityError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise InfinityError(message=raw_response.text, status_code=raw_response.status_code) _billed_units = RerankBilledUnits(**raw_response_json.get("usage", {})) _tokens = RerankTokens( @@ -108,9 +102,7 @@ class InfinityRerankConfig(CohereRerankConfig): relevance_score=result.get("relevance_score"), ) if result.get("document"): - _rerank_response["document"] = RerankResponseDocument( - text=result.get("document") - ) + _rerank_response["document"] = RerankResponseDocument(text=result.get("document")) cohere_results.append(_rerank_response) if cohere_results is None: raise ValueError(f"No results found in the response={raw_response_json}") diff --git a/litellm/llms/jina_ai/embedding/transformation.py b/litellm/llms/jina_ai/embedding/transformation.py index 7a634903005..80927a59a64 100644 --- a/litellm/llms/jina_ai/embedding/transformation.py +++ b/litellm/llms/jina_ai/embedding/transformation.py @@ -80,9 +80,7 @@ class JinaAIEmbeddingConfig(BaseEmbeddingConfig): - api_base: str - dynamic_api_key: str """ - api_base = ( - api_base or get_secret_str("JINA_AI_API_BASE") or "https://api.jina.ai/v1" - ) # type: ignore + api_base = api_base or get_secret_str("JINA_AI_API_BASE") or "https://api.jina.ai/v1" # type: ignore dynamic_api_key = api_key or ( get_secret_str("JINA_AI_API_KEY") or get_secret_str("JINA_AI_API_KEY") @@ -100,11 +98,7 @@ class JinaAIEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - return ( - f"{api_base}/embeddings" - if api_base - else "https://api.jina.ai/v1/embeddings" - ) + return f"{api_base}/embeddings" if api_base else "https://api.jina.ai/v1/embeddings" def transform_embedding_request( self, diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index cabaf079edc..7f4c0709bdd 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -141,9 +141,7 @@ class JinaAIRerankConfig(BaseRerankConfig): optional_params: dict | None = None, ) -> Dict: if api_key is None: - raise ValueError( - "api_key is required. Set via `api_key` parameter or `JINA_API_KEY` environment variable." - ) + raise ValueError("api_key is required. Set via `api_key` parameter or `JINA_API_KEY` environment variable.") return { "accept": "application/json", "content-type": "application/json", diff --git a/litellm/llms/lambda_ai/chat/transformation.py b/litellm/llms/lambda_ai/chat/transformation.py index 262a189428d..96d1dad1416 100644 --- a/litellm/llms/lambda_ai/chat/transformation.py +++ b/litellm/llms/lambda_ai/chat/transformation.py @@ -23,9 +23,7 @@ class LambdaAIChatConfig(OpenAILikeChatConfig): ) -> Tuple[Optional[str], Optional[str]]: # Lambda AI is openai compatible, we just need to set the api_base api_base = ( - api_base - or get_secret_str("LAMBDA_API_BASE") - or "https://api.lambda.ai/v1" # Default Lambda API base URL + api_base or get_secret_str("LAMBDA_API_BASE") or "https://api.lambda.ai/v1" # Default Lambda API base URL ) # type: ignore dynamic_api_key = api_key or get_secret_str("LAMBDA_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/langflow/chat/transformation.py b/litellm/llms/langflow/chat/transformation.py index f898163ad02..73fa49f492b 100644 --- a/litellm/llms/langflow/chat/transformation.py +++ b/litellm/llms/langflow/chat/transformation.py @@ -50,9 +50,7 @@ class LangFlowConfig(BaseConfig): ) -> Tuple[Optional[str], Optional[str]]: from litellm.secret_managers.main import get_secret_str - api_base = ( - api_base or get_secret_str("LANGFLOW_API_BASE") or "http://localhost:7860" - ) + api_base = api_base or get_secret_str("LANGFLOW_API_BASE") or "http://localhost:7860" api_key = api_key or get_secret_str("LANGFLOW_API_KEY") return api_base, api_key @@ -78,10 +76,7 @@ class LangFlowConfig(BaseConfig): if optional_params.get("flow_id") is not None: raise LangFlowError( status_code=400, - message=( - "flow_id cannot be set via request parameters; " - "use model langflow/{flow_id}" - ), + message=("flow_id cannot be set via request parameters; use model langflow/{flow_id}"), ) flow_id = (model.split("/", 1)[1] if "/" in model else model).strip() @@ -264,9 +259,7 @@ class LangFlowConfig(BaseConfig): from litellm.utils import token_counter prompt_tokens = token_counter(model=model, messages=messages) - completion_tokens = token_counter( - model=model, text=content, count_response_tokens=True - ) + completion_tokens = token_counter(model=model, text=content, count_response_tokens=True) usage = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index 9808b665b54..77b5cfbc3fa 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -65,9 +65,7 @@ class LangGraphConfig(BaseConfig): """ from litellm.secret_managers.main import get_secret_str - api_base = ( - api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" - ) + api_base = api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" api_key = api_key or get_secret_str("LANGGRAPH_API_KEY") @@ -137,9 +135,7 @@ class LangGraphConfig(BaseConfig): return parts[1] return model - def _convert_messages_to_langgraph_format( - self, messages: List[AllMessageValues] - ) -> List[Dict[str, Any]]: + def _convert_messages_to_langgraph_format(self, messages: List[AllMessageValues]) -> List[Dict[str, Any]]: """ Convert OpenAI-format messages to LangGraph format. @@ -265,9 +261,7 @@ class LangGraphConfig(BaseConfig): return msg.get("content", "") # Fallback: try to serialize the whole response - verbose_logger.warning( - "Could not extract content from LangGraph response, returning raw" - ) + verbose_logger.warning("Could not extract content from LangGraph response, returning raw") return json.dumps(response_json) def get_streaming_response( @@ -317,14 +311,10 @@ class LangGraphConfig(BaseConfig): ) if response.status_code != 200: - raise LangGraphError( - status_code=response.status_code, message=str(response.read()) - ) + raise LangGraphError(status_code=response.status_code, message=str(response.read())) # Create iterator for SSE stream - completion_stream = self.get_streaming_response( - model=model, raw_response=response - ) + completion_stream = self.get_streaming_response(model=model, raw_response=response) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -366,9 +356,7 @@ class LangGraphConfig(BaseConfig): from litellm.utils import CustomStreamWrapper if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client( - llm_provider=cast(Any, "langgraph"), params={} - ) + client = get_async_httpx_client(llm_provider=cast(Any, "langgraph"), params={}) verbose_logger.debug(f"Making async streaming request to: {api_base}") @@ -382,14 +370,10 @@ class LangGraphConfig(BaseConfig): ) if response.status_code != 200: - raise LangGraphError( - status_code=response.status_code, message=str(await response.aread()) - ) + raise LangGraphError(status_code=response.status_code, message=str(await response.aread())) # Create iterator for SSE stream - completion_stream = self.get_streaming_response( - model=model, raw_response=response - ) + completion_stream = self.get_streaming_response(model=model, raw_response=response) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -459,9 +443,7 @@ class LangGraphConfig(BaseConfig): from litellm.utils import token_counter prompt_tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - completion_tokens = token_counter( - model="gpt-3.5-turbo", text=content, count_response_tokens=True - ) + completion_tokens = token_counter(model="gpt-3.5-turbo", text=content, count_response_tokens=True) total_tokens = prompt_tokens + completion_tokens usage = Usage( diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index fa546f9e147..f10dbf49f66 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -78,9 +78,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): Returns: List of model names prefixed with "lemonade/" """ - api_base, api_key = self._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + api_base, api_key = self._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) if api_base is None: raise ValueError( @@ -173,9 +171,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): if model.startswith("lemonade/"): model = model.split("/", 1)[1] - api_base, api_key = self._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + api_base, api_key = self._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) encoded_model = quote(model, safe="") try: @@ -211,19 +207,10 @@ class LemonadeChatConfig(OpenAILikeChatConfig): ) -> Tuple[Optional[str], Optional[str]]: # lemonade is openai compatible, we just need to set this to custom_openai and have the api_base be lemonade's endpoint passed_api_base = api_base - api_base = ( - api_base - or get_secret_str("LEMONADE_API_BASE") - or "http://localhost:8000/api/v1" - ) # type: ignore + api_base = api_base or get_secret_str("LEMONADE_API_BASE") or "http://localhost:8000/api/v1" # type: ignore key = self._DEFAULT_API_KEY if passed_api_base is None or api_key: - key = ( - api_key - or litellm.lemonade_key - or get_secret_str("LEMONADE_API_KEY") - or self._DEFAULT_API_KEY - ) + key = api_key or litellm.lemonade_key or get_secret_str("LEMONADE_API_KEY") or self._DEFAULT_API_KEY return api_base, key def _get_auth_headers(self, api_key: Optional[str]) -> dict: diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index d27ae038f9e..a68231fa867 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -22,9 +22,7 @@ class _LinkupSearchRequestRequired(TypedDict): q: str # Required - The natural language question for which you want to retrieve context depth: Literal["deep", "standard"] # Required - Defines the precision of the search - outputType: Literal[ - "searchResults", "sourcedAnswer", "structured" - ] # Required - The type of output + outputType: Literal["searchResults", "sourcedAnswer", "structured"] # Required - The type of output class LinkupSearchRequest(_LinkupSearchRequestRequired, total=False): @@ -69,9 +67,7 @@ class LinkupSearchConfig(BaseSearchConfig): default_api_base=self.LINKUP_API_BASE, ) if not api_key: - raise ValueError( - "LINKUP_API_KEY is not set. Set `LINKUP_API_KEY` environment variable." - ) + raise ValueError("LINKUP_API_KEY is not set. Set `LINKUP_API_KEY` environment variable.") headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -141,10 +137,7 @@ class LinkupSearchConfig(BaseSearchConfig): # pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value return result_data diff --git a/litellm/llms/litellm_proxy/chat/transformation.py b/litellm/llms/litellm_proxy/chat/transformation.py index cf6a6ed7a54..eee0ec6fa08 100644 --- a/litellm/llms/litellm_proxy/chat/transformation.py +++ b/litellm/llms/litellm_proxy/chat/transformation.py @@ -42,14 +42,10 @@ class LiteLLMProxyChatConfig(OpenAIGPTConfig): dynamic_api_key = api_key or get_secret_str("LITELLM_PROXY_API_KEY") return api_base, dynamic_api_key - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: api_base, api_key = self._get_openai_compatible_provider_info(api_base, api_key) if api_base is None: - raise ValueError( - "api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`" - ) + raise ValueError("api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`") models = super().get_models(api_key=api_key, api_base=api_base) return [f"litellm_proxy/{model}" for model in models] @@ -111,9 +107,7 @@ class LiteLLMProxyChatConfig(OpenAIGPTConfig): ( api_base, api_key, - ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) return model, custom_llm_provider, api_key, api_base diff --git a/litellm/llms/litellm_proxy/image_edit/transformation.py b/litellm/llms/litellm_proxy/image_edit/transformation.py index 79cd6e15c68..94825cffeae 100644 --- a/litellm/llms/litellm_proxy/image_edit/transformation.py +++ b/litellm/llms/litellm_proxy/image_edit/transformation.py @@ -19,13 +19,9 @@ class LiteLLMProxyImageEditConfig(OpenAIImageEditConfig): headers.update({"Authorization": f"Bearer {api_key}"}) return headers - def get_complete_url( - self, model: str, api_base: Optional[str], litellm_params: dict - ) -> str: + def get_complete_url(self, model: str, api_base: Optional[str], litellm_params: dict) -> str: api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") if api_base is None: - raise ValueError( - "api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`" - ) + raise ValueError("api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`") api_base = api_base.rstrip("/") return f"{api_base}/images/edits" diff --git a/litellm/llms/litellm_proxy/image_generation/transformation.py b/litellm/llms/litellm_proxy/image_generation/transformation.py index 3932070e964..5fad663d126 100644 --- a/litellm/llms/litellm_proxy/image_generation/transformation.py +++ b/litellm/llms/litellm_proxy/image_generation/transformation.py @@ -34,8 +34,6 @@ class LiteLLMProxyImageGenerationConfig(GPTImageGenerationConfig): ) -> str: api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") if api_base is None: - raise ValueError( - "api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`" - ) + raise ValueError("api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`") api_base = api_base.rstrip("/") return f"{api_base}/images/generations" diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index 2b567f03760..4ac3311921d 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -44,9 +44,7 @@ def get_litellm_code_execution_tool() -> Dict[str, Any]: "description": "Execute Python code in a sandboxed environment. Use this to run code that generates files, processes data, or performs computations. Generated files will be returned directly.", "parameters": { "type": "object", - "properties": { - "code": {"type": "string", "description": "Python code to execute"} - }, + "properties": {"code": {"type": "string", "description": "Python code to execute"}}, "required": ["code"], }, }, @@ -65,9 +63,7 @@ def get_litellm_code_execution_tool_anthropic() -> Dict[str, Any]: "description": "Execute Python code in a sandboxed environment. Use this to run code that generates files, processes data, or performs computations. Generated files will be returned directly.", "input_schema": { "type": "object", - "properties": { - "code": {"type": "string", "description": "Python code to execute"} - }, + "properties": {"code": {"type": "string", "description": "Python code to execute"}}, "required": ["code"], }, } @@ -145,9 +141,7 @@ class CodeExecutionHandler: response: Any = None # Initialize to avoid possibly unbound error for iteration in range(self.max_iterations): - verbose_logger.debug( - f"CodeExecutionHandler: Iteration {iteration + 1}/{self.max_iterations}" - ) + verbose_logger.debug(f"CodeExecutionHandler: Iteration {iteration + 1}/{self.max_iterations}") # Make LLM call response = await litellm.acompletion( @@ -181,9 +175,7 @@ class CodeExecutionHandler: # Check if we're done (no tool calls or not tool_calls finish reason) if stop_reason != "tool_calls" or not assistant_message.tool_calls: - verbose_logger.debug( - f"CodeExecutionHandler: Completed after {iteration + 1} iterations" - ) + verbose_logger.debug(f"CodeExecutionHandler: Completed after {iteration + 1} iterations") return { "response": response, "files": generated_files, # Files returned directly with base64 content @@ -201,18 +193,14 @@ class CodeExecutionHandler: args = json.loads(tool_call.function.arguments) code = args.get("code", "") - verbose_logger.debug( - f"CodeExecutionHandler: Executing code ({len(code)} chars)" - ) + verbose_logger.debug(f"CodeExecutionHandler: Executing code ({len(code)} chars)") exec_result = executor.execute( code=code, skill_files=skill_files, ) - verbose_logger.debug( - f"CodeExecutionHandler: Execution result: {exec_result}" - ) + verbose_logger.debug(f"CodeExecutionHandler: Execution result: {exec_result}") execution_results.append( { @@ -241,9 +229,7 @@ class CodeExecutionHandler: "size": len(file_content), } ) - tool_result += ( - f"\n- {f['name']} ({len(file_content)} bytes)" - ) + tool_result += f"\n- {f['name']} ({len(file_content)} bytes)" verbose_logger.debug( f"CodeExecutionHandler: Generated file {f['name']} ({len(file_content)} bytes)" @@ -282,9 +268,7 @@ class CodeExecutionHandler: ) # Max iterations reached - verbose_logger.warning( - f"CodeExecutionHandler: Max iterations ({self.max_iterations}) reached" - ) + verbose_logger.warning(f"CodeExecutionHandler: Max iterations ({self.max_iterations}) reached") return { "response": response, "files": generated_files, diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 9138b9a712f..6f5ae261d2e 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -55,10 +55,7 @@ class LiteLLMSkillsHandler: from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise ValueError( - "Prisma client is not initialized. " - "Database connection required for LiteLLM skills." - ) + raise ValueError("Prisma client is not initialized. Database connection required for LiteLLM skills.") return prisma_client @staticmethod @@ -77,9 +74,7 @@ class LiteLLMSkillsHandler: # Stamping a placeholder would let any two such callers see # each other's skills via the shared owner. ValueError keeps # this module FastAPI-free per the project layering rule. - raise ValueError( - "Unable to record skill ownership: caller has no identity scope." - ) + raise ValueError("Unable to record skill ownership: caller has no identity scope.") skill_data: Dict[str, Any] = { "skill_id": skill_id, @@ -105,9 +100,7 @@ class LiteLLMSkillsHandler: if data.file_type is not None: skill_data["file_type"] = data.file_type - verbose_logger.debug( - f"LiteLLMSkillsHandler: Creating skill {skill_id} with title={data.display_title}" - ) + verbose_logger.debug(f"LiteLLMSkillsHandler: Creating skill {skill_id} with title={data.display_title}") new_skill = await SkillsRepository(prisma_client).table.create(data=skill_data) return _prisma_skill_to_litellm(new_skill) @@ -120,9 +113,7 @@ class LiteLLMSkillsHandler: ) -> List[LiteLLM_SkillsTable]: prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - verbose_logger.debug( - f"LiteLLMSkillsHandler: Listing skills with limit={limit}, offset={offset}" - ) + verbose_logger.debug(f"LiteLLMSkillsHandler: Listing skills with limit={limit}, offset={offset}") find_many_kwargs: Dict[str, Any] = { "take": limit, @@ -135,9 +126,7 @@ class LiteLLMSkillsHandler: return [] find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} - skills = await SkillsRepository(prisma_client).table.find_many( - **find_many_kwargs - ) + skills = await SkillsRepository(prisma_client).table.find_many(**find_many_kwargs) return [_prisma_skill_to_litellm(s) for s in skills] @staticmethod @@ -152,12 +141,8 @@ class LiteLLMSkillsHandler: return cached prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - skill = await SkillsRepository(prisma_client).table.find_unique( - where={"skill_id": skill_id} - ) - _SKILL_CACHE.set_cache( - skill_id, skill if skill is not None else _NEGATIVE_SKILL_SENTINEL - ) + skill = await SkillsRepository(prisma_client).table.find_unique(where={"skill_id": skill_id}) + _SKILL_CACHE.set_cache(skill_id, skill if skill is not None else _NEGATIVE_SKILL_SENTINEL) return skill @staticmethod @@ -170,9 +155,7 @@ class LiteLLMSkillsHandler: skill = await LiteLLMSkillsHandler._load_skill(skill_id) # Same "not found" message for both "missing" and "cross-tenant" # so callers can't enumerate skill IDs they don't own. - if skill is None or not user_can_access_resource_owner( - getattr(skill, "created_by", None), user_api_key_dict - ): + if skill is None or not user_can_access_resource_owner(getattr(skill, "created_by", None), user_api_key_dict): raise ValueError(f"Skill not found: {skill_id}") return _prisma_skill_to_litellm(skill) @@ -186,9 +169,7 @@ class LiteLLMSkillsHandler: verbose_logger.debug(f"LiteLLMSkillsHandler: Deleting skill {skill_id}") skill = await LiteLLMSkillsHandler._load_skill(skill_id) - if skill is None or not user_can_access_resource_owner( - getattr(skill, "created_by", None), user_api_key_dict - ): + if skill is None or not user_can_access_resource_owner(getattr(skill, "created_by", None), user_api_key_dict): raise ValueError(f"Skill not found: {skill_id}") await SkillsRepository(prisma_client).table.delete(where={"skill_id": skill_id}) @@ -204,13 +185,9 @@ class LiteLLMSkillsHandler: """Skills-injection-hook helper: returns None instead of raising on not-found / not-authorized so the hook can silently skip.""" try: - return await LiteLLMSkillsHandler.get_skill( - skill_id, user_api_key_dict=user_api_key_dict - ) + return await LiteLLMSkillsHandler.get_skill(skill_id, user_api_key_dict=user_api_key_dict) except ValueError: return None except Exception as e: - verbose_logger.warning( - f"LiteLLMSkillsHandler: Error fetching skill {skill_id}: {e}" - ) + verbose_logger.warning(f"LiteLLMSkillsHandler: Error fetching skill {skill_id}: {e}") return None diff --git a/litellm/llms/litellm_proxy/skills/prompt_injection.py b/litellm/llms/litellm_proxy/skills/prompt_injection.py index 86b6e223512..8be6f105845 100644 --- a/litellm/llms/litellm_proxy/skills/prompt_injection.py +++ b/litellm/llms/litellm_proxy/skills/prompt_injection.py @@ -147,9 +147,7 @@ class SkillPromptInjectionHandler: return data # Build the skill injection text - skill_section = "\n\n---\n\n# Available Skills\n\n" + "\n\n---\n\n".join( - skill_contents - ) + skill_section = "\n\n---\n\n# Available Skills\n\n" + "\n\n---\n\n".join(skill_contents) if use_anthropic_format: # Anthropic messages API: use top-level 'system' parameter @@ -243,12 +241,7 @@ class SkillPromptInjectionHandler: func_name = skill.skill_id.replace("-", "_").replace(" ", "_") # Use instructions as description, fall back to description or title - description = ( - skill.instructions - or skill.description - or skill.display_title - or f"Skill: {skill.skill_id}" - ) + description = skill.instructions or skill.description or skill.display_title or f"Skill: {skill.skill_id}" # Truncate description if too long (OpenAI has limits) max_desc_length = 1024 @@ -276,9 +269,7 @@ class SkillPromptInjectionHandler: return tool - def convert_skill_to_anthropic_tool( - self, skill: LiteLLM_SkillsTable - ) -> Dict[str, Any]: + def convert_skill_to_anthropic_tool(self, skill: LiteLLM_SkillsTable) -> Dict[str, Any]: """ Convert a LiteLLM skill to an Anthropic-style tool (messages API format). @@ -290,12 +281,7 @@ class SkillPromptInjectionHandler: """ func_name = skill.skill_id.replace("-", "_").replace(" ", "_") - description = ( - skill.instructions - or skill.description - or skill.display_title - or f"Skill: {skill.skill_id}" - ) + description = skill.instructions or skill.description or skill.display_title or f"Skill: {skill.skill_id}" max_desc_length = 1024 if len(description) > max_desc_length: diff --git a/litellm/llms/litellm_proxy/skills/sandbox_executor.py b/litellm/llms/litellm_proxy/skills/sandbox_executor.py index 4514512fc59..5f1f129032c 100644 --- a/litellm/llms/litellm_proxy/skills/sandbox_executor.py +++ b/litellm/llms/litellm_proxy/skills/sandbox_executor.py @@ -67,10 +67,7 @@ class SkillsSandboxExecutor: try: from llm_sandbox import SandboxSession except ImportError: - verbose_logger.error( - "SkillsSandboxExecutor: llm-sandbox not installed. " - "Install `llm-sandbox`." - ) + verbose_logger.error("SkillsSandboxExecutor: llm-sandbox not installed. Install `llm-sandbox`.") return { "success": False, "output": "", @@ -99,9 +96,7 @@ class SkillsSandboxExecutor: # Create the file in temp directory local_path = os.path.abspath(os.path.join(tmpdir, path)) if not local_path.startswith(tmpdir_abs + os.sep): - verbose_logger.warning( - f"SkillsSandboxExecutor: Skipping file with invalid path: {path}" - ) + verbose_logger.warning(f"SkillsSandboxExecutor: Skipping file with invalid path: {path}") continue os.makedirs(os.path.dirname(local_path), exist_ok=True) with open(local_path, "wb") as f: @@ -111,9 +106,7 @@ class SkillsSandboxExecutor: sandbox_path = f"/sandbox/{path}" session.copy_to_runtime(local_path, sandbox_path) - verbose_logger.debug( - f"SkillsSandboxExecutor: Copied {len(skill_files)} files to sandbox" - ) + verbose_logger.debug(f"SkillsSandboxExecutor: Copied {len(skill_files)} files to sandbox") # 2. Install requirements if present. Let pip parse the # requirements file inside the sandbox so standard syntax like @@ -149,18 +142,14 @@ subprocess.run( """ install_result = session.run(pip_code) if install_result.exit_code != 0: - verbose_logger.debug( - "SkillsSandboxExecutor: Requirements installation failed" - ) + verbose_logger.debug("SkillsSandboxExecutor: Requirements installation failed") return { "success": False, "output": install_result.stdout or "", "error": install_result.stderr or "", "files": [], } - verbose_logger.debug( - "SkillsSandboxExecutor: Installed requirements" - ) + verbose_logger.debug("SkillsSandboxExecutor: Installed requirements") # 3. Execute the code # Wrap code to run from /sandbox directory @@ -179,19 +168,13 @@ sys.path.insert(0, '/sandbox') error = result.stderr or "" if success: - verbose_logger.debug( - "SkillsSandboxExecutor: Code execution succeeded" - ) + verbose_logger.debug("SkillsSandboxExecutor: Code execution succeeded") else: verbose_logger.debug( f"SkillsSandboxExecutor: Code execution failed with exit code {result.exit_code}" ) - verbose_logger.debug( - f"SkillsSandboxExecutor: stderr: {error[:500] if error else 'No stderr'}" - ) - verbose_logger.debug( - f"SkillsSandboxExecutor: stdout: {output[:500] if output else 'No stdout'}" - ) + verbose_logger.debug(f"SkillsSandboxExecutor: stderr: {error[:500] if error else 'No stderr'}") + verbose_logger.debug(f"SkillsSandboxExecutor: stdout: {output[:500] if output else 'No stdout'}") # 4. Collect generated files generated_files = self._collect_generated_files(session, skill_files) @@ -287,21 +270,15 @@ print(json.dumps(files)) } ) - verbose_logger.debug( - f"SkillsSandboxExecutor: Collected generated file: {rel_path}" - ) + verbose_logger.debug(f"SkillsSandboxExecutor: Collected generated file: {rel_path}") except Exception as e: - verbose_logger.warning( - f"SkillsSandboxExecutor: Error copying file {filepath}: {e}" - ) + verbose_logger.warning(f"SkillsSandboxExecutor: Error copying file {filepath}: {e}") finally: if os.path.exists(tmp_path): os.unlink(tmp_path) except Exception as e: - verbose_logger.warning( - f"SkillsSandboxExecutor: Error collecting generated files: {e}" - ) + verbose_logger.warning(f"SkillsSandboxExecutor: Error collecting generated files: {e}") return generated_files diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index 199f13191fe..7fa58ad9df2 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -87,9 +87,7 @@ class LiteLLMSkillsTransformationHandler: if isinstance(first_file, tuple) and len(first_file) >= 2: file_name = first_file[0] file_content = first_file[1] - file_type = ( - first_file[2] if len(first_file) > 2 else "application/zip" - ) + file_type = first_file[2] if len(first_file) > 2 else "application/zip" if _is_async: return self._async_create_skill( diff --git a/litellm/llms/llamafile/chat/transformation.py b/litellm/llms/llamafile/chat/transformation.py index 223f5503c9c..78cc58708ad 100644 --- a/litellm/llms/llamafile/chat/transformation.py +++ b/litellm/llms/llamafile/chat/transformation.py @@ -15,9 +15,7 @@ class LlamafileChatConfig(OpenAIGPTConfig): If both are None, a fake API key is returned. """ - return ( - api_key or get_secret_str("LLAMAFILE_API_KEY") or "fake-api-key" - ) # llamafile does not require an API key + return api_key or get_secret_str("LLAMAFILE_API_KEY") or "fake-api-key" # llamafile does not require an API key @staticmethod def _resolve_api_base(api_base: Optional[str] = None) -> Optional[str]: @@ -27,11 +25,7 @@ class LlamafileChatConfig(OpenAIGPTConfig): If both are None, a default Llamafile server URL is returned. See: https://github.com/Mozilla-Ocho/llamafile/blob/bd1bbe9aabb1ee12dbdcafa8936db443c571eb9d/README.md#L61 """ - return ( - api_base - or get_secret_str("LLAMAFILE_API_BASE") - or "http://127.0.0.1:8080/v1" - ) # type: ignore + return api_base or get_secret_str("LLAMAFILE_API_BASE") or "http://127.0.0.1:8080/v1" # type: ignore def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] diff --git a/litellm/llms/lm_studio/embed/transformation.py b/litellm/llms/lm_studio/embed/transformation.py index 87f4f6e73d5..f0357b9428c 100644 --- a/litellm/llms/lm_studio/embed/transformation.py +++ b/litellm/llms/lm_studio/embed/transformation.py @@ -44,7 +44,5 @@ class LmStudioEmbeddingConfig: def get_supported_openai_params(self) -> List[str]: return [] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: return optional_params diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index 34166161390..4a65fac709b 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -92,9 +92,7 @@ class ManusFilesConfig(BaseFilesConfig): ) return headers - def get_supported_openai_params( - self, model: str - ) -> List[OpenAICreateFileRequestOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAICreateFileRequestOptionalParams]: """ Return supported OpenAI file creation parameters for Manus. Manus supports the standard 'purpose' parameter. @@ -129,12 +127,7 @@ class ManusFilesConfig(BaseFilesConfig): Returns: str: The full URL for the Manus /v1/files endpoint """ - api_base = ( - api_base - or litellm.api_base - or get_secret_str("MANUS_API_BASE") - or MANUS_API_BASE - ) + api_base = api_base or litellm.api_base or get_secret_str("MANUS_API_BASE") or MANUS_API_BASE # Remove trailing slashes api_base = api_base.rstrip("/") @@ -193,11 +186,7 @@ class ManusFilesConfig(BaseFilesConfig): ) # Get API key - api_key = ( - litellm_params.get("api_key") - or litellm.api_key - or get_secret_str("MANUS_API_KEY") - ) + api_key = litellm_params.get("api_key") or litellm.api_key or get_secret_str("MANUS_API_KEY") if not api_key: raise ValueError( diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index b3a0073a5c2..0db53f90330 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -75,18 +75,14 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): # If no slash, assume the model name itself is the agent profile return model - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate environment and set up headers for Manus API. Manus uses `API_KEY` header instead of `Authorization: Bearer`. """ litellm_params = litellm_params or GenericLiteLLMParams() - api_key = ( - litellm_params.api_key or litellm.api_key or get_secret_str("MANUS_API_KEY") - ) + api_key = litellm_params.api_key or litellm.api_key or get_secret_str("MANUS_API_KEY") if not api_key: raise ValueError( @@ -114,12 +110,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): Returns: str: The full URL for the Manus /v1/responses endpoint """ - api_base = ( - api_base - or litellm.api_base - or get_secret_str("MANUS_API_BASE") - or MANUS_API_BASE - ) + api_base = api_base or litellm.api_base or get_secret_str("MANUS_API_BASE") or MANUS_API_BASE # Remove trailing slashes api_base = api_base.rstrip("/") @@ -166,9 +157,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): if extra_body: base_request.update(extra_body) - verbose_logger.debug( - f"Manus: Using agent_profile={agent_profile}, task_mode=agent" - ) + verbose_logger.debug(f"Manus: Using agent_profile={agent_profile}, task_mode=agent") return base_request @@ -191,32 +180,20 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): raw_response_json = raw_response.json() # Manus uses camelCase "createdAt" instead of snake_case "created_at" - if ( - "createdAt" in raw_response_json - and "created_at" not in raw_response_json - ): - raw_response_json["created_at"] = _safe_convert_created_field( - raw_response_json["createdAt"] - ) + if "createdAt" in raw_response_json and "created_at" not in raw_response_json: + raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["createdAt"]) # Ensure created_at is set if "created_at" in raw_response_json: - raw_response_json["created_at"] = _safe_convert_created_field( - raw_response_json["created_at"] - ) + raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["created_at"]) except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) # Ensure reasoning is an empty dict if not present, OpenAI SDK does not allow None - if ( - "reasoning" not in raw_response_json - or raw_response_json.get("reasoning") is None - ): + if "reasoning" not in raw_response_json or raw_response_json.get("reasoning") is None: raw_response_json["reasoning"] = {} if "text" not in raw_response_json or raw_response_json.get("text") is None: @@ -242,9 +219,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): try: response = ResponsesAPIResponse(**raw_response_json) except Exception: - verbose_logger.debug( - f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" - ) + verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct") response = ResponsesAPIResponse.model_construct(**raw_response_json) # Store processed headers in additional_headers so they get returned to the client @@ -271,9 +246,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): Reference: https://open.manus.im/docs/openai-compatibility """ - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -297,32 +270,20 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): raw_response_json = raw_response.json() # Manus uses camelCase "createdAt" instead of snake_case "created_at" - if ( - "createdAt" in raw_response_json - and "created_at" not in raw_response_json - ): - raw_response_json["created_at"] = _safe_convert_created_field( - raw_response_json["createdAt"] - ) + if "createdAt" in raw_response_json and "created_at" not in raw_response_json: + raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["createdAt"]) # Ensure created_at is set if "created_at" in raw_response_json: - raw_response_json["created_at"] = _safe_convert_created_field( - raw_response_json["created_at"] - ) + raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["created_at"]) except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) # Ensure reasoning, text, output, and usage are present with defaults - if ( - "reasoning" not in raw_response_json - or raw_response_json.get("reasoning") is None - ): + if "reasoning" not in raw_response_json or raw_response_json.get("reasoning") is None: raw_response_json["reasoning"] = {} if "text" not in raw_response_json or raw_response_json.get("text") is None: @@ -346,9 +307,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): try: response = ResponsesAPIResponse(**raw_response_json) except Exception: - verbose_logger.debug( - f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" - ) + verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct") response = ResponsesAPIResponse.model_construct(**raw_response_json) # Store processed headers in additional_headers so they get returned to the client diff --git a/litellm/llms/maritalk.py b/litellm/llms/maritalk.py index 418d13b3448..4b3a569357f 100644 --- a/litellm/llms/maritalk.py +++ b/litellm/llms/maritalk.py @@ -57,9 +57,5 @@ class MaritalkConfig(OpenAIGPTConfig): "tool_choice", ] - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return MaritalkError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return MaritalkError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/meta_llama/chat/transformation.py b/litellm/llms/meta_llama/chat/transformation.py index 6c9b79005f5..d9ffbc46f21 100644 --- a/litellm/llms/meta_llama/chat/transformation.py +++ b/litellm/llms/meta_llama/chat/transformation.py @@ -33,9 +33,7 @@ class LlamaAPIConfig(OpenAIGPTConfig): model: str, drop_params: bool, ) -> dict: - mapped_openai_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + mapped_openai_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) # Only json_schema is working for response_format if ( diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 867f6d4b1f5..a53075ba1d6 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -47,9 +47,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): def __init__(self): super().__init__() - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: api_key: Optional[str] = None if litellm_params is not None: api_key = litellm_params.api_key or get_secret_str("MILVUS_API_KEY") @@ -63,9 +61,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): return headers - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: api_key = litellm_params.get("api_key") if not api_key: raise ValueError( @@ -90,9 +86,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): ], } - def map_openai_params( - self, non_default_params: dict, optional_params: dict, drop_params: bool - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict, drop_params: bool) -> dict: for param, value in non_default_params.items(): if param in MILVUS_OPTIONAL_PARAMS: optional_params[param] = value @@ -218,17 +212,15 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): results = response_json.get("data", []) # Try to get text_field from optional_params first, then litellm_params - optional_params = litellm_logging_obj.model_call_details.get( - "optional_params", {} - ) + optional_params = litellm_logging_obj.model_call_details.get("optional_params", {}) text_field = optional_params.get("milvus_text_field", "") # Fallback to litellm_params if not in optional_params if not text_field: - text_field = litellm_logging_obj.model_call_details.get( - "litellm_params", {} - ).get("milvus_text_field", "") + text_field = litellm_logging_obj.model_call_details.get("litellm_params", {}).get( + "milvus_text_field", "" + ) # Transform results to standard format search_results: List[VectorStoreSearchResult] = [] @@ -282,7 +274,5 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): ) -> Tuple[str, Dict]: raise NotImplementedError - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: raise NotImplementedError diff --git a/litellm/llms/minimax/chat/transformation.py b/litellm/llms/minimax/chat/transformation.py index 69f228160f6..512c162658c 100644 --- a/litellm/llms/minimax/chat/transformation.py +++ b/litellm/llms/minimax/chat/transformation.py @@ -39,11 +39,7 @@ class MinimaxChatConfig(OpenAIGPTConfig): Defaults to international endpoint: https://api.minimax.io/v1 For China, set to: https://api.minimaxi.com/v1 """ - return ( - api_base - or get_secret_str("MINIMAX_API_BASE") - or "https://api.minimax.io/v1" - ) + return api_base or get_secret_str("MINIMAX_API_BASE") or "https://api.minimax.io/v1" def get_complete_url( self, diff --git a/litellm/llms/minimax/messages/transformation.py b/litellm/llms/minimax/messages/transformation.py index 57cfcbf0621..3f46aae1aaa 100644 --- a/litellm/llms/minimax/messages/transformation.py +++ b/litellm/llms/minimax/messages/transformation.py @@ -47,11 +47,7 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): Defaults to international endpoint: https://api.minimax.io/anthropic For China, set to: https://api.minimaxi.com/anthropic """ - return ( - api_base - or get_secret_str("MINIMAX_API_BASE") - or "https://api.minimax.io/anthropic/v1/messages" - ) + return api_base or get_secret_str("MINIMAX_API_BASE") or "https://api.minimax.io/anthropic/v1/messages" def get_complete_url( self, diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index 2a7d6897edc..70ce2e71731 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -202,12 +202,8 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): return headers - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return MinimaxException( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return MinimaxException(message=error_message, status_code=status_code, headers=headers) def transform_text_to_speech_request( self, @@ -240,9 +236,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): # Extract audio settings sample_rate = params.pop("sample_rate", 32000) # 16000, 24000, 32000 - bitrate = params.pop( - "bitrate", 128000 - ) # For MP3: 64000, 128000, 192000, 256000 + bitrate = params.pop("bitrate", 128000) # For MP3: 64000, 128000, 192000, 256000 channel = params.pop("channel", 1) # 1 for mono, 2 for stereo # Output format: 'url' or 'hex' (default is 'hex') diff --git a/litellm/llms/mistral/audio_transcription/transformation.py b/litellm/llms/mistral/audio_transcription/transformation.py index 8c6d604acb4..53d1428e1f1 100644 --- a/litellm/llms/mistral/audio_transcription/transformation.py +++ b/litellm/llms/mistral/audio_transcription/transformation.py @@ -27,9 +27,7 @@ class MistralAudioTranscriptionException(BaseLLMException): class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: return [ "language", "temperature", @@ -59,9 +57,7 @@ class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = ( - "https://api.mistral.ai/v1" if api_base is None else api_base.rstrip("/") - ) + api_base = "https://api.mistral.ai/v1" if api_base is None else api_base.rstrip("/") return f"{api_base}/audio/transcriptions" def get_error_class( @@ -119,9 +115,7 @@ class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig): openai_params=self.get_supported_openai_params(model), ) for key, value in provider_specific_params.items(): - form_fields[key] = ( - str(value).lower() if isinstance(value, bool) else str(value) - ) + form_fields[key] = str(value).lower() if isinstance(value, bool) else str(value) files = { "file": ( diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 8d0cf993814..0f202a22c96 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -161,9 +161,7 @@ class MistralConfig(OpenAIGPTConfig): for param, value in non_default_params.items(): if param == "max_tokens": optional_params["max_tokens"] = value - if ( - param == "max_completion_tokens" - ): # max_completion_tokens should take priority + if param == "max_completion_tokens": # max_completion_tokens should take priority optional_params["max_tokens"] = value if param == "tools": # Clean tools to remove problematic schema fields for Mistral API @@ -177,9 +175,7 @@ class MistralConfig(OpenAIGPTConfig): if param == "stop": optional_params["stop"] = value if param == "tool_choice" and isinstance(value, str): - optional_params["tool_choice"] = self._map_tool_choice( - tool_choice=value - ) + optional_params["tool_choice"] = self._map_tool_choice(tool_choice=value) if param == "seed": optional_params["extra_body"] = {"random_seed": value} if param == "response_format": @@ -205,9 +201,7 @@ class MistralConfig(OpenAIGPTConfig): ) # type: ignore # if api_base does not end with /v1 we add it - if api_base is not None and not api_base.endswith( - "/v1" - ): # Mistral always needs a /v1 at the end + if api_base is not None and not api_base.endswith("/v1"): # Mistral always needs a /v1 at the end api_base = api_base + "/v1" dynamic_api_key = ( api_key @@ -278,9 +272,7 @@ class MistralConfig(OpenAIGPTConfig): else: return super()._transform_messages(new_messages, model, False) - async def _transform_messages_async( - self, messages: List[AllMessageValues], model: str - ) -> List[AllMessageValues]: + async def _transform_messages_async(self, messages: List[AllMessageValues], model: str) -> List[AllMessageValues]: """ Handle modification of messages for Mistral API in an async context. """ @@ -290,9 +282,7 @@ class MistralConfig(OpenAIGPTConfig): messages = self._handle_message_with_file(messages) return messages - def _transform_messages_sync( - self, messages: List[AllMessageValues], model: str - ) -> List[AllMessageValues]: + def _transform_messages_sync(self, messages: List[AllMessageValues], model: str) -> List[AllMessageValues]: """Handle modification of messages for Mistral API in a sync context.""" # Call parent sync method to handle basic transformations # and then apply Mistral-specific handling for files @@ -301,9 +291,7 @@ class MistralConfig(OpenAIGPTConfig): messages = self._handle_message_with_file(messages) return messages - def _handle_message_with_file( - self, messages: List[AllMessageValues] - ) -> List[AllMessageValues]: + def _handle_message_with_file(self, messages: List[AllMessageValues]) -> List[AllMessageValues]: """ Mistral API supports only 'file_id' in message content with type 'file'. """ @@ -313,9 +301,7 @@ class MistralConfig(OpenAIGPTConfig): if any(c.get("type") == "file" for c in _content_block): # If file content is present, we get file_id from 'file' attribute of content block # then replace 'file' with 'file_id' and assign the value of 'file_id' attribute to it. - file_contents = [ - c for c in _content_block if c.get("type") == "file" - ] + file_contents = [c for c in _content_block if c.get("type") == "file"] for file_content in file_contents: file_id = file_content.get("file", {}).get("file_id") if file_id: @@ -346,21 +332,15 @@ class MistralConfig(OpenAIGPTConfig): # Handle both string and list content, preserving original format if isinstance(existing_content, str): # String content - prepend reasoning prompt - new_content: Union[str, list] = ( - f"{reasoning_prompt}\n\n{existing_content}" - ) + new_content: Union[str, list] = f"{reasoning_prompt}\n\n{existing_content}" elif isinstance(existing_content, list): # List content - prepend reasoning prompt as text block - new_content = [ - {"type": "text", "text": reasoning_prompt + "\n\n"} - ] + existing_content + new_content = [{"type": "text", "text": reasoning_prompt + "\n\n"}] + existing_content else: # Fallback for any other type - convert to string new_content = f"{reasoning_prompt}\n\n{str(existing_content)}" - messages[i] = cast( - AllMessageValues, {**msg, "content": new_content} - ) + messages[i] = cast(AllMessageValues, {**msg, "content": new_content}) break else: # Add new system message with reasoning instructions @@ -405,9 +385,7 @@ class MistralConfig(OpenAIGPTConfig): cleaned_tools = copy.deepcopy(tools) # Apply all cleaning functions with max_depth protection - cleaned_tools = _remove_json_schema_refs( - cleaned_tools, max_depth=DEFAULT_MAX_RECURSE_DEPTH - ) + cleaned_tools = _remove_json_schema_refs(cleaned_tools, max_depth=DEFAULT_MAX_RECURSE_DEPTH) return cleaned_tools @@ -423,11 +401,7 @@ class MistralConfig(OpenAIGPTConfig): return message return cast( AllMessageValues, - { - k: v - for k, v in message.items() - if k not in ("reasoning_content", "thinking_blocks") - }, + {k: v for k, v in message.items() if k not in ("reasoning_content", "thinking_blocks")}, ) @classmethod @@ -514,9 +488,7 @@ class MistralConfig(OpenAIGPTConfig): """ Convert Mistral thinking blocks to reasoning content. """ - return "\n".join( - [block.get("text", "") for block in thinking_blocks["thinking"]] - ) + return "\n".join([block.get("text", "") for block in thinking_blocks["thinking"]]) @staticmethod def _handle_content_list_to_str_conversion(response_data: dict) -> dict: @@ -545,9 +517,7 @@ class MistralConfig(OpenAIGPTConfig): thinking_texts = [] for thinking_block in thinking_blocks: if thinking_block.get("type") == "text": - thinking_texts.append( - thinking_block.get("text", "") - ) + thinking_texts.append(thinking_block.get("text", "")) thinking_content = "\n".join(thinking_texts) elif block.get("type") == "text": text_content = block.get("text", "") @@ -575,12 +545,8 @@ class MistralConfig(OpenAIGPTConfig): dict: The transformed request. Sent as the body of the API call. """ # Add reasoning system prompt if needed (for magistral models) - if "magistral" in model.lower() and optional_params.get( - "_add_reasoning_prompt", False - ): - messages = self._add_reasoning_system_prompt_if_needed( - messages, optional_params - ) + if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False): + messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params) # Call parent transform_request which handles _transform_messages return super().transform_request( @@ -701,7 +667,5 @@ class MistralChatResponseIterator(OpenAIChatCompletionStreamingHandler): text_segments.append(block.get("text", "")) normalized_text = "".join(text_segments) if text_segments else None - reasoning_content = ( - "\n".join(reasoning_segments) if reasoning_segments else None - ) + reasoning_content = "\n".join(reasoning_segments) if reasoning_segments else None return normalized_text, thinking_blocks, reasoning_content diff --git a/litellm/llms/mistral/ocr/guardrail_translation/handler.py b/litellm/llms/mistral/ocr/guardrail_translation/handler.py index 7d3797a1dbe..9144c71f70a 100644 --- a/litellm/llms/mistral/ocr/guardrail_translation/handler.py +++ b/litellm/llms/mistral/ocr/guardrail_translation/handler.py @@ -51,9 +51,7 @@ class OCRHandler(BaseTranslation): """ document = data.get("document") if document is None or not isinstance(document, dict): - verbose_proxy_logger.debug( - "OCR guardrail: No valid document found in request data" - ) + verbose_proxy_logger.debug("OCR guardrail: No valid document found in request data") return data # Extract the document URL for guardrail checking @@ -135,9 +133,7 @@ class OCRHandler(BaseTranslation): # Add user metadata if available if user_api_key_dict is not None: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: # Preserve original behavior: inject metadata into inputs for # third-party guardrail providers that read it from there diff --git a/litellm/llms/modelscope/chat/transformation.py b/litellm/llms/modelscope/chat/transformation.py index 162ef1a236c..1a54be6e1c8 100644 --- a/litellm/llms/modelscope/chat/transformation.py +++ b/litellm/llms/modelscope/chat/transformation.py @@ -54,20 +54,14 @@ class ModelScopeChatConfig(OpenAIGPTConfig): message["content"] = "".join(item.get("text") or "" for item in content) if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = ( - api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL - ) # type: ignore + api_base = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL # type: ignore dynamic_api_key = api_key or get_secret_str("MODELSCOPE_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 0d85f7796fb..a3d890734d1 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -41,9 +41,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://api-inference.modelscope.cn/v1" - def get_supported_openai_params( - self, model: str - ) -> list[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: """ Return list of OpenAI params supported by ModelScope. @@ -70,9 +68,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): """ supported_params = self.get_supported_openai_params(model) if drop_params: - non_default_params = { - k: v for k, v in non_default_params.items() if k in supported_params - } + non_default_params = {k: v for k, v in non_default_params.items() if k in supported_params} optional_params.update(non_default_params) return optional_params @@ -89,9 +85,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the ModelScope image generation API request. """ - base_url: str = ( - api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL - ) + base_url: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL base_url = base_url.rstrip("/") # Return the images endpoint @@ -115,8 +109,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): if not final_api_key: raise ValueError( - "MODELSCOPE_API_KEY is not set. " - "Please set it via environment variable or pass api_key parameter." + "MODELSCOPE_API_KEY is not set. Please set it via environment variable or pass api_key parameter." ) default_headers = { @@ -185,9 +178,7 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): # Check for errors in response if "error" in response_data: - error_msg = response_data["error"].get( - "message", str(response_data["error"]) - ) + error_msg = response_data["error"].get("message", str(response_data["error"])) raise self.get_error_class( error_message=f"ModelScope error: {error_msg}", status_code=raw_response.status_code, diff --git a/litellm/llms/moonshot/chat/transformation.py b/litellm/llms/moonshot/chat/transformation.py index a5ac696aae2..07a963e95fd 100644 --- a/litellm/llms/moonshot/chat/transformation.py +++ b/litellm/llms/moonshot/chat/transformation.py @@ -53,22 +53,14 @@ class MoonshotChatConfig(OpenAIGPTConfig): messages = handle_messages_with_content_list_to_str_conversion(messages) if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = ( - api_base - or get_secret_str("MOONSHOT_API_BASE") - or "https://api.moonshot.ai/v1" - ) # type: ignore + api_base = api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("MOONSHOT_API_KEY") return api_base, dynamic_api_key @@ -153,9 +145,7 @@ class MoonshotChatConfig(OpenAIGPTConfig): optional_params["temperature"] = 0.3 return optional_params - def fill_reasoning_content( - self, messages: List[AllMessageValues] - ) -> List[AllMessageValues]: + def fill_reasoning_content(self, messages: List[AllMessageValues]) -> List[AllMessageValues]: """ Moonshot reasoning models require `reasoning_content` on every assistant message that contains tool_calls (multi-turn tool-calling flows). diff --git a/litellm/llms/morph/chat/transformation.py b/litellm/llms/morph/chat/transformation.py index 93bd7e16aef..97ddc12920f 100644 --- a/litellm/llms/morph/chat/transformation.py +++ b/litellm/llms/morph/chat/transformation.py @@ -25,9 +25,7 @@ class MorphChatConfig(OpenAILikeChatConfig): self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: api_base = ( - api_base - or get_secret_str("MORPH_API_BASE") - or "https://api.morphllm.com/v1" # default api base + api_base or get_secret_str("MORPH_API_BASE") or "https://api.morphllm.com/v1" # default api base ) dynamic_api_key = api_key or get_secret_str("MORPH_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/nlp_cloud/chat/transformation.py b/litellm/llms/nlp_cloud/chat/transformation.py index 8037a458321..5aafc4cd45c 100644 --- a/litellm/llms/nlp_cloud/chat/transformation.py +++ b/litellm/llms/nlp_cloud/chat/transformation.py @@ -146,9 +146,7 @@ class NLPCloudConfig(BaseConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return NLPCloudError( - status_code=status_code, message=error_message, headers=headers - ) + return NLPCloudError(status_code=status_code, message=error_message, headers=headers) def transform_request( self, @@ -193,9 +191,7 @@ class NLPCloudConfig(BaseConfig): try: completion_response = raw_response.json() except Exception: - raise NLPCloudError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise NLPCloudError(message=raw_response.text, status_code=raw_response.status_code) if "error" in completion_response: raise NLPCloudError( message=completion_response["error"], diff --git a/litellm/llms/nscale/chat/transformation.py b/litellm/llms/nscale/chat/transformation.py index 6103b8e3c49..1b032fab2ac 100644 --- a/litellm/llms/nscale/chat/transformation.py +++ b/litellm/llms/nscale/chat/transformation.py @@ -23,9 +23,7 @@ class NscaleConfig(OpenAIGPTConfig): @staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: - return ( - api_base or get_secret_str("NSCALE_API_BASE") or NscaleConfig.API_BASE_URL - ) + return api_base or get_secret_str("NSCALE_API_BASE") or NscaleConfig.API_BASE_URL def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index 7ae3f913297..2d72d52f991 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -156,9 +156,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): api_key = get_secret_str("NVIDIA_NIM_API_KEY") or litellm.api_key if api_key is None: - raise ValueError( - "Nvidia NIM API key is required. Please set 'NVIDIA_NIM_API_KEY' in your environment" - ) + raise ValueError("Nvidia NIM API key is required. Please set 'NVIDIA_NIM_API_KEY' in your environment") default_headers = { "Authorization": f"Bearer {api_key}", @@ -235,18 +233,12 @@ class NvidiaNimRerankConfig(BaseRerankConfig): } # Add optional top_k parameter if provided (already mapped from top_n in map_cohere_rerank_params) - if ( - "top_k" in optional_rerank_params - and optional_rerank_params.get("top_k") is not None - ): # type: ignore + if "top_k" in optional_rerank_params and optional_rerank_params.get("top_k") is not None: # type: ignore request_data["top_k"] = optional_rerank_params.get("top_k") # type: ignore # Add Nvidia-specific truncate parameter if provided # This is passed through from non_default_params, not in base OptionalRerankParams - if ( - "truncate" in optional_rerank_params - and optional_rerank_params.get("truncate") is not None - ): # type: ignore + if "truncate" in optional_rerank_params and optional_rerank_params.get("truncate") is not None: # type: ignore truncate_value = optional_rerank_params.get("truncate") # type: ignore if truncate_value in ["NONE", "END"]: request_data["truncate"] = truncate_value # type: ignore @@ -305,9 +297,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): rankings = nvidia_response.get("rankings", []) # Get original documents from request if we need to include them - original_passages: List[NvidiaNimPassageObject] = request_data.get( - "passages", [] - ) + original_passages: List[NvidiaNimPassageObject] = request_data.get("passages", []) for ranking in rankings: result_item: RerankResponseResult = { @@ -327,9 +317,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): usage = raw_response_json.get("usage", {}) total_tokens = usage.get("total_tokens", 0) - billed_units: RerankBilledUnits = { - "total_tokens": total_tokens if total_tokens > 0 else len(results) - } + billed_units: RerankBilledUnits = {"total_tokens": total_tokens if total_tokens > 0 else len(results)} meta: RerankResponseMeta = {"billed_units": billed_units} diff --git a/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py b/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py index 253d6d2f73f..7ec679c858d 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py +++ b/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py @@ -30,10 +30,7 @@ from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException FloatArray = Any -_INSTALL_HINT = ( - "Install Riva STT extras to enable automatic audio resampling: " - "`pip install 'litellm[stt-nvidia-riva]'`" -) +_INSTALL_HINT = "Install Riva STT extras to enable automatic audio resampling: `pip install 'litellm[stt-nvidia-riva]'`" @dataclass @@ -69,9 +66,7 @@ def resample_to_riva_pcm(file_bytes: bytes) -> ResampledAudio: samples_float = np.asarray(samples_float, dtype=np.float32).ravel() if source_rate != RIVA_TARGET_SAMPLE_RATE_HZ: - samples_float = _resample( - samples_float, source_rate, RIVA_TARGET_SAMPLE_RATE_HZ - ) + samples_float = _resample(samples_float, source_rate, RIVA_TARGET_SAMPLE_RATE_HZ) # Clip + convert float [-1, 1] to int16 little-endian PCM. np.clip(samples_float, -1.0, 1.0, out=samples_float) @@ -167,9 +162,7 @@ def _decode_to_float32(file_bytes: bytes) -> Tuple["FloatArray", int]: pass -def _resample( - samples: "FloatArray", source_rate: int, target_rate: int -) -> "FloatArray": +def _resample(samples: "FloatArray", source_rate: int, target_rate: int) -> "FloatArray": """ Resample mono float32 ``samples`` from ``source_rate`` to ``target_rate``. @@ -189,9 +182,7 @@ def _resample( return cast( "FloatArray", - np.asarray( - soxr.resample(samples, source_rate, target_rate), dtype=np.float32 - ), + np.asarray(soxr.resample(samples, source_rate, target_rate), dtype=np.float32), ) except ImportError: pass @@ -204,18 +195,14 @@ def _resample( g = gcd(int(source_rate), int(target_rate)) up = int(target_rate) // g down = int(source_rate) // g - return cast( - "FloatArray", np.asarray(resample_poly(samples, up, down), dtype=np.float32) - ) + return cast("FloatArray", np.asarray(resample_poly(samples, up, down), dtype=np.float32)) except ImportError: pass return _linear_resample(samples, source_rate, target_rate) -def _linear_resample( - samples: "FloatArray", source_rate: int, target_rate: int -) -> "FloatArray": +def _linear_resample(samples: "FloatArray", source_rate: int, target_rate: int) -> "FloatArray": """Linear-interpolation fallback. See :func:`_resample` for caveats.""" import numpy as np # type: ignore diff --git a/litellm/llms/nvidia_riva/audio_transcription/handler.py b/litellm/llms/nvidia_riva/audio_transcription/handler.py index 9740162ba1c..eab5abd475b 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/handler.py +++ b/litellm/llms/nvidia_riva/audio_transcription/handler.py @@ -59,10 +59,7 @@ _DEFAULT_CHUNK_SAMPLES = 1600 _DEFAULT_CHUNK_BYTES = _DEFAULT_CHUNK_SAMPLES * 2 # int16 = 2 bytes/sample -_RIVA_INSTALL_HINT = ( - "NVIDIA Riva client is not installed. " - "Install with `pip install 'litellm[stt-nvidia-riva]'`." -) +_RIVA_INSTALL_HINT = "NVIDIA Riva client is not installed. Install with `pip install 'litellm[stt-nvidia-riva]'`." class NvidiaRivaAudioTranscription: @@ -209,9 +206,7 @@ class NvidiaRivaAudioTranscription: riva_asr_module=riva_asr_module, recognition_config_dict=recognition_config_dict, ) - streaming_config = riva_asr_module.StreamingRecognitionConfig( - config=recognition_config, interim_results=False - ) + streaming_config = riva_asr_module.StreamingRecognitionConfig(config=recognition_config, interim_results=False) logging_obj.pre_call( input=None, @@ -221,9 +216,7 @@ class NvidiaRivaAudioTranscription: "atranscription": atranscription, "complete_input_dict": { "recognition_config": recognition_config_dict, - "nvcf_function_id_set": bool( - optional_params.get("nvcf_function_id") - ), + "nvcf_function_id_set": bool(optional_params.get("nvcf_function_id")), "use_ssl": optional_params.get("use_ssl"), }, }, @@ -239,9 +232,7 @@ class NvidiaRivaAudioTranscription: # Forward the deadline so the stream cannot block forever if the # server stalls. Older riva-client versions do not accept a # ``timeout`` kwarg, so pass it only when supported. - if timeout is not None and self._supports_timeout_kwarg( - asr_service.streaming_response_generator - ): + if timeout is not None and self._supports_timeout_kwarg(asr_service.streaming_response_generator): stream_kwargs["timeout"] = float(timeout) stream = asr_service.streaming_response_generator(**stream_kwargs) final_results = self._collect_final_results(stream) @@ -300,11 +291,7 @@ class NvidiaRivaAudioTranscription: """ nvcf_function_id = optional_params.get("nvcf_function_id") use_ssl_override = optional_params.get("use_ssl") - use_ssl = ( - bool(use_ssl_override) - if use_ssl_override is not None - else bool(nvcf_function_id) - ) + use_ssl = bool(use_ssl_override) if use_ssl_override is not None else bool(nvcf_function_id) metadata: List[Tuple[str, str]] = [] if nvcf_function_id: @@ -313,19 +300,13 @@ class NvidiaRivaAudioTranscription: metadata.append(("authorization", f"Bearer {api_key}")) try: - return riva_module.Auth( - uri=api_base, use_ssl=use_ssl, metadata_args=metadata - ) + return riva_module.Auth(uri=api_base, use_ssl=use_ssl, metadata_args=metadata) except TypeError: # Older riva-client signatures used positional-only args. return riva_module.Auth(None, use_ssl, api_base, metadata) - def _build_recognition_config_proto( - self, riva_asr_module: Any, recognition_config_dict: Dict[str, Any] - ): - encoding_name = ( - recognition_config_dict.get("encoding") or "LINEAR_PCM" - ).upper() + def _build_recognition_config_proto(self, riva_asr_module: Any, recognition_config_dict: Dict[str, Any]): + encoding_name = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper() encoding_enum = getattr( riva_asr_module.AudioEncoding, encoding_name, @@ -337,20 +318,12 @@ class NvidiaRivaAudioTranscription: sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]), language_code=recognition_config_dict["language_code"], audio_channel_count=int(recognition_config_dict["audio_channel_count"]), - enable_automatic_punctuation=bool( - recognition_config_dict.get("enable_automatic_punctuation", True) - ), - enable_word_time_offsets=bool( - recognition_config_dict.get("enable_word_time_offsets", False) - ), + enable_automatic_punctuation=bool(recognition_config_dict.get("enable_automatic_punctuation", True)), + enable_word_time_offsets=bool(recognition_config_dict.get("enable_word_time_offsets", False)), max_alternatives=int(recognition_config_dict.get("max_alternatives", 1)), model=recognition_config_dict.get("model", "") or "", - verbatim_transcripts=bool( - recognition_config_dict.get("verbatim_transcripts", False) - ), - profanity_filter=bool( - recognition_config_dict.get("profanity_filter", False) - ), + verbatim_transcripts=bool(recognition_config_dict.get("verbatim_transcripts", False)), + profanity_filter=bool(recognition_config_dict.get("profanity_filter", False)), ) endpointing = recognition_config_dict.get("endpointing_config") @@ -437,8 +410,6 @@ def _import_riva(): riva_asr_module = riva_asr_pb2 except ImportError as e: - raise NvidiaRivaException( - status_code=500, message=_RIVA_INSTALL_HINT - ) from e + raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e return riva_client, riva_asr_module diff --git a/litellm/llms/nvidia_riva/audio_transcription/transformation.py b/litellm/llms/nvidia_riva/audio_transcription/transformation.py index c2dfc25d945..43185cb2f7a 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/transformation.py +++ b/litellm/llms/nvidia_riva/audio_transcription/transformation.py @@ -43,9 +43,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): optional TLS via ``use_ssl``). """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: # Riva natively understands language + word timestamps. # `response_format` is honored at response-shaping time in the handler. return ["language", "response_format", "timestamp_granularities"] @@ -79,12 +77,8 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return NvidiaRivaException( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return NvidiaRivaException(message=error_message, status_code=status_code, headers=headers) def transform_audio_transcription_request( self, @@ -141,9 +135,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # gRPC auth is constructed in the handler, not via HTTP headers. return headers - def _build_recognition_config_dict( - self, model: str, optional_params: dict - ) -> Dict[str, Any]: + def _build_recognition_config_dict(self, model: str, optional_params: dict) -> Dict[str, Any]: """ Build the Riva ``RecognitionConfig`` shape as a plain dict. @@ -156,28 +148,18 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """ return { "language_code": optional_params.get("language_code", "en-US"), - "sample_rate_hertz": optional_params.get( - "sample_rate_hertz", RIVA_TARGET_SAMPLE_RATE_HZ - ), + "sample_rate_hertz": optional_params.get("sample_rate_hertz", RIVA_TARGET_SAMPLE_RATE_HZ), "encoding": optional_params.get("encoding", RIVA_TARGET_ENCODING), - "audio_channel_count": optional_params.get( - "audio_channel_count", RIVA_TARGET_NUM_CHANNELS - ), - "enable_automatic_punctuation": optional_params.get( - "enable_automatic_punctuation", True - ), - "enable_word_time_offsets": bool( - optional_params.get("enable_word_time_offsets", False) - ), + "audio_channel_count": optional_params.get("audio_channel_count", RIVA_TARGET_NUM_CHANNELS), + "enable_automatic_punctuation": optional_params.get("enable_automatic_punctuation", True), + "enable_word_time_offsets": bool(optional_params.get("enable_word_time_offsets", False)), "max_alternatives": optional_params.get("max_alternatives", 1), "model": optional_params.get("riva_model_name", ""), "verbatim_transcripts": optional_params.get("verbatim_transcripts", False), "profanity_filter": optional_params.get("profanity_filter", False), } - def _build_endpointing_config_dict( - self, optional_params: dict - ) -> Optional[Dict[str, Any]]: + def _build_endpointing_config_dict(self, optional_params: dict) -> Optional[Dict[str, Any]]: """ Translate an OpenAI-style ``chunking_strategy`` into Riva's ``EndpointingConfig`` shape, or pass through an explicit @@ -257,9 +239,7 @@ class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig): only ``result.is_final`` entries (empty/non-final chunks are ignored). """ - full_transcript = "".join( - (item.get("transcript") or "") for item in final_results - ).strip() + full_transcript = "".join((item.get("transcript") or "") for item in final_results).strip() response = TranscriptionResponse(text=full_transcript) response["task"] = "transcribe" diff --git a/litellm/llms/nvidia_riva/common_utils.py b/litellm/llms/nvidia_riva/common_utils.py index a3071cf7060..4206fc91cc6 100644 --- a/litellm/llms/nvidia_riva/common_utils.py +++ b/litellm/llms/nvidia_riva/common_utils.py @@ -84,9 +84,5 @@ def grpc_error_to_litellm_exception(error: Exception) -> NvidiaRivaException: http_status = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500) detail = _extract_grpc_details(error) or str(error) - message = ( - f"NVIDIA Riva gRPC error ({status_name}): {detail}" - if status_name - else f"NVIDIA Riva error: {detail}" - ) + message = f"NVIDIA Riva gRPC error ({status_name}): {detail}" if status_name else f"NVIDIA Riva error: {detail}" return NvidiaRivaException(status_code=http_status, message=message) diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index de8a09b7a3b..3661ac908d2 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -54,9 +54,7 @@ def _extract_text_content(content: Any) -> str: return content if isinstance(content, list): return "".join( - item.get("text", "") - for item in content - if isinstance(item, dict) and item.get("type") == "text" + item.get("text", "") for item in content if isinstance(item, dict) and item.get("type") == "text" ) return str(content) @@ -88,9 +86,7 @@ def adapt_messages_to_cohere_standard( tc_id = tc.get("id", "") raw_args: Any = tc.get("function", {}).get("arguments", "{}") try: - params: Dict[str, Any] = ( - json.loads(raw_args) if isinstance(raw_args, str) else raw_args - ) + params: Dict[str, Any] = json.loads(raw_args) if isinstance(raw_args, str) else raw_args except json.JSONDecodeError: params = {} tool_call_lookup[tc_id] = CohereToolCall( @@ -99,17 +95,11 @@ def adapt_messages_to_cohere_standard( ) last_user_index = next( - ( - i - for i in range(len(messages) - 1, -1, -1) - if messages[i].get("role") == "user" - ), + (i for i in range(len(messages) - 1, -1, -1) if messages[i].get("role") == "user"), None, ) history_source = ( - messages - if last_user_index is None - else [m for i, m in enumerate(messages) if i != last_user_index] + messages if last_user_index is None else [m for i, m in enumerate(messages) if i != last_user_index] ) chat_history: List[CohereMessage] = [] @@ -139,14 +129,10 @@ def adapt_messages_to_cohere_standard( if role == "user": chat_history.append(CohereMessage(role="USER", message=content)) elif role == "assistant": - chat_history.append( - CohereMessage(role="CHATBOT", message=content, toolCalls=tool_calls) - ) + chat_history.append(CohereMessage(role="CHATBOT", message=content, toolCalls=tool_calls)) elif role == "tool": tool_call_id = str(msg.get("tool_call_id", "") or "") - cohere_call = tool_call_lookup.get( - tool_call_id, CohereToolCall(name="", parameters={}) - ) + cohere_call = tool_call_lookup.get(tool_call_id, CohereToolCall(name="", parameters={})) tool_result = CohereToolResult( call=cohere_call, outputs=[{"output": content}], @@ -179,9 +165,7 @@ def adapt_tool_definitions_to_cohere_standard( function_def = tool.get("function", {}) raw_params = function_def.get("parameters", {}) - resolved = sanitize_oci_schema( - resolve_oci_schema_anyof(resolve_oci_schema_refs(raw_params)) - ) + resolved = sanitize_oci_schema(resolve_oci_schema_anyof(resolve_oci_schema_refs(raw_params))) properties = resolved.get("properties", {}) required = resolved.get("required", []) @@ -190,9 +174,7 @@ def adapt_tool_definitions_to_cohere_standard( json_type = param_schema.get("type", "string") python_type = OCI_JSON_TO_PYTHON_TYPES.get(json_type, json_type) parameter_definitions[param_name] = CohereParameterDefinition( - description=enrich_cohere_param_description( - param_schema.get("description", ""), param_schema - ), + description=enrich_cohere_param_description(param_schema.get("description", ""), param_schema), type=python_type, isRequired=param_name in required, ) @@ -227,17 +209,13 @@ def handle_cohere_response( model_response.created = int(datetime.datetime.now().timestamp()) response_text = cohere_response.chatResponse.text - finish_reason = _normalize_oci_finish_reason( - cohere_response.chatResponse.finishReason - ) + finish_reason = _normalize_oci_finish_reason(cohere_response.chatResponse.finishReason) tool_calls: Optional[List[Dict[str, Any]]] = None if cohere_response.chatResponse.toolCalls: tool_calls = [ { - "id": _synthesize_oci_tool_call_id( - i, tc.name, json.dumps(tc.parameters, sort_keys=True) - ), + "id": _synthesize_oci_tool_call_id(i, tc.name, json.dumps(tc.parameters, sort_keys=True)), "type": "function", "function": { "name": tc.name, @@ -274,9 +252,7 @@ def handle_cohere_response( total_tokens=usage_info.totalTokens, ) else: - model_response.usage = Usage( - prompt_tokens=0, completion_tokens=0, total_tokens=0 - ) # type: ignore[attr-defined] + model_response.usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) # type: ignore[attr-defined] return model_response @@ -319,9 +295,7 @@ def handle_cohere_stream_chunk( # already-streamed deltas. We require both signals to be present so that a # future API change which adds `chatHistory` to intermediate chunks (or a # rare early-populated case) doesn't silently drop legitimate token deltas. - is_terminal_consolidation = ( - typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None - ) + is_terminal_consolidation = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None # On non-terminal text-free chunks (e.g. tool-call-only or keep-alive # chunks) emit ``content=None`` rather than ``content=""`` so downstream # stream-mergers that distinguish "no text in this delta" from "an @@ -331,9 +305,7 @@ def handle_cohere_stream_chunk( # confirmed that text deltas were already emitted earlier — otherwise # (e.g. a degenerate stream that delivers the whole response in a # single SSE event), passing it through is the only chance to surface it. - text: Optional[str] = ( - None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text - ) + text: Optional[str] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text # Tool calls on the terminal consolidation chunk (whether from # `typed_chunk.toolCalls` or from `chatHistory`) typically restate what @@ -343,11 +315,7 @@ def handle_cohere_stream_chunk( # tool calls were already emitted earlier — otherwise (e.g. a short # response that delivers tool calls exclusively on the terminal chunk), # passing them through is the only chance to surface them. - cohere_tool_calls = ( - None - if (is_terminal_consolidation and prior_tool_calls_emitted) - else typed_chunk.toolCalls - ) + cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls tool_calls: Optional[List[Dict[str, Any]]] = None if cohere_tool_calls: @@ -357,9 +325,7 @@ def handle_cohere_stream_chunk( # deterministically from the call's content/position. A random # uuid4 per chunk would cause downstream stream-mergers to # treat each chunk as a distinct tool call. - "id": _synthesize_oci_tool_call_id( - i, tc.name, json.dumps(tc.parameters, sort_keys=True) - ), + "id": _synthesize_oci_tool_call_id(i, tc.name, json.dumps(tc.parameters, sort_keys=True)), "type": "function", "function": { "name": tc.name, diff --git a/litellm/llms/oci/chat/generic.py b/litellm/llms/oci/chat/generic.py index 2cc1ac77a40..02ec762488d 100644 --- a/litellm/llms/oci/chat/generic.py +++ b/litellm/llms/oci/chat/generic.py @@ -55,9 +55,7 @@ open_ai_to_generic_oci_role_map: Dict[str, OCIRoles] = { # --------------------------------------------------------------------------- -def adapt_messages_to_generic_oci_standard_content_message( - role: str, content: Union[str, list] -) -> OCIMessage: +def adapt_messages_to_generic_oci_standard_content_message(role: str, content: Union[str, list]) -> OCIMessage: """Convert a plain-text or multipart content message to OCI format.""" new_content: List[OCIContentPartUnion] = [] if isinstance(content, str): @@ -70,9 +68,7 @@ def adapt_messages_to_generic_oci_standard_content_message( for content_item in content: if not isinstance(content_item, dict): - raise OCIError( - status_code=400, message="Each content item must be a dictionary" - ) + raise OCIError(status_code=400, message="Each content item must be a dictionary") item_type = content_item.get("type") if not isinstance(item_type, str): @@ -114,20 +110,14 @@ def adapt_messages_to_generic_oci_standard_content_message( ) -def adapt_messages_to_generic_oci_standard_tool_call( - role: str, tool_calls: list -) -> OCIMessage: +def adapt_messages_to_generic_oci_standard_tool_call(role: str, tool_calls: list) -> OCIMessage: """Convert an assistant tool-call message to OCI format.""" tool_calls_formatted = [] for tool_call in tool_calls: if not isinstance(tool_call, dict): - raise OCIError( - status_code=400, message="Each tool call must be a dictionary" - ) + raise OCIError(status_code=400, message="Each tool call must be a dictionary") if tool_call.get("type") != "function": - raise OCIError( - status_code=400, message="OCI only supports function tool calls" - ) + raise OCIError(status_code=400, message="OCI only supports function tool calls") tool_call_id = tool_call.get("id") if not isinstance(tool_call_id, str): @@ -135,15 +125,11 @@ def adapt_messages_to_generic_oci_standard_tool_call( tool_function = tool_call.get("function") if not isinstance(tool_function, dict): - raise OCIError( - status_code=400, message="Tool call `function` must be a dictionary" - ) + raise OCIError(status_code=400, message="Tool call `function` must be a dictionary") function_name = tool_function.get("name") if not isinstance(function_name, str): - raise OCIError( - status_code=400, message="Tool call `function.name` must be a string" - ) + raise OCIError(status_code=400, message="Tool call `function.name` must be a string") arguments = tool_call["function"].get("arguments", "{}") if not isinstance(arguments, str): @@ -169,9 +155,7 @@ def adapt_messages_to_generic_oci_standard_tool_call( ) -def adapt_messages_to_generic_oci_standard_tool_response( - role: str, tool_call_id: str, content: str -) -> OCIMessage: +def adapt_messages_to_generic_oci_standard_tool_response(role: str, tool_call_id: str, content: str) -> OCIMessage: """Convert a tool-result message to OCI format.""" return OCIMessage( role=open_ai_to_generic_oci_role_map[role], @@ -194,12 +178,8 @@ def adapt_messages_to_generic_oci_standard( if role == "assistant" and tool_calls is not None: if not isinstance(tool_calls, list): - raise OCIError( - status_code=400, message="Message `tool_calls` must be a list" - ) - new_messages.append( - adapt_messages_to_generic_oci_standard_tool_call(role, tool_calls) - ) + raise OCIError(status_code=400, message="Message `tool_calls` must be a list") + new_messages.append(adapt_messages_to_generic_oci_standard_tool_call(role, tool_calls)) elif role in ["system", "user", "assistant"] and content is not None: if not isinstance(content, (str, list)): @@ -207,9 +187,7 @@ def adapt_messages_to_generic_oci_standard( status_code=400, message="Message `content` must be a string or list of content parts", ) - new_messages.append( - adapt_messages_to_generic_oci_standard_content_message(role, content) - ) + new_messages.append(adapt_messages_to_generic_oci_standard_content_message(role, content)) elif role == "tool": if not isinstance(tool_call_id, str): @@ -222,11 +200,7 @@ def adapt_messages_to_generic_oci_standard( status_code=400, message="Tool result message `content` must be a string", ) - new_messages.append( - adapt_messages_to_generic_oci_standard_tool_response( - role, tool_call_id, content - ) - ) + new_messages.append(adapt_messages_to_generic_oci_standard_tool_response(role, tool_call_id, content)) return new_messages @@ -236,9 +210,7 @@ def adapt_messages_to_generic_oci_standard( # --------------------------------------------------------------------------- -def adapt_tool_definition_to_oci_standard( - tools: List[Dict], vendor: OCIVendors -) -> List[OCIToolDefinition]: +def adapt_tool_definition_to_oci_standard(tools: List[Dict], vendor: OCIVendors) -> List[OCIToolDefinition]: """Convert OpenAI-format tool definitions to OCI GENERIC format. Resolves ``$ref``/``$defs`` and ``anyOf`` that the OCI endpoint rejects. @@ -250,14 +222,10 @@ def adapt_tool_definition_to_oci_standard( tool_function = tool.get("function") if not isinstance(tool_function, dict): - raise OCIError( - status_code=400, message="Tool `function` must be a dictionary" - ) + raise OCIError(status_code=400, message="Tool `function` must be a dictionary") raw_params = tool_function.get("parameters", {}) - resolved_params = sanitize_oci_schema( - resolve_oci_schema_anyof(resolve_oci_schema_refs(raw_params)) - ) + resolved_params = sanitize_oci_schema(resolve_oci_schema_anyof(resolve_oci_schema_refs(raw_params))) new_tools.append( OCIToolDefinition( @@ -370,9 +338,7 @@ def handle_generic_response( if text is not None: message.content = text if response_message.toolCalls: - message.tool_calls = adapt_tools_to_openai_standard( - response_message.toolCalls - ) + message.tool_calls = adapt_tools_to_openai_standard(response_message.toolCalls) model_response.choices[0].finish_reason = _normalize_oci_finish_reason( # type: ignore[union-attr,assignment] response_choice.finishReason @@ -380,10 +346,7 @@ def handle_generic_response( oci_usage = completion_response.chatResponse.usage reasoning_tokens: Optional[int] = None - if ( - oci_usage.completionTokensDetails - and oci_usage.completionTokensDetails.reasoningTokens is not None - ): + if oci_usage.completionTokensDetails and oci_usage.completionTokensDetails.reasoningTokens is not None: reasoning_tokens = oci_usage.completionTokensDetails.reasoningTokens model_response.usage = Usage( # type: ignore[attr-defined] prompt_tokens=oci_usage.promptTokens, @@ -456,9 +419,7 @@ def handle_generic_stream_chunk(dict_chunk: dict) -> ModelResponseStream: for i, tc in enumerate(typed_chunk.message.toolCalls) ] - finish_reason: Optional[str] = _normalize_oci_finish_reason( - typed_chunk.finishReason - ) + finish_reason: Optional[str] = _normalize_oci_finish_reason(typed_chunk.finishReason) return ModelResponseStream( choices=[ diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 35d5aefeacb..496656dd5ac 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -154,9 +154,7 @@ def _normalize_tool_choice(selected_params: Dict) -> None: "required": {"type": "REQUIRED"}, "any": {"type": "REQUIRED"}, } - selected_params["toolChoice"] = tc_map.get( - tc.lower(), {"type": "FUNCTION", "name": tc} - ) + selected_params["toolChoice"] = tc_map.get(tc.lower(), {"type": "FUNCTION", "name": tc}) return if isinstance(tc, dict): raw_type = tc.get("type") @@ -188,10 +186,7 @@ def _normalize_tool_choice(selected_params: Dict) -> None: return raise OCIError( status_code=400, - message=( - f"Invalid tool_choice for OCI: expected str or dict, got " - f"{type(tc).__name__}" - ), + message=(f"Invalid tool_choice for OCI: expected str or dict, got {type(tc).__name__}"), ) @@ -239,9 +234,7 @@ def _normalize_response_format(selected_params: Dict, vendor: OCIVendors) -> Non return fmt = rf_type.upper() - selected_params["responseFormat"] = { - "type": "JSON_OBJECT" if fmt == "JSON" else fmt - } + selected_params["responseFormat"] = {"type": "JSON_OBJECT" if fmt == "JSON" else fmt} def get_vendor_from_model(model: str) -> OCIVendors: @@ -305,8 +298,7 @@ class OCIChatConfig(BaseConfig): # ``map_openai_params`` either drops them (under drop_params) or raises # a clear error, rather than silently passing them through. self.openai_to_oci_cohere_param_map = { - k: ("stopSequences" if k == "stop" else v) - for k, v in self.openai_to_oci_generic_param_map.items() + k: ("stopSequences" if k == "stop" else v) for k, v in self.openai_to_oci_generic_param_map.items() } self.openai_to_oci_cohere_param_map["tool_choice"] = False self.openai_to_oci_cohere_param_map["n"] = False @@ -350,9 +342,7 @@ class OCIChatConfig(BaseConfig): adapted_params = {} vendor = get_vendor_from_model(model) param_map = ( - self.openai_to_oci_cohere_param_map - if vendor == OCIVendors.COHERE - else self.openai_to_oci_generic_param_map + self.openai_to_oci_cohere_param_map if vendor == OCIVendors.COHERE else self.openai_to_oci_generic_param_map ) for key, value in {**non_default_params, **optional_params}.items(): @@ -464,13 +454,9 @@ class OCIChatConfig(BaseConfig): base = get_oci_base_url(optional_params, api_base or litellm.api_base) return f"{base}/{OCI_API_VERSION}/actions/chat" - def _get_optional_params( - self, vendor: OCIVendors, optional_params: dict, model: str = "" - ) -> Dict: + def _get_optional_params(self, vendor: OCIVendors, optional_params: dict, model: str = "") -> Dict: param_map = ( - self.openai_to_oci_cohere_param_map - if vendor == OCIVendors.COHERE - else self.openai_to_oci_generic_param_map + self.openai_to_oci_cohere_param_map if vendor == OCIVendors.COHERE else self.openai_to_oci_generic_param_map ) selected_params: Dict = {} @@ -480,9 +466,7 @@ class OCIChatConfig(BaseConfig): # endpoint uses "maxTokens" regardless, so the override is GENERIC-only. max_tokens_key = ( "maxCompletionTokens" - if vendor != OCIVendors.COHERE - and model - and _model_uses_max_completion_tokens(model) + if vendor != OCIVendors.COHERE and model and _model_uses_max_completion_tokens(model) else "maxTokens" ) @@ -589,18 +573,14 @@ class OCIChatConfig(BaseConfig): system_messages = [m for m in messages if m.get("role") == "system"] preamble_override = None if system_messages: - preamble = "\n".join( - _extract_text_content(m["content"]) for m in system_messages - ) + preamble = "\n".join(_extract_text_content(m["content"]) for m in system_messages) if preamble: preamble_override = preamble chat_request = CohereChatRequest( apiFormat="COHERE", message=_extract_text_content(user_messages[-1]["content"]), - chatHistory=adapt_messages_to_cohere_standard( - [m for m in messages if m.get("role") != "system"] - ), + chatHistory=adapt_messages_to_cohere_standard([m for m in messages if m.get("role") != "system"]), preambleOverride=preamble_override, **self._get_optional_params(OCIVendors.COHERE, optional_params, model), ) @@ -652,13 +632,9 @@ class OCIChatConfig(BaseConfig): vendor = get_vendor_from_model(model) if vendor == OCIVendors.COHERE: - model_response = handle_cohere_response( - response_json, model, model_response, raw_response - ) + model_response = handle_cohere_response(response_json, model, model_response, raw_response) else: - model_response = handle_generic_response( - response_json, model, model_response, raw_response - ) + model_response = handle_generic_response(response_json, model, model_response, raw_response) model_response._hidden_params["additional_headers"] = raw_response.headers return model_response @@ -684,11 +660,7 @@ class OCIChatConfig(BaseConfig): response = client.post( api_base, headers=headers, - data=( - signed_json_body - if signed_json_body is not None - else json.dumps(data) - ), + data=(signed_json_body if signed_json_body is not None else json.dumps(data)), stream=True, logging_obj=logging_obj, timeout=STREAMING_TIMEOUT, @@ -727,11 +699,7 @@ class OCIChatConfig(BaseConfig): response = await client.post( api_base, headers=headers, - data=( - signed_json_body - if signed_json_body is not None - else json.dumps(data) - ), + data=(signed_json_body if signed_json_body is not None else json.dumps(data)), stream=True, logging_obj=logging_obj, timeout=STREAMING_TIMEOUT, diff --git a/litellm/llms/oci/common_utils.py b/litellm/llms/oci/common_utils.py index 29c88cbd50f..4ecbcbfb656 100644 --- a/litellm/llms/oci/common_utils.py +++ b/litellm/llms/oci/common_utils.py @@ -33,8 +33,7 @@ OCI_API_VERSION = "20231130" def _require_cryptography() -> None: if not _CRYPTOGRAPHY_AVAILABLE: raise ImportError( - "cryptography package is required for OCI authentication. " - "Please install it with: pip install cryptography" + "cryptography package is required for OCI authentication. Please install it with: pip install cryptography" ) @@ -65,9 +64,7 @@ class OCISignerProtocol(Protocol): See: https://docs.oracle.com/en-us/iaas/tools/python/latest/api/signing.html """ - def do_request_sign( - self, request: Any, *, enforce_content_headers: bool = False - ) -> None: + def do_request_sign(self, request: Any, *, enforce_content_headers: bool = False) -> None: pass @@ -105,9 +102,7 @@ def sha256_base64(data: bytes) -> str: return base64.b64encode(digest).decode() -def build_signature_string( - method: str, path: str, headers: dict, signed_headers: list -) -> str: +def build_signature_string(method: str, path: str, headers: dict, signed_headers: list) -> str: lines = [] for header in signed_headers: if header == "(request-target)": @@ -125,9 +120,7 @@ def load_private_key_from_str(key_str: str) -> Any: password=None, ) if not isinstance(key, rsa.RSAPrivateKey): # type: ignore[union-attr] - raise TypeError( - "The provided private key is not an RSA key, which is required for OCI signing." - ) + raise TypeError("The provided private key is not an RSA key, which is required for OCI signing.") return key @@ -170,19 +163,13 @@ def resolve_oci_credentials(optional_params: dict) -> dict: oci_key, oci_key_file, oci_compartment_id """ return { - "oci_region": optional_params.get("oci_region") - or os.environ.get(_OCI_REGION_ENV) - or "us-ashburn-1", + "oci_region": optional_params.get("oci_region") or os.environ.get(_OCI_REGION_ENV) or "us-ashburn-1", "oci_user": optional_params.get("oci_user") or os.environ.get(_OCI_USER_ENV), - "oci_fingerprint": optional_params.get("oci_fingerprint") - or os.environ.get(_OCI_FINGERPRINT_ENV), - "oci_tenancy": optional_params.get("oci_tenancy") - or os.environ.get(_OCI_TENANCY_ENV), + "oci_fingerprint": optional_params.get("oci_fingerprint") or os.environ.get(_OCI_FINGERPRINT_ENV), + "oci_tenancy": optional_params.get("oci_tenancy") or os.environ.get(_OCI_TENANCY_ENV), "oci_key": optional_params.get("oci_key") or os.environ.get(_OCI_KEY_ENV), - "oci_key_file": optional_params.get("oci_key_file") - or os.environ.get(_OCI_KEY_FILE_ENV), - "oci_compartment_id": optional_params.get("oci_compartment_id") - or os.environ.get(_OCI_COMPARTMENT_ID_ENV), + "oci_key_file": optional_params.get("oci_key_file") or os.environ.get(_OCI_KEY_FILE_ENV), + "oci_compartment_id": optional_params.get("oci_compartment_id") or os.environ.get(_OCI_COMPARTMENT_ID_ENV), } @@ -205,8 +192,7 @@ def get_oci_base_url(optional_params: dict, api_base: Optional[str] = None) -> s raise OCIError( status_code=400, message=( - f"Invalid OCI region {region!r}: must match " - "^[a-z][a-z0-9-]{0,30}[a-z0-9]$ (e.g. 'us-ashburn-1')." + f"Invalid OCI region {region!r}: must match ^[a-z][a-z0-9-]{{0,30}}[a-z0-9]$ (e.g. 'us-ashburn-1')." ), ) return f"https://inference.generativeai.{region}.oci.oraclecloud.com" @@ -235,9 +221,7 @@ def sign_with_oci_signer( prepared_headers.setdefault("content-type", "application/json") prepared_headers.setdefault("content-length", str(len(body))) - request_wrapper = OCIRequestWrapper( - method=method, url=api_base, headers=prepared_headers, body=body - ) + request_wrapper = OCIRequestWrapper(method=method, url=api_base, headers=prepared_headers, body=body) if oci_signer is None: raise ValueError("oci_signer cannot be None when calling sign_with_oci_signer") @@ -273,12 +257,7 @@ def sign_with_manual_credentials( oci_key = creds["oci_key"] oci_key_file = creds["oci_key_file"] - if ( - not oci_user - or not oci_fingerprint - or not oci_tenancy - or not (oci_key or oci_key_file) - ): + if not oci_user or not oci_fingerprint or not oci_tenancy or not (oci_key or oci_key_file): raise OCIError( status_code=401, message=( @@ -317,9 +296,7 @@ def sign_with_manual_credentials( "content-type", "x-content-sha256", ] - signing_string = build_signature_string( - method, path, headers_to_sign, signed_header_names - ) + signing_string = build_signature_string(method, path, headers_to_sign, signed_header_names) _require_cryptography() @@ -401,9 +378,7 @@ def sign_oci_request( """ if optional_params.get("oci_signer") is not None: return sign_with_oci_signer(headers, optional_params, request_data, api_base) - return sign_with_manual_credentials( - headers, optional_params, request_data, api_base - ) + return sign_with_manual_credentials(headers, optional_params, request_data, api_base) def validate_oci_environment( @@ -485,11 +460,7 @@ def resolve_oci_schema_anyof(obj: Any) -> Any: """ if isinstance(obj, dict): if "anyOf" in obj and "type" not in obj: - non_null = [ - t - for t in obj["anyOf"] - if not (isinstance(t, dict) and t.get("type") == "null") - ] + non_null = [t for t in obj["anyOf"] if not (isinstance(t, dict) and t.get("type") == "null")] if non_null: resolved = {**obj, **non_null[0]} resolved.pop("anyOf", None) @@ -535,18 +506,14 @@ def sanitize_oci_schema(schema: Any) -> Any: properties = sanitized.get("properties") if "required" in sanitized: if isinstance(required, list) and isinstance(properties, dict): - sanitized["required"] = [ - f for f in required if isinstance(f, str) and f in properties - ] + sanitized["required"] = [f for f in required if isinstance(f, str) and f in properties] elif not isinstance(required, list): sanitized["required"] = [] return sanitized -def enrich_cohere_param_description( - description: str, param_schema: Dict[str, Any] -) -> str: +def enrich_cohere_param_description(description: str, param_schema: Dict[str, Any]) -> str: """Embed schema constraints into a Cohere parameter description. ``CohereParameterDefinition`` only has ``type``, ``description``, and diff --git a/litellm/llms/oci/embed/transformation.py b/litellm/llms/oci/embed/transformation.py index 6cfa85b4bc4..44f5d941db4 100644 --- a/litellm/llms/oci/embed/transformation.py +++ b/litellm/llms/oci/embed/transformation.py @@ -226,9 +226,7 @@ class OCIEmbedConfig(BaseEmbeddingConfig): if serving_mode_type == "DEDICATED": endpoint_id = optional_params.get("oci_endpoint_id", model) - serving_mode = OCIServingMode( - servingType="DEDICATED", endpointId=endpoint_id - ) + serving_mode = OCIServingMode(servingType="DEDICATED", endpointId=endpoint_id) else: serving_mode = OCIServingMode(servingType="ON_DEMAND", modelId=model) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 5c0624fe5c9..3152ded2367 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -172,17 +172,9 @@ class OllamaChatConfig(BaseConfig): optional_params["repeat_penalty"] = value if param == "stop": optional_params["stop"] = value - if ( - param == "response_format" - and isinstance(value, dict) - and value.get("type") == "json_object" - ): + if param == "response_format" and isinstance(value, dict) and value.get("type") == "json_object": optional_params["format"] = "json" - if ( - param == "response_format" - and isinstance(value, dict) - and value.get("type") == "json_schema" - ): + if param == "response_format" and isinstance(value, dict) and value.get("type") == "json_schema": if value.get("json_schema") and value["json_schema"].get("schema"): optional_params["format"] = value["json_schema"]["schema"] if param == "reasoning_effort" and value is not None: @@ -281,9 +273,7 @@ class OllamaChatConfig(BaseConfig): ) ) new_tools.append(ollama_tool_call) - reasoning_content, parsed_content = _extract_reasoning_content( - cast(dict, m) - ) + reasoning_content, parsed_content = _extract_reasoning_content(cast(dict, m)) content_str = convert_content_list_to_str(cast(AllMessageValues, m)) images = extract_images_from_message(cast(AllMessageValues, m)) @@ -361,9 +351,7 @@ class OllamaChatConfig(BaseConfig): if response_json_message is not None: if "thinking" in response_json_message: # remap 'thinking' to 'reasoning_content' - response_json_message["reasoning_content"] = response_json_message[ - "thinking" - ] + response_json_message["reasoning_content"] = response_json_message["thinking"] del response_json_message["thinking"] elif response_json_message.get("content") is not None: # parse reasoning content from content @@ -371,16 +359,11 @@ class OllamaChatConfig(BaseConfig): _parse_content_for_reasoning, ) - reasoning_content, content = _parse_content_for_reasoning( - response_json_message["content"] - ) + reasoning_content, content = _parse_content_for_reasoning(response_json_message["content"]) response_json_message["reasoning_content"] = reasoning_content response_json_message["content"] = content - if ( - request_data.get("format", "") == "json" - and litellm_params.get("function_name") is not None - ): + if request_data.get("format", "") == "json" and litellm_params.get("function_name") is not None: function_call = json.loads(response_json_message["content"]) message = litellm.Message( content=None, @@ -388,12 +371,8 @@ class OllamaChatConfig(BaseConfig): { "id": f"call_{str(uuid.uuid4())}", "function": { - "name": function_call.get( - "name", litellm_params.get("function_name") - ), - "arguments": json.dumps( - function_call.get("arguments", function_call) - ), + "name": function_call.get("name", litellm_params.get("function_name")), + "arguments": json.dumps(function_call.get("arguments", function_call)), }, "type": "function", } @@ -411,9 +390,7 @@ class OllamaChatConfig(BaseConfig): model_response.choices[0].finish_reason = "tool_calls" model_response.created = int(time.time()) model_response.model = "ollama_chat/" + model - prompt_tokens = response_json.get( - "prompt_eval_count", litellm.token_counter(messages=messages) - ) # type: ignore + prompt_tokens = response_json.get("prompt_eval_count", litellm.token_counter(messages=messages)) # type: ignore completion_tokens = response_json.get( "eval_count", litellm.token_counter(text=response_json["message"]["content"]), @@ -429,12 +406,8 @@ class OllamaChatConfig(BaseConfig): ) return model_response - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return OllamaError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return OllamaError(status_code=status_code, message=error_message, headers=headers) def get_model_response_iterator( self, @@ -500,9 +473,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): for tool_call in tool_calls: function_args = tool_call.get("function").get("arguments") if function_args is not None and len(function_args) > 0: - is_function_call_complete = self._is_function_call_complete( - function_args - ) + is_function_call_complete = self._is_function_call_complete(function_args) if is_function_call_complete: tool_call["id"] = str(uuid.uuid4()) @@ -513,10 +484,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): reasoning_content = chunk["message"].get("thinking") self.started_reasoning_content = True if chunk["message"].get("content"): - if ( - self.started_reasoning_content - and not self.finished_reasoning_content - ): + if self.started_reasoning_content and not self.finished_reasoning_content: self.finished_reasoning_content = True message_content = chunk["message"].get("content") @@ -529,10 +497,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): message_content = message_content.replace("", "") self.finished_reasoning_content = True - if ( - self.started_reasoning_content - and not self.finished_reasoning_content - ): + if self.started_reasoning_content and not self.finished_reasoning_content: reasoning_content = message_content else: content = message_content @@ -565,8 +530,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): usage = ChatCompletionUsageBlock( prompt_tokens=chunk.get("prompt_eval_count", 0), completion_tokens=chunk.get("eval_count", 0), - total_tokens=chunk.get("prompt_eval_count", 0) - + chunk.get("eval_count", 0), + total_tokens=chunk.get("prompt_eval_count", 0) + chunk.get("eval_count", 0), ) return ModelResponseStream( diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 7d52ef14dd9..21ff3612a49 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -7,9 +7,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException class OllamaError(BaseLLMException): - def __init__( - self, status_code: int, message: str, headers: Union[dict, httpx.Headers] - ): + def __init__(self, status_code: int, message: str, headers: Union[dict, httpx.Headers]): super().__init__(status_code=status_code, message=message, headers=headers) @@ -27,9 +25,7 @@ def _convert_image(image): try: from PIL import Image except Exception: - raise Exception( - "ollama image conversion failed please run `pip install Pillow`" - ) + raise Exception("ollama image conversion failed please run `pip install Pillow`") orig = image if image.startswith("data:"): @@ -101,9 +97,7 @@ class OllamaModelInfo(BaseLLMModelInfo): passed_api_base = api_base base = self.get_server_api_base(api_base) - api_key = ( - self.get_api_key(api_key) if passed_api_base is None or api_key else None - ) + api_key = self.get_api_key(api_key) if passed_api_base is None or api_key else None headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} names: set[str] = set() @@ -113,11 +107,7 @@ class OllamaModelInfo(BaseLLMModelInfo): data = resp.json() # Expecting a dict with a 'models' list models_list = [] - if ( - isinstance(data, dict) - and "models" in data - and isinstance(data["models"], list) - ): + if isinstance(data, dict) and "models" in data and isinstance(data["models"], list): models_list = data["models"] elif isinstance(data, list): models_list = data @@ -137,9 +127,7 @@ class OllamaModelInfo(BaseLLMModelInfo): static = models_by_provider.get("ollama", []) or [] return [f"ollama/{m}" for m in static] except Exception as e1: - verbose_logger.warning( - f"Error retrieving static ollama models as fallback: {e1}" - ) + verbose_logger.warning(f"Error retrieving static ollama models as fallback: {e1}") return [] # assemble full model names result = sorted(names) @@ -190,9 +178,7 @@ class OllamaModelInfo(BaseLLMModelInfo): model = self._strip_ollama_model_prefix(model) passed_api_base = api_base api_base = self.get_server_api_base(api_base) - api_key = ( - self.get_api_key(api_key) if passed_api_base is None or api_key else None - ) + api_key = self.get_api_key(api_key) if passed_api_base is None or api_key else None headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} try: @@ -238,9 +224,7 @@ class OllamaModelInfo(BaseLLMModelInfo): ) -> Optional[dict[str, Any]]: if self._is_static_ollama_model(model): return None - return self.get_runtime_model_info( - model=model, api_base=api_base, api_key=api_key - ) + return self.get_runtime_model_info(model=model, api_base=api_base, api_key=api_key) def validate_environment( self, diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 97e4f13b560..7f229be53ae 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -46,15 +46,11 @@ def _process_ollama_embedding_response( if encoding is not None: input_tokens = len(encoding.encode("".join(prompts))) if logging_obj: - logging_obj.debug( - "Ollama response missing prompt_eval_count; estimated with encoding." - ) + logging_obj.debug("Ollama response missing prompt_eval_count; estimated with encoding.") else: input_tokens = 0 if logging_obj: - logging_obj.warning( - "Missing prompt_eval_count and no encoding provided; defaulted to 0." - ) + logging_obj.warning("Missing prompt_eval_count and no encoding provided; defaulted to 0.") model_response.object = "list" model_response.data = output_data diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index a8cdbff87d0..204b0d15c03 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -91,9 +91,7 @@ class OllamaConfig(BaseConfig): repeat_penalty: Optional[float] = None temperature: Optional[float] = None seed: Optional[int] = None - stop: Optional[list] = ( - None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 - ) + stop: Optional[list] = None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 tfs_z: Optional[float] = None num_predict: Optional[int] = None top_k: Optional[int] = None @@ -232,16 +230,10 @@ class OllamaConfig(BaseConfig): "name": "mistral" }' """ - return OllamaModelInfo().get_model_info( - model=model, api_base=api_base, api_key=api_key - ) + return OllamaModelInfo().get_model_info(model=model, api_base=api_base, api_key=api_key) - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return OllamaError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return OllamaError(status_code=status_code, message=error_message, headers=headers) def transform_response( self, @@ -292,9 +284,7 @@ class OllamaConfig(BaseConfig): "id": f"call_{str(uuid.uuid4())}", "function": { "name": function_call["name"], - "arguments": json.dumps( - function_call["arguments"] - ), + "arguments": json.dumps(function_call["arguments"]), }, "type": "function", } @@ -315,12 +305,8 @@ class OllamaConfig(BaseConfig): reasoning_content: Optional[str] = None content: Optional[str] = None if response_text is not None: - reasoning_content, content = _parse_content_for_reasoning( - response_text - ) - message = litellm.Message( - content=content, reasoning_content=reasoning_content - ) + reasoning_content, content = _parse_content_for_reasoning(response_text) + message = litellm.Message(content=content, reasoning_content=reasoning_content) model_response.choices[0].message = message # type: ignore model_response.choices[0].finish_reason = "stop" else: @@ -362,9 +348,7 @@ class OllamaConfig(BaseConfig): litellm_params: dict, headers: dict, ) -> dict: - custom_prompt_dict = ( - litellm_params.get("custom_prompt_dict") or litellm.custom_prompt_dict - ) + custom_prompt_dict = litellm_params.get("custom_prompt_dict") or litellm.custom_prompt_dict text_completion_request = litellm_params.get("text_completion") if model in custom_prompt_dict: @@ -402,9 +386,7 @@ class OllamaConfig(BaseConfig): if format is not None: data["format"] = format if images is not None: - data["images"] = [ - _convert_image(convert_to_ollama_image(image)) for image in images - ] + data["images"] = [_convert_image(convert_to_ollama_image(image)) for image in images] if think is not None: data["think"] = think @@ -461,21 +443,15 @@ class OllamaConfig(BaseConfig): class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): - def __init__( - self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False - ): + def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): super().__init__(streaming_response, sync_stream, json_mode) self.started_reasoning_content: bool = False self.finished_reasoning_content: bool = False - def _handle_string_chunk( - self, str_line: str - ) -> Union[GenericStreamingChunk, ModelResponseStream]: + def _handle_string_chunk(self, str_line: str) -> Union[GenericStreamingChunk, ModelResponseStream]: return self.chunk_parser(json.loads(str_line)) - def chunk_parser( - self, chunk: dict - ) -> Union[GenericStreamingChunk, ModelResponseStream]: + def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]: try: if "error" in chunk: raise Exception(f"Ollama Error - {chunk}") @@ -515,10 +491,7 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): text = text.replace("", "") self.finished_reasoning_content = True - if ( - self.started_reasoning_content - and not self.finished_reasoning_content - ): + if self.started_reasoning_content and not self.finished_reasoning_content: reasoning_content = text else: content = text @@ -527,9 +500,7 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): choices=[ StreamingChoices( index=0, - delta=Delta( - reasoning_content=reasoning_content, content=content - ), + delta=Delta(reasoning_content=reasoning_content, content=content), ) ], finish_reason=finish_reason, diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index 5eb68a03d4b..fe2bb9dc6d1 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -113,9 +113,7 @@ def embedding( # Logging before API call if logging_obj: - logging_obj.pre_call( - input=input, api_key=api_key, additional_args={"complete_input_dict": data} - ) + logging_obj.pre_call(input=input, api_key=api_key, additional_args={"complete_input_dict": data}) # Send POST request headers = oobabooga_config.validate_environment( @@ -126,9 +124,7 @@ def embedding( optional_params=optional_params, litellm_params={}, ) - response = litellm.module_level_client.post( - embeddings_url, headers=headers, json=data - ) + response = litellm.module_level_client.post(embeddings_url, headers=headers, json=data) completion_response = response.json() # Check for errors in response diff --git a/litellm/llms/oobabooga/chat/transformation.py b/litellm/llms/oobabooga/chat/transformation.py index 0c118efbc35..608fbc5cb35 100644 --- a/litellm/llms/oobabooga/chat/transformation.py +++ b/litellm/llms/oobabooga/chat/transformation.py @@ -25,9 +25,7 @@ class OobaboogaConfig(OpenAIGPTConfig): status_code: int, headers: Optional[Union[dict, httpx.Headers]] = None, ) -> BaseLLMException: - return OobaboogaError( - status_code=status_code, message=error_message, headers=headers - ) + return OobaboogaError(status_code=status_code, message=error_message, headers=headers) def transform_response( self, @@ -55,9 +53,7 @@ class OobaboogaConfig(OpenAIGPTConfig): try: completion_response = raw_response.json() except Exception: - raise OobaboogaError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OobaboogaError(message=raw_response.text, status_code=raw_response.status_code) if "error" in completion_response: raise OobaboogaError( message=completion_response["error"], @@ -65,9 +61,7 @@ class OobaboogaConfig(OpenAIGPTConfig): ) else: try: - model_response.choices[0].message.content = completion_response[ - "choices" - ][0]["message"]["content"] # type: ignore + model_response.choices[0].message.content = completion_response["choices"][0]["message"]["content"] # type: ignore except Exception as e: raise OobaboogaError( message=str(e), diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 9ccb2e1c267..f0a859deba0 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -129,9 +129,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): ) @classmethod - def _is_reasoning_effort_level_explicitly_disabled( - cls, model: str, level: str - ) -> bool: + def _is_reasoning_effort_level_explicitly_disabled(cls, model: str, level: str) -> bool: """Return True only when the model map explicitly sets the capability to False. Unlike ``_supports_reasoning_effort_level`` (which requires an explicit True), @@ -188,11 +186,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): if not self._supports_reasoning_effort_level(model, "none"): non_supported_params.extend(["logprobs", "top_p", "top_logprobs"]) - return [ - param - for param in base_gpt_series_params - if param not in non_supported_params - ] + return [param for param in base_gpt_series_params if param not in non_supported_params] def map_openai_params( self, @@ -203,9 +197,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): ) -> dict: if self.is_model_gpt_5_search_model(model): if "max_tokens" in non_default_params: - optional_params["max_completion_tokens"] = non_default_params.pop( - "max_tokens" - ) + optional_params["max_completion_tokens"] = non_default_params.pop("max_tokens") return super()._map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -217,17 +209,13 @@ class OpenAIGPT5Config(OpenAIGPTConfig): # Use effective_effort (extracted string) for xhigh validation, "none" checks, and # tool/sampling guards — dict inputs like {"effort": "none", "summary": "detailed"} # must be treated as effort="none" to avoid incorrect tool-drop or sampling errors. - raw_reasoning_effort = non_default_params.get( - "reasoning_effort" - ) or optional_params.get("reasoning_effort") + raw_reasoning_effort = non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort") effective_effort = _get_effort_level(raw_reasoning_effort) # Normalize dict reasoning_effort to string for Chat Completions API. # Example: {"effort": "high", "summary": "detailed"} -> "high" if isinstance(raw_reasoning_effort, dict) and "effort" in raw_reasoning_effort: - normalized = _normalize_reasoning_effort_for_chat_completion( - raw_reasoning_effort - ) + normalized = _normalize_reasoning_effort_for_chat_completion(raw_reasoning_effort) if normalized is not None: if "reasoning_effort" in non_default_params: non_default_params["reasoning_effort"] = normalized @@ -242,9 +230,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): optional_params.pop("reasoning_effort", None) else: raise litellm.utils.UnsupportedParamsError( - message=( - f"reasoning_effort={effective_effort} is not supported for this model." - ), + message=(f"reasoning_effort={effective_effort} is not supported for this model."), status_code=400, ) elif effective_effort in ("minimal", "low"): @@ -252,17 +238,13 @@ class OpenAIGPT5Config(OpenAIGPTConfig): # the model map explicitly sets supports_{level}_reasoning_effort=false. # Example: gpt-5.5-pro only accepts {medium, high, xhigh}, so it sets # supports_low_reasoning_effort=false (and supports_minimal=false). - if self._is_reasoning_effort_level_explicitly_disabled( - model, effective_effort - ): + if self._is_reasoning_effort_level_explicitly_disabled(model, effective_effort): if litellm.drop_params or drop_params: non_default_params.pop("reasoning_effort", None) optional_params.pop("reasoning_effort", None) else: raise litellm.utils.UnsupportedParamsError( - message=( - f"reasoning_effort={effective_effort} is not supported for this model." - ), + message=(f"reasoning_effort={effective_effort} is not supported for this model."), status_code=400, ) @@ -271,9 +253,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): # Relevant issue: https://github.com/BerriAI/litellm/issues/13381 ################################################################ if "max_tokens" in non_default_params: - optional_params["max_completion_tokens"] = non_default_params.pop( - "max_tokens" - ) + optional_params["max_completion_tokens"] = non_default_params.pop("max_tokens") # gpt-5.1/5.2 support logprobs, top_p, top_logprobs only when reasoning_effort="none" supports_none = self._supports_reasoning_effort_level(model, "none") @@ -298,9 +278,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): temperature_value: Optional[float] = non_default_params.pop("temperature") if temperature_value is not None: # models supporting reasoning_effort="none" also support flexible temperature - if supports_none and ( - effective_effort == "none" or effective_effort is None - ): + if supports_none and (effective_effort == "none" or effective_effort is None): optional_params["temperature"] = temperature_value elif temperature_value == 1: optional_params["temperature"] = temperature_value diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index e4d55404743..396ad5b105e 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -172,15 +172,11 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ] # works across all models model_specific_params = [] - if ( - model != "gpt-3.5-turbo-16k" and model != "gpt-4" - ): # gpt-4 does not support 'response_format' + if model != "gpt-3.5-turbo-16k" and model != "gpt-4": # gpt-4 does not support 'response_format' model_specific_params.append("response_format") # Normalize model name for responses API (e.g., "responses/gpt-4.1" -> "gpt-4.1") - model_for_check = ( - model.split("responses/", 1)[1] if "responses/" in model else model - ) + model_for_check = model.split("responses/", 1)[1] if "responses/" in model else model if ( model_for_check in litellm.open_ai_chat_completion_models ) or model_for_check in litellm.open_ai_text_completion_models: @@ -230,15 +226,11 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): def contains_pdf_url(self, content_item: ChatCompletionFileObjectFile) -> bool: potential_pdf_url_starts = ["https://", "http://", "www."] file_id = content_item.get("file_id") - if file_id and any( - file_id.startswith(start) for start in potential_pdf_url_starts - ): + if file_id and any(file_id.startswith(start) for start in potential_pdf_url_starts): return True return False - def _handle_pdf_url( - self, content_item: ChatCompletionFileObjectFile - ) -> ChatCompletionFileObjectFile: + def _handle_pdf_url(self, content_item: ChatCompletionFileObjectFile) -> ChatCompletionFileObjectFile: content_copy = content_item.copy() file_id = content_copy.get("file_id") if file_id is not None: @@ -248,9 +240,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): content_copy.pop("file_id") return content_copy - async def _async_handle_pdf_url( - self, content_item: ChatCompletionFileObjectFile - ) -> ChatCompletionFileObjectFile: + async def _async_handle_pdf_url(self, content_item: ChatCompletionFileObjectFile) -> ChatCompletionFileObjectFile: file_id = content_item.get("file_id") if file_id is not None: # check for file id being url done in _handle_pdf_url base64_data = await async_convert_url_to_base64(file_id) @@ -259,9 +249,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): content_item.pop("file_id") return content_item - def _common_file_data_check( - self, content_item: ChatCompletionFileObjectFile - ) -> ChatCompletionFileObjectFile: + def _common_file_data_check(self, content_item: ChatCompletionFileObjectFile) -> ChatCompletionFileObjectFile: file_data = content_item.get("file_data") filename = content_item.get("filename") if file_data is not None and filename is None: @@ -282,9 +270,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): elif isinstance(content_item["image_url"], dict): new_image_url_obj = ChatCompletionImageUrlObject( **{ # type: ignore - k: v - for k, v in content_item["image_url"].items() - if k not in litellm_specific_params + k: v for k, v in content_item["image_url"].items() if k not in litellm_specific_params } ) content_item["image_url"] = new_image_url_obj @@ -299,9 +285,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) new_file_obj = ChatCompletionFileObjectFile( **{ # type: ignore - k: v - for k, v in file_obj.items() - if k not in litellm_specific_params + k: v for k, v in file_obj.items() if k not in litellm_specific_params } ) content_item["file"] = new_file_obj @@ -370,18 +354,10 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): message_content = message.get("content") message_role = message.get("role") - if ( - message_role == "user" - and message_content - and isinstance(message_content, list) - ): - message_content_types = cast( - List[OpenAIMessageContentListBlock], message_content - ) + if message_role == "user" and message_content and isinstance(message_content, list): + message_content_types = cast(List[OpenAIMessageContentListBlock], message_content) for i, content_item in enumerate(message_content_types): - message_content_types[ - i - ] = await self._async_transform_content_item( + message_content_types[i] = await self._async_transform_content_item( cast(OpenAIMessageContentListBlock, content_item), ) return messages @@ -392,14 +368,8 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): for message in messages: message_content = message.get("content") message_role = message.get("role") - if ( - message_role == "user" - and message_content - and isinstance(message_content, list) - ): - message_content_types = cast( - List[OpenAIMessageContentListBlock], message_content - ) + if message_role == "user" and message_content and isinstance(message_content, list): + message_content_types = cast(List[OpenAIMessageContentListBlock], message_content) for i, content_item in enumerate(message_content): message_content_types[i] = self._transform_content_item( cast(OpenAIMessageContentListBlock, content_item) @@ -443,12 +413,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ if custom_llm_provider != "openai": return False - resolved_api_base = ( - api_base - or litellm.api_base - or os.getenv("OPENAI_BASE_URL") - or os.getenv("OPENAI_API_BASE") - ) + resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") if not resolved_api_base: return False hostname = urlparse(resolved_api_base).hostname @@ -496,9 +461,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): litellm_params: dict, headers: dict, ) -> dict: - transformed_messages = await self._transform_messages( - messages=messages, model=model, is_async=True - ) + transformed_messages = await self._transform_messages(messages=messages, model=model, is_async=True) if not self._should_preserve_cache_control_for_endpoint( litellm_params.get("custom_llm_provider"), litellm_params.get("api_base") ): @@ -520,9 +483,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): } else: ## allow for any object specific behaviour to be handled - return self.transform_request( - model, messages, optional_params, litellm_params, headers - ) + return self.transform_request(model, messages, optional_params, litellm_params, headers) def _passed_in_tools(self, optional_params: dict) -> bool: return optional_params.get("tools", None) is not None @@ -540,10 +501,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): tool_call_names = get_tool_call_names(optional_params.get("tools", [])) try: json_content = json.loads(content) - if ( - json_content.get("type") == "function" - and json_content.get("name") in tool_call_names - ): + if json_content.get("type") == "function" and json_content.get("name") in tool_call_names: return ChatCompletionMessageToolCall( function=Function( name=json_content.get("name"), @@ -579,20 +537,12 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): for _tc in tool_calls: _openai_tc = ChatCompletionMessageToolCall(**_tc) # type: ignore _openai_tool_calls.append(_openai_tc) - fixed_tool_calls = _handle_invalid_parallel_tool_calls( - _openai_tool_calls - ) + fixed_tool_calls = _handle_invalid_parallel_tool_calls(_openai_tool_calls) if fixed_tool_calls is not None: new_tool_calls = fixed_tool_calls - elif ( - optional_params is not None - and message_content - and isinstance(message_content, str) - ): - new_tool_call = self._check_and_fix_if_content_is_tool_call( - message_content, optional_params - ) + elif optional_params is not None and message_content and isinstance(message_content, str): + new_tool_call = self._check_and_fix_if_content_is_tool_call(message_content, optional_params) if new_tool_call is not None: choice["message"]["content"] = None # remove the content new_tool_calls = [new_tool_call] @@ -604,9 +554,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): convert_tool_call_to_json_mode=json_mode, ): # to support response_format on claude models - json_mode_content_str: Optional[str] = ( - str(new_tool_calls[0]["function"].get("arguments", "")) or None - ) + json_mode_content_str: Optional[str] = str(new_tool_calls[0]["function"].get("arguments", "")) or None if json_mode_content_str is not None: translated_message = Message(content=json_mode_content_str) finish_reason = "stop" @@ -679,9 +627,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise OpenAIError( - message="Unable to get json response - {}, Original Response: {}".format( - str(e), raw_response.text - ), + message="Unable to get json response - {}, Original Response: {}".format(str(e), raw_response.text), status_code=raw_response.status_code, headers=response_headers, ) @@ -751,9 +697,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): return headers - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: """ Calls OpenAI's `/v1/models` endpoint and returns the list of models. """ @@ -782,12 +726,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: - return ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") @staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a62dd939a71..c7a49a2e47f 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -107,13 +107,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): structured_messages = self.get_structured_messages(data) if structured_messages: if skip_system: - structured_messages = openai_messages_without_system( - structured_messages - ) + structured_messages = openai_messages_without_system(structured_messages) if skip_tool: - structured_messages = openai_messages_without_tool( - structured_messages - ) + structured_messages = openai_messages_without_tool(structured_messages) inputs["structured_messages"] = structured_messages # Pass tools (function definitions) to the guardrail tools = data.get("tools") @@ -138,9 +134,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if guardrailed_tools is not None: data["tools"] = guardrailed_tools - guardrailed_structured_messages = guardrailed_inputs.get( - "structured_messages" - ) + guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") if ( guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages @@ -268,9 +262,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content - messages[msg_idx]["content"][content_idx_optional]["text"] = ( - guardrail_response - ) + messages[msg_idx]["content"][content_idx_optional]["text"] = guardrail_response async def _apply_guardrail_responses_to_input_tool_calls( self, @@ -290,9 +282,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if task_idx < len(tool_calls): guardrailed_tool_call = tool_calls[task_idx] message_tool_calls = messages[msg_idx].get("tool_calls", None) - if message_tool_calls is not None and isinstance( - message_tool_calls, list - ): + if message_tool_calls is not None and isinstance(message_tool_calls, list): if tool_call_idx < len(message_tool_calls): # Replace the tool call with the guardrailed version message_tool_calls[tool_call_idx] = guardrailed_tool_call @@ -324,9 +314,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Step 0: Check if response has any text content to process if not self._has_text_content(response): - verbose_proxy_logger.warning( - "OpenAI Chat Completions: No text content in response, skipping guardrail" - ) + verbose_proxy_logger.warning("OpenAI Chat Completions: No text content in response, skipping guardrail") return response texts_to_check: List[str] = [] @@ -362,9 +350,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Add user API key metadata with prefixed keys if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata @@ -388,8 +374,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): returned_tool_calls = guardrailed_inputs.get("tool_calls") guardrailed_tool_calls: List[Dict[str, Any]] = ( cast(List[Dict[str, Any]], returned_tool_calls) - if isinstance(returned_tool_calls, list) - and len(returned_tool_calls) == len(tool_calls_to_check) + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) else tool_calls_to_check ) @@ -409,9 +394,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=tool_call_task_mappings, ) - verbose_proxy_logger.debug( - "OpenAI Chat Completions: Processed output response: %s", response - ) + verbose_proxy_logger.debug("OpenAI Chat Completions: Processed output response: %s", response) return response @@ -450,9 +433,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # convert to model response model_response = cast( ModelResponse, - stream_chunk_builder( - chunks=responses_so_far, logging_obj=litellm_logging_obj - ), + stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), ) # run process_output_response await self.process_output_response( @@ -505,9 +486,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Add user API key metadata with prefixed keys if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata @@ -515,11 +494,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if images_to_check: inputs["images"] = images_to_check # Include model information from the first response if available - if ( - responses_so_far - and hasattr(responses_so_far[0], "model") - and responses_so_far[0].model - ): + if responses_so_far and hasattr(responses_so_far[0], "model") and responses_so_far[0].model: inputs["model"] = responses_so_far[0].model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -595,9 +570,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return combined_texts - def _has_text_content( - self, response: Union["ModelResponse", "ModelResponseStream"] - ) -> bool: + def _has_text_content(self, response: Union["ModelResponse", "ModelResponseStream"]) -> bool: """ Check if response has any text content or tool calls to process. @@ -609,28 +582,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for choice in response.choices: if isinstance(choice, litellm.Choices): # Check for text content - if choice.message.content and isinstance( - choice.message.content, str - ): + if choice.message.content and isinstance(choice.message.content, str): return True # Check for tool calls - if choice.message.tool_calls and isinstance( - choice.message.tool_calls, list - ): + if choice.message.tool_calls and isinstance(choice.message.tool_calls, list): if len(choice.message.tool_calls) > 0: return True elif isinstance(response, ModelResponseStream): for streaming_choice in response.choices: if isinstance(streaming_choice, litellm.StreamingChoices): # Check for text content - if streaming_choice.delta.content and isinstance( - streaming_choice.delta.content, str - ): + if streaming_choice.delta.content and isinstance(streaming_choice.delta.content, str): return True # Check for tool calls - if streaming_choice.delta.tool_calls and isinstance( - streaming_choice.delta.tool_calls, list - ): + if streaming_choice.delta.tool_calls and isinstance(streaming_choice.delta.tool_calls, list): if len(streaming_choice.delta.tool_calls) > 0: return True return False @@ -650,9 +615,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize text/image/tool call extraction logic. """ - verbose_proxy_logger.debug( - "OpenAI Chat Completions: Processing choice: %s", choice - ) + verbose_proxy_logger.debug("OpenAI Chat Completions: Processing choice: %s", choice) # Determine content source and tool calls based on choice type content = None @@ -699,9 +662,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_calls_to_check.append(tool_call_dict) tool_call_task_mappings.append((choice_idx, int(tool_call_idx))) - def _convert_tool_call_to_dict( - self, tool_call: Union[Dict[str, Any], Any] - ) -> Optional[Dict[str, Any]]: + def _convert_tool_call_to_dict(self, tool_call: Union[Dict[str, Any], Any]) -> Optional[Dict[str, Any]]: """ Convert a tool call object to dictionary format. @@ -755,9 +716,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content - choice.message.content[content_idx_optional]["text"] = ( - guardrail_response # type: ignore - ) + choice.message.content[content_idx_optional]["text"] = guardrail_response # type: ignore async def _apply_guardrail_responses_to_output_tool_calls( self, @@ -780,9 +739,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): choice = cast(Choices, response.choices[choice_idx]) choice_tool_calls = choice.message.tool_calls - if choice_tool_calls is not None and isinstance( - choice_tool_calls, list - ): + if choice_tool_calls is not None and isinstance(choice_tool_calls, list): if tool_call_idx < len(choice_tool_calls): # Update the tool call with guardrailed version existing_tool_call = choice_tool_calls[tool_call_idx] @@ -790,9 +747,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "function" in guardrailed_tool_call: func_dict = guardrailed_tool_call["function"] if "arguments" in func_dict: - existing_tool_call.function.arguments = func_dict[ - "arguments" - ] + existing_tool_call.function.arguments = func_dict["arguments"] if "name" in func_dict: existing_tool_call.function.name = func_dict["name"] diff --git a/litellm/llms/openai/chat/o_series_transformation.py b/litellm/llms/openai/chat/o_series_transformation.py index 8db7ecf7b3a..78a5b3512a4 100644 --- a/litellm/llms/openai/chat/o_series_transformation.py +++ b/litellm/llms/openai/chat/o_series_transformation.py @@ -36,9 +36,7 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): def get_config(cls): return super().get_config() - def translate_developer_role_to_system_role( - self, messages: List[AllMessageValues] - ) -> List[AllMessageValues]: + def translate_developer_role_to_system_role(self, messages: List[AllMessageValues]) -> List[AllMessageValues]: """ O-series models support `developer` role. """ @@ -64,22 +62,16 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): all_openai_params.extend(o_series_only_param) try: - model, custom_llm_provider, api_base, api_key = get_llm_provider( - model=model - ) + model, custom_llm_provider, api_base, api_key = get_llm_provider(model=model) except Exception: verbose_logger.debug( f"Unable to infer model provider for model={model}, defaulting to openai for o1 supported param check" ) custom_llm_provider = "openai" - _supports_function_calling = supports_function_calling( - model, custom_llm_provider - ) + _supports_function_calling = supports_function_calling(model, custom_llm_provider) _supports_response_schema = supports_response_schema(model, custom_llm_provider) - _supports_parallel_tool_calls = supports_parallel_function_calling( - model, custom_llm_provider - ) + _supports_parallel_tool_calls = supports_parallel_function_calling(model, custom_llm_provider) if not _supports_function_calling: non_supported_params.append("tools") @@ -93,9 +85,7 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): if not _supports_response_schema: non_supported_params.append("response_format") - return [ - param for param in all_openai_params if param not in non_supported_params - ] + return [param for param in all_openai_params if param not in non_supported_params] def map_openai_params( self, @@ -105,9 +95,7 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): drop_params: bool, ): if "max_tokens" in non_default_params: - optional_params["max_completion_tokens"] = non_default_params.pop( - "max_tokens" - ) + optional_params["max_completion_tokens"] = non_default_params.pop("max_tokens") if "temperature" in non_default_params: temperature_value: Optional[float] = non_default_params.pop("temperature") if temperature_value is not None: @@ -125,9 +113,7 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): status_code=400, ) - return super()._map_openai_params( - non_default_params, optional_params, model, drop_params - ) + return super()._map_openai_params(non_default_params, optional_params, model, drop_params) def is_model_o_series_model(self, model: str) -> bool: model = model.split("/")[-1] # could be "openai/o3" or "o3" @@ -162,16 +148,10 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): _supports_system_messages = supports_system_messages(model, "openai") for i, message in enumerate(messages): if message["role"] == "system" and not _supports_system_messages: - new_message = ChatCompletionUserMessage( - content=message["content"], role="user" - ) + new_message = ChatCompletionUserMessage(content=message["content"], role="user") messages[i] = new_message # Replace the old message with the new one if is_async: - return super()._transform_messages( - messages, model, is_async=cast(Literal[True], True) - ) + return super()._transform_messages(messages, model, is_async=cast(Literal[True], True)) else: - return super()._transform_messages( - messages, model, is_async=cast(Literal[False], False) - ) + return super()._transform_messages(messages, model, is_async=cast(Literal[False], False)) diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 381f215a13f..6731d4a6a4a 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -64,9 +64,7 @@ class OpenAIError(BaseLLMException): if response: self.response = response else: - self.response = httpx.Response( - status_code=status_code, request=self.request - ) + self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, message=self.message, @@ -100,11 +98,7 @@ def drop_params_from_unprocessable_entity_error( error_body = error_message else: error_body = e.body - if ( - error_body is not None - and isinstance(error_body, dict) - and error_body.get("message") - ): + if error_body is not None and isinstance(error_body, dict) and error_body.get("message"): message = error_body.get("message", {}) if isinstance(message, str): try: @@ -162,15 +156,11 @@ class BaseOpenAILLM: ) @staticmethod - def get_openai_client_cache_key( - client_initialization_params: dict, client_type: Literal["openai", "azure"] - ) -> str: + def get_openai_client_cache_key(client_initialization_params: dict, client_type: Literal["openai", "azure"]) -> str: """Creates a cache key for the OpenAI client based on the client initialization parameters""" hashed_api_key = None if client_initialization_params.get("api_key") is not None: - hash_object = hashlib.sha256( - client_initialization_params.get("api_key", "").encode() - ) + hash_object = hashlib.sha256(client_initialization_params.get("api_key", "").encode()) # Hexadecimal representation of the hash hashed_api_key = hash_object.hexdigest() @@ -187,9 +177,7 @@ class BaseOpenAILLM: "api_base", ) openai_client_fields = ( - BaseOpenAILLM.get_openai_client_initialization_param_fields( - client_type=client_type - ) + BaseOpenAILLM.get_openai_client_initialization_param_fields(client_type=client_type) + LITELLM_CLIENT_SPECIFIC_PARAMS ) @@ -227,9 +215,7 @@ class BaseOpenAILLM: return httpx.AsyncClient( verify=ssl_config, transport=AsyncHTTPHandler._create_async_transport( - ssl_context=( - ssl_config if isinstance(ssl_config, ssl.SSLContext) else None - ), + ssl_context=(ssl_config if isinstance(ssl_config, ssl.SSLContext) else None), ssl_verify=ssl_config if isinstance(ssl_config, bool) else None, shared_session=shared_session, ), @@ -274,15 +260,8 @@ def get_openai_credentials( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - resolved_organization = ( - organization - or litellm.organization - or os.getenv("OPENAI_ORGANIZATION", None) - or None - ) - resolved_api_key = ( - api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY") - ) + resolved_organization = organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None + resolved_api_key = api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY") return OpenAICredentials( api_base=resolved_api_base, api_key=resolved_api_key, diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index 593ab0ed2e5..8537fefe1e2 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -47,9 +47,7 @@ class OpenAITextCompletionHandler(BaseTranslation): """ prompt = data.get("prompt") if prompt is None: - verbose_proxy_logger.debug( - "OpenAI Text Completion: No prompt found in request data" - ) + verbose_proxy_logger.debug("OpenAI Text Completion: No prompt found in request data") return data if isinstance(prompt, str): @@ -69,8 +67,7 @@ class OpenAITextCompletionHandler(BaseTranslation): data["prompt"] = guardrailed_texts[0] if guardrailed_texts else prompt verbose_proxy_logger.debug( - "OpenAI Text Completion: Applied guardrail to string prompt. " - "Original length: %d, New length: %d", + "OpenAI Text Completion: Applied guardrail to string prompt. Original length: %d, New length: %d", len(prompt), len(data["prompt"]), ) @@ -140,9 +137,7 @@ class OpenAITextCompletionHandler(BaseTranslation): Modified response with guardrails applied to completion text """ if not hasattr(response, "choices") or not response.choices: - verbose_proxy_logger.debug( - "OpenAI Text Completion: No choices in response to process" - ) + verbose_proxy_logger.debug("OpenAI Text Completion: No choices in response to process") return response # Collect all texts to check @@ -166,9 +161,7 @@ class OpenAITextCompletionHandler(BaseTranslation): # Add user API key metadata with prefixed keys if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index f08ef844bb3..376d2636ba7 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -136,9 +136,7 @@ class OpenAITextCompletion(BaseLLM): else: openai_client = client - raw_response = openai_client.completions.with_raw_response.create( - **data - ) # type: ignore + raw_response = openai_client.completions.with_raw_response.create(**data) # type: ignore response = raw_response.parse() response_json = response.model_dump() @@ -161,9 +159,7 @@ class OpenAITextCompletion(BaseLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) async def acompletion( self, @@ -192,9 +188,7 @@ class OpenAITextCompletion(BaseLLM): else: openai_aclient = client - raw_response = await openai_aclient.completions.with_raw_response.create( - **data - ) + raw_response = await openai_aclient.completions.with_raw_response.create(**data) response = raw_response.parse() response_json = response.model_dump() @@ -218,9 +212,7 @@ class OpenAITextCompletion(BaseLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) def streaming( self, @@ -258,9 +250,7 @@ class OpenAITextCompletion(BaseLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) streamwrapper = CustomStreamWrapper( completion_stream=response, model=model, @@ -279,9 +269,7 @@ class OpenAITextCompletion(BaseLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) async def async_streaming( self, @@ -329,6 +317,4 @@ class OpenAITextCompletion(BaseLLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) diff --git a/litellm/llms/openai/completion/utils.py b/litellm/llms/openai/completion/utils.py index 8b3efb4cda8..a7b7e7a67ce 100644 --- a/litellm/llms/openai/completion/utils.py +++ b/litellm/llms/openai/completion/utils.py @@ -16,8 +16,7 @@ def is_tokens_or_list_of_tokens(value: List): return True # Check if it's a list of lists of integers (list of tokens) if isinstance(value, list) and all( - isinstance(item, list) and all(isinstance(i, int) for i in item) - for item in value + isinstance(item, list) and all(isinstance(i, int) for i in item) for item in value ): return True return False @@ -28,11 +27,7 @@ def _transform_prompt( ) -> AllPromptValues: if len(messages) == 1: # base case message_content = messages[0].get("content") - if ( - message_content - and isinstance(message_content, list) - and is_tokens_or_list_of_tokens(message_content) - ): + if message_content and isinstance(message_content, list) and is_tokens_or_list_of_tokens(message_content): openai_prompt: AllPromptValues = cast(AllPromptValues, message_content) else: openai_prompt = "" diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 7f874ffd3b1..b5f4334af0a 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -60,12 +60,7 @@ class OpenAIContainerConfig(BaseContainerConfig): headers: dict, api_key: Optional[str] = None, ) -> dict: - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.update( { "Authorization": f"Bearer {api_key}", @@ -99,9 +94,7 @@ class OpenAIContainerConfig(BaseContainerConfig): """Transform the container creation request for OpenAI API.""" # Remove extra_headers from optional params as they're handled separately container_create_optional_request_params = { - k: v - for k, v in container_create_optional_request_params.items() - if k not in ["extra_headers"] + k: v for k, v in container_create_optional_request_params.items() if k not in ["extra_headers"] } # Create the request data @@ -131,16 +124,11 @@ class OpenAIContainerConfig(BaseContainerConfig): provider="openai", ) - if ( - not hasattr(container_obj, "_hidden_params") - or container_obj._hidden_params is None - ): + if not hasattr(container_obj, "_hidden_params") or container_obj._hidden_params is None: container_obj._hidden_params = {} if "additional_headers" not in container_obj._hidden_params: container_obj._hidden_params["additional_headers"] = {} - container_obj._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = container_cost + container_obj._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = container_cost return container_obj @@ -199,9 +187,7 @@ class OpenAIContainerConfig(BaseContainerConfig): ) -> Tuple[str, Dict]: """Transform the OpenAI container retrieve request.""" # For container retrieve, we just need to construct the URL - encoded_container_id = encode_url_path_segment( - container_id, field_name="container_id" - ) + encoded_container_id = encode_url_path_segment(container_id, field_name="container_id") url = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request @@ -234,9 +220,7 @@ class OpenAIContainerConfig(BaseContainerConfig): - DELETE /v1/containers/{container_id} """ # Construct the URL for container delete - encoded_container_id = encode_url_path_segment( - container_id, field_name="container_id" - ) + encoded_container_id = encode_url_path_segment(container_id, field_name="container_id") url = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request @@ -274,9 +258,7 @@ class OpenAIContainerConfig(BaseContainerConfig): - GET /v1/containers/{container_id}/files """ # Construct the URL for container files - encoded_container_id = encode_url_path_segment( - container_id, field_name="container_id" - ) + encoded_container_id = encode_url_path_segment(container_id, field_name="container_id") url = join_container_api_base_path(api_base, f"/{encoded_container_id}/files") # Prepare query parameters @@ -321,13 +303,9 @@ class OpenAIContainerConfig(BaseContainerConfig): - GET /v1/containers/{container_id}/files/{file_id}/content """ # Construct the URL for container file content - encoded_container_id = encode_url_path_segment( - container_id, field_name="container_id" - ) + encoded_container_id = encode_url_path_segment(container_id, field_name="container_id") encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") - url = join_container_api_base_path( - api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content" - ) + url = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content") # No query parameters needed params: Dict[str, Any] = {} diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index 6935cafd0d9..25376419b9e 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/litellm/llms/openai/cost_calculation.py @@ -90,9 +90,7 @@ def cost_per_token( # return prompt_cost, completion_cost -def cost_per_second( - model: str, custom_llm_provider: Optional[str], duration: float = 0.0 -) -> Tuple[float, float]: +def cost_per_second(model: str, custom_llm_provider: Optional[str], duration: float = 0.0) -> Tuple[float, float]: """ Calculates the cost per second for a given model, prompt tokens, and completion tokens. @@ -106,25 +104,17 @@ def cost_per_second( """ ## GET MODEL INFO - model_info = get_model_info( - model=model, custom_llm_provider=custom_llm_provider or "openai" - ) + model_info = get_model_info(model=model, custom_llm_provider=custom_llm_provider or "openai") prompt_cost = 0.0 completion_cost = 0.0 ## Speech / Audio cost calculation - if ( - "output_cost_per_second" in model_info - and model_info["output_cost_per_second"] is not None - ): + if "output_cost_per_second" in model_info and model_info["output_cost_per_second"] is not None: verbose_logger.debug( f"For model={model} - output_cost_per_second: {model_info.get('output_cost_per_second')}; duration: {duration}" ) ## COST PER SECOND ## completion_cost = model_info["output_cost_per_second"] * duration - elif ( - "input_cost_per_second" in model_info - and model_info["input_cost_per_second"] is not None - ): + elif "input_cost_per_second" in model_info and model_info["input_cost_per_second"] is not None: verbose_logger.debug( f"For model={model} - input_cost_per_second: {model_info.get('input_cost_per_second')}; duration: {duration}" ) @@ -202,9 +192,7 @@ def video_generation_cost( """ ## GET MODEL INFO if model_info is None: - model_info = get_model_info( - model=model, custom_llm_provider=custom_llm_provider or "openai" - ) + model_info = get_model_info(model=model, custom_llm_provider=custom_llm_provider or "openai") # Check for video-specific cost per second video_cost_per_second = model_info.get("output_cost_per_video_per_second") diff --git a/litellm/llms/openai/data_residency.py b/litellm/llms/openai/data_residency.py index 7162f70ca5f..db3c49d7583 100644 --- a/litellm/llms/openai/data_residency.py +++ b/litellm/llms/openai/data_residency.py @@ -20,9 +20,7 @@ _OPENAI_REGIONAL_HOSTS: Dict[str, str] = { } -def infer_openai_data_residency( - custom_llm_provider: Optional[str], api_base: Optional[str] -) -> Optional[str]: +def infer_openai_data_residency(custom_llm_provider: Optional[str], api_base: Optional[str]) -> Optional[str]: """ Derive the OpenAI data-residency region from an api_base URL. diff --git a/litellm/llms/openai/embeddings/guardrail_translation/handler.py b/litellm/llms/openai/embeddings/guardrail_translation/handler.py index ff5021b8ce0..d208c98b0e4 100644 --- a/litellm/llms/openai/embeddings/guardrail_translation/handler.py +++ b/litellm/llms/openai/embeddings/guardrail_translation/handler.py @@ -50,19 +50,13 @@ class OpenAIEmbeddingsHandler(BaseTranslation): """ input_data = data.get("input") if input_data is None: - verbose_proxy_logger.debug( - "OpenAI Embeddings: No input found in request data" - ) + verbose_proxy_logger.debug("OpenAI Embeddings: No input found in request data") return data if isinstance(input_data, str): - data = await self._process_string_input( - data, input_data, guardrail_to_apply, litellm_logging_obj - ) + data = await self._process_string_input(data, input_data, guardrail_to_apply, litellm_logging_obj) elif isinstance(input_data, list): - data = await self._process_list_input( - data, input_data, guardrail_to_apply, litellm_logging_obj - ) + data = await self._process_list_input(data, input_data, guardrail_to_apply, litellm_logging_obj) else: verbose_proxy_logger.warning( "OpenAI Embeddings: Unexpected input type: %s. Expected string or list.", @@ -93,8 +87,7 @@ class OpenAIEmbeddingsHandler(BaseTranslation): if guardrailed_texts := guardrailed_inputs.get("texts"): data["input"] = guardrailed_texts[0] verbose_proxy_logger.debug( - "OpenAI Embeddings: Applied guardrail to string input. " - "Original length: %d, New length: %d", + "OpenAI Embeddings: Applied guardrail to string input. Original length: %d, New length: %d", len(input_data), len(data["input"]), ) @@ -116,9 +109,7 @@ class OpenAIEmbeddingsHandler(BaseTranslation): # Skip non-text inputs (token IDs) if isinstance(first_item, (int, list)): - verbose_proxy_logger.debug( - "OpenAI Embeddings: Input is token IDs, skipping guardrail processing" - ) + verbose_proxy_logger.debug("OpenAI Embeddings: Input is token IDs, skipping guardrail processing") return data if not isinstance(first_item, str): @@ -174,7 +165,6 @@ class OpenAIEmbeddingsHandler(BaseTranslation): Unmodified response (embeddings don't have text output to guard) """ verbose_proxy_logger.debug( - "OpenAI Embeddings: Output response processing skipped - " - "embeddings contain vectors, not text" + "OpenAI Embeddings: Output response processing skipped - embeddings contain vectors, not text" ) return response diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index 66537e56a6f..8a55fec58a6 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -38,9 +38,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.OPENAI - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """Add OpenAI-specific headers""" import litellm from litellm.secret_managers.main import get_secret_str @@ -50,12 +48,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): if litellm_params: api_key = litellm_params.api_key - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") if not api_key: raise ValueError("OPENAI_API_KEY is required for Evals API") @@ -158,9 +151,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform get eval request for OpenAI""" - url = self.get_complete_url( - api_base=api_base, endpoint="evals", eval_id=eval_id - ) + url = self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id) verbose_logger.debug("Get eval request - URL: %s", url) @@ -186,16 +177,12 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict, Dict]: """Transform update eval request for OpenAI""" - url = self.get_complete_url( - api_base=api_base, endpoint="evals", eval_id=eval_id - ) + url = self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id) # Build request body request_body = {k: v for k, v in update_request.items() if v is not None} - verbose_logger.debug( - "Update eval request - URL: %s, body: %s", url, request_body - ) + verbose_logger.debug("Update eval request - URL: %s, body: %s", url, request_body) return url, headers, request_body @@ -218,9 +205,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform delete eval request for OpenAI""" - url = self.get_complete_url( - api_base=api_base, endpoint="evals", eval_id=eval_id - ) + url = self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id) verbose_logger.debug("Delete eval request - URL: %s", url) @@ -284,9 +269,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): # Build request body request_body = {k: v for k, v in create_request.items() if v is not None} - verbose_logger.debug( - "Create run request - URL: %s, body: %s", url, request_body - ) + verbose_logger.debug("Create run request - URL: %s, body: %s", url, request_body) return url, request_body diff --git a/litellm/llms/openai/fine_tuning/handler.py b/litellm/llms/openai/fine_tuning/handler.py index ca93622d9de..e0914a9ff0d 100644 --- a/litellm/llms/openai/fine_tuning/handler.py +++ b/litellm/llms/openai/fine_tuning/handler.py @@ -19,9 +19,7 @@ _AZURE_STATUS_MAP = { # because LiteLLMFineTuningJob schema has no intermediate cancellation state. -def _normalize_fine_tuning_job_dict( - data: Dict[str, Any], is_azure: bool = False -) -> Dict[str, Any]: +def _normalize_fine_tuning_job_dict(data: Dict[str, Any], is_azure: bool = False) -> Dict[str, Any]: """ Normalize Azure OpenAI FineTuningJob response to match OpenAI schema. @@ -48,12 +46,8 @@ def _normalize_fine_tuning_job_dict( return normalized -def _litellm_fine_tuning_job_from_response( - response: Any, is_azure: bool = False -) -> LiteLLMFineTuningJob: - return LiteLLMFineTuningJob( - **_normalize_fine_tuning_job_dict(response.model_dump(), is_azure=is_azure) - ) +def _litellm_fine_tuning_job_from_response(response: Any, is_azure: bool = False) -> LiteLLMFineTuningJob: + return LiteLLMFineTuningJob(**_normalize_fine_tuning_job_dict(response.model_dump(), is_azure=is_azure)) class OpenAIFineTuningAPI: @@ -71,9 +65,7 @@ class OpenAIFineTuningAPI: timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, _is_async: bool = False, api_version: Optional[str] = None, litellm_params: Optional[dict] = None, @@ -86,9 +78,7 @@ class OpenAIFineTuningAPI: ] ]: received_args = locals() - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None if client is None: data = {} for k, v in received_args.items(): @@ -112,9 +102,7 @@ class OpenAIFineTuningAPI: create_fine_tuning_job_data: dict, openai_client: Union[AsyncOpenAI, AsyncAzureOpenAI], ) -> LiteLLMFineTuningJob: - response = await openai_client.fine_tuning.jobs.create( - **create_fine_tuning_job_data - ) + response = await openai_client.fine_tuning.jobs.create(**create_fine_tuning_job_data) return _litellm_fine_tuning_job_from_response(response) @@ -128,13 +116,9 @@ class OpenAIFineTuningAPI: timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = self.get_openai_client( + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, timeout=timeout, @@ -158,12 +142,8 @@ class OpenAIFineTuningAPI: create_fine_tuning_job_data=create_fine_tuning_job_data, openai_client=openai_client, ) - verbose_logger.debug( - "creating fine tuning job, args= %s", create_fine_tuning_job_data - ) - response = cast(OpenAI, openai_client).fine_tuning.jobs.create( - **create_fine_tuning_job_data - ) + verbose_logger.debug("creating fine tuning job, args= %s", create_fine_tuning_job_data) + response = cast(OpenAI, openai_client).fine_tuning.jobs.create(**create_fine_tuning_job_data) return _litellm_fine_tuning_job_from_response(response) async def acancel_fine_tuning_job( @@ -171,9 +151,7 @@ class OpenAIFineTuningAPI: fine_tuning_job_id: str, openai_client: Union[AsyncOpenAI, AsyncAzureOpenAI], ) -> LiteLLMFineTuningJob: - response = await openai_client.fine_tuning.jobs.cancel( - fine_tuning_job_id=fine_tuning_job_id - ) + response = await openai_client.fine_tuning.jobs.cancel(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response) def cancel_fine_tuning_job( @@ -186,13 +164,9 @@ class OpenAIFineTuningAPI: timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = self.get_openai_client( + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, timeout=timeout, @@ -217,9 +191,7 @@ class OpenAIFineTuningAPI: openai_client=openai_client, ) verbose_logger.debug("canceling fine tuning job, args= %s", fine_tuning_job_id) - response = cast(OpenAI, openai_client).fine_tuning.jobs.cancel( - fine_tuning_job_id=fine_tuning_job_id - ) + response = cast(OpenAI, openai_client).fine_tuning.jobs.cancel(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response) async def alist_fine_tuning_jobs( @@ -240,15 +212,11 @@ class OpenAIFineTuningAPI: timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, after: Optional[str] = None, limit: Optional[int] = None, ): - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = self.get_openai_client( + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, timeout=timeout, @@ -282,9 +250,7 @@ class OpenAIFineTuningAPI: fine_tuning_job_id: str, openai_client: Union[AsyncOpenAI, AsyncAzureOpenAI], ) -> LiteLLMFineTuningJob: - response = await openai_client.fine_tuning.jobs.retrieve( - fine_tuning_job_id=fine_tuning_job_id - ) + response = await openai_client.fine_tuning.jobs.retrieve(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response) def retrieve_fine_tuning_job( @@ -297,13 +263,9 @@ class OpenAIFineTuningAPI: timeout: Union[float, httpx.Timeout], max_retries: Optional[int], organization: Optional[str], - client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = None, + client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: - openai_client: Optional[ - Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI] - ] = self.get_openai_client( + openai_client: Optional[Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, timeout=timeout, @@ -328,7 +290,5 @@ class OpenAIFineTuningAPI: openai_client=openai_client, ) verbose_logger.debug("retrieving fine tuning job, id= %s", fine_tuning_job_id) - response = cast(OpenAI, openai_client).fine_tuning.jobs.retrieve( - fine_tuning_job_id=fine_tuning_job_id - ) + response = cast(OpenAI, openai_client).fine_tuning.jobs.retrieve(fine_tuning_job_id=fine_tuning_job_id) return _litellm_fine_tuning_job_from_response(response) diff --git a/litellm/llms/openai/image_edit/dalle2_transformation.py b/litellm/llms/openai/image_edit/dalle2_transformation.py index 04995ce9514..ac08d056a34 100644 --- a/litellm/llms/openai/image_edit/dalle2_transformation.py +++ b/litellm/llms/openai/image_edit/dalle2_transformation.py @@ -58,16 +58,12 @@ class DallE2ImageEditConfig(OpenAIImageEditConfig): ######################################################### _image_list = request_dict.get("image") _mask = request_dict.get("mask") - data_without_files = { - k: v for k, v in request_dict.items() if k not in ["image", "mask"] - } + data_without_files = {k: v for k, v in request_dict.items() if k not in ["image", "mask"]} files_list: List[Tuple[str, Any]] = [] # Handle image parameter - DALL-E-2 only supports single image if _image_list is not None: - image_list = ( - [_image_list] if not isinstance(_image_list, list) else _image_list - ) + image_list = [_image_list] if not isinstance(_image_list, list) else _image_list # Validate only one image is provided if len(image_list) > 1: @@ -93,9 +89,7 @@ class DallE2ImageEditConfig(OpenAIImageEditConfig): _mask = _mask[0] if _mask else None if _mask is not None: - mask_content_type: str = ImageEditRequestUtils.get_image_content_type( - _mask - ) + mask_content_type: str = ImageEditRequestUtils.get_image_content_type(_mask) if isinstance(_mask, BufferedReader): files_list.append(("mask", (_mask.name, _mask, mask_content_type))) else: diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index 9c0daca8022..f53c1731f58 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -110,16 +110,12 @@ class OpenAIImageEditConfig(BaseImageEditConfig): ######################################################### _image_list = request_dict.get("image") _mask = request_dict.get("mask") - data_without_files = { - k: v for k, v in request_dict.items() if k not in ["image", "mask"] - } + data_without_files = {k: v for k, v in request_dict.items() if k not in ["image", "mask"]} files_list: List[Tuple[str, Any]] = [] # Handle image parameter if _image_list is not None: - image_list = ( - [_image_list] if not isinstance(_image_list, list) else _image_list - ) + image_list = [_image_list] if not isinstance(_image_list, list) else _image_list for _image in image_list: if _image is not None: @@ -135,9 +131,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): _mask = _mask[0] if _mask else None if _mask is not None: - mask_content_type: str = ImageEditRequestUtils.get_image_content_type( - _mask - ) + mask_content_type: str = ImageEditRequestUtils.get_image_content_type(_mask) if isinstance(_mask, BufferedReader): files_list.append(("mask", (_mask.name, _mask, mask_content_type))) else: @@ -155,9 +149,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): try: raw_response_json = raw_response.json() except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) return ImageResponse(**raw_response_json) def validate_environment( @@ -168,12 +160,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): litellm_params: Optional[dict] = None, api_base: Optional[str] = None, ) -> dict: - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.update( { "Authorization": f"Bearer {api_key}", diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py index dab277a7ba8..effda2fa3ee 100644 --- a/litellm/llms/openai/image_generation/cost_calculator.py +++ b/litellm/llms/openai/image_generation/cost_calculator.py @@ -22,18 +22,14 @@ def cost_calculator( """Calculate cost for OpenAI gpt-image models (token-based pricing).""" usage = getattr(image_response, "usage", None) if usage is None: - verbose_logger.debug( - f"No usage data available for {model}, cannot calculate token-based cost" - ) + verbose_logger.debug(f"No usage data available for {model}, cannot calculate token-based cost") return 0.0 provider = custom_llm_provider or "openai" # A chat Usage with an explicit output breakdown: cost via generic_cost_per_token. if isinstance(usage, Usage) and usage.completion_tokens_details is not None: - prompt_cost, completion_cost = generic_cost_per_token( - model=model, usage=usage, custom_llm_provider=provider - ) + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider) return prompt_cost + completion_cost # ImageUsage / ResponseAPIUsage: reuse the shared helper (same path as @@ -50,9 +46,7 @@ def cost_calculator( # Fallback: a Usage with no output breakdown that the image helper can't read — # cost via generic_cost_per_token (text rate) instead of returning 0.0. if isinstance(usage, Usage): - prompt_cost, completion_cost = generic_cost_per_token( - model=model, usage=usage, custom_llm_provider=provider - ) + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider) return prompt_cost + completion_cost return 0.0 diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py index 22c2349a837..fbc2e8dec3d 100644 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_2_transformation.py @@ -18,9 +18,7 @@ class DallE2ImageGenerationConfig(BaseImageGenerationConfig): OpenAI dall-e-2 image generation config """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: return ["n", "response_format", "quality", "size", "user"] def map_openai_params( @@ -74,14 +72,8 @@ class DallE2ImageGenerationConfig(BaseImageGenerationConfig): ) # set optional params - image_response.size = optional_params.get( - "size", "1024x1024" - ) # default is always 1024x1024 - image_response.quality = optional_params.get( - "quality", "standard" - ) # always standard for dall-e-2 - image_response.output_format = optional_params.get( - "output_format", "png" - ) # always png for dall-e-2 + image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 + image_response.quality = optional_params.get("quality", "standard") # always standard for dall-e-2 + image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-2 return image_response diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py index 9e2bdabc3a1..3434c708113 100644 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_3_transformation.py @@ -18,9 +18,7 @@ class DallE3ImageGenerationConfig(BaseImageGenerationConfig): OpenAI dall-e-3 image generation config """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: return ["n", "response_format", "quality", "size", "user", "style"] def map_openai_params( @@ -74,14 +72,8 @@ class DallE3ImageGenerationConfig(BaseImageGenerationConfig): ) # set optional params - image_response.size = optional_params.get( - "size", "1024x1024" - ) # default is always 1024x1024 - image_response.quality = optional_params.get( - "quality", "hd" - ) # always hd for dall-e-3 - image_response.output_format = optional_params.get( - "output_format", "png" - ) # always png for dall-e-3 + image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 + image_response.quality = optional_params.get("quality", "hd") # always hd for dall-e-3 + image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-3 return image_response diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 68f799e5747..b9c2368d4be 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -18,9 +18,7 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): OpenAI gpt-image image generation config """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: return [ "background", "moderation", @@ -83,14 +81,8 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): ) # set optional params - image_response.size = optional_params.get( - "size", "1024x1024" - ) # default is always 1024x1024 - image_response.quality = optional_params.get( - "quality", "high" - ) # always hd for dall-e-3 - image_response.output_format = optional_params.get( - "response_format", "png" - ) # always png for dall-e-3 + image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 + image_response.quality = optional_params.get("quality", "high") # always hd for dall-e-3 + image_response.output_format = optional_params.get("response_format", "png") # always png for dall-e-3 return image_response diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index 76610088d0c..56bc00f319c 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -46,9 +46,7 @@ class OpenAIImageGenerationHandler(BaseTranslation): """ prompt = data.get("prompt") if prompt is None: - verbose_proxy_logger.debug( - "OpenAI Image Generation: No prompt found in request data" - ) + verbose_proxy_logger.debug("OpenAI Image Generation: No prompt found in request data") return data # Apply guardrail to the prompt @@ -68,8 +66,7 @@ class OpenAIImageGenerationHandler(BaseTranslation): data["prompt"] = guardrailed_texts[0] if guardrailed_texts else prompt verbose_proxy_logger.debug( - "OpenAI Image Generation: Applied guardrail to prompt. " - "Original length: %d, New length: %d", + "OpenAI Image Generation: Applied guardrail to prompt. Original length: %d, New length: %d", len(prompt), len(data["prompt"]), ) @@ -105,7 +102,5 @@ class OpenAIImageGenerationHandler(BaseTranslation): Returns: Unmodified response (images don't need text guardrails) """ - verbose_proxy_logger.debug( - "OpenAI Image Generation: Output processing not needed for image responses" - ) + verbose_proxy_logger.debug("OpenAI Image Generation: Output processing not needed for image responses") return response diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index dae3fa9d457..00cbb87e31d 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -30,9 +30,7 @@ class OpenAIImageVariationsHandler: openai_client = client return openai_client - def get_async_client( - self, client: Optional[AsyncOpenAI], init_client_params: dict - ) -> AsyncOpenAI: + def get_async_client(self, client: Optional[AsyncOpenAI], init_client_params: dict) -> AsyncOpenAI: if client is None: openai_client = AsyncOpenAI( **init_client_params, @@ -69,13 +67,9 @@ class OpenAIImageVariationsHandler: "organization": organization, } - client = self.get_async_client( - client=client, init_client_params=init_client_params - ) + client = self.get_async_client(client=client, init_client_params=init_client_params) - raw_response = await client.images.with_raw_response.create_variation( - **data - ) # type: ignore + raw_response = await client.images.with_raw_response.create_variation(**data) # type: ignore response = raw_response.parse() response_json = response.model_dump() @@ -95,9 +89,7 @@ class OpenAIImageVariationsHandler: model_response=ImageResponse(**response_json), raw_response=httpx.Response( status_code=200, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), # mock request object + request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ), logging_obj=logging_obj, request_data=data, @@ -114,9 +106,7 @@ class OpenAIImageVariationsHandler: error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) def image_variations( self, @@ -143,9 +133,7 @@ class OpenAIImageVariationsHandler: ) if provider_config is None: - raise ValueError( - f"image variation provider not found: {custom_llm_provider}." - ) + raise ValueError(f"image variation provider not found: {custom_llm_provider}.") max_retries = optional_params.pop("max_retries", 2) @@ -157,9 +145,7 @@ class OpenAIImageVariationsHandler: ) json_data = data.get("data") if not json_data: - raise ValueError( - f"data field is required, for openai image variations. Got={data}" - ) + raise ValueError(f"data field is required, for openai image variations. Got={data}") ## LOGGING logging_obj.pre_call( input="", @@ -198,9 +184,7 @@ class OpenAIImageVariationsHandler: "organization": organization, } - client = self.get_sync_client( - client=client, init_client_params=init_client_params - ) + client = self.get_sync_client(client=client, init_client_params=init_client_params) raw_response = client.images.with_raw_response.create_variation(**json_data) # type: ignore response = raw_response.parse() @@ -222,9 +206,7 @@ class OpenAIImageVariationsHandler: model_response=ImageResponse(**response_json), raw_response=httpx.Response( status_code=200, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), # mock request object + request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ), logging_obj=logging_obj, request_data=json_data, @@ -241,6 +223,4 @@ class OpenAIImageVariationsHandler: error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) diff --git a/litellm/llms/openai/image_variations/transformation.py b/litellm/llms/openai/image_variations/transformation.py index 96d1a302761..2f16c6f3d23 100644 --- a/litellm/llms/openai/image_variations/transformation.py +++ b/litellm/llms/openai/image_variations/transformation.py @@ -13,9 +13,7 @@ from ..common_utils import OpenAIError class OpenAIImageVariationConfig(BaseImageVariationConfig): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageVariationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageVariationOptionalParams]: return ["n", "size", "response_format", "user"] def map_openai_params( @@ -72,9 +70,7 @@ class OpenAIImageVariationConfig(BaseImageVariationConfig): ) -> ImageResponse: return model_response - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: return OpenAIError( status_code=status_code, message=error_message, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 0b90381ba59..6b191144a11 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -198,18 +198,14 @@ class OpenAIConfig(BaseConfig): else: return litellm.openAIGPTConfig.get_supported_openai_params(model=model) - def _map_openai_params( - self, non_default_params: dict, optional_params: dict, model: str - ) -> dict: + def _map_openai_params(self, non_default_params: dict, optional_params: dict, model: str) -> dict: supported_openai_params = self.get_supported_openai_params(model) for param, value in non_default_params.items(): if param in supported_openai_params: optional_params[param] = value return optional_params - def _transform_messages( - self, messages: List[AllMessageValues], model: str - ) -> List[AllMessageValues]: + def _transform_messages(self, messages: List[AllMessageValues], model: str) -> List[AllMessageValues]: return messages def map_openai_params( @@ -368,9 +364,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): if not isinstance(max_retries, int): raise OpenAIError( status_code=422, - message="max retries must be an int. Passed in value: {}".format( - max_retries - ), + message="max retries must be an int. Passed in value: {}".format(max_retries), ) cached_client = self.get_cached_openai_client( client_initialization_params=client_initialization_params, @@ -378,17 +372,13 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) if cached_client: - if isinstance(cached_client, OpenAI) or isinstance( - cached_client, AsyncOpenAI - ): + if isinstance(cached_client, OpenAI) or isinstance(cached_client, AsyncOpenAI): return cached_client if is_async: _new_client: Union[OpenAI, AsyncOpenAI] = AsyncOpenAI( api_key=api_key, base_url=api_base, - http_client=OpenAIChatCompletion._get_async_http_client( - shared_session=shared_session - ), + http_client=OpenAIChatCompletion._get_async_http_client(shared_session=shared_session), timeout=timeout, max_retries=max_retries, organization=organization, @@ -434,11 +424,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ start_time = time.time() try: - raw_response = ( - await openai_aclient.chat.completions.with_raw_response.create( - **data, timeout=timeout - ) - ) + raw_response = await openai_aclient.chat.completions.with_raw_response.create(**data, timeout=timeout) end_time = time.time() if hasattr(raw_response, "headers"): @@ -475,9 +461,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ raw_response = None try: - raw_response = openai_client.chat.completions.with_raw_response.create( - **data, timeout=timeout - ) + raw_response = openai_client.chat.completions.with_raw_response.create(**data, timeout=timeout) if hasattr(raw_response, "headers"): headers = dict(raw_response.headers) @@ -539,9 +523,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): try: if isinstance(callback, CustomLogger): # Check if the callback has the chat completion agentic loop methods - if not hasattr( - callback, "async_should_run_chat_completion_agentic_loop" - ): + if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"): continue # First: Check if agentic loop should run (using chat completion method) @@ -560,25 +542,19 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): if should_run: # Second: Execute agentic loop - kwargs_with_provider = ( - litellm_params.copy() if litellm_params else {} - ) - kwargs_with_provider["custom_llm_provider"] = ( - custom_llm_provider - ) + kwargs_with_provider = litellm_params.copy() if litellm_params else {} + kwargs_with_provider["custom_llm_provider"] = custom_llm_provider # For OpenAI Chat Completions, use the chat completion agentic loop method - agentic_response = ( - await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) + agentic_response = await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, ) # First hook that runs agentic loop wins return agentic_response @@ -637,9 +613,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): try: fake_stream: bool = False inference_params = optional_params.copy() - stream_options: Optional[dict] = inference_params.pop( - "stream_options", None - ) + stream_options: Optional[dict] = inference_params.pop("stream_options", None) stream: Optional[bool] = inference_params.pop("stream", False) provider_config: Optional[BaseConfig] = None @@ -665,9 +639,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): if model is None or messages is None: raise OpenAIError(status_code=422, message="Missing model or messages") - if not isinstance(timeout, float) and not isinstance( - timeout, httpx.Timeout - ): + if not isinstance(timeout, float) and not isinstance(timeout, httpx.Timeout): raise OpenAIError( status_code=422, message="Timeout needs to be a float or httpx.Timeout", @@ -676,9 +648,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): if custom_llm_provider is not None and custom_llm_provider != "openai": model_response.model = f"{custom_llm_provider}/{model}" - for _ in range( - 2 - ): # if call fails due to alternating messages, retry with reformatted message + for _ in range(2): # if call fails due to alternating messages, retry with reformatted message try: max_retries = inference_params.pop("max_retries", 2) if acompletion is True: @@ -748,9 +718,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) else: if not isinstance(max_retries, int): - raise OpenAIError( - status_code=422, message="max retries must be an int" - ) + raise OpenAIError(status_code=422, message="max retries must be an int") openai_client: OpenAI = self._get_openai_client( # type: ignore is_async=False, api_key=api_key, @@ -785,11 +753,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) logging_obj.model_call_details["response_headers"] = headers - stringified_response = ( - provider_config.transform_parsed_response_dict( - response.model_dump() - ) - ) + stringified_response = provider_config.transform_parsed_response_dict(response.model_dump()) logging_obj.post_call( input=messages, api_key=api_key, @@ -814,9 +778,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): except openai.UnprocessableEntityError as e: ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 if litellm.drop_params is True or drop_params is True: - inference_params = drop_params_from_unprocessable_entity_error( - e, inference_params - ) + inference_params = drop_params_from_unprocessable_entity_error(e, inference_params) else: raise e # e.message @@ -835,22 +797,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): new_messages.append(messages[i]) if messages[i]["role"] == messages[i + 1]["role"]: if messages[i]["role"] == "user": - new_messages.append( - {"role": "assistant", "content": ""} - ) + new_messages.append({"role": "assistant", "content": ""}) else: new_messages.append({"role": "user", "content": ""}) new_messages.append(messages[-1]) messages = new_messages - elif ( - "Last message must have role `user`" in str(e) - ) and messages is not None: + elif ("Last message must have role `user`" in str(e)) and messages is not None: new_messages = messages new_messages.append({"role": "user", "content": ""}) messages = new_messages - elif "unknown field: parameter index is not a valid field" in str( - e - ): + elif "unknown field: parameter index is not a valid field" in str(e): litellm.remove_index_from_tool_calls(messages=messages) else: raise e @@ -901,9 +857,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): litellm_params=litellm_params, headers=headers or {}, ) - for _ in range( - 2 - ): # if call fails due to alternating messages, retry with reformatted message + for _ in range(2): # if call fails due to alternating messages, retry with reformatted message try: openai_aclient: AsyncOpenAI = self._get_openai_client( # type: ignore is_async=True, @@ -922,9 +876,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): input=data["messages"], api_key=openai_aclient.api_key, additional_args={ - "headers": { - "Authorization": f"Bearer {openai_aclient.api_key}" - }, + "headers": {"Authorization": f"Bearer {openai_aclient.api_key}"}, "api_base": openai_aclient._base_url._uri_reference, "acompletion": True, "complete_input_dict": data, @@ -937,9 +889,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): timeout=timeout, logging_obj=logging_obj, ) - stringified_response = provider_config.transform_parsed_response_dict( - response.model_dump() - ) + stringified_response = provider_config.transform_parsed_response_dict(response.model_dump()) logging_obj.post_call( input=data["messages"], api_key=api_key, @@ -1016,9 +966,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): stream_options: Optional[dict] = None, ): data["stream"] = True - data.update( - self.get_stream_options(stream_options=stream_options, api_base=api_base) - ) + data.update(self.get_stream_options(stream_options=stream_options, api_base=api_base)) openai_client: OpenAI = self._get_openai_client( # type: ignore is_async=False, @@ -1088,9 +1036,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): headers=headers or {}, ) data["stream"] = True - data.update( - self.get_stream_options(stream_options=stream_options, api_base=api_base) - ) + data.update(self.get_stream_options(stream_options=stream_options, api_base=api_base)) for _ in range(2): try: openai_aclient: AsyncOpenAI = self._get_openai_client( # type: ignore @@ -1138,7 +1084,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data = drop_params_from_unprocessable_entity_error(e, data) else: raise e - except Exception as e: # need to exception handle here. async exceptions don't get caught in sync functions. + except ( + Exception + ) as e: # need to exception handle here. async exceptions don't get caught in sync functions. if isinstance(e, OpenAIError): raise e @@ -1178,9 +1126,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): body=exception_body, ) - def get_stream_options( - self, stream_options: Optional[dict], api_base: Optional[str] - ) -> dict: + def get_stream_options(self, stream_options: Optional[dict], api_base: Optional[str]) -> dict: """ Pass `stream_options` to the data dict for OpenAI requests """ @@ -1207,9 +1153,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): - call embeddings.create by default """ try: - raw_response = await openai_aclient.embeddings.with_raw_response.create( - **data, timeout=timeout - ) # type: ignore + raw_response = await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) # type: ignore headers = dict(raw_response.headers) response = raw_response.parse() return headers, response @@ -1230,9 +1174,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): - call embeddings.create by default """ try: - raw_response = openai_client.embeddings.with_raw_response.create( - **data, timeout=timeout - ) # type: ignore + raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout) # type: ignore headers = dict(raw_response.headers) response = raw_response.parse() @@ -1308,9 +1250,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) def embedding( # type: ignore self, @@ -1396,9 +1336,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): error_response = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) - raise OpenAIError( - status_code=status_code, message=error_text, headers=error_headers - ) + raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers) async def aimage_generation( self, @@ -1546,9 +1484,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): original_response=str(e), ) if hasattr(e, "status_code"): - raise OpenAIError( - status_code=getattr(e, "status_code", 500), message=str(e) - ) + raise OpenAIError(status_code=getattr(e, "status_code", 500), message=str(e)) else: raise OpenAIError(status_code=500, message=str(e)) @@ -1746,9 +1682,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: Optional[int], organization: Optional[str], client: Optional[Union[OpenAI, AsyncOpenAI]] = None, - ) -> Union[ - HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] - ]: + ) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]: openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -1782,9 +1716,7 @@ class OpenAIFilesAPI(BaseLLM): openai_client: AsyncOpenAI, chunk_size: int = 1024 * 1024, ) -> FileContentStreamingResult: - response_cm = openai_client.files.with_streaming_response.content( - **file_content_request - ) + response_cm = openai_client.files.with_streaming_response.content(**file_content_request) response = await response_cm.__aenter__() headers = dict(response.headers) @@ -1841,9 +1773,7 @@ class OpenAIFilesAPI(BaseLLM): chunk_size=chunk_size, ) - response_cm = cast(OpenAI, openai_client).files.with_streaming_response.content( - **file_content_request - ) + response_cm = cast(OpenAI, openai_client).files.with_streaming_response.content(**file_content_request) response = response_cm.__enter__() headers = dict(response.headers) @@ -2185,9 +2115,7 @@ class OpenAIBatchesAPI(BaseLLM): # At this point, openai_client is guaranteed to be a sync OpenAI client if not isinstance(openai_client, OpenAI): - raise ValueError( - "OpenAI client is not an instance of OpenAI. Make sure you passed a sync OpenAI client." - ) + raise ValueError("OpenAI client is not an instance of OpenAI. Make sure you passed a sync OpenAI client.") response = openai_client.batches.cancel(**cancel_batch_data) return LiteLLMBatch(**response.model_dump()) @@ -2541,11 +2469,9 @@ class OpenAIAssistantsAPI(BaseLLM): client=client, ) - thread_message: OpenAIMessage = ( - await openai_client.beta.threads.messages.create( # type: ignore - thread_id, - **message_data, # type: ignore - ) + thread_message: OpenAIMessage = await openai_client.beta.threads.messages.create( # type: ignore + thread_id, + **message_data, # type: ignore ) response_obj: Optional[OpenAIMessage] = None diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index cebab20adde..626d2f3a28e 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -142,9 +142,7 @@ class OpenAIRealtime(OpenAIChatCompletion): "If your client expects beta event names, add 'OpenAI-Beta: realtime=v1' " "to the WebSocket headers sent to the LiteLLM proxy." ) - headers = self._get_additional_headers( - api_key, openai_beta_realtime=openai_beta_realtime - ) + headers = self._get_additional_headers(api_key, openai_beta_realtime=openai_beta_realtime) # Log a masked request preview consistent with other endpoints. logging_obj.pre_call( @@ -170,9 +168,7 @@ class OpenAIRealtime(OpenAIChatCompletion): user_api_key_dict=user_api_key_dict, request_data={"litellm_metadata": litellm_metadata or {}}, force_transcription_model=( - model - if (query_params or {}).get("intent") == "transcription" - else None + model if (query_params or {}).get("intent") == "transcription" else None ), event_normalizer=self._make_event_normalizer(), ) @@ -182,17 +178,11 @@ class OpenAIRealtime(OpenAIChatCompletion): await websocket.close(code=e.status_code, reason=_redact_string(str(e))) except Exception as e: try: - await websocket.close( - code=1011, reason=_redact_string(f"Internal server error: {str(e)}") - ) + await websocket.close(code=1011, reason=_redact_string(f"Internal server error: {str(e)}")) except RuntimeError as close_error: - if "already completed" in str(close_error) or "websocket.close" in str( - close_error - ): + if "already completed" in str(close_error) or "websocket.close" in str(close_error): # The WebSocket is already closed or the response is completed, so we can ignore this error pass else: # If it's a different RuntimeError, we might want to log it or handle it differently - raise Exception( - f"Unexpected error while closing WebSocket: {close_error}" - ) + raise Exception(f"Unexpected error while closing WebSocket: {close_error}") diff --git a/litellm/llms/openai/realtime/http_transformation.py b/litellm/llms/openai/realtime/http_transformation.py index 7a6af39ba65..0a7e65dfea2 100644 --- a/litellm/llms/openai/realtime/http_transformation.py +++ b/litellm/llms/openai/realtime/http_transformation.py @@ -9,33 +9,18 @@ from litellm.secret_managers.main import get_secret_str class OpenAIRealtimeHTTPConfig(BaseRealtimeHTTPConfig): def get_api_base(self, api_base: Optional[str], **kwargs) -> str: - return ( - api_base - or litellm.api_base - or get_secret_str("OPENAI_API_BASE") - or "https://api.openai.com" - ) + return api_base or litellm.api_base or get_secret_str("OPENAI_API_BASE") or "https://api.openai.com" def get_api_key(self, api_key: Optional[str], **kwargs) -> str: - return ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - or "" - ) + return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") or "" - def get_complete_url( - self, api_base: Optional[str], model: str, api_version: Optional[str] = None - ) -> str: + def get_complete_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: base = self.get_api_base(api_base).rstrip("/") if base.endswith("/v1"): base = base[:-3] return f"{base}/v1/realtime/client_secrets" - def get_realtime_calls_url( - self, api_base: Optional[str], model: str, api_version: Optional[str] = None - ) -> str: + def get_realtime_calls_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: base = self.get_api_base(api_base).rstrip("/") if base.endswith("/v1"): base = base[:-3] diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 7fb5f6dad78..3dded042de8 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -45,9 +45,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): try: self.validate_request(model, input) - verbose_logger.debug( - f"Processing OpenAI CountTokens request for model: {model}" - ) + verbose_logger.debug(f"Processing OpenAI CountTokens request for model: {model}") request_body = self.transform_request_to_count_tokens( model=model, @@ -62,13 +60,9 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): headers = self.get_required_headers(api_key) - async_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.OPENAI - ) + async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.OPENAI) - request_timeout = ( - timeout if timeout is not None else litellm.request_timeout - ) + request_timeout = timeout if timeout is not None else litellm.request_timeout response = await async_client.post( endpoint_url, diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 3d3a659075e..8e700ecafa1 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -60,9 +60,7 @@ class OpenAITokenCounter(BaseTokenCounter): api_base = litellm_params.get("api_base") # Convert chat messages to Responses API input format - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( - messages - ) + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) # Use system param if instructions not extracted from messages if instructions is None and system is not None: @@ -91,9 +89,7 @@ class OpenAITokenCounter(BaseTokenCounter): original_response=result, ) except OpenAIError as e: - verbose_logger.warning( - f"OpenAI CountTokens API error: status={e.status_code}, message={e.message}" - ) + verbose_logger.warning(f"OpenAI CountTokens API error: status={e.status_code}, message={e.message}") return TokenCountResponse( total_tokens=0, request_model=request_model, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 83ec69a9d7a..6ac33ffa44a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -80,11 +80,9 @@ class OpenAIResponsesHandler(BaseTranslation): input_data = data.get("input") if input_data is None: return None - messages = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=input_data, - responses_api_request=data, - ) + messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=input_data, + responses_api_request=data, ) return cast(List[AllMessageValues], messages) if messages else None @@ -132,9 +130,7 @@ class OpenAIResponsesHandler(BaseTranslation): ) 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, guardrailed_inputs.get("tools") - ) + self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") return data @@ -195,9 +191,7 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) - verbose_proxy_logger.debug( - "OpenAI Responses API: Processed input messages: %s", input_data - ) + verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", input_data) return data @@ -232,13 +226,9 @@ class OpenAIResponsesHandler(BaseTranslation): ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( tools # type: ignore ) - tools_to_check.extend( - cast(List[ChatCompletionToolParam], transformed_tools) - ) + tools_to_check.extend(cast(List[ChatCompletionToolParam], transformed_tools)) - def _remap_tools_to_responses_api_format( - self, guardrailed_tools: List[Any] - ) -> List[Dict[str, Any]]: + def _remap_tools_to_responses_api_format(self, guardrailed_tools: List[Any]) -> List[Dict[str, Any]]: """ Remap guardrail-returned tools (Chat Completion format) back to Responses API request tool format. @@ -350,9 +340,7 @@ class OpenAIResponsesHandler(BaseTranslation): elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content if isinstance(messages[msg_idx]["content"][content_idx_optional], dict): - messages[msg_idx]["content"][content_idx_optional]["text"] = ( - guardrail_response - ) + messages[msg_idx]["content"][content_idx_optional]["text"] = guardrail_response async def process_output_response( self, @@ -394,9 +382,7 @@ class OpenAIResponsesHandler(BaseTranslation): elif hasattr(response, "output"): response_output = response.output or [] else: - verbose_proxy_logger.debug( - "OpenAI Responses API: No output found in response" - ) + verbose_proxy_logger.debug("OpenAI Responses API: No output found in response") return response if not response_output: @@ -426,9 +412,7 @@ class OpenAIResponsesHandler(BaseTranslation): # Add user API key metadata with prefixed keys if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata @@ -462,9 +446,7 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) - verbose_proxy_logger.debug( - "OpenAI Responses API: Processed output response: %s", response - ) + verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response) return response @@ -527,17 +509,13 @@ class OpenAIResponsesHandler(BaseTranslation): if "response" not in request_data: request_data["response"] = response_obj if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if tool_calls_to_check: - inputs["tool_calls"] = cast( - List[ChatCompletionToolCallChunk], tool_calls_to_check - ) + inputs["tool_calls"] = cast(List[ChatCompletionToolCallChunk], tool_calls_to_check) response_model = response_obj.get("model") if response_model: inputs["model"] = response_model @@ -566,19 +544,14 @@ class OpenAIResponsesHandler(BaseTranslation): # Case 2: response.output_item.done — extract tool calls only. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": - model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( - final_chunk + model_response_stream = ( + OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) ) tool_calls = model_response_stream.choices[0].delta.tool_calls if tool_calls: inputs = GenericGuardrailAPIInputs() - inputs["tool_calls"] = cast( - List[ChatCompletionToolCallChunk], tool_calls - ) - if ( - hasattr(model_response_stream, "model") - and model_response_stream.model - ): + inputs["tool_calls"] = cast(List[ChatCompletionToolCallChunk], tool_calls) + if hasattr(model_response_stream, "model") and model_response_stream.model: inputs["model"] = model_response_stream.model await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -597,9 +570,7 @@ class OpenAIResponsesHandler(BaseTranslation): if string_so_far: fallback_inputs = GenericGuardrailAPIInputs(texts=[string_so_far]) response_model = ( - final_chunk.get("response", {}).get("model") - if isinstance(final_chunk.get("response"), dict) - else None + final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None ) if response_model: fallback_inputs["model"] = response_model @@ -615,10 +586,7 @@ class OpenAIResponsesHandler(BaseTranslation): """ Check if the streaming has ended. """ - return all( - response.choices[0].finish_reason is not None - for response in responses_so_far - ) + return all(response.choices[0].finish_reason is not None for response in responses_so_far) def get_streaming_string_so_far(self, responses_so_far: List[Any]) -> str: """ @@ -638,11 +606,7 @@ class OpenAIResponsesHandler(BaseTranslation): for output_item in response.output: if isinstance(output_item, BaseModel): try: - generic_response_output_item = ( - GenericResponseOutputItem.model_validate( - output_item.model_dump() - ) - ) + generic_response_output_item = GenericResponseOutputItem.model_validate(output_item.model_dump()) if generic_response_output_item.content: output_item = generic_response_output_item except Exception: @@ -682,13 +646,13 @@ class OpenAIResponsesHandler(BaseTranslation): # Check if this is a tool call (OutputFunctionToolCall) if isinstance(output_item, OutputFunctionToolCall): if tool_calls_to_check is not None: - tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=output_item, - index=output_idx, - ) - tool_calls_to_check.append( - cast(ChatCompletionToolCallChunk, tool_call_dict) + tool_call_dict = ( + LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( + tool_call_item=output_item, + index=output_idx, + ) ) + tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) return elif ( isinstance(output_item, BaseModel) @@ -696,17 +660,15 @@ class OpenAIResponsesHandler(BaseTranslation): and getattr(output_item, "type") == "function_call" ): if tool_calls_to_check is not None: - tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=output_item, - index=output_idx, - ) - tool_calls_to_check.append( - cast(ChatCompletionToolCallChunk, tool_call_dict) + tool_call_dict = ( + LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( + tool_call_item=output_item, + index=output_idx, + ) ) + tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) return - elif ( - isinstance(output_item, dict) and output_item.get("type") == "function_call" - ): + elif isinstance(output_item, dict) and output_item.get("type") == "function_call": # Handle dict representation of tool call if tool_calls_to_check is not None: # Convert dict to ResponseFunctionToolCall for processing @@ -716,9 +678,7 @@ class OpenAIResponsesHandler(BaseTranslation): tool_call_item=tool_call_obj, index=output_idx, ) - tool_calls_to_check.append( - cast(ChatCompletionToolCallChunk, tool_call_dict) - ) + tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict)) except Exception: pass return @@ -728,9 +688,7 @@ class OpenAIResponsesHandler(BaseTranslation): if isinstance(output_item, BaseModel): try: output_item_dump = output_item.model_dump() - generic_response_output_item = GenericResponseOutputItem.model_validate( - output_item_dump - ) + generic_response_output_item = GenericResponseOutputItem.model_validate(output_item_dump) if generic_response_output_item.content: content = generic_response_output_item.content except Exception: @@ -747,9 +705,7 @@ class OpenAIResponsesHandler(BaseTranslation): if not content: return - verbose_proxy_logger.debug( - "OpenAI Responses API: Processing output item: %s", output_item - ) + verbose_proxy_logger.debug("OpenAI Responses API: Processing output item: %s", output_item) # Iterate through content items (list of OutputText objects) for content_idx, content_item in enumerate(content): @@ -805,9 +761,7 @@ class OpenAIResponsesHandler(BaseTranslation): elif isinstance(output_item, BaseModel): # Handle other Pydantic models by converting to GenericResponseOutputItem try: - generic_item = GenericResponseOutputItem.model_validate( - output_item.model_dump() - ) + generic_item = GenericResponseOutputItem.model_validate(output_item.model_dump()) if generic_item.content and content_idx < len(generic_item.content): content_item = generic_item.content[content_idx] if isinstance(content_item, OutputText): diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index c18f2216f61..d107ca7a0d7 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -96,9 +96,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): temperature = params.get("temperature") if temperature is not None and temperature != 1: reasoning = params.get("reasoning") or {} - effort = ( - reasoning.get("effort") if isinstance(reasoning, dict) else None - ) + effort = reasoning.get("effort") if isinstance(reasoning, dict) else None supports_none = self._supports_reasoning_effort_none(model=model) if supports_none and (effort == "none" or effort is None): pass # flexible temperature allowed @@ -136,15 +134,11 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): input = self._validate_input_param(input) tools = response_api_optional_request_params.get("tools") - input, tools = self.remove_cache_control_flag_from_input_and_tools( - model=model, input=input, tools=tools - ) + input, tools = self.remove_cache_control_flag_from_input_and_tools(model=model, input=input, tools=tools) if tools is not None: response_api_optional_request_params["tools"] = tools final_request_params = dict( - ResponsesAPIRequestParams( - model=model, input=input, **response_api_optional_request_params - ) + ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params) ) return final_request_params @@ -181,9 +175,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return input, tools - def _validate_input_param( - self, input: Union[str, ResponseInputParam] - ) -> Union[str, ResponseInputParam]: + def _validate_input_param(self, input: Union[str, ResponseInputParam]) -> Union[str, ResponseInputParam]: """ Ensure all input fields if pydantic are converted to dict @@ -241,15 +233,12 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return dict_reasoning_item except Exception as e: - verbose_logger.debug( - f"Failed to create ResponseReasoningItem, falling back to manual filtering: {e}" - ) + verbose_logger.debug(f"Failed to create ResponseReasoningItem, falling back to manual filtering: {e}") # Fallback: manually filter out known None fields filtered_item = { k: v for k, v in item.items() - if v is not None - or k not in {"status", "content", "encrypted_content"} + if v is not None or k not in {"status", "content", "encrypted_content"} } return filtered_item return item @@ -267,21 +256,15 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): additional_args={"complete_input_dict": {}}, ) raw_response_json = raw_response.json() - raw_response_json["created_at"] = _safe_convert_created_field( - raw_response_json["created_at"] - ) + raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["created_at"]) except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) try: response = ResponsesAPIResponse(**raw_response_json) except Exception: - verbose_logger.debug( - f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" - ) + verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct") response = ResponsesAPIResponse.model_construct(**raw_response_json) # Store processed headers in additional_headers so they get returned to the client @@ -289,16 +272,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): response._hidden_params["headers"] = raw_response_headers return response - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() - api_key = ( - litellm_params.api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.setdefault("Content-Type", "application/json") headers["Authorization"] = f"Bearer {api_key}" return headers @@ -336,9 +312,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): # Convert the dictionary to a properly typed ResponsesAPIStreamingResponse verbose_logger.debug("Raw OpenAI Chunk=%s", parsed_chunk) event_type = str(parsed_chunk.get("type")) - event_pydantic_model = OpenAIResponsesAPIConfig.get_event_model_class( - event_type=event_type - ) + event_pydantic_model = OpenAIResponsesAPIConfig.get_event_model_class(event_type=event_type) # Some OpenAI-compatible providers send error.code: null; coalesce so validation succeeds. try: error_obj = parsed_chunk.get("error") @@ -353,8 +327,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return event_pydantic_model(**parsed_chunk) except ValidationError: verbose_logger.debug( - "Pydantic validation failed for %s with chunk %s, " - "falling back to model_construct", + "Pydantic validation failed for %s with chunk %s, falling back to model_construct", event_pydantic_model.__name__, parsed_chunk, ) @@ -438,9 +411,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ): return True except Exception as e: - verbose_logger.debug( - f"Error getting model info in OpenAIResponsesAPIConfig: {e}" - ) + verbose_logger.debug(f"Error getting model info in OpenAIResponsesAPIConfig: {e}") return False def supports_native_websocket(self) -> bool: @@ -463,9 +434,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - DELETE /v1/responses/{response_id} """ - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -481,9 +450,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): try: raw_response_json = raw_response.json() except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) return DeleteResponseResult(**raw_response_json) ######################################################### @@ -502,9 +469,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - GET /v1/responses/{response_id} """ - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -520,9 +485,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): try: raw_response_json = raw_response.json() except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) response = ResponsesAPIResponse(**raw_response_json) @@ -546,9 +509,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): limit: int = 20, order: Literal["asc", "desc"] = "desc", ) -> Tuple[str, Dict]: - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}/input_items" params: Dict[str, Any] = {} if after is not None: @@ -571,9 +532,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): try: return raw_response.json() except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) ######################################################### ########## CANCEL RESPONSE API TRANSFORMATION ########## @@ -591,9 +550,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - POST /v1/responses/{response_id}/cancel """ - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}/cancel" data: Dict = {} return url, data @@ -609,9 +566,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): try: raw_response_json = raw_response.json() except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) @@ -646,16 +601,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): input = self._validate_input_param(input) tools = response_api_optional_request_params.get("tools") - input, tools = self.remove_cache_control_flag_from_input_and_tools( - model=model, input=input, tools=tools - ) + input, tools = self.remove_cache_control_flag_from_input_and_tools(model=model, input=input, tools=tools) if tools is not None: response_api_optional_request_params["tools"] = tools - data = dict( - ResponsesAPIRequestParams( - model=model, input=input, **response_api_optional_request_params - ) - ) + data = dict(ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params)) return url, data @@ -673,22 +622,16 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): additional_args={"complete_input_dict": {}}, ) raw_response_json = raw_response.json() - raw_response_json["created_at"] = _safe_convert_created_field( - raw_response_json["created_at"] - ) + raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["created_at"]) except Exception: - raise OpenAIError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) try: response = ResponsesAPIResponse(**raw_response_json) except Exception: - verbose_logger.debug( - f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" - ) + verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct") response = ResponsesAPIResponse.model_construct(**raw_response_json) response._hidden_params["additional_headers"] = processed_headers diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py index f0c3149d0ae..3f29a8055d8 100644 --- a/litellm/llms/openai/speech/guardrail_translation/handler.py +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -45,9 +45,7 @@ class OpenAITextToSpeechHandler(BaseTranslation): """ input_text = data.get("input") if input_text is None: - verbose_proxy_logger.debug( - "OpenAI Text-to-Speech: No input text found in request data" - ) + verbose_proxy_logger.debug("OpenAI Text-to-Speech: No input text found in request data") return data if isinstance(input_text, str): @@ -66,8 +64,7 @@ class OpenAITextToSpeechHandler(BaseTranslation): data["input"] = guardrailed_texts[0] if guardrailed_texts else input_text verbose_proxy_logger.debug( - "OpenAI Text-to-Speech: Applied guardrail to input text. " - "Original length: %d, New length: %d", + "OpenAI Text-to-Speech: Applied guardrail to input text. Original length: %d, New length: %d", len(input_text), len(data["input"]), ) @@ -103,7 +100,6 @@ class OpenAITextToSpeechHandler(BaseTranslation): Unmodified response (audio data doesn't need text guardrails) """ verbose_proxy_logger.debug( - "OpenAI Text-to-Speech: Output processing not applicable " - "(output is audio data, not text)" + "OpenAI Text-to-Speech: Output processing not applicable (output is audio data, not text)" ) return response diff --git a/litellm/llms/openai/transcriptions/gpt_transformation.py b/litellm/llms/openai/transcriptions/gpt_transformation.py index 34621c44e22..56a1e39ecef 100644 --- a/litellm/llms/openai/transcriptions/gpt_transformation.py +++ b/litellm/llms/openai/transcriptions/gpt_transformation.py @@ -10,9 +10,7 @@ from .whisper_transformation import OpenAIWhisperAudioTranscriptionConfig class OpenAIGPTAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: """ Get the supported OpenAI params for the `gpt-4o-transcribe` models """ diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py index 92cf4398f05..fc1cae75b80 100644 --- a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py +++ b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py @@ -47,8 +47,7 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): Unmodified data (audio files don't need text guardrails) """ verbose_proxy_logger.debug( - "OpenAI Audio Transcription: Input processing not applicable " - "(input is audio file, not text)" + "OpenAI Audio Transcription: Input processing not applicable (input is audio file, not text)" ) return data @@ -73,9 +72,7 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): Modified response with guardrails applied to transcribed text """ if not hasattr(response, "text") or response.text is None: - verbose_proxy_logger.debug( - "OpenAI Audio Transcription: No text in response to process" - ) + verbose_proxy_logger.debug("OpenAI Audio Transcription: No text in response to process") return response if isinstance(response.text, str): @@ -90,9 +87,7 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): # Add user API key metadata with prefixed keys if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata diff --git a/litellm/llms/openai/transcriptions/handler.py b/litellm/llms/openai/transcriptions/handler.py index 44fb5da8590..76178051ca1 100644 --- a/litellm/llms/openai/transcriptions/handler.py +++ b/litellm/llms/openai/transcriptions/handler.py @@ -37,11 +37,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): - call openai_aclient.audio.transcriptions.create by default """ try: - raw_response = ( - await openai_aclient.audio.transcriptions.with_raw_response.create( - **data, timeout=timeout - ) - ) # type: ignore + raw_response = await openai_aclient.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) # type: ignore headers = dict(raw_response.headers) response = raw_response.parse() @@ -62,18 +58,12 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): """ try: if litellm.return_response_headers is True: - raw_response = ( - openai_client.audio.transcriptions.with_raw_response.create( - **data, timeout=timeout - ) - ) # type: ignore + raw_response = openai_client.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) # type: ignore headers = dict(raw_response.headers) response = raw_response.parse() return headers, response else: - response = openai_client.audio.transcriptions.create( - **data, timeout=timeout - ) # type: ignore + response = openai_client.audio.transcriptions.create(**data, timeout=timeout) # type: ignore return None, response except Exception as e: raise e diff --git a/litellm/llms/openai/transcriptions/whisper_transformation.py b/litellm/llms/openai/transcriptions/whisper_transformation.py index 2c01156fe05..ae7d0bb30b2 100644 --- a/litellm/llms/openai/transcriptions/whisper_transformation.py +++ b/litellm/llms/openai/transcriptions/whisper_transformation.py @@ -47,9 +47,7 @@ class OpenAIWhisperAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return api_base or "" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: """ Get the supported OpenAI params for the `whisper-1` models """ @@ -109,17 +107,13 @@ class OpenAIWhisperAudioTranscriptionConfig(BaseAudioTranscriptionConfig): data = {"model": model, "file": audio_file, **optional_params} if "response_format" not in data: - data["response_format"] = ( - "verbose_json" # ensures 'duration' is received - used for cost calculation - ) + data["response_format"] = "verbose_json" # ensures 'duration' is received - used for cost calculation return AudioTranscriptionRequestData( data=data, ) - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: return OpenAIError( status_code=status_code, message=error_message, @@ -138,10 +132,7 @@ class OpenAIWhisperAudioTranscriptionConfig(BaseAudioTranscriptionConfig): raise return TranscriptionResponse(text=raw_response.text) - if any( - key in raw_response_json - for key in TranscriptionResponse.model_fields.keys() - ): + if any(key in raw_response_json for key in TranscriptionResponse.model_fields.keys()): return TranscriptionResponse(**raw_response_json) else: raise ValueError( diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index 52202f57fd3..653a31f2e80 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -30,9 +30,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): ASSISTANTS_HEADER_KEY = "OpenAI-Beta" ASSISTANTS_HEADER_VALUE = "assistants=v2" - def get_auth_credentials( - self, litellm_params: Dict[str, Any] - ) -> VectorStoreFileAuthCredentials: + def get_auth_credentials(self, litellm_params: Dict[str, Any]) -> VectorStoreFileAuthCredentials: api_key = litellm_params.get("api_key") if api_key is None: raise ValueError("api_key is required") @@ -68,12 +66,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): litellm_params: Optional[GenericLiteLLMParams], ) -> Dict[str, str]: litellm_params = litellm_params or GenericLiteLLMParams() - api_key = ( - litellm_params.api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.update( { "Authorization": f"Bearer {api_key}", @@ -99,9 +92,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): or "https://api.openai.com/v1" ) base_url = base_url.rstrip("/") - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") return f"{base_url}/vector_stores/{encoded_vector_store_id}/files" def transform_create_vector_store_file_request( diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index bd095a0a1b7..6ccf8e271e5 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -31,9 +31,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): ASSISTANTS_HEADER_KEY = "OpenAI-Beta" ASSISTANTS_HEADER_VALUE = "assistants=v2" - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: api_key = litellm_params.get("api_key") if api_key is None: raise ValueError("api_key is required") @@ -49,16 +47,9 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): "write": [("POST", "/vector_stores")], } - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() - api_key = ( - litellm_params.api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.update( { "Authorization": f"Bearer {api_key}", @@ -109,22 +100,14 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}/search" typed_request_body = VectorStoreSearchRequest( query=query, filters=vector_store_search_optional_params.get("filters", None), - max_num_results=vector_store_search_optional_params.get( - "max_num_results", None - ), - ranking_options=vector_store_search_optional_params.get( - "ranking_options", None - ), - rewrite_query=vector_store_search_optional_params.get( - "rewrite_query", None - ), + max_num_results=vector_store_search_optional_params.get("max_num_results", None), + ranking_options=vector_store_search_optional_params.get("ranking_options", None), + rewrite_query=vector_store_search_optional_params.get("rewrite_query", None), ) dict_request_body = cast(dict, typed_request_body) @@ -155,21 +138,15 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): typed_request_body = VectorStoreCreateRequest( name=vector_store_create_optional_params.get("name", None), file_ids=vector_store_create_optional_params.get("file_ids", None), - expires_after=vector_store_create_optional_params.get( - "expires_after", None - ), - chunking_strategy=vector_store_create_optional_params.get( - "chunking_strategy", None - ), + expires_after=vector_store_create_optional_params.get("expires_after", None), + chunking_strategy=vector_store_create_optional_params.get("chunking_strategy", None), metadata=metadata_payload, ) dict_request_body = cast(dict, typed_request_body) return url, dict_request_body - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: try: response_json = response.json() return VectorStoreCreateResponse(**response_json) diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 520a42e9dd1..684601367b6 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -79,12 +79,7 @@ class OpenAIVideoConfig(BaseVideoConfig): if litellm_params and litellm_params.api_key: api_key = api_key or litellm_params.api_key - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.update( { "Authorization": f"Bearer {api_key}", @@ -126,17 +121,13 @@ class OpenAIVideoConfig(BaseVideoConfig): } # Create the request data - video_create_request = CreateVideoRequest( - model=model, prompt=prompt, **video_create_optional_request_params - ) + video_create_request = CreateVideoRequest(model=model, prompt=prompt, **video_create_optional_request_params) request_dict = cast(Dict, video_create_request) request_dict = self._decode_character_ids_in_create_video_request(request_dict) # Handle input_reference parameter if provided _input_reference = video_create_optional_request_params.get("input_reference") - data_without_files = { - k: v for k, v in request_dict.items() if k not in ["input_reference"] - } + data_without_files = {k: v for k, v in request_dict.items() if k not in ["input_reference"]} files_list: List[Tuple[str, Any]] = [] # Handle input_reference parameter @@ -191,9 +182,7 @@ class OpenAIVideoConfig(BaseVideoConfig): video_obj = VideoObject(**response_data) # type: ignore[arg-type] if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider( - video_obj.id, custom_llm_provider, model - ) + video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) usage_data = {} if video_obj: @@ -222,9 +211,7 @@ class OpenAIVideoConfig(BaseVideoConfig): - GET /v1/videos/{video_id}/content?variant=thumbnail """ original_video_id = extract_original_video_id(video_id) - encoded_video_id = encode_url_path_segment( - original_video_id, field_name="video_id" - ) + encoded_video_id = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for video content download url = f"{api_base.rstrip('/')}/{encoded_video_id}/content" @@ -256,9 +243,7 @@ class OpenAIVideoConfig(BaseVideoConfig): - POST /v1/videos/{video_id}/remix """ original_video_id = extract_original_video_id(video_id) - encoded_video_id = encode_url_path_segment( - original_video_id, field_name="video_id" - ) + encoded_video_id = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for video remix url = f"{api_base.rstrip('/')}/{encoded_video_id}/remix" @@ -295,9 +280,7 @@ class OpenAIVideoConfig(BaseVideoConfig): video_obj = VideoObject(**response_data) # type: ignore[arg-type] if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider( - video_obj.id, custom_llm_provider, None - ) + video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) # Create usage object with duration information for cost calculation # Video remix API doesn't provide usage, so we create one with duration @@ -403,9 +386,7 @@ class OpenAIVideoConfig(BaseVideoConfig): - DELETE /v1/videos/{video_id} """ original_video_id = extract_original_video_id(video_id) - encoded_video_id = encode_url_path_segment( - original_video_id, field_name="video_id" - ) + encoded_video_id = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for video delete url = f"{api_base.rstrip('/')}/{encoded_video_id}" @@ -442,9 +423,7 @@ class OpenAIVideoConfig(BaseVideoConfig): """ # Extract the original video_id (remove provider encoding if present) original_video_id = extract_original_video_id(video_id) - encoded_video_id = encode_url_path_segment( - original_video_id, field_name="video_id" - ) + encoded_video_id = encode_url_path_segment(original_video_id, field_name="video_id") # For video retrieve, we just need to construct the URL url = f"{api_base.rstrip('/')}/{encoded_video_id}" @@ -468,9 +447,7 @@ class OpenAIVideoConfig(BaseVideoConfig): video_obj = VideoObject(**response_data) # type: ignore[arg-type] if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider( - video_obj.id, custom_llm_provider, None - ) + video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj @@ -513,9 +490,7 @@ class OpenAIVideoConfig(BaseVideoConfig): headers: dict, ) -> Tuple[str, Dict]: original_character_id = extract_original_character_id(character_id) - encoded_character_id = encode_url_path_segment( - original_character_id, field_name="character_id" - ) + encoded_character_id = encode_url_path_segment(original_character_id, field_name="character_id") url = f"{api_base.rstrip('/')}/characters/{encoded_character_id}" return url, {} @@ -552,9 +527,7 @@ class OpenAIVideoConfig(BaseVideoConfig): ) -> VideoObject: video_obj = VideoObject(**raw_response.json()) if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider( - video_obj.id, custom_llm_provider, None - ) + video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj def transform_video_extension_request( @@ -586,9 +559,7 @@ class OpenAIVideoConfig(BaseVideoConfig): ) -> VideoObject: video_obj = VideoObject(**raw_response.json()) if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider( - video_obj.id, custom_llm_provider, None - ) + video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj def _add_image_to_files( @@ -603,9 +574,7 @@ class OpenAIVideoConfig(BaseVideoConfig): if isinstance(image, BufferedReader): files_list.append((field_name, (image.name, image, image_content_type))) else: - files_list.append( - (field_name, ("input_reference.png", image, image_content_type)) - ) + files_list.append((field_name, ("input_reference.png", image, image_content_type))) def _add_video_to_files( self, diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 821fc9b7f15..0da0f3d90f0 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -37,21 +37,15 @@ async def make_call( if client is None: client = litellm.module_level_aclient - response = await client.post( - api_base, headers=headers, data=data, stream=not fake_stream - ) + response = await client.post(api_base, headers=headers, data=data, stream=not fake_stream) if streaming_decoder is not None: - completion_stream: Any = streaming_decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) - ) + completion_stream: Any = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) elif fake_stream: model_response = ModelResponse(**response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator( - streaming_response=response.aiter_lines(), sync_stream=False - ) + completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False) # LOGGING logging_obj.post_call( input=messages, @@ -78,24 +72,18 @@ def make_sync_call( if client is None: client = litellm.module_level_client # Create a new client if none provided - response = client.post( - api_base, headers=headers, data=data, stream=not fake_stream, timeout=timeout - ) + response = client.post(api_base, headers=headers, data=data, stream=not fake_stream, timeout=timeout) if response.status_code != 200: raise OpenAILikeError(status_code=response.status_code, message=response.read()) if streaming_decoder is not None: - completion_stream = streaming_decoder.iter_bytes( - response.iter_bytes(chunk_size=1024) - ) + completion_stream = streaming_decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) elif fake_stream: model_response = ModelResponse(**response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator( - streaming_response=response.iter_lines(), sync_stream=True - ) + completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True) # LOGGING logging_obj.post_call( @@ -184,9 +172,7 @@ class OpenAILikeChatHandler(OpenAILikeBase): client = litellm.module_level_aclient try: - response = await client.post( - api_base, headers=headers, data=json.dumps(data), timeout=timeout - ) + response = await client.post(api_base, headers=headers, data=json.dumps(data), timeout=timeout) response.raise_for_status() except httpx.HTTPStatusError as e: raise OpenAILikeError( @@ -241,9 +227,7 @@ class OpenAILikeChatHandler(OpenAILikeBase): ] = None, # if openai-compatible api needs custom stream decoder - e.g. sagemaker fake_stream: bool = False, ): - custom_endpoint = custom_endpoint or optional_params.pop( - "custom_endpoint", None - ) + custom_endpoint = custom_endpoint or optional_params.pop("custom_endpoint", None) base_model: Optional[str] = optional_params.pop("base_model", None) api_base, headers = self._validate_environment( api_base=api_base, @@ -264,12 +248,8 @@ class OpenAILikeChatHandler(OpenAILikeBase): provider_config = ProviderConfigManager.get_provider_chat_config( model=model, provider=LlmProviders(custom_llm_provider) ) - if isinstance(provider_config, OpenAIGPTConfig) or isinstance( - provider_config, OpenAIConfig - ): - messages = provider_config._transform_messages( - messages=messages, model=model - ) + if isinstance(provider_config, OpenAIGPTConfig) or isinstance(provider_config, OpenAIConfig): + messages = provider_config._transform_messages(messages=messages, model=model) data = { "model": model, @@ -343,11 +323,7 @@ class OpenAILikeChatHandler(OpenAILikeBase): ## COMPLETION CALL if stream is True: completion_stream = make_sync_call( - client=( - client - if client is not None and isinstance(client, HTTPHandler) - else None - ), + client=(client if client is not None and isinstance(client, HTTPHandler) else None), api_base=api_base, headers=headers, data=json.dumps(data), @@ -369,9 +345,7 @@ class OpenAILikeChatHandler(OpenAILikeBase): if client is None or not isinstance(client, HTTPHandler): client = HTTPHandler(timeout=timeout) # type: ignore try: - response = client.post( - url=api_base, headers=headers, data=json.dumps(data) - ) + response = client.post(url=api_base, headers=headers, data=json.dumps(data)) response.raise_for_status() except httpx.HTTPStatusError as e: @@ -380,9 +354,7 @@ class OpenAILikeChatHandler(OpenAILikeBase): message=e.response.text, ) except httpx.TimeoutException: - raise OpenAILikeError( - status_code=408, message="Timeout error occurred." - ) + raise OpenAILikeError(status_code=408, message="Timeout error occurred.") except Exception as e: raise OpenAILikeError(status_code=500, message=str(e)) return OpenAILikeChatConfig._transform_response( diff --git a/litellm/llms/openai_like/chat/transformation.py b/litellm/llms/openai_like/chat/transformation.py index 1c8cd574c01..a2c847a410f 100644 --- a/litellm/llms/openai_like/chat/transformation.py +++ b/litellm/llms/openai_like/chat/transformation.py @@ -27,9 +27,7 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): api_key: Optional[str], ) -> Tuple[Optional[str], Optional[str]]: api_base = api_base or get_secret_str("OPENAI_LIKE_API_BASE") # type: ignore - dynamic_api_key = ( - api_key or get_secret_str("OPENAI_LIKE_API_KEY") or "" - ) # vllm does not require an api key + dynamic_api_key = api_key or get_secret_str("OPENAI_LIKE_API_KEY") or "" # vllm does not require an api key return api_base, dynamic_api_key @staticmethod @@ -107,19 +105,15 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): if json_mode: for choice in response_json["choices"]: - message = ( - OpenAILikeChatConfig._json_mode_convert_tool_response_to_message( - choice.get("message"), json_mode - ) + message = OpenAILikeChatConfig._json_mode_convert_tool_response_to_message( + choice.get("message"), json_mode ) choice["message"] = message returned_response = ModelResponse(**response_json) if custom_llm_provider is not None: - returned_response.model = ( - custom_llm_provider + "/" + (returned_response.model or "") - ) + returned_response.model = custom_llm_provider + "/" + (returned_response.model or "") if base_model is not None: returned_response._hidden_params["model"] = base_model @@ -164,13 +158,8 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): drop_params: bool, replace_max_completion_tokens_with_max_tokens: bool = True, ) -> dict: - mapped_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) - if ( - "max_completion_tokens" in non_default_params - and replace_max_completion_tokens_with_max_tokens - ): + mapped_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) + if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: mapped_params["max_tokens"] = non_default_params[ "max_completion_tokens" ] # most openai-compatible providers support 'max_tokens' not 'max_completion_tokens' diff --git a/litellm/llms/openai_like/common_utils.py b/litellm/llms/openai_like/common_utils.py index 116277b6dd3..40f2e5c3f5c 100644 --- a/litellm/llms/openai_like/common_utils.py +++ b/litellm/llms/openai_like/common_utils.py @@ -9,9 +9,7 @@ class OpenAILikeError(Exception): self.message = message self.request = httpx.Request(method="POST", url="https://www.litellm.ai") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class OpenAILikeBase: diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index 9ed9734edae..3c763ed9b9b 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -20,9 +20,7 @@ def create_config_class(provider: SimpleProviderConfig): """Generate config class dynamically from JSON configuration""" # Choose base class - base_class: type = ( - OpenAIGPTConfig if provider.base_class == "openai_gpt" else OpenAILikeChatConfig - ) + base_class: type = OpenAIGPTConfig if provider.base_class == "openai_gpt" else OpenAILikeChatConfig class JSONProviderConfig(base_class): # type: ignore[valid-type,misc] @overload @@ -48,13 +46,9 @@ def create_config_class(provider: SimpleProviderConfig): messages = handle_messages_with_content_list_to_str_conversion(messages) if is_async: - return super()._transform_messages( - messages=messages, model=model, is_async=True - ) + return super()._transform_messages(messages=messages, model=model, is_async=True) else: - return super()._transform_messages( - messages=messages, model=model, is_async=False - ) + return super()._transform_messages(messages=messages, model=model, is_async=False) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] @@ -101,9 +95,7 @@ def create_config_class(provider: SimpleProviderConfig): supported_params = super().get_supported_openai_params(model=model) - _supports_fc = supports_function_calling( - model=model, custom_llm_provider=provider.slug - ) + _supports_fc = supports_function_calling(model=model, custom_llm_provider=provider.slug) if not _supports_fc: tool_params = [ diff --git a/litellm/llms/openai_like/embedding/handler.py b/litellm/llms/openai_like/embedding/handler.py index 285595d2791..52eafc05b2c 100644 --- a/litellm/llms/openai_like/embedding/handler.py +++ b/litellm/llms/openai_like/embedding/handler.py @@ -59,9 +59,7 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): message=e.response.text if e.response else str(e), ) except httpx.TimeoutException: - raise OpenAILikeError( - status_code=408, message="Timeout error occurred." - ) + raise OpenAILikeError(status_code=408, message="Timeout error occurred.") except Exception as e: raise OpenAILikeError(status_code=500, message=str(e)) @@ -105,9 +103,7 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): custom_endpoint=custom_endpoint, ) model = model - filtered_optional_params = { - k: v for k, v in optional_params.items() if v not in (None, "") - } + filtered_optional_params = {k: v for k, v in optional_params.items() if v not in (None, "")} data = {"model": model, "input": input, **filtered_optional_params} ## LOGGING diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index c6ff0f7a394..4640bb8a422 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -52,9 +52,7 @@ class JSONProviderRegistry: cls._loaded = True except Exception as e: - verbose_logger.warning( - f"Warning: Failed to load JSON provider configs: {e}" - ) + verbose_logger.warning(f"Warning: Failed to load JSON provider configs: {e}") cls._loaded = True @classmethod diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index 107d5c25e6d..ca287f5de04 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -39,9 +39,9 @@ class OpenrouterConfig(OpenAIGPTConfig): """ supported_params = super().get_supported_openai_params(model=model) try: - if litellm.supports_reasoning( - model=model, custom_llm_provider="openrouter" - ) or litellm.supports_reasoning(model=model): + if litellm.supports_reasoning(model=model, custom_llm_provider="openrouter") or litellm.supports_reasoning( + model=model + ): supported_params.append("reasoning_effort") supported_params.append("thinking") except Exception: @@ -59,9 +59,7 @@ class OpenrouterConfig(OpenAIGPTConfig): if non_default_params.get("reasoning_effort") == "max": non_default_params = {**non_default_params, "reasoning_effort": "xhigh"} - mapped_openai_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + mapped_openai_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) # OpenRouter-only parameters extra_body = {} @@ -74,9 +72,7 @@ class OpenrouterConfig(OpenAIGPTConfig): extra_body["models"] = models if route is not None: extra_body["route"] = route - mapped_openai_params["extra_body"] = ( - extra_body # openai client supports `extra_body` param - ) + mapped_openai_params["extra_body"] = extra_body # openai client supports `extra_body` param return mapped_openai_params def _supports_cache_control_in_content(self, model: str) -> bool: @@ -87,10 +83,7 @@ class OpenrouterConfig(OpenAIGPTConfig): bool: True if model supports cache_control (Claude or Gemini models) """ model_lower = model.lower() - return any( - supported_model.value in model_lower - for supported_model in CacheControlSupportedModels - ) + return any(supported_model.value in model_lower for supported_model in CacheControlSupportedModels) def remove_cache_control_flag_from_messages_and_tools( self, @@ -101,13 +94,9 @@ class OpenrouterConfig(OpenAIGPTConfig): if self._supports_cache_control_in_content(model): return messages, tools else: - return super().remove_cache_control_flag_from_messages_and_tools( - model, messages, tools - ) + return super().remove_cache_control_flag_from_messages_and_tools(model, messages, tools) - def _move_cache_control_to_content( - self, messages: List[AllMessageValues] - ) -> List[AllMessageValues]: + def _move_cache_control_to_content(self, messages: List[AllMessageValues]) -> List[AllMessageValues]: """ Move cache_control from message level to content blocks. OpenRouter requires cache_control to be inside content blocks, not at message level. @@ -167,9 +156,7 @@ class OpenrouterConfig(OpenAIGPTConfig): messages = self._move_cache_control_to_content(messages) extra_body = optional_params.pop("extra_body", {}) - response = super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + response = super().transform_request(model, messages, optional_params, litellm_params, headers) response.update(extra_body) # ALWAYS add usage parameter to get cost data from OpenRouter @@ -228,9 +215,9 @@ class OpenrouterConfig(OpenAIGPTConfig): model_response._hidden_params = {} if "additional_headers" not in model_response._hidden_params: model_response._hidden_params["additional_headers"] = {} - model_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(response_cost) + model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float( + response_cost + ) except Exception: # If we can't extract cost, continue without it - don't fail the response pass diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index 8b836e8e5d2..c6c3df083a1 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -170,9 +170,7 @@ class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Any - ) -> Any: + def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any: """ Get the error class for OpenRouter errors. """ diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index 0d96b62425f..f4531932f96 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -97,9 +97,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): if key == "size": if "image_config" not in mapped_params: mapped_params["image_config"] = {} - mapped_params["image_config"]["aspect_ratio"] = ( - self._map_size_to_aspect_ratio(cast(str, value)) - ) + mapped_params["image_config"]["aspect_ratio"] = self._map_size_to_aspect_ratio(cast(str, value)) elif key == "quality": image_size = self._map_quality_to_image_size(cast(str, value)) if image_size: @@ -139,11 +137,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): api_base: Optional[str], litellm_params: dict, ) -> str: - base_url = ( - api_base - or get_secret_str("OPENROUTER_API_BASE") - or "https://openrouter.ai/api/v1" - ) + base_url = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" base_url = base_url.rstrip("/") if not base_url.endswith("/chat/completions"): return f"{base_url}/chat/completions" @@ -344,17 +338,15 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): model_response._hidden_params = {} if "additional_headers" not in model_response._hidden_params: model_response._hidden_params["additional_headers"] = {} - model_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(cost) + model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float( + cost + ) cost_details = usage_data.get("cost_details", {}) if cost_details: if "response_cost_details" not in model_response._hidden_params: model_response._hidden_params["response_cost_details"] = {} - model_response._hidden_params["response_cost_details"].update( - cost_details - ) + model_response._hidden_params["response_cost_details"].update(cost_details) model_response._hidden_params["model"] = response_json.get("model", model) diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 9c2293eb3f1..eabb76f00c0 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -64,9 +64,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): and extract images from chat responses. """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for OpenRouter image generation. @@ -224,17 +222,15 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): model_response._hidden_params = {} if "additional_headers" not in model_response._hidden_params: model_response._hidden_params["additional_headers"] = {} - model_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(cost) + model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float( + cost + ) cost_details = usage_data.get("cost_details", {}) if cost_details: if "response_cost_details" not in model_response._hidden_params: model_response._hidden_params["response_cost_details"] = {} - model_response._hidden_params["response_cost_details"].update( - cost_details - ) + model_response._hidden_params["response_cost_details"].update(cost_details) model_response._hidden_params["model"] = response_json.get("model", model) diff --git a/litellm/llms/openrouter/responses/transformation.py b/litellm/llms/openrouter/responses/transformation.py index 864e1549274..217a419ed22 100644 --- a/litellm/llms/openrouter/responses/transformation.py +++ b/litellm/llms/openrouter/responses/transformation.py @@ -49,8 +49,7 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig): if not api_key: raise ValueError( - "OpenRouter API key is required. Set OPENROUTER_API_KEY " - "environment variable or pass api_key parameter." + "OpenRouter API key is required. Set OPENROUTER_API_KEY environment variable or pass api_key parameter." ) headers.update( @@ -66,10 +65,7 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: dict, ) -> str: api_base = ( - api_base - or litellm.api_base - or get_secret_str("OPENROUTER_API_BASE") - or "https://openrouter.ai/api/v1" + api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" ) api_base = api_base.rstrip("/") diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py index dc9f8440d30..60266c988df 100644 --- a/litellm/llms/opensandbox/sandbox/transformation.py +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -71,12 +71,8 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): ) -> ContainerHandle: key = self.validate_environment(api_key=api_key) base = self._api_base(api_base) - ready_timeout_seconds = ( - float(ready_timeout) if ready_timeout is not None else DEFAULT_READY_TIMEOUT - ) - poll_interval_seconds = ( - float(poll_interval) if poll_interval is not None else DEFAULT_POLL_INTERVAL - ) + ready_timeout_seconds = float(ready_timeout) if ready_timeout is not None else DEFAULT_READY_TIMEOUT + poll_interval_seconds = float(poll_interval) if poll_interval is not None else DEFAULT_POLL_INTERVAL body = self._create_body( template=template, timeout=timeout, @@ -150,25 +146,13 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): api_key=api_key, api_base=api_base, use_server_proxy=use_server_proxy, - ready_timeout=( - float(ready_timeout) - if ready_timeout is not None - else DEFAULT_READY_TIMEOUT - ), - poll_interval=( - float(poll_interval) - if poll_interval is not None - else DEFAULT_POLL_INTERVAL - ), + ready_timeout=(float(ready_timeout) if ready_timeout is not None else DEFAULT_READY_TIMEOUT), + poll_interval=(float(poll_interval) if poll_interval is not None else DEFAULT_POLL_INTERVAL), client=client, ) endpoint = str(handle._hidden_params["execd_endpoint"]) endpoint_headers = self._as_str_dict(handle._hidden_params.get("execd_headers")) - base = str( - handle._hidden_params.get("api_base") - or handle.domain - or self._api_base(api_base) - ) + base = str(handle._hidden_params.get("api_base") or handle.domain or self._api_base(api_base)) lines = await self._post_code( url=f"{self._endpoint_base_url(endpoint, base)}/code", headers={ @@ -228,9 +212,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): base = str(handle._hidden_params.get("api_base") or self._api_base(api_base)) key = self._api_key(api_key=api_key, handle=handle) - resolved_use_server_proxy = bool( - handle._hidden_params.get("use_server_proxy", use_server_proxy) - ) + resolved_use_server_proxy = bool(handle._hidden_params.get("use_server_proxy", use_server_proxy)) endpoint, endpoint_headers = await self._wait_for_execd_endpoint( sandbox_id=handle.id, api_base=base, @@ -277,10 +259,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): if state in {"Failed", "Stopping", "Terminated"}: raise ValueError(f"OpenSandbox sandbox {sandbox_id} entered {state}") if time.monotonic() >= deadline: - raise TimeoutError( - f"OpenSandbox sandbox {sandbox_id} was not Running within " - f"{ready_timeout} seconds" - ) + raise TimeoutError(f"OpenSandbox sandbox {sandbox_id} was not Running within {ready_timeout} seconds") await asyncio.sleep(poll_interval) async def _wait_for_execd_endpoint( @@ -314,8 +293,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): if time.monotonic() >= deadline: raise TimeoutError( - f"OpenSandbox execd endpoint for {sandbox_id} was not ready within " - f"{ready_timeout} seconds" + f"OpenSandbox execd endpoint for {sandbox_id} was not ready within {ready_timeout} seconds" ) from last_error await asyncio.sleep(poll_interval) @@ -339,9 +317,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): data = response.json() endpoint = data.get("endpoint") if not endpoint: - raise ValueError( - f"OpenSandbox did not return an execd endpoint for {sandbox_id}" - ) + raise ValueError(f"OpenSandbox did not return an execd endpoint for {sandbox_id}") return str(endpoint), self._as_str_dict(data.get("headers")) async def _post_code( @@ -390,8 +366,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): "image": {"uri": template or OPEN_SANDBOX_DEFAULT_TEMPLATE}, "entrypoint": list(entrypoint or OPEN_SANDBOX_DEFAULT_ENTRYPOINT), "timeout": timeout if timeout is not None else DEFAULT_SANDBOX_TIMEOUT, - "resourceLimits": resource_limits - or OpenSandboxSandboxConfig._default_resource_limits(), + "resourceLimits": resource_limits or OpenSandboxSandboxConfig._default_resource_limits(), } if metadata: body["metadata"] = metadata @@ -434,10 +409,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): def _api_base(api_base: str | None) -> str: base = api_base or get_secret_str(OPEN_SANDBOX_API_BASE_ENV_VAR) if not base: - raise ValueError( - "OpenSandbox api_base is required. Pass api_base or set " - f"{OPEN_SANDBOX_API_BASE_ENV_VAR}." - ) + raise ValueError(f"OpenSandbox api_base is required. Pass api_base or set {OPEN_SANDBOX_API_BASE_ENV_VAR}.") return str(base).rstrip("/") @staticmethod @@ -456,9 +428,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): return f"{protocol}://{normalized_endpoint}" @staticmethod - def _as_handle( - container: Union[ContainerHandle, str], *, api_base: str | None - ) -> ContainerHandle: + def _as_handle(container: Union[ContainerHandle, str], *, api_base: str | None) -> ContainerHandle: if isinstance(container, ContainerHandle): return container handle = ContainerHandle( @@ -472,9 +442,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): @staticmethod def _parse_lines(lines: list[str]) -> CodeExecutionResult: messages = tuple( - event - for line in lines - if (event := OpenSandboxSandboxConfig._parse_sse_line(line)) is not None + event for line in lines if (event := OpenSandboxSandboxConfig._parse_sse_line(line)) is not None ) def of_type(message_type: str): @@ -488,8 +456,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): ( OpenSandboxSandboxConfig._as_int(m.get("execution_count")) for m in of_type("execution_count") - if OpenSandboxSandboxConfig._as_int(m.get("execution_count")) - is not None + if OpenSandboxSandboxConfig._as_int(m.get("execution_count")) is not None ), None, ) @@ -497,9 +464,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): return CodeExecutionResult( stdout="".join(str(m.get("text", "")) for m in of_type("stdout")), stderr="".join(str(m.get("text", "")) for m in of_type("stderr")), - results=[ - OpenSandboxSandboxConfig._normalize_result(m) for m in of_type("result") - ], + results=[OpenSandboxSandboxConfig._normalize_result(m) for m in of_type("result")], error=error, execution_count=execution_count, ) @@ -541,40 +506,24 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): results = message.get("results") if isinstance(results, dict): return {str(k): v for k, v in results.items()} - return { - str(k): v - for k, v in message.items() - if k not in {"type", "timestamp", "execution_count"} - } + return {str(k): v for k, v in message.items() if k not in {"type", "timestamp", "execution_count"}} @staticmethod def _normalize_error(message: dict[str, object]) -> dict[str, object]: raw_error = message.get("error") if isinstance(raw_error, dict): - name = OpenSandboxSandboxConfig._first_non_none_value( - raw_error, "ename", "name", default="" - ) - value = OpenSandboxSandboxConfig._first_non_none_value( - raw_error, "evalue", "value", default="" - ) - traceback = OpenSandboxSandboxConfig._first_non_none_value( - raw_error, "traceback", default=[] - ) + name = OpenSandboxSandboxConfig._first_non_none_value(raw_error, "ename", "name", default="") + value = OpenSandboxSandboxConfig._first_non_none_value(raw_error, "evalue", "value", default="") + traceback = OpenSandboxSandboxConfig._first_non_none_value(raw_error, "traceback", default=[]) return { "name": name, "value": value, "traceback": traceback, } return { - "name": OpenSandboxSandboxConfig._first_non_none_value( - message, "name", default="" - ), - "value": OpenSandboxSandboxConfig._first_non_none_value( - message, "value", "text", default="" - ), - "traceback": OpenSandboxSandboxConfig._first_non_none_value( - message, "traceback", default=[] - ), + "name": OpenSandboxSandboxConfig._first_non_none_value(message, "name", default=""), + "value": OpenSandboxSandboxConfig._first_non_none_value(message, "value", "text", default=""), + "traceback": OpenSandboxSandboxConfig._first_non_none_value(message, "traceback", default=[]), } @staticmethod @@ -589,9 +538,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): return None @staticmethod - def _first_non_none_value( - values: dict[str, object], *keys: str, default: object - ) -> object: + def _first_non_none_value(values: dict[str, object], *keys: str, default: object) -> object: return next( (values[key] for key in keys if key in values and values[key] is not None), default, diff --git a/litellm/llms/ovhcloud/audio_transcription/transformation.py b/litellm/llms/ovhcloud/audio_transcription/transformation.py index f49f31d7ecd..43b68c6503d 100644 --- a/litellm/llms/ovhcloud/audio_transcription/transformation.py +++ b/litellm/llms/ovhcloud/audio_transcription/transformation.py @@ -26,9 +26,7 @@ from ..utils import OVHCloudException class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: # OVHCloud implements the OpenAI-compatible Whisper interface. # We pass through the same optional params as the OpenAI Whisper API. return [ @@ -61,11 +59,7 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = ( - "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" - if api_base is None - else api_base.rstrip("/") - ) + api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") complete_url = f"{api_base}/audio/transcriptions" return complete_url diff --git a/litellm/llms/ovhcloud/chat/transformation.py b/litellm/llms/ovhcloud/chat/transformation.py index 62f51f1e9da..0090ae168f7 100644 --- a/litellm/llms/ovhcloud/chat/transformation.py +++ b/litellm/llms/ovhcloud/chat/transformation.py @@ -31,11 +31,7 @@ class OVHCloudChatConfig(OpenAIGPTConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = ( - "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" - if api_base is None - else api_base.rstrip("/") - ) + api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") complete_url = f"{api_base}/chat/completions" return complete_url @@ -55,9 +51,7 @@ class OVHCloudChatConfig(OpenAIGPTConfig): model: str, drop_params: bool, ) -> dict: - mapped_openai_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + mapped_openai_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) return mapped_openai_params def transform_request( @@ -69,9 +63,7 @@ class OVHCloudChatConfig(OpenAIGPTConfig): headers: dict, ) -> dict: extra_body = optional_params.pop("extra_body", {}) - response = super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + response = super().transform_request(model, messages, optional_params, litellm_params, headers) response.update(extra_body) return response @@ -88,9 +80,7 @@ class OVHCloudChatCompletionStreamingHandler(BaseModelResponseIterator): try: if "error" in chunk: error_chunk = chunk["error"] - error_message = "OVHCloud Error: {}".format( - error_chunk.get("message", "Unknown error") - ) + error_message = "OVHCloud Error: {}".format(error_chunk.get("message", "Unknown error")) raise OVHCloudException( message=error_message, status_code=error_chunk.get("code", 400), diff --git a/litellm/llms/ovhcloud/embedding/transformation.py b/litellm/llms/ovhcloud/embedding/transformation.py index 6b5c43e2d06..006f2a2349b 100644 --- a/litellm/llms/ovhcloud/embedding/transformation.py +++ b/litellm/llms/ovhcloud/embedding/transformation.py @@ -30,11 +30,7 @@ class OVHCloudEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = ( - "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" - if api_base is None - else api_base.rstrip("/") - ) + api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") complete_url = f"{api_base}/embeddings" return complete_url @@ -122,6 +118,4 @@ class OVHCloudEmbeddingConfig(BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return OVHCloudException( - message=error_message, status_code=status_code, headers=headers - ) + return OVHCloudException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index 35a0d84df40..56566aea0b1 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -75,9 +75,7 @@ class ParallelAISearchConfig(BaseSearchConfig): default_api_base=self.PARALLEL_AI_API_BASE, ) if not api_key: - raise ValueError( - "PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable." - ) + raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.") headers["x-api-key"] = api_key headers["Content-Type"] = "application/json" return headers @@ -89,11 +87,7 @@ class ParallelAISearchConfig(BaseSearchConfig): data: Optional[Union[Dict, List[Dict]]] = None, **kwargs, ) -> str: - api_base = ( - api_base - or get_secret_str("PARALLEL_AI_API_BASE") - or self.PARALLEL_AI_API_BASE - ) + api_base = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE api_base = api_base.rstrip("/") if not api_base.endswith("/v1/search"): @@ -155,9 +149,7 @@ class ParallelAISearchConfig(BaseSearchConfig): advanced_settings["location"] = params.pop("country") if "max_chars_per_result" in params: - advanced_settings["excerpt_settings"] = { - "max_chars_per_result": params.pop("max_chars_per_result") - } + advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")} source_policy: _ParallelAISourcePolicy = {} diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index db8d519d9be..8ca600b0bcf 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -43,15 +43,11 @@ class PassThroughEndpointHandler(BaseTranslation): if litellm_logging_obj is None: return None - passthrough_config = getattr( - litellm_logging_obj, "passthrough_guardrails_config", None - ) + passthrough_config = getattr(litellm_logging_obj, "passthrough_guardrails_config", None) if not passthrough_config or not guardrail_name: return None - return PassthroughGuardrailHandler.get_settings( - passthrough_config, guardrail_name - ) + return PassthroughGuardrailHandler.get_settings(passthrough_config, guardrail_name) def _extract_text_for_guardrail( self, @@ -83,13 +79,9 @@ class PassThroughEndpointHandler(BaseTranslation): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps payload_to_check = { - k: v - for k, v in data.items() - if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj") + k: v for k, v in data.items() if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj") } - verbose_proxy_logger.debug( - "PassThroughEndpointHandler: Using full payload for guardrail" - ) + verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail") return safe_dumps(payload_to_check) async def process_input_messages( @@ -115,9 +107,7 @@ class PassThroughEndpointHandler(BaseTranslation): text_to_check = self._extract_text_for_guardrail(data, field_expressions) if not text_to_check: - verbose_proxy_logger.debug( - "PassThroughEndpointHandler: No text to check, skipping guardrail" - ) + verbose_proxy_logger.debug("PassThroughEndpointHandler: No text to check, skipping guardrail") return data # Apply guardrail (pass-through doesn't modify the text, just checks it) @@ -153,9 +143,7 @@ class PassThroughEndpointHandler(BaseTranslation): user_api_key_dict: User API key metadata to pass to guardrails """ if not isinstance(response, dict): - verbose_proxy_logger.debug( - "PassThroughEndpointHandler: Response is not a dict, skipping" - ) + verbose_proxy_logger.debug("PassThroughEndpointHandler: Response is not a dict, skipping") return response guardrail_name = guardrail_to_apply.guardrail_name @@ -177,22 +165,14 @@ class PassThroughEndpointHandler(BaseTranslation): # Use the real request_data if provided (proxy path), otherwise # create a standalone dict (SDK / direct-call path). if request_data is None: - request_data = ( - {"response": response} - if not isinstance(response, dict) - else response.copy() - ) + request_data = {"response": response} if not isinstance(response, dict) else response.copy() else: if "response" not in request_data: - request_data["response"] = ( - response if not isinstance(response, dict) else response.copy() - ) + request_data["response"] = response if not isinstance(response, dict) else response.copy() # Add user API key metadata with prefixed keys if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: request_data["litellm_metadata"] = user_metadata @@ -300,13 +280,9 @@ class LlmPassthroughRouteHandler(BaseTranslation): return getattr(handler_cls, "de_anonymize_event_stream", None) @staticmethod - def supports_event_stream_de_anonymization( - provider: Optional[str], endpoint: Optional[str] - ) -> bool: + def supports_event_stream_de_anonymization(provider: Optional[str], endpoint: Optional[str]) -> bool: handler_cls = _get_provider_handlers().get(provider or "") - endpoint_check = getattr( - handler_cls, "event_stream_endpoint_is_de_anonymizable", None - ) + endpoint_check = getattr(handler_cls, "event_stream_endpoint_is_de_anonymizable", None) if endpoint_check is None: return False return endpoint_check(endpoint or "") @@ -319,13 +295,10 @@ class LlmPassthroughRouteHandler(BaseTranslation): data: dict, ) -> bytes: provider = data.get("custom_llm_provider") - de_anonymize = LlmPassthroughRouteHandler._resolve_event_stream_de_anonymizer( - provider - ) + de_anonymize = LlmPassthroughRouteHandler._resolve_event_stream_de_anonymizer(provider) if de_anonymize is None: verbose_proxy_logger.debug( - "LlmPassthroughRouteHandler: no event-stream handler for provider=%s, " - "leaving stream unmodified", + "LlmPassthroughRouteHandler: no event-stream handler for provider=%s, leaving stream unmodified", provider, ) return body_bytes diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index cc0738697d1..93afccd5c9d 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -25,16 +25,8 @@ class PerplexityChatConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = ( - api_base - or get_secret_str("PERPLEXITY_API_BASE") - or "https://api.perplexity.ai" - ) # type: ignore - dynamic_api_key = ( - api_key - or get_secret_str("PERPLEXITYAI_API_KEY") - or get_secret_str("PERPLEXITY_API_KEY") - ) + api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" # type: ignore + dynamic_api_key = api_key or get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str("PERPLEXITY_API_KEY") return api_base, dynamic_api_key def get_supported_openai_params(self, model: str) -> list: @@ -59,17 +51,13 @@ class PerplexityChatConfig(OpenAIGPTConfig): ] try: - if litellm.supports_reasoning( - model=model, custom_llm_provider=self.custom_llm_provider - ): + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): base_openai_params.append("reasoning_effort") except Exception as e: verbose_logger.debug(f"Error checking if model supports reasoning: {e}") try: - if litellm.supports_web_search( - model=model, custom_llm_provider=self.custom_llm_provider - ): + if litellm.supports_web_search(model=model, custom_llm_provider=self.custom_llm_provider): base_openai_params.append("web_search_options") except Exception as e: verbose_logger.debug(f"Error checking if model supports web search: {e}") @@ -108,20 +96,14 @@ class PerplexityChatConfig(OpenAIGPTConfig): # Extract and enhance usage with Perplexity-specific fields try: raw_response_json = raw_response.json() - self._enhance_usage_with_perplexity_fields( - model_response, raw_response_json - ) + self._enhance_usage_with_perplexity_fields(model_response, raw_response_json) self._add_citations_as_annotations(model_response, raw_response_json) except Exception as e: - verbose_logger.debug( - f"Error extracting Perplexity-specific usage fields: {e}" - ) + verbose_logger.debug(f"Error extracting Perplexity-specific usage fields: {e}") return model_response - def _enhance_usage_with_perplexity_fields( - self, model_response: ModelResponse, raw_response_json: dict - ) -> None: + def _enhance_usage_with_perplexity_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None: """ Extract citation tokens and search queries from Perplexity API response and add them to the usage object using standard LiteLLM fields. @@ -140,9 +122,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): if citations: # Count total characters in citations as a proxy for citation tokens # This is an estimation - in practice, you might want to use proper tokenization - total_citation_chars = sum( - len(str(citation)) for citation in citations if citation - ) + total_citation_chars = sum(len(str(citation)) for citation in citations if citation) # Rough estimation: ~4 characters per token (OpenAI's general rule) if total_citation_chars > 0: citation_tokens = max(1, total_citation_chars // 4) @@ -161,9 +141,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): num_search_queries = raw_response_json.get("search_queries") # Create or update prompt_tokens_details to include web search requests and citation tokens - if citation_tokens > 0 or ( - num_search_queries is not None and num_search_queries > 0 - ): + if citation_tokens > 0 or (num_search_queries is not None and num_search_queries > 0): if usage.prompt_tokens_details is None: usage.prompt_tokens_details = PromptTokensDetailsWrapper() @@ -175,9 +153,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): if num_search_queries is not None and num_search_queries > 0: usage.prompt_tokens_details.web_search_requests = num_search_queries - def _add_citations_as_annotations( - self, model_response: ModelResponse, raw_response_json: dict - ) -> None: + def _add_citations_as_annotations(self, model_response: ModelResponse, raw_response_json: dict) -> None: """ Extract citations and search_results from Perplexity API response and add them as ChatCompletionAnnotation objects to the message. diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index ec7ec397ea6..c9574f3be80 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -34,9 +34,7 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: ## GET MODEL INFO model_info = get_model_info(model=model, custom_llm_provider="perplexity") - def _safe_float_cast( - value: Union[str, int, float, None, object], default: float = 0.0 - ) -> float: + def _safe_float_cast(value: Union[str, int, float, None, object], default: float = 0.0) -> float: """Safely cast a value to float with proper type handling for mypy.""" if value is None: return default @@ -60,14 +58,8 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: output_cost_per_token = _safe_float_cast(model_info.get("output_cost_per_token")) reasoning_tokens = getattr(usage, "reasoning_tokens", 0) or 0 - if ( - reasoning_tokens == 0 - and hasattr(usage, "completion_tokens_details") - and usage.completion_tokens_details - ): - reasoning_tokens = ( - getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0 - ) + if reasoning_tokens == 0 and hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: + reasoning_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0 reasoning_cost_value = model_info.get("output_cost_per_reasoning_token") @@ -76,9 +68,7 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: # configured we subtract before the output-rate multiplication so the reasoning # tokens are not billed twice. if reasoning_tokens > 0 and reasoning_cost_value is not None: - non_reasoning_completion_tokens = max( - 0, (usage.completion_tokens or 0) - reasoning_tokens - ) + non_reasoning_completion_tokens = max(0, (usage.completion_tokens or 0) - reasoning_tokens) completion_cost: float = non_reasoning_completion_tokens * output_cost_per_token completion_cost += reasoning_tokens * _safe_float_cast(reasoning_cost_value) else: @@ -87,23 +77,19 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: ## ADD SEARCH QUERIES COST (if present) num_search_queries = 0 if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - num_search_queries = ( - getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0 - ) + num_search_queries = getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0 # Check both possible keys for search cost (legacy and current) - search_cost_value = model_info.get( - "search_queries_cost_per_query" - ) or model_info.get("search_context_cost_per_query") + search_cost_value = model_info.get("search_queries_cost_per_query") or model_info.get( + "search_context_cost_per_query" + ) if num_search_queries > 0 and search_cost_value is not None: # Handle both dict and float formats if isinstance(search_cost_value, dict): # search_context_cost_per_query stores the per-request price in USD # (e.g. sonar low = $0.005/request). Use it directly, matching the # gemini cost calculator which reads the same field per request. - search_cost_per_query = _safe_float_cast( - search_cost_value.get("search_context_size_low", 0) - ) + search_cost_per_query = _safe_float_cast(search_cost_value.get("search_context_size_low", 0)) else: search_cost_per_query = _safe_float_cast(search_cost_value) search_cost = num_search_queries * search_cost_per_query diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index 24881ccebf8..a52eab34c08 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -34,9 +34,7 @@ class PerplexityEmbeddingError(BaseLLMException): ): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://api.perplexity.ai/v1/embeddings" - ) + self.request = httpx.Request(method="POST", url="https://api.perplexity.ai/v1/embeddings") self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, @@ -99,9 +97,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): api_base: Optional[str] = None, ) -> dict: if api_key is None: - api_key = get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str( - "PERPLEXITY_API_KEY" - ) + api_key = get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str("PERPLEXITY_API_KEY") return { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", @@ -152,9 +148,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): try: raw_response_json = raw_response.json() except Exception: - raise PerplexityEmbeddingError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise PerplexityEmbeddingError(message=raw_response.text, status_code=raw_response.status_code) model_response.model = raw_response_json.get("model", model) model_response.object = raw_response_json.get("object", "list") @@ -163,16 +157,13 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): decoded_data: List[Dict[str, Any]] = [] for item in raw_data: decoded_item = dict(item) - decoded_item["embedding"] = self._decode_base64_embedding( - item.get("embedding") - ) + decoded_item["embedding"] = self._decode_base64_embedding(item.get("embedding")) decoded_data.append(decoded_item) model_response.data = decoded_data usage_data = raw_response_json.get("usage", {}) usage = Usage( - prompt_tokens=usage_data.get("prompt_tokens", 0) - or usage_data.get("total_tokens", 0), + prompt_tokens=usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0), total_tokens=usage_data.get("total_tokens", 0), ) model_response.usage = usage @@ -184,6 +175,4 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): status_code: int, headers: Union[dict, httpx.Headers], ) -> BaseLLMException: - return PerplexityEmbeddingError( - message=error_message, status_code=status_code, headers=headers - ) + return PerplexityEmbeddingError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index e09dc01f1c1..dd5517f6c33 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -40,30 +40,20 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.PERPLEXITY - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() api_key = ( - litellm_params.api_key - or get_secret_str("PERPLEXITYAI_API_KEY") - or get_secret_str("PERPLEXITY_API_KEY") + litellm_params.api_key or get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str("PERPLEXITY_API_KEY") ) if api_key: headers["Authorization"] = f"Bearer {api_key}" return headers def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str: - api_base = ( - api_base - or get_secret_str("PERPLEXITY_API_BASE") - or "https://api.perplexity.ai" - ) + api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" return f"{api_base.rstrip('/')}/v1/responses" - def _ensure_message_type( - self, input: Union[str, ResponseInputParam] - ) -> Union[str, ResponseInputParam]: + def _ensure_message_type(self, input: Union[str, ResponseInputParam]) -> Union[str, ResponseInputParam]: """Ensure list input items have type='message' (required by Perplexity).""" if isinstance(input, str): return input @@ -71,9 +61,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): result: List[Any] = [] for item in input: if isinstance(item, dict) and "type" not in item: - new_item = dict( - item - ) # convert to plain dict to avoid TypedDict checking + new_item = dict(item) # convert to plain dict to avoid TypedDict checking new_item["type"] = "message" result.append(new_item) else: @@ -119,10 +107,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): except Exception: raw_response_json = None - if ( - isinstance(raw_response_json, dict) - and raw_response_json.get("status") == "failed" - ): + if isinstance(raw_response_json, dict) and raw_response_json.get("status") == "failed": error = raw_response_json.get("error", {}) raise BaseLLMException( status_code=raw_response.status_code, diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index 55de52c5384..8ed165de742 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -58,9 +58,7 @@ class PerplexitySearchConfig(BaseSearchConfig): default_api_base=self.PERPLEXITY_API_BASE, ) if not api_key: - raise ValueError( - "PERPLEXITYAI_API_KEY is not set. Set `PERPLEXITYAI_API_KEY` environment variable." - ) + raise ValueError("PERPLEXITYAI_API_KEY is not set. Set `PERPLEXITYAI_API_KEY` environment variable.") headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -75,11 +73,7 @@ class PerplexitySearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - api_base = ( - api_base - or get_secret_str("PERPLEXITY_API_BASE") - or self.PERPLEXITY_API_BASE - ) + api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or self.PERPLEXITY_API_BASE # append "/search" to the api base if it's not already there if not api_base.endswith("/search"): diff --git a/litellm/llms/petals/completion/handler.py b/litellm/llms/petals/completion/handler.py index ae38baecf22..4a4a820d56d 100644 --- a/litellm/llms/petals/completion/handler.py +++ b/litellm/llms/petals/completion/handler.py @@ -97,9 +97,7 @@ def completion( model = model - tokenizer = AutoTokenizer.from_pretrained( - model, use_fast=False, add_bos_token=False - ) + tokenizer = AutoTokenizer.from_pretrained(model, use_fast=False, add_bos_token=False) model_obj = AutoDistributedModelForCausalLM.from_pretrained(model) ## LOGGING @@ -129,9 +127,7 @@ def completion( model_response.choices[0].message.content = output_text # type: ignore prompt_tokens = len(encoding.encode(prompt)) - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content")) - ) + completion_tokens = len(encoding.encode(model_response["choices"][0]["message"].get("content"))) model_response.created = int(time.time()) model_response.model = model diff --git a/litellm/llms/petals/completion/transformation.py b/litellm/llms/petals/completion/transformation.py index d50afc4625a..ae6415680b1 100644 --- a/litellm/llms/petals/completion/transformation.py +++ b/litellm/llms/petals/completion/transformation.py @@ -37,9 +37,7 @@ class PetalsConfig(BaseConfig): """ max_length: Optional[int] = None - max_new_tokens: Optional[int] = ( - litellm.max_tokens - ) # petals requires max tokens to be set + max_new_tokens: Optional[int] = litellm.max_tokens # petals requires max tokens to be set do_sample: Optional[bool] = None temperature: Optional[float] = None top_k: Optional[int] = None @@ -49,9 +47,7 @@ class PetalsConfig(BaseConfig): def __init__( self, max_length: Optional[int] = None, - max_new_tokens: Optional[ - int - ] = litellm.max_tokens, # petals requires max tokens to be set + max_new_tokens: Optional[int] = litellm.max_tokens, # petals requires max tokens to be set do_sample: Optional[bool] = None, temperature: Optional[float] = None, top_k: Optional[int] = None, @@ -67,12 +63,8 @@ class PetalsConfig(BaseConfig): def get_config(cls): return super().get_config() - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return PetalsError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return PetalsError(status_code=status_code, message=error_message, headers=headers) def get_supported_openai_params(self, model: str) -> List: return ["max_tokens", "temperature", "top_p", "stream"] diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index fc4cfc7b083..b58b6e7f498 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -27,9 +27,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): - api_key: API key for authentication with the PG vector service """ - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate environment and set headers for PG vector service authentication """ @@ -83,9 +81,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: - encoded_vector_store_id = encode_url_path_segment( - vector_store_id, field_name="vector_store_id" - ) + encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py index 07f2738aa96..fe8ee508fcd 100644 --- a/litellm/llms/predibase/chat/handler.py +++ b/litellm/llms/predibase/chat/handler.py @@ -27,9 +27,7 @@ async def make_call( logging_obj, timeout: Optional[Union[float, httpx.Timeout]], ): - response = await client.post( - api_base, headers=headers, data=data, stream=True, timeout=timeout - ) + response = await client.post(api_base, headers=headers, data=data, stream=True, timeout=timeout) if response.status_code != 200: raise PredibaseError(status_code=response.status_code, message=response.text) @@ -216,9 +214,7 @@ class PredibaseChatCompletion: params={"timeout": timeout}, ) try: - response = await async_handler.post( - api_base, headers=headers, data=json.dumps(data) - ) + response = await async_handler.post(api_base, headers=headers, data=json.dumps(data)) except httpx.HTTPStatusError as e: raise PredibaseError( status_code=e.response.status_code, diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index ce004f60bfc..fcb21272be2 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -35,13 +35,9 @@ class PredibaseConfig(BaseConfig): best_of: Optional[int] = None decoder_input_details: Optional[bool] = None details: bool = True # enables returning logprobs + best of - max_new_tokens: int = ( - DEFAULT_MAX_TOKENS # openai default - requests hang if max_new_tokens not given - ) + max_new_tokens: int = DEFAULT_MAX_TOKENS # openai default - requests hang if max_new_tokens not given repetition_penalty: Optional[float] = None - return_full_text: Optional[bool] = ( - False # by default don't return the input as part of the output - ) + return_full_text: Optional[bool] = False # by default don't return the input as part of the output seed: Optional[int] = None stop: Optional[List[str]] = None temperature: Optional[float] = None @@ -108,9 +104,7 @@ class PredibaseConfig(BaseConfig): optional_params["top_p"] = value if param == "n": optional_params["best_of"] = value - optional_params["do_sample"] = ( - True # Need to sample if you want best of for hf inference endpoints - ) + optional_params["do_sample"] = True # Need to sample if you want best of for hf inference endpoints if param == "stream": optional_params["stream"] = value if param == "stop": @@ -175,13 +169,8 @@ class PredibaseConfig(BaseConfig): completion_response["generated_text"] ) - if ( - "details" in completion_response - and "tokens" in completion_response["details"] - ): - model_response.choices[0].finish_reason = map_finish_reason( - completion_response["details"]["finish_reason"] - ) + if "details" in completion_response and "tokens" in completion_response["details"]: + model_response.choices[0].finish_reason = map_finish_reason(completion_response["details"]["finish_reason"]) sum_logprob = 0 for token in completion_response["details"]["tokens"]: if token["logprob"] is not None: @@ -201,14 +190,9 @@ class PredibaseConfig(BaseConfig): best_of_value = 0 if best_of_value > 1: - if ( - "details" in completion_response - and "best_of_sequences" in completion_response["details"] - ): + if "details" in completion_response and "best_of_sequences" in completion_response["details"]: choices_list = [] - for idx, item in enumerate( - completion_response["details"]["best_of_sequences"] - ): + for idx, item in enumerate(completion_response["details"]["best_of_sequences"]): sum_logprob = 0 for token in item["tokens"]: if token["logprob"] is not None: @@ -238,11 +222,7 @@ class PredibaseConfig(BaseConfig): if output_text is not None and len(output_text) > 0: completion_tokens = 0 try: - completion_tokens = len( - encoding.encode( - model_response["choices"][0]["message"].get("content", "") - ) - ) + completion_tokens = len(encoding.encode(model_response["choices"][0]["message"].get("content", ""))) except Exception: # Keep usage calculation non-blocking if encoding fails. pass @@ -332,9 +312,7 @@ class PredibaseConfig(BaseConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - tenant_id = litellm_params.get("predibase_tenant_id") or litellm_params.get( - "tenant_id" - ) + tenant_id = litellm_params.get("predibase_tenant_id") or litellm_params.get("tenant_id") if tenant_id is None: raise ValueError( "Missing Predibase Tenant ID - Required for making the request. Set dynamically (e.g. `completion(..tenant_id=)`) or in env - `PREDIBASE_TENANT_ID`." @@ -347,21 +325,15 @@ class PredibaseConfig(BaseConfig): base_url = os.getenv("PREDIBASE_API_BASE", "") completion_url = f"{base_url}/{tenant_id}/deployments/v2/llms/{model}" - should_stream = ( - stream if stream is not None else optional_params.get("stream", False) - ) + should_stream = stream if stream is not None else optional_params.get("stream", False) if should_stream is True: completion_url += "/generate_stream" else: completion_url += "/generate" return completion_url - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return PredibaseError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return PredibaseError(status_code=status_code, message=error_message, headers=headers) def validate_environment( self, diff --git a/litellm/llms/ragflow/chat/transformation.py b/litellm/llms/ragflow/chat/transformation.py index 990fc2b2e61..be3417d1aad 100644 --- a/litellm/llms/ragflow/chat/transformation.py +++ b/litellm/llms/ragflow/chat/transformation.py @@ -49,20 +49,14 @@ class RAGFlowConfig(OpenAIConfig): ) if parts[0] != "ragflow": - raise ValueError( - f"Invalid RAGFlow model format: {model}. Must start with 'ragflow/'" - ) + raise ValueError(f"Invalid RAGFlow model format: {model}. Must start with 'ragflow/'") endpoint_type = parts[1] if endpoint_type not in ["chat", "agent"]: - raise ValueError( - f"Invalid RAGFlow endpoint type: {endpoint_type}. Must be 'chat' or 'agent'" - ) + raise ValueError(f"Invalid RAGFlow endpoint type: {endpoint_type}. Must be 'chat' or 'agent'") entity_id = parts[2] - model_name = "/".join( - parts[3:] - ) # Handle model names that might contain slashes + model_name = "/".join(parts[3:]) # Handle model names that might contain slashes return endpoint_type, entity_id, model_name @@ -94,19 +88,10 @@ class RAGFlowConfig(OpenAIConfig): Complete URL for the API call """ # Get api_base from multiple sources: input param, litellm_params, environment, or global litellm setting - if ( - litellm_params - and hasattr(litellm_params, "api_base") - and litellm_params.api_base - ): + if litellm_params and hasattr(litellm_params, "api_base") and litellm_params.api_base: api_base = api_base or litellm_params.api_base - api_base = ( - api_base - or litellm.api_base - or get_secret("RAGFLOW_API_BASE") - or get_secret_str("RAGFLOW_API_BASE") - ) + api_base = api_base or litellm.api_base or get_secret("RAGFLOW_API_BASE") or get_secret_str("RAGFLOW_API_BASE") if api_base is None: raise ValueError( @@ -164,16 +149,11 @@ class RAGFlowConfig(OpenAIConfig): # Get api_base from multiple sources: input param, environment, or global litellm setting dynamic_api_base = ( - api_base - or litellm.api_base - or get_secret("RAGFLOW_API_BASE") - or get_secret_str("RAGFLOW_API_BASE") + api_base or litellm.api_base or get_secret("RAGFLOW_API_BASE") or get_secret_str("RAGFLOW_API_BASE") ) # Get api_key from multiple sources: input param, environment, or global litellm setting - dynamic_api_key = ( - api_key or litellm.api_key or get_secret_str("RAGFLOW_API_KEY") - ) + dynamic_api_key = api_key or litellm.api_key or get_secret_str("RAGFLOW_API_KEY") return dynamic_api_base, dynamic_api_key, custom_llm_provider @@ -203,11 +183,7 @@ class RAGFlowConfig(OpenAIConfig): Updated headers dictionary """ # Use api_key from litellm_params if available, otherwise fall back to other sources - if ( - litellm_params - and hasattr(litellm_params, "api_key") - and litellm_params.api_key - ): + if litellm_params and hasattr(litellm_params, "api_key") and litellm_params.api_key: api_key = api_key or litellm_params.api_key # Get api_key from multiple sources: input param, litellm_params, environment, or global litellm setting @@ -266,6 +242,4 @@ class RAGFlowConfig(OpenAIConfig): actual_model = model # Use parent's transform_request with the actual model name - return super().transform_request( - actual_model, messages, optional_params, litellm_params, headers - ) + return super().transform_request(actual_model, messages, optional_params, litellm_params, headers) diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index 3238d3e9c14..d8bdd981425 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -24,17 +24,13 @@ else: class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): """Vector store configuration for RAGFlow datasets.""" - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: api_key = litellm_params.get("api_key") if api_key is None: # Try to get from environment variable api_key = get_secret_str("RAGFLOW_API_KEY") if api_key is None: - raise ValueError( - "api_key is required (set RAGFLOW_API_KEY env var or pass in litellm_params)" - ) + raise ValueError("api_key is required (set RAGFLOW_API_KEY env var or pass in litellm_params)") return { "headers": { "Authorization": f"Bearer {api_key}", @@ -48,17 +44,13 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): "write": [], } - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """Validate environment and set headers for RAGFlow API.""" litellm_params = litellm_params or GenericLiteLLMParams() api_key = litellm_params.api_key or get_secret_str("RAGFLOW_API_KEY") if api_key is None: - raise ValueError( - "RAGFLOW_API_KEY is required (set env var or pass in litellm_params)" - ) + raise ValueError("RAGFLOW_API_KEY is required (set env var or pass in litellm_params)") headers.update( { @@ -82,10 +74,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): - Default: http://localhost:9380 """ api_base = ( - api_base - or litellm_params.get("api_base") - or get_secret_str("RAGFLOW_API_BASE") - or "http://localhost:9380" + api_base or litellm_params.get("api_base") or get_secret_str("RAGFLOW_API_BASE") or "http://localhost:9380" ) # Remove trailing slashes @@ -105,17 +94,13 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """RAGFlow vector stores are management-only, search is not supported.""" - raise NotImplementedError( - "RAGFlow vector stores support dataset management only, not search/retrieval" - ) + raise NotImplementedError("RAGFlow vector stores support dataset management only, not search/retrieval") def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: """RAGFlow vector stores are management-only, search is not supported.""" - raise NotImplementedError( - "RAGFlow vector stores support dataset management only, not search/retrieval" - ) + raise NotImplementedError("RAGFlow vector stores support dataset management only, not search/retrieval") def transform_create_vector_store_request( self, @@ -172,9 +157,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): return url, request_body - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: """ Transform RAGFlow response to VectorStoreCreateResponse format. diff --git a/litellm/llms/recraft/cost_calculator.py b/litellm/llms/recraft/cost_calculator.py index 27b9108e5fe..5ab47e9395e 100644 --- a/litellm/llms/recraft/cost_calculator.py +++ b/litellm/llms/recraft/cost_calculator.py @@ -22,6 +22,4 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index 1dccd406058..61c669b50c0 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -68,9 +68,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): Some providers need `model` in `api_base` """ - complete_url: str = ( - api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url: str = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") complete_url = f"{complete_url}/{self.IMAGE_EDIT_ENDPOINT}" @@ -109,9 +107,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): request_params = { "model": model, - "strength": image_edit_optional_request_params.pop( - "strength", self.DEFAULT_STRENGTH - ), + "strength": image_edit_optional_request_params.pop("strength", self.DEFAULT_STRENGTH), **image_edit_optional_request_params, } if prompt is not None: @@ -122,9 +118,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): ######################################################### # Reuse OpenAI logic: Separate images as `files` and send other parameters as `data` ######################################################### - files_list = ( - self._get_image_files_for_request(image=image) if image is not None else [] - ) + files_list = self._get_image_files_for_request(image=image) if image is not None else [] data_without_images = {k: v for k, v in request_dict.items() if k != "image"} return data_without_images, files_list @@ -144,17 +138,11 @@ class RecraftImageEditConfig(BaseImageEditConfig): _image = image if _image is not None: - image_content_type: str = ImageEditRequestUtils.get_image_content_type( - _image - ) + image_content_type: str = ImageEditRequestUtils.get_image_content_type(_image) if isinstance(_image, BufferedReader): - files_list.append( - ("image", (_image.name, _image, image_content_type)) - ) + files_list.append(("image", (_image.name, _image, image_content_type))) else: - files_list.append( - ("image", ("image.png", _image, image_content_type)) - ) + files_list.append(("image", ("image.png", _image, image_content_type))) return files_list diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index 4a00512dfb9..9f48273c306 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -25,9 +25,7 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ https://www.recraft.ai/docs#generate-image """ @@ -68,9 +66,7 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = ( - api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url: str = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" diff --git a/litellm/llms/reducto/common.py b/litellm/llms/reducto/common.py index 4e7d96dbe87..364f269feb1 100644 --- a/litellm/llms/reducto/common.py +++ b/litellm/llms/reducto/common.py @@ -51,9 +51,7 @@ def extract_file_id_or_bytes( _raise_bad_request("Invalid Reducto data URI provided.", model=model) if ";base64" not in header: - _raise_bad_request( - "Reducto only supports base64-encoded data URIs.", model=model - ) + _raise_bad_request("Reducto only supports base64-encoded data URIs.", model=model) mime = header.removeprefix("data:").split(";")[0] or "application/octet-stream" try: @@ -68,16 +66,10 @@ def _extract_file_id_from_upload_response(response: Any) -> str: try: payload = response.json() except ValueError as exc: - raise ValueError( - "Reducto /upload returned a non-JSON 200 response: {}".format(response.text) - ) from exc + raise ValueError("Reducto /upload returned a non-JSON 200 response: {}".format(response.text)) from exc file_id = (payload or {}).get("file_id") if isinstance(payload, dict) else None if not isinstance(file_id, str) or not file_id: - raise ValueError( - "Reducto /upload returned 200 without a file_id; got payload={}".format( - payload - ) - ) + raise ValueError("Reducto /upload returned 200 without a file_id; got payload={}".format(payload)) return file_id @@ -135,18 +127,14 @@ def build_pages_from_reducto(result: Dict[str, Any]) -> List["OCRPage"]: blocks_by_page[normalized_page].append(block) if not blocks_by_page: - fallback_markdown = "\n\n".join( - chunk.get("content", "") for chunk in chunks if chunk.get("content") - ) + fallback_markdown = "\n\n".join(chunk.get("content", "") for chunk in chunks if chunk.get("content")) if fallback_markdown == "": return [] return [OCRPage(index=0, markdown=fallback_markdown)] pages: List["OCRPage"] = [] for page_no, blocks in sorted(blocks_by_page.items()): - markdown = "\n\n".join( - block.get("content", "") for block in blocks if block.get("content") - ) + markdown = "\n\n".join(block.get("content", "") for block in blocks if block.get("content")) page_index = max(page_no - 1, 0) page = OCRPage( index=page_index, diff --git a/litellm/llms/reducto/ocr/transformation.py b/litellm/llms/reducto/ocr/transformation.py index cc338ecc484..e8bfcceea2a 100644 --- a/litellm/llms/reducto/ocr/transformation.py +++ b/litellm/llms/reducto/ocr/transformation.py @@ -69,16 +69,12 @@ class _BaseReductoOCRConfig(BaseOCRConfig): source_url = document.get("document_url") or document.get("image_url") if source_url is None: raise ValueError( - "Reducto expected OCR preprocessing to produce document_url or image_url for model={}".format( - model - ) + "Reducto expected OCR preprocessing to produce document_url or image_url for model={}".format(model) ) return source_url @staticmethod - def _resolve_credentials( - api_key: Optional[str], api_base: Optional[str] - ) -> Tuple[str, str]: + def _resolve_credentials(api_key: Optional[str], api_base: Optional[str]) -> Tuple[str, str]: from litellm.secret_managers.main import get_secret_str resolved_key = api_key or get_secret_str("REDUCTO_API_KEY") @@ -213,9 +209,7 @@ class ReductoParseLegacyConfig(_BaseReductoOCRConfig): api_base=kwargs.get("api_base"), ) return OCRRequestData( - data=self._build_legacy_body( - file_id=file_id, optional_params=optional_params - ), + data=self._build_legacy_body(file_id=file_id, optional_params=optional_params), files=None, ) @@ -234,8 +228,6 @@ class ReductoParseLegacyConfig(_BaseReductoOCRConfig): api_base=kwargs.get("api_base"), ) return OCRRequestData( - data=self._build_legacy_body( - file_id=file_id, optional_params=optional_params - ), + data=self._build_legacy_body(file_id=file_id, optional_params=optional_params), files=None, ) diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py index a2eddb65a54..57381e57dab 100644 --- a/litellm/llms/replicate/chat/handler.py +++ b/litellm/llms/replicate/chat/handler.py @@ -29,9 +29,7 @@ def handle_prediction_response_streaming( status = "" while True and (status not in ["succeeded", "failed", "canceled"]): - time.sleep( - REPLICATE_POLLING_DELAY_SECONDS - ) # prevent being rate limited by replicate + time.sleep(REPLICATE_POLLING_DELAY_SECONDS) # prevent being rate limited by replicate print_verbose(f"replicate: polling endpoint: {prediction_url}") response = http_client.get(prediction_url, headers=headers) if response.status_code == 200: @@ -43,9 +41,7 @@ def handle_prediction_response_streaming( except Exception: raise ReplicateError( status_code=422, - message="Unable to parse response. Got={}".format( - response_data["output"] - ), + message="Unable to parse response. Got={}".format(response_data["output"]), headers=response.headers, ) new_output = output_string[len(previous_output) :] @@ -80,17 +76,13 @@ async def async_handle_prediction_response_streaming( status = "" while True and (status not in ["succeeded", "failed", "canceled"]): - await asyncio.sleep( - REPLICATE_POLLING_DELAY_SECONDS - ) # prevent being rate limited by replicate + await asyncio.sleep(REPLICATE_POLLING_DELAY_SECONDS) # prevent being rate limited by replicate response = await http_client.get(prediction_url, headers=headers) if response.status_code == 200: response_data = response.json() status = response_data.get("status", "") # Check that "output" exists and is not None or empty - output_present = ( - "output" in response_data and response_data["output"] is not None - ) + output_present = "output" in response_data and response_data["output"] is not None if output_present: try: # If output is None or not a list, treat as empty string @@ -104,9 +96,7 @@ async def async_handle_prediction_response_streaming( except Exception: raise ReplicateError( status_code=422, - message="Unable to parse response. Got={}".format( - response_data.get("output", None) - ), + message="Unable to parse response. Got={}".format(response_data.get("output", None)), headers=response.headers, ) new_output = output_string[len(previous_output) :] @@ -180,9 +170,7 @@ def completion( headers=headers, ) # type: ignore ## COMPLETION CALL - model_response.created = int( - time.time() - ) # for pricing this must remain right before calling api + model_response.created = int(time.time()) # for pricing this must remain right before calling api prediction_url = replicate_config.get_complete_url( api_base=api_base, @@ -214,9 +202,7 @@ def completion( headers=headers, http_client=httpx_client, ) - return CustomStreamWrapper( - _response, model, logging_obj=logging_obj, custom_llm_provider="replicate" - ) # type: ignore + return CustomStreamWrapper(_response, model, logging_obj=logging_obj, custom_llm_provider="replicate") # type: ignore else: for retry in range(litellm.DEFAULT_REPLICATE_POLLING_RETRIES): time.sleep( @@ -274,9 +260,7 @@ async def async_completion( llm_provider=litellm.LlmProviders.REPLICATE, params={"timeout": 600.0}, ) - response = await async_handler.post( - url=prediction_url, headers=headers, data=json.dumps(input_data) - ) + response = await async_handler.post(url=prediction_url, headers=headers, data=json.dumps(input_data)) prediction_url = replicate_config.get_prediction_url(response) if "stream" in optional_params and optional_params["stream"] is True: @@ -287,9 +271,7 @@ async def async_completion( headers=headers, http_client=async_handler, ) - return CustomStreamWrapper( - _response, model, logging_obj=logging_obj, custom_llm_provider="replicate" - ) # type: ignore + return CustomStreamWrapper(_response, model, logging_obj=logging_obj, custom_llm_provider="replicate") # type: ignore for retry in range(litellm.DEFAULT_REPLICATE_POLLING_RETRIES): await asyncio.sleep( diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py index 4c610868018..6da26b966f3 100644 --- a/litellm/llms/replicate/chat/transformation.py +++ b/litellm/llms/replicate/chat/transformation.py @@ -133,9 +133,7 @@ class ReplicateConfig(BaseConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return ReplicateError( - status_code=status_code, message=error_message, headers=headers - ) + return ReplicateError(status_code=status_code, message=error_message, headers=headers) def get_complete_url( self, @@ -191,9 +189,7 @@ class ReplicateConfig(BaseConfig): model_prompt_details = litellm.custom_prompt_dict[model] prompt = custom_prompt( role_dict=model_prompt_details.get("roles", {}), - initial_prompt_value=model_prompt_details.get( - "initial_prompt_value", "" - ), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), final_prompt_value=model_prompt_details.get("final_prompt_value", ""), bos_token=model_prompt_details.get("bos_token", ""), eos_token=model_prompt_details.get("eos_token", ""), @@ -225,8 +221,7 @@ class ReplicateConfig(BaseConfig): if ":" in version_id and len(version_id) > REPLICATE_MODEL_NAME_WITH_ID_LENGTH: model_parts = version_id.split(":") if ( - len(model_parts) > 1 - and len(model_parts[1]) == REPLICATE_MODEL_NAME_WITH_ID_LENGTH + len(model_parts) > 1 and len(model_parts[1]) == REPLICATE_MODEL_NAME_WITH_ID_LENGTH ): ## checks if model name has a 64 digit code - e.g. "meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3" request_data["version"] = model_parts[1] @@ -256,9 +251,7 @@ class ReplicateConfig(BaseConfig): if raw_response_json.get("status") != "succeeded": raise ReplicateError( status_code=422, - message="LiteLLM Error - prediction not succeeded - {}".format( - raw_response_json - ), + message="LiteLLM Error - prediction not succeeded - {}".format(raw_response_json), headers=raw_response.headers, ) outputs = raw_response_json.get("output", []) @@ -299,9 +292,7 @@ class ReplicateConfig(BaseConfig): if prediction_url is None: raise ReplicateError( status_code=400, - message="LiteLLM Error - prediction url is None - {}".format( - response_json - ), + message="LiteLLM Error - prediction url is None - {}".format(response_json), headers=response.headers, ) return prediction_url diff --git a/litellm/llms/runwayml/cost_calculator.py b/litellm/llms/runwayml/cost_calculator.py index 35b6086f196..564f4814eec 100644 --- a/litellm/llms/runwayml/cost_calculator.py +++ b/litellm/llms/runwayml/cost_calculator.py @@ -25,6 +25,4 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError( - f"image_response must be of type ImageResponse, got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse, got type={type(image_response)}") diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index 448dcd4a67b..fddd0b1350b 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -49,9 +49,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = ( - api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url: str = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") if self.IMAGE_GENERATION_ENDPOINT: @@ -69,9 +67,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): api_base: Optional[str] = None, ) -> dict: final_api_key: Optional[str] = ( - api_key - or get_secret_str("RUNWAYML_API_SECRET") - or get_secret_str("RUNWAYML_API_KEY") + api_key or get_secret_str("RUNWAYML_API_SECRET") or get_secret_str("RUNWAYML_API_KEY") ) if not final_api_key: raise ValueError("RUNWAYML_API_SECRET or RUNWAYML_API_KEY is not set") @@ -154,9 +150,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): TimeoutError: If operation has exceeded timeout """ if time.time() - start_time > timeout_secs: - raise TimeoutError( - f"RunwayML task polling timed out after {timeout_secs} seconds" - ) + raise TimeoutError(f"RunwayML task polling timed out after {timeout_secs} seconds") @staticmethod def _check_task_status(response_data: Dict[str, Any]) -> str: @@ -183,9 +177,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): elif status == "FAILED": failure_reason = response_data.get("failure", "Unknown error") failure_code = response_data.get("failureCode", "unknown") - raise ValueError( - f"RunwayML image generation failed: {failure_reason} (code: {failure_code})" - ) + raise ValueError(f"RunwayML image generation failed: {failure_reason} (code: {failure_code})") elif status == "CANCELLED": raise ValueError("RunwayML image generation was cancelled") elif status in ["PENDING", "RUNNING", "THROTTLED"]: @@ -346,9 +338,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), - "X-Runway-Version": raw_response.request.headers.get( - "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION - ), + "X-Runway-Version": raw_response.request.headers.get("X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION), } # Poll until task completes @@ -408,9 +398,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), - "X-Runway-Version": raw_response.request.headers.get( - "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION - ), + "X-Runway-Version": raw_response.request.headers.get("X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION), } # Poll until task completes (async) @@ -424,9 +412,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): # Update response_data with polled result response_data = raw_response.json() - verbose_logger.debug( - "RunwayML polling complete (async), transforming to OpenAI format" - ) + verbose_logger.debug("RunwayML polling complete (async), transforming to OpenAI format") # Transform RunwayML response to OpenAI format return self._transform_runwayml_response_to_openai( @@ -434,9 +420,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): model_response=model_response, ) - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for RunwayML image generation """ diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index 314a538f7c5..0f3da5f7ac5 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -200,11 +200,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ validated_headers = headers.copy() - final_api_key = ( - api_key - or get_secret_str("RUNWAYML_API_SECRET") - or get_secret_str("RUNWAYML_API_KEY") - ) + final_api_key = api_key or get_secret_str("RUNWAYML_API_SECRET") or get_secret_str("RUNWAYML_API_KEY") if not final_api_key: raise ValueError("RUNWAYML_API_SECRET or RUNWAYML_API_KEY is not set") @@ -224,9 +220,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Get the complete URL for RunwayML TTS request """ - complete_url = ( - api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL - ) + complete_url = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL complete_url = complete_url.rstrip("/") return f"{complete_url}/{self.TTS_ENDPOINT_PATH}" @@ -244,9 +238,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): TimeoutError: If operation has exceeded timeout """ if time.time() - start_time > timeout_secs: - raise TimeoutError( - f"RunwayML TTS task polling timed out after {timeout_secs} seconds" - ) + raise TimeoutError(f"RunwayML TTS task polling timed out after {timeout_secs} seconds") @staticmethod def _check_task_status(response_data: Dict[str, Any]) -> str: @@ -273,9 +265,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): elif status == "FAILED": failure_reason = response_data.get("failure", "Unknown error") failure_code = response_data.get("failureCode", "unknown") - raise ValueError( - f"RunwayML TTS failed: {failure_reason} (code: {failure_code})" - ) + raise ValueError(f"RunwayML TTS failed: {failure_reason} (code: {failure_code})") elif status == "CANCELLED": raise ValueError("RunwayML TTS was cancelled") elif status in ["PENDING", "RUNNING", "THROTTLED"]: @@ -480,9 +470,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), - "X-Runway-Version": raw_response.request.headers.get( - "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION - ), + "X-Runway-Version": raw_response.request.headers.get("X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION), } # Poll until task completes @@ -551,9 +539,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), - "X-Runway-Version": raw_response.request.headers.get( - "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION - ), + "X-Runway-Version": raw_response.request.headers.get("X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION), } # Poll until task completes (async) diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index b1723f494ec..b11671c9431 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -97,11 +97,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): seconds = video_create_optional_params["seconds"] if seconds is not None: try: - mapped_params["duration"] = ( - int(float(seconds)) - if isinstance(seconds, str) - else int(seconds) - ) + mapped_params["duration"] = int(float(seconds)) if isinstance(seconds, str) else int(seconds) except (ValueError, TypeError): # If conversion fails, use default duration pass @@ -130,16 +126,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): api_key = api_key or litellm_params.api_key api_key = ( - api_key - or litellm.api_key - or get_secret_str("RUNWAYML_API_SECRET") - or get_secret_str("RUNWAYML_API_KEY") + api_key or litellm.api_key or get_secret_str("RUNWAYML_API_SECRET") or get_secret_str("RUNWAYML_API_KEY") ) if api_key is None: raise ValueError( - "RunwayML API key is required. Set RUNWAYML_API_SECRET environment variable " - "or pass api_key parameter." + "RunwayML API key is required. Set RUNWAYML_API_SECRET environment variable or pass api_key parameter." ) headers.update( @@ -238,15 +230,11 @@ class RunwayMLVideoConfig(BaseVideoConfig): if "output" in response_data and response_data["output"]: # RunwayML returns output as array of URLs when task succeeds video_data["output_url"] = ( - response_data["output"][0] - if isinstance(response_data["output"], list) - else response_data["output"] + response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] ) if "completedAt" in response_data: - video_data["completed_at"] = self._parse_runway_timestamp( - response_data.get("completedAt") - ) + video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) if "failureCode" in response_data or "failure" in response_data: video_data["error"] = { @@ -269,9 +257,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_obj = VideoObject(**video_data) # type: ignore[arg-type] if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider( - video_obj.id, custom_llm_provider, model - ) + video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) # Add usage data for cost tracking usage_data = {} @@ -335,9 +321,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): We'll retrieve the task and extract the video URL. """ original_video_id = extract_original_video_id(video_id) - encoded_video_id = encode_url_path_segment( - original_video_id, field_name="video_id" - ) + encoded_video_id = encode_url_path_segment(original_video_id, field_name="video_id") # Get task status to retrieve video URL url = f"{api_base}/tasks/{encoded_video_id}" @@ -361,16 +345,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Check if the video generation failed or is still processing status = response_data.get("status", "UNKNOWN") if status in ["PENDING", "RUNNING", "THROTTLED"]: - raise ValueError( - f"Video is still processing (status: {status}). Please wait and try again." - ) + raise ValueError(f"Video is still processing (status: {status}). Please wait and try again.") elif status == "FAILED": failure_reason = response_data.get("failure", "Unknown error") raise ValueError(f"Video generation failed: {failure_reason}") else: - raise ValueError( - "Video URL not found in response. Video may not be ready yet." - ) + raise ValueError("Video URL not found in response. Video may not be ready yet.") return video_url @@ -499,9 +479,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): RunwayML uses task cancellation. """ original_video_id = extract_original_video_id(video_id) - encoded_video_id = encode_url_path_segment( - original_video_id, field_name="video_id" - ) + encoded_video_id = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for task cancellation url = f"{api_base}/tasks/{encoded_video_id}/cancel" @@ -540,9 +518,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): RunwayML uses GET /v1/tasks/{task_id} to retrieve task status. """ original_video_id = extract_original_video_id(video_id) - encoded_video_id = encode_url_path_segment( - original_video_id, field_name="video_id" - ) + encoded_video_id = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the full URL for task status retrieval url = f"{api_base}/tasks/{encoded_video_id}" @@ -574,15 +550,11 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Add optional fields if present if "output" in response_data and response_data["output"]: video_data["output_url"] = ( - response_data["output"][0] - if isinstance(response_data["output"], list) - else response_data["output"] + response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] ) if "completedAt" in response_data: - video_data["completed_at"] = self._parse_runway_timestamp( - response_data.get("completedAt") - ) + video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) if "progress" in response_data: video_data["progress"] = response_data["progress"] @@ -596,27 +568,17 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_obj = VideoObject(**video_data) # type: ignore[arg-type] if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider( - video_obj.id, custom_llm_provider, None - ) + video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj - def transform_video_create_character_request( - self, name, video, api_base, litellm_params, headers - ): - raise NotImplementedError( - "video create character is not supported for RunwayML" - ) + def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + raise NotImplementedError("video create character is not supported for RunwayML") def transform_video_create_character_response(self, raw_response, logging_obj): - raise NotImplementedError( - "video create character is not supported for RunwayML" - ) + raise NotImplementedError("video create character is not supported for RunwayML") - def transform_video_get_character_request( - self, character_id, api_base, litellm_params, headers - ): + def transform_video_get_character_request(self, character_id, api_base, litellm_params, headers): raise NotImplementedError("video get character is not supported for RunwayML") def transform_video_get_character_response(self, raw_response, logging_obj): @@ -655,9 +617,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): ): raise NotImplementedError("video extension is not supported for RunwayML") - def transform_video_extension_response( - self, raw_response, logging_obj, custom_llm_provider=None - ): + def transform_video_extension_response(self, raw_response, logging_obj, custom_llm_provider=None): raise NotImplementedError("video extension is not supported for RunwayML") def get_error_class( diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 8270e99d456..b31e6f4511a 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -29,9 +29,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): BaseVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: return {} def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: @@ -40,9 +38,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "write": [], } - def get_supported_openai_params( - self, model: str - ) -> List[VECTOR_STORE_OPENAI_PARAMS]: + def get_supported_openai_params(self, model: str) -> List[VECTOR_STORE_OPENAI_PARAMS]: return ["max_num_results"] def map_openai_params( @@ -56,9 +52,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): optional_params["maxResults"] = value return optional_params - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: headers = headers or {} headers.setdefault("Content-Type", "application/json") return headers @@ -92,9 +86,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): else: # Try to get bucket_name from litellm_params bucket_name_from_params = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance( - bucket_name_from_params, str - ): + if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): raise ValueError( "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " "or vector_bucket_name must be provided in litellm_params" @@ -106,15 +98,11 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): query = " ".join(query) # Generate embedding for the query - embedding_model = litellm_params.get( - "embedding_model", "text-embedding-3-small" - ) + embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") import litellm as litellm_module - embedding_response = litellm_module.embedding( - model=embedding_model, input=[query] - ) + embedding_response = litellm_module.embedding(model=embedding_model, input=[query]) query_embedding = embedding_response.data[0]["embedding"] url = f"{api_base}/QueryVectors" @@ -123,9 +111,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "vectorBucketName": bucket_name, "indexName": index_name, "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get( - "max_num_results", 5 - ), # Default to 5 + "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 "returnDistance": True, "returnMetadata": True, } @@ -154,9 +140,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): else: # Try to get bucket_name from litellm_params bucket_name_from_params = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance( - bucket_name_from_params, str - ): + if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): raise ValueError( "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " "or vector_bucket_name must be provided in litellm_params" @@ -168,15 +152,11 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): query = " ".join(query) # Generate embedding for the query asynchronously - embedding_model = litellm_params.get( - "embedding_model", "text-embedding-3-small" - ) + embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") import litellm as litellm_module - embedding_response = await litellm_module.aembedding( - model=embedding_model, input=[query] - ) + embedding_response = await litellm_module.aembedding(model=embedding_model, input=[query]) query_embedding = embedding_response.data[0]["embedding"] url = f"{api_base}/QueryVectors" @@ -185,9 +165,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "vectorBucketName": bucket_name, "indexName": index_name, "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get( - "max_num_results", 5 - ), # Default to 5 + "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 "returnDistance": True, "returnMetadata": True, } @@ -246,9 +224,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): results.append( VectorStoreSearchResult( score=score, - content=[ - VectorStoreResultContent(text=source_text, type="text") - ], + content=[VectorStoreResultContent(text=source_text, type="text")], file_id=file_id, filename=filename, attributes=metadata, diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index b86cda7aeaf..c01e93c4bf4 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -30,9 +30,7 @@ class SagemakerChatHandler(BaseAWSLLM): aws_role_name = optional_params.pop("aws_role_name", None) aws_session_name = optional_params.pop("aws_session_name", None) aws_profile_name = optional_params.pop("aws_profile_name", None) - optional_params.pop( - "aws_bedrock_runtime_endpoint", None - ) # https://bedrock-runtime.{region_name}.amazonaws.com + optional_params.pop("aws_bedrock_runtime_endpoint", None) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None) @@ -41,15 +39,11 @@ class SagemakerChatHandler(BaseAWSLLM): # check env # litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) - if litellm_aws_region_name is not None and isinstance( - litellm_aws_region_name, str - ): + if litellm_aws_region_name is not None and isinstance(litellm_aws_region_name, str): aws_region_name = litellm_aws_region_name standard_aws_region_name = get_secret("AWS_REGION", None) - if standard_aws_region_name is not None and isinstance( - standard_aws_region_name, str - ): + if standard_aws_region_name is not None and isinstance(standard_aws_region_name, str): aws_region_name = standard_aws_region_name if aws_region_name is None: @@ -97,9 +91,7 @@ class SagemakerChatHandler(BaseAWSLLM): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - request = AWSRequest( - method="POST", url=api_base, data=encoded_data, headers=headers - ) + request = AWSRequest(method="POST", url=api_base, data=encoded_data, headers=headers) sigv4.add_auth(request) if ( extra_headers is not None and "Authorization" in extra_headers diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 3e42c1e8c15..4e4e088f491 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -41,12 +41,8 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): OpenAIGPTConfig.__init__(self, **kwargs) BaseAWSLLM.__init__(self, **kwargs) - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return SagemakerError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return SagemakerError(status_code=status_code, message=error_message, headers=headers) def validate_environment( self, @@ -79,9 +75,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): else: api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" - sagemaker_base_url = cast( - Optional[str], optional_params.get("sagemaker_base_url") - ) + sagemaker_base_url = cast(Optional[str], optional_params.get("sagemaker_base_url")) if sagemaker_base_url is not None: api_base = sagemaker_base_url @@ -143,19 +137,13 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): logging_obj=logging_obj, ) except httpx.HTTPStatusError as e: - raise SagemakerError( - status_code=e.response.status_code, message=e.response.text - ) + raise SagemakerError(status_code=e.response.status_code, message=e.response.text) if response.status_code != 200: - raise SagemakerError( - status_code=response.status_code, message=response.text - ) + raise SagemakerError(status_code=response.status_code, message=response.text) custom_stream_decoder = AWSEventStreamDecoder(model="", is_messages_api=True) - completion_stream = custom_stream_decoder.iter_bytes( - response.iter_bytes(chunk_size=1024) - ) + completion_stream = custom_stream_decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -195,19 +183,13 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): logging_obj=logging_obj, ) except httpx.HTTPStatusError as e: - raise SagemakerError( - status_code=e.response.status_code, message=e.response.text - ) + raise SagemakerError(status_code=e.response.status_code, message=e.response.text) if response.status_code != 200: - raise SagemakerError( - status_code=response.status_code, message=response.text - ) + raise SagemakerError(status_code=response.status_code, message=response.text) custom_stream_decoder = AWSEventStreamDecoder(model="", is_messages_api=True) - completion_stream = custom_stream_decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) - ) + completion_stream = custom_stream_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 6c15d642f8c..2fddde291f4 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -18,9 +18,7 @@ def _load_sagemaker_response_stream_shape(): loader = Loader() service_dict = loader.load_service_model("sagemaker-runtime", "service-2") - return ServiceModel(service_dict).shape_for( - "InvokeEndpointWithResponseStreamOutput" - ) + return ServiceModel(service_dict).shape_for("InvokeEndpointWithResponseStreamOutput") except Exception as e: verbose_logger.warning( "litellm: could not load sagemaker-runtime response stream shape " @@ -60,12 +58,8 @@ class AWSEventStreamDecoder: self.content_blocks: List = [] self.is_messages_api = is_messages_api - def _chunk_parser_messages_api( - self, chunk_data: dict - ) -> StreamingChatCompletionChunk: - openai_chunk = StreamingChatCompletionChunk( - **{"model": self.model, **chunk_data} - ) + def _chunk_parser_messages_api(self, chunk_data: dict) -> StreamingChatCompletionChunk: + openai_chunk = StreamingChatCompletionChunk(**{"model": self.model, **chunk_data}) return openai_chunk @@ -94,9 +88,7 @@ class AWSEventStreamDecoder: usage=None, ) - def iter_bytes( - self, iterator: Iterator[bytes] - ) -> Iterator[Optional[Union[GChunk, StreamingChatCompletionChunk]]]: + def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[Optional[Union[GChunk, StreamingChatCompletionChunk]]]: """Given an iterator that yields lines, iterate over it & yield every event encountered""" from botocore.eventstream import EventStreamBuffer @@ -109,10 +101,7 @@ class AWSEventStreamDecoder: message = self._parse_message_from_event(event) if message: # remove data: prefix and "\n\n" at the end - message = ( - litellm.CustomStreamWrapper._strip_sse_data_from_chunk(message) - or "" - ) + message = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(message) or "" message = message.replace("\n\n", "") # Accumulate JSON data @@ -141,9 +130,7 @@ class AWSEventStreamDecoder: yield self._chunk_parser(chunk_data=_data) except json.JSONDecodeError: # Handle or log any unparseable data at the end - verbose_logger.error( - f"Warning: Unparseable JSON data remained: {accumulated_json}" - ) + verbose_logger.error(f"Warning: Unparseable JSON data remained: {accumulated_json}") yield None async def aiter_bytes( @@ -161,16 +148,9 @@ class AWSEventStreamDecoder: try: message = self._parse_message_from_event(event) if message: - verbose_logger.debug( - "sagemaker parsed chunk bytes %s", message - ) + verbose_logger.debug("sagemaker parsed chunk bytes %s", message) # remove data: prefix and "\n\n" at the end - message = ( - litellm.CustomStreamWrapper._strip_sse_data_from_chunk( - message - ) - or "" - ) + message = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(message) or "" message = message.replace("\n\n", "") # Accumulate JSON data @@ -188,14 +168,10 @@ class AWSEventStreamDecoder: # If it's not valid JSON yet, continue to the next event continue except UnicodeDecodeError as e: - verbose_logger.warning( - f"UnicodeDecodeError: {e}. Attempting to combine with next event." - ) + verbose_logger.warning(f"UnicodeDecodeError: {e}. Attempting to combine with next event.") continue except Exception as e: - verbose_logger.error( - f"Error parsing message: {e}. Attempting to combine with next event." - ) + verbose_logger.error(f"Error parsing message: {e}. Attempting to combine with next event.") continue # Handle any remaining data after the iterator is exhausted @@ -208,9 +184,7 @@ class AWSEventStreamDecoder: yield self._chunk_parser(chunk_data=_data) except json.JSONDecodeError: # Handle or log any unparseable data at the end - verbose_logger.error( - f"Warning: Unparseable JSON data remained: {accumulated_json}" - ) + verbose_logger.error(f"Warning: Unparseable JSON data remained: {accumulated_json}") yield None except Exception as e: verbose_logger.error(f"Final error parsing accumulated JSON: {e}") diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index aa4663666c2..4b87271fd44 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -52,9 +52,7 @@ class SagemakerLLM(BaseAWSLLM): aws_role_name = optional_params.pop("aws_role_name", None) aws_session_name = optional_params.pop("aws_session_name", None) aws_profile_name = optional_params.pop("aws_profile_name", None) - optional_params.pop( - "aws_bedrock_runtime_endpoint", None - ) # https://bedrock-runtime.{region_name}.amazonaws.com + optional_params.pop("aws_bedrock_runtime_endpoint", None) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None) @@ -63,15 +61,11 @@ class SagemakerLLM(BaseAWSLLM): # check env # litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) - if litellm_aws_region_name is not None and isinstance( - litellm_aws_region_name, str - ): + if litellm_aws_region_name is not None and isinstance(litellm_aws_region_name, str): aws_region_name = litellm_aws_region_name standard_aws_region_name = get_secret("AWS_REGION", None) - if standard_aws_region_name is not None and isinstance( - standard_aws_region_name, str - ): + if standard_aws_region_name is not None and isinstance(standard_aws_region_name, str): aws_region_name = standard_aws_region_name if aws_region_name is None: @@ -125,9 +119,7 @@ class SagemakerLLM(BaseAWSLLM): optional_params=optional_params, litellm_params=litellm_params, ) - request = AWSRequest( - method="POST", url=api_base, data=encoded_data, headers=headers - ) + request = AWSRequest(method="POST", url=api_base, data=encoded_data, headers=headers) sigv4.add_auth(request) if ( extra_headers is not None and "Authorization" in extra_headers @@ -207,9 +199,7 @@ class SagemakerLLM(BaseAWSLLM): if model_id is not None: # Add model_id as InferenceComponentName header # boto3 doc: https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html - prepared_request.headers.update( - {"X-Amzn-SageMaker-Inference-Component": model_id} - ) + prepared_request.headers.update({"X-Amzn-SageMaker-Inference-Component": model_id}) sync_handler = _get_httpx_client() sync_response = sync_handler.post( url=prepared_request.url, @@ -226,9 +216,7 @@ class SagemakerLLM(BaseAWSLLM): decoder = AWSEventStreamDecoder(model="") - completion_stream = decoder.iter_bytes( - sync_response.iter_bytes(chunk_size=1024) - ) + completion_stream = decoder.iter_bytes(sync_response.iter_bytes(chunk_size=1024)) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, @@ -287,9 +275,7 @@ class SagemakerLLM(BaseAWSLLM): if model_id is not None: # Add model_id as InferenceComponentName header # boto3 doc: https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html - prepared_request.headers.update( - {"X-Amzn-SageMaker-Inference-Component": model_id} - ) + prepared_request.headers.update({"X-Amzn-SageMaker-Inference-Component": model_id}) ## LOGGING timeout = 300.0 @@ -330,14 +316,8 @@ class SagemakerLLM(BaseAWSLLM): raise e except Exception as e: verbose_logger.error("Sagemaker error %s", str(e)) - status_code = ( - getattr(e, "response", {}) - .get("ResponseMetadata", {}) - .get("HTTPStatusCode", 500) - ) - error_message = ( - getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) - ) + status_code = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500) + error_message = getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) if "Inference Component Name header is required" in error_message: error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`" raise SagemakerError(status_code=status_code, message=error_message) @@ -375,14 +355,10 @@ class SagemakerLLM(BaseAWSLLM): ) if response.status_code != 200: - raise SagemakerError( - status_code=response.status_code, message=response.text - ) + raise SagemakerError(status_code=response.status_code, message=response.text) decoder = AWSEventStreamDecoder(model="") - completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) - ) + completion_stream = decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) return completion_stream @@ -437,9 +413,7 @@ class SagemakerLLM(BaseAWSLLM): } prepared_request = await asyncified_prepare_request(**prepared_request_args) if model_id is not None: # Fixes https://github.com/BerriAI/litellm/issues/8889 - prepared_request.headers.update( - {"X-Amzn-SageMaker-Inference-Component": model_id} - ) + prepared_request.headers.update({"X-Amzn-SageMaker-Inference-Component": model_id}) if not prepared_request.body: raise ValueError("Prepared request body is empty") @@ -484,9 +458,7 @@ class SagemakerLLM(BaseAWSLLM): litellm_params: dict, ): timeout = 300.0 - async_handler = get_async_httpx_client( - llm_provider=litellm.LlmProviders.SAGEMAKER - ) + async_handler = get_async_httpx_client(llm_provider=litellm.LlmProviders.SAGEMAKER) data = await sagemaker_config.async_transform_request( model=model, @@ -522,9 +494,7 @@ class SagemakerLLM(BaseAWSLLM): if model_id is not None: # Add model_id as InferenceComponentName header # boto3 doc: https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html - prepared_request.headers.update( - {"X-Amzn-SageMaker-Inference-Component": model_id} - ) + prepared_request.headers.update({"X-Amzn-SageMaker-Inference-Component": model_id}) # make async httpx post request here try: response = await async_handler.post( @@ -535,9 +505,7 @@ class SagemakerLLM(BaseAWSLLM): ) if response.status_code != 200: - raise SagemakerError( - status_code=response.status_code, message=response.text - ) + raise SagemakerError(status_code=response.status_code, message=response.text) except Exception as e: ## LOGGING logging_obj.post_call( @@ -610,9 +578,7 @@ class SagemakerLLM(BaseAWSLLM): #### EMBEDDING LOGIC # Transform request based on model type provider_config = SagemakerEmbeddingConfig.get_model_config(model) - request_data = provider_config.transform_embedding_request( - model, input, optional_params, {} - ) + request_data = provider_config.transform_embedding_request(model, input, optional_params, {}) data = json.dumps(request_data).encode("utf-8") ## LOGGING @@ -637,14 +603,8 @@ class SagemakerLLM(BaseAWSLLM): CustomAttributes="accept_eula=true", ) except Exception as e: - status_code = ( - getattr(e, "response", {}) - .get("ResponseMetadata", {}) - .get("HTTPStatusCode", 500) - ) - error_message = ( - getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) - ) + status_code = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 500) + error_message = getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) raise SagemakerError(status_code=status_code, message=error_message) response = json.loads(response["Body"].read().decode("utf8")) diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index 8fd32bc4460..918af7f586d 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -60,12 +60,8 @@ class SagemakerConfig(BaseConfig): def get_config(cls): return super().get_config() - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return SagemakerError( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return SagemakerError(message=error_message, status_code=status_code, headers=headers) def get_supported_openai_params(self, model: str) -> List: return [ @@ -90,9 +86,7 @@ class SagemakerConfig(BaseConfig): if value == 0.0 or value == 0: # hugging face exception raised when temp==0 # Failed: Error occurred: HuggingfaceException - Input validation error: `temperature` must be strictly positive - if not non_default_params.get( - "aws_sagemaker_allow_zero_temp", False - ): + if not non_default_params.get("aws_sagemaker_allow_zero_temp", False): value = 0.01 optional_params["temperature"] = value @@ -100,9 +94,7 @@ class SagemakerConfig(BaseConfig): optional_params["top_p"] = value if param == "n": optional_params["best_of"] = value - optional_params["do_sample"] = ( - True # Need to sample if you want best of for hf inference endpoints - ) + optional_params["do_sample"] = True # Need to sample if you want best of for hf inference endpoints if param == "stream": optional_params["stream"] = value if param == "stop": @@ -130,9 +122,7 @@ class SagemakerConfig(BaseConfig): model_prompt_details = custom_prompt_dict[model] prompt = custom_prompt( role_dict=model_prompt_details.get("roles", None), - initial_prompt_value=model_prompt_details.get( - "initial_prompt_value", "" - ), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), final_prompt_value=model_prompt_details.get("final_prompt_value", ""), messages=messages, ) @@ -141,9 +131,7 @@ class SagemakerConfig(BaseConfig): model_prompt_details = custom_prompt_dict[hf_model_name] prompt = custom_prompt( role_dict=model_prompt_details.get("roles", None), - initial_prompt_value=model_prompt_details.get( - "initial_prompt_value", "" - ), + initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), final_prompt_value=model_prompt_details.get("final_prompt_value", ""), messages=messages, ) @@ -175,9 +163,7 @@ class SagemakerConfig(BaseConfig): if stream is True: data["stream"] = True - custom_prompt_dict = ( - litellm_params.get("custom_prompt_dict", None) or litellm.custom_prompt_dict - ) + custom_prompt_dict = litellm_params.get("custom_prompt_dict", None) or litellm.custom_prompt_dict hf_model_name = litellm_params.get("hf_model_name", None) @@ -199,9 +185,7 @@ class SagemakerConfig(BaseConfig): litellm_params: dict, headers: dict, ) -> dict: - return await asyncify(self.transform_request)( - model, messages, optional_params, litellm_params, headers - ) + return await asyncify(self.transform_request)(model, messages, optional_params, litellm_params, headers) def transform_response( self, diff --git a/litellm/llms/sagemaker/embedding/cohere_transformation.py b/litellm/llms/sagemaker/embedding/cohere_transformation.py index fdb67202ebb..126f153222d 100644 --- a/litellm/llms/sagemaker/embedding/cohere_transformation.py +++ b/litellm/llms/sagemaker/embedding/cohere_transformation.py @@ -55,12 +55,8 @@ class SagemakerCohereEmbeddingConfig(BaseEmbeddingConfig): optional_params["input_type"] = non_default_params["input_type"] return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return SagemakerError( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return SagemakerError(message=error_message, status_code=status_code, headers=headers) def transform_embedding_request( self, @@ -109,10 +105,7 @@ class SagemakerCohereEmbeddingConfig(BaseEmbeddingConfig): invoking this transform. """ input_value = ( - logging_obj.model_call_details.get("input") - or request_data.get("texts") - or request_data.get("images") - or [] + logging_obj.model_call_details.get("input") or request_data.get("texts") or request_data.get("images") or [] ) if isinstance(input_value, str): input_value = [input_value] diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 5e2aa99534f..fce2bfd22e7 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -63,12 +63,8 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): ) -> dict: return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return SagemakerError( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return SagemakerError(message=error_message, status_code=status_code, headers=headers) def transform_embedding_request( self, @@ -126,9 +122,7 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): output_data = [] for idx, embedding in enumerate(embeddings): - output_data.append( - {"object": "embedding", "index": idx, "embedding": embedding} - ) + output_data.append({"object": "embedding", "index": idx, "embedding": embedding}) model_response.object = "list" model_response.data = output_data diff --git a/litellm/llms/sambanova/embedding/transformation.py b/litellm/llms/sambanova/embedding/transformation.py index 5c88188b84e..611507bcf0d 100644 --- a/litellm/llms/sambanova/embedding/transformation.py +++ b/litellm/llms/sambanova/embedding/transformation.py @@ -135,6 +135,4 @@ class SambaNovaEmbeddingConfig(BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return SambaNovaError( - message=error_message, status_code=status_code, headers=headers - ) + return SambaNovaError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/sap/chat/handler.py b/litellm/llms/sap/chat/handler.py index 713143d895f..a679e4cf704 100755 --- a/litellm/llms/sap/chat/handler.py +++ b/litellm/llms/sap/chat/handler.py @@ -139,9 +139,7 @@ class SAPStreamIterator: if not line: continue - payload = ( - line[len(self._prefix) :] if line.startswith(self._prefix) else line - ) + payload = line[len(self._prefix) :] if line.startswith(self._prefix) else line if payload == self._final: self._safe_close() raise StopIteration @@ -213,9 +211,7 @@ class AsyncSAPStreamIterator: continue # now = lambda: int(time.time() * 1000) - payload = ( - line[len(self._prefix) :] if line.startswith(self._prefix) else line - ) + payload = line[len(self._prefix) :] if line.startswith(self._prefix) else line if payload == self._final: await self._aclose() raise StopAsyncIteration @@ -250,9 +246,7 @@ class AsyncSAPStreamIterator: # LLM handler # ------------------------------- class GenAIHubOrchestration(BaseLLMHTTPHandler): - def _add_stream_param_to_request_body( - self, data: dict, provider_config: BaseConfig, fake_stream: bool - ): + def _add_stream_param_to_request_body(self, data: dict, provider_config: BaseConfig, fake_stream: bool): if data.get("config", {}).get("stream", None) is not None: data["config"]["stream"]["enabled"] = True else: diff --git a/litellm/llms/sap/chat/models.py b/litellm/llms/sap/chat/models.py index a107901de76..2756dd0e67e 100644 --- a/litellm/llms/sap/chat/models.py +++ b/litellm/llms/sap/chat/models.py @@ -90,16 +90,12 @@ class SAPMessage(BaseModel): role: Literal["system", "developer"] = "system" content: str - _content_validator = field_validator("content", mode="before")( - validate_different_content - ) + _content_validator = field_validator("content", mode="before")(validate_different_content) class SAPUserMessage(BaseModel): role: Literal["user"] = "user" - content: Union[ - str, TextContent, ImageContent, list[Union[TextContent, ImageContent]] - ] + content: Union[str, TextContent, ImageContent, list[Union[TextContent, ImageContent]]] class SAPAssistantMessage(BaseModel): @@ -108,9 +104,7 @@ class SAPAssistantMessage(BaseModel): refusal: str = "" tool_calls: list[MessageToolCall] = [] - _content_validator = field_validator("content", mode="before")( - validate_different_content - ) + _content_validator = field_validator("content", mode="before")(validate_different_content) class SAPToolChatMessage(BaseModel): @@ -118,9 +112,7 @@ class SAPToolChatMessage(BaseModel): tool_call_id: str content: str - _content_validator = field_validator("content", mode="before")( - validate_different_content - ) + _content_validator = field_validator("content", mode="before")(validate_different_content) ChatMessage = Union[SAPMessage, SAPUserMessage, SAPAssistantMessage, SAPToolChatMessage] @@ -184,9 +176,7 @@ class DocumentGroundingConfig(BaseModel): class GroundingModuleConfig(BaseModel): - type_: Literal["document_grounding_service"] = Field( - default="document_grounding_service", alias="type" - ) + type_: Literal["document_grounding_service"] = Field(default="document_grounding_service", alias="type") config: DocumentGroundingConfig @@ -329,9 +319,7 @@ class DPIStandardEntity(BaseModel): """ type_: SAPMaskingProfileEntity = Field(..., alias="type") - replacement_strategy: Optional[ - Union[DPIMethodConstant, DPIMethodFabricatedData] - ] = None + replacement_strategy: Optional[Union[DPIMethodConstant, DPIMethodFabricatedData]] = None class MaskGroundingInput(BaseModel): @@ -361,9 +349,7 @@ class MaskingProviderConfig(BaseModel): mask_grounding_input: A flag indicating whether to mask input to the grounding module. """ - type_: Literal["sap_data_privacy_integration"] = Field( - default="sap_data_privacy_integration", alias="type" - ) + type_: Literal["sap_data_privacy_integration"] = Field(default="sap_data_privacy_integration", alias="type") method: Literal["anonymization", "pseudonymization"] entities: list[Union[DPIStandardEntity, DPICustomEntity]] allowlist: Optional[list[str]] = None @@ -382,9 +368,7 @@ class MaskingModuleConfig(BaseModel): """ providers: Optional[list[MaskingProviderConfig]] = Field(min_length=1, default=None) - masking_providers: Optional[list[MaskingProviderConfig]] = Field( - min_length=1, default=None - ) + masking_providers: Optional[list[MaskingProviderConfig]] = Field(min_length=1, default=None) @model_validator(mode="after") def enforce_exactly_one_provider_list(self): @@ -392,9 +376,7 @@ class MaskingModuleConfig(BaseModel): has_masking_providers = self.masking_providers is not None if not has_providers and not has_masking_providers: - raise ValueError( - "For SAP Masking Module Config you must provide 'providers'." - ) + raise ValueError("For SAP Masking Module Config you must provide 'providers'.") if has_providers and has_masking_providers: raise ValueError( "For SAP Masking Module Config you must set exactly one of: 'providers' or 'masking_providers', not both." @@ -556,16 +538,12 @@ class LlamaGuard38bFilterConfig(BaseModel): class AzureContentSafetyInputFilterConfig(BaseModel): - type_: Literal["azure_content_safety"] = Field( - default="azure_content_safety", alias="type" - ) + type_: Literal["azure_content_safety"] = Field(default="azure_content_safety", alias="type") config: Optional[AzureContentSafetyInput] = None class AzureContentSafetyOutputFilterConfig(BaseModel): - type_: Literal["azure_content_safety"] = Field( - default="azure_content_safety", alias="type" - ) + type_: Literal["azure_content_safety"] = Field(default="azure_content_safety", alias="type") config: Optional[AzureContentSafetyOutput] = None @@ -585,9 +563,7 @@ class InputFiltering(BaseModel): filters: List of ContentFilter objects to be applied to input content. """ - filters: list[ - Union[AzureContentSafetyInputFilterConfig, LlamaGuard38bFilterConfig] - ] = Field(min_length=1) + filters: list[Union[AzureContentSafetyInputFilterConfig, LlamaGuard38bFilterConfig]] = Field(min_length=1) class OutputFiltering(BaseModel): @@ -599,9 +575,7 @@ class OutputFiltering(BaseModel): stream_options: Module-specific streaming options. """ - filters: list[ - Union[AzureContentSafetyOutputFilterConfig, LlamaGuard38bFilterConfig] - ] = Field(min_length=1) + filters: list[Union[AzureContentSafetyOutputFilterConfig, LlamaGuard38bFilterConfig]] = Field(min_length=1) stream_options: Optional[FilteringStreamOptions] = None @@ -677,9 +651,7 @@ class SAPDocumentTranslationInput(BaseModel): config: Configuration object for the translation module. """ - type_: Literal["sap_document_translation"] = Field( - default="sap_document_translation", alias="type" - ) + type_: Literal["sap_document_translation"] = Field(default="sap_document_translation", alias="type") translate_messages_history: Optional[bool] = None config: InputTranslationConfig @@ -694,9 +666,7 @@ class SAPDocumentTranslationOutput(BaseModel): config: Configuration object for the translation module. """ - type_: Literal["sap_document_translation"] = Field( - default="sap_document_translation", alias="type" - ) + type_: Literal["sap_document_translation"] = Field(default="sap_document_translation", alias="type") config: OutputTranslationConfig @@ -716,9 +686,7 @@ class TranslationModuleConfig(BaseModel): @model_validator(mode="after") def enforce_min_properties(self) -> "TranslationModuleConfig": if self.input is None and self.output is None: - raise ValueError( - "TranslationModuleConfig requires at least one of 'input' or 'output'." - ) + raise ValueError("TranslationModuleConfig requires at least one of 'input' or 'output'.") return self diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index 9899e3be9ad..d9d0f9bc236 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -79,9 +79,7 @@ def _messages_to_sap_template(messages: List[Dict[str, str]]) -> list: # type: return template -def _tools_response_format_and_stream( - optional_params: dict, model_params: dict -) -> Tuple[dict, dict, dict]: +def _tools_response_format_and_stream(optional_params: dict, model_params: dict) -> Tuple[dict, dict, dict]: tools_ = optional_params.pop("tools", []) tools_ = [validate_dict(tool, ChatCompletionTool) for tool in tools_] tools: dict = {"tools": tools_} if tools_ else {} @@ -149,9 +147,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): def run_env_setup(self, service_key: Optional[str] = None) -> None: try: - self.token_creator, self._base_url, self._resource_group = ( - get_token_creator(service_key) - ) # type: ignore + self.token_creator, self._base_url, self._resource_group = get_token_creator(service_key) # type: ignore except ValueError as err: raise GenAIHubOrchestrationError(status_code=400, message=err.args[0]) @@ -184,9 +180,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): # Keep a short, tight client lifecycle here to avoid fd leaks client = litellm.module_level_client # with httpx.Client(timeout=30) as client: - deployments = client.get( - f"{self.base_url}/lm/deployments", headers=self.headers - ).json() + deployments = client.get(f"{self.base_url}/lm/deployments", headers=self.headers).json() valid: List[Tuple[str, str]] = [] for dep in deployments.get("resources", []): if dep.get("scenarioId") == "orchestration": @@ -289,9 +283,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): resp_type = response_format.get("type", None) if resp_type: if resp_type == "json_schema": - response_format = validate_dict( - response_format, ResponseFormatJSONSchema - ) + response_format = validate_dict(response_format, ResponseFormatJSONSchema) else: response_format = validate_dict(response_format, ResponseFormat) response_format = {"response_format": response_format} @@ -299,9 +291,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): response_format = {} placeholder_defaults = params.pop("placeholder_defaults", {}) - placeholder_defaults = ( - {"defaults": placeholder_defaults} if placeholder_defaults else {} - ) + placeholder_defaults = {"defaults": placeholder_defaults} if placeholder_defaults else {} optional_modules = {} optional_modules_lst = ["grounding", "masking", "filtering", "translation"] @@ -365,9 +355,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): modules_dict = dict(modules_dict) fallback_model = modules_dict.pop("model", None) if fallback_model is None: - raise ValueError( - "Each entry in `fallback_sap_modules` must include a 'model' key." - ) + raise ValueError("Each entry in `fallback_sap_modules` must include a 'model' key.") if fallback_model.startswith("sap/"): fallback_model = fallback_model[4:] fallback_template = modules_dict.pop("messages", []) diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index 8cb19f195f2..54e6b1af50e 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -37,9 +37,7 @@ def _get_nested(d: Union[Dict[str, Any], str], path: Sequence[str]) -> Any: try: cur = json.loads(cur) except json.JSONDecodeError: - verbose_logger.warning( - "SAP service key or VCAP service is a string but not valid JSON." - ) + verbose_logger.warning("SAP service key or VCAP service is a string but not valid JSON.") return None for k in path: if not isinstance(cur, dict): @@ -102,36 +100,24 @@ CREDENTIAL_VALUES: Final[List[CredentialsValue]] = [ CredentialsValue( "auth_url", ("url",), - transform_fn=lambda url: ( - url.rstrip("/") - + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX) - ), + transform_fn=lambda url: url.rstrip("/") + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX), ), CredentialsValue( "base_url", ("serviceurls", "AI_API_URL"), - transform_fn=lambda url: ( - url.rstrip("/") + ("" if url.endswith("/v2") else "/v2") - ), + transform_fn=lambda url: url.rstrip("/") + ("" if url.endswith("/v2") else "/v2"), ), CredentialsValue( "cert_url", ("certurl",), - transform_fn=lambda url: ( - url.rstrip("/") - + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX) - ), + transform_fn=lambda url: url.rstrip("/") + ("" if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX), ), # file paths (kept for config compatibility) CredentialsValue("cert_file_path"), CredentialsValue("key_file_path"), # inline PEMs from VCAP - CredentialsValue( - "cert_str", ("certificate",), transform_fn=lambda s: s.replace("\\n", "\n") - ), - CredentialsValue( - "key_str", ("key",), transform_fn=lambda s: s.replace("\\n", "\n") - ), + CredentialsValue("cert_str", ("certificate",), transform_fn=lambda s: s.replace("\\n", "\n")), + CredentialsValue("key_str", ("key",), transform_fn=lambda s: s.replace("\\n", "\n")), ] @@ -148,14 +134,7 @@ def init_conf(profile: Optional[str] = None) -> Dict[str, Any]: cfg_path = ( Path(cfg_env) if cfg_env - else ( - home - / ( - "config.json" - if profile in (None, "", "default") - else f"config_{profile}.json" - ) - ) + else (home / ("config.json" if profile in (None, "", "default") else f"config_{profile}.json")) ) if cfg_path and cfg_path.exists(): @@ -167,9 +146,7 @@ def init_conf(profile: Optional[str] = None) -> Dict[str, Any]: # If an explicit non-default profile was requested but not found, raise. if cfg_env or (profile not in (None, "", "default")): - raise FileNotFoundError( - f"Unable to locate profile config file at '{cfg_path}' in AICORE_HOME '{home}'" - ) + raise FileNotFoundError(f"Unable to locate profile config file at '{cfg_path}' in AICORE_HOME '{home}'") return {} @@ -204,9 +181,7 @@ def resolve_resource_group(sources: List[Source]) -> Optional[str]: for source in sources: value = source.get(rg_cred) if value is not None: - verbose_logger.debug( - f"Resolved GEN AI Hub resource_group from source {source.name}" - ) + verbose_logger.debug(f"Resolved GEN AI Hub resource_group from source {source.name}") return value return rg_cred.default @@ -227,9 +202,7 @@ def _parse_service_key_once( try: return json.loads(service_key) except json.JSONDecodeError: - verbose_logger.warning( - "SAP service key is a string but not valid JSON. Skipping this source." - ) + verbose_logger.warning("SAP service key is a string but not valid JSON. Skipping this source.") return None verbose_logger.warning( f"SAP service key has unexpected type '{type(service_key).__name__}'. Expected str or dict. Ignoring." @@ -242,15 +215,9 @@ def _resolve_credential_from_service_key( ) -> Optional[str]: if service_key is None: return None - val = _str_or_none( - _get_nested( - service_key, (("credentials",) + cv.vcap_key) if cv.vcap_key else (cv.name,) - ) - ) + val = _str_or_none(_get_nested(service_key, (("credentials",) + cv.vcap_key) if cv.vcap_key else (cv.name,))) if val is None: - return _str_or_none( - _get_nested(service_key, cv.vcap_key if cv.vcap_key else (cv.name,)) - ) + return _str_or_none(_get_nested(service_key, cv.vcap_key if cv.vcap_key else (cv.name,))) return val @@ -280,9 +247,7 @@ def fetch_credentials( """ config = init_conf(profile) - service_key = _parse_service_key_once( - service_key or litellm.sap_service_key or os.environ.get(SERVICE_KEY_ENV_VAR) - ) + service_key = _parse_service_key_once(service_key or litellm.sap_service_key or os.environ.get(SERVICE_KEY_ENV_VAR)) vcap_service = _get_vcap_service(VCAP_AICORE_SERVICE_NAME) sources = [ @@ -437,9 +402,7 @@ def get_token_creator( """ # Resolve credentials using your helper - credentials: Dict[str, str] = fetch_credentials( - service_key=service_key, profile=profile, **overrides - ) + credentials: Dict[str, str] = fetch_credentials(service_key=service_key, profile=profile, **overrides) auth_url = credentials.get("auth_url") base_url = credentials.get("base_url") @@ -501,19 +464,13 @@ def get_token_creator( cert_pair=(cert_file_path, key_file_path), ) # Defensive guard: should never reach here due to validate_credentials() - raise ValueError( - "Invalid authentication configuration: no valid credentials found. " - ) + raise ValueError("Invalid authentication configuration: no valid credentials found. ") def get_token() -> str: nonlocal token, token_expiry with lock: now = datetime.now(timezone.utc) - if ( - token is None - or token_expiry is None - or token_expiry - now < timedelta(minutes=expiry_buffer_minutes) - ): + if token is None or token_expiry is None or token_expiry - now < timedelta(minutes=expiry_buffer_minutes): token, token_expiry = _fetch_token() return token diff --git a/litellm/llms/sap/embed/transformation.py b/litellm/llms/sap/embed/transformation.py index 4344c2cc545..8368be718ad 100644 --- a/litellm/llms/sap/embed/transformation.py +++ b/litellm/llms/sap/embed/transformation.py @@ -27,9 +27,7 @@ class Usage(BaseModel): class EmbeddingItem(BaseModel): object: Literal["embedding"] - embedding: List[float] = Field( - ..., description="Vector of floats (length varies by model)." - ) + embedding: List[float] = Field(..., description="Vector of floats (length varies by model).") index: int @@ -102,20 +100,15 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig): def deployment_url(self) -> str: with httpx.Client(timeout=30) as client: valid_deployments = [] - deployments = client.get( - self.base_url + "/lm/deployments", headers=self.headers - ).json() + deployments = client.get(self.base_url + "/lm/deployments", headers=self.headers).json() for deployment in deployments.get("resources", []): if deployment["scenarioId"] == "orchestration": config_details = client.get( - self.base_url - + f"/lm/configurations/{deployment['configurationId']}", + self.base_url + f"/lm/configurations/{deployment['configurationId']}", headers=self.headers, ).json() if config_details["executableId"] == "orchestration": - valid_deployments.append( - (deployment["deploymentUrl"], deployment["createdAt"]) - ) + valid_deployments.append((deployment["deploymentUrl"], deployment["createdAt"])) return sorted(valid_deployments, key=lambda x: x[1], reverse=True)[0][0] def get_error_class(self, error_message, status_code, headers): diff --git a/litellm/llms/scaleway/audio_transcription/transformation.py b/litellm/llms/scaleway/audio_transcription/transformation.py index b45f287afb4..d5438cbf930 100644 --- a/litellm/llms/scaleway/audio_transcription/transformation.py +++ b/litellm/llms/scaleway/audio_transcription/transformation.py @@ -27,9 +27,7 @@ class ScalewayAudioTranscriptionException(BaseLLMException): class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: return [ "language", "prompt", @@ -60,9 +58,7 @@ class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = ( - "https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/") - ) + api_base = "https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/") return f"{api_base}/audio/transcriptions" def get_error_class( @@ -90,8 +86,7 @@ class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig): if not api_key: raise ScalewayAudioTranscriptionException( message=( - "Scaleway API key not found. Pass `api_key=...` or set the " - "SCW_SECRET_KEY environment variable." + "Scaleway API key not found. Pass `api_key=...` or set the SCW_SECRET_KEY environment variable." ), status_code=401, headers={}, diff --git a/litellm/llms/searchapi/search/transformation.py b/litellm/llms/searchapi/search/transformation.py index ae8413684cc..5f3e535d7fd 100644 --- a/litellm/llms/searchapi/search/transformation.py +++ b/litellm/llms/searchapi/search/transformation.py @@ -83,9 +83,7 @@ class SearchAPIConfig(BaseSearchConfig): ) if not api_key: - raise ValueError( - "SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable." - ) + raise ValueError("SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable.") headers["Content-Type"] = "application/json" @@ -103,9 +101,7 @@ class SearchAPIConfig(BaseSearchConfig): SearchAPI.io uses GET requests and includes api_key in query params. """ - api_base = ( - api_base or get_secret_str("SEARCHAPI_API_BASE") or self.SEARCHAPI_API_BASE - ) + api_base = api_base or get_secret_str("SEARCHAPI_API_BASE") or self.SEARCHAPI_API_BASE # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_searchapi_params" in data: @@ -155,9 +151,7 @@ class SearchAPIConfig(BaseSearchConfig): default_api_base=self.SEARCHAPI_API_BASE, ) if not api_key: - raise ValueError( - "SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable." - ) + raise ValueError("SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable.") request_data: SearchAPIRequest = { "engine": "google", @@ -178,9 +172,7 @@ class SearchAPIConfig(BaseSearchConfig): # Convert to multiple "site:domain" clauses domains = optional_params["search_domain_filter"] if isinstance(domains, list) and len(domains) > 0: - result_data["q"] = self._append_domain_filters( - str(result_data["q"]), domains - ) + result_data["q"] = self._append_domain_filters(str(result_data["q"]), domains) if "country" in optional_params: # Map to gl parameter @@ -188,10 +180,7 @@ class SearchAPIConfig(BaseSearchConfig): # Pass through all other SearchAPI.io-specific parameters for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value # Store params in special key for URL building (GET request) diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py index ff68be5709e..b5f41015112 100644 --- a/litellm/llms/searxng/search/transformation.py +++ b/litellm/llms/searxng/search/transformation.py @@ -174,10 +174,7 @@ class SearXNGSearchConfig(BaseSearchConfig): # Pass through all other SearXNG-specific parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value # Store params in special key for GET request URL building diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py index dd43f2d2dc9..31a0d3f2bac 100644 --- a/litellm/llms/serper/search/transformation.py +++ b/litellm/llms/serper/search/transformation.py @@ -63,9 +63,7 @@ class SerperSearchConfig(BaseSearchConfig): default_api_base=self.SERPER_API_BASE, ) if not api_key: - raise ValueError( - "SERPER_API_KEY is not set. Set `SERPER_API_KEY` environment variable." - ) + raise ValueError("SERPER_API_KEY is not set. Set `SERPER_API_KEY` environment variable.") headers["X-API-KEY"] = api_key headers["Content-Type"] = "application/json" return headers @@ -131,10 +129,7 @@ class SerperSearchConfig(BaseSearchConfig): # pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value return result_data diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index ed30522876a..8b23ae135b5 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -146,9 +146,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): anthropic_tools.append(tool) return anthropic_tools - def _extract_system_and_messages( - self, messages: List[AllMessageValues] - ) -> tuple[Optional[str], List[Dict]]: + def _extract_system_and_messages(self, messages: List[AllMessageValues]) -> tuple[Optional[str], List[Dict]]: """ Split messages into system prompt and conversation turns for Anthropic format. @@ -171,50 +169,22 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): if isinstance(content, str) and content: system_parts.append(content) elif isinstance(content, list): - system_parts.append( - "\n".join( - b.get("text", "") - for b in content - if b.get("type") == "text" - ) - ) + system_parts.append("\n".join(b.get("text", "") for b in content if b.get("type") == "text")) elif role == "assistant": - tool_calls = ( - msg.get("tool_calls") - if isinstance(msg, dict) - else getattr(msg, "tool_calls", None) - ) + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) if tool_calls: # type: ignore[truthy-bool] content_blocks: List[Dict[str, Any]] = [] if content: content_blocks.append({"type": "text", "text": content}) for tc in tool_calls: # type: ignore[attr-defined] - func = ( - tc.get("function", {}) - if isinstance(tc, dict) - else getattr(tc, "function", {}) - ) - tc_id = ( - tc.get("id", "") - if isinstance(tc, dict) - else getattr(tc, "id", "") - ) - func_name = ( - func.get("name", "") - if isinstance(func, dict) - else getattr(func, "name", "") - ) + func = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", {}) + tc_id = tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", "") + func_name = func.get("name", "") if isinstance(func, dict) else getattr(func, "name", "") func_args = ( - func.get("arguments", "{}") - if isinstance(func, dict) - else getattr(func, "arguments", "{}") + func.get("arguments", "{}") if isinstance(func, dict) else getattr(func, "arguments", "{}") ) try: - input_data = ( - json.loads(func_args) - if isinstance(func_args, str) - else func_args - ) + input_data = json.loads(func_args) if isinstance(func_args, str) else func_args except (json.JSONDecodeError, TypeError): input_data = {} content_blocks.append( @@ -225,20 +195,14 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): "input": input_data, } ) - conversation.append( - {"role": "assistant", "content": content_blocks} - ) + conversation.append({"role": "assistant", "content": content_blocks}) else: conversation.append({"role": "assistant", "content": content}) elif role == "tool": tool_call_id = ( - msg.get("tool_call_id", "") - if isinstance(msg, dict) - else getattr(msg, "tool_call_id", "") - ) - tool_content = ( - content if isinstance(content, str) else json.dumps(content) + msg.get("tool_call_id", "") if isinstance(msg, dict) else getattr(msg, "tool_call_id", "") ) + tool_content = content if isinstance(content, str) else json.dumps(content) tool_result_block = { "type": "tool_result", "tool_use_id": tool_call_id, @@ -253,9 +217,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): ): conversation[-1]["content"].append(tool_result_block) else: - conversation.append( - {"role": "user", "content": [tool_result_block]} - ) + conversation.append({"role": "user", "content": [tool_result_block]}) else: conversation.append({"role": role, "content": content}) @@ -274,12 +236,8 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): extra_body = optional_params.pop("extra_body", {}) if _is_claude_model(model): - return self._transform_request_anthropic( - model, messages, optional_params, stream, extra_body - ) - return self._transform_request_openai( - model, messages, optional_params, stream, extra_body - ) + return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body) + return self._transform_request_openai(model, messages, optional_params, stream, extra_body) def _transform_request_openai( self, @@ -341,14 +299,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): system, conversation = self._extract_system_and_messages(messages) if "tools" in optional_params: - optional_params["tools"] = self._transform_tools_to_anthropic( - optional_params["tools"] - ) + optional_params["tools"] = self._transform_tools_to_anthropic(optional_params["tools"]) if "tool_choice" in optional_params: - optional_params["tool_choice"] = self._transform_tool_choice_to_anthropic( - optional_params["tool_choice"] - ) + optional_params["tool_choice"] = self._transform_tool_choice_to_anthropic(optional_params["tool_choice"]) max_completion_tokens = optional_params.pop("max_completion_tokens", None) if max_completion_tokens and "max_tokens" not in optional_params: @@ -368,9 +322,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): body["system"] = system if "max_tokens" not in body: - body["max_tokens"] = ( - 4096 # reasonable default; Anthropic API max varies by model - ) + body["max_tokens"] = 4096 # reasonable default; Anthropic API max varies by model return body @@ -392,9 +344,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): return self._transform_response_anthropic( model, raw_response, model_response, logging_obj, request_data, messages ) - return self._transform_response_openai( - model, raw_response, model_response, logging_obj, request_data, messages - ) + return self._transform_response_openai(model, raw_response, model_response, logging_obj, request_data, messages) def _transform_response_openai( self, @@ -466,9 +416,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): "tool_use": "tool_calls", "stop_sequence": "stop", } - finish_reason = _stop_reason_map.get( - response_json.get("stop_reason", "end_turn"), "stop" - ) + finish_reason = _stop_reason_map.get(response_json.get("stop_reason", "end_turn"), "stop") message = Message(content=text_content or None, role="assistant") if tool_calls: @@ -484,8 +432,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): usage = Usage( prompt_tokens=usage_data.get("input_tokens", 0), completion_tokens=usage_data.get("output_tokens", 0), - total_tokens=usage_data.get("input_tokens", 0) - + usage_data.get("output_tokens", 0), + total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0), ) model_response.choices = [choice] diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index 83716f3ef26..44abb66b900 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -64,6 +64,4 @@ class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return SnowflakeException( - message=error_message, status_code=status_code, headers=headers - ) + return SnowflakeException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/soniox/audio_transcription/handler.py b/litellm/llms/soniox/audio_transcription/handler.py index d4774fea460..88fa8f10580 100644 --- a/litellm/llms/soniox/audio_transcription/handler.py +++ b/litellm/llms/soniox/audio_transcription/handler.py @@ -168,15 +168,9 @@ class SonioxAudioTranscriptionHandler: # Pull handler-only kwargs out of params so they aren't sent # to Soniox. - poll_interval = float( - params.pop("soniox_polling_interval", SONIOX_DEFAULT_POLL_INTERVAL) - ) + poll_interval = float(params.pop("soniox_polling_interval", SONIOX_DEFAULT_POLL_INTERVAL)) try: - max_attempts = int( - params.pop( - "soniox_max_polling_attempts", SONIOX_DEFAULT_MAX_POLL_ATTEMPTS - ) - ) + max_attempts = int(params.pop("soniox_max_polling_attempts", SONIOX_DEFAULT_MAX_POLL_ATTEMPTS)) except (ValueError, OverflowError): max_attempts = SONIOX_DEFAULT_MAX_POLL_ATTEMPTS cleanup_raw = params.pop("soniox_cleanup", SONIOX_DEFAULT_CLEANUP) @@ -195,9 +189,7 @@ class SonioxAudioTranscriptionHandler: # SONIOX_MAX_POLL_ATTEMPTS * SONIOX_MAX_POLL_INTERVAL. if not math.isfinite(poll_interval): poll_interval = SONIOX_DEFAULT_POLL_INTERVAL - clamped_poll_interval = max( - SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL) - ) + clamped_poll_interval = max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)) clamped_max_attempts = max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)) handler_opts: Dict[str, Any] = { @@ -273,9 +265,7 @@ class SonioxAudioTranscriptionHandler: additional_args={ "api_base": f"{api_base}/v1/transcriptions", "atranscription": True, - "complete_input_dict": SonioxAudioTranscriptionHandler._redact_body_for_logging( - body - ), + "complete_input_dict": SonioxAudioTranscriptionHandler._redact_body_for_logging(body), }, ) except Exception: @@ -295,11 +285,7 @@ class SonioxAudioTranscriptionHandler: logging_obj.post_call( input=get_audio_file_name(audio_file) if audio_file else None, api_key=api_key, - additional_args={ - "complete_input_dict": SonioxAudioTranscriptionHandler._redact_body_for_logging( - body - ) - }, + additional_args={"complete_input_dict": SonioxAudioTranscriptionHandler._redact_body_for_logging(body)}, original_response=original_response, ) except Exception: @@ -316,11 +302,7 @@ class SonioxAudioTranscriptionHandler: if response.status_code >= 400: try: payload = response.json() - message = ( - payload.get("error_message") - or payload.get("error") - or response.text - ) + message = payload.get("error_message") or payload.get("error") or response.text except Exception: message = response.text raise provider_config.get_error_class( @@ -403,9 +385,7 @@ class SonioxAudioTranscriptionHandler: json=body, timeout=timeout, ) - self._raise_for_response( - create_resp, provider_config, "create transcription" - ) + self._raise_for_response(create_resp, provider_config, "create transcription") transcription_id = create_resp.json()["id"] transcription_meta = self._sync_poll_until_completed( @@ -424,9 +404,7 @@ class SonioxAudioTranscriptionHandler: headers=auth_headers, timeout=timeout, ) - self._raise_for_response( - transcript_resp, provider_config, "fetch transcript" - ) + self._raise_for_response(transcript_resp, provider_config, "fetch transcript") transcript = transcript_resp.json() payload = {"transcription": transcription_meta, "transcript": transcript} @@ -444,9 +422,7 @@ class SonioxAudioTranscriptionHandler: "model": model, "custom_llm_provider": "soniox", "audio_transcription_duration": ( - float(audio_duration_ms) / 1000.0 - if audio_duration_ms is not None - else None + float(audio_duration_ms) / 1000.0 if audio_duration_ms is not None else None ), } ) @@ -641,9 +617,7 @@ class SonioxAudioTranscriptionHandler: json=body, timeout=timeout, ) - self._raise_for_response( - create_resp, provider_config, "create transcription" - ) + self._raise_for_response(create_resp, provider_config, "create transcription") transcription_id = create_resp.json()["id"] transcription_meta = await self._async_poll_until_completed( @@ -662,9 +636,7 @@ class SonioxAudioTranscriptionHandler: headers=auth_headers, timeout=timeout, ) - self._raise_for_response( - transcript_resp, provider_config, "fetch transcript" - ) + self._raise_for_response(transcript_resp, provider_config, "fetch transcript") transcript = transcript_resp.json() payload = {"transcription": transcription_meta, "transcript": transcript} @@ -682,9 +654,7 @@ class SonioxAudioTranscriptionHandler: "model": model, "custom_llm_provider": "soniox", "audio_transcription_duration": ( - float(audio_duration_ms) / 1000.0 - if audio_duration_ms is not None - else None + float(audio_duration_ms) / 1000.0 if audio_duration_ms is not None else None ), } ) diff --git a/litellm/llms/soniox/audio_transcription/transformation.py b/litellm/llms/soniox/audio_transcription/transformation.py index 681d4352dfe..7160d2548df 100644 --- a/litellm/llms/soniox/audio_transcription/transformation.py +++ b/litellm/llms/soniox/audio_transcription/transformation.py @@ -61,9 +61,7 @@ SONIOX_HANDLER_ONLY_PARAMS: List[str] = [ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """Configuration for Soniox async speech-to-text transcription.""" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: # `language` is mapped onto Soniox's `language_hints`. # `response_format` is handled by LiteLLM (Soniox doesn't support # SRT/VTT natively but we synthesize them from token timestamps). @@ -96,12 +94,8 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return SonioxException( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return SonioxException(message=error_message, status_code=status_code, headers=headers) def validate_environment( self, @@ -165,9 +159,7 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): if value is not None: body[key] = value - return AudioTranscriptionRequestData( - data=body, files=None, content_type="application/json" - ) + return AudioTranscriptionRequestData(data=body, files=None, content_type="application/json") def transform_audio_transcription_response( self, @@ -242,9 +234,7 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # Best-effort metadata fields matching OpenAI's verbose_json shape. if transcription_meta.get("audio_duration_ms") is not None: try: - response["duration"] = ( - float(transcription_meta["audio_duration_ms"]) / 1000.0 - ) + response["duration"] = float(transcription_meta["audio_duration_ms"]) / 1000.0 except (TypeError, ValueError): pass diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index 01f8062fc96..76aa25522d0 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -178,9 +178,7 @@ def _group_tokens_into_cues( cues.append( { "start_ms": current_start, - "end_ms": ( - current_end if current_end is not None else current_start - ), + "end_ms": (current_end if current_end is not None else current_start), "text": text, } ) @@ -209,11 +207,7 @@ def _group_tokens_into_cues( should_break = False if len(current_tokens) >= _CUE_MAX_TOKENS: should_break = True - elif ( - current_start is not None - and start_ms is not None - and (start_ms - current_start) >= _CUE_MAX_DURATION_MS - ): + elif current_start is not None and start_ms is not None and (start_ms - current_start) >= _CUE_MAX_DURATION_MS: should_break = True if should_break: diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 9c325c3cdda..05a200246a1 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -80,9 +80,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): if k in param_mapping: # Map param if mapping exists and value is valid if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO: - mapped_params[param_mapping[k]] = ( - OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] - ) # type: ignore + mapped_params[param_mapping[k]] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] # type: ignore # Don't copy "size" itself to final dict elif k == "n": # Store for logic but do not add to outgoing params @@ -161,8 +159,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): if not final_api_key: raise ValueError( - "STABILITY_API_KEY is not set. " - "Please set it via environment variable or pass api_key parameter." + "STABILITY_API_KEY is not set. Please set it via environment variable or pass api_key parameter." ) headers["Authorization"] = f"Bearer {final_api_key}" @@ -312,9 +309,9 @@ class StabilityImageEditConfig(BaseImageEditConfig): model_info = get_model_info(model, custom_llm_provider="stability") cost_per_image = model_info.get("output_cost_per_image", 0) if cost_per_image is not None: - model_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(cost_per_image) + model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float( + cost_per_image + ) return model_response def use_multipart_form_data(self) -> bool: diff --git a/litellm/llms/stability/image_generation/transformation.py b/litellm/llms/stability/image_generation/transformation.py index c8c2a16fcd1..a5b18b0f325 100644 --- a/litellm/llms/stability/image_generation/transformation.py +++ b/litellm/llms/stability/image_generation/transformation.py @@ -45,9 +45,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://api.stability.ai" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Return list of OpenAI params supported by Stability AI. @@ -80,9 +78,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: # Map size to aspect_ratio if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO: - optional_params["aspect_ratio"] = ( - OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] - ) + optional_params["aspect_ratio"] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] elif k == "n": # Store n for later, but don't pass to Stability optional_params["_n"] = v @@ -131,9 +127,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the Stability AI API request. """ - base_url: str = ( - api_base or get_secret_str("STABILITY_API_BASE") or self.DEFAULT_BASE_URL - ) + base_url: str = api_base or get_secret_str("STABILITY_API_BASE") or self.DEFAULT_BASE_URL base_url = base_url.rstrip("/") endpoint = self._get_model_endpoint(model) @@ -156,8 +150,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): if not final_api_key: raise ValueError( - "STABILITY_API_KEY is not set. " - "Please set it via environment variable or pass api_key parameter." + "STABILITY_API_KEY is not set. Please set it via environment variable or pass api_key parameter." ) headers["Authorization"] = f"Bearer {final_api_key}" diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index 647cfb5fa84..51b897d93b2 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -33,9 +33,7 @@ class TavilySearchRequest(_TavilySearchRequestRequired, total=False): include_domains: List[str] # Optional - list of domains to include (max 300) exclude_domains: List[str] # Optional - list of domains to exclude (max 150) topic: str # Optional - category of search ('general', 'news', 'finance'), default 'general' - search_depth: ( - str # Optional - depth of search ('basic', 'advanced'), default 'basic' - ) + search_depth: str # Optional - depth of search ('basic', 'advanced'), default 'basic' include_answer: Union[bool, str] # Optional - include LLM-generated answer include_raw_content: Union[bool, str] # Optional - include raw HTML content include_images: bool # Optional - perform image search @@ -72,9 +70,7 @@ class TavilySearchConfig(BaseSearchConfig): default_api_base=self.TAVILY_API_BASE, ) if not api_key: - raise ValueError( - "TAVILY_API_KEY is not set. Set `TAVILY_API_KEY` environment variable." - ) + raise ValueError("TAVILY_API_KEY is not set. Set `TAVILY_API_KEY` environment variable.") headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -151,10 +147,7 @@ class TavilySearchConfig(BaseSearchConfig): # pass through all other parameters as-is for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value return result_data @@ -189,9 +182,7 @@ class TavilySearchConfig(BaseSearchConfig): search_result = SearchResult( title=result.get("title", ""), url=result.get("url", ""), - snippet=result.get( - "content", "" - ), # Tavily uses "content" instead of "snippet" + snippet=result.get("content", ""), # Tavily uses "content" instead of "snippet" date=None, # Tavily doesn't provide date in response last_updated=None, # Tavily doesn't provide last_updated in response ) diff --git a/litellm/llms/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index 4a95f519646..4b7d38e3661 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -75,9 +75,7 @@ class TinyfishSearchConfig(BaseSearchConfig): default_api_base=self.TINYFISH_API_BASE, ) if not resolved_key: - raise ValueError( - "TINYFISH_API_KEY is not set. Set `TINYFISH_API_KEY` environment variable." - ) + raise ValueError("TINYFISH_API_KEY is not set. Set `TINYFISH_API_KEY` environment variable.") return {**headers, "X-API-Key": resolved_key, "Accept": "application/json"} def get_complete_url( @@ -87,13 +85,9 @@ class TinyfishSearchConfig(BaseSearchConfig): data: dict[str, object] | list[dict[str, object]] | None = None, **kwargs: object, ) -> str: - resolved_base = ( - api_base or get_secret_str("TINYFISH_API_BASE") or self.TINYFISH_API_BASE - ) + resolved_base = api_base or get_secret_str("TINYFISH_API_BASE") or self.TINYFISH_API_BASE if isinstance(data, dict) and _TINYFISH_PARAMS_KEY in data: - validated_params = _UrlEncodableParams.validate_python( - data[_TINYFISH_PARAMS_KEY] - ) + validated_params = _UrlEncodableParams.validate_python(data[_TINYFISH_PARAMS_KEY]) return f"{resolved_base}?{urlencode(validated_params, doseq=True)}" return resolved_base @@ -116,15 +110,11 @@ class TinyfishSearchConfig(BaseSearchConfig): request_data["max_results"] = max(1, min(int(raw_max), 20)) try: - domains = _StrList.validate_python( - optional_params.get("search_domain_filter") - ) + domains = _StrList.validate_python(optional_params.get("search_domain_filter")) except (ValidationError, TypeError): domains = [] if domains: - request_data["query"] = _append_domain_filters( - request_data["query"], domains - ) + request_data["query"] = _append_domain_filters(request_data["query"], domains) result_data: dict[str, object] = dict(request_data) @@ -156,8 +146,7 @@ class TinyfishSearchConfig(BaseSearchConfig): max_results: int = min(int(max_results_str), 20) results = [ - SearchResult(title=item.title, url=item.url, snippet=item.snippet) - for item in parsed.results[:max_results] + SearchResult(title=item.title, url=item.url, snippet=item.snippet) for item in parsed.results[:max_results] ] return SearchResponse(results=results, object="search") diff --git a/litellm/llms/together_ai/chat.py b/litellm/llms/together_ai/chat.py index 238849cc1ec..a78b023f287 100644 --- a/litellm/llms/together_ai/chat.py +++ b/litellm/llms/together_ai/chat.py @@ -29,9 +29,7 @@ class TogetherAIConfig(OpenAIGPTConfig): # exception in _get_model_info_helper is hit (~332 deep calls). supports_fc: Optional[bool] = None try: - supports_fc = supports_function_calling( - model, custom_llm_provider="together_ai" - ) + supports_fc = supports_function_calling(model, custom_llm_provider="together_ai") except Exception as e: verbose_logger.debug(f"Error getting supported openai params: {e}") pass @@ -54,12 +52,8 @@ class TogetherAIConfig(OpenAIGPTConfig): model: str, drop_params: bool, ) -> dict: - mapped_openai_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + mapped_openai_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) - if "response_format" in mapped_openai_params and mapped_openai_params[ - "response_format" - ] == {"type": "text"}: + if "response_format" in mapped_openai_params and mapped_openai_params["response_format"] == {"type": "text"}: mapped_openai_params.pop("response_format") return mapped_openai_params diff --git a/litellm/llms/together_ai/completion/transformation.py b/litellm/llms/together_ai/completion/transformation.py index 8b9dc750c63..6e0b862c183 100644 --- a/litellm/llms/together_ai/completion/transformation.py +++ b/litellm/llms/together_ai/completion/transformation.py @@ -29,15 +29,9 @@ class TogetherAITextCompletionConfig(OpenAITextCompletionConfig): """ initial_prompt: AllPromptValues = _transform_prompt(messages) ## TOGETHER AI SPECIFIC VALIDATION ## - if isinstance(initial_prompt, list) and is_tokens_or_list_of_tokens( - value=initial_prompt - ): + if isinstance(initial_prompt, list) and is_tokens_or_list_of_tokens(value=initial_prompt): raise ValueError("TogetherAI does not support integers as input") - if ( - isinstance(initial_prompt, list) - and len(initial_prompt) == 1 - and isinstance(initial_prompt[0], str) - ): + if isinstance(initial_prompt, list) and len(initial_prompt) == 1 and isinstance(initial_prompt[0], str): together_prompt = initial_prompt[0] elif isinstance(initial_prompt, list): raise ValueError("TogetherAI does not support multiple prompts.") diff --git a/litellm/llms/together_ai/cost_calculator.py b/litellm/llms/together_ai/cost_calculator.py index a1be097bc86..191521266e7 100644 --- a/litellm/llms/together_ai/cost_calculator.py +++ b/litellm/llms/together_ai/cost_calculator.py @@ -29,9 +29,7 @@ def get_model_params_and_category(model_name, call_type: CallTypes) -> str: if call_type == CallTypes.embedding or call_type == CallTypes.aembedding: return get_model_params_and_category_embeddings(model_name=model_name) model_name = model_name.lower() - re_params_match = re.search( - r"(\d+b)", model_name - ) # catch all decimals like 3b, 70b, etc + re_params_match = re.search(r"(\d+b)", model_name) # catch all decimals like 3b, 70b, etc category = None if re_params_match is not None: params_match = str(re_params_match.group(1)) @@ -67,9 +65,7 @@ def get_model_params_and_category_embeddings(model_name) -> str: - str - model pricing category if mapped else received model name """ model_name = model_name.lower() - re_params_match = re.search( - r"(\d+m)", model_name - ) # catch all decimals like 100m, 200m, etc. + re_params_match = re.search(r"(\d+m)", model_name) # catch all decimals like 100m, 200m, etc. category = None if re_params_match is not None: params_match = str(re_params_match.group(1)) diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py index c5b02731e1e..08acdead386 100644 --- a/litellm/llms/together_ai/rerank/handler.py +++ b/litellm/llms/together_ai/rerank/handler.py @@ -70,9 +70,7 @@ class TogetherAIRerank(BaseLLM): request_data_dict: Dict[str, Any], api_key: str, ) -> RerankResponse: - client = get_async_httpx_client( - llm_provider=litellm.LlmProviders.TOGETHER_AI - ) # Use async client + client = get_async_httpx_client(llm_provider=litellm.LlmProviders.TOGETHER_AI) # Use async client response = await client.post( "https://api.together.xyz/v1/rerank", diff --git a/litellm/llms/together_ai/rerank/transformation.py b/litellm/llms/together_ai/rerank/transformation.py index f4d642bd25a..3610a5853ac 100644 --- a/litellm/llms/together_ai/rerank/transformation.py +++ b/litellm/llms/together_ai/rerank/transformation.py @@ -37,11 +37,7 @@ class TogetherAIRerankConfig: # Get document data if it exists document_data = result.get("document", {}) - document = ( - RerankResponseDocument(text=str(document_data.get("text", ""))) - if document_data - else None - ) + document = RerankResponseDocument(text=str(document_data.get("text", ""))) if document_data else None # Create typed result rerank_result = RerankResponseResult( diff --git a/litellm/llms/topaz/common_utils.py b/litellm/llms/topaz/common_utils.py index 95fe2914934..27603b3b401 100644 --- a/litellm/llms/topaz/common_utils.py +++ b/litellm/llms/topaz/common_utils.py @@ -23,18 +23,14 @@ class TopazModelInfo(BaseLLMModelInfo): api_base: Optional[str] = None, ) -> dict: if api_key is None: - raise ValueError( - "API key is required for Topaz image variations. Set via `TOPAZ_API_KEY` or `api_key=..`" - ) + raise ValueError("API key is required for Topaz image variations. Set via `TOPAZ_API_KEY` or `api_key=..`") return { # "Content-Type": "multipart/form-data", "Accept": "image/jpeg", "X-API-Key": api_key, } - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: return [ "topaz/Standard V2", "topaz/Low Resolution V2", @@ -49,9 +45,7 @@ class TopazModelInfo(BaseLLMModelInfo): @staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: - return ( - api_base or get_secret_str("TOPAZ_API_BASE") or "https://api.topazlabs.com" - ) + return api_base or get_secret_str("TOPAZ_API_BASE") or "https://api.topazlabs.com" @staticmethod def get_base_model(model: str) -> str: diff --git a/litellm/llms/topaz/image_variations/transformation.py b/litellm/llms/topaz/image_variations/transformation.py index 41b51a558c5..01239d600b6 100644 --- a/litellm/llms/topaz/image_variations/transformation.py +++ b/litellm/llms/topaz/image_variations/transformation.py @@ -23,9 +23,7 @@ from ..common_utils import TopazException, TopazModelInfo class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageVariationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageVariationOptionalParams]: return ["response_format", "size"] def get_complete_url( @@ -144,9 +142,7 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): response_ms = logging_obj.get_response_ms() - return self._common_transform_response_image_variation( - image_content, response_ms - ) + return self._common_transform_response_image_variation(image_content, response_ms) def transform_response_image_variation( self, @@ -163,17 +159,11 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): ) -> ImageResponse: image_content = raw_response.content - response_ms = ( - raw_response.elapsed.total_seconds() * 1000 - ) # Convert to milliseconds + response_ms = raw_response.elapsed.total_seconds() * 1000 # Convert to milliseconds - return self._common_transform_response_image_variation( - image_content, response_ms - ) + return self._common_transform_response_image_variation(image_content, response_ms) - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: return TopazException( status_code=status_code, message=error_message, diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index 0db83b2d3de..44fe32e2e5d 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -35,12 +35,8 @@ class TritonConfig(BaseConfig): Handles routing between /infer and /generate triton completion llms """ - def get_error_class( - self, error_message: str, status_code: int, headers: Union[Dict, Headers] - ) -> BaseLLMException: - return TritonError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[Dict, Headers]) -> BaseLLMException: + return TritonError(status_code=status_code, message=error_message, headers=headers) def validate_environment( self, @@ -198,9 +194,7 @@ class TritonGenerateConfig(TritonConfig): data_for_triton: Dict[str, Any] = { "text_input": prompt_factory(model=model, messages=messages), "parameters": { - "max_tokens": int( - optional_params.get("max_tokens", DEFAULT_MAX_TOKENS_FOR_TRITON) - ), + "max_tokens": int(optional_params.get("max_tokens", DEFAULT_MAX_TOKENS_FOR_TRITON)), }, "stream": bool(stream), } @@ -224,12 +218,8 @@ class TritonGenerateConfig(TritonConfig): try: raw_response_json = raw_response.json() except Exception: - raise TritonError( - message=raw_response.text, status_code=raw_response.status_code - ) - model_response.choices = [ - Choices(index=0, message=Message(content=raw_response_json["text_output"])) - ] + raise TritonError(message=raw_response.text, status_code=raw_response.status_code) + model_response.choices = [Choices(index=0, message=Message(content=raw_response_json["text_output"]))] return model_response @@ -263,9 +253,7 @@ class TritonInferConfig(TritonConfig): if not (k == "stream" or k == "max_retries"): datatype = "INT32" if isinstance(v, int) else "BYTES" datatype = "FP32" if isinstance(v, float) else datatype - data_for_triton["inputs"].append( - {"name": k, "shape": [1], "datatype": datatype, "data": [v]} - ) + data_for_triton["inputs"].append({"name": k, "shape": [1], "datatype": datatype, "data": [v]}) if "max_tokens" not in optional_params: data_for_triton["inputs"].append( @@ -295,9 +283,7 @@ class TritonInferConfig(TritonConfig): try: raw_response_json = raw_response.json() except Exception: - raise TritonError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise TritonError(message=raw_response.text, status_code=raw_response.status_code) _triton_response_data = raw_response_json["outputs"][0]["data"] triton_response_data: Optional[str] = None diff --git a/litellm/llms/triton/embedding/transformation.py b/litellm/llms/triton/embedding/transformation.py index 93d1c25f169..2426520e630 100644 --- a/litellm/llms/triton/embedding/transformation.py +++ b/litellm/llms/triton/embedding/transformation.py @@ -81,9 +81,7 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): try: raw_response_json = raw_response.json() except Exception: - raise TritonError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise TritonError(message=raw_response.text, status_code=raw_response.status_code) _embedding_output = [] @@ -104,9 +102,7 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): model_response.model = raw_response_json.get("model_name", "None") model_response.data = _embedding_output - model_response.usage = self._build_embedding_usage( - model=model, request_data=request_data - ) + model_response.usage = self._build_embedding_usage(model=model, request_data=request_data) return model_response def _build_embedding_usage(self, model: str, request_data: dict) -> Usage: @@ -137,17 +133,11 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return TritonError( - message=error_message, status_code=status_code, headers=headers - ) + return TritonError(message=error_message, status_code=status_code, headers=headers) @staticmethod - def split_embedding_by_shape( - data: List[float], shape: List[int] - ) -> List[List[float]]: + def split_embedding_by_shape(data: List[float], shape: List[int]) -> List[List[float]]: if len(shape) != 2: raise ValueError("Shape must be of length 2.") embedding_size = shape[1] - return [ - data[i * embedding_size : (i + 1) * embedding_size] for i in range(shape[0]) - ] + return [data[i * embedding_size : (i + 1) * embedding_size] for i in range(shape[0])] diff --git a/litellm/llms/v0/chat/transformation.py b/litellm/llms/v0/chat/transformation.py index 7b65cec9d39..5e029512471 100644 --- a/litellm/llms/v0/chat/transformation.py +++ b/litellm/llms/v0/chat/transformation.py @@ -23,9 +23,7 @@ class V0ChatConfig(OpenAILikeChatConfig): ) -> Tuple[Optional[str], Optional[str]]: # v0 is openai compatible, we just need to set the api_base api_base = ( - api_base - or get_secret_str("V0_API_BASE") - or "https://api.v0.dev/v1" # Default v0 API base URL + api_base or get_secret_str("V0_API_BASE") or "https://api.v0.dev/v1" # Default v0 API base URL ) # type: ignore dynamic_api_key = api_key or get_secret_str("V0_API_KEY") return api_base, dynamic_api_key diff --git a/litellm/llms/vercel_ai_gateway/chat/transformation.py b/litellm/llms/vercel_ai_gateway/chat/transformation.py index fda1c4a77cb..1c2e29234e6 100644 --- a/litellm/llms/vercel_ai_gateway/chat/transformation.py +++ b/litellm/llms/vercel_ai_gateway/chat/transformation.py @@ -33,16 +33,8 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = ( - api_base - or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") - or "https://ai-gateway.vercel.sh/v1" - ) - user_api_key = ( - api_key - or get_secret_str("VERCEL_AI_GATEWAY_API_KEY") - or get_secret_str("VERCEL_OIDC_TOKEN") - ) + api_base = api_base or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") or "https://ai-gateway.vercel.sh/v1" + user_api_key = api_key or get_secret_str("VERCEL_AI_GATEWAY_API_KEY") or get_secret_str("VERCEL_OIDC_TOKEN") return api_base, user_api_key def map_openai_params( @@ -52,9 +44,7 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): model: str, drop_params: bool, ) -> dict: - mapped_openai_params = super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + mapped_openai_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) # Vercel AI Gateway-only parameters extra_body = {} @@ -63,9 +53,7 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): if provider_options is not None: extra_body["providerOptions"] = provider_options - mapped_openai_params["extra_body"] = ( - extra_body # openai client supports `extra_body` param - ) + mapped_openai_params["extra_body"] = extra_body # openai client supports `extra_body` param return mapped_openai_params def transform_request( @@ -82,9 +70,7 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): Returns: dict: The transformed request. Sent as the body of the API call. """ - return super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + return super().transform_request(model, messages, optional_params, litellm_params, headers) def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] @@ -95,9 +81,7 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): headers=headers, ) - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: api_base, _ = self._get_openai_compatible_provider_info(api_base, api_key) if api_base is None: diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py index 7238b05f10d..e4036f415a9 100644 --- a/litellm/llms/vercel_ai_gateway/embedding/transformation.py +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -78,10 +78,7 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): if api_base: api_base = api_base.rstrip("/") else: - api_base = ( - get_secret_str("VERCEL_AI_GATEWAY_API_BASE") - or "https://ai-gateway.vercel.sh/v1" - ) + api_base = get_secret_str("VERCEL_AI_GATEWAY_API_BASE") or "https://ai-gateway.vercel.sh/v1" return f"{api_base}/embeddings" @@ -163,9 +160,7 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class( - self, error_message: str, status_code: int, headers: Any - ) -> Any: + def get_error_class(self, error_message: str, status_code: int, headers: Any) -> Any: """ Get the error class for Vercel AI Gateway errors. """ diff --git a/litellm/llms/vertex_ai/agent_engine/sse_iterator.py b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py index 06fb55e1848..d3e95f46be9 100644 --- a/litellm/llms/vertex_ai/agent_engine/sse_iterator.py +++ b/litellm/llms/vertex_ai/agent_engine/sse_iterator.py @@ -27,9 +27,7 @@ class VertexAgentEngineResponseIterator(BaseModelResponseIterator): def __init__(self, streaming_response: Any, sync_stream: bool) -> None: super().__init__(streaming_response=streaming_response, sync_stream=sync_stream) - def chunk_parser( - self, chunk: dict - ) -> Union[GenericStreamingChunk, ModelResponseStream]: + def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]: """ Parse a Vertex Agent Engine response chunk into ModelResponseStream. diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 0707a7b4c26..20c86a25f82 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -120,9 +120,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): # Get project and location from litellm_params or environment vertex_project = self.safe_get_vertex_ai_project(litellm_params) - vertex_location = ( - self.safe_get_vertex_ai_location(litellm_params) or "us-central1" - ) + vertex_location = self.safe_get_vertex_ai_location(litellm_params) or "us-central1" # Build the full resource path if only engine_id was provided if not resource_path.startswith("projects/"): @@ -158,9 +156,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): project_id=vertex_project, ) - verbose_logger.debug( - f"Vertex Agent Engine: Authenticated for project {project_id}" - ) + verbose_logger.debug(f"Vertex Agent Engine: Authenticated for project {project_id}") return { "Authorization": f"Bearer {access_token}", @@ -260,17 +256,13 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): return "" - def _calculate_usage( - self, model: str, messages: List[AllMessageValues], content: str - ) -> Optional[Usage]: + def _calculate_usage(self, model: str, messages: List[AllMessageValues], content: str) -> Optional[Usage]: """Calculate token usage using LiteLLM's token counter.""" try: from litellm.utils import token_counter prompt_tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - completion_tokens = token_counter( - model="gpt-3.5-turbo", text=content, count_response_tokens=True - ) + completion_tokens = token_counter(model="gpt-3.5-turbo", text=content, count_response_tokens=True) total_tokens = prompt_tokens + completion_tokens return Usage( @@ -304,9 +296,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): """ try: content_type = raw_response.headers.get("content-type", "").lower() - verbose_logger.debug( - f"Vertex Agent Engine response Content-Type: {content_type}" - ) + verbose_logger.debug(f"Vertex Agent Engine response Content-Type: {content_type}") # Parse the SSE response response_text = raw_response.text @@ -346,9 +336,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): return model_response except Exception as e: - verbose_logger.error( - f"Error processing Vertex Agent Engine response: {str(e)}" - ) + verbose_logger.error(f"Error processing Vertex Agent Engine response: {str(e)}") raise VertexAgentEngineError( message=f"Error processing response: {str(e)}", status_code=raw_response.status_code, @@ -401,14 +389,10 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): ) if response.status_code != 200: - raise VertexAgentEngineError( - status_code=response.status_code, message=str(response.read()) - ) + raise VertexAgentEngineError(status_code=response.status_code, message=str(response.read())) # Create iterator for SSE stream - completion_stream = self.get_streaming_response( - model=model, raw_response=response - ) + completion_stream = self.get_streaming_response(model=model, raw_response=response) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -448,9 +432,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): from litellm.utils import CustomStreamWrapper if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client( - llm_provider=cast(Any, "vertex_ai"), params={} - ) + client = get_async_httpx_client(llm_provider=cast(Any, "vertex_ai"), params={}) # Avoid logging sensitive api_base directly verbose_logger.debug("Making async streaming request to Vertex AI endpoint.") @@ -465,9 +447,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): ) if response.status_code != 200: - raise VertexAgentEngineError( - status_code=response.status_code, message=str(await response.aread()) - ) + raise VertexAgentEngineError(status_code=response.status_code, message=str(await response.aread())) # Create iterator for SSE stream (async) completion_stream = VertexAgentEngineResponseIterator( diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 00222faf274..ada1356fb6b 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -78,8 +78,10 @@ class VertexAIBatchPrediction(VertexLLM): "Authorization": f"Bearer {access_token}", } - vertex_batch_request: VertexAIBatchPredictionJob = VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( - request=create_batch_data + vertex_batch_request: VertexAIBatchPredictionJob = ( + VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( + request=create_batch_data + ) ) if _is_async is True: @@ -366,8 +368,10 @@ class VertexAIBatchPrediction(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") _json_response = response.json() - vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response( - response=_json_response + vertex_batch_response = ( + VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response( + response=_json_response + ) ) return vertex_batch_response @@ -389,8 +393,10 @@ class VertexAIBatchPrediction(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") _json_response = response.json() - vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response( - response=_json_response + vertex_batch_response = ( + VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response( + response=_json_response + ) ) return vertex_batch_response @@ -482,9 +488,7 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception( - f"Error: {retrieve_response.status_code} {retrieve_response.text}" - ) + raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -530,9 +534,7 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception( - f"Error: {retrieve_response.status_code} {retrieve_response.text}" - ) + raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index c1144654908..c75efdb43e8 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -28,15 +28,11 @@ class VertexAIBatchTransformation: input_file_id = request.get("input_file_id") if input_file_id is None: raise ValueError("input_file_id is required, but not provided") - input_config: InputConfig = InputConfig( - gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl" - ) + input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl") model: str = cls._get_model_from_gcs_file(input_file_id) output_config: OutputConfig = OutputConfig( predictionsFormat="jsonl", - gcsDestination=GcsDestination( - outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id) - ), + gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), ) return VertexAIBatchPredictionJob( inputConfig=input_config, @@ -52,19 +48,13 @@ class VertexAIBatchTransformation: return LiteLLMBatch( id=cls._get_batch_id_from_vertex_ai_batch_response(response), completion_window="24hrs", - created_at=_convert_vertex_datetime_to_openai_datetime( - vertex_datetime=response.get("createTime", "") - ), + created_at=_convert_vertex_datetime_to_openai_datetime(vertex_datetime=response.get("createTime", "")), endpoint="", - input_file_id=cls._get_input_file_id_from_vertex_ai_batch_response( - response - ), + input_file_id=cls._get_input_file_id_from_vertex_ai_batch_response(response), object="batch", status=cls._get_batch_job_status_from_vertex_ai_batch_response(response), error_file_id=None, # Vertex AI doesn't seem to have a direct equivalent - output_file_id=cls._get_output_file_id_from_vertex_ai_batch_response( - response - ), + output_file_id=cls._get_output_file_id_from_vertex_ai_batch_response(response), ) @classmethod @@ -76,10 +66,7 @@ class VertexAIBatchTransformation: """ batch_jobs = response.get("batchPredictionJobs", []) or [] - data = [ - cls.transform_vertex_ai_batch_response_to_openai_batch_response(job) - for job in batch_jobs - ] + data = [cls.transform_vertex_ai_batch_response_to_openai_batch_response(job) for job in batch_jobs] first_id = data[0].id if len(data) > 0 else None last_id = data[-1].id if len(data) > 0 else None @@ -95,9 +82,7 @@ class VertexAIBatchTransformation: } @classmethod - def _get_batch_id_from_vertex_ai_batch_response( - cls, response: VertexBatchPredictionResponse - ) -> str: + def _get_batch_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: """ Gets the batch id from the Vertex AI Batch response safely @@ -113,9 +98,7 @@ class VertexAIBatchTransformation: return parts[-1] if parts else _name @classmethod - def _get_input_file_id_from_vertex_ai_batch_response( - cls, response: VertexBatchPredictionResponse - ) -> str: + def _get_input_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: """ Gets the input file id from the Vertex AI Batch response """ @@ -135,16 +118,12 @@ class VertexAIBatchTransformation: return uris[0] @classmethod - def _get_output_file_id_from_vertex_ai_batch_response( - cls, response: VertexBatchPredictionResponse - ) -> str: + def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: """ Gets the output file id from the Vertex AI Batch response """ - output_file_id: str = response.get("outputInfo", OutputInfo()).get( - "gcsOutputDirectory", "" - ) + output_file_id: str = response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "") if output_file_id: output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" if output_file_id and output_file_id != "/predictions.jsonl": diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 96f016da94d..36522dfe396 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -135,9 +135,7 @@ class VertexAIModelRoute(str, Enum): VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute] -def get_vertex_ai_model_route( - model: str, litellm_params: Optional[dict] = None -) -> VertexAIModelRoute: +def get_vertex_ai_model_route(model: str, litellm_params: Optional[dict] = None) -> VertexAIModelRoute: """ Determine which handler to use for a Vertex AI model based on the model name. @@ -216,9 +214,7 @@ def get_supports_system_message( _custom_llm_provider = custom_llm_provider if custom_llm_provider == "vertex_ai_beta": _custom_llm_provider = "vertex_ai" - supports_system_message = supports_system_messages( - model=model, custom_llm_provider=_custom_llm_provider - ) + supports_system_message = supports_system_messages(model=model, custom_llm_provider=_custom_llm_provider) # Vertex Models called in the `/gemini` request/response format also support system messages if litellm.VertexGeminiConfig._is_model_gemini_spec_model(model): @@ -241,9 +237,7 @@ def get_supports_response_schema( if custom_llm_provider == "vertex_ai_beta": _custom_llm_provider = "vertex_ai" - _supports_response_schema = supports_response_schema( - model=model, custom_llm_provider=_custom_llm_provider - ) + _supports_response_schema = supports_response_schema(model=model, custom_llm_provider=_custom_llm_provider) return _supports_response_schema @@ -278,9 +272,7 @@ def supports_response_json_schema(model: str) -> bool: from typing import Literal, Optional -all_gemini_url_modes = Literal[ - "chat", "embedding", "batch_embedding", "image_generation", "count_tokens" -] +all_gemini_url_modes = Literal["chat", "embedding", "batch_embedding", "image_generation", "count_tokens"] def get_vertex_base_model_name(model: str) -> str: @@ -453,9 +445,7 @@ def _get_gemini_url( ) _gemini_model_name = "models/{}".format(model) - api_version = ( - "v1alpha" if VertexGeminiConfig._is_gemini_3_or_newer(model) else "v1beta" - ) + api_version = "v1alpha" if VertexGeminiConfig._is_gemini_3_or_newer(model) else "v1beta" if mode == "chat": endpoint = "generateContent" @@ -465,24 +455,16 @@ def _get_gemini_url( api_version, _gemini_model_name, endpoint ) else: - url = "https://generativelanguage.googleapis.com/{}/{}:{}".format( - api_version, _gemini_model_name, endpoint - ) + url = "https://generativelanguage.googleapis.com/{}/{}:{}".format(api_version, _gemini_model_name, endpoint) elif mode == "embedding": endpoint = "embedContent" - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format( - _gemini_model_name, endpoint - ) + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(_gemini_model_name, endpoint) elif mode == "batch_embedding": endpoint = "batchEmbedContents" - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format( - _gemini_model_name, endpoint - ) + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(_gemini_model_name, endpoint) elif mode == "count_tokens": endpoint = "countTokens" - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format( - _gemini_model_name, endpoint - ) + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(_gemini_model_name, endpoint) elif mode == "image_generation": raise ValueError( "LiteLLM's `gemini/` route does not support image generation yet. Let us know if you need this feature by opening an issue at https://github.com/BerriAI/litellm/issues" @@ -511,9 +493,7 @@ def _check_text_in_content(parts: List[PartType]) -> bool: def _fix_enum_empty_strings(schema, depth=0): """Fix empty strings in enum values by replacing them with None. Gemini doesn't accept empty strings in enums.""" if depth > DEFAULT_MAX_RECURSE_DEPTH: - raise ValueError( - f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema." - ) + raise ValueError(f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema.") if "enum" in schema and isinstance(schema["enum"], list): schema["enum"] = [None if value == "" else value for value in schema["enum"]] @@ -537,9 +517,7 @@ def _fix_enum_types(schema, depth=0): include a string type), remove the enum to avoid provider validation errors. """ if depth > DEFAULT_MAX_RECURSE_DEPTH: - raise ValueError( - f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema." - ) + raise ValueError(f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema.") if not isinstance(schema, dict): return @@ -672,11 +650,7 @@ def _filter_anyof_fields(schema_dict: Dict[str, Any]) -> Dict[str, Any]: if isinstance(schema_dict, dict) and schema_dict.get("anyOf"): any_of = schema_dict["anyOf"] - if ( - (title or description) - and isinstance(any_of, list) - and all(isinstance(item, dict) for item in any_of) - ): + if (title or description) and isinstance(any_of, list) and all(isinstance(item, dict) for item in any_of): for item in any_of: if title: item["title"] = title @@ -712,9 +686,7 @@ def process_items(schema, depth=0): process_items(item, depth + 1) -def set_schema_property_ordering( - schema: Dict[str, Any], depth: int = 0 -) -> Dict[str, Any]: +def set_schema_property_ordering(schema: Dict[str, Any], depth: int = 0) -> Dict[str, Any]: """ vertex ai and generativeai apis order output of fields alphabetically, unless you specify the order. python dicts retain order, so we just use that. Note that this field only applies to structured outputs, and not tools. @@ -741,9 +713,7 @@ def set_schema_property_ordering( return schema -def filter_schema_fields( - schema_dict: Dict[str, Any], valid_fields: Set[str], processed=None -) -> Dict[str, Any]: +def filter_schema_fields(schema_dict: Dict[str, Any], valid_fields: Set[str], processed=None) -> Dict[str, Any]: """ Recursively filter a schema dictionary to keep only valid fields. """ @@ -766,10 +736,7 @@ def filter_schema_fields( continue if key == "properties" and isinstance(value, dict): - result[key] = { - k: filter_schema_fields(v, valid_fields, processed) - for k, v in value.items() - } + result[key] = {k: filter_schema_fields(v, valid_fields, processed) for k, v in value.items()} elif key == "format": if value in {"enum", "date-time"}: result[key] = value @@ -809,8 +776,7 @@ def convert_anyof_null_to_nullable(schema, depth=0): if len(anyof) == 0: # Edge case: response schema with only null type present is invalid in Vertex AI raise ValueError( - "Invalid input: AnyOf schema with only null type is not supported. " - "Please provide a non-null type." + "Invalid input: AnyOf schema with only null type is not supported. Please provide a non-null type." ) if contains_null: @@ -834,12 +800,7 @@ def convert_anyof_null_to_nullable(schema, depth=0): def add_object_type(schema): # Gemini requires all function parameters to be type OBJECT # Handle case where schema has no properties and no type (e.g. tools with no arguments) - if ( - "type" not in schema - and "anyOf" not in schema - and "oneOf" not in schema - and "allOf" not in schema - ): + if "type" not in schema and "anyOf" not in schema and "oneOf" not in schema and "allOf" not in schema: schema["type"] = "object" properties = schema.get("properties", None) @@ -951,9 +912,7 @@ def _convert_schema_types(schema, depth=0): any_of.append({"type": t}) # Remove type-specific fields from parent if we moved them into anyOf - has_object_or_array = any( - t in ("object", "array") for t in type_val if isinstance(t, str) - ) + has_object_or_array = any(t in ("object", "array") for t in type_val if isinstance(t, str)) if has_object_or_array: for field in type_specific_fields: schema.pop(field, None) @@ -1009,9 +968,7 @@ def get_vertex_model_id_from_url(url: str) -> Optional[str]: return match.group(1) if match else None -def replace_project_and_location_in_route( - requested_route: str, vertex_project: str, vertex_location: str -) -> str: +def replace_project_and_location_in_route(requested_route: str, vertex_project: str, vertex_location: str) -> str: """ Replace project and location values in the route with the provided values """ @@ -1044,9 +1001,7 @@ def construct_target_url( new_base_url = httpx.URL(base_url) if "locations" in requested_route: # contains the target project id + location if vertex_project and vertex_location: - requested_route = replace_project_and_location_in_route( - requested_route, vertex_project, vertex_location - ) + requested_route = replace_project_and_location_in_route(requested_route, vertex_project, vertex_location) return new_base_url.copy_with(path=requested_route) """ @@ -1067,9 +1022,7 @@ def construct_target_url( vertex_version = "v1beta1" requested_route = requested_route.replace("/v1beta1/", "/", 1) - base_requested_route = "{}/projects/{}/locations/{}".format( - vertex_version, vertex_project, vertex_location - ) + base_requested_route = "{}/projects/{}/locations/{}".format(vertex_version, vertex_project, vertex_location) updated_requested_route = "/" + base_requested_route + requested_route @@ -1100,9 +1053,7 @@ class VertexAIModelInfo(BaseLLMModelInfo): ) -> dict: raise NotImplementedError("Vertex AI models are not supported yet") - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: """ Returns a list of models supported by this provider. """ @@ -1157,9 +1108,7 @@ class VertexAITokenCounter(BaseTokenCounter): ) deployment = deployment or {} - count_tokens_params_request = copy.deepcopy( - deployment.get("litellm_params", {}) - ) + count_tokens_params_request = copy.deepcopy(deployment.get("litellm_params", {})) # Check if this is a partner model (Claude, Mistral, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model_to_use): @@ -1167,19 +1116,16 @@ class VertexAITokenCounter(BaseTokenCounter): partner_models_handler = VertexAIPartnerModels() # Extract vertex-specific params from litellm_params - vertex_project = count_tokens_params_request.get( - "vertex_project" - ) or count_tokens_params_request.get("vertex_ai_project") + vertex_project = count_tokens_params_request.get("vertex_project") or count_tokens_params_request.get( + "vertex_ai_project" + ) - vertex_location = count_tokens_params_request.get( - "vertex_location" - ) or count_tokens_params_request.get("vertex_ai_location") + vertex_location = count_tokens_params_request.get("vertex_location") or count_tokens_params_request.get( + "vertex_ai_location" + ) # Count tokens not available on global location: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens - vertex_location = ( - count_tokens_params_request.get("vertex_count_tokens_location") - or vertex_location - ) + vertex_location = count_tokens_params_request.get("vertex_count_tokens_location") or vertex_location vertex_credentials = count_tokens_params_request.get( "vertex_credentials" diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index f73eb220cc6..f0ce3323ef6 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -145,9 +145,7 @@ def separate_cached_messages( last_cached_idx = filtered_messages[last_continuous_block_idx][0] cached_messages = messages[first_cached_idx : last_cached_idx + 1] - non_cached_messages = ( - messages[:first_cached_idx] + messages[last_cached_idx + 1 :] - ) + non_cached_messages = messages[:first_cached_idx] + messages[last_cached_idx + 1 :] else: non_cached_messages = messages @@ -165,9 +163,7 @@ def transform_openai_messages_to_gemini_context_caching( # Extract TTL from cached messages BEFORE system message transformation ttl = extract_ttl_from_cached_messages(messages) - supports_system_message = get_supports_system_message( - model=model, custom_llm_provider=custom_llm_provider - ) + supports_system_message = get_supports_system_message(model=model, custom_llm_provider=custom_llm_provider) transformed_system_messages, new_messages = _transform_system_message( supports_system_message=supports_system_message, messages=messages diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index d29734c0294..0bf3715f798 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -26,9 +26,7 @@ from .transformation import ( transform_openai_messages_to_gemini_context_caching, ) -local_cache_obj = Cache( - type=LiteLLMCacheType.LOCAL -) # only used for calling 'get_cache_key' function +local_cache_obj = Cache(type=LiteLLMCacheType.LOCAL) # only used for calling 'get_cache_key' function MAX_PAGINATION_PAGES = 100 # Reasonable upper bound for pagination @@ -88,9 +86,7 @@ class ContextCachingEndpoints(VertexBase): model=model, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_api_version=( - "v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1" - ), + vertex_api_version=("v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1"), ) def check_cache( @@ -156,9 +152,7 @@ class ContextCachingEndpoints(VertexBase): except httpx.HTTPStatusError as e: if e.response.status_code == 403: return None - raise VertexAIError( - status_code=e.response.status_code, message=e.response.text - ) + raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -250,9 +244,7 @@ class ContextCachingEndpoints(VertexBase): except httpx.HTTPStatusError as e: if e.response.status_code == 403: return None - raise VertexAIError( - status_code=e.response.status_code, message=e.response.text - ) + raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -311,9 +303,7 @@ class ContextCachingEndpoints(VertexBase): if cached_content is not None: return messages, optional_params, cached_content - cached_messages, non_cached_messages = separate_cached_messages( - messages=messages - ) + cached_messages, non_cached_messages = separate_cached_messages(messages=messages) if len(cached_messages) == 0: return messages, optional_params, None @@ -387,15 +377,13 @@ class ContextCachingEndpoints(VertexBase): return non_cached_messages, optional_params, google_cache_name ## TRANSFORM REQUEST - cached_content_request_body = ( - transform_openai_messages_to_gemini_context_caching( - model=model, - messages=cached_messages, - cache_key=generated_cache_key, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - ) + cached_content_request_body = transform_openai_messages_to_gemini_context_caching( + model=model, + messages=cached_messages, + cache_key=generated_cache_key, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, ) cached_content_request_body["tools"] = tools @@ -466,9 +454,7 @@ class ContextCachingEndpoints(VertexBase): if cached_content is not None: return messages, optional_params, cached_content - cached_messages, non_cached_messages = separate_cached_messages( - messages=messages - ) + cached_messages, non_cached_messages = separate_cached_messages(messages=messages) if len(cached_messages) == 0: return messages, optional_params, None @@ -512,9 +498,7 @@ class ContextCachingEndpoints(VertexBase): headers.update(extra_headers) if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client( - params={"timeout": timeout}, llm_provider=litellm.LlmProviders.VERTEX_AI - ) + client = get_async_httpx_client(params={"timeout": timeout}, llm_provider=litellm.LlmProviders.VERTEX_AI) else: client = client @@ -540,15 +524,13 @@ class ContextCachingEndpoints(VertexBase): return non_cached_messages, optional_params, google_cache_name ## TRANSFORM REQUEST - cached_content_request_body = ( - transform_openai_messages_to_gemini_context_caching( - model=model, - messages=cached_messages, - cache_key=generated_cache_key, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - ) + cached_content_request_body = transform_openai_messages_to_gemini_context_caching( + model=model, + messages=cached_messages, + cache_key=generated_cache_key, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, ) cached_content_request_body["tools"] = tools diff --git a/litellm/llms/vertex_ai/cost_calculator.py b/litellm/llms/vertex_ai/cost_calculator.py index 758fc0bae87..84c9108847b 100644 --- a/litellm/llms/vertex_ai/cost_calculator.py +++ b/litellm/llms/vertex_ai/cost_calculator.py @@ -47,9 +47,7 @@ def cost_router( or "gemma" in model ): return "cost_per_token" - elif custom_llm_provider == "vertex_ai" and ( - call_type == "embedding" or call_type == "aembedding" - ): + elif custom_llm_provider == "vertex_ai" and (call_type == "embedding" or call_type == "aembedding"): return "cost_per_token" elif custom_llm_provider == "vertex_ai" and ("gemini-2" in model): return "cost_per_token" @@ -78,14 +76,10 @@ def cost_per_character( Raises: Exception if model requires >128k pricing, but model cost not mapped """ - model_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) ## GET MODEL INFO - model_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) ## CALCULATE INPUT COST if prompt_characters is None: @@ -103,25 +97,18 @@ def cost_per_character( ## check if character pricing, else default to token pricing assert ( "input_cost_per_character_above_128k_tokens" in model_info - and model_info["input_cost_per_character_above_128k_tokens"] - is not None + and model_info["input_cost_per_character_above_128k_tokens"] is not None ), ( "model info for model={} does not have 'input_cost_per_character_above_128k_tokens'-pricing for > 128k tokens\nmodel_info={}".format( model, model_info ) ) - prompt_cost = ( - prompt_characters - * model_info["input_cost_per_character_above_128k_tokens"] - ) + prompt_cost = prompt_characters * model_info["input_cost_per_character_above_128k_tokens"] else: assert ( - "input_cost_per_character" in model_info - and model_info["input_cost_per_character"] is not None - ), ( - "model info for model={} does not have 'input_cost_per_character'-pricing\nmodel_info={}".format( - model, model_info - ) + "input_cost_per_character" in model_info and model_info["input_cost_per_character"] is not None + ), "model info for model={} does not have 'input_cost_per_character'-pricing\nmodel_info={}".format( + model, model_info ) prompt_cost = prompt_characters * model_info["input_cost_per_character"] except Exception as e: @@ -152,29 +139,20 @@ def cost_per_character( ): assert ( "output_cost_per_character_above_128k_tokens" in model_info - and model_info["output_cost_per_character_above_128k_tokens"] - is not None + and model_info["output_cost_per_character_above_128k_tokens"] is not None ), ( "model info for model={} does not have 'output_cost_per_character_above_128k_tokens' pricing\nmodel_info={}".format( model, model_info ) ) - completion_cost = ( - completion_tokens - * model_info["output_cost_per_character_above_128k_tokens"] - ) + completion_cost = completion_tokens * model_info["output_cost_per_character_above_128k_tokens"] else: assert ( - "output_cost_per_character" in model_info - and model_info["output_cost_per_character"] is not None - ), ( - "model info for model={} does not have 'output_cost_per_character'-pricing\nmodel_info={}".format( - model, model_info - ) - ) - completion_cost = ( - completion_characters * model_info["output_cost_per_character"] + "output_cost_per_character" in model_info and model_info["output_cost_per_character"] is not None + ), "model info for model={} does not have 'output_cost_per_character'-pricing\nmodel_info={}".format( + model, model_info ) + completion_cost = completion_characters * model_info["output_cost_per_character"] except Exception as e: verbose_logger.debug( "litellm.litellm_core_utils.llm_cost_calc.google.py::cost_per_character(): Exception occured - {}\nDefaulting to None".format( @@ -195,37 +173,23 @@ def _handle_128k_pricing( usage: Usage, ) -> Tuple[float, float]: ## CALCULATE INPUT COST - input_cost_per_token_above_128k_tokens = model_info.get( - "input_cost_per_token_above_128k_tokens" - ) - output_cost_per_token_above_128k_tokens = model_info.get( - "output_cost_per_token_above_128k_tokens" - ) + input_cost_per_token_above_128k_tokens = model_info.get("input_cost_per_token_above_128k_tokens") + output_cost_per_token_above_128k_tokens = model_info.get("output_cost_per_token_above_128k_tokens") prompt_tokens = usage.prompt_tokens completion_tokens = usage.completion_tokens - if ( - _is_above_128k(tokens=prompt_tokens) - and input_cost_per_token_above_128k_tokens is not None - ): + if _is_above_128k(tokens=prompt_tokens) and input_cost_per_token_above_128k_tokens is not None: prompt_cost = prompt_tokens * input_cost_per_token_above_128k_tokens else: prompt_cost = prompt_tokens * (model_info["input_cost_per_token"] or 0.0) ## CALCULATE OUTPUT COST - output_cost_per_token_above_128k_tokens = model_info.get( - "output_cost_per_token_above_128k_tokens" - ) - if ( - _is_above_128k(tokens=completion_tokens) - and output_cost_per_token_above_128k_tokens is not None - ): + output_cost_per_token_above_128k_tokens = model_info.get("output_cost_per_token_above_128k_tokens") + if _is_above_128k(tokens=completion_tokens) and output_cost_per_token_above_128k_tokens is not None: completion_cost = completion_tokens * output_cost_per_token_above_128k_tokens else: - completion_cost = completion_tokens * ( - model_info["output_cost_per_token"] or 0.0 - ) + completion_cost = completion_tokens * (model_info["output_cost_per_token"] or 0.0) return prompt_cost, completion_cost @@ -255,21 +219,12 @@ def cost_per_token( """ ## GET MODEL INFO - model_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) ## HANDLE 128k+ PRICING - input_cost_per_token_above_128k_tokens = model_info.get( - "input_cost_per_token_above_128k_tokens" - ) - output_cost_per_token_above_128k_tokens = model_info.get( - "output_cost_per_token_above_128k_tokens" - ) - if ( - input_cost_per_token_above_128k_tokens is not None - or output_cost_per_token_above_128k_tokens is not None - ): + input_cost_per_token_above_128k_tokens = model_info.get("input_cost_per_token_above_128k_tokens") + output_cost_per_token_above_128k_tokens = model_info.get("output_cost_per_token_above_128k_tokens") + if input_cost_per_token_above_128k_tokens is not None or output_cost_per_token_above_128k_tokens is not None: return _handle_128k_pricing( model_info=model_info, usage=usage, diff --git a/litellm/llms/vertex_ai/count_tokens/handler.py b/litellm/llms/vertex_ai/count_tokens/handler.py index 9a175371a27..9f2826a4bb4 100644 --- a/litellm/llms/vertex_ai/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/count_tokens/handler.py @@ -17,9 +17,7 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): Returns a Tuple of headers and url for the Vertex AI countTokens endpoint. """ litellm_params = litellm_params or {} - vertex_credentials = self.get_vertex_ai_credentials( - litellm_params=litellm_params - ) + vertex_credentials = self.get_vertex_ai_credentials(litellm_params=litellm_params) vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params) vertex_location = self.get_vertex_ai_location(litellm_params=litellm_params) should_use_v1beta1_features = self.is_using_v1beta1_features(litellm_params) diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index 176cfe98411..3bc09139f8f 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -60,9 +60,7 @@ class VertexAIFilesHandler(GCSBucketBase): scheme="gs://", configured_bucket_name=configured_bucket_name, allowed_object_prefixes=(VERTEX_AI_MANAGED_GCS_PREFIX,), - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( - litellm_params - ), + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) async def afile_content( @@ -93,9 +91,7 @@ class VertexAIFilesHandler(GCSBucketBase): if not file_id: raise ValueError("file_id is required in file_content_request") - gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( - kwargs={} - ) + gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(kwargs={}) bucket_name, object_path = self._extract_bucket_and_object_from_file_id( file_id=file_id, configured_bucket_name=gcs_logging_config["bucket_name"], @@ -109,9 +105,7 @@ class VertexAIFilesHandler(GCSBucketBase): } } - file_content = await self.download_gcs_object( - object_name=object_path, **download_kwargs - ) + file_content = await self.download_gcs_object(object_name=object_path, **download_kwargs) decoded_file_id = unquote(file_id) if file_content is None: @@ -156,9 +150,7 @@ class VertexAIFilesHandler(GCSBucketBase): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], litellm_params: Optional[dict] = None, - ) -> Union[ - HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] - ]: + ) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]: """ Download file content from GCS bucket for VertexAI files. Supports both sync and async operations. diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index d5164d8c1c2..fcb7617d97e 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -93,9 +93,7 @@ def _sanitize_gcp_label_value(value: str) -> str: def _encode_gcp_label_value_chunks(value: str) -> List[str]: """Encode arbitrary text across one or more GCP-label-safe values.""" max_encoded_len = _GCP_LABEL_VALUE_MAX_LEN - len(_CUSTOM_ID_RAW_LABEL_PREFIX) - encoded = ( - base64.b32encode(value.encode("utf-8")).decode("ascii").rstrip("=").lower() - ) + encoded = base64.b32encode(value.encode("utf-8")).decode("ascii").rstrip("=").lower() return [ f"{_CUSTOM_ID_RAW_LABEL_PREFIX}{encoded[i : i + max_encoded_len]}" for i in range(0, len(encoded), max_encoded_len) @@ -143,10 +141,7 @@ def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: for key, value in labels.items(): if key.startswith(chunk_prefix) and key[len(chunk_prefix) :].isdigit(): indexed_chunks.append((int(key[len(chunk_prefix) :]), str(value))) - raw_chunks.extend( - raw_label_chunk - for _, raw_label_chunk in sorted(indexed_chunks, key=lambda item: item[0]) - ) + raw_chunks.extend(raw_label_chunk for _, raw_label_chunk in sorted(indexed_chunks, key=lambda item: item[0])) decoded = _decode_gcp_label_value_chunks(raw_chunks) if decoded is not None: return decoded @@ -279,9 +274,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True for entry in _iter_openai_jsonl_entries(self._openai_file_content): - wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request( - entry, self._map_openai_to_vertex_params - ) + wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, self._map_openai_to_vertex_params) prefix = b"" if first else b"\n" first = False yield prefix + json.dumps(wrapped).encode("utf-8") @@ -361,9 +354,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_configured_bucket_name(self, litellm_params: Dict) -> str: bucket_name = ( - litellm_params.get("gcs_bucket_name") - or litellm_params.get("bucket_name") - or os.getenv("GCS_BUCKET_NAME") + litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") ) if not bucket_name: raise ValueError("GCS bucket_name is required") @@ -398,9 +389,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # large uploads); everything else is a single simple-media upload. upload_type = ( "resumable" - if FilesAPIUtils.is_batch_jsonl_request( - create_file_data=data, content_type=content_type - ) + if FilesAPIUtils.is_batch_jsonl_request(create_file_data=data, content_type=content_type) else "media" ) endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType={upload_type}&name={encoded_object_name}" @@ -410,9 +399,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): return f"{api_base}/{endpoint}" - def get_supported_openai_params( - self, model: str - ) -> List[OpenAICreateFileRequestOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAICreateFileRequestOptionalParams]: return [] def map_openai_params( @@ -521,16 +508,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): object="file", ) - def get_error_class( - self, error_message: str, status_code: int, headers: Union[Dict, Headers] - ) -> BaseLLMException: - return VertexAIError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[Dict, Headers]) -> BaseLLMException: + return VertexAIError(status_code=status_code, message=error_message, headers=headers) - def _parse_gcs_uri( - self, file_id: str, litellm_params: Optional[Dict] = None - ) -> Tuple[str, str]: + def _parse_gcs_uri(self, file_id: str, litellm_params: Optional[Dict] = None) -> Tuple[str, str]: """ Validate a managed GCS file_id and return (bucket, url-encoded-object-path). """ @@ -540,9 +521,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): scheme="gs://", configured_bucket_name=configured_bucket_name, allowed_object_prefixes=(VERTEX_AI_MANAGED_GCS_PREFIX,), - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( - litellm_params - ), + allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) return bucket_name, encode_gcs_object_name_for_url(object_path) @@ -761,13 +740,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): output = bytearray() for line in itertools.chain([first_line], lines): try: - openai_output = ( - self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), - vertex_gemini_config=vertex_gemini_config, - logging_obj=batch_transform_logging_obj, - mock_httpx_response=mock_httpx_response, - ) + openai_output = self._transform_single_vertex_batch_output_to_openai( + vertex_output=json.loads(line), + vertex_gemini_config=vertex_gemini_config, + logging_obj=batch_transform_logging_obj, + mock_httpx_response=mock_httpx_response, ) except Exception: return content diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index a5971de0e94..b220b1544b5 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -38,9 +38,7 @@ class VertexFineTuningAPI(VertexLLM): def convert_response_created_at(self, response: ResponseTuningJob): try: create_time_str = response.get("createTime", "") or "" - create_time_datetime = datetime.fromisoformat( - create_time_str.replace("Z", "+00:00") - ) + create_time_datetime = datetime.fromisoformat(create_time_str.replace("Z", "+00:00")) # Convert to Unix timestamp (seconds since epoch) created_at = int(create_time_datetime.timestamp()) @@ -65,16 +63,12 @@ class VertexFineTuningAPI(VertexLLM): ) if create_fine_tuning_job_data.validation_file: - supervised_tuning_spec["validation_dataset"] = ( - create_fine_tuning_job_data.validation_file - ) + supervised_tuning_spec["validation_dataset"] = create_fine_tuning_job_data.validation_file - _vertex_hyperparameters = ( - self._transform_openai_hyperparameters_to_vertex_hyperparameters( - create_fine_tuning_job_data=create_fine_tuning_job_data, - kwargs=kwargs, - original_hyperparameters=original_hyperparameters, - ) + _vertex_hyperparameters = self._transform_openai_hyperparameters_to_vertex_hyperparameters( + create_fine_tuning_job_data=create_fine_tuning_job_data, + kwargs=kwargs, + original_hyperparameters=original_hyperparameters, ) if _vertex_hyperparameters and len(_vertex_hyperparameters) > 0: @@ -98,9 +92,7 @@ class VertexFineTuningAPI(VertexLLM): _vertex_hyperparameters = FineTuneHyperparameters() if _oai_hyperparameters: if _oai_hyperparameters.n_epochs: - _vertex_hyperparameters["epoch_count"] = int( - _oai_hyperparameters.n_epochs - ) + _vertex_hyperparameters["epoch_count"] = int(_oai_hyperparameters.n_epochs) if _oai_hyperparameters.learning_rate_multiplier: _vertex_hyperparameters["learning_rate_multiplier"] = float( _oai_hyperparameters.learning_rate_multiplier @@ -112,12 +104,8 @@ class VertexFineTuningAPI(VertexLLM): return _vertex_hyperparameters - def convert_vertex_response_to_open_ai_response( - self, response: ResponseTuningJob - ) -> LiteLLMFineTuningJob: - status: Literal[ - "validating_files", "queued", "running", "succeeded", "failed", "cancelled" - ] = "queued" + def convert_vertex_response_to_open_ai_response(self, response: ResponseTuningJob) -> LiteLLMFineTuningJob: + status: Literal["validating_files", "queued", "running", "succeeded", "failed", "cancelled"] = "queued" if response["state"] == "JOB_STATE_PENDING": status = "queued" if response["state"] == "JOB_STATE_SUCCEEDED": @@ -131,9 +119,7 @@ class VertexFineTuningAPI(VertexLLM): created_at = self.convert_response_created_at(response) - _supervisedTuningSpec: ResponseSupervisedTuningSpec = ( - response.get("supervisedTuningSpec", None) or {} - ) + _supervisedTuningSpec: ResponseSupervisedTuningSpec = response.get("supervisedTuningSpec", None) or {} training_uri: str = _supervisedTuningSpec.get("trainingDatasetUri", "") or "" return LiteLLMFineTuningJob( id=response.get("name", "") or "", @@ -141,10 +127,7 @@ class VertexFineTuningAPI(VertexLLM): fine_tuned_model=response.get("tunedModelDisplayName", ""), finished_at=None, hyperparameters=self._translate_vertex_response_hyperparameters( - vertex_hyper_parameters=_supervisedTuningSpec.get( - "hyperParameters", FineTuneHyperparameters() - ) - or {} + vertex_hyper_parameters=_supervisedTuningSpec.get("hyperParameters", FineTuneHyperparameters()) or {} ), model=response.get("baseModel", "") or "", object="fine_tuning.job", @@ -184,9 +167,7 @@ class VertexFineTuningAPI(VertexLLM): json.dumps(request_data, indent=4), ) if self.async_handler is None: - raise ValueError( - "VertexAI Fine Tuning - async_handler is not initialized" - ) + raise ValueError("VertexAI Fine Tuning - async_handler is not initialized") response = await self.async_handler.post( headers=headers, url=fine_tuning_url, @@ -198,18 +179,14 @@ class VertexFineTuningAPI(VertexLLM): f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}" ) - verbose_logger.debug( - "got response from creating fine tuning job: %s", response.json() - ) + verbose_logger.debug("got response from creating fine tuning job: %s", response.json()) vertex_response = ResponseTuningJob( # type: ignore **response.json(), ) verbose_logger.debug("vertex_response %s", vertex_response) - open_ai_response = self.convert_vertex_response_to_open_ai_response( - vertex_response - ) + open_ai_response = self.convert_vertex_response_to_open_ai_response(vertex_response) return open_ai_response except Exception as e: @@ -230,9 +207,7 @@ class VertexFineTuningAPI(VertexLLM): kwargs: Optional[dict] = None, original_hyperparameters: Optional[dict] = {}, ) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: - verbose_logger.debug( - "creating fine tuning job, args= %s", create_fine_tuning_job_data - ) + verbose_logger.debug("creating fine tuning job, args= %s", create_fine_tuning_job_data) _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -288,17 +263,13 @@ class VertexFineTuningAPI(VertexLLM): f"Error creating fine tuning job. Status code: {response.status_code}. Response: {response.text}" ) - verbose_logger.debug( - "got response from creating fine tuning job: %s", response.json() - ) + verbose_logger.debug("got response from creating fine tuning job: %s", response.json()) vertex_response = ResponseTuningJob( # type: ignore **response.json(), ) verbose_logger.debug("vertex_response %s", vertex_response) - open_ai_response = self.convert_vertex_response_to_open_ai_response( - vertex_response - ) + open_ai_response = self.convert_vertex_response_to_open_ai_response(vertex_response) return open_ai_response async def pass_through_vertex_ai_POST_request( diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 3bdcbd25949..0db1118a7b4 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -121,9 +121,7 @@ def _convert_detail_to_media_resolution_enum( return None -def _get_highest_media_resolution( - current: Optional[str], new_detail: Optional[str] -) -> Optional[str]: +def _get_highest_media_resolution(current: Optional[str], new_detail: Optional[str]) -> Optional[str]: """ Compare two media resolution values and return the highest one. Resolution hierarchy: ultra_high > high > medium > low > None @@ -169,9 +167,7 @@ def _extract_max_media_resolution_from_messages( if isinstance(file_obj, dict): detail = file_obj.get("detail") if detail: - max_resolution = _get_highest_media_resolution( - max_resolution, detail - ) + max_resolution = _get_highest_media_resolution(max_resolution, detail) return max_resolution @@ -194,9 +190,7 @@ def _apply_gemini_metadata( part_dict = dict(part) - if media_resolution_enum is not None and VertexGeminiConfig._is_gemini_3_or_newer( - model - ): + if media_resolution_enum is not None and VertexGeminiConfig._is_gemini_3_or_newer(model): part_dict["media_resolution"] = media_resolution_enum if video_metadata is not None: @@ -231,9 +225,7 @@ def _is_valid_gcs_bucket_name(bucket: str) -> bool: max_bucket_length = 222 if "." in bucket else 63 if bucket_length < 3 or bucket_length > max_bucket_length: return False - if "." in bucket and any( - len(label) == 0 or len(label) > 63 for label in bucket.split(".") - ): + if "." in bucket and any(len(label) == 0 or len(label) > 63 for label in bucket.split(".")): return False if not re.fullmatch(r"[a-z0-9][a-z0-9._-]*[a-z0-9]", bucket): return False @@ -269,11 +261,7 @@ def _image_url_payload_may_need_sync_gcs_metadata_fetch( url = raw_image_url.get("url") # type: ignore[assignment] if not isinstance(url, str): return False - fmt = ( - raw_image_url.get("format") - or raw_image_url.get("mime_type") - or raw_image_url.get("content_type") - ) + fmt = raw_image_url.get("format") or raw_image_url.get("mime_type") or raw_image_url.get("content_type") elif isinstance(raw_image_url, str): url = raw_image_url else: @@ -305,9 +293,7 @@ def _openai_messages_may_need_sync_gcs_metadata_fetch( for image_item in images_field: if not isinstance(image_item, dict): continue - if _image_url_payload_may_need_sync_gcs_metadata_fetch( - image_item.get("image_url") - ): + if _image_url_payload_may_need_sync_gcs_metadata_fetch(image_item.get("image_url")): return True content = msg.get("content") @@ -318,19 +304,13 @@ def _openai_messages_may_need_sync_gcs_metadata_fetch( continue itype = item.get("type") if itype == "image_url": - if _image_url_payload_may_need_sync_gcs_metadata_fetch( - item.get("image_url") - ): + if _image_url_payload_may_need_sync_gcs_metadata_fetch(item.get("image_url")): return True elif itype == "file": file_obj = item.get("file") if not isinstance(file_obj, dict): continue - fmt = ( - file_obj.get("format") - or file_obj.get("mime_type") - or file_obj.get("content_type") - ) + fmt = file_obj.get("format") or file_obj.get("mime_type") or file_obj.get("content_type") passed = file_obj.get("file_id") or file_obj.get("file_data") if ( isinstance(passed, str) @@ -365,9 +345,7 @@ def _get_gcs_object_content_type( return None headers: Dict[str, str] = {} - explicit_vertex_auth_provided = ( - vertex_project is not None or vertex_credentials is not None - ) + explicit_vertex_auth_provided = vertex_project is not None or vertex_credentials is not None if explicit_vertex_auth_provided: try: access_token, _ = _get_vertex_base().get_access_token( @@ -378,8 +356,7 @@ def _get_gcs_object_content_type( except Exception as e: raise litellm.BadRequestError( message=( - "Unable to fetch GCS metadata with provided Vertex credentials/project. " - f"Original error: {str(e)}" + f"Unable to fetch GCS metadata with provided Vertex credentials/project. Original error: {str(e)}" ), model=None, llm_provider="vertex_ai", @@ -470,9 +447,7 @@ def _get_gcs_object_content_type( return None -def _normalize_and_validate_gemini_mime_type( - mime_type: str, model: Optional[str] -) -> str: +def _normalize_and_validate_gemini_mime_type(mime_type: str, model: Optional[str]) -> str: # Import lazily to avoid a module-level cyclic-import alert with # litellm.types.files. from litellm.types.files import get_file_extension_from_mime_type @@ -581,12 +556,8 @@ def _process_gemini_media( ) file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} - return _apply_gemini_metadata( - part, model, media_resolution_enum, video_metadata - ) - elif image_url.startswith( - "https://generativelanguage.googleapis.com/v1beta/files/" - ): + return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) + elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"): # Gemini Files API URIs — the file is already uploaded to Google's # servers; pass the URI through as file_data without fetching it. # These URLs return 403 when accessed directly, so we must not try @@ -597,26 +568,16 @@ def _process_gemini_media( # Gemini Files API references can be passed through as URI-only. file_data = cast(FileDataType, {"file_uri": image_url}) part = {"file_data": file_data} - return _apply_gemini_metadata( - part, model, media_resolution_enum, video_metadata - ) - elif ( - "https://" in image_url - and (image_type := format or _get_image_mime_type_from_url(image_url)) - is not None - ): + return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) + elif "https://" in image_url and (image_type := format or _get_image_mime_type_from_url(image_url)) is not None: file_data = FileDataType(mime_type=image_type, file_uri=image_url) part = {"file_data": file_data} - return _apply_gemini_metadata( - part, model, media_resolution_enum, video_metadata - ) + return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) elif "http://" in image_url or "https://" in image_url or "base64" in image_url: image = convert_to_anthropic_image_obj(image_url, format=format) _blob: BlobType = {"data": image["data"], "mime_type": image["media_type"]} part = {"inline_data": cast(BlobType, _blob)} - return _apply_gemini_metadata( - part, model, media_resolution_enum, video_metadata - ) + return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) raise Exception("Invalid image received - {}".format(image_url)) except Exception as e: raise e @@ -653,9 +614,7 @@ def _get_equivalent_key(key: str, available_keys: set) -> Optional[str]: return None -def check_if_part_exists_in_parts( - parts: List[PartType], part: PartType, excluded_keys: List[str] = [] -) -> bool: +def check_if_part_exists_in_parts(parts: List[PartType], part: PartType, excluded_keys: List[str] = []) -> bool: """ Check if a part exists in a list of parts Handles both camelCase and snake_case key variations (e.g., function_call vs functionCall) @@ -667,9 +626,7 @@ def check_if_part_exists_in_parts( match_found = True for key in keys_to_compare: equivalent_key = _get_equivalent_key(key, p_keys) - if equivalent_key is None or p.get(equivalent_key, None) != part.get( - key, None - ): + if equivalent_key is None or p.get(equivalent_key, None) != part.get(key, None): match_found = False break @@ -701,30 +658,20 @@ def _gemini_convert_messages_with_history( vertex_project = None vertex_credentials = None if litellm_params: - vertex_project = litellm_params.get("vertex_project") or litellm_params.get( - "vertex_ai_project" - ) - vertex_credentials = litellm_params.get( - "vertex_credentials" - ) or litellm_params.get("vertex_ai_credentials") + vertex_project = litellm_params.get("vertex_project") or litellm_params.get("vertex_ai_project") + vertex_credentials = litellm_params.get("vertex_credentials") or litellm_params.get("vertex_ai_credentials") try: while msg_i < len(messages): user_content: List[PartType] = [] init_msg_i = msg_i ## MERGE CONSECUTIVE USER CONTENT ## - while ( - msg_i < len(messages) and messages[msg_i]["role"] in user_message_types - ): + while msg_i < len(messages) and messages[msg_i]["role"] in user_message_types: _message_content = messages[msg_i].get("content") if _message_content is not None and isinstance(_message_content, list): _parts: List[PartType] = [] for element_idx, element in enumerate(_message_content): - if ( - element["type"] == "text" - and "text" in element - and len(element["text"]) > 0 - ): + if element["type"] == "text" and "text" in element and len(element["text"]) > 0: element = cast(ChatCompletionTextObject, element) _part = PartType(text=element["text"]) _parts.append(_part) @@ -757,9 +704,7 @@ def _gemini_convert_messages_with_history( or image_url_dict.get("content_type") ) detail = image_url_dict.get("detail") - media_resolution_enum = ( - _convert_detail_to_media_resolution_enum(detail) - ) + media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) else: image_url = raw_image_url _part = _process_gemini_media( @@ -781,13 +726,11 @@ def _gemini_convert_messages_with_history( if audio_format.startswith("audio/") is False else audio_format ) # Gemini expects audio/wav, audio/mp3, etc. - openai_image_str = ( - convert_generic_image_chunk_to_openai_image_obj( - image_chunk=GenericImageParsingChunk( - type="base64", - media_type=audio_format_modified, - data=audio_data, - ) + openai_image_str = convert_generic_image_chunk_to_openai_image_obj( + image_chunk=GenericImageParsingChunk( + type="base64", + media_type=audio_format_modified, + data=audio_data, ) ) _part = _process_gemini_media( @@ -812,23 +755,17 @@ def _gemini_convert_messages_with_history( file_dict = cast(Dict[str, Any], _file_field) file_id = file_dict.get("file_id") format = ( - file_dict.get("format") - or file_dict.get("mime_type") - or file_dict.get("content_type") + file_dict.get("format") or file_dict.get("mime_type") or file_dict.get("content_type") ) file_data = file_dict.get("file_data") detail = file_dict.get("detail") video_metadata = file_dict.get("video_metadata") passed_file = file_id or file_data if passed_file is None: - raise Exception( - "Unknown file type. Please pass in a file_id or file_data" - ) + raise Exception("Unknown file type. Please pass in a file_id or file_data") # Convert detail to media_resolution_enum - media_resolution_enum = ( - _convert_detail_to_media_resolution_enum(detail) - ) + media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) try: _part = _process_gemini_media( @@ -881,9 +818,7 @@ def _gemini_convert_messages_with_history( ## MERGE CONSECUTIVE ASSISTANT CONTENT ## while msg_i < len(messages) and messages[msg_i]["role"] == "assistant": if isinstance(messages[msg_i], BaseModel): - msg_dict: Union[ChatCompletionAssistantMessage, dict] = messages[ - msg_i - ].model_dump() # type: ignore + msg_dict: Union[ChatCompletionAssistantMessage, dict] = messages[msg_i].model_dump() # type: ignore else: msg_dict = messages[msg_i] # type: ignore assistant_msg = ChatCompletionAssistantMessage(**msg_dict) # type: ignore @@ -891,18 +826,13 @@ def _gemini_convert_messages_with_history( reasoning_content = assistant_msg.get("reasoning_content", None) thinking_blocks = assistant_msg.get("thinking_blocks") if reasoning_content is not None: - assistant_content.append( - PartType(thought=True, text=reasoning_content) - ) + assistant_content.append(PartType(thought=True, text=reasoning_content)) if thinking_blocks is not None: for block in thinking_blocks: if block["type"] == "thinking": block_thinking_str = block.get("thinking") block_signature = block.get("signature") - if ( - block_thinking_str is not None - and block_signature is not None - ): + if block_thinking_str is not None and block_signature is not None: try: assistant_content.append( PartType( @@ -929,23 +859,13 @@ def _gemini_convert_messages_with_history( elif _message_content is not None and isinstance(_message_content, str): assistant_text = _message_content # Check if message has thought_signatures in provider_specific_fields - provider_specific_fields = assistant_msg.get( - "provider_specific_fields" - ) + provider_specific_fields = assistant_msg.get("provider_specific_fields") thought_signatures = None - if provider_specific_fields and isinstance( - provider_specific_fields, dict - ): - thought_signatures = provider_specific_fields.get( - "thought_signatures" - ) + if provider_specific_fields and isinstance(provider_specific_fields, dict): + thought_signatures = provider_specific_fields.get("thought_signatures") # If we have thought signatures, add them to the part - if ( - thought_signatures - and isinstance(thought_signatures, list) - and len(thought_signatures) > 0 - ): + if thought_signatures and isinstance(thought_signatures, list) and len(thought_signatures) > 0: # Use the first signature for the text part (Gemini expects one signature per part) assistant_content.append( PartType( @@ -971,9 +891,7 @@ def _gemini_convert_messages_with_history( or image_url_obj.get("content_type") ) detail = image_url_obj.get("detail") - media_resolution_enum = ( - _convert_detail_to_media_resolution_enum(detail) - ) + media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) if assistant_image_url: _part = _process_gemini_media( image_url=assistant_image_url, @@ -987,8 +905,7 @@ def _gemini_convert_messages_with_history( ## HANDLE ASSISTANT FUNCTION CALL if ( - assistant_msg.get("tool_calls", []) is not None - or assistant_msg.get("function_call") is not None + assistant_msg.get("tool_calls", []) is not None or assistant_msg.get("function_call") is not None ): # support assistant tool invoke conversion gemini_tool_call_parts = convert_to_gemini_tool_call_invoke( assistant_msg, @@ -1011,10 +928,7 @@ def _gemini_convert_messages_with_history( # reference. The following tool result would then be matched against # an assistant message that has no tool_calls, raising "Missing # corresponding tool call for tool response message". - if ( - assistant_msg.get("tool_calls") - or assistant_msg.get("function_call") is not None - ): + if assistant_msg.get("tool_calls") or assistant_msg.get("function_call") is not None: last_message_with_tool_calls = assistant_msg ## HANDLE SERVER-SIDE TOOL INVOCATIONS (context circulation) @@ -1032,9 +946,7 @@ def _gemini_convert_messages_with_history( } } if "thought_signature" in invocation: - tc_part["thoughtSignature"] = invocation[ - "thought_signature" - ] + tc_part["thoughtSignature"] = invocation["thought_signature"] assistant_content.append(tc_part) # type: ignore # Re-inject toolResponse part if response is present @@ -1047,9 +959,7 @@ def _gemini_convert_messages_with_history( tr_dict["toolType"] = invocation["tool_type"] tr_part: Dict[str, Any] = {"toolResponse": tr_dict} if "response_thought_signature" in invocation: - tr_part["thoughtSignature"] = invocation[ - "response_thought_signature" - ] + tr_part["thoughtSignature"] = invocation["response_thought_signature"] assistant_content.append(tr_part) # type: ignore msg_i += 1 @@ -1059,10 +969,7 @@ def _gemini_convert_messages_with_history( ## APPEND TOOL CALL MESSAGES ## tool_call_message_roles = ["tool", "function"] - if ( - msg_i < len(messages) - and messages[msg_i]["role"] in tool_call_message_roles - ): + if msg_i < len(messages) and messages[msg_i]["role"] in tool_call_message_roles: _part = convert_to_gemini_tool_call_result( messages[msg_i], # type: ignore last_message_with_tool_calls, # type: ignore @@ -1075,9 +982,7 @@ def _gemini_convert_messages_with_history( tool_call_responses.extend(_part) else: tool_call_responses.append(_part) - if msg_i < len(messages) and ( - messages[msg_i]["role"] not in tool_call_message_roles - ): + if msg_i < len(messages) and (messages[msg_i]["role"] not in tool_call_message_roles): if len(tool_call_responses) > 0: contents.append(ContentType(role="user", parts=tool_call_responses)) tool_call_responses = [] @@ -1118,11 +1023,7 @@ def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None: for k, v in extra_body.items(): if k in _LITELLM_INTERNAL_EXTRA_BODY_KEYS: continue - if ( - k in data_dict - and isinstance(data_dict[k], dict) - and isinstance(v, dict) - ): + if k in data_dict and isinstance(data_dict[k], dict) and isinstance(v, dict): data_dict[k].update(v) else: data_dict[k] = v @@ -1132,9 +1033,7 @@ def _has_google_maps_tool(tools: Optional[Any]) -> bool: """Return True if any tool object in the list has a 'googleMaps' key.""" if not isinstance(tools, list): return False - return any( - isinstance(t, dict) and VertexToolName.GOOGLE_MAPS.value in t for t in tools - ) + return any(isinstance(t, dict) and VertexToolName.GOOGLE_MAPS.value in t for t in tools) def _rewrite_mime_type_to_response_format(generation_config: GenerationConfig) -> None: @@ -1195,17 +1094,13 @@ def _transform_request_body( Common transformation logic across sync + async Gemini /generateContent calls. """ # Separate system prompt from rest of message - supports_system_message = get_supports_system_message( - model=model, custom_llm_provider=custom_llm_provider - ) + supports_system_message = get_supports_system_message(model=model, custom_llm_provider=custom_llm_provider) system_instructions, messages = _transform_system_message( supports_system_message=supports_system_message, messages=messages ) # Checks for 'response_schema' support - if passed in if "response_schema" in optional_params: - supports_response_schema = get_supports_response_schema( - model=model, custom_llm_provider=custom_llm_provider - ) + supports_response_schema = get_supports_response_schema(model=model, custom_llm_provider=custom_llm_provider) if supports_response_schema is False: user_response_schema_message = response_schema_prompt( model=model, @@ -1235,12 +1130,8 @@ def _transform_request_body( ) tools: Optional[Tools] = optional_params.pop("tools", None) tool_choice: Optional[ToolConfig] = optional_params.pop("tool_choice", None) - include_server_side_tool_invocations: bool = optional_params.pop( - "include_server_side_tool_invocations", False - ) - safety_settings: Optional[List[SafetSettingsConfig]] = optional_params.pop( - "safety_settings", None - ) # type: ignore + include_server_side_tool_invocations: bool = optional_params.pop("include_server_side_tool_invocations", False) + safety_settings: Optional[List[SafetSettingsConfig]] = optional_params.pop("safety_settings", None) # type: ignore # Drop output_config as it's not supported by Vertex AI optional_params.pop("output_config", None) config_fields = GenerationConfig.__annotations__.keys() @@ -1248,15 +1139,9 @@ def _transform_request_body( # labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata) labels = pop_vertex_request_labels(optional_params, litellm_params) - filtered_params = { - k: v - for k, v in optional_params.items() - if _get_equivalent_key(k, set(config_fields)) - } + filtered_params = {k: v for k, v in optional_params.items() if _get_equivalent_key(k, set(config_fields))} - generation_config: Optional[GenerationConfig] = GenerationConfig( - **filtered_params - ) + generation_config: Optional[GenerationConfig] = GenerationConfig(**filtered_params) # For Gemini 2.x models, also add media_resolution to generation_config (global) # as a fallback, since some 2.x versions may not support per-part media_resolution. @@ -1264,20 +1149,14 @@ def _transform_request_body( if "gemini-2" in model: max_media_resolution = _extract_max_media_resolution_from_messages(messages) if max_media_resolution: - media_resolution_value = _convert_detail_to_media_resolution_enum( - max_media_resolution - ) + media_resolution_value = _convert_detail_to_media_resolution_enum(max_media_resolution) if media_resolution_value and generation_config is not None: - generation_config["mediaResolution"] = media_resolution_value[ - "level" - ] + generation_config["mediaResolution"] = media_resolution_value["level"] data = RequestBody(contents=content) # Vertex rejects system_instruction/tools/toolConfig alongside cachedContent. # Treat dropping these fields as a request mutation guarded by modify_params. - can_send_cache_incompatible_fields = ( - cached_content is None or litellm.modify_params is False - ) + can_send_cache_incompatible_fields = cached_content is None or litellm.modify_params is False if can_send_cache_incompatible_fields: if system_instructions is not None: data["system_instruction"] = system_instructions diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 32bdb17840d..678877c0721 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -256,9 +256,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if isinstance(response_format, dict): return response_format - if isinstance(response_format, type) and issubclass( - response_format, _BaseModel - ): + if isinstance(response_format, type) and issubclass(response_format, _BaseModel): schema = response_format.model_json_schema() return { "type": "json_schema", @@ -291,9 +289,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return False @staticmethod - def _forward_gemini_function_call_id( - model: str, custom_llm_provider: Optional[str] = None - ) -> bool: + def _forward_gemini_function_call_id(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Whether to include `id` on function_call / function_response parts. @@ -348,9 +344,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): supported_params.append("thinking") return supported_params - def map_tool_choice_values( - self, model: str, tool_choice: Union[str, dict] - ) -> Optional[ToolConfig]: + def map_tool_choice_values(self, model: str, tool_choice: Union[str, dict]) -> Optional[ToolConfig]: if tool_choice == "none": return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="NONE")) elif tool_choice == "required": @@ -360,11 +354,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif isinstance(tool_choice, dict): # only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html name = tool_choice.get("function", {}).get("name", "") - return ToolConfig( - functionCallingConfig=FunctionCallingConfig( - mode="ANY", allowed_function_names=[name] - ) - ) + return ToolConfig(functionCallingConfig=FunctionCallingConfig(mode="ANY", allowed_function_names=[name])) else: raise litellm.utils.UnsupportedParamsError( message="VertexAI doesn't support tool_choice={}. Supported tool_choice values=['auto', 'required', json object]. To drop it from the call, set `litellm.drop_params = True.".format( @@ -412,16 +402,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return search_tool_keys = cls._search_tool_keys() - has_function_declarations = any( - isinstance(tool, dict) and tool.get("function_declarations") - for tool in tools - ) + has_function_declarations = any(isinstance(tool, dict) and tool.get("function_declarations") for tool in tools) if not has_function_declarations: return has_search_tools = any( - isinstance(tool, dict) and any(key in tool for key in search_tool_keys) - for tool in tools + isinstance(tool, dict) and any(key in tool for key in search_tool_keys) for tool in tools ) if not has_search_tools: return @@ -434,11 +420,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "send a request without function calling tools." ) optional_params["tools"] = [ - tool - for tool in tools - if not ( - isinstance(tool, dict) and any(key in tool for key in search_tool_keys) - ) + tool for tool in tools if not (isinstance(tool, dict) and any(key in tool for key in search_tool_keys)) ] def _map_service_tier_param(self, value: str, optional_params: dict) -> None: @@ -482,19 +464,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Transform excluded_predefined_functions to camelCase if "excluded_predefined_functions" in computer_use_config: - transformed_config["excludedPredefinedFunctions"] = computer_use_config[ - "excluded_predefined_functions" - ] + transformed_config["excludedPredefinedFunctions"] = computer_use_config["excluded_predefined_functions"] elif "excludedPredefinedFunctions" in computer_use_config: - transformed_config["excludedPredefinedFunctions"] = computer_use_config[ - "excludedPredefinedFunctions" - ] + transformed_config["excludedPredefinedFunctions"] = computer_use_config["excludedPredefinedFunctions"] return transformed_config - def _extract_google_maps_retrieval_config( - self, google_maps_config: dict - ) -> Tuple[dict, Optional[dict]]: + def _extract_google_maps_retrieval_config(self, google_maps_config: dict) -> Tuple[dict, Optional[dict]]: """ Extract location configuration from googleMaps tool for Vertex AI toolConfig. @@ -527,9 +503,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Remove location fields from tool definition cleaned_config = { - k: v - for k, v in google_maps_config.items() - if k not in ["latitude", "longitude", "languageCode"] + k: v for k, v in google_maps_config.items() if k not in ["latitude", "longitude", "languageCode"] } return cleaned_config, retrieval_config @@ -546,9 +520,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Optional[dict]: The tool value if found, None otherwise """ # Convert camelCase to underscore_case - underscore_name = "".join( - ["_" + c.lower() if c.isupper() else c for c in tool_name] - ).lstrip("_") + underscore_name = "".join(["_" + c.lower() if c.isupper() else c for c in tool_name]).lstrip("_") # Try both camelCase and underscore_case variants if tool.get(tool_name) is not None: @@ -592,14 +564,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): urlContext, ] ) - server_side_tool_invocations = optional_params.get( - "include_server_side_tool_invocations", False - ) - if ( - gtool_func_declarations - and has_search_tools - and not server_side_tool_invocations - ): + server_side_tool_invocations = optional_params.get("include_server_side_tool_invocations", False) + if gtool_func_declarations and has_search_tools and not server_side_tool_invocations: verbose_logger.warning( "Vertex AI does not support mixing function declarations with " "search tools (googleSearch, enterpriseWebSearch, urlContext, " @@ -644,9 +610,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): value = _remove_strict_from_schema(value) for tool in value: - openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = ( - None - ) + openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = None if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore **tool["function"] @@ -657,9 +621,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): and _openai_function_object["parameters"] is not None and isinstance(_openai_function_object["parameters"], dict) ): # OPENAI accepts JSON Schema, Google accepts OpenAPI schema. - _openai_function_object["parameters"] = _build_vertex_schema( - _openai_function_object["parameters"] - ) + _openai_function_object["parameters"] = _build_vertex_schema(_openai_function_object["parameters"]) openai_function_object = _openai_function_object @@ -675,68 +637,43 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "web_search", "web_search_preview", ): - verbose_logger.info( - f"Gemini: Transforming OpenAI-style '{tool['type']}' tool to googleSearch" - ) + verbose_logger.info(f"Gemini: Transforming OpenAI-style '{tool['type']}' tool to googleSearch") tool = {VertexToolName.GOOGLE_SEARCH.value: {}} # Handle tools with 'type' field (OpenAI spec compliance) Ignore this field -> https://github.com/BerriAI/litellm/issues/14644#issuecomment-3342061838 elif "type" in tool: tool = {k: tool[k] for k in tool if k != "type"} tool_name = list(tool.keys())[0] if len(tool.keys()) == 1 else None if tool_name and ( - tool_name == "codeExecution" - or tool_name == VertexToolName.CODE_EXECUTION.value + tool_name == "codeExecution" or tool_name == VertexToolName.CODE_EXECUTION.value ): # code_execution maintained for backwards compatibility code_execution = self.get_tool_value(tool, "codeExecution") - elif tool_name and ( - tool_name == VertexToolName.GOOGLE_SEARCH.value - or tool_name == "google_search" - ): + elif tool_name and (tool_name == VertexToolName.GOOGLE_SEARCH.value or tool_name == "google_search"): googleSearch = self.get_tool_value(tool, tool_name) elif tool_name and ( - tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value - or tool_name == "google_search_retrieval" + tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value or tool_name == "google_search_retrieval" ): googleSearchRetrieval = self.get_tool_value(tool, tool_name) elif tool_name and ( - tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value - or tool_name == "enterprise_web_search" + tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value or tool_name == "enterprise_web_search" ): enterpriseWebSearch = self.get_tool_value(tool, tool_name) - elif tool_name and ( - tool_name == VertexToolName.URL_CONTEXT.value - or tool_name == "urlContext" - ): + elif tool_name and (tool_name == VertexToolName.URL_CONTEXT.value or tool_name == "urlContext"): urlContext = self.get_tool_value(tool, tool_name) - elif tool_name and ( - tool_name == VertexToolName.GOOGLE_MAPS.value - or tool_name == "google_maps" - ): - google_maps_value = self.get_tool_value( - tool, VertexToolName.GOOGLE_MAPS.value - ) + elif tool_name and (tool_name == VertexToolName.GOOGLE_MAPS.value or tool_name == "google_maps"): + google_maps_value = self.get_tool_value(tool, VertexToolName.GOOGLE_MAPS.value) # Extract and transform location configuration for toolConfig if google_maps_value is not None: ( googleMaps, google_maps_retrieval_config, - ) = self._extract_google_maps_retrieval_config( - google_maps_config=google_maps_value - ) - elif tool_name and ( - tool_name == VertexToolName.COMPUTER_USE.value - or tool_name == "computer_use" - ): - computer_use_value = self.get_tool_value( - tool, VertexToolName.COMPUTER_USE.value - ) + ) = self._extract_google_maps_retrieval_config(google_maps_config=google_maps_value) + elif tool_name and (tool_name == VertexToolName.COMPUTER_USE.value or tool_name == "computer_use"): + computer_use_value = self.get_tool_value(tool, VertexToolName.COMPUTER_USE.value) # Transform Computer Use configuration to Gemini API format if computer_use_value is not None: - computerUse = self._transform_computer_use_config( - computer_use_config=computer_use_value - ) + computerUse = self._transform_computer_use_config(computer_use_config=computer_use_value) else: # Empty config - Gemini will use defaults computerUse = {} @@ -792,15 +729,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tools_list.append(search_tool) if googleSearchRetrieval is not None: retrieval_tool = Tools() - retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = ( - googleSearchRetrieval - ) + retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval _tools_list.append(retrieval_tool) if enterpriseWebSearch is not None: enterprise_tool = Tools() - enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = ( - enterpriseWebSearch - ) + enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch _tools_list.append(enterprise_tool) if code_execution is not None: code_tool = Tools() @@ -823,9 +756,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if google_maps_retrieval_config is not None: if "toolConfig" not in optional_params: optional_params["toolConfig"] = {} - optional_params["toolConfig"]["retrievalConfig"] = ( - google_maps_retrieval_config - ) + optional_params["toolConfig"]["retrievalConfig"] = google_maps_retrieval_config return _tools_list @@ -834,19 +765,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if isinstance(old_schema, list): for item in old_schema: if isinstance(item, dict): - item = _build_vertex_schema( - parameters=item, add_property_ordering=True - ) + item = _build_vertex_schema(parameters=item, add_property_ordering=True) elif isinstance(old_schema, dict): - old_schema = _build_vertex_schema( - parameters=old_schema, add_property_ordering=True - ) + old_schema = _build_vertex_schema(parameters=old_schema, add_property_ordering=True) return old_schema - def apply_response_schema_transformation( - self, value: dict, optional_params: dict, model: str - ): + def apply_response_schema_transformation(self, value: dict, optional_params: dict, model: str): new_value = deepcopy(value) # remove 'strict' from json schema (not supported by Gemini) new_value = _remove_strict_from_schema(new_value) @@ -882,17 +807,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # - Standard JSON Schema format (lowercase types) # - Supports additionalProperties # - No propertyOrdering needed - optional_params["response_json_schema"] = _build_json_schema( - deepcopy(schema) - ) + optional_params["response_json_schema"] = _build_json_schema(deepcopy(schema)) else: # Use responseSchema (default, backwards compatible) # - OpenAPI-style format (uppercase types) # - No additionalProperties support # - Requires propertyOrdering - optional_params["response_schema"] = self._map_response_schema( - value=schema - ) + optional_params["response_schema"] = self._map_response_schema(value=schema) @staticmethod def _map_reasoning_effort_to_thinking_budget( @@ -907,9 +828,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif model and "gemini-2.5-pro" in model.lower(): budget = DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO elif model and "gemini-2.5-flash" in model.lower(): - budget = ( - DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH - ) + budget = DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH else: budget = DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET @@ -962,9 +881,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Check if this is gemini-3-flash which supports MINIMAL thinking level # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, # gemini-3.5-flash, and any future 3.x-flash variants. - is_gemini3flash = model and ( - "flash" in model.lower() and "gemini-3" in model.lower() - ) + is_gemini3flash = model and ("flash" in model.lower() and "gemini-3" in model.lower()) is_gemini31pro = model and ("gemini-3.1-pro-preview" in model.lower()) if reasoning_effort == "minimal": if is_gemini3flash: @@ -1058,20 +975,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): params["includeThoughts"] = True # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: - is_gemini3flash = ( - "gemini-3" in model.lower() and "flash" in model.lower() - ) - params["thinkingLevel"] = ( - "minimal" if is_gemini3flash else "low" - ) + is_gemini3flash = "gemini-3" in model.lower() and "flash" in model.lower() + params["thinkingLevel"] = "minimal" if is_gemini3flash else "low" else: # Thinking disabled params["includeThoughts"] = False else: # For older Gemini models, use thinkingBudget - if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero( - thinking_budget - ): + if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero(thinking_budget): params["includeThoughts"] = True if thinking_budget is not None and isinstance(thinking_budget, int): params["thinkingBudget"] = thinking_budget @@ -1178,9 +1089,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model: str, drop_params: bool, ) -> Dict: - self._apply_include_server_side_tool_invocations( - non_default_params, optional_params - ) + self._apply_include_server_side_tool_invocations(non_default_params, optional_params) gemini_sampling_params_warned: bool = False for param, value in non_default_params.items(): if param == "temperature": @@ -1201,10 +1110,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): gemini_sampling_params_warned = True optional_params["temperature"] = value elif param == "top_p": - if ( - VertexGeminiConfig._is_gemini_3_or_newer(model) - and not gemini_sampling_params_warned - ): + if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned: verbose_logger.warning( "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to " f"function for Gemini 3+ ({model}) but are planned for removal in a " @@ -1214,10 +1120,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): gemini_sampling_params_warned = True optional_params["top_p"] = value elif param == "top_k": - if ( - VertexGeminiConfig._is_gemini_3_or_newer(model) - and not gemini_sampling_params_warned - ): + if VertexGeminiConfig._is_gemini_3_or_newer(model) and not gemini_sampling_params_warned: verbose_logger.warning( "DeprecationWarning: `temperature`, `top_p`, and `top_k` continue to " f"function for Gemini 3+ ({model}) but are planned for removal in a " @@ -1242,9 +1145,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif param == "max_tokens" or param == "max_completion_tokens": optional_params["max_output_tokens"] = value elif param == "response_format" and isinstance(value, dict): # type: ignore - self.apply_response_schema_transformation( - value=value, optional_params=optional_params, model=model - ) + self.apply_response_schema_transformation(value=value, optional_params=optional_params, model=model) elif param == "frequency_penalty": if self._supports_penalty_parameters(model): optional_params["frequency_penalty"] = value @@ -1255,21 +1156,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params["responseLogprobs"] = value elif param == "top_logprobs": optional_params["logprobs"] = value - elif ( - (param == "tools" or param == "functions") - and isinstance(value, list) - and value - ): + elif (param == "tools" or param == "functions") and isinstance(value, list) and value: # Pass optional_params so _map_function can add toolConfig if needed - mapped_tools = self._map_function( - value=value, optional_params=optional_params - ) - optional_params = self._add_tools_to_optional_params( - optional_params, mapped_tools - ) - elif param == "tool_choice" and ( - isinstance(value, str) or isinstance(value, dict) - ): + mapped_tools = self._map_function(value=value, optional_params=optional_params) + optional_params = self._add_tools_to_optional_params(optional_params, mapped_tools) + elif param == "tool_choice" and (isinstance(value, str) or isinstance(value, dict)): _tool_choice_value = self.map_tool_choice_values( model=model, tool_choice=value, # type: ignore @@ -1277,9 +1168,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value elif param == "parallel_tool_calls": - tools_list = non_default_params.get( - "tools", non_default_params.get("functions") - ) + tools_list = non_default_params.get("tools", non_default_params.get("functions")) num_tools = len(tools_list) if isinstance(tools_list, list) else 0 # Gemini does not support parallel_tool_calls=False with multiple # tools. Drop the param instead of failing — Responses API clients @@ -1305,16 +1194,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_description="thinking_budget", ) if VertexGeminiConfig._is_gemini_3_or_newer(model): - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - effort_value, model - ) + optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + effort_value, model ) else: - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - effort_value, model - ) + optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( + effort_value, model ) elif param == "thinking": # Validate no conflict with thinking_level @@ -1323,20 +1208,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_name="thinking", param_description="thinking_budget", ) - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_thinking_param( - cast(AnthropicThinkingParam, value), - model=model, - ) + optional_params["thinkingConfig"] = VertexGeminiConfig._map_thinking_param( + cast(AnthropicThinkingParam, value), + model=model, ) elif param == "modalities" and isinstance(value, list): response_modalities = self.map_response_modalities(value) optional_params["responseModalities"] = response_modalities elif param == "web_search_options" and isinstance(value, dict): _tools = self._map_web_search_options(value) - optional_params = self._add_tools_to_optional_params( - optional_params, [_tools] - ) + optional_params = self._add_tools_to_optional_params(optional_params, [_tools]) elif param == "service_tier" and isinstance(value, str): self._map_service_tier_param(value, optional_params) elif param == "include_server_side_tool_invocations" and value is True: @@ -1479,11 +1360,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ from litellm.litellm_core_utils.core_helpers import _FINISH_REASON_MAP - return { - k: v - for k, v in _FINISH_REASON_MAP.items() - if k in VertexGeminiConfig._GEMINI_FINISH_REASON_KEYS - } + return {k: v for k, v in _FINISH_REASON_MAP.items() if k in VertexGeminiConfig._GEMINI_FINISH_REASON_KEYS} def translate_exception_str(self, exception_string: str): if ( @@ -1495,9 +1372,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) return exception_string - def get_assistant_content_message( - self, parts: List[HttpxPartType] - ) -> Tuple[Optional[str], Optional[str]]: + def get_assistant_content_message(self, parts: List[HttpxPartType]) -> Tuple[Optional[str], Optional[str]]: content_str: Optional[str] = None reasoning_content_str: Optional[str] = None @@ -1509,9 +1384,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if text_content.startswith("data:audio") and ";base64," in text_content: try: if is_base64_encoded(text_content): - media_type, _ = text_content.split("data:")[1].split( - ";base64," - ) + media_type, _ = text_content.split("data:")[1].split(";base64,") if media_type.startswith("audio/"): continue except (ValueError, IndexError): @@ -1540,9 +1413,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return content_str, reasoning_content_str - def _extract_thinking_blocks_from_parts( - self, parts: List[HttpxPartType] - ) -> List[ChatCompletionThinkingBlock]: + def _extract_thinking_blocks_from_parts(self, parts: List[HttpxPartType]) -> List[ChatCompletionThinkingBlock]: """Extract thinking blocks from parts if present. Per Google's docs (https://ai.google.dev/gemini-api/docs/thinking): @@ -1565,9 +1436,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): thinking_blocks.append(block) return thinking_blocks - def _extract_thought_signatures_from_parts( - self, parts: List[HttpxPartType] - ) -> Optional[List[str]]: + def _extract_thought_signatures_from_parts(self, parts: List[HttpxPartType]) -> Optional[List[str]]: """Extract thoughtSignature values from parts. Per Google's docs, thoughtSignature is returned for multi-turn context preservation @@ -1641,16 +1510,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Any orphan responses (shouldn't happen, but be safe) for resp_id, resp_entry in tool_responses_by_id.items(): if "thought_signature" in resp_entry: - resp_entry["response_thought_signature"] = resp_entry[ - "thought_signature" - ] + resp_entry["response_thought_signature"] = resp_entry["thought_signature"] invocations.append(resp_entry) return invocations if invocations else None - def _extract_image_response_from_parts( - self, parts: List[HttpxPartType] - ) -> Optional[List[ImageURLListItem]]: + def _extract_image_response_from_parts(self, parts: List[HttpxPartType]) -> Optional[List[ImageURLListItem]]: """Extract image response from parts if present""" images: List[ImageURLListItem] = [] for part in parts: @@ -1670,9 +1535,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) return images - def _extract_audio_response_from_parts( - self, parts: List[HttpxPartType] - ) -> Optional[ChatCompletionAudioResponse]: + def _extract_audio_response_from_parts(self, parts: List[HttpxPartType]) -> Optional[ChatCompletionAudioResponse]: """Extract audio response from parts if present""" for part in parts: if "text" in part: @@ -1681,9 +1544,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if text_content.startswith("data:audio") and ";base64," in text_content: try: if is_base64_encoded(text_content): - media_type, audio_data = text_content.split("data:")[ - 1 - ].split(";base64,") + media_type, audio_data = text_content.split("data:")[1].split(";base64,") if media_type.startswith("audio/"): expires_at = int(time.time()) + (24 * 60 * 60) @@ -1706,9 +1567,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): expires_at = int(time.time()) + (24 * 60 * 60) transcript = "" # Gemini doesn't provide transcript - return ChatCompletionAudioResponse( - data=data, expires_at=expires_at, transcript=transcript - ) + return ChatCompletionAudioResponse(data=data, expires_at=expires_at, transcript=transcript) return None @@ -1728,9 +1587,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if "functionCall" in part: _function_chunk: ChatCompletionToolCallFunctionChunk = { "name": part["functionCall"]["name"], - "arguments": json.dumps( - part["functionCall"]["args"], ensure_ascii=False - ), + "arguments": json.dumps(part["functionCall"]["args"], ensure_ascii=False), } # Extract thought signature if present thought_signature = part.get("thoughtSignature") @@ -1744,9 +1601,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if thought_signature: if "provider_specific_fields" not in function_dict: function_dict["provider_specific_fields"] = {} - function_dict["provider_specific_fields"][ - "thought_signature" - ] = thought_signature + function_dict["provider_specific_fields"]["thought_signature"] = thought_signature function = cast(ChatCompletionToolCallFunctionChunk, function_dict) else: _tool_response_chunk: ChatCompletionToolCallChunk = { @@ -1765,10 +1620,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tool_response_chunk["provider_specific_fields"] = { # type: ignore "thought_signature": thought_signature } - _tool_response_chunk["id"] = ( - _encode_tool_call_id_with_signature( - _tool_response_chunk["id"] or "", thought_signature - ) + _tool_response_chunk["id"] = _encode_tool_call_id_with_signature( + _tool_response_chunk["id"] or "", thought_signature ) _tools.append(_tool_response_chunk) cumulative_tool_call_idx += 1 @@ -1789,19 +1642,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): logprobs_list: List[ChatCompletionTokenLogprob] = [] for index, candidate in enumerate(logprobs_result["chosenCandidates"]): top_logprobs: List[TopLogprob] = [] - if "topCandidates" in logprobs_result and index < len( - logprobs_result["topCandidates"] - ): - top_candidates_for_index = logprobs_result["topCandidates"][index][ - "candidates" - ] + if "topCandidates" in logprobs_result and index < len(logprobs_result["topCandidates"]): + top_candidates_for_index = logprobs_result["topCandidates"][index]["candidates"] for options in top_candidates_for_index: - top_logprobs.append( - TopLogprob( - token=options["token"], logprob=options["logProbability"] - ) - ) + top_logprobs.append(TopLogprob(token=options["token"], logprob=options["logProbability"])) logprobs_list.append( ChatCompletionTokenLogprob( token=candidate["token"], @@ -1836,12 +1681,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## GET USAGE ## usage = Usage( - prompt_tokens=completion_response["usageMetadata"].get( - "promptTokenCount", 0 - ), - completion_tokens=completion_response["usageMetadata"].get( - "candidatesTokenCount", 0 - ), + prompt_tokens=completion_response["usageMetadata"].get("promptTokenCount", 0), + completion_tokens=completion_response["usageMetadata"].get("candidatesTokenCount", 0), total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0), ) @@ -1874,12 +1715,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## GET USAGE ## usage = Usage( - prompt_tokens=completion_response["usageMetadata"].get( - "promptTokenCount", 0 - ), - completion_tokens=completion_response["usageMetadata"].get( - "candidatesTokenCount", 0 - ), + prompt_tokens=completion_response["usageMetadata"].get("promptTokenCount", 0), + completion_tokens=completion_response["usageMetadata"].get("candidatesTokenCount", 0), total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0), ) @@ -1907,17 +1744,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _calculate_usage( - completion_response: Union[ - GenerateContentResponseBody, BidiGenerateContentServerMessage - ], + completion_response: Union[GenerateContentResponseBody, BidiGenerateContentServerMessage], ) -> Usage: - if ( - completion_response is not None - and "usageMetadata" not in completion_response - ): - raise ValueError( - f"usageMetadata not found in completion_response. Got={completion_response}" - ) + if completion_response is not None and "usageMetadata" not in completion_response: + raise ValueError(f"usageMetadata not found in completion_response. Got={completion_response}") cached_tokens: Optional[int] = None # Separate variables for prompt tokens by modality prompt_audio_tokens: Optional[int] = None @@ -1946,17 +1776,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): modality = str(detail.get("modality", "")).upper() token_count = _get_token_count(detail) if modality == "TEXT": - response_tokens_details.text_tokens = ( - response_tokens_details.text_tokens or 0 - ) + token_count + response_tokens_details.text_tokens = (response_tokens_details.text_tokens or 0) + token_count elif modality == "AUDIO": - response_tokens_details.audio_tokens = ( - response_tokens_details.audio_tokens or 0 - ) + token_count + response_tokens_details.audio_tokens = (response_tokens_details.audio_tokens or 0) + token_count elif modality == "DOCUMENT": - response_tokens_details.text_tokens = ( - response_tokens_details.text_tokens or 0 - ) + token_count + response_tokens_details.text_tokens = (response_tokens_details.text_tokens or 0) + token_count ######################################################### @@ -1968,25 +1792,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): modality = str(detail.get("modality", "")).upper() token_count = _get_token_count(detail) if modality == "TEXT": - response_tokens_details.text_tokens = ( - response_tokens_details.text_tokens or 0 - ) + token_count + response_tokens_details.text_tokens = (response_tokens_details.text_tokens or 0) + token_count elif modality == "AUDIO": - response_tokens_details.audio_tokens = ( - response_tokens_details.audio_tokens or 0 - ) + token_count + response_tokens_details.audio_tokens = (response_tokens_details.audio_tokens or 0) + token_count elif modality == "IMAGE": - response_tokens_details.image_tokens = ( - response_tokens_details.image_tokens or 0 - ) + token_count + response_tokens_details.image_tokens = (response_tokens_details.image_tokens or 0) + token_count elif modality == "VIDEO": - response_tokens_details.video_tokens = ( - response_tokens_details.video_tokens or 0 - ) + token_count + response_tokens_details.video_tokens = (response_tokens_details.video_tokens or 0) + token_count elif modality == "DOCUMENT": - response_tokens_details.text_tokens = ( - response_tokens_details.text_tokens or 0 - ) + token_count + response_tokens_details.text_tokens = (response_tokens_details.text_tokens or 0) + token_count # Calculate text_tokens if not explicitly provided in candidatesTokensDetails # candidatesTokenCount includes all modalities, so: text = total - (image + audio + video) @@ -1999,10 +1813,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): completion_audio_tokens = response_tokens_details.audio_tokens or 0 completion_video_tokens = response_tokens_details.video_tokens or 0 calculated_text_tokens = ( - candidates_token_count - - completion_image_tokens - - completion_audio_tokens - - completion_video_tokens + candidates_token_count - completion_image_tokens - completion_audio_tokens - completion_video_tokens ) response_tokens_details.text_tokens = calculated_text_tokens ######################################################### @@ -2083,13 +1894,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): video_tokens=prompt_video_tokens, ) - completion_tokens = response_tokens or completion_response["usageMetadata"].get( - "candidatesTokenCount", 0 - ) - if ( - not VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) - and reasoning_tokens - ): + completion_tokens = response_tokens or completion_response["usageMetadata"].get("candidatesTokenCount", 0) + if not VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) and reasoning_tokens: completion_tokens = reasoning_tokens + completion_tokens ## GET USAGE ## usage = Usage( @@ -2170,11 +1976,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _calculate_web_search_requests(grounding_metadata: List[dict]) -> Optional[int]: web_search_requests: Optional[int] = None - if ( - grounding_metadata - and isinstance(grounding_metadata, list) - and len(grounding_metadata) > 0 - ): + if grounding_metadata and isinstance(grounding_metadata, list) and len(grounding_metadata) > 0: for grounding_metadata_item in grounding_metadata: web_search_queries = grounding_metadata_item.get("webSearchQueries") if web_search_queries and web_search_requests: @@ -2199,9 +2001,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): from litellm.types.utils import Delta, StreamingChoices annotations = chat_completion_message.get("annotations") # type: ignore - provider_specific_fields = chat_completion_message.get( - "provider_specific_fields" - ) # type: ignore + provider_specific_fields = chat_completion_message.get("provider_specific_fields") # type: ignore # create a streaming choice object choice = StreamingChoices( finish_reason=VertexGeminiConfig._check_finish_reason( @@ -2290,14 +2090,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) -> None: setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore if grounding_metadata: - model_response._hidden_params["vertex_ai_grounding_metadata"] = ( - grounding_metadata - ) + model_response._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore if url_context_metadata: - model_response._hidden_params["vertex_ai_url_context_metadata"] = ( - url_context_metadata - ) + model_response._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore setattr(model_response, "vertex_ai_safety_results", safety_ratings) # type: ignore if safety_ratings: @@ -2305,9 +2101,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model_response._hidden_params["vertex_ai_safety_results"] = safety_ratings setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore if citation_metadata: - model_response._hidden_params["vertex_ai_citation_metadata"] = ( - citation_metadata - ) + model_response._hidden_params["vertex_ai_citation_metadata"] = citation_metadata def apply_assembled_streaming_response_metadata( self, @@ -2440,51 +2234,35 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ( content, reasoning_content, - ) = VertexGeminiConfig().get_assistant_content_message( + ) = VertexGeminiConfig().get_assistant_content_message(parts=candidate["content"]["parts"]) + + audio_response = VertexGeminiConfig()._extract_audio_response_from_parts( + parts=candidate["content"]["parts"] + ) + image_response = VertexGeminiConfig()._extract_image_response_from_parts( parts=candidate["content"]["parts"] ) - audio_response = ( - VertexGeminiConfig()._extract_audio_response_from_parts( - parts=candidate["content"]["parts"] - ) - ) - image_response = ( - VertexGeminiConfig()._extract_image_response_from_parts( - parts=candidate["content"]["parts"] - ) - ) - - thinking_blocks = ( - VertexGeminiConfig()._extract_thinking_blocks_from_parts( - parts=candidate["content"]["parts"] - ) + thinking_blocks = VertexGeminiConfig()._extract_thinking_blocks_from_parts( + parts=candidate["content"]["parts"] ) # Extract thoughtSignatures from parts (can exist without thought: true) - thought_signatures = ( - VertexGeminiConfig()._extract_thought_signatures_from_parts( - parts=candidate["content"]["parts"] - ) + thought_signatures = VertexGeminiConfig()._extract_thought_signatures_from_parts( + parts=candidate["content"]["parts"] ) # Extract server-side tool invocations (context circulation) - server_side_tool_invocations = ( - VertexGeminiConfig._extract_server_side_tool_invocations( - parts=candidate["content"]["parts"] - ) + server_side_tool_invocations = VertexGeminiConfig._extract_server_side_tool_invocations( + parts=candidate["content"]["parts"] ) if audio_response is not None: - cast(Dict[str, Any], chat_completion_message)["audio"] = ( - audio_response - ) + cast(Dict[str, Any], chat_completion_message)["audio"] = audio_response chat_completion_message["content"] = None # OpenAI spec if image_response is not None: # Handle image response - combine with text content into structured format - cast(Dict[str, Any], chat_completion_message)["images"] = ( - image_response - ) + cast(Dict[str, Any], chat_completion_message)["images"] = image_response if content is not None: chat_completion_message["content"] = content @@ -2492,11 +2270,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): chat_completion_message["reasoning_content"] = reasoning_content if candidate_grounding_metadata: - annotations = ( - VertexGeminiConfig._convert_grounding_metadata_to_annotations( - grounding_metadata=candidate_grounding_metadata, - content_text=content, - ) + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( + grounding_metadata=candidate_grounding_metadata, + content_text=content, ) if annotations: chat_completion_message["annotations"] = annotations # type: ignore @@ -2526,10 +2302,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Convert thinking_blocks to reasoning_content for streaming # This ensures reasoning_content is available in streaming responses - if ( - isinstance(model_response, ModelResponseStream) - and reasoning_content is None - ): + if isinstance(model_response, ModelResponseStream) and reasoning_content is None: reasoning_content_parts = [] for block in thinking_blocks: thinking_text = block.get("thinking") @@ -2544,17 +2317,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if thought_signatures is not None: if "provider_specific_fields" not in chat_completion_message: chat_completion_message["provider_specific_fields"] = {} - chat_completion_message["provider_specific_fields"][ - "thought_signatures" - ] = thought_signatures # type: ignore + chat_completion_message["provider_specific_fields"]["thought_signatures"] = thought_signatures # type: ignore # Store server-side tool invocations in provider_specific_fields if server_side_tool_invocations is not None: if "provider_specific_fields" not in chat_completion_message: chat_completion_message["provider_specific_fields"] = {} - chat_completion_message["provider_specific_fields"][ - "server_side_tool_invocations" - ] = server_side_tool_invocations # type: ignore + chat_completion_message["provider_specific_fields"]["server_side_tool_invocations"] = ( + server_side_tool_invocations # type: ignore + ) if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( @@ -2647,10 +2418,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model_response.model = model ## CHECK IF RESPONSE FLAGGED - if ( - "promptFeedback" in completion_response - and "blockReason" in completion_response["promptFeedback"] - ): + if "promptFeedback" in completion_response and "blockReason" in completion_response["promptFeedback"]: return self._handle_blocked_response( model_response=model_response, completion_response=completion_response, @@ -2658,13 +2426,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _candidates = completion_response.get("candidates") if _candidates and len(_candidates) > 0: - content_policy_violations = ( - VertexGeminiConfig().get_flagged_finish_reasons() - ) - if ( - "finishReason" in _candidates[0] - and _candidates[0]["finishReason"] in content_policy_violations.keys() - ): + content_policy_violations = VertexGeminiConfig().get_flagged_finish_reasons() + if "finishReason" in _candidates[0] and _candidates[0]["finishReason"] in content_policy_violations.keys(): return self._handle_content_policy_violation( model_response=model_response, completion_response=completion_response, @@ -2686,38 +2449,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): safety_ratings, citation_metadata, _, # cumulative_tool_call_index not needed in non-streaming - ) = VertexGeminiConfig._process_candidates( - _candidates, model_response, logging_obj.optional_params - ) + ) = VertexGeminiConfig._process_candidates(_candidates, model_response, logging_obj.optional_params) - usage = VertexGeminiConfig._calculate_usage( - completion_response=completion_response - ) + usage = VertexGeminiConfig._calculate_usage(completion_response=completion_response) - web_search_requests = VertexGeminiConfig._calculate_web_search_requests( - grounding_metadata - ) + web_search_requests = VertexGeminiConfig._calculate_web_search_requests(grounding_metadata) if web_search_requests is not None: - cast( - PromptTokensDetailsWrapper, usage.prompt_tokens_details - ).web_search_requests = web_search_requests + cast(PromptTokensDetailsWrapper, usage.prompt_tokens_details).web_search_requests = web_search_requests setattr(model_response, "usage", usage) ## ADD METADATA TO RESPONSE ## setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) - model_response._hidden_params["vertex_ai_grounding_metadata"] = ( - grounding_metadata - ) + model_response._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata - setattr( - model_response, "vertex_ai_url_context_metadata", url_context_metadata - ) + setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) - model_response._hidden_params["vertex_ai_url_context_metadata"] = ( - url_context_metadata - ) + model_response._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata setattr(model_response, "vertex_ai_safety_results", safety_ratings) model_response._hidden_params["vertex_ai_safety_results"] = ( @@ -2731,13 +2480,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) ## ADD TRAFFIC TYPE ## - traffic_type = completion_response.get("usageMetadata", {}).get( - "trafficType" - ) + traffic_type = completion_response.get("usageMetadata", {}).get("trafficType") if traffic_type: - model_response._hidden_params.setdefault( - "provider_specific_fields", {} - )["traffic_type"] = traffic_type + model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type ## ADD SERVICE TIER ## if getattr(raw_response, "headers", None): @@ -2774,9 +2519,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers] ) -> BaseLLMException: - return VertexAIError( - message=error_message, status_code=status_code, headers=headers - ) + return VertexAIError(message=error_message, status_code=status_code, headers=headers) def transform_request( self, @@ -2786,9 +2529,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): litellm_params: Dict, headers: Dict, ) -> Dict: - raise NotImplementedError( - "Vertex AI has a custom implementation of transform_request. Needs sync + async." - ) + raise NotImplementedError("Vertex AI has a custom implementation of transform_request. Needs sync + async.") def validate_environment( self, @@ -2831,9 +2572,7 @@ async def make_call( ) try: - response = await client.post( - api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj - ) + response = await client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) response.raise_for_status() except httpx.HTTPStatusError as e: exception_string = str(await e.response.aread()) @@ -2882,9 +2621,7 @@ def make_sync_call( if client is None: client = HTTPHandler() # Create a new client if none provided - response = client.post( - api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj - ) + response = client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj) if response.status_code != 200 and response.status_code != 201: raise VertexAIError( @@ -2941,9 +2678,7 @@ class VertexLLM(VertexBase): gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, ) -> CustomStreamWrapper: - should_use_v1beta1_features = self.is_using_v1beta1_features( - optional_params=optional_params - ) + should_use_v1beta1_features = self.is_using_v1beta1_features(optional_params=optional_params) _auth_header, vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, @@ -3000,11 +2735,7 @@ class VertexLLM(VertexBase): completion_stream=None, make_call=partial( make_call, - gemini_client=( - client - if client is not None and isinstance(client, AsyncHTTPHandler) - else None - ), + gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), api_base=api_base, headers=headers, data=request_body_str, @@ -3043,9 +2774,7 @@ class VertexLLM(VertexBase): gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, ) -> Union[ModelResponse, CustomStreamWrapper]: - should_use_v1beta1_features = self.is_using_v1beta1_features( - optional_params=optional_params - ) + should_use_v1beta1_features = self.is_using_v1beta1_features(optional_params=optional_params) _auth_header, vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, @@ -3090,9 +2819,7 @@ class VertexLLM(VertexBase): if timeout: _async_client_params["timeout"] = timeout if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client( - params=_async_client_params, llm_provider=litellm.LlmProviders.VERTEX_AI - ) + client = get_async_httpx_client(params=_async_client_params, llm_provider=litellm.LlmProviders.VERTEX_AI) else: client = client # type: ignore ## LOGGING @@ -3231,9 +2958,7 @@ class VertexLLM(VertexBase): extra_headers=extra_headers, ) - should_use_v1beta1_features = self.is_using_v1beta1_features( - optional_params=optional_params - ) + should_use_v1beta1_features = self.is_using_v1beta1_features(optional_params=optional_params) _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, @@ -3292,11 +3017,7 @@ class VertexLLM(VertexBase): completion_stream=None, make_call=partial( make_sync_call, - gemini_client=( - client - if client is not None and isinstance(client, HTTPHandler) - else None - ), + gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None), api_base=url, data=request_data_str, model=model, @@ -3323,9 +3044,7 @@ class VertexLLM(VertexBase): client = client try: - response = client.post( - url=url, headers=headers, json=data, logging_obj=logging_obj - ) # type: ignore + response = client.post(url=url, headers=headers, json=data, logging_obj=logging_obj) # type: ignore response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -3428,11 +3147,7 @@ class ModelResponseIterator: # to correctly set finish_reason="tool_calls" per the OpenAI spec. if not self.has_seen_tool_calls: for choice in model_response.choices: - if ( - hasattr(choice, "delta") - and choice.delta - and choice.delta.tool_calls - ): + if hasattr(choice, "delta") and choice.delta and choice.delta.tool_calls: self.has_seen_tool_calls = True break @@ -3452,9 +3167,7 @@ class ModelResponseIterator: if self.has_seen_tool_calls: mapped_finish_reason = "tool_calls" else: - mapped_finish_reason = VertexGeminiConfig._check_finish_reason( - None, finish_reason_str - ) + mapped_finish_reason = VertexGeminiConfig._check_finish_reason(None, finish_reason_str) choice = StreamingChoices( finish_reason=mapped_finish_reason, index=candidate.get("index", 0), @@ -3502,19 +3215,13 @@ class ModelResponseIterator: completion_response=processed_chunk, ) - web_search_requests = VertexGeminiConfig._calculate_web_search_requests( - grounding_metadata - ) + web_search_requests = VertexGeminiConfig._calculate_web_search_requests(grounding_metadata) if web_search_requests is not None: - cast( - PromptTokensDetailsWrapper, usage.prompt_tokens_details - ).web_search_requests = web_search_requests + cast(PromptTokensDetailsWrapper, usage.prompt_tokens_details).web_search_requests = web_search_requests traffic_type = processed_chunk.get("usageMetadata", {}).get("trafficType") if traffic_type: - model_response._hidden_params.setdefault("provider_specific_fields", {})[ - "traffic_type" - ] = traffic_type + model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type service_tier = self.response_headers.get("x-gemini-service-tier") if service_tier: @@ -3561,9 +3268,7 @@ class ModelResponseIterator: citation_metadata, ) = self._apply_stream_candidates(_candidates, model_response) - usage = self._apply_stream_usage_metadata( - processed_chunk, model_response, grounding_metadata - ) + usage = self._apply_stream_usage_metadata(processed_chunk, model_response, grounding_metadata) setattr(model_response, "usage", usage) # type: ignore @@ -3595,9 +3300,7 @@ class ModelResponseIterator: return self.chunk_parser(chunk=json_chunk) - def handle_accumulated_json_chunk( - self, chunk: str - ) -> Optional["ModelResponseStream"]: + def handle_accumulated_json_chunk(self, chunk: str) -> Optional["ModelResponseStream"]: chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" message = chunk.replace("\n\n", "") @@ -3619,9 +3322,7 @@ class ModelResponseIterator: except json.JSONDecodeError: return None - def _common_chunk_parsing_logic( - self, chunk: str - ) -> Optional["ModelResponseStream"]: + def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]: try: chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" if len(chunk) > 0: @@ -3688,16 +3389,12 @@ class ModelResponseIterator: try: await iterator.aclose() except Exception as e: # noqa: BLE001 - verbose_logger.debug( - "ModelResponseIterator.aclose: error closing iterator: %s", e - ) + verbose_logger.debug("ModelResponseIterator.aclose: error closing iterator: %s", e) if self.response is not None: try: await self.response.aclose() except Exception as e: # noqa: BLE001 - verbose_logger.debug( - "ModelResponseIterator.aclose: error closing response: %s", e - ) + verbose_logger.debug("ModelResponseIterator.aclose: error closing response: %s", e) def close(self) -> None: iterator = getattr(self, "response_iterator", self.streaming_response) @@ -3705,13 +3402,9 @@ class ModelResponseIterator: try: iterator.close() except Exception as e: # noqa: BLE001 - verbose_logger.debug( - "ModelResponseIterator.close: error closing iterator: %s", e - ) + verbose_logger.debug("ModelResponseIterator.close: error closing iterator: %s", e) if self.response is not None: try: self.response.close() except Exception as e: # noqa: BLE001 - verbose_logger.debug( - "ModelResponseIterator.close: error closing response: %s", e - ) + verbose_logger.debug("ModelResponseIterator.close: error closing response: %s", e) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 53f5fb464df..d989750a5f3 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -38,10 +38,7 @@ class GoogleBatchEmbeddings(VertexLLM): """Flatten nested input lists and detect file references.""" input_list = [input] if isinstance(input, str) else input flat_elements = [ - e - for item in input_list - for e in (item if isinstance(item, list) else [item]) - if isinstance(e, str) + e for item in input_list for e in (item if isinstance(item, list) else [item]) if isinstance(e, str) ] has_file_refs = any(_is_file_reference(e) for e in flat_elements) return flat_elements, has_file_refs @@ -73,9 +70,7 @@ class GoogleBatchEmbeddings(VertexLLM): response = sync_handler.get(url=url, headers=headers) if response.status_code != 200: - raise Exception( - f"Error fetching file {element}: {response.status_code} {response.text}" - ) + raise Exception(f"Error fetching file {element}: {response.status_code} {response.text}") file_data = response.json() resolved_files[element] = { @@ -112,9 +107,7 @@ class GoogleBatchEmbeddings(VertexLLM): response = await async_handler.get(url=url, headers=headers) if response.status_code != 200: - raise Exception( - f"Error fetching file {element}: {response.status_code} {response.text}" - ) + raise Exception(f"Error fetching file {element}: {response.status_code} {response.text}") file_data = response.json() resolved_files[element] = { @@ -218,9 +211,7 @@ class GoogleBatchEmbeddings(VertexLLM): if use_embed_content: resolved_files = {} if api_key: - resolved_files = self._resolve_file_references( - input=input, api_key=api_key, sync_handler=sync_handler - ) + resolved_files = self._resolve_file_references(input=input, api_key=api_key, sync_handler=sync_handler) request_data = transform_openai_input_gemini_embed_content( input=input, model=model, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 27ca1bd92a4..fd08fdf4c8c 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -140,9 +140,7 @@ def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool: for element in input: if isinstance(element, list): - if any( - _is_multimodal_element(sub) for sub in element if isinstance(sub, str) - ): + if any(_is_multimodal_element(sub) for sub in element if isinstance(sub, str)): return True elif isinstance(element, str) and _is_multimodal_element(element): return True @@ -242,13 +240,8 @@ def transform_openai_input_gemini_content( raise ValueError("Nested input list must not be empty") for sub in element: if not isinstance(sub, str): - raise ValueError( - f"Elements inside a nested input list must be strings, got {type(sub)}" - ) - parts = [ - _build_part_for_input(sub, resolved_files=resolved_files) - for sub in element - ] + raise ValueError(f"Elements inside a nested input list must be strings, got {type(sub)}") + parts = [_build_part_for_input(sub, resolved_files=resolved_files) for sub in element] else: parts = [_build_part_for_input(element, resolved_files=resolved_files)] request = EmbedContentRequest( @@ -322,11 +315,7 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> Optional[UsageMetadata] def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: if isinstance(input, str): return (input,) - return tuple( - sub - for element in input - for sub in (element if isinstance(element, list) else [element]) - ) + return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element])) def _is_image_element( @@ -354,17 +343,11 @@ def _count_input_images( input: GeminiEmbeddingInput, resolved_files: Mapping[str, Mapping[str, str]], ) -> int: - return sum( - 1 - for element in _flatten_input(input) - if _is_image_element(element, resolved_files) - ) + return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files)) def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: - return sum( - detail["tokenCount"] for detail in details if detail["modality"] == modality - ) + return sum(detail["tokenCount"] for detail in details if detail["modality"] == modality) def _fallback_usage(input: GeminiEmbeddingInput, model: str) -> Usage: @@ -388,28 +371,20 @@ def _usage_from_embed_content_response( prompt_tokens = usage_metadata.get("promptTokenCount", 0) total_tokens = usage_metadata.get("totalTokenCount") or prompt_tokens - details: Sequence[PromptTokensDetails] = ( - usage_metadata.get("promptTokensDetails") or () - ) + details: Sequence[PromptTokensDetails] = usage_metadata.get("promptTokensDetails") or () text_tokens = _tokens_for_modality(details, "TEXT") audio_tokens = _tokens_for_modality(details, "AUDIO") video_tokens = _tokens_for_modality(details, "VIDEO") image_count = _count_input_images(input, resolved_files) - video_length_seconds = ( - video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 - ) - audio_length_seconds = ( - audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0 - ) + video_length_seconds = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 + audio_length_seconds = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0 # generic_cost_per_token rewrites text_tokens to the full prompt minus # other modalities when both text_tokens and image_count are zero. For # video, that misallocates video tokens to text; a 1-token floor sidesteps # the rewrite and keeps billing on input_cost_per_video_per_second. - needs_video_text_floor = ( - video_length_seconds > 0 and text_tokens == 0 and image_count == 0 - ) + needs_video_text_floor = video_length_seconds > 0 and text_tokens == 0 and image_count == 0 resolved_text_tokens = 1 if needs_video_text_floor else text_tokens return Usage( @@ -447,9 +422,7 @@ def process_embed_content_response( EmbeddingResponse with single embedding """ if "embedding" not in response_json: - raise ValueError( - f"embedContent response missing 'embedding' field: {response_json}" - ) + raise ValueError(f"embedContent response missing 'embedding' field: {response_json}") embedding_data = response_json["embedding"] @@ -495,25 +468,17 @@ def process_response( text_elements: List[str] = [] for e in input_list: if isinstance(e, list): - text_elements.extend( - sub - for sub in e - if isinstance(sub, str) and not _is_multimodal_element(sub) - ) + text_elements.extend(sub for sub in e if isinstance(sub, str) and not _is_multimodal_element(sub)) elif isinstance(e, str) and not _is_multimodal_element(e): text_elements.append(e) if text_elements: - input_text = get_formatted_prompt( - data={"input": text_elements}, call_type="embedding" - ) + input_text = get_formatted_prompt(data={"input": text_elements}, call_type="embedding") prompt_tokens = token_counter(model=model, text=input_text) else: prompt_tokens = 0 else: input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") prompt_tokens = token_counter(model=model, text=input_text) - model_response.usage = Usage( - prompt_tokens=prompt_tokens, total_tokens=prompt_tokens - ) + model_response.usage = Usage(prompt_tokens=prompt_tokens, total_tokens=prompt_tokens) return model_response diff --git a/litellm/llms/vertex_ai/image_edit/cost_calculator.py b/litellm/llms/vertex_ai/image_edit/cost_calculator.py index b346622a336..6e951624081 100644 --- a/litellm/llms/vertex_ai/image_edit/cost_calculator.py +++ b/litellm/llms/vertex_ai/image_edit/cost_calculator.py @@ -26,9 +26,7 @@ def cost_calculator( output_cost_per_image: float = model_info.get("output_cost_per_image") or 0.0 if not isinstance(image_response, ImageResponse): - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") num_images = len(image_response.data or []) return output_cost_per_image * num_images diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index de7f234a861..a2020149ef2 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -48,11 +48,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): drop_params: bool, ) -> Dict[str, Any]: supported_params = self.get_supported_openai_params(model) - filtered_params = { - key: value - for key, value in image_edit_optional_params.items() - if key in supported_params - } + filtered_params = {key: value for key, value in image_edit_optional_params.items() if key in supported_params} mapped_params: Dict[str, Any] = {} @@ -109,14 +105,8 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed) # then fall back to environment variables and other sources - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_credentials = ( - self.safe_get_vertex_ai_credentials(litellm_params) - or self._resolve_vertex_credentials() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials() access_token, _ = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -145,19 +135,11 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_location are passed) # then fall back to environment variables and other sources - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_location = ( - self.safe_get_vertex_ai_location(litellm_params) - or self._resolve_vertex_location() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() if not vertex_project or not vertex_location: - raise ValueError( - "vertex_project and vertex_location are required for Vertex AI" - ) + raise ValueError("vertex_project and vertex_location are required for Vertex AI") base_url = get_vertex_base_url(vertex_location) @@ -192,9 +174,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): # Add image-specific configuration image_config: Dict[str, Any] = {} if "aspectRatio" in image_edit_optional_request_params: - image_config["aspect_ratio"] = image_edit_optional_request_params[ - "aspectRatio" - ] + image_config["aspect_ratio"] = image_edit_optional_request_params["aspectRatio"] if image_config: generation_config["image_config"] = image_config @@ -203,9 +183,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): payload: Any = json.dumps(request_body) empty_files = cast(RequestFiles, []) - return cast( - Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files) - ) + return cast(Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files)) def transform_image_edit_response( self, @@ -253,9 +231,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): } return aspect_ratio_map.get(size, "1:1") - def _prepare_inline_image_parts( - self, image: Union[FileTypes, List[FileTypes]] - ) -> List[Dict[str, Any]]: + def _prepare_inline_image_parts(self, image: Union[FileTypes, List[FileTypes]]) -> List[Dict[str, Any]]: images: List[FileTypes] if isinstance(image, list): images = image diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index 3eb039614fd..d9127a1929f 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -49,11 +49,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): drop_params: bool, ) -> Dict[str, Any]: supported_params = self.get_supported_openai_params(model) - filtered_params = { - key: value - for key, value in image_edit_optional_params.items() - if key in supported_params - } + filtered_params = {key: value for key, value in image_edit_optional_params.items() if key in supported_params} mapped_params: Dict[str, Any] = {} @@ -113,14 +109,8 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if _api_base is not None: return headers - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_credentials = ( - self.safe_get_vertex_ai_credentials(litellm_params) - or self._resolve_vertex_credentials() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials() access_token, _ = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -137,19 +127,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ Get the complete URL for Vertex AI Imagen predict API """ - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_location = ( - self.safe_get_vertex_ai_location(litellm_params) - or self._resolve_vertex_location() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() if not vertex_project or not vertex_location: - raise ValueError( - "vertex_project and vertex_location are required for Vertex AI" - ) + raise ValueError("vertex_project and vertex_location are required for Vertex AI") # Use the model name as provided, handling vertex_ai prefix model_name = model @@ -174,16 +156,10 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): ) -> Tuple[Dict[str, Any], Optional[RequestFiles]]: # Prepare reference images in the correct Imagen format if image is None: - raise ValueError( - "Vertex AI Imagen image edit requires at least one reference image." - ) - reference_images = self._prepare_reference_images( - image, image_edit_optional_request_params - ) + raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") + reference_images = self._prepare_reference_images(image, image_edit_optional_request_params) if not reference_images: - raise ValueError( - "Vertex AI Imagen image edit requires at least one reference image." - ) + raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") if prompt is None: raise ValueError("Vertex AI Imagen image edit requires a prompt.") @@ -215,9 +191,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): payload: Any = json.dumps(request_body) empty_files = cast(RequestFiles, []) - return cast( - Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files) - ) + return cast(Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files)) def transform_image_edit_response( self, @@ -313,9 +287,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): return reference_images - def _read_all_bytes( - self, image: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH - ) -> bytes: + def _read_all_bytes(self, image: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bytes: if depth > max_depth: raise ValueError( f"Max recursion depth {max_depth} reached while reading image bytes for Vertex AI Imagen image edit." @@ -324,9 +296,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if isinstance(image, (list, tuple)): for item in image: if item is not None: - return self._read_all_bytes( - item, depth=depth + 1, max_depth=max_depth - ) + return self._read_all_bytes(item, depth=depth + 1, max_depth=max_depth) raise ValueError("Unsupported image type for Vertex AI Imagen image edit.") if isinstance(image, dict): @@ -338,13 +308,9 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): return base64.b64decode(value) except Exception: continue - return self._read_all_bytes( - value, depth=depth + 1, max_depth=max_depth - ) + return self._read_all_bytes(value, depth=depth + 1, max_depth=max_depth) if "path" in image: - return self._read_all_bytes( - image["path"], depth=depth + 1, max_depth=max_depth - ) + return self._read_all_bytes(image["path"], depth=depth + 1, max_depth=max_depth) if isinstance(image, bytes): return image @@ -383,6 +349,4 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if isinstance(data, str): data = data.encode("utf-8") return data - raise ValueError( - f"Unsupported image type for Vertex AI Imagen image edit. Got type={type(image)}" - ) + raise ValueError(f"Unsupported image type for Vertex AI Imagen image edit. Got type={type(image)}") diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index e14cfe3be0b..d265352ca0a 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -131,9 +131,7 @@ class VertexImageGeneration(VertexLLM): should_use_v1beta1_features=False, mode="image_generation", ) - optional_params = optional_params or { - "sampleCount": 1 - } # default optional params + optional_params = optional_params or {"sampleCount": 1} # default optional params # Transform optional params to camelCase format optional_params = self.transform_optional_params(optional_params) @@ -165,9 +163,7 @@ class VertexImageGeneration(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") json_response = response.json() - return self.process_image_generation_response( - json_response, model_response, model - ) + return self.process_image_generation_response(json_response, model_response, model) async def aimage_generation( self, @@ -271,9 +267,7 @@ class VertexImageGeneration(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") json_response = response.json() - return self.process_image_generation_response( - json_response, model_response, model - ) + return self.process_image_generation_response(json_response, model_response, model) def is_image_generation_response(self, json_response: Dict[str, Any]) -> bool: if "predictions" in json_response: diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 103c7b2a28a..39503bd78dd 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -153,19 +153,11 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_location are passed) # then fall back to environment variables and other sources - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_location = ( - self.safe_get_vertex_ai_location(litellm_params) - or self._resolve_vertex_location() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() if not vertex_project or not vertex_location: - raise ValueError( - "vertex_project and vertex_location are required for Vertex AI" - ) + raise ValueError("vertex_project and vertex_location are required for Vertex AI") base_url = get_vertex_base_url(vertex_location) @@ -191,14 +183,8 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed) # then fall back to environment variables and other sources - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_credentials = ( - self.safe_get_vertex_ai_credentials(litellm_params) - or self._resolve_vertex_credentials() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials() access_token, _ = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -325,11 +311,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ImageObject( b64_json=inline_data["data"], url=None, - provider_specific_fields=( - {"thought_signature": thought_sig} - if thought_sig - else None - ), + provider_specific_fields=({"thought_signature": thought_sig} if thought_sig else None), ) ) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 05ebd685d91..2cd3df010d6 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -39,9 +39,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): BaseImageGenerationConfig.__init__(self) VertexLLM.__init__(self) - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: """ Imagen API supported parameters """ @@ -135,19 +133,11 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_location are passed) # then fall back to environment variables and other sources - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_location = ( - self.safe_get_vertex_ai_location(litellm_params) - or self._resolve_vertex_location() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() if not vertex_project or not vertex_location: - raise ValueError( - "vertex_project and vertex_location are required for Vertex AI" - ) + raise ValueError("vertex_project and vertex_location are required for Vertex AI") base_url = get_vertex_base_url(vertex_location) @@ -173,14 +163,8 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed) # then fall back to environment variables and other sources - vertex_project = ( - self.safe_get_vertex_ai_project(litellm_params) - or self._resolve_vertex_project() - ) - vertex_credentials = ( - self.safe_get_vertex_ai_credentials(litellm_params) - or self._resolve_vertex_credentials() - ) + vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() + vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials() access_token, _ = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index f1d121099f9..4bcfdee2d17 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -75,11 +75,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): if self._is_gcs_uri(input_str): return InstanceImage(gcsUri=input_str) else: - return InstanceImage( - bytesBase64Encoded=( - input_str.split(",")[1] if "," in input_str else input_str - ) - ) + return InstanceImage(bytesBase64Encoded=(input_str.split(",")[1] if "," in input_str else input_str)) def _create_video_instance(self, input_str: str) -> InstanceVideo: """Create an InstanceVideo from a GCS URI.""" @@ -108,9 +104,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): else: return Instance(text=input_element) - def _try_merge_text_with_media( - self, text_str: str, next_elem: Optional[str] - ) -> tuple[Instance, bool]: + def _try_merge_text_with_media(self, text_str: str, next_elem: Optional[str]) -> tuple[Instance, bool]: """ Try to merge a text element with a following media element into a single instance. @@ -133,9 +127,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): return instance_args, False - def process_openai_embedding_input( - self, _input: Union[list, str] - ) -> List[Instance]: + def process_openai_embedding_input(self, _input: Union[list, str]) -> List[Instance]: """ Process the input for multimodal embedding requests. @@ -160,9 +152,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): i += 1 else: # Current element is text - try to merge with next media element - instance, consumed_next = self._try_merge_text_with_media( - text_str=current, next_elem=next_elem - ) + instance, consumed_next = self._try_merge_text_with_media(text_str=current, next_elem=next_elem) processed_instances.append(instance) i += 2 if consumed_next else 1 elif isinstance(current, dict): @@ -187,9 +177,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): if "instances" in optional_params: request_data["instances"] = optional_params["instances"] elif isinstance(input, list): - vertex_instances: List[Instance] = self.process_openai_embedding_input( - _input=input - ) + vertex_instances: List[Instance] = self.process_openai_embedding_input(_input=input) request_data["instances"] = vertex_instances else: @@ -202,9 +190,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): request_data["instances"] = [vertex_request_instance] if "outputDimensionality" in optional_params: - request_data["parameters"] = { - "dimension": optional_params["outputDimensionality"] - } + request_data["parameters"] = {"dimension": optional_params["outputDimensionality"]} return cast(dict, request_data) @@ -231,9 +217,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): ) _predictions = _json_response["predictions"] vertex_predictions = MultimodalPredictions(predictions=_predictions) - model_response.data = self.transform_embedding_response_to_openai( - predictions=vertex_predictions - ) + model_response.data = self.transform_embedding_response_to_openai(predictions=vertex_predictions) model_response.model = model model_response.usage = self.calculate_usage( @@ -291,9 +275,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): prompt_tokens_details=prompt_tokens_details, ) - def transform_embedding_response_to_openai( - self, predictions: MultimodalPredictions - ) -> List[Embedding]: + def transform_embedding_response_to_openai(self, predictions: MultimodalPredictions) -> List[Embedding]: openai_embeddings: List[Embedding] = [] if "predictions" in predictions: for idx, _prediction in enumerate(predictions["predictions"]): @@ -322,9 +304,5 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): openai_embeddings.append(openai_embedding_object) return openai_embeddings - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, Headers] - ) -> BaseLLMException: - return VertexAIError( - status_code=status_code, message=error_message, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException: + return VertexAIError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py index a98311d04eb..68836a64027 100644 --- a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py +++ b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py @@ -66,12 +66,8 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - vertex_project = VertexBase.safe_get_vertex_ai_project( - litellm_params=litellm_params - ) - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( - litellm_params=litellm_params - ) + vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=litellm_params) # Get access token from Vertex credentials access_token, project_id = self.vertex_base.get_access_token( @@ -110,12 +106,8 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - vertex_project = VertexBase.safe_get_vertex_ai_project( - litellm_params=litellm_params - ) - vertex_location = VertexBase.safe_get_vertex_ai_location( - litellm_params=litellm_params - ) + vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) + vertex_location = VertexBase.safe_get_vertex_ai_location(litellm_params=litellm_params) if vertex_project is None: raise ValueError( @@ -159,9 +151,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): Returns: OCRRequestData with JSON data for the DeepSeek OCR endpoint """ - verbose_logger.debug( - "Vertex AI DeepSeek OCR transform_ocr_request (sync) called" - ) + verbose_logger.debug("Vertex AI DeepSeek OCR transform_ocr_request (sync) called") if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") @@ -176,9 +166,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): elif doc_type == "document_url": document_url = document.get("document_url", "") else: - raise ValueError( - f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'" - ) + raise ValueError(f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'") # Build DeepSeek OCR message content content_item = {} @@ -317,17 +305,11 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): "pages": [ { "index": 0, - "markdown": ( - content - if isinstance(content, str) - else json.dumps(content) - ), + "markdown": (content if isinstance(content, str) else json.dumps(content)), } ], "model": ocr_data.get("model", model), - "usage_info": ocr_data.get( - "usage_info", response_json.get("usage", {}) - ), + "usage_info": ocr_data.get("usage_info", response_json.get("usage", {})), } # Convert usage info if present @@ -352,11 +334,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): if not pages: # Create a default page if none exist - pages = [ - OCRPage( - index=0, markdown=content if isinstance(content, str) else "" - ) - ] + pages = [OCRPage(index=0, markdown=content if isinstance(content, str) else "")] return OCRResponse( pages=pages, diff --git a/litellm/llms/vertex_ai/ocr/transformation.py b/litellm/llms/vertex_ai/ocr/transformation.py index a725762b3c5..d67c5f2b089 100644 --- a/litellm/llms/vertex_ai/ocr/transformation.py +++ b/litellm/llms/vertex_ai/ocr/transformation.py @@ -62,12 +62,8 @@ class VertexAIOCRConfig(MistralOCRConfig): # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - vertex_project = VertexBase.safe_get_vertex_ai_project( - litellm_params=litellm_params - ) - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( - litellm_params=litellm_params - ) + vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=litellm_params) # Get access token from Vertex credentials access_token, project_id = self.vertex_base.get_access_token( @@ -109,12 +105,8 @@ class VertexAIOCRConfig(MistralOCRConfig): # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - vertex_project = VertexBase.safe_get_vertex_ai_project( - litellm_params=litellm_params - ) - vertex_location = VertexBase.safe_get_vertex_ai_location( - litellm_params=litellm_params - ) + vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) + vertex_location = VertexBase.safe_get_vertex_ai_location(litellm_params=litellm_params) if vertex_project is None: raise ValueError( @@ -148,17 +140,13 @@ class VertexAIOCRConfig(MistralOCRConfig): Returns: Base64 data URI string """ - verbose_logger.debug( - f"Vertex AI OCR: Converting URL to base64 data URI (sync): {url}" - ) + verbose_logger.debug(f"Vertex AI OCR: Converting URL to base64 data URI (sync): {url}") # Fetch and convert to base64 data URI # convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = convert_url_to_base64(url=url) - verbose_logger.debug( - f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})" - ) + verbose_logger.debug(f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})") return data_uri @@ -175,17 +163,13 @@ class VertexAIOCRConfig(MistralOCRConfig): Returns: Base64 data URI string """ - verbose_logger.debug( - f"Vertex AI OCR: Converting URL to base64 data URI (async): {url}" - ) + verbose_logger.debug(f"Vertex AI OCR: Converting URL to base64 data URI (async): {url}") # Fetch and convert to base64 data URI asynchronously # async_convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = await async_convert_url_to_base64(url=url) - verbose_logger.debug( - f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})" - ) + verbose_logger.debug(f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})") return data_uri @@ -226,18 +210,14 @@ class VertexAIOCRConfig(MistralOCRConfig): document_url = document.get("document_url", "") # If it's not already a data URI, convert it if document_url and not document_url.startswith("data:"): - verbose_logger.debug( - "Vertex AI OCR: Converting document URL to base64 data URI (sync)" - ) + verbose_logger.debug("Vertex AI OCR: Converting document URL to base64 data URI (sync)") data_uri = self._convert_url_to_data_uri_sync(url=document_url) transformed_document["document_url"] = data_uri elif doc_type == "image_url": image_url = document.get("image_url", "") # If it's not already a data URI, convert it if image_url and not image_url.startswith("data:"): - verbose_logger.debug( - "Vertex AI OCR: Converting image URL to base64 data URI (sync)" - ) + verbose_logger.debug("Vertex AI OCR: Converting image URL to base64 data URI (sync)") data_uri = self._convert_url_to_data_uri_sync(url=image_url) transformed_document["image_url"] = data_uri @@ -274,9 +254,7 @@ class VertexAIOCRConfig(MistralOCRConfig): Returns: OCRRequestData with JSON data """ - verbose_logger.debug( - f"Vertex AI OCR async_transform_ocr_request - model: {model}" - ) + verbose_logger.debug(f"Vertex AI OCR async_transform_ocr_request - model: {model}") if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") @@ -289,18 +267,14 @@ class VertexAIOCRConfig(MistralOCRConfig): document_url = document.get("document_url", "") # If it's not already a data URI, convert it if document_url and not document_url.startswith("data:"): - verbose_logger.debug( - "Vertex AI OCR: Converting document URL to base64 data URI (async)" - ) + verbose_logger.debug("Vertex AI OCR: Converting document URL to base64 data URI (async)") data_uri = await self._convert_url_to_data_uri_async(url=document_url) transformed_document["document_url"] = data_uri elif doc_type == "image_url": image_url = document.get("image_url", "") # If it's not already a data URI, convert it if image_url and not image_url.startswith("data:"): - verbose_logger.debug( - "Vertex AI OCR: Converting image URL to base64 data URI (async)" - ) + verbose_logger.debug("Vertex AI OCR: Converting image URL to base64 data URI (async)") data_uri = await self._convert_url_to_data_uri_async(url=image_url) transformed_document["image_url"] = data_uri diff --git a/litellm/llms/vertex_ai/rag_engine/ingestion.py b/litellm/llms/vertex_ai/rag_engine/ingestion.py index 2ec61667795..d9e0035aa99 100644 --- a/litellm/llms/vertex_ai/rag_engine/ingestion.py +++ b/litellm/llms/vertex_ai/rag_engine/ingestion.py @@ -79,20 +79,14 @@ class VertexAIRAGIngestion(BaseRAGIngestion): ) # GCP config - self.vertex_project = self.vector_store_config.get( - "vertex_project" - ) or get_secret_str("VERTEXAI_PROJECT") + self.vertex_project = self.vector_store_config.get("vertex_project") or get_secret_str("VERTEXAI_PROJECT") self.vertex_location = ( - self.vector_store_config.get("vertex_location") - or get_secret_str("VERTEXAI_LOCATION") - or "us-central1" + self.vector_store_config.get("vertex_location") or get_secret_str("VERTEXAI_LOCATION") or "us-central1" ) self.vertex_credentials = self.vector_store_config.get("vertex_credentials") # GCS bucket for file uploads - self.gcs_bucket = self.vector_store_config.get("gcs_bucket") or os.environ.get( - "GCS_BUCKET_NAME" - ) + self.gcs_bucket = self.vector_store_config.get("gcs_bucket") or os.environ.get("GCS_BUCKET_NAME") if not self.gcs_bucket: raise ValueError( "gcs_bucket is required for Vertex AI RAG ingestion. " @@ -101,9 +95,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): # Import settings self.wait_for_import = self.vector_store_config.get("wait_for_import", True) - self.import_timeout = _get_int( - self.vector_store_config.get("import_timeout"), 600 - ) + self.import_timeout = _get_int(self.vector_store_config.get("import_timeout"), 600) # Validate required config if not self.vertex_project: @@ -141,8 +133,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): file_tuple = (filename, file_content, content_type) verbose_logger.debug( - f"Uploading file to GCS via litellm.files.acreate_file: {filename} " - f"(bucket: {self.gcs_bucket})" + f"Uploading file to GCS via litellm.files.acreate_file: {filename} (bucket: {self.gcs_bucket})" ) # Upload to GCS using LiteLLM's file upload @@ -204,9 +195,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): transformation_config=transformation_config, timeout=self.import_timeout, ) - verbose_logger.info( - f"Import complete: {response.imported_rag_files_count} files imported" - ) + verbose_logger.info(f"Import complete: {response.imported_rag_files_count} files imported") else: # Async import - don't wait _ = rag.import_files_async( @@ -290,9 +279,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion): Tuple of (corpus_id, gcs_uri) """ if not file_content or not filename: - verbose_logger.warning( - "No file content or filename provided for Vertex AI ingestion" - ) + verbose_logger.warning("No file content or filename provided for Vertex AI ingestion") return _get_str_or_none(self.corpus_id), None # Step 1: Upload file to GCS diff --git a/litellm/llms/vertex_ai/rag_engine/transformation.py b/litellm/llms/vertex_ai/rag_engine/transformation.py index ed5154bbdff..4aa2fcb49be 100644 --- a/litellm/llms/vertex_ai/rag_engine/transformation.py +++ b/litellm/llms/vertex_ai/rag_engine/transformation.py @@ -39,7 +39,9 @@ class VertexAIRAGTransformation(VertexBase): Vertex AI RAG Engine primarily uses gRPC-based SDK. """ base_url = get_vertex_base_url(vertex_location) - return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{corpus_id}:importRagFiles" + return ( + f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{corpus_id}:importRagFiles" + ) def get_retrieve_contexts_url( self, @@ -89,8 +91,7 @@ class VertexAIRAGTransformation(VertexBase): # Log if separators are provided (not supported by Vertex AI) if chunking_strategy.get("separators"): verbose_logger.warning( - "Vertex AI RAG Engine does not support custom separators. " - "The 'separators' parameter will be ignored." + "Vertex AI RAG Engine does not support custom separators. The 'separators' parameter will be ignored." ) return { @@ -115,9 +116,7 @@ class VertexAIRAGTransformation(VertexBase): Returns: Request payload dict for importRagFiles API """ - transformation_config = self.transform_chunking_strategy_to_vertex_format( - chunking_strategy - ) + transformation_config = self.transform_chunking_strategy_to_vertex_format(chunking_strategy) return { "import_rag_files_config": { @@ -136,9 +135,7 @@ class VertexAIRAGTransformation(VertexBase): Uses the base class method to get credentials. """ - credentials = self.get_vertex_ai_credentials( - {"vertex_credentials": vertex_credentials} - ) + credentials = self.get_vertex_ai_credentials({"vertex_credentials": vertex_credentials}) project = vertex_project or self.get_vertex_ai_project({}) access_token, _ = self._ensure_access_token( diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index 9f339d3dc29..beb8bc0be6f 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -118,11 +118,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): from litellm.types.llms.vertex_ai import GeminiResponseModalities response_modalities: list[GeminiResponseModalities] = ["AUDIO"] - full_model_path = ( - f"projects/{self._project}" - f"/locations/{self._location}" - f"/publishers/google/models/{model}" - ) + full_model_path = f"projects/{self._project}/locations/{self._location}/publishers/google/models/{model}" setup_config: BidiGenerateContentSetup = { "model": full_model_path, "generationConfig": {"responseModalities": response_modalities}, @@ -146,11 +142,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): def _vertex_model_path(self, model: str) -> str: """Return the fully-qualified Vertex AI model resource path.""" - return ( - f"projects/{self._project}" - f"/locations/{self._location}" - f"/publishers/google/models/{model}" - ) + return f"projects/{self._project}/locations/{self._location}/publishers/google/models/{model}" def _build_vertex_ai_setup_config(self, model: str, session_params: dict) -> dict: """Build Vertex AI setup configuration with proper model path and defaults.""" @@ -161,9 +153,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): # settings would be silently dropped because ``map_openai_params`` only # recognises the flat OpenAI-beta key names. session_params = self._normalize_session_payload_for_mapping(session_params) - setup_config = self.map_openai_params( - optional_params={}, non_default_params=session_params - ) + setup_config = self.map_openai_params(optional_params={}, non_default_params=session_params) # Use full Vertex AI model path setup_config["model"] = self._vertex_model_path(model) @@ -184,13 +174,10 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): # that need that behaviour must accept that VAD is off. client_turn_detection = session_params.get("turn_detection") client_disabled_auto_response = ( - isinstance(client_turn_detection, dict) - and client_turn_detection.get("create_response") is False + isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False ) realtime_input_config = setup_config.setdefault("realtimeInputConfig", {}) - automatic_detection = realtime_input_config.setdefault( - "automaticActivityDetection", {} - ) + automatic_detection = realtime_input_config.setdefault("automaticActivityDetection", {}) if not client_disabled_auto_response: automatic_detection["disabled"] = False automatic_detection.setdefault("silenceDurationMs", 800) @@ -220,14 +207,10 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): if msg_type == "session.update": if session_configuration_request is None: - setup_config = self._build_vertex_ai_setup_config( - model, json_message.get("session") or {} - ) + setup_config = self._build_vertex_ai_setup_config(model, json_message.get("session") or {}) gemini_setup_msg = json.dumps({"setup": setup_config}) - verbose_logger.debug( - "Vertex AI Realtime: Sending initial setup with tools to backend" - ) + verbose_logger.debug("Vertex AI Realtime: Sending initial setup with tools to backend") return [gemini_setup_msg] # A follow-up session.update can't be forwarded as a second setup @@ -235,13 +218,8 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): # silencing the audio-transcription guardrail's create_response # disable, surface a warning so operators know the model will # auto-respond before the guardrail can gate it on Vertex AI. - client_turn_detection = GeminiRealtimeConfig._extract_turn_detection( - json_message.get("session") or {} - ) - if ( - isinstance(client_turn_detection, dict) - and client_turn_detection.get("create_response") is False - ): + client_turn_detection = GeminiRealtimeConfig._extract_turn_detection(json_message.get("session") or {}) + if isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False: verbose_logger.warning( "Vertex AI Realtime: Dropping subsequent session.update " "(turn_detection.create_response=False) — Vertex Live " @@ -250,11 +228,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): "Vertex AI in non-deferred mode." ) else: - verbose_logger.debug( - "Vertex AI Realtime: Ignoring session.update (setup already sent)" - ) + verbose_logger.debug("Vertex AI Realtime: Ignoring session.update (setup already sent)") return [] - return super().transform_realtime_request( - message, model, session_configuration_request - ) + return super().transform_realtime_request(message, model, session_configuration_request) diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 382a8498d40..b9680af20cc 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -59,11 +59,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): ) # Fallback to environment or litellm config - project_id = ( - vertex_project - or get_secret_str("VERTEXAI_PROJECT") - or litellm.vertex_project - ) + project_id = vertex_project or get_secret_str("VERTEXAI_PROJECT") or litellm.vertex_project if not project_id: raise ValueError( @@ -207,19 +203,13 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): rerank_results = [] for result in results: rerank_results.append( - RerankResponseResult( - index=result["index"], relevance_score=result["relevance_score"] - ) + RerankResponseResult(index=result["index"], relevance_score=result["relevance_score"]) ) # Create meta object - meta = RerankResponseMeta( - billed_units=RerankBilledUnits(search_units=len(records)) - ) + meta = RerankResponseMeta(billed_units=RerankBilledUnits(search_units=len(records))) - return RerankResponse( - id=f"vertex_ai_rerank_{model}", results=rerank_results, meta=meta - ) + return RerankResponse(id=f"vertex_ai_rerank_{model}", results=rerank_results, meta=meta) def get_supported_cohere_rerank_params(self, model: str) -> list: return [ diff --git a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py index b835ad7d8fa..e27df956c9d 100644 --- a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py +++ b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py @@ -150,9 +150,7 @@ class VertexTextToSpeechAPI(VertexLLM): json=request, # type: ignore ) if response.status_code != 200: - raise Exception( - f"Request failed with status code {response.status_code}, {response.text}" - ) + raise Exception(f"Request failed with status code {response.status_code}, {response.text}") ############ Process the response ############ _json_response = response.json() @@ -180,9 +178,7 @@ class VertexTextToSpeechAPI(VertexLLM): ) -> HttpxBinaryResponseContent: import base64 - async_handler = get_async_httpx_client( - llm_provider=litellm.LlmProviders.VERTEX_AI - ) + async_handler = get_async_httpx_client(llm_provider=litellm.LlmProviders.VERTEX_AI) response = await async_handler.post( url=url, @@ -191,9 +187,7 @@ class VertexTextToSpeechAPI(VertexLLM): ) if response.status_code != 200: - raise Exception( - f"Request did not return a 200 status code: {response.status_code}, {response.text}" - ) + raise Exception(f"Request did not return a 200 status code: {response.status_code}, {response.text}") _json_response = response.json() @@ -213,9 +207,7 @@ class VertexTextToSpeechAPI(VertexLLM): return http_binary_response -def validate_vertex_input( - input_data: VertexInput, kwargs: dict, optional_params: dict -) -> None: +def validate_vertex_input(input_data: VertexInput, kwargs: dict, optional_params: dict) -> None: # Remove None values if input_data.get("text") is None: input_data.pop("text", None) diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index be7bcfcadd7..a003409f7a6 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -330,9 +330,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): if not input_data: raise ValueError("Either 'text' or 'ssml' must be provided.") if "text" in input_data and "ssml" in input_data: - raise ValueError( - "Only one of 'text' or 'ssml' should be provided, not both." - ) + raise ValueError("Only one of 'text' or 'ssml' should be provided, not both.") return input_data @@ -357,9 +355,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): TextToSpeechRequestData: Contains dict_body and headers """ # Get Vertex AI credentials from litellm_params - vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = litellm_params.get( - "vertex_credentials" - ) + vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = litellm_params.get("vertex_credentials") vertex_project: Optional[str] = litellm_params.get("vertex_project") ####### Authenticate with Vertex AI ######## @@ -393,9 +389,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): # Check for voice dict stored in: # 1. litellm_params by dispatch method # 2. optional_params by map_openai_params - voice_dict = litellm_params.get("vertex_voice_dict") or optional_params.get( - "vertex_voice_dict" - ) + voice_dict = litellm_params.get("vertex_voice_dict") or optional_params.get("vertex_voice_dict") if voice_dict is not None and isinstance(voice_dict, dict): vertex_voice = VertexTextToSpeechVoice(**voice_dict) elif voice is not None and isinstance(voice, str): @@ -417,16 +411,12 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): ) # Build audio configuration - audio_encoding = optional_params.get( - "audioEncoding", self.DEFAULT_AUDIO_ENCODING - ) + audio_encoding = optional_params.get("audioEncoding", self.DEFAULT_AUDIO_ENCODING) speaking_rate = optional_params.get("speakingRate", self.DEFAULT_SPEAKING_RATE) # Check for full audioConfig in optional_params if "audioConfig" in optional_params: - vertex_audio_config = VertexTextToSpeechAudioConfig( - **optional_params["audioConfig"] - ) + vertex_audio_config = VertexTextToSpeechAudioConfig(**optional_params["audioConfig"]) else: vertex_audio_config = VertexTextToSpeechAudioConfig( audioEncoding=audio_encoding, diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index d31e1f6c8f2..47a81fc07bf 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -35,9 +35,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): def __init__(self): super().__init__() - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: # Get credentials and project info vertex_credentials = self.get_vertex_ai_credentials(dict(litellm_params)) vertex_project = self.get_vertex_ai_project(dict(litellm_params)) @@ -62,9 +60,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): "write": [("POST", "/ragCorpora")], } - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate and set up authentication for Vertex AI RAG API """ @@ -183,9 +179,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Generate file_id from source URI or use display name as fallback file_id = source_uri if source_uri else source_display_name - filename = ( - source_display_name if source_display_name else "Unknown Document" - ) + filename = source_display_name if source_display_name else "Unknown Document" # Build attributes with available metadata attributes = {} @@ -233,9 +227,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Build the request body for Vertex AI RAG Corpus creation request_body: Dict[str, Any] = { - "display_name": vector_store_create_optional_params.get( - "name", "litellm-vector-store" - ), + "display_name": vector_store_create_optional_params.get("name", "litellm-vector-store"), "description": "Vector store created via LiteLLM", } @@ -246,9 +238,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): return url, request_body - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: """ Transform Vertex AI RAG Corpus creation response to standard vector store response """ @@ -257,9 +247,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Extract the corpus ID from the response name corpus_name = response_json.get("name", "") - corpus_id = ( - corpus_name.split("/")[-1] if "/" in corpus_name else corpus_name - ) + corpus_id = corpus_name.split("/")[-1] if "/" in corpus_name else corpus_name # Handle createTime conversion create_time = response_json.get("createTime", 0) diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 46dedb3d0a4..958839d4a48 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -45,13 +45,9 @@ VERTEX_SEARCH_TARGET_SELECTING_FIELDS = frozenset( # via extra_body, derived from the TypedDicts so the type is the source of truth. # Engine/app mode is a superset (adds dataStoreSpecs, numResultsPerDataStore), # since an app fans out across multiple member data stores. -VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS = frozenset( - VertexSearchDataStoreExtraBody.__annotations__ -) +VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS = frozenset(VertexSearchDataStoreExtraBody.__annotations__) -VERTEX_SEARCH_ENGINE_EXTRA_BODY_FIELDS = frozenset( - VertexSearchEngineExtraBody.__annotations__ -) +VERTEX_SEARCH_ENGINE_EXTRA_BODY_FIELDS = frozenset(VertexSearchEngineExtraBody.__annotations__) class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): @@ -79,9 +75,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): return VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS @classmethod - def _filter_extra_body( - cls, extra_body: Dict[str, Any], is_engine: bool = False - ) -> Dict[str, Any]: + def _filter_extra_body(cls, extra_body: Dict[str, Any], is_engine: bool = False) -> Dict[str, Any]: """ Validate ``extra_body`` against the supported-field allowlist for the active serving config (engine/app vs data store). @@ -94,9 +88,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): data-store mode where they are meaningless. """ supported = cls.get_supported_extra_body_fields(is_engine=is_engine) - filtered = { - key: value for key, value in extra_body.items() if value is not None - } + filtered = {key: value for key, value in extra_body.items() if value is not None} target_selecting = set(filtered) & VERTEX_SEARCH_TARGET_SELECTING_FIELDS if target_selecting: @@ -124,9 +116,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): return filtered - def get_auth_credentials( - self, litellm_params: dict - ) -> BaseVectorStoreAuthCredentials: + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: # Get credentials and project info vertex_credentials = self.get_vertex_ai_credentials(dict(litellm_params)) vertex_project = self.get_vertex_ai_project(dict(litellm_params)) @@ -151,9 +141,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): "write": [], } - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate and set up authentication for Vertex AI RAG API """ @@ -181,12 +169,8 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vertex_location = self.get_vertex_ai_location(litellm_params) vertex_project = self.get_vertex_ai_project(litellm_params) - collection_id = ( - litellm_params.get("vertex_collection_id") or "default_collection" - ) - encoded_collection_id = encode_url_path_segment( - collection_id, field_name="vertex_collection_id" - ) + collection_id = litellm_params.get("vertex_collection_id") or "default_collection" + encoded_collection_id = encode_url_path_segment(collection_id, field_name="vertex_collection_id") base = ( f"https://discoveryengine.googleapis.com/v1/" f"projects/{vertex_project}/locations/{vertex_location}/" @@ -195,19 +179,13 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): engine_id = litellm_params.get("vertex_engine_id") if engine_id: - encoded_engine_id = encode_url_path_segment( - engine_id, field_name="vertex_engine_id" - ) + encoded_engine_id = encode_url_path_segment(engine_id, field_name="vertex_engine_id") return f"{base}/engines/{encoded_engine_id}/servingConfigs/default_serving_config" datastore_id = litellm_params.get("vector_store_id") if not datastore_id: - raise ValueError( - "vector_store_id is required when vertex_engine_id is not set" - ) - encoded_datastore_id = encode_url_path_segment( - datastore_id, field_name="vector_store_id" - ) + raise ValueError("vector_store_id is required when vertex_engine_id is not set") + encoded_datastore_id = encode_url_path_segment(datastore_id, field_name="vector_store_id") return f"{base}/dataStores/{encoded_datastore_id}/servingConfigs/default_config" def transform_search_vector_store_request( @@ -249,13 +227,9 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): if max_num_results is not None: request_body["pageSize"] = max_num_results if isinstance(extra_body, dict): - request_body.update( - self._filter_extra_body(extra_body, is_engine=is_engine) - ) + request_body.update(self._filter_extra_body(extra_body, is_engine=is_engine)) - litellm_logging_obj.model_call_details["query"] = request_body.get( - "query", query - ) + litellm_logging_obj.model_call_details["query"] = request_body.get("query", query) return url, request_body @@ -299,10 +273,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): if snippets: # Combine all snippets into one text - text_parts = [ - snippet.get("snippet", snippet.get("htmlSnippet", "")) - for snippet in snippets - ] + text_parts = [snippet.get("snippet", snippet.get("htmlSnippet", "")) for snippet in snippets] text_content = " ".join(text_parts) # If no snippets, use title as fallback @@ -347,9 +318,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Note: Search API doesn't provide explicit scores in the response # You can use the position/rank as an implicit score - score = 1.0 / ( - float(search_results.__len__() + 1) - ) # Decreasing score based on position + score = 1.0 / (float(search_results.__len__() + 1)) # Decreasing score based on position result_obj = VectorStoreSearchResult( score=score, @@ -380,9 +349,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): ) -> Tuple[str, Dict]: raise NotImplementedError - def transform_create_vector_store_response( - self, response: httpx.Response - ) -> VectorStoreCreateResponse: + def transform_create_vector_store_response(self, response: httpx.Response) -> VectorStoreCreateResponse: raise NotImplementedError def calculate_vector_store_cost( diff --git a/litellm/llms/vertex_ai/vertex_ai_aws_wif.py b/litellm/llms/vertex_ai/vertex_ai_aws_wif.py index a03a4e37a21..da95ac72c2f 100644 --- a/litellm/llms/vertex_ai/vertex_ai_aws_wif.py +++ b/litellm/llms/vertex_ai/vertex_ai_aws_wif.py @@ -12,8 +12,7 @@ AwsSecurityCredentialsSupplier for google-auth. from typing import Dict GOOGLE_IMPORT_ERROR_MESSAGE = ( - "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' " - "or pip install google-cloud-aiplatform" + "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform" ) # AWS params recognized in WIF credential JSON for explicit auth. @@ -107,9 +106,7 @@ class VertexAIAwsWifAuth: token_url=json_obj.get("token_url"), credential_source=None, # Not using metadata endpoints aws_security_credentials_supplier=supplier, - service_account_impersonation_url=json_obj.get( - "service_account_impersonation_url" - ), + service_account_impersonation_url=json_obj.get("service_account_impersonation_url"), ) # Forward universe_domain if present (defaults to googleapis.com) if "universe_domain" in json_obj: diff --git a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py index c134dee7ad4..33606013d5c 100644 --- a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py +++ b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py @@ -17,13 +17,9 @@ class VertexAIError(Exception): def __init__(self, status_code, message): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url=" https://cloud.google.com/vertex-ai/" - ) + self.request = httpx.Request(method="POST", url=" https://cloud.google.com/vertex-ai/") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class TextStreamer: @@ -58,9 +54,7 @@ class TextStreamer: raise StopAsyncIteration # once we run out of data to stream, we raise this error -def _get_client_cache_key( - model: str, vertex_project: Optional[str], vertex_location: Optional[str] -): +def _get_client_cache_key(model: str, vertex_project: Optional[str], vertex_location: Optional[str]): _cache_key = f"{model}-{vertex_project}-{vertex_location}" return _cache_key @@ -108,9 +102,7 @@ def completion( message="vertexai import failed please run `pip install google-cloud-aiplatform`. This is required for the 'vertex_ai/' route on LiteLLM", ) - if not ( - hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models") - ): + if not (hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")): raise VertexAIError( status_code=400, message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", @@ -128,13 +120,9 @@ def completion( from vertexai.preview.language_models import ChatModel, CodeChatModel ## Load credentials with the correct quota project ref: https://github.com/googleapis/python-aiplatform/issues/2557#issuecomment-1709284744 - print_verbose( - f"VERTEX AI: vertex_project={vertex_project}; vertex_location={vertex_location}" - ) + print_verbose(f"VERTEX AI: vertex_project={vertex_project}; vertex_location={vertex_location}") - _cache_key = _get_client_cache_key( - model=model, vertex_project=vertex_project, vertex_location=vertex_location - ) + _cache_key = _get_client_cache_key(model=model, vertex_project=vertex_project, vertex_location=vertex_location) _vertex_llm_model_object = _get_client_from_cache(client_cache_key=_cache_key) # Load credentials - needed for both vertexai.init() and PredictionServiceClient @@ -176,18 +164,12 @@ def completion( raise ValueError("safety_settings must be a list") if len(safety_settings) > 0 and not isinstance(safety_settings[0], dict): raise ValueError("safety_settings must be a list of dicts") - safety_settings = [ - gapic_content_types.SafetySetting(x) for x in safety_settings - ] + safety_settings = [gapic_content_types.SafetySetting(x) for x in safety_settings] # vertexai does not use an API key, it looks for credentials.json in the environment prompt = " ".join( - [ - message.get("content") - for message in messages - if isinstance(message.get("content", None), str) - ] + [message.get("content") for message in messages if isinstance(message.get("content", None), str)] ) mode = "" @@ -195,14 +177,9 @@ def completion( request_str = "" response_obj = None instances = None - client_options = { - "api_endpoint": f"{vertex_location}-aiplatform.googleapis.com" - } + client_options = {"api_endpoint": f"{vertex_location}-aiplatform.googleapis.com"} fake_stream = False - if ( - model in litellm.vertex_language_models - or model in litellm.vertex_vision_models - ): + if model in litellm.vertex_language_models or model in litellm.vertex_vision_models: llm_model: Any = _vertex_llm_model_object or GenerativeModel(model) mode = "vision" request_str += f"llm_model = GenerativeModel({model})\n" @@ -211,15 +188,11 @@ def completion( mode = "chat" request_str += f"llm_model = ChatModel.from_pretrained({model})\n" elif model in litellm.vertex_text_models: - llm_model = _vertex_llm_model_object or TextGenerationModel.from_pretrained( - model - ) + llm_model = _vertex_llm_model_object or TextGenerationModel.from_pretrained(model) mode = "text" request_str += f"llm_model = TextGenerationModel.from_pretrained({model})\n" elif model in litellm.vertex_code_text_models: - llm_model = _vertex_llm_model_object or CodeGenerationModel.from_pretrained( - model - ) + llm_model = _vertex_llm_model_object or CodeGenerationModel.from_pretrained(model) mode = "text" request_str += f"llm_model = CodeGenerationModel.from_pretrained({model})\n" fake_stream = True @@ -280,9 +253,7 @@ def completion( completion_response = None - stream = optional_params.pop( - "stream", None - ) # See note above on handling streaming for vertex ai + stream = optional_params.pop("stream", None) # See note above on handling streaming for vertex ai if mode == "chat": chat = llm_model.start_chat() request_str += "chat = llm_model.start_chat()\n" @@ -291,13 +262,9 @@ def completion( # NOTE: VertexAI does not accept stream=True as a param and raises an error, # we handle this by removing 'stream' from optional params and sending the request # after we get the response we add optional_params["stream"] = True, since main.py needs to know it's a streaming response to then transform it for the OpenAI format - optional_params.pop( - "stream", None - ) # vertex ai raises an error when passing stream in optional params + optional_params.pop("stream", None) # vertex ai raises an error when passing stream in optional params - request_str += ( - f"chat.send_message_streaming({prompt}, **{optional_params})\n" - ) + request_str += f"chat.send_message_streaming({prompt}, **{optional_params})\n" ## LOGGING logging_obj.pre_call( input=prompt, @@ -325,9 +292,7 @@ def completion( completion_response = chat.send_message(prompt, **optional_params).text elif mode == "text": if fake_stream is not True and stream is True: - request_str += ( - f"llm_model.predict_streaming({prompt}, **{optional_params})\n" - ) + request_str += f"llm_model.predict_streaming({prompt}, **{optional_params})\n" ## LOGGING logging_obj.pre_call( input=prompt, @@ -358,9 +323,7 @@ def completion( """ if vertex_project is None or vertex_location is None: - raise ValueError( - "Vertex project and location are required for custom endpoint" - ) + raise ValueError("Vertex project and location are required for custom endpoint") ## LOGGING logging_obj.pre_call( @@ -376,21 +339,12 @@ def completion( credentials=creds, # type: ignore[arg-type] ) request_str += f"llm_model = aiplatform.gapic.PredictionServiceClient(client_options={client_options}, credentials=...)\n" - endpoint_path = llm_model.endpoint_path( - project=vertex_project, location=vertex_location, endpoint=model - ) - request_str += ( - f"llm_model.predict(endpoint={endpoint_path}, instances={instances})\n" - ) - response = llm_model.predict( - endpoint=endpoint_path, instances=instances - ).predictions + endpoint_path = llm_model.endpoint_path(project=vertex_project, location=vertex_location, endpoint=model) + request_str += f"llm_model.predict(endpoint={endpoint_path}, instances={instances})\n" + response = llm_model.predict(endpoint=endpoint_path, instances=instances).predictions completion_response = response[0] - if ( - isinstance(completion_response, str) - and "\nOutput:\n" in completion_response - ): + if isinstance(completion_response, str) and "\nOutput:\n" in completion_response: completion_response = completion_response.split("\nOutput:\n", 1)[1] if stream is True: response = TextStreamer(completion_response) @@ -416,19 +370,14 @@ def completion( response = llm_model.predict(instances=instances).predictions completion_response = response[0] - if ( - isinstance(completion_response, str) - and "\nOutput:\n" in completion_response - ): + if isinstance(completion_response, str) and "\nOutput:\n" in completion_response: completion_response = completion_response.split("\nOutput:\n", 1)[1] if stream is True: response = TextStreamer(completion_response) return response ## LOGGING - logging_obj.post_call( - input=prompt, api_key=None, original_response=completion_response - ) + logging_obj.post_call(input=prompt, api_key=None, original_response=completion_response) ## RESPONSE OBJECT if isinstance(completion_response, litellm.Message): @@ -456,16 +405,10 @@ def completion( response_obj.usage_metadata, "prompt_token_count" ): prompt_tokens = response_obj.usage_metadata.prompt_token_count - completion_tokens = ( - response_obj.usage_metadata.candidates_token_count - ) + completion_tokens = response_obj.usage_metadata.candidates_token_count else: prompt_tokens = len(encoding.encode(prompt)) - completion_tokens = len( - encoding.encode( - model_response["choices"][0]["message"].get("content", "") - ) - ) + completion_tokens = len(encoding.encode(model_response["choices"][0]["message"].get("content", ""))) usage = Usage( prompt_tokens=prompt_tokens, @@ -480,9 +423,7 @@ def completion( except Exception as e: if isinstance(e, VertexAIError): raise e - raise litellm.APIConnectionError( - message=str(e), llm_provider="vertex_ai", model=model - ) + raise litellm.APIConnectionError(message=str(e), llm_provider="vertex_ai", model=model) async def async_completion( @@ -545,9 +486,7 @@ async def async_completion( from google.cloud import aiplatform # type: ignore if vertex_project is None or vertex_location is None: - raise ValueError( - "Vertex project and location are required for custom endpoint" - ) + raise ValueError("Vertex project and location are required for custom endpoint") ## LOGGING logging_obj.pre_call( @@ -564,22 +503,15 @@ async def async_completion( credentials=vertex_credentials, ) request_str += f"llm_model = aiplatform.gapic.PredictionServiceAsyncClient(client_options={client_options}, credentials=...)\n" - endpoint_path = llm_model.endpoint_path( - project=vertex_project, location=vertex_location, endpoint=model - ) - request_str += ( - f"llm_model.predict(endpoint={endpoint_path}, instances={instances})\n" - ) + endpoint_path = llm_model.endpoint_path(project=vertex_project, location=vertex_location, endpoint=model) + request_str += f"llm_model.predict(endpoint={endpoint_path}, instances={instances})\n" response_obj = await llm_model.predict( endpoint=endpoint_path, instances=instances, ) response = response_obj.predictions completion_response = response[0] - if ( - isinstance(completion_response, str) - and "\nOutput:\n" in completion_response - ): + if isinstance(completion_response, str) and "\nOutput:\n" in completion_response: completion_response = completion_response.split("\nOutput:\n", 1)[1] elif mode == "private": @@ -590,16 +522,11 @@ async def async_completion( response = response_obj.predictions completion_response = response[0] - if ( - isinstance(completion_response, str) - and "\nOutput:\n" in completion_response - ): + if isinstance(completion_response, str) and "\nOutput:\n" in completion_response: completion_response = completion_response.split("\nOutput:\n", 1)[1] ## LOGGING - logging_obj.post_call( - input=prompt, api_key=None, original_response=completion_response - ) + logging_obj.post_call(input=prompt, api_key=None, original_response=completion_response) ## RESPONSE OBJECT if isinstance(completion_response, litellm.Message): @@ -625,18 +552,13 @@ async def async_completion( # this block attempts to get usage from response_obj if it exists, if not it uses the litellm token counter prompt_tokens, completion_tokens, _ = 0, 0, 0 if response_obj is not None and ( - hasattr(response_obj, "usage_metadata") - and hasattr(response_obj.usage_metadata, "prompt_token_count") + hasattr(response_obj, "usage_metadata") and hasattr(response_obj.usage_metadata, "prompt_token_count") ): prompt_tokens = response_obj.usage_metadata.prompt_token_count completion_tokens = response_obj.usage_metadata.candidates_token_count else: prompt_tokens = len(encoding.encode(prompt)) - completion_tokens = len( - encoding.encode( - model_response["choices"][0]["message"].get("content", "") - ) - ) + completion_tokens = len(encoding.encode(model_response["choices"][0]["message"].get("content", ""))) # set usage usage = Usage( @@ -675,12 +597,8 @@ async def async_streaming( response: Any = None if mode == "chat": chat = llm_model.start_chat() - optional_params.pop( - "stream", None - ) # vertex ai raises an error when passing stream in optional params - request_str += ( - f"chat.send_message_streaming_async({prompt}, **{optional_params})\n" - ) + optional_params.pop("stream", None) # vertex ai raises an error when passing stream in optional params + request_str += f"chat.send_message_streaming_async({prompt}, **{optional_params})\n" ## LOGGING logging_obj.pre_call( input=prompt, @@ -693,12 +611,8 @@ async def async_streaming( response = chat.send_message_streaming_async(prompt, **optional_params) elif mode == "text": - optional_params.pop( - "stream", None - ) # See note above on handling streaming for vertex ai - request_str += ( - f"llm_model.predict_streaming_async({prompt}, **{optional_params})\n" - ) + optional_params.pop("stream", None) # See note above on handling streaming for vertex ai + request_str += f"llm_model.predict_streaming_async({prompt}, **{optional_params})\n" ## LOGGING logging_obj.pre_call( input=prompt, @@ -713,9 +627,7 @@ async def async_streaming( from google.cloud import aiplatform # type: ignore if vertex_project is None or vertex_location is None: - raise ValueError( - "Vertex project and location are required for custom endpoint" - ) + raise ValueError("Vertex project and location are required for custom endpoint") stream = optional_params.pop("stream", None) @@ -733,12 +645,8 @@ async def async_streaming( credentials=vertex_credentials, ) request_str += f"llm_model = aiplatform.gapic.PredictionServiceAsyncClient(client_options={client_options}, credentials=...)\n" - endpoint_path = llm_model.endpoint_path( - project=vertex_project, location=vertex_location, endpoint=model - ) - request_str += ( - f"client.predict(endpoint={endpoint_path}, instances={instances})\n" - ) + endpoint_path = llm_model.endpoint_path(project=vertex_project, location=vertex_location, endpoint=model) + request_str += f"client.predict(endpoint={endpoint_path}, instances={instances})\n" response_obj = await llm_model.predict( endpoint=endpoint_path, instances=instances, @@ -746,10 +654,7 @@ async def async_streaming( response = response_obj.predictions completion_response = response[0] - if ( - isinstance(completion_response, str) - and "\nOutput:\n" in completion_response - ): + if isinstance(completion_response, str) and "\nOutput:\n" in completion_response: completion_response = completion_response.split("\nOutput:\n", 1)[1] if stream: response = TextStreamer(completion_response) @@ -765,10 +670,7 @@ async def async_streaming( ) response = response_obj.predictions completion_response = response[0] - if ( - isinstance(completion_response, str) - and "\nOutput:\n" in completion_response - ): + if isinstance(completion_response, str) and "\nOutput:\n" in completion_response: completion_response = completion_response.split("\nOutput:\n", 1)[1] if stream: response = TextStreamer(completion_response) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py index cc0ecc2e3c6..a9c1e5819f2 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py @@ -1,9 +1,7 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig -def get_vertex_ai_partner_model_config( - model: str, vertex_publisher_or_api_spec: str -) -> BaseConfig: +def get_vertex_ai_partner_model_config(model: str, vertex_publisher_or_api_spec: str) -> BaseConfig: """Return config for handling response transformation for vertex ai partner models""" if vertex_publisher_or_api_spec == "anthropic": from .anthropic.transformation import VertexAIAnthropicConfig @@ -13,10 +11,7 @@ def get_vertex_ai_partner_model_config( from .ai21.transformation import VertexAIAi21Config return VertexAIAi21Config() - elif ( - vertex_publisher_or_api_spec == "openapi" - or vertex_publisher_or_api_spec == "mistralai" - ): + elif vertex_publisher_or_api_spec == "openapi" or vertex_publisher_or_api_spec == "mistralai": from .llama3.transformation import VertexAILlama3Config return VertexAILlama3Config() diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py index 8ffc00cc957..c8163708574 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/ai21/transformation.py @@ -49,9 +49,7 @@ class VertexAIAi21Config(OpenAIGPTConfig): drop_params: bool, ): if "max_completion_tokens" in non_default_params: - non_default_params["max_tokens"] = non_default_params.pop( - "max_completion_tokens" - ) + non_default_params["max_tokens"] = non_default_params.pop("max_completion_tokens") return litellm.OpenAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index a633ef3298a..5d338198c19 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -93,18 +93,12 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert # Add context management header if any other edits exist if has_other: - beta_values.add( - ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value - ) + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) # Check for web search tool for tool in tools: - if isinstance(tool, dict) and tool.get("type", "").startswith( - ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value - ): - beta_values.add( - ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value - ) + if isinstance(tool, dict) and tool.get("type", "").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value): + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value) break # Check for tool search tools - Vertex AI uses different beta header @@ -127,9 +121,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert stream: Optional[bool] = None, ) -> str: if api_base is None: - raise ValueError( - "api_base is required. Unable to determine the correct api_base for the request." - ) + raise ValueError("api_base is required. Unable to determine the correct api_base for the request.") return api_base # no transformation is needed - handled in validate_environment def transform_anthropic_messages_request( @@ -152,9 +144,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert anthropic_messages_request["anthropic_version"] = "vertex-2023-10-16" - anthropic_messages_request.pop( - "model", None - ) # do not pass model in request body to vertex ai + anthropic_messages_request.pop("model", None) # do not pass model in request body to vertex ai sanitize_vertex_anthropic_output_params(anthropic_messages_request, model) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index ae8bdc55443..c8d91be359b 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -17,13 +17,9 @@ class VertexAIError(Exception): def __init__(self, status_code, message): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url=" https://cloud.google.com/vertex-ai/" - ) + self.request = httpx.Request(method="POST", url=" https://cloud.google.com/vertex-ai/") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class VertexAIAnthropicConfig(AnthropicConfig): @@ -55,9 +51,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): def should_strip_billing_metadata(self) -> bool: return True - def _add_context_management_beta_headers( - self, beta_set: set, context_management: dict - ) -> None: + def _add_context_management_beta_headers(self, beta_set: set, context_management: dict) -> None: """ Add context_management beta headers to the beta_set. @@ -87,9 +81,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): # Add context management header if any other edits exist if has_other: - beta_set.add( - ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value - ) + beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) def transform_request( self, @@ -124,9 +116,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): beta_set = set(auto_betas) if tool_search_used: - beta_set.add( - "tool-search-tool-2025-10-19" - ) # Vertex requires this header for tool search + beta_set.add("tool-search-tool-2025-10-19") # Vertex requires this header for tool search # Add context_management beta headers (compact and/or context-management) context_management = optional_params.get("context_management") @@ -218,10 +208,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): """ Check if the model is supported by the VertexAI Anthropic API. """ - if ( - custom_llm_provider != "vertex_ai" - and custom_llm_provider != "vertex_ai_beta" - ): + if custom_llm_provider != "vertex_ai" and custom_llm_provider != "vertex_ai_beta": return False if "claude" in model.lower(): return True diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 3a3ab2e2465..d3edf2e9848 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -130,9 +130,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): vertex_project = self.get_vertex_ai_project(litellm_params) # Check for count_tokens specific location override - vertex_count_tokens_location = litellm_params.get( - "vertex_count_tokens_location" - ) + vertex_count_tokens_location = litellm_params.get("vertex_count_tokens_location") vertex_location_raw = self.get_vertex_ai_location(litellm_params) # Determine final location with precedence: @@ -185,9 +183,7 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): # Check for errors if response.status_code != 200: error_text = response.text - raise ValueError( - f"Token counting request failed with status {response.status_code}: {error_text}" - ) + raise ValueError(f"Token counting request failed with status {response.status_code}: {error_text}") # Parse response result = response.json() diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py index 47c388f0a54..13cb09dc22c 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py @@ -28,9 +28,7 @@ class VertexAIGPTOSSTransformation(OpenAIGPTConfig): "functions", ] base_gpt_series_params = [ - param - for param in base_gpt_series_params - if param not in TOOL_CALLING_PARAMS_TO_REMOVE + param for param in base_gpt_series_params if param not in TOOL_CALLING_PARAMS_TO_REMOVE ] return base_gpt_series_params diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index 3031f159d87..411a2a1cb0d 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -78,9 +78,7 @@ class VertexAILlama3Config(OpenAIGPTConfig): drop_params: bool, ): if "max_completion_tokens" in non_default_params: - non_default_params["max_tokens"] = non_default_params.pop( - "max_completion_tokens" - ) + non_default_params["max_tokens"] = non_default_params.pop("max_completion_tokens") return super().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -128,9 +126,7 @@ class VertexAILlama3Config(OpenAIGPTConfig): except Exception as e: response_headers = getattr(raw_response, "headers", None) raise VertexAIError( - message="Unable to get json response - {}, Original Response: {}".format( - str(e), raw_response.text - ), + message="Unable to get json response - {}, Original Response: {}".format(str(e), raw_response.text), status_code=raw_response.status_code, headers=response_headers, ) @@ -202,9 +198,7 @@ class VertexAILlama3StreamingHandler(OpenAIChatCompletionStreamingHandler): elif delta.role is None: delta.role = "assistant" # If the first chunk has empty content, ensure it's still emitted - if ( - delta.content == "" or delta.content is None - ) and delta.provider_specific_fields is None: + if (delta.content == "" or delta.content is None) and delta.provider_specific_fields is None: delta.provider_specific_fields = {} self.sent_role = True return result diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 78669f1e789..097928508a5 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -20,13 +20,9 @@ class VertexAIError(Exception): def __init__(self, status_code, message): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url=" https://cloud.google.com/vertex-ai/" - ) + self.request = httpx.Request(method="POST", url=" https://cloud.google.com/vertex-ai/") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs class PartnerModelPrefixes(str, Enum): @@ -125,9 +121,7 @@ class VertexAIPartnerModels(VertexBase): message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""", ) - if not ( - hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models") - ): + if not (hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")): raise VertexAIError( status_code=400, message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", @@ -174,9 +168,7 @@ class VertexAIPartnerModels(VertexBase): if "codestral" in model and litellm_params.get("text_completion") is True: optional_params["model"] = model - text_completion_model_response = litellm.TextCompletionResponse( - stream=stream - ) + text_completion_model_response = litellm.TextCompletionResponse(stream=stream) return codestral_fim_completions.completion( model=model, messages=messages, diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py index e3f25b425ff..6525d3342f5 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -57,9 +57,7 @@ class VertexBGEConfig: return model_lower.startswith("bge/") or "bge" in model_lower @staticmethod - def transform_request( - input: Union[list, str], optional_params: dict, model: str - ) -> VertexEmbeddingRequest: + def transform_request(input: Union[list, str], optional_params: dict, model: str) -> VertexEmbeddingRequest: """ Transforms an OpenAI request to a Vertex BGE embedding request. @@ -82,9 +80,7 @@ class VertexBGEConfig: input = [input] for text in input: - embedding_input = VertexBGEConfig._create_embedding_input( - prompt=text, task_type=task_type, title=title - ) + embedding_input = VertexBGEConfig._create_embedding_input(prompt=text, task_type=task_type, title=title) vertex_text_embedding_input_list.append(embedding_input) vertex_request["instances"] = vertex_text_embedding_input_list @@ -119,9 +115,7 @@ class VertexBGEConfig: return text_embedding_input @staticmethod - def transform_response( - response: dict, model: str, model_response: EmbeddingResponse - ) -> EmbeddingResponse: + def transform_response(response: dict, model: str, model_response: EmbeddingResponse) -> EmbeddingResponse: """ Transforms a Vertex BGE embedding response to OpenAI format. @@ -151,9 +145,7 @@ class VertexBGEConfig: _predictions = response["predictions"] if not isinstance(_predictions, list): - raise ValueError( - f"Expected 'predictions' to be a list, got {type(_predictions)}" - ) + raise ValueError(f"Expected 'predictions' to be a list, got {type(_predictions)}") embedding_response = [] # BGE models don't return token counts, so we estimate or set to 0 @@ -161,9 +153,7 @@ class VertexBGEConfig: for idx, embedding_values in enumerate(_predictions): if not isinstance(embedding_values, list): - raise ValueError( - f"Expected embedding at index {idx} to be a list, got {type(embedding_values)}" - ) + raise ValueError(f"Expected embedding at index {idx} to be a list, got {type(embedding_values)}") embedding_response.append( { @@ -176,8 +166,6 @@ class VertexBGEConfig: model_response.object = "list" model_response.data = embedding_response model_response.model = model - usage = Usage( - prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens - ) + usage = Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens) setattr(model_response, "usage", usage) return model_response diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index 729cc9c3ead..0e7afd5da3f 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -65,9 +65,7 @@ class VertexEmbedding(VertexBase): litellm_params=litellm_params, ) - should_use_v1beta1_features = self.is_using_v1beta1_features( - optional_params=optional_params - ) + should_use_v1beta1_features = self.is_using_v1beta1_features(optional_params=optional_params) _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, @@ -92,11 +90,13 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) - vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, - optional_params=optional_params, - model=model, - litellm_params=litellm_params, + vertex_request: VertexEmbeddingRequest = ( + litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, + ) ) _client_params = {} @@ -128,14 +128,10 @@ class VertexEmbedding(VertexBase): _json_response = response.json() ## LOGGING POST-CALL - logging_obj.post_call( - input=input, api_key=None, original_response=_json_response - ) + logging_obj.post_call(input=input, api_key=None, original_response=_json_response) - model_response = ( - litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( - response=_json_response, model=model, model_response=model_response - ) + model_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( + response=_json_response, model=model, model_response=model_response ) return model_response @@ -164,9 +160,7 @@ class VertexEmbedding(VertexBase): """ Async embedding implementation """ - should_use_v1beta1_features = self.is_using_v1beta1_features( - optional_params=optional_params - ) + should_use_v1beta1_features = self.is_using_v1beta1_features(optional_params=optional_params) _auth_header, vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, @@ -190,20 +184,20 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) - vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, - optional_params=optional_params, - model=model, - litellm_params=litellm_params, + vertex_request: VertexEmbeddingRequest = ( + litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, + ) ) _async_client_params = {} if timeout: _async_client_params["timeout"] = timeout if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client( - params=_async_client_params, llm_provider=litellm.LlmProviders.VERTEX_AI - ) + client = get_async_httpx_client(params=_async_client_params, llm_provider=litellm.LlmProviders.VERTEX_AI) else: client = client # type: ignore ## LOGGING @@ -228,14 +222,10 @@ class VertexEmbedding(VertexBase): _json_response = response.json() ## LOGGING POST-CALL - logging_obj.post_call( - input=input, api_key=None, original_response=_json_response - ) + logging_obj.post_call(input=input, api_key=None, original_response=_json_response) - model_response = ( - litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( - response=_json_response, model=model, model_response=model_response - ) + model_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( + response=_json_response, model=model, model_response=model_response ) return model_response diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 24396628dbd..6b7e6c036c0 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -75,9 +75,7 @@ class VertexAITextEmbeddingConfig(BaseModel): def get_supported_openai_params(self): return ["dimensions"] - def map_openai_params( - self, non_default_params: dict, optional_params: dict, kwargs: dict - ): + def map_openai_params(self, non_default_params: dict, optional_params: dict, kwargs: dict): for param, value in non_default_params.items(): if param == "dimensions": optional_params["outputDimensionality"] = value @@ -116,10 +114,8 @@ class VertexAITextEmbeddingConfig(BaseModel): labels = pop_vertex_request_labels(optional_params, litellm_params) if model.isdigit(): - vertex_request = ( - self._transform_openai_request_to_fine_tuned_embedding_request( - input, optional_params, model - ) + vertex_request = self._transform_openai_request_to_fine_tuned_embedding_request( + input, optional_params, model ) if labels: vertex_request["labels"] = labels @@ -141,9 +137,7 @@ class VertexAITextEmbeddingConfig(BaseModel): input = [input] # Convert single string to list for uniform processing for text in input: - embedding_input = self.create_embedding_input( - content=text, task_type=task_type, title=title - ) + embedding_input = self.create_embedding_input(content=text, task_type=task_type, title=title) vertex_text_embedding_input_list.append(embedding_input) vertex_request["instances"] = vertex_text_embedding_input_list @@ -188,14 +182,9 @@ class VertexAITextEmbeddingConfig(BaseModel): vertex_text_embedding_input_list.append(embedding_input) vertex_request["instances"] = vertex_text_embedding_input_list - vertex_request["parameters"] = TextEmbeddingFineTunedParameters( - **optional_params - ) + vertex_request["parameters"] = TextEmbeddingFineTunedParameters(**optional_params) # Remove 'shared_session' from parameters if present - if ( - vertex_request["parameters"] is not None - and "shared_session" in vertex_request["parameters"] - ): + if vertex_request["parameters"] is not None and "shared_session" in vertex_request["parameters"]: del vertex_request["parameters"]["shared_session"] # type: ignore[typeddict-item] return vertex_request @@ -233,17 +222,13 @@ class VertexAITextEmbeddingConfig(BaseModel): Transforms a vertex embedding response to an openai response. """ if model.isdigit(): - return self._transform_vertex_response_to_openai_for_fine_tuned_models( - response, model, model_response - ) + return self._transform_vertex_response_to_openai_for_fine_tuned_models(response, model, model_response) # Import here to avoid circular import issues with litellm.__init__ from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig if VertexBGEConfig.is_bge_model(model): - return VertexBGEConfig.transform_response( - response=response, model=model, model_response=model_response - ) + return VertexBGEConfig.transform_response(response=response, model=model, model_response=model_response) _predictions = response["predictions"] @@ -263,9 +248,7 @@ class VertexAITextEmbeddingConfig(BaseModel): model_response.object = "list" model_response.data = embedding_response model_response.model = model - usage = Usage( - prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens - ) + usage = Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens) setattr(model_response, "usage", usage) return model_response @@ -286,17 +269,13 @@ class VertexAITextEmbeddingConfig(BaseModel): { "object": "embedding", "index": idx, - "embedding": embedding_values[ - 0 - ], # The embedding values are nested one level deeper + "embedding": embedding_values[0], # The embedding values are nested one level deeper } ) model_response.object = "list" model_response.data = embedding_response model_response.model = model - usage = Usage( - prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens - ) + usage = Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens) setattr(model_response, "usage", usage) return model_response diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index b6bf2f73b72..9622a93c0d8 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -71,9 +71,7 @@ class VertexAIGemmaModels(VertexBase): message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""", ) - if not ( - hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models") - ): + if not (hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")): raise VertexAIError( status_code=400, message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 35cd54d65f6..567c8c6a3ee 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -87,9 +87,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): # Remove params not needed/supported by Vertex Gemma openai_request.pop("model", None) - openai_request.pop( - "stream", None - ) # Streaming not supported, will be faked client-side + openai_request.pop("stream", None) # Streaming not supported, will be faked client-side openai_request.pop("stream_options", None) # Stream options not supported # Vertex Gemma's chatCompletions wrapper does not understand # `context_management` (an Anthropic/Responses API concept). Strip it @@ -264,9 +262,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) # Return fake stream iterator if streaming was requested - return self._handle_fake_stream_response( - model_response=model_response, stream=stream - ) + return self._handle_fake_stream_response(model_response=model_response, stream=stream) async def _async_completion( self, @@ -359,6 +355,4 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) # Return fake stream iterator if streaming was requested - return self._handle_fake_stream_response( - model_response=model_response, stream=stream - ) + return self._handle_fake_stream_response(model_response=model_response, stream=stream) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 18a9c98c315..d57d7bf17df 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -26,8 +26,7 @@ from .common_utils import ( ) GOOGLE_IMPORT_ERROR_MESSAGE = ( - "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' " - "or pip install google-cloud-aiplatform" + "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform" ) if TYPE_CHECKING: @@ -73,9 +72,7 @@ class VertexBase: # Try to get supported_regions directly from model_cost # Check both with and without vertex_ai/ prefix - model_key = ( - f"vertex_ai/{model}" if not model.startswith("vertex_ai/") else model - ) + model_key = f"vertex_ai/{model}" if not model.startswith("vertex_ai/") else model model_info = litellm.model_cost.get(model_key, {}) supported_regions = model_info.get("supported_regions") @@ -86,8 +83,7 @@ class VertexBase: # If user specified a region not supported by this model, override it if vertex_region not in supported_regions: verbose_logger.warning( - "Vertex AI model '%s' does not support region '%s' " - "(supported: %s). Routing to '%s'.", + "Vertex AI model '%s' does not support region '%s' (supported: %s). Routing to '%s'.", model, vertex_region, supported_regions, @@ -128,18 +124,14 @@ class VertexBase: elif isinstance(credentials, dict): json_obj = credentials else: - raise ValueError( - "Invalid credentials type: {}".format(type(credentials)) - ) + raise ValueError("Invalid credentials type: {}".format(type(credentials))) # Check if the JSON object contains Workload Identity Federation configuration if "type" in json_obj and json_obj["type"] == "external_account": # If environment_id key contains "aws" value it corresponds to an AWS config file credential_source = json_obj.get("credential_source", {}) environment_id = ( - credential_source.get("environment_id", "") - if isinstance(credential_source, dict) - else "" + credential_source.get("environment_id", "") if isinstance(credential_source, dict) else "" ) if isinstance(environment_id, str) and "aws" in environment_id: # Check if explicit AWS params are in the JSON (bypasses metadata) @@ -159,10 +151,7 @@ class VertexBase: json_obj, scopes=["https://www.googleapis.com/auth/cloud-platform"], ) - elif ( - isinstance(credential_source, dict) - and "executable" in credential_source - ): + elif isinstance(credential_source, dict) and "executable" in credential_source: creds = self._credentials_from_pluggable( json_obj, scopes=["https://www.googleapis.com/auth/cloud-platform"], @@ -179,7 +168,9 @@ class VertexBase: scopes=["https://www.googleapis.com/auth/cloud-platform"], ) if project_id is None: - project_id = creds.quota_project_id # authorized user credentials don't have a project_id, only quota_project_id + project_id = ( + creds.quota_project_id + ) # authorized user credentials don't have a project_id, only quota_project_id else: creds = self._credentials_from_service_account( json_obj, @@ -201,9 +192,7 @@ class VertexBase: raise ValueError("Could not resolve project_id") if not isinstance(project_id, str): - raise TypeError( - f"Expected project_id to be a str but got {type(project_id)}" - ) + raise TypeError(f"Expected project_id to be a str but got {type(project_id)}") return creds, project_id @@ -247,9 +236,7 @@ class VertexBase: except ImportError: raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE) - return google.oauth2.credentials.Credentials.from_authorized_user_info( - json_obj, scopes=scopes - ) + return google.oauth2.credentials.Credentials.from_authorized_user_info(json_obj, scopes=scopes) def _credentials_from_service_account(self, json_obj, scopes): try: @@ -257,9 +244,7 @@ class VertexBase: except ImportError: raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE) - return google.oauth2.service_account.Credentials.from_service_account_info( - json_obj, scopes=scopes - ) + return google.oauth2.service_account.Credentials.from_service_account_info(json_obj, scopes=scopes) def _credentials_from_default_auth(self, scopes): try: @@ -272,14 +257,10 @@ class VertexBase: def get_default_vertex_location(self) -> str: return "us-central1" - def get_api_base( - self, api_base: Optional[str], vertex_location: Optional[str] - ) -> str: + def get_api_base(self, api_base: Optional[str], vertex_location: Optional[str]) -> str: if api_base: return api_base - return get_vertex_base_url( - vertex_location or self.get_default_vertex_location() - ) + return get_vertex_base_url(vertex_location or self.get_default_vertex_location()) @staticmethod def create_vertex_url( @@ -324,9 +305,7 @@ class VertexBase: ) -> str: # Use get_vertex_region to handle global-only models resolved_location = self.get_vertex_region(vertex_location, model) - api_base = self.get_api_base( - api_base=custom_api_base, vertex_location=resolved_location - ) + api_base = self.get_api_base(api_base=custom_api_base, vertex_location=resolved_location) default_api_base = VertexBase.create_vertex_url( vertex_location=resolved_location, vertex_project=vertex_project or project_id, @@ -383,17 +362,13 @@ class VertexBase: caller is done with the lock so the entry can be pruned when no other coroutine is holding or waiting on it. """ - lock = self._async_refresh_locks.setdefault( - credential_cache_key, asyncio.Lock() - ) + lock = self._async_refresh_locks.setdefault(credential_cache_key, asyncio.Lock()) self._async_refresh_lock_refcounts[credential_cache_key] = ( self._async_refresh_lock_refcounts.get(credential_cache_key, 0) + 1 ) return lock - def _release_async_refresh_lock( - self, credential_cache_key: tuple, lock: asyncio.Lock - ) -> None: + def _release_async_refresh_lock(self, credential_cache_key: tuple, lock: asyncio.Lock) -> None: """Decrement the refcount and drop the lock entry when it reaches zero. Must be called only after the caller has released ``lock`` (i.e. once @@ -459,9 +434,7 @@ class VertexBase: return None return creds.token, resolved_project, token_state, creds, cached_project_id - def _unpack_cached_credentials( - self, credential_cache_key: tuple - ) -> Tuple[Any, Optional[str]]: + def _unpack_cached_credentials(self, credential_cache_key: tuple) -> Tuple[Any, Optional[str]]: """ Return (credentials, project_id) from the cache, or (None, None) if not cached. Handles both tuple and legacy cache formats. @@ -471,9 +444,7 @@ class VertexBase: cached_entry = self._credentials_project_mapping[credential_cache_key] if isinstance(cached_entry, tuple): return cached_entry - return cached_entry, cached_entry.quota_project_id or getattr( - cached_entry, "project_id", None - ) + return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None) def _get_token_state(self, credentials: Any) -> "TokenState": """ @@ -550,9 +521,7 @@ class VertexBase: exc_info=True, ) - async def _await_in_flight_background_refresh( - self, credential_cache_key: tuple - ) -> None: + async def _await_in_flight_background_refresh(self, credential_cache_key: tuple) -> None: """Wait for an in-flight background refresh to finish, if any. google-auth's ``Credentials.refresh()`` is not safe to invoke @@ -588,9 +557,7 @@ class VertexBase: return self._background_refresh_tasks.pop(credential_cache_key, None) task = asyncio.create_task( - self._background_refresh_credentials( - credentials, credential_cache_key, credential_project_id - ) + self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id) ) def _drop_background_refresh_task(_fut: asyncio.Future[Any]) -> None: @@ -664,9 +631,7 @@ class VertexBase: if custom_llm_provider == "gemini": # For Gemini (Google AI Studio), construct the full path like other providers if model is None: - raise ValueError( - "Model parameter is required for Gemini custom API base URLs" - ) + raise ValueError("Model parameter is required for Gemini custom API base URLs") url = "{}/models/{}:{}".format(api_base, model, endpoint) if gemini_api_key is None: raise ValueError( @@ -795,8 +760,7 @@ class VertexBase: The original error if reauthentication fails """ verbose_logger.debug( - f"Handling reauthentication for project_id: {project_id}. " - f"Clearing cache and retrying once." + f"Handling reauthentication for project_id: {project_id}. Clearing cache and retrying once." ) # Clear the cached credentials @@ -829,8 +793,7 @@ class VertexBase: Async reauthentication retry that stays within the per-key async lock. """ verbose_logger.debug( - f"Handling async reauthentication for project_id: {project_id}. " - f"Clearing cache and retrying once." + f"Handling async reauthentication for project_id: {project_id}. Clearing cache and retrying once." ) self._credentials_project_mapping.pop(credential_cache_key, None) @@ -846,11 +809,7 @@ class VertexBase: ) if project_id is None and isinstance(credential_project_id, str): project_id = credential_project_id - cache_credentials = ( - json.dumps(credentials) - if isinstance(credentials, dict) - else credentials - ) + cache_credentials = json.dumps(credentials) if isinstance(credentials, dict) else credentials resolved_cache_key = (cache_credentials, project_id) # Always overwrite — any pre-existing entry at the resolved key # references the OLD credentials object we just replaced, and @@ -903,20 +862,14 @@ class VertexBase: """ # Convert dict credentials to string for caching - cache_credentials = ( - json.dumps(credentials) if isinstance(credentials, dict) else credentials - ) + cache_credentials = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key = (cache_credentials, project_id) _credentials: Optional[GoogleCredentialsObject] = None - verbose_logger.debug( - f"Checking cached credentials for project_id: {project_id}" - ) + verbose_logger.debug(f"Checking cached credentials for project_id: {project_id}") if credential_cache_key in self._credentials_project_mapping: - verbose_logger.debug( - f"Cached credentials found for project_id: {project_id}." - ) + verbose_logger.debug(f"Cached credentials found for project_id: {project_id}.") # Retrieve both credentials and cached project_id cached_entry = self._credentials_project_mapping[credential_cache_key] verbose_logger.debug("cached_entry: %s", cached_entry) @@ -925,9 +878,7 @@ class VertexBase: else: # Backward compatibility with old cache format _credentials = cached_entry - credential_project_id = _credentials.quota_project_id or getattr( - _credentials, "project_id", None - ) + credential_project_id = _credentials.quota_project_id or getattr(_credentials, "project_id", None) verbose_logger.debug( "Using cached credentials for project_id: %s", credential_project_id, @@ -939,9 +890,7 @@ class VertexBase: ) try: - _credentials, credential_project_id = self.load_auth( - credentials=credentials, project_id=project_id - ) + _credentials, credential_project_id = self.load_auth(credentials=credentials, project_id=project_id) except Exception as e: verbose_logger.exception( f"Failed to load vertex credentials. Check to see if credentials containing partial/invalid information. Error: {str(e)}" @@ -962,11 +911,7 @@ class VertexBase: ## VALIDATE CREDENTIALS verbose_logger.debug("Validating credentials") - if ( - project_id is None - and credential_project_id is not None - and isinstance(credential_project_id, str) - ): + if project_id is None and credential_project_id is not None and isinstance(credential_project_id, str): project_id = credential_project_id # Update cache with resolved project_id for future lookups resolved_cache_key = (cache_credentials, project_id) @@ -1030,9 +975,7 @@ class VertexBase: """ from google.auth.credentials import TokenState - cache_credentials = ( - json.dumps(credentials) if isinstance(credentials, dict) else credentials - ) + cache_credentials = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key = (cache_credentials, project_id) # === FAST PATH (no lock) === @@ -1043,13 +986,9 @@ class VertexBase: # callers on the lock just to schedule that refresh. usable = self._try_get_usable_cached_token(credential_cache_key, project_id) if usable is not None: - cached_token, resolved_project, token_state, creds, cached_project_id = ( - usable - ) + cached_token, resolved_project, token_state, creds, cached_project_id = usable if token_state == TokenState.STALE: - self._schedule_background_refresh( - creds, credential_cache_key, cached_project_id - ) + self._schedule_background_refresh(creds, credential_cache_key, cached_project_id) return cached_token, resolved_project # === SLOW PATH (per-key lock) === @@ -1061,18 +1000,14 @@ class VertexBase: if cached is not None: return cached - _credentials, credential_project_id = self._unpack_cached_credentials( - credential_cache_key - ) + _credentials, credential_project_id = self._unpack_cached_credentials(credential_cache_key) # Load credentials if not cached if _credentials is None: ( _credentials, credential_project_id, - ) = await self._load_and_cache_credentials( - credentials, project_id, credential_cache_key - ) + ) = await self._load_and_cache_credentials(credentials, project_id, credential_cache_key) # Resolve project_id from credentials if not provided if project_id is None and isinstance(credential_project_id, str): @@ -1117,9 +1052,7 @@ class VertexBase: # on the same credentials object, and the background task # runs outside this lock. await self._await_in_flight_background_refresh(credential_cache_key) - cached = self._try_get_cached_token( - credential_cache_key, project_id - ) + cached = self._try_get_cached_token(credential_cache_key, project_id) if cached is not None: return cached @@ -1133,9 +1066,7 @@ class VertexBase: ) except Exception as e: if "Reauthentication is needed" in str(e): - verbose_logger.debug( - "Reauthentication needed, clearing cache and retrying" - ) + verbose_logger.debug("Reauthentication needed, clearing cache and retrying") return await self._handle_reauthentication_async( credentials=credentials, project_id=project_id, @@ -1145,9 +1076,7 @@ class VertexBase: raise # Final validation - if _credentials.token is None or not isinstance( - _credentials.token, str - ): + if _credentials.token is None or not isinstance(_credentials.token, str): raise ValueError( "Could not resolve credentials token. Got None or non-string token (type={})".format( type(_credentials.token).__name__ @@ -1179,9 +1108,7 @@ class VertexBase: project_id=project_id, ) - def set_headers( - self, auth_header: Optional[str], extra_headers: Optional[dict] - ) -> dict: + def set_headers(self, auth_header: Optional[str], extra_headers: Optional[dict]) -> dict: headers = { "Content-Type": "application/json", } diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index f54b8d93500..bd9d95e6d04 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -48,10 +48,7 @@ def create_vertex_url( """Return the api base for vertex model garden (without /chat/completions).""" base_url = get_vertex_base_url(vertex_location) if _vertex_model_garden_model_id_in_json_body(model): - return ( - f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}" - "/endpoints/openapi" - ) + return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi" return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}" @@ -95,9 +92,7 @@ class VertexAIModelGardenModels(VertexBase): message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""", ) - if not ( - hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models") - ): + if not (hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")): raise VertexAIError( status_code=400, message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index b84966354b8..98af8ca30ea 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -49,9 +49,7 @@ def _build_vertex_video_usage_from_request_data( return usage_data parameters = request_data.get("parameters", {}) - duration = ( - parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS - ) + duration = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS if duration is not None: try: usage_data["duration_seconds"] = float(duration) @@ -203,16 +201,10 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict # Ensure litellm_params is a dict for type checking - params_dict: Dict[str, Any] = ( - cast(Dict[str, Any], litellm_params) if litellm_params is not None else {} - ) + params_dict: Dict[str, Any] = cast(Dict[str, Any], litellm_params) if litellm_params is not None else {} - vertex_project = VertexBase.safe_get_vertex_ai_project( - litellm_params=params_dict - ) - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( - litellm_params=params_dict - ) + vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=params_dict) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=params_dict) # Get access token from Vertex credentials access_token, project_id = self.get_access_token( @@ -261,7 +253,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): else: base_url = get_vertex_base_url(vertex_location) - url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}" + url = ( + f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}" + ) return url @@ -376,15 +370,11 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): raise ValueError(f"No operation name in Veo response: {response_data}") if custom_llm_provider: - video_id = encode_video_id_with_provider( - operation_name, custom_llm_provider, model - ) + video_id = encode_video_id_with_provider(operation_name, custom_llm_provider, model) else: video_id = operation_name - video_obj = VideoObject( - id=video_id, object="video", status="processing", model=model - ) + video_obj = VideoObject(id=video_id, object="video", status="processing", model=model) video_obj.usage = _build_vertex_video_usage_from_request_data(request_data) return video_obj @@ -461,9 +451,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): model = self.extract_model_from_operation_name(operation_name) if custom_llm_provider: - video_id = encode_video_id_with_provider( - operation_name, custom_llm_provider, model - ) + video_id = encode_video_id_with_provider(operation_name, custom_llm_provider, model) else: video_id = operation_name @@ -471,9 +459,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): create_time_str = response_data.get("metadata", {}).get("createTime") if create_time_str: try: - created_at = _convert_vertex_datetime_to_openai_datetime( - create_time_str - ) + created_at = _convert_vertex_datetime_to_openai_datetime(create_time_str) except Exception: created_at = int(time.time()) else: @@ -515,9 +501,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): Since we need to make an HTTP call here, we'll use the same fetchPredictOperation approach as status retrieval. """ - return self.transform_video_status_retrieve_request( - video_id, api_base, litellm_params, headers - ) + return self.transform_video_status_retrieve_request(video_id, api_base, litellm_params, headers) def transform_video_content_response( self, @@ -533,8 +517,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): if not response_data.get("done", False): raise ValueError( - "Video generation is not complete yet. " - "Please check status with video_status() before downloading." + "Video generation is not complete yet. Please check status with video_status() before downloading." ) try: @@ -571,8 +554,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): Video remix is not supported by Veo API. """ raise NotImplementedError( - "Video remix is not supported by Vertex AI Veo. " - "Please use video_generation() to create new videos." + "Video remix is not supported by Vertex AI Veo. Please use video_generation() to create new videos." ) def transform_video_remix_response( @@ -622,8 +604,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): Video delete is not supported by Veo API. """ raise NotImplementedError( - "Video delete is not supported by Vertex AI Veo. " - "Videos are automatically cleaned up by Google." + "Video delete is not supported by Vertex AI Veo. Videos are automatically cleaned up by Google." ) def transform_video_delete_response( @@ -634,21 +615,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): """Video delete is not supported.""" raise NotImplementedError("Video delete is not supported by Vertex AI Veo.") - def transform_video_create_character_request( - self, name, video, api_base, litellm_params, headers - ): - raise NotImplementedError( - "video create character is not supported for Vertex AI" - ) + def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + raise NotImplementedError("video create character is not supported for Vertex AI") def transform_video_create_character_response(self, raw_response, logging_obj): - raise NotImplementedError( - "video create character is not supported for Vertex AI" - ) + raise NotImplementedError("video create character is not supported for Vertex AI") - def transform_video_get_character_request( - self, character_id, api_base, litellm_params, headers - ): + def transform_video_get_character_request(self, character_id, api_base, litellm_params, headers): raise NotImplementedError("video get character is not supported for Vertex AI") def transform_video_get_character_response(self, raw_response, logging_obj): @@ -692,10 +665,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): ) if not prefetched_source_data.get("done", False): - raise ValueError( - "Source video generation is not complete yet. " - "Check the video status before editing." - ) + raise ValueError("Source video generation is not complete yet. Check the video status before editing.") videos = prefetched_source_data.get("response", {}).get("videos", []) if not videos: @@ -709,9 +679,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): video_input["bytesBase64Encoded"] = source_video["bytesBase64Encoded"] video_input["mimeType"] = source_video.get("mimeType", "video/mp4") else: - raise ValueError( - "Source video has neither gcsUri nor bytesBase64Encoded. Cannot edit." - ) + raise ValueError("Source video has neither gcsUri nor bytesBase64Encoded. Cannot edit.") operation_name = extract_original_video_id(video_id) model = self.extract_model_from_operation_name(operation_name) or "" @@ -757,9 +725,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): model = self.extract_model_from_operation_name(operation_name) or "" if custom_llm_provider: - video_id = encode_video_id_with_provider( - operation_name, custom_llm_provider, model - ) + video_id = encode_video_id_with_provider(operation_name, custom_llm_provider, model) else: video_id = operation_name @@ -784,9 +750,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): ): raise NotImplementedError("video extension is not supported for Vertex AI") - def transform_video_extension_response( - self, raw_response, logging_obj, custom_llm_provider=None - ): + def transform_video_extension_response(self, raw_response, logging_obj, custom_llm_provider=None): raise NotImplementedError("video extension is not supported for Vertex AI") def get_error_class( diff --git a/litellm/llms/vllm/common_utils.py b/litellm/llms/vllm/common_utils.py index e2ed0daafe4..1d6b8d7897e 100644 --- a/litellm/llms/vllm/common_utils.py +++ b/litellm/llms/vllm/common_utils.py @@ -60,9 +60,7 @@ class VLLMModelInfo(BaseLLMModelInfo): def get_base_model(model: str) -> Optional[str]: return model - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: api_base = VLLMModelInfo.get_api_base(api_base) api_key = VLLMModelInfo.get_api_key(api_key) endpoint = "/v1/models" @@ -85,6 +83,4 @@ class VLLMModelInfo(BaseLLMModelInfo): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return VLLMError( - status_code=status_code, message=error_message, headers=headers - ) + return VLLMError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/vllm/completion/handler.py b/litellm/llms/vllm/completion/handler.py index 1f13082917f..cb352b599f9 100644 --- a/litellm/llms/vllm/completion/handler.py +++ b/litellm/llms/vllm/completion/handler.py @@ -18,9 +18,7 @@ class VLLMError(Exception): self.message = message self.request = httpx.Request(method="POST", url="http://0.0.0.0:8000") self.response = httpx.Response(status_code=status_code, request=self.request) - super().__init__( - self.message - ) # Call the base class constructor with the parameters it needs + super().__init__(self.message) # Call the base class constructor with the parameters it needs # check if vllm is installed @@ -76,9 +74,7 @@ def completion( if llm: outputs = llm.generate(prompt, sampling_params) else: - raise VLLMError( - status_code=0, message="Need to pass in a model name to initialize vllm" - ) + raise VLLMError(status_code=0, message="Need to pass in a model name to initialize vllm") ## COMPLETION CALL if "stream" in optional_params and optional_params["stream"] is True: @@ -110,9 +106,7 @@ def completion( return model_response -def batch_completions( - model: str, messages: list, optional_params=None, custom_prompt_dict={} -): +def batch_completions(model: str, messages: list, optional_params=None, custom_prompt_dict={}): """ Example usage: import litellm @@ -164,9 +158,7 @@ def batch_completions( if llm: outputs = llm.generate(prompts, sampling_params) else: - raise VLLMError( - status_code=0, message="Need to pass in a model name to initialize vllm" - ) + raise VLLMError(status_code=0, message="Need to pass in a model name to initialize vllm") final_outputs = [] for output in outputs: diff --git a/litellm/llms/volcengine/chat/transformation.py b/litellm/llms/volcengine/chat/transformation.py index d60f54615fa..c6dbbdbce60 100644 --- a/litellm/llms/volcengine/chat/transformation.py +++ b/litellm/llms/volcengine/chat/transformation.py @@ -96,13 +96,10 @@ class VolcEngineChatConfig(OpenAILikeChatConfig): if ( thinking_value is not None and isinstance(thinking_value, dict) - and thinking_value.get("type", None) - in ["enabled", "disabled", "auto"] # legal values, see docs + and thinking_value.get("type", None) in ["enabled", "disabled", "auto"] # legal values, see docs ): # Add thinking parameter to extra_body for all legal cases - optional_params.setdefault("extra_body", {})["thinking"] = ( - thinking_value - ) + optional_params.setdefault("extra_body", {})["thinking"] = thinking_value else: # Skip adding thinking parameter when it's not set or has invalid value pass diff --git a/litellm/llms/volcengine/common_utils.py b/litellm/llms/volcengine/common_utils.py index 0c8d3daebdc..be639086437 100644 --- a/litellm/llms/volcengine/common_utils.py +++ b/litellm/llms/volcengine/common_utils.py @@ -14,15 +14,11 @@ class VolcEngineError(BaseLLMException): Custom exception class for Volcengine provider errors. """ - def __init__( - self, status_code: int, message: str, headers: Optional[httpx.Headers] = None - ): + def __init__(self, status_code: int, message: str, headers: Optional[httpx.Headers] = None): self.status_code = status_code self.message = message self.headers = headers or httpx.Headers() - super().__init__( - status_code=status_code, message=message, headers=dict(self.headers) - ) + super().__init__(status_code=status_code, message=message, headers=dict(self.headers)) def get_volcengine_base_url(api_base: Optional[str] = None) -> str: diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 99e0a958ef1..56950151969 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -92,20 +92,14 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> VolcEngineError: - typed_headers: httpx.Headers = ( - headers - if isinstance(headers, httpx.Headers) - else httpx.Headers(headers or {}) - ) + typed_headers: httpx.Headers = headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers or {}) return VolcEngineError( status_code=status_code, message=error_message, headers=typed_headers, ) - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Build auth headers for Volcengine Responses API. """ @@ -122,9 +116,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): ) if api_key is None: - raise ValueError( - "Volcengine API key is required. Set ARK_API_KEY / VOLCENGINE_API_KEY or pass api_key." - ) + raise ValueError("Volcengine API key is required. Set ARK_API_KEY / VOLCENGINE_API_KEY or pass api_key.") return get_volcengine_headers(api_key=api_key, extra_headers=headers) @@ -173,9 +165,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): # Volcengine docs do not list parallel_tool_calls; drop it to avoid backend errors. if "parallel_tool_calls" in params: - verbose_logger.debug( - "Volcengine Responses API: dropping unsupported 'parallel_tool_calls' param." - ) + verbose_logger.debug("Volcengine Responses API: dropping unsupported 'parallel_tool_calls' param.") params.pop("parallel_tool_calls", None) return params @@ -195,11 +185,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): """ allowed = set(self._SUPPORTED_OPTIONAL_PARAMS) - sanitized_optional = { - k: v - for k, v in response_api_optional_request_params.items() - if k in allowed - } + sanitized_optional = {k: v for k, v in response_api_optional_request_params.items() if k in allowed} # Ensure metadata never reaches provider sanitized_optional.pop("metadata", None) sanitized_optional.pop("parallel_tool_calls", None) @@ -207,11 +193,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): # If extra_body is provided, filter its keys against the same allowlist to avoid # leaking unsupported params to the provider. if isinstance(sanitized_optional.get("extra_body"), dict): - filtered_body = { - k: v - for k, v in sanitized_optional["extra_body"].items() - if k in allowed - } + filtered_body = {k: v for k, v in sanitized_optional["extra_body"].items() if k in allowed} if filtered_body: sanitized_optional["extra_body"] = filtered_body else: @@ -247,9 +229,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): chunk = patched_chunk event_type = str(chunk.get("type")) if isinstance(chunk, dict) else None - event_pydantic_model = OpenAIResponsesAPIConfig.get_event_model_class( - event_type=event_type - ) + event_pydantic_model = OpenAIResponsesAPIConfig.get_event_model_class(event_type=event_type) patched_chunk = self._fill_missing_fields(chunk, event_pydantic_model) @@ -268,13 +248,9 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): ) raw_response_json = raw_response.json() if "created_at" in raw_response_json: - raw_response_json["created_at"] = _safe_convert_created_field( - raw_response_json["created_at"] - ) + raw_response_json["created_at"] = _safe_convert_created_field(raw_response_json["created_at"]) except Exception: - raise VolcEngineError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) @@ -282,9 +258,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): try: response = ResponsesAPIResponse(**raw_response_json) except Exception: - verbose_logger.debug( - "Volcengine Responses API: falling back to model_construct for response parsing." - ) + verbose_logger.debug("Volcengine Responses API: falling back to model_construct for response parsing.") response = ResponsesAPIResponse.model_construct(**raw_response_json) response._hidden_params["additional_headers"] = processed_headers @@ -301,9 +275,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -316,9 +288,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): try: raw_response_json = raw_response.json() except Exception: - raise VolcEngineError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) try: return DeleteResponseResult(**raw_response_json) except Exception: @@ -337,9 +307,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -352,9 +320,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): try: raw_response_json = raw_response.json() except Exception: - raise VolcEngineError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) @@ -379,9 +345,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): limit: int = 20, order: Literal["asc", "desc"] = "desc", ) -> Tuple[str, Dict]: - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}/input_items" params: Dict[str, Any] = {} if after is not None: @@ -404,9 +368,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): try: return raw_response.json() except Exception: - raise VolcEngineError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) ######################################################### ########## CANCEL RESPONSE API TRANSFORMATION ########## @@ -418,9 +380,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - encoded_response_id = encode_url_path_segment( - response_id, field_name="response_id" - ) + encoded_response_id = encode_url_path_segment(response_id, field_name="response_id") url = f"{api_base}/{encoded_response_id}/cancel" data: Dict = {} return url, data @@ -433,9 +393,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): try: raw_response_json = raw_response.json() except Exception: - raise VolcEngineError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VolcEngineError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) @@ -471,29 +429,19 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): for name, field in fields_map.items(): if name in patched: - patched[name] = VolcEngineResponsesAPIConfig._maybe_fill_nested( - patched[name], field.annotation - ) + patched[name] = VolcEngineResponsesAPIConfig._maybe_fill_nested(patched[name], field.annotation) continue # Explicit default or factory - if ( - field.default is not pyd_fields.PydanticUndefined - and field.default is not None - ): + if field.default is not pyd_fields.PydanticUndefined and field.default is not None: patched[name] = field.default continue - if ( - field.default_factory is not None - and field.default_factory is not pyd_fields.PydanticUndefined - ): + if field.default_factory is not None and field.default_factory is not pyd_fields.PydanticUndefined: patched[name] = field.default_factory() continue # Heuristic defaults for missing required fields - patched[name] = VolcEngineResponsesAPIConfig._default_for_annotation( - field.annotation - ) + patched[name] = VolcEngineResponsesAPIConfig._default_for_annotation(field.annotation) return patched @@ -533,10 +481,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): # Attempt to fill list elements if we know the element annotation elem_ann: Any = args[0] if args else None if elem_ann is not None: - return [ - VolcEngineResponsesAPIConfig._maybe_fill_nested(v, elem_ann) - for v in value - ] + return [VolcEngineResponsesAPIConfig._maybe_fill_nested(v, elem_ann) for v in value] return value diff --git a/litellm/llms/voyage/embedding/transformation.py b/litellm/llms/voyage/embedding/transformation.py index 91811e03927..7193fd2f10a 100644 --- a/litellm/llms/voyage/embedding/transformation.py +++ b/litellm/llms/voyage/embedding/transformation.py @@ -19,9 +19,7 @@ class VoyageError(BaseLLMException): ): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://api.voyageai.com/v1/embeddings" - ) + self.request = httpx.Request(method="POST", url="https://api.voyageai.com/v1/embeddings") self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, @@ -124,9 +122,7 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig): try: raw_response_json = raw_response.json() except Exception: - raise VoyageError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VoyageError(message=raw_response.text, status_code=raw_response.status_code) # model_response.usage model_response.model = raw_response_json.get("model") @@ -143,6 +139,4 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return VoyageError( - message=error_message, status_code=status_code, headers=headers - ) + return VoyageError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 1f5ca99f47d..d7cca3c87a8 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -24,9 +24,7 @@ class VoyageError(BaseLLMException): ): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings" - ) + self.request = httpx.Request(method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings") self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, @@ -126,9 +124,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): try: raw_response_json = raw_response.json() except Exception: - raise VoyageError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VoyageError(message=raw_response.text, status_code=raw_response.status_code) # model_response.usage model_response.model = raw_response_json.get("model") @@ -145,9 +141,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return VoyageError( - message=error_message, status_code=status_code, headers=headers - ) + return VoyageError(message=error_message, status_code=status_code, headers=headers) @staticmethod def is_contextualized_embeddings(model: str) -> bool: diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py index 55e221b065b..916037054ef 100644 --- a/litellm/llms/voyage/embedding/transformation_multimodal.py +++ b/litellm/llms/voyage/embedding/transformation_multimodal.py @@ -27,9 +27,7 @@ class VoyageMultimodalEmbeddingError(BaseLLMException): ): self.status_code = status_code self.message = message - self.request = httpx.Request( - method="POST", url="https://api.voyageai.com/v1/multimodalembeddings" - ) + self.request = httpx.Request(method="POST", url="https://api.voyageai.com/v1/multimodalembeddings") self.response = httpx.Response(status_code=status_code, request=self.request) super().__init__( status_code=status_code, @@ -124,10 +122,7 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): content = item.get("content") or [] return { **item, - "content": [ - self._normalize_content_item(content_item) - for content_item in content - ], + "content": [self._normalize_content_item(content_item) for content_item in content], } return item @@ -159,9 +154,7 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): try: raw_response_json = raw_response.json() except Exception: - raise VoyageMultimodalEmbeddingError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VoyageMultimodalEmbeddingError(message=raw_response.text, status_code=raw_response.status_code) model_response.model = raw_response_json.get("model") model_response.data = raw_response_json.get("data") @@ -178,6 +171,4 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return VoyageMultimodalEmbeddingError( - message=error_message, status_code=status_code, headers=headers - ) + return VoyageMultimodalEmbeddingError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 05991d4cc8c..e426e39962b 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -88,9 +88,7 @@ class VoyageRerankConfig(BaseRerankConfig): litellm_params: Dict = {}, ) -> RerankResponse: if raw_response.status_code != 200: - raise VoyageError( - message=raw_response.text, status_code=raw_response.status_code - ) + raise VoyageError(message=raw_response.text, status_code=raw_response.status_code) logging_obj.post_call(original_response=raw_response.text) @@ -141,13 +139,9 @@ class VoyageRerankConfig(BaseRerankConfig): optional_params: dict | None = None, ) -> Dict: if api_key is None: - api_key = get_secret_str("VOYAGE_API_KEY") or get_secret_str( - "VOYAGE_AI_API_KEY" - ) + api_key = get_secret_str("VOYAGE_API_KEY") or get_secret_str("VOYAGE_AI_API_KEY") if api_key is None: - raise ValueError( - "Voyage AI API key is required. Set via `api_key` parameter or `VOYAGE_API_KEY` env var." - ) + raise ValueError("Voyage AI API key is required. Set via `api_key` parameter or `VOYAGE_API_KEY` env var.") return { "Authorization": f"Bearer {api_key}", "content-type": "application/json", @@ -172,9 +166,5 @@ class VoyageRerankConfig(BaseRerankConfig): return 0.0, 0.0 return model_info["input_cost_per_token"] * total_tokens, 0.0 - def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] - ): - return VoyageError( - message=error_message, status_code=status_code, headers=headers - ) + def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]): + return VoyageError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 5944705258e..6d28790b8d1 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -25,9 +25,7 @@ from ...openai.transcriptions.whisper_transformation import ( from ..common_utils import IBMWatsonXMixin -class IBMWatsonXAudioTranscriptionConfig( - IBMWatsonXMixin, OpenAIWhisperAudioTranscriptionConfig -): +class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTranscriptionConfig): """ IBM WatsonX Audio Transcription Config @@ -65,9 +63,7 @@ class IBMWatsonXAudioTranscriptionConfig( result.pop("Content-Type", None) return result - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIAudioTranscriptionOptionalParams]: """ Get the supported OpenAI params for WatsonX audio transcription. """ @@ -98,9 +94,7 @@ class IBMWatsonXAudioTranscriptionConfig( """ # Use common utility to process the audio file processed_audio = process_audio_file(audio_file) - project_id = optional_params.get("project_id") or optional_params.get( - "watsonx_project" - ) + project_id = optional_params.get("project_id") or optional_params.get("watsonx_project") space_id = optional_params.get("space_id") # api_params = _get_api_params(params=optional_params, model=model) @@ -157,10 +151,7 @@ class IBMWatsonXAudioTranscriptionConfig( url = f"{url}/ml/v1/audio/transcriptions" # Add version parameter (only version in query string, not project_id) - api_version = ( - optional_params.get("api_version", None) - or litellm.WATSONX_DEFAULT_API_VERSION - ) + api_version = optional_params.get("api_version", None) or litellm.WATSONX_DEFAULT_API_VERSION url = f"{url}?version={api_version}" return url @@ -178,9 +169,7 @@ class IBMWatsonXAudioTranscriptionConfig( try: raw_response_json = raw_response.json() except Exception as e: - raise ValueError( - f"Error transforming response to json: {str(e)}\nResponse: {raw_response.text}" - ) + raise ValueError(f"Error transforming response to json: {str(e)}\nResponse: {raw_response.text}") # Extract only valid fields for TranscriptionResponse.__init__() # TranscriptionResponse only accepts 'text' and 'usage' in __init__() diff --git a/litellm/llms/watsonx/chat/transformation.py b/litellm/llms/watsonx/chat/transformation.py index 157493a4ce8..8c938e8dc4d 100644 --- a/litellm/llms/watsonx/chat/transformation.py +++ b/litellm/llms/watsonx/chat/transformation.py @@ -69,17 +69,13 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): optional_params["tool_choice_option"] = _tool_choice elif _tool_choice is not None: optional_params["tool_choice"] = _tool_choice - return super().map_openai_params( - non_default_params, optional_params, model, drop_params - ) + return super().map_openai_params(non_default_params, optional_params, model, drop_params) def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: api_base = api_base or get_secret_str("HOSTED_VLLM_API_BASE") # type: ignore - dynamic_api_key = ( - api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "" - ) # vllm does not require an api key + dynamic_api_key = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "" # vllm does not require an api key return api_base, dynamic_api_key def get_complete_url( @@ -95,29 +91,19 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): if model.startswith("deployment/"): deployment_id = "/".join(model.split("/")[1:]) endpoint = ( - WatsonXAIEndpoint.DEPLOYMENT_CHAT_STREAM.value - if stream - else WatsonXAIEndpoint.DEPLOYMENT_CHAT.value + WatsonXAIEndpoint.DEPLOYMENT_CHAT_STREAM.value if stream else WatsonXAIEndpoint.DEPLOYMENT_CHAT.value ) endpoint = endpoint.format(deployment_id=deployment_id) else: - endpoint = ( - WatsonXAIEndpoint.CHAT_STREAM.value - if stream - else WatsonXAIEndpoint.CHAT.value - ) + endpoint = WatsonXAIEndpoint.CHAT_STREAM.value if stream else WatsonXAIEndpoint.CHAT.value url = url.rstrip("/") + endpoint ## add api version - url = self._add_api_version_to_url( - url=url, api_version=optional_params.pop("api_version", None) - ) + url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) return url @staticmethod - def _apply_prompt_template_core( - model: str, messages: List[Dict[str, str]], hf_template_fn - ) -> Optional[str]: + def _apply_prompt_template_core(model: str, messages: List[Dict[str, str]], hf_template_fn) -> Optional[str]: """Core logic for applying prompt templates""" from litellm.litellm_core_utils.prompt_templates.factory import ( custom_prompt, @@ -169,9 +155,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): return None @staticmethod - async def aapply_prompt_template( - model: str, messages: List[Dict[str, str]] - ) -> Optional[str]: + async def aapply_prompt_template(model: str, messages: List[Dict[str, str]]) -> Optional[str]: """Apply prompt template (async version)""" import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -208,9 +192,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): # Log the exception for debugging but don't raise it # The caller will fall back to default prompt factory try: - verbose_logger.debug( - f"Failed to apply HuggingFace template for model {hf_model}: {e}" - ) + verbose_logger.debug(f"Failed to apply HuggingFace template for model {hf_model}: {e}") except Exception: # If logging fails, silently continue - don't break the flow pass @@ -237,9 +219,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): return None @staticmethod - def apply_prompt_template( - model: str, messages: List[Dict[str, str]] - ) -> Optional[str]: + def apply_prompt_template(model: str, messages: List[Dict[str, str]]) -> Optional[str]: """Apply prompt template (sync version)""" from litellm.litellm_core_utils.prompt_templates.factory import ( hf_chat_template, diff --git a/litellm/llms/watsonx/common_utils.py b/litellm/llms/watsonx/common_utils.py index 2aead929c8d..d1b065dbc6d 100644 --- a/litellm/llms/watsonx/common_utils.py +++ b/litellm/llms/watsonx/common_utils.py @@ -26,9 +26,7 @@ iam_token_cache = InMemoryCache() def get_watsonx_iam_url(): - return ( - get_secret_str("WATSONX_IAM_URL") or "https://iam.cloud.ibm.com/identity/token" - ) + return get_secret_str("WATSONX_IAM_URL") or "https://iam.cloud.ibm.com/identity/token" def generate_iam_token(api_key=None, **params) -> str: @@ -58,9 +56,7 @@ def generate_iam_token(api_key=None, **params) -> str: headers, data, ) - response = litellm.module_level_client.post( - url=iam_token_url, data=data, headers=headers - ) + response = litellm.module_level_client.post(url=iam_token_url, data=data, headers=headers) response.raise_for_status() json_data = response.json() @@ -99,16 +95,10 @@ def _get_api_params(params: dict, model: Optional[str] = None) -> WatsonXAPIPara # Load auth variables from environment variables if project_id is None: project_id = ( - get_secret_str("WATSONX_PROJECT_ID") - or get_secret_str("WX_PROJECT_ID") - or get_secret_str("PROJECT_ID") + get_secret_str("WATSONX_PROJECT_ID") or get_secret_str("WX_PROJECT_ID") or get_secret_str("PROJECT_ID") ) if region_name is None: - region_name = ( - get_secret_str("WATSONX_REGION") - or get_secret_str("WX_REGION") - or get_secret_str("REGION") - ) + region_name = get_secret_str("WATSONX_REGION") or get_secret_str("WX_REGION") or get_secret_str("REGION") if space_id is None: space_id = ( get_secret_str("WATSONX_DEPLOYMENT_SPACE_ID") @@ -117,12 +107,7 @@ def _get_api_params(params: dict, model: Optional[str] = None) -> WatsonXAPIPara or get_secret_str("SPACE_ID") ) - if ( - project_id is None - and space_id is None - and model is not None - and not model.startswith("deployment/") - ): + if project_id is None and space_id is None and model is not None and not model.startswith("deployment/"): raise WatsonXAIError( status_code=401, message="Error: Watsonx project_id and space_id not set. Set WX_PROJECT_ID or WX_SPACE_ID in environment variables or pass in as a parameter.", @@ -150,9 +135,7 @@ async def _aconvert_watsonx_messages_core( model_prompt_dict = custom_prompt_dict[model] return ptf.custom_prompt( messages=messages, - role_dict=model_prompt_dict.get( - "role_dict", model_prompt_dict.get("roles") - ), + role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")), initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""), final_prompt_value=model_prompt_dict.get("final_prompt_value", ""), bos_token=model_prompt_dict.get("bos_token", ""), @@ -166,9 +149,7 @@ async def _aconvert_watsonx_messages_core( if result: return result # Fallback to default - return ptf.prompt_factory( - model=model, messages=messages, custom_llm_provider="watsonx" - ) # type: ignore + return ptf.prompt_factory(model=model, messages=messages, custom_llm_provider="watsonx") # type: ignore def _convert_watsonx_messages_core( @@ -186,9 +167,7 @@ def _convert_watsonx_messages_core( model_prompt_dict = custom_prompt_dict[model] return ptf.custom_prompt( messages=messages, - role_dict=model_prompt_dict.get( - "role_dict", model_prompt_dict.get("roles") - ), + role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")), initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""), final_prompt_value=model_prompt_dict.get("final_prompt_value", ""), bos_token=model_prompt_dict.get("bos_token", ""), @@ -202,9 +181,7 @@ def _convert_watsonx_messages_core( if result: return result # Fallback to default - return ptf.prompt_factory( - model=model, messages=messages, custom_llm_provider="watsonx" - ) # type: ignore + return ptf.prompt_factory(model=model, messages=messages, custom_llm_provider="watsonx") # type: ignore async def aconvert_watsonx_messages_to_prompt( @@ -268,8 +245,7 @@ class IBMWatsonXMixin: ) zen_api_key = cast( Optional[str], - optional_params.pop("zen_api_key", None) - or get_secret_str("WATSONX_ZENAPIKEY"), + optional_params.pop("zen_api_key", None) or get_secret_str("WATSONX_ZENAPIKEY"), ) if token: headers["Authorization"] = f"Bearer {token}" @@ -306,9 +282,7 @@ class IBMWatsonXMixin: def get_error_class( self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers] ) -> BaseLLMException: - return WatsonXAIError( - status_code=status_code, message=error_message, headers=headers - ) + return WatsonXAIError(status_code=status_code, message=error_message, headers=headers) @staticmethod def get_watsonx_credentials( @@ -337,18 +311,14 @@ class IBMWatsonXMixin: wx_credentials = optional_params.pop( "wx_credentials", - optional_params.pop( - "watsonx_credentials", None - ), # follow {provider}_credentials, same as vertex ai + optional_params.pop("watsonx_credentials", None), # follow {provider}_credentials, same as vertex ai ) token: Optional[str] = None if wx_credentials is not None: api_base = wx_credentials.get("url", api_base) - api_key = wx_credentials.get( - "apikey", wx_credentials.get("api_key", api_key) - ) + api_key = wx_credentials.get("apikey", wx_credentials.get("api_key", api_key)) token = wx_credentials.get( "token", wx_credentials.get( @@ -365,9 +335,7 @@ class IBMWatsonXMixin: status_code=401, message="Error: Watsonx API base not set. Set WATSONX_API_BASE in environment variables or pass in as parameter - 'api_base='.", ) - return WatsonXCredentials( - api_key=api_key, api_base=api_base, token=cast(Optional[str], token) - ) + return WatsonXCredentials(api_key=api_key, api_base=api_base, token=cast(Optional[str], token)) def _prepare_payload(self, model: str, api_params: WatsonXAPIParams) -> dict: payload: dict = {} diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 8f418567371..190e2f7e93d 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -228,16 +228,12 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): "us-south", ] - def _build_request_payload( - self, model: str, prompt: str, optional_params: Dict - ) -> Dict: + def _build_request_payload(self, model: str, prompt: str, optional_params: Dict) -> Dict: """Shared logic to build request payload""" extra_body_params = optional_params.pop("extra_body", {}) optional_params.update(extra_body_params) watsonx_api_params = _get_api_params(params=optional_params, model=model) - watsonx_auth_payload = self._prepare_payload( - model=model, api_params=watsonx_api_params - ) + watsonx_auth_payload = self._prepare_payload(model=model, api_params=watsonx_api_params) return { "input": prompt, @@ -263,9 +259,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): prompt = await aconvert_watsonx_messages_to_prompt( model=model, messages=messages, provider=provider, custom_prompt_dict={} ) - return self._build_request_payload( - model=model, prompt=prompt, optional_params=optional_params - ) + return self._build_request_payload(model=model, prompt=prompt, optional_params=optional_params) def transform_request( self, @@ -280,9 +274,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): prompt = convert_watsonx_messages_to_prompt( model=model, messages=messages, provider=provider, custom_prompt_dict={} ) - return self._build_request_payload( - model=model, prompt=prompt, optional_params=optional_params - ) + return self._build_request_payload(model=model, prompt=prompt, optional_params=optional_params) def transform_response( self, @@ -318,17 +310,13 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): prompt_tokens = json_resp["results"][0]["input_token_count"] completion_tokens = json_resp["results"][0]["generated_token_count"] model_response.choices[0].message.content = generated_text # type: ignore - model_response.choices[0].finish_reason = map_finish_reason( - json_resp["results"][0]["stop_reason"] - ) + model_response.choices[0].finish_reason = map_finish_reason(json_resp["results"][0]["stop_reason"]) if json_resp.get("created_at"): try: created_datetime = datetime.fromisoformat(json_resp["created_at"]) except ValueError: # datetime.fromisoformat cannot handle 'Z' in Python 3.10 - created_datetime = datetime.fromisoformat( - f"{json_resp['created_at'].rstrip('Z')}+00:00" - ) + created_datetime = datetime.fromisoformat(f"{json_resp['created_at'].rstrip('Z')}+00:00") model_response.created = int(created_datetime.timestamp()) else: model_response.created = int(time.time()) @@ -360,17 +348,11 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): ) endpoint = endpoint.format(deployment_id=deployment_id) else: - endpoint = ( - WatsonXAIEndpoint.TEXT_GENERATION_STREAM - if stream - else WatsonXAIEndpoint.TEXT_GENERATION - ) + endpoint = WatsonXAIEndpoint.TEXT_GENERATION_STREAM if stream else WatsonXAIEndpoint.TEXT_GENERATION url = url.rstrip("/") + endpoint ## add api version - url = self._add_api_version_to_url( - url=url, api_version=optional_params.pop("api_version", None) - ) + url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) return url def get_model_response_iterator( diff --git a/litellm/llms/watsonx/embed/transformation.py b/litellm/llms/watsonx/embed/transformation.py index ae873d63fe9..a841ba9d3ad 100644 --- a/litellm/llms/watsonx/embed/transformation.py +++ b/litellm/llms/watsonx/embed/transformation.py @@ -47,9 +47,7 @@ class IBMWatsonXEmbeddingConfig(IBMWatsonXMixin, BaseEmbeddingConfig): inputs: list[str] = [input] elif isinstance(input, list): if len(input) > 0 and isinstance(input[0], (list, int)): - raise ValueError( - "WatsonX embeddings require a string or list of strings" - ) + raise ValueError("WatsonX embeddings require a string or list of strings") inputs = input else: inputs = [input] @@ -77,9 +75,7 @@ class IBMWatsonXEmbeddingConfig(IBMWatsonXMixin, BaseEmbeddingConfig): url = url.rstrip("/") + endpoint ## add api version - url = self._add_api_version_to_url( - url=url, api_version=optional_params.pop("api_version", None) - ) + url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) return url def transform_embedding_response( diff --git a/litellm/llms/watsonx/passthrough/transformation.py b/litellm/llms/watsonx/passthrough/transformation.py index 9162eef0e03..a89c72dbe10 100644 --- a/litellm/llms/watsonx/passthrough/transformation.py +++ b/litellm/llms/watsonx/passthrough/transformation.py @@ -54,16 +54,14 @@ class WatsonxPassthroughConfig(IBMWatsonXMixin, BasePassthroughConfig): ) -> Optional[str]: return ( api_key - or IBMWatsonXMixin.get_watsonx_credentials( - optional_params=dict(), api_base=None, api_key=api_key - )["api_key"] + or IBMWatsonXMixin.get_watsonx_credentials(optional_params=dict(), api_base=None, api_key=api_key)[ + "api_key" + ] ) @staticmethod def get_base_model(model: str) -> Optional[str]: return model - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: return super().get_models(api_key, api_base) diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 790606c7e6d..25b593f1c0a 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -42,9 +42,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): params = optional_params or {} - complete_url = self._add_api_version_to_url( - url=url, api_version=(params.get("api_version", None)) - ) + complete_url = self._add_api_version_to_url(url=url, api_version=(params.get("api_version", None))) return complete_url def get_supported_cohere_rerank_params(self, model: str) -> list: @@ -78,8 +76,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): ) zen_api_key = cast( str | None, - optional_params.pop("zen_api_key", None) - or get_secret_str("WATSONX_ZENAPIKEY"), + optional_params.pop("zen_api_key", None) or get_secret_str("WATSONX_ZENAPIKEY"), ) if token: headers["Authorization"] = f"Bearer {token}" @@ -115,21 +112,13 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): if k == "query" and v is not None: optional_rerank_params["query"] = v elif k == "documents" and v is not None: - optional_rerank_params["inputs"] = [ - {"text": el} if isinstance(el, str) else el for el in v - ] + optional_rerank_params["inputs"] = [{"text": el} if isinstance(el, str) else el for el in v] elif k == "top_n" and v is not None: - optional_rerank_params.setdefault("parameters", {}).setdefault( - "return_options", {} - )["top_n"] = v + optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["top_n"] = v elif k == "return_documents" and v is not None and isinstance(v, bool): - optional_rerank_params.setdefault("parameters", {}).setdefault( - "return_options", {} - )["inputs"] = v + optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["inputs"] = v elif k == "max_tokens_per_doc" and v is not None: - optional_rerank_params.setdefault("parameters", {})[ - "truncate_input_tokens" - ] = v + optional_rerank_params.setdefault("parameters", {})["truncate_input_tokens"] = v # IBM watsonx.ai require one of below parameters elif k == "project_id" and v is not None: @@ -200,11 +189,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): transformed_results.append(transformed_result) - response_id = ( - raw_response_json.get("id") - or raw_response_json.get("model_id") - or str(uuid.uuid4()) - ) + response_id = raw_response_json.get("id") or raw_response_json.get("model_id") or str(uuid.uuid4()) # Extract usage information _tokens = RerankTokens( diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 87b5757ef35..0e689549421 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -58,9 +58,7 @@ class XAIChatConfig(OpenAIGPTConfig): dynamic_api_key = XAIModelInfo.get_api_key(api_key) if should_use_xai_oauth(litellm_params) and not dynamic_api_key: try: - headers["Authorization"] = ( - f"Bearer {XAIOAuthAuthenticator().get_access_token()}" - ) + headers["Authorization"] = f"Bearer {XAIOAuthAuthenticator().get_access_token()}" except XAIOAuthError as exc: raise AuthenticationError( model=model, @@ -142,9 +140,7 @@ class XAIChatConfig(OpenAIGPTConfig): # reasoning check ######################################################### try: - if litellm.supports_reasoning( - model=model, custom_llm_provider=self.custom_llm_provider - ): + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): base_openai_params.append("reasoning_effort") except Exception as e: verbose_logger.debug(f"Error checking if model supports reasoning: {e}") @@ -222,9 +218,7 @@ class XAIChatConfig(OpenAIGPTConfig): Filter out 'name' from messages """ messages = strip_name_from_messages(messages) - return super().transform_request( - model, messages, optional_params, litellm_params, headers - ) + return super().transform_request(model, messages, optional_params, litellm_params, headers) @staticmethod def _fix_choice_finish_reason_for_tool_calls(choice: Choices) -> None: @@ -234,11 +228,7 @@ class XAIChatConfig(OpenAIGPTConfig): XAI API returns empty string for finish_reason when using tools, so we need to set it to "tool_calls" when tool_calls are present. """ - if ( - choice.finish_reason == "" - and choice.message.tool_calls - and len(choice.message.tool_calls) > 0 - ): + if choice.finish_reason == "" and choice.message.tool_calls and len(choice.message.tool_calls) > 0: choice.finish_reason = "tool_calls" def transform_response( @@ -345,9 +335,7 @@ class XAIChatConfig(OpenAIGPTConfig): return details = getattr(usage, "completion_tokens_details", None) - reasoning_tokens = ( - int(getattr(details, "reasoning_tokens", 0) or 0) if details else 0 - ) + reasoning_tokens = int(getattr(details, "reasoning_tokens", 0) or 0) if details else 0 if reasoning_tokens <= 0: return @@ -364,9 +352,7 @@ class XAIChatConfig(OpenAIGPTConfig): usage.completion_tokens = completion_tokens + reasoning_tokens - def _enhance_usage_with_xai_web_search_fields( - self, model_response: ModelResponse, raw_response_json: dict - ) -> None: + def _enhance_usage_with_xai_web_search_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None: """ Extract num_sources_used from X.AI response and map it to web_search_requests. """ diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index adc857894c5..0e499e33ed1 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -59,12 +59,7 @@ class XAIModelInfo(BaseLLMModelInfo): the provider-specific litellm.xai_key takes precedence over fallbacks. """ if legacy_generic_before_env: - return ( - api_key - or litellm.xai_key - or litellm.api_key - or get_secret_str("XAI_API_KEY") - ) + return api_key or litellm.xai_key or litellm.api_key or get_secret_str("XAI_API_KEY") return api_key or litellm.xai_key or get_secret_str("XAI_API_KEY") @@ -72,9 +67,7 @@ class XAIModelInfo(BaseLLMModelInfo): def get_base_model(model: str) -> Optional[str]: return model.replace("xai/", "") - def get_models( - self, api_key: Optional[str] = None, api_base: Optional[str] = None - ) -> List[str]: + def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> List[str]: api_base = self.get_api_base(api_base) api_key = self.get_api_key(api_key) if api_base is None or api_key is None: diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index 8edfd0c27ad..284400b0824 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -34,16 +34,10 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: total_tokens = int(getattr(usage, "total_tokens", 0) or 0) reasoning_tokens = 0 if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: - reasoning_tokens = int( - getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0 - ) + reasoning_tokens = int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0) already_normalised = total_tokens == prompt_tokens + completion_tokens - total_completion_tokens = ( - completion_tokens - if already_normalised - else completion_tokens + reasoning_tokens - ) + total_completion_tokens = completion_tokens if already_normalised else completion_tokens + reasoning_tokens modified_usage = Usage( prompt_tokens=usage.prompt_tokens, @@ -53,9 +47,7 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: completion_tokens_details=None, ) - prompt_cost, completion_cost = generic_cost_per_token( - model=model, usage=modified_usage, custom_llm_provider="xai" - ) + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=modified_usage, custom_llm_provider="xai") return prompt_cost, completion_cost diff --git a/litellm/llms/xai/oauth.py b/litellm/llms/xai/oauth.py index 30c717b7ca0..064e0ff77d6 100644 --- a/litellm/llms/xai/oauth.py +++ b/litellm/llms/xai/oauth.py @@ -62,9 +62,7 @@ class _CallbackHandler(BaseHTTPRequestHandler): self.send_response(400) self.send_header("Content-Type", "text/html; charset=utf-8") self.end_headers() - self.wfile.write( - b"

xAI authorization state mismatch.

" - ) + self.wfile.write(b"

xAI authorization state mismatch.

") return self.send_response(200) @@ -87,30 +85,18 @@ class _CallbackServer(HTTPServer): class XAIOAuthAuthenticator: - def __init__( - self, http_client: Optional[Union[httpx.Client, HTTPHandler]] = None - ) -> None: - self.token_dir = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser( - "~/.config/litellm/xai_oauth" - ) - self.auth_file = os.path.join( - self.token_dir, get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json" - ) + def __init__(self, http_client: Optional[Union[httpx.Client, HTTPHandler]] = None) -> None: + self.token_dir = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser("~/.config/litellm/xai_oauth") + self.auth_file = os.path.join(self.token_dir, get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json") self.http_client = http_client def get_api_base(self) -> str: - return ( - get_secret_str("XAI_OAUTH_API_BASE") - or get_secret_str("XAI_API_BASE") - or XAI_API_BASE - ) + return get_secret_str("XAI_OAUTH_API_BASE") or get_secret_str("XAI_API_BASE") or XAI_API_BASE def get_access_token(self) -> str: auth_data = self._read_auth_file() if not auth_data: - raise XAIOAuthLoginRequiredError( - "xAI OAuth login required. Run `litellm xai-oauth login`." - ) + raise XAIOAuthLoginRequiredError("xAI OAuth login required. Run `litellm xai-oauth login`.") access_token = auth_data.get("access_token") if access_token and not self._is_expired(auth_data): @@ -118,9 +104,7 @@ class XAIOAuthAuthenticator: refresh_token = auth_data.get("refresh_token") if not refresh_token: - raise XAIOAuthLoginRequiredError( - "xAI OAuth refresh token missing. Run `litellm xai-oauth login`." - ) + raise XAIOAuthLoginRequiredError("xAI OAuth refresh token missing. Run `litellm xai-oauth login`.") with _XAI_OAUTH_REFRESH_LOCK: locked_auth_data = self._read_auth_file() or auth_data @@ -156,9 +140,7 @@ class XAIOAuthAuthenticator: ) if no_browser or not webbrowser.open(authorize_url): - sys.stdout.write( - f"Open this URL to authenticate with xAI:\n{authorize_url}\n" - ) + sys.stdout.write(f"Open this URL to authenticate with xAI:\n{authorize_url}\n") sys.stdout.flush() result = self._wait_for_callback(server) @@ -245,9 +227,7 @@ class XAIOAuthAuthenticator: def _discover(self) -> Dict[str, str]: try: - response = self._client().get( - XAI_OAUTH_DISCOVERY_URL, headers={"Accept": "application/json"} - ) + response = self._client().get(XAI_OAUTH_DISCOVERY_URL, headers={"Accept": "application/json"}) response.raise_for_status() except httpx.HTTPStatusError as exc: raise XAIOAuthError( @@ -256,17 +236,13 @@ class XAIOAuthAuthenticator: try: data = response.json() except ValueError as exc: - raise XAIOAuthError( - "xAI OAuth discovery response was not valid JSON" - ) from exc + raise XAIOAuthError("xAI OAuth discovery response was not valid JSON") from exc authorization_endpoint = data.get("authorization_endpoint") token_endpoint = data.get("token_endpoint") if not authorization_endpoint or not token_endpoint: raise XAIOAuthError("xAI OAuth discovery missing endpoints") return { - "authorization_endpoint": self._validate_xai_endpoint( - authorization_endpoint - ), + "authorization_endpoint": self._validate_xai_endpoint(authorization_endpoint), "token_endpoint": self._validate_xai_endpoint(token_endpoint), } @@ -274,29 +250,19 @@ class XAIOAuthAuthenticator: parsed = urlparse(url) host = (parsed.hostname or "").lower() if parsed.scheme != "https" or (host != "x.ai" and not host.endswith(".x.ai")): - raise XAIOAuthError( - f"xAI OAuth discovery returned unexpected endpoint: {url}" - ) + raise XAIOAuthError(f"xAI OAuth discovery returned unexpected endpoint: {url}") return url def _pkce_pair(self) -> Tuple[str, str]: - verifier = ( - base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode() - ) - challenge = ( - base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) - .rstrip(b"=") - .decode() - ) + verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode() + challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() return verifier, challenge def _start_callback_server(self, state: str) -> Tuple[_CallbackServer, str]: last_error: Optional[OSError] = None for port in (XAI_OAUTH_REDIRECT_PORT, 0): try: - server = _CallbackServer( - (XAI_OAUTH_REDIRECT_HOST, port), _CallbackHandler - ) + server = _CallbackServer((XAI_OAUTH_REDIRECT_HOST, port), _CallbackHandler) server.expected_state = state server.callback_result = None actual_port = server.server_address[1] @@ -338,9 +304,7 @@ class XAIOAuthAuthenticator: server.server_close() raise XAIOAuthError("Timed out waiting for xAI OAuth callback") - def _exchange_token( - self, token_endpoint: str, data: Dict[str, str] - ) -> Dict[str, Any]: + def _exchange_token(self, token_endpoint: str, data: Dict[str, str]) -> Dict[str, Any]: try: response = self._client().post( token_endpoint, @@ -396,9 +360,7 @@ class XAIOAuthAuthenticator: token_endpoint = self._validate_xai_endpoint(token_endpoint) refresh_token = auth_data.get("refresh_token") if not refresh_token: - raise XAIOAuthLoginRequiredError( - "xAI OAuth refresh token missing. Run `litellm xai-oauth login`." - ) + raise XAIOAuthLoginRequiredError("xAI OAuth refresh token missing. Run `litellm xai-oauth login`.") token_payload = self._exchange_token( token_endpoint, diff --git a/litellm/llms/xai/realtime/transformation.py b/litellm/llms/xai/realtime/transformation.py index f7981dc547a..6d8a8948f06 100644 --- a/litellm/llms/xai/realtime/transformation.py +++ b/litellm/llms/xai/realtime/transformation.py @@ -105,10 +105,7 @@ class XAIRealtimeNormalizer: @staticmethod def _ensure_server_vad_create_response(turn_detection: dict) -> None: - if ( - turn_detection.get("type") == "server_vad" - and "create_response" not in turn_detection - ): + if turn_detection.get("type") == "server_vad" and "create_response" not in turn_detection: turn_detection["create_response"] = True # --------------------------------------------------------------------------- @@ -128,9 +125,7 @@ class XAIRealtimeNormalizer: if isinstance(part, dict): self._content_part_by_key[self._content_part_key(event)] = part - def _update_content_part_field( - self, event: dict, *, part_type: str, field: str, value: object - ) -> None: + def _update_content_part_field(self, event: dict, *, part_type: str, field: str, value: object) -> None: if value is None: return key = self._content_part_key(event) @@ -164,9 +159,7 @@ class XAIRealtimeNormalizer: return event if event_type == "response.output_text.done": - self._update_content_part_field( - event, part_type="text", field="text", value=event.get("text") - ) + self._update_content_part_field(event, part_type="text", field="text", value=event.get("text")) return event if event_type == "response.output_audio_transcript.done": @@ -250,9 +243,7 @@ class XAIRealtimeNormalizer: } @staticmethod - def _normalize_usage( - usage: object, *, empty_as_null: bool - ) -> Optional[dict[str, Any]]: + def _normalize_usage(usage: object, *, empty_as_null: bool) -> Optional[dict[str, Any]]: """Coerce a usage object into the full OpenAI GA shape. ``empty_as_null=True`` for ``response.created`` (usage optional). diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index f81e860a8ce..2773444bce9 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -51,9 +51,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return supported_params - def _transform_web_search_tool( - self, tool: Dict[str, Any] - ) -> Union[XAIWebSearchTool, Dict[str, Any]]: + def _transform_web_search_tool(self, tool: Dict[str, Any]) -> Union[XAIWebSearchTool, Dict[str, Any]]: """ Transform web_search tool to XAI format. @@ -92,9 +90,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return xai_tool - def _transform_x_search_tool( - self, tool: Dict[str, Any] - ) -> Union[XAIXSearchTool, Dict[str, Any]]: + def _transform_x_search_tool(self, tool: Dict[str, Any]) -> Union[XAIXSearchTool, Dict[str, Any]]: """ Transform x_search tool to XAI format. @@ -154,15 +150,11 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Drop instructions parameter (not supported by XAI) if "instructions" in params: - verbose_logger.debug( - "XAI Responses API does not support 'instructions' parameter. Dropping it." - ) + verbose_logger.debug("XAI Responses API does not support 'instructions' parameter. Dropping it.") params.pop("instructions") if "metadata" in params: - verbose_logger.debug( - "XAI Responses API does not support 'metadata' parameter. Dropping it." - ) + verbose_logger.debug("XAI Responses API does not support 'metadata' parameter. Dropping it.") params.pop("metadata") # Transform tools @@ -179,23 +171,17 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): if tool_type == "code_interpreter": # XAI supports code_interpreter but doesn't use the container field - verbose_logger.debug( - "XAI: Transforming code_interpreter tool, removing container field" - ) + verbose_logger.debug("XAI: Transforming code_interpreter tool, removing container field") transformed_tools.append({"type": "code_interpreter"}) elif tool_type == "web_search": # Transform web_search to XAI format - verbose_logger.debug( - "XAI: Transforming web_search tool to XAI format" - ) + verbose_logger.debug("XAI: Transforming web_search tool to XAI format") transformed_tools.append(self._transform_web_search_tool(tool)) elif tool_type == "x_search": # Transform x_search to XAI format - verbose_logger.debug( - "XAI: Transforming x_search tool to XAI format" - ) + verbose_logger.debug("XAI: Transforming x_search tool to XAI format") transformed_tools.append(self._transform_x_search_tool(tool)) else: @@ -208,18 +194,14 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return params - def validate_environment( - self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] - ) -> dict: + def validate_environment(self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]) -> dict: """ Validate environment and set up headers for XAI API. Uses the shared xAI key resolver with Responses API legacy precedence. """ litellm_params = litellm_params or GenericLiteLLMParams() - api_key = XAIModelInfo.get_api_key( - litellm_params.api_key, legacy_generic_before_env=True - ) + api_key = XAIModelInfo.get_api_key(litellm_params.api_key, legacy_generic_before_env=True) if not api_key: from litellm.llms.xai.oauth import ( @@ -264,18 +246,11 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): """ from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth - api_key = XAIModelInfo.get_api_key( - litellm_params.get("api_key"), legacy_generic_before_env=True - ) + api_key = XAIModelInfo.get_api_key(litellm_params.get("api_key"), legacy_generic_before_env=True) if should_use_xai_oauth(litellm_params) and not api_key: api_base = XAIOAuthAuthenticator().get_api_base() else: - api_base = ( - api_base - or litellm.api_base - or get_secret_str("XAI_API_BASE") - or XAI_API_BASE - ) + api_base = api_base or litellm.api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE # Remove trailing slashes api_base = api_base.rstrip("/") diff --git a/litellm/llms/xinference/image_generation/transformation.py b/litellm/llms/xinference/image_generation/transformation.py index 6ff70d0642d..0d2d890ddf4 100644 --- a/litellm/llms/xinference/image_generation/transformation.py +++ b/litellm/llms/xinference/image_generation/transformation.py @@ -13,9 +13,7 @@ class XInferenceImageGenerationConfig(BaseImageGenerationConfig): https://inference.readthedocs.io/en/v1.1.1/reference/generated/xinference.client.handlers.ImageModelHandle.text_to_image.html#xinference.client.handlers.ImageModelHandle.text_to_image """ - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + def get_supported_openai_params(self, model: str) -> List[OpenAIImageGenerationOptionalParams]: return ["n", "response_format", "size", "response_format"] def map_openai_params( diff --git a/litellm/llms/you_com/search/transformation.py b/litellm/llms/you_com/search/transformation.py index 0c7916e4c05..0cd825c3ab8 100644 --- a/litellm/llms/you_com/search/transformation.py +++ b/litellm/llms/you_com/search/transformation.py @@ -108,9 +108,7 @@ class YouComSearchConfig(BaseSearchConfig): api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/search") and not api_base.endswith( - "/v1/agents/search" - ): + if not api_base.endswith("/v1/search") and not api_base.endswith("/v1/agents/search"): api_base = f"{api_base}/v1/search" return api_base @@ -150,10 +148,7 @@ class YouComSearchConfig(BaseSearchConfig): result_data = dict(request_data) for param, value in optional_params.items(): - if ( - param not in self.get_supported_perplexity_optional_params() - and param not in result_data - ): + if param not in self.get_supported_perplexity_optional_params() and param not in result_data: result_data[param] = value return result_data diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index c932dcd2e03..fb1d67df357 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/litellm/llms/zai/chat/transformation.py @@ -48,9 +48,7 @@ class ZAIChatConfig(OpenAIGPTConfig): import litellm try: - if litellm.supports_reasoning( - model=model, custom_llm_provider=self.custom_llm_provider - ): + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): base_params.append("thinking") except Exception: pass diff --git a/litellm/main.py b/litellm/main.py index cdea7dbd684..18d2c367f8d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -374,9 +374,7 @@ class Completions: self.params[k] = v model = model or self.params.get("model") if self.router_obj is not None: - response = self.router_obj.completion( - model=model, messages=messages, **self.params - ) + response = self.router_obj.completion(model=model, messages=messages, **self.params) else: response = completion(model=model, messages=messages, **self.params) return response @@ -392,9 +390,7 @@ class AsyncCompletions: self.params[k] = v model = model or self.params.get("model") if self.router_obj is not None: - response = await self.router_obj.acompletion( - model=model, messages=messages, **self.params - ) + response = await self.router_obj.acompletion(model=model, messages=messages, **self.params) else: response = await acompletion(model=model, messages=messages, **self.params) return response @@ -433,9 +429,7 @@ async def acompletion( logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, deployment_id=None, - reasoning_effort: Optional[ - Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] - ] = None, + reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"]] = None, verbosity: Optional[Literal["low", "medium", "high"]] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, @@ -547,13 +541,9 @@ async def acompletion( # Log shared session usage if shared_session is not None: - verbose_logger.debug( - f"🔄 SHARED SESSION: acompletion called with shared_session (ID: {id(shared_session)})" - ) + verbose_logger.debug(f"🔄 SHARED SESSION: acompletion called with shared_session (ID: {id(shared_session)})") else: - verbose_logger.debug( - "🔄 NO SHARED SESSION: acompletion called without shared_session" - ) + verbose_logger.debug("🔄 NO SHARED SESSION: acompletion called without shared_session") # Adjusted to use explicit arguments instead of *args and **kwargs completion_kwargs = { @@ -609,9 +599,7 @@ async def acompletion( fallbacks = fallbacks or litellm.model_fallbacks if fallbacks is not None: - response = await async_completion_with_fallbacks( - **completion_kwargs, kwargs={"fallbacks": fallbacks, **kwargs} - ) + response = await async_completion_with_fallbacks(**completion_kwargs, kwargs={"fallbacks": fallbacks, **kwargs}) if response is None: raise Exception( "No response from fallbacks. Got none. Turn on `litellm.set_verbose=True` to see more details." @@ -641,9 +629,7 @@ async def acompletion( func_with_context = partial(ctx.run, func) init_response = await loop.run_in_executor(None, func_with_context) - if isinstance(init_response, dict) or isinstance( - init_response, ModelResponse - ): ## CACHING SCENARIO + if isinstance(init_response, dict) or isinstance(init_response, ModelResponse): ## CACHING SCENARIO if isinstance(init_response, dict): response = ModelResponse(**init_response) response = init_response @@ -738,45 +724,29 @@ def _handle_mock_potential_exceptions( raise litellm.MockException( status_code=getattr(mock_response, "status_code", 500), # type: ignore message=getattr(mock_response, "text", str(mock_response)), - llm_provider=getattr( - mock_response, "llm_provider", custom_llm_provider or "openai" - ), # type: ignore + llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"), # type: ignore model=model, # type: ignore request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), ) elif isinstance(mock_response, str) and mock_response == "litellm.RateLimitError": raise litellm.RateLimitError( message="this is a mock rate limit error", - llm_provider=getattr( - mock_response, "llm_provider", custom_llm_provider or "openai" - ), # type: ignore + llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"), # type: ignore model=model, ) - elif ( - isinstance(mock_response, str) - and mock_response == "litellm.ContextWindowExceededError" - ): + elif isinstance(mock_response, str) and mock_response == "litellm.ContextWindowExceededError": raise litellm.ContextWindowExceededError( message="this is a mock context window exceeded error", - llm_provider=getattr( - mock_response, "llm_provider", custom_llm_provider or "openai" - ), # type: ignore + llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"), # type: ignore model=model, ) - elif ( - isinstance(mock_response, str) - and mock_response == "litellm.InternalServerError" - ): + elif isinstance(mock_response, str) and mock_response == "litellm.InternalServerError": raise litellm.InternalServerError( message="this is a mock internal server error", - llm_provider=getattr( - mock_response, "llm_provider", custom_llm_provider or "openai" - ), # type: ignore + llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"), # type: ignore model=model, ) - elif isinstance(mock_response, str) and mock_response.startswith( - "Exception: content_filter_policy" - ): + elif isinstance(mock_response, str) and mock_response.startswith("Exception: content_filter_policy"): raise litellm.MockException( status_code=400, message=mock_response, @@ -892,9 +862,7 @@ def mock_completion( mock_response = cast( Union[str, dict, ModelResponse, ModelResponseStream], mock_response ) # after this point, mock_response is a string, dict, ModelResponse, or ModelResponseStream - if isinstance(mock_response, str) and mock_response.startswith( - "Exception: mock_streaming_error" - ): + if isinstance(mock_response, str) and mock_response.startswith("Exception: mock_streaming_error"): mock_response = litellm.MockException( message="This is a mock error raised mid-stream", llm_provider="anthropic", @@ -948,9 +916,7 @@ def mock_completion( for i in range(n): _choice = litellm.utils.Choices( index=i, - message=litellm.utils.Message( - content=mock_response, role="assistant" - ), + message=litellm.utils.Message(content=mock_response, role="assistant"), ) _all_choices.append(_choice) model_response.choices = _all_choices # type: ignore @@ -959,8 +925,7 @@ def mock_completion( if mock_tool_calls: model_response.choices[0].message.tool_calls = [ # type: ignore - ChatCompletionMessageToolCall(**tool_call) - for tool_call in mock_tool_calls + ChatCompletionMessageToolCall(**tool_call) for tool_call in mock_tool_calls ] setattr( @@ -969,8 +934,7 @@ def mock_completion( Usage( prompt_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, - total_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT - + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + total_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, ), ) @@ -1016,9 +980,7 @@ def responses_api_bridge_check( try: model_info = cast( dict, - _get_model_info_helper( - model=model, custom_llm_provider=custom_llm_provider - ), + _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider), ) if model_info.get("mode") is None and model.startswith("responses/"): model = model.replace("responses/", "") @@ -1032,9 +994,7 @@ def responses_api_bridge_check( except Exception as e: verbose_logger.debug("Error getting model info: {}".format(e)) - if model.startswith( - "responses/" - ): # handle azure models - `azure/responses/` + if model.startswith("responses/"): # handle azure models - `azure/responses/` model = model.replace("responses/", "") mode = "responses" model_info["mode"] = mode @@ -1052,10 +1012,7 @@ def responses_api_bridge_check( and OpenAIGPT5Config.is_model_gpt_5_model(model) and not OpenAIGPT5Config.is_model_gpt_5_search_model(model) and reasoning_effort is not None - and ( - reasoning_summary is not None - or (OpenAIGPT5Config.is_model_gpt_5_4_plus_model(model) and tools) - ) + and (reasoning_summary is not None or (OpenAIGPT5Config.is_model_gpt_5_4_plus_model(model) and tools)) ): model_info["mode"] = "responses" model = model.replace("responses/", "") @@ -1063,16 +1020,10 @@ def responses_api_bridge_check( return model_info, model -def _should_allow_input_examples( - custom_llm_provider: Optional[str], model: str -) -> bool: +def _should_allow_input_examples(custom_llm_provider: Optional[str], model: str) -> bool: if custom_llm_provider == "anthropic": return True - if ( - custom_llm_provider == "azure_ai" - or custom_llm_provider == "bedrock" - or custom_llm_provider == "vertex_ai" - ): + if custom_llm_provider == "azure_ai" or custom_llm_provider == "bedrock" or custom_llm_provider == "vertex_ai": return "claude" in model.lower() return False @@ -1149,10 +1100,7 @@ def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul timeout = ctx.timeout dynamic_params = False - if client is not None and ( - isinstance(client, openai.AzureOpenAI) - or isinstance(client, openai.AsyncAzureOpenAI) - ): + if client is not None and (isinstance(client, openai.AzureOpenAI) or isinstance(client, openai.AsyncAzureOpenAI)): dynamic_params = _check_dynamic_azure_params( azure_client_params={"api_version": api_version}, azure_client=client, @@ -1163,10 +1111,7 @@ def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - or litellm.AZURE_DEFAULT_API_VERSION + api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") or litellm.AZURE_DEFAULT_API_VERSION ) api_key = ( @@ -1177,9 +1122,9 @@ def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul or get_secret_str("AZURE_API_KEY") ) - azure_ad_token = optional_params.get("extra_body", {}).pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") + azure_ad_token = optional_params.get("extra_body", {}).pop("azure_ad_token", None) or get_secret_str( + "AZURE_AD_TOKEN" + ) azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) @@ -1293,9 +1238,7 @@ def _complete_azure_text(ctx: _CompletionDispatchContext) -> _CompletionDispatch "api_base is required for Azure OpenAI LLM provider. Either set it dynamically or set the AZURE_API_BASE environment variable." ) - api_version = ( - api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") - ) + api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") api_key = ( api_key @@ -1305,9 +1248,9 @@ def _complete_azure_text(ctx: _CompletionDispatchContext) -> _CompletionDispatch or get_secret_str("AZURE_API_KEY") ) - azure_ad_token = optional_params.get("extra_body", {}).pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") + azure_ad_token = optional_params.get("extra_body", {}).pop("azure_ad_token", None) or get_secret_str( + "AZURE_AD_TOKEN" + ) azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) @@ -1443,8 +1386,7 @@ def _complete_azure_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe api_base = AzureFoundryModelInfo.get_api_base(api_base) if api_base is None: raise ValueError( - "Azure AI Agents requests require an api_base. " - "Set `api_base` or the AZURE_AI_API_BASE env var." + "Azure AI Agents requests require an api_base. Set `api_base` or the AZURE_AI_API_BASE env var." ) api_key = AzureFoundryModelInfo.get_api_key(api_key) @@ -1469,8 +1411,7 @@ def _complete_azure_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe api_base = AzureFoundryModelInfo.get_api_base(api_base) if api_base is None: raise ValueError( - "Azure Anthropic requests require an api_base. " - "Set `api_base` or the AZURE_AI_API_BASE env var." + "Azure Anthropic requests require an api_base. Set `api_base` or the AZURE_AI_API_BASE env var." ) api_key = AzureFoundryModelInfo.get_api_key(api_key) @@ -1601,9 +1542,7 @@ def _complete_text_completion_openai( openai.api_version = None # set API KEY - api_key = ( - api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") headers = headers or litellm.headers @@ -1636,16 +1575,10 @@ def _complete_text_completion_openai( timeout=timeout, # type: ignore ) - if ( - optional_params.get("stream", False) is False - and acompletion is False - and text_completion is False - ): + if optional_params.get("stream", False) is False and acompletion is False and text_completion is False: # convert to chat completion response - _response = ( - litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( - response_object=_response, model_response_object=model_response - ) + _response = litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( + response_object=_response, model_response_object=model_response ) if optional_params.get("stream", False) or acompletion is True: @@ -2217,19 +2150,9 @@ def _complete_cometapi(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe stream = ctx.stream timeout = ctx.timeout - api_key = ( - api_key - or litellm.cometapi_key - or get_secret_str("COMETAPI_KEY") - or litellm.api_key - ) + api_key = api_key or litellm.cometapi_key or get_secret_str("COMETAPI_KEY") or litellm.api_key - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COMETAPI_API_BASE") - or "https://api.cometapi.com/v1" - ) + api_base = api_base or litellm.api_base or get_secret_str("COMETAPI_API_BASE") or "https://api.cometapi.com/v1" ## COMPLETION CALL response = base_llm_http_handler.completion( @@ -2278,12 +2201,7 @@ def _complete_minimax(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key - api_base = ( - api_base - or litellm.api_base - or get_secret_str("MINIMAX_API_BASE") - or "https://api.minimax.io/v1" - ) + api_base = api_base or litellm.api_base or get_secret_str("MINIMAX_API_BASE") or "https://api.minimax.io/v1" response = base_llm_http_handler.completion( model=model, @@ -2419,9 +2337,7 @@ def _complete_custom_openai( if extra_headers is not None: optional_params["extra_headers"] = extra_headers - if ( - litellm.enable_preview_features and metadata is not None - ): # [PREVIEW] allow metadata to be passed to OPENAI + if litellm.enable_preview_features and metadata is not None: # [PREVIEW] allow metadata to be passed to OPENAI openai_metadata = get_requester_metadata(metadata) if openai_metadata is not None: optional_params["metadata"] = openai_metadata @@ -2435,9 +2351,7 @@ def _complete_custom_openai( optional_params[k] = v ## COMPLETION CALL - use_base_llm_http_handler = get_secret_bool( - "EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER" - ) + use_base_llm_http_handler = get_secret_bool("EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER") try: if use_base_llm_http_handler: @@ -2522,12 +2436,7 @@ def _complete_mistral(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes timeout = ctx.timeout api_key = api_key or litellm.api_key or get_secret("MISTRAL_API_KEY") - api_base = ( - api_base - or litellm.api_base - or get_secret("MISTRAL_API_BASE") - or "https://api.mistral.ai/v1" - ) + api_base = api_base or litellm.api_base or get_secret("MISTRAL_API_BASE") or "https://api.mistral.ai/v1" return base_llm_http_handler.completion( model=model, @@ -2572,12 +2481,7 @@ def _complete_replicate(ctx: _CompletionDispatchContext) -> _CompletionDispatchR or get_secret("REPLICATE_API_TOKEN") ) - api_base = ( - api_base - or litellm.api_base - or get_secret("REPLICATE_API_BASE") - or "https://api.replicate.com/v1" - ) + api_base = api_base or litellm.api_base or get_secret("REPLICATE_API_BASE") or "https://api.replicate.com/v1" custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict @@ -2627,12 +2531,7 @@ def _complete_anthropic_text( stream = ctx.stream timeout = ctx.timeout - api_key = ( - api_key - or litellm.anthropic_key - or litellm.api_key - or os.environ.get("ANTHROPIC_API_KEY") - ) + api_key = api_key or litellm.anthropic_key or litellm.api_key or os.environ.get("ANTHROPIC_API_KEY") custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict api_base = cast( Optional[str], @@ -2645,16 +2544,10 @@ def _complete_anthropic_text( # Check if we should disable automatic URL suffix appending disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if ( - api_base is not None - and not disable_url_suffix - and not api_base.endswith("/v1/complete") - ): + if api_base is not None and not disable_url_suffix and not api_base.endswith("/v1/complete"): api_base += "/v1/complete" elif disable_url_suffix: - verbose_logger.debug( - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix" - ) + verbose_logger.debug("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix") return base_llm_http_handler.completion( model=model, @@ -2692,12 +2585,7 @@ def _complete_anthropic(ctx: _CompletionDispatchContext) -> _CompletionDispatchR optional_params = ctx.optional_params timeout = ctx.timeout - api_key = ( - api_key - or litellm.anthropic_key - or litellm.api_key - or os.environ.get("ANTHROPIC_API_KEY") - ) + api_key = api_key or litellm.anthropic_key or litellm.api_key or os.environ.get("ANTHROPIC_API_KEY") custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict # call /messages # default route for all anthropic models @@ -2712,16 +2600,10 @@ def _complete_anthropic(ctx: _CompletionDispatchContext) -> _CompletionDispatchR # Check if we should disable automatic URL suffix appending disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if ( - api_base is not None - and not disable_url_suffix - and not api_base.endswith("/v1/messages") - ): + if api_base is not None and not disable_url_suffix and not api_base.endswith("/v1/messages"): api_base += "/v1/messages" elif disable_url_suffix: - verbose_logger.debug( - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix" - ) + verbose_logger.debug("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix") response = anthropic_chat_completions.completion( model=model, @@ -2764,19 +2646,9 @@ def _complete_nlp_cloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchR model_response = ctx.model_response optional_params = ctx.optional_params - nlp_cloud_key = ( - api_key - or litellm.nlp_cloud_key - or get_secret("NLP_CLOUD_API_KEY") - or litellm.api_key - ) + nlp_cloud_key = api_key or litellm.nlp_cloud_key or get_secret("NLP_CLOUD_API_KEY") or litellm.api_key - api_base = ( - api_base - or litellm.api_base - or get_secret("NLP_CLOUD_API_BASE") - or "https://api.nlpcloud.io/v1/gpu/" - ) + api_base = api_base or litellm.api_base or get_secret("NLP_CLOUD_API_BASE") or "https://api.nlpcloud.io/v1/gpu/" response = nlp_cloud_chat_completion( model=model, @@ -2832,10 +2704,7 @@ def _complete_aleph_alpha(ctx: _CompletionDispatchContext) -> _CompletionDispatc ) api_base = ( - api_base - or litellm.api_base - or get_secret("ALEPH_ALPHA_API_BASE") - or "https://api.aleph-alpha.com/complete" + api_base or litellm.api_base or get_secret("ALEPH_ALPHA_API_BASE") or "https://api.aleph-alpha.com/complete" ) model_response = aleph_alpha.completion( @@ -2893,22 +2762,12 @@ def _complete_cohere_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatc verbose_logger.debug(f"Cohere route: {cohere_route}") # Set API base based on route if cohere_route == "v2": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COHERE_API_BASE") - or "https://api.cohere.com/v2/chat" - ) + api_base = api_base or litellm.api_base or get_secret_str("COHERE_API_BASE") or "https://api.cohere.com/v2/chat" # Remove v2/ prefix from model name for the actual API call if "v2/" in model: model = model.replace("v2/", "") else: - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COHERE_API_BASE") - or "https://api.cohere.ai/v1/chat" - ) + api_base = api_base or litellm.api_base or get_secret_str("COHERE_API_BASE") or "https://api.cohere.ai/v1/chat" headers = headers or litellm.headers or {} if headers is None: @@ -2951,19 +2810,9 @@ def _complete_maritalk(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe model_response = ctx.model_response optional_params = ctx.optional_params - maritalk_key = ( - api_key - or litellm.maritalk_key - or get_secret("MARITALK_API_KEY") - or litellm.api_key - ) + maritalk_key = api_key or litellm.maritalk_key or get_secret("MARITALK_API_KEY") or litellm.api_key - api_base = ( - api_base - or litellm.api_base - or get_secret("MARITALK_API_BASE") - or "https://chat.maritaca.ai/api" - ) + api_base = api_base or litellm.api_base or get_secret("MARITALK_API_BASE") or "https://chat.maritaca.ai/api" return openai_like_chat_completion.completion( model=model, @@ -2996,17 +2845,9 @@ def _complete_amazon_nova(ctx: _CompletionDispatchContext) -> _CompletionDispatc optional_params = ctx.optional_params timeout = ctx.timeout - api_key = ( - api_key - or litellm.amazon_nova_api_key - or get_secret_str("AMAZON_NOVA_API_KEY") - or litellm.api_key - ) + api_key = api_key or litellm.amazon_nova_api_key or get_secret_str("AMAZON_NOVA_API_KEY") or litellm.api_key api_base = ( - api_base - or litellm.api_base - or get_secret_str("AMAZON_NOVA_API_BASE") - or "https://api.nova.amazon.com/v1" + api_base or litellm.api_base or get_secret_str("AMAZON_NOVA_API_BASE") or "https://api.nova.amazon.com/v1" ) return openai_like_chat_completion.completion( model=model, @@ -3305,12 +3146,7 @@ def _complete_openrouter(ctx: _CompletionDispatchContext) -> _CompletionDispatch stream = ctx.stream timeout = ctx.timeout - api_base = ( - api_base - or litellm.api_base - or get_secret_str("OPENROUTER_API_BASE") - or "https://openrouter.ai/api/v1" - ) + api_base = api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" api_key = ( api_key @@ -3366,9 +3202,7 @@ def _complete_openrouter(ctx: _CompletionDispatchContext) -> _CompletionDispatch client=client, ) ## LOGGING - logging.post_call( - input=messages, api_key=openai.api_key, original_response=response - ) + logging.post_call(input=messages, api_key=openai.api_key, original_response=response) return response @@ -3446,9 +3280,7 @@ def _complete_vercel_ai_gateway( client=client, ) ## LOGGING - logging.post_call( - input=messages, api_key=openai.api_key, original_response=response - ) + logging.post_call(input=messages, api_key=openai.api_key, original_response=response) return response @@ -3696,11 +3528,7 @@ def _complete_vertex_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion=acompletion, ) - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): + if "stream" in optional_params and optional_params["stream"] is True and acompletion is False: return CustomStreamWrapper( model_response, model, @@ -3744,12 +3572,7 @@ def _complete_predibase(ctx: _CompletionDispatchContext) -> _CompletionDispatchR or get_secret("PREDIBASE_API_BASE") ) - api_key = ( - api_key - or litellm.api_key - or litellm.predibase_key - or get_secret("PREDIBASE_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.predibase_key or get_secret("PREDIBASE_API_KEY") _model_response = predibase_chat_completions.completion( model=model, @@ -3769,11 +3592,7 @@ def _complete_predibase(ctx: _CompletionDispatchContext) -> _CompletionDispatchR timeout=timeout, ) - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): + if "stream" in optional_params and optional_params["stream"] is True and acompletion is False: return _model_response return _model_response @@ -3823,11 +3642,7 @@ def _complete_text_completion_codestral( timeout=timeout, ) - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): + if "stream" in optional_params and optional_params["stream"] is True and acompletion is False: return _model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract return _model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract @@ -3850,16 +3665,8 @@ def _complete_text_completion_inception( text_completion = ctx.text_completion timeout = ctx.timeout - passed_api_base = ( - api_base - or optional_params.pop("api_base", None) - or optional_params.pop("base_url", None) - ) - api_base = ( - passed_api_base - or get_secret_str("INCEPTION_API_BASE") - or "https://api.inceptionlabs.ai/v1" - ) + passed_api_base = api_base or optional_params.pop("api_base", None) or optional_params.pop("base_url", None) + api_base = passed_api_base or get_secret_str("INCEPTION_API_BASE") or "https://api.inceptionlabs.ai/v1" # FIM is served at `/v1/fim/completions`; the OpenAI client appends # `/completions`, so point it at the `/v1/fim` base. api_base = api_base.rstrip("/") @@ -3870,9 +3677,7 @@ def _complete_text_completion_inception( # api_base; only resolve it for the default/server base, or when the # caller passes their own key. if passed_api_base is None or api_key: - api_key = ( - api_key or litellm.inception_key or get_secret_str("INCEPTION_API_KEY") - ) + api_key = api_key or litellm.inception_key or get_secret_str("INCEPTION_API_KEY") _response = openai_text_completions.completion( model=model, @@ -3891,15 +3696,9 @@ def _complete_text_completion_inception( timeout=timeout, # type: ignore ) - if ( - optional_params.get("stream", False) is False - and acompletion is False - and text_completion is False - ): - _response = ( - litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( - response_object=_response, model_response_object=model_response - ) + if optional_params.get("stream", False) is False and acompletion is False and text_completion is False: + _response = litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( + response_object=_response, model_response_object=model_response ) if optional_params.get("stream", False) or acompletion is True: @@ -4012,10 +3811,7 @@ def _complete_bedrock(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes optional_params["aws_secret_access_key"] = creds.secret_key if creds.token: optional_params["aws_session_token"] = creds.token - if ( - "aws_region_name" not in optional_params - or optional_params["aws_region_name"] is None - ): + if "aws_region_name" not in optional_params or optional_params["aws_region_name"] is None: optional_params["aws_region_name"] = aws_bedrock_client.meta.region_name bedrock_route = BedrockModelInfo.get_bedrock_route(model) @@ -4181,9 +3977,7 @@ def _complete_watsonx_text( wx_credentials = optional_params.pop( "wx_credentials", - optional_params.pop( - "watsonx_credentials", None - ), # follow {provider}_credentials, same as vertex ai + optional_params.pop("watsonx_credentials", None), # follow {provider}_credentials, same as vertex ai ) token: Optional[str] = None @@ -4273,12 +4067,7 @@ def _complete_ollama(ctx: _CompletionDispatchContext) -> _CompletionDispatchResu stream = ctx.stream timeout = ctx.timeout - api_base = ( - litellm.api_base - or api_base - or get_secret("OLLAMA_API_BASE") - or "http://localhost:11434" - ) + api_base = litellm.api_base or api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" if api_key is not None and "Authorization" not in headers: headers["Authorization"] = f"Bearer {api_key}" @@ -4318,19 +4107,9 @@ def _complete_ollama_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatc stream = ctx.stream timeout = ctx.timeout - api_base = ( - litellm.api_base - or api_base - or get_secret("OLLAMA_API_BASE") - or "http://localhost:11434" - ) + api_base = litellm.api_base or api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" - api_key = ( - api_key - or litellm.ollama_key - or os.environ.get("OLLAMA_API_KEY") - or litellm.api_key - ) + api_key = api_key or litellm.ollama_key or os.environ.get("OLLAMA_API_KEY") or litellm.api_key if api_key is not None and "Authorization" not in headers: headers["Authorization"] = f"Bearer {api_key}" @@ -4406,12 +4185,7 @@ def _complete_cloudflare(ctx: _CompletionDispatchContext) -> _CompletionDispatch stream = ctx.stream timeout = ctx.timeout - api_key = ( - api_key - or litellm.cloudflare_api_key - or litellm.api_key - or get_secret("CLOUDFLARE_API_KEY") - ) + api_key = api_key or litellm.cloudflare_api_key or litellm.api_key or get_secret("CLOUDFLARE_API_KEY") api_base = api_base or litellm.api_base or get_secret("CLOUDFLARE_API_BASE") custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict @@ -4578,12 +4352,7 @@ def _complete_bytez(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul stream = ctx.stream timeout = ctx.timeout - api_key = ( - api_key - or litellm.bytez_key - or get_secret_str("BYTEZ_API_KEY") - or litellm.api_key - ) + api_key = api_key or litellm.bytez_key or get_secret_str("BYTEZ_API_KEY") or litellm.api_key response = base_llm_http_handler.completion( model=model, @@ -4625,12 +4394,7 @@ def _complete_lemonade(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe stream = ctx.stream timeout = ctx.timeout - api_key = ( - api_key - or litellm.lemonade_key - or get_secret_str("LEMONADE_API_KEY") - or litellm.api_key - ) + api_key = api_key or litellm.lemonade_key or get_secret_str("LEMONADE_API_KEY") or litellm.api_key response = base_llm_http_handler.completion( model=model, @@ -4672,12 +4436,7 @@ def _complete_ovhcloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe stream = ctx.stream timeout = ctx.timeout - api_key = ( - api_key - or litellm.ovhcloud_key - or get_secret_str("OVHCLOUD_API_KEY") - or litellm.api_key - ) + api_key = api_key or litellm.ovhcloud_key or get_secret_str("OVHCLOUD_API_KEY") or litellm.api_key api_base = ( api_base @@ -4723,9 +4482,7 @@ def _complete_custom(ctx: _CompletionDispatchContext) -> _CompletionDispatchResu url = litellm.api_base or api_base or "" if url is None or url == "": - raise ValueError( - "api_base not set. Set api_base or litellm.api_base for custom endpoints" - ) + raise ValueError("api_base not set. Set api_base or litellm.api_base for custom endpoints") """ assume input to custom LLM api bases follow this format: @@ -4808,14 +4565,10 @@ def _complete_custom_providers( custom_handler = item["custom_handler"] if custom_handler is None: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) + raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) ## ROUTE LLM CALL ## - handler_fn = custom_chat_llm_router( - async_fn=acompletion, stream=stream, custom_llm=custom_handler - ) + handler_fn = custom_chat_llm_router(async_fn=acompletion, stream=stream, custom_llm=custom_handler) headers = headers or litellm.headers or {} @@ -4970,9 +4723,7 @@ def completion( # type: ignore logit_bias: Optional[dict] = None, user: Optional[str] = None, # openai v1.0+ new params - reasoning_effort: Optional[ - Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] - ] = None, + reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"]] = None, verbosity: Optional[Literal["low", "medium", "high"]] = None, response_format: Optional[Union[dict, Type[BaseModel]]] = None, seed: Optional[int] = None, @@ -5081,9 +4832,7 @@ def completion( # type: ignore # Check if MCP tools are present (following responses pattern) # Cast tools to Optional[Iterable[ToolParam]] for type checking tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools) - if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( - tools=tools_for_mcp - ): + if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools_for_mcp): return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it model=model, messages=messages, @@ -5142,17 +4891,11 @@ def completion( # type: ignore model_info = kwargs.get("model_info", None) proxy_server_request = kwargs.get("proxy_server_request", None) fallbacks = kwargs.get("fallbacks", None) - provider_specific_header = cast( - Optional[ProviderSpecificHeader], kwargs.get("provider_specific_header", None) - ) + provider_specific_header = cast(Optional[ProviderSpecificHeader], kwargs.get("provider_specific_header", None)) headers = kwargs.get("headers", None) or extra_headers - ensure_alternating_roles: Optional[bool] = kwargs.get( - "ensure_alternating_roles", None - ) - user_continue_message: Optional[ChatCompletionUserMessage] = kwargs.get( - "user_continue_message", None - ) + ensure_alternating_roles: Optional[bool] = kwargs.get("ensure_alternating_roles", None) + user_continue_message: Optional[ChatCompletionUserMessage] = kwargs.get("user_continue_message", None) assistant_continue_message: Optional[ChatCompletionAssistantMessage] = kwargs.get( "assistant_continue_message", None ) @@ -5194,9 +4937,7 @@ def completion( # type: ignore model_info.get("base_model") if isinstance(model_info, dict) else None ) ### DISABLE FLAGS ### - disable_add_transform_inline_image_block = kwargs.get( - "disable_add_transform_inline_image_block", None - ) + disable_add_transform_inline_image_block = kwargs.get("disable_add_transform_inline_image_block", None) ### TEXT COMPLETION CALLS ### text_completion = kwargs.get("text_completion", False) atext_completion = kwargs.get("atext_completion", False) @@ -5260,9 +5001,7 @@ def completion( # type: ignore **args ) if model_list is not None: - deployments = [ - m["litellm_params"] for m in model_list if m["model_name"] == model - ] + deployments = [m["litellm_params"] for m in model_list if m["model_name"] == model] return litellm.batch_completion_models( # pyright: ignore[reportReturnType] # batch path returns a list of responses, outside completion()'s single-response return type deployments=deployments, **args ) @@ -5279,18 +5018,14 @@ def completion( # type: ignore if deployment_id is not None: # azure llms model = deployment_id custom_llm_provider = "azure" - _supplemental_provider_params = { - k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs - } + _supplemental_provider_params = {k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs} model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key, litellm_params=( - GenericLiteLLMParams(**_supplemental_provider_params) - if _supplemental_provider_params - else None + GenericLiteLLMParams(**_supplemental_provider_params) if _supplemental_provider_params else None ), ) @@ -5301,9 +5036,7 @@ def completion( # type: ignore web_search_options=web_search_options, ) - if not _should_allow_input_examples( - custom_llm_provider=custom_llm_provider, model=model - ): + if not _should_allow_input_examples(custom_llm_provider=custom_llm_provider, model=model): tools = _drop_input_examples_from_tools(tools=tools) if provider_specific_header is not None: @@ -5344,13 +5077,7 @@ def completion( # type: ignore ) ### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ### custom_prompt_dict = {} # type: ignore - if ( - initial_prompt_value - or roles - or final_prompt_value - or bos_token - or eos_token - ): + if initial_prompt_value or roles or final_prompt_value or bos_token or eos_token: custom_prompt_dict = {model: {}} if initial_prompt_value: custom_prompt_dict[model]["initial_prompt_value"] = initial_prompt_value @@ -5373,9 +5100,7 @@ def completion( # type: ignore ) provider_config: Optional[BaseConfig] = None - if custom_llm_provider is not None and custom_llm_provider in [ - provider.value for provider in LlmProviders - ]: + if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: provider_config = ProviderConfigManager.get_provider_chat_config( model=model, provider=LlmProviders(custom_llm_provider), @@ -5383,9 +5108,7 @@ def completion( # type: ignore ) if provider_config is not None: - messages = provider_config.translate_developer_role_to_system_role( - messages=messages - ) + messages = provider_config.translate_developer_role_to_system_role(messages=messages) if ( supports_system_message is not None @@ -5441,9 +5164,7 @@ def completion( # type: ignore "allowed_openai_params": kwargs.get("allowed_openai_params"), "base_model": base_model, } - optional_params = get_optional_params( - **optional_param_args, **non_default_params - ) + optional_params = get_optional_params(**optional_param_args, **non_default_params) processed_non_default_params = pre_process_non_default_params( model=model, passed_params=optional_param_args, @@ -5455,16 +5176,11 @@ def completion( # type: ignore provider_config=provider_config, ) - if ( - litellm.add_function_to_prompt - and optional_params.get("functions_unsupported_model", None) + if litellm.add_function_to_prompt and optional_params.get( + "functions_unsupported_model", None ): # if user opts to add it to prompt, when API doesn't support function calling - functions_unsupported_model = optional_params.pop( - "functions_unsupported_model" - ) - messages = function_call_prompt( - messages=messages, functions=functions_unsupported_model - ) + functions_unsupported_model = optional_params.pop("functions_unsupported_model") + messages = function_call_prompt(messages=messages, functions=functions_unsupported_model) # For logging - save the values of the litellm-specific params passed in litellm_params = get_litellm_params( @@ -5502,9 +5218,7 @@ def completion( # type: ignore prompt_id=prompt_id, prompt_variables=prompt_variables, ssl_verify=ssl_verify, - merge_reasoning_content_in_choices=kwargs.get( - "merge_reasoning_content_in_choices", None - ), + merge_reasoning_content_in_choices=kwargs.get("merge_reasoning_content_in_choices", None), use_litellm_proxy=kwargs.get("use_litellm_proxy", False), api_version=api_version, azure_ad_token=kwargs.get("azure_ad_token"), @@ -5567,15 +5281,10 @@ def completion( # type: ignore # detection when the deployment name differs from the model name. _azure_detection_model = base_model or model - if ( - responses_api_model_info.get("mode") == "responses" - and not skip_responses_api_bridge - ): + if responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge: from litellm.completion_extras import responses_api_bridge - optional_params, rs_val = ( - strip_reasoning_summary_aliases_from_optional_params(optional_params) - ) + optional_params, rs_val = strip_reasoning_summary_aliases_from_optional_params(optional_params) if isinstance(reasoning_effort, dict) and "summary" in reasoning_effort: optional_params["reasoning_effort"] = reasoning_effort @@ -5608,18 +5317,11 @@ def completion( # type: ignore encoding=_get_encoding(), stream=stream, ) - elif ( - custom_llm_provider == "openai" - and OpenAIGPT5Config.is_model_gpt_5_model(model) - ) or ( + elif (custom_llm_provider == "openai" and OpenAIGPT5Config.is_model_gpt_5_model(model)) or ( custom_llm_provider == "azure" - and litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( - _azure_detection_model - ) + and litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(_azure_detection_model) ): - optional_params, _ = strip_reasoning_summary_aliases_from_optional_params( - optional_params - ) + optional_params, _ = strip_reasoning_summary_aliases_from_optional_params(optional_params) _dispatch_ctx = _CompletionDispatchContext( _azure_detection_model=_azure_detection_model, @@ -5671,8 +5373,7 @@ def completion( # type: ignore custom_llm_provider == "text-completion-openai" or "ft:babbage-002" in model or "ft:davinci-002" in model # support for finetuned completion models - or custom_llm_provider - in litellm.openai_text_completion_compatible_providers + or custom_llm_provider in litellm.openai_text_completion_compatible_providers and kwargs.get("text_completion") is True ): response = _complete_text_completion_openai(_dispatch_ctx) @@ -5728,9 +5429,7 @@ def completion( # type: ignore or custom_llm_provider == "wandb" or custom_llm_provider == "clarifai" or custom_llm_provider in litellm.openai_compatible_providers - or JSONProviderRegistry.exists( - custom_llm_provider - ) # JSON-configured providers + or JSONProviderRegistry.exists(custom_llm_provider) # JSON-configured providers or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo ): # allow user to make an openai call with a custom base # note: if a user sets a custom base - we should ensure this works @@ -5739,18 +5438,10 @@ def completion( # type: ignore elif custom_llm_provider == "mistral": response = _complete_mistral(_dispatch_ctx) - elif ( - "replicate" in model - or custom_llm_provider == "replicate" - or model in litellm.replicate_models - ): + elif "replicate" in model or custom_llm_provider == "replicate" or model in litellm.replicate_models: # Setting the relevant API KEY for replicate, replicate defaults to using os.environ.get("REPLICATE_API_TOKEN") response = _complete_replicate(_dispatch_ctx) - elif ( - "clarifai" in model - or custom_llm_provider == "clarifai" - or model in litellm.clarifai_models - ): + elif "clarifai" in model or custom_llm_provider == "clarifai" or model in litellm.clarifai_models: pass # Deprecated - handled in the openai compatible provider section above elif custom_llm_provider == "anthropic_text": response = _complete_anthropic_text(_dispatch_ctx) @@ -5853,9 +5544,7 @@ def completion( # type: ignore elif custom_llm_provider == "custom": response = _complete_custom(_dispatch_ctx) - elif ( - custom_llm_provider in litellm._custom_providers - ): # Assume custom LLM provider + elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider # Get the Custom Handler response = _complete_custom_providers(_dispatch_ctx) @@ -5868,9 +5557,7 @@ def completion( # type: ignore response = _complete_langflow(_dispatch_ctx) else: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) + raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) return response except Exception as e: ## Map to OpenAI Exception @@ -5890,9 +5577,7 @@ def completion_with_retries(*args, **kwargs): try: import tenacity except Exception as e: - raise Exception( - f"tenacity import failed please run `pip install tenacity`. Error{e}" - ) + raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") num_retries = kwargs.pop("num_retries", 3) # reset retries in .completion() @@ -5909,9 +5594,7 @@ def completion_with_retries(*args, **kwargs): reraise=True, ) else: - retryer = tenacity.Retrying( - stop=tenacity.stop_after_attempt(num_retries), reraise=True - ) + retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) return retryer(original_function, *args, **kwargs) @@ -5923,9 +5606,7 @@ async def acompletion_with_retries(*args, **kwargs): try: import tenacity except Exception as e: - raise Exception( - f"tenacity import failed please run `pip install tenacity`. Error{e}" - ) + raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") num_retries = kwargs.pop("num_retries", 3) kwargs["max_retries"] = 0 @@ -5939,9 +5620,7 @@ async def acompletion_with_retries(*args, **kwargs): reraise=True, ) else: - retryer = tenacity.AsyncRetrying( - stop=tenacity.stop_after_attempt(num_retries), reraise=True - ) + retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) return await retryer(original_function, *args, **kwargs) @@ -5952,9 +5631,7 @@ def responses_with_retries(*args, **kwargs): try: import tenacity except Exception as e: - raise Exception( - f"tenacity import failed please run `pip install tenacity`. Error{e}" - ) + raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") from litellm.responses.main import responses @@ -5973,9 +5650,7 @@ def responses_with_retries(*args, **kwargs): reraise=True, ) else: - retryer = tenacity.Retrying( - stop=tenacity.stop_after_attempt(num_retries), reraise=True - ) + retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) return retryer(original_function, *args, **kwargs) @@ -5986,9 +5661,7 @@ async def aresponses_with_retries(*args, **kwargs): try: import tenacity except Exception as e: - raise Exception( - f"tenacity import failed please run `pip install tenacity`. Error{e}" - ) + raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") from litellm.responses.main import aresponses @@ -6004,9 +5677,7 @@ async def aresponses_with_retries(*args, **kwargs): reraise=True, ) else: - retryer = tenacity.AsyncRetrying( - stop=tenacity.stop_after_attempt(num_retries), reraise=True - ) + retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) return await retryer(original_function, *args, **kwargs) @@ -6052,17 +5723,11 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: response = init_response elif asyncio.iscoroutine(init_response): response = await init_response # type: ignore - if ( - response is not None - and isinstance(response, EmbeddingResponse) - and hasattr(response, "_hidden_params") - ): + if response is not None and isinstance(response, EmbeddingResponse) and hasattr(response, "_hidden_params"): response._hidden_params["custom_llm_provider"] = custom_llm_provider if response is None: - raise ValueError( - "Unable to get Embedding Response. Please pass a valid llm_provider." - ) + raise ValueError("Unable to get Embedding Response. Please pass a valid llm_provider.") return response except Exception as e: custom_llm_provider = custom_llm_provider or "openai" @@ -6236,9 +5901,7 @@ def embedding( if dynamic_api_key is not None: api_key = dynamic_api_key - allowed_openai_params: Optional[List[str]] = kwargs.get( - "allowed_openai_params", None - ) + allowed_openai_params: Optional[List[str]] = kwargs.get("allowed_openai_params", None) optional_params = get_optional_params_embeddings( model=model, user=user, @@ -6250,9 +5913,7 @@ def embedding( ) ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### - if ( - input_cost_per_token is not None and output_cost_per_token is not None - ) or input_cost_per_second is not None: + if (input_cost_per_token is not None and output_cost_per_token is not None) or input_cost_per_second is not None: litellm.register_model( { f"{custom_llm_provider}/{model}": _build_custom_pricing_entry( @@ -6277,9 +5938,7 @@ def embedding( if mock_response is not None: return mock_embedding(model=model, mock_response=mock_response) try: - response: Optional[ - Union[EmbeddingResponse, Coroutine[Any, Any, EmbeddingResponse]] - ] = None + response: Optional[Union[EmbeddingResponse, Coroutine[Any, Any, EmbeddingResponse]]] = None if azure is True or custom_llm_provider == "azure": # azure configs @@ -6293,21 +5952,12 @@ def embedding( or litellm.AZURE_DEFAULT_API_VERSION ) - azure_ad_token = optional_params.pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") + azure_ad_token = optional_params.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") - api_key = ( - api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.azure_key or get_secret_str("AZURE_API_KEY") if api_base is None: - raise ValueError( - "No API Base provided for Azure OpenAI LLM provider. Set 'AZURE_API_BASE' in .env" - ) + raise ValueError("No API Base provided for Azure OpenAI LLM provider. Set 'AZURE_API_BASE' in .env") ## EMBEDDING CALL response = azure_chat_completions.embedding( @@ -6349,10 +5999,7 @@ def embedding( or custom_llm_provider == "together_ai" or custom_llm_provider == "nvidia_nim" or custom_llm_provider == "litellm_proxy" - or ( - model in litellm.open_ai_embedding_models - and custom_llm_provider is None - ) + or (model in litellm.open_ai_embedding_models and custom_llm_provider is None) ): api_base = ( api_base @@ -6367,12 +6014,7 @@ def embedding( or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 ) # set API KEY - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") if headers is not None and headers != {}: optional_params["extra_headers"] = headers @@ -6384,9 +6026,7 @@ def embedding( if env_fmt is not None and env_fmt.strip().lower() == "none": optional_params.pop("encoding_format", None) else: - _default_fmt = ( - optional_params.get("encoding_format") or env_fmt or "float" - ) + _default_fmt = optional_params.get("encoding_format") or env_fmt or "float" if _default_fmt.strip().lower() == "none": optional_params.pop("encoding_format", None) else: @@ -6413,12 +6053,7 @@ def embedding( api_base = api_base or litellm.api_base or get_secret("DATABRICKS_API_BASE") # type: ignore # set API KEY - api_key = ( - api_key - or litellm.api_key - or litellm.databricks_key - or get_secret("DATABRICKS_API_KEY") - ) # type: ignore + api_key = api_key or litellm.api_key or litellm.databricks_key or get_secret("DATABRICKS_API_KEY") # type: ignore ## EMBEDDING CALL response = databricks_embedding.embedding( @@ -6434,9 +6069,7 @@ def embedding( aembedding=aembedding, ) elif custom_llm_provider == "hosted_vllm": - api_base = ( - api_base or litellm.api_base or get_secret_str("HOSTED_VLLM_API_BASE") - ) + api_base = api_base or litellm.api_base or get_secret_str("HOSTED_VLLM_API_BASE") # set API KEY if api_key is None: @@ -6462,18 +6095,11 @@ def embedding( or custom_llm_provider == "llamafile" or custom_llm_provider == "lm_studio" ): - api_base = ( - api_base or litellm.api_base or get_secret_str("OPENAI_LIKE_API_BASE") - ) + api_base = api_base or litellm.api_base or get_secret_str("OPENAI_LIKE_API_BASE") # set API KEY if api_key is None: - api_key = ( - api_key - or litellm.api_key - or litellm.openai_like_key - or get_secret_str("OPENAI_LIKE_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_like_key or get_secret_str("OPENAI_LIKE_API_KEY") if headers is not None and headers != {}: optional_params["extra_headers"] = headers @@ -6540,10 +6166,7 @@ def embedding( ) elif custom_llm_provider == "openrouter": api_base = ( - api_base - or litellm.api_base - or get_secret_str("OPENROUTER_API_BASE") - or "https://openrouter.ai/api/v1" + api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" ) api_key = ( @@ -6614,12 +6237,7 @@ def embedding( headers=headers, ) elif custom_llm_provider == "huggingface": - api_key = ( - api_key - or litellm.huggingface_key - or get_secret("HUGGINGFACE_API_KEY") - or litellm.api_key - ) # type: ignore + api_key = api_key or litellm.huggingface_key or get_secret("HUGGINGFACE_API_KEY") or litellm.api_key # type: ignore response = huggingface_embed.embedding( model=model, input=input, @@ -6657,9 +6275,7 @@ def embedding( ) elif custom_llm_provider == "triton": if api_base is None: - raise ValueError( - "api_base is required for triton. Please pass `api_base`" - ) + raise ValueError("api_base is required for triton. Please pass `api_base`") response = base_llm_http_handler.embedding( model=model, input=input, @@ -6721,16 +6337,11 @@ def embedding( ) api_base = ( - api_base - or litellm.api_base - or get_secret_str("VERTEXAI_API_BASE") - or get_secret_str("VERTEX_API_BASE") + api_base or litellm.api_base or get_secret_str("VERTEXAI_API_BASE") or get_secret_str("VERTEX_API_BASE") ) try: - model_info = get_model_info( - model=model, custom_llm_provider="vertex_ai" - ) + model_info = get_model_info(model=model, custom_llm_provider="vertex_ai") uses_embed_content = model_info.get("uses_embed_content", False) except Exception: uses_embed_content = False @@ -6757,8 +6368,7 @@ def embedding( elif ( "image" in optional_params or "video" in optional_params - or model - in vertex_multimodal_embedding.SUPPORTED_MULTIMODAL_EMBEDDING_MODELS + or model in vertex_multimodal_embedding.SUPPORTED_MULTIMODAL_EMBEDDING_MODELS ): response = vertex_multimodal_embedding.multimodal_embedding( model=model, @@ -6809,12 +6419,7 @@ def embedding( api_key=api_key, ) elif custom_llm_provider == "ollama": - api_base = ( - litellm.api_base - or api_base - or get_secret_str("OLLAMA_API_BASE") - or "http://localhost:11434" - ) # type: ignore + api_base = litellm.api_base or api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" # type: ignore if isinstance(input, str): input = [input] @@ -6824,11 +6429,7 @@ def embedding( model=model, # type: ignore llm_provider="ollama", # type: ignore ) - ollama_embeddings_fn = ( - ollama.ollama_aembeddings - if aembedding is True - else ollama.ollama_embeddings - ) + ollama_embeddings_fn = ollama.ollama_aembeddings if aembedding is True else ollama.ollama_embeddings response = ollama_embeddings_fn( # type: ignore api_base=api_base, model=model, @@ -6863,9 +6464,7 @@ def embedding( aembedding=aembedding, ) elif custom_llm_provider == "fireworks_ai": - api_key = ( - api_key or litellm.api_key or get_secret_str("FIREWORKS_AI_API_KEY") - ) + api_key = api_key or litellm.api_key or get_secret_str("FIREWORKS_AI_API_KEY") response = openai_chat_completions.embedding( model=model, input=input, @@ -6880,12 +6479,7 @@ def embedding( ) elif custom_llm_provider == "nebius": api_key = api_key or litellm.api_key or get_secret_str("NEBIUS_API_KEY") - api_base = ( - api_base - or litellm.api_base - or get_secret_str("NEBIUS_API_BASE") - or "api.studio.nebius.ai/v1" - ) + api_base = api_base or litellm.api_base or get_secret_str("NEBIUS_API_BASE") or "api.studio.nebius.ai/v1" response = openai_chat_completions.embedding( model=model, @@ -6902,10 +6496,7 @@ def embedding( elif custom_llm_provider == "wandb": api_key = api_key or litellm.api_key or get_secret_str("WANDB_API_KEY") api_base = ( - api_base - or litellm.api_base - or get_secret_str("WANDB_API_BASE") - or "https://api.inference.wandb.ai/v1" + api_base or litellm.api_base or get_secret_str("WANDB_API_BASE") or "https://api.inference.wandb.ai/v1" ) response = openai_chat_completions.embedding( @@ -6923,10 +6514,7 @@ def embedding( elif custom_llm_provider == "sambanova": api_key = api_key or litellm.api_key or get_secret_str("SAMBANOVA_API_KEY") api_base = ( - api_base - or litellm.api_base - or get_secret_str("SAMBANOVA_API_BASE") - or "https://api.sambanova.ai/v1" + api_base or litellm.api_base or get_secret_str("SAMBANOVA_API_BASE") or "https://api.sambanova.ai/v1" ) response = base_llm_http_handler.embedding( model=model, @@ -6999,16 +6587,10 @@ def embedding( ) elif custom_llm_provider == "xinference": api_key = ( - api_key - or litellm.api_key - or get_secret_str("XINFERENCE_API_KEY") - or "stub-xinference-key" + api_key or litellm.api_key or get_secret_str("XINFERENCE_API_KEY") or "stub-xinference-key" ) # xinference does not need an api key, pass a stub key if user did not set one api_base = ( - api_base - or litellm.api_base - or get_secret_str("XINFERENCE_API_BASE") - or "http://127.0.0.1:9997/v1" + api_base or litellm.api_base or get_secret_str("XINFERENCE_API_BASE") or "http://127.0.0.1:9997/v1" ) response = openai_chat_completions.embedding( model=model, @@ -7085,10 +6667,7 @@ def embedding( ) elif custom_llm_provider == "volcengine": volcengine_key = ( - api_key - or litellm.api_key - or get_secret_str("ARK_API_KEY") - or get_secret_str("VOLCENGINE_API_KEY") + api_key or litellm.api_key or get_secret_str("ARK_API_KEY") or get_secret_str("VOLCENGINE_API_KEY") ) if volcengine_key is None: raise ValueError( @@ -7114,9 +6693,7 @@ def embedding( headers=headers, ) elif custom_llm_provider == "dashscope": - dashscope_key = ( - api_key or litellm.api_key or get_secret_str("DASHSCOPE_API_KEY") - ) + dashscope_key = api_key or litellm.api_key or get_secret_str("DASHSCOPE_API_KEY") if dashscope_key is None: raise ValueError( "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter." @@ -7163,17 +6740,9 @@ def embedding( litellm_params={}, ) elif custom_llm_provider == "cometapi": - api_key = ( - api_key - or litellm.cometapi_key - or get_secret_str("COMETAPI_KEY") - or litellm.api_key - ) + api_key = api_key or litellm.cometapi_key or get_secret_str("COMETAPI_KEY") or litellm.api_key api_base = ( - api_base - or litellm.api_base - or get_secret_str("COMETAPI_API_BASE") - or "https://api.cometapi.com/v1" + api_base or litellm.api_base or get_secret_str("COMETAPI_API_BASE") or "https://api.cometapi.com/v1" ) response = base_llm_http_handler.embedding( model=model, @@ -7196,15 +6765,9 @@ def embedding( custom_handler = item["custom_handler"] if custom_handler is None: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) + raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) - handler_fn = ( - custom_handler.embedding - if not aembedding - else custom_handler.aembedding - ) + handler_fn = custom_handler.embedding if not aembedding else custom_handler.aembedding response = handler_fn( model=model, @@ -7272,20 +6835,12 @@ def embedding( litellm_params={}, ) else: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) - if ( - response is not None - and hasattr(response, "_hidden_params") - and isinstance(response, EmbeddingResponse) - ): + raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) + if response is not None and hasattr(response, "_hidden_params") and isinstance(response, EmbeddingResponse): response._hidden_params["custom_llm_provider"] = custom_llm_provider if response is None: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) + raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) return response except Exception as e: ## LOGGING @@ -7305,9 +6860,7 @@ def embedding( ###### Text Completion ################ @client -async def atext_completion( - *args, **kwargs -) -> Union[TextCompletionResponse, TextCompletionStreamWrapper]: +async def atext_completion(*args, **kwargs) -> Union[TextCompletionResponse, TextCompletionStreamWrapper]: """ Implemented to handle async streaming for the text completion endpoint """ @@ -7325,9 +6878,7 @@ async def atext_completion( func_with_context = partial(ctx.run, func) init_response = await loop.run_in_executor(None, func_with_context) - if isinstance(init_response, dict) or isinstance( - init_response, TextCompletionResponse - ): ## CACHING SCENARIO + if isinstance(init_response, dict) or isinstance(init_response, TextCompletionResponse): ## CACHING SCENARIO if isinstance(init_response, dict): response = TextCompletionResponse(**init_response) else: @@ -7384,27 +6935,13 @@ def text_completion( str, List[Union[str, List[Union[str, List[int]]]]] ], # Required: The prompt(s) to generate completions for. model: Optional[str] = None, # Optional: either `model` or `engine` can be set - best_of: Optional[ - int - ] = None, # Optional: Generates best_of completions server-side. - echo: Optional[ - bool - ] = None, # Optional: Echo back the prompt in addition to the completion. - frequency_penalty: Optional[ - float - ] = None, # Optional: Penalize new tokens based on their existing frequency. - logit_bias: Optional[ - Dict[int, int] - ] = None, # Optional: Modify the likelihood of specified tokens. - logprobs: Optional[ - int - ] = None, # Optional: Include the log probabilities on the most likely tokens. - max_tokens: Optional[ - int - ] = None, # Optional: The maximum number of tokens to generate in the completion. - n: Optional[ - int - ] = None, # Optional: How many completions to generate for each prompt. + best_of: Optional[int] = None, # Optional: Generates best_of completions server-side. + echo: Optional[bool] = None, # Optional: Echo back the prompt in addition to the completion. + frequency_penalty: Optional[float] = None, # Optional: Penalize new tokens based on their existing frequency. + logit_bias: Optional[Dict[int, int]] = None, # Optional: Modify the likelihood of specified tokens. + logprobs: Optional[int] = None, # Optional: Include the log probabilities on the most likely tokens. + max_tokens: Optional[int] = None, # Optional: The maximum number of tokens to generate in the completion. + n: Optional[int] = None, # Optional: How many completions to generate for each prompt. presence_penalty: Optional[ float ] = None, # Optional: Penalize new tokens based on whether they appear in the text so far. @@ -7413,14 +6950,10 @@ def text_completion( ] = None, # Optional: Sequences where the API will stop generating further tokens. stream: Optional[bool] = None, # Optional: Whether to stream back partial progress. stream_options: Optional[dict] = None, - suffix: Optional[ - str - ] = None, # Optional: The suffix that comes after a completion of inserted text. + suffix: Optional[str] = None, # Optional: The suffix that comes after a completion of inserted text. temperature: Optional[float] = None, # Optional: Sampling temperature to use. top_p: Optional[float] = None, # Optional: Nucleus sampling parameter. - user: Optional[ - str - ] = None, # Optional: A unique identifier representing your end-user. + user: Optional[str] = None, # Optional: A unique identifier representing your end-user. # set api_base, api_version, api_key api_base: Optional[str] = None, api_version: Optional[str] = None, @@ -7554,9 +7087,7 @@ def text_completion( executor.submit(process_prompt, i, individual_prompt) for i, individual_prompt in enumerate(prompt) ] - for i, future in enumerate( - concurrent.futures.as_completed(completed_futures) - ): + for i, future in enumerate(concurrent.futures.as_completed(completed_futures)): responses[i] = future.result() text_completion_response.choices = responses # type: ignore @@ -7614,11 +7145,7 @@ def text_completion( ) if kwargs.get("acompletion", False) is True: return response - if ( - stream is True - or kwargs.get("stream", False) is True - or isinstance(response, CustomStreamWrapper) - ): + if stream is True or kwargs.get("stream", False) is True or isinstance(response, CustomStreamWrapper): response = TextCompletionStreamWrapper( completion_stream=response, model=model, @@ -7633,11 +7160,9 @@ def text_completion( if isinstance(response, TextCompletionResponse): return response - text_completion_response = ( - litellm.utils.LiteLLMResponseObjectHandler.convert_chat_to_text_completion( - response=response, - text_completion_response=text_completion_response, - ) + text_completion_response = litellm.utils.LiteLLMResponseObjectHandler.convert_chat_to_text_completion( + response=response, + text_completion_response=text_completion_response, ) return text_completion_response @@ -7667,21 +7192,13 @@ async def aadapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) - response: Union[ModelResponse, CustomStreamWrapper] = await acompletion( - **new_kwargs - ) # type: ignore - translated_response: Optional[ - Union[BaseModel, AdapterCompletionStreamWrapper] - ] = None + response: Union[ModelResponse, CustomStreamWrapper] = await acompletion(**new_kwargs) # type: ignore + translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = None if isinstance(response, ModelResponse): - translated_response = translation_obj.translate_completion_output_params( - response=response - ) + translated_response = translation_obj.translate_completion_output_params(response=response) if isinstance(response, CustomStreamWrapper): - translated_response = ( - translation_obj.translate_completion_output_params_streaming( - completion_stream=response - ) + translated_response = translation_obj.translate_completion_output_params_streaming( + completion_stream=response ) return translated_response @@ -7696,16 +7213,12 @@ async def aadapter_generate_content( coro = cast( Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]], - GenerateContentToCompletionHandler.generate_content_handler( - **kwargs, _is_async=True - ), + GenerateContentToCompletionHandler.generate_content_handler(**kwargs, _is_async=True), ) return await coro -def adapter_completion( - *, adapter_id: str, **kwargs -) -> Optional[Union[BaseModel, AdapterCompletionStreamWrapper]]: +def adapter_completion(*, adapter_id: str, **kwargs) -> Optional[Union[BaseModel, AdapterCompletionStreamWrapper]]: translation_obj: Optional[CustomLogger] = None for item in litellm.adapters: if item["id"] == adapter_id: @@ -7721,19 +7234,11 @@ def adapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore - translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = ( - None - ) + translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = None if isinstance(response, ModelResponse): - translated_response = translation_obj.translate_completion_output_params( - response=response - ) + translated_response = translation_obj.translate_completion_output_params(response=response) elif isinstance(response, CustomStreamWrapper) or inspect.isgenerator(response): - translated_response = ( - translation_obj.translate_completion_output_params_streaming( - completion_stream=response - ) - ) + translated_response = translation_obj.translate_completion_output_params_streaming(completion_stream=response) return translated_response @@ -7745,12 +7250,7 @@ def moderation( input: str, model: Optional[str] = None, api_key: Optional[str] = None, **kwargs ) -> OpenAIModerationResponse: # only supports open ai for now - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") # Extract api_base from kwargs api_base = kwargs.get("api_base", None) @@ -7784,16 +7284,9 @@ async def amoderation( from openai import AsyncOpenAI # only supports open ai for now - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") optional_params = GenericLiteLLMParams(**kwargs) - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj", None - ) + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) _dynamic_api_base = None try: ( @@ -7870,9 +7363,7 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) - _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) - ) + _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None)) # Await normally init_response = await loop.run_in_executor(None, func_with_context) @@ -7894,18 +7385,12 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: # exposing it in the response body. Adding duration to the response # tricks the OpenAI SDK's "best match deserialization" into thinking # a plain Transcription is a TranscriptionVerbose/Diarized type. - if ( - response is not None - and not isinstance(response, Coroutine) - and file is not None - ): + if response is not None and not isinstance(response, Coroutine) and file is not None: existing_duration = getattr(response, "duration", None) if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params["audio_transcription_duration"] = ( - calculated_duration - ) + response._hidden_params["audio_transcription_duration"] = calculated_duration return response except Exception as e: @@ -7926,9 +7411,7 @@ def transcription( ## OPTIONAL OPENAI PARAMS ## language: Optional[str] = None, prompt: Optional[str] = None, - response_format: Optional[ - Literal["json", "text", "srt", "verbose_json", "vtt"] - ] = None, + response_format: Optional[Literal["json", "text", "srt", "verbose_json", "vtt"]] = None, timestamp_granularities: Optional[List[Literal["word", "segment"]]] = None, temperature: Optional[int] = None, # openai defaults this to 0 ## LITELLM PARAMS ## @@ -8012,9 +7495,7 @@ def transcription( custom_llm_provider=custom_llm_provider, ) - response: Optional[ - Union[TranscriptionResponse, Coroutine[Any, Any, TranscriptionResponse]] - ] = None + response: Optional[Union[TranscriptionResponse, Coroutine[Any, Any, TranscriptionResponse]]] = None provider_config = ProviderConfigManager.get_provider_audio_transcription_config( model=model, @@ -8025,20 +7506,11 @@ def transcription( # azure configs api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") - api_version = ( - api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") - ) + api_version = api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") - azure_ad_token = kwargs.pop("azure_ad_token", None) or get_secret_str( - "AZURE_AD_TOKEN" - ) + azure_ad_token = kwargs.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") - api_key = ( - api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_API_KEY") - ) + api_key = api_key or litellm.api_key or litellm.azure_key or get_secret_str("AZURE_API_KEY") optional_params["extra_headers"] = extra_headers @@ -8058,9 +7530,7 @@ def transcription( max_retries=max_retries, litellm_params=litellm_params_dict, ) - elif custom_llm_provider == "openai" or ( - custom_llm_provider in litellm.openai_compatible_providers - ): + elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers): api_base = ( api_base or litellm.api_base @@ -8075,12 +7545,7 @@ def transcription( ) # set API KEY - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret("OPENAI_API_KEY") - ) # type: ignore + api_key = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") # type: ignore response = openai_audio_transcriptions.audio_transcriptions( model=model, audio_file=file, @@ -8112,9 +7577,7 @@ def transcription( api_base=api_base, api_key=api_key, provider_config=( - provider_config - if isinstance(provider_config, NvidiaRivaAudioTranscriptionConfig) - else None + provider_config if isinstance(provider_config, NvidiaRivaAudioTranscriptionConfig) else None ), ) elif custom_llm_provider == "soniox": @@ -8131,11 +7594,7 @@ def transcription( atranscription=atranscription, client=( client - if client is not None - and ( - isinstance(client, HTTPHandler) - or isinstance(client, AsyncHTTPHandler) - ) + if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler)) else None ), timeout=timeout, @@ -8156,11 +7615,7 @@ def transcription( atranscription=atranscription, client=( client - if client is not None - and ( - isinstance(client, HTTPHandler) - or isinstance(client, AsyncHTTPHandler) - ) + if client is not None and (isinstance(client, HTTPHandler) or isinstance(client, AsyncHTTPHandler)) else None ), timeout=timeout, @@ -8181,9 +7636,7 @@ def transcription( if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params["audio_transcription_duration"] = ( - calculated_duration - ) + response._hidden_params["audio_transcription_duration"] = calculated_duration if response is None: raise ValueError("Unmapped provider passed in. Unable to get the response.") @@ -8208,9 +7661,7 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent: ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) - _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) - ) + _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None)) # Await normally init_response = await loop.run_in_executor(None, func_with_context) @@ -8280,11 +7731,9 @@ def speech( litellm_params_dict = get_litellm_params(**kwargs) # Get provider-specific text-to-speech config and map parameters - text_to_speech_provider_config = ( - ProviderConfigManager.get_provider_text_to_speech_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + text_to_speech_provider_config = ProviderConfigManager.get_provider_text_to_speech_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) # Map OpenAI params to provider-specific params if config exists @@ -8297,9 +7746,7 @@ def speech( kwargs=kwargs, ) - logging_obj: LiteLLMLoggingObj = cast( - LiteLLMLoggingObj, kwargs.get("litellm_logging_obj") - ) + logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj")) logging_obj.update_environment_variables( model=model, user=user, @@ -8320,10 +7767,7 @@ def speech( Coroutine[Any, Any, 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: 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", @@ -8392,9 +7836,7 @@ def speech( ) # Cast to specific Azure config type to access dispatch method - azure_config = cast( - AzureAVATextToSpeechConfig, text_to_speech_provider_config - ) + azure_config = cast(AzureAVATextToSpeechConfig, text_to_speech_provider_config) response = azure_config.dispatch_text_to_speech( # type: ignore model=model, @@ -8421,9 +7863,7 @@ def speech( ) api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version = ( - api_version or litellm.api_version or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version = api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( api_key @@ -8466,9 +7906,7 @@ def speech( if text_to_speech_provider_config is None: text_to_speech_provider_config = ElevenLabsTextToSpeechConfig() - elevenlabs_config = cast( - ElevenLabsTextToSpeechConfig, text_to_speech_provider_config - ) + elevenlabs_config = cast(ElevenLabsTextToSpeechConfig, text_to_speech_provider_config) voice_id = voice if isinstance(voice, str) else None if voice_id is None or not voice_id.strip(): @@ -8479,17 +7917,11 @@ def speech( ) voice_id = voice_id.strip() - query_params = kwargs.pop( - ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY, None - ) + query_params = kwargs.pop(ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY, None) if isinstance(query_params, dict): - litellm_params_dict[ - ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY - ] = query_params + litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY] = query_params - litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = ( - voice_id - ) + litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = voice_id if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -8594,9 +8026,7 @@ def speech( ) # Cast to specific RunwayML config type to access dispatch method - runwayml_config = cast( - RunwayMLTextToSpeechConfig, text_to_speech_provider_config - ) + runwayml_config = cast(RunwayMLTextToSpeechConfig, text_to_speech_provider_config) response = runwayml_config.dispatch_text_to_speech( # type: ignore model=model, @@ -8661,9 +8091,7 @@ def speech( text_to_speech_provider_config = AWSPollyTextToSpeechConfig() # Cast to specific AWS Polly config type to access dispatch method - aws_polly_config = cast( - AWSPollyTextToSpeechConfig, text_to_speech_provider_config - ) + aws_polly_config = cast(AWSPollyTextToSpeechConfig, text_to_speech_provider_config) response = aws_polly_config.dispatch_text_to_speech( model=model, @@ -8730,10 +8158,8 @@ async def ahealth_check( log_raw_request_response=True, ) model_params["litellm_logging_obj"] = litellm_logging_obj - model_params = ( - HealthCheckHelpers._update_model_params_with_health_check_tracking_information( - model_params=model_params - ) + model_params = HealthCheckHelpers._update_model_params_with_health_check_tracking_information( + model_params=model_params ) ######################################################### try: @@ -8757,9 +8183,7 @@ async def ahealth_check( if model in litellm.model_cost and mode is None: mode = litellm.model_cost[model].get("mode") - model_params["cache"] = { - "no-cache": True - } # don't used cached responses for making health check calls + model_params["cache"] = {"no-cache": True} # don't used cached responses for making health check calls mode = mode or "chat" if "*" in model: return await HealthCheckHelpers.ahealth_check_wildcard_models( @@ -8780,14 +8204,10 @@ async def ahealth_check( if mode in mode_handlers: _response = await mode_handlers[mode]() # Only process headers for chat mode - _response_headers: dict = ( - getattr(_response, "_hidden_params", {}).get("headers", {}) or {} - ) + _response_headers: dict = getattr(_response, "_hidden_params", {}).get("headers", {}) or {} return _create_health_check_response(_response_headers) else: - raise Exception( - f"Mode {mode} not supported. See modes here: https://docs.litellm.ai/docs/proxy/health" - ) + raise Exception(f"Mode {mode} not supported. See modes here: https://docs.litellm.ai/docs/proxy/health") except Exception as e: stack_trace = _redact_string(traceback.format_exc()) if isinstance(stack_trace, str): @@ -8801,9 +8221,7 @@ async def ahealth_check( error_to_return = str(e) + "\nstack trace: " + stack_trace - raw_request_typed_dict = litellm_logging_obj.model_call_details.get( - "raw_request_typed_dict" - ) + raw_request_typed_dict = litellm_logging_obj.model_call_details.get("raw_request_typed_dict") return { "error": error_to_return, @@ -8834,9 +8252,7 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion( - chunks: list, messages: Optional[List] = None -) -> TextCompletionResponse: +def stream_chunk_builder_text_completion(chunks: list, messages: Optional[List] = None) -> TextCompletionResponse: id = chunks[0]["id"] object = chunks[0]["object"] created = chunks[0]["created"] @@ -8869,11 +8285,7 @@ def stream_chunk_builder_text_completion( for chunk in chunks: choices = chunk["choices"] for choice in choices: - if ( - choice is not None - and hasattr(choice, "text") - and choice.get("text") is not None - ): + if choice is not None and hasattr(choice, "text") and choice.get("text") is not None: _choice = choice.get("text") content_list.append(_choice) @@ -8889,9 +8301,7 @@ def stream_chunk_builder_text_completion( pass # # Update usage information if needed try: - response["usage"]["prompt_tokens"] = token_counter( - model=model, messages=messages - ) + response["usage"]["prompt_tokens"] = token_counter(model=model, messages=messages) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") response["usage"]["prompt_tokens"] = 0 @@ -8900,9 +8310,7 @@ def stream_chunk_builder_text_completion( text=combined_content, count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages ) - response["usage"]["total_tokens"] = ( - response["usage"]["prompt_tokens"] + response["usage"]["completion_tokens"] - ) + response["usage"]["total_tokens"] = response["usage"]["prompt_tokens"] + response["usage"]["completion_tokens"] return TextCompletionResponse(**response) @@ -8935,9 +8343,7 @@ def stream_chunk_builder( if first_chunk_with_choices is not None and isinstance( first_chunk_with_choices["choices"][0], litellm.utils.TextChoices ): # route to the text completion logic - return stream_chunk_builder_text_completion( - chunks=chunks, messages=messages - ) + return stream_chunk_builder_text_completion(chunks=chunks, messages=messages) model = chunks[0]["model"] # Initialize the response dictionary @@ -8952,11 +8358,7 @@ def stream_chunk_builder( continue choice = chunk["choices"][0] - delta_obj = ( - choice.get("delta", {}) - if isinstance(choice, dict) - else getattr(choice, "delta", {}) - ) + delta_obj = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) if isinstance(delta_obj, dict): delta = delta_obj elif hasattr(delta_obj, "model_dump"): @@ -8983,9 +8385,7 @@ def stream_chunk_builder( if is_simple_text_stream: if simple_content_parts: - response["choices"][0]["message"]["content"] = "".join( - simple_content_parts - ) + response["choices"][0]["message"]["content"] = "".join(simple_content_parts) completion_output = get_content_from_model_response(response) usage = processor.calculate_usage( chunks=chunks, @@ -9003,9 +8403,9 @@ def stream_chunk_builder( else: hidden = getattr(chunk, "_hidden_params", None) if isinstance(hidden, dict) and "provider_specific_fields" in hidden: - response._hidden_params.setdefault( - "provider_specific_fields", {} - ).update(hidden["provider_specific_fields"]) + response._hidden_params.setdefault("provider_specific_fields", {}).update( + hidden["provider_specific_fields"] + ) break if litellm.include_cost_in_streaming_usage and logging_obj is not None: @@ -9014,9 +8414,7 @@ def stream_chunk_builder( "cost", logging_obj._response_cost_calculator(result=response), ) - processor.apply_provider_assembled_streaming_metadata( - response, chunks, logging_obj - ) + processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response tool_call_chunks = [ @@ -9044,9 +8442,7 @@ def stream_chunk_builder( if len(function_call_chunks) > 0: _choice = cast(Choices, response.choices[0]) _choice.message.content = None - _choice.message.function_call = ( - processor.get_combined_function_call_content(function_call_chunks) - ) + _choice.message.function_call = processor.get_combined_function_call_content(function_call_chunks) content_chunks = [ chunk @@ -9057,9 +8453,7 @@ def stream_chunk_builder( ] if len(content_chunks) > 0: - response["choices"][0]["message"]["content"] = ( - processor.get_combined_content(content_chunks) - ) + response["choices"][0]["message"]["content"] = processor.get_combined_content(content_chunks) thinking_blocks = [ chunk @@ -9070,8 +8464,8 @@ def stream_chunk_builder( ] if len(thinking_blocks) > 0: - response["choices"][0]["message"]["thinking_blocks"] = ( - processor.get_combined_thinking_content(thinking_blocks) + response["choices"][0]["message"]["thinking_blocks"] = processor.get_combined_thinking_content( + thinking_blocks ) reasoning_chunks = [ @@ -9083,8 +8477,8 @@ def stream_chunk_builder( ] if len(reasoning_chunks) > 0: - response["choices"][0]["message"]["reasoning_content"] = ( - processor.get_combined_reasoning_content(reasoning_chunks) + response["choices"][0]["message"]["reasoning_content"] = processor.get_combined_reasoning_content( + reasoning_chunks ) annotation_chunks = [ @@ -9151,9 +8545,7 @@ def stream_chunk_builder( for key, value in fields.items(): if key not in combined_provider_fields: combined_provider_fields[key] = value - elif isinstance(value, list) and isinstance( - combined_provider_fields[key], list - ): + elif isinstance(value, list) and isinstance(combined_provider_fields[key], list): # For lists like web_search_results, take the last (most complete) one combined_provider_fields[key] = value else: @@ -9185,27 +8577,19 @@ def stream_chunk_builder( else: hidden = getattr(chunk, "_hidden_params", None) if isinstance(hidden, dict) and "provider_specific_fields" in hidden: - response._hidden_params.setdefault( - "provider_specific_fields", {} - ).update(hidden["provider_specific_fields"]) + response._hidden_params.setdefault("provider_specific_fields", {}).update( + hidden["provider_specific_fields"] + ) break # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and logging_obj is not None: - setattr( - usage, "cost", logging_obj._response_cost_calculator(result=response) - ) + setattr(usage, "cost", logging_obj._response_cost_calculator(result=response)) - processor.apply_provider_assembled_streaming_metadata( - response, chunks, logging_obj - ) + processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response except Exception as e: - verbose_logger.exception( - "litellm.main.py::stream_chunk_builder() - Exception occurred - {}".format( - str(e) - ) - ) + verbose_logger.exception("litellm.main.py::stream_chunk_builder() - Exception occurred - {}".format(str(e))) raise litellm.APIError( status_code=500, message="Error building chunks for logging/streaming usage calculation", @@ -9276,17 +8660,12 @@ async def acount_tokens( # Try to get provider-specific token counter try: llm_provider_enum = LlmProviders(custom_llm_provider) - provider_model_info = ProviderConfigManager.get_provider_model_info( - model=model, provider=llm_provider_enum - ) + provider_model_info = ProviderConfigManager.get_provider_model_info(model=model, provider=llm_provider_enum) if provider_model_info is not None: token_counter_instance = provider_model_info.get_token_counter() - if ( - token_counter_instance is not None - and token_counter_instance.should_use_token_counting_api( - custom_llm_provider - ) + if token_counter_instance is not None and token_counter_instance.should_use_token_counting_api( + custom_llm_provider ): result = await token_counter_instance.count_tokens( model_to_use=resolved_model, @@ -9300,9 +8679,7 @@ async def acount_tokens( if result is not None and not result.error: return result except Exception as e: - verbose_logger.debug( - f"Provider token counting failed for model={model}, falling back to local: {e}" - ) + verbose_logger.debug(f"Provider token counting failed for model={model}, falling back to local: {e}") # Fallback to local tiktoken-based token counting fallback_messages = messages or [] diff --git a/litellm/models/budget.py b/litellm/models/budget.py index e7dfe2f8fbc..8c35aebd208 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -29,9 +29,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None model_max_budget: Optional[dict] = None budget_duration: Optional[str] = None - allowed_models: Optional[List[str]] = ( - None # per-member model scope; empty = inherit team models - ) + allowed_models: Optional[List[str]] = None # per-member model scope; empty = inherit team models model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/models/team.py b/litellm/models/team.py index aa0798955f2..f11c21a078e 100644 --- a/litellm/models/team.py +++ b/litellm/models/team.py @@ -118,10 +118,7 @@ class LiteLLM_TeamTable(TeamBase): if isinstance(values, BaseModel): values = values.model_dump() - if ( - isinstance(values.get("members_with_roles"), dict) - and not values["members_with_roles"] - ): + if isinstance(values.get("members_with_roles"), dict) and not values["members_with_roles"]: values["members_with_roles"] = [] for field in dict_fields: diff --git a/litellm/models/team_membership.py b/litellm/models/team_membership.py index d0a1308ce7c..e79b64977d4 100644 --- a/litellm/models/team_membership.py +++ b/litellm/models/team_membership.py @@ -17,9 +17,7 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None spend: Optional[float] = 0.0 total_spend: Optional[float] = 0.0 - litellm_budget_table: Optional[ - Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable] - ] = None + litellm_budget_table: Optional[Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable]] = None def safe_get_team_member_rpm_limit(self) -> Optional[int]: if self.litellm_budget_table is not None: diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 93b3a892659..5716155361d 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -72,9 +72,7 @@ def _prepare_ocr_request( litellm_call_id = cast(str | None, kwargs.get("litellm_call_id", None)) if not isinstance(document, dict): - raise ValueError( - f"document must be a dict with 'type' and URL/file field, got {type(document)}" - ) + raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}") doc_type = document.get("type") @@ -83,10 +81,7 @@ def _prepare_ocr_request( doc_type = document.get("type") if doc_type not in ["document_url", "image_url"]: - raise ValueError( - f"Invalid document type: {doc_type}. " - "Must be 'document_url', 'image_url', or 'file'" - ) + raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") ( model, @@ -226,9 +221,7 @@ def _prepare_rust_ocr_call( litellm_params=prepared_request.litellm_params, ) rust_api_base = _rust_bridge_api_base(prepared_request, resolve_api_key) - rust_optional_params = _rust_bridge_optional_params( - prepared_request, resolve_api_key - ) + rust_optional_params = _rust_bridge_optional_params(prepared_request, resolve_api_key) prepared_request.litellm_logging_obj.pre_call( input="OCR document processing", api_key=resolved_api_key, @@ -392,9 +385,7 @@ async def aocr( ) model = prepared.model custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update( - {"model": model, "custom_llm_provider": custom_llm_provider} - ) + completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): from litellm.secret_managers.main import get_secret_str @@ -404,9 +395,7 @@ async def aocr( resolve_api_key=get_secret_str, ) if rust_response is None: - verbose_logger.debug( - "Async Rust OCR bridge unavailable; falling back to Python path" - ) + verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") else: return rust_response @@ -429,9 +418,7 @@ async def aocr( response = await response if response is None: - raise ValueError( - f"Got an unexpected None response from the OCR API: {response}" - ) + raise ValueError(f"Got an unexpected None response from the OCR API: {response}") return response except Exception as e: @@ -539,8 +526,7 @@ def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, file_bytes = file_bytes.encode("utf-8") else: raise ValueError( - f"Unsupported file input type: {type(file_input)}. " - "Expected pathlib.Path, bytes, or a file-like object." + f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object." ) if not file_bytes: @@ -667,9 +653,7 @@ def ocr( ) model = prepared.model custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update( - {"model": model, "custom_llm_provider": custom_llm_provider} - ) + completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): from litellm.secret_managers.main import get_secret_str @@ -679,9 +663,7 @@ def ocr( resolve_api_key=get_secret_str, ) if rust_response is None: - verbose_logger.debug( - "Rust OCR bridge unavailable; falling back to Python path" - ) + verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") else: return rust_response diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index b1d4864cea7..66367513062 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -264,9 +264,7 @@ def llm_passthrough_route( # [TODO: Refactor to bedrockpassthroughconfig] need to encode the id of application-inference-profile for bedrock if custom_llm_provider == "bedrock" and "application-inference-profile" in endpoint: - encoded_url_str = CommonUtils.encode_bedrock_runtime_modelid_arn( - str(updated_url) - ) + encoded_url_str = CommonUtils.encode_bedrock_runtime_modelid_arn(str(updated_url)) updated_url = httpx.URL(encoded_url_str) # Add or update query parameters @@ -345,9 +343,7 @@ def llm_passthrough_route( ) else: # Sync path - client.client.send returns Response directly - response: httpx.Response = client.client.send( - request=request, stream=is_streaming_request - ) # type: ignore + response: httpx.Response = client.client.send(request=request, stream=is_streaming_request) # type: ignore response.raise_for_status() if ( diff --git a/litellm/passthrough/timeout_utils.py b/litellm/passthrough/timeout_utils.py index a423db2aa91..84ec89b7e2a 100644 --- a/litellm/passthrough/timeout_utils.py +++ b/litellm/passthrough/timeout_utils.py @@ -21,9 +21,7 @@ def resolve_pass_through_request_timeout( try: proxy_server = sys.modules.get("litellm.proxy.proxy_server") if proxy_server is not None: - global_timeout = getattr(proxy_server, "general_settings", {}).get( - "pass_through_request_timeout" - ) + global_timeout = getattr(proxy_server, "general_settings", {}).get("pass_through_request_timeout") if global_timeout is not None: return float(global_timeout) except Exception: diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index 9484922833a..706beb7dc5e 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -36,9 +36,7 @@ class BasePassthroughUtils: existing_query_params = parse_qs(existing_query_string) # parse_qs returns a dict where each value is a list, so let's flatten it - updated_existing_query_params = { - k: v[0] if len(v) == 1 else v for k, v in existing_query_params.items() - } + updated_existing_query_params = {k: v[0] if len(v) == 1 else v for k, v in existing_query_params.items()} # Start with default query params (lowest priority) merged_params = {} @@ -84,12 +82,9 @@ class BasePassthroughUtils: for header_name, header_value in request_headers.items(): if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX): # Strip the 'x-pass-' prefix and normalize to lowercase - actual_header_name = header_name[ - len(PASS_THROUGH_HEADER_PREFIX) : - ].lower() + actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower() if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any( - actual_header_name.startswith(p) - for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES + actual_header_name.startswith(p) for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES ): verbose_logger.debug( "x-pass- header %s maps to a protected header name; skipping", diff --git a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py b/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py index 97a16ad3e15..aa42074f2c7 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/auth/token_exchange.py @@ -49,9 +49,7 @@ class TokenExchangeHandler: ) # WeakValueDictionary so locks are GC'd once no coroutine holds a reference, # preventing unbounded growth with many rotating user tokens. - self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( - weakref.WeakValueDictionary() - ) + self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() def _get_lock(self, cache_key: str) -> asyncio.Lock: lock = self._locks.get(cache_key) @@ -118,8 +116,7 @@ class TokenExchangeHandler: data: Dict[str, str] = { "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, "subject_token": subject_token, - "subject_token_type": server.subject_token_type - or DEFAULT_SUBJECT_TOKEN_TYPE, + "subject_token_type": server.subject_token_type or DEFAULT_SUBJECT_TOKEN_TYPE, "client_id": server.client_id, "client_secret": server.client_secret, } @@ -146,8 +143,7 @@ class TokenExchangeHandler: exc.response.status_code, ) raise ValueError( - f"Token exchange for MCP server '{server.server_id}' " - f"failed with status {exc.response.status_code}" + f"Token exchange for MCP server '{server.server_id}' failed with status {exc.response.status_code}" ) from exc body = response.json() @@ -159,18 +155,11 @@ class TokenExchangeHandler: access_token = body.get("access_token") if not access_token: - raise ValueError( - f"Token exchange response for MCP server '{server.server_id}' " - f"missing 'access_token'" - ) + raise ValueError(f"Token exchange response for MCP server '{server.server_id}' missing 'access_token'") raw_expires_in = body.get("expires_in") try: - expires_in = ( - int(raw_expires_in) - if raw_expires_in is not None - else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL - ) + expires_in = int(raw_expires_in) if raw_expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL except (TypeError, ValueError): expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index bd4c45d5bcc..2520c7e82a1 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -23,9 +23,7 @@ from litellm.repositories.table_repositories import ( ) -def _parse_mcp_server_names_from_path( - path: str, mcp_servers_header: Optional[List[str]] = None -) -> Optional[List[str]]: +def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: Optional[List[str]] = None) -> Optional[List[str]]: """Resolve the single MCP server name a cold-start passthrough bypass may target. Delegates parsing to :meth:`MCPRequestHandler._extract_target_server_names_from_path` so the @@ -59,9 +57,7 @@ def _parse_mcp_server_names_from_path( return servers -def _is_mcp_passthrough_cold_start( - mcp_servers: Optional[List[str]], client_ip: Optional[str] -) -> bool: +def _is_mcp_passthrough_cold_start(mcp_servers: Optional[List[str]], client_ip: Optional[str]) -> bool: """True only when EVERY targeted server is a pass-through server with no auth headers — the cold-start OAuth discovery case per RFC 9728 / MCP Authorization spec. Lets the route handler's 401 emitter produce the @@ -79,9 +75,7 @@ def _is_mcp_passthrough_cold_start( ) for name in mcp_servers: - server = global_mcp_server_manager.get_mcp_server_by_name( - name, client_ip=client_ip - ) + server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip) if server is None or not getattr(server, "is_oauth_passthrough", False): return False return True @@ -161,44 +155,31 @@ class MCPRequestHandler: headers = MCPRequestHandler._safe_get_headers_from_scope(scope) # Check if there is an explicit LiteLLM API key (primary header) - has_explicit_litellm_key = ( - headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) - is not None - ) + has_explicit_litellm_key = headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) is not None - litellm_api_key = ( - MCPRequestHandler.get_litellm_api_key_from_headers(headers) or "" - ) + litellm_api_key = MCPRequestHandler.get_litellm_api_key_from_headers(headers) or "" # Get the old mcp_auth_header for backward compatibility mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers) # Get the new server-specific auth headers - mcp_server_auth_headers = ( - MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - ) + mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) # Get the oauth2 headers oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) # Parse MCP servers from header - mcp_servers_header = headers.get( - MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME - ) + mcp_servers_header = headers.get(MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME) verbose_logger.debug(f"Raw MCP servers header: {mcp_servers_header}") mcp_servers = None if mcp_servers_header is not None: try: - mcp_servers = [ - s.strip() for s in mcp_servers_header.split(",") if s.strip() - ] + mcp_servers = [s.strip() for s in mcp_servers_header.split(",") if s.strip()] verbose_logger.debug(f"Parsed MCP servers: {mcp_servers}") except Exception as e: verbose_logger.debug(f"Error parsing mcp_servers header: {e}") mcp_servers = None - if mcp_servers_header == "" or ( - mcp_servers is not None and len(mcp_servers) == 0 - ): + if mcp_servers_header == "" or (mcp_servers is not None and len(mcp_servers) == 0): mcp_servers = [] # Create a proper Request object with mock body method to avoid ASGI receive channel issues request = Request(scope=scope) @@ -220,9 +201,7 @@ class MCPRequestHandler: # An explicit x-litellm-api-key is always a LiteLLM credential, even # for a delegated server, so validate it: identity / spend / rate # limits resolve and any stored upstream token can be forwarded. - validated_user_api_key_auth = await user_api_key_auth( - api_key=litellm_api_key, request=request - ) + validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) elif MCPRequestHandler._target_servers_delegate_auth_to_upstream( path=request_route, mcp_servers=mcp_servers, @@ -247,17 +226,13 @@ class MCPRequestHandler: # so a recognized-but-forbidden key still fails closed. client_ip = IPAddressUtils.get_mcp_client_ip(request) try: - validated_user_api_key_auth = await user_api_key_auth( - api_key=litellm_api_key, request=request - ) + validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as e: # ProxyException.code is normalized to str (possibly "None"), so # compare both int and str forms rather than coercing. status = e.status_code if isinstance(e, HTTPException) else e.code is_unauthenticated = status in (401, "401") - mcp_servers_from_path = _parse_mcp_server_names_from_path( - request_route, mcp_servers - ) + mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) if ( is_unauthenticated and mcp_servers_from_path is not None @@ -265,30 +240,23 @@ class MCPRequestHandler: mcp_auth_header, mcp_server_auth_headers, ) - and _is_mcp_passthrough_cold_start( - mcp_servers_from_path, client_ip=client_ip - ) + and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) ): verbose_logger.debug( - "MCP pass-through return: forwarding Authorization as " - "upstream OAuth token for delegated auth" + "MCP pass-through return: forwarding Authorization as upstream OAuth token for delegated auth" ) validated_user_api_key_auth = UserAPIKeyAuth() else: raise else: try: - validated_user_api_key_auth = await user_api_key_auth( - api_key=litellm_api_key, request=request - ) + validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as exc: # Cold-start MCP OAuth discovery: RFC 9728 / MCP Authorization spec # require unauthenticated requests to protected resources to receive # 401 + WWW-Authenticate. Defer to _raise_preemptive_401_for_unauthenticated_servers # for pass-through servers instead of surfacing a generic admission error. - mcp_servers_from_path = _parse_mcp_server_names_from_path( - request_route, mcp_servers - ) + mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) client_ip = IPAddressUtils.get_mcp_client_ip(request) if ( mcp_servers_from_path is not None @@ -297,13 +265,9 @@ class MCPRequestHandler: mcp_server_auth_headers, ) and _is_litellm_auth_admission_error(exc) - and _is_mcp_passthrough_cold_start( - mcp_servers_from_path, client_ip=client_ip - ) + and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) ): - verbose_logger.debug( - "MCP pass-through cold start: deferring admission to route 401 emitter" - ) + verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") validated_user_api_key_auth = UserAPIKeyAuth() else: raise @@ -370,9 +334,7 @@ class MCPRequestHandler: return [s.strip() for s in servers_part.split(",") if s.strip()] # Single-server case — server name may contain at most one slash. - single_server_match = re.match( - r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path - ) + single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path) if single_server_match: return [single_server_match.group(1)] return [servers_and_path] @@ -402,16 +364,12 @@ class MCPRequestHandler: # (``extract_mcp_auth_context``) or an attacker could set # ``x-mcp-servers`` to a delegate-enabled server while the URL path # targets a non-delegate server, skipping LiteLLM auth for it. - target_names = MCPRequestHandler._resolve_target_server_names( - path=path, mcp_servers_header=mcp_servers - ) + target_names = MCPRequestHandler._resolve_target_server_names(path=path, mcp_servers_header=mcp_servers) if not target_names: return False for name in target_names: - server = global_mcp_server_manager.get_mcp_server_by_name( - name, client_ip=client_ip - ) + server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip) if server is None or server.auth_type != MCPAuth.oauth2: return False # `is True` is intentional: opt-in must be an explicit boolean @@ -428,9 +386,7 @@ class MCPRequestHandler: return True @staticmethod - def _resolve_target_server_names( - path: str, mcp_servers_header: Optional[List[str]] - ) -> List[str]: + def _resolve_target_server_names(path: str, mcp_servers_header: Optional[List[str]]) -> List[str]: """ Resolve the target MCP server names exactly as downstream routing does (``server.py::extract_mcp_auth_context``). @@ -464,9 +420,7 @@ class MCPRequestHandler: DEPRECATED: This method is deprecated in favor of server-specific auth headers using the format x-mcp-{{server_alias}}-{{header_name}} instead. """ - mcp_client_side_auth_header_name: str = ( - MCPRequestHandler._get_mcp_client_side_auth_header_name() - ) + mcp_client_side_auth_header_name: str = MCPRequestHandler._get_mcp_client_side_auth_header_name() auth_header = headers.get(mcp_client_side_auth_header_name) if auth_header: verbose_logger.warning( @@ -498,10 +452,8 @@ class MCPRequestHandler: if header_name.lower().startswith(prefix): # Skip the access groups header as it's not a server auth header if ( - header_name.lower() - == MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower() - or header_name.lower() - == MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower() + header_name.lower() == MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower() + or header_name.lower() == MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower() ): continue @@ -521,9 +473,7 @@ class MCPRequestHandler: if server_alias not in server_auth_headers: server_auth_headers[server_alias] = {} - server_auth_headers[server_alias][auth_header_name] = ( - header_value - ) + server_auth_headers[server_alias][auth_header_name] = header_value verbose_logger.debug( f"Found server auth header: {server_alias} -> {auth_header_name}: {header_value[:10]}..." ) @@ -553,18 +503,14 @@ class MCPRequestHandler: from litellm.proxy.proxy_server import general_settings from litellm.secret_managers.main import get_secret_str - MCP_CLIENT_SIDE_AUTH_HEADER_NAME: str = ( - MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME - ) + MCP_CLIENT_SIDE_AUTH_HEADER_NAME: str = MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME if get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") is not None: MCP_CLIENT_SIDE_AUTH_HEADER_NAME = ( - get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") - or MCP_CLIENT_SIDE_AUTH_HEADER_NAME + get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") or MCP_CLIENT_SIDE_AUTH_HEADER_NAME ) elif general_settings.get("mcp_client_side_auth_header_name") is not None: MCP_CLIENT_SIDE_AUTH_HEADER_NAME = ( - general_settings.get("mcp_client_side_auth_header_name") - or MCP_CLIENT_SIDE_AUTH_HEADER_NAME + general_settings.get("mcp_client_side_auth_header_name") or MCP_CLIENT_SIDE_AUTH_HEADER_NAME ) return MCP_CLIENT_SIDE_AUTH_HEADER_NAME @@ -584,9 +530,7 @@ class MCPRequestHandler: if api_key: return api_key - auth_header = headers.get( - MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_SECONDARY - ) + auth_header = headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_SECONDARY) if auth_header: return auth_header @@ -605,10 +549,7 @@ class MCPRequestHandler: # ASGI headers are list of [name: bytes, value: bytes] pairs raw_headers = scope.get("headers", []) # Convert bytes to strings and create dict for Headers constructor - headers_dict = { - name.decode("latin-1"): value.decode("latin-1") - for name, value in raw_headers - } + headers_dict = {name.decode("latin-1"): value.decode("latin-1") for name, value in raw_headers} return Headers(headers_dict) except (UnicodeDecodeError, AttributeError, TypeError) as e: verbose_logger.exception(f"Error getting headers from scope: {e}") @@ -640,31 +581,16 @@ class MCPRequestHandler: try: # Get allowed servers from key and team - allowed_mcp_servers_for_key = ( - await MCPRequestHandler._get_allowed_mcp_servers_for_key( - user_api_key_auth - ) - ) + allowed_mcp_servers_for_key = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth) # The key explicitly opted out of every MCP server. This overrides # team inheritance and additive grants (mirrors no-default-models). - if ( - SpecialMCPServerNames.no_mcp_servers.value - in allowed_mcp_servers_for_key - ): + if SpecialMCPServerNames.no_mcp_servers.value in allowed_mcp_servers_for_key: return [] - allowed_mcp_servers_for_team = ( - await MCPRequestHandler._get_allowed_mcp_servers_for_team( - user_api_key_auth - ) - ) + allowed_mcp_servers_for_team = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_api_key_auth) - key_access_group_grants = ( - await MCPRequestHandler._get_key_access_group_mcp_server_extras( - user_api_key_auth - ) - ) + key_access_group_grants = await MCPRequestHandler._get_key_access_group_mcp_server_extras(user_api_key_auth) ######################################################### # Calculate key/team allowed servers using inheritance and intersection logic @@ -684,11 +610,7 @@ class MCPRequestHandler: # team is a ceiling rather than a default, so the key must # grant servers explicitly (or via an access group) to reach # any — it inherits none. - base = ( - set() - if general_settings.get("require_key_mcp_access_defined", False) - else team_set - ) + base = set() if general_settings.get("require_key_mcp_access_defined", False) else team_set else: base = key_set & team_set # both restrict → intersect @@ -701,10 +623,8 @@ class MCPRequestHandler: # Check end_user permissions if end_user_id is set ######################################################### if user_api_key_auth and user_api_key_auth.end_user_id: - allowed_mcp_servers_for_end_user = ( - await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( - user_api_key_auth - ) + allowed_mcp_servers_for_end_user = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( + user_api_key_auth ) # If end_user has explicit MCP server permissions, apply intersection @@ -735,19 +655,13 @@ class MCPRequestHandler: # Check agent permissions if agent_id is set on the key ######################################################### if user_api_key_auth and user_api_key_auth.agent_id: - allowed_mcp_servers_for_agent = ( - await MCPRequestHandler._get_allowed_mcp_servers_for_agent( - user_api_key_auth - ) + allowed_mcp_servers_for_agent = await MCPRequestHandler._get_allowed_mcp_servers_for_agent( + user_api_key_auth ) if len(allowed_mcp_servers_for_agent) > 0: has_lower_level_mcp_restrictions = True # Intersect: agent can only use servers allowed by BOTH key/team AND agent config - allowed_mcp_servers = [ - s - for s in allowed_mcp_servers - if s in allowed_mcp_servers_for_agent - ] + allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_agent] verbose_logger.debug( f"Applied agent intersection filter. Final allowed servers: {allowed_mcp_servers}" ) @@ -756,25 +670,17 @@ class MCPRequestHandler: # Apply org-level ceiling if org_id is set ######################################################### if user_api_key_auth and user_api_key_auth.org_id: - allowed_mcp_servers_for_org = ( - await MCPRequestHandler._get_allowed_mcp_servers_for_org( - user_api_key_auth - ) + allowed_mcp_servers_for_org = await MCPRequestHandler._get_allowed_mcp_servers_for_org( + user_api_key_auth ) if len(allowed_mcp_servers_for_org) > 0: if has_lower_level_mcp_restrictions: # Lower-level restrictions exist, so org can only cap them. - allowed_mcp_servers = [ - s - for s in allowed_mcp_servers - if s in allowed_mcp_servers_for_org - ] + allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_org] else: # No lower-level restrictions → org list becomes the ceiling allowed_mcp_servers = allowed_mcp_servers_for_org - verbose_logger.debug( - f"Applied org ceiling filter. Final allowed servers: {allowed_mcp_servers}" - ) + verbose_logger.debug(f"Applied org ceiling filter. Final allowed servers: {allowed_mcp_servers}") return list(set(allowed_mcp_servers)) except Exception as e: @@ -854,12 +760,8 @@ class MCPRequestHandler: try: # Get key and team object permissions (already loaded in main auth flow) - key_obj_perm = MCPRequestHandler._get_key_object_permission( - user_api_key_auth - ) - team_obj_perm = await MCPRequestHandler._get_team_object_permission( - user_api_key_auth - ) + key_obj_perm = MCPRequestHandler._get_key_object_permission(user_api_key_auth) + team_obj_perm = await MCPRequestHandler._get_team_object_permission(user_api_key_auth) # Extract tool permissions for this server. Dict keys may be # server_ids OR names/aliases; normalize to server_id-keyed form @@ -870,16 +772,12 @@ class MCPRequestHandler: ) key_tools = ( - global_mcp_server_manager.expand_tool_permissions( - key_obj_perm.mcp_tool_permissions - ).get(server_id) + global_mcp_server_manager.expand_tool_permissions(key_obj_perm.mcp_tool_permissions).get(server_id) if key_obj_perm else None ) team_tools = ( - global_mcp_server_manager.expand_tool_permissions( - team_obj_perm.mcp_tool_permissions - ).get(server_id) + global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id) if team_obj_perm else None ) @@ -899,15 +797,11 @@ class MCPRequestHandler: # Intersect with agent's tool permissions if agent_id is set if user_api_key_auth.agent_id: # Pre-fetch agent object_permission once to avoid duplicate DB query - agent_obj_perm = await MCPRequestHandler._get_agent_object_permission( - user_api_key_auth - ) - agent_tools = ( - await MCPRequestHandler._get_agent_tool_permissions_for_server( - server_id=server_id, - user_api_key_auth=user_api_key_auth, - agent_object_permission=agent_obj_perm, - ) + agent_obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + agent_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_id=server_id, + user_api_key_auth=user_api_key_auth, + agent_object_permission=agent_obj_perm, ) if agent_tools is not None: if allowed_tools is not None: @@ -919,13 +813,9 @@ class MCPRequestHandler: if user_api_key_auth.org_id: # _get_org_object_permission uses user_api_key_cache, so this is not a # fresh DB round-trip when get_allowed_mcp_servers was already called. - org_obj_perm = await MCPRequestHandler._get_org_object_permission( - user_api_key_auth - ) + org_obj_perm = await MCPRequestHandler._get_org_object_permission(user_api_key_auth) org_tools = ( - global_mcp_server_manager.expand_tool_permissions( - org_obj_perm.mcp_tool_permissions - ).get(server_id) + global_mcp_server_manager.expand_tool_permissions(org_obj_perm.mcp_tool_permissions).get(server_id) if org_obj_perm and org_obj_perm.mcp_tool_permissions else None ) @@ -1027,9 +917,7 @@ class MCPRequestHandler: # Permission entries may be server_ids OR names/aliases — expand to ids. return global_mcp_server_manager.expand_permission_list(raw_server_ids) except Exception as e: - verbose_logger.warning( - f"Failed to get key access group MCP server grants: {str(e)}" - ) + verbose_logger.warning(f"Failed to get key access group MCP server grants: {str(e)}") return [] @staticmethod @@ -1061,14 +949,8 @@ class MCPRequestHandler: ) # Get key object permission (already loaded in main auth flow, or fetch from DB) - key_object_permission = MCPRequestHandler._get_key_object_permission( - user_api_key_auth - ) - if ( - key_object_permission is None - and user_api_key_auth.object_permission_id - and prisma_client is not None - ): + key_object_permission = MCPRequestHandler._get_key_object_permission(user_api_key_auth) + if key_object_permission is None and user_api_key_auth.object_permission_id and prisma_client is not None: key_object_permission = await get_object_permission( object_permission_id=user_api_key_auth.object_permission_id, prisma_client=prisma_client, @@ -1081,9 +963,7 @@ class MCPRequestHandler: # Sentinel opt-out: surface it unexpanded so the caller can short-circuit # to zero servers instead of inheriting the team. - if SpecialMCPServerNames.no_mcp_servers.value in ( - key_object_permission.mcp_servers or [] - ): + if SpecialMCPServerNames.no_mcp_servers.value in (key_object_permission.mcp_servers or []): return [SpecialMCPServerNames.no_mcp_servers.value] # Permission entries may be server_ids OR names/aliases — expand to ids. @@ -1092,26 +972,20 @@ class MCPRequestHandler: ) # Get MCP servers from access groups - access_group_servers = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - key_object_permission.mcp_access_groups or [] - ) + access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( + key_object_permission.mcp_access_groups or [] ) # servers referenced in tool permissions should also be accessible tool_perm_servers = list( - global_mcp_server_manager.expand_tool_permissions( - key_object_permission.mcp_tool_permissions - ).keys() + global_mcp_server_manager.expand_tool_permissions(key_object_permission.mcp_tool_permissions).keys() ) # Combine all lists all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers return list(set(all_servers)) except Exception as e: - verbose_logger.warning( - f"Failed to get allowed MCP servers for key: {str(e)}" - ) + verbose_logger.warning(f"Failed to get allowed MCP servers for key: {str(e)}") return [] @staticmethod @@ -1143,11 +1017,7 @@ class MCPRequestHandler: user_api_key_cache, ) - if ( - user_api_key_auth is None - or not user_api_key_auth.team_id - or prisma_client is None - ): + if user_api_key_auth is None or not user_api_key_auth.team_id or prisma_client is None: return [] team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( @@ -1171,33 +1041,22 @@ class MCPRequestHandler: if object_permissions is None: return list(set(team_access_group_servers)) - direct_mcp_servers = global_mcp_server_manager.expand_permission_list( - object_permissions.mcp_servers or [] - ) + direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) - legacy_access_group_servers = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - object_permissions.mcp_access_groups or [] - ) + legacy_access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( + object_permissions.mcp_access_groups or [] ) tool_perm_servers = list( - global_mcp_server_manager.expand_tool_permissions( - object_permissions.mcp_tool_permissions - ).keys() + global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys() ) all_servers = ( - direct_mcp_servers - + legacy_access_group_servers - + tool_perm_servers - + team_access_group_servers + direct_mcp_servers + legacy_access_group_servers + tool_perm_servers + team_access_group_servers ) return list(set(all_servers)) except Exception as e: - verbose_logger.warning( - f"Failed to get allowed MCP servers for team: {str(e)}" - ) + verbose_logger.warning(f"Failed to get allowed MCP servers for team: {str(e)}") return [] @staticmethod @@ -1257,9 +1116,7 @@ class MCPRequestHandler: An empty result means the org places no restriction (allow-all from this level). """ try: - object_permissions = await MCPRequestHandler._get_org_object_permission( - user_api_key_auth - ) + object_permissions = await MCPRequestHandler._get_org_object_permission(user_api_key_auth) if object_permissions is None: return [] @@ -1269,28 +1126,20 @@ class MCPRequestHandler: ) # Expand names/aliases to canonical server IDs (consistent with key/team/end-user path) - direct_mcp_servers = global_mcp_server_manager.expand_permission_list( - object_permissions.mcp_servers or [] - ) + direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) - access_group_servers = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - object_permissions.mcp_access_groups or [] - ) + access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( + object_permissions.mcp_access_groups or [] ) tool_perm_servers = list( - global_mcp_server_manager.expand_tool_permissions( - object_permissions.mcp_tool_permissions - ).keys() + global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys() ) all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers return list(set(all_servers)) except Exception as e: - verbose_logger.warning( - f"Failed to get allowed MCP servers for org: {str(e)}" - ) + verbose_logger.warning(f"Failed to get allowed MCP servers for org: {str(e)}") return [] @staticmethod @@ -1340,10 +1189,8 @@ class MCPRequestHandler: ) # Get MCP servers from access groups - access_group_servers = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - end_user_obj.object_permission.mcp_access_groups or [] - ) + access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( + end_user_obj.object_permission.mcp_access_groups or [] ) # servers referenced in tool permissions should also be accessible @@ -1357,9 +1204,7 @@ class MCPRequestHandler: all_servers = direct_mcp_servers + access_group_servers + tool_perm_servers return list(set(all_servers)) except Exception as e: - verbose_logger.warning( - f"Failed to get allowed MCP servers for end_user: {str(e)}" - ) + verbose_logger.warning(f"Failed to get allowed MCP servers for end_user: {str(e)}") return [] # Sentinel stored in cache when an agent has no object_permission, so we @@ -1394,9 +1239,7 @@ class MCPRequestHandler: cache_key = f"agent_object_permission_id:{agent_id}" try: - object_permission_id: Optional[ - str - ] = await user_api_key_cache.async_get_cache(key=cache_key) + object_permission_id: Optional[str] = await user_api_key_cache.async_get_cache(key=cache_key) if object_permission_id == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: return None @@ -1406,14 +1249,11 @@ class MCPRequestHandler: where={"agent_id": agent_id}, ) object_permission_id = ( - getattr(agent_row, "object_permission_id", None) - if agent_row is not None - else None + getattr(agent_row, "object_permission_id", None) if agent_row is not None else None ) await user_api_key_cache.async_set_cache( key=cache_key, - value=object_permission_id - or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, + value=object_permission_id or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, ttl=get_management_object_ttl(user_api_key_cache), ) if not object_permission_id: @@ -1452,9 +1292,7 @@ class MCPRequestHandler: try: obj_perm = agent_object_permission if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission( - user_api_key_auth - ) + obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) if obj_perm is None: return [] @@ -1470,21 +1308,13 @@ class MCPRequestHandler: global_mcp_server_manager, ) - expanded_direct_servers = global_mcp_server_manager.expand_permission_list( - list(direct_mcp_servers) - ) + expanded_direct_servers = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers)) - access_group_servers = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - mcp_access_groups - ) - ) + access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups) all_servers = expanded_direct_servers + access_group_servers return list(set(all_servers)) except Exception as e: - verbose_logger.warning( - f"Failed to get allowed MCP servers for agent: {str(e)}" - ) + verbose_logger.warning(f"Failed to get allowed MCP servers for agent: {str(e)}") return [] @staticmethod @@ -1509,9 +1339,7 @@ class MCPRequestHandler: try: obj_perm = agent_object_permission if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission( - user_api_key_auth - ) + obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) if obj_perm is None: return None @@ -1523,20 +1351,14 @@ class MCPRequestHandler: global_mcp_server_manager, ) - tools = global_mcp_server_manager.expand_tool_permissions( - mcp_tool_permissions - ).get(server_id) + tools = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id) return list(tools) if tools else None except Exception as e: - verbose_logger.warning( - f"Failed to get agent tool permissions for server: {str(e)}" - ) + verbose_logger.warning(f"Failed to get agent tool permissions for server: {str(e)}") return None @staticmethod - def _get_config_server_ids_for_access_groups( - config_mcp_servers, access_groups: List[str] - ) -> Set[str]: + def _get_config_server_ids_for_access_groups(config_mcp_servers, access_groups: List[str]) -> Set[str]: """ Helper to get server_ids from config-loaded servers that match any of the given access groups. """ @@ -1548,9 +1370,7 @@ class MCPRequestHandler: return server_ids @staticmethod - async def _get_db_server_ids_for_access_groups( - prisma_client, access_groups: List[str] - ) -> Set[str]: + async def _get_db_server_ids_for_access_groups(prisma_client, access_groups: List[str]) -> Set[str]: """ Helper to get server_ids from DB servers that match any of the given access groups. """ @@ -1563,9 +1383,7 @@ class MCPRequestHandler: for server in mcp_servers: server_ids.add(server.server_id) except Exception as e: - verbose_logger.debug( - f"Error getting MCP servers from access groups: {e}" - ) + verbose_logger.debug(f"Error getting MCP servers from access groups: {e}") return server_ids @staticmethod @@ -1589,18 +1407,12 @@ class MCPRequestHandler: ) # Use the new helper for DB servers - db_server_ids = ( - await MCPRequestHandler._get_db_server_ids_for_access_groups( - prisma_client, access_groups - ) - ) + db_server_ids = await MCPRequestHandler._get_db_server_ids_for_access_groups(prisma_client, access_groups) server_ids.update(db_server_ids) return list(server_ids) except Exception as e: - verbose_logger.warning( - f"Failed to get MCP servers from access groups: {str(e)}" - ) + verbose_logger.warning(f"Failed to get MCP servers from access groups: {str(e)}") return [] @staticmethod @@ -1611,12 +1423,8 @@ class MCPRequestHandler: Get list of MCP access groups for the given user/key based on permissions """ access_groups: List[str] = [] - access_groups_for_key = await MCPRequestHandler._get_mcp_access_groups_for_key( - user_api_key_auth - ) - access_groups_for_team = ( - await MCPRequestHandler._get_mcp_access_groups_for_team(user_api_key_auth) - ) + access_groups_for_key = await MCPRequestHandler._get_mcp_access_groups_for_key(user_api_key_auth) + access_groups_for_team = await MCPRequestHandler._get_mcp_access_groups_for_team(user_api_key_auth) ######################################################### # If team has access groups, then key must have a subset of the team's access groups @@ -1709,9 +1517,7 @@ class MCPRequestHandler: return object_permissions.mcp_access_groups or [] except Exception as e: - verbose_logger.warning( - f"Failed to get MCP access groups for team: {str(e)}" - ) + verbose_logger.warning(f"Failed to get MCP access groups for team: {str(e)}") return [] @staticmethod @@ -1719,14 +1525,10 @@ class MCPRequestHandler: """ Extract and parse the x-mcp-access-groups header as a list of strings. """ - mcp_access_groups_header = headers.get( - MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME - ) + mcp_access_groups_header = headers.get(MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME) if mcp_access_groups_header is not None: try: - return [ - s.strip() for s in mcp_access_groups_header.split(",") if s.strip() - ] + return [s.strip() for s in mcp_access_groups_header.split(",") if s.strip()] except Exception: return None return None diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 2f5973ca371..4f58f4bdbb3 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -79,9 +79,7 @@ def _oauth_token_error(code: str, status: int = 400) -> JSONResponse: FastAPI's default ``HTTPException`` renders ``{"detail": ...}`` which spec-compliant OAuth clients parsing the ``error`` field won't recognize. """ - return JSONResponse( - status_code=status, content={"error": code}, headers=TOKEN_NO_CACHE_HEADERS - ) + return JSONResponse(status_code=status, content={"error": code}, headers=TOKEN_NO_CACHE_HEADERS) def _user_id_from_session_cookie(request: Request) -> Optional[str]: @@ -160,8 +158,7 @@ def _build_authorize_html( # Build access checklist rows access_rows = "".join( - f'
✓{e(item)}
' - for item in access_items + f'
✓{e(item)}
' for item in access_items ) access_section = "" if access_rows: @@ -177,7 +174,9 @@ def _build_authorize_html( # Help link for step 2 help_link_html = "" if help_url: - help_link_html = f'Where do I find my API key? ↗' + help_link_html = ( + f'Where do I find my API key? ↗' + ) return f""" @@ -722,14 +721,10 @@ async def byok_authorize_post( # Reject new codes if the store is at capacity (prevents memory exhaustion # from a burst of abandoned OAuth flows). if len(_byok_auth_codes) >= _AUTH_CODES_MAX_SIZE: - raise HTTPException( - status_code=503, detail="Too many pending authorization flows" - ) + raise HTTPException(status_code=503, detail="Too many pending authorization flows") if code_challenge_method != "S256": - raise HTTPException( - status_code=400, detail="Only S256 code_challenge_method is supported" - ) + raise HTTPException(status_code=400, detail="Only S256 code_challenge_method is supported") # Identity comes from the authenticated session, not the OAuth client_id # form field (RFC 6749 §2.2: client_id identifies the client application, @@ -806,11 +801,7 @@ async def byok_token( # actually submitted a value, so we stay RFC 6749-backward-compatible # without breaking OAuth 2.1 clients. PKCE + client_id binding # (checked below) cover the security role redirect_uri played. - if ( - record.get("redirect_uri") - and redirect_uri - and redirect_uri != record["redirect_uri"] - ): + if record.get("redirect_uri") and redirect_uri and redirect_uri != record["redirect_uri"]: return _oauth_token_error("invalid_grant") # RFC 6749 §4.1.3: if the client was identified at /authorize, the @@ -865,9 +856,7 @@ async def byok_token( ) return _oauth_token_error("server_error", status=500) else: - verbose_proxy_logger.warning( - "byok_token: prisma_client is None — credential not persisted" - ) + verbose_proxy_logger.warning("byok_token: prisma_client is None — credential not persisted") now = int(time.time()) payload = { diff --git a/litellm/proxy/_experimental/mcp_server/cost_calculator.py b/litellm/proxy/_experimental/mcp_server/cost_calculator.py index b8fdba23d92..9b6f89bc7bd 100644 --- a/litellm/proxy/_experimental/mcp_server/cost_calculator.py +++ b/litellm/proxy/_experimental/mcp_server/cost_calculator.py @@ -32,9 +32,7 @@ class MCPCostCalculator: # Get the response cost from logging object model_call_details # This is set when a user modifies the response in a post_mcp_tool_call_hook ######################################################### - response_cost = litellm_logging_obj.model_call_details.get( - "response_cost", None - ) + response_cost = litellm_logging_obj.model_call_details.get("response_cost", None) if response_cost is not None: return response_cost @@ -44,9 +42,7 @@ class MCPCostCalculator: mcp_tool_call_metadata: StandardLoggingMCPToolCall = ( cast( StandardLoggingMCPToolCall, - litellm_logging_obj.model_call_details.get( - "mcp_tool_call_metadata", {} - ), + litellm_logging_obj.model_call_details.get("mcp_tool_call_metadata", {}), ) or {} ) @@ -56,12 +52,8 @@ class MCPCostCalculator: ######################################################### # User defined cost per query ######################################################### - default_cost_per_query = mcp_server_cost_info.get( - "default_cost_per_query", None - ) - tool_name_to_cost_per_query: dict = ( - mcp_server_cost_info.get("tool_name_to_cost_per_query", {}) or {} - ) + default_cost_per_query = mcp_server_cost_info.get("default_cost_per_query", None) + tool_name_to_cost_per_query: dict = mcp_server_cost_info.get("tool_name_to_cost_per_query", {}) or {} tool_name = mcp_tool_call_metadata.get("name", "") ######################################################### diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 668e936b8a2..2dd046ceada 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -169,14 +169,11 @@ def _reencrypt_global_env_var_values( ) if decrypted is None: verbose_proxy_logger.warning( - "rotate_mcp_server_credentials_master_key: could not decrypt " - "global env var %s, skipping", + "rotate_mcp_server_credentials_master_key: could not decrypt global env var %s, skipping", entry.get("name"), ) continue - entry["value"] = encrypt_value_helper( - decrypted, new_encryption_key=new_encryption_key - ) + entry["value"] = encrypt_value_helper(decrypted, new_encryption_key=new_encryption_key) rotated = True return rebuilt if rotated else None @@ -240,9 +237,7 @@ def _prepare_mcp_server_data( # Handle credentials serialization credentials = data_dict.get("credentials") if credentials is not None: - data_dict["credentials"] = encrypt_credentials( - credentials=credentials, encryption_key=_get_salt_key() - ) + data_dict["credentials"] = encrypt_credentials(credentials=credentials, encryption_key=_get_salt_key()) data_dict["credentials"] = safe_dumps(data_dict["credentials"]) # Serialize JSON fields from ``data_dict`` (not ``data``) so the @@ -268,13 +263,9 @@ def _prepare_mcp_server_data( data_dict["env"] = safe_dumps(data_dict["env"]) if "tool_name_to_display_name" in data_dict: - data_dict["tool_name_to_display_name"] = safe_dumps( - data_dict["tool_name_to_display_name"] or {} - ) + data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {}) if "tool_name_to_description" in data_dict: - data_dict["tool_name_to_description"] = safe_dumps( - data_dict["tool_name_to_description"] or {} - ) + data_dict["tool_name_to_description"] = safe_dumps(data_dict["tool_name_to_description"] or {}) # mcp_access_groups is already List[str], no serialization needed @@ -286,9 +277,7 @@ def _prepare_mcp_server_data( return data_dict -def encrypt_credentials( - credentials: MCPCredentials, encryption_key: Optional[str] -) -> MCPCredentials: +def encrypt_credentials(credentials: MCPCredentials, encryption_key: Optional[str]) -> MCPCredentials: auth_value = credentials.get("auth_value") if auth_value is not None: credentials["auth_value"] = encrypt_value_helper( @@ -366,35 +355,24 @@ async def get_all_mcp_servers( where: Dict[str, Any] = {} if approval_status is not None: where["approval_status"] = approval_status - mcp_servers = await MCPServerRepository(prisma_client).table.find_many( - where=where if where else {} - ) + mcp_servers = await MCPServerRepository(prisma_client).table.find_many(where=where if where else {}) - tables = [ - LiteLLM_MCPServerTable(**mcp_server.model_dump()) - for mcp_server in mcp_servers - ] + tables = [LiteLLM_MCPServerTable(**mcp_server.model_dump()) for mcp_server in mcp_servers] for table in tables: decrypt_global_env_var_values(table.env_vars) return tables except Exception as e: verbose_proxy_logger.debug( - "litellm.proxy._experimental.mcp_server.db.py::get_all_mcp_servers - {}".format( - str(e) - ) + "litellm.proxy._experimental.mcp_server.db.py::get_all_mcp_servers - {}".format(str(e)) ) return [] -async def get_mcp_server( - prisma_client: PrismaClient, server_id: str -) -> Optional[LiteLLM_MCPServerTable]: +async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> Optional[LiteLLM_MCPServerTable]: """ Returns the matching mcp server from the db iff exists """ - mcp_server: Optional[LiteLLM_MCPServerTable] = await MCPServerRepository( - prisma_client - ).table.find_unique( + mcp_server: Optional[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_unique( where={ "server_id": server_id, } @@ -406,15 +384,11 @@ async def get_mcp_server( return table -async def get_mcp_servers( - prisma_client: PrismaClient, server_ids: Iterable[str] -) -> List[LiteLLM_MCPServerTable]: +async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]) -> List[LiteLLM_MCPServerTable]: """ Returns the matching mcp servers from the db with the server_ids """ - _mcp_servers: List[LiteLLM_MCPServerTable] = await MCPServerRepository( - prisma_client - ).table.find_many( + _mcp_servers: List[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.find_many( where={ "server_id": {"in": server_ids}, } @@ -428,15 +402,11 @@ async def get_mcp_servers( return final_mcp_servers -async def get_mcp_servers_by_verificationtoken( - prisma_client: PrismaClient, token: str -) -> List[str]: +async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> List[str]: """ Returns the mcp servers from the db for the verification token """ - verification_token_record: LiteLLM_TeamTable = await VerificationTokenRepository( - prisma_client - ).table.find_unique( + verification_token_record: LiteLLM_TeamTable = await VerificationTokenRepository(prisma_client).table.find_unique( where={ "token": token, }, @@ -446,23 +416,16 @@ async def get_mcp_servers_by_verificationtoken( ) mcp_servers: Optional[List[str]] = [] - if ( - verification_token_record is not None - and verification_token_record.object_permission is not None - ): + if verification_token_record is not None and verification_token_record.object_permission is not None: mcp_servers = verification_token_record.object_permission.mcp_servers return mcp_servers or [] -async def get_mcp_servers_by_team( - prisma_client: PrismaClient, team_id: str -) -> List[str]: +async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> List[str]: """ Returns the mcp servers from the db for the team id """ - team_record: LiteLLM_TeamTable = await TeamRepository( - prisma_client - ).table.find_unique( + team_record: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.find_unique( where={ "team_id": team_id, }, @@ -492,19 +455,12 @@ async def get_all_mcp_servers_for_user( # Get the mcp servers for the key if user.api_key: - token_mcp_servers = await get_mcp_servers_by_verificationtoken( - prisma_client, user.api_key - ) + token_mcp_servers = await get_mcp_servers_by_verificationtoken(prisma_client, user.api_key) mcp_server_ids.update(token_mcp_servers) # check for special team membership - if ( - SpecialMCPServerName.all_team_servers in mcp_server_ids - and user.team_id is not None - ): - team_mcp_servers = await get_mcp_servers_by_team( - prisma_client, user.team_id - ) + if SpecialMCPServerName.all_team_servers in mcp_server_ids and user.team_id is not None: + team_mcp_servers = await get_mcp_servers_by_team(prisma_client, user.team_id) mcp_server_ids.update(team_mcp_servers) if len(mcp_server_ids) > 0: @@ -519,9 +475,7 @@ async def get_objectpermissions_for_mcp_server( """ Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server """ - object_permission_records = await ObjectPermissionRepository( - prisma_client - ).table.find_many( + object_permission_records = await ObjectPermissionRepository(prisma_client).table.find_many( where={ "mcp_servers": {"has": mcp_server_id}, }, @@ -534,9 +488,7 @@ async def get_objectpermissions_for_mcp_server( return object_permission_records -async def get_virtualkeys_for_mcp_server( - prisma_client: PrismaClient, server_id: str -) -> List: +async def get_virtualkeys_for_mcp_server(prisma_client: PrismaClient, server_id: str) -> List: """ Get all the virtual keys that have access to the mcp server """ @@ -565,9 +517,7 @@ async def delete_mcp_server_from_virtualkey(): pass -async def delete_mcp_server( - prisma_client: PrismaClient, server_id: str -) -> Optional[LiteLLM_MCPServerTable]: +async def delete_mcp_server(prisma_client: PrismaClient, server_id: str) -> Optional[LiteLLM_MCPServerTable]: """ Delete the mcp server from the db by server_id @@ -644,19 +594,13 @@ async def update_mcp_server( # exclude_unset=True makes this a true partial update: fields the caller did # not provide are not written, so they keep their existing DB value instead # of being reset to a schema default (transport=sse, allow_all_keys=False...). - data_dict = _prepare_mcp_server_data( - data, exclude_unset=True, fields_set=fields_set - ) + data_dict = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set) # Pre-fetch existing record once if we need it for auth_type or credential logic existing = None - has_credentials = ( - "credentials" in data_dict and data_dict["credentials"] is not None - ) + has_credentials = "credentials" in data_dict and data_dict["credentials"] is not None if data.auth_type or has_credentials: - existing = await MCPServerRepository(prisma_client).table.find_unique( - where={"server_id": data.server_id} - ) + existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id}) # Clear stale credentials when auth_type changes but no new credentials provided if ( @@ -676,9 +620,7 @@ async def update_mcp_server( # Only merge when auth_type is unchanged. Switching auth types # (e.g. oauth2 → api_key) should replace credentials entirely # to avoid stale secrets from the previous auth type lingering. - auth_type_unchanged = ( - data.auth_type is None or data.auth_type == existing.auth_type - ) + auth_type_unchanged = data.auth_type is None or data.auth_type == existing.auth_type if auth_type_unchanged: existing_creds = ( json.loads(existing.credentials) @@ -706,9 +648,7 @@ async def update_mcp_server( return updated_mcp_server -async def rotate_mcp_server_credentials_master_key( - prisma_client: PrismaClient, touched_by: str, new_master_key: str -): +async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps mcp_servers = await MCPServerRepository(prisma_client).table.find_many() @@ -729,9 +669,7 @@ async def rotate_mcp_server_credentials_master_key( ) update_data["credentials"] = safe_dumps(encrypted_credentials) - rotated_env_vars = _reencrypt_global_env_var_values( - mcp_server.env_vars, new_master_key - ) + rotated_env_vars = _reencrypt_global_env_var_values(mcp_server.env_vars, new_master_key) if rotated_env_vars is not None: update_data["env_vars"] = safe_dumps(rotated_env_vars) @@ -791,9 +729,7 @@ def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: return None -async def rotate_mcp_user_credentials_master_key( - prisma_client: PrismaClient, new_master_key: str -): +async def rotate_mcp_user_credentials_master_key(prisma_client: PrismaClient, new_master_key: str): """Re-encrypt every ``LiteLLM_MCPUserCredentials`` row with ``new_master_key``. Reads each ``credential_b64`` with the current salt key (falling back to @@ -815,9 +751,7 @@ async def rotate_mcp_user_credentials_master_key( ) skipped += 1 continue - re_encrypted = encrypt_value_helper( - plaintext, new_encryption_key=new_master_key - ) + re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key) await MCPUserCredentialsRepository(prisma_client).table.update( where={ "user_id_server_id": { @@ -835,9 +769,7 @@ async def rotate_mcp_user_credentials_master_key( ) -async def rotate_mcp_user_env_vars_master_key( - prisma_client: PrismaClient, new_master_key: str -): +async def rotate_mcp_user_env_vars_master_key(prisma_client: PrismaClient, new_master_key: str): """Re-encrypt every ``LiteLLM_MCPUserEnvVars`` row with ``new_master_key``. Reads each ``values_b64`` blob with the current salt key and writes it back @@ -857,16 +789,13 @@ async def rotate_mcp_user_env_vars_master_key( ) if plaintext is None: verbose_proxy_logger.warning( - "rotate_mcp_user_env_vars_master_key: could not decrypt env vars " - "for user_id=%s server_id=%s, skipping", + "rotate_mcp_user_env_vars_master_key: could not decrypt env vars for user_id=%s server_id=%s, skipping", row.user_id, row.server_id, ) skipped += 1 continue - re_encrypted = encrypt_value_helper( - plaintext, new_encryption_key=new_master_key - ) + re_encrypted = encrypt_value_helper(plaintext, new_encryption_key=new_master_key) await prisma_client.db.litellm_mcpuserenvvars.update( where={ "user_id_server_id": { @@ -966,9 +895,7 @@ async def store_user_oauth_credential( expires_at: Optional[str] = None if expires_in is not None: - expires_at = ( - datetime.now(timezone.utc) + timedelta(seconds=expires_in) - ).isoformat() + expires_at = (datetime.now(timezone.utc) + timedelta(seconds=expires_in)).isoformat() payload: Dict[str, Any] = { "type": "oauth2", @@ -989,10 +916,7 @@ async def store_user_oauth_credential( existing = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - if ( - existing is not None - and _decode_oauth_payload(existing.credential_b64) is None - ): + if existing is not None and _decode_oauth_payload(existing.credential_b64) is None: # Existing row is either a BYOK secret or an OAuth2 row that no # longer decrypts (e.g. after a salt-key rotation). In either # case, refuse to overwrite — the caller would clobber data @@ -1059,9 +983,7 @@ async def list_user_oauth_credentials( ) -> List[Dict[str, Any]]: """Return all OAuth2 credential payloads for a user, tagged with server_id.""" - rows = await MCPUserCredentialsRepository(prisma_client).table.find_many( - where={"user_id": user_id} - ) + rows = await MCPUserCredentialsRepository(prisma_client).table.find_many(where={"user_id": user_id}) results: List[Dict[str, Any]] = [] for row in rows: payload = _decode_oauth_payload(row.credential_b64) @@ -1118,9 +1040,7 @@ async def refresh_user_oauth_token( token_data["client_secret"] = client_secret try: - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.Oauth2Check - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response = await async_client.post( token_url, headers={"Accept": "application/json"}, @@ -1140,8 +1060,7 @@ async def refresh_user_oauth_token( access_token: Optional[str] = body.get("access_token") if not access_token: verbose_proxy_logger.warning( - "refresh_user_oauth_token: token response missing access_token for " - "user=%s server=%s", + "refresh_user_oauth_token: token response missing access_token for user=%s server=%s", user_id, server_id, ) @@ -1158,9 +1077,9 @@ async def refresh_user_oauth_token( new_refresh_token: Optional[str] = body.get("refresh_token") or refresh_token raw_scope = body.get("scope") - scopes: Optional[List[str]] = ( - raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None - ) or cred.get("scopes") + scopes: Optional[List[str]] = (raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None) or cred.get( + "scopes" + ) await store_user_oauth_credential( prisma_client=prisma_client, @@ -1202,18 +1121,14 @@ async def resolve_valid_user_oauth_token( """ if not cred or not cred.get("access_token"): return None - if not is_oauth_credential_expired( - cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS - ): + if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): return cred if not cred.get("refresh_token"): return None if prisma_client is None: from litellm.proxy.utils import get_prisma_client_or_throw - prisma_client = get_prisma_client_or_throw( - "Database not connected. Cannot refresh OAuth token." - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Cannot refresh OAuth token.") refreshed = await refresh_user_oauth_token( prisma_client=prisma_client, user_id=user_id, @@ -1281,9 +1196,7 @@ async def resolve_user_oauth_access_token( access_token: str = cred["access_token"] if prefetched_creds is None: - ttl = _compute_per_user_token_ttl( - server, _remaining_token_seconds(cred.get("expires_at")) - ) + ttl = _compute_per_user_token_ttl(server, _remaining_token_seconds(cred.get("expires_at"))) await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl) return access_token except Exception as e: @@ -1371,9 +1284,7 @@ async def get_mcp_submissions( for item in items: decrypt_global_env_var_values(item.env_vars) - pending = sum( - 1 for i in items if i.approval_status == MCPApprovalStatus.pending_review - ) + pending = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review) active = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active) rejected = sum(1 for i in items if i.approval_status == MCPApprovalStatus.rejected) @@ -1440,9 +1351,7 @@ async def get_user_env_vars_bulk( ids = list(server_ids) if not ids: return {} - rows = await prisma_client.db.litellm_mcpuserenvvars.find_many( - where={"user_id": user_id, "server_id": {"in": ids}} - ) + rows = await prisma_client.db.litellm_mcpuserenvvars.find_many(where={"user_id": user_id, "server_id": {"in": ids}}) return {row.server_id: _decode_user_env_vars(row.values_b64) for row in rows} @@ -1499,6 +1408,4 @@ async def delete_user_env_vars( Uses ``delete_many`` so a missing row is a no-op; real DB errors still propagate to the caller instead of being silently swallowed. """ - await prisma_client.db.litellm_mcpuserenvvars.delete_many( - where={"user_id": user_id, "server_id": server_id} - ) + await prisma_client.db.litellm_mcpuserenvvars.delete_many(where={"user_id": user_id, "server_id": server_id}) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3beddd2c435..7a8df83f9f9 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -50,9 +50,7 @@ router = APIRouter( def _prune_oauth_metadata_cache(now: Optional[float] = None) -> None: now = now if now is not None else time.time() expired_cache_keys = [ - cache_key - for cache_key, (expires_at, _payload) in _OAUTH_METADATA_CACHE.items() - if expires_at <= now + cache_key for cache_key, (expires_at, _payload) in _OAUTH_METADATA_CACHE.items() if expires_at <= now ] for cache_key in expired_cache_keys: _OAUTH_METADATA_CACHE.pop(cache_key, None) @@ -130,9 +128,7 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data -def _get_validated_client_redirect_uri( - request: Request, state_data: Dict[str, Any] -) -> str: +def _get_validated_client_redirect_uri(request: Request, state_data: Dict[str, Any]) -> str: """Return a trusted (same-origin, loopback, or ops-allowlisted) client redirect URI from OAuth state. """ @@ -217,24 +213,17 @@ def _validate_token_response( "error": "token_validation_failed", "server_id": server_id, "field": key, - "message": ( - f"OAuth token rejected: required field '{key}' is absent" - ), + "message": (f"OAuth token rejected: required field '{key}' is absent"), }, ) - if _normalize_for_token_comparison(actual) != _normalize_for_token_comparison( - expected - ): + if _normalize_for_token_comparison(actual) != _normalize_for_token_comparison(expected): raise HTTPException( status_code=403, detail={ "error": "token_validation_failed", "server_id": server_id, "field": key, - "message": ( - f"OAuth token rejected: '{key}' = '{actual}', " - f"expected '{expected}'" - ), + "message": (f"OAuth token rejected: '{key}' = '{actual}', expected '{expected}'"), }, ) @@ -247,9 +236,7 @@ async def _extract_user_id_from_request(request: Request) -> Optional[str]: auth pipeline (which has side effects such as rate-limit increments and spend logging). Returns ``None`` if no cached credential is found. """ - auth_header = request.headers.get("Authorization") or request.headers.get( - "authorization" - ) + auth_header = request.headers.get("Authorization") or request.headers.get("authorization") if not auth_header: return None lower = auth_header.lower() @@ -289,22 +276,16 @@ async def _store_per_user_token_server_side( raw_expires = token_response.get("expires_in") try: - expires_in: Optional[int] = ( - int(raw_expires) if raw_expires is not None else None - ) + expires_in: Optional[int] = int(raw_expires) if raw_expires is not None else None except (TypeError, ValueError): expires_in = None refresh_token: Optional[str] = token_response.get("refresh_token") or None raw_scope = token_response.get("scope") - scopes: Optional[list] = ( - raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None - ) + scopes: Optional[list] = raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None try: - prisma_client = get_prisma_client_or_throw( - "Database not connected. Cannot store per-user OAuth token." - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Cannot store per-user OAuth token.") from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 store_user_oauth_credential, ) @@ -356,9 +337,7 @@ async def authorize_with_server( if mcp_server.auth_type != "oauth2": raise HTTPException(status_code=400, detail="MCP server is not OAuth2") if mcp_server.authorization_url is None: - raise HTTPException( - status_code=400, detail="MCP server authorization url is not set" - ) + raise HTTPException(status_code=400, detail="MCP server authorization url is not set") # Trusted redirect_uri: same-origin, loopback, or ops-allowlisted. # The URI is encrypted into the OAuth state and decoded on @@ -418,9 +397,7 @@ async def exchange_token_with_server( raise HTTPException(status_code=400, detail="MCP server token url is not set") resolved_client_id = mcp_server.client_id if mcp_server.client_id else client_id - resolved_client_secret = ( - mcp_server.client_secret if mcp_server.client_secret else client_secret - ) + resolved_client_secret = mcp_server.client_secret if mcp_server.client_secret else client_secret if grant_type == "refresh_token": if not refresh_token: @@ -494,8 +471,7 @@ async def exchange_token_with_server( ) except Exception as exc: verbose_logger.warning( - "exchange_token_with_server: server-side storage failed " - "for user=%s server=%s: %s", + "exchange_token_with_server: server-side storage failed for user=%s server=%s: %s", user_id, mcp_server.server_id, exc, @@ -545,9 +521,7 @@ async def register_client_with_server( return dummy_return if mcp_server.authorization_url is None: - raise HTTPException( - status_code=400, detail="MCP server authorization url is not set" - ) + raise HTTPException(status_code=400, detail="MCP server authorization url is not set") if mcp_server.registration_url is None: return dummy_return @@ -564,9 +538,7 @@ async def register_client_with_server( "Accept": "application/json", } - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.Oauth2Register - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Register) response = await async_client.post( mcp_server.registration_url, headers=headers, @@ -605,11 +577,7 @@ async def authorize( lookup_name: Optional[str] = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = ( - global_mcp_server_manager.get_mcp_server_by_name( - lookup_name, client_ip=client_ip - ) - if lookup_name - else None + global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None ) if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) @@ -670,9 +638,7 @@ async def token_endpoint( lookup_name = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( - lookup_name, client_ip=client_ip - ) + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: @@ -781,9 +747,7 @@ async def callback( # 2. Neither success nor error parameters present — most likely a stray # GET / dropped SSO redirect chain. Surface a 400 instead of 422. if not code or not state: - missing = [ - name for name, value in (("code", code), ("state", state)) if not value - ] + missing = [name for name, value in (("code", code), ("state", state)) if not value] return _render_oauth_error_html( "invalid_request", f"Missing authorization {' and '.join(repr(m) for m in missing)} parameter(s).", @@ -811,9 +775,7 @@ async def callback( # a generic "authentication incomplete" redirect. raise except Exception: - return HTMLResponse( - "Authentication incomplete. You can close this window." - ) + return HTMLResponse("Authentication incomplete. You can close this window.") # ------------------------------ @@ -880,14 +842,9 @@ async def fetch_upstream_oauth_protected_resource( candidates = [f"{host_base}/.well-known/oauth-protected-resource"] # RFC 9728 §3.1 path fallback if upstream.path and upstream.path not in ("", "/"): - candidates.append( - f"{host_base}/.well-known/oauth-protected-resource" - f"{upstream.path.rstrip('/')}" - ) + candidates.append(f"{host_base}/.well-known/oauth-protected-resource{upstream.path.rstrip('/')}") - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.Oauth2Check - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) network_errors: list[Exception] = [] for candidate in candidates: @@ -988,9 +945,7 @@ async def _build_oauth_protected_resource_response( mcp_server: Optional[MCPServer] = None if mcp_server_name: - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( - mcp_server_name, client_ip=client_ip - ) + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) # Build resource URL based on the pattern if mcp_server_name: @@ -1007,9 +962,7 @@ async def _build_oauth_protected_resource_response( # directs the client at the real IdP (Okta, Keycloak, …) instead of us. if mcp_server is not None and mcp_server.is_oauth_passthrough: try: - upstream_metadata = await fetch_upstream_oauth_protected_resource( - mcp_server - ) + upstream_metadata = await fetch_upstream_oauth_protected_resource(mcp_server) except Exception as exc: verbose_logger.warning( "Failed to fetch upstream oauth-protected-resource metadata " @@ -1018,8 +971,7 @@ async def _build_oauth_protected_resource_response( raise HTTPException( status_code=502, detail=( - "Failed to fetch upstream oauth-protected-resource " - f"metadata for MCP server {mcp_server.name!r}" + f"Failed to fetch upstream oauth-protected-resource metadata for MCP server {mcp_server.name!r}" ), ) @@ -1032,29 +984,19 @@ async def _build_oauth_protected_resource_response( # so we must not fall through to the default gateway metadata — # that would point clients at the wrong IdP. verbose_logger.warning( - "Upstream oauth-protected-resource metadata unavailable for " - f"pass-through MCP server {mcp_server.name!r}" + f"Upstream oauth-protected-resource metadata unavailable for pass-through MCP server {mcp_server.name!r}" ) raise HTTPException( status_code=502, - detail=( - "Upstream oauth-protected-resource metadata unavailable " - f"for MCP server {mcp_server.name!r}" - ), + detail=(f"Upstream oauth-protected-resource metadata unavailable for MCP server {mcp_server.name!r}"), ) return { "authorization_servers": [ - ( - f"{request_base_url}/{mcp_server_name}" - if mcp_server_name - else f"{request_base_url}" - ) + (f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}") ], "resource": resource_url, - "scopes_supported": ( - mcp_server.scopes if mcp_server and mcp_server.scopes else [] - ), + "scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []), } @@ -1086,9 +1028,7 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp" ) @router.get("/.well-known/oauth-protected-resource") -async def oauth_protected_resource_mcp( - request: Request, mcp_server_name: Optional[str] = None -): +async def oauth_protected_resource_mcp(request: Request, mcp_server_name: Optional[str] = None): """ OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern. @@ -1129,38 +1069,26 @@ def _build_oauth_authorization_server_response( mcp_server_name = resolved.server_name or resolved.name authorization_endpoint = ( - f"{request_base_url}/{mcp_server_name}/authorize" - if mcp_server_name - else f"{request_base_url}/authorize" - ) - token_endpoint = ( - f"{request_base_url}/{mcp_server_name}/token" - if mcp_server_name - else f"{request_base_url}/token" + f"{request_base_url}/{mcp_server_name}/authorize" if mcp_server_name else f"{request_base_url}/authorize" ) + token_endpoint = f"{request_base_url}/{mcp_server_name}/token" if mcp_server_name else f"{request_base_url}/token" mcp_server: Optional[MCPServer] = None if mcp_server_name: - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( - mcp_server_name, client_ip=client_ip - ) + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) return { "issuer": request_base_url, # point to your proxy "authorization_endpoint": authorization_endpoint, "token_endpoint": token_endpoint, "response_types_supported": ["code"], - "scopes_supported": ( - mcp_server.scopes if mcp_server and mcp_server.scopes else [] - ), + "scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []), "grant_types_supported": ["authorization_code", "refresh_token"], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["client_secret_post"], # Claude expects a registration endpoint, even if we just fake it "registration_endpoint": ( - f"{request_base_url}/{mcp_server_name}/register" - if mcp_server_name - else f"{request_base_url}/register" + f"{request_base_url}/{mcp_server_name}/register" if mcp_server_name else f"{request_base_url}/register" ), } @@ -1169,9 +1097,7 @@ def _build_oauth_authorization_server_response( @router.get( f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" ) -async def oauth_authorization_server_mcp_standard( - request: Request, mcp_server_name: str -): +async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_name: str): """ OAuth authorization server discovery endpoint using standard MCP URL pattern. @@ -1189,9 +1115,7 @@ async def oauth_authorization_server_mcp_standard( f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}" ) @router.get("/.well-known/oauth-authorization-server") -async def oauth_authorization_server_mcp( - request: Request, mcp_server_name: Optional[str] = None -): +async def oauth_authorization_server_mcp(request: Request, mcp_server_name: Optional[str] = None): """ OAuth authorization server discovery endpoint. @@ -1307,9 +1231,7 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non ) return dummy_return - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( - mcp_server_name, client_ip=client_ip - ) + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) if mcp_server is None: return dummy_return return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py index e42270bf10b..030f4dfeca6 100644 --- a/litellm/proxy/_experimental/mcp_server/elicitation_handler.py +++ b/litellm/proxy/_experimental/mcp_server/elicitation_handler.py @@ -70,9 +70,7 @@ async def handle_elicitation_request( ) # No downstream session — we're in Tool Bridge mode # or the client doesn't support elicitation - verbose_logger.info( - "MCP elicitation: no downstream session available, declining" - ) + verbose_logger.info("MCP elicitation: no downstream session available, declining") return ElicitResult( action="decline", ) @@ -105,23 +103,17 @@ async def _relay_elicitation_to_downstream( if downstream_capabilities is not None: elicit_caps = getattr(downstream_capabilities, "elicitation", None) if elicit_caps is None: - verbose_logger.info( - "MCP elicitation: downstream client does not support elicitation" - ) + verbose_logger.info("MCP elicitation: downstream client does not support elicitation") return ElicitResult(action="decline") if mode == "url": url_cap = getattr(elicit_caps, "url", None) if url_cap is None: - verbose_logger.info( - "MCP elicitation: downstream client does not support URL mode" - ) + verbose_logger.info("MCP elicitation: downstream client does not support URL mode") return ElicitResult(action="decline") if mode == "form": form_cap = getattr(elicit_caps, "form", None) if form_cap is None: - verbose_logger.info( - "MCP elicitation: downstream client does not support form mode" - ) + verbose_logger.info("MCP elicitation: downstream client does not support form mode") return ElicitResult(action="decline") try: if mode == "url" and isinstance(params, ElicitRequestURLParams): @@ -145,9 +137,7 @@ async def _relay_elicitation_to_downstream( else: # Fallback for generic ElicitRequestParams — pass an empty schema # since elicit() requires requestedSchema as a positional arg. - verbose_logger.info( - "MCP elicitation: relaying generic elicitation to downstream" - ) + verbose_logger.info("MCP elicitation: relaying generic elicitation to downstream") result = await downstream_session.elicit( message=getattr(params, "message", ""), requestedSchema=getattr(params, "requestedSchema", {}), diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index 430634e802e..d2500894000 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -80,15 +80,9 @@ class MCPUpstreamAuthError(Exception): if challenge is None and self.status_code == 401 and base_url: prefix = base_url.rstrip("/") if request_path and request_path.startswith(f"/{self.server_name}/mcp"): - resource_metadata_url = ( - f"{prefix}/.well-known/oauth-protected-resource/" - f"{self.server_name}/mcp" - ) + resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/{self.server_name}/mcp" else: - resource_metadata_url = ( - f"{prefix}/.well-known/oauth-protected-resource/" - f"mcp/{self.server_name}" - ) + resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/mcp/{self.server_name}" challenge = f'Bearer resource_metadata="{resource_metadata_url}"' detail = "Forbidden" if self.status_code == 403 else "Unauthorized" return HTTPException( diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 6997f5241de..b668833e638 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -41,9 +41,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): ) -> Dict[str, Any]: mcp_tool_name = data.get("mcp_tool_name") or data.get("name") mcp_arguments = data.get("mcp_arguments") or data.get("arguments") - mcp_tool_description = data.get("mcp_tool_description") or data.get( - "description" - ) + mcp_tool_description = data.get("mcp_tool_description") or data.get("description") if mcp_arguments is None or not isinstance(mcp_arguments, dict): mcp_arguments = {} diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py index 51918509441..8a85c0c516b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_context.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -11,9 +11,7 @@ from typing import Optional # Set server-side in proxy_server.py route handlers when a request arrives via # /toolset/{name}/mcp or the toolset fallback in dynamic_mcp_route. # Never populated from client-supplied headers. -_mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar( - "_mcp_active_toolset_id", default=None -) +_mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar("_mcp_active_toolset_id", default=None) # Per-request merged InitializeResult.instructions; set in MCP HTTP/SSE handlers. _mcp_gateway_initialize_instructions: ContextVar[Optional[str]] = ContextVar( @@ -22,6 +20,4 @@ _mcp_gateway_initialize_instructions: ContextVar[Optional[str]] = ContextVar( # Per-request scoped server name; set in MCP HTTP/SSE handlers when the path # identifies exactly one upstream server. Never populated from client-supplied headers. -_mcp_gateway_server_name: ContextVar[Optional[str]] = ContextVar( - "_mcp_gateway_server_name", default=None -) +_mcp_gateway_server_name: ContextVar[Optional[str]] = ContextVar("_mcp_gateway_server_name", default=None) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 254f208e231..42e2b17d697 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -162,8 +162,7 @@ class MCPDebug: has_server_specific = bool( mcp_server_auth_headers and ( - mcp_server_auth_headers.get(server.alias or "") - or mcp_server_auth_headers.get(server.server_name or "") + mcp_server_auth_headers.get(server.alias or "") or mcp_server_auth_headers.get(server.server_name or "") ) ) if has_server_specific or mcp_auth_header: @@ -219,9 +218,7 @@ class MCPDebug: if k.lower() == hdr_name: inbound_parts.append(f"{hdr_name}={MCPDebug._mask(v)}") break - debug[f"{_RESPONSE_HEADER_PREFIX}-inbound-auth"] = ( - "; ".join(inbound_parts) if inbound_parts else "(none)" - ) + debug[f"{_RESPONSE_HEADER_PREFIX}-inbound-auth"] = "; ".join(inbound_parts) if inbound_parts else "(none)" # --- OAuth2 token --- oauth2_token = (oauth2_headers or {}).get("Authorization") @@ -230,26 +227,19 @@ class MCPDebug: litellm_raw = litellm_api_key.removeprefix("Bearer ").strip() if oauth2_raw == litellm_raw: debug[f"{_RESPONSE_HEADER_PREFIX}-oauth2-token"] = ( - f"{MCPDebug._mask(oauth2_token)} " - f"(SAME_AS_LITELLM_KEY - likely misconfigured)" + f"{MCPDebug._mask(oauth2_token)} (SAME_AS_LITELLM_KEY - likely misconfigured)" ) else: - debug[f"{_RESPONSE_HEADER_PREFIX}-oauth2-token"] = MCPDebug._mask( - oauth2_token - ) + debug[f"{_RESPONSE_HEADER_PREFIX}-oauth2-token"] = MCPDebug._mask(oauth2_token) else: - debug[f"{_RESPONSE_HEADER_PREFIX}-oauth2-token"] = MCPDebug._mask( - oauth2_token - ) + debug[f"{_RESPONSE_HEADER_PREFIX}-oauth2-token"] = MCPDebug._mask(oauth2_token) # --- Auth resolution --- debug[f"{_RESPONSE_HEADER_PREFIX}-auth-resolution"] = auth_resolution # --- Server info --- debug[f"{_RESPONSE_HEADER_PREFIX}-outbound-url"] = server_url or "(unknown)" - debug[f"{_RESPONSE_HEADER_PREFIX}-server-auth-type"] = ( - server_auth_type or "(none)" - ) + debug[f"{_RESPONSE_HEADER_PREFIX}-server-auth-type"] = server_auth_type or "(none)" return debug @@ -301,9 +291,7 @@ class MCPDebug: auth_resolution = "no-auth" for server_name in mcp_servers or []: - server = global_mcp_server_manager.get_mcp_server_by_name( - server_name, client_ip=client_ip - ) + server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) if server: server_url = server.url server_auth_type = server.auth_type diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 835c505a164..cb9f4685bfd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -166,9 +166,7 @@ def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None: _user_env_vars_cache.pop((user_id, server_id), None) -def _write_user_env_vars_cache( - user_id: str, server_id: str, values: Dict[str, str] -) -> None: +def _write_user_env_vars_cache(user_id: str, server_id: str, values: Dict[str, str]) -> None: cache_key = (user_id, server_id) # Re-insert at the tail so eviction drops the oldest-written entry, not a # freshly refreshed one, and only sheds a single entry instead of wiping the @@ -210,10 +208,7 @@ def _should_strip_caller_authorization( """ if mcp_server.has_client_credentials: return True - if ( - mcp_server.auth_type == MCPAuth.oauth2 - and to_server_spec(mcp_server) is not None - ): + if mcp_server.auth_type == MCPAuth.oauth2 and to_server_spec(mcp_server) is not None: # Migrated per-user OAuth (authorization_code): the v2 resolver injects the # stored token, so a caller-forwarded Authorization must not be forwarded # upstream — it would override another user's stored credential. Delegate and @@ -222,12 +217,8 @@ def _should_strip_caller_authorization( if not mcp_server.is_oauth_passthrough: return False - normalized_raw_headers = { - str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str) - } - has_explicit_litellm_admission_header = ( - normalized_raw_headers.get("x-litellm-api-key") is not None - ) + normalized_raw_headers = {str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str)} + has_explicit_litellm_admission_header = normalized_raw_headers.get("x-litellm-api-key") is not None admission_consumed_authorization_as_litellm_key = ( user_api_key_auth is not None and bool(getattr(user_api_key_auth, "api_key", None)) @@ -292,10 +283,7 @@ def _extract_upstream_auth_failure( if current.__cause__ is not None: stack.append(current.__cause__) - if ( - current.__context__ is not None - and current.__context__ is not current.__cause__ - ): + if current.__context__ is not None and current.__context__ is not current.__cause__: stack.append(current.__context__) return None @@ -314,9 +302,7 @@ def _warn_on_server_name_fields( if result.is_valid: return - warning_text = ( - "; ".join(result.warnings) if result.warnings else "Validation failed" - ) + warning_text = "; ".join(result.warnings) if result.warnings else "Validation failed" verbose_logger.warning( "MCP server '%s' has invalid %s '%s': %s", server_id, @@ -329,9 +315,7 @@ def _warn_on_server_name_fields( _warn("server_name", server_name) -def _warn_internal_delegate_pkce_if_applicable( - server: MCPServer, *, source: str -) -> None: +def _warn_internal_delegate_pkce_if_applicable(server: MCPServer, *, source: str) -> None: """Surface internal + upstream PKCE delegate in logs for operators.""" if server.auth_type != MCPAuth.oauth2: return @@ -393,10 +377,7 @@ def _deserialize_json_list(data: Any) -> Optional[List[Dict[str, Any]]]: data = parsed if not isinstance(data, list): return None - return [ - item.model_dump(mode="json") if hasattr(item, "model_dump") else item - for item in data - ] + return [item.model_dump(mode="json") if hasattr(item, "model_dump") else item for item in data] def _normalize_mcp_server_cost_info(mcp_info: MCPInfo) -> None: @@ -464,9 +445,7 @@ def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): ) auth_context = get_active_auth_context() - resolved_auth = user_api_key_auth or ( - auth_context.user_api_key_auth if auth_context else None - ) + resolved_auth = user_api_key_auth or (auth_context.user_api_key_auth if auth_context else None) # Forward original HTTP headers and client IP so that # header-dependent guardrails, tag-based routing, trace # correlation, and forward_llm_provider_auth_headers work @@ -505,11 +484,7 @@ def _create_elicitation_callback(): # In Gateway mode, we relay the elicitation request to the downstream client # that triggered the current operation. downstream_session = get_active_mcp_session() - downstream_capabilities = ( - getattr(downstream_session, "capabilities", None) - if downstream_session - else None - ) + downstream_capabilities = getattr(downstream_session, "capabilities", None) if downstream_session else None return await handle_elicitation_request( context=context, @@ -541,9 +516,7 @@ class MCPServerManager: unless authorization_url is present (interactive OAuth). """ if oauth2_flow in ("client_credentials", "authorization_code"): - return cast( - Literal["client_credentials", "authorization_code"], oauth2_flow - ) + return cast(Literal["client_credentials", "authorization_code"], oauth2_flow) if oauth2_flow: # Ignore unknown/untyped values and continue legacy inference. return None @@ -589,18 +562,12 @@ class MCPServerManager: # not return instructions, and to apply a short cooldown after failures. self._upstream_initialize_instructions_probed_at: Dict[str, float] = {} - def _remember_upstream_initialize_instructions( - self, server: MCPServer, client: MCPClient - ) -> None: + def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient) -> None: raw = getattr(client, "_last_initialize_instructions", None) if raw and str(raw).strip(): - self._upstream_initialize_instructions_by_server_id[server.server_id] = str( - raw - ).strip() + self._upstream_initialize_instructions_by_server_id[server.server_id] = str(raw).strip() - async def _ensure_upstream_initialize_instructions_cached( - self, server: MCPServer - ) -> None: + async def _ensure_upstream_initialize_instructions_cached(self, server: MCPServer) -> None: """ Open one upstream session and cache InitializeResult.instructions if missing. @@ -635,20 +602,13 @@ class MCPServerManager: ): return - last_probed_at = self._upstream_initialize_instructions_probed_at.get( - server.server_id - ) - if ( - last_probed_at is not None - and (time.monotonic() - last_probed_at) < MCP_HEALTH_CHECK_TIMEOUT - ): + last_probed_at = self._upstream_initialize_instructions_probed_at.get(server.server_id) + if last_probed_at is not None and (time.monotonic() - last_probed_at) < MCP_HEALTH_CHECK_TIMEOUT: return # Record the attempt up-front so that a failure / empty response does not # cause every subsequent initialize request to re-open the upstream session. - self._upstream_initialize_instructions_probed_at[server.server_id] = ( - time.monotonic() - ) + self._upstream_initialize_instructions_probed_at[server.server_id] = time.monotonic() try: resolved_static_headers = await self._resolve_static_headers_with_env_vars( @@ -656,9 +616,7 @@ class MCPServerManager: user_api_key_auth=None, raise_on_missing=False, ) - extra_headers: Optional[Dict[str, str]] = ( - dict(resolved_static_headers) if resolved_static_headers else None - ) + extra_headers: Optional[Dict[str, str]] = dict(resolved_static_headers) if resolved_static_headers else None client = await self._create_mcp_client( server=server, mcp_auth_header=None, @@ -669,9 +627,7 @@ class MCPServerManager: async def _noop(_session): return "ok" - await asyncio.wait_for( - client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT - ) + await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT) self._remember_upstream_initialize_instructions(server, client) except Exception as e: verbose_logger.debug( @@ -724,15 +680,10 @@ class MCPServerManager: if mcp_aliases and alias is None: # Check if this server_name has an alias in mcp_aliases for alias_name, target_server_name in mcp_aliases.items(): - if ( - target_server_name == server_name - and alias_name not in used_aliases - ): + if target_server_name == server_name and alias_name not in used_aliases: alias = alias_name used_aliases.add(alias_name) - verbose_logger.debug( - f"Mapped alias '{alias_name}' to server '{server_name}'" - ) + verbose_logger.debug(f"Mapped alias '{alias_name}' to server '{server_name}'") break # Create a temporary server object to use with get_server_prefix utility @@ -767,9 +718,7 @@ class MCPServerManager: else: mcp_oauth_metadata = None - resolved_scopes = server_config.get("scopes") or ( - mcp_oauth_metadata.scopes if mcp_oauth_metadata else None - ) + resolved_scopes = server_config.get("scopes") or (mcp_oauth_metadata.scopes if mcp_oauth_metadata else None) resolved_authorization_url = server_config.get("authorization_url") or ( mcp_oauth_metadata.authorization_url if mcp_oauth_metadata else None ) @@ -808,9 +757,7 @@ class MCPServerManager: # TODO: utility fn the default values transport=server_config.get("transport", MCPTransport.http), auth_type=auth_type, - authentication_token=server_config.get( - "authentication_token", server_config.get("auth_value", None) - ), + authentication_token=server_config.get("authentication_token", server_config.get("auth_value", None)), mcp_info=mcp_info, extra_headers=server_config.get("extra_headers", None), allowed_tools=server_config.get("allowed_tools", None), @@ -820,12 +767,8 @@ class MCPServerManager: static_headers=server_config.get("static_headers", None), env_vars=server_config.get("env_vars", None), allow_all_keys=bool(server_config.get("allow_all_keys", False)), - available_on_public_internet=bool( - server_config.get("available_on_public_internet", True) - ), - delegate_auth_to_upstream=bool( - server_config.get("delegate_auth_to_upstream", False) - ), + available_on_public_internet=bool(server_config.get("available_on_public_internet", True)), + delegate_auth_to_upstream=bool(server_config.get("delegate_auth_to_upstream", False)), oauth_passthrough=bool(server_config.get("oauth_passthrough", False)), # AWS SigV4 fields aws_access_key_id=server_config.get("aws_access_key_id", None), @@ -837,9 +780,7 @@ class MCPServerManager: aws_session_name=server_config.get("aws_session_name", None), instructions=server_config.get("instructions", None), # Token Exchange (OBO) fields - token_exchange_endpoint=server_config.get( - "token_exchange_endpoint", None - ), + token_exchange_endpoint=server_config.get("token_exchange_endpoint", None), audience=server_config.get("audience", None), subject_token_type=server_config.get( "subject_token_type", @@ -856,24 +797,18 @@ class MCPServerManager: # Check if this is an OpenAPI-based server spec_path = server_config.get("spec_path", None) if spec_path: - verbose_logger.info( - f"Loading OpenAPI spec from {spec_path} for server {server_name}" - ) + verbose_logger.info(f"Loading OpenAPI spec from {spec_path} for server {server_name}") await self._register_openapi_tools( spec_path=spec_path, server=new_server, base_url=server_config.get("url", ""), ) - verbose_logger.debug( - f"Loaded MCP Servers: {json.dumps(self.config_mcp_servers, indent=4, default=str)}" - ) + verbose_logger.debug(f"Loaded MCP Servers: {json.dumps(self.config_mcp_servers, indent=4, default=str)}") self.initialize_tool_name_to_mcp_server_name_mapping() - async def _register_openapi_tools( - self, spec_path: str, server: MCPServer, base_url: str - ): + async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_url: str): """ Register tools from an OpenAPI specification for a given server. @@ -909,9 +844,7 @@ class MCPServerManager: # Use base_url from config if provided, otherwise extract from spec if not base_url: base_url = get_openapi_base_url(spec, spec_path) - verbose_logger.info( - f"Registering OpenAPI tools for server {server.name} with base URL: {base_url}" - ) + verbose_logger.info(f"Registering OpenAPI tools for server {server.name} with base URL: {base_url}") # Get server prefix for tool naming server_prefix = get_server_prefix(server) @@ -965,20 +898,14 @@ class MCPServerManager: operation = path_item[method] # Resolve $ref params and merge path-level params into the operation. - resolved_operation = resolve_operation_params( - operation, path_item, components - ) + resolved_operation = resolve_operation_params(operation, path_item, components) # Generate tool name (without prefix initially) - operation_id = operation.get( - "operationId", f"{method}_{path.replace('/', '_')}" - ) + operation_id = operation.get("operationId", f"{method}_{path.replace('/', '_')}") base_tool_name = operation_id.replace(" ", "_").lower() # Add server prefix to tool name - prefixed_tool_name = add_server_prefix_to_name( - base_tool_name, server_prefix - ) + prefixed_tool_name = add_server_prefix_to_name(base_tool_name, server_prefix) # Get description description = operation.get( @@ -990,9 +917,7 @@ class MCPServerManager: input_schema = build_input_schema(resolved_operation) # Create tool function with headers using imported function - tool_func = create_tool_function( - path, method, resolved_operation, base_url, headers=headers - ) + tool_func = create_tool_function(path, method, resolved_operation, base_url, headers=headers) tool_func.__name__ = prefixed_tool_name tool_func.__doc__ = description @@ -1005,26 +930,16 @@ class MCPServerManager: ) # Update tool name to server name mapping (for both prefixed and base names) - self.tool_name_to_mcp_server_name_mapping[base_tool_name] = ( - server_prefix - ) - self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = ( - server_prefix - ) + self.tool_name_to_mcp_server_name_mapping[base_tool_name] = server_prefix + self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = server_prefix registered_count += 1 - verbose_logger.debug( - f"Registered OpenAPI tool: {prefixed_tool_name} for server {server.name}" - ) + verbose_logger.debug(f"Registered OpenAPI tool: {prefixed_tool_name} for server {server.name}") - verbose_logger.info( - f"Successfully registered {registered_count} OpenAPI tools for server {server.name}" - ) + verbose_logger.info(f"Successfully registered {registered_count} OpenAPI tools for server {server.name}") except Exception as e: - verbose_logger.error( - f"Failed to register OpenAPI tools for server {server.name}: {str(e)}" - ) + verbose_logger.error(f"Failed to register OpenAPI tools for server {server.name}: {str(e)}") raise e def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None: @@ -1055,9 +970,7 @@ class MCPServerManager: owned_normalized = {normalize_server_name(x) for x in owned_raw} stale_mapping_keys: List[str] = [] - for tool_name, mapped_server in list( - self.tool_name_to_mcp_server_name_mapping.items() - ): + for tool_name, mapped_server in list(self.tool_name_to_mcp_server_name_mapping.items()): if mapped_server in owned_raw: stale_mapping_keys.append(tool_name) elif normalize_server_name(str(mapped_server)) in owned_normalized: @@ -1074,14 +987,10 @@ class MCPServerManager: if evicted is None and mcp_server.server_name: evicted = self.registry.pop(mcp_server.server_name, None) if evicted is not None: - verbose_logger.debug( - "Removed MCP Server: %s", mcp_server.server_id or mcp_server.server_name - ) + verbose_logger.debug("Removed MCP Server: %s", mcp_server.server_id or mcp_server.server_name) self._cleanup_server_tool_routing_artifacts(evicted) else: - verbose_logger.warning( - f"Server ID {mcp_server.server_id} not found in registry" - ) + verbose_logger.warning(f"Server ID {mcp_server.server_id} not found in registry") def _resolve_env_vars_list( self, @@ -1107,20 +1016,14 @@ class MCPServerManager: ) -> MCPServer: _mcp_info: MCPInfo = mcp_server.mcp_info or {} env_dict = _deserialize_json_dict(getattr(mcp_server, "env", None)) - static_headers_dict = _deserialize_json_dict( - getattr(mcp_server, "static_headers", None) - ) + static_headers_dict = _deserialize_json_dict(getattr(mcp_server, "static_headers", None)) env_vars_list = self._resolve_env_vars_list( mcp_server, env_vars_are_encrypted=( - credentials_are_encrypted - if env_vars_are_encrypted is None - else env_vars_are_encrypted + credentials_are_encrypted if env_vars_are_encrypted is None else env_vars_are_encrypted ), ) - credentials_dict = _deserialize_json_dict( - getattr(mcp_server, "credentials", None) - ) + credentials_dict = _deserialize_json_dict(getattr(mcp_server, "credentials", None)) encrypted_auth_value: Optional[str] = None encrypted_client_id: Optional[str] = None @@ -1167,9 +1070,7 @@ class MCPServerManager: client_secret_value = encrypted_client_secret # AWS SigV4 credential fields - aws_creds = self._extract_aws_credentials( - credentials_dict, credentials_are_encrypted - ) + aws_creds = self._extract_aws_credentials(credentials_dict, credentials_are_encrypted) scopes: Optional[List[str]] = None if credentials_dict: @@ -1177,9 +1078,7 @@ class MCPServerManager: if scopes_value is not None: scopes = self._extract_scopes(scopes_value) - name_for_prefix = ( - mcp_server.alias or mcp_server.server_name or mcp_server.server_id - ) + name_for_prefix = mcp_server.alias or mcp_server.server_name or mcp_server.server_id mcp_info: MCPInfo = _mcp_info.copy() if "server_name" not in mcp_info: @@ -1190,20 +1089,14 @@ class MCPServerManager: auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url - needs_discovery = ( - bool(server_url) - and auth_type == MCPAuth.oauth2 - and not mcp_server.authorization_url - ) + needs_discovery = bool(server_url) and auth_type == MCPAuth.oauth2 and not mcp_server.authorization_url mcp_oauth_metadata = ( await self._descovery_metadata(server_url=server_url) # type: ignore[arg-type] if needs_discovery else None ) - resolved_scopes = scopes or ( - mcp_oauth_metadata.scopes if mcp_oauth_metadata else None - ) + resolved_scopes = scopes or (mcp_oauth_metadata.scopes if mcp_oauth_metadata else None) new_server = MCPServer( server_id=mcp_server.server_id, @@ -1220,26 +1113,20 @@ class MCPServerManager: static_headers=static_headers_dict, env_vars=env_vars_list, client_id=client_id_value or getattr(mcp_server, "client_id", None), - client_secret=client_secret_value - or getattr(mcp_server, "client_secret", None), + client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), oauth2_flow=self._resolve_oauth2_flow( auth_type=auth_type, oauth2_flow=getattr(mcp_server, "oauth2_flow", None), - token_url=mcp_server.token_url - or getattr(mcp_oauth_metadata, "token_url", None), + token_url=mcp_server.token_url or getattr(mcp_oauth_metadata, "token_url", None), authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None), client_id=client_id_value or getattr(mcp_server, "client_id", None), - client_secret=client_secret_value - or getattr(mcp_server, "client_secret", None), + client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), ), scopes=resolved_scopes, - authorization_url=mcp_server.authorization_url - or getattr(mcp_oauth_metadata, "authorization_url", None), - token_url=mcp_server.token_url - or getattr(mcp_oauth_metadata, "token_url", None), - registration_url=mcp_server.registration_url - or getattr(mcp_oauth_metadata, "registration_url", None), + authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None), + token_url=mcp_server.token_url or getattr(mcp_oauth_metadata, "token_url", None), + registration_url=mcp_server.registration_url or getattr(mcp_oauth_metadata, "registration_url", None), command=getattr(mcp_server, "command", None), args=getattr(mcp_server, "args", None) or [], env=env_dict, @@ -1247,21 +1134,13 @@ class MCPServerManager: allowed_tools=getattr(mcp_server, "allowed_tools", None), disallowed_tools=getattr(mcp_server, "disallowed_tools", None), allow_all_keys=mcp_server.allow_all_keys, - available_on_public_internet=bool( - getattr(mcp_server, "available_on_public_internet", True) - ), - delegate_auth_to_upstream=bool( - getattr(mcp_server, "delegate_auth_to_upstream", False) - ), + available_on_public_internet=bool(getattr(mcp_server, "available_on_public_internet", True)), + delegate_auth_to_upstream=bool(getattr(mcp_server, "delegate_auth_to_upstream", False)), oauth_passthrough=bool(getattr(mcp_server, "oauth_passthrough", False)), created_at=getattr(mcp_server, "created_at", None), updated_at=getattr(mcp_server, "updated_at", None), - tool_name_to_display_name=_deserialize_json_dict( - getattr(mcp_server, "tool_name_to_display_name", None) - ), - tool_name_to_description=_deserialize_json_dict( - getattr(mcp_server, "tool_name_to_description", None) - ), + tool_name_to_display_name=_deserialize_json_dict(getattr(mcp_server, "tool_name_to_display_name", None)), + tool_name_to_description=_deserialize_json_dict(getattr(mcp_server, "tool_name_to_description", None)), is_byok=bool(getattr(mcp_server, "is_byok", False)), byok_description=getattr(mcp_server, "byok_description", None) or [], byok_api_key_help_url=getattr(mcp_server, "byok_api_key_help_url", None), @@ -1276,29 +1155,19 @@ class MCPServerManager: aws_session_name=aws_creds.get("aws_session_name"), instructions=mcp_server.instructions, # Token Exchange (OBO) fields — read from credentials JSON blob - token_exchange_endpoint=( - credentials_dict.get("token_exchange_endpoint") - if credentials_dict - else None - ), + token_exchange_endpoint=(credentials_dict.get("token_exchange_endpoint") if credentials_dict else None), audience=(credentials_dict.get("audience") if credentials_dict else None), - subject_token_type=( - credentials_dict.get("subject_token_type") if credentials_dict else None - ) + subject_token_type=(credentials_dict.get("subject_token_type") if credentials_dict else None) or "urn:ietf:params:oauth:token-type:access_token", timeout=getattr(mcp_server, "timeout", None), ) _warn_internal_delegate_pkce_if_applicable(new_server, source="database") return new_server - async def _maybe_register_openapi_tools( - self, server: MCPServer, *, initialize_mapping: bool = True - ): + async def _maybe_register_openapi_tools(self, server: MCPServer, *, initialize_mapping: bool = True): """Register OpenAPI tools if the server has a spec_path configured.""" if server.spec_path: - verbose_logger.info( - f"Loading OpenAPI spec from {server.spec_path} for server {server.name}" - ) + verbose_logger.info(f"Loading OpenAPI spec from {server.spec_path} for server {server.name}") await self._register_openapi_tools( spec_path=server.spec_path, server=server, @@ -1322,9 +1191,7 @@ class MCPServerManager: # `credentials` field is the only one still encrypted here). # Re-decrypting plaintext would zero the values, so build with # env_vars_are_encrypted=False. - new_server = await self.build_mcp_server_from_table( - mcp_server, env_vars_are_encrypted=False - ) + new_server = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False) self._assign_unique_short_prefix(new_server) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) @@ -1349,9 +1216,7 @@ class MCPServerManager: if mcp_server.server_id in self.registry: # See add_server: db.py helpers already decrypted env var # values, so don't decrypt them a second time here. - new_server = await self.build_mcp_server_from_table( - mcp_server, env_vars_are_encrypted=False - ) + new_server = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False) # Carry the previously-resolved short prefix across so the # tool names stay stable for clients holding cached lists. existing_prefix = self.registry[mcp_server.server_id].short_prefix @@ -1375,15 +1240,9 @@ class MCPServerManager: def get_allow_all_keys_server_ids(self) -> List[str]: """Return server IDs that bypass per-key restrictions.""" - return [ - server.server_id - for server in self.get_registry().values() - if server.allow_all_keys is True - ] + return [server.server_id for server in self.get_registry().values() if server.allow_all_keys is True] - async def get_allowed_mcp_servers( - self, user_api_key_auth: Optional[UserAPIKeyAuth] = None - ) -> List[str]: + async def get_allowed_mcp_servers(self, user_api_key_auth: Optional[UserAPIKeyAuth] = None) -> List[str]: """ Get the allowed MCP Servers for the user. @@ -1399,12 +1258,9 @@ class MCPServerManager: try: # The key explicitly opted out of every MCP server. Return zero before # layering on allow_all_keys servers so the opt-out is absolute. - key_object_permission = ( - user_api_key_auth.object_permission if user_api_key_auth else None - ) + key_object_permission = user_api_key_auth.object_permission if user_api_key_auth else None if key_object_permission is not None and ( - SpecialMCPServerNames.no_mcp_servers.value - in (key_object_permission.mcp_servers or []) + SpecialMCPServerNames.no_mcp_servers.value in (key_object_permission.mcp_servers or []) ): return [] @@ -1419,23 +1275,13 @@ class MCPServerManager: ) # If admin but NO explicit object permission, get all servers - if ( - user_api_key_auth - and _user_has_admin_view(user_api_key_auth) - and not has_explicit_object_permission - ): - verbose_logger.debug( - "Admin user without explicit object_permission - returning all servers" - ) + if user_api_key_auth and _user_has_admin_view(user_api_key_auth) and not has_explicit_object_permission: + verbose_logger.debug("Admin user without explicit object_permission - returning all servers") return list(self.get_registry().keys()) # Get allowed servers from object permissions (respects object_permission even for admins) - allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers( - user_api_key_auth - ) - verbose_logger.debug( - f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}" - ) + allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + verbose_logger.debug(f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}") combined_servers = set(allowed_mcp_servers) # Only skip allow_all_keys servers when the request is inside a toolset # scope. toolset_mcp_route / dynamic_mcp_route set _mcp_active_toolset_id @@ -1476,9 +1322,7 @@ class MCPServerManager: combined_servers.update(delegate_server_ids) if len(combined_servers) == 0: - verbose_logger.debug( - "No allowed MCP Servers found for user api key auth." - ) + verbose_logger.debug("No allowed MCP Servers found for user api key auth.") return list(combined_servers) except Exception: # noqa: BLE001 verbose_logger.exception( @@ -1558,15 +1402,12 @@ class MCPServerManager: keys_to_remove = [ k for k in cache_dict - if (k.startswith("toolset_perms:") and toolset_id in k) - or k.startswith("toolset_name:") + if (k.startswith("toolset_perms:") and toolset_id in k) or k.startswith("toolset_name:") ] for k in keys_to_remove: cache_dict.pop(k, None) except Exception as e: - verbose_logger.warning( - f"invalidate_toolset_cache: failed to evict in-memory entries: {e}" - ) + verbose_logger.warning(f"invalidate_toolset_cache: failed to evict in-memory entries: {e}") async def get_toolset_by_name_cached( self, @@ -1603,18 +1444,12 @@ class MCPServerManager: toolset = await get_mcp_toolset_by_name(prisma_client, toolset_name) await user_api_key_cache.async_set_cache( key=cache_key, - value=( - toolset.model_dump(mode="json") - if toolset is not None - else "__not_found__" - ), + value=(toolset.model_dump(mode="json") if toolset is not None else "__not_found__"), ttl=get_management_object_ttl(user_api_key_cache), ) return toolset - def filter_server_ids_by_ip( - self, server_ids: List[str], client_ip: Optional[str] - ) -> List[str]: + def filter_server_ids_by_ip(self, server_ids: List[str], client_ip: Optional[str]) -> List[str]: """ Filter server IDs by client IP — external callers only see public servers. @@ -1656,9 +1491,7 @@ class MCPServerManager: return [] return await self._get_tools_from_server(server) except Exception as e: - verbose_logger.warning( - f"Failed to get tools from server {server_id}: {str(e)}" - ) + verbose_logger.warning(f"Failed to get tools from server {server_id}: {str(e)}") return [] async def list_tools( @@ -1727,9 +1560,7 @@ class MCPServerManager: # Flatten results into single list list_tools_result: List[MCPTool] = [tool for tools in results for tool in tools] - verbose_logger.info( - f"Successfully fetched {len(list_tools_result)} tools total from all servers" - ) + verbose_logger.info(f"Successfully fetched {len(list_tools_result)} tools total from all servers") return list_tools_result ######################################################### @@ -1851,9 +1682,7 @@ class MCPServerManager: # the user hasn't filled it in -- only vars without a global fallback do. referenced = collect_env_var_references(strings=(static_headers or {}).values()) referenced_user_vars = referenced & user_var_names - required_user_vars = { - name for name in referenced_user_vars if name not in global_values - } + required_user_vars = {name for name in referenced_user_vars if name not in global_values} user_values: Dict[str, str] = {} if required_user_vars: @@ -1867,28 +1696,21 @@ class MCPServerManager: if raise_on_missing: raise verbose_logger.warning( - "MCPServerManager: best-effort user env var load failed for " - "server=%s: %s", + "MCPServerManager: best-effort user env var load failed for server=%s: %s", server.server_id, exc, ) if raise_on_missing: - missing = sorted( - name for name in required_user_vars if not user_values.get(name) - ) + missing = sorted(name for name in required_user_vars if not user_values.get(name)) if missing: # A cached negative must never produce a 412: cache # invalidation is process-local, so a user who just stored # values on another worker would otherwise be told their # credentials are missing until the entry expires. Confirm # against the DB before raising. - user_values = await self._load_user_env_vars( - server, user_api_key_auth, force_refresh=True - ) - missing = sorted( - name for name in required_user_vars if not user_values.get(name) - ) + user_values = await self._load_user_env_vars(server, user_api_key_auth, force_refresh=True) + missing = sorted(name for name in required_user_vars if not user_values.get(name)) if missing: raise MCPMissingUserEnvVarsError( server_id=server.server_id, @@ -1900,9 +1722,7 @@ class MCPServerManager: # Only honor stored user values for currently user-scoped vars, and let # admin globals win, so a stale row from when a var was user-scoped can # never override the global value the admin set after switching it. - scoped_user_values = { - name: value for name, value in user_values.items() if name in user_var_names - } + scoped_user_values = {name: value for name, value in user_values.items() if name in user_var_names} merged_vars: Dict[str, str] = {**scoped_user_values, **global_values} if not static_headers: return static_headers @@ -1997,34 +1817,20 @@ class MCPServerManager: # caller must not be able to substitute another user's stored credential, so we keep the v2 # spec and ignore the override there; the REST tools preview supplies its not-yet-persisted # token through the resolver (cred_provider), never this path. - if ( - spec is not None - and mcp_auth_header - and not isinstance(spec.config, AuthorizationCodeConfig) - ): + if spec is not None and mcp_auth_header and not isinstance(spec.config, AuthorizationCodeConfig): spec = None auth_value = ( - await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token) - if spec is None - else None + await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token) if spec is None else None ) # Create sampling and elicitation callbacks for this client - sampling_cb = ( - _create_sampling_callback(user_api_key_auth=user_api_key_auth) - if server.allow_sampling - else None - ) - elicitation_cb = ( - _create_elicitation_callback() if server.allow_elicitation else None - ) + sampling_cb = _create_sampling_callback(user_api_key_auth=user_api_key_auth) if server.allow_sampling else None + elicitation_cb = _create_elicitation_callback() if server.allow_elicitation else None # Handle stdio transport if transport == MCPTransport.stdio: resolved_env = ( - stdio_env - if stdio_env is not None - else (dict(server.env) if server.env is not None else None) + stdio_env if stdio_env is not None else (dict(server.env) if server.env is not None else None) ) # Ensure npm-based STDIO MCP servers have a writable cache dir. @@ -2066,9 +1872,7 @@ class MCPServerManager: transport_type=transport, auth_type=server.auth_type, auth_value=auth_value, - timeout=( - server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT - ), + timeout=(server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT), stdio_config=stdio_config, extra_headers=extra_headers, sampling_callback=sampling_cb, @@ -2079,9 +1883,7 @@ class MCPServerManager: server_url = server.url or "" if spec is not None: - match await provider.resolve_credentials( - to_subject(user_api_key_auth, subject_token), spec - ): + match await provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(auth): resolved_auth = auth # Do not override an Authorization already supplied via extra_headers @@ -2092,10 +1894,7 @@ class MCPServerManager: if ( header_name and extra_headers - and any( - key.lower() == header_name.lower() - for key in extra_headers - ) + and any(key.lower() == header_name.lower() for key in extra_headers) ): resolved_auth = None case Error(err): @@ -2108,11 +1907,7 @@ class MCPServerManager: server_url=server_url, transport_type=transport, auth_type=server.auth_type, - timeout=( - server.timeout - if server.timeout is not None - else MCP_CLIENT_TIMEOUT - ), + timeout=(server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT), extra_headers=extra_headers, resolved_auth=resolved_auth, sampling_callback=sampling_cb, @@ -2137,9 +1932,7 @@ class MCPServerManager: transport_type=transport, auth_type=server.auth_type, auth_value=auth_value, - timeout=( - server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT - ), + timeout=(server.timeout if server.timeout is not None else MCP_CLIENT_TIMEOUT), extra_headers=extra_headers, aws_auth=aws_auth, sampling_callback=sampling_cb, @@ -2206,12 +1999,10 @@ class MCPServerManager: static_headers = server.static_headers or {} has_static_authorization = any( - isinstance(k, str) and k.lower() == "authorization" - for k in static_headers.keys() + isinstance(k, str) and k.lower() == "authorization" for k in static_headers.keys() ) has_extra_authorization = bool(extra_headers) and any( - isinstance(k, str) and k.lower() == "authorization" - for k in (extra_headers or {}).keys() + isinstance(k, str) and k.lower() == "authorization" for k in (extra_headers or {}).keys() ) if ( @@ -2241,12 +2032,8 @@ class MCPServerManager: if server.spec_path: # OpenAPI tools were stored in the registry under the prefix # active at registration time — fetch by that same prefix. - _tools = global_mcp_tool_registry.list_tools( - tool_prefix=get_server_prefix(server) - ) - tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type( - _tools - ) + _tools = global_mcp_tool_registry.list_tools(tool_prefix=get_server_prefix(server)) + tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type(_tools) # OpenAPI tools are stored in the registry with their prefix already # applied (e.g. "test_petstore-getinventory"). Do NOT pass them # through _create_prefixed_tools — that would add the prefix a second @@ -2256,9 +2043,7 @@ class MCPServerManager: sep = MCP_TOOL_PREFIX_SEPARATOR tools = [ ( - t.model_copy( - update={"name": t.name[len(prefix) + len(sep) :]} - ) + t.model_copy(update={"name": t.name[len(prefix) + len(sep) :]}) if t.name.startswith(f"{prefix}{sep}") else t ) @@ -2266,14 +2051,10 @@ class MCPServerManager: ] return tools else: - tools = await self._fetch_tools_with_timeout( - client, server.name, server=server - ) + tools = await self._fetch_tools_with_timeout(client, server.name, server=server) self._remember_upstream_initialize_instructions(server, client) - prefixed_or_original_tools = self._create_prefixed_tools( - tools, server, add_prefix=add_prefix - ) + prefixed_or_original_tools = self._create_prefixed_tools(tools, server, add_prefix=add_prefix) return prefixed_or_original_tools @@ -2283,9 +2064,7 @@ class MCPServerManager: # aggregator catches this explicitly to keep absorbing. raise except Exception as e: - verbose_logger.warning( - f"Failed to get tools from server {server.name}: {str(e)}" - ) + verbose_logger.warning(f"Failed to get tools from server {server.name}: {str(e)}") return [] async def get_prompts_from_server( @@ -2329,16 +2108,12 @@ class MCPServerManager: prompts = await client.list_prompts() - prefixed_or_original_prompts = self._create_prefixed_prompts( - prompts, server, add_prefix=add_prefix - ) + prefixed_or_original_prompts = self._create_prefixed_prompts(prompts, server, add_prefix=add_prefix) return prefixed_or_original_prompts except Exception as e: - verbose_logger.warning( - f"Failed to get prompts from server {server.name}: {str(e)}" - ) + verbose_logger.warning(f"Failed to get prompts from server {server.name}: {str(e)}") return [] async def get_resources_from_server( @@ -2373,16 +2148,12 @@ class MCPServerManager: resources = await client.list_resources() - prefixed_resources = self._create_prefixed_resources( - resources, server, add_prefix=add_prefix - ) + prefixed_resources = self._create_prefixed_resources(resources, server, add_prefix=add_prefix) return prefixed_resources except Exception as e: - verbose_logger.warning( - f"Failed to get resources from server {server.name}: {str(e)}" - ) + verbose_logger.warning(f"Failed to get resources from server {server.name}: {str(e)}") return [] async def get_resource_templates_from_server( @@ -2424,9 +2195,7 @@ class MCPServerManager: return prefixed_templates except Exception as e: - verbose_logger.warning( - f"Failed to get resource templates from server {server.name}: {str(e)}" - ) + verbose_logger.warning(f"Failed to get resource templates from server {server.name}: {str(e)}") return [] async def read_resource_from_server( @@ -2547,15 +2316,8 @@ class MCPServerManager: authorization_servers, resource_scopes, ) = await self._attempt_well_known_discovery(server_url) - metadata = await self._fetch_authorization_server_metadata( - authorization_servers, server_url - ) - if ( - metadata is None - and not resource_scopes - and authorization_servers - and response.status_code == 200 - ): + metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url) + if metadata is None and not resource_scopes and authorization_servers and response.status_code == 200: verbose_logger.warning( "MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.", server_url, @@ -2574,13 +2336,11 @@ class MCPServerManager: header_value: Optional[str] = None if exc.response is not None: - header_value = exc.response.headers.get( - "WWW-Authenticate" - ) or exc.response.headers.get("www-authenticate") + header_value = exc.response.headers.get("WWW-Authenticate") or exc.response.headers.get( + "www-authenticate" + ) - resource_metadata_url, scopes = self._parse_www_authenticate_header( - header_value - ) + resource_metadata_url, scopes = self._parse_www_authenticate_header(header_value) authorization_servers = [] resource_scopes = None @@ -2588,9 +2348,7 @@ class MCPServerManager: ( authorization_servers, resource_scopes, - ) = await self._fetch_oauth_metadata_from_resource( - resource_metadata_url, server_url - ) + ) = await self._fetch_oauth_metadata_from_resource(resource_metadata_url, server_url) else: ( authorization_servers, @@ -2602,16 +2360,12 @@ class MCPServerManager: try: parsed_url = urlparse(server_url) if parsed_url.scheme and parsed_url.netloc: - authorization_servers = [ - f"{parsed_url.scheme}://{parsed_url.netloc}" - ] + authorization_servers = [f"{parsed_url.scheme}://{parsed_url.netloc}"] except Exception: authorization_servers = [] if authorization_servers: - metadata = await self._fetch_authorization_server_metadata( - authorization_servers, server_url - ) + metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url) preferred_scopes = scopes or resource_scopes if metadata is None and preferred_scopes: @@ -2621,14 +2375,10 @@ class MCPServerManager: return metadata except Exception as exc: # pragma: no cover - network/transient issues - verbose_logger.debug( - "MCP OAuth discovery failed for %s: %s", server_url, exc - ) + verbose_logger.debug("MCP OAuth discovery failed for %s: %s", server_url, exc) return None - def _parse_www_authenticate_header( - self, header_value: Optional[str] - ) -> Tuple[Optional[str], Optional[List[str]]]: + def _parse_www_authenticate_header(self, header_value: Optional[str]) -> Tuple[Optional[str], Optional[List[str]]]: if not header_value: return None, None @@ -2637,8 +2387,7 @@ class MCPServerManager: param_pattern = re.compile(r"([a-zA-Z0-9_]+)\s*=\s*\"?([^\",]+)\"?") params: Dict[str, str] = { - match.group(1).lower(): match.group(2).strip() - for match in param_pattern.finditer(params_section) + match.group(1).lower(): match.group(2).strip() for match in param_pattern.finditer(params_section) } resource_metadata_url = params.get("resource_metadata") @@ -2656,9 +2405,7 @@ class MCPServerManager: return [], None try: - response = await self._fetch_oauth_discovery_url( - resource_metadata_url, server_url - ) + response = await self._fetch_oauth_discovery_url(resource_metadata_url, server_url) response.raise_for_status() data = response.json() except SSRFError as exc: @@ -2680,23 +2427,15 @@ class MCPServerManager: raw_servers = data.get("authorization_servers") if isinstance(raw_servers, list): - authorization_servers = [ - entry - for entry in raw_servers - if isinstance(entry, str) and entry.strip() != "" - ] + authorization_servers = [entry for entry in raw_servers if isinstance(entry, str) and entry.strip() != ""] else: authorization_servers = [] - scopes = self._extract_scopes( - data.get("scopes_supported") or data.get("scopes") - ) + scopes = self._extract_scopes(data.get("scopes_supported") or data.get("scopes")) return authorization_servers, scopes - async def _attempt_well_known_discovery( - self, server_url: str - ) -> Tuple[List[str], Optional[List[str]]]: + async def _attempt_well_known_discovery(self, server_url: str) -> Tuple[List[str], Optional[List[str]]]: try: parsed = urlparse(server_url) except Exception: @@ -2728,9 +2467,7 @@ class MCPServerManager: self, authorization_servers: List[str], server_url: str ) -> Optional[MCPOAuthMetadata]: for issuer in authorization_servers: - metadata = await self._fetch_single_authorization_server_metadata( - issuer, server_url - ) + metadata = await self._fetch_single_authorization_server_metadata(issuer, server_url) if metadata is not None: return metadata return None @@ -2751,13 +2488,9 @@ class MCPServerManager: candidate_urls: List[str] = [] if path: - candidate_urls.append( - f"{base}/.well-known/oauth-authorization-server/{path}" - ) + candidate_urls.append(f"{base}/.well-known/oauth-authorization-server/{path}") candidate_urls.append(f"{base}/.well-known/openid-configuration/{path}") - candidate_urls.append( - f"{issuer_url.rstrip('/')}/.well-known/openid-configuration" - ) + candidate_urls.append(f"{issuer_url.rstrip('/')}/.well-known/openid-configuration") candidate_urls.append(f"{base}/.well-known/oauth-authorization-server") candidate_urls.append(f"{base}/.well-known/openid-configuration") candidate_urls.append(issuer_url.rstrip("/")) @@ -2808,14 +2541,8 @@ class MCPServerManager: def _build_azure_authorization_server_metadata( parsed_issuer_url: Any, ) -> Optional[MCPOAuthMetadata]: - path_parts = [ - part for part in (parsed_issuer_url.path or "").split("/") if part - ] - if ( - parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS - or len(path_parts) != 2 - or path_parts[1] != "v2.0" - ): + path_parts = [part for part in (parsed_issuer_url.path or "").split("/") if part] + if parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS or len(path_parts) != 2 or path_parts[1] != "v2.0": return None tenant = path_parts[0] @@ -2925,33 +2652,24 @@ class MCPServerManager: ) try: with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): - tools = await client.list_tools( - raise_on_error=should_surface_upstream_auth - ) + tools = await client.list_tools(raise_on_error=should_surface_upstream_auth) verbose_logger.debug(f"Tools from {server_name}: {tools}") return tools except TimeoutError: verbose_logger.warning(f"Timeout while listing tools from {server_name}") return [] except asyncio.CancelledError: - verbose_logger.warning( - f"Task cancelled while listing tools from {server_name}" - ) + verbose_logger.warning(f"Task cancelled while listing tools from {server_name}") return [] except ConnectionError as e: - verbose_logger.warning( - f"Connection error while listing tools from {server_name}: {str(e)}" - ) + verbose_logger.warning(f"Connection error while listing tools from {server_name}: {str(e)}") return [] except Exception as e: if should_surface_upstream_auth: auth_info = _extract_upstream_auth_failure(e) if auth_info is not None: status_code, www_authenticate = auth_info - verbose_logger.info( - f"Upstream auth failure from MCP server " - f"{server_name}: HTTP {status_code}" - ) + verbose_logger.info(f"Upstream auth failure from MCP server {server_name}: HTTP {status_code}") raise MCPUpstreamAuthError( status_code=status_code, www_authenticate=www_authenticate, @@ -3022,9 +2740,7 @@ class MCPServerManager: "attempts; the 3-character prefix space is too crowded." ) - def _create_prefixed_tools( - self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True - ) -> List[MCPTool]: + def _create_prefixed_tools(self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True) -> List[MCPTool]: """ Create prefixed tools and update tool mapping. @@ -3058,9 +2774,7 @@ class MCPServerManager: qualified = add_server_prefix_to_name(original_name, known_prefix) self.tool_name_to_mcp_server_name_mapping[qualified] = prefix - verbose_logger.info( - f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}" - ) + verbose_logger.info(f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}") return prefixed_tools def _create_prefixed_prompts( @@ -3087,9 +2801,7 @@ class MCPServerManager: prompt.name = name_to_use prefixed_prompts.append(prompt) - verbose_logger.info( - f"Successfully fetched {len(prefixed_prompts)} prompts from server {server.name}" - ) + verbose_logger.info(f"Successfully fetched {len(prefixed_prompts)} prompts from server {server.name}") return prefixed_prompts def _create_prefixed_resources( @@ -3101,17 +2813,11 @@ class MCPServerManager: prefix = get_server_prefix(server) for resource in resources: - name_to_use = ( - add_server_prefix_to_name(resource.name, prefix) - if add_prefix - else resource.name - ) + name_to_use = add_server_prefix_to_name(resource.name, prefix) if add_prefix else resource.name resource.name = name_to_use prefixed_resources.append(resource) - verbose_logger.info( - f"Successfully fetched {len(prefixed_resources)} resources from server {server.name}" - ) + verbose_logger.info(f"Successfully fetched {len(prefixed_resources)} resources from server {server.name}") return prefixed_resources def _create_prefixed_resource_templates( @@ -3127,9 +2833,7 @@ class MCPServerManager: for resource_template in resource_templates: name_to_use = ( - add_server_prefix_to_name(resource_template.name, prefix) - if add_prefix - else resource_template.name + add_server_prefix_to_name(resource_template.name, prefix) if add_prefix else resource_template.name ) resource_template.name = name_to_use prefixed_templates.append(resource_template) @@ -3150,20 +2854,14 @@ class MCPServerManager: if server_applies_tool_allowlist(server): if not server.allowed_tools: return False - return ( - tool_name in server.allowed_tools - or f"{server.name}-{tool_name}" in server.allowed_tools - ) + return tool_name in server.allowed_tools or f"{server.name}-{tool_name}" in server.allowed_tools if server.disallowed_tools: return ( - tool_name not in server.disallowed_tools - and f"{server.name}-{tool_name}" not in server.disallowed_tools + tool_name not in server.disallowed_tools and f"{server.name}-{tool_name}" not in server.disallowed_tools ) return True - def validate_allowed_params( - self, tool_name: str, arguments: Dict[str, Any], server: MCPServer - ) -> None: + def validate_allowed_params(self, tool_name: str, arguments: Dict[str, Any], server: MCPServer) -> None: """ Filter arguments to only include allowed parameters for the given tool. @@ -3190,18 +2888,14 @@ class MCPServerManager: unprefixed_tool_name, _ = split_server_prefix_from_name(tool_name) # Check both prefixed and unprefixed tool names - allowed_params_list = server.allowed_params.get( - tool_name - ) or server.allowed_params.get(unprefixed_tool_name) + allowed_params_list = server.allowed_params.get(tool_name) or server.allowed_params.get(unprefixed_tool_name) # If this tool doesn't have allowed_params specified, allow all params if allowed_params_list is None: return None # Filter arguments to only include allowed parameters - disallowed_params = [ - param for param in arguments.keys() if param not in allowed_params_list - ] + disallowed_params = [param for param in arguments.keys() if param not in allowed_params_list] if disallowed_params: raise HTTPException( @@ -3364,42 +3058,22 @@ class MCPServerManager: "name": name, "arguments": arguments, "server_name": server_name, - "mcp_rate_limit_server_name": server.alias - or server.server_name - or server.name, + "mcp_rate_limit_server_name": server.alias or server.server_name or server.name, "user_api_key_auth": user_api_key_auth, - "user_api_key_user_id": ( - getattr(user_api_key_auth, "user_id", None) - if user_api_key_auth - else None - ), - "user_api_key_team_id": ( - getattr(user_api_key_auth, "team_id", None) - if user_api_key_auth - else None - ), + "user_api_key_user_id": (getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None), + "user_api_key_team_id": (getattr(user_api_key_auth, "team_id", None) if user_api_key_auth else None), "user_api_key_end_user_id": ( - getattr(user_api_key_auth, "end_user_id", None) - if user_api_key_auth - else None - ), - "user_api_key_hash": ( - getattr(user_api_key_auth, "api_key_hash", None) - if user_api_key_auth - else None + getattr(user_api_key_auth, "end_user_id", None) if user_api_key_auth else None ), + "user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None), "incoming_bearer_token": incoming_bearer_token, } # Create MCP request object for processing - mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs( - pre_hook_kwargs - ) + mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs) # Convert to LLM format for existing guardrail compatibility - synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format( - mcp_request_obj, pre_hook_kwargs - ) + synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs) hook_result: Dict[str, Any] = {} try: @@ -3411,11 +3085,7 @@ class MCPServerManager: ) if modified_data: # Convert response back to MCP format and apply modifications - modified_kwargs = ( - proxy_logging_obj._convert_mcp_hook_response_to_kwargs( - modified_data, pre_hook_kwargs - ) - ) + modified_kwargs = proxy_logging_obj._convert_mcp_hook_response_to_kwargs(modified_data, pre_hook_kwargs) if modified_kwargs.get("arguments") != arguments: hook_result["arguments"] = modified_kwargs["arguments"] if modified_kwargs.get("extra_headers"): @@ -3460,9 +3130,7 @@ class MCPServerManager: "user_api_key_auth": user_api_key_auth, } - synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format( - request_obj, during_hook_kwargs - ) + synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, during_hook_kwargs) return asyncio.create_task( proxy_logging_obj.during_call_hook( @@ -3558,9 +3226,7 @@ class MCPServerManager: if extra_headers is None: extra_headers = {} - normalized_raw_headers = { - str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str) - } + normalized_raw_headers = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} strip_caller_authorization = _should_strip_caller_authorization( mcp_server=mcp_server, raw_headers=raw_headers, @@ -3581,9 +3247,7 @@ class MCPServerManager: # MCPMissingUserEnvVarsError when the calling user has not filled in # a required per-user variable — the REST layer converts that into # a friendly 412 with a setup URL. - resolved_static_headers = await self._resolve_static_headers_with_env_vars( - mcp_server, user_api_key_auth - ) + resolved_static_headers = await self._resolve_static_headers_with_env_vars(mcp_server, user_api_key_auth) if resolved_static_headers: if extra_headers is None: extra_headers = {} @@ -3635,21 +3299,13 @@ class MCPServerManager: ) async def _call_tool_via_client(client, params): - return await client.call_tool( - params, host_progress_callback=host_progress_callback - ) + return await client.call_tool(params, host_progress_callback=host_progress_callback) - tasks.append( - asyncio.create_task(_call_tool_via_client(client, call_tool_params)) - ) + tasks.append(asyncio.create_task(_call_tool_via_client(client, call_tool_params))) - _timeout = ( - mcp_server.timeout if mcp_server.timeout is not None else MCP_CLIENT_TIMEOUT - ) + _timeout = mcp_server.timeout if mcp_server.timeout is not None else MCP_CLIENT_TIMEOUT try: - mcp_responses = await asyncio.wait_for( - asyncio.gather(*tasks), timeout=_timeout - ) + mcp_responses = await asyncio.wait_for(asyncio.gather(*tasks), timeout=_timeout) except asyncio.TimeoutError: raise HTTPException( status_code=504, @@ -3663,9 +3319,7 @@ class MCPServerManager: GuardrailRaisedException, HTTPException, ) as e: - verbose_logger.error( - f"Guardrail blocked MCP tool call during result check: {str(e)}" - ) + verbose_logger.error(f"Guardrail blocked MCP tool call during result check: {str(e)}") raise e # If proxy_logging_obj is None, the tool call result is at index 0 @@ -3693,9 +3347,7 @@ class MCPServerManager: candidate.server_name, candidate.name, ): - if identifier and normalize_server_name(identifier) == ( - normalized_server_name - ): + if identifier and normalize_server_name(identifier) == (normalized_server_name): return True return False @@ -3707,9 +3359,7 @@ class MCPServerManager: break if mcp_server is None: fallback = self._get_mcp_server_from_tool_name(name) - if fallback is not None and ( - not server_name or _candidate_matches_server_name(fallback) - ): + if fallback is not None and (not server_name or _candidate_matches_server_name(fallback)): mcp_server = fallback if mcp_server is None: raise ValueError(f"Tool {name} not found") @@ -3724,9 +3374,7 @@ class MCPServerManager: return mcp_server - async def has_user_oauth_token( - self, server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth] - ) -> bool: + async def has_user_oauth_token(self, server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth]) -> bool: """Whether the v2 resolver can produce a per-user token for this server right now. This is the preemptive 401's existence check, routed through the same resolver that drives @@ -3736,9 +3384,7 @@ class MCPServerManager: spec = to_server_spec(server) if spec is None: return False - return await self._cred_provider.has_user_token( - to_subject(user_api_key_auth, None), spec - ) + return await self._cred_provider.has_user_token(to_subject(user_api_key_auth, None), spec) async def _resolve_oauth2_headers_for_tool_call( self, @@ -3747,11 +3393,7 @@ class MCPServerManager: user_api_key_auth: Optional[UserAPIKeyAuth], ) -> Optional[Dict[str, str]]: """Look up per-user OAuth headers when the client did not supply a token.""" - if ( - not mcp_server.needs_user_oauth_token - or oauth2_headers - or user_api_key_auth is None - ): + if not mcp_server.needs_user_oauth_token or oauth2_headers or user_api_key_auth is None: return oauth2_headers if to_server_spec(mcp_server) is not None: @@ -3799,9 +3441,7 @@ class MCPServerManager: GuardrailRaisedException, HTTPException, ) as e: - verbose_logger.error( - f"Guardrail blocked MCP tool call during result check: {str(e)}" - ) + verbose_logger.error(f"Guardrail blocked MCP tool call during result check: {str(e)}") raise e async def call_tool( @@ -3868,15 +3508,11 @@ class MCPServerManager: ) tasks.append(during_hook_task) - oauth2_headers = await self._resolve_oauth2_headers_for_tool_call( - mcp_server, oauth2_headers, user_api_key_auth - ) + oauth2_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, oauth2_headers, user_api_key_auth) # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: - verbose_logger.debug( - "Calling OpenAPI tool %s directly via HTTP handler", name - ) + verbose_logger.debug("Calling OpenAPI tool %s directly via HTTP handler", name) if hook_result.get("extra_headers"): verbose_logger.warning( "pre_mcp_call hook returned extra_headers for OpenAPI-backed " @@ -3885,11 +3521,7 @@ class MCPServerManager: "transport to enable hook header injection.", server_name, ) - tasks.append( - asyncio.create_task( - self._call_openapi_tool_handler(mcp_server, name, arguments) - ) - ) + tasks.append(asyncio.create_task(self._call_openapi_tool_handler(mcp_server, name, arguments))) else: return await self._call_regular_mcp_tool( mcp_server=mcp_server, @@ -3918,9 +3550,7 @@ class MCPServerManager: """ try: if asyncio.get_running_loop(): - asyncio.create_task( - self._initialize_tool_name_to_mcp_server_name_mapping() - ) + asyncio.create_task(self._initialize_tool_name_to_mcp_server_name_mapping()) except RuntimeError as e: # no running event loop verbose_logger.exception( f"No running event loop - skipping tool name to MCP server name mapping initialization: {str(e)}" @@ -3942,14 +3572,12 @@ class MCPServerManager: # at startup we have none, so an upstream 401 is normal. # Swallow it so we keep mapping the remaining servers. verbose_logger.debug( - f"Skipping tool name mapping for server {server.name} " - f"due to upstream auth error: {str(e)}" + f"Skipping tool name mapping for server {server.name} due to upstream auth error: {str(e)}" ) continue except Exception as e: verbose_logger.warning( - f"Failed to get tools from server {server.name} during " - f"tool name mapping initialization: {str(e)}" + f"Failed to get tools from server {server.name} during tool name mapping initialization: {str(e)}" ) continue for tool in tools: @@ -3992,9 +3620,7 @@ class MCPServerManager: # If not found and tool name is prefixed, extract the prefix and # match against any known form. - if is_tool_name_prefixed( - tool_name, known_server_prefixes=set(prefix_to_server.keys()) - ): + if is_tool_name_prefixed(tool_name, known_server_prefixes=set(prefix_to_server.keys())): ( original_tool_name, server_name_from_prefix, @@ -4020,9 +3646,7 @@ class MCPServerManager: self._upstream_initialize_instructions_probed_at.clear() # perform authz check to filter the mcp servers user has access to - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") # Load only "active", legacy "approved", and NULL (no approval workflow) rows. # Pending/rejected servers are excluded at the DB level so we never load them. from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable @@ -4064,16 +3688,12 @@ class MCPServerManager: alias=getattr(server, "alias", None), server_name=getattr(server, "server_name", None), ) - verbose_logger.debug( - f"Building server from DB: {server.server_id} ({server.server_name})" - ) + verbose_logger.debug(f"Building server from DB: {server.server_id} ({server.server_name})") # raw_rows come straight from the DB, so their global env var # values (like credentials) are still encrypted here, unlike the # already-decrypted records add_server/update_server are handed. # Decrypt them while building the registry entry. - new_server = await self.build_mcp_server_from_table( - server, env_vars_are_encrypted=True - ) + new_server = await self.build_mcp_server_from_table(server, env_vars_are_encrypted=True) # Carry the cached short_prefix from the previous registry entry # (if any) so the prefix is stable across reloads. if existing_server is not None and existing_server.short_prefix: @@ -4097,9 +3717,7 @@ class MCPServerManager: # Register OpenAPI tools *after* the final short prefix is assigned # so the tools are stored in the global registry under the same # prefix that lookups will use. - await self._maybe_register_openapi_tools( - new_server, initialize_mapping=False - ) + await self._maybe_register_openapi_tools(new_server, initialize_mapping=False) registered_registry[server_id] = new_server if new_server.spec_path: registered_openapi_tools = True @@ -4115,9 +3733,7 @@ class MCPServerManager: if registered_openapi_tools: self.initialize_tool_name_to_mcp_server_name_mapping() - verbose_logger.debug( - "MCP registry refreshed (%s servers in registry)", len(registered_registry) - ) + verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) def get_mcp_servers_from_ids(self, server_ids: List[str]) -> List[MCPServer]: servers = [] @@ -4139,9 +3755,7 @@ class MCPServerManager: # Fallback if proxy_server not available return {} - def _is_server_accessible_from_ip( - self, server: MCPServer, client_ip: Optional[str] - ) -> bool: + def _is_server_accessible_from_ip(self, server: MCPServer, client_ip: Optional[str]) -> bool: """ Check if a server is accessible from the given client IP. @@ -4159,9 +3773,7 @@ class MCPServerManager: return True # Non-public server: only accessible from internal IPs general_settings = self._get_general_settings() - internal_networks = IPAddressUtils.parse_internal_networks( - general_settings.get("mcp_internal_ip_ranges") - ) + internal_networks = IPAddressUtils.parse_internal_networks(general_settings.get("mcp_internal_ip_ranges")) return IPAddressUtils.is_internal_ip(client_ip, internal_networks) def get_mcp_server_by_id(self, server_id: str) -> Optional[MCPServer]: @@ -4195,11 +3807,7 @@ class MCPServerManager: if litellm.public_mcp_servers is None: return [] public_ids = set(litellm.public_mcp_servers) - return [ - server - for server in self.get_registry().values() - if server.server_id in public_ids - ] + return [server for server in self.get_registry().values() if server.server_id in public_ids] public_ids = set(litellm.public_mcp_servers or []) return [ @@ -4232,9 +3840,7 @@ class MCPServerManager: matches: List[str] = [ server_id for server_id, server in registry.items() - if server.alias == identifier - or server.server_name == identifier - or server.name == identifier + if server.alias == identifier or server.server_name == identifier or server.name == identifier ] if matches: expanded.update(matches) @@ -4274,9 +3880,7 @@ class MCPServerManager: result.setdefault(server_id, []).extend(tools or []) return result - def get_mcp_server_by_name( - self, server_name: str, client_ip: Optional[str] = None - ) -> Optional[MCPServer]: + def get_mcp_server_by_name(self, server_name: str, client_ip: Optional[str] = None) -> Optional[MCPServer]: """ Get the MCP Server from the server name. @@ -4311,9 +3915,7 @@ class MCPServerManager: return server return None - def get_filtered_registry( - self, client_ip: Optional[str] = None - ) -> Dict[str, MCPServer]: + def get_filtered_registry(self, client_ip: Optional[str] = None) -> Dict[str, MCPServer]: """ Get registry filtered by client IP access control. @@ -4324,11 +3926,7 @@ class MCPServerManager: registry = self.get_registry() if client_ip is None: return registry - return { - k: v - for k, v in registry.items() - if self._is_server_accessible_from_ip(v, client_ip) - } + return {k: v for k, v in registry.items() if self._is_server_accessible_from_ip(v, client_ip)} def _generate_stable_server_id( self, @@ -4357,9 +3955,7 @@ class MCPServerManager: A deterministic server ID string """ # Create a string from all the identifying parameters - params_string = ( - f"{server_name}|{url}|{transport}|{auth_type or ''}|{alias or ''}" - ) + params_string = f"{server_name}|{url}|{transport}|{auth_type or ''}|{alias or ''}" # Generate SHA-256 hash hash_object = hashlib.sha256(params_string.encode("utf-8")) @@ -4425,9 +4021,7 @@ class MCPServerManager: user_api_key_auth=None, raise_on_missing=False, ) - extra_headers = ( - dict(resolved_static_headers) if resolved_static_headers else {} - ) + extra_headers = dict(resolved_static_headers) if resolved_static_headers else {} client = await self._create_mcp_client( server=server, @@ -4442,15 +4036,11 @@ class MCPServerManager: return "ok" # Add timeout wrapper to prevent hanging - await asyncio.wait_for( - client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT - ) + await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT) self._remember_upstream_initialize_instructions(server, client) status = "healthy" except asyncio.TimeoutError: - health_check_error = ( - f"Health check timed out after {MCP_HEALTH_CHECK_TIMEOUT} seconds" - ) + health_check_error = f"Health check timed out after {MCP_HEALTH_CHECK_TIMEOUT} seconds" status = "unhealthy" except asyncio.CancelledError: health_check_error = "Health check was cancelled" @@ -4463,9 +4053,7 @@ class MCPServerManager: server_id=server.server_id, server_name=server.server_name, alias=server.alias, - description=( - server.mcp_info.get("description") if server.mcp_info else None - ), + description=(server.mcp_info.get("description") if server.mcp_info else None), url=server.url, transport=server.transport, auth_type=server.auth_type, @@ -4564,9 +4152,7 @@ class MCPServerManager: server_id=server.server_id, server_name=server.server_name, alias=server.alias, - description=( - server.mcp_info.get("description") if server.mcp_info else None - ), + description=(server.mcp_info.get("description") if server.mcp_info else None), url=server.url, spec_path=server.spec_path, transport=server.transport, @@ -4632,9 +4218,7 @@ class MCPServerManager: return await self._run_health_checks(target_server_ids) - async def _run_health_checks( - self, target_server_ids: List[str] - ) -> List[LiteLLM_MCPServerTable]: + async def _run_health_checks(self, target_server_ids: List[str]) -> List[LiteLLM_MCPServerTable]: if not target_server_ids: return [] diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 92ef57d8cd5..f18ff04c4e8 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -135,19 +135,12 @@ class MCPOAuth2TokenCache(InMemoryCache): access_token = body.get("access_token") if not access_token: - raise ValueError( - f"OAuth2 token response for MCP server '{server.server_id}' " - f"missing 'access_token'" - ) + raise ValueError(f"OAuth2 token response for MCP server '{server.server_id}' missing 'access_token'") # Safely parse expires_in — providers may return null or non-numeric values raw_expires_in = body.get("expires_in") try: - expires_in = ( - int(raw_expires_in) - if raw_expires_in is not None - else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL - ) + expires_in = int(raw_expires_in) if raw_expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL except (TypeError, ValueError): expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL @@ -289,9 +282,7 @@ async def resolve_mcp_auth( return mcp_auth_header if server.has_token_exchange_config: if subject_token: - return await token_exchange.mcp_token_exchange_handler.exchange_token( - subject_token, server - ) + return await token_exchange.mcp_token_exchange_handler.exchange_token(subject_token, server) # No subject_token — fall back to client_credentials using the same client # credentials and token_url so M2M scenarios still work. if server.client_id and server.client_secret and server.token_url: diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index e8b591c39cf..4d5813dbc5b 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -379,10 +379,8 @@ def _trusted_redirect_uri_is_allowed( ) -> bool: if proxy_base: proxy_parsed = urlparse(proxy_base) - if ( - parsed.scheme == proxy_parsed.scheme - and redirect_netloc - == _strip_default_port(proxy_parsed.scheme, proxy_parsed.netloc) + if parsed.scheme == proxy_parsed.scheme and redirect_netloc == _strip_default_port( + proxy_parsed.scheme, proxy_parsed.netloc ): return True @@ -418,9 +416,7 @@ def _build_trusted_redirect_rejection_message( redirect_origin = _origin_label(parsed.scheme, redirect_netloc) proxy_parsed = urlparse(proxy_base) if proxy_base else None proxy_netloc_norm = ( - _strip_default_port(proxy_parsed.scheme, proxy_parsed.netloc) - if proxy_parsed and proxy_parsed.netloc - else "" + _strip_default_port(proxy_parsed.scheme, proxy_parsed.netloc) if proxy_parsed and proxy_parsed.netloc else "" ) mismatch_parts: List[str] = [] @@ -433,16 +429,10 @@ def _build_trusted_redirect_rejection_message( "or trust X-Forwarded-Proto from your ingress)" ) if redirect_netloc != proxy_netloc_norm: - mismatch_parts.append( - f"host/port: redirect_uri {redirect_netloc!r} does not match " - "the proxy origin" - ) + mismatch_parts.append(f"host/port: redirect_uri {redirect_netloc!r} does not match the proxy origin") if mismatch_parts: - return ( - f"redirect_uri origin ({redirect_origin}) does not match the proxy " - "origin. " + "; ".join(mismatch_parts) - ) + return f"redirect_uri origin ({redirect_origin}) does not match the proxy origin. " + "; ".join(mismatch_parts) return ( f"redirect_uri ({redirect_uri!r}) is not allowed: not same-origin with " f"the proxy origin, not loopback, and not listed in " @@ -457,9 +447,7 @@ def _raise_trusted_redirect_uri_rejected( redirect_netloc: str, proxy_base: Optional[str], ) -> NoReturn: - description = _build_trusted_redirect_rejection_message( - redirect_uri, parsed, redirect_netloc, proxy_base - ) + description = _build_trusted_redirect_rejection_message(redirect_uri, parsed, redirect_netloc, proxy_base) hint = ( "Align the proxy public URL with the browser URL. Set PROXY_BASE_URL to your " @@ -525,6 +513,4 @@ def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: proxy_base = _resolve_proxy_base_for_redirect(request) if _trusted_redirect_uri_is_allowed(parsed, redirect_netloc, proxy_base): return - _raise_trusted_redirect_uri_rejected( - request, redirect_uri, parsed, redirect_netloc, proxy_base - ) + _raise_trusted_redirect_uri_rejected(request, redirect_uri, parsed, redirect_netloc, proxy_base) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index de70fe1331e..1ee300be718 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -58,8 +58,8 @@ _request_auth_header: contextvars.ContextVar[Optional[str]] = contextvars.Contex # Per-request extra headers forwarded from the client request. # Populated from MCPServer.extra_headers names matched against raw request # headers in server.py before dispatching to a local/OpenAPI tool handler. -_request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = ( - contextvars.ContextVar("_request_extra_headers", default=None) +_request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = contextvars.ContextVar( + "_request_extra_headers", default=None ) @@ -74,14 +74,10 @@ def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: normalized_value = value_str.replace("\\", "/") if "/" in normalized_value: - raise ValueError( - f"Path parameter '{param_name}' must not contain path separators" - ) + raise ValueError(f"Path parameter '{param_name}' must not contain path separators") if any(part in {".", ".."} for part in PurePosixPath(normalized_value).parts): - raise ValueError( - f"Path parameter '{param_name}' cannot include '.' or '..' segments" - ) + raise ValueError(f"Path parameter '{param_name}' cannot include '.' or '..' segments") return quote(value_str, safe="") @@ -149,9 +145,7 @@ def get_base_url(spec: Dict[str, Any], spec_path: Optional[str] = None) -> str: return f"{scheme}://{spec['host']}{base_path}" # Fallback: derive base URL from spec_path if it's a URL - if spec_path and ( - spec_path.startswith("http://") or spec_path.startswith("https://") - ): + if spec_path and (spec_path.startswith("http://") or spec_path.startswith("https://")): for suffix in [ "/openapi.json", "/openapi.yaml", @@ -160,24 +154,18 @@ def get_base_url(spec: Dict[str, Any], spec_path: Optional[str] = None) -> str: ]: if spec_path.endswith(suffix): base_url = spec_path[: -len(suffix)] - verbose_logger.info( - f"No server info in OpenAPI spec. Using derived base URL: {base_url}" - ) + verbose_logger.info(f"No server info in OpenAPI spec. Using derived base URL: {base_url}") return base_url if spec_path.split("/")[-1].endswith((".json", ".yaml", ".yml")): base_url = "/".join(spec_path.split("/")[:-1]) - verbose_logger.info( - f"No server info in OpenAPI spec. Using derived base URL: {base_url}" - ) + verbose_logger.info(f"No server info in OpenAPI spec. Using derived base URL: {base_url}") return base_url return "" -def _resolve_ref( - param: Dict[str, Any], component_params: Dict[str, Any] -) -> Optional[Dict[str, Any]]: +def _resolve_ref(param: Dict[str, Any], component_params: Dict[str, Any]) -> Optional[Dict[str, Any]]: """Resolve a single parameter, following a $ref if present. Returns the resolved param dict, or None if the $ref target is absent from @@ -190,9 +178,7 @@ def _resolve_ref( return component_params.get(ref.split("/")[-1]) -def _resolve_param_list( - raw: List[Dict[str, Any]], component_params: Dict[str, Any] -) -> List[Dict[str, Any]]: +def _resolve_param_list(raw: List[Dict[str, Any]], component_params: Dict[str, Any]) -> List[Dict[str, Any]]: """Resolve $refs in a parameter list, dropping any unresolvable entries.""" result = [] for p in raw: @@ -225,9 +211,7 @@ def resolve_operation_params( path_level = _resolve_param_list(path_item.get("parameters", []), component_params) op_level = _resolve_param_list(operation.get("parameters", []), component_params) op_keys = {(p["name"], p.get("in")) for p in op_level} - merged = [ - p for p in path_level if (p["name"], p.get("in")) not in op_keys - ] + op_level + merged = [p for p in path_level if (p["name"], p.get("in")) not in op_keys] + op_level result = dict(operation) result["parameters"] = merged return result @@ -330,9 +314,7 @@ def _merge_openapi_tool_request_headers( static = static_headers or {} static_lower_names = {k.lower() for k in static} - effective_headers: Dict[str, str] = { - k: v for k, v in request_extra.items() if k.lower() not in static_lower_names - } + effective_headers: Dict[str, str] = {k: v for k, v in request_extra.items() if k.lower() not in static_lower_names} effective_headers.update(static) override_auth = _request_auth_header.get() @@ -424,11 +406,7 @@ def create_tool_function( elif body_value: # If it's a string, try to parse as JSON try: - json_body = ( - json.loads(body_value) - if isinstance(body_value, str) - else {"data": body_value} - ) + json_body = json.loads(body_value) if isinstance(body_value, str) else {"data": body_value} except (json.JSONDecodeError, TypeError): json_body = {"data": body_value} @@ -437,21 +415,13 @@ def create_tool_function( if original_method == "get": response = await client.get(url, params=params, headers=effective_headers) elif original_method == "post": - response = await client.post( - url, params=params, json=json_body, headers=effective_headers - ) + response = await client.post(url, params=params, json=json_body, headers=effective_headers) elif original_method == "put": - response = await client.put( - url, params=params, json=json_body, headers=effective_headers - ) + response = await client.put(url, params=params, json=json_body, headers=effective_headers) elif original_method == "delete": - response = await client.delete( - url, params=params, headers=effective_headers - ) + response = await client.delete(url, params=params, headers=effective_headers) elif original_method == "patch": - response = await client.patch( - url, params=params, json=json_body, headers=effective_headers - ) + response = await client.patch(url, params=params, json=json_body, headers=effective_headers) else: return f"Unsupported HTTP method: {original_method}" @@ -488,16 +458,12 @@ def register_tools_from_openapi(spec: Dict[str, Any], base_url: str): while unique in used_names: n += 1 suffix = f"_{n}" - unique = ( - tool_name[: _OPENAPI_TOOL_NAME_MAX_LEN - len(suffix)] + suffix - ) + unique = tool_name[: _OPENAPI_TOOL_NAME_MAX_LEN - len(suffix)] + suffix tool_name = unique used_names.add(tool_name) # Get description - description = operation.get( - "summary", operation.get("description", f"{method.upper()} {path}") - ) + description = operation.get("summary", operation.get("description", f"{method.upper()} {path}")) # Build input schema input_schema = build_input_schema(operation) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index b3e1cd844d8..815fc2ba29d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -34,9 +34,7 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer -def to_subject( - user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optional[str] -) -> Subject: +def to_subject(user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optional[str]) -> Subject: """Map v1's authenticated principal onto the resolver's Subject. tenant_id / subject_id are empty for an unauthenticated caller; the per-user arms must reject @@ -67,18 +65,14 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: (M2M), delegated/passthrough oauth2, token exchange, and SigV4 return None and stay on v1. """ if server.is_byok: - return ( - None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) - ) + return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) resource = server.url or server.server_id auth_type = server.auth_type match auth_type: case None | MCPAuth.none: if server.is_oauth_passthrough: return None # passthrough is not migrated yet -> defer to v1 - return ServerSpec( - server_id=server.server_id, resource=resource, config=NoneConfig() - ) + return ServerSpec(server_id=server.server_id, resource=resource, config=NoneConfig()) case MCPAuth.api_key: return _shared_key_spec(server, resource, "X-API-Key", "") case MCPAuth.bearer_token: @@ -88,9 +82,7 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: case MCPAuth.authorization: return _shared_key_spec(server, resource, "Authorization", "") case MCPAuth.basic: - return _shared_key_spec( - server, resource, "Authorization", "Basic", encode=True - ) + return _shared_key_spec(server, resource, "Authorization", "Basic", encode=True) case MCPAuth.oauth2: if server.needs_user_oauth_token and not server.delegate_auth_to_upstream: return ServerSpec( @@ -141,11 +133,7 @@ def raise_public(error: CredError) -> NoReturn: raise HTTPException( status_code=401, detail=challenge.body if challenge.body is not None else error.summary, - headers=( - {"WWW-Authenticate": challenge.www_authenticate} - if challenge.www_authenticate - else None - ), + headers=({"WWW-Authenticate": challenge.www_authenticate} if challenge.www_authenticate else None), ) case "misconfigured": raise HTTPException(status_code=500, detail=error.summary) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py index 04b7a54aaa0..2504ff67e3e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -22,9 +22,7 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer ServerLookup = Callable[[str], "MCPServer | None"] -TokenEndpointPost = Callable[ - [str, dict[str, str]], Awaitable["dict[str, object] | None"] -] +TokenEndpointPost = Callable[[str, dict[str, str]], Awaitable["dict[str, object] | None"]] class CredentialPersist(Protocol): @@ -81,9 +79,7 @@ class AuthorizationCodeRefresher: self._persist = persist self._clock = clock - async def refresh( - self, user_id: str, server_id: str, token: OAuthToken - ) -> OAuthToken | None: + async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: if token.refresh_token is None: return None server = self._server_lookup(server_id) @@ -104,15 +100,11 @@ class AuthorizationCodeRefresher: return None rotated = body.get("refresh_token") - new_refresh = ( - rotated if isinstance(rotated, str) and rotated else token.refresh_token - ) + new_refresh = rotated if isinstance(rotated, str) and rotated else token.refresh_token expires_in = _parse_expires_in(body.get("expires_in")) scopes = _parse_scopes(body.get("scope")) or token.scopes - await self._persist( - user_id, server_id, access_token, new_refresh, expires_in, scopes or None - ) + await self._persist(user_id, server_id, access_token, new_refresh, expires_in, scopes or None) return OAuthToken( access_token=access_token, expires_at=self._clock() + expires_in if expires_in is not None else None, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py index 2345fa98123..e4d8fd25748 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py @@ -19,9 +19,7 @@ from pydantic import SecretStr class NoOpAuth(httpx.Auth): """Attaches nothing — the `none` mode (and the seam-level default).""" - def auth_flow( - self, request: httpx.Request - ) -> Generator[httpx.Request, httpx.Response, None]: + def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: yield request @@ -38,8 +36,6 @@ class StaticHeaderAuth(httpx.Auth): self.header_name = header_name self._header_value = SecretStr(header_value) - def auth_flow( - self, request: httpx.Request - ) -> Generator[httpx.Request, httpx.Response, None]: + def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: request.headers[self.header_name] = self._header_value.get_secret_value() yield request diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py index 1c089f2a931..9c7ab326f29 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py @@ -82,9 +82,7 @@ class TokenRefresher(Protocol): are not derivable from ``token``, so the seam threads them alongside it. """ - async def refresh( - self, user_id: str, server_id: str, token: OAuthToken - ) -> OAuthToken | None: ... + async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: ... class TokenCacheBackend(Protocol): @@ -96,9 +94,7 @@ class TokenCacheBackend(Protocol): async def get(self, user_id: str, server_id: str) -> OAuthToken | None: ... - async def set( - self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float - ) -> None: ... + async def set(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None: ... async def delete(self, user_id: str, server_id: str) -> None: ... @@ -106,9 +102,7 @@ class TokenCacheBackend(Protocol): class InMemoryTokenCacheBackend: """Per-process token cache: a bounded dict with wall-clock TTLs (the default backend).""" - def __init__( - self, *, max_size: int = 4096, clock: Callable[[], float] = time.time - ) -> None: + def __init__(self, *, max_size: int = 4096, clock: Callable[[], float] = time.time) -> None: self._max_size = max_size self._clock = clock self._cache: dict[tuple[str, str], tuple[OAuthToken, float]] = {} @@ -124,9 +118,7 @@ class InMemoryTokenCacheBackend: self._cache.pop(key, None) return None - async def set( - self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float - ) -> None: + async def set(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None: key = (user_id, server_id) if key not in self._cache and len(self._cache) >= self._max_size: # Evict the oldest entry (insertion order), rather than clearing the whole cache and @@ -164,15 +156,11 @@ class CachedOAuthTokenStore: self._default_ttl_seconds = default_ttl_seconds self._expiry_skew_seconds = expiry_skew_seconds self._clock = clock - self._backend: TokenCacheBackend = backend or InMemoryTokenCacheBackend( - max_size=max_size, clock=clock - ) + self._backend: TokenCacheBackend = backend or InMemoryTokenCacheBackend(max_size=max_size, clock=clock) def _ttl(self, token: OAuthToken) -> float: if token.expires_at is not None: - return max( - 0.0, token.expires_at - self._expiry_skew_seconds - self._clock() - ) + return max(0.0, token.expires_at - self._expiry_skew_seconds - self._clock()) return self._default_ttl_seconds async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: @@ -268,15 +256,10 @@ class RefreshingTokenStore: self._refresher = refresher self._expiry_skew_seconds = expiry_skew_seconds self._clock = clock - self._coordinator: RefreshCoordinator = ( - coordinator or InProcessRefreshCoordinator() - ) + self._coordinator: RefreshCoordinator = coordinator or InProcessRefreshCoordinator() def _is_expired(self, token: OAuthToken) -> bool: - return ( - token.expires_at is not None - and self._clock() >= token.expires_at - self._expiry_skew_seconds - ) + return token.expires_at is not None and self._clock() >= token.expires_at - self._expiry_skew_seconds async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: token = await self._inner.fetch(user_id, server_id) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 7ffde55c2c1..cedd13e7b2f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -74,9 +74,7 @@ async def _persist_credential( ) -async def _post_token_endpoint( - url: str, form: dict[str, str] -) -> dict[str, object] | None: +async def _post_token_endpoint(url: str, form: dict[str, str]) -> dict[str, object] | None: from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 get_async_httpx_client, # pyright: ignore ) @@ -102,9 +100,7 @@ async def _post_token_endpoint( def build_per_user_oauth_token_store( server_lookup: ServerLookup, ) -> CachedOAuthTokenStore: - refresher = AuthorizationCodeRefresher( - server_lookup, _post_token_endpoint, _persist_credential - ) + refresher = AuthorizationCodeRefresher(server_lookup, _post_token_endpoint, _persist_credential) # Cache and refresh coordinator use the foundation's in-process defaults (a single replica needs # no shared cache or lock); the cross-replica path is layered on separately. refreshing = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 87fa66aeab9..f9a9fa00b23 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -64,13 +64,9 @@ class UpstreamCredentialProvider: """ def __init__(self, oauth_token_store: OAuthTokenStore | None = None) -> None: - self._oauth_token_store: OAuthTokenStore = ( - oauth_token_store or _NullOAuthTokenStore() - ) + self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore() - async def resolve_credentials( - self, subject: Subject, server: ServerSpec - ) -> Result[httpx.Auth, CredError]: + async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: case NoneConfig(): return Ok(NoOpAuth()) @@ -101,52 +97,30 @@ class UpstreamCredentialProvider: def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]: match config.key_source: case SharedKey() as source: - header_name, header_value = config.header( - source.value.get_secret_value() - ) + header_name, header_value = config.header(source.value.get_secret_value()) return Ok(StaticHeaderAuth(header_value, header_name=header_name)) case Byok(): # Per-user key pulled from the credential store; lands with that seam. - return Error( - CredError.of_not_implemented( - "api_key BYOK source not implemented yet" - ) - ) + return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet")) assert_never(config.key_source) - async def _authorization_code( - self, subject: Subject, server: ServerSpec - ) -> Result[StaticHeaderAuth, CredError]: + async def _authorization_code(self, subject: Subject, server: ServerSpec) -> Result[StaticHeaderAuth, CredError]: token = await self._authz_token(subject, server) if token is None: - return Error( - CredError.of_unauthorized( - "Authorization required: complete the OAuth flow for this server." - ) - ) - return Ok( - StaticHeaderAuth( - f"Bearer {token.access_token}", header_name="Authorization" - ) - ) + return Error(CredError.of_unauthorized("Authorization required: complete the OAuth flow for this server.")) + return Ok(StaticHeaderAuth(f"Bearer {token.access_token}", header_name="Authorization")) - async def _authz_token( - self, subject: Subject, server: ServerSpec - ) -> OAuthToken | None: + async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None: """The user's authorization_code token, or None when absent or the store is unreachable. A store outage is mapped to None (the OAuth challenge), not raised, so a transient outage does not 500; it is the store, not this resolver, that declines to cache the failure. """ try: - return await self._oauth_token_store.fetch( - subject.subject_id, server.server_id - ) + return await self._oauth_token_store.fetch(subject.subject_id, server.server_id) except TokenStoreUnavailable: return None def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: - return Error( - CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet") - ) + return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet")) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 8e589d3a24b..671de63eabe 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -91,24 +91,12 @@ class CredError: "not_implemented", ] = tag() - unauthorized: Unauthorized = ( - case() - ) # no usable credential for this (subject, server) -> 401 challenge - misconfigured: str = ( - case() - ) # the declared mode is missing required config -> 5xx (operator) - upstream_unavailable: str = ( - case() - ) # the IdP / token endpoint could not be reached -> 503 - unsupported_mode: str = ( - case() - ) # a raw mode string did not parse into AuthSpecKind (boundary) - precondition_required: str = ( - case() - ) # a required per-user value (e.g. an env var) has not been provided -> 412 - not_implemented: str = ( - case() - ) # the declared mode's resolver arm is not built yet -> 501 (not operator error) + unauthorized: Unauthorized = case() # no usable credential for this (subject, server) -> 401 challenge + misconfigured: str = case() # the declared mode is missing required config -> 5xx (operator) + upstream_unavailable: str = case() # the IdP / token endpoint could not be reached -> 503 + unsupported_mode: str = case() # a raw mode string did not parse into AuthSpecKind (boundary) + precondition_required: str = case() # a required per-user value (e.g. an env var) has not been provided -> 412 + not_implemented: str = case() # the declared mode's resolver arm is not built yet -> 501 (not operator error) @staticmethod def of_unauthorized( @@ -117,11 +105,7 @@ class CredError: www_authenticate: str | None = None, body: Mapping[str, str] | None = None, ) -> CredError: - return CredError( - unauthorized=Unauthorized( - detail=detail, www_authenticate=www_authenticate, body=body - ) - ) + return CredError(unauthorized=Unauthorized(detail=detail, www_authenticate=www_authenticate, body=body)) @staticmethod def of_misconfigured(detail: str) -> CredError: @@ -292,9 +276,7 @@ class Ambient(BaseModel): source: Literal["ambient"] = "ambient" -AwsCredentialSource = Annotated[ - StaticKeys | AssumeRole | Ambient, Field(discriminator="source") -] +AwsCredentialSource = Annotated[StaticKeys | AssumeRole | Ambient, Field(discriminator="source")] class AwsSigV4Config(BaseModel): diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 0d71174e9d7..d30d8af2af2 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -55,16 +55,12 @@ def _connection_error_message(exc: BaseException) -> str: ) if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout)): return ( - "Failed to connect to MCP server: the server is unreachable. " - "Check the URL and that the server is running." + "Failed to connect to MCP server: the server is unreachable. Check the URL and that the server is running." ) if isinstance(exc, httpx.TimeoutException): return "Failed to connect to MCP server: the connection timed out." if isinstance(exc, httpx.HTTPStatusError): - return ( - f"Failed to connect to MCP server: it returned HTTP " - f"{exc.response.status_code}." - ) + return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}." return "Failed to connect to MCP server. Check proxy logs for details." @@ -117,10 +113,7 @@ if MCP_AVAILABLE: return { sid for sid in allowed_server_ids - if getattr( - global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None - ) - == MCPAuth.oauth2 + if getattr(global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None) == MCPAuth.oauth2 } async def _get_user_oauth_extra_headers( @@ -159,9 +152,7 @@ if MCP_AVAILABLE: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to use OAuth2 MCP tools." ) - cred = await get_user_oauth_credential( - prisma_client, user_id, server_id - ) + cred = await get_user_oauth_credential(prisma_client, user_id, server_id) cred = await resolve_valid_user_oauth_token( user_id=user_id, server=server, @@ -200,9 +191,7 @@ if MCP_AVAILABLE: creds = await list_user_oauth_credentials(prisma_client, user_id) return {c["server_id"]: c for c in creds if "server_id" in c} except Exception as e: - verbose_logger.warning( - f"_prefetch_user_oauth_creds: failed to prefetch for user={user_id}: {e}" - ) + verbose_logger.warning(f"_prefetch_user_oauth_creds: failed to prefetch for user={user_id}: {e}") return {} async def _get_bulk_user_oauth_headers( @@ -234,9 +223,7 @@ if MCP_AVAILABLE: if c.get("access_token") and c.get("server_id") } except Exception: - verbose_logger.debug( - "Failed to bulk-fetch OAuth credentials", exc_info=True - ) + verbose_logger.debug("Failed to bulk-fetch OAuth credentials", exc_info=True) return {} def _create_tool_response_objects(tools, server: MCPServer): @@ -273,12 +260,8 @@ if MCP_AVAILABLE: """ headers = request.headers raw_headers = dict(headers) - mcp_auth_header = mcp_request_handler_cls._get_mcp_auth_header_from_headers( - headers - ) - mcp_server_auth_headers = ( - mcp_request_handler_cls._get_mcp_server_auth_headers_from_headers(headers) - ) + mcp_auth_header = mcp_request_handler_cls._get_mcp_auth_header_from_headers(headers) + mcp_server_auth_headers = mcp_request_handler_cls._get_mcp_server_auth_headers_from_headers(headers) return mcp_auth_header, mcp_server_auth_headers, raw_headers def _resolve_mcp_server_id_for_rest( @@ -295,9 +278,7 @@ if MCP_AVAILABLE: allowed = set(allowed_server_ids) if server_id in allowed: return server_id - by_name = global_mcp_server_manager.get_mcp_server_by_name( - server_id, client_ip=client_ip - ) + by_name = global_mcp_server_manager.get_mcp_server_by_name(server_id, client_ip=client_ip) if by_name is not None and by_name.server_id in allowed: return by_name.server_id return server_id @@ -334,14 +315,10 @@ if MCP_AVAILABLE: allowed_server_ids_set.update(servers) allowed_server_ids_set = set( - global_mcp_server_manager.filter_server_ids_by_ip( - list(allowed_server_ids_set), _rest_client_ip - ) + global_mcp_server_manager.filter_server_ids_by_ip(list(allowed_server_ids_set), _rest_client_ip) ) - canonical_server_id = _resolve_mcp_server_id_for_rest( - server_id, allowed_server_ids_set, _rest_client_ip - ) + canonical_server_id = _resolve_mcp_server_id_for_rest(server_id, allowed_server_ids_set, _rest_client_ip) if canonical_server_id not in allowed_server_ids_set: _server = global_mcp_server_manager.get_mcp_server_by_id( @@ -350,9 +327,7 @@ if MCP_AVAILABLE: if ( _server is not None and _rest_client_ip is not None - and not global_mcp_server_manager._is_server_accessible_from_ip( - _server, _rest_client_ip - ) + and not global_mcp_server_manager._is_server_accessible_from_ip(_server, _rest_client_ip) ): raise HTTPException( status_code=403, @@ -431,21 +406,12 @@ if MCP_AVAILABLE: ): # Dict keys may be server_ids OR names/aliases; normalize so lookup # by concrete server_id resolves name-keyed restrictions too. - allowed_tools_for_server = ( - global_mcp_server_manager.expand_tool_permissions( - user_api_key_auth.object_permission.mcp_tool_permissions - ).get(server.server_id) - ) - if ( - allowed_tools_for_server is not None - and len(allowed_tools_for_server) > 0 - ): + allowed_tools_for_server = global_mcp_server_manager.expand_tool_permissions( + user_api_key_auth.object_permission.mcp_tool_permissions + ).get(server.server_id) + if allowed_tools_for_server is not None and len(allowed_tools_for_server) > 0: # Filter tools to only include those in the allowed list - tools = [ - tool - for tool in tools - if _tool_name_matches(tool.name, allowed_tools_for_server) - ] + tools = [tool for tool in tools if _tool_name_matches(tool.name, allowed_tools_for_server)] return _create_tool_response_objects(tools, server) @@ -457,9 +423,7 @@ if MCP_AVAILABLE: auth_contexts = await build_effective_auth_contexts(user_api_key_dict) allowed_server_ids_set = set() for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth=auth_context - ) + servers = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth=auth_context) allowed_server_ids_set.update(servers) if server_id not in allowed_server_ids_set: raise HTTPException( @@ -491,22 +455,15 @@ if MCP_AVAILABLE: _name_resolved = None if server_id not in allowed_server_ids: _name_resolved = global_mcp_server_manager.get_mcp_server_by_name(server_id) - if _name_resolved is not None and _name_resolved.server_id in set( - allowed_server_ids - ): + if _name_resolved is not None and _name_resolved.server_id in set(allowed_server_ids): server_id = _name_resolved.server_id if server_id not in allowed_server_ids: - _server = ( - global_mcp_server_manager.get_mcp_server_by_id(server_id) - or _name_resolved - ) + _server = global_mcp_server_manager.get_mcp_server_by_id(server_id) or _name_resolved if ( _server is not None and rest_client_ip is not None - and not global_mcp_server_manager._is_server_accessible_from_ip( - _server, rest_client_ip - ) + and not global_mcp_server_manager._is_server_accessible_from_ip(_server, rest_client_ip) ): raise HTTPException( status_code=403, @@ -535,12 +492,8 @@ if MCP_AVAILABLE: "message": f"Server with id {server_id} not found", } - server_auth_header = _get_server_auth_header( - server, mcp_server_auth_headers, mcp_auth_header - ) - user_oauth_extra_headers = await _get_user_oauth_extra_headers( - server, user_api_key_dict - ) + server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header) + user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict) try: list_tools_result = await _get_tools_for_single_server( @@ -572,9 +525,7 @@ if MCP_AVAILABLE: @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)]) async def list_tool_rest_api( request: Request, - server_id: Optional[str] = Query( - None, description="The server id to list tools for" - ), + server_id: Optional[str] = Query(None, description="The server id to list tools for"), include_disabled_tools: bool = Query( False, description=( @@ -615,19 +566,14 @@ if MCP_AVAILABLE: # The full catalog (allowlist filter skipped) is admin-only so the # REST endpoint can't be used to enumerate deliberately-disabled tools. apply_tool_filters = not ( - include_disabled_tools - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + include_disabled_tools and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ) # Extract auth headers from request headers = request.headers raw_headers_from_request = dict(headers) - mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( - headers - ) - mcp_server_auth_headers = ( - MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) - ) + mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers) + mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) auth_contexts = await build_effective_auth_contexts(user_api_key_dict) @@ -696,15 +642,11 @@ if MCP_AVAILABLE: # Query all servers the user has access to errors = [] for allowed_server_id in allowed_server_ids: - server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_server_id - ) + server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) if server is None: continue - server_auth_header = _get_server_auth_header( - server, mcp_server_auth_headers, mcp_auth_header - ) + server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header) user_oauth_extra_headers = await _get_user_oauth_extra_headers( server, user_api_key_dict, @@ -722,23 +664,17 @@ if MCP_AVAILABLE: ) list_tools_result.extend(tools_result) except Exception as e: - verbose_logger.exception( - f"Error getting tools from {server.name}: {e}" - ) + verbose_logger.exception(f"Error getting tools from {server.name}: {e}") errors.append(f"{server.name}: {str(e)}") continue if errors and not list_tools_result: - error_message = "Failed to get tools from servers: " + "; ".join( - errors - ) + error_message = "Failed to get tools from servers: " + "; ".join(errors) return { "tools": list_tools_result, "error": "partial_failure" if error_message else None, - "message": ( - error_message if error_message else "Successfully retrieved tools" - ), + "message": (error_message if error_message else "Successfully retrieved tools"), } except MCPUpstreamAuthError as e: @@ -751,18 +687,14 @@ if MCP_AVAILABLE: except HTTPException as http_exc: # Internal access/IP 403s keep the legacy error-dict response shape # so the existing contract stays intact. - verbose_logger.exception( - "HTTPException in list_tool_rest_api: %s", str(http_exc) - ) + verbose_logger.exception("HTTPException in list_tool_rest_api: %s", str(http_exc)) return { "tools": [], "error": "unexpected_error", "message": (f"An unexpected error occurred: {http_exc.detail}"), } except Exception as e: - verbose_logger.exception( - "Unexpected error in list_tool_rest_api: %s", str(e) - ) + verbose_logger.exception("Unexpected error in list_tool_rest_api: %s", str(e)) return { "tools": [], "error": "unexpected_error", @@ -852,9 +784,7 @@ if MCP_AVAILABLE: ( allowed_mcp_servers, canonical_server_id, - ) = await _resolve_allowed_mcp_servers_with_ip_filter( - request, user_api_key_dict, server_id - ) + ) = await _resolve_allowed_mcp_servers_with_ip_filter(request, user_api_key_dict, server_id) # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). user_oauth_extra_headers: Optional[Dict[str, str]] = None @@ -863,9 +793,7 @@ if MCP_AVAILABLE: None, ) if target_server is not None: - user_oauth_extra_headers = await _get_user_oauth_extra_headers( - target_server, user_api_key_dict - ) + user_oauth_extra_headers = await _get_user_oauth_extra_headers(target_server, user_api_key_dict) # Call execute_mcp_tool directly (permission checks already done) result = await execute_mcp_tool( @@ -958,9 +886,7 @@ if MCP_AVAILABLE: client_id: Optional[str] = creds.get("client_id") client_secret: Optional[str] = creds.get("client_secret") scopes_raw = creds.get("scopes") - scopes: Optional[List[str]] = ( - scopes_raw if isinstance(scopes_raw, list) else None - ) + scopes: Optional[List[str]] = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes async def _execute_with_mcp_client( @@ -991,12 +917,8 @@ if MCP_AVAILABLE: try: client_id, client_secret, scopes = _extract_credentials(request) - _oauth2_flow: Optional[ - Literal["client_credentials", "authorization_code"] - ] = request.oauth2_flow or ( - "client_credentials" - if client_id and client_secret and request.token_url - else None + _oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = request.oauth2_flow or ( + "client_credentials" if client_id and client_secret and request.token_url else None ) # client_credentials requires token_url to fetch a token; without it the # incoming auth header would be dropped with nothing to replace it. @@ -1024,15 +946,11 @@ if MCP_AVAILABLE: instructions=request.instructions, ) - stdio_env = global_mcp_server_manager._build_stdio_env( - server_model, raw_headers - ) + stdio_env = global_mcp_server_manager._build_stdio_env(server_model, raw_headers) # For M2M OAuth servers, drop the incoming Authorization header so that # resolve_mcp_auth can auto-fetch a token via client_credentials. - effective_oauth2_headers = ( - None if server_model.has_client_credentials else oauth2_headers - ) + effective_oauth2_headers = None if server_model.has_client_credentials else oauth2_headers # Interactive authorization_code tools preview: the operator holds a just-authorized # token but it is not persisted yet. Resolve it through the v2 resolver via a one-shot @@ -1055,9 +973,7 @@ if MCP_AVAILABLE: ) forwarded_authorization = ( - effective_oauth2_headers.get("Authorization") - if effective_oauth2_headers - else None + effective_oauth2_headers.get("Authorization") if effective_oauth2_headers else None ) is_interactive_authz_code = ( server_model.auth_type == MCPAuth.oauth2 @@ -1079,9 +995,7 @@ if MCP_AVAILABLE: ) merged_headers = merge_mcp_headers( - extra_headers=( - None if preview_cred_provider else effective_oauth2_headers - ), + extra_headers=(None if preview_cred_provider else effective_oauth2_headers), static_headers=request.static_headers, ) @@ -1127,9 +1041,7 @@ if MCP_AVAILABLE: if operation is None: continue - resolved_op = resolve_operation_params( - operation, path_item, components - ) + resolved_op = resolve_operation_params(operation, path_item, components) raw_op_id = operation.get("operationId", f"{method}_{path}") # Match what register_tools_from_openapi does so the preview @@ -1143,9 +1055,7 @@ if MCP_AVAILABLE: while unique in used_names: n += 1 suffix = f"_{n}" - unique = ( - op_id[: _OPENAPI_TOOL_NAME_MAX_LEN - len(suffix)] + suffix - ) + unique = op_id[: _OPENAPI_TOOL_NAME_MAX_LEN - len(suffix)] + suffix op_id = unique used_names.add(op_id) summary = operation.get("summary", "") @@ -1154,9 +1064,7 @@ if MCP_AVAILABLE: tools.append( { "name": op_id, - "description": description - or summary - or f"{method.upper()} {path}", + "description": description or summary or f"{method.upper()} {path}", "inputSchema": input_schema, } ) @@ -1220,9 +1128,7 @@ if MCP_AVAILABLE: }, ) - new_mcp_server_request = _inherit_credentials_from_existing_server( - new_mcp_server_request - ) + new_mcp_server_request = _inherit_credentials_from_existing_server(new_mcp_server_request) # For OpenAPI spec servers, generate tools from the spec directly if new_mcp_server_request.spec_path: @@ -1253,13 +1159,9 @@ if MCP_AVAILABLE: async def _list_tools_session_operation(session): return await session.list_tools() - list_tools_response = await client.run_with_session( - _list_tools_session_operation - ) + list_tools_response = await client.run_with_session(_list_tools_session_operation) list_tools_result: List[MCPTool] = list_tools_response.tools - model_dumped_tools: List[dict] = [ - tool.model_dump() for tool in list_tools_result - ] + model_dumped_tools: List[dict] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, "error": None, diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index b659ba6f813..65630f74e90 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -97,25 +97,19 @@ def _resolve_model_from_preferences( for model_name in available_model_names: if hint_name.lower() in model_name.lower(): verbose_logger.debug( - "MCP sampling model resolution: substring hint match " - "'%s' -> '%s'", + "MCP sampling model resolution: substring hint match '%s' -> '%s'", hint_name, model_name, ) return model_name verbose_logger.debug( - "MCP sampling model resolution: no hint matched from %s " - "against %d available models", + "MCP sampling model resolution: no hint matched from %s against %d available models", [getattr(h, "name", None) for h in model_preferences.hints], len(available_model_names), ) # 2. Priority-based selection (cost/speed/intelligence) - if ( - model_preferences - and available_model_names - and _has_priorities(model_preferences) - ): + if model_preferences and available_model_names and _has_priorities(model_preferences): best = _select_model_by_priority(available_model_names, model_preferences) if best is not None: verbose_logger.debug( @@ -134,8 +128,7 @@ def _resolve_model_from_preferences( # Fall back to first available model if available_model_names: verbose_logger.debug( - "MCP sampling model resolution: no default configured, " - "falling back to first available model '%s'", + "MCP sampling model resolution: no default configured, falling back to first available model '%s'", available_model_names[0], ) return available_model_names[0] @@ -247,14 +240,9 @@ def _select_model_by_priority( best_name = None best_score = -1.0 for i, entry in enumerate(scored): - score = ( - cost_weight * cost_scores[i] - + speed_weight * speed_scores[i] - + intel_weight * intel_scores[i] - ) + score = cost_weight * cost_scores[i] + speed_weight * speed_scores[i] + intel_weight * intel_scores[i] verbose_logger.debug( - "MCP priority scoring: model=%s cost_score=%.3f speed_score=%.3f " - "intel_score=%.3f → weighted=%.3f", + "MCP priority scoring: model=%s cost_score=%.3f speed_score=%.3f intel_score=%.3f → weighted=%.3f", entry["name"], cost_scores[i], speed_scores[i], @@ -353,11 +341,7 @@ def _convert_single_content( tool_use_id = getattr(content, "toolUseId", "") nested_content = getattr(content, "content", []) if isinstance(nested_content, list): - text_parts = [ - getattr(c, "text", str(c)) - for c in nested_content - if getattr(c, "type", None) == "text" - ] + text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"] result_text = "\n".join(text_parts) if text_parts else "" else: result_text = str(nested_content) @@ -417,9 +401,7 @@ def _convert_mcp_messages_to_openai( # above (e.g. unexpected role, single non-list content). converted = _convert_mcp_content_to_openai(content) converted_parts = ( - converted - if isinstance(converted, list) - else ([converted] if isinstance(converted, dict) else []) + converted if isinstance(converted, list) else ([converted] if isinstance(converted, dict) else []) ) # Separate marker items from regular content parts @@ -488,9 +470,7 @@ def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]: "type": "function", "function": { "name": getattr(item, "name", ""), - "arguments": json.dumps( - getattr(item, "input", {}), default=str - ), + "arguments": json.dumps(getattr(item, "input", {}), default=str), }, } ) @@ -517,11 +497,7 @@ def _extract_tool_results(content: Any) -> List[Dict[str, Any]]: # Extract text from nested content nested_content = getattr(item, "content", []) if isinstance(nested_content, list): - text_parts = [ - getattr(c, "text", str(c)) - for c in nested_content - if getattr(c, "type", None) == "text" - ] + text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"] result_text = "\n".join(text_parts) if text_parts else "" else: result_text = str(nested_content) @@ -597,8 +573,7 @@ def _convert_openai_response_to_mcp_result( """ if not response.choices: verbose_logger.warning( - "MCP sampling: LLM returned empty choices list for model=%s " - "(possible content filter or provider error)", + "MCP sampling: LLM returned empty choices list for model=%s (possible content filter or provider error)", model_name, ) return ErrorData( @@ -661,9 +636,7 @@ def _convert_openai_response_to_mcp_result( ) -async def _check_model_access( - model: str, user_api_key_auth: Any -) -> Optional["ErrorData"]: +async def _check_model_access(model: str, user_api_key_auth: Any) -> Optional["ErrorData"]: """Enforce model-permission checks for MCP sampling requests. Runs the same authorization checks as ``/chat/completions``: @@ -681,9 +654,7 @@ async def _check_model_access( _user_role = getattr(user_api_key_auth, "user_role", None) _has_real_credential = bool(_api_key) or bool(_token) - _is_admin = ( - _user_role in ("proxy_admin", "proxy_admin_viewer") if _user_role else False - ) + _is_admin = _user_role in ("proxy_admin", "proxy_admin_viewer") if _user_role else False if not _has_real_credential and not _is_admin: verbose_logger.warning( @@ -760,9 +731,7 @@ async def _check_model_access( model=model, team_object=team_obj, llm_router=_llm_router, - team_model_aliases=getattr( - user_api_key_auth, "team_model_aliases", None - ), + team_model_aliases=getattr(user_api_key_auth, "team_model_aliases", None), ) if _user_id and _proxy_logging_obj: await _check_team_member_model_access( @@ -824,10 +793,7 @@ async def _check_model_access( ) return ErrorData( code=-1, - message=( - f"Model access denied: the API key is not authorized " - f"to use model '{model}'. {access_err}" - ), + message=(f"Model access denied: the API key is not authorized to use model '{model}'. {access_err}"), ) @@ -859,9 +825,7 @@ async def _run_budget_checks( ) import litellm except ImportError as import_err: - verbose_logger.warning( - "MCP sampling: budget check imports unavailable: %s", import_err - ) + verbose_logger.warning("MCP sampling: budget check imports unavailable: %s", import_err) return None # Can't enforce budgets without the modules _team_id = getattr(user_api_key_auth, "team_id", None) @@ -1102,9 +1066,7 @@ async def _build_completion_kwargs( from litellm.proxy.proxy_server import proxy_config completion_kwargs["user"] = getattr(user_api_key_auth, "user_id", None) - _dummy_request = _build_sampling_request( - raw_headers=raw_headers, client_ip=client_ip - ) + _dummy_request = _build_sampling_request(raw_headers=raw_headers, client_ip=client_ip) completion_kwargs = await add_litellm_data_to_request( data=completion_kwargs, request=_dummy_request, @@ -1236,9 +1198,7 @@ async def handle_sampling_create_message( user_api_key_auth=user_api_key_auth, ) - result = _convert_openai_response_to_mcp_result( - response=response, model_name=model - ) + result = _convert_openai_response_to_mcp_result(response=response, model_name=model) verbose_logger.info( "MCP sampling: completed successfully, model=%s, stopReason=%s", getattr(result, "model", "unknown"), diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index a9c4d2ece46..f24d5715e83 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -62,14 +62,10 @@ class SemanticMCPToolFilter: all_tools = [] for server_id, server in registry.items(): try: - tools = await global_mcp_server_manager.get_tools_for_server( - server_id - ) + tools = await global_mcp_server_manager.get_tools_for_server(server_id) all_tools.extend(tools) except Exception as e: - verbose_logger.warning( - f"Failed to fetch tools from server {server_id}: {e}" - ) + verbose_logger.warning(f"Failed to fetch tools from server {server_id}: {e}") continue if not all_tools: @@ -77,9 +73,7 @@ class SemanticMCPToolFilter: self.tool_router = None return - verbose_logger.info( - f"Fetched {len(all_tools)} tools from {len(registry)} MCP servers" - ) + verbose_logger.info(f"Fetched {len(all_tools)} tools from {len(registry)} MCP servers") self._build_router(all_tools) except Exception as e: @@ -180,9 +174,7 @@ class SemanticMCPToolFilter: # Router should be built on startup - if not, something went wrong if self.tool_router is None: - verbose_logger.warning( - "Router not initialized - was build_router_from_mcp_registry() called on startup?" - ) + verbose_logger.warning("Router not initialized - was build_router_from_mcp_registry() called on startup?") return available_tools # Run semantic filtering @@ -252,9 +244,7 @@ class SemanticMCPToolFilter: separator = client_name[-len(canonical) - 1] return separator in ("_", "-") - def _get_tools_by_names( - self, tool_names: List[str], available_tools: List[Any] - ) -> List[Any]: + def _get_tools_by_names(self, tool_names: List[str], available_tools: List[Any]) -> List[Any]: """ Get tools from available_tools by their names, preserving the semantic router's ordering. diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a78e43a0226..49a1fd99e3a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -112,9 +112,7 @@ def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: _byok_cred_cache.pop((user_id, server_id), None) -def _write_byok_cred_cache( - user_id: str, server_id: str, credential: Optional[str] -) -> None: +def _write_byok_cred_cache(user_id: str, server_id: str, credential: Optional[str]) -> None: """Write a credential value to the cache, evicting all entries if at capacity.""" if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE: _byok_cred_cache.clear() @@ -143,8 +141,8 @@ try: # Robust auth lookup keyed by session_object. _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() - active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = ( - contextvars.ContextVar("active_mcp_session", default=None) + active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = contextvars.ContextVar( + "active_mcp_session", default=None ) except ImportError as e: verbose_logger.debug(f"MCP module not found: {e}") @@ -441,10 +439,7 @@ if MCP_AVAILABLE: for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): if _stateful_session_active_request_counts.get(session_id, 0) > 0: continue - if ( - now - last_seen >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS - or session_id not in server_instances - ): + if now - last_seen >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS or session_id not in server_instances: expired_session_ids.append(session_id) for session_id in expired_session_ids: @@ -514,9 +509,7 @@ if MCP_AVAILABLE: try: await _purge_expired_stateful_session_auth_contexts() except Exception as e: - verbose_logger.exception( - f"Error cleaning up expired MCP stateful sessions: {e}" - ) + verbose_logger.exception(f"Error cleaning up expired MCP stateful sessions: {e}") async def initialize_session_managers(): """Initialize the session managers. Can be called from main app lifespan.""" @@ -543,14 +536,10 @@ if MCP_AVAILABLE: await _session_manager_cm.__aenter__() await _session_manager_stateful_cm.__aenter__() await _sse_session_manager_cm.__aenter__() - _stateful_auth_context_cleanup_task = asyncio.create_task( - _cleanup_expired_stateful_session_auth_contexts() - ) + _stateful_auth_context_cleanup_task = asyncio.create_task(_cleanup_expired_stateful_session_auth_contexts()) _SESSION_MANAGERS_INITIALIZED = True - verbose_logger.info( - "MCP Server started with StreamableHTTP and SSE session managers!" - ) + verbose_logger.info("MCP Server started with StreamableHTTP and SSE session managers!") async def shutdown_session_managers(): """Shutdown the session managers.""" @@ -621,12 +610,8 @@ if MCP_AVAILABLE: raw_headers, _client_ip, ) = await get_or_extract_auth_context() - verbose_logger.debug( - f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}" - ) - verbose_logger.debug( - f"MCP list_tools - MCP servers from context: {mcp_servers}" - ) + verbose_logger.debug(f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}") + verbose_logger.debug(f"MCP list_tools - MCP servers from context: {mcp_servers}") verbose_logger.debug( f"MCP list_tools - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) @@ -642,9 +627,7 @@ if MCP_AVAILABLE: log_list_tools_to_spendlogs=True, list_tools_log_source="mcp_protocol", ) - verbose_logger.info( - f"MCP list_tools - Successfully returned {len(tools)} tools" - ) + verbose_logger.info(f"MCP list_tools - Successfully returned {len(tools)} tools") return tools except Exception as e: verbose_logger.exception(f"Error in list_tools endpoint: {str(e)}") @@ -656,9 +639,7 @@ if MCP_AVAILABLE: active_mcp_session_var.reset(_session_reset_token) @server.call_tool() - async def mcp_server_tool_call( - name: str, arguments: Dict[str, Any] | None - ) -> CallToolResult: + async def mcp_server_tool_call(name: str, arguments: Dict[str, Any] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: @@ -696,9 +677,7 @@ if MCP_AVAILABLE: f"MCP mcp_server_tool_call - user_api_key_auth={user_api_key_auth}, user_role={getattr(user_api_key_auth, 'user_role', 'N/A')}" ) - verbose_logger.debug( - f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" - ) + verbose_logger.debug(f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}") host_progress_callback = None try: host_ctx = server.request_context @@ -707,9 +686,7 @@ if MCP_AVAILABLE: if host_token and hasattr(host_ctx, "session") and host_ctx.session: host_session = host_ctx.session - async def forward_progress( - progress: float, total: Optional[float] - ): + async def forward_progress(progress: float, total: Optional[float]): """Forward progress notifications from external MCP to Host""" try: await host_session.send_progress_notification( @@ -717,18 +694,12 @@ if MCP_AVAILABLE: progress=progress, total=total, ) - verbose_logger.debug( - f"Forwarded progress {progress}/{total} to Host" - ) + verbose_logger.debug(f"Forwarded progress {progress}/{total} to Host") except Exception as e: - verbose_logger.error( - f"Failed to forward progress to Host: {e}" - ) + verbose_logger.error(f"Failed to forward progress to Host: {e}") host_progress_callback = forward_progress - verbose_logger.debug( - f"Host progressToken captured: {host_token[:8]}..." - ) + verbose_logger.debug(f"Host progressToken captured: {host_token[:8]}...") except Exception as e: verbose_logger.warning(f"Could not capture host progress context: {e}") try: @@ -779,9 +750,7 @@ if MCP_AVAILABLE: isError=True, ) except BlockedPiiEntityError as e: - verbose_logger.error( - f"BlockedPiiEntityError in MCP tool call: {str(e)}" - ) + verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}") return CallToolResult( content=[ TextContent( @@ -792,15 +761,9 @@ if MCP_AVAILABLE: isError=True, ) except GuardrailRaisedException as e: - verbose_logger.error( - f"GuardrailRaisedException in MCP tool call: {str(e)}" - ) + verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}") return CallToolResult( - content=[ - TextContent( - text=f"Error: Guardrail violation - {str(e)}", type="text" - ) - ], + content=[TextContent(text=f"Error: Guardrail violation - {str(e)}", type="text")], isError=True, ) except HTTPException as e: @@ -844,12 +807,8 @@ if MCP_AVAILABLE: raw_headers, _client_ip, ) = await get_or_extract_auth_context() - verbose_logger.debug( - f"MCP list_prompts - User API Key Auth from context: {user_api_key_auth}" - ) - verbose_logger.debug( - f"MCP list_prompts - MCP servers from context: {mcp_servers}" - ) + verbose_logger.debug(f"MCP list_prompts - User API Key Auth from context: {user_api_key_auth}") + verbose_logger.debug(f"MCP list_prompts - MCP servers from context: {mcp_servers}") verbose_logger.debug( f"MCP list_prompts - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) @@ -863,9 +822,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - verbose_logger.info( - f"MCP list_prompts - Successfully returned {len(prompts)} prompts" - ) + verbose_logger.info(f"MCP list_prompts - Successfully returned {len(prompts)} prompts") return prompts except Exception as e: verbose_logger.exception(f"Error in list_prompts endpoint: {str(e)}") @@ -877,9 +834,7 @@ if MCP_AVAILABLE: active_mcp_session_var.reset(_session_reset_token) @server.get_prompt() - async def get_prompt( - name: str, arguments: Optional[Dict[str, str]] - ) -> GetPromptResult: + async def get_prompt(name: str, arguments: Optional[Dict[str, str]]) -> GetPromptResult: """ Get a specific prompt with the provided arguments @@ -910,9 +865,7 @@ if MCP_AVAILABLE: _client_ip, ) = await get_or_extract_auth_context() - verbose_logger.debug( - f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" - ) + verbose_logger.debug(f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}") return await mcp_get_prompt( name=name, arguments=arguments, @@ -947,12 +900,8 @@ if MCP_AVAILABLE: raw_headers, _client_ip, ) = await get_or_extract_auth_context() - verbose_logger.debug( - f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}" - ) - verbose_logger.debug( - f"MCP list_resources - MCP servers from context: {mcp_servers}" - ) + verbose_logger.debug(f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}") + verbose_logger.debug(f"MCP list_resources - MCP servers from context: {mcp_servers}") verbose_logger.debug( f"MCP list_resources - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) @@ -965,9 +914,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - verbose_logger.info( - f"MCP list_resources - Successfully returned {len(resources)} resources" - ) + verbose_logger.info(f"MCP list_resources - Successfully returned {len(resources)} resources") return resources except Exception as e: verbose_logger.exception(f"Error in list_resources endpoint: {str(e)}") @@ -996,12 +943,8 @@ if MCP_AVAILABLE: raw_headers, _client_ip, ) = await get_or_extract_auth_context() - verbose_logger.debug( - f"MCP list_resource_templates - User API Key Auth from context: {user_api_key_auth}" - ) - verbose_logger.debug( - f"MCP list_resource_templates - MCP servers from context: {mcp_servers}" - ) + verbose_logger.debug(f"MCP list_resource_templates - User API Key Auth from context: {user_api_key_auth}") + verbose_logger.debug(f"MCP list_resource_templates - MCP servers from context: {mcp_servers}") verbose_logger.debug( f"MCP list_resource_templates - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) @@ -1019,9 +962,7 @@ if MCP_AVAILABLE: ) return resource_templates except Exception as e: - verbose_logger.exception( - f"Error in list_resource_templates endpoint: {str(e)}" - ) + verbose_logger.exception(f"Error in list_resource_templates endpoint: {str(e)}") return [] finally: if _session_reset_token is not None: @@ -1093,9 +1034,7 @@ if MCP_AVAILABLE: for server in allowed_mcp_servers: if server: - match_list = [ - s.lower() for s in iter_known_server_prefixes(server) if s - ] + match_list = [s.lower() for s in iter_known_server_prefixes(server) if s] if server_or_group.lower() in match_list: filtered_server[server.server_id] = server @@ -1104,10 +1043,8 @@ if MCP_AVAILABLE: if not server_name_matched: try: - access_group_server_ids = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - [server_or_group] - ) + access_group_server_ids = await MCPRequestHandler._get_mcp_servers_from_access_groups( + [server_or_group] ) # Only include servers that the user has access to for server_id in access_group_server_ids: @@ -1115,9 +1052,7 @@ if MCP_AVAILABLE: if server_id == server.server_id: filtered_server[server.server_id] = server except Exception as e: - verbose_logger.debug( - f"Could not resolve '{server_or_group}' as access group: {e}" - ) + verbose_logger.debug(f"Could not resolve '{server_or_group}' as access group: {e}") if filtered_server: return list(filtered_server.values()) @@ -1127,8 +1062,7 @@ if MCP_AVAILABLE: # closed so URL/header namespacing cannot silently fall back to # the caller's full allowed-server set. verbose_logger.debug( - "MCP scope filter resolved to no servers for requested names %s; " - "returning empty list (fail-closed).", + "MCP scope filter resolved to no servers for requested names %s; returning empty list (fail-closed).", mcp_servers, ) return [] @@ -1192,18 +1126,12 @@ if MCP_AVAILABLE: if server_applies_tool_allowlist(mcp_server): if not mcp_server.allowed_tools: return [] - tools_to_return = [ - tool - for tool in tools - if _tool_name_matches(tool.name, mcp_server.allowed_tools) - ] + tools_to_return = [tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools)] # Filter by disallowed_tools (blacklist) if mcp_server.disallowed_tools: tools_to_return = [ - tool - for tool in tools_to_return - if not _tool_name_matches(tool.name, mcp_server.disallowed_tools) + tool for tool in tools_to_return if not _tool_name_matches(tool.name, mcp_server.disallowed_tools) ] return tools_to_return @@ -1269,15 +1197,11 @@ if MCP_AVAILABLE: "IP filtering will be skipped. This is expected for internal calls." ) - allowed_mcp_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) - ) + allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) ( allowed_mcp_server_ids, _ip_blocked, - ) = global_mcp_server_manager.filter_server_ids_by_ip_with_info( - allowed_mcp_server_ids, client_ip - ) + ) = global_mcp_server_manager.filter_server_ids_by_ip_with_info(allowed_mcp_server_ids, client_ip) verbose_logger.debug( "MCP IP filter: client_ip=%s, allowed_server_ids=%s", client_ip, @@ -1295,9 +1219,7 @@ if MCP_AVAILABLE: ) allowed_mcp_servers: List[MCPServer] = [] for allowed_mcp_server_id in allowed_mcp_server_ids: - mcp_server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_mcp_server_id - ) + mcp_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) if mcp_server is not None: # Apply oauth2_flow resolution for legacy DB rows where it may be NULL resolved_flow = MCPServerManager._resolve_oauth2_flow( @@ -1310,9 +1232,7 @@ if MCP_AVAILABLE: ) if resolved_flow and resolved_flow != mcp_server.oauth2_flow: # Create a new instance with the resolved flow for this request - mcp_server = mcp_server.model_copy( - update={"oauth2_flow": resolved_flow} - ) + mcp_server = mcp_server.model_copy(update={"oauth2_flow": resolved_flow}) allowed_mcp_servers.append(mcp_server) if mcp_servers is not None: @@ -1387,9 +1307,7 @@ if MCP_AVAILABLE: Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ - user_id = ( - getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None - ) + user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None if not user_id: return {} try: @@ -1404,9 +1322,7 @@ if MCP_AVAILABLE: creds = await list_user_oauth_credentials(prisma_client, user_id) return {c["server_id"]: c for c in creds if "server_id" in c} except Exception as e: - verbose_logger.warning( - f"_prefetch_oauth_creds_for_user: failed to prefetch for user={user_id}: {e}" - ) + verbose_logger.warning(f"_prefetch_oauth_creds_for_user: failed to prefetch for user={user_id}: {e}") return {} def _prepare_mcp_server_headers( @@ -1454,9 +1370,7 @@ if MCP_AVAILABLE: if extra_headers is None: extra_headers = {} - normalized_raw_headers = { - str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str) - } + normalized_raw_headers = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} # Centralized strip decision shared with # ``MCPServerManager._call_regular_mcp_tool`` so the two @@ -1496,21 +1410,13 @@ if MCP_AVAILABLE: texts: List[Tuple[str, str]] = [] for server in allowed_mcp_servers: - label = ( - server.alias - or server.server_name - or server.name - or server.server_id - or "mcp" - ) + label = server.alias or server.server_name or server.name or server.server_id or "mcp" if server.instructions and server.instructions.strip(): texts.append((label, server.instructions.strip())) continue if server.spec_path: continue - cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get( - server.server_id - ) + cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get(server.server_id) if cached and cached.strip(): texts.append((label, cached.strip())) @@ -1538,9 +1444,7 @@ if MCP_AVAILABLE: # cancel sibling probes or 500 the gateway initialize request. await asyncio.gather( *[ - global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( - s - ) + global_mcp_server_manager._ensure_upstream_initialize_instructions_cached(s) for s in allowed if s is not None ], @@ -1551,10 +1455,7 @@ if MCP_AVAILABLE: if scoped_server_endpoint and len(allowed) == 1: scoped_server = allowed[0] scoped_server_name = ( - scoped_server.alias - or scoped_server.server_name - or scoped_server.name - or scoped_server.server_id + scoped_server.alias or scoped_server.server_name or scoped_server.name or scoped_server.server_id ) instructions_token = _mcp_gateway_initialize_instructions.set(merged) server_name_token = _mcp_gateway_server_name.set(scoped_server_name) @@ -1600,9 +1501,7 @@ if MCP_AVAILABLE: rules_obj = Rules() list_tools_call_id = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) - effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers( - raw_headers - ) + effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers(raw_headers) spend_logs_metadata: Dict[str, Any] = { "mcp_operation": "list_tools", } @@ -1639,9 +1538,9 @@ if MCP_AVAILABLE: _metadata_variable_name="metadata", ) - user_identifier = getattr( - user_api_key_auth, "end_user_id", None - ) or getattr(user_api_key_auth, "user_id", None) + user_identifier = getattr(user_api_key_auth, "end_user_id", None) or getattr( + user_api_key_auth, "user_id", None + ) if user_identifier: list_tools_request_data["user"] = user_identifier @@ -1656,9 +1555,7 @@ if MCP_AVAILABLE: litellm_logging_obj.call_type = CallTypes.list_mcp_tools.value litellm_logging_obj.model = "MCP: list_tools" except Exception as logging_error: - verbose_logger.debug( - "Failed to initialize logging for MCP list_tools: %s", logging_error - ) + verbose_logger.debug("Failed to initialize logging for MCP list_tools: %s", logging_error) litellm_logging_obj = None try: @@ -1669,14 +1566,9 @@ if MCP_AVAILABLE: # Pre-fetch OAuth credentials only when at least one server uses OAuth2, # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. - _has_oauth2_server = any( - getattr(s, "auth_type", None) == MCPAuth.oauth2 - for s in allowed_mcp_servers - ) + _has_oauth2_server = any(getattr(s, "auth_type", None) == MCPAuth.oauth2 for s in allowed_mcp_servers) _prefetched_oauth_creds = ( - await _prefetch_oauth_creds_for_user(user_api_key_auth) - if _has_oauth2_server - else {} + await _prefetch_oauth_creds_for_user(user_api_key_auth) if _has_oauth2_server else {} ) async def _fetch_and_filter_server_tools( @@ -1721,11 +1613,7 @@ if MCP_AVAILABLE: extra_headers = db_headers # If still no OAuth2 token, fall back to pre-fetched creds (non-stale-client path) - elif ( - not migrated_to_v2 - and extra_headers is None - and server.auth_type == MCPAuth.oauth2 - ): + elif not migrated_to_v2 and extra_headers is None and server.auth_type == MCPAuth.oauth2: extra_headers = await _get_user_oauth_extra_headers_from_db( server, user_api_key_auth, @@ -1765,15 +1653,11 @@ if MCP_AVAILABLE: # swallow the auth error. raise except Exception as e: - verbose_logger.exception( - f"Error getting tools from server {server.name}: {str(e)}" - ) + verbose_logger.exception(f"Error getting tools from server {server.name}: {str(e)}") return [] # Fetch tools from all servers in parallel - tasks = [ - _fetch_and_filter_server_tools(server) for server in allowed_mcp_servers - ] + tasks = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers] results = await asyncio.gather(*tasks) # Flatten results into single list @@ -1818,9 +1702,7 @@ if MCP_AVAILABLE: log_exc, ) - verbose_logger.info( - f"Successfully fetched {len(all_tools)} tools total from all MCP servers" - ) + verbose_logger.info(f"Successfully fetched {len(all_tools)} tools total from all MCP servers") return all_tools except Exception as e: @@ -1830,9 +1712,7 @@ if MCP_AVAILABLE: from litellm.proxy.proxy_server import proxy_logging_obj if proxy_logging_obj: - traceback_str = traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG - ) + traceback_str = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) await proxy_logging_obj.post_call_failure_hook( request_data=list_tools_request_data or {}, original_exception=e, @@ -1841,9 +1721,7 @@ if MCP_AVAILABLE: traceback_str=traceback_str, ) except Exception: - verbose_logger.debug( - "Failed to log MCP list_tools failure via post_call_failure_hook" - ) + verbose_logger.debug("Failed to log MCP list_tools failure via post_call_failure_hook") raise async def _get_prompts_from_mcp_servers( @@ -1901,18 +1779,12 @@ if MCP_AVAILABLE: all_prompts.extend(prompts) - verbose_logger.debug( - f"Successfully fetched {len(prompts)} prompts from server {server.name}" - ) + verbose_logger.debug(f"Successfully fetched {len(prompts)} prompts from server {server.name}") except Exception as e: - verbose_logger.exception( - f"Error getting prompts from server {server.name}: {str(e)}" - ) + verbose_logger.exception(f"Error getting prompts from server {server.name}: {str(e)}") # Continue with other servers instead of failing completely - verbose_logger.info( - f"Successfully fetched {len(all_prompts)} prompts total from all MCP servers" - ) + verbose_logger.info(f"Successfully fetched {len(all_prompts)} prompts total from all MCP servers") return all_prompts @@ -1958,17 +1830,11 @@ if MCP_AVAILABLE: ) all_resources.extend(resources) - verbose_logger.debug( - f"Successfully fetched {len(resources)} resources from server {server.name}" - ) + verbose_logger.debug(f"Successfully fetched {len(resources)} resources from server {server.name}") except Exception as e: - verbose_logger.exception( - f"Error getting resources from server {server.name}: {str(e)}" - ) + verbose_logger.exception(f"Error getting resources from server {server.name}: {str(e)}") - verbose_logger.info( - f"Successfully fetched {len(all_resources)} resources total from all MCP servers" - ) + verbose_logger.info(f"Successfully fetched {len(all_resources)} resources total from all MCP servers") return all_resources @@ -2005,14 +1871,12 @@ if MCP_AVAILABLE: ) try: - resource_templates = ( - await global_mcp_server_manager.get_resource_templates_from_server( - server=server, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - add_prefix=True, # Always add server prefix - raw_headers=raw_headers, - ) + resource_templates = await global_mcp_server_manager.get_resource_templates_from_server( + server=server, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + add_prefix=True, # Always add server prefix + raw_headers=raw_headers, ) all_resource_templates.extend(resource_templates) verbose_logger.debug( @@ -2058,11 +1922,7 @@ if MCP_AVAILABLE: # prefix (resolved from the server) rather than the first separator, so a # prefix containing the separator still reduces to the stored bare name. server = global_mcp_server_manager.get_mcp_server_by_id(server_id) - return [ - t - for t in tools - if strip_known_server_prefix(t.name, server) in allowed_tool_names - ] + return [t for t in tools if strip_known_server_prefix(t.name, server) in allowed_tool_names] async def _merge_toolset_permissions( user_api_key_auth: Optional[UserAPIKeyAuth], @@ -2082,11 +1942,7 @@ if MCP_AVAILABLE: if not toolset_ids: return user_api_key_auth - toolset_perms = ( - await global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=toolset_ids - ) - ) + toolset_perms = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=toolset_ids) if not toolset_perms: return user_api_key_auth @@ -2102,9 +1958,7 @@ if MCP_AVAILABLE: # filtering doesn't silently drop servers that the toolset references but that # aren't already in the key's explicit mcp_servers list. merged_servers = list(set(op.mcp_servers or []) | set(existing.keys())) - updated_op = op.model_copy( - update={"mcp_servers": merged_servers, "mcp_tool_permissions": existing} - ) + updated_op = op.model_copy(update={"mcp_servers": merged_servers, "mcp_tool_permissions": existing}) return user_api_key_auth.model_copy(update={"object_permission": updated_op}) async def _list_mcp_tools( @@ -2149,13 +2003,9 @@ if MCP_AVAILABLE: log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, list_tools_log_source=list_tools_log_source, ) - verbose_logger.debug( - f"Successfully fetched {len(managed_tools)} tools from managed MCP servers" - ) + verbose_logger.debug(f"Successfully fetched {len(managed_tools)} tools from managed MCP servers") except Exception as e: - verbose_logger.exception( - f"Error getting tools from managed MCP servers: {str(e)}" - ) + verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}") # Continue with empty managed tools list instead of failing completely return managed_tools @@ -2193,13 +2043,9 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - verbose_logger.debug( - f"Successfully fetched {len(managed_prompts)} prompts from managed MCP servers" - ) + verbose_logger.debug(f"Successfully fetched {len(managed_prompts)} prompts from managed MCP servers") except Exception as e: - verbose_logger.exception( - f"Error getting tools from managed MCP servers: {str(e)}" - ) + verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}") # Continue with empty managed tools list instead of failing completely return managed_prompts @@ -2227,13 +2073,9 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - verbose_logger.debug( - f"Successfully fetched {len(managed_resources)} resources from managed MCP servers" - ) + verbose_logger.debug(f"Successfully fetched {len(managed_resources)} resources from managed MCP servers") except Exception as e: - verbose_logger.exception( - f"Error getting resources from managed MCP servers: {str(e)}" - ) + verbose_logger.exception(f"Error getting resources from managed MCP servers: {str(e)}") return managed_resources @@ -2287,9 +2129,7 @@ if MCP_AVAILABLE: display_map = server.tool_name_to_display_name or {} for unprefixed_name, display_name in display_map.items(): if display_name == name: - return add_server_prefix_to_name( - unprefixed_name, get_server_prefix(server) - ) + return add_server_prefix_to_name(unprefixed_name, get_server_prefix(server)) return name async def _get_byok_credential( @@ -2350,9 +2190,7 @@ if MCP_AVAILABLE: "server_name": mcp_server.server_name or mcp_server.name, "message": "User identity is required for BYOK servers", }, - headers={ - "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' - }, + headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, ) # Check shared credential cache before hitting the DB. @@ -2414,9 +2252,7 @@ if MCP_AVAILABLE: "Complete the OAuth authorization flow to provide your API key." ), }, - headers={ - "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' - }, + headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, ) async def execute_mcp_tool( @@ -2476,9 +2312,7 @@ if MCP_AVAILABLE: for registry_server in global_mcp_server_manager.get_registry().values(): for known_prefix in iter_known_server_prefixes(registry_server): all_registry_prefixes.add(normalize_server_name(known_prefix)) - name_is_prefixed = is_tool_name_prefixed( - name, known_server_prefixes=all_registry_prefixes - ) + name_is_prefixed = is_tool_name_prefixed(name, known_server_prefixes=all_registry_prefixes) if requested_server is not None and not name_is_prefixed: # REST callers may pass server_id with the upstream tool name (no @@ -2494,10 +2328,8 @@ if MCP_AVAILABLE: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) if mcp_server is None and requested_server is not None: for known_prefix in iter_known_server_prefixes(requested_server): - candidate = ( - global_mcp_server_manager._get_mcp_server_from_tool_name( - add_server_prefix_to_name(name, known_prefix) - ) + candidate = global_mcp_server_manager._get_mcp_server_from_tool_name( + add_server_prefix_to_name(name, known_prefix) ) if candidate is not None: mcp_server = candidate @@ -2506,10 +2338,7 @@ if MCP_AVAILABLE: server_name = mcp_server.name if requested_server is not None: - if ( - mcp_server is not None - and mcp_server.server_id != requested_server.server_id - ): + if mcp_server is not None and mcp_server.server_id != requested_server.server_id: raise HTTPException( status_code=403, detail={ @@ -2536,21 +2365,15 @@ if MCP_AVAILABLE: detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}", ) - standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = ( - _get_standard_logging_mcp_tool_call( - name=original_tool_name, # Use original name for logging - arguments=arguments, - server_name=server_name, - session_id=_mcp_session_id_from_headers(raw_headers), - ) - ) - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj", None + standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = _get_standard_logging_mcp_tool_call( + name=original_tool_name, # Use original name for logging + arguments=arguments, + server_name=server_name, + session_id=_mcp_session_id_from_headers(raw_headers), ) + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) if litellm_logging_obj: - litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( - standard_logging_mcp_tool_call - ) + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call litellm_logging_obj.model = f"MCP: {name}" litellm_logging_obj.model_call_details["model"] = f"MCP: {name}" # Resolve the MCP server early so BYOK checks and credential injection @@ -2559,13 +2382,11 @@ if MCP_AVAILABLE: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) if mcp_server: - standard_logging_mcp_tool_call["mcp_server_cost_info"] = ( - mcp_server.mcp_info or {} - ).get("mcp_server_cost_info") + standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get( + "mcp_server_cost_info" + ) if litellm_logging_obj: - litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( - standard_logging_mcp_tool_call - ) + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call # BYOK: retrieve the stored per-user credential. A single DB call # both checks existence and fetches the value, avoiding a double query. @@ -2645,9 +2466,7 @@ if MCP_AVAILABLE: # configured auth_type so the generator doesn't need to know the prefix. auth_header_value: Optional[str] = None if mcp_auth_header: - server_auth_type = ( - getattr(mcp_server, "auth_type", None) if mcp_server else None - ) + server_auth_type = getattr(mcp_server, "auth_type", None) if mcp_server else None if server_auth_type == MCPAuth.api_key: auth_header_value = f"ApiKey {mcp_auth_header}" elif server_auth_type == MCPAuth.basic: @@ -2661,19 +2480,12 @@ if MCP_AVAILABLE: # _prepare_mcp_server_headers for managed MCP). forwarded_headers: Optional[Dict[str, str]] = None if mcp_server and mcp_server.extra_headers and raw_headers: - normalized_raw = { - str(k).lower(): v - for k, v in raw_headers.items() - if isinstance(k, str) - } + normalized_raw = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} skip_caller_authorization = bool(mcp_server.has_client_credentials) for header_name in mcp_server.extra_headers: if not isinstance(header_name, str): continue - if ( - skip_caller_authorization - and header_name.lower() == "authorization" - ): + if skip_caller_authorization and header_name.lower() == "authorization": continue value = normalized_raw.get(header_name.lower()) if value is not None: @@ -2753,28 +2565,20 @@ if MCP_AVAILABLE: Call a specific tool with the provided arguments (handles prefixed tool names). """ start_time = datetime.now() - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj", None - ) + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) try: if arguments is None: - raise HTTPException( - status_code=400, detail="Request arguments are required" - ) + raise HTTPException(status_code=400, detail="Request arguments are required") ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL - allowed_mcp_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - ) + allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, ) allowed_mcp_servers: List[MCPServer] = [] for allowed_mcp_server_id in allowed_mcp_server_ids: - allowed_server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_mcp_server_id - ) + allowed_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) if allowed_server is not None: allowed_mcp_servers.append(allowed_server) @@ -3056,21 +2860,15 @@ if MCP_AVAILABLE: # Path found at the end, remove it from servers path_part = "/" + path_match.group(1) servers_part = servers_and_path[: -len(path_part)] - mcp_servers_from_path = [ - s.strip() for s in servers_part.split(",") if s.strip() - ] + mcp_servers_from_path = [s.strip() for s in servers_part.split(",") if s.strip()] else: # No path, just comma-separated servers - mcp_servers_from_path = [ - s.strip() for s in servers_and_path.split(",") if s.strip() - ] + mcp_servers_from_path = [s.strip() for s in servers_and_path.split(",") if s.strip()] else: # Single server case - use regex approach for server/path separation # This handles cases like "custom_solutions/user_123/chat/completions" # where we want to extract "custom_solutions/user_123" as the server name - single_server_match = re.match( - r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path - ) + single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path) if single_server_match: server_name = single_server_match.group(1) mcp_servers_from_path = [server_name] @@ -3118,15 +2916,9 @@ if MCP_AVAILABLE: Returns None if not present. """ for header_name, header_value in scope.get("headers", []): - name = ( - header_name if isinstance(header_name, bytes) else header_name.encode() - ) + name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": - return ( - header_value.decode() - if isinstance(header_value, bytes) - else str(header_value) - ) + return header_value.decode() if isinstance(header_value, bytes) else str(header_value) return None def _owner_fingerprint_for( @@ -3177,9 +2969,7 @@ if MCP_AVAILABLE: user_id_hash = hashlib.sha256(uid_material).hexdigest() return f"user:{user_id_hash}" if oauth2_headers: - authz = oauth2_headers.get("Authorization") or oauth2_headers.get( - "authorization" - ) + authz = oauth2_headers.get("Authorization") or oauth2_headers.get("authorization") authz_bytes = _bytes_for_hash(authz) if authz_bytes: return f"oauth:{hashlib.sha256(authz_bytes).hexdigest()}" @@ -3330,11 +3120,7 @@ if MCP_AVAILABLE: "Stripping stale header to force new session creation.", _session_id, ) - scope["headers"] = [ - (k, v) - for k, v in _headers - if _normalize_header_name(k) != _mcp_session_header - ] + scope["headers"] = [(k, v) for k, v in _headers if _normalize_header_name(k) != _mcp_session_header] return False async def _apply_toolset_scope( @@ -3359,9 +3145,7 @@ if MCP_AVAILABLE: # drop the sentinel. Checked before the admin branch, mirroring # get_allowed_mcp_servers. original_op = user_api_key_auth.object_permission - if original_op is not None and SpecialMCPServerNames.no_mcp_servers.value in ( - original_op.mcp_servers or [] - ): + if original_op is not None and SpecialMCPServerNames.no_mcp_servers.value in (original_op.mcp_servers or []): raise HTTPException( status_code=403, detail="API key is scoped to no MCP servers; toolset access is denied.", @@ -3382,11 +3166,7 @@ if MCP_AVAILABLE: detail=f"API key does not have access to toolset '{toolset_id}'.", ) - tool_permissions = ( - await global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=[toolset_id] - ) - ) + tool_permissions = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=[toolset_id]) server_ids = list(tool_permissions.keys()) existing_op = user_api_key_auth.object_permission if existing_op is not None: @@ -3452,14 +3232,8 @@ if MCP_AVAILABLE: a server it will be 403'd on immediately after authentication. """ for server_name in mcp_servers or []: - server = global_mcp_server_manager.get_mcp_server_by_name( - server_name, client_ip=client_ip - ) - if ( - server is not None - and allowed_server_ids is not None - and server.server_id not in allowed_server_ids - ): + server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) + if server is not None and allowed_server_ids is not None and server.server_id not in allowed_server_ids: # Caller's narrowed scope excludes this server — skip the # preemptive challenge and let downstream authorization # return 403. @@ -3486,9 +3260,7 @@ if MCP_AVAILABLE: ) # The v2 resolver owns the existence check, so every authorization_code # resolution (egress and this discovery challenge) runs through it. - if await global_mcp_server_manager.has_user_oauth_token( - server, user_api_key_auth - ): + if await global_mcp_server_manager.has_user_oauth_token(server, user_api_key_auth): continue request = StarletteRequest(scope) @@ -3518,9 +3290,7 @@ if MCP_AVAILABLE: if ( server and server.is_oauth_passthrough - and not _client_has_passthrough_authorization( - server, oauth2_headers, mcp_server_auth_headers - ) + and not _client_has_passthrough_authorization(server, oauth2_headers, mcp_server_auth_headers) ): www_authenticate = _get_passthrough_www_authenticate( scope=scope, @@ -3604,13 +3374,9 @@ if MCP_AVAILABLE: # AsyncHTTPHandler.post() calls raise_for_status(); a 401/403 from # upstream lands here. Return its status so the caller can map it # to the appropriate response. - return exc.response.status_code, exc.response.headers.get( - "www-authenticate" - ) + return exc.response.status_code, exc.response.headers.get("www-authenticate") except Exception as exc: - verbose_logger.debug( - f"_probe_upstream_auth: probe to {url} failed ({exc}), allowing request through" - ) + verbose_logger.debug(f"_probe_upstream_auth: probe to {url} failed ({exc}), allowing request through") return 200, None async def _check_passthrough_upstream_auth( @@ -3658,10 +3424,7 @@ if MCP_AVAILABLE: return probe_results = await asyncio.gather( - *[ - _probe_upstream_auth(srv.url or "", forwarded_auth) - for srv in passthrough_servers - ] + *[_probe_upstream_auth(srv.url or "", forwarded_auth) for srv in passthrough_servers] ) for srv, (probe_status, _) in zip(passthrough_servers, probe_results): if probe_status == 401: @@ -3687,9 +3450,7 @@ if MCP_AVAILABLE: detail="Forbidden", ) - async def handle_streamable_http_mcp( - scope: Scope, receive: Receive, send: Send - ) -> None: + async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: path = scope.get("path", "") @@ -3706,28 +3467,20 @@ if MCP_AVAILABLE: # Extract client IP for MCP access control _client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) - verbose_logger.debug( - f"MCP request mcp_servers (header/path): {mcp_servers}" - ) + verbose_logger.debug(f"MCP request mcp_servers (header/path): {mcp_servers}") verbose_logger.debug( f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [ - (k, v) - for k, v in scope.get("headers", []) - if k.lower() != b"x-mcp-toolset-id" - ] + scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar (set by # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). active_toolset_id = _mcp_active_toolset_id.get() toolset_allowed_server_ids: Optional[Set[str]] = None if active_toolset_id and user_api_key_auth is not None: - user_api_key_auth = await _apply_toolset_scope( - user_api_key_auth, active_toolset_id - ) + user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id) op = user_api_key_auth.object_permission toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set() @@ -3749,9 +3502,7 @@ if MCP_AVAILABLE: # Pre-flight auth check for pass-through servers. Must run after # toolset scoping so the probe list is derived from the fully-authorized # server set, not the raw user-supplied names. - await _check_passthrough_upstream_auth( - scope, user_api_key_auth, mcp_servers, _client_ip - ) + await _check_passthrough_upstream_auth(scope, user_api_key_auth, mcp_servers, _client_ip) # Inject masked debug headers when client sends x-litellm-mcp-debug: true _debug_headers = MCPDebug.maybe_build_debug_headers( @@ -3790,9 +3541,7 @@ if MCP_AVAILABLE: # response sees a pristine ``receive`` channel. if session_id: expected_owner = _stateful_session_owners.get(session_id) - request_owner = _owner_fingerprint_for( - user_api_key_auth, oauth2_headers, _client_ip - ) + request_owner = _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip) if expected_owner is not None and expected_owner != request_owner: verbose_logger.warning( "Rejecting MCP request: session '%s' owner mismatch.", @@ -3812,9 +3561,7 @@ if MCP_AVAILABLE: # non-DELETE requests have their session header stripped and should # be routed as no-session requests. if session_id: - handled = await _handle_stale_mcp_session( - scope, receive, send, session_manager_stateful - ) + handled = await _handle_stale_mcp_session(scope, receive, send, session_manager_stateful) if handled: # Request was fully handled (e.g., DELETE on non-existent session) return @@ -3826,9 +3573,7 @@ if MCP_AVAILABLE: is_initialize = _is_initialize_request(body) use_stateful = bool(session_id or is_initialize) - target_manager = ( - session_manager_stateful if use_stateful else session_manager_stateless - ) + target_manager = session_manager_stateful if use_stateful else session_manager_stateless verbose_logger.debug( f"MCP routing to {'stateful' if use_stateful else 'stateless'} manager" @@ -3840,9 +3585,7 @@ if MCP_AVAILABLE: # session. Cap how many a single caller can hold so an authenticated # client cannot spam `initialize` and exhaust memory. if is_initialize and not session_id: - request_owner = _owner_fingerprint_for( - user_api_key_auth, oauth2_headers, _client_ip - ) + request_owner = _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip) if not await _enforce_stateful_session_cap_for_owner(request_owner): verbose_logger.warning( "Rejecting MCP initialize: caller already holds the maximum number of active stateful sessions." @@ -3921,15 +3664,8 @@ if MCP_AVAILABLE: ) session_lock: Optional[asyncio.Lock] = None - if ( - use_stateful - and session_id - and request_method in ("POST", "DELETE") - and not is_jsonrpc_response - ): - session_lock = _stateful_session_locks.setdefault( - session_id, asyncio.Lock() - ) + if use_stateful and session_id and request_method in ("POST", "DELETE") and not is_jsonrpc_response: + session_lock = _stateful_session_locks.setdefault(session_id, asyncio.Lock()) active_request_session_ids: List[str] = [] @@ -3938,8 +3674,7 @@ if MCP_AVAILABLE: return active_request_session_ids.append(session_id_to_track) _stateful_session_active_request_counts[session_id_to_track] = ( - _stateful_session_active_request_counts.get(session_id_to_track, 0) - + 1 + _stateful_session_active_request_counts.get(session_id_to_track, 0) + 1 ) if use_stateful and session_id: @@ -3968,9 +3703,7 @@ if MCP_AVAILABLE: local_send = _wrap_send_with_stateful_session_auth_context( local_send, auth_user, - _owner_fingerprint_for( - user_api_key_auth, oauth2_headers, _client_ip - ), + _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip), _track_initialized_stateful_session, ) @@ -3992,36 +3725,18 @@ if MCP_AVAILABLE: await _dispatch() finally: for active_request_session_id in active_request_session_ids: - active_request_count = ( - _stateful_session_active_request_counts.get( - active_request_session_id, 0 - ) - - 1 - ) + active_request_count = _stateful_session_active_request_counts.get(active_request_session_id, 0) - 1 if active_request_count > 0: - _stateful_session_active_request_counts[ - active_request_session_id - ] = active_request_count + _stateful_session_active_request_counts[active_request_session_id] = active_request_count else: - _stateful_session_active_request_counts.pop( - active_request_session_id, None - ) + _stateful_session_active_request_counts.pop(active_request_session_id, None) - if ( - scope.get("method") != "DELETE" - and active_request_session_id in _stateful_session_auth_contexts - ): - _stateful_session_auth_context_last_seen[ - active_request_session_id - ] = time.monotonic() + if scope.get("method") != "DELETE" and active_request_session_id in _stateful_session_auth_contexts: + _stateful_session_auth_context_last_seen[active_request_session_id] = time.monotonic() # Periodic cleanup iterates _stateful_session_auth_context_last_seen, # so locks for untracked sessions must be dropped here. - if ( - active_request_count <= 0 - and active_request_session_id - not in _stateful_session_auth_contexts - ): + if active_request_count <= 0 and active_request_session_id not in _stateful_session_auth_contexts: _stateful_session_locks.pop(active_request_session_id, None) except MCPUpstreamAuthError as e: # Upstream delegated auth returned 401; surface it to the client so @@ -4051,9 +3766,7 @@ if MCP_AVAILABLE: ) await error_response(scope, receive, send) except Exception as response_error: - verbose_logger.exception( - f"Failed to send error response: {response_error}" - ) + verbose_logger.exception(f"Failed to send error response: {response_error}") # If we can't send a proper response, re-raise the original error raise e @@ -4074,19 +3787,13 @@ if MCP_AVAILABLE: # Extract client IP for MCP access control _sse_client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) - verbose_logger.debug( - f"MCP request mcp_servers (header/path): {mcp_servers}" - ) + verbose_logger.debug(f"MCP request mcp_servers (header/path): {mcp_servers}") verbose_logger.debug( f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [ - (k, v) - for k, v in scope.get("headers", []) - if k.lower() != b"x-mcp-toolset-id" - ] + scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar so the # downstream probe list matches the fully-authorized server set @@ -4094,9 +3801,7 @@ if MCP_AVAILABLE: active_toolset_id = _mcp_active_toolset_id.get() toolset_allowed_server_ids: Optional[Set[str]] = None if active_toolset_id and user_api_key_auth is not None: - user_api_key_auth = await _apply_toolset_scope( - user_api_key_auth, active_toolset_id - ) + user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id) op = user_api_key_auth.object_permission toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set() @@ -4121,9 +3826,7 @@ if MCP_AVAILABLE: # being stuck with a silently empty tool list. Must run after # toolset scoping so the probe list is derived from the fully- # authorized server set, not the raw user-supplied names. - await _check_passthrough_upstream_auth( - scope, user_api_key_auth, mcp_servers, _sse_client_ip - ) + await _check_passthrough_upstream_auth(scope, user_api_key_auth, mcp_servers, _sse_client_ip) set_auth_context( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -4176,9 +3879,7 @@ if MCP_AVAILABLE: ) await error_response(scope, receive, send) except Exception as response_error: - verbose_logger.exception( - f"Failed to send error response: {response_error}" - ) + verbose_logger.exception(f"Failed to send error response: {response_error}") # If we can't send a proper response, re-raise the original error raise e @@ -4272,9 +3973,7 @@ if MCP_AVAILABLE: touch_last_seen: bool = True, copy_existing_session_auth_context: bool = False, ) -> MCPAuthenticatedUser: - auth_user = ( - _stateful_session_auth_contexts.get(session_id) if session_id else None - ) + auth_user = _stateful_session_auth_contexts.get(session_id) if session_id else None if auth_user is not None and session_id is not None: if touch_last_seen: _stateful_session_auth_context_last_seen[session_id] = time.monotonic() @@ -4321,16 +4020,12 @@ if MCP_AVAILABLE: for key, value in message.get("headers", []): header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": - session_id = ( - value.decode() if isinstance(value, bytes) else str(value) - ) + session_id = value.decode() if isinstance(value, bytes) else str(value) if on_session_registered is not None: on_session_registered(session_id) auth_context_var.set(auth_user) _stateful_session_auth_contexts[session_id] = auth_user - _stateful_session_auth_context_last_seen[session_id] = ( - time.monotonic() - ) + _stateful_session_auth_context_last_seen[session_id] = time.monotonic() _stateful_session_owners[session_id] = owner_fingerprint break await send(message) diff --git a/litellm/proxy/_experimental/mcp_server/sse_transport.py b/litellm/proxy/_experimental/mcp_server/sse_transport.py index 63ffd403c66..0a896328dde 100644 --- a/litellm/proxy/_experimental/mcp_server/sse_transport.py +++ b/litellm/proxy/_experimental/mcp_server/sse_transport.py @@ -35,9 +35,7 @@ class SseServerTransport: """ _endpoint: str - _read_stream_writers: dict[ - UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception] - ] + _read_stream_writers: dict[UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]] def __init__(self, endpoint: str) -> None: """ @@ -48,9 +46,7 @@ class SseServerTransport: super().__init__() self._endpoint = endpoint self._read_stream_writers = {} - verbose_logger.debug( - f"SseServerTransport initialized with endpoint: {endpoint}" - ) + verbose_logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}") @asynccontextmanager async def connect_sse(self, request: Request): @@ -75,9 +71,7 @@ class SseServerTransport: sse_stream_writer: MemoryObjectSendStream[dict[str, Any]] sse_stream_reader: MemoryObjectReceiveStream[dict[str, Any]] - sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream( - 0, dict[str, Any] - ) + sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0, dict[str, Any]) async def sse_writer(): verbose_logger.debug("Starting SSE writer") @@ -90,25 +84,19 @@ class SseServerTransport: await sse_stream_writer.send( { "event": "message", - "data": message.model_dump_json( - by_alias=True, exclude_none=True - ), + "data": message.model_dump_json(by_alias=True, exclude_none=True), } ) async with anyio.create_task_group() as tg: - response = EventSourceResponse( - content=sse_stream_reader, data_sender_callable=sse_writer - ) + response = EventSourceResponse(content=sse_stream_reader, data_sender_callable=sse_writer) verbose_logger.debug("Starting SSE response task") tg.start_soon(response, request.scope, request.receive, request._send) verbose_logger.debug("Yielding read and write streams") yield (read_stream, write_stream) - async def handle_post_message( - self, scope: Scope, receive: Receive, send: Send - ) -> Response: + async def handle_post_message(self, scope: Scope, receive: Receive, send: Send) -> Response: verbose_logger.debug("Handling POST message") request = Request(scope, receive) diff --git a/litellm/proxy/_experimental/mcp_server/tool_registry.py b/litellm/proxy/_experimental/mcp_server/tool_registry.py index bb30ff55c5c..2da22671c91 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_registry.py +++ b/litellm/proxy/_experimental/mcp_server/tool_registry.py @@ -52,11 +52,7 @@ class MCPToolRegistry: List all registered tools """ if tool_prefix: - return [ - tool - for tool in self.tools.values() - if tool.name.startswith(tool_prefix) - ] + return [tool for tool in self.tools.values() if tool.name.startswith(tool_prefix)] return list(self.tools.values()) def unregister_tools_with_prefix(self, prefix: str) -> int: @@ -75,13 +71,9 @@ class MCPToolRegistry: verbose_logger.debug("Unregistered MCP tool %s", name) return removed - def convert_tools_to_mcp_sdk_tool_type( - self, tools: List[MCPTool] - ) -> List["MCPToolSDKTool"]: + def convert_tools_to_mcp_sdk_tool_type(self, tools: List[MCPTool]) -> List["MCPToolSDKTool"]: if MCPToolSDKTool is None: - raise ImportError( - "MCP SDK is not installed. Please install it with: pip install 'litellm[proxy]'" - ) + raise ImportError("MCP SDK is not installed. Please install it with: pip install 'litellm[proxy]'") return [ MCPToolSDKTool( name=tool.name, @@ -108,9 +100,7 @@ class MCPToolRegistry: fires. """ if mcp_tools_config is None: - raise ValueError( - "mcp_tools_config is required, please set `mcp_tools` in your proxy config" - ) + raise ValueError("mcp_tools_config is required, please set `mcp_tools` in your proxy config") for tool_config in mcp_tools_config: if not isinstance(tool_config, dict): @@ -131,9 +121,7 @@ class MCPToolRegistry: handler = get_instance_fn(handler_name, config_file_path) if handler is None: - verbose_logger.warning( - f"Warning: Could not find handler {handler_name} for tool {name}" - ) + verbose_logger.warning(f"Warning: Could not find handler {handler_name} for tool {name}") continue # Register the tool @@ -148,9 +136,7 @@ class MCPToolRegistry: input_schema=input_schema, handler=handler, ) - verbose_logger.debug( - "all registered tools: %s", json.dumps(self.tools, indent=4, default=str) - ) + verbose_logger.debug("all registered tools: %s", json.dumps(self.tools, indent=4, default=str)) global_mcp_tool_registry = MCPToolRegistry() diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index a996131653f..9652a3a2888 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -39,9 +39,7 @@ async def get_mcp_toolset( prisma_client: PrismaClient, toolset_id: str, ) -> Optional[MCPToolset]: - row = await MCPToolsetRepository(prisma_client).table.find_unique( - where={"toolset_id": toolset_id} - ) + row = await MCPToolsetRepository(prisma_client).table.find_unique(where={"toolset_id": toolset_id}) if row is None: return None return _toolset_from_row(row) @@ -59,9 +57,7 @@ async def list_mcp_toolsets( return [_toolset_from_row(r) for r in rows] except Exception as e: verbose_proxy_logger.warning( - "litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - {}".format( - str(e) - ) + "litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - {}".format(str(e)) ) return [] @@ -70,9 +66,7 @@ async def get_mcp_toolset_by_name( prisma_client: PrismaClient, toolset_name: str, ) -> Optional[MCPToolset]: - row = await MCPToolsetRepository(prisma_client).table.find_first( - where={"toolset_name": toolset_name} - ) + row = await MCPToolsetRepository(prisma_client).table.find_first(where={"toolset_name": toolset_name}) if row is None: return None return _toolset_from_row(row) @@ -106,9 +100,7 @@ async def delete_mcp_toolset( toolset_id: str, ) -> Optional[MCPToolset]: try: - row = await MCPToolsetRepository(prisma_client).table.delete( - where={"toolset_id": toolset_id} - ) + row = await MCPToolsetRepository(prisma_client).table.delete(where={"toolset_id": toolset_id}) except Exception as e: from prisma.errors import RecordNotFoundError diff --git a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py index 37a3228ebf0..1b37b884987 100644 --- a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py +++ b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py @@ -28,10 +28,7 @@ async def resolve_ui_session_team_ids( ) -> List[str]: """Resolve the real team ids backing a UI session token.""" - if ( - user_api_key_auth.team_id != UI_SESSION_TOKEN_TEAM_ID - or not user_api_key_auth.user_id - ): + if user_api_key_auth.team_id != UI_SESSION_TOKEN_TEAM_ID or not user_api_key_auth.user_id: return [] from litellm.proxy.auth.auth_checks import get_user_object @@ -78,8 +75,5 @@ async def build_effective_auth_contexts( resolved_team_ids = await resolve_ui_session_team_ids(user_api_key_auth) if resolved_team_ids: - return [ - clone_user_api_key_auth_with_team(user_api_key_auth, team_id) - for team_id in resolved_team_ids - ] + return [clone_user_api_key_auth_with_team(user_api_key_auth, team_id) for team_id in resolved_team_ids] return [user_api_key_auth] diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 3418417a8f8..9cb6d404b01 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -30,13 +30,9 @@ from urllib.parse import quote # module is reloaded (e.g. ``importlib.reload``). Tests that override these # variables must reload this module — see # ``tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py``. -LITELLM_MCP_SERVER_NAME = os.environ.get( - "LITELLM_MCP_SERVER_NAME", "litellm-mcp-server" -) +LITELLM_MCP_SERVER_NAME = os.environ.get("LITELLM_MCP_SERVER_NAME", "litellm-mcp-server") LITELLM_MCP_SERVER_VERSION = "1.0.0" -LITELLM_MCP_SERVER_DESCRIPTION = os.environ.get( - "LITELLM_MCP_SERVER_DESCRIPTION", "MCP Server for LiteLLM" -) +LITELLM_MCP_SERVER_DESCRIPTION = os.environ.get("LITELLM_MCP_SERVER_DESCRIPTION", "MCP Server for LiteLLM") MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-") MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}" @@ -391,9 +387,7 @@ def is_tool_name_prefixed( return True -def validate_mcp_server_name( - server_name: str, raise_http_exception: bool = False -) -> None: +def validate_mcp_server_name(server_name: str, raise_http_exception: bool = False) -> None: """ Validate that MCP server name does not contain 'MCP_TOOL_PREFIX_SEPARATOR'. @@ -410,9 +404,7 @@ def validate_mcp_server_name( from fastapi import HTTPException from starlette import status - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, detail={"error": error_message} - ) + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail={"error": error_message}) else: raise Exception(error_message) @@ -527,9 +519,7 @@ def interpolate_env_vars(value: str, variables: Mapping[str, str]) -> str: return _ENV_VAR_PATTERN.sub(_sub, value) -def interpolate_headers( - headers: Mapping[str, str], variables: Mapping[str, str] -) -> Dict[str, str]: +def interpolate_headers(headers: Mapping[str, str], variables: Mapping[str, str]) -> Dict[str, str]: """Return a copy of ``headers`` with every value passed through ``interpolate_env_vars``.""" return {k: interpolate_env_vars(v, variables) for k, v in headers.items()} diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 73b8a5538de..1dda1f29fb9 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -27,9 +27,7 @@ def _include_router(attr_name: str = "router") -> Callable[["FastAPI", object], return _register -def _mount_app( - prefix: str, attr_name: str = "app" -) -> Callable[["FastAPI", object], None]: +def _mount_app(prefix: str, attr_name: str = "app") -> Callable[["FastAPI", object], None]: def _register(app: "FastAPI", module: object) -> None: app.mount(path=prefix, app=getattr(module, attr_name)) @@ -41,9 +39,7 @@ class LazyFeature: name: str module_path: str path_prefixes: Tuple[str, ...] - register_fn: Callable[["FastAPI", object], None] = field( - default_factory=lambda: _include_router("router") - ) + register_fn: Callable[["FastAPI", object], None] = field(default_factory=lambda: _include_router("router")) # For routes whose path has a leading parameter (e.g. /{server}/authorize) # — startswith can't match those, so the matcher also checks endswith. path_suffixes: Tuple[str, ...] = () @@ -52,9 +48,7 @@ class LazyFeature: persistent_swagger_stub: bool = False def matches(self, path: str) -> bool: - return any(path.startswith(p) for p in self.path_prefixes) or any( - path.endswith(s) for s in self.path_suffixes - ) + return any(path.startswith(p) for p in self.path_prefixes) or any(path.endswith(s) for s in self.path_suffixes) LAZY_FEATURES: Tuple[LazyFeature, ...] = ( @@ -295,9 +289,7 @@ class LazyFeatureMiddleware: async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: # Short-circuit once every feature has loaded. - if scope["type"] in ("http", "websocket") and len(self._loaded) < len( - self._features - ): + if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features): path = scope.get("path", "") # Strip SERVER_ROOT_PATH so prefix matching works under a server # root path. Without this, requests like /api/v1/policies/... never @@ -330,9 +322,7 @@ async def _force_load(app: "FastAPI", feat: LazyFeature) -> bool: # Import on a thread (heavy modules take 1-3 s). register_fn # mutates app.router.routes, so it stays on the loop thread. loop = asyncio.get_running_loop() - module = await loop.run_in_executor( - None, importlib.import_module, feat.module_path - ) + module = await loop.run_in_executor(None, importlib.import_module, feat.module_path) feat.register_fn(app, module) app.state.lazy_loaded.add(feat.module_path) app.openapi_schema = None @@ -424,9 +414,7 @@ def inject_lazy_stubs(schema: Dict) -> Dict: if fragment: for p, ops in fragment.get("paths", {}).items(): paths.setdefault(p, ops) - for name, sch in ( - fragment.get("components", {}).get("schemas", {}).items() - ): + for name, sch in fragment.get("components", {}).get("schemas", {}).items(): schemas.setdefault(name, sch) continue @@ -453,8 +441,4 @@ def lazy_tag_to_prefix() -> Dict[str, str]: if load_snapshot(): return {} - return { - feat.name: feat.path_prefixes[0] - for feat in LAZY_FEATURES - if not feat.persistent_swagger_stub - } + return {feat.name: feat.path_prefixes[0] for feat in LAZY_FEATURES if not feat.persistent_swagger_stub} diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index c63ff8d0733..818232d650e 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -76,9 +76,7 @@ def _normalize_operation_ids(paths: Dict[str, Dict]) -> None: for suffix in methods: suffix_token = f"_{suffix}" if operation_id.endswith(suffix_token): - operation["operationId"] = ( - operation_id[: -len(suffix_token)] + f"_{method}" - ) + operation["operationId"] = operation_id[: -len(suffix_token)] + f"_{method}" break @@ -102,11 +100,7 @@ def generate_snapshot() -> Dict[str, Dict]: fragments: Dict[str, Dict] = {} used_operation_ids: Set[str] = set() for feat in LAZY_FEATURES: - feat_routes = [ - r - for r in app.routes - if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes) - ] + feat_routes = [r for r in app.routes if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)] if not feat_routes: continue _stabilize_multi_method_route_ids(feat_routes) @@ -121,9 +115,7 @@ def generate_snapshot() -> Dict[str, Dict]: if isinstance(operation_id, str): for suffix in HTTP_METHOD_SUFFIXES: if operation_id.endswith(f"_{suffix}"): - op["operationId"] = ( - operation_id[: -len(suffix)] + method - ) + op["operationId"] = operation_id[: -len(suffix)] + method break op["tags"] = [feat.name] full = ensure_unique_openapi_operation_ids(full, used_operation_ids) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 559720eb64e..d84588a4c24 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -850,12 +850,7 @@ class LiteLLMRoutes(enum.Enum): ) # All routes accesible by an Org Admin - org_admin_allowed_routes = ( - org_admin_only_routes - + management_routes - + self_managed_routes - + admin_viewer_routes - ) + org_admin_allowed_routes = org_admin_only_routes + management_routes + self_managed_routes + admin_viewer_routes class LiteLLMPromptInjectionParams(LiteLLMPydanticObjectBase): @@ -876,23 +871,11 @@ class LiteLLMPromptInjectionParams(LiteLLMPydanticObjectBase): llm_api_check = values.get("llm_api_check") if llm_api_check is True: if "llm_api_name" not in values or not values["llm_api_name"]: - raise ValueError( - "If llm_api_check is set to True, llm_api_name must be provided" - ) - if ( - "llm_api_system_prompt" not in values - or not values["llm_api_system_prompt"] - ): - raise ValueError( - "If llm_api_check is set to True, llm_api_system_prompt must be provided" - ) - if ( - "llm_api_fail_call_string" not in values - or not values["llm_api_fail_call_string"] - ): - raise ValueError( - "If llm_api_check is set to True, llm_api_fail_call_string must be provided" - ) + raise ValueError("If llm_api_check is set to True, llm_api_name must be provided") + if "llm_api_system_prompt" not in values or not values["llm_api_system_prompt"]: + raise ValueError("If llm_api_check is set to True, llm_api_system_prompt must be provided") + if "llm_api_fail_call_string" not in values or not values["llm_api_fail_call_string"]: + raise ValueError("If llm_api_check is set to True, llm_api_fail_call_string must be provided") return values @@ -1005,9 +988,7 @@ class ModelParams(LiteLLMPydanticObjectBase): @classmethod def set_model_info(cls, values): if values.get("model_info") is None: - values.update( - {"model_info": ModelInfo(id=None, mode="chat", base_model=None)} - ) + values.update({"model_info": ModelInfo(id=None, mode="chat", base_model=None)}) return values @@ -1046,15 +1027,11 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None - budget_limits: Optional[List[BudgetLimitEntry]] = ( - None # multiple concurrent budget windows - ) + budget_limits: Optional[List[BudgetLimitEntry]] = None # multiple concurrent budget windows allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[ - dict - ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[dict] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1089,12 +1066,12 @@ class KeyRequestBase(GenerateRequestBase): allowed_routes: Optional[list] = [] allowed_passthrough_routes: Optional[list] = None allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None - rpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm - tpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm + rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]] = ( + None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm + ) + tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]] = ( + None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm + ) router_settings: Optional[UpdateRouterConfig] = None access_group_ids: Optional[List[str]] = None @@ -1117,9 +1094,7 @@ class GenerateKeyRequest(KeyRequestBase): default=LiteLLMKeyType.DEFAULT, description="Type of key that determines default allowed routes.", ) - auto_rotate: Optional[bool] = Field( - default=False, description="Whether this key should be automatically rotated" - ) + auto_rotate: Optional[bool] = Field(default=False, description="Whether this key should be automatically rotated") rotation_interval: Optional[str] = Field( default=None, description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True", @@ -1185,9 +1160,7 @@ class UpdateKeyRequest(KeyRequestBase): def validate_temp_budget(self) -> "UpdateKeyRequest": if self.temp_budget_increase is not None or self.temp_budget_expiry is not None: if self.temp_budget_increase is None or self.temp_budget_expiry is None: - raise ValueError( - "temp_budget_increase and temp_budget_expiry must be set together" - ) + raise ValueError("temp_budget_increase and temp_budget_expiry must be set together") return self @@ -1199,9 +1172,7 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None - grace_period: Optional[str] = ( - None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke - ) + grace_period: Optional[str] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke class ResetSpendRequest(LiteLLMPydanticObjectBase): @@ -1216,9 +1187,7 @@ class KeyRequest(LiteLLMPydanticObjectBase): @classmethod def validate_at_least_one(cls, values): if not values.get("keys") and not values.get("key_aliases"): - raise ValueError( - "At least one of 'keys' or 'key_aliases' must be provided." - ) + raise ValueError("At least one of 'keys' or 'key_aliases' must be provided.") return values @@ -1319,9 +1288,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): ) elif transport in [MCPTransport.http, MCPTransport.sse]: if not values.get("url") and not values.get("spec_path"): - raise ValueError( - "url or spec_path is required for HTTP/SSE transport" - ) + raise ValueError("url or spec_path is required for HTTP/SSE transport") return values @model_validator(mode="before") @@ -1392,9 +1359,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): ) elif transport in [MCPTransport.http, MCPTransport.sse]: if not values.get("url") and not values.get("spec_path"): - raise ValueError( - "url or spec_path is required for HTTP/SSE transport" - ) + raise ValueError("url or spec_path is required for HTTP/SSE transport") return values @@ -1551,9 +1516,7 @@ class NewUserRequest(GenerateRequestBase): ] ] = None teams: Optional[Union[List[str], List[NewUserRequestTeam]]] = None - auto_create_key: bool = ( - True # flag used for returning a key as part of the /user/new response - ) + auto_create_key: bool = True # flag used for returning a key as part of the /user/new response send_invite_email: Optional[bool] = None sso_user_id: Optional[str] = None organizations: Optional[List[str]] = None @@ -1577,9 +1540,7 @@ class NewUserResponse(GenerateKeyResponse): updated_at: Optional[datetime] = None -class UpdateUserRequestNoUserIDorEmail( - GenerateRequestBase -): # shared with BulkUpdateUserRequest +class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest password: Optional[str] = None spend: Optional[float] = None metadata: Optional[dict] = None @@ -1629,12 +1590,8 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase): max_parallel_requests: Optional[int] = Field( default=None, description="Max concurrent requests allowed for this budget id." ) - tpm_limit: Optional[int] = Field( - default=None, description="Max tokens per minute, allowed for this budget id." - ) - rpm_limit: Optional[int] = Field( - default=None, description="Max requests per minute, allowed for this budget id." - ) + tpm_limit: Optional[int] = Field(default=None, description="Max tokens per minute, allowed for this budget id.") + rpm_limit: Optional[int] = Field(default=None, description="Max requests per minute, allowed for this budget id.") budget_duration: Optional[str] = Field( default=None, description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')", @@ -1681,9 +1638,7 @@ class NewCustomerRequest(BudgetNewRequest): allowed_model_region: Optional[AllowedModelRegion] = ( None # require all user requests to use models in this specific region ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + default_model: Optional[str] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1709,9 +1664,7 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): allowed_model_region: Optional[AllowedModelRegion] = ( None # require all user requests to use models in this specific region ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + default_model: Optional[str] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -1749,24 +1702,18 @@ class NewTeamRequest(TeamBase): disable_global_guardrails: Optional[bool] = None secret_manager_settings: Optional[dict] = None model_rpm_limit: Optional[Dict[str, int]] = None - rpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm - tpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput"] - ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm + rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] = ( + None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm + ) + tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] = ( + None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm + ) model_tpm_limit: Optional[Dict[str, int]] = None mcp_rpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[float] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[int] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[int] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1833,12 +1780,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): enforced_file_expires_after: Optional[dict] = None router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None - budget_limits: Optional[List[BudgetLimitEntry]] = ( - None # multiple concurrent budget windows - ) - default_team_member_models: Optional[List[str]] = ( - None # default allowed_models seeded onto new team members - ) + budget_limits: Optional[List[BudgetLimitEntry]] = None # multiple concurrent budget windows + default_team_member_models: Optional[List[str]] = None # default allowed_models seeded onto new team members class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -1875,9 +1818,7 @@ class BlockModelRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1887,13 +1828,9 @@ class AddTeamCallback(LiteLLMPydanticObjectBase): valid_keys = set(StandardCallbackDynamicParams.__annotations__.keys()) for key, value in callback_vars.items(): if key not in valid_keys: - raise ValueError( - f"Invalid callback variable: {key}. Must be one of {valid_keys}" - ) + raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}") callback_vars[key] = str(value) - validate_no_callback_env_reference( - key, callback_vars[key], source="key/team callback metadata" - ) + validate_no_callback_env_reference(key, callback_vars[key], source="key/team callback metadata") return values @@ -1931,9 +1868,7 @@ class TeamCallbackMetadata(LiteLLMPydanticObjectBase): if callback_vars is not None: for key in callback_vars: if key not in valid_keys: - raise ValueError( - f"Invalid callback variable: {key}. Must be one of {valid_keys}" - ) + raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}") return values @@ -2043,9 +1978,7 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): description="Optional unique identifier for the pass-through endpoint. If not provided, endpoints will be identified by path for backwards compatibility.", ) path: str = Field(description="The route to be added to the LiteLLM Proxy Server.") - target: str = Field( - description="The URL to which requests for this path should be forwarded." - ) + target: str = Field(description="The URL to which requests for this path should be forwarded.") headers: dict = Field( default={}, description="Key-value pairs of headers to be forwarded with the request. You can set any key value pair here and it will be forwarded to your target endpoint", @@ -2119,9 +2052,7 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[List[FieldDetail]] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2147,9 +2078,7 @@ class PluginConfig(LiteLLMPydanticObjectBase): """A single external service registered as an embeddable UI plugin.""" name: str = Field(description="unique plugin identifier (kebab-case)") - display_name: str | None = Field( - None, description="human-readable label shown in the UI view switcher" - ) + display_name: str | None = Field(None, description="human-readable label shown in the UI view switcher") url: str = Field(description="base URL of the plugin service") plugin_key: str | None = Field( None, @@ -2162,24 +2091,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ - completion_model: Optional[str] = Field( - None, description="proxy level default model for all chat completion calls" - ) + completion_model: Optional[str] = Field(None, description="proxy level default model for all chat completion calls") plugins: list[PluginConfig] | None = Field( None, description="external services registered as embeddable UI plugins" ) key_management_system: Optional[KeyManagementSystem] = Field( None, description="key manager to load keys from / decrypt keys with" ) - use_google_kms: Optional[bool] = Field( - None, description="decrypt keys with google kms" - ) - use_azure_key_vault: Optional[bool] = Field( - None, description="load keys from azure key vault" - ) - master_key: Optional[str] = Field( - None, description="require a key for all calls to proxy" - ) + use_google_kms: Optional[bool] = Field(None, description="decrypt keys with google kms") + use_azure_key_vault: Optional[bool] = Field(None, description="load keys from azure key vault") + master_key: Optional[str] = Field(None, description="require a key for all calls to proxy") allow_cli_sso_verification_uri_complete: bool | None = Field( None, description="opt-in to RFC 8628 verification_uri_complete for the CLI SSO device flow, pre-filling the user_code in the browser. Off by default; intended for same-host clients where the device that starts the flow and the browser run on the same machine", @@ -2230,9 +2151,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "takes precedence." ), ) - database_type: Optional[Literal["dynamo_db"]] = Field( - None, description="to use dynamodb instead of postgres db" - ) + database_type: Optional[Literal["dynamo_db"]] = Field(None, description="to use dynamodb instead of postgres db") database_args: Optional[DynamoDBArgs] = Field( None, description="custom args for instantiating dynamodb client - e.g. billing provision", @@ -2268,12 +2187,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="for `/models` endpoint, infers available model based on environment keys (e.g. OPENAI_API_KEY)", ) - background_health_checks: Optional[bool] = Field( - None, description="run health checks in background" - ) - health_check_interval: int = Field( - 300, description="background health check interval in seconds" - ) + background_health_checks: Optional[bool] = Field(None, description="run health checks in background") + health_check_interval: int = Field(300, description="background health check interval in seconds") health_check_concurrency: Optional[int] = Field( None, description=( @@ -2306,12 +2221,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="sends alerts if requests hang for 5min+", ) - ui_access_mode: Optional[Literal["admin_only", "all"]] = Field( - "all", description="Control access to the Proxy UI" - ) - allowed_routes: Optional[List] = Field( - None, description="Proxy API Endpoints you want users to be able to access" - ) + ui_access_mode: Optional[Literal["admin_only", "all"]] = Field("all", description="Control access to the Proxy UI") + allowed_routes: Optional[List] = Field(None, description="Proxy API Endpoints you want users to be able to access") reject_clientside_metadata_tags: Optional[bool] = Field( None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", @@ -2493,9 +2404,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): current = getattr(self, attr_name, None) # Apply budget value when key has no value, or for model_max_budget when key has empty dict should_apply = current is None or ( - attr_name == "model_max_budget" - and isinstance(current, dict) - and len(current) == 0 + attr_name == "model_max_budget" and isinstance(current, dict) and len(current) == 0 ) if should_apply: kwargs[attr_name] = value @@ -2508,9 +2417,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): super().__init__(**kwargs) -class UserAPIKeyAuth( - LiteLLM_VerificationTokenView -): # the expected response object for user api key auth +class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response object for user api key auth """ Return the row in the db """ @@ -2530,9 +2437,7 @@ class UserAPIKeyAuth( is_session_token: bool = False budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True) user: Optional[Any] = None # Expanded user object when expand=user is used - created_by_user: Optional[Any] = ( - None # Expanded created_by user when expand=user is used - ) + created_by_user: Optional[Any] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None # Team object_permission preloaded in auth (e.g. get_team_object) to avoid # per-request object_permission fetches in downstream checks (vector stores, etc.) @@ -2550,13 +2455,9 @@ class UserAPIKeyAuth( if not isinstance(values, dict): return values if values.get("api_key") is not None: - values.update( - {"token": cls._safe_hash_litellm_api_key(values.get("api_key"))} - ) + values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): - values.update( - {"api_key": cls._safe_hash_litellm_api_key(values.get("api_key"))} - ) + values.update({"api_key": cls._safe_hash_litellm_api_key(values.get("api_key"))}) return values @classmethod @@ -2765,9 +2666,7 @@ class NewProjectRequest(LiteLLM_BudgetTable): def set_model_info(cls, values): if "tags" in values and values["tags"] is not None: if not isinstance(values["tags"], list): - raise ValueError( - f"tags must be a list of strings, got {type(values['tags']).__name__}" - ) + raise ValueError(f"tags must be a list of strings, got {type(values['tags']).__name__}") for field in LiteLLM_ManagementEndpoint_MetadataFields: if values.get(field) is not None: if values.get("metadata") is None: @@ -2800,9 +2699,7 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): def set_model_info(cls, values): if "tags" in values and values["tags"] is not None: if not isinstance(values["tags"], list): - raise ValueError( - f"tags must be a list of strings, got {type(values['tags']).__name__}" - ) + raise ValueError(f"tags must be a list of strings, got {type(values['tags']).__name__}") for field in LiteLLM_ManagementEndpoint_MetadataFields: if values.get(field) is not None: if values.get("metadata") is None: @@ -2860,9 +2757,7 @@ from litellm.models.spend_logs import ( # noqa: E402 ) from litellm.models.tag import LiteLLM_TagTable as LiteLLM_TagTable # noqa: E402 -AUDIT_ACTIONS = Literal[ - "created", "updated", "deleted", "blocked", "unblocked", "rotated" -] +AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "unblocked", "rotated"] class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): @@ -3167,9 +3062,7 @@ class SpendLogsMetadata(TypedDict): Specific metadata k,v pairs logged to spendlogs for easier cost tracking """ - additional_usage_values: Optional[ - dict - ] # covers provider-specific usage information - e.g. prompt caching + additional_usage_values: Optional[dict] # covers provider-specific usage information - e.g. prompt caching user_api_key: Optional[str] user_api_key_alias: Optional[str] user_api_key_team_id: Optional[str] @@ -3178,9 +3071,7 @@ class SpendLogsMetadata(TypedDict): user_api_key_org_id: Optional[str] user_api_key_user_id: Optional[str] user_api_key_team_alias: Optional[str] - spend_logs_metadata: Optional[ - dict - ] # special param to log k,v pairs to spendlogs for a call + spend_logs_metadata: Optional[dict] # special param to log k,v pairs to spendlogs for a call requester_ip_address: Optional[str] litellm_call_id: Optional[str] applied_guardrails: Optional[List[str]] @@ -3194,17 +3085,11 @@ class SpendLogsMetadata(TypedDict): error_information: Optional[StandardLoggingPayloadErrorInformation] usage_object: Optional[dict] model_map_information: Optional[StandardLoggingModelInformation] - cold_storage_object_key: Optional[ - str - ] # S3/GCS object key for cold storage retrieval + cold_storage_object_key: Optional[str] # S3/GCS object key for cold storage retrieval litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds - attempted_retries: Optional[ - int - ] # Number of retries attempted (0 = first attempt succeeded) + attempted_retries: Optional[int] # Number of retries attempted (0 = first attempt succeeded) max_retries: Optional[int] # Max retries configured for this request - cost_breakdown: Optional[ - CostBreakdown - ] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) + cost_breakdown: Optional[CostBreakdown] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) class SpendLogsPayload(TypedDict): @@ -3342,10 +3227,7 @@ class ProxyException(Exception): # rules for proxyExceptions # Litellm router.py returns "No healthy deployment available" when there are no deployments available # Should map to 429 errors https://github.com/BerriAI/litellm/issues/2487 - if ( - "No healthy deployment available" in self.message - or "No deployments available" in self.message - ): + if "No healthy deployment available" in self.message or "No deployments available" in self.message: self.code = "429" elif RouterErrors.no_deployments_with_tag_routing.value in self.message: self.code = "401" @@ -3372,13 +3254,9 @@ class CommonProxyErrors(str, enum.Enum): no_llm_router = "No models configured on proxy" not_allowed_access = "Admin-only endpoint. Not allowed to access this." not_premium_user = "You must be a LiteLLM Enterprise user to use this feature. If you have a license please set `LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/enterprise#trial. \nPricing: https://www.litellm.ai/#pricing" - max_parallel_request_limit_reached = ( - "Crossed TPM / RPM / Max Parallel Request Limit" - ) + max_parallel_request_limit_reached = "Crossed TPM / RPM / Max Parallel Request Limit" missing_enterprise_package = "Missing litellm-enterprise package. Please install it to use this feature. Run `pip install litellm-enterprise`" - missing_enterprise_package_docker = ( - "This uses the enterprise folder - only available on the Docker image." - ) + missing_enterprise_package_docker = "This uses the enterprise folder - only available on the Docker image." class SpendCalculateRequest(LiteLLMPydanticObjectBase): @@ -3570,10 +3448,7 @@ class MemberAddRequest(LiteLLMPydanticObjectBase): member_data = data.get("member") if isinstance(member_data, list): # If member is a list of dictionaries, convert each dictionary to a Member object - members = [ - Member(**item) if isinstance(item, dict) else item - for item in member_data - ] + members = [Member(**item) if isinstance(item, dict) else item for item in member_data] # Replace member_data with the list of Member objects data["member"] = members elif isinstance(member_data, dict): @@ -3675,12 +3550,8 @@ class TeamMemberDeleteRequest(MemberDeleteRequest): class TeamMemberUpdateRequest(TeamMemberDeleteRequest): max_budget_in_team: Optional[float] = None role: Optional[Literal["admin", "user"]] = None - tpm_limit: Optional[int] = Field( - default=None, description="Tokens per minute limit for this team member" - ) - rpm_limit: Optional[int] = Field( - default=None, description="Requests per minute limit for this team member" - ) + tpm_limit: Optional[int] = Field(default=None, description="Tokens per minute limit for this team member") + rpm_limit: Optional[int] = Field(default=None, description="Requests per minute limit for this team member") budget_duration: Optional[str] = Field( default=None, description="Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.", @@ -3717,9 +3588,7 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[float] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -3738,13 +3607,9 @@ class OrganizationMemberUpdateRequest(OrganizationMemberDeleteRequest): role: Optional[LitellmUserRoles] = None @field_validator("role") - def validate_role( - cls, value: Optional[LitellmUserRoles] - ) -> Optional[LitellmUserRoles]: + def validate_role(cls, value: Optional[LitellmUserRoles]) -> Optional[LitellmUserRoles]: if value is not None and value not in ROLES_WITHIN_ORG: - raise ValueError( - f"Invalid role. Must be one of: {[role.value for role in ROLES_WITHIN_ORG]}" - ) + raise ValueError(f"Invalid role. Must be one of: {[role.value for role in ROLES_WITHIN_ORG]}") return value @@ -3898,43 +3763,37 @@ class JWKUrlResponse(TypedDict, total=False): class UserManagementEndpointParamDocStringEnums(str, enum.Enum): - user_id_doc_str = ( - "Optional[str] - Specify a user id. If not set, a unique id will be generated." - ) - user_alias_doc_str = ( - "Optional[str] - A descriptive name for you to know who this user id refers to." - ) + user_id_doc_str = "Optional[str] - Specify a user id. If not set, a unique id will be generated." + user_alias_doc_str = "Optional[str] - A descriptive name for you to know who this user id refers to." teams_doc_str = "Optional[list] - specify a list of team id's a user belongs to." user_email_doc_str = "Optional[str] - Specify a user email." - send_invite_email_doc_str = ( - "Optional[bool] - Specify if an invite email should be sent." - ) + send_invite_email_doc_str = "Optional[bool] - Specify if an invite email should be sent." user_role_doc_str = """Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`""" max_budget_doc_str = """Optional[float] - Specify max budget for a given user.""" budget_duration_doc_str = """Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo").""" - models_doc_str = """Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)""" - tpm_limit_doc_str = ( - """Optional[int] - Specify tpm limit for a given user (Tokens per minute)""" - ) - rpm_limit_doc_str = ( - """Optional[int] - Specify rpm limit for a given user (Requests per minute)""" + models_doc_str = ( + """Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)""" ) + tpm_limit_doc_str = """Optional[int] - Specify tpm limit for a given user (Tokens per minute)""" + rpm_limit_doc_str = """Optional[int] - Specify rpm limit for a given user (Requests per minute)""" auto_create_key_doc_str = """bool - Default=True. Flag used for returning a key as part of the /user/new response""" aliases_doc_str = """Optional[dict] - Model aliases for the user - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)""" config_doc_str = """Optional[dict] - [DEPRECATED PARAM] User-specific config.""" allowed_cache_controls_doc_str = """Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request-""" - blocked_doc_str = ( - """Optional[bool] - [Not Implemented Yet] Whether the user is blocked.""" - ) + blocked_doc_str = """Optional[bool] - [Not Implemented Yet] Whether the user is blocked.""" guardrails_doc_str = """Optional[List[str]] - [Not Implemented Yet] List of active guardrails for the user""" - permissions_doc_str = """Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.""" + permissions_doc_str = ( + """Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.""" + ) metadata_doc_str = """Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }""" max_parallel_requests_doc_str = """Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.""" soft_budget_doc_str = """Optional[float] - Get alerts when user crosses given budget, doesn't block requests.""" model_max_budget_doc_str = """Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)""" model_rpm_limit_doc_str = """Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" model_tpm_limit_doc_str = """Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" - spend_doc_str = """Optional[float] - Amount spent by user. Default is 0. Will be updated by proxy whenever user is used.""" + spend_doc_str = ( + """Optional[float] - Amount spent by user. Default is 0. Will be updated by proxy whenever user is used.""" + ) team_id_doc_str = """Optional[str] - [DEPRECATED PARAM] The team id of the user. Default is None.""" duration_doc_str = """Optional[str] - Duration for the key auto-created on `/user/new`. Default is None.""" @@ -4251,18 +4110,14 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): user_allowed_email_domain: Optional[str] = None user_roles_jwt_field: Optional[str] = None user_allowed_roles: Optional[List[str]] = None - user_id_upsert: bool = Field( - default=False, description="If user doesn't exist, upsert them into the db." - ) + user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.") end_user_id_jwt_field: Optional[str] = None public_key_ttl: float = 600 public_allowed_routes: List[str] = ["public_routes"] enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[str] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False @@ -4364,9 +4219,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): if (user_roles_jwt_field is not None and user_allowed_roles is None) or ( user_roles_jwt_field is None and user_allowed_roles is not None ): - raise ValueError( - "user_allowed_roles must be provided if user_roles_jwt_field is set." - ) + raise ValueError("user_allowed_roles must be provided if user_roles_jwt_field is set.") if object_id_jwt_field is not None and role_mappings is None: raise ValueError( @@ -4374,9 +4227,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): ) if scope_mappings is not None and not enforce_scope_based_access: - raise ValueError( - "scope_mappings must be set if enforce_scope_based_access is true." - ) + raise ValueError("scope_mappings must be set if enforce_scope_based_access is true.") super().__init__(**kwargs) @@ -4423,9 +4274,7 @@ class DefaultInternalUserParams(LiteLLMPydanticObjectBase): default=None, description="Default budget duration for new users (e.g. 'daily', 'weekly', 'monthly')", ) - models: Optional[List[str]] = Field( - default=None, description="Default list of models that new users can access" - ) + models: Optional[List[str]] = Field(default=None, description="Default list of models that new users can access") teams: Optional[Union[List[str], List[NewUserRequestTeam]]] = Field( default=None, @@ -4543,12 +4392,8 @@ class CostEstimateRequest(LiteLLMPydanticObjectBase): model: str = Field(description="Model name (from /model_group/info)") input_tokens: int = Field(description="Expected input tokens per request", ge=0) output_tokens: int = Field(description="Expected output tokens per request", ge=0) - num_requests_per_day: Optional[int] = Field( - default=None, description="Number of requests per day", ge=0 - ) - num_requests_per_month: Optional[int] = Field( - default=None, description="Number of requests per month", ge=0 - ) + num_requests_per_day: Optional[int] = Field(default=None, description="Number of requests per day", ge=0) + num_requests_per_month: Optional[int] = Field(default=None, description="Number of requests per month", ge=0) class CostEstimateResponse(LiteLLMPydanticObjectBase): @@ -4560,44 +4405,20 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase): num_requests_per_day: Optional[int] = None num_requests_per_month: Optional[int] = None # Per-request costs - cost_per_request: float = Field( - description="Total cost per request (includes margin)" - ) - input_cost_per_request: float = Field( - description="Input token cost per request (before margin)" - ) - output_cost_per_request: float = Field( - description="Output token cost per request (before margin)" - ) - margin_cost_per_request: float = Field( - default=0.0, description="Margin/fee added per request" - ) + cost_per_request: float = Field(description="Total cost per request (includes margin)") + input_cost_per_request: float = Field(description="Input token cost per request (before margin)") + output_cost_per_request: float = Field(description="Output token cost per request (before margin)") + margin_cost_per_request: float = Field(default=0.0, description="Margin/fee added per request") # Daily costs (if num_requests_per_day provided) - daily_cost: Optional[float] = Field( - default=None, description="Total daily cost (includes margin)" - ) - daily_input_cost: Optional[float] = Field( - default=None, description="Daily input token cost" - ) - daily_output_cost: Optional[float] = Field( - default=None, description="Daily output token cost" - ) - daily_margin_cost: Optional[float] = Field( - default=None, description="Daily margin/fee" - ) + daily_cost: Optional[float] = Field(default=None, description="Total daily cost (includes margin)") + daily_input_cost: Optional[float] = Field(default=None, description="Daily input token cost") + daily_output_cost: Optional[float] = Field(default=None, description="Daily output token cost") + daily_margin_cost: Optional[float] = Field(default=None, description="Daily margin/fee") # Monthly costs (if num_requests_per_month provided) - monthly_cost: Optional[float] = Field( - default=None, description="Total monthly cost (includes margin)" - ) - monthly_input_cost: Optional[float] = Field( - default=None, description="Monthly input token cost" - ) - monthly_output_cost: Optional[float] = Field( - default=None, description="Monthly output token cost" - ) - monthly_margin_cost: Optional[float] = Field( - default=None, description="Monthly margin/fee" - ) + monthly_cost: Optional[float] = Field(default=None, description="Total monthly cost (includes margin)") + monthly_input_cost: Optional[float] = Field(default=None, description="Monthly input token cost") + monthly_output_cost: Optional[float] = Field(default=None, description="Monthly output token cost") + monthly_margin_cost: Optional[float] = Field(default=None, description="Monthly margin/fee") # Pricing info input_cost_per_token: Optional[float] = None output_cost_per_token: Optional[float] = None diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index 57d360ab5af..79638341ac1 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -97,9 +97,7 @@ def _filter_capabilities(upstream_capabilities: Any) -> Dict[str, Any]: if not isinstance(upstream_capabilities, dict): return {} return { - key: value - for key, value in upstream_capabilities.items() - if key in _ALLOWED_CAPABILITY_KEYS and bool(value) + key: value for key, value in upstream_capabilities.items() if key in _ALLOWED_CAPABILITY_KEYS and bool(value) } diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py index c7ac8dc0bee..a95f1d2dfeb 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/litellm/proxy/a2a/discovery.py @@ -66,16 +66,12 @@ def _build_langgraph_platform_paths( """ assistant_id = (params or {}).get("assistant_id") if not assistant_id: - raise AgentCardDiscoveryError( - "langgraph_platform discovery requires params.assistant_id" - ) + raise AgentCardDiscoveryError("langgraph_platform discovery requires params.assistant_id") query = urlencode({"assistant_id": str(assistant_id)}) return tuple(f"{path}?{query}" for path in AGENT_CARD_WELL_KNOWN_PATHS) -def _paths_for_mode( - mode: DiscoveryMode, params: Optional[Dict[str, Any]] -) -> Tuple[str, ...]: +def _paths_for_mode(mode: DiscoveryMode, params: Optional[Dict[str, Any]]) -> Tuple[str, ...]: if mode == DiscoveryMode.WELL_KNOWN_FALLBACK: return AGENT_CARD_WELL_KNOWN_PATHS if mode == DiscoveryMode.LANGGRAPH_PLATFORM: @@ -127,9 +123,7 @@ async def fetch_well_known_card( response = await async_safe_get(client, url, headers=headers or {}) except SSRFError as exc: last_error = f"{url}: {exc!s}" - verbose_proxy_logger.debug( - "A2A discovery blocked by SSRF guard for %s: %s", url, exc - ) + verbose_proxy_logger.debug("A2A discovery blocked by SSRF guard for %s: %s", url, exc) continue except Exception as exc: last_error = f"{url}: {exc!s}" @@ -138,9 +132,7 @@ async def fetch_well_known_card( if response.status_code >= 400: last_error = f"{url}: HTTP {response.status_code}" - verbose_proxy_logger.debug( - "A2A discovery HTTP %s for %s", response.status_code, url - ) + verbose_proxy_logger.debug("A2A discovery HTTP %s for %s", response.status_code, url) continue try: @@ -157,6 +149,5 @@ async def fetch_well_known_card( return card raise AgentCardDiscoveryError( - f"Could not fetch agent card from {base_url} (mode={discovery_mode.value}). " - f"Last error: {last_error}" + f"Could not fetch agent card from {base_url} (mode={discovery_mode.value}). Last error: {last_error}" ) diff --git a/litellm/proxy/a2a/endpoints.py b/litellm/proxy/a2a/endpoints.py index 520fdc9d8c5..a46de73fabc 100644 --- a/litellm/proxy/a2a/endpoints.py +++ b/litellm/proxy/a2a/endpoints.py @@ -91,10 +91,7 @@ async def discover_agent_card( if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=403, - detail=( - "Only proxy admins can discover agent cards. " - f"Your role={user_api_key_dict.user_role}" - ), + detail=(f"Only proxy admins can discover agent cards. Your role={user_api_key_dict.user_role}"), ) try: diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index b1fb72619b1..78be6231ae0 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -66,11 +66,7 @@ def _forwarding_headers( agent_extra_headers: Optional[Dict[str, str]], ) -> Optional[Dict[str, str]]: sanitized = ( - { - k: v - for k, v in agent_extra_headers.items() - if not k.lower().startswith("x-litellm-") - } + {k: v for k, v in agent_extra_headers.items() if not k.lower().startswith("x-litellm-")} if agent_extra_headers else None ) @@ -123,10 +119,7 @@ def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: if not trace_id: raise HTTPException( status_code=400, - detail=( - f"Agent '{agent.agent_id}' requires x-litellm-trace-id header " - "on all inbound requests." - ), + detail=(f"Agent '{agent.agent_id}' requires x-litellm-trace-id header on all inbound requests."), ) @@ -188,9 +181,7 @@ async def _a2a_sse_event_source( try: parsed = json.loads(error_body) if isinstance(parsed, dict) and "error" in parsed: - error_event = _normalize_a2a_jsonrpc_response( - parsed, request_id=request_id - ) + error_event = _normalize_a2a_jsonrpc_response(parsed, request_id=request_id) except Exception: error_event = None yield error_event or { @@ -223,9 +214,7 @@ async def _forward_jsonrpc_sse( user_api_key_dict: Optional[Any] = None, request_data: Optional[dict] = None, ) -> StreamingResponse: - event_source = _a2a_sse_event_source( - agent_url, body, request_id=request_id, extra_headers=extra_headers - ) + event_source = _a2a_sse_event_source(agent_url, body, request_id=request_id, extra_headers=extra_headers) def _serialize_chunk(chunk: Any) -> str: return f"data: {json.dumps(chunk)}\n\n" @@ -246,11 +235,7 @@ async def _forward_jsonrpc_sse( + "\n\n" ) - if ( - proxy_logging_obj is not None - and user_api_key_dict is not None - and request_data is not None - ): + if proxy_logging_obj is not None and user_api_key_dict is not None and request_data is not None: # Route streamed events through the shared streaming generator so the # post-call streaming hook (and therefore agent guardrails) inspects # tasks/resubscribe output the same way message/stream does. @@ -258,15 +243,13 @@ async def _forward_jsonrpc_sse( ProxyBaseLLMRequestProcessing, ) - generator: AsyncGenerator[str, None] = ( - ProxyBaseLLMRequestProcessing.async_streaming_data_generator( - response=event_source, - user_api_key_dict=user_api_key_dict, - request_data=request_data, - proxy_logging_obj=proxy_logging_obj, - serialize_chunk=_serialize_chunk, - serialize_error=_serialize_error, - ) + generator: AsyncGenerator[str, None] = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=event_source, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + proxy_logging_obj=proxy_logging_obj, + serialize_chunk=_serialize_chunk, + serialize_error=_serialize_error, ) else: @@ -323,11 +306,7 @@ async def _handle_stream_message( from a2a.types import MessageSendParams, SendStreamingMessageRequest - use_proxy_hooks = ( - user_api_key_dict is not None - and request_data is not None - and proxy_logging_obj is not None - ) + use_proxy_hooks = user_api_key_dict is not None and request_data is not None and proxy_logging_obj is not None async def stream_response(): try: @@ -381,9 +360,7 @@ async def _handle_stream_message( + "\n" ) - async for ( - line - ) in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + async for line in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=a2a_stream, user_api_key_dict=user_api_key_dict, request_data=request_data, @@ -395,10 +372,7 @@ async def _handle_stream_message( else: async for chunk in a2a_stream: if hasattr(chunk, "model_dump"): - yield ( - json.dumps(chunk.model_dump(mode="json", exclude_none=True)) - + "\n" - ) + yield (json.dumps(chunk.model_dump(mode="json", exclude_none=True)) + "\n") else: yield json.dumps(chunk) + "\n" except Exception as e: @@ -492,9 +466,7 @@ async def get_agent_card( "url": f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}", } - verbose_proxy_logger.debug( - f"Returning agent card for '{agent_id}' with proxy URL: {agent_card['url']}" - ) + verbose_proxy_logger.debug(f"Returning agent card for '{agent_id}' with proxy URL: {agent_card['url']}") return JSONResponse(content=agent_card) except HTTPException: @@ -552,9 +524,7 @@ async def invoke_agent_a2a( # Validate JSON-RPC format if body.get("jsonrpc") != "2.0": - return _jsonrpc_error( - body.get("id"), -32600, "Invalid Request: jsonrpc must be '2.0'" - ) + return _jsonrpc_error(body.get("id"), -32600, "Invalid Request: jsonrpc must be '2.0'") request_id: Optional[Any] = body.get("id") method: Optional[str] = body.get("method") @@ -581,9 +551,7 @@ async def invoke_agent_a2a( # Find the agent agent = _get_agent(agent_id) if agent is None: - return _jsonrpc_error( - request_id, -32000, f"Agent '{agent_id}' not found", 404 - ) + return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404) is_allowed = await AgentRequestHandler.is_agent_allowed( agent_id=agent.agent_id, @@ -622,13 +590,9 @@ async def invoke_agent_a2a( # URL is required unless using completion bridge with a provider that derives endpoint from model # (e.g., bedrock/agentcore derives endpoint from ARN in model string) if not agent_url and not custom_llm_provider: - return _jsonrpc_error( - request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500 - ) + return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500) - verbose_proxy_logger.info( - f"Proxying A2A request to agent '{agent_id}' at {agent_url or 'completion-bridge'}" - ) + verbose_proxy_logger.info(f"Proxying A2A request to agent '{agent_id}' at {agent_url or 'completion-bridge'}") # Set up data dict for litellm processing if "metadata" not in body: @@ -713,9 +677,7 @@ async def invoke_agent_a2a( _existing_guardrails: List = data.get("guardrails") or [] if not isinstance(_existing_guardrails, list): _existing_guardrails = [_existing_guardrails] - data["guardrails"] = _existing_guardrails + [ - g for g in _agent_guardrails if g not in _existing_guardrails - ] + data["guardrails"] = _existing_guardrails + [g for g in _agent_guardrails if g not in _existing_guardrails] # Route through SDK functions if method == "message/send": @@ -794,9 +756,7 @@ async def invoke_agent_a2a( "agent/getAuthenticatedExtendedCard", }: if not agent_url: - return _jsonrpc_error( - request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500 - ) + return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500) if method == "tasks/pushNotificationConfig/set": if not isinstance(params, dict): raise HTTPException( @@ -804,9 +764,7 @@ async def invoke_agent_a2a( detail="params must be an object", ) push_config = params.get("pushNotificationConfig", {}) - if "pushNotificationConfig" in params and not isinstance( - push_config, dict - ): + if "pushNotificationConfig" in params and not isinstance(push_config, dict): raise HTTPException( status_code=400, detail="pushNotificationConfig must be an object", @@ -831,19 +789,13 @@ async def invoke_agent_a2a( request_data=data, agent_extra_headers=agent_extra_headers, ) - result = await _forward_jsonrpc( - agent_url, forward_body, extra_headers=caller_headers - ) + result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=caller_headers) if method == "agent/getAuthenticatedExtendedCard": if isinstance(result.get("result"), dict) and "url" in result["result"]: - result["result"]["url"] = ( - f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}" - ) + result["result"]["url"] = f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}" from litellm.types.agents import LiteLLMSendMessageResponse - response = LiteLLMSendMessageResponse.from_dict( - result, request_id=request_id - ) + response = LiteLLMSendMessageResponse.from_dict(result, request_id=request_id) response = await proxy_logging_obj.post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, @@ -851,17 +803,13 @@ async def invoke_agent_a2a( ) return JSONResponse( content=( - response.model_dump(mode="json", exclude_none=True) - if hasattr(response, "model_dump") - else response + response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response ) ) elif method == "tasks/resubscribe": if not agent_url: - return _jsonrpc_error( - request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500 - ) + return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500) forward_body = { "jsonrpc": "2.0", "id": request_id, diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index c1f9c89529b..1373d055d4f 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -24,15 +24,11 @@ class AgentRegistry: self.agent_list.append(agent_config) def deregister_agent(self, agent_name: str): - self.agent_list = [ - agent for agent in self.agent_list if agent.agent_name != agent_name - ] + self.agent_list = [agent for agent in self.agent_list if agent.agent_name != agent_name] def get_agent_list(self, agent_names: Optional[List[str]] = None): if agent_names is not None: - return [ - agent for agent in self.agent_list if agent.agent_name in agent_names - ] + return [agent for agent in self.agent_list if agent.agent_name in agent_names] return self.agent_list def get_public_agent_list(self) -> List[AgentResponse]: @@ -45,9 +41,7 @@ class AgentRegistry: return public_agent_list def _create_agent_id(self, agent_config: AgentConfig) -> str: - return hashlib.sha256( - json.dumps(agent_config, sort_keys=True).encode() - ).hexdigest() + return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() def load_agents_from_config(self, agent_config: Optional[List[AgentConfig]] = None): if agent_config is None: @@ -65,9 +59,7 @@ class AgentRegistry: # create a stable hash id for config item config_hash = self._create_agent_id(agent_config_item) - self.register_agent( - agent_config=AgentResponse(agent_id=config_hash, **agent_config_item) - ) # type: ignore + self.register_agent(agent_config=AgentResponse(agent_id=config_hash, **agent_config_item)) # type: ignore def load_agents_from_db_and_config( self, @@ -122,9 +114,7 @@ class AgentRegistry: if hasattr(litellm_params_obj, "model_dump"): litellm_params_dict = litellm_params_obj.model_dump() else: - litellm_params_dict = ( - dict(litellm_params_obj) if litellm_params_obj else {} - ) + litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} litellm_params: str = safe_dumps(litellm_params_dict) # Serialize agent_card_params @@ -132,24 +122,18 @@ class AgentRegistry: if hasattr(agent_card_params_obj, "model_dump"): agent_card_params_dict = agent_card_params_obj.model_dump() else: - agent_card_params_dict = ( - dict(agent_card_params_obj) if agent_card_params_obj else {} - ) + agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} agent_card_params: str = safe_dumps(agent_card_params_dict) # Handle object_permission (MCP tool access for agent) object_permission_id: Optional[str] = None if agent.get("object_permission") is not None: agent_copy = dict(agent) - object_permission_id = await handle_update_object_permission_common( - agent_copy, None, prisma_client - ) + object_permission_id = await handle_update_object_permission_common(agent_copy, None, prisma_client) # Serialize static_headers static_headers_obj = agent.get("static_headers") - static_headers_val: Optional[str] = ( - safe_dumps(dict(static_headers_obj)) if static_headers_obj else None - ) + static_headers_val: Optional[str] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None extra_headers_val: Optional[List[str]] = agent.get("extra_headers") @@ -190,27 +174,19 @@ class AgentRegistry: created_agent_dict = created_agent.model_dump() if created_agent.object_permission is not None: try: - created_agent_dict["object_permission"] = ( - created_agent.object_permission.model_dump() - ) + created_agent_dict["object_permission"] = created_agent.object_permission.model_dump() except Exception: - created_agent_dict["object_permission"] = ( - created_agent.object_permission.dict() - ) + created_agent_dict["object_permission"] = created_agent.object_permission.dict() return AgentResponse(**created_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") - async def delete_agent_from_db( - self, agent_id: str, prisma_client: PrismaClient - ) -> Dict[str, Any]: + async def delete_agent_from_db(self, agent_id: str, prisma_client: PrismaClient) -> Dict[str, Any]: """ Delete an agent from the database """ try: - deleted_agent = await AgentsRepository(prisma_client).table.delete( - where={"agent_id": agent_id} - ) + deleted_agent = await AgentsRepository(prisma_client).table.delete(where={"agent_id": agent_id}) return dict(deleted_agent) except Exception as e: raise Exception(f"Error deleting agent from DB: {str(e)}") @@ -237,9 +213,7 @@ class AgentRegistry: The patched agent """ try: - existing_agent = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id} - ) + existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) if existing_agent is not None: existing_agent = dict(existing_agent) @@ -251,13 +225,9 @@ class AgentRegistry: if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): - update_data["litellm_params"] = safe_dumps( - augment_agent.get("litellm_params") - ) + update_data["litellm_params"] = safe_dumps(augment_agent.get("litellm_params")) if augment_agent.get("agent_card_params"): - update_data["agent_card_params"] = safe_dumps( - augment_agent.get("agent_card_params") - ) + update_data["agent_card_params"] = safe_dumps(augment_agent.get("agent_card_params")) for rate_field in ( "tpm_limit", @@ -269,19 +239,13 @@ class AgentRegistry: update_data[rate_field] = agent.get(rate_field) if "static_headers" in agent: headers_value = agent.get("static_headers") - update_data["static_headers"] = safe_dumps( - dict(headers_value) if headers_value is not None else {} - ) + update_data["static_headers"] = safe_dumps(dict(headers_value) if headers_value is not None else {}) if "extra_headers" in agent: extra_headers_value = agent.get("extra_headers") - update_data["extra_headers"] = ( - extra_headers_value if extra_headers_value is not None else [] - ) + update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else [] if agent.get("object_permission") is not None: agent_copy = dict(augment_agent) - existing_object_permission_id = existing_agent.get( - "object_permission_id" - ) + existing_object_permission_id = existing_agent.get("object_permission_id") object_permission_id = await handle_update_object_permission_common( agent_copy, existing_object_permission_id, @@ -302,13 +266,9 @@ class AgentRegistry: patched_agent_dict = patched_agent.model_dump() if patched_agent.object_permission is not None: try: - patched_agent_dict["object_permission"] = ( - patched_agent.object_permission.model_dump() - ) + patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump() except Exception: - patched_agent_dict["object_permission"] = ( - patched_agent.object_permission.dict() - ) + patched_agent_dict["object_permission"] = patched_agent.object_permission.dict() return AgentResponse(**patched_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error patching agent in DB: {str(e)}") @@ -331,9 +291,7 @@ class AgentRegistry: if hasattr(litellm_params_obj, "model_dump"): litellm_params_dict = litellm_params_obj.model_dump() else: - litellm_params_dict = ( - dict(litellm_params_obj) if litellm_params_obj else {} - ) + litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} litellm_params: str = safe_dumps(litellm_params_dict) # Serialize agent_card_params @@ -341,17 +299,13 @@ class AgentRegistry: if hasattr(agent_card_params_obj, "model_dump"): agent_card_params_dict = agent_card_params_obj.model_dump() else: - agent_card_params_dict = ( - dict(agent_card_params_obj) if agent_card_params_obj else {} - ) + agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} agent_card_params: str = safe_dumps(agent_card_params_dict) # Serialize static_headers for update static_headers_obj_u = agent.get("static_headers") static_headers_val_u: str = ( - safe_dumps(dict(static_headers_obj_u)) - if static_headers_obj_u is not None - else safe_dumps({}) + safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({}) ) extra_headers_val_u: List[str] = agent.get("extra_headers") or [] @@ -376,13 +330,9 @@ class AgentRegistry: update_data[rate_field] = _val if agent.get("object_permission") is not None: - existing_agent = await AgentsRepository( - prisma_client - ).table.find_unique(where={"agent_id": agent_id}) + existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) existing_object_permission_id = ( - existing_agent.object_permission_id - if existing_agent is not None - else None + existing_agent.object_permission_id if existing_agent is not None else None ) agent_copy = dict(agent) object_permission_id = await handle_update_object_permission_common( @@ -403,13 +353,9 @@ class AgentRegistry: updated_agent_dict = updated_agent.model_dump() if updated_agent.object_permission is not None: try: - updated_agent_dict["object_permission"] = ( - updated_agent.object_permission.model_dump() - ) + updated_agent_dict["object_permission"] = updated_agent.object_permission.model_dump() except Exception: - updated_agent_dict["object_permission"] = ( - updated_agent.object_permission.dict() - ) + updated_agent_dict["object_permission"] = updated_agent.object_permission.dict() return AgentResponse(**updated_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error updating agent in DB: {str(e)}") @@ -433,9 +379,7 @@ class AgentRegistry: # object_permission is eagerly loaded via include above if agent.object_permission is not None: try: - agent_dict["object_permission"] = ( - agent.object_permission.model_dump() - ) + agent_dict["object_permission"] = agent.object_permission.model_dump() except Exception: agent_dict["object_permission"] = agent.object_permission.dict() agents.append(agent_dict) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 2577615fc8e..24f70351ddc 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -43,14 +43,8 @@ class AgentRequestHandler: """ try: allowed_agents: List[str] = [] - allowed_agents_for_key = ( - await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) - ) - allowed_agents_for_team = ( - await AgentRequestHandler._get_allowed_agents_for_team( - user_api_key_auth - ) - ) + allowed_agents_for_key = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) + allowed_agents_for_team = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) # If team has agent restrictions, handle inheritance and intersection logic if len(allowed_agents_for_team) > 0: @@ -161,18 +155,14 @@ class AgentRequestHandler: all_agents: List[str] = [] # 1. Get agents from object_permission (native permissions) - key_object_permission = AgentRequestHandler._get_key_object_permission( - user_api_key_auth - ) + key_object_permission = AgentRequestHandler._get_key_object_permission(user_api_key_auth) if key_object_permission is not None: # Get direct agents direct_agents = key_object_permission.agents or [] # Get agents from access groups - access_group_agents = ( - await AgentRequestHandler._get_agents_from_access_groups( - key_object_permission.agent_access_groups or [] - ) + access_group_agents = await AgentRequestHandler._get_agents_from_access_groups( + key_object_permission.agent_access_groups or [] ) all_agents = direct_agents + access_group_agents @@ -244,10 +234,8 @@ class AgentRequestHandler: direct_agents = object_permissions.agents or [] # Get agents from access groups - access_group_agents = ( - await AgentRequestHandler._get_agents_from_access_groups( - object_permissions.agent_access_groups or [] - ) + access_group_agents = await AgentRequestHandler._get_agents_from_access_groups( + object_permissions.agent_access_groups or [] ) all_agents = direct_agents + access_group_agents @@ -269,15 +257,11 @@ class AgentRequestHandler: # litellm-dashboard is the default UI team and will never have agents; # skip noisy warnings for it. if user_api_key_auth.team_id != UI_TEAM_ID: - verbose_logger.warning( - f"Failed to get allowed agents for team: {str(e)}" - ) + verbose_logger.warning(f"Failed to get allowed agents for team: {str(e)}") return [] @staticmethod - def _get_config_agent_ids_for_access_groups( - config_agents: List, access_groups: List[str] - ) -> Set[str]: + def _get_config_agent_ids_for_access_groups(config_agents: List, access_groups: List[str]) -> Set[str]: """ Helper to get agent_ids from config-loaded agents that match any of the given access groups. """ @@ -290,9 +274,7 @@ class AgentRequestHandler: return server_ids @staticmethod - async def _get_db_agent_ids_for_access_groups( - prisma_client, access_groups: List[str] - ) -> Set[str]: + async def _get_db_agent_ids_for_access_groups(prisma_client, access_groups: List[str]) -> Set[str]: """ Helper to get agent_ids from DB agents that match any of the given access groups. """ @@ -325,11 +307,7 @@ class AgentRequestHandler: ) # Use the helper for DB agents - db_agent_ids = ( - await AgentRequestHandler._get_db_agent_ids_for_access_groups( - prisma_client, access_groups - ) - ) + db_agent_ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups(prisma_client, access_groups) agent_ids.update(db_agent_ids) return list(agent_ids) @@ -345,16 +323,8 @@ class AgentRequestHandler: Get list of agent access groups for the given user/key based on permissions. """ access_groups: List[str] = [] - access_groups_for_key = ( - await AgentRequestHandler._get_agent_access_groups_for_key( - user_api_key_auth - ) - ) - access_groups_for_team = ( - await AgentRequestHandler._get_agent_access_groups_for_team( - user_api_key_auth - ) - ) + access_groups_for_key = await AgentRequestHandler._get_agent_access_groups_for_key(user_api_key_auth) + access_groups_for_team = await AgentRequestHandler._get_agent_access_groups_for_team(user_api_key_auth) # If team has access groups, then key must have a subset of the team's access groups if len(access_groups_for_team) > 0: @@ -401,9 +371,7 @@ class AgentRequestHandler: return key_object_permission.agent_access_groups or [] except Exception as e: - verbose_logger.warning( - f"Failed to get agent access groups for key: {str(e)}" - ) + verbose_logger.warning(f"Failed to get agent access groups for key: {str(e)}") return [] @staticmethod @@ -446,7 +414,5 @@ class AgentRequestHandler: return object_permissions.agent_access_groups or [] except Exception as e: - verbose_logger.warning( - f"Failed to get agent access groups for team: {str(e)}" - ) + verbose_logger.warning(f"Failed to get agent access groups for team: {str(e)}") return [] diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 1c1f5a2b4c4..1c01f789916 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -91,10 +91,7 @@ def parse_databricks_oauth_config( if raw is None: return None if not isinstance(raw, dict): - raise ValueError( - f"'{DATABRICKS_OAUTH_PARAM}' must be a mapping of OAuth settings, " - f"got {type(raw).__name__}" - ) + raise ValueError(f"'{DATABRICKS_OAUTH_PARAM}' must be a mapping of OAuth settings, got {type(raw).__name__}") client_id = _resolve_secret(raw.get("client_id")) client_secret = _resolve_secret(raw.get("client_secret")) @@ -110,10 +107,7 @@ def parse_databricks_oauth_config( if not value ] if missing: - raise ValueError( - f"Databricks App OAuth config is missing required field(s): " - f"{', '.join(missing)}" - ) + raise ValueError(f"Databricks App OAuth config is missing required field(s): {', '.join(missing)}") scope = _resolve_secret(raw.get("scope")) or _DEFAULT_SCOPE @@ -176,13 +170,9 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): async def _fetch_token(self, config: DatabricksAppOAuthConfig) -> Tuple[str, int]: client = get_async_httpx_client(llm_provider=httpxSpecialProvider.A2A) - verbose_logger.debug( - "Fetching Databricks App OAuth token from %s", config.token_url - ) + verbose_logger.debug("Fetching Databricks App OAuth token from %s", config.token_url) - basic_auth = base64.b64encode( - f"{config.client_id}:{config.client_secret}".encode() - ).decode() + basic_auth = base64.b64encode(f"{config.client_id}:{config.client_secret}".encode()).decode() try: response = await client.post( config.token_url, @@ -197,34 +187,24 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): ) except httpx.HTTPStatusError as exc: raise ValueError( - "Databricks App OAuth token request failed with status " - f"{exc.response.status_code}" + f"Databricks App OAuth token request failed with status {exc.response.status_code}" ) from exc except httpx.HTTPError as exc: - raise ValueError( - f"Databricks App OAuth token request failed: {exc}" - ) from exc + raise ValueError(f"Databricks App OAuth token request failed: {exc}") from exc body = response.json() if not isinstance(body, dict): raise ValueError( - "Databricks App OAuth token response returned non-object JSON " - f"(got {type(body).__name__})" + f"Databricks App OAuth token response returned non-object JSON (got {type(body).__name__})" ) access_token = body.get("access_token") if not access_token: - raise ValueError( - "Databricks App OAuth token response missing 'access_token'" - ) + raise ValueError("Databricks App OAuth token response missing 'access_token'") raw_expires_in = body.get("expires_in") try: - expires_in = ( - int(raw_expires_in) - if raw_expires_in is not None - else _DEFAULT_TTL_SECONDS - ) + expires_in = int(raw_expires_in) if raw_expires_in is not None else _DEFAULT_TTL_SECONDS except (TypeError, ValueError): expires_in = _DEFAULT_TTL_SECONDS diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index d19008856bd..ebbc26e6a9a 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -107,12 +107,8 @@ def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> Non ) -AGENT_HEALTH_CHECK_TIMEOUT_SECONDS = float( - os.environ.get("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", "5.0") -) -AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS = float( - os.environ.get("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", "30.0") -) +AGENT_HEALTH_CHECK_TIMEOUT_SECONDS = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", "5.0")) +AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", "30.0")) async def _check_agent_url_health( @@ -199,9 +195,7 @@ async def get_agents( returned_agents = global_agent_registry.get_agent_list() else: # Get allowed agents from object_permission (key/team level) - allowed_agent_ids = await AgentRequestHandler.get_allowed_agents( - user_api_key_auth=user_api_key_dict - ) + allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) # If no restrictions (empty list), return all agents if len(allowed_agent_ids) == 0: @@ -209,9 +203,7 @@ async def get_agents( else: # Filter agents by allowed IDs all_agents = global_agent_registry.get_agent_list() - returned_agents = [ - agent for agent in all_agents if agent.agent_id in allowed_agent_ids - ] + returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] # Fetch current spend from DB for all returned agents from litellm.proxy.proxy_server import prisma_client @@ -231,9 +223,8 @@ async def get_agents( for agent in returned_agents: if agent.litellm_params is None: agent.litellm_params = {} - agent.litellm_params["is_public"] = ( - litellm.public_agent_groups is not None - and (agent.agent_id in litellm.public_agent_groups) + agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and ( + agent.agent_id in litellm.public_agent_groups ) # Redact sensitive fields for non-admin users @@ -245,21 +236,11 @@ async def get_agents( returned_agents = _redact_sensitive_agent_fields(returned_agents) if health_check: - agents_with_url = [ - agent - for agent in returned_agents - if (agent.agent_card_params or {}).get("url") - ] - agents_without_url = [ - agent - for agent in returned_agents - if not (agent.agent_card_params or {}).get("url") - ] + agents_with_url = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")] + agents_without_url = [agent for agent in returned_agents if not (agent.agent_card_params or {}).get("url")] try: health_results = await asyncio.wait_for( - asyncio.gather( - *[_check_agent_url_health(agent) for agent in agents_with_url] - ), + asyncio.gather(*[_check_agent_url_health(agent) for agent in agents_with_url]), timeout=AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, ) except asyncio.TimeoutError: @@ -275,25 +256,17 @@ async def get_agents( } for agent in agents_with_url ] - healthy_ids = { - result["agent_id"] for result in health_results if result["healthy"] - } - returned_agents = [ - agent for agent in agents_with_url if agent.agent_id in healthy_ids - ] + agents_without_url + healthy_ids = {result["agent_id"] for result in health_results if result["healthy"]} + returned_agents = [agent for agent in agents_with_url if agent.agent_id in healthy_ids] + agents_without_url return returned_agents except HTTPException: raise except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.agent_endpoints.get_agents(): Exception occurred - {}".format( - str(e) - ) - ) - raise HTTPException( - status_code=500, detail={"error": f"Internal server error: {str(e)}"} + "litellm.proxy.agent_endpoints.get_agents(): Exception occurred - {}".format(str(e)) ) + raise HTTPException(status_code=500, detail={"error": f"Internal server error: {str(e)}"}) #### CRUD ENDPOINTS FOR AGENTS #### @@ -408,9 +381,7 @@ async def create_agent( # Also register in memory try: AGENT_REGISTRY.register_agent(agent_config=result) - verbose_proxy_logger.info( - f"Successfully registered agent '{agent_name}' (ID: {agent_id}) in memory" - ) + verbose_proxy_logger.info(f"Successfully registered agent '{agent_name}' (ID: {agent_id}) in memory") except Exception as reg_error: verbose_proxy_logger.warning( f"Failed to register agent '{agent_name}' (ID: {agent_id}) in memory: {reg_error}" @@ -455,9 +426,7 @@ async def get_agent_by_id( AgentRequestHandler, ) - is_allowed = await AgentRequestHandler.is_agent_allowed( - agent_id=agent_id, user_api_key_auth=user_api_key_dict - ) + is_allowed = await AgentRequestHandler.is_agent_allowed(agent_id=agent_id, user_api_key_auth=user_api_key_dict) if not is_allowed: raise HTTPException( status_code=403, @@ -480,26 +449,18 @@ async def get_agent_by_id( agent_dict = agent_row.model_dump() if agent_row.object_permission is not None: try: - agent_dict["object_permission"] = ( - agent_row.object_permission.model_dump() - ) + agent_dict["object_permission"] = agent_row.object_permission.model_dump() except Exception: - agent_dict["object_permission"] = ( - agent_row.object_permission.dict() - ) + agent_dict["object_permission"] = agent_row.object_permission.dict() agent = AgentResponse(**agent_dict) # type: ignore else: # Agent found in memory — refresh spend from DB - db_row = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id} - ) + db_row = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) if db_row is not None: agent.spend = db_row.spend if agent is None: - raise HTTPException( - status_code=404, detail=f"Agent with ID {agent_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") # Redact sensitive fields for non-admin users is_admin = ( @@ -565,22 +526,16 @@ async def update_agent( _check_agent_management_permission(user_api_key_dict) if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Check if agent exists - existing_agent = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id} - ) + existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) if existing_agent is not None: existing_agent = dict(existing_agent) if existing_agent is None: - raise HTTPException( - status_code=404, detail=f"Agent with ID {agent_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") # Get the user ID from the API key auth updated_by = user_api_key_dict.user_id or "unknown" @@ -673,22 +628,16 @@ async def patch_agent( _check_agent_management_permission(user_api_key_dict) if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Check if agent exists - existing_agent = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id} - ) + existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) if existing_agent is not None: existing_agent = dict(existing_agent) if existing_agent is None: - raise HTTPException( - status_code=404, detail=f"Agent with ID {agent_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") # Get the user ID from the API key auth updated_by = user_api_key_dict.user_id or "unknown" @@ -770,20 +719,14 @@ async def delete_agent( try: # Check if agent exists - existing_agent = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id} - ) + existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) if existing_agent is not None: existing_agent = dict[Any, Any](existing_agent) if existing_agent is None: - raise HTTPException( - status_code=404, detail=f"Agent with ID {agent_id} not found in DB." - ) + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.") - await AGENT_REGISTRY.delete_agent_from_db( - agent_id=agent_id, prisma_client=prisma_client - ) + await AGENT_REGISTRY.delete_agent_from_db(agent_id=agent_id, prisma_client=prisma_client) AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) # type: ignore @@ -834,9 +777,7 @@ async def make_agent_public( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Update the public model groups @@ -860,16 +801,12 @@ async def make_agent_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id} - ) + agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) if agent is not None: agent = AgentResponse(**agent.model_dump()) # type: ignore if agent is None: - raise HTTPException( - status_code=404, detail=f"Agent with ID {agent_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") if litellm.public_agent_groups is None: litellm.public_agent_groups = [] @@ -951,9 +888,7 @@ async def make_agents_public( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Update the public model groups @@ -983,16 +918,12 @@ async def make_agents_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": agent_id} - ) + agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) if agent is not None: agent = AgentResponse(**agent.model_dump()) # type: ignore if agent is None: - raise HTTPException( - status_code=404, detail=f"Agent with ID {agent_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") litellm.public_agent_groups = request.agent_ids @@ -1054,9 +985,7 @@ async def get_agent_daily_activity( agent_ids_list = agent_ids.split(",") if agent_ids else None exclude_agent_ids_list: Optional[List[str]] = None if exclude_agent_ids: - exclude_agent_ids_list = ( - exclude_agent_ids.split(",") if exclude_agent_ids else None - ) + exclude_agent_ids_list = exclude_agent_ids.split(",") if exclude_agent_ids else None # Without scoping, an empty `agent_ids` query returned every agent's # spend/token rows on the proxy. Restrict non-admin callers to the @@ -1069,9 +998,7 @@ async def get_agent_daily_activity( where_condition: Dict[str, Any] = {} if not _user_has_admin_view(user_api_key_dict): - permitted_agent_ids = await AgentRequestHandler.get_allowed_agents( - user_api_key_auth=user_api_key_dict - ) + permitted_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) # `get_allowed_agents` returns an empty list when the caller's key # and team carry no agent restrictions. For activity scoping that's # not "see everything" — fall back to the agents the caller @@ -1090,9 +1017,7 @@ async def get_agent_daily_activity( if agent_ids_list: permitted_agent_id_set = set(permitted_agent_ids) - agent_ids_list = [ - aid for aid in agent_ids_list if aid in permitted_agent_id_set - ] + agent_ids_list = [aid for aid in agent_ids_list if aid in permitted_agent_id_set] else: agent_ids_list = list(permitted_agent_ids) @@ -1119,12 +1044,8 @@ async def get_agent_daily_activity( if agent_ids_list: where_condition["agent_id"] = {"in": list(agent_ids_list)} - agent_records = await AgentsRepository(prisma_client).table.find_many( - where=where_condition - ) - agent_metadata = { - agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records - } + agent_records = await AgentsRepository(prisma_client).table.find_many(where=where_condition) + agent_metadata = {agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records} return await get_daily_activity( prisma_client=prisma_client, diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index b88c602ac34..d8e2639521f 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -29,9 +29,7 @@ async def append_agents_to_model_group( AgentRequestHandler, ) - allowed_agent_ids = await AgentRequestHandler.get_allowed_agents( - user_api_key_auth=user_api_key_dict - ) + allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) @@ -65,9 +63,7 @@ async def append_agents_to_model_info( AgentRequestHandler, ) - allowed_agent_ids = await AgentRequestHandler.get_allowed_agents( - user_api_key_auth=user_api_key_dict - ) + allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) diff --git a/litellm/proxy/analytics_endpoints/analytics_endpoints.py b/litellm/proxy/analytics_endpoints/analytics_endpoints.py index 6835f0c9095..4c1ff31e5a1 100644 --- a/litellm/proxy/analytics_endpoints/analytics_endpoints.py +++ b/litellm/proxy/analytics_endpoints/analytics_endpoints.py @@ -58,9 +58,7 @@ async def get_global_activity( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -93,9 +91,7 @@ async def get_global_activity( sl."call_type", sl."model" """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) if db_response is None: return [] diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index dd7350e13ce..7d69e09b8da 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -72,25 +72,19 @@ async def get_marketplace(): try: prisma_client = await _get_prisma_client() - plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( - where={"enabled": True} - ) + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many(where={"enabled": True}) plugin_list = [] for plugin in plugins: try: manifest = json.loads(plugin.manifest_json) except json.JSONDecodeError: - verbose_proxy_logger.warning( - f"Plugin {plugin.name} has invalid manifest JSON, skipping" - ) + verbose_proxy_logger.warning(f"Plugin {plugin.name} has invalid manifest JSON, skipping") continue # Source must be specified for URL-based marketplaces if "source" not in manifest: - verbose_proxy_logger.warning( - f"Plugin {plugin.name} has no source field, skipping" - ) + verbose_proxy_logger.warning(f"Plugin {plugin.name} has no source field, skipping") continue entry: Dict[str, Any] = { @@ -135,9 +129,7 @@ async def get_marketplace(): # Each segment must start with an alphanumeric character and contain only # alphanumeric characters, dots, hyphens, and underscores. # This implicitly blocks '..', leading '/', backslashes, and percent-encoded sequences. -_VALID_GIT_SUBDIR_PATH_RE = re.compile( - r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-zA-Z0-9][a-zA-Z0-9._-]*)*$" -) +_VALID_GIT_SUBDIR_PATH_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-zA-Z0-9][a-zA-Z0-9._-]*)*$") def _validate_plugin_source(source: Dict[str, Any]) -> None: @@ -147,17 +139,13 @@ def _validate_plugin_source(source: Dict[str, Any]) -> None: if "repo" not in source: raise HTTPException( status_code=400, - detail={ - "error": "GitHub source must include 'repo' field (e.g., 'org/repo')" - }, + detail={"error": "GitHub source must include 'repo' field (e.g., 'org/repo')"}, ) elif source_type == "url": if "url" not in source: raise HTTPException( status_code=400, - detail={ - "error": "URL source must include 'url' field (e.g., 'https://github.com/org/repo.git')" - }, + detail={"error": "URL source must include 'url' field (e.g., 'https://github.com/org/repo.git')"}, ) elif source_type == "git-subdir": if not source.get("url"): @@ -170,9 +158,7 @@ def _validate_plugin_source(source: Dict[str, Any]) -> None: if not source.get("path"): raise HTTPException( status_code=400, - detail={ - "error": "git-subdir source must include 'path' field (e.g., 'plugins/plugin-name')" - }, + detail={"error": "git-subdir source must include 'path' field (e.g., 'plugins/plugin-name')"}, ) if not _VALID_GIT_SUBDIR_PATH_RE.match(source["path"]): raise HTTPException( @@ -237,9 +223,7 @@ async def register_plugin( if not re.match(r"^[a-z0-9-]+$", request.name): raise HTTPException( status_code=400, - detail={ - "error": "Plugin name must be kebab-case (lowercase letters, numbers, hyphens)" - }, + detail={"error": "Plugin name must be kebab-case (lowercase letters, numbers, hyphens)"}, ) # Validate source format @@ -269,9 +253,7 @@ async def register_plugin( manifest["namespace"] = request.namespace # Check if plugin exists - existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique( - where={"name": request.name} - ) + existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": request.name}) if existing: plugin = await ClaudeCodePluginRepository(prisma_client).table.update( @@ -349,9 +331,7 @@ async def list_plugins( prisma_client = await _get_prisma_client() where = {"enabled": True} if enabled_only else {} - plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( - where=where - ) + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many(where=where) plugin_list = [] for p in plugins: @@ -416,9 +396,7 @@ async def get_plugin( try: prisma_client = await _get_prisma_client() - plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( - where={"name": plugin_name} - ) + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) if not plugin: raise HTTPException( @@ -472,9 +450,7 @@ async def enable_plugin( try: prisma_client = await _get_prisma_client() - plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( - where={"name": plugin_name} - ) + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) if not plugin: raise HTTPException( status_code=404, @@ -517,9 +493,7 @@ async def disable_plugin( try: prisma_client = await _get_prisma_client() - plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( - where={"name": plugin_name} - ) + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) if not plugin: raise HTTPException( status_code=404, @@ -562,18 +536,14 @@ async def delete_plugin( try: prisma_client = await _get_prisma_client() - plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( - where={"name": plugin_name} - ) + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) if not plugin: raise HTTPException( status_code=404, detail={"error": f"Plugin '{plugin_name}' not found"}, ) - await ClaudeCodePluginRepository(prisma_client).table.delete( - where={"name": plugin_name} - ) + await ClaudeCodePluginRepository(prisma_client).table.delete(where={"name": plugin_name}) verbose_proxy_logger.info(f"Plugin {plugin_name} deleted") return {"status": "success", "message": f"Plugin '{plugin_name}' deleted"} diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 856b788b54b..71acc1f3106 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -149,13 +149,11 @@ async def anthropic_response( async def _passthrough_stream_generator(): yield _anthropic_response - selected_data_generator = ( - ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=_passthrough_stream_generator(), - user_api_key_dict=user_api_key_dict, - request_data=_data, - proxy_logging_obj=proxy_logging_obj, - ) + selected_data_generator = ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=_passthrough_stream_generator(), + user_api_key_dict=user_api_key_dict, + request_data=_data, + proxy_logging_obj=proxy_logging_obj, ) return await create_response( @@ -185,9 +183,7 @@ async def anthropic_response( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.anthropic_response(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.anthropic_response(): Exception occured - {}".format(str(e)) ) # Extract model_id from request metadata (same as success path) @@ -258,14 +254,10 @@ async def count_tokens( messages = data.get("messages", []) if not model_name: - raise HTTPException( - status_code=400, detail={"error": "model parameter is required"} - ) + raise HTTPException(status_code=400, detail={"error": "model parameter is required"}) if not messages: - raise HTTPException( - status_code=400, detail={"error": "messages parameter is required"} - ) + raise HTTPException(status_code=400, detail={"error": "messages parameter is required"}) # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest @@ -305,13 +297,9 @@ async def count_tokens( ) except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - {}".format( - str(e) - ) - ) - raise HTTPException( - status_code=500, detail={"error": f"Internal server error: {str(e)}"} + "litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - {}".format(str(e)) ) + raise HTTPException(status_code=500, detail={"error": f"Internal server error: {str(e)}"}) @router.post( diff --git a/litellm/proxy/anthropic_endpoints/skills_endpoints.py b/litellm/proxy/anthropic_endpoints/skills_endpoints.py index cd19e7731f0..772006a7fa1 100644 --- a/litellm/proxy/anthropic_endpoints/skills_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/skills_endpoints.py @@ -84,11 +84,7 @@ async def create_skill( data = await convert_upload_files_to_file_data(form_data) # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -191,11 +187,7 @@ async def list_skills( data["before_id"] = before_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -292,11 +284,7 @@ async def get_skill( data["skill_id"] = skill_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -395,11 +383,7 @@ async def delete_skill( data["skill_id"] = skill_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b16b2defdab..8858289ad44 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -124,13 +124,8 @@ def _log_budget_lookup_failure(entity: str, error: Exception) -> None: return err_str = str(error).lower() hint = "" - if any( - x in err_str - for x in ("column", "schema", "does not exist", "prisma", "migrate") - ): - hint = ( - " Run `prisma db push` or `prisma migrate deploy` to fix schema mismatches." - ) + if any(x in err_str for x in ("column", "schema", "does not exist", "prisma", "migrate")): + hint = " Run `prisma db push` or `prisma migrate deploy` to fix schema mismatches." verbose_proxy_logger.error( f"Budget lookup failed for {entity}; cache will not be populated. " f"Each request will hit the database. Error: {error}.{hint}" @@ -153,9 +148,7 @@ def _get_router_zero_cost_cache(llm_router: Router) -> Optional[Dict[str, bool]] return cache if isinstance(cache, dict) else None -def _is_model_cost_zero( - model: Optional[Union[str, List[str]]], llm_router: Optional[Router] -) -> bool: +def _is_model_cost_zero(model: Optional[Union[str, List[str]]], llm_router: Optional[Router]) -> bool: """ Check if a model has zero cost (no configured pricing). @@ -190,9 +183,7 @@ def _is_model_cost_zero( if model_group_info is None: # Model not found or no pricing info available # Conservative approach: assume it has cost - verbose_proxy_logger.debug( - f"No model group info found for {model_name}, assuming it has cost" - ) + verbose_proxy_logger.debug(f"No model group info found for {model_name}, assuming it has cost") if zero_cost_cache is not None: zero_cost_cache[model_name] = False return False @@ -246,9 +237,7 @@ def _is_model_cost_zero( except Exception as e: # If we can't determine the cost, assume it has cost (conservative approach) - verbose_proxy_logger.debug( - f"Error checking cost for model {model_name}: {str(e)}, assuming it has cost" - ) + verbose_proxy_logger.debug(f"Error checking cost for model {model_name}: {str(e)}, assuming it has cost") return False # All models checked have zero cost @@ -322,36 +311,22 @@ async def _run_project_checks( ) -def _enforce_user_param_check( - general_settings: dict, request: Request, request_body: dict, route: str -) -> None: +def _enforce_user_param_check(general_settings: dict, request: Request, request_body: dict, route: str) -> None: if not general_settings.get("enforce_user_param", False): return http_method = request.method if hasattr(request, "method") else None is_post_method = http_method and http_method.upper() == "POST" is_openai_route = RouteChecks.is_llm_api_route(route=route) - is_mcp_route = ( - route in LiteLLMRoutes.mcp_routes.value - or RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value - ) + is_mcp_route = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value ) - if ( - is_post_method - and is_openai_route - and not is_mcp_route - and "user" not in request_body - ): - raise Exception( - f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" - ) + if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body: + raise Exception(f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}") -def _reject_clientside_metadata_tags_check( - general_settings: dict, request_body: dict, route: str -) -> None: +def _reject_clientside_metadata_tags_check(general_settings: dict, request_body: dict, route: str) -> None: if not general_settings.get("reject_clientside_metadata_tags", False): return @@ -369,9 +344,7 @@ def _reject_clientside_metadata_tags_check( ) -def _global_proxy_budget_check( - global_proxy_spend: Optional[float], skip_budget_checks: bool, route: str -) -> None: +def _global_proxy_budget_check(global_proxy_spend: Optional[float], skip_budget_checks: bool, route: str) -> None: if ( litellm.max_budget > 0 and not skip_budget_checks @@ -380,13 +353,8 @@ def _global_proxy_budget_check( and route != "/v1/models" and route != "/models" ): - if ( - math.isfinite(litellm.max_budget) - and global_proxy_spend > litellm.max_budget - ): - raise litellm.BudgetExceededError( - current_cost=global_proxy_spend, max_budget=litellm.max_budget - ) + if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget: + raise litellm.BudgetExceededError(current_cost=global_proxy_spend, max_budget=litellm.max_budget) _GUARDRAIL_MODIFICATION_KEYS: tuple = ( @@ -397,9 +365,7 @@ _GUARDRAIL_MODIFICATION_KEYS: tuple = ( ) -def _guardrail_modification_check( - request_body: dict, team_object: Optional[LiteLLM_TeamTable] -) -> None: +def _guardrail_modification_check(request_body: dict, team_object: Optional[LiteLLM_TeamTable]) -> None: """ Reject user-supplied metadata flags that would modify guardrail behavior unless the team has explicit permission. Checked keys include the plural @@ -449,9 +415,7 @@ def _guardrail_modification_check( if not can_modify_guardrails(team_object): raise HTTPException( status_code=403, - detail={ - "error": "Your team does not have permission to modify guardrails." - }, + detail={"error": "Your team does not have permission to modify guardrails."}, ) @@ -473,28 +437,16 @@ async def check_tools_allowlist( if valid_token is None: return call_types = get_call_types_for_route(route) - if not call_types or not any( - ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types - ): + if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types): return tool_names = extract_request_tool_names(route, request_body) if not tool_names: return - key_meta = ( - (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} - ) - team_meta = ( - (valid_token.team_metadata or {}) - if isinstance(valid_token.team_metadata, dict) - else {} - ) + key_meta = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} + team_meta = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} key_allowed = key_meta.get("allowed_tools") team_allowed = team_meta.get("allowed_tools") - effective = ( - key_allowed - if (isinstance(key_allowed, list) and len(key_allowed) > 0) - else team_allowed - ) + effective = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed if not isinstance(effective, list) or len(effective) == 0: return allowed_set = {str(t) for t in effective} @@ -571,9 +523,7 @@ async def common_checks( # 1. If team is blocked if team_object is not None and team_object.blocked is True: - raise Exception( - f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin." - ) + raise Exception(f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin.") # 2. If team can call model (or key's access_group_ids grant it) if _model and team_object: @@ -583,9 +533,7 @@ async def common_checks( model=_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=( - valid_token.team_model_aliases if valid_token else None - ), + team_model_aliases=(valid_token.team_model_aliases if valid_token else None), ) except ProxyException as team_denial: if team_denial.type != ProxyErrorTypes.team_model_access_denied: @@ -600,9 +548,7 @@ async def common_checks( # 2.2. If team member has per-member model scope, enforce it if _model and team_object and valid_token and valid_token.user_id: - with tracer.trace( - "litellm.proxy.auth.common_checks.check_team_member_model_access" - ): + with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_model_access"): await _check_team_member_model_access( model=_model, team_object=team_object, @@ -620,9 +566,7 @@ async def common_checks( agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) if agent is not None: - require_trace_id = (agent.litellm_params or {}).get( - "require_trace_id_on_calls_by_agent" - ) + require_trace_id = (agent.litellm_params or {}).get("require_trace_id_on_calls_by_agent") if require_trace_id: headers_dict = dict(request.headers) trace_id = get_chain_id_from_headers(headers_dict) @@ -672,9 +616,7 @@ async def common_checks( await _team_multi_budget_check(team_object=team_object) # 3.2. Multi-window budget check for key - with tracer.trace( - "litellm.proxy.auth.common_checks.virtual_key_multi_budget_check" - ): + with tracer.trace("litellm.proxy.auth.common_checks.virtual_key_multi_budget_check"): if valid_token is not None: await _virtual_key_multi_budget_check(valid_token=valid_token) @@ -687,9 +629,7 @@ async def common_checks( ) # 3.1. If organization is in budget - with tracer.trace( - "litellm.proxy.auth.common_checks.organization_max_budget_check" - ): + with tracer.trace("litellm.proxy.auth.common_checks.organization_max_budget_check"): await _organization_max_budget_check( valid_token=valid_token, team_object=team_object, @@ -754,10 +694,7 @@ async def common_checks( ) # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget - if ( - end_user_object is not None - and end_user_object.litellm_budget_table is not None - ): + if end_user_object is not None and end_user_object.litellm_budget_table is not None: await _check_end_user_budget(end_user_obj=end_user_object, route=route) _enforce_user_param_check(general_settings, request, request_body, route) @@ -765,9 +702,7 @@ async def common_checks( _guardrail_modification_check(request_body, team_object) # 10 [OPTIONAL] Organization RBAC checks - organization_role_based_access_check( - user_object=user_object, route=route, request_body=request_body - ) + organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body) _is_route_allowed = _is_api_route_allowed( route=route, @@ -845,16 +780,10 @@ def _is_user_proxy_admin(user_obj: Optional[LiteLLM_UserTable]): if user_obj is None: return False - if ( - user_obj.user_role is not None - and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value - ): + if user_obj.user_role is not None and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value: return True - if ( - user_obj.user_role is not None - and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value - ): + if user_obj.user_role is not None and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value: return True return False @@ -902,9 +831,7 @@ def allowed_routes_check( """ By default allow a team to call openai + info routes """ - is_allowed = _allowed_routes_check( - user_route=user_route, allowed_routes=["openai_routes", "info_routes"] - ) + is_allowed = _allowed_routes_check(user_route=user_route, allowed_routes=["openai_routes", "info_routes"]) return is_allowed elif litellm_proxy_roles.team_allowed_routes is not None: is_allowed = _allowed_routes_check( @@ -1039,14 +966,10 @@ async def get_team_member_default_budget( return LiteLLM_BudgetTable(**cached_budget) try: - budget_record = await BudgetRepository(prisma_client).table.find_unique( - where={"budget_id": budget_id} - ) + budget_record = await BudgetRepository(prisma_client).table.find_unique(where={"budget_id": budget_id}) if budget_record is None: - verbose_proxy_logger.warning( - f"Team-default member budget not found in database: {budget_id}" - ) + verbose_proxy_logger.warning(f"Team-default member budget not found in database: {budget_id}") return None await user_api_key_cache.async_set_cache( @@ -1058,9 +981,7 @@ async def get_team_member_default_budget( return LiteLLM_BudgetTable(**budget_record.dict()) except Exception: - verbose_proxy_logger.exception( - f"Error fetching team-default member budget {budget_id}" - ) + verbose_proxy_logger.exception(f"Error fetching team-default member budget {budget_id}") return None @@ -1290,11 +1211,7 @@ async def resolve_and_validate_end_user_id( await user_api_key_cache.async_set_cache( key=cache_key, value="valid" if is_valid else "invalid", - ttl=( - _END_USER_VALIDATION_POSITIVE_TTL - if is_valid - else _END_USER_VALIDATION_NEGATIVE_TTL - ), + ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL), ) if is_valid: @@ -1326,9 +1243,7 @@ async def _end_user_id_exists_in_db( if end_user_obj is not None: return True except Exception as e: - verbose_proxy_logger.debug( - f"end_user validation: get_end_user_object lookup failed: {e}" - ) + verbose_proxy_logger.debug(f"end_user validation: get_end_user_object lookup failed: {e}") try: user_obj = await get_user_object( @@ -1344,9 +1259,7 @@ async def _end_user_id_exists_in_db( if user_obj is not None: return True except Exception as e: - verbose_proxy_logger.debug( - f"end_user validation: get_user_object lookup failed: {e}" - ) + verbose_proxy_logger.debug(f"end_user validation: get_user_object lookup failed: {e}") return False @@ -1519,9 +1432,7 @@ async def get_team_membership( return None -def model_in_access_group( - model: str, team_models: Optional[List[str]], llm_router: Optional[Router] -) -> bool: +def model_in_access_group(model: str, team_models: Optional[List[str]], llm_router: Optional[Router]) -> bool: from collections import defaultdict if team_models is None: @@ -1549,9 +1460,7 @@ def model_in_access_group( return False -def _should_check_db( - key: str, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int -) -> bool: +def _should_check_db(key: str, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int) -> bool: """ Prevent calling db repeatedly for items that don't exist in the db. """ @@ -1559,9 +1468,7 @@ def _should_check_db( # if key doesn't exist in last_db_access_time -> check db if key not in last_db_access_time: return True - elif ( - last_db_access_time[key][0] is not None - ): # check db for non-null values (for refresh operations) + elif last_db_access_time[key][0] is not None: # check db for non-null values (for refresh operations) return True elif last_db_access_time[key][0] is None: if current_time - last_db_access_time[key] >= db_cache_expiry: @@ -1569,9 +1476,7 @@ def _should_check_db( return False -def _update_last_db_access_time( - key: str, value: Optional[Any], last_db_access_time: LimitedSizeOrderedDict -): +def _update_last_db_access_time(key: str, value: Optional[Any], last_db_access_time: LimitedSizeOrderedDict): last_db_access_time[key] = (value, time.time()) @@ -1743,10 +1648,7 @@ async def get_user_object( else: raise Exception - if ( - response.organization_memberships is not None - and len(response.organization_memberships) > 0 - ): + if response.organization_memberships is not None and len(response.organization_memberships) > 0: # dump each organization membership to type LiteLLM_OrganizationMembershipTable _dumped_memberships = [ LiteLLM_OrganizationMembershipTable(**membership.model_dump()) @@ -1837,9 +1739,7 @@ async def _cache_team_object( alias_key = "team_alias:{}".format(team_table.team_alias) user_api_key_cache.delete_cache(key=alias_key) if proxy_logging_obj is not None: - await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( - key=alias_key - ) + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=alias_key) async def _cache_key_object( @@ -1853,9 +1753,7 @@ async def _cache_key_object( ## CACHE REFRESH TIME user_api_key_obj.last_refreshed_at = time.time() - cached_key_obj = _copy_user_api_key_auth_for_cache( - user_api_key_obj=user_api_key_obj - ) + cached_key_obj = _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_obj) await _cache_management_object( key=key, value=cached_key_obj, @@ -1876,18 +1774,12 @@ async def _delete_cache_key_object( ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: - await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( - key=key - ) + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) @log_db_metrics -async def _get_team_db_check( - team_id: str, prisma_client: PrismaClient, team_id_upsert: Optional[bool] = None -): - response = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) +async def _get_team_db_check(team_id: str, prisma_client: PrismaClient, team_id_upsert: Optional[bool] = None): + response = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if response is None and team_id_upsert: from litellm.proxy.management_endpoints.team_endpoints import new_team @@ -1907,9 +1799,7 @@ async def _get_team_db_check( async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): - return await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + return await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) async def _get_team_object_from_user_api_key_cache( @@ -1929,9 +1819,7 @@ async def _get_team_object_from_user_api_key_cache( db_cache_expiry=db_cache_expiry, ) if should_check_db: - response = await _get_team_db_check( - team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert - ) + response = await _get_team_db_check(team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert) else: response = None @@ -1980,19 +1868,12 @@ async def _get_team_object_from_cache( parent_otel_span: Optional[Span], ) -> Optional[LiteLLM_TeamTableCachedObj]: ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ## - if ( - proxy_logging_obj is not None - and proxy_logging_obj.internal_usage_cache.dual_cache - ): - cached_raw = ( - await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( - key=key, parent_otel_span=parent_otel_span - ) + if proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache: + cached_raw = await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( + key=key, parent_otel_span=parent_otel_span ) if cached_raw is not None: - from_internal = CacheCodec.deserialize( - cached_raw, LiteLLM_TeamTableCachedObj - ) + from_internal = CacheCodec.deserialize(cached_raw, LiteLLM_TeamTableCachedObj) if from_internal is not None: return from_internal @@ -2023,9 +1904,7 @@ async def get_team_object( - HTTPException: If team doesn't exist in db or cache (status_code=404) """ if prisma_client is None: - raise Exception( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") # check if in cache key = "team_id:{}".format(team_id) @@ -2044,9 +1923,7 @@ async def get_team_object( if check_cache_only: raise HTTPException( status_code=404, - detail={ - "error": f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}." - }, + detail={"error": f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}."}, ) # else, check db @@ -2064,9 +1941,7 @@ async def get_team_object( except Exception: raise HTTPException( status_code=404, - detail={ - "error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call." - }, + detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."}, ) @@ -2096,9 +1971,7 @@ async def _delete_cache_access_object( ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: - await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( - key=key - ) + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) @log_db_metrics @@ -2121,9 +1994,7 @@ async def get_access_object( - HTTPException: If access group doesn't exist in db or cache (status_code=404) """ if prisma_client is None: - raise Exception( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") key = "access_group_id:{}".format(access_group_id) @@ -2143,9 +2014,7 @@ async def get_access_object( if response is None: raise HTTPException( status_code=404, - detail={ - "error": f"Access group doesn't exist in db. Access group={access_group_id}." - }, + detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."}, ) _response = LiteLLM_AccessGroupTable(**response.dict()) @@ -2168,9 +2037,7 @@ async def get_access_object( ) raise HTTPException( status_code=404, - detail={ - "error": f"Access group doesn't exist in db. Access group={access_group_id}. Error: {e}" - }, + detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}. Error: {e}"}, ) @@ -2199,9 +2066,7 @@ async def get_team_object_by_alias( HTTPException: If team doesn't exist or multiple teams have the same alias """ if prisma_client is None: - raise Exception( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") # Check cache first (keyed by alias) cache_key = "team_alias:{}".format(team_alias) @@ -2218,9 +2083,7 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams = await TeamRepository(prisma_client).table.find_many( - where={"team_alias": team_alias} - ) + teams = await TeamRepository(prisma_client).table.find_many(where={"team_alias": team_alias}) if not teams: raise HTTPException( @@ -2280,9 +2143,7 @@ async def get_team_object_by_alias( verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias) raise HTTPException( status_code=500, - detail={ - "error": f"Error looking up team by alias '{team_alias}': {str(e)}" - }, + detail={"error": f"Error looking up team by alias '{team_alias}': {str(e)}"}, ) @@ -2311,9 +2172,7 @@ async def get_org_object_by_alias( HTTPException: If organization not found or multiple orgs have the same alias """ if prisma_client is None: - raise Exception( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") # Check cache first (keyed by alias) cache_key = "org_alias:{}".format(org_alias) @@ -2326,9 +2185,7 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await OrganizationRepository(prisma_client).table.find_many( - where={"organization_alias": org_alias} - ) + orgs = await OrganizationRepository(prisma_client).table.find_many(where={"organization_alias": org_alias}) if not orgs: raise HTTPException( @@ -2369,14 +2226,10 @@ async def get_org_object_by_alias( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - "Error looking up organization by alias: %s", org_alias - ) + verbose_proxy_logger.exception("Error looking up organization by alias: %s", org_alias) raise HTTPException( status_code=500, - detail={ - "error": f"Error looking up organization by alias '{org_alias}': {str(e)}" - }, + detail={"error": f"Error looking up organization by alias '{org_alias}': {str(e)}"}, ) @@ -2488,17 +2341,13 @@ class ExperimentalUIJWTToken: decrypt_value_helper, ) - decrypted_token = decrypt_value_helper( - hashed_token, key="ui_hash_key", exception_type="debug" - ) + decrypted_token = decrypt_value_helper(hashed_token, key="ui_hash_key", exception_type="debug") if decrypted_token is None: return None try: return UserAPIKeyAuth(**json.loads(decrypted_token)) except Exception as e: - raise Exception( - f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}" - ) + raise Exception(f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}") async def _fetch_key_object_from_db_with_reconnect( @@ -2521,14 +2370,10 @@ async def _fetch_key_object_from_db_with_reconnect( if PrismaDBExceptionHandler.is_database_transport_error(e): did_reconnect = False if hasattr(prisma_client, "attempt_db_reconnect"): - auth_reconnect_timeout = getattr( - prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0 - ) + auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0) if not isinstance(auth_reconnect_timeout, (int, float)): auth_reconnect_timeout = 2.0 - auth_reconnect_lock_timeout = getattr( - prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1 - ) + auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1) if not isinstance(auth_reconnect_lock_timeout, (int, float)): auth_reconnect_lock_timeout = 0.1 did_reconnect = await prisma_client.attempt_db_reconnect( @@ -2584,9 +2429,7 @@ async def get_key_object( - if not, then raise an error """ if prisma_client is None: - raise Exception( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") # check if in cache key = hashed_token @@ -2601,9 +2444,7 @@ async def get_key_object( return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth) if check_cache_only: - raise Exception( - f"Key doesn't exist in cache + check_cache_only=True. key={key}." - ) + raise Exception(f"Key doesn't exist in cache + check_cache_only=True. key={key}.") # else, check db _valid_token: Optional[BaseModel] = await _fetch_key_object_from_db_with_reconnect( @@ -2675,9 +2516,7 @@ async def get_object_permission( - if not, then raise an error """ if prisma_client is None: - raise Exception( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") # check if in cache key = "object_permission_id:{}".format(object_permission_id) @@ -2752,11 +2591,7 @@ async def get_managed_vector_store_rows_by_uuids( ) for row in rows: - row_dict = ( - row.model_dump() - if hasattr(row, "model_dump") - else (row.dict() if hasattr(row, "dict") else None) - ) + row_dict = row.model_dump() if hasattr(row, "model_dump") else (row.dict() if hasattr(row, "dict") else None) if not isinstance(row_dict, dict) or not row_dict: row_dict = dict(row) if hasattr(row, "__dict__") else {} if not row_dict: @@ -2797,9 +2632,7 @@ async def get_org_object( include_budget_table: If True, includes litellm_budget_table in the query """ if prisma_client is None: - raise Exception( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") if not isinstance(org_id, str): return None @@ -2821,9 +2654,7 @@ async def get_org_object( if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response = await OrganizationRepository(prisma_client).table.find_unique( - **query_kwargs - ) + response = await OrganizationRepository(prisma_client).table.find_unique(**query_kwargs) if response is None: raise Exception @@ -2846,9 +2677,7 @@ async def get_org_object( async def _get_resources_from_access_groups( access_group_ids: List[str], - resource_field: Literal[ - "access_model_names", "access_mcp_server_ids", "access_agent_ids" - ], + resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], prisma_client: Optional[PrismaClient] = None, user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, @@ -2967,16 +2796,10 @@ def _resolve_all_team_model_sentinel_for_auth_check( llm_router: Optional[Router], team_id: Optional[str], ) -> List[str]: - if ( - SpecialModelNames.all_team_models.value not in models - or team_id is None - or llm_router is None - ): + if SpecialModelNames.all_team_models.value not in models or team_id is None or llm_router is None: return models proxy_models = llm_router.get_model_names() - non_sentinel_models = [ - model for model in models if model != SpecialModelNames.all_team_models.value - ] + non_sentinel_models = [model for model in models if model != SpecialModelNames.all_team_models.value] if not proxy_models: return non_sentinel_models or models return list(dict.fromkeys(non_sentinel_models + proxy_models)) @@ -2995,9 +2818,7 @@ def _check_model_access_helper( access_groups: Dict[str, List[str]] = defaultdict(list) if llm_router: - access_groups = llm_router.get_model_access_groups( - model_name=model, team_id=team_id - ) + access_groups = llm_router.get_model_access_groups(model_name=model, team_id=team_id) models = _resolve_all_team_model_sentinel_for_auth_check( models=models, @@ -3005,12 +2826,8 @@ def _check_model_access_helper( team_id=team_id, ) - if ( - len(access_groups) > 0 and llm_router is not None - ): # check if token contains any model access groups - for idx, m in enumerate( - models - ): # loop token models, if any of them are an access group add the access group + if len(access_groups) > 0 and llm_router is not None: # check if token contains any model access groups + for idx, m in enumerate(models): # loop token models, if any of them are an access group add the access group if m in access_groups: return True @@ -3020,9 +2837,7 @@ def _check_model_access_helper( if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases): return True - if _model_matches_any_wildcard_pattern_in_list( - model=model, allowed_model_list=filtered_models - ): + if _model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models): return True all_model_access: bool = False @@ -3064,11 +2879,7 @@ def _can_object_call_model( - Exception: If token not allowed to call model """ if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: - raise Exception( - "Unable to parse model, max fallback depth exceeded - received model: {}".format( - model - ) - ) + raise Exception("Unable to parse model, max fallback depth exceeded - received model: {}".format(model)) if isinstance(model, list): for m in model: _can_object_call_model( @@ -3103,17 +2914,13 @@ def _can_object_call_model( raise ProxyException( message=f"{object_type} not allowed to access model. This {object_type} can only access models={models}. Tried to access {model}", - type=ProxyErrorTypes.get_model_access_error_type_for_object( - object_type=object_type - ), + type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type=object_type), param="model", code=status.HTTP_403_FORBIDDEN, ) -def _model_in_team_aliases( - model: str, team_model_aliases: Optional[Dict[str, str]] = None -) -> bool: +def _model_in_team_aliases(model: str, team_model_aliases: Optional[Dict[str, str]] = None) -> bool: """ Returns True if `model` being accessed is an alias of a team model @@ -3211,8 +3018,7 @@ async def can_key_call_resolved_model( ) skip_key_model_check = valid_token.config or ( - isinstance(valid_token.models, list) - and SpecialModelNames.all_team_models.value in valid_token.models + isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models ) if not skip_key_model_check: await can_key_call_model( @@ -3332,9 +3138,7 @@ async def can_team_access_model( ) except ProxyException: # Fallback: check team's access_group_ids - team_access_group_ids = ( - (team_object.access_group_ids or []) if team_object else [] - ) + team_access_group_ids = (team_object.access_group_ids or []) if team_object else [] if team_access_group_ids: models_from_groups = await _get_models_from_access_groups( access_group_ids=team_access_group_ids, @@ -3354,9 +3158,7 @@ async def can_team_access_model( async def get_authorized_resources_from_key_access_groups( valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], - resource_field: Literal[ - "access_model_names", "access_mcp_server_ids", "access_agent_ids" - ], + resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], ) -> List[str]: """ For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row @@ -3379,9 +3181,7 @@ async def get_authorized_resources_from_key_access_groups( if _prisma_client is None or _user_api_key_cache is None: return [] - key_team_id = valid_token.team_id or ( - team_object.team_id if team_object is not None else None - ) + key_team_id = valid_token.team_id or (team_object.team_id if team_object is not None else None) key_token = valid_token.token authorized_resources: List[str] = [] @@ -3395,9 +3195,7 @@ async def get_authorized_resources_from_key_access_groups( ) except Exception: continue - team_authorized = bool( - key_team_id and key_team_id in (ag.assigned_team_ids or []) - ) + team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or [])) key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) if team_authorized or key_authorized: authorized_resources.extend(getattr(ag, resource_field, []) or []) @@ -3551,9 +3349,7 @@ async def can_key_call_search_tool( """ return _can_object_call_search_tool( search_tool_name=search_tool_name, - allowed_search_tools=_search_tool_names_from_object_permission( - valid_token.object_permission - ), + allowed_search_tools=_search_tool_names_from_object_permission(valid_token.object_permission), object_type="key", ) @@ -3582,9 +3378,7 @@ async def can_team_call_search_tool( return _can_object_call_search_tool( search_tool_name=search_tool_name, - allowed_search_tools=_search_tool_names_from_object_permission( - team_object.object_permission - ), + allowed_search_tools=_search_tool_names_from_object_permission(team_object.object_permission), object_type="team", ) @@ -3700,11 +3494,7 @@ async def _virtual_key_max_budget_check( # name the key in the error so operators don't have to reverse-map # spend back to a key; key_name is the masked form (last 4 chars) key_label = valid_token.key_alias or "key" - key_descriptor = ( - f"{key_label} ({valid_token.key_name})" - if valid_token.key_name - else key_label - ) + key_descriptor = f"{key_label} ({valid_token.key_name})" if valid_token.key_name else key_label raise litellm.BudgetExceededError( current_cost=spend, max_budget=valid_token.max_budget, @@ -3829,12 +3619,7 @@ def _merge_budget_alert_email_configs( return None thresholds = set(global_cfg_normalized) | set(per_key_cfg_normalized) return { - t: list( - dict.fromkeys( - global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []) - ) - ) - for t in thresholds + t: list(dict.fromkeys(global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []))) for t in thresholds } @@ -3850,17 +3635,11 @@ async def _virtual_key_max_budget_alert_check( """ - if ( - valid_token.max_budget is not None - and valid_token.spend is not None - and valid_token.spend > 0 - ): + if valid_token.max_budget is not None and valid_token.spend is not None and valid_token.spend > 0: owner_email = user_obj.user_email if user_obj else None - alert_email_config: Optional[Dict[str, List[str]]] = ( - _merge_budget_alert_email_configs( - global_cfg=litellm.default_key_max_budget_alert_emails, - per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"), - ) + alert_email_config: Optional[Dict[str, List[str]]] = _merge_budget_alert_email_configs( + global_cfg=litellm.default_key_max_budget_alert_emails, + per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"), ) if isinstance(alert_email_config, dict) and alert_email_config: @@ -3869,9 +3648,7 @@ async def _virtual_key_max_budget_alert_check( (int(k) for k in alert_email_config if k.isdigit()), default=None, ) - if min_pct is None or valid_token.spend < valid_token.max_budget * ( - min_pct / 100.0 - ): + if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0): return call_info = CallInfo( @@ -3896,14 +3673,9 @@ async def _virtual_key_max_budget_alert_check( ) else: # Old path: existing single 80% threshold — completely unchanged - alert_threshold = ( - valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - ) + alert_threshold = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE - if ( - valid_token.spend >= alert_threshold - and valid_token.spend < valid_token.max_budget - ): + if valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget: verbose_proxy_logger.debug( "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", valid_token.token, @@ -3966,9 +3738,7 @@ async def _check_team_member_budget( ): team_member_budget = team_membership.litellm_budget_table.max_budget else: - default_budget_id = (team_object.metadata or {}).get( - "team_member_budget_id" - ) + default_budget_id = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): default_budget = await get_team_member_default_budget( budget_id=default_budget_id, @@ -3985,9 +3755,7 @@ async def _check_team_member_budget( team_member_budget = default_budget.max_budget if team_member_budget is not None: - team_member_spend = ( - team_membership.spend if team_membership is not None else 0.0 - ) or 0.0 + team_member_spend = (team_membership.spend if team_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend @@ -3998,10 +3766,7 @@ async def _check_team_member_budget( max_budget=team_member_budget, ) - if ( - math.isfinite(team_member_budget) - and team_member_spend >= team_member_budget - ): + if math.isfinite(team_member_budget) and team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, @@ -4042,9 +3807,7 @@ async def _check_team_member_model_access( ): return # no per-member restriction — inherit team-level check - member_allowed_models: List[str] = ( - team_membership.litellm_budget_table.allowed_models - ) + member_allowed_models: List[str] = team_membership.litellm_budget_table.allowed_models try: _can_object_call_model( model=model, @@ -4170,26 +3933,16 @@ async def _team_soft_budget_check( if valid_token: # Extract alert emails from team metadata alert_emails: Optional[List[str]] = None - if team_object.metadata is not None and isinstance( - team_object.metadata, dict - ): - soft_budget_alert_emails = team_object.metadata.get( - "soft_budget_alerting_emails" - ) + if team_object.metadata is not None and isinstance(team_object.metadata, dict): + soft_budget_alert_emails = team_object.metadata.get("soft_budget_alerting_emails") if soft_budget_alert_emails is not None: if isinstance(soft_budget_alert_emails, list): alert_emails = [ - email - for email in soft_budget_alert_emails - if isinstance(email, str) and email.strip() + email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip() ] elif isinstance(soft_budget_alert_emails, str): # Handle comma-separated string - alert_emails = [ - email.strip() - for email in soft_budget_alert_emails.split(",") - if email.strip() - ] + alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()] # Filter out empty strings if alert_emails: alert_emails = [email for email in alert_emails if email] @@ -4295,11 +4048,7 @@ async def _project_soft_budget_check( if project_object.litellm_budget_table is not None: soft_budget = project_object.litellm_budget_table.soft_budget - if ( - soft_budget is not None - and project_object.spend is not None - and project_object.spend >= soft_budget - ): + if soft_budget is not None and project_object.spend is not None and project_object.spend >= soft_budget: verbose_proxy_logger.debug( "Crossed Soft Budget for project %s, spend %s, soft_budget %s", project_object.project_id, @@ -4509,10 +4258,7 @@ async def _tag_max_budget_check( continue # Check if tag has budget limits - if ( - tag_object.litellm_budget_table is not None - and tag_object.litellm_budget_table.max_budget is not None - ): + if tag_object.litellm_budget_table is not None and tag_object.litellm_budget_table.max_budget is not None: from litellm.proxy.proxy_server import get_current_spend tag_spend = await get_current_spend( @@ -4549,9 +4295,7 @@ def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: return False -def _model_matches_any_wildcard_pattern_in_list( - model: str, allowed_model_list: list -) -> bool: +def _model_matches_any_wildcard_pattern_in_list(model: str, allowed_model_list: list) -> bool: """ Returns True if a model matches any wildcard pattern in a list. @@ -4563,9 +4307,7 @@ def _model_matches_any_wildcard_pattern_in_list( if any( _is_wildcard_pattern(allowed_model_pattern) - and is_model_allowed_by_pattern( - model=model, allowed_model_pattern=allowed_model_pattern - ) + and is_model_allowed_by_pattern(model=model, allowed_model_pattern=allowed_model_pattern) for allowed_model_pattern in allowed_model_list ): return True @@ -4582,9 +4324,7 @@ def _model_matches_any_wildcard_pattern_in_list( return False -def _model_custom_llm_provider_matches_wildcard_pattern( - model: str, allowed_model_pattern: str -) -> bool: +def _model_custom_llm_provider_matches_wildcard_pattern(model: str, allowed_model_pattern: str) -> bool: """ Returns True for this scenario: - `model=gpt-4o` @@ -4630,24 +4370,18 @@ async def vector_store_access_check( # Get the vector store the user is trying to access ######################################################### if prisma_client is None: - verbose_proxy_logger.debug( - "Prisma client not found, skipping vector store access check" - ) + verbose_proxy_logger.debug("Prisma client not found, skipping vector store access check") return True if litellm.vector_store_registry is None: - verbose_proxy_logger.debug( - "Vector store registry not found, skipping vector store access check" - ) + verbose_proxy_logger.debug("Vector store registry not found, skipping vector store access check") return True vector_store_ids_to_run = litellm.vector_store_registry.get_vector_store_ids_to_run( non_default_params=request_body, tools=request_body.get("tools", None) ) if vector_store_ids_to_run is None: - verbose_proxy_logger.debug( - "Vector store to run not found, skipping vector store access check" - ) + verbose_proxy_logger.debug("Vector store to run not found, skipping vector store access check") return True ######################################################### @@ -4655,9 +4389,7 @@ async def vector_store_access_check( ######################################################### # Check if the key can access the vector store if valid_token is not None and valid_token.object_permission_id is not None: - key_object_permission = await ObjectPermissionRepository( - prisma_client - ).table.find_unique( + key_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": valid_token.object_permission_id}, ) if key_object_permission is not None: @@ -4669,9 +4401,7 @@ async def vector_store_access_check( # Check if the team can access the vector store if team_object is not None and team_object.object_permission_id is not None: - team_object_permission = await ObjectPermissionRepository( - prisma_client - ).table.find_unique( + team_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": team_object.object_permission_id}, ) if team_object_permission is not None: @@ -4705,9 +4435,7 @@ def _can_object_call_vector_stores( if vector_store_id not in object_permissions.vector_stores: raise ProxyException( message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}", - type=ProxyErrorTypes.get_vector_store_access_error_type_for_object( - object_type - ), + type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type), param="vector_store", code=status.HTTP_401_UNAUTHORIZED, ) diff --git a/litellm/proxy/auth/auth_checks_organization.py b/litellm/proxy/auth/auth_checks_organization.py index 00aac0d48f9..44c1d158cbe 100644 --- a/litellm/proxy/auth/auth_checks_organization.py +++ b/litellm/proxy/auth/auth_checks_organization.py @@ -65,9 +65,7 @@ def organization_role_based_access_check( code=status.HTTP_401_UNAUTHORIZED, ) - user_role: Optional[LitellmUserRoles] = _user_organization_role_mapping.get( - passed_organization_id - ) + user_role: Optional[LitellmUserRoles] = _user_organization_role_mapping.get(passed_organization_id) if user_role is None: raise ProxyException( message=f"You do not have a role within the selected organization. Passed organization_id: {passed_organization_id}. Please contact the organization admin to request access.", @@ -89,10 +87,7 @@ def organization_role_based_access_check( _user_organizations, _user_organization_role_mapping, ) = get_user_organization_info(user_object) - if ( - user_object.organization_memberships is not None - and len(user_object.organization_memberships) > 0 - ): + if user_object.organization_memberships is not None and len(user_object.organization_memberships) > 0: if passed_organization_id is None: raise ProxyException( message=f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {_user_organizations}", @@ -101,9 +96,7 @@ def organization_role_based_access_check( code=status.HTTP_401_UNAUTHORIZED, ) - _user_role_in_passed_org = _user_organization_role_mapping.get( - passed_organization_id - ) + _user_role_in_passed_org = _user_organization_role_mapping.get(passed_organization_id) if _user_role_in_passed_org != LitellmUserRoles.ORG_ADMIN.value: raise ProxyException( message=f"You do not have the required role to call {route}. Your role is {_user_role_in_passed_org} in Organization {passed_organization_id}", @@ -134,9 +127,7 @@ def get_user_organization_info( for _membership in user_object.organization_memberships: if _membership.organization_id is not None: _user_organizations.append(_membership.organization_id) - _user_organization_role_mapping[_membership.organization_id] = ( - _membership.user_role - ) # type: ignore + _user_organization_role_mapping[_membership.organization_id] = _membership.user_role # type: ignore return _user_organizations, _user_organization_role_mapping @@ -174,8 +165,7 @@ def _user_is_org_admin( admin_org_ids = { _membership.organization_id for _membership in user_object.organization_memberships - if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value - and _membership.organization_id is not None + if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value and _membership.organization_id is not None } # User must be admin of ALL requested orgs, not just any one diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 83f18173182..e19bb3f9f5f 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -107,16 +107,10 @@ class UserAPIKeyAuthExceptionHandler: # is labeled — a fresh UserAPIKeyAuth here would drop everything auth had # already looked up (e.g. an expired key whose team/user is known). Copy # so the handler is side-effect-free for the caller's identity object. - user_api_key_dict = ( - resolved_identity.model_copy() - if resolved_identity is not None - else UserAPIKeyAuth() - ) + user_api_key_dict = resolved_identity.model_copy() if resolved_identity is not None else UserAPIKeyAuth() user_api_key_dict.parent_otel_span = parent_otel_span user_api_key_dict.request_route = route - user_api_key_dict.api_key = ( - user_api_key_dict.api_key or UserAPIKeyAuth(api_key=api_key).api_key - ) + user_api_key_dict.api_key = user_api_key_dict.api_key or UserAPIKeyAuth(api_key=api_key).api_key # Stamp identity onto the request's server span now, before the request # is rejected; the OTEL failure hooks don't touch the server span, so @@ -136,9 +130,7 @@ class UserAPIKeyAuthExceptionHandler: resolve_llm_provider_for_rate_limit, ) - _, e.llm_provider = resolve_llm_provider_for_rate_limit( - request_data.get("model") - ) + _, e.llm_provider = resolve_llm_provider_for_rate_limit(request_data.get("model")) # Allow callbacks to transform the error response transformed_exception = await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 3a2f2221ee3..b1bce352784 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -18,9 +18,7 @@ from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS from litellm.types.utils import CustomPricingLiteLLMParams -def _get_request_ip_address( - request: Request, use_x_forwarded_for: Optional[bool] = False -) -> Optional[str]: +def _get_request_ip_address(request: Request, use_x_forwarded_for: Optional[bool] = False) -> Optional[str]: client_ip = None if use_x_forwarded_for is True and "x-forwarded-for" in request.headers: client_ip = request.headers["x-forwarded-for"] @@ -44,9 +42,7 @@ def _check_valid_ip( return True, None # if general_settings.get("use_x_forwarded_for") is True then use x-forwarded-for - client_ip = _get_request_ip_address( - request=request, use_x_forwarded_for=use_x_forwarded_for - ) + client_ip = _get_request_ip_address(request=request, use_x_forwarded_for=use_x_forwarded_for) # Check if IP address is allowed if client_ip not in allowed_ips: @@ -97,8 +93,7 @@ def check_complete_credentials(request_body: dict) -> bool: validate_url(url_value) except SSRFError as e: raise ValueError( - f"Rejected request: client-side {url_field}={url_value!r} " - f"is rejected by the SSRF guard ({e})." + f"Rejected request: client-side {url_field}={url_value!r} is rejected by the SSRF guard ({e})." ) return True @@ -153,9 +148,7 @@ def _allow_model_level_clientside_configurable_parameters( if model_info is None: # check if wildcard model is set if model.split("/", 1)[0] in provider_list: - model_info = llm_router.get_model_group_info( - model_group=model.split("/", 1)[0] - ) + model_info = llm_router.get_model_group_info(model_group=model.split("/", 1)[0]) if model_info is None: return False @@ -339,9 +332,7 @@ def _check_banned_params( ) -def is_request_body_safe( - request_body: dict, general_settings: dict, llm_router: Optional[Router], model: str -) -> bool: +def is_request_body_safe(request_body: dict, general_settings: dict, llm_router: Optional[Router], model: str) -> bool: """ Check if the request body is safe. @@ -368,9 +359,7 @@ def is_request_body_safe( recursion-depth DoS surface. """ if "model_list" in request_body: - raise ValueError( - "Rejected Request: model_list is not allowed in the request body." - ) + raise ValueError("Rejected Request: model_list is not allowed in the request body.") _check_banned_params(request_body, general_settings, llm_router, model) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) @@ -430,9 +419,7 @@ async def pre_db_read_auth_checks( request_body=request_data, general_settings=general_settings, llm_router=llm_router, - model=request_data.get( - "model", "" - ), # [TODO] use model passed in url as well (azure openai routes) + model=request_data.get("model", ""), # [TODO] use model passed in url as well (azure openai routes) ) # Check 3. Check if IP address is allowed @@ -456,9 +443,7 @@ async def pre_db_read_auth_checks( f"Trying to set allowed_routes. This is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}" ) if route not in _allowed_routes: - verbose_proxy_logger.error( - f"Route {route} not in allowed_routes={_allowed_routes}" - ) + verbose_proxy_logger.error(f"Route {route} not in allowed_routes={_allowed_routes}") raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"Access forbidden: Route {route} not allowed", @@ -503,9 +488,7 @@ def route_in_additonal_public_routes(current_route: str): # Check wildcard patterns for route_pattern in routes_defined: - if RouteChecks._route_matches_wildcard_pattern( - route=current_route, pattern=route_pattern - ): + if RouteChecks._route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): return True return False @@ -534,18 +517,14 @@ def get_request_route(request: Request) -> str: if not isinstance(scope, dict): return str(request.url.path) raw_path: str = str(scope.get("path", request.url.path)) - root_path: str = str( - scope.get("app_root_path", scope.get("root_path", "")) - ).rstrip("/") + root_path: str = str(scope.get("app_root_path", scope.get("root_path", ""))).rstrip("/") if not isinstance(raw_path, str): return str(request.url.path) # Strip root_path only when it matches whole path segments — guarding # against sibling paths like "/apifoo" being truncated under # root_path="/api". Trailing slashes on root_path are stripped above, # so bare "/" or "/prefix/" still leave the leading "/" intact. - if root_path and ( - raw_path == root_path or raw_path.startswith(root_path + "/") - ): + if root_path and (raw_path == root_path or raw_path.startswith(root_path + "/")): stripped = raw_path[len(root_path) :] return stripped or "/" return raw_path @@ -723,9 +702,7 @@ async def check_if_request_size_is_safe(request: Request) -> bool: if content_length: header_size = int(content_length) header_size_mb = bytes_to_mb(bytes_value=header_size) - verbose_proxy_logger.debug( - f"content_length request size in MB={header_size_mb}" - ) + verbose_proxy_logger.debug(f"content_length request size in MB={header_size_mb}") if header_size_mb > max_request_size_mb: raise ProxyException( @@ -740,9 +717,7 @@ async def check_if_request_size_is_safe(request: Request) -> bool: body_size = len(body) request_size_mb = bytes_to_mb(bytes_value=body_size) - verbose_proxy_logger.debug( - f"request body request size in MB={request_size_mb}" - ) + verbose_proxy_logger.debug(f"request body request size in MB={request_size_mb}") if request_size_mb > max_request_size_mb: raise ProxyException( message=f"Request size is too large. Request size is {request_size_mb} MB. Max size is {max_request_size_mb} MB", @@ -926,9 +901,7 @@ def get_key_model_tpm_limit( def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, - metadata_accessor_key: Literal[ - "team_metadata", "organization_metadata", "project_metadata" - ], + metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], ) -> Optional[Dict[str, int]]: if getattr(user_api_key_dict, metadata_accessor_key): @@ -1063,11 +1036,7 @@ def _has_user_setup_sso(): google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) - sso_setup = ( - (microsoft_client_id is not None) - or (google_client_id is not None) - or (generic_client_id is not None) - ) + sso_setup = (microsoft_client_id is not None) or (google_client_id is not None) or (generic_client_id is not None) return sso_setup @@ -1156,17 +1125,13 @@ def _coerce_user_id_to_str(value: Any) -> Optional[str]: return None -def get_end_user_id_from_request_body( - request_body: dict, request_headers: Optional[dict] = None -) -> Optional[str]: +def get_end_user_id_from_request_body(request_body: dict, request_headers: Optional[dict] = None) -> Optional[str]: # Import general_settings here to avoid potential circular import issues at module level # and to ensure it's fetched at runtime. from litellm.proxy.proxy_server import general_settings # Check 1: Standard customer ID headers (always checked, no configuration required) - customer_id = _get_customer_id_from_standard_headers( - request_headers=request_headers - ) + customer_id = _get_customer_id_from_standard_headers(request_headers=request_headers) if customer_id is not None: return customer_id @@ -1179,9 +1144,7 @@ def get_end_user_id_from_request_body( # Prefer user mappings (new behavior) user_id_mapping = general_settings.get("user_header_mappings", None) if user_id_mapping: - custom_header_name_to_check = get_customer_user_header_from_mapping( - user_id_mapping - ) + custom_header_name_to_check = get_customer_user_header_from_mapping(user_id_mapping) # Fallback to deprecated user_header_name if mapping did not specify if not custom_header_name_to_check: @@ -1319,9 +1282,7 @@ def _dedupe_model_candidates(candidates: List[str]) -> List[str]: return deduped -def _get_case_insensitive_mapping_value( - mapping: Optional[Mapping[str, Any]], key: str -) -> Any: +def _get_case_insensitive_mapping_value(mapping: Optional[Mapping[str, Any]], key: str) -> Any: if not mapping: return None if key in mapping: @@ -1360,9 +1321,7 @@ def _extract_models_from_managed_resource_id( get_models_from_unified_file_id, ) - _append_model_candidates( - candidates=candidates, value=decode_model_from_file_id(resource_id) - ) + _append_model_candidates(candidates=candidates, value=decode_model_from_file_id(resource_id)) unified_file_id = _is_base64_encoded_unified_file_id(resource_id) if unified_file_id: _append_model_candidates( @@ -1374,25 +1333,17 @@ def _extract_models_from_managed_resource_id( value=get_model_id_from_unified_batch_id(unified_file_id), ) except Exception as e: - verbose_proxy_logger.debug( - "Unable to extract model from managed file/batch ID: %s", str(e) - ) + verbose_proxy_logger.debug("Unable to extract model from managed file/batch ID: %s", str(e)) try: from litellm.llms.base_llm.managed_resources.utils import parse_unified_id parsed_id = parse_unified_id(resource_id) if parsed_id: - _append_model_candidates( - candidates=candidates, value=parsed_id.get("model_id") - ) - _append_model_candidates( - candidates=candidates, value=parsed_id.get("target_model_names") - ) + _append_model_candidates(candidates=candidates, value=parsed_id.get("model_id")) + _append_model_candidates(candidates=candidates, value=parsed_id.get("target_model_names")) except Exception as e: - verbose_proxy_logger.debug( - "Unable to extract model from unified managed resource ID: %s", str(e) - ) + verbose_proxy_logger.debug("Unable to extract model from unified managed resource ID: %s", str(e)) if resource_id_field in ("video_id", "character_id"): try: @@ -1408,32 +1359,24 @@ def _extract_models_from_managed_resource_id( value=_resolve_model_id_with_router(model_id, llm_router), ) else: - model_id = decode_character_id_with_provider(resource_id).get( - "model_id" - ) + model_id = decode_character_id_with_provider(resource_id).get("model_id") _append_model_candidates( candidates=candidates, value=_resolve_model_id_with_router(model_id, llm_router), ) except Exception as e: - verbose_proxy_logger.debug( - "Unable to extract model from managed video/character ID: %s", str(e) - ) + verbose_proxy_logger.debug("Unable to extract model from managed video/character ID: %s", str(e)) return _dedupe_model_candidates(candidates) -def _resolve_model_id_with_router( - model_id: Optional[str], llm_router: Optional[Router] -) -> Optional[str]: +def _resolve_model_id_with_router(model_id: Optional[str], llm_router: Optional[Router]) -> Optional[str]: if model_id is None or llm_router is None: return model_id try: return llm_router.resolve_model_name_from_model_id(model_id) or model_id except Exception as e: - verbose_proxy_logger.debug( - "Unable to resolve model_id from managed resource ID: %s", str(e) - ) + verbose_proxy_logger.debug("Unable to resolve model_id from managed resource ID: %s", str(e)) return model_id @@ -1463,15 +1406,11 @@ def _extract_model_candidates_from_request( _append_model_candidates(candidates, body_model) if uses_body_target_model_sources or not body_model: _append_model_candidates(candidates, request_data.get("target_model_names")) - if _route_matches_any_marker( - route=route, markers=_MODEL_ROUTING_SESSION_MODEL_ROUTE_MARKERS - ): + if _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_SESSION_MODEL_ROUTE_MARKERS): session = request_data.get("session") if isinstance(session, dict): _append_model_candidates(candidates, session.get("model")) - if uses_completion_model_sources and isinstance( - request_data.get("completion"), dict - ): + if uses_completion_model_sources and isinstance(request_data.get("completion"), dict): _append_model_candidates(candidates, request_data["completion"].get("model")) if uses_model_routing_sources: @@ -1482,16 +1421,12 @@ def _extract_model_candidates_from_request( ) _append_model_candidates( candidates, - _get_case_insensitive_mapping_value( - request_headers, MODEL_ROUTING_HEADER_NAME - ), + _get_case_insensitive_mapping_value(request_headers, MODEL_ROUTING_HEADER_NAME), ) if uses_query_target_model_sources: _append_model_candidates( candidates, - _get_case_insensitive_mapping_value( - request_query_params, "target_model_names" - ), + _get_case_insensitive_mapping_value(request_query_params, "target_model_names"), ) for field in _MODEL_ROUTING_ID_FIELDS: diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 1b73bb16db3..9db9b970d88 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -160,9 +160,7 @@ class JWTHandler: return claims return None except Exception as e: - verbose_proxy_logger.debug( - "Failed to decode unverified JWT claims for routing: %s", e - ) + verbose_proxy_logger.debug("Failed to decode unverified JWT claims for routing: %s", e) return None def _rbac_role_from_role_mapping(self, token: dict) -> Optional[RBAC_ROLES]: @@ -223,9 +221,7 @@ class JWTHandler: return LitellmUserRoles.TEAM elif self.get_user_id(token=token, default_value=None) is not None: return LitellmUserRoles.INTERNAL_USER - elif user_roles is not None and self.is_allowed_user_role( - user_roles=user_roles - ): + elif user_roles is not None and self.is_allowed_user_role(user_roles=user_roles): return LitellmUserRoles.INTERNAL_USER elif rbac_role := self._rbac_role_from_role_mapping(token=token): return rbac_role @@ -250,9 +246,7 @@ class JWTHandler: return self._is_trusted_issuer_normalized_token(token=token) and claim in token def get_team_ids_from_jwt(self, token: dict) -> List[str]: - if self._has_trusted_issuer_normalized_claim( - token=token, claim=self.LITELLM_TEAM_IDS_CLAIM - ): + if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_TEAM_IDS_CLAIM): issuer_team_ids = token.get(self.LITELLM_TEAM_IDS_CLAIM) if isinstance(issuer_team_ids, list): return issuer_team_ids @@ -293,9 +287,7 @@ class JWTHandler: """ team_ids: List[str] = list(self.get_team_ids_from_jwt(token)) singular: Any = None - if self._has_trusted_issuer_normalized_claim( - token=token, claim=self.LITELLM_TEAM_ID_CLAIM - ): + if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_TEAM_ID_CLAIM): singular = token.get(self.LITELLM_TEAM_ID_CLAIM) elif self.litellm_jwtauth.team_id_jwt_field is not None: singular = get_nested_value( @@ -315,12 +307,8 @@ class JWTHandler: team_ids.append(str(singular)) return team_ids - def get_end_user_id( - self, token: dict, default_value: Optional[str] - ) -> Optional[str]: - if self._has_trusted_issuer_normalized_claim( - token=token, claim=self.LITELLM_END_USER_ID_CLAIM - ): + def get_end_user_id(self, token: dict, default_value: Optional[str]) -> Optional[str]: + if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_END_USER_ID_CLAIM): return token.get(self.LITELLM_END_USER_ID_CLAIM) try: @@ -343,10 +331,7 @@ class JWTHandler: - True: if 'team_id_jwt_field' or 'team_alias_jwt_field' is set - False: if neither is set """ - if ( - self.litellm_jwtauth.team_id_jwt_field is None - and self.litellm_jwtauth.team_alias_jwt_field is None - ): + if self.litellm_jwtauth.team_id_jwt_field is None and self.litellm_jwtauth.team_alias_jwt_field is None: return False return True @@ -364,9 +349,7 @@ class JWTHandler: return False def get_team_id(self, token: dict, default_value: Optional[str]) -> Optional[str]: - if self._has_trusted_issuer_normalized_claim( - token=token, claim=self.LITELLM_TEAM_ID_CLAIM - ): + if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_TEAM_ID_CLAIM): team_id = token.get(self.LITELLM_TEAM_ID_CLAIM) if isinstance(team_id, list): return team_id[0] if team_id else default_value @@ -407,9 +390,7 @@ class JWTHandler: team_id = default_value return team_id - def get_team_alias( - self, token: dict, default_value: Optional[str] - ) -> Optional[str]: + def get_team_alias(self, token: dict, default_value: Optional[str]) -> Optional[str]: """ Extract team name/alias from JWT token using the configured team_alias_jwt_field. @@ -445,9 +426,7 @@ class JWTHandler: return self.litellm_jwtauth.user_id_upsert def get_user_id(self, token: dict, default_value: Optional[str]) -> Optional[str]: - if self._has_trusted_issuer_normalized_claim( - token=token, claim=self.LITELLM_USER_ID_CLAIM - ): + if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_USER_ID_CLAIM): return token.get(self.LITELLM_USER_ID_CLAIM) try: @@ -463,9 +442,7 @@ class JWTHandler: user_id = default_value return user_id - def get_user_roles( - self, token: dict, default_value: Optional[List[str]] - ) -> Optional[List[str]]: + def get_user_roles(self, token: dict, default_value: Optional[List[str]]) -> Optional[List[str]]: """ Returns the user role from the token. @@ -499,9 +476,7 @@ class JWTHandler: return mapping.litellm_role return None - def get_jwt_role( - self, token: dict, default_value: Optional[List[str]] - ) -> Optional[List[str]]: + def get_jwt_role(self, token: dict, default_value: Optional[List[str]]) -> Optional[List[str]]: """ Generic implementation of `get_user_roles` that can be used for both user and team roles. @@ -531,19 +506,13 @@ class JWTHandler: if ( user_roles is not None and self.litellm_jwtauth.user_allowed_roles is not None - and any( - role in self.litellm_jwtauth.user_allowed_roles for role in user_roles - ) + and any(role in self.litellm_jwtauth.user_allowed_roles for role in user_roles) ): return True return False - def get_user_email( - self, token: dict, default_value: Optional[str] - ) -> Optional[str]: - if self._has_trusted_issuer_normalized_claim( - token=token, claim=self.LITELLM_USER_EMAIL_CLAIM - ): + def get_user_email(self, token: dict, default_value: Optional[str]) -> Optional[str]: + if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_USER_EMAIL_CLAIM): return token.get(self.LITELLM_USER_EMAIL_CLAIM) try: @@ -574,9 +543,7 @@ class JWTHandler: return object_id def get_org_id(self, token: dict, default_value: Optional[str]) -> Optional[str]: - if self._has_trusted_issuer_normalized_claim( - token=token, claim=self.LITELLM_ORG_ID_CLAIM - ): + if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_ORG_ID_CLAIM): return token.get(self.LITELLM_ORG_ID_CLAIM) try: @@ -625,9 +592,7 @@ class JWTHandler: elif isinstance(token["scope"], list): scopes = token["scope"] else: - raise Exception( - f"Unmapped scope type - {type(token['scope'])}. Supported types - list, str." - ) + raise Exception(f"Unmapped scope type - {type(token['scope'])}. Supported types - list, str.") except KeyError: scopes = [] return scopes @@ -647,9 +612,7 @@ class JWTHandler: if cached_jwks_uri is not None: return cached_jwks_uri - verbose_proxy_logger.debug( - f"JWT Auth: Fetching OIDC discovery document from {url}" - ) + verbose_proxy_logger.debug(f"JWT Auth: Fetching OIDC discovery document from {url}") response = await self.http_handler.get(url) if response.status_code != 200: raise Exception( @@ -658,19 +621,13 @@ class JWTHandler: try: discovery = response.json() except Exception as e: - raise Exception( - f"JWT Auth: Failed to parse OIDC discovery document at {url}: {e}" - ) + raise Exception(f"JWT Auth: Failed to parse OIDC discovery document at {url}: {e}") jwks_uri = discovery.get("jwks_uri") if not jwks_uri: - raise Exception( - f"JWT Auth: OIDC discovery document at {url} does not contain a 'jwks_uri' field." - ) + raise Exception(f"JWT Auth: OIDC discovery document at {url} does not contain a 'jwks_uri' field.") - verbose_proxy_logger.debug( - f"JWT Auth: Resolved OIDC discovery {url} -> jwks_uri={jwks_uri}" - ) + verbose_proxy_logger.debug(f"JWT Auth: Resolved OIDC discovery {url} -> jwks_uri={jwks_uri}") await self.user_api_key_cache.async_set_cache( key=cache_key, value=jwks_uri, @@ -684,9 +641,7 @@ class JWTHandler: return 600 return litellm_jwtauth.public_key_ttl - async def _get_public_key_from_jwks_url( - self, jwks_url: str, kid: Optional[str] - ) -> dict: + async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: Optional[str]) -> dict: resolved_jwks_url = await self._resolve_jwks_url(jwks_url) cache_key = f"litellm_jwt_auth_keys_{resolved_jwks_url}" @@ -698,12 +653,8 @@ class JWTHandler: try: response_json = response.json() except Exception as e: - verbose_proxy_logger.error( - f"Error parsing response: {e}. Original Response: {response.text}" - ) - raise Exception( - f"Error parsing response: {e}. Check server logs for original response." - ) + verbose_proxy_logger.error(f"Error parsing response: {e}. Original Response: {response.text}") + raise Exception(f"Error parsing response: {e}. Check server logs for original response.") if "keys" in response_json: keys: JWKKeyValue = response_json["keys"] @@ -722,9 +673,7 @@ class JWTHandler: if public_key is not None: return cast(dict, public_key) - raise NoMatchingJWTPublicKeyError( - f"No matching public key found. keys={resolved_jwks_url}, kid={kid}" - ) + raise NoMatchingJWTPublicKeyError(f"No matching public key found. keys={resolved_jwks_url}, kid={kid}") async def get_public_key(self, kid: Optional[str]) -> dict: keys_url = os.getenv("JWT_PUBLIC_KEY_URL") @@ -736,26 +685,18 @@ class JWTHandler: for key_url in keys_url_list: try: - return await self._get_public_key_from_jwks_url( - jwks_url=key_url, kid=kid - ) + return await self._get_public_key_from_jwks_url(jwks_url=key_url, kid=kid) except NoMatchingJWTPublicKeyError as e: - verbose_proxy_logger.debug( - "JWT Auth: No matching public key found at %s: %s", key_url, e - ) + verbose_proxy_logger.debug("JWT Auth: No matching public key found at %s: %s", key_url, e) - raise NoMatchingJWTPublicKeyError( - f"No matching public key found. keys={keys_url_list}, kid={kid}" - ) + raise NoMatchingJWTPublicKeyError(f"No matching public key found. keys={keys_url_list}, kid={kid}") def parse_keys(self, keys: JWKKeyValue, kid: Optional[str]) -> Optional[JWTKeyItem]: public_key: Optional[JWTKeyItem] = None if len(keys) == 1: if isinstance(keys, dict) and (keys.get("kid", None) == kid or kid is None): public_key = keys - elif isinstance(keys, list) and ( - keys[0].get("kid", None) == kid or kid is None - ): + elif isinstance(keys, list) and (keys[0].get("kid", None) == kid or kid is None): public_key = keys[0] elif len(keys) > 1: for key in keys: @@ -763,12 +704,7 @@ class JWTHandler: key_kid = key.get("kid", None) else: key_kid = None - if ( - kid is not None - and isinstance(key, dict) - and key_kid is not None - and key_kid == kid - ): + if kid is not None and isinstance(key, dict) and key_kid is not None and key_kid == kid: public_key = key return public_key @@ -801,9 +737,7 @@ class JWTHandler: Exception: If UserInfo endpoint is not configured or request fails """ if not self.litellm_jwtauth.oidc_userinfo_endpoint: - raise Exception( - "OIDC UserInfo endpoint not configured. Set 'oidc_userinfo_endpoint' in JWT auth config." - ) + raise Exception("OIDC UserInfo endpoint not configured. Set 'oidc_userinfo_endpoint' in JWT auth config.") # Check cache first cache_key = f"oidc_userinfo_{hashlib.sha256(token.encode()).hexdigest()}" @@ -813,9 +747,7 @@ class JWTHandler: verbose_proxy_logger.debug("Returning cached OIDC UserInfo") return cached_userinfo - verbose_proxy_logger.debug( - f"Calling OIDC UserInfo endpoint: {self.litellm_jwtauth.oidc_userinfo_endpoint}" - ) + verbose_proxy_logger.debug(f"Calling OIDC UserInfo endpoint: {self.litellm_jwtauth.oidc_userinfo_endpoint}") try: # Call the UserInfo endpoint with the access token @@ -828,9 +760,7 @@ class JWTHandler: ) if response.status_code != 200: - raise Exception( - f"OIDC UserInfo endpoint returned status {response.status_code}: {response.text}" - ) + raise Exception(f"OIDC UserInfo endpoint returned status {response.status_code}: {response.text}") userinfo = response.json() verbose_proxy_logger.debug(f"Received OIDC UserInfo: {userinfo}") @@ -868,11 +798,7 @@ class JWTHandler: audience = os.getenv("JWT_AUDIENCE") issuer = os.getenv("JWT_ISSUER") - if ( - audience is None - and issuer is None - and not cls._unscoped_jwt_warning_emitted - ): + if audience is None and issuer is None and not cls._unscoped_jwt_warning_emitted: verbose_proxy_logger.warning( "JWT auth is enabled but neither JWT_AUDIENCE nor JWT_ISSUER " "is configured. Tokens minted by any application that shares " @@ -940,12 +866,8 @@ class JWTHandler: return None return claim_value - def _apply_issuer_claim_mappings( - self, token: dict, issuer_config: JWTIssuerConfig - ) -> dict: - normalized: dict = { - k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS - } + def _apply_issuer_claim_mappings(self, token: dict, issuer_config: JWTIssuerConfig) -> dict: + normalized: dict = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS} normalized[self.LITELLM_JWT_ISSUER_CLAIM] = issuer_config.issuer claim_mappings = [ (issuer_config.user_id_jwt_field, self.LITELLM_USER_ID_CLAIM), @@ -988,9 +910,7 @@ class JWTHandler: # validation. Require callers to opt in via # ``disable_audience_validation=True``. if audience is None and not disable_audience_validation: - raise ValueError( - "audience must be provided unless disable_audience_validation=True" - ) + raise ValueError("audience must be provided unless disable_audience_validation=True") options: dict = {} if audience is None: options["verify_aud"] = False @@ -1018,9 +938,7 @@ class JWTHandler: ) if isinstance(public_key, dict): - public_key_obj = PyJWK.from_dict( - self._get_jwk_from_public_key(public_key=public_key) - ).key + public_key_obj = PyJWK.from_dict(self._get_jwk_from_public_key(public_key=public_key)).key return jwt.decode( token, public_key_obj, # type: ignore @@ -1046,9 +964,7 @@ class JWTHandler: leeway=self.leeway, ) - async def _auth_jwt_with_issuer( - self, token: str, issuer_config: JWTIssuerConfig, kid: Optional[str] - ) -> dict: + async def _auth_jwt_with_issuer(self, token: str, issuer_config: JWTIssuerConfig, kid: Optional[str]) -> dict: public_key = await self._get_public_key_from_jwks_url( jwks_url=self._get_jwks_url_for_issuer(issuer_config=issuer_config), kid=kid, @@ -1104,11 +1020,7 @@ class JWTHandler: issuer=decode_kwargs["issuer"], options=decode_kwargs["options"], ) - return { - k: v - for k, v in payload.items() - if k not in self.LITELLM_INTERNAL_CLAIMS - } + return {k: v for k, v in payload.items() if k not in self.LITELLM_INTERNAL_CLAIMS} except jwt.ExpiredSignatureError: raise ProxyException( @@ -1138,9 +1050,7 @@ class JWTAuthManager: """ Checks if user is allowed to access the route, based on their role. """ - role_based_routes = get_role_based_routes( - rbac_role=rbac_role, general_settings=general_settings - ) + role_based_routes = get_role_based_routes(rbac_role=rbac_role, general_settings=general_settings) if role_based_routes is None or route is None: return True @@ -1167,9 +1077,7 @@ class JWTAuthManager: """ Checks if user is allowed to access the model, based on their role. """ - role_based_models = get_role_based_models( - rbac_role=rbac_role, general_settings=general_settings - ) + role_based_models = get_role_based_models(rbac_role=rbac_role, general_settings=general_settings) if role_based_models is None or model is None: return True @@ -1207,11 +1115,7 @@ class JWTAuthManager: if requested_model not in allowed_models: raise HTTPException( status_code=403, - detail={ - "error": "model={} not allowed. Allowed_models={}".format( - requested_model, allowed_models - ) - }, + detail={"error": "model={} not allowed. Allowed_models={}".format(requested_model, allowed_models)}, ) return None @@ -1264,9 +1168,7 @@ class JWTAuthManager: if not is_allowed: allowed_routes: List[Any] = jwt_handler.litellm_jwtauth.admin_allowed_routes actual_routes = get_actual_routes(allowed_routes=allowed_routes) - raise Exception( - f"Admin not allowed to access this route. Route={route}, Allowed Routes={actual_routes}" - ) + raise Exception(f"Admin not allowed to access this route. Route={route}, Allowed Routes={actual_routes}") return JWTAuthBuilderResult( is_proxy_admin=True, @@ -1293,9 +1195,7 @@ class JWTAuthManager: proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]: """Find and validate specific team ID from team_id_jwt_field or team_alias_jwt_field""" - individual_team_id = jwt_handler.get_team_id( - token=jwt_valid_token, default_value=None - ) + individual_team_id = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None) team_object: Optional[LiteLLM_TeamTable] = None @@ -1312,10 +1212,7 @@ class JWTAuthManager: ) return individual_team_id, team_object except HTTPException as e: - if ( - e.status_code != 404 - or not jwt_handler.litellm_jwtauth.team_claim_fallback - ): + if e.status_code != 404 or not jwt_handler.litellm_jwtauth.team_claim_fallback: raise # Claim doesn't map to a known team — defer to fallback. verbose_proxy_logger.debug( @@ -1326,13 +1223,9 @@ class JWTAuthManager: return None, None # If no team_id found, try to resolve via team_alias_jwt_field - team_alias = jwt_handler.get_team_alias( - token=jwt_valid_token, default_value=None - ) + team_alias = jwt_handler.get_team_alias(token=jwt_valid_token, default_value=None) if team_alias: - verbose_proxy_logger.info( - f"JWT Auth: Resolving team by alias: '{team_alias}'" - ) + verbose_proxy_logger.info(f"JWT Auth: Resolving team by alias: '{team_alias}'") team_object = await get_team_object_by_alias( team_alias=team_alias, prisma_client=prisma_client, @@ -1394,9 +1287,7 @@ class JWTAuthManager: route: str, request_method: Optional[str] = None, ) -> bool: - normalized_request_method = ( - request_method.upper() if isinstance(request_method, str) else None - ) + normalized_request_method = request_method.upper() if isinstance(request_method, str) else None if not RouteChecks.is_auth_enforced_pass_through_route( route=route, method=normalized_request_method, @@ -1407,9 +1298,7 @@ class JWTAuthManager: # so passthrough access is granted only by the selected team's metadata. return RouteChecks.check_passthrough_route_access( route=route, - user_api_key_dict=UserAPIKeyAuth( - team_metadata=(team_object.metadata or {}) if team_object else {} - ), + user_api_key_dict=UserAPIKeyAuth(team_metadata=(team_object.metadata or {}) if team_object else {}), ) @staticmethod @@ -1477,13 +1366,10 @@ class JWTAuthManager: user_route=route, litellm_proxy_roles=jwt_handler.litellm_jwtauth, ) - if ( - is_allowed - and not JWTAuthManager._team_has_passthrough_route_access( - team_object=team_object, - route=route, - request_method=request_method, - ) + if is_allowed and not JWTAuthManager._team_has_passthrough_route_access( + team_object=team_object, + route=route, + request_method=request_method, ): is_allowed = False denied_auth_enforced_pass_through_route = True @@ -1498,10 +1384,7 @@ class JWTAuthManager: if denied_auth_enforced_pass_through_route: JWTAuthManager._raise_team_passthrough_route_denial(route=route) - if requested_model and ( - any_claim_team_resolved - or not jwt_handler.litellm_jwtauth.team_claim_fallback - ): + if requested_model and (any_claim_team_resolved or not jwt_handler.litellm_jwtauth.team_claim_fallback): # Claim resolved but no model access, or fallback disabled — deny. raise HTTPException( status_code=403, @@ -1517,19 +1400,11 @@ class JWTAuthManager: jwt_valid_token: dict, ) -> Tuple[Optional[str], Optional[str], Optional[bool]]: """Get user email and validation status""" - user_email = jwt_handler.get_user_email( - token=jwt_valid_token, default_value=None - ) + user_email = jwt_handler.get_user_email(token=jwt_valid_token, default_value=None) valid_user_email = None if jwt_handler.is_enforced_email_domain(): - valid_user_email = ( - False - if user_email is None - else jwt_handler.is_allowed_domain(user_email=user_email) - ) - user_id = jwt_handler.get_user_id( - token=jwt_valid_token, default_value=user_email - ) + valid_user_email = False if user_email is None else jwt_handler.is_allowed_domain(user_email=user_email) + user_id = jwt_handler.get_user_id(token=jwt_valid_token, default_value=user_email) return user_id, user_email, valid_user_email @staticmethod @@ -1590,9 +1465,7 @@ class JWTAuthManager: else None ) elif org_alias: - verbose_proxy_logger.info( - f"JWT Auth: Resolving org by alias: '{org_alias}'" - ) + verbose_proxy_logger.info(f"JWT Auth: Resolving org by alias: '{org_alias}'") org_object = await get_org_object_by_alias( org_alias=org_alias, prisma_client=prisma_client, @@ -1621,9 +1494,7 @@ class JWTAuthManager: user_id=user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - user_id_upsert=jwt_handler.is_upsert_user_id( - valid_user_email=valid_user_email - ), + user_id_upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email), parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, user_email=user_email, @@ -1649,9 +1520,7 @@ class JWTAuthManager: ) # Rebind to resolved DB user_id for team_membership + auth_builder (GH #26789). - effective_user_id = JWTAuthManager._canonical_user_id_from_db( - user_id=user_id, user_object=user_object - ) + effective_user_id = JWTAuthManager._canonical_user_id_from_db(user_id=user_id, user_object=user_object) if effective_user_id != user_id: verbose_proxy_logger.debug( "JWT Auth: rebinding user_id %r -> DB user_id %r (email/sso match)", @@ -1734,9 +1603,7 @@ class JWTAuthManager: detail=f"Team '{header_team_id}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}", ) - verbose_proxy_logger.debug( - f"Using team_id from x-litellm-team-id header: {header_team_id}" - ) + verbose_proxy_logger.debug(f"Using team_id from x-litellm-team-id header: {header_team_id}") return header_team_id @staticmethod @@ -1777,9 +1644,7 @@ class JWTAuthManager: user_role=LitellmUserRoles.PROXY_ADMIN ), # [TODO]: expose an internal service role, for better tracking ) - verbose_proxy_logger.debug( - f"Successfully added user {user_object.user_id} to team {team_object.team_id}" - ) + verbose_proxy_logger.debug(f"Successfully added user {user_object.user_id} to team {team_object.team_id}") except ProxyException as e: if e.type == ProxyErrorTypes.team_member_already_in_team: verbose_proxy_logger.debug( @@ -1871,9 +1736,7 @@ class JWTAuthManager: are enforced when admins act on behalf of a team. Admin management routes ignore the header to preserve pre-existing bypass behavior. """ - header_team_id = ( - request_headers.get("x-litellm-team-id") if request_headers else None - ) + header_team_id = request_headers.get("x-litellm-team-id") if request_headers else None if not header_team_id or not RouteChecks.is_llm_api_route(route=route): return try: @@ -1972,12 +1835,8 @@ class JWTAuthManager: """Main authentication and authorization builder""" # Check if OIDC UserInfo endpoint is enabled, but fall back to standard # JWT auth if the token itself is a well-formed JWT (3-part structure). - if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt( - token=api_key - ): - verbose_proxy_logger.debug( - "OIDC UserInfo is enabled. Fetching user info from UserInfo endpoint." - ) + if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(token=api_key): + verbose_proxy_logger.debug("OIDC UserInfo is enabled. Fetching user info from UserInfo endpoint.") # Use the access token to fetch user info from OIDC UserInfo endpoint jwt_valid_token: dict = await jwt_handler.get_oidc_userinfo(token=api_key) else: @@ -2005,10 +1864,7 @@ class JWTAuthManager: # Check Scope Based Access scopes = jwt_handler.get_scopes(token=jwt_valid_token) - if ( - jwt_handler.litellm_jwtauth.enforce_scope_based_access - and jwt_handler.litellm_jwtauth.scope_mappings - ): + if jwt_handler.litellm_jwtauth.enforce_scope_based_access and jwt_handler.litellm_jwtauth.scope_mappings: JWTAuthManager.check_scope_based_access( scope_mappings=jwt_handler.litellm_jwtauth.scope_mappings, scopes=scopes, @@ -2019,15 +1875,11 @@ class JWTAuthManager: object_id = jwt_handler.get_object_id(token=jwt_valid_token, default_value=None) # Get basic user info - user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info( - jwt_handler, jwt_valid_token - ) + user_id, user_email, valid_user_email = await JWTAuthManager.get_user_info(jwt_handler, jwt_valid_token) # Get IDs org_id = jwt_handler.get_org_id(token=jwt_valid_token, default_value=None) - end_user_id = jwt_handler.get_end_user_id( - token=jwt_valid_token, default_value=None - ) + end_user_id = jwt_handler.get_end_user_id(token=jwt_valid_token, default_value=None) team_id: Optional[str] = None team_object: Optional[LiteLLM_TeamTable] = None object_id = jwt_handler.get_object_id(token=jwt_valid_token, default_value=None) @@ -2058,9 +1910,7 @@ class JWTAuthManager: # Get team with model access ## Check if team_id is specified via x-litellm-team-id header all_team_ids = JWTAuthManager.get_all_team_ids(jwt_handler, jwt_valid_token) - specific_team_id = jwt_handler.get_team_id( - token=jwt_valid_token, default_value=None - ) + specific_team_id = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None) if specific_team_id: all_team_ids.add(specific_team_id) @@ -2114,9 +1964,7 @@ class JWTAuthManager: and team_object is None and RouteChecks.is_auth_enforced_pass_through_route( route=route, - method=( - request_method.upper() if isinstance(request_method, str) else None - ), + method=(request_method.upper() if isinstance(request_method, str) else None), ) ): team_object = await get_team_object( @@ -2203,9 +2051,7 @@ class JWTAuthManager: ) # check if user is proxy admin - is_proxy_admin = bool( - user_object and user_object.user_role == LitellmUserRoles.PROXY_ADMIN - ) + is_proxy_admin = bool(user_object and user_object.user_role == LitellmUserRoles.PROXY_ADMIN) return JWTAuthBuilderResult( is_proxy_admin=is_proxy_admin, diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index f1ef05c7ad4..bd1cd596f9c 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -72,9 +72,7 @@ class IPAddressUtils: try: networks.append(ipaddress.ip_network(cidr, strict=False)) except ValueError: - verbose_proxy_logger.warning( - "Invalid CIDR in mcp_internal_ip_ranges: %s, skipping", cidr - ) + verbose_proxy_logger.warning("Invalid CIDR in mcp_internal_ip_ranges: %s, skipping", cidr) return networks if networks else IPAddressUtils._DEFAULT_INTERNAL_NETWORKS @staticmethod @@ -92,9 +90,7 @@ class IPAddressUtils: try: networks.append(ipaddress.ip_network(cidr, strict=False)) except ValueError: - verbose_proxy_logger.warning( - "Invalid CIDR in mcp_trusted_proxy_ranges: %s, skipping", cidr - ) + verbose_proxy_logger.warning("Invalid CIDR in mcp_trusted_proxy_ranges: %s, skipping", cidr) return networks @staticmethod @@ -114,9 +110,7 @@ class IPAddressUtils: @staticmethod def is_internal_ip( client_ip: Optional[str], - internal_networks: Optional[ - List[Union[ipaddress.IPv4Network, ipaddress.IPv6Network]] - ] = None, + internal_networks: Optional[List[Union[ipaddress.IPv4Network, ipaddress.IPv6Network]]] = None, ) -> bool: """ Check if a client IP is from an internal/private network. @@ -312,24 +306,18 @@ class IPAddressUtils: # If XFF is enabled, validate the request comes from a trusted proxy if use_xff and "x-forwarded-for" in request.headers: - if not IPAddressUtils.is_request_from_trusted_proxy( - request, general_settings=general_settings - ): + if not IPAddressUtils.is_request_from_trusted_proxy(request, general_settings=general_settings): direct_ip = request.client.host if request.client else None if general_settings.get("mcp_trusted_proxy_ranges"): # Direct connection isn't in any configured trusted CIDR. - verbose_proxy_logger.warning( - "XFF header from untrusted IP %s, ignoring", direct_ip - ) + verbose_proxy_logger.warning("XFF header from untrusted IP %s, ignoring", direct_ip) return direct_ip # XFF enabled but no trusted proxy ranges configured: the direct # peer is typically the reverse proxy's own (private) IP, so # returning it would mis-classify external callers as internal. # Fail closed for access control. return "" - match IPAddressUtils._resolve_num_trusted_hops( - general_settings.get("mcp_xff_num_trusted_hops") - ): + match IPAddressUtils._resolve_num_trusted_hops(general_settings.get("mcp_xff_num_trusted_hops")): case _HopCountInvalid(): return "" case _HopCount(value=num_trusted_hops): diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index ec2c1eb8e19..2bb4ae0b354 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -81,9 +81,7 @@ class LicenseCheck: assert isinstance(premium, bool) verbose_proxy_logger.debug( - "litellm.proxy.auth.litellm_license.py::_verify - License={} is premium={}".format( - license_str, premium - ) + "litellm.proxy.auth.litellm_license.py::_verify - License={} is premium={}".format(license_str, premium) ) return premium except Exception as e: @@ -121,9 +119,7 @@ class LicenseCheck: if self.license_str is None: return False elif ( - self.verify_license_without_api_request( - public_key=self.public_key, license_key=self.license_str - ) + self.verify_license_without_api_request(public_key=self.public_key, license_key=self.license_str) is True ): return True @@ -152,12 +148,8 @@ class LicenseCheck: if self.airgapped_license_data is None: return False - _max_teams_in_license: Optional[int] = self.airgapped_license_data.get( - "max_teams" - ) - if "max_teams" not in self.airgapped_license_data or not isinstance( - _max_teams_in_license, int - ): + _max_teams_in_license: Optional[int] = self.airgapped_license_data.get("max_teams") + if "max_teams" not in self.airgapped_license_data or not isinstance(_max_teams_in_license, int): return False return team_count > _max_teams_in_license @@ -197,9 +189,7 @@ class LicenseCheck: verbose_proxy_logger.debug("License data: %s", license_data) # Check expiration date - expiration_date = datetime.strptime( - license_data["expiration_date"], "%Y-%m-%d" - ) + expiration_date = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") if expiration_date < datetime.now(): return False, "License has expired" diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index bd2e7560430..11f12e597b9 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -162,9 +162,9 @@ async def authenticate_user( - Login with UI_USERNAME and UI_PASSWORD - Login with Invite Link `user_email` and `password` combination """ - if secrets.compare_digest( - username.encode("utf-8"), ui_username.encode("utf-8") - ) and secrets.compare_digest(password.encode("utf-8"), ui_password.encode("utf-8")): + if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest( + password.encode("utf-8"), ui_password.encode("utf-8") + ): # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin user_role = LitellmUserRoles.PROXY_ADMIN user_id = LITELLM_PROXY_ADMIN_NAME @@ -172,8 +172,7 @@ async def authenticate_user( # we want the key created to have PROXY_ADMIN_PERMISSIONS key_user_id = LITELLM_PROXY_ADMIN_NAME if ( - os.getenv("PROXY_ADMIN_ID", None) is not None - and os.environ["PROXY_ADMIN_ID"] == user_id + os.getenv("PROXY_ADMIN_ID", None) is not None and os.environ["PROXY_ADMIN_ID"] == user_id ) or user_id == LITELLM_PROXY_ADMIN_NAME: # checks if user is admin key_user_id = os.getenv("PROXY_ADMIN_ID", LITELLM_PROXY_ADMIN_NAME) @@ -222,9 +221,7 @@ async def authenticate_user( user_info: Optional[LiteLLM_UserTable] = None if _user_row is not None: user_info = _user_row - elif ( - user_id is not None - ): # if user_id is not None, we are using the UI_USERNAME and UI_PASSWORD + elif user_id is not None: # if user_id is not None, we are using the UI_USERNAME and UI_PASSWORD user_info = LiteLLM_UserTable( user_id=user_id, user_role=user_role, @@ -234,14 +231,10 @@ async def authenticate_user( if user_info is None: raise HTTPException( status_code=401, - detail={ - "error": "User Information is required for experimental UI login" - }, + detail={"error": "User Information is required for experimental UI login"}, ) - key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - user_info - ) + key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(user_info) return LoginResult( user_id=user_id, @@ -258,9 +251,7 @@ async def authenticate_user( -> if the user has no role in the DB assume they are only a viewer """ user_id = getattr(_user_row, "user_id", "unknown") - user_role = getattr( - _user_row, "user_role", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY - ) + user_role = getattr(_user_row, "user_role", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) user_email = getattr(_user_row, "user_email", "unknown") _password = getattr(_user_row, "password", "unknown") @@ -338,9 +329,7 @@ def create_ui_token_object( Returns: ReturnedUITokenObject: Token object ready for JWT encoding """ - disabled_non_admin_personal_key_creation = ( - get_disabled_non_admin_personal_key_creation() - ) + disabled_non_admin_personal_key_creation = get_disabled_non_admin_personal_key_creation() return ReturnedUITokenObject( user_id=login_result.user_id, @@ -349,9 +338,7 @@ def create_ui_token_object( user_role=login_result.user_role, login_method=login_result.login_method, premium_user=premium_user, - auth_header_name=general_settings.get( - "litellm_key_header_name", "Authorization" - ), + auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"), disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, server_root_path=get_server_root_path(), ) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 5d5ab4f224f..0f3b8b5a412 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -34,9 +34,7 @@ def _check_wildcard_routing(model: str) -> bool: return False -def get_provider_models( - provider: str, litellm_params: Optional[LiteLLM_Params] = None -) -> Optional[List[str]]: +def get_provider_models(provider: str, litellm_params: Optional[LiteLLM_Params] = None) -> Optional[List[str]]: """ Returns the list of known models by provider """ @@ -44,9 +42,7 @@ def get_provider_models( return get_valid_models(litellm_params=litellm_params) if provider in litellm.models_by_provider: - provider_models = get_valid_models( - custom_llm_provider=provider, litellm_params=litellm_params - ) + provider_models = get_valid_models(custom_llm_provider=provider, litellm_params=litellm_params) return provider_models return None @@ -60,9 +56,7 @@ def _get_models_from_access_groups( new_models = [] for idx, model in enumerate(all_models): if model in model_access_groups: - if ( - not include_model_access_groups - ): # remove access group, unless requested - e.g. when creating a key + if not include_model_access_groups: # remove access group, unless requested - e.g. when creating a key idx_to_remove.append(idx) new_models.extend(model_access_groups[model]) @@ -116,20 +110,11 @@ def get_key_models( """ all_models: List[str] = [] if len(user_api_key_dict.models) > 0: - all_models = list( - user_api_key_dict.models - ) # copy to avoid mutating cached objects - if ( - SpecialModelNames.all_team_models.value in all_models - and user_api_key_dict.team_id is not None - ): + all_models = list(user_api_key_dict.models) # copy to avoid mutating cached objects + if SpecialModelNames.all_team_models.value in all_models and user_api_key_dict.team_id is not None: all_models = list(user_api_key_dict.team_models) if SpecialModelNames.all_team_models.value in all_models: - all_models = [ - model - for model in all_models - if model != SpecialModelNames.all_team_models.value - ] + all_models = [model for model in all_models if model != SpecialModelNames.all_team_models.value] all_models.extend(proxy_model_list) if include_model_access_groups: all_models.extend(model_access_groups.keys()) @@ -262,26 +247,19 @@ def _hydrate_litellm_credential_name( if litellm_params is None or litellm_params.litellm_credential_name is None: return litellm_params - credential_values = CredentialAccessor.get_credential_values( - litellm_params.litellm_credential_name - ) + credential_values = CredentialAccessor.get_credential_values(litellm_params.litellm_credential_name) if not credential_values: return litellm_params litellm_params = litellm_params.model_copy() for key, value in credential_values.items(): - if ( - key in _CREDENTIAL_LITELLM_PARAM_FIELDS - and getattr(litellm_params, key, None) is None - ): + if key in _CREDENTIAL_LITELLM_PARAM_FIELDS and getattr(litellm_params, key, None) is None: setattr(litellm_params, key, value) litellm_params.litellm_credential_name = None return litellm_params -def get_known_models_from_wildcard( - wildcard_model: str, litellm_params: Optional[LiteLLM_Params] = None -) -> List[str]: +def get_known_models_from_wildcard(wildcard_model: str, litellm_params: Optional[LiteLLM_Params] = None) -> List[str]: wildcard_model_to_expand = ( litellm_params.model if wildcard_model == "*" @@ -291,9 +269,7 @@ def get_known_models_from_wildcard( else wildcard_model ) try: - wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split( - "/", 1 - ) + wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split("/", 1) except ValueError: # safely fail return [] @@ -309,9 +285,7 @@ def get_known_models_from_wildcard( litellm_params = _hydrate_litellm_credential_name(litellm_params) - wildcard_models = get_provider_models( - provider=provider, litellm_params=litellm_params - ) + wildcard_models = get_provider_models(provider=provider, litellm_params=litellm_params) if wildcard_models is None: return [] @@ -319,15 +293,9 @@ def get_known_models_from_wildcard( ## CHECK IF PARTIAL FILTER e.g. `gemini-*` model_prefix = wildcard_suffix.replace("*", "") - is_partial_filter = any( - wc_model.startswith(model_prefix) for wc_model in wildcard_models - ) + is_partial_filter = any(wc_model.startswith(model_prefix) for wc_model in wildcard_models) if is_partial_filter: - filtered_wildcard_models = [ - wc_model - for wc_model in wildcard_models - if wc_model.startswith(model_prefix) - ] + filtered_wildcard_models = [wc_model for wc_model in wildcard_models if wc_model.startswith(model_prefix)] wildcard_models = filtered_wildcard_models else: # add model prefix to wildcard models @@ -366,9 +334,7 @@ def expand_wildcard_deployments_for_model_info( for deployment in deployments: model_name = str(deployment.get("model_name") or "") raw_params = deployment.get("litellm_params") - litellm_params_dict: dict[str, Any] = ( - raw_params if isinstance(raw_params, dict) else {} - ) + litellm_params_dict: dict[str, Any] = raw_params if isinstance(raw_params, dict) else {} litellm_model = str(litellm_params_dict.get("model") or "") # Determine the wildcard pattern to expand. @@ -376,9 +342,7 @@ def expand_wildcard_deployments_for_model_info( # also a wildcard, so a concrete model_name is never overwritten. if _check_wildcard_routing(model_name) and "/" in model_name: wildcard_pattern = model_name - elif _check_wildcard_routing(model_name) and _check_wildcard_routing( - litellm_model - ): + elif _check_wildcard_routing(model_name) and _check_wildcard_routing(litellm_model): wildcard_pattern = litellm_model elif _check_wildcard_routing(model_name): wildcard_pattern = model_name @@ -387,11 +351,7 @@ def expand_wildcard_deployments_for_model_info( continue try: - litellm_params = ( - LiteLLM_Params.model_validate(litellm_params_dict) - if litellm_params_dict - else None - ) + litellm_params = LiteLLM_Params.model_validate(litellm_params_dict) if litellm_params_dict else None except Exception: expanded.append(deployment) continue @@ -424,16 +384,12 @@ def _get_wildcard_models( all_wildcard_models = [] for model in unique_models: if _check_wildcard_routing(model=model): - if ( - return_wildcard_routes - ): # will add the wildcard route to the list eg: anthropic/*. + if return_wildcard_routes: # will add the wildcard route to the list eg: anthropic/*. all_wildcard_models.append(model) ## get litellm params from model if llm_router is not None: - model_list = llm_router.get_model_list( - model_name=model, team_id=team_id - ) + model_list = llm_router.get_model_list(model_name=model, team_id=team_id) if model_list: for router_model in model_list: wildcard_models = get_known_models_from_wildcard( @@ -446,17 +402,13 @@ def _get_wildcard_models( else: # Router has no deployment for this wildcard (e.g., BYOK team models) # Fall back to expanding from known provider models - wildcard_models = get_known_models_from_wildcard( - wildcard_model=model, litellm_params=None - ) + wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None) if wildcard_models: models_to_remove.add(model) all_wildcard_models.extend(wildcard_models) else: # get all known provider models - wildcard_models = get_known_models_from_wildcard( - wildcard_model=model, litellm_params=None - ) + wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None) if wildcard_models: models_to_remove.add(model) @@ -504,9 +456,7 @@ def get_all_fallbacks( try: # Use existing function to get fallback model group - fallback_model_group, _ = get_fallback_model_group( - fallbacks=fallbacks_config, model_group=model - ) + fallback_model_group, _ = get_fallback_model_group(fallbacks=fallbacks_config, model_group=model) if fallback_model_group is None: return [] diff --git a/litellm/proxy/auth/network.py b/litellm/proxy/auth/network.py index 4eb6f1dcec2..1ff1b11de96 100644 --- a/litellm/proxy/auth/network.py +++ b/litellm/proxy/auth/network.py @@ -22,9 +22,7 @@ class TrustedProxyConfig(BaseModel): trusted_proxy_cidrs: list[str] = Field(default_factory=list) -def normalize_cidr_ranges( - configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs" -) -> list[str]: +def normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs") -> list[str]: if not configured_ranges: return [] if isinstance(configured_ranges, str): @@ -47,9 +45,7 @@ def parse_trusted_proxy_ranges( try: networks.append(ipaddress.ip_network(cidr, strict=False)) except ValueError: - verbose_proxy_logger.warning( - "Invalid CIDR in %s: %s, skipping", setting_name, cidr - ) + verbose_proxy_logger.warning("Invalid CIDR in %s: %s, skipping", setting_name, cidr) return networks @@ -71,9 +67,7 @@ def _is_valid_ip(value: str) -> bool: return False -def resolve_client_ip( - request: Request, config: TrustedProxyConfig -) -> tuple[str | None, bool]: +def resolve_client_ip(request: Request, config: TrustedProxyConfig) -> tuple[str | None, bool]: """Resolve the real client IP, trusting X-Forwarded-For only when the direct peer is itself a configured trusted proxy. Walks the header right-to-left and returns the first hop that is not a trusted proxy, so a forged left-most entry @@ -90,9 +84,7 @@ def resolve_client_ip( return peer, True -def resolve_network_context( - request: Request, config: TrustedProxyConfig -) -> NetworkContext: +def resolve_network_context(request: Request, config: TrustedProxyConfig) -> NetworkContext: ip, via_proxy = resolve_client_ip(request, config) return NetworkContext( client_ip=ip, diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index 10b1759b77e..a77fc510bb7 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -63,9 +63,7 @@ class Oauth2Handler: # Add client authentication if credentials are provided if oauth_client_id and oauth_client_secret: # Use HTTP Basic authentication for client credentials - credentials = base64.b64encode( - f"{oauth_client_id}:{oauth_client_secret}".encode() - ).decode() + credentials = base64.b64encode(f"{oauth_client_id}:{oauth_client_secret}".encode()).decode() headers["Authorization"] = f"Basic {credentials}" elif oauth_client_id: # For public clients, include client_id in the request body @@ -132,30 +130,23 @@ class Oauth2Handler: if premium_user is not True: raise ValueError( - "Oauth2 token validation is only available for premium users" - + CommonProxyErrors.not_premium_user.value + "Oauth2 token validation is only available for premium users" + CommonProxyErrors.not_premium_user.value ) - verbose_proxy_logger.debug( - "Oauth2 token validation for token=[set=%s]", token is not None - ) + verbose_proxy_logger.debug("Oauth2 token validation for token=[set=%s]", token is not None) # Get the token info endpoint from environment variable token_info_endpoint = os.getenv("OAUTH_TOKEN_INFO_ENDPOINT") user_id_field_name = os.environ.get("OAUTH_USER_ID_FIELD_NAME", "sub") user_role_field_name = os.environ.get("OAUTH_USER_ROLE_FIELD_NAME", "role") - user_team_id_field_name = os.environ.get( - "OAUTH_USER_TEAM_ID_FIELD_NAME", "team_id" - ) + user_team_id_field_name = os.environ.get("OAUTH_USER_TEAM_ID_FIELD_NAME", "team_id") # OAuth2 client credentials for introspection endpoint authentication oauth_client_id = os.environ.get("OAUTH_CLIENT_ID") oauth_client_secret = os.environ.get("OAUTH_CLIENT_SECRET") if not token_info_endpoint: - raise ValueError( - "OAUTH_TOKEN_INFO_ENDPOINT environment variable is not set" - ) + raise ValueError("OAUTH_TOKEN_INFO_ENDPOINT environment variable is not set") client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) @@ -177,9 +168,7 @@ class Oauth2Handler: oauth_client_secret=oauth_client_secret, ) - response = await client.post( - token_info_endpoint, headers=headers, data=data - ) + response = await client.post(token_info_endpoint, headers=headers, data=data) else: # Generic token info endpoint - uses GET with Bearer token verbose_proxy_logger.debug("Using generic token info endpoint (GET)") diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index 9fc4c4fb531..2b0593d3618 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -64,17 +64,13 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: feature_name="OAuth2 proxy auth", ) - oauth2_config_mappings: Dict[str, str] = ( - general_settings.get("oauth2_config_mappings") or {} - ) + oauth2_config_mappings: Dict[str, str] = general_settings.get("oauth2_config_mappings") or {} verbose_proxy_logger.debug(f"Oauth2 config mappings: {oauth2_config_mappings}") if not oauth2_config_mappings: raise ValueError("Oauth2 config mappings not found in general_settings") - disallowed = sorted( - set(oauth2_config_mappings.keys()) - ALLOWED_OAUTH2_PROXY_FIELDS - ) + disallowed = sorted(set(oauth2_config_mappings.keys()) - ALLOWED_OAUTH2_PROXY_FIELDS) if disallowed: raise ValueError( "Oauth2 proxy auth refuses to map non-identity UserAPIKeyAuth " diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py index 053cdb91f17..2ef66d3b7ff 100644 --- a/litellm/proxy/auth/rds_iam_token.py +++ b/litellm/proxy/auth/rds_iam_token.py @@ -71,11 +71,7 @@ def init_rds_client( config = boto3.session.Config() # type: ignore ### CHECK STS ### - if ( - aws_web_identity_token is not None - and aws_role_name is not None - and aws_session_name is not None - ): + if aws_web_identity_token is not None and aws_role_name is not None and aws_session_name is not None: try: oidc_token = open(aws_web_identity_token).read() # check if filepath except Exception: @@ -114,9 +110,7 @@ def init_rds_client( aws_secret_access_key=aws_secret_access_key, ) - sts_response = sts_client.assume_role( - RoleArn=aws_role_name, RoleSessionName=aws_session_name - ) + sts_response = sts_client.assume_role(RoleArn=aws_role_name, RoleSessionName=aws_session_name) client = boto3.client( service_name="rds", @@ -159,9 +153,7 @@ def init_rds_client( return client -def generate_iam_auth_token( - db_host, db_port, db_user, client: Optional[Any] = None -) -> str: +def generate_iam_auth_token(db_host, db_port, db_user, client: Optional[Any] = None) -> str: from urllib.parse import quote if client is None: @@ -172,16 +164,12 @@ def generate_iam_auth_token( aws_session_name=os.getenv("AWS_SESSION_NAME"), aws_profile_name=os.getenv("AWS_PROFILE_NAME"), aws_role_name=os.getenv("AWS_ROLE_NAME", os.getenv("AWS_ROLE_ARN")), - aws_web_identity_token=os.getenv( - "AWS_WEB_IDENTITY_TOKEN", os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") - ), + aws_web_identity_token=os.getenv("AWS_WEB_IDENTITY_TOKEN", os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")), ) else: boto_client = client - token = boto_client.generate_db_auth_token( - DBHostname=db_host, Port=db_port, DBUsername=db_user - ) + token = boto_client.generate_db_auth_token(DBHostname=db_host, Port=db_port, DBUsername=db_user) cleaned_token = quote(token, safe="") return cleaned_token diff --git a/litellm/proxy/auth/resolvers/exceptions.py b/litellm/proxy/auth/resolvers/exceptions.py index dd953e66659..a2ece57209b 100644 --- a/litellm/proxy/auth/resolvers/exceptions.py +++ b/litellm/proxy/auth/resolvers/exceptions.py @@ -11,16 +11,12 @@ class IdentityResolutionError(Exception): class NoDatabaseConnectionError(IdentityResolutionError): def __init__(self) -> None: - super().__init__( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + super().__init__("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") class KeyNotInCacheError(IdentityResolutionError): def __init__(self, hashed_token: str) -> None: - super().__init__( - f"Key doesn't exist in cache + check_cache_only=True. key={hashed_token}." - ) + super().__init__(f"Key doesn't exist in cache + check_cache_only=True. key={hashed_token}.") class KeyNotFoundError(IdentityResolutionError, ProxyException): @@ -44,7 +40,4 @@ class KeyNotFoundError(IdentityResolutionError, ProxyException): class PrincipalMissingSourceKeyError(IdentityResolutionError): def __init__(self) -> None: - super().__init__( - "Principal carries no source key; it was not produced by " - "IdentityStore.resolve" - ) + super().__init__("Principal carries no source key; it was not produced by IdentityStore.resolve") diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py index c43d4c84ca4..ad9fd234163 100644 --- a/litellm/proxy/auth/resolvers/store.py +++ b/litellm/proxy/auth/resolvers/store.py @@ -102,9 +102,7 @@ class IdentityStore: if self._prisma is None: raise NoDatabaseConnectionError() - cached = await self._cache.async_get_cache( - key=hashed_token, model_type=UserAPIKeyAuth - ) + cached = await self._cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth) if cached is not None: return _copy_user_api_key_auth_for_cache(user_api_key_obj=cached) @@ -163,33 +161,17 @@ class IdentityStore: """ teams: list[TeamIdentity] = [] if key.team_id is not None: - role = ( - team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER - ) + role = team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER teams.append(TeamIdentity(id=key.team_id, name=key.team_alias, role=role)) organization = ( - OrganizationIdentity(id=key.org_id, name=key.organization_alias) - if key.org_id is not None - else None - ) - user = ( - UserIdentity(id=key.user_id, email=key.user_email) - if key.user_id is not None - else None - ) - project = ( - ProjectIdentity(id=key.project_id, name=key.project_alias) - if key.project_id is not None - else None - ) - end_user = ( - EndUserIdentity(id=key.end_user_id) if key.end_user_id is not None else None + OrganizationIdentity(id=key.org_id, name=key.organization_alias) if key.org_id is not None else None ) + user = UserIdentity(id=key.user_id, email=key.user_email) if key.user_id is not None else None + project = ProjectIdentity(id=key.project_id, name=key.project_alias) if key.project_id is not None else None + end_user = EndUserIdentity(id=key.end_user_id) if key.end_user_id is not None else None mapped = map_role(key.user_role) return Principal( - principal_type=( - PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT - ), + principal_type=(PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT), subject=key.user_id or key.key_alias or subject_fallback or "", issuer=issuer, user=user, diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 4be2a185ce4..dd0a34a7898 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -59,9 +59,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES = frozenset( # paths directly because the request route carries the resolved key id. _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES = ("/regenerate", "/reset_spend") -_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS = frozenset( - ("openai_routes", "llm_api_routes") -) +_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS = frozenset(("openai_routes", "llm_api_routes")) class RouteChecks: @@ -84,9 +82,7 @@ class RouteChecks: pass # Check if Virtual Key is allowed to call the route - Applies to all Roles - RouteChecks.is_virtual_key_allowed_to_call_route( - route=route, valid_token=valid_token, request=request - ) + RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token, request=request) return True @staticmethod @@ -111,16 +107,11 @@ class RouteChecks: # explicit check for allowed routes (exact match or prefix match) for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_allowed_route( - route=route, allowed_route=allowed_route - ): + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): return True ## check if 'allowed_route' is a field name in LiteLLMRoutes - if any( - allowed_route in LiteLLMRoutes._member_names_ - for allowed_route in valid_token.allowed_routes - ): + if any(allowed_route in LiteLLMRoutes._member_names_ for allowed_route in valid_token.allowed_routes): for allowed_route in valid_token.allowed_routes: if allowed_route in LiteLLMRoutes._member_names_: if RouteChecks.check_route_access( @@ -134,9 +125,7 @@ class RouteChecks: method=RouteChecks._get_request_method(request=request), ) ): - if RouteChecks.check_passthrough_route_access( - route=route, user_api_key_dict=valid_token - ): + if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): return True denied_auth_enforced_pass_through_route = True else: @@ -150,9 +139,7 @@ class RouteChecks: InitPassThroughEndpointHelpers, ) - if InitPassThroughEndpointHelpers.is_registered_pass_through_route( - route=route - ): + if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): if RouteChecks.is_auth_enforced_pass_through_route( route=route, method=RouteChecks._get_request_method(request=request), @@ -175,16 +162,12 @@ class RouteChecks: # headers, env, credentials). POST/PUT/DELETE on # these paths are admin-only management writes and # are intentionally not covered. - if RouteChecks._is_get_mcp_server_discovery_route( - route=route, request=request - ): + if RouteChecks._is_get_mcp_server_discovery_route(route=route, request=request): return True # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_wildcard_pattern( - route=route, pattern=allowed_route - ): + if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): return True if denied_auth_enforced_pass_through_route: @@ -264,9 +247,7 @@ class RouteChecks: route=route, method=RouteChecks._get_request_method(request=request), ): - RouteChecks._require_auth_pass_through_access( - route=route, valid_token=valid_token - ) + RouteChecks._require_auth_pass_through_access(route=route, valid_token=valid_token) elif RouteChecks.is_llm_api_route(route=route): pass elif RouteChecks.is_info_route(route=route): @@ -278,9 +259,7 @@ class RouteChecks: # check if user can access this route query_params = request.query_params user_id = query_params.get("user_id") - verbose_proxy_logger.debug( - f"user_id: {user_id} & valid_token.user_id: {valid_token.user_id}" - ) + verbose_proxy_logger.debug(f"user_id: {user_id} & valid_token.user_id: {valid_token.user_id}") if user_id and user_id != valid_token.user_id: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -309,25 +288,17 @@ class RouteChecks: request_data=request_data, request=request, ) - elif ( - _user_role == LitellmUserRoles.INTERNAL_USER.value - and RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value - ) + elif _user_role == LitellmUserRoles.INTERNAL_USER.value and RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value ): pass - elif _user_is_org_admin( - request_data=request_data, user_object=user_obj - ) and RouteChecks.check_route_access( + elif _user_is_org_admin(request_data=request_data, user_object=user_obj) and RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value ): pass - elif ( - _user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - and RouteChecks.check_route_access( - route=route, - allowed_routes=LiteLLMRoutes.internal_user_view_only_routes.value, - ) + elif _user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value and RouteChecks.check_route_access( + route=route, + allowed_routes=LiteLLMRoutes.internal_user_view_only_routes.value, ): pass elif RouteChecks.check_route_access( @@ -336,34 +307,24 @@ class RouteChecks: pass elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"): pass # authN/authZ handled by api itself - elif RouteChecks.check_passthrough_route_access( - route=route, user_api_key_dict=valid_token - ): + elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): pass elif valid_token.allowed_routes is not None: # check if route is in allowed_routes (exact match or prefix match) route_allowed = False for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_allowed_route( - route=route, allowed_route=allowed_route - ): + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): route_allowed = True break - if RouteChecks._route_matches_wildcard_pattern( - route=route, pattern=allowed_route - ): + if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): route_allowed = True break if not route_allowed: - RouteChecks._raise_admin_only_route_exception( - user_obj=user_obj, route=route - ) + RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) else: - RouteChecks._raise_admin_only_route_exception( - user_obj=user_obj, route=route - ) + RouteChecks._raise_admin_only_route_exception(user_obj=user_obj, route=route) @staticmethod def custom_admin_only_route_check(route: str): @@ -405,14 +366,10 @@ class RouteChecks: if route in LiteLLMRoutes.google_routes.value: return True - if RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value - ): + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value): return True - if RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.agent_routes.value - ): + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_routes.value): return True if route in LiteLLMRoutes.litellm_native_routes.value: @@ -424,47 +381,35 @@ class RouteChecks: # Replace placeholders with regex pattern # placeholders are written as "/threads/{thread_id}" if "{" in openai_route: - if RouteChecks._route_matches_pattern( - route=route, pattern=openai_route - ): + if RouteChecks._route_matches_pattern(route=route, pattern=openai_route): return True # Check for wildcard patterns like "/containers/*" if RouteChecks._is_wildcard_pattern(pattern=openai_route): - if RouteChecks._route_matches_wildcard_pattern( - route=route, pattern=openai_route - ): + if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=openai_route): return True # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" for google_route in LiteLLMRoutes.google_routes.value: if "{" in google_route: - if RouteChecks._route_matches_pattern( - route=route, pattern=google_route - ): + if RouteChecks._route_matches_pattern(route=route, pattern=google_route): return True # Check for Anthropic routes with placeholders for anthropic_route in LiteLLMRoutes.anthropic_routes.value: if "{" in anthropic_route: - if RouteChecks._route_matches_pattern( - route=route, pattern=anthropic_route - ): + if RouteChecks._route_matches_pattern(route=route, pattern=anthropic_route): return True if RouteChecks._is_azure_openai_route(route=route): return True for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value: - if route == _llm_passthrough_route or route.startswith( - _llm_passthrough_route + "/" - ): + if route == _llm_passthrough_route or route.startswith(_llm_passthrough_route + "/"): return True return False @staticmethod - def _is_get_mcp_server_discovery_route( - route: str, request: Optional[Request] - ) -> bool: + def _is_get_mcp_server_discovery_route(route: str, request: Optional[Request]) -> bool: """ Returns True if `request` is a GET against one of the two read-only MCP-server discovery paths: @@ -491,9 +436,7 @@ class RouteChecks: """ Check if route is a management route """ - return RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.management_routes.value - ) + return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value) @staticmethod def is_info_route(route: str) -> bool: @@ -638,15 +581,9 @@ class RouteChecks: # wildcard match route is in allowed_routes # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* ######################################################### - wildcard_allowed_routes = [ - route - for route in allowed_routes - if RouteChecks._is_wildcard_pattern(pattern=route) - ] + wildcard_allowed_routes = [route for route in allowed_routes if RouteChecks._is_wildcard_pattern(pattern=route)] for allowed_route in wildcard_allowed_routes: - if RouteChecks._route_matches_wildcard_pattern( - route=route, pattern=allowed_route - ): + if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): return True ######################################################### @@ -656,8 +593,7 @@ class RouteChecks: # returns: True ######################################################### if any( # Check pattern match - RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) - for allowed_route in allowed_routes + RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes ): return True @@ -678,9 +614,7 @@ class RouteChecks: return method.upper() @staticmethod - def is_auth_enforced_pass_through_route( - route: str, method: Optional[str] = None - ) -> bool: + def is_auth_enforced_pass_through_route(route: str, method: Optional[str] = None) -> bool: """ True for config/DB pass-through endpoints registered with auth=true. @@ -692,9 +626,7 @@ class RouteChecks: InitPassThroughEndpointHelpers, ) - route_info = InitPassThroughEndpointHelpers.get_registered_pass_through_route( - route=route, method=method - ) + route_info = InitPassThroughEndpointHelpers.get_registered_pass_through_route(route=route, method=method) if route_info is None: return False return route_info.get("auth") is True @@ -717,16 +649,12 @@ class RouteChecks: """ Require an explicit ``allowed_passthrough_routes`` match for auth=true pass-through. """ - if RouteChecks.check_passthrough_route_access( - route=route, user_api_key_dict=valid_token - ): + if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): return raise RouteChecks._auth_pass_through_denied_exception(route=route) @staticmethod - def check_passthrough_route_access( - route: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: + def check_passthrough_route_access(route: str, user_api_key_dict: UserAPIKeyAuth) -> bool: """ Check if route is a passthrough route. Supports both exact match and prefix match. @@ -735,10 +663,7 @@ class RouteChecks: team_metadata = user_api_key_dict.team_metadata or {} if metadata is None and team_metadata is None: return False - if ( - "allowed_passthrough_routes" not in metadata - and "allowed_passthrough_routes" not in team_metadata - ): + if "allowed_passthrough_routes" not in metadata and "allowed_passthrough_routes" not in team_metadata: return False if ( metadata.get("allowed_passthrough_routes") is None @@ -747,16 +672,12 @@ class RouteChecks: return False allowed_passthrough_routes = ( - metadata.get("allowed_passthrough_routes") - or team_metadata.get("allowed_passthrough_routes") - or [] + metadata.get("allowed_passthrough_routes") or team_metadata.get("allowed_passthrough_routes") or [] ) # Check if route matches any allowed passthrough route (exact or prefix match) for allowed_route in allowed_passthrough_routes: - if RouteChecks._route_matches_allowed_route( - route=route, allowed_route=allowed_route - ): + if RouteChecks._route_matches_allowed_route(route=route, allowed_route=allowed_route): return True return False @@ -862,9 +783,7 @@ class RouteChecks: ) # Check if this is a write operation on management routes - if RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.management_routes.value - ): + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value): # For management routes, only allow read operations or specific allowed updates if route == "/user/update": # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY @@ -877,8 +796,7 @@ class RouteChecks: detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated", ) elif route in _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES or ( - route.startswith("/key/") - and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) + route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES) ): # Block write operations for PROXY_ADMIN_VIEW_ONLY raise HTTPException( @@ -923,13 +841,9 @@ class RouteChecks: # Legacy explicit-allow sets (kept for routes that are POST but # semantically read-only, e.g. /spend/calculate). Both admin_viewer_routes # and global_spend_tracking_routes are reads/listings. - if RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value - ): + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value): return - if RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value - ): + if RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value): return # NOTE: We intentionally do NOT fall back to allowing all diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py index 35bb79e7efe..86eb22aee84 100644 --- a/litellm/proxy/auth/trusted_proxy_utils.py +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -63,9 +63,7 @@ def require_trusted_proxy_request( if general_settings is None: general_settings = _get_proxy_general_settings() - trusted_networks = parse_trusted_proxy_ranges( - general_settings.get(setting_name), setting_name=setting_name - ) + trusted_networks = parse_trusted_proxy_ranges(general_settings.get(setting_name), setting_name=setting_name) if not trusted_networks: raise ValueError( f"{feature_name} requires general_settings.{setting_name} before " diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 184d91fb103..73f9def822b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -106,9 +106,7 @@ def _normalize_public_auth_route(route: str) -> str: return route -def _route_requires_auth_despite_public( - route: str, general_settings: Optional[dict] -) -> bool: +def _route_requires_auth_despite_public(route: str, general_settings: Optional[dict]) -> bool: normalized_route = _normalize_public_auth_route(route) if normalized_route == "/metrics": return litellm.require_auth_for_metrics_endpoint is not False @@ -206,9 +204,7 @@ def _routing_selector_matches_claim( return True selector_list: List[str] = ( - [str(v) for v in selector_value] - if isinstance(selector_value, list) - else [str(selector_value)] + [str(v) for v in selector_value] if isinstance(selector_value, list) else [str(selector_value)] ) if claim_value is None: @@ -216,11 +212,7 @@ def _routing_selector_matches_claim( if isinstance(claim_value, list): claim_list = [str(v) for v in claim_value] - elif ( - split_space_delimited - and isinstance(claim_value, str) - and " " in claim_value.strip() - ): + elif split_space_delimited and isinstance(claim_value, str) and " " in claim_value.strip(): # OAuth/OIDC often sends scope as a single space-delimited string. Only split # for the scope selector: iss/aud/client_id must stay exact full-string match # on unverified claims (see routing override security review). The elif guard @@ -240,21 +232,13 @@ def _routing_selector_matches_claim( return fnmatch.fnmatchcase(claim, selector) return selector == claim - return any( - _selector_matches_claim(selector=s, claim=c) - for s in selector_list - for c in claim_list - ) + return any(_selector_matches_claim(selector=s, claim=c) for s in selector_list for c in claim_list) -def _matches_routing_override( - token_claims: dict, override: "JWTRoutingOverride" -) -> bool: +def _matches_routing_override(token_claims: dict, override: "JWTRoutingOverride") -> bool: return ( _routing_selector_matches_claim(override.iss, token_claims.get("iss")) - and _routing_selector_matches_claim( - override.client_id, token_claims.get("client_id") - ) + and _routing_selector_matches_claim(override.client_id, token_claims.get("client_id")) and _routing_selector_matches_claim( override.scope, token_claims.get("scope"), @@ -274,12 +258,8 @@ def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> return False for override in routing_overrides: - if override.path == "oauth2" and _matches_routing_override( - token_claims=token_claims, override=override - ): - verbose_proxy_logger.debug( - "JWT routing override matched. Routing token to OAuth2 introspection." - ) + if override.path == "oauth2" and _matches_routing_override(token_claims=token_claims, override=override): + verbose_proxy_logger.debug("JWT routing override matched. Routing token to OAuth2 introspection.") return True return False @@ -378,9 +358,7 @@ async def user_api_key_auth_websocket(websocket: WebSocket): api_key = websocket.headers.get("api-key") if not api_key: # Try extracting from WebSocket subprotocol (browser clients) - for protocol in websocket.headers.get("sec-websocket-protocol", "").split( - "," - ): + for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): protocol = protocol.strip() if protocol.startswith("openai-insecure-api-key."): api_key = protocol[len("openai-insecure-api-key.") :] @@ -392,9 +370,7 @@ async def user_api_key_auth_websocket(websocket: WebSocket): # Extract the API key from the Bearer token if not authorization.startswith("Bearer "): await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException( - status_code=403, detail="Invalid Authorization header format" - ) + raise HTTPException(status_code=403, detail="Invalid Authorization header format") api_key = authorization[len("Bearer ") :].strip() @@ -408,9 +384,7 @@ async def user_api_key_auth_websocket(websocket: WebSocket): raise HTTPException(status_code=403, detail=str(e)) -def update_valid_token_with_end_user_params( - valid_token: UserAPIKeyAuth, end_user_params: dict -) -> UserAPIKeyAuth: +def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_user_params: dict) -> UserAPIKeyAuth: valid_token.end_user_id = end_user_params.get("end_user_id") # Only overwrite token fields when the DB-derived value is not None. # This prevents DB lookups (where the budget table has no value set) @@ -423,9 +397,7 @@ def update_valid_token_with_end_user_params( if end_user_params.get("allowed_model_region") is not None: valid_token.allowed_model_region = end_user_params["allowed_model_region"] if end_user_params.get("end_user_model_max_budget") is not None: - valid_token.end_user_model_max_budget = end_user_params[ - "end_user_model_max_budget" - ] + valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"] return valid_token @@ -464,9 +436,7 @@ async def get_global_proxy_spend( proxy_logging_obj: ProxyLogging, ) -> Optional[float]: global_proxy_spend = None - if ( - litellm.max_budget > 0 and prisma_client is not None - ): # user set proxy max budget + if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget # Use event-driven coordination to prevent cache stampede cache_key = "{}:spend".format(litellm_proxy_admin_name) global_proxy_spend = await _fetch_global_spend_with_event_coordination( @@ -588,10 +558,7 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( return UserAPIKeyAuth() ## IF AUTH ENABLED ### IF CUSTOM PARSER REQUIRED - if ( - endpoint.get("custom_auth_parser") is not None - and endpoint.get("custom_auth_parser") == "langfuse" - ): + if endpoint.get("custom_auth_parser") is not None and endpoint.get("custom_auth_parser") == "langfuse": # langfuse returns {'Authorization': 'Basic '} # check the langfuse public key if it contains the litellm api key import base64 @@ -605,8 +572,7 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( if headers is not None: header_key = headers.get("litellm_user_api_key", "") if ( - isinstance(request.headers, dict) - and request.headers.get(key=header_key) is not None # type: ignore + isinstance(request.headers, dict) and request.headers.get(key=header_key) is not None # type: ignore ): api_key = request.headers.get(key=header_key) # type: ignore return api_key @@ -719,9 +685,7 @@ async def _auto_register_jwt_mapping( claim_value, ) try: - await prisma_client.db.litellm_verificationtoken.delete( - where={"token": token_hash} - ) + await prisma_client.db.litellm_verificationtoken.delete(where={"token": token_hash}) except Exception as delete_err: # Don't fail the request if cleanup fails — the orphan is # unmapped and inert. Log so an operator can prune it later. @@ -810,9 +774,7 @@ async def _resolve_jwt_to_virtual_key( ) if claim_value is None: - verbose_proxy_logger.debug( - f"JWT Key Mapping: Claim field '{virtual_key_claim_field}' not found in JWT claims." - ) + verbose_proxy_logger.debug(f"JWT Key Mapping: Claim field '{virtual_key_claim_field}' not found in JWT claims.") # A missing claim is an unmapped client — apply the no-match policy # rather than returning early. Otherwise a caller can bypass REJECT # simply by presenting a JWT that omits the configured field. For @@ -974,9 +936,7 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None: ) # Under V2 the FastAPI instrumentor stamps http.route / url.path on the server # span; only the legacy logger needs these set explicitly. - set_route_attrs = getattr( - open_telemetry_logger, "set_proxy_request_route_attributes", None - ) + set_route_attrs = getattr(open_telemetry_logger, "set_proxy_request_route_attributes", None) if not is_otel_v2_enabled() and set_route_attrs is not None: set_route_attrs( parent_otel_span, @@ -1031,9 +991,7 @@ async def _user_api_key_auth_builder( request=request, route=route, ) - pass_through_endpoints: Optional[List[dict]] = general_settings.get( - "pass_through_endpoints", None - ) + pass_through_endpoints: Optional[List[dict]] = general_settings.get("pass_through_endpoints", None) ## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header api_key, passed_in_key = get_api_key( custom_litellm_key_header=custom_litellm_key_header, @@ -1107,9 +1065,7 @@ async def _user_api_key_auth_builder( """ ######## Route Checks Before Reading DB / Cache for "token" ################ - if not _route_requires_auth_despite_public( - route=route, general_settings=general_settings - ) and ( + if not _route_requires_auth_despite_public(route=route, general_settings=general_settings) and ( route in LiteLLMRoutes.public_routes.value # type: ignore or route_in_additonal_public_routes(current_route=route) ): @@ -1124,20 +1080,16 @@ async def _user_api_key_auth_builder( # Routing uses unverified JWT claims only to choose auth path. # Final authentication is enforced by the selected validator. - route_jwt_to_oauth2 = is_jwt and _should_route_jwt_to_oauth2_override( - token=api_key, jwt_handler=jwt_handler - ) + route_jwt_to_oauth2 = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler) # OAuth2 applies for: # 1) when global OAuth2 auth is enabled on LLM + info routes # 2) JWT tokens that explicitly match routing_overrides on LLM + info routes should_apply_override_oauth2 = route_jwt_to_oauth2 and ( - RouteChecks.is_llm_api_route(route=route) - or RouteChecks.is_info_route(route=route) + RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) ) should_apply_global_oauth2 = enable_oauth2_auth and ( - RouteChecks.is_llm_api_route(route=route) - or RouteChecks.is_info_route(route=route) + RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) ) if (should_apply_global_oauth2 and not is_jwt) or should_apply_override_oauth2: from litellm.proxy.proxy_server import premium_user @@ -1157,9 +1109,7 @@ async def _user_api_key_auth_builder( from litellm.proxy.proxy_server import premium_user if premium_user is not True: - raise ValueError( - f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}") is_jwt = jwt_handler.is_jwt(token=api_key) verbose_proxy_logger.debug("is_jwt: %s", is_jwt) if is_jwt: @@ -1210,9 +1160,7 @@ async def _user_api_key_auth_builder( proxy_logging_obj=proxy_logging_obj, parent_otel_span=parent_otel_span, request_headers=_safe_get_request_headers(request), - request_method=RouteChecks._get_request_method( - request=request - ), + request_method=RouteChecks._get_request_method(request=request), ) is_proxy_admin = result["is_proxy_admin"] @@ -1222,9 +1170,7 @@ async def _user_api_key_auth_builder( user_object = result["user_object"] end_user_id = result["end_user_id"] org_id = result["org_id"] - team_membership: Optional[LiteLLM_TeamMembership] = result.get( - "team_membership", None - ) + team_membership: Optional[LiteLLM_TeamMembership] = result.get("team_membership", None) jwt_claims = result.get("jwt_claims", None) if is_proxy_admin: @@ -1246,29 +1192,11 @@ async def _user_api_key_auth_builder( user_role=LitellmUserRoles.PROXY_ADMIN, user_id=user_id, team_id=team_id, - team_alias=( - team_object.team_alias - if team_object is not None - else None - ), - team_tpm_limit=( - team_object.tpm_limit - if team_object is not None - else None - ), - team_rpm_limit=( - team_object.rpm_limit - if team_object is not None - else None - ), - team_models=( - team_object.models if team_object is not None else [] - ), - team_metadata=( - team_object.metadata - if team_object is not None - else None - ), + team_alias=(team_object.team_alias if team_object is not None else None), + team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), + team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), + team_models=(team_object.models if team_object is not None else []), + team_metadata=(team_object.metadata if team_object is not None else None), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, @@ -1278,53 +1206,32 @@ async def _user_api_key_auth_builder( valid_token = UserAPIKeyAuth( api_key=None, team_id=team_id, - team_alias=( - team_object.team_alias if team_object is not None else None - ), - team_tpm_limit=( - team_object.tpm_limit if team_object is not None else None - ), - team_rpm_limit=( - team_object.rpm_limit if team_object is not None else None - ), - team_models=( - team_object.models if team_object is not None else [] - ), + team_alias=(team_object.team_alias if team_object is not None else None), + team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), + team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), + team_models=(team_object.models if team_object is not None else []), user_role=( LitellmUserRoles(user_object.user_role) - if user_object is not None - and user_object.user_role is not None + if user_object is not None and user_object.user_role is not None else LitellmUserRoles.INTERNAL_USER ), user_id=user_id, org_id=org_id, parent_otel_span=parent_otel_span, end_user_id=end_user_id, - user_tpm_limit=( - user_object.tpm_limit if user_object is not None else None - ), - user_rpm_limit=( - user_object.rpm_limit if user_object is not None else None - ), + user_tpm_limit=(user_object.tpm_limit if user_object is not None else None), + user_rpm_limit=(user_object.rpm_limit if user_object is not None else None), team_member_rpm_limit=( - team_membership.safe_get_team_member_rpm_limit() - if team_membership is not None - else None + team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None ), team_member_tpm_limit=( - team_membership.safe_get_team_member_tpm_limit() - if team_membership is not None - else None - ), - team_metadata=( - team_object.metadata if team_object is not None else None + team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None ), + team_metadata=(team_object.metadata if team_object is not None else None), jwt_claims=jwt_claims, ) valid_token.team_object_permission = ( - team_object.object_permission - if team_object is not None - else None + team_object.object_permission if team_object is not None else None ) # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. @@ -1364,13 +1271,9 @@ async def _user_api_key_auth_builder( if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero - skip_budget_checks = _is_model_cost_zero( - model=model, llm_router=llm_router - ) + skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) if skip_budget_checks: - verbose_proxy_logger.info( - f"Skipping all budget checks for zero-cost model: {model}" - ) + verbose_proxy_logger.info(f"Skipping all budget checks for zero-cost model: {model}") # Fetch project object for JWT path if project_id is set _jwt_project_obj = None @@ -1416,9 +1319,7 @@ async def _user_api_key_auth_builder( raise Exception("No api key passed in.") elif api_key == "": # missing 'Bearer ' prefix - raise Exception( - "Malformed API Key passed in. Ensure Key has `Bearer ` prefix." - ) + raise Exception("Malformed API Key passed in. Ensure Key has `Bearer ` prefix.") if route == "/user/auth": if general_settings.get("allow_user_auth", False) is True: @@ -1433,9 +1334,7 @@ async def _user_api_key_auth_builder( _end_user_object = None end_user_params = {} - raw_end_user_id = get_end_user_id_from_request_body( - request_data, _safe_get_request_headers(request) - ) + raw_end_user_id = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) end_user_id = await resolve_and_validate_end_user_id( raw_end_user_id=raw_end_user_id, prisma_client=prisma_client, @@ -1458,9 +1357,7 @@ async def _user_api_key_auth_builder( route=route, ) if _end_user_object is not None: - end_user_params["allowed_model_region"] = ( - _end_user_object.allowed_model_region - ) + end_user_params["allowed_model_region"] = _end_user_object.allowed_model_region if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -1487,9 +1384,7 @@ async def _user_api_key_auth_builder( except Exception as e: if isinstance(e, litellm.BudgetExceededError): raise e - verbose_proxy_logger.debug( - "Unable to find user in db. Error - {}".format(str(e)) - ) + verbose_proxy_logger.debug("Unable to find user in db. Error - {}".format(str(e))) pass ### CHECK IF ADMIN ### @@ -1523,9 +1418,7 @@ async def _user_api_key_auth_builder( and not api_key.startswith("sk-") and get_secret_bool("EXPERIMENTAL_UI_LOGIN") is not False ): - valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key( - api_key - ) + valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) if ( valid_token is not None @@ -1538,10 +1431,7 @@ async def _user_api_key_auth_builder( expiry_time = valid_token.expires else: expiry_time = datetime.fromisoformat(valid_token.expires) - if ( - expiry_time.tzinfo is None - or expiry_time.tzinfo.utcoffset(expiry_time) is None - ): + if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: expiry_time = expiry_time.replace(tzinfo=timezone.utc) if expiry_time < current_time: await _delete_cache_key_object( @@ -1560,17 +1450,11 @@ async def _user_api_key_auth_builder( ) valid_token.parent_otel_span = parent_otel_span if _end_user_object is not None: - valid_token.end_user_object_permission = ( - _end_user_object.object_permission - ) + valid_token.end_user_object_permission = _end_user_object.object_permission return valid_token - if ( - valid_token is not None - and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.team_id is not None - ): + if valid_token is not None and isinstance(valid_token, UserAPIKeyAuth) and valid_token.team_id is not None: ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token try: team_obj: LiteLLM_TeamTableCachedObj = await get_team_object( @@ -1594,9 +1478,7 @@ async def _user_api_key_auth_builder( if field_name in valid_token.__fields__: setattr(valid_token, field_name, v) except Exception as e: - verbose_logger.debug( - e - ) # moving from .warning to .debug as it spams logs when team missing from cache. + verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. try: is_master_key_valid = secrets.compare_digest(api_key, master_key) # type: ignore @@ -1607,11 +1489,7 @@ async def _user_api_key_auth_builder( if not isinstance(master_key, str): raise HTTPException( status_code=500, - detail={ - "Master key must be a valid string. Current type={}".format( - type(master_key) - ) - }, + detail={"Master key must be a valid string. Current type={}".format(type(master_key))}, ) if is_master_key_valid: @@ -1649,9 +1527,7 @@ async def _user_api_key_auth_builder( ## IF it's not a master key ## Route should not be in master_key_only_routes if route in LiteLLMRoutes.master_key_only_routes.value: # type: ignore - raise Exception( - f"Tried to access route={route}, which is only for MASTER KEY" - ) + raise Exception(f"Tried to access route={route}, which is only for MASTER KEY") ## Check DB @@ -1666,14 +1542,8 @@ async def _user_api_key_auth_builder( ) if valid_token is None: - if isinstance( - api_key, str - ): # if generated token, make sure it starts with sk-. - _masked_key = ( - "{}****{}".format(api_key[:4], api_key[-4:]) - if len(api_key) > 8 - else "****" - ) + if isinstance(api_key, str): # if generated token, make sure it starts with sk-. + _masked_key = "{}****{}".format(api_key[:4], api_key[-4:]) if len(api_key) > 8 else "****" if not api_key.startswith("sk-"): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, @@ -1714,9 +1584,7 @@ async def _user_api_key_auth_builder( valid_token.end_user_id = end_user_params.get("end_user_id") valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") - valid_token.allowed_model_region = end_user_params.get( - "allowed_model_region" - ) + valid_token.allowed_model_region = end_user_params.get("allowed_model_region") # update key budget with temp budget increase valid_token = _update_key_budget_with_temp_budget_increase( valid_token @@ -1740,9 +1608,7 @@ async def _user_api_key_auth_builder( ## base case ## key is disabled if valid_token.blocked is True: - raise Exception( - "Key is blocked. Update via `/key/unblock` if you're an admin." - ) + raise Exception("Key is blocked. Update via `/key/unblock` if you're an admin.") await _enforce_key_and_fallback_model_access( valid_token=valid_token, request_data=request_data, @@ -1792,13 +1658,9 @@ async def _user_api_key_auth_builder( if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero - skip_budget_checks = _is_model_cost_zero( - model=model, llm_router=llm_router - ) + skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) if skip_budget_checks: - verbose_proxy_logger.info( - f"Skipping all budget checks for zero-cost model: {model}" - ) + verbose_proxy_logger.info(f"Skipping all budget checks for zero-cost model: {model}") # Check 3. Check if user is in their team budget if not skip_budget_checks and valid_token.team_member_spend is not None: @@ -1815,9 +1677,7 @@ async def _user_api_key_auth_builder( _team_id = valid_token.team_id if _user_id is not None and _team_id is not None: - _db_member = await TeamMembershipRepository( - prisma_client - ).table.find_first( + _db_member = await TeamMembershipRepository(prisma_client).table.find_first( where={ "user_id": _user_id, "team_id": _team_id, @@ -1825,9 +1685,7 @@ async def _user_api_key_auth_builder( include={"litellm_budget_table": True}, ) if _db_member is not None: - team_member_info = LiteLLM_TeamMembership( - **_db_member.dict() - ) + team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) await user_api_key_cache.async_set_cache( key=_cache_key, value=team_member_info, @@ -1835,22 +1693,14 @@ async def _user_api_key_auth_builder( ttl=5, ) - if ( - team_member_info is not None - and team_member_info.litellm_budget_table is not None - ): - team_member_budget = ( - team_member_info.litellm_budget_table.max_budget - ) + if team_member_info is not None and team_member_info.litellm_budget_table is not None: + team_member_budget = team_member_info.litellm_budget_table.max_budget if team_member_budget is not None and team_member_budget > 0: # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend team_member_spend = valid_token.team_member_spend - if ( - valid_token.user_id is not None - and valid_token.team_id is not None - ): + if valid_token.user_id is not None and valid_token.team_id is not None: team_member_spend = await get_current_spend( counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", fallback_spend=team_member_spend, @@ -1869,10 +1719,7 @@ async def _user_api_key_auth_builder( expiry_time = valid_token.expires else: expiry_time = datetime.fromisoformat(valid_token.expires) - if ( - expiry_time.tzinfo is None - or expiry_time.tzinfo.utcoffset(expiry_time) is None - ): + if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: expiry_time = expiry_time.replace(tzinfo=timezone.utc) verbose_proxy_logger.debug( f"Checking if token expired, expiry time {expiry_time} and current time {current_time}" @@ -1920,9 +1767,7 @@ async def _user_api_key_auth_builder( request=request, llm_router=llm_router, ) - current_models = _get_model_names_for_budget_checks( - model=current_model - ) + current_models = _get_model_names_for_budget_checks(model=current_model) if ( max_budget_per_model is not None @@ -2016,17 +1861,13 @@ async def _user_api_key_auth_builder( valid_token.project_alias = _project_obj.project_alias global_proxy_spend = None - if ( - litellm.max_budget > 0 and prisma_client is not None - ): # user set proxy max budget + if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget cache_key = "{}:spend".format(litellm_proxy_admin_name) with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): - global_proxy_spend = ( - await _fetch_global_spend_with_event_coordination( - cache_key=cache_key, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) + global_proxy_spend = await _fetch_global_spend_with_event_coordination( + cache_key=cache_key, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, ) if global_proxy_spend is not None: @@ -2066,9 +1907,7 @@ async def _user_api_key_auth_builder( if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict["end_user_object_permission"] = ( - _end_user_object.object_permission - ) + valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions @@ -2199,11 +2038,7 @@ async def _run_centralized_common_checks( pass_through_endpoints = general_settings.get("pass_through_endpoints", None) if pass_through_endpoints is not None: for endpoint in pass_through_endpoints: - if ( - isinstance(endpoint, dict) - and endpoint.get("path", "") == route - and endpoint.get("auth") is not True - ): + if isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True: return # No-auth dev mode: master_key unset AND no JWT/OAuth2 auth @@ -2219,9 +2054,7 @@ async def _run_centralized_common_checks( ): return - if user_custom_auth is not None and not general_settings.get( - "custom_auth_run_common_checks", False - ): + if user_custom_auth is not None and not general_settings.get("custom_auth_run_common_checks", False): return parent_otel_span = user_api_key_auth_obj.parent_otel_span @@ -2231,9 +2064,7 @@ async def _run_centralized_common_checks( # function is invoked in isolation (e.g. in direct unit tests). end_user_id = user_api_key_auth_obj.end_user_id if end_user_id is None: - raw_end_user_id = get_end_user_id_from_request_body( - request_data, _safe_get_request_headers(request) - ) + raw_end_user_id = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) end_user_id = await resolve_and_validate_end_user_id( raw_end_user_id=raw_end_user_id, prisma_client=prisma_client, @@ -2360,17 +2191,11 @@ async def _run_centralized_common_checks( if isinstance(team_result, BaseException): # Token-derived fallback only valid when a team_id is set; # _team_obj_from_token asserts that precondition. - team_object = ( - _team_obj_from_token(user_api_key_auth_obj) - if user_api_key_auth_obj.team_id is not None - else None - ) + team_object = _team_obj_from_token(user_api_key_auth_obj) if user_api_key_auth_obj.team_id is not None else None else: team_object = team_result - user_object: Optional[LiteLLM_UserTable] = ( - None if isinstance(user_result, BaseException) else user_result - ) + user_object: Optional[LiteLLM_UserTable] = None if isinstance(user_result, BaseException) else user_result project_object: Optional[LiteLLM_ProjectTableCachedObj] = ( None if isinstance(project_result, BaseException) else project_result ) @@ -2532,9 +2357,7 @@ def _should_skip_budget_checks( return False -def _resolve_request_principal( - request: Request, valid_token: UserAPIKeyAuth -) -> Principal: +def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> Principal: """Project the resolved identity into one per-request Principal, off the key object the builder already fetched, and stamp the request network context onto it once. X-Forwarded-For is only trusted when the operator configured @@ -2548,9 +2371,7 @@ def _resolve_request_principal( request, TrustedProxyConfig(use_forwarded_for=bool(cidrs), trusted_proxy_cidrs=cidrs), ) - auth_method = ( - AuthMethod.BEARER_JWT if valid_token.jwt_claims else AuthMethod.API_KEY - ) + auth_method = AuthMethod.BEARER_JWT if valid_token.jwt_claims else AuthMethod.API_KEY return IdentityStore._principal_from_key( valid_token, auth_method=auth_method, @@ -2565,16 +2386,10 @@ async def user_api_key_auth( request: Request, api_key: str = fastapi.Security(api_key_header), azure_api_key_header: str = fastapi.Security(azure_api_key_header), - anthropic_api_key_header: Optional[str] = fastapi.Security( - anthropic_api_key_header - ), - google_ai_studio_api_key_header: Optional[str] = fastapi.Security( - google_ai_studio_api_key_header - ), + anthropic_api_key_header: Optional[str] = fastapi.Security(anthropic_api_key_header), + google_ai_studio_api_key_header: Optional[str] = fastapi.Security(google_ai_studio_api_key_header), azure_apim_header: Optional[str] = fastapi.Security(azure_apim_header), - custom_litellm_key_header: Optional[str] = fastapi.Security( - custom_litellm_key_header - ), + custom_litellm_key_header: Optional[str] = fastapi.Security(custom_litellm_key_header), ) -> UserAPIKeyAuth: """ Parent function to authenticate user api key / jwt token. @@ -2587,9 +2402,7 @@ async def user_api_key_auth( _ensure_parent_otel_span_on_request_state(request) request_data = await _read_request_body(request=request) - request_data = populate_request_with_path_params( - request_data=request_data, request=request - ) + request_data = populate_request_with_path_params(request_data=request_data, request=request) route: str = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED @@ -2610,9 +2423,7 @@ async def user_api_key_auth( user_api_key_auth_obj.budget_reservation = None ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## - RouteChecks.should_call_route( - route=route, valid_token=user_api_key_auth_obj, request=request - ) + RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) # Single authorization point. Builder paths MUST NOT call common_checks. # Route through the same exception handler the builder uses so @@ -2650,9 +2461,7 @@ async def user_api_key_auth( user_api_key_cache, ) - raw_end_user_id = get_end_user_id_from_request_body( - request_data, _safe_get_request_headers(request) - ) + raw_end_user_id = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) if raw_end_user_id is not None: resolved_end_user_id = await resolve_and_validate_end_user_id( raw_end_user_id=raw_end_user_id, @@ -2682,13 +2491,9 @@ async def user_api_key_auth( # reject an already-authenticated request, so it is left unset on failure; # any future consumer must treat a missing principal as deny, not allow. try: - request.state.principal = _resolve_request_principal( - request, user_api_key_auth_obj - ) + request.state.principal = _resolve_request_principal(request, user_api_key_auth_obj) except Exception as e: - verbose_proxy_logger.warning( - "Principal projection at auth seam failed (non-fatal): %s", e - ) + verbose_proxy_logger.warning("Principal projection at auth seam failed (non-fatal): %s", e) return user_api_key_auth_obj @@ -2715,9 +2520,7 @@ async def _return_user_api_key_auth_obj( ) ) - retrieved_user_role = ( - user_role or _get_user_role(user_obj=user_obj) or LitellmUserRoles.INTERNAL_USER - ) + retrieved_user_role = user_role or _get_user_role(user_obj=user_obj) or LitellmUserRoles.INTERNAL_USER user_api_key_kwargs = { "api_key": api_key, @@ -2742,9 +2545,7 @@ async def _return_user_api_key_auth_obj( return UserAPIKeyAuth.model_validate(user_api_key_kwargs) -def get_api_key_from_custom_header( - request: Request, custom_litellm_key_header_name: str -) -> str: +def get_api_key_from_custom_header(request: Request, custom_litellm_key_header_name: str) -> str: """ Get API key from custom header @@ -2781,10 +2582,7 @@ def get_api_key_from_custom_header( def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): valid_token_metadata = valid_token.metadata - if ( - "temp_budget_increase" in valid_token_metadata - and "temp_budget_expiry" in valid_token_metadata - ): + if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata: expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"]) if expiry > datetime.now(): return valid_token_metadata["temp_budget_increase"] @@ -2878,9 +2676,7 @@ async def _enforce_key_and_fallback_model_access( model_list = config.get("model_list", []) new_model_list = model_list verbose_proxy_logger.debug(f"\n new llm router model list {new_model_list}") - elif ( - isinstance(valid_token.models, list) and "all-team-models" in valid_token.models - ): + elif isinstance(valid_token.models, list) and "all-team-models" in valid_token.models: pass else: model = _get_model_from_request_context( @@ -2908,13 +2704,9 @@ async def _enforce_key_and_fallback_model_access( fallback_names: List[str] = [] override_settings = request_data.get("router_settings_override") for _fb_key in ROUTER_FALLBACK_FIELDS: - fallback_names.extend( - iter_router_fallback_model_names(request_data.get(_fb_key)) - ) + fallback_names.extend(iter_router_fallback_model_names(request_data.get(_fb_key))) if isinstance(override_settings, dict): - fallback_names.extend( - iter_router_fallback_model_names(override_settings.get(_fb_key)) - ) + fallback_names.extend(iter_router_fallback_model_names(override_settings.get(_fb_key))) for _name in dict.fromkeys(fallback_names): # dedupe, preserve order await can_key_call_model( @@ -2994,9 +2786,7 @@ async def _run_post_custom_auth_checks( # gate skips it for custom-auth deployments unless # custom_auth_run_common_checks is set. Enforce it here on that path # so an over-budget end user can't keep making requests. - if end_user_object is not None and not general_settings.get( - "custom_auth_run_common_checks", False - ): + if end_user_object is not None and not general_settings.get("custom_auth_run_common_checks", False): await _check_end_user_budget(end_user_obj=end_user_object, route=route) # 2. Check token expiry @@ -3006,21 +2796,14 @@ async def _run_post_custom_auth_checks( expiry_time = valid_token.expires else: expiry_time = datetime.fromisoformat(valid_token.expires) - if ( - expiry_time.tzinfo is None - or expiry_time.tzinfo.utcoffset(expiry_time) is None - ): + if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: expiry_time = expiry_time.replace(tzinfo=timezone.utc) if expiry_time < current_time: raise ProxyException( message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=status.HTTP_401_UNAUTHORIZED, - param=( - abbreviate_api_key(api_key=valid_token.token) - if valid_token.token - else "" - ), + param=(abbreviate_api_key(api_key=valid_token.token) if valid_token.token else ""), ) if general_settings.get("custom_auth_run_common_checks", False): diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 6db75eeb3d9..02e4d9c170e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -129,10 +129,7 @@ async def create_batch( team_metadata = user_api_key_dict.team_metadata or {} enforced_batch_expiry = team_metadata.get("enforced_batch_output_expires_after") if enforced_batch_expiry is not None: - if ( - "anchor" not in enforced_batch_expiry - or "seconds" not in enforced_batch_expiry - ): + if "anchor" not in enforced_batch_expiry or "seconds" not in enforced_batch_expiry: raise HTTPException( status_code=500, detail={ @@ -207,17 +204,11 @@ async def create_batch( response.input_file_id = input_file_id - elif ( - litellm.enable_loadbalancing_on_batch_endpoints is True - and is_router_model - and router_model is not None - ): + elif litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model and router_model is not None: if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) response = await llm_router.acreate_batch(**_create_batch_data) # type: ignore @@ -229,20 +220,14 @@ async def create_batch( if len(target_model_names) != 1: raise HTTPException( status_code=400, - detail={ - "error": "Expected 1 model, got {}".format( - len(target_model_names) - ) - }, + detail={"error": "Expected 1 model, got {}".format(len(target_model_names))}, ) model = target_model_names[0] _create_batch_data["model"] = model if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) response = await llm_router.acreate_batch(**_create_batch_data) @@ -251,9 +236,7 @@ async def create_batch( else: # Check if model specified via header/query/body param model_param = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") + data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") ) # SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback @@ -293,9 +276,7 @@ async def create_batch( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -322,9 +303,7 @@ async def create_batch( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.create_batch(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.create_batch(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -349,9 +328,7 @@ async def retrieve_batch( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), provider: Optional[str] = None, - batch_id: str = Path( - title="Batch ID to retrieve", description="The ID of the batch to retrieve" - ), + batch_id: str = Path(title="Batch ID to retrieve", description="The ID of the batch to retrieve"), ): """ Retrieves a batch. @@ -495,23 +472,17 @@ async def retrieve_batch( f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}" ) - elif ( - litellm.enable_loadbalancing_on_batch_endpoints is True or unified_batch_id - ): + elif litellm.enable_loadbalancing_on_batch_endpoints is True or unified_batch_id: if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) response = await llm_router.aretrieve_batch(**data) # type: ignore response._hidden_params["unified_batch_id"] = unified_batch_id if unified_batch_id: - model_id_from_batch = get_model_id_from_unified_batch_id( - unified_batch_id - ) + model_id_from_batch = get_model_id_from_unified_batch_id(unified_batch_id) if model_id_from_batch: response._hidden_params["model_id"] = model_id_from_batch @@ -553,9 +524,7 @@ async def retrieve_batch( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -582,9 +551,7 @@ async def retrieve_batch( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -660,9 +627,7 @@ async def list_batches( # Try to use managed objects table for listing batches (returns encoded IDs) managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_obj is not None and hasattr( - managed_files_obj, "list_user_batches" - ): + if managed_files_obj is not None and hasattr(managed_files_obj, "list_user_batches"): verbose_proxy_logger.debug("Using managed objects table for batch listing") response = await cast(Any, managed_files_obj).list_user_batches( user_api_key_dict=user_api_key_dict, @@ -673,9 +638,7 @@ async def list_batches( llm_router=llm_router, ) elif model_param := ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") + data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") ): # SCENARIO 2: Use model-based routing from header/query/body credentials = get_credentials_for_model( @@ -704,13 +667,9 @@ async def list_batches( # SCENARIO 2 (alternative): target_model_names based routing elif target_model_names or data.get("target_model_names", None): - target_model_names = target_model_names or data.get( - "target_model_names", None - ) + target_model_names = target_model_names or data.get("target_model_names", None) if target_model_names is None: - raise ValueError( - "target_model_names is required for this routing scenario" - ) + raise ValueError("target_model_names is required for this routing scenario") model = target_model_names.split(",")[0] data.pop("model", None) response = await llm_router.alist_batches( @@ -768,11 +727,7 @@ async def list_batches( original_exception=e, request_data={"after": after, "limit": limit}, ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - {}".format(str(e))) raise handle_exception_on_proxy(e) @@ -893,18 +848,14 @@ async def cancel_batch( if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) model_id_from_batch = get_model_id_from_unified_batch_id(unified_batch_id) if model_id_from_batch is None: raise HTTPException( status_code=400, - detail={ - "error": "Invalid LiteLLM managed batch ID. Missing model_id." - }, + detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."}, ) data["model"] = model_id_from_batch data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id) @@ -953,9 +904,7 @@ async def cancel_batch( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -982,9 +931,7 @@ async def cancel_batch( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.create_batch(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.create_batch(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/caching_routes.py b/litellm/proxy/caching_routes.py index f0d8ddf97d6..09d68ee3a0c 100644 --- a/litellm/proxy/caching_routes.py +++ b/litellm/proxy/caching_routes.py @@ -40,9 +40,7 @@ def _extract_cache_params() -> Dict[str, Any]: return {} try: cache_params = vars(litellm.cache.cache) - cleaned_params = ( - HealthCheckCacheParams(**cache_params).model_dump() if cache_params else {} - ) + cleaned_params = HealthCheckCacheParams(**cache_params).model_dump() if cache_params else {} return masker.mask_dict(cleaned_params) except (AttributeError, TypeError) as e: verbose_proxy_logger.debug(f"Error extracting cache params: {str(e)}") @@ -81,9 +79,7 @@ async def cache_ping(): if litellm.cache.type == "redis": ping_response = await litellm.cache.ping() - verbose_proxy_logger.debug( - "/cache/ping: ping_response: " + str(ping_response) - ) + verbose_proxy_logger.debug("/cache/ping: ping_response: " + str(ping_response)) # add cache does not return anything await litellm.cache.async_add_cache( result="test_key", @@ -144,9 +140,7 @@ async def cache_delete(request: Request): """ try: if litellm.cache is None: - raise HTTPException( - status_code=503, detail="Cache not initialized. litellm.cache is None" - ) + raise HTTPException(status_code=503, detail="Cache not initialized. litellm.cache is None") request_data = await request.json() keys = request_data.get("keys", None) @@ -179,9 +173,7 @@ def _get_redis_client_info(cache_instance) -> Tuple[List, int]: client_list = cache_instance.client_list() return client_list, len(client_list) except Exception as e: - verbose_proxy_logger.warning( - f"CLIENT LIST command failed (likely restricted on managed Redis): {str(e)}" - ) + verbose_proxy_logger.warning(f"CLIENT LIST command failed (likely restricted on managed Redis): {str(e)}") return ["CLIENT LIST command not available on this Redis instance"], -1 @@ -195,14 +187,9 @@ async def cache_redis_info(): """ try: if litellm.cache is None: - raise HTTPException( - status_code=503, detail="Cache not initialized. litellm.cache is None" - ) + raise HTTPException(status_code=503, detail="Cache not initialized. litellm.cache is None") - if not ( - litellm.cache.type == "redis" - and isinstance(litellm.cache.cache, RedisCache) - ): + if not (litellm.cache.type == "redis" and isinstance(litellm.cache.cache, RedisCache)): raise HTTPException( status_code=500, detail=f"Cache type {litellm.cache.type} does not support redis info", @@ -244,12 +231,8 @@ async def cache_flushall(): """ try: if litellm.cache is None: - raise HTTPException( - status_code=503, detail="Cache not initialized. litellm.cache is None" - ) - if litellm.cache.type == "redis" and isinstance( - litellm.cache.cache, RedisCache - ): + raise HTTPException(status_code=503, detail="Cache not initialized. litellm.cache is None") + if litellm.cache.type == "redis" and isinstance(litellm.cache.cache, RedisCache): litellm.cache.cache.flushall() return { "status": "success", diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index 064c6162b0a..41015ffd341 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -160,9 +160,7 @@ class ChatClient: # Make streaming request session = requests.Session() try: - response = session.post( - url, headers=self._get_headers(), json=data, stream=True - ) + response = session.post(url, headers=self._get_headers(), json=data, stream=True) response.raise_for_status() # Parse SSE stream diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index f39ffb3e864..dfcedb70686 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -228,26 +228,20 @@ def _resolve_api_key(ctx: click.Context) -> str: ctx.invoke(login) api_key = get_stored_api_key(expected_base_url=base_url) if not api_key: - raise click.ClickException( - "Login did not produce an API key; cannot start the agent." - ) + raise click.ClickException("Login did not produce an API key; cannot start the agent.") return api_key _SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy." -def _launch( - ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool -) -> None: +def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None: base_url = ctx.obj["base_url"] started_interactive = _is_interactive() api_key = _resolve_api_key(ctx) display_name, _ = agent_profile(binary) - click.echo( - f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}" - ) + click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}") try: run_agent( @@ -255,9 +249,7 @@ def _launch( api_key, [binary, *args], skip_verify=skip_verify, - reattach_terminal=( - _restore_controlling_terminal if started_interactive else None - ), + reattach_terminal=(_restore_controlling_terminal if started_interactive else None), ) except AgentRunError as e: raise click.ClickException(str(e)) @@ -286,10 +278,7 @@ def _make_agent_command(binary: str, display_name: str) -> click.Command: def agent_commands() -> List[click.Command]: """Build one top-level command per known agent, e.g. `lite claude`.""" - return [ - _make_agent_command(binary, name) - for binary, (name, _profiles) in _KNOWN_AGENTS.items() - ] + return [_make_agent_command(binary, name) for binary, (name, _profiles) in _KNOWN_AGENTS.items()] __all__ = [ diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b06d86d5965..b258664ee16 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -154,9 +154,7 @@ def get_key_input(): return None -def display_interactive_team_selection( - teams: List[Dict[str, Any]], selected_index: int = 0 -) -> None: +def display_interactive_team_selection(teams: List[Dict[str, Any]], selected_index: int = 0) -> None: """Display teams with one highlighted for selection""" console = Console() @@ -270,9 +268,7 @@ def prompt_team_selection_fallback( ) return selected_team else: - click.echo( - f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}" - ) + click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}") except ValueError: click.echo("❌ Invalid input. Please enter a number or 'skip'") except KeyboardInterrupt: @@ -307,25 +303,14 @@ def _poll_for_ready_data( if status == "ready": return data if status == "pending": - if ( - pending_message - and pending_log_every > 0 - and attempt % pending_log_every == 0 - ): + if pending_message and pending_log_every > 0 and attempt % pending_log_every == 0: click.echo(pending_message) - elif ( - other_status_message - and other_status_log_every > 0 - and attempt % other_status_log_every == 0 - ): + elif other_status_message and other_status_log_every > 0 and attempt % other_status_log_every == 0: click.echo(other_status_message) elif http_error_log_every > 0 and attempt % http_error_log_every == 0: click.echo(f"Polling error: HTTP {response.status_code}") except requests.RequestException as e: - if ( - connection_error_log_every > 0 - and attempt % connection_error_log_every == 0 - ): + if connection_error_log_every > 0 and attempt % connection_error_log_every == 0: click.echo(f"Connection error (will retry): {e}") time.sleep(poll_interval) return None @@ -369,9 +354,7 @@ def _get_cli_sso_poll_headers(poll_secret: str) -> Dict[str, str]: return {"x-litellm-cli-poll-secret": poll_secret} -def _poll_for_authentication( - base_url: str, key_id: str, poll_secret: str -) -> Optional[dict]: +def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Optional[dict]: """ Poll the server for authentication completion and handle team selection. @@ -449,9 +432,7 @@ def _handle_team_selection_during_polling( The JWT token with the selected team, or None if selection was skipped """ if not teams: - click.echo( - "ℹ️ No teams found. You can create or join teams using the web interface." - ) + click.echo("ℹ️ No teams found. You can create or join teams using the web interface.") return None click.echo("\n" + "=" * 60) @@ -528,9 +509,7 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option click.echo(f"\n✅ Selected team: {team_alias} ({team_id})") return team_id - click.echo( - f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}" - ) + click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}") except ValueError: click.echo("❌ Invalid input. Please enter a number or 'skip'") except KeyboardInterrupt: @@ -553,9 +532,7 @@ def login(ctx: click.Context): poll_secret = cli_sso_flow["poll_secret"] user_code = cli_sso_flow["user_code"] - sso_url = f"{base_url}/sso/key/generate?" + urlencode( - {"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id} - ) + sso_url = f"{base_url}/sso/key/generate?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id}) click.echo(f"Opening browser to: {sso_url}") click.echo("Please complete the SSO authentication in your browser...") @@ -568,9 +545,7 @@ def login(ctx: click.Context): # Poll for authentication completion click.echo("Waiting for authentication...") - auth_result = _poll_for_authentication( - base_url=base_url, key_id=key_id, poll_secret=poll_secret - ) + auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id, poll_secret=poll_secret) if auth_result: api_key = auth_result["api_key"] @@ -638,9 +613,7 @@ def whoami(): click.echo(f"Token age: {age_hours:.1f} hours") if age_hours > CLI_JWT_EXPIRATION_HOURS: - click.echo( - f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired." - ) + click.echo(f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") # Export functions for use by other CLI commands diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index 696e34c3ecd..9d1a1aa7a30 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -28,14 +28,10 @@ def _get_available_models(ctx: click.Context) -> List[Dict[str, Any]]: return [] -def _select_model( - console: Console, available_models: List[Dict[str, Any]] -) -> Optional[str]: +def _select_model(console: Console, available_models: List[Dict[str, Any]]) -> Optional[str]: """Interactive model selection""" if not available_models: - console.print( - "[yellow]No models available or could not fetch models list.[/yellow]" - ) + console.print("[yellow]No models available or could not fetch models list.[/yellow]") model_name = Prompt.ask("Please enter a model name") return model_name if model_name.strip() else None @@ -48,14 +44,10 @@ def _select_model( models_to_display: List[Dict[str, Any]] = available_models[:MAX_MODELS_TO_DISPLAY] for i, model in enumerate(models_to_display): # Limit to first 200 models - table.add_row( - str(i + 1), str(model.get("id", "")), str(model.get("owned_by", "")) - ) + table.add_row(str(i + 1), str(model.get("id", "")), str(model.get("owned_by", ""))) if len(available_models) > MAX_MODELS_TO_DISPLAY: - console.print( - f"\n[dim]... and {len(available_models) - MAX_MODELS_TO_DISPLAY} more models[/dim]" - ) + console.print(f"\n[dim]... and {len(available_models) - MAX_MODELS_TO_DISPLAY} more models[/dim]") console.print(table) @@ -172,9 +164,7 @@ def chat( break # Handle special commands - should_exit, messages, new_model = _handle_special_commands( - console, user_input, messages, system, ctx - ) + should_exit, messages, new_model = _handle_special_commands(console, user_input, messages, system, ctx) if should_exit: break @@ -261,9 +251,7 @@ def _show_history(console: Console, messages: List[Dict[str, Any]]): content = message["content"] if role == "system": - console.print( - f"[dim]{i}. [bold magenta]System:[/bold magenta] {content}[/dim]" - ) + console.print(f"[dim]{i}. [bold magenta]System:[/bold magenta] {content}[/dim]") elif role == "user": console.print(f"{i}. [bold cyan]You:[/bold cyan] {content}") elif role == "assistant": @@ -291,9 +279,7 @@ def _save_conversation(console: Console, messages: List[Dict[str, Any]], command console.print(f"[red]Error saving conversation: {e}[/red]") -def _load_conversation( - console: Console, command: str, system: Optional[str] -) -> List[Dict[str, Any]]: +def _load_conversation(console: Console, command: str, system: Optional[str]) -> List[Dict[str, Any]]: """Load conversation from a file""" parts = command.split() if len(parts) < 2: @@ -395,9 +381,7 @@ def _stream_response( console.print(f"\n[red]Error: HTTP {e.response.status_code}[/red]") try: error_body = e.response.json() - console.print( - f"[red]{error_body.get('error', {}).get('message', 'Unknown error')}[/red]" - ) + console.print(f"[red]{error_body.get('error', {}).get('message', 'Unknown error')}[/red]") except json.JSONDecodeError: console.print(f"[red]{e.response.text}[/red]") return None diff --git a/litellm/proxy/client/cli/commands/http.py b/litellm/proxy/client/cli/commands/http.py index b724f3cf2c7..dba36f9d92c 100644 --- a/litellm/proxy/client/cli/commands/http.py +++ b/litellm/proxy/client/cli/commands/http.py @@ -61,9 +61,7 @@ def request( key, value = h.split(":", 1) headers[key.strip()] = value.strip() except ValueError: - raise click.BadParameter( - f"Invalid header format: {h}. Expected format: 'key:value'" - ) + raise click.BadParameter(f"Invalid header format: {h}. Expected format: 'key:value'") # Parse JSON data if provided json_data = None diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py index a007d260199..45e27442708 100644 --- a/litellm/proxy/client/cli/commands/keys.py +++ b/litellm/proxy/client/cli/commands/keys.py @@ -30,9 +30,7 @@ def keys(): default=True, help="Return the full key object", ) -@click.option( - "--include-team-keys", is_flag=True, help="Include team keys in the response" -) +@click.option("--include-team-keys", is_flag=True, help="Include team keys in the response") @click.option( "--format", "output_format", @@ -72,9 +70,7 @@ def list( if output_format == "json": rich.print_json(data=response) else: - rich.print( - f"Showing {len(response.get('keys', []))} keys out of {response.get('total_count', 0)}" - ) + rich.print(f"Showing {len(response.get('keys', []))} keys out of {response.get('total_count', 0)}") table = Table(title="API Keys") table.add_column("Key Hash", style="cyan") table.add_column("Alias", style="green") @@ -105,9 +101,7 @@ def list( @click.option("--team-id", type=str, help="Team ID to associate the key with") @click.option("--user-id", type=str, help="User ID to associate the key with") @click.option("--budget-id", type=str, help="Budget ID to associate the key with") -@click.option( - "--config", type=str, help="JSON string of additional configuration parameters" -) +@click.option("--config", type=str, help="JSON string of additional configuration parameters") @click.pass_context def generate( ctx: click.Context, @@ -154,9 +148,7 @@ def generate( @keys.command() @click.option("--keys", type=str, help="Comma-separated list of API keys to delete") -@click.option( - "--key-aliases", type=str, help="Comma-separated list of key aliases to delete" -) +@click.option("--key-aliases", type=str, help="Comma-separated list of key aliases to delete") @click.pass_context def delete(ctx: click.Context, keys: Optional[str], key_aliases: Optional[str]): """Delete API keys by key or alias""" @@ -195,9 +187,7 @@ def _parse_created_since_filter(created_since: Optional[str]) -> Optional[dateti raise click.Abort() -def _fetch_all_keys_with_pagination( - source_client: KeysManagementClient, source_base_url: str -) -> List[Dict[str, Any]]: +def _fetch_all_keys_with_pagination(source_client: KeysManagementClient, source_base_url: str) -> List[Dict[str, Any]]: """Fetch all keys from source instance using pagination.""" click.echo(f"Fetching keys from source server: {source_base_url}") source_keys = [] @@ -205,9 +195,7 @@ def _fetch_all_keys_with_pagination( page_size = 100 # Use a larger page size to minimize API calls while True: - source_response = source_client.list( - return_full_object=True, page=page, size=page_size - ) + source_response = source_client.list(return_full_object=True, page=page, size=page_size) # source_client.list() returns Dict[str, Any] when return_request is False (default) assert isinstance(source_response, dict), "Expected dict response from list API" page_keys = source_response.get("keys", []) @@ -243,9 +231,7 @@ def _filter_keys_by_created_since( # Parse the key's created_at timestamp if isinstance(key_created_at, str): if "T" in key_created_at: - key_dt = datetime.fromisoformat( - key_created_at.replace("Z", "+00:00") - ) + key_dt = datetime.fromisoformat(key_created_at.replace("Z", "+00:00")) else: key_dt = datetime.fromisoformat(key_created_at) @@ -256,9 +242,7 @@ def _filter_keys_by_created_since( if key_dt >= created_since_dt: filtered_keys.append(key) - click.echo( - f"Filtered {len(source_keys)} keys to {len(filtered_keys)} keys created since {created_since}" - ) + click.echo(f"Filtered {len(source_keys)} keys to {len(filtered_keys)} keys created since {created_since}") return filtered_keys @@ -281,9 +265,7 @@ def _display_dry_run_table(source_keys: List[Dict[str, Any]]) -> None: dt = datetime.fromisoformat(created_at.replace("Z", "+00:00")) created_at = dt.strftime("%Y-%m-%d %H:%M") - table.add_row( - str(key.get("key_alias", "")), str(key.get("user_id", "")), str(created_at) - ) + table.add_row(str(key.get("key_alias", "")), str(key.get("user_id", "")), str(created_at)) rich.print(table) @@ -343,9 +325,7 @@ def _import_keys_to_destination( required=True, help="Base URL of the source LiteLLM proxy server to import keys from", ) -@click.option( - "--source-api-key", help="API key for authentication to the source server" -) +@click.option("--source-api-key", help="API key for authentication to the source server") @click.option( "--dry-run", is_flag=True, @@ -377,9 +357,7 @@ def import_keys( # Filter keys by created_since if specified if created_since: - source_keys = _filter_keys_by_created_since( - source_keys, created_since_dt, created_since - ) + source_keys = _filter_keys_by_created_since(source_keys, created_since_dt, created_since) if not source_keys: click.echo("No keys found in source instance.") @@ -392,9 +370,7 @@ def import_keys( return # Import each key - imported_count, failed_count = _import_keys_to_destination( - source_keys, dest_client - ) + imported_count, failed_count = _import_keys_to_destination(source_keys, dest_client) # Summary click.echo("\nImport completed:") diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index 387979a69a0..15266488c84 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -129,11 +129,7 @@ def list_models(ctx: click.Context, output_format: Literal["table", "json"]) -> table.add_row( str(model.get("id", "")), str(model.get("object", "model")), - ( - format_timestamp(created) - if isinstance(created, int) - else format_iso_datetime_str(created) - ), + (format_timestamp(created) if isinstance(created, int) else format_iso_datetime_str(created)), str(model.get("owned_by", "")), ) @@ -155,9 +151,7 @@ def list_models(ctx: click.Context, output_format: Literal["table", "json"]) -> help="Model info in key=value format (can be specified multiple times)", ) @click.pass_context -def add_model( - ctx: click.Context, model_name: str, param: tuple[str, ...], info: tuple[str, ...] -) -> None: +def add_model(ctx: click.Context, model_name: str, param: tuple[str, ...], info: tuple[str, ...]) -> None: """Add a new model to the proxy""" # Convert parameters from key=value format to dict model_params = dict(p.split("=", 1) for p in param) @@ -186,9 +180,7 @@ def delete_model(ctx: click.Context, model_id: str) -> None: @click.option("--id", "model_id", help="ID of the model to retrieve") @click.option("--name", "model_name", help="Name of the model to retrieve") @click.pass_context -def get_model( - ctx: click.Context, model_id: Optional[str], model_name: Optional[str] -) -> None: +def get_model(ctx: click.Context, model_id: Optional[str], model_name: Optional[str]) -> None: """Get information about a specific model""" if not model_id and not model_name: raise click.UsageError("Either --id or --name must be provided") @@ -213,9 +205,7 @@ def get_model( help="Comma-separated list of columns to display. Valid columns: public_model, upstream_model, credential_name, created_at, updated_at, id, input_cost, output_cost. Default: public_model,upstream_model,updated_at", ) @click.pass_context -def get_models_info( - ctx: click.Context, output_format: Literal["table", "json"], columns: str -) -> None: +def get_models_info(ctx: click.Context, output_format: Literal["table", "json"], columns: str) -> None: """Get detailed information about all models""" client = create_client(ctx) models_info = client.models.info() @@ -236,30 +226,22 @@ def get_models_info( "upstream_model": { "header": "Upstream Model", "style": "green", - "get_value": lambda m: str( - m.get("litellm_params", {}).get("model", "") - ), + "get_value": lambda m: str(m.get("litellm_params", {}).get("model", "")), }, "credential_name": { "header": "Credential Name", "style": "yellow", - "get_value": lambda m: str( - m.get("litellm_params", {}).get("litellm_credential_name", "") - ), + "get_value": lambda m: str(m.get("litellm_params", {}).get("litellm_credential_name", "")), }, "created_at": { "header": "Created At", "style": "magenta", - "get_value": lambda m: format_iso_datetime_str( - m.get("model_info", {}).get("created_at") - ), + "get_value": lambda m: format_iso_datetime_str(m.get("model_info", {}).get("created_at")), }, "updated_at": { "header": "Updated At", "style": "magenta", - "get_value": lambda m: format_iso_datetime_str( - m.get("model_info", {}).get("updated_at") - ), + "get_value": lambda m: format_iso_datetime_str(m.get("model_info", {}).get("updated_at")), }, "id": { "header": "ID", @@ -270,17 +252,13 @@ def get_models_info( "header": "Input Cost", "style": "green", "justify": "right", - "get_value": lambda m: format_cost_per_1k_tokens( - m.get("model_info", {}).get("input_cost_per_token") - ), + "get_value": lambda m: format_cost_per_1k_tokens(m.get("model_info", {}).get("input_cost_per_token")), }, "output_cost": { "header": "Output Cost", "style": "green", "justify": "right", - "get_value": lambda m: format_cost_per_1k_tokens( - m.get("model_info", {}).get("output_cost_per_token") - ), + "get_value": lambda m: format_cost_per_1k_tokens(m.get("model_info", {}).get("output_cost_per_token")), }, } @@ -324,9 +302,7 @@ def get_models_info( help="Model info in key=value format (can be specified multiple times)", ) @click.pass_context -def update_model( - ctx: click.Context, model_id: str, param: tuple[str, ...], info: tuple[str, ...] -) -> None: +def update_model(ctx: click.Context, model_id: str, param: tuple[str, ...], info: tuple[str, ...]) -> None: """Update an existing model's configuration""" # Convert parameters from key=value format to dict model_params = dict(p.split("=", 1) for p in param) @@ -356,10 +332,7 @@ def _filter_model(model, model_regex, access_group_regex): if access_group_regex: if not isinstance(access_groups, list): return False - if not any( - isinstance(group, str) and access_group_regex.search(group) - for group in access_groups - ): + if not any(isinstance(group, str) and access_group_regex.search(group) for group in access_groups): return False return True @@ -395,32 +368,18 @@ def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, Any]]: with open(yaml_file, "r") as f: data = yaml.safe_load(f) if not data or "model_list" not in data: - raise click.ClickException( - "YAML file must contain a 'model_list' key with a list of models." - ) + raise click.ClickException("YAML file must contain a 'model_list' key with a list of models.") model_list = data["model_list"] if not isinstance(model_list, list): raise click.ClickException("'model_list' must be a list of model definitions.") return model_list -def _get_filtered_model_list( - model_list, only_models_matching_regex, only_access_groups_matching_regex -): +def _get_filtered_model_list(model_list, only_models_matching_regex, only_access_groups_matching_regex): """Return a list of models that pass the filter criteria.""" - model_regex = ( - re.compile(only_models_matching_regex) if only_models_matching_regex else None - ) - access_group_regex = ( - re.compile(only_access_groups_matching_regex) - if only_access_groups_matching_regex - else None - ) - return [ - model - for model in model_list - if _filter_model(model, model_regex, access_group_regex) - ] + model_regex = re.compile(only_models_matching_regex) if only_models_matching_regex else None + access_group_regex = re.compile(only_access_groups_matching_regex) if only_access_groups_matching_regex else None + return [model for model in model_list if _filter_model(model, model_regex, access_group_regex)] def _import_models_get_table_title(dry_run: bool) -> str: @@ -431,9 +390,7 @@ def _import_models_get_table_title(dry_run: bool) -> str: @models.command("import") -@click.argument( - "yaml_file", type=click.Path(exists=True, dir_okay=False, readable=True) -) +@click.argument("yaml_file", type=click.Path(exists=True, dir_okay=False, readable=True)) @click.option( "--dry-run", is_flag=True, diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index 51a3250162a..c0d45544b11 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -52,11 +52,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None: # Try to determine role (this might vary based on API response structure) role = "Member" # Default role - if ( - isinstance(team, dict) - and "members_with_roles" in team - and team["members_with_roles"] - ): + if isinstance(team, dict) and "members_with_roles" in team and team["members_with_roles"]: # This would need to be implemented based on actual API response structure pass diff --git a/litellm/proxy/client/cli/commands/users.py b/litellm/proxy/client/cli/commands/users.py index 36b29b8fe6b..1671b92c089 100644 --- a/litellm/proxy/client/cli/commands/users.py +++ b/litellm/proxy/client/cli/commands/users.py @@ -13,9 +13,7 @@ def users(): @click.pass_context def list_users(ctx: click.Context): """List all users""" - client = UsersManagementClient( - base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] - ) + client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) users = client.list_users() if isinstance(users, dict) and "users" in users: users = users["users"] @@ -46,9 +44,7 @@ def list_users(ctx: click.Context): @click.pass_context def get_user(ctx: click.Context, user_id: str): """Get information about a specific user""" - client = UsersManagementClient( - base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] - ) + client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) result = client.get_user(user_id=user_id) rich.print_json(data=result) @@ -62,9 +58,7 @@ def get_user(ctx: click.Context, user_id: str): @click.pass_context def create_user(ctx: click.Context, email, role, alias, team, max_budget): """Create a new user""" - client = UsersManagementClient( - base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] - ) + client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) user_data = { "user_email": email, "user_role": role, @@ -84,8 +78,6 @@ def create_user(ctx: click.Context, email, role, alias, team, max_budget): @click.pass_context def delete_user(ctx: click.Context, user_ids): """Delete one or more users by user_id""" - client = UsersManagementClient( - base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] - ) + client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) result = client.delete_user(list(user_ids)) rich.print_json(data=result) diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index a32d60aadd9..33f1f4a4480 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -57,9 +57,7 @@ def styled_prompt(): # Now move cursor up to the input line and get input click.echo("\033[2A", nl=False) # Move cursor up 2 lines - click.echo( - f"\r{left_border} {prompt_text}", nl=False - ) # Position at start of input line + click.echo(f"\r{left_border} {prompt_text}", nl=False) # Position at start of input line try: # Get user input diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index 929ad46a77c..f481a61c328 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -30,26 +30,14 @@ class Client: """ self._base_url = base_url.rstrip("/") # Only use the stored CLI key when it was issued for this server. - self._api_key = api_key or get_litellm_gateway_api_key( - expected_base_url=self._base_url - ) + self._api_key = api_key or get_litellm_gateway_api_key(expected_base_url=self._base_url) # Initialize resource clients - self.http = HTTPClient( - base_url=base_url, api_key=self._api_key, timeout=timeout - ) - self.models = ModelsManagementClient( - base_url=self._base_url, api_key=self._api_key - ) - self.model_groups = ModelGroupsManagementClient( - base_url=self._base_url, api_key=self._api_key - ) + self.http = HTTPClient(base_url=base_url, api_key=self._api_key, timeout=timeout) + self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key) + self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key) self.chat = ChatClient(base_url=self._base_url, api_key=self._api_key) self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key) - self.credentials = CredentialsManagementClient( - base_url=self._base_url, api_key=self._api_key - ) - self.teams = TeamsManagementClient( - base_url=self._base_url, api_key=self._api_key - ) + self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key) + self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key) diff --git a/litellm/proxy/client/exceptions.py b/litellm/proxy/client/exceptions.py index c4089381e30..34884c0b995 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -9,9 +9,7 @@ def _redact_orig_exception( orig_exception: Union[requests.exceptions.HTTPError, str], ) -> Union[requests.exceptions.HTTPError, str]: if isinstance(orig_exception, requests.exceptions.HTTPError): - return requests.exceptions.HTTPError( - redact_string(str(orig_exception)), response=orig_exception.response - ) + return requests.exceptions.HTTPError(redact_string(str(orig_exception)), response=orig_exception.response) return redact_string(str(orig_exception)) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index 845b49d1581..dcf915d812a 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -91,9 +91,7 @@ class KeysManagementClient: if include_team_keys is not None: params["include_team_keys"] = str(include_team_keys).lower() - request = requests.Request( - "GET", url, headers=self._get_headers(), params=params - ) + request = requests.Request("GET", url, headers=self._get_headers(), params=params) if return_request: return request @@ -287,9 +285,7 @@ class KeysManagementClient: except Exception: raise Exception(f"Error updating key: {response_text}") - def info( - self, key: str, return_request: bool = False - ) -> Union[Dict[str, Any], requests.Request]: + def info(self, key: str, return_request: bool = False) -> Union[Dict[str, Any], requests.Request]: """ Get information about API keys. @@ -319,6 +315,4 @@ class KeysManagementClient: redacted_message = redact_string(str(e)) if e.response.status_code == 401: raise UnauthorizedError(e) from None - raise requests.exceptions.HTTPError( - redacted_message, response=e.response - ) from None + raise requests.exceptions.HTTPError(redacted_message, response=e.response) from None diff --git a/litellm/proxy/client/model_groups.py b/litellm/proxy/client/model_groups.py index 03bc3eae466..2be6e10e542 100644 --- a/litellm/proxy/client/model_groups.py +++ b/litellm/proxy/client/model_groups.py @@ -27,9 +27,7 @@ class ModelGroupsManagementClient: headers["Authorization"] = f"Bearer {self._api_key}" return headers - def info( - self, return_request: bool = False - ) -> Union[List[Dict[str, Any]], requests.Request]: + def info(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]: """ Get detailed information about all model groups from the server. diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index d2f5eead284..bb375426295 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -27,9 +27,7 @@ class ModelsManagementClient: headers["Authorization"] = f"Bearer {self._api_key}" return headers - def list( - self, return_request: bool = False - ) -> Union[List[Dict[str, Any]], requests.Request]: + def list(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]: """ Get the list of models supported by the server. @@ -111,9 +109,7 @@ class ModelsManagementClient: raise UnauthorizedError(e) raise - def delete( - self, model_id: str, return_request: bool = False - ) -> Union[Dict[str, Any], requests.Request]: + def delete(self, model_id: str, return_request: bool = False) -> Union[Dict[str, Any], requests.Request]: """ Delete a model from the proxy. @@ -175,9 +171,7 @@ class ModelsManagementClient: NotFoundError: If the model is not found requests.exceptions.RequestException: If the request fails with any other error """ - if (model_id is None and model_name is None) or ( - model_id is not None and model_name is not None - ): + if (model_id is None and model_name is None) or (model_id is not None and model_name is not None): raise ValueError("Exactly one of model_id or model_name must be provided") # If return_request is True, delegate to info @@ -211,9 +205,7 @@ class ModelsManagementClient: ) ) - def info( - self, return_request: bool = False - ) -> Union[List[Dict[str, Any]], requests.Request]: + def info(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]: """ Get detailed information about all models from the server. diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index 9f80e171914..9f2d53c6d7d 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -14,9 +14,7 @@ class UsersManagementClient: headers["Authorization"] = f"Bearer {self.api_key}" return headers - def list_users( - self, params: Optional[Dict[str, Any]] = None - ) -> List[Dict[str, Any]]: + def list_users(self, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: """List users (GET /user/list)""" url = f"{self.base_url}/user/list" response = requests.get(url, headers=self._get_headers(), params=params) @@ -61,9 +59,7 @@ class UsersManagementClient: def delete_user(self, user_ids: List[str]) -> Dict[str, Any]: """Delete users (POST /user/delete)""" url = f"{self.base_url}/user/delete" - response = requests.post( - url, headers=self._get_headers(), json={"user_ids": user_ids} - ) + response = requests.post(url, headers=self._get_headers(), json={"user_ids": user_ids}) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b5e98dbc2af..4d931a47e9d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -114,9 +114,7 @@ async def _record_streaming_client_disconnect_if_needed( if logging_obj is not None: litellm_params = logging_obj.model_call_details.setdefault("litellm_params", {}) _apply_client_disconnect_metadata(litellm_params.setdefault("metadata", {})) - _apply_client_disconnect_metadata( - logging_obj.model_call_details.setdefault("metadata", {}) - ) + _apply_client_disconnect_metadata(logging_obj.model_call_details.setdefault("metadata", {})) _apply_client_disconnect_metadata(request_data.setdefault("metadata", {})) litellm_params = request_data.setdefault("litellm_params", {}) @@ -184,17 +182,13 @@ def _collect_response_file_search_vector_store_ids(data: Dict[str, Any]) -> set[ if not isinstance(ids, list): raise HTTPException( status_code=400, - detail={ - "error": "file_search.vector_store_ids must be a list of strings" - }, + detail={"error": "file_search.vector_store_ids must be a list of strings"}, ) for vector_store_id in ids: if not isinstance(vector_store_id, str) or not vector_store_id: raise HTTPException( status_code=400, - detail={ - "error": "file_search.vector_store_ids must be a list of strings" - }, + detail={"error": "file_search.vector_store_ids must be a list of strings"}, ) vector_store_ids.add(vector_store_id) @@ -222,20 +216,14 @@ async def _authorize_response_file_search_vector_stores( async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]: """Parses an event line and returns an error code if present, else None.""" - event_line = ( - event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line - ) + event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line if event_line.startswith("data: "): json_str = event_line[len("data: ") :].strip() if not json_str or json_str == "[DONE]": # handle empty data or [DONE] message return None try: data = orjson.loads(json_str) - if ( - isinstance(data, dict) - and "error" in data - and isinstance(data["error"], dict) - ): + if isinstance(data, dict) and "error" in data and isinstance(data["error"], dict): error_code_raw = data["error"].get("code") error_code: Optional[int] = None @@ -254,12 +242,8 @@ async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional # Ensure error_code is a valid HTTP status code if error_code is not None and 100 <= error_code <= 599: return error_code - elif ( - error_code_raw is not None - ): # Log if original code was present but not valid - verbose_proxy_logger.warning( - f"Error has invalid or non-convertible code: {error_code_raw}" - ) + elif error_code_raw is not None: # Log if original code was present but not valid + verbose_proxy_logger.warning(f"Error has invalid or non-convertible code: {error_code_raw}") except (orjson.JSONDecodeError, json.JSONDecodeError): # not a known error chunk pass @@ -276,9 +260,7 @@ def _extract_error_from_sse_chunk(event_line: Union[str, bytes]) -> dict: Returns: Error dictionary in OpenAI API format """ - event_line = ( - event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line - ) + event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line # Default error format default_error = { @@ -327,9 +309,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse): status_code: int = status.HTTP_200_OK, upstream_generator: Optional[AsyncGenerator[str, None]] = None, ) -> None: - super().__init__( - content, status_code=status_code, headers=headers, media_type=media_type - ) + super().__init__(content, status_code=status_code, headers=headers, media_type=media_type) self._upstream_generator = upstream_generator async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: @@ -344,9 +324,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse): try: await aclose() except BaseException as e: - verbose_proxy_logger.debug( - "error closing streaming generator: %s", e - ) + verbose_proxy_logger.debug("error closing streaming generator: %s", e) class _ClientDisconnectedBeforeFirstChunk(Exception): @@ -366,8 +344,7 @@ async def _wait_for_http_disconnect(request: Request) -> None: raise except Exception as exc: # noqa: BLE001 verbose_proxy_logger.warning( - "create_response: request.receive() raised %s; first-chunk disconnect " - "monitoring disabled for this request", + "create_response: request.receive() raised %s; first-chunk disconnect monitoring disabled for this request", exc, ) # A receive() failure must not masquerade as a disconnect. @@ -393,13 +370,9 @@ async def _buffer_first_chunk_honoring_disconnect( return await generator.__anext__() chunk_task: asyncio.Task[str] = asyncio.ensure_future(generator.__anext__()) - disconnect_task: asyncio.Task[None] = asyncio.ensure_future( - _wait_for_http_disconnect(request) - ) + disconnect_task: asyncio.Task[None] = asyncio.ensure_future(_wait_for_http_disconnect(request)) try: - await asyncio.wait( - {chunk_task, disconnect_task}, return_when=asyncio.FIRST_COMPLETED - ) + await asyncio.wait({chunk_task, disconnect_task}, return_when=asyncio.FIRST_COMPLETED) # A completed disconnect_task has already consumed the http.disconnect # message, so Starlette's later listen_for_disconnect would never see it. # Take the cancellation path whenever a disconnect was observed, even if @@ -424,13 +397,8 @@ async def _buffer_first_chunk_honoring_disconnect( try: await generator.aclose() except BaseException as exc: # noqa: BLE001 - verbose_proxy_logger.debug( - "create_response: error closing generator on disconnect: %s", exc - ) - verbose_proxy_logger.info( - "create_response: client disconnected before first chunk, " - "upstream LLM request cancelled" - ) + verbose_proxy_logger.debug("create_response: error closing generator on disconnect: %s", exc) + verbose_proxy_logger.info("create_response: client disconnected before first chunk, upstream LLM request cancelled") raise _ClientDisconnectedBeforeFirstChunk() @@ -462,15 +430,11 @@ async def create_response( generator = await generator # Now get the first chunk from the actual generator - first_chunk_value = await _buffer_first_chunk_honoring_disconnect( - generator, request - ) + first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) if first_chunk_value is not None: try: - error_code_from_chunk = await _parse_event_data_for_error( - first_chunk_value - ) + error_code_from_chunk = await _parse_event_data_for_error(first_chunk_value) if error_code_from_chunk is not None: # First chunk is an error, stream hasn't really started yet # Should return standard JSON error response instead of SSE format @@ -526,9 +490,7 @@ async def create_response( ) except Exception as e: # Unexpected error consuming first chunk. - verbose_proxy_logger.exception( - f"Error consuming first chunk from generator: {e}" - ) + verbose_proxy_logger.exception(f"Error consuming first chunk from generator: {e}") # Preserve status code from HTTPException (e.g., guardrail blocks) error_status = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -642,9 +604,7 @@ def _override_openai_response_model( if isinstance(hidden_params, dict): # Check if a fallback occurred - if so, preserve the actual model used fallback_headers = hidden_params.get("additional_headers", {}) or {} - attempted_fallbacks = fallback_headers.get( - "x-litellm-attempted-fallbacks", None - ) + attempted_fallbacks = fallback_headers.get("x-litellm-attempted-fallbacks", None) if attempted_fallbacks is not None and attempted_fallbacks > 0: verbose_proxy_logger.debug( "%s: fallback detected (attempted_fallbacks=%d), preserving actual model used instead of overriding to requested model.", @@ -773,10 +733,7 @@ _CLIENT_DISCONNECT_DETAIL = "Client disconnected the request" def _log_llm_api_exception(e: Exception) -> None: - if ( - getattr(e, "status_code", None) == 499 - and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL - ): + if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL: verbose_proxy_logger.info( "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled" ) @@ -800,8 +757,7 @@ async def _cancel_llm_call_on_client_disconnect( return except Exception as exc: verbose_proxy_logger.warning( - "cancel_on_disconnect: request.receive() raised %s; " - "upstream LLM call will not be cancelled on disconnect", + "cancel_on_disconnect: request.receive() raised %s; upstream LLM call will not be cancelled on disconnect", exc, ) @@ -811,9 +767,7 @@ async def _await_llm_call_cancelling_on_disconnect( llm_api_call: "asyncio.Future[Any]", ) -> Any: disconnect_event = asyncio.Event() - monitor = asyncio.create_task( - _cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event) - ) + monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event)) try: return await llm_api_call except asyncio.CancelledError: @@ -858,9 +812,7 @@ class ProxyBaseLLMRequestProcessing: discount_amount, margin_total_amount, margin_percent, - ) = _get_cost_breakdown_from_logging_obj( - litellm_logging_obj=litellm_logging_obj - ) + ) = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=litellm_logging_obj) # Calculate updated spend for header (include current response_cost) current_spend = user_api_key_dict.spend or 0.0 @@ -868,11 +820,7 @@ class ProxyBaseLLMRequestProcessing: if response_cost is not None: try: # Convert response_cost to float if it's a string - cost_value = ( - float(response_cost) - if isinstance(response_cost, str) - else response_cost - ) + cost_value = float(response_cost) if isinstance(response_cost, str) else response_cost if cost_value > 0: updated_spend = current_spend + cost_value except (ValueError, TypeError): @@ -889,61 +837,37 @@ class ProxyBaseLLMRequestProcessing: "x-litellm-version": version, "x-litellm-model-region": model_region, "x-litellm-response-cost": str(response_cost), - "x-litellm-response-cost-original": ( - str(original_cost) if original_cost is not None else None - ), - "x-litellm-response-cost-discount-amount": ( - str(discount_amount) if discount_amount is not None else None - ), + "x-litellm-response-cost-original": (str(original_cost) if original_cost is not None else None), + "x-litellm-response-cost-discount-amount": (str(discount_amount) if discount_amount is not None else None), "x-litellm-response-cost-margin-amount": ( str(margin_total_amount) if margin_total_amount is not None else None ), - "x-litellm-response-cost-margin-percent": ( - str(margin_percent) if margin_percent is not None else None - ), + "x-litellm-response-cost-margin-percent": (str(margin_percent) if margin_percent is not None else None), "x-litellm-key-tpm-limit": str(user_api_key_dict.tpm_limit), "x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit), "x-litellm-key-max-budget": str(user_api_key_dict.max_budget), "x-litellm-key-spend": str(updated_spend), - "x-litellm-response-duration-ms": str( - hidden_params.get("_response_ms", None) - ), - "x-litellm-overhead-duration-ms": str( - hidden_params.get("litellm_overhead_time_ms", None) - ), - "x-litellm-callback-duration-ms": str( - hidden_params.get("callback_duration_ms", None) - ), + "x-litellm-response-duration-ms": str(hidden_params.get("_response_ms", None)), + "x-litellm-overhead-duration-ms": str(hidden_params.get("litellm_overhead_time_ms", None)), + "x-litellm-callback-duration-ms": str(hidden_params.get("callback_duration_ms", None)), **( { - "x-litellm-timing-pre-processing-ms": str( - hidden_params.get("timing_pre_processing_ms", None) - ), - "x-litellm-timing-llm-api-ms": str( - hidden_params.get("timing_llm_api_ms", None) - ), - "x-litellm-timing-post-processing-ms": str( - hidden_params.get("timing_post_processing_ms", None) - ), - "x-litellm-timing-message-copy-ms": str( - hidden_params.get("timing_message_copy_ms", None) - ), + "x-litellm-timing-pre-processing-ms": str(hidden_params.get("timing_pre_processing_ms", None)), + "x-litellm-timing-llm-api-ms": str(hidden_params.get("timing_llm_api_ms", None)), + "x-litellm-timing-post-processing-ms": str(hidden_params.get("timing_post_processing_ms", None)), + "x-litellm-timing-message-copy-ms": str(hidden_params.get("timing_message_copy_ms", None)), } if LITELLM_DETAILED_TIMING else {} ), "x-litellm-fastest_response_batch_completion": ( - str(fastest_response_batch_completion) - if fastest_response_batch_completion is not None - else None + str(fastest_response_batch_completion) if fastest_response_batch_completion is not None else None ), "x-litellm-timeout": str(timeout) if timeout is not None else None, **{k: str(v) for k, v in kwargs.items()}, } if request_data: - remaining_tokens_header = ( - get_remaining_tokens_and_requests_from_request_data(request_data) - ) + remaining_tokens_header = get_remaining_tokens_and_requests_from_request_data(request_data) headers.update(remaining_tokens_header) logging_caching_headers = get_logging_caching_headers(request_data) @@ -951,11 +875,7 @@ class ProxyBaseLLMRequestProcessing: headers.update(logging_caching_headers) try: - return { - key: str(value) - for key, value in headers.items() - if value not in exclude_values - } + return {key: str(value) for key, value in headers.items() if value not in exclude_values} except Exception as e: verbose_proxy_logger.error(f"Error setting custom headers: {e}") return {} @@ -982,16 +902,12 @@ class ProxyBaseLLMRequestProcessing: if not isinstance(hidden_params, dict): hidden_params = {} - model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response( - hidden_params, request_data - ) + model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response(hidden_params, request_data) cache_key = hidden_params.get("cache_key", None) or "" api_base = hidden_params.get("api_base", None) or "" response_cost = hidden_params.get("response_cost", None) or "" - fastest_response_batch_completion = hidden_params.get( - "fastest_response_batch_completion", None - ) + fastest_response_batch_completion = hidden_params.get("fastest_response_batch_completion", None) additional_headers = hidden_params.get("additional_headers", {}) or {} custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -1164,9 +1080,7 @@ class ProxyBaseLLMRequestProcessing: self.data[_metadata_variable_name] = {} if not isinstance(self.data[_metadata_variable_name], dict): self.data[_metadata_variable_name] = {} - self.data[_metadata_variable_name]["queue_time_seconds"] = ( - queue_time_seconds - ) + self.data[_metadata_variable_name]["queue_time_seconds"] = queue_time_seconds self.data["model"] = ( general_settings.get("completion_model", None) # server default @@ -1188,10 +1102,7 @@ class ProxyBaseLLMRequestProcessing: ### MODEL ALIAS MAPPING ### # check if model name in model alias map # get the actual model name - if ( - isinstance(self.data["model"], str) - and self.data["model"] in litellm.model_alias_map - ): + if isinstance(self.data["model"], str) and self.data["model"] in litellm.model_alias_map: self.data["model"] = litellm.model_alias_map[self.data["model"]] # Check key-specific aliases @@ -1203,9 +1114,7 @@ class ProxyBaseLLMRequestProcessing: ): self.data["model"] = user_api_key_dict.aliases[self.data["model"]] - self.data["litellm_call_id"] = request.headers.get( - "x-litellm-call-id", str(uuid.uuid4()) - ) + self.data["litellm_call_id"] = request.headers.get("x-litellm-call-id", str(uuid.uuid4())) DDSpanTagger.tag_call_id(self.data.get("litellm_call_id")) DDSpanTagger.tag_request( user_api_key_dict=user_api_key_dict, @@ -1222,10 +1131,7 @@ class ProxyBaseLLMRequestProcessing: # Only set if stream_options is not already provided by the client if "stream_options" not in self.data: self.data["stream_options"] = {"include_usage": True} - elif ( - isinstance(self.data["stream_options"], dict) - and "include_usage" not in self.data["stream_options"] - ): + elif isinstance(self.data["stream_options"], dict) and "include_usage" not in self.data["stream_options"]: self.data["stream_options"]["include_usage"] = True ### CALL HOOKS ### - modify/reject incoming data before calling the model @@ -1288,11 +1194,7 @@ class ProxyBaseLLMRequestProcessing: "Request received by LiteLLM: payload too large to log (%d bytes, limit %d). Keys: %s", len(_payload_str), MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, - ( - list(self.data.keys()) - if isinstance(self.data, dict) - else type(self.data).__name__ - ), + (list(self.data.keys()) if isinstance(self.data, dict) else type(self.data).__name__), ) else: verbose_proxy_logger.debug( @@ -1493,15 +1395,11 @@ class ProxyBaseLLMRequestProcessing: llm_call_task = asyncio.create_task(llm_call) tasks.append(llm_call_task) - llm_responses = asyncio.gather( - *tasks - ) # run the moderation check in parallel to the actual llm api call + llm_responses = asyncio.gather(*tasks) # run the moderation check in parallel to the actual llm api call try: if general_settings.get("cancel_on_disconnect", False): - responses = await _await_llm_call_cancelling_on_disconnect( - request, llm_responses - ) + responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses) else: responses = await llm_responses finally: @@ -1535,9 +1433,7 @@ class ProxyBaseLLMRequestProcessing: ) if self._is_streaming_request( data=self.data, is_streaming_request=is_streaming_request - ) or self._is_streaming_response( - response - ): # use generate_responses to stream responses + ) or self._is_streaming_response(response): # use generate_responses to stream responses custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, call_id=logging_obj.litellm_call_id, @@ -1555,13 +1451,11 @@ class ProxyBaseLLMRequestProcessing: ) # Call response headers hook for streaming success - callback_headers = ( - await proxy_logging_obj.post_call_response_headers_hook( - data=self.data, - user_api_key_dict=user_api_key_dict, - response=response, - request_headers=dict(request.headers), - ) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=self.data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), ) if callback_headers: custom_headers.update(callback_headers) @@ -1571,9 +1465,7 @@ class ProxyBaseLLMRequestProcessing: # aliasing/routing, but the OpenAI-compatible response `model` field should reflect # what the client sent. if requested_model_from_client: - self.data["_litellm_client_requested_model"] = ( - requested_model_from_client - ) + self.data["_litellm_client_requested_model"] = requested_model_from_client # Streaming: attach a closure that fires after all guardrail # end-of-stream blocks complete. CSW.__anext__ stores the @@ -1588,9 +1480,7 @@ class ProxyBaseLLMRequestProcessing: CustomStreamWrapper, ) - if _post_call_guardrails_active and isinstance( - response, CustomStreamWrapper - ): + if _post_call_guardrails_active and isinstance(response, CustomStreamWrapper): # Intentionally a live reference (not a copy) — mirrors # ProxyLogging.post_call_success_hook which also mutates # data["guardrail_to_apply"] during iteration. @@ -1598,9 +1488,7 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict = user_api_key_dict _captured_logging_obj = logging_obj - async def _on_deferred_stream_complete( - assembled_response, cache_hit - ): + async def _on_deferred_stream_complete(assembled_response, cache_hit): await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( captured_data=_captured_data, captured_user_api_key_dict=_captured_user_api_key_dict, @@ -1609,9 +1497,7 @@ class ProxyBaseLLMRequestProcessing: cache_hit=cache_hit, ) - logging_obj._on_deferred_stream_complete = ( - _on_deferred_stream_complete # type: ignore[union-attr] - ) + logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[union-attr] if route_type == "allm_passthrough_route": # Check if response is an async generator @@ -1628,17 +1514,13 @@ class ProxyBaseLLMRequestProcessing: body_bytes = b"".join( [chunk async for chunk in generator] # type: ignore[union-attr] ) - modified_bytes = ( - await self._handle_event_stream_allm_passthrough_route( - body_bytes=body_bytes, - proxy_logging_obj=proxy_logging_obj, - user_api_key_dict=user_api_key_dict, - ) + modified_bytes = await self._handle_event_stream_allm_passthrough_route( + body_bytes=body_bytes, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, ) response_headers = { - k: v - for k, v in custom_headers.items() - if k.lower() != "content-length" + k: v for k, v in custom_headers.items() if k.lower() != "content-length" } return Response( content=modified_bytes, @@ -1655,14 +1537,12 @@ class ProxyBaseLLMRequestProcessing: headers=custom_headers, ) else: - _early = ( - await self._handle_non_streaming_allm_passthrough_route( - response=response, - proxy_logging_obj=proxy_logging_obj, - user_api_key_dict=user_api_key_dict, - custom_headers=custom_headers, - request_headers=dict(request.headers), - ) + _early = await self._handle_non_streaming_allm_passthrough_route( + response=response, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, + custom_headers=custom_headers, + request_headers=dict(request.headers), ) if _early is not None: return _early @@ -1677,14 +1557,12 @@ class ProxyBaseLLMRequestProcessing: # This handles cases like websearch_interception agentic loop # which returns a non-streaming dict even for streaming requests if self._is_streaming_response(response): - selected_data_generator = ( - ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=self.data, - proxy_logging_obj=proxy_logging_obj, - request=request, - ) + selected_data_generator = ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=self.data, + proxy_logging_obj=proxy_logging_obj, + request=request, ) return await create_response( generator=selected_data_generator, @@ -1710,10 +1588,12 @@ class ProxyBaseLLMRequestProcessing: # stream stay unregistered and follow-up file API # calls 403. Covers the background-polling path # too, which loops ``body_iterator`` end-to-end. - selected_data_generator = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( - original_stream_response=response, - wrapped_generator=selected_data_generator, - user_api_key_dict=user_api_key_dict, + selected_data_generator = ( + ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( + original_stream_response=response, + wrapped_generator=selected_data_generator, + user_api_key_dict=user_api_key_dict, + ) ) return await create_response( generator=selected_data_generator, @@ -1779,9 +1659,7 @@ class ProxyBaseLLMRequestProcessing: # it at stream end. _exception_raised is function-scoped and # immune to outer exception context, avoiding false positives. if _exception_raised: - _deferred_fn = getattr( - logging_obj, "_on_deferred_stream_complete", None - ) + _deferred_fn = getattr(logging_obj, "_on_deferred_stream_complete", None) if _deferred_fn is not None: logging_obj._on_deferred_stream_complete = None # type: ignore[union-attr] try: @@ -1795,9 +1673,7 @@ class ProxyBaseLLMRequestProcessing: ) ) except Exception as e: - verbose_proxy_logger.exception( - "Error in orphaned streaming async logging: %s", e - ) + verbose_proxy_logger.exception("Error in orphaned streaming async logging: %s", e) # Always return the client-requested model name (not provider-prefixed internal identifiers) # for OpenAI-compatible responses. @@ -1808,9 +1684,7 @@ class ProxyBaseLLMRequestProcessing: log_context=f"litellm_call_id={logging_obj.litellm_call_id}", ) - hidden_params = ( - getattr(response, "_hidden_params", {}) or {} - ) # get any updated response headers + hidden_params = getattr(response, "_hidden_params", {}) or {} # get any updated response headers additional_headers = hidden_params.get("additional_headers", {}) or {} fastapi_response.headers.update( @@ -1914,10 +1788,8 @@ class ProxyBaseLLMRequestProcessing: yield chunk finally: try: - completed_obj = ( - ProxyBaseLLMRequestProcessing._extract_completed_responses_response( - original_stream_response - ) + completed_obj = ProxyBaseLLMRequestProcessing._extract_completed_responses_response( + original_stream_response ) if completed_obj is not None: await ProxyBaseLLMRequestProcessing._record_container_owners_from_responses_if_needed( @@ -2027,9 +1899,7 @@ class ProxyBaseLLMRequestProcessing: return False - def _is_streaming_request( - self, data: dict, is_streaming_request: Optional[bool] = False - ) -> bool: + def _is_streaming_request(self, data: dict, is_streaming_request: Optional[bool] = False) -> bool: """ Check if the request is a streaming request. @@ -2072,9 +1942,7 @@ class ProxyBaseLLMRequestProcessing: from litellm.proxy.proxy_server import llm_router from litellm.proxy.utils import _check_and_merge_model_level_guardrails - guardrail_data = _check_and_merge_model_level_guardrails( - data=self.data, llm_router=llm_router - ) + guardrail_data = _check_and_merge_model_level_guardrails(data=self.data, llm_router=llm_router) for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue @@ -2115,9 +1983,7 @@ class ProxyBaseLLMRequestProcessing: LlmPassthroughRouteHandler, ) - return LlmPassthroughRouteHandler.event_stream_media_type( - self.data.get("custom_llm_provider") - ) + return LlmPassthroughRouteHandler.event_stream_media_type(self.data.get("custom_llm_provider")) async def _handle_non_streaming_allm_passthrough_route( self, @@ -2200,8 +2066,7 @@ class ProxyBaseLLMRequestProcessing: content = _json.dumps(processed).encode() else: verbose_proxy_logger.debug( - "allm_passthrough_route: post_call_success_hook returned %s, " - "leaving JSON response unmodified", + "allm_passthrough_route: post_call_success_hook returned %s, leaving JSON response unmodified", type(processed).__name__, ) content = body_bytes @@ -2298,9 +2163,7 @@ class ProxyBaseLLMRequestProcessing: from litellm.proxy.proxy_server import llm_router as _global_llm_router from litellm.proxy.utils import _check_and_merge_model_level_guardrails - guardrail_data = _check_and_merge_model_level_guardrails( - data=captured_data, llm_router=_global_llm_router - ) + guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue @@ -2338,12 +2201,8 @@ class ProxyBaseLLMRequestProcessing: getattr(cb, "guardrail_name", type(cb).__name__), e, ) - if isinstance(e, HTTPException) and hasattr( - captured_logging_obj, "model_call_details" - ): - captured_logging_obj.model_call_details.setdefault( - "metadata", {} - )["guardrail_blocked"] = True + if isinstance(e, HTTPException) and hasattr(captured_logging_obj, "model_call_details"): + captured_logging_obj.model_call_details.setdefault("metadata", {})["guardrail_blocked"] = True except Exception as e: verbose_proxy_logger.exception( "Error in deferred streaming guardrail initialization: %s", @@ -2403,9 +2262,7 @@ class ProxyBaseLLMRequestProcessing: timeout = getattr( e, "timeout", None ) # returns the timeout set by the wrapper. Used for testing if model-specific timeout are set correctly - _litellm_logging_obj: Optional[LiteLLMLoggingObj] = self.data.get( - "litellm_logging_obj", None - ) + _litellm_logging_obj: Optional[LiteLLMLoggingObj] = self.data.get("litellm_logging_obj", None) # Attempt to get model_id from logging object # @@ -2416,9 +2273,7 @@ class ProxyBaseLLMRequestProcessing: custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, call_id=( - _litellm_logging_obj.litellm_call_id - if _litellm_logging_obj - else self.data.get("litellm_call_id") + _litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id") ), model_id=model_id, version=version, @@ -2445,9 +2300,7 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, response=None, - request_headers=(self.data.get("proxy_server_request") or {}).get( - "headers", {} - ), + request_headers=(self.data.get("proxy_server_request") or {}).get("headers", {}), ) if callback_headers: headers.update(callback_headers) @@ -2511,11 +2364,7 @@ class ProxyBaseLLMRequestProcessing: # the upstream API response. Use it to return the correct HTTP code # instead of defaulting to 500. _exc_status_code = getattr(e, "status_code", None) - if ( - _exc_status_code is not None - and isinstance(_exc_status_code, int) - and 400 <= _exc_status_code <= 599 - ): + if _exc_status_code is not None and isinstance(_exc_status_code, int) and 400 <= _exc_status_code <= 599: _code = _exc_status_code else: _code = status.HTTP_500_INTERNAL_SERVER_ERROR @@ -2560,17 +2409,13 @@ class ProxyBaseLLMRequestProcessing: client_disconnected: bool = False, ) -> None: with anyio.CancelScope(shield=True): - should_record_client_disconnect = client_disconnected or ( - not stream_completed - ) + should_record_client_disconnect = client_disconnected or (not stream_completed) recorded_client_disconnect = False if should_record_client_disconnect: - recorded_client_disconnect = ( - await _record_streaming_client_disconnect_if_needed( - request, - request_data, - client_disconnected, - ) + recorded_client_disconnect = await _record_streaming_client_disconnect_if_needed( + request, + request_data, + client_disconnected, ) if recorded_client_disconnect: ProxyLogging._fire_deferred_stream_logging(request_data) @@ -2609,23 +2454,15 @@ class ProxyBaseLLMRequestProcessing: # await, response-string materialization, and cost-injection call are # pure overhead on the streaming hot path (the default config). caps = ProxyLogging._callback_capabilities() - cost_injection_enabled = bool( - getattr(litellm, "include_cost_in_streaming_usage", False) - ) - fast_path = ( - not caps.has_streaming_chunk_override - and not caps.has_guardrail - and not cost_injection_enabled - ) + cost_injection_enabled = bool(getattr(litellm, "include_cost_in_streaming_usage", False)) + fast_path = not caps.has_streaming_chunk_override and not caps.has_guardrail and not cost_injection_enabled debug_enabled = verbose_proxy_logger.isEnabledFor(logging.DEBUG) stream_completed = False client_disconnected = False delivered_chunk = False try: str_so_far = "" - async for ( - chunk - ) in proxy_logging_obj.async_post_call_streaming_iterator_hook( + async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=response, request_data=request_data, @@ -2633,9 +2470,7 @@ class ProxyBaseLLMRequestProcessing: # ``.format(chunk)`` was previously evaluated for every chunk # regardless of log level; gate it behind the level check. if debug_enabled: - verbose_proxy_logger.debug( - "async_data_generator: received streaming chunk - %s", chunk - ) + verbose_proxy_logger.debug("async_data_generator: received streaming chunk - %s", chunk) if not fast_path: chunk = await proxy_logging_obj.async_post_call_streaming_hook( @@ -2659,9 +2494,7 @@ class ProxyBaseLLMRequestProcessing: str_so_far += str(chunk.get("content", "")) model_name = request_data.get("model", "") - chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - chunk, model_name - ) + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) # Set before the yield: an async generator suspends at the yield, # so a GeneratorExit on client disconnect is raised there and any @@ -2679,24 +2512,18 @@ class ProxyBaseLLMRequestProcessing: # on disconnect, so the nested iterator hook (which only sees # GeneratorExit on GC) cannot own the refund. if not stream_completed: - proxy_logging_obj._release_max_parallel_requests_on_disconnect( - user_api_key_dict - ) + proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict) client_disconnected = True if not delivered_chunk: from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, ) - await release_budget_reservation_on_cancel( - getattr(user_api_key_dict, "budget_reservation", None) - ) + await release_budget_reservation_on_cancel(getattr(user_api_key_dict, "budget_reservation", None)) raise except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format(str(e)) ) transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -2775,42 +2602,24 @@ class ProxyBaseLLMRequestProcessing: try: if isinstance(chunk, dict): - maybe_modified = ( - ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( - chunk, model_name - ) - ) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) if maybe_modified is not None: return maybe_modified elif isinstance(chunk, (bytes, bytearray)): # Decode to str, inject, and rebuild as bytes try: s = chunk.decode("utf-8", errors="ignore") - maybe_mod = ( - ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( - s, model_name - ) - ) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) if maybe_mod is not None: - return ( - maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n") - ).encode("utf-8") + return (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") except Exception: pass elif isinstance(chunk, str): # Try to parse SSE frame and inject cost into the data line - maybe_mod = ( - ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( - chunk, model_name - ) - ) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) if maybe_mod is not None: # Ensure trailing frame separator - return ( - maybe_mod - if maybe_mod.endswith("\n\n") - else (maybe_mod + "\n\n") - ) + return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") except Exception: # Never break streaming on optional cost injection pass @@ -2818,9 +2627,7 @@ class ProxyBaseLLMRequestProcessing: return chunk @staticmethod - def _inject_cost_into_sse_frame_str( - frame_str: str, model_name: str - ) -> Optional[str]: + def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> Optional[str]: """ Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. @@ -2840,11 +2647,7 @@ class ProxyBaseLLMRequestProcessing: json_part = stripped_ln.split("data:", 1)[1].strip() if json_part and json_part != "[DONE]": obj = json.loads(json_part) - maybe_modified = ( - ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( - obj, model_name - ) - ) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) if maybe_modified is not None: # Replace just this line with updated JSON using safe_dumps lines[idx] = f"data: {safe_dumps(maybe_modified)}" @@ -2870,8 +2673,7 @@ class ProxyBaseLLMRequestProcessing: prompt_tokens = int(_usage.get("input_tokens", 0) or 0) completion_tokens = int(_usage.get("output_tokens", 0) or 0) total_tokens = int( - _usage.get("total_tokens", prompt_tokens + completion_tokens) - or (prompt_tokens + completion_tokens) + _usage.get("total_tokens", prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens) ) # Extract additional usage fields @@ -2895,15 +2697,11 @@ class ProxyBaseLLMRequestProcessing: # Handle web_search_requests by wrapping in ServerToolUse if web_search_requests is not None: - usage_kwargs["server_tool_use"] = ServerToolUse( - web_search_requests=web_search_requests - ) + usage_kwargs["server_tool_use"] = ServerToolUse(web_search_requests=web_search_requests) # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs["cache_creation_input_tokens"] = ( - cache_creation_input_tokens - ) + usage_kwargs["cache_creation_input_tokens"] = cache_creation_input_tokens if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens @@ -2922,9 +2720,7 @@ class ProxyBaseLLMRequestProcessing: return obj return None - def maybe_get_model_id( - self, _logging_obj: Optional[LiteLLMLoggingObj] - ) -> Optional[str]: + def maybe_get_model_id(self, _logging_obj: Optional[LiteLLMLoggingObj]) -> Optional[str]: """ Get model_id from logging object or request metadata. diff --git a/litellm/proxy/common_utils/admin_ui_utils.py b/litellm/proxy/common_utils/admin_ui_utils.py index 41fd3a1a76f..07cd6c540a6 100644 --- a/litellm/proxy/common_utils/admin_ui_utils.py +++ b/litellm/proxy/common_utils/admin_ui_utils.py @@ -5,15 +5,11 @@ def show_missing_vars_in_env(): if prisma_client is None and master_key is None: return HTMLResponse( - content=missing_keys_form( - missing_key_names="DATABASE_URL, LITELLM_MASTER_KEY" - ), + content=missing_keys_form(missing_key_names="DATABASE_URL, LITELLM_MASTER_KEY"), status_code=200, ) if prisma_client is None: - return HTMLResponse( - content=missing_keys_form(missing_key_names="DATABASE_URL"), status_code=200 - ) + return HTMLResponse(content=missing_keys_form(missing_key_names="DATABASE_URL"), status_code=200) if master_key is None: return HTMLResponse( diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py index abb0402d3b9..2f09e92f1ba 100644 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -63,17 +63,13 @@ class EventDrivenCacheCoordinator: self._query_in_progress = False self._log_prefix = log_prefix - async def _get_cached( - self, cache_key: str, cache: AsyncCacheProtocol - ) -> Optional[Any]: + async def _get_cached(self, cache_key: str, cache: AsyncCacheProtocol) -> Optional[Any]: """Return value from cache if present, else None.""" return await cache.async_get_cache(key=cache_key) def _log_cache_hit(self, value: T) -> None: if self._log_prefix: - verbose_proxy_logger.debug( - "%s Cache hit, value: %s", self._log_prefix, value - ) + verbose_proxy_logger.debug("%s Cache hit, value: %s", self._log_prefix, value) def _log_cache_miss(self) -> None: if self._log_prefix: @@ -86,9 +82,7 @@ class EventDrivenCacheCoordinator: async with self._lock: if self._query_in_progress and self._event is not None: if self._log_prefix: - verbose_proxy_logger.debug( - "%s Load in flight, waiting for signal", self._log_prefix - ) + verbose_proxy_logger.debug("%s Load in flight, waiting for signal", self._log_prefix) return self._event self._query_in_progress = True self._event = asyncio.Event() @@ -108,9 +102,7 @@ class EventDrivenCacheCoordinator: """Wait for loader to finish, then read from cache.""" await event.wait() if self._log_prefix: - verbose_proxy_logger.debug( - "%s Signal received, reading from cache", self._log_prefix - ) + verbose_proxy_logger.debug("%s Signal received, reading from cache", self._log_prefix) value: Optional[T] = await cache.async_get_cache(key=cache_key) if value is not None and self._log_prefix: verbose_proxy_logger.debug( @@ -119,9 +111,7 @@ class EventDrivenCacheCoordinator: value, ) elif value is None and self._log_prefix: - verbose_proxy_logger.debug( - "%s Signal received but cache still empty", self._log_prefix - ) + verbose_proxy_logger.debug("%s Signal received but cache still empty", self._log_prefix) return value async def _load_and_cache( @@ -165,9 +155,7 @@ class EventDrivenCacheCoordinator: self._query_in_progress = False if self._event is not None: if self._log_prefix: - verbose_proxy_logger.debug( - "%s Signaling all waiting requests", self._log_prefix - ) + verbose_proxy_logger.debug("%s Signaling all waiting requests", self._log_prefix) self._event.set() self._event = None diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py index 80a8d6281a1..25d33a0aa52 100644 --- a/litellm/proxy/common_utils/cache_pydantic_utils.py +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -56,9 +56,7 @@ class CacheCodec: # Already the right type: dump directly, skip re-validation copy. return value.model_dump(mode="json", exclude_none=True) if isinstance(value, (dict, BaseModel)): - return model_type.model_validate(value).model_dump( - mode="json", exclude_none=True - ) + return model_type.model_validate(value).model_dump(mode="json", exclude_none=True) return value if isinstance(value, BaseModel): return value.model_dump(mode="json", exclude_none=True) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index cd41bce97d6..573763d6627 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -51,9 +51,7 @@ def initialize_callbacks_on_proxy( ) from litellm.proxy.proxy_server import prisma_client - verbose_proxy_logger.debug( - f"{blue_color_code}initializing callbacks={value} on proxy{reset_color_code}" - ) + verbose_proxy_logger.debug(f"{blue_color_code}initializing callbacks={value} on proxy{reset_color_code}") if isinstance(value, list): imported_list: List[Any] = [] for callback in value: # ["presidio", ] @@ -62,59 +60,41 @@ def initialize_callbacks_on_proxy( CompressionInterceptionLogger, ) - compression_interception_obj = ( - CompressionInterceptionLogger.initialize_from_proxy_config( - litellm_settings=litellm_settings, - callback_specific_params=callback_specific_params, - ) + compression_interception_obj = CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, ) imported_list.append(compression_interception_obj) continue - if ( - isinstance(callback, str) - and callback == "code_interpreter_interception" - ): + if isinstance(callback, str) and callback == "code_interpreter_interception": from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, ) - code_interpreter_interception_obj = ( - CodeInterpreterInterceptionLogger.initialize_from_proxy_config( - litellm_settings=litellm_settings, - callback_specific_params=callback_specific_params, - ) + code_interpreter_interception_obj = CodeInterpreterInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, ) imported_list.append(code_interpreter_interception_obj) continue # check if callback is a custom logger compatible callback if isinstance(callback, str): - callback = LoggingCallbackManager._add_custom_callback_generic_api_str( - callback - ) - if ( - isinstance(callback, str) - and callback in litellm._known_custom_logger_compatible_callbacks - ): + callback = LoggingCallbackManager._add_custom_callback_generic_api_str(callback) + if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks: imported_list.append(callback) elif isinstance(callback, str) and callback == "presidio": from litellm.proxy.guardrails.guardrail_hooks.presidio import ( _OPTIONAL_PresidioPIIMasking, ) - presidio_logging_only: Optional[bool] = litellm_settings.get( - "presidio_logging_only", None - ) + presidio_logging_only: Optional[bool] = litellm_settings.get("presidio_logging_only", None) if presidio_logging_only is not None: - presidio_logging_only = bool( - presidio_logging_only - ) # validate boolean given + presidio_logging_only = bool(presidio_logging_only) # validate boolean given _presidio_params = {} - if "presidio" in callback_specific_params and isinstance( - callback_specific_params["presidio"], dict - ): + if "presidio" in callback_specific_params and isinstance(callback_specific_params["presidio"], dict): _presidio_params = callback_specific_params["presidio"] params: Dict[str, Any] = { @@ -130,15 +110,11 @@ def initialize_callbacks_on_proxy( ) except ImportError: raise Exception( - "MissingTrying to use Llama Guard" - + CommonProxyErrors.missing_enterprise_package.value + "MissingTrying to use Llama Guard" + CommonProxyErrors.missing_enterprise_package.value ) if premium_user is not True: - raise Exception( - "Trying to use Llama Guard" - + CommonProxyErrors.not_premium_user.value - ) + raise Exception("Trying to use Llama Guard" + CommonProxyErrors.not_premium_user.value) llama_guard_object = _ENTERPRISE_LlamaGuard() imported_list.append(llama_guard_object) @@ -149,15 +125,11 @@ def initialize_callbacks_on_proxy( ) except ImportError: raise Exception( - "Trying to use Secret Detection" - + CommonProxyErrors.missing_enterprise_package.value + "Trying to use Secret Detection" + CommonProxyErrors.missing_enterprise_package.value ) if premium_user is not True: - raise Exception( - "Trying to use secret hiding" - + CommonProxyErrors.not_premium_user.value - ) + raise Exception("Trying to use secret hiding" + CommonProxyErrors.not_premium_user.value) _secret_detection_object = _ENTERPRISE_SecretDetection() imported_list.append(_secret_detection_object) @@ -173,10 +145,7 @@ def initialize_callbacks_on_proxy( ) if premium_user is not True: - raise Exception( - "Trying to use OpenAI Moderations Check" - + CommonProxyErrors.not_premium_user.value - ) + raise Exception("Trying to use OpenAI Moderations Check" + CommonProxyErrors.not_premium_user.value) openai_moderations_object = _ENTERPRISE_OpenAI_Moderation() imported_list.append(openai_moderations_object) @@ -211,10 +180,7 @@ def initialize_callbacks_on_proxy( ) if premium_user is not True: - raise Exception( - "Trying to use Google Text Moderation" - + CommonProxyErrors.not_premium_user.value - ) + raise Exception("Trying to use Google Text Moderation" + CommonProxyErrors.not_premium_user.value) google_text_moderation_obj = _ENTERPRISE_GoogleTextModeration() imported_list.append(google_text_moderation_obj) @@ -224,16 +190,10 @@ def initialize_callbacks_on_proxy( _ENTERPRISE_LLMGuard, ) except ImportError: - raise Exception( - "Trying to use Llm Guard" - + CommonProxyErrors.missing_enterprise_package.value - ) + raise Exception("Trying to use Llm Guard" + CommonProxyErrors.missing_enterprise_package.value) if premium_user is not True: - raise Exception( - "Trying to use Llm Guard" - + CommonProxyErrors.not_premium_user.value - ) + raise Exception("Trying to use Llm Guard" + CommonProxyErrors.not_premium_user.value) llm_guard_moderation_obj = _ENTERPRISE_LLMGuard() imported_list.append(llm_guard_moderation_obj) @@ -244,19 +204,13 @@ def initialize_callbacks_on_proxy( ) except ImportError: raise Exception( - "Trying to use Blocked User List" - + CommonProxyErrors.missing_enterprise_package_docker.value + "Trying to use Blocked User List" + CommonProxyErrors.missing_enterprise_package_docker.value ) if premium_user is not True: - raise Exception( - "Trying to use ENTERPRISE BlockedUser" - + CommonProxyErrors.not_premium_user.value - ) + raise Exception("Trying to use ENTERPRISE BlockedUser" + CommonProxyErrors.not_premium_user.value) - blocked_user_list = _ENTERPRISE_BlockedUserList( - prisma_client=prisma_client - ) + blocked_user_list = _ENTERPRISE_BlockedUserList(prisma_client=prisma_client) imported_list.append(blocked_user_list) elif isinstance(callback, str) and callback == "banned_keywords": try: @@ -265,15 +219,11 @@ def initialize_callbacks_on_proxy( ) except ImportError: raise Exception( - "Trying to use Banned Keywords" - + CommonProxyErrors.missing_enterprise_package_docker.value + "Trying to use Banned Keywords" + CommonProxyErrors.missing_enterprise_package_docker.value ) if premium_user is not True: - raise Exception( - "Trying to use ENTERPRISE BannedKeyword" - + CommonProxyErrors.not_premium_user.value - ) + raise Exception("Trying to use ENTERPRISE BannedKeyword" + CommonProxyErrors.not_premium_user.value) banned_keywords_obj = _ENTERPRISE_BannedKeywords() imported_list.append(banned_keywords_obj) @@ -284,12 +234,8 @@ def initialize_callbacks_on_proxy( prompt_injection_params = None if "prompt_injection_params" in litellm_settings: - prompt_injection_params_in_config = litellm_settings[ - "prompt_injection_params" - ] - prompt_injection_params = LiteLLMPromptInjectionParams( - **prompt_injection_params_in_config - ) + prompt_injection_params_in_config = litellm_settings["prompt_injection_params"] + prompt_injection_params = LiteLLMPromptInjectionParams(**prompt_injection_params_in_config) prompt_injection_detection_obj = _OPTIONAL_PromptInjectionDetection( prompt_injection_params=prompt_injection_params, @@ -307,15 +253,9 @@ def initialize_callbacks_on_proxy( _PROXY_AzureContentSafety, ) - azure_content_safety_params = litellm_settings[ - "azure_content_safety_params" - ] + azure_content_safety_params = litellm_settings["azure_content_safety_params"] for k, v in azure_content_safety_params.items(): - if ( - v is not None - and isinstance(v, str) - and v.startswith("os.environ/") - ): + if v is not None and isinstance(v, str) and v.startswith("os.environ/"): azure_content_safety_params[k] = get_secret(v) azure_content_safety_obj = _PROXY_AzureContentSafety( @@ -327,11 +267,9 @@ def initialize_callbacks_on_proxy( WebSearchInterceptionLogger, ) - websearch_interception_obj = ( - WebSearchInterceptionLogger.initialize_from_proxy_config( - litellm_settings=litellm_settings, - callback_specific_params=callback_specific_params, - ) + websearch_interception_obj = WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, ) imported_list.append(websearch_interception_obj) elif isinstance(callback, str) and callback == "datadog_cost_management": @@ -374,16 +312,12 @@ def initialize_callbacks_on_proxy( config_file_path=config_file_path, ) ] - verbose_proxy_logger.debug( - f"{blue_color_code} Initialized Callbacks - {litellm.callbacks} {reset_color_code}" - ) + verbose_proxy_logger.debug(f"{blue_color_code} Initialized Callbacks - {litellm.callbacks} {reset_color_code}") def get_model_group_from_litellm_kwargs(kwargs: dict) -> Optional[str]: _litellm_params = kwargs.get("litellm_params", None) or {} - _metadata = ( - _litellm_params.get(get_metadata_variable_name_from_kwargs(kwargs)) or {} - ) + _metadata = _litellm_params.get(get_metadata_variable_name_from_kwargs(kwargs)) or {} _model_group = _metadata.get("model_group", None) if _model_group is not None: return _model_group @@ -412,25 +346,19 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str, model_group = get_model_group_from_request_data(data) # The h11 package considers "/" or ":" invalid and raise a LocalProtocolError - h11_model_group_name = ( - model_group.replace("/", "-").replace(":", "-") if model_group else None - ) + h11_model_group_name = model_group.replace("/", "-").replace(":", "-") if model_group else None # Remaining Requests remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_requests = _metadata.get(remaining_requests_variable_name, None) if remaining_requests: - headers[f"x-litellm-key-remaining-requests-{h11_model_group_name}"] = ( - remaining_requests - ) + headers[f"x-litellm-key-remaining-requests-{h11_model_group_name}"] = remaining_requests # Remaining Tokens remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" remaining_tokens = _metadata.get(remaining_tokens_variable_name, None) if remaining_tokens: - headers[f"x-litellm-key-remaining-tokens-{h11_model_group_name}"] = ( - remaining_tokens - ) + headers[f"x-litellm-key-remaining-tokens-{h11_model_group_name}"] = remaining_tokens return headers @@ -447,9 +375,7 @@ def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]: _metadata.update(litellm_metadata_bucket) headers = {} if "applied_guardrails" in _metadata: - headers["x-litellm-applied-guardrails"] = ",".join( - _metadata["applied_guardrails"] - ) + headers["x-litellm-applied-guardrails"] = ",".join(_metadata["applied_guardrails"]) if "applied_policies" in _metadata: headers["x-litellm-applied-policies"] = ",".join(_metadata["applied_policies"]) @@ -458,16 +384,12 @@ def get_logging_caching_headers(request_data: Dict) -> Optional[Dict]: sources = _metadata["policy_sources"] if isinstance(sources, dict) and sources: # Use ';' as delimiter — matched_via reasons may contain commas - headers["x-litellm-policy-sources"] = "; ".join( - f"{name}={reason}" for name, reason in sources.items() - ) + headers["x-litellm-policy-sources"] = "; ".join(f"{name}={reason}" for name, reason in sources.items()) if "semantic-similarity" in _metadata: headers["x-litellm-semantic-similarity"] = str(_metadata["semantic-similarity"]) - is_trusted_pillar_metadata = ( - _metadata.get(TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY) is True - ) + is_trusted_pillar_metadata = _metadata.get(TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY) is True pillar_headers = _metadata.get("pillar_response_headers") if is_trusted_pillar_metadata and isinstance(pillar_headers, dict): headers.update( @@ -570,9 +492,7 @@ def sanitize_openai_provider_metadata( return sanitized or None -def add_guardrail_to_applied_guardrails_header( - request_data: Dict, guardrail_name: Optional[str] -): +def add_guardrail_to_applied_guardrails_header(request_data: Dict, guardrail_name: Optional[str]): if guardrail_name is None: return _, _metadata = _get_or_create_proxy_metadata_bucket(request_data) @@ -583,9 +503,7 @@ def add_guardrail_to_applied_guardrails_header( _metadata["applied_guardrails"] = [guardrail_name] -def add_policy_to_applied_policies_header( - request_data: Dict, policy_name: Optional[str] -): +def add_policy_to_applied_policies_header(request_data: Dict, policy_name: Optional[str]): """ Add a policy name to the applied_policies list in request metadata. @@ -626,8 +544,8 @@ def add_guardrail_response_to_standard_logging_object( ): if litellm_logging_obj is None: return - standard_logging_object: Optional[StandardLoggingPayload] = ( - litellm_logging_obj.model_call_details.get("standard_logging_object") + standard_logging_object: Optional[StandardLoggingPayload] = litellm_logging_obj.model_call_details.get( + "standard_logging_object" ) if standard_logging_object is None: return @@ -640,9 +558,7 @@ def add_guardrail_response_to_standard_logging_object( return standard_logging_object -def process_callback( - _callback: str, callback_type: str, environment_variables: dict -) -> dict: +def process_callback(_callback: str, callback_type: str, environment_variables: dict) -> dict: """Process a single callback and return its data with environment variables""" env_vars = CustomLogger.get_callback_env_vars(_callback) @@ -680,9 +596,7 @@ def decrypt_callback_vars(metadata: Any) -> Any: return _transform_callback_vars(metadata, _decrypt_or_passthrough) -def _transform_callback_vars( - metadata: Any, transform: Callable[[str, Any], Any] -) -> Any: +def _transform_callback_vars(metadata: Any, transform: Callable[[str, Any], Any]) -> Any: if not isinstance(metadata, dict): return metadata out = copy.deepcopy(metadata) @@ -690,16 +604,10 @@ def _transform_callback_vars( if isinstance(logging_entries, list): for entry in logging_entries: if isinstance(entry, dict) and isinstance(entry.get("callback_vars"), dict): - entry["callback_vars"] = { - k: transform(k, v) for k, v in entry["callback_vars"].items() - } + entry["callback_vars"] = {k: transform(k, v) for k, v in entry["callback_vars"].items()} callback_settings = out.get("callback_settings") - if isinstance(callback_settings, dict) and isinstance( - callback_settings.get("callback_vars"), dict - ): - callback_settings["callback_vars"] = { - k: transform(k, v) for k, v in callback_settings["callback_vars"].items() - } + if isinstance(callback_settings, dict) and isinstance(callback_settings.get("callback_vars"), dict): + callback_settings["callback_vars"] = {k: transform(k, v) for k, v in callback_settings["callback_vars"].items()} return out @@ -739,7 +647,5 @@ def _decrypt_or_passthrough(key: str, value: Any) -> Any: # Legacy plaintext rows or non-credential fields — return as-is. return value inner = value[len(_CALLBACK_VAR_ENCRYPTED_PREFIX) :] - decrypted = decrypt_value_helper( - value=inner, key=key, exception_type="debug", return_original_value=False - ) + decrypted = decrypt_value_helper(value=inner, key=key, exception_type="debug", return_original_value=False) return decrypted if decrypted is not None else value diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index ce5d2539fb3..92cd8741b32 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -49,15 +49,11 @@ class CustomOpenAPISpec: except Exception as e: # FastAPI 0.120+ may fail schema generation for certain types (e.g., openai.Timeout) # Log the error and return None to skip schema generation for this model - verbose_proxy_logger.debug( - f"Failed to generate schema for {model_class}: {e}" - ) + verbose_proxy_logger.debug(f"Failed to generate schema for {model_class}: {e}") return None @staticmethod - def add_schema_to_components( - openapi_schema: Dict[str, Any], schema_name: str, schema_def: Dict[str, Any] - ) -> None: + def add_schema_to_components(openapi_schema: Dict[str, Any], schema_name: str, schema_def: Dict[str, Any]) -> None: """ Add a schema definition to the OpenAPI components/schemas section. @@ -73,14 +69,10 @@ class CustomOpenAPISpec: openapi_schema["components"]["schemas"] = {} # Add the schema - CustomOpenAPISpec._move_defs_to_components( - openapi_schema, {schema_name: schema_def} - ) + CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def}) @staticmethod - def add_request_body_to_paths( - openapi_schema: Dict[str, Any], paths: List[str], schema_ref: str - ) -> None: + def add_request_body_to_paths(openapi_schema: Dict[str, Any], paths: List[str], schema_ref: str) -> None: """ Add request body with expanded form fields for better Swagger UI display. This keeps the request body but expands it to show individual fields in the UI. @@ -91,28 +83,17 @@ class CustomOpenAPISpec: schema_ref: Reference to the schema component (e.g., "#/components/schemas/ModelName") """ for path in paths: - if ( - path in openapi_schema.get("paths", {}) - and "post" in openapi_schema["paths"][path] - ): + if path in openapi_schema.get("paths", {}) and "post" in openapi_schema["paths"][path]: # Get the actual schema to extract ALL field definitions - schema_name = schema_ref.split("/")[ - -1 - ] # Extract "ProxyChatCompletionRequest" from the ref - actual_schema = ( - openapi_schema.get("components", {}) - .get("schemas", {}) - .get(schema_name, {}) - ) + schema_name = schema_ref.split("/")[-1] # Extract "ProxyChatCompletionRequest" from the ref + actual_schema = openapi_schema.get("components", {}).get("schemas", {}).get(schema_name, {}) schema_properties = actual_schema.get("properties", {}) required_fields = actual_schema.get("required", []) # Extract $defs and add them to components/schemas # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI if "$defs" in actual_schema: - CustomOpenAPISpec._move_defs_to_components( - openapi_schema, actual_schema["$defs"] - ) + CustomOpenAPISpec._move_defs_to_components(openapi_schema, actual_schema["$defs"]) # Create an expanded inline schema instead of just a $ref # This makes Swagger UI show all individual fields in the request body editor @@ -124,20 +105,14 @@ class CustomOpenAPISpec: # Add all properties with their full definitions for field_name, field_def in schema_properties.items(): - expanded_field = CustomOpenAPISpec._expand_field_definition( - field_def - ) + expanded_field = CustomOpenAPISpec._expand_field_definition(field_def) # Rewrite $defs references to use components/schemas instead - expanded_field = CustomOpenAPISpec._rewrite_defs_refs( - expanded_field - ) + expanded_field = CustomOpenAPISpec._rewrite_defs_refs(expanded_field) # Add a simple example for the messages field if field_name == "messages": - expanded_field["example"] = [ - {"role": "user", "content": "Hello, how are you?"} - ] + expanded_field["example"] = [{"role": "user", "content": "Hello, how are you?"}] expanded_schema["properties"][field_name] = expanded_field @@ -149,21 +124,13 @@ class CustomOpenAPISpec: # Keep any existing parameters (like path parameters) but remove conflicting query params if "parameters" in openapi_schema["paths"][path]["post"]: - existing_params = openapi_schema["paths"][path]["post"][ - "parameters" - ] + existing_params = openapi_schema["paths"][path]["post"]["parameters"] # Only keep path parameters, remove query params that conflict with request body - filtered_params = [ - param for param in existing_params if param.get("in") == "path" - ] - openapi_schema["paths"][path]["post"]["parameters"] = ( - filtered_params - ) + filtered_params = [param for param in existing_params if param.get("in") == "path"] + openapi_schema["paths"][path]["post"]["parameters"] = filtered_params @staticmethod - def _move_defs_to_components( - openapi_schema: Dict[str, Any], defs: Dict[str, Any] - ) -> None: + def _move_defs_to_components(openapi_schema: Dict[str, Any], defs: Dict[str, Any]) -> None: """ Move $defs from Pydantic v2 schema to OpenAPI components/schemas. This makes the definitions resolvable in Swagger/OpenAPI viewers. @@ -189,9 +156,7 @@ class CustomOpenAPISpec: # If this definition also has $defs, process them recursively if "$defs" in def_schema: - CustomOpenAPISpec._move_defs_to_components( - openapi_schema, def_schema["$defs"] - ) + CustomOpenAPISpec._move_defs_to_components(openapi_schema, def_schema["$defs"]) @staticmethod def _rewrite_defs_refs(schema: Any) -> Any: @@ -208,11 +173,7 @@ class CustomOpenAPISpec: if isinstance(schema, dict): result = {} for key, value in schema.items(): - if ( - key == "$ref" - and isinstance(value, str) - and value.startswith("#/$defs/") - ): + if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"): # Rewrite the reference to use components/schemas def_name = value.replace("#/$defs/", "") result[key] = f"#/components/schemas/{def_name}" @@ -299,26 +260,20 @@ class CustomOpenAPISpec: # Only proceed if we successfully got the schema if request_schema is not None: # Add schema to components - CustomOpenAPISpec.add_schema_to_components( - openapi_schema, schema_name, request_schema - ) + CustomOpenAPISpec.add_schema_to_components(openapi_schema, schema_name, request_schema) # Add request body to specified endpoints CustomOpenAPISpec.add_request_body_to_paths( openapi_schema, paths, f"#/components/schemas/{schema_name}" ) - verbose_proxy_logger.debug( - f"Successfully added {schema_name} schema to OpenAPI spec" - ) + verbose_proxy_logger.debug(f"Successfully added {schema_name} schema to OpenAPI spec") else: verbose_proxy_logger.debug(f"Could not get schema for {schema_name}") except Exception as e: # If schema addition fails, continue without it - verbose_proxy_logger.debug( - f"Failed to add {operation_name} request schema: {str(e)}" - ) + verbose_proxy_logger.debug(f"Failed to add {operation_name} request schema: {str(e)}") return openapi_schema @@ -347,9 +302,7 @@ class CustomOpenAPISpec: operation_name="chat completion", ) except ImportError as e: - verbose_proxy_logger.debug( - f"Failed to import ProxyChatCompletionRequest: {str(e)}" - ) + verbose_proxy_logger.debug(f"Failed to import ProxyChatCompletionRequest: {str(e)}") return openapi_schema @staticmethod @@ -403,9 +356,7 @@ class CustomOpenAPISpec: operation_name="responses API", ) except ImportError as e: - verbose_proxy_logger.debug( - f"Failed to import ResponsesAPIRequestParams: {str(e)}" - ) + verbose_proxy_logger.debug(f"Failed to import ResponsesAPIRequestParams: {str(e)}") return openapi_schema @staticmethod @@ -422,16 +373,12 @@ class CustomOpenAPISpec: OpenAPI schema with added request body schemas """ # Add chat completion request schema - openapi_schema = CustomOpenAPISpec.add_chat_completion_request_schema( - openapi_schema - ) + openapi_schema = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema) # Add embedding request schema openapi_schema = CustomOpenAPISpec.add_embedding_request_schema(openapi_schema) # Add responses API request schema - openapi_schema = CustomOpenAPISpec.add_responses_api_request_schema( - openapi_schema - ) + openapi_schema = CustomOpenAPISpec.add_responses_api_request_schema(openapi_schema) return openapi_schema diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 4cc62e1adbd..aed3e1228f1 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -35,9 +35,7 @@ def configure_gc_thresholds(): f"GC threshold not set: {gc_threshold_env}. Expected format: 'gen0,gen1,gen2'" ) except ValueError as e: - verbose_proxy_logger.warning( - f"Failed to parse GC threshold: {gc_threshold_env}. Error: {e}" - ) + verbose_proxy_logger.warning(f"Failed to parse GC threshold: {gc_threshold_env}. Error: {e}") # Log current thresholds current_thresholds = gc.get_threshold() @@ -75,11 +73,7 @@ async def get_active_tasks_stats(): break coro = task.get_coro() # Derive a human‐readable name from the coroutine: - name = ( - getattr(coro, "__qualname__", None) - or getattr(coro, "__name__", None) - or repr(coro) - ) + name = getattr(coro, "__qualname__", None) or getattr(coro, "__name__", None) or repr(coro) counter[name] += 1 return { @@ -100,9 +94,7 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true": print("\n\nLeaking objects") # noqa: T201 objgraph.show_most_common_types(objects=roots) except ImportError: - raise ImportError( - "objgraph not found. Please install objgraph to use this feature." - ) + raise ImportError("objgraph not found. Please install objgraph to use this feature.") tracemalloc.start(10) @@ -146,13 +138,13 @@ async def memory_usage_in_mem_cache( if llm_router is None: num_items_in_llm_router_cache = 0 else: - num_items_in_llm_router_cache = len( - llm_router.cache.in_memory_cache.cache_dict - ) + len(llm_router.cache.in_memory_cache.ttl_dict) + num_items_in_llm_router_cache = len(llm_router.cache.in_memory_cache.cache_dict) + len( + llm_router.cache.in_memory_cache.ttl_dict + ) - num_items_in_user_api_key_cache = len( - user_api_key_cache.in_memory_cache.cache_dict - ) + len(user_api_key_cache.in_memory_cache.ttl_dict) + num_items_in_user_api_key_cache = len(user_api_key_cache.in_memory_cache.cache_dict) + len( + user_api_key_cache.in_memory_cache.ttl_dict + ) num_items_in_proxy_logging_obj_cache = len( proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict @@ -252,9 +244,7 @@ async def get_memory_summary( health_status = "healthy" except ImportError: - process_memory["error"] = ( - "Install psutil for memory monitoring: pip install psutil" - ) + process_memory["error"] = "Install psutil for memory monitoring: pip install psutil" except Exception as e: process_memory["error"] = str(e) @@ -283,9 +273,7 @@ async def get_memory_summary( } # Proxy logging cache - logging_cache_items = len( - proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict - ) + logging_cache_items = len(proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict) total_cache_items += logging_cache_items caches["usage_tracking"] = { "count": logging_cache_items, @@ -308,9 +296,7 @@ async def get_memory_summary( # Add warning if garbage collection issues detected if uncollectable > 0: - gc_info["warning"] = ( - f"{uncollectable} uncollectable objects (possible memory leak)" - ) + gc_info["warning"] = f"{uncollectable} uncollectable objects (possible memory leak)" return { "worker_pid": os.getpid(), @@ -384,9 +370,7 @@ def _get_uncollectable_objects_info() -> Dict[str, Any]: } -def _get_cache_memory_stats( - user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache -) -> Dict[str, Any]: +def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) -> Dict[str, Any]: """Calculate memory usage for all caches.""" cache_stats: Dict[str, Any] = {} try: @@ -397,42 +381,28 @@ def _get_cache_memory_stats( "num_items": len(user_api_key_cache.in_memory_cache.cache_dict), "cache_dict_size_bytes": user_cache_size, "ttl_dict_size_bytes": user_ttl_size, - "total_size_mb": round( - (user_cache_size + user_ttl_size) / (1024 * 1024), 2 - ), + "total_size_mb": round((user_cache_size + user_ttl_size) / (1024 * 1024), 2), } # Router cache if llm_router is not None: - router_cache_size = sys.getsizeof( - llm_router.cache.in_memory_cache.cache_dict - ) + router_cache_size = sys.getsizeof(llm_router.cache.in_memory_cache.cache_dict) router_ttl_size = sys.getsizeof(llm_router.cache.in_memory_cache.ttl_dict) cache_stats["llm_router_cache"] = { "num_items": len(llm_router.cache.in_memory_cache.cache_dict), "cache_dict_size_bytes": router_cache_size, "ttl_dict_size_bytes": router_ttl_size, - "total_size_mb": round( - (router_cache_size + router_ttl_size) / (1024 * 1024), 2 - ), + "total_size_mb": round((router_cache_size + router_ttl_size) / (1024 * 1024), 2), } # Proxy logging cache - logging_cache_size = sys.getsizeof( - proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict - ) - logging_ttl_size = sys.getsizeof( - proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.ttl_dict - ) + logging_cache_size = sys.getsizeof(proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict) + logging_ttl_size = sys.getsizeof(proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.ttl_dict) cache_stats["proxy_logging_cache"] = { - "num_items": len( - proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict - ), + "num_items": len(proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict), "cache_dict_size_bytes": logging_cache_size, "ttl_dict_size_bytes": logging_ttl_size, - "total_size_mb": round( - (logging_cache_size + logging_ttl_size) / (1024 * 1024), 2 - ), + "total_size_mb": round((logging_cache_size + logging_ttl_size) / (1024 * 1024), 2), } # Redis cache info @@ -443,22 +413,15 @@ def _get_cache_memory_stats( } # Try to get Redis connection pool info if available try: - if ( - hasattr(redis_usage_cache, "redis_client") - and redis_usage_cache.redis_client - ): + if hasattr(redis_usage_cache, "redis_client") and redis_usage_cache.redis_client: if hasattr(redis_usage_cache.redis_client, "connection_pool"): pool_info = redis_usage_cache.redis_client.connection_pool # type: ignore cache_stats["redis_usage_cache"]["connection_pool"] = { "max_connections": ( - pool_info.max_connections - if hasattr(pool_info, "max_connections") - else None + pool_info.max_connections if hasattr(pool_info, "max_connections") else None ), "connection_class": ( - pool_info.connection_class.__name__ - if hasattr(pool_info, "connection_class") - else None + pool_info.connection_class.__name__ if hasattr(pool_info, "connection_class") else None ), } except Exception as e: @@ -506,10 +469,7 @@ def _get_router_memory_stats(llm_router) -> Dict[str, Any]: } # Deployment latency map - if ( - hasattr(llm_router, "deployment_latency_map") - and llm_router.deployment_latency_map - ): + if hasattr(llm_router, "deployment_latency_map") and llm_router.deployment_latency_map: latency_map_size = sys.getsizeof(llm_router.deployment_latency_map) litellm_router_memory["deployment_latency_map"] = { "num_tracked_deployments": len(llm_router.deployment_latency_map), @@ -542,9 +502,7 @@ def _get_router_memory_stats(llm_router) -> Dict[str, Any]: return litellm_router_memory -def _get_process_memory_info( - worker_pid: int, include_process_info: bool -) -> Optional[Dict[str, Any]]: +def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> Optional[Dict[str, Any]]: """Get process-level memory information using psutil.""" if not include_process_info: return None @@ -574,11 +532,7 @@ def _get_process_memory_info( "description": "Percentage of total system RAM being used", }, "open_file_handles": { - "count": ( - process.num_fds() - if hasattr(process, "num_fds") - else "N/A (Windows)" - ), + "count": (process.num_fds() if hasattr(process, "num_fds") else "N/A (Windows)"), "description": "Number of open file descriptors/handles", }, "threads": { @@ -636,9 +590,7 @@ async def get_memory_details( gc_stats = _get_gc_statistics() total_objects, top_object_types = _get_object_type_counts(top_n) uncollectable_info = _get_uncollectable_objects_info() - cache_stats = _get_cache_memory_stats( - user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache - ) + cache_stats = _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) litellm_router_memory = _get_router_memory_stats(llm_router) process_info = _get_process_memory_info(worker_pid, include_process_info) @@ -697,14 +649,11 @@ async def configure_gc_thresholds_endpoint( try: gc.set_threshold(generation_0, generation_1, generation_2) verbose_proxy_logger.info( - f"GC thresholds updated from {old_thresholds} to " - f"({generation_0}, {generation_1}, {generation_2})" + f"GC thresholds updated from {old_thresholds} to ({generation_0}, {generation_1}, {generation_2})" ) except Exception as e: verbose_proxy_logger.error(f"Failed to set GC thresholds: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to set GC thresholds: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Failed to set GC thresholds: {str(e)}") # Get current object count to show immediate impact current_count = gc.get_count()[0] @@ -791,9 +740,7 @@ def init_verbose_loggers(): # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS verbose_logger.setLevel(level=logging.INFO) # sets package logs to info - verbose_router_logger.setLevel( - level=logging.INFO - ) # set router logs to info + verbose_router_logger.setLevel(level=logging.INFO) # set router logs to info verbose_proxy_logger.setLevel(level=logging.INFO) # set proxy logs to info if detailed_debug is True: import logging @@ -805,12 +752,8 @@ def init_verbose_loggers(): ) verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug - verbose_router_logger.setLevel( - level=logging.DEBUG - ) # set router logs to debug - verbose_proxy_logger.setLevel( - level=logging.DEBUG - ) # set proxy logs to debug + verbose_router_logger.setLevel(level=logging.DEBUG) # set router logs to debug + verbose_proxy_logger.setLevel(level=logging.DEBUG) # set proxy logs to debug elif debug is False and detailed_debug is False: # users can control proxy debugging using env variable = 'LITELLM_LOG' litellm_log_setting = os.environ.get("LITELLM_LOG", "") @@ -825,12 +768,8 @@ def init_verbose_loggers(): # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS - verbose_router_logger.setLevel( - level=logging.INFO - ) # set router logs to info - verbose_proxy_logger.setLevel( - level=logging.INFO - ) # set proxy logs to info + verbose_router_logger.setLevel(level=logging.INFO) # set router logs to info + verbose_proxy_logger.setLevel(level=logging.INFO) # set proxy logs to info elif litellm_log_setting.upper() == "DEBUG": import logging @@ -839,12 +778,8 @@ def init_verbose_loggers(): verbose_router_logger, ) - verbose_router_logger.setLevel( - level=logging.DEBUG - ) # set router logs to info - verbose_proxy_logger.setLevel( - level=logging.DEBUG - ) # set proxy logs to debug + verbose_router_logger.setLevel(level=logging.DEBUG) # set router logs to info + verbose_proxy_logger.setLevel(level=logging.DEBUG) # set proxy logs to debug except Exception as e: import logging diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index a5da5798f47..6be56de1260 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -65,9 +65,7 @@ def decrypt_value_helper( verbose_proxy_logger.debug(error_message) return value if return_original_value else None - verbose_proxy_logger.debug( - f"Unable to decrypt value={value} for key: {key}, returning None" - ) + verbose_proxy_logger.debug(f"Unable to decrypt value={value} for key: {key}, returning None") if return_original_value: return value else: diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index 4f3e26ab5fb..85ec6b839a4 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -92,26 +92,19 @@ class ExpiredUISessionKeyCleanupManager: tokens=tokens, response=response, ) - verbose_proxy_logger.info( - "Deleted %s expired UI session key(s)", deleted_count - ) + verbose_proxy_logger.info("Deleted %s expired UI session key(s)", deleted_count) return deleted_count except Exception as e: if getattr(e, "status_code", None) == 404: verbose_proxy_logger.debug( - "Expired UI session key cleanup skipped because selected keys " - "were already deleted: %s", + "Expired UI session key cleanup skipped because selected keys were already deleted: %s", e, ) return 0 verbose_proxy_logger.error(f"Expired UI session key cleanup failed: {e}") return 0 finally: - if ( - lock_acquired - and self.pod_lock_manager - and self.pod_lock_manager.redis_cache - ): + if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: await self.pod_lock_manager.release_lock( cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, ) diff --git a/litellm/proxy/common_utils/get_routes.py b/litellm/proxy/common_utils/get_routes.py index 743c3b6e9d9..64a1332ccb0 100644 --- a/litellm/proxy/common_utils/get_routes.py +++ b/litellm/proxy/common_utils/get_routes.py @@ -23,9 +23,7 @@ class GetRoutes: "path": getattr(route, "path", None), "methods": getattr(route, "methods", None), "name": getattr(route, "name", None), - "endpoint": ( - endpoint_route.__name__ if getattr(route, "endpoint", None) else None - ), + "endpoint": (endpoint_route.__name__ if getattr(route, "endpoint", None) else None), } routes.append(route_info) return routes @@ -43,9 +41,7 @@ class GetRoutes: if sub_app and hasattr(sub_app, "routes"): for sub_route in sub_app.routes: # Get endpoint - either from endpoint attribute or app attribute - endpoint_func = getattr(sub_route, "endpoint", None) or getattr( - sub_route, "app", None - ) + endpoint_func = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) if endpoint_func is not None: sub_route_path = getattr(sub_route, "path", "") @@ -69,14 +65,10 @@ class GetRoutes: try: if hasattr(endpoint_function, "__name__"): return getattr(endpoint_function, "__name__") - elif hasattr(endpoint_function, "__class__") and hasattr( - endpoint_function.__class__, "__name__" - ): + elif hasattr(endpoint_function, "__class__") and hasattr(endpoint_function.__class__, "__name__"): return getattr(endpoint_function.__class__, "__name__") else: return None except Exception: - verbose_logger.exception( - f"Error getting endpoint name for route: {endpoint_function}" - ) + verbose_logger.exception(f"Error getting endpoint name for route: {endpoint_function}") return None diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index efb4dbc60b4..719c39642c6 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -13,9 +13,7 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.types.router import Deployment -_FORM_CONTENT_TYPES: frozenset[str] = frozenset( - {"application/x-www-form-urlencoded", "multipart/form-data"} -) +_FORM_CONTENT_TYPES: frozenset[str] = frozenset({"application/x-www-form-urlencoded", "multipart/form-data"}) def _normalize_media_type(content_type: str) -> str: @@ -56,9 +54,7 @@ async def _read_request_body(request: Optional[Request]) -> Dict: return {} # Check if we already read and parsed the body - _cached_request_body: Optional[dict] = _safe_get_request_parsed_body( - request=request - ) + _cached_request_body: Optional[dict] = _safe_get_request_parsed_body(request=request) if _cached_request_body is not None: return _cached_request_body @@ -99,13 +95,9 @@ async def _read_request_body(request: Optional[Request]) -> Dict: # The surrogate-repair fallback below runs two full-body re.sub # passes, which block the event loop on multi-MB malformed bodies. # Above the configured size, skip the repair and raise the 400 now. - repair_limit_bytes = ( - MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB * 1024 * 1024 - ) + repair_limit_bytes = MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB * 1024 * 1024 if repair_limit_bytes > 0 and len(body) > repair_limit_bytes: - verbose_proxy_logger.error( - f"Invalid JSON payload received: {str(e)}" - ) + verbose_proxy_logger.error(f"Invalid JSON payload received: {str(e)}") raise ProxyException( message=f"Invalid JSON payload: {str(e)}", type="invalid_request_error", @@ -119,21 +111,15 @@ async def _read_request_body(request: Optional[Request]) -> Dict: # Replace invalid surrogate pairs # This regex finds incomplete surrogate pairs - body_str = re.sub( - r"[\uD800-\uDBFF](?![\uDC00-\uDFFF])", "", body_str - ) + body_str = re.sub(r"[\uD800-\uDBFF](?![\uDC00-\uDFFF])", "", body_str) # This regex finds low surrogates without high surrogates - body_str = re.sub( - r"(? Dict: raise except Exception as e: # Catch unexpected errors to avoid crashes - verbose_proxy_logger.exception( - "Unexpected error reading request body - {}".format(e) - ) + verbose_proxy_logger.exception("Unexpected error reading request body - {}".format(e)) return {} def _safe_get_request_parsed_body(request: Optional[Request]) -> Optional[dict]: if request is None: return None - if ( - hasattr(request, "scope") - and "parsed_body" in request.scope - and isinstance(request.scope["parsed_body"], tuple) - ): + if hasattr(request, "scope") and "parsed_body" in request.scope and isinstance(request.scope["parsed_body"], tuple): accepted_keys, parsed_body = request.scope["parsed_body"] return {key: parsed_body[key] for key in accepted_keys} return None @@ -178,9 +158,7 @@ def _safe_get_request_query_params(request: Optional[Request]) -> Dict: return dict(request.query_params) return {} except Exception as e: - verbose_proxy_logger.debug( - "Unexpected error reading request query params - {}".format(e) - ) + verbose_proxy_logger.debug("Unexpected error reading request query params - {}".format(e)) return {} @@ -193,9 +171,7 @@ def _safe_set_request_parsed_body( return request.scope["parsed_body"] = (tuple(parsed_body.keys()), parsed_body) except Exception as e: - verbose_proxy_logger.debug( - "Unexpected error setting request parsed body - {}".format(e) - ) + verbose_proxy_logger.debug("Unexpected error setting request parsed body - {}".format(e)) def _safe_get_request_headers(request: Optional[Request]) -> dict: @@ -213,15 +189,11 @@ def _safe_get_request_headers(request: Optional[Request]) -> dict: if isinstance(cached, dict): return cached if cached is not None: - verbose_proxy_logger.debug( - "Unexpected cached request headers type - {}".format(type(cached)) - ) + verbose_proxy_logger.debug("Unexpected cached request headers type - {}".format(type(cached))) try: headers = dict(request.headers) except Exception as e: - verbose_proxy_logger.debug( - "Unexpected error reading request headers - {}".format(e) - ) + verbose_proxy_logger.debug("Unexpected error reading request headers - {}".format(e)) headers = {} try: if state is not None: @@ -258,10 +230,8 @@ def check_file_size_under_limit( if llm_router is not None and request_data["model"] in router_model_names: try: - deployment: Optional[Deployment] = ( - llm_router.get_deployment_by_model_group_name( - model_group_name=request_data["model"] - ) + deployment: Optional[Deployment] = llm_router.get_deployment_by_model_group_name( + model_group_name=request_data["model"] ) if ( deployment @@ -270,9 +240,7 @@ def check_file_size_under_limit( ): max_file_size_mb = deployment.litellm_params.max_file_size_mb except Exception as e: - verbose_proxy_logger.error( - "Got error when checking file size: %s", (str(e)) - ) + verbose_proxy_logger.error("Got error when checking file size: %s", (str(e))) if max_file_size_mb is not None: verbose_proxy_logger.debug( @@ -377,9 +345,7 @@ async def get_request_body(request: Request) -> Dict[str, Any]: return {} -def extract_nested_form_metadata( - form_data: Dict[str, Any], prefix: str = "litellm_metadata[" -) -> Dict[str, Any]: +def extract_nested_form_metadata(form_data: Dict[str, Any], prefix: str = "litellm_metadata[") -> Dict[str, Any]: """ Extract nested metadata from form data with bracket notation. @@ -426,9 +392,7 @@ def extract_nested_form_metadata( # Skip UploadFile objects - they should not be in metadata if isinstance(value, UploadFile): - verbose_proxy_logger.warning( - f"Skipping UploadFile in metadata extraction for key: {key}" - ) + verbose_proxy_logger.warning(f"Skipping UploadFile in metadata extraction for key: {key}") continue # Extract the nested path from bracket notation @@ -441,9 +405,7 @@ def extract_nested_form_metadata( parts = path_string.split("][") if not parts or not parts[0]: - verbose_proxy_logger.warning( - f"Invalid metadata key format (empty path): {key}" - ) + verbose_proxy_logger.warning(f"Invalid metadata key format (empty path): {key}") continue # Navigate/create nested dictionary structure @@ -460,9 +422,7 @@ def extract_nested_form_metadata( if isinstance(current, dict): current[parts[-1]] = value else: - verbose_proxy_logger.warning( - f"Cannot set value - parent is not a dict for key: {key}" - ) + verbose_proxy_logger.warning(f"Cannot set value - parent is not a dict for key: {key}") except Exception as e: verbose_proxy_logger.error(f"Error parsing metadata key '{key}': {str(e)}") @@ -585,6 +545,4 @@ def _add_vector_store_id_from_path(request_data: dict, request: Request) -> None f"populate_request_with_path_params: Updated request_data with vector_store_ids={request_data.get('vector_store_ids')}" ) else: - verbose_proxy_logger.debug( - f"populate_request_with_path_params: No vector_store_id present in path={path}" - ) + verbose_proxy_logger.debug(f"populate_request_with_path_params: No vector_store_id present in path={path}") diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index d622f612494..51f86e514ca 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -85,37 +85,23 @@ class KeyRotationManager: verbose_proxy_logger.debug("No keys are due for rotation at this time") return - verbose_proxy_logger.info( - f"Found {len(keys_to_rotate)} keys due for rotation" - ) + verbose_proxy_logger.info(f"Found {len(keys_to_rotate)} keys due for rotation") # Rotate each key for key in keys_to_rotate: try: await self._rotate_key(key) - key_identifier = key.key_name or ( - key.token[:8] + "..." if key.token else "unknown" - ) - verbose_proxy_logger.info( - f"Successfully rotated key: {key_identifier}" - ) + key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") + verbose_proxy_logger.info(f"Successfully rotated key: {key_identifier}") except Exception as e: - key_identifier = key.key_name or ( - key.token[:8] + "..." if key.token else "unknown" - ) - verbose_proxy_logger.error( - f"Failed to rotate key {key_identifier}: {e}" - ) + key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") + verbose_proxy_logger.error(f"Failed to rotate key {key_identifier}: {e}") except Exception as e: verbose_proxy_logger.error(f"Key rotation process failed: {e}") finally: # Only release the lock if it was actually acquired - if ( - lock_acquired - and self.pod_lock_manager - and self.pod_lock_manager.redis_cache - ): + if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: await self.pod_lock_manager.release_lock( cronjob_id=KEY_ROTATION_JOB_NAME, ) @@ -130,18 +116,12 @@ class KeyRotationManager: """ now = datetime.now(timezone.utc) - keys_with_rotation = await VerificationTokenRepository( - self.prisma_client - ).table.find_many( + keys_with_rotation = await VerificationTokenRepository(self.prisma_client).table.find_many( where={ "auto_rotate": True, # Only keys marked for auto rotation "OR": [ - { - "key_rotation_at": None - }, # Keys that need initial rotation time setup - { - "key_rotation_at": {"lte": now} - }, # Keys where rotation time has passed + {"key_rotation_at": None}, # Keys that need initial rotation time setup + {"key_rotation_at": {"lte": now}}, # Keys where rotation time has passed ], } ) @@ -154,17 +134,13 @@ class KeyRotationManager: """ try: now = datetime.now(timezone.utc) - result = await DeprecatedVerificationTokenRepository( - self.prisma_client - ).table.delete_many(where={"revoke_at": {"lt": now}}) - if result > 0: - verbose_proxy_logger.debug( - "Cleaned up %s expired deprecated key(s)", result - ) - except Exception as e: - verbose_proxy_logger.debug( - "Deprecated key cleanup skipped (table may not exist): %s", e + result = await DeprecatedVerificationTokenRepository(self.prisma_client).table.delete_many( + where={"revoke_at": {"lt": now}} ) + if result > 0: + verbose_proxy_logger.debug("Cleaned up %s expired deprecated key(s)", result) + except Exception as e: + verbose_proxy_logger.debug("Deprecated key cleanup skipped (table may not exist): %s", e) def _should_rotate_key(self, key: LiteLLM_VerificationToken, now: datetime) -> bool: """ @@ -204,11 +180,7 @@ class KeyRotationManager: ) # Update the NEW key with rotation info (regenerate_key_fn creates a new token) - if ( - isinstance(response, GenerateKeyResponse) - and response.token_id - and key.rotation_interval - ): + if isinstance(response, GenerateKeyResponse) and response.token_id and key.rotation_interval: # Calculate next rotation time using helper function now = datetime.now(timezone.utc) next_rotation_time = _calculate_key_rotation_time(key.rotation_interval) diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 3c7329c2c0c..71e8238db27 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -20,9 +20,7 @@ def get_file_contents_from_s3(bucket_name, object_key): aws_secret_access_key=credentials.secret_key, aws_session_token=credentials.token, # Optional, if using temporary credentials ) - verbose_proxy_logger.debug( - f"Retrieving {object_key} from S3 bucket: {bucket_name}" - ) + verbose_proxy_logger.debug(f"Retrieving {object_key} from S3 bucket: {bucket_name}") response = s3_client.get_object(Bucket=bucket_name, Key=object_key) verbose_proxy_logger.debug(f"Response: {response}") @@ -96,9 +94,7 @@ def download_python_file_from_s3( aws_session_token=credentials.token, ) - verbose_proxy_logger.debug( - f"Downloading Python file {object_key} from S3 bucket: {bucket_name}" - ) + verbose_proxy_logger.debug(f"Downloading Python file {object_key} from S3 bucket: {bucket_name}") response = s3_client.get_object(Bucket=bucket_name, Key=object_key) # Read the file contents @@ -112,9 +108,7 @@ def download_python_file_from_s3( with open(local_file_path, "w") as f: f.write(file_contents) - verbose_proxy_logger.debug( - f"Python file downloaded successfully to {local_file_path}" - ) + verbose_proxy_logger.debug(f"Python file downloaded successfully to {local_file_path}") return True except ImportError as e: @@ -161,15 +155,11 @@ async def download_python_file_from_gcs( with open(local_file_path, "w") as f: f.write(file_contents) - verbose_proxy_logger.debug( - f"Python file downloaded successfully to {local_file_path}" - ) + verbose_proxy_logger.debug(f"Python file downloaded successfully to {local_file_path}") return True except Exception as e: - verbose_proxy_logger.exception( - f"Error downloading Python file from GCS: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error downloading Python file from GCS: {str(e)}") return False diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 3a70377037d..2835ecfc751 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -61,19 +61,14 @@ class TeamModelNameTranslator: Empty when disabled via the legacy flag, the router is absent, or the router model list is malformed. """ - if llm_router is None or not TeamModelNameTranslator._is_enabled( - general_settings - ): + if llm_router is None or not TeamModelNameTranslator._is_enabled(general_settings): return {} router_model_list = llm_router.get_model_list() if not isinstance(router_model_list, list): return {} return dict( pair - for pair in ( - TeamModelNameTranslator._internal_public_pair(model) - for model in router_model_list - ) + for pair in (TeamModelNameTranslator._internal_public_pair(model) for model in router_model_list) if pair is not None ) @@ -111,16 +106,10 @@ class TeamModelNameTranslator: internal key. Both ids are identical for unmapped names (globals, access-group keys). """ - internal_to_public = TeamModelNameTranslator.build_internal_to_public_map( - llm_router, general_settings - ) + internal_to_public = TeamModelNameTranslator.build_internal_to_public_map(llm_router, general_settings) if not internal_to_public: return [(name, name) for name in model_names] - return list( - TeamModelNameTranslator._response_to_lookup_map( - model_names, internal_to_public - ).items() - ) + return list(TeamModelNameTranslator._response_to_lookup_map(model_names, internal_to_public).items()) @staticmethod def translate_listing( @@ -133,10 +122,7 @@ class TeamModelNameTranslator: while preserving order; unmapped names pass through. """ return [ - entry[0] - for entry in TeamModelNameTranslator.listing_entries( - model_names, llm_router, general_settings - ) + entry[0] for entry in TeamModelNameTranslator.listing_entries(model_names, llm_router, general_settings) ] @staticmethod @@ -157,11 +143,9 @@ class TeamModelNameTranslator: unchanged when it is not an accessible public team name (already-internal names and globals pass through). """ - internal_to_public = TeamModelNameTranslator.build_internal_to_public_map( - llm_router, general_settings - ) + internal_to_public = TeamModelNameTranslator.build_internal_to_public_map(llm_router, general_settings) if not internal_to_public: return model_id - return TeamModelNameTranslator._response_to_lookup_map( - available_models, internal_to_public - ).get(model_id, model_id) + return TeamModelNameTranslator._response_to_lookup_map(available_models, internal_to_public).get( + model_id, model_id + ) diff --git a/litellm/proxy/common_utils/openapi_schema_compat.py b/litellm/proxy/common_utils/openapi_schema_compat.py index 83b1751c947..cd96bd2d2a4 100644 --- a/litellm/proxy/common_utils/openapi_schema_compat.py +++ b/litellm/proxy/common_utils/openapi_schema_compat.py @@ -49,9 +49,7 @@ def get_openapi_schema_with_compat( obj_module = getattr(obj, "__module__", "") if (obj_module == "openai" and "Timeout" in obj_str) or ( - hasattr(obj, "__name__") - and obj.__name__ == "Timeout" - and obj_module == "openai" + hasattr(obj, "__name__") and obj.__name__ == "Timeout" and obj_module == "openai" ): # Return a simple string schema for Timeout types return core_schema.str_schema() @@ -76,17 +74,13 @@ def get_openapi_schema_with_compat( ) finally: # Restore original method - setattr( - GenerateSchema, "_unknown_type_schema", original_unknown_type_schema - ) + setattr(GenerateSchema, "_unknown_type_schema", original_unknown_type_schema) return openapi_schema except (ImportError, AttributeError) as e: # If patching fails, try normal generation with error handling - verbose_proxy_logger.debug( - f"Could not patch Pydantic schema generation: {e}. Trying normal generation." - ) + verbose_proxy_logger.debug(f"Could not patch Pydantic schema generation: {e}. Trying normal generation.") try: return get_openapi_func( title=title, @@ -98,9 +92,8 @@ def get_openapi_schema_with_compat( # Check if it's a PydanticSchemaGenerationError by checking the error type name # This avoids import issues if PydanticSchemaGenerationError is not available error_type_name = type(pydantic_error).__name__ - if ( - error_type_name == "PydanticSchemaGenerationError" - or "PydanticSchemaGenerationError" in str(type(pydantic_error)) + if error_type_name == "PydanticSchemaGenerationError" or "PydanticSchemaGenerationError" in str( + type(pydantic_error) ): # If we still get the error, log it and return minimal schema verbose_proxy_logger.warning( diff --git a/litellm/proxy/common_utils/performance_utils.py b/litellm/proxy/common_utils/performance_utils.py index 6853a86d1df..6a67fbd93e8 100644 --- a/litellm/proxy/common_utils/performance_utils.py +++ b/litellm/proxy/common_utils/performance_utils.py @@ -54,9 +54,7 @@ def _start_profiling(profile_sampling_rate: float) -> None: if _profiler is None: _profiler = cProfile.Profile() _profiler.enable() - verbose_proxy_logger.info( - f"Profiling started with sampling rate: {profile_sampling_rate}" - ) + verbose_proxy_logger.info(f"Profiling started with sampling rate: {profile_sampling_rate}") def _start_profiling_for_request(profile_sampling_rate: float) -> bool: @@ -179,9 +177,7 @@ def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool: try: original_function = getattr(module, function_name, None) if original_function is None: - verbose_proxy_logger.warning( - f"Function {function_name} not found in module {module.__name__}" - ) + verbose_proxy_logger.warning(f"Function {function_name} not found in module {module.__name__}") return False # Store original function if not already wrapped @@ -192,14 +188,10 @@ def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool: profiled_function = _line_profiler(original_function) setattr(module, function_name, profiled_function) - verbose_proxy_logger.info( - f"Wrapped {module.__name__}.{function_name} with line_profiler" - ) + verbose_proxy_logger.info(f"Wrapped {module.__name__}.{function_name} with line_profiler") return True except Exception as e: - verbose_proxy_logger.error( - f"Error wrapping {function_name} with line_profiler: {e}" - ) + verbose_proxy_logger.error(f"Error wrapping {function_name} with line_profiler: {e}") return False @@ -228,9 +220,7 @@ def wrap_function_directly(func: Callable) -> Callable: # Suppress warnings about __wrapped__ - we intentionally want to profile the wrapper with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", message=".*__wrapped__.*", category=UserWarning - ) + warnings.filterwarnings("ignore", message=".*__wrapped__.*", category=UserWarning) # Add function to line_profiler and wrap it _line_profiler.add_function(func) profiled_function = _line_profiler(func) @@ -291,6 +281,4 @@ def register_shutdown_handler(output_file: Optional[str] = None) -> None: collect_line_profiler_stats(output_file=output_file) atexit.register(shutdown_handler) - verbose_proxy_logger.debug( - f"Registered line_profiler shutdown handler for {output_file}" - ) + verbose_proxy_logger.debug(f"Registered line_profiler shutdown handler for {output_file}") diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index 24e5c991794..8e7a568b14f 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -144,9 +144,7 @@ class ProxyRateLimitError(HTTPException, RateLimitError): # type: ignore[misc] self, detail: Any, headers: Optional[Mapping[str, Any]] = None, - category: Union[ - str, RateLimitErrorCategory - ] = RateLimitErrorCategory.LITELLM_RATE_LIMIT, + category: Union[str, RateLimitErrorCategory] = RateLimitErrorCategory.LITELLM_RATE_LIMIT, rate_limit_type: Optional[Union[str, RateLimitType]] = None, model: Optional[str] = None, llm_provider: Optional[str] = "litellm_proxy", @@ -159,9 +157,7 @@ class ProxyRateLimitError(HTTPException, RateLimitError): # type: ignore[misc] model = model or "" llm_provider = llm_provider or "litellm_proxy" message = _coerce_message(detail) - stringified_headers: Optional[Dict[str, str]] = ( - {k: str(v) for k, v in headers.items()} if headers else None - ) + stringified_headers: Optional[Dict[str, str]] = {k: str(v) for k, v in headers.items()} if headers else None # Initialize the FastAPI HTTPException portion first so its attributes # (status_code, detail, headers) are already on the instance before diff --git a/litellm/proxy/common_utils/rbac_utils.py b/litellm/proxy/common_utils/rbac_utils.py index 0c3ddf8e6fe..1e842610820 100644 --- a/litellm/proxy/common_utils/rbac_utils.py +++ b/litellm/proxy/common_utils/rbac_utils.py @@ -64,9 +64,7 @@ async def check_feature_access_for_user( raise HTTPException( status_code=403, - detail={ - "error": f"Access to {feature_name} is disabled for your role. Contact your proxy admin." - }, + detail={"error": f"Access to {feature_name} is disabled for your role. Contact your proxy admin."}, ) @@ -94,7 +92,5 @@ async def check_org_admin_can_generate_keys( raise HTTPException( status_code=403, - detail={ - "error": "key generation is disabled for org admins. Contact your proxy admin." - }, + detail={"error": "key generation is disabled for org admins. Contact your proxy admin."}, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 7c1dfe8dc90..e758420ee37 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -73,14 +73,10 @@ class ResetBudgetJob: try: from litellm.proxy.proxy_server import spend_counter_cache - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0, ttl=60 - ) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0, ttl=60) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0, ttl=60 - ) + await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=0.0, ttl=60) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to reset spend counter %s in Redis: %s. " @@ -89,9 +85,7 @@ class ResetBudgetJob: redis_err, ) except Exception as e: - verbose_proxy_logger.warning( - "Failed to reset spend counter %s: %s", counter_key, e - ) + verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: @@ -144,9 +138,7 @@ class ResetBudgetJob: rows = await table.find_many(where=where) except Exception as e: rows = [] - verbose_proxy_logger.warning( - "Failed to fetch %s for counter invalidation: %s", log_subject, e - ) + verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e) update_result = await table.update_many(where=where, data={"spend": 0}) @@ -161,9 +153,7 @@ class ResetBudgetJob: return update_result - async def reset_budget_for_litellm_team_members( - self, budgets_to_reset: List[LiteLLM_BudgetTableFull] - ): + async def reset_budget_for_litellm_team_members(self, budgets_to_reset: List[LiteLLM_BudgetTableFull]): """ Resets the budget for all LiteLLM Team Members if their budget has expired """ @@ -175,9 +165,7 @@ class ResetBudgetJob: cache_key_fn=lambda m: f"{m.team_id}_{m.user_id}", ) - async def reset_budget_for_keys_linked_to_budgets( - self, budgets_to_reset: List[LiteLLM_BudgetTableFull] - ): + async def reset_budget_for_keys_linked_to_budgets(self, budgets_to_reset: List[LiteLLM_BudgetTableFull]): """ Resets the spend for keys linked to budget tiers that are being reset. @@ -193,9 +181,7 @@ class ResetBudgetJob: cache_key_fn=lambda k: k.token, ) - async def reset_budget_for_orgs_linked_to_budgets( - self, budgets_to_reset: List[LiteLLM_BudgetTableFull] - ): + async def reset_budget_for_orgs_linked_to_budgets(self, budgets_to_reset: List[LiteLLM_BudgetTableFull]): """ Resets the spend for orgs linked to budget tiers that are being reset. """ @@ -211,9 +197,7 @@ class ResetBudgetJob: ], ) - async def reset_budget_for_tags_linked_to_budgets( - self, budgets_to_reset: List[LiteLLM_BudgetTableFull] - ): + async def reset_budget_for_tags_linked_to_budgets(self, budgets_to_reset: List[LiteLLM_BudgetTableFull]): """ Resets the spend for tags linked to budget tiers that are being reset. @@ -253,9 +237,7 @@ class ResetBudgetJob: if budgets_to_reset is not None and len(budgets_to_reset) > 0: for budget in budgets_to_reset: - budget = await ResetBudgetJob._reset_budget_reset_at_date( - budget, now - ) + budget = await ResetBudgetJob._reset_budget_reset_at_date(budget, now) await self.prisma_client.update_data( query_type="update_many", @@ -263,11 +245,7 @@ class ResetBudgetJob: table_name="budget", ) - budget_ids_to_reset = [ - budget.budget_id - for budget in budgets_to_reset - if budget.budget_id is not None - ] + budget_ids_to_reset = [budget.budget_id for budget in budgets_to_reset if budget.budget_id is not None] endusers_to_reset = await self.prisma_client.get_data( table_name="enduser", @@ -279,43 +257,26 @@ class ResetBudgetJob: # default budget via litellm.max_end_user_budget_id. These # users are enforced in-memory but never had budget_id # persisted, so the query above misses them. - if ( - litellm.max_end_user_budget_id is not None - and litellm.max_end_user_budget_id in budget_ids_to_reset - ): - default_budget_endusers = ( - await self._get_endusers_with_no_budget_id() - ) + if litellm.max_end_user_budget_id is not None and litellm.max_end_user_budget_id in budget_ids_to_reset: + default_budget_endusers = await self._get_endusers_with_no_budget_id() if default_budget_endusers: if endusers_to_reset is None: endusers_to_reset = default_budget_endusers else: endusers_to_reset.extend(default_budget_endusers) - await self.reset_budget_for_litellm_team_members( - budgets_to_reset=budgets_to_reset - ) + await self.reset_budget_for_litellm_team_members(budgets_to_reset=budgets_to_reset) - await self.reset_budget_for_keys_linked_to_budgets( - budgets_to_reset=budgets_to_reset - ) + await self.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset) - await self.reset_budget_for_orgs_linked_to_budgets( - budgets_to_reset=budgets_to_reset - ) + await self.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=budgets_to_reset) - await self.reset_budget_for_tags_linked_to_budgets( - budgets_to_reset=budgets_to_reset - ) + await self.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=budgets_to_reset) if endusers_to_reset is not None and len(endusers_to_reset) > 0: for enduser in endusers_to_reset: try: - updated_enduser = ( - await ResetBudgetJob._reset_budget_for_enduser( - enduser=enduser - ) - ) + updated_enduser = await ResetBudgetJob._reset_budget_for_enduser(enduser=enduser) if updated_enduser is not None: updated_endusers.append(updated_enduser) else: @@ -327,9 +288,7 @@ class ResetBudgetJob: ) except Exception as e: failed_endusers.append({"enduser": enduser, "error": str(e)}) - verbose_proxy_logger.exception( - "Failed to reset budget for enduser: %s", enduser - ) + verbose_proxy_logger.exception("Failed to reset budget for enduser: %s", enduser) verbose_proxy_logger.debug( "Updated users %s", @@ -356,26 +315,14 @@ class ResetBudgetJob: start_time=start_time, end_time=end_time, event_metadata={ - "num_budgets_found": ( - len(budgets_to_reset) if budgets_to_reset else 0 - ), - "budgets_found": json.dumps( - budgets_to_reset, indent=4, default=str - ), - "num_endusers_found": ( - len(endusers_to_reset) if endusers_to_reset else 0 - ), - "endusers_found": json.dumps( - endusers_to_reset, indent=4, default=str - ), + "num_budgets_found": (len(budgets_to_reset) if budgets_to_reset else 0), + "budgets_found": json.dumps(budgets_to_reset, indent=4, default=str), + "num_endusers_found": (len(endusers_to_reset) if endusers_to_reset else 0), + "endusers_found": json.dumps(endusers_to_reset, indent=4, default=str), "num_endusers_updated": len(updated_endusers), - "endusers_updated": json.dumps( - updated_endusers, indent=4, default=str - ), + "endusers_updated": json.dumps(updated_endusers, indent=4, default=str), "num_endusers_failed": len(failed_endusers), - "endusers_failed": json.dumps( - failed_endusers, indent=4, default=str - ), + "endusers_failed": json.dumps(failed_endusers, indent=4, default=str), }, ) ) @@ -390,18 +337,10 @@ class ResetBudgetJob: start_time=start_time, end_time=end_time, event_metadata={ - "num_budgets_found": ( - len(budgets_to_reset) if budgets_to_reset else 0 - ), - "budgets_found": json.dumps( - budgets_to_reset, indent=4, default=str - ), - "num_endusers_found": ( - len(endusers_to_reset) if endusers_to_reset else 0 - ), - "endusers_found": json.dumps( - endusers_to_reset, indent=4, default=str - ), + "num_budgets_found": (len(budgets_to_reset) if budgets_to_reset else 0), + "budgets_found": json.dumps(budgets_to_reset, indent=4, default=str), + "num_endusers_found": (len(endusers_to_reset) if endusers_to_reset else 0), + "endusers_found": json.dumps(endusers_to_reset, indent=4, default=str), }, ) ) @@ -424,9 +363,7 @@ class ResetBudgetJob: ) return [LiteLLM_EndUserTable(**row.dict()) for row in rows] - async def _write_key_reset_updates( - self, updated_keys: List[LiteLLM_VerificationToken] - ) -> None: + async def _write_key_reset_updates(self, updated_keys: List[LiteLLM_VerificationToken]) -> None: """ Write per-row {spend, budget_reset_at} updates for keys. @@ -448,9 +385,7 @@ class ResetBudgetJob: ) await batcher.commit() - async def _write_user_reset_updates( - self, updated_users: List[LiteLLM_UserTable] - ) -> None: + async def _write_user_reset_updates(self, updated_users: List[LiteLLM_UserTable]) -> None: """ Write per-row {spend, budget_reset_at} updates for users. @@ -469,9 +404,7 @@ class ResetBudgetJob: ) await batcher.commit() - async def _write_team_reset_updates( - self, updated_teams: List[LiteLLM_TeamTable] - ) -> None: + async def _write_team_reset_updates(self, updated_teams: List[LiteLLM_TeamTable]) -> None: """ Write per-row {spend, budget_reset_at} updates for teams. @@ -503,32 +436,22 @@ class ResetBudgetJob: keys_to_reset = await self.prisma_client.get_data( table_name="key", query_type="find_all", expires=now, reset_at=now ) - verbose_proxy_logger.debug( - "Keys to reset %s", json.dumps(keys_to_reset, indent=4, default=str) - ) + verbose_proxy_logger.debug("Keys to reset %s", json.dumps(keys_to_reset, indent=4, default=str)) updated_keys: List[LiteLLM_VerificationToken] = [] failed_keys = [] if keys_to_reset is not None and len(keys_to_reset) > 0: for key in keys_to_reset: try: - updated_key = await ResetBudgetJob._reset_budget_for_key( - key=key, current_time=now - ) + updated_key = await ResetBudgetJob._reset_budget_for_key(key=key, current_time=now) if updated_key is not None: updated_keys.append(updated_key) else: - failed_keys.append( - {"key": key, "error": "Returned None without exception"} - ) + failed_keys.append({"key": key, "error": "Returned None without exception"}) except Exception as e: failed_keys.append({"key": key, "error": str(e)}) - verbose_proxy_logger.exception( - "Failed to reset budget for key: %s", key - ) + verbose_proxy_logger.exception("Failed to reset budget for key: %s", key) - verbose_proxy_logger.debug( - "Updated keys %s", json.dumps(updated_keys, indent=4, default=str) - ) + verbose_proxy_logger.debug("Updated keys %s", json.dumps(updated_keys, indent=4, default=str)) if updated_keys: await self._write_key_reset_updates(updated_keys=updated_keys) @@ -539,9 +462,7 @@ class ResetBudgetJob: end_time = time.time() if len(failed_keys) > 0: # If any keys failed to reset - raise Exception( - f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}" - ) + raise Exception(f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}") asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( @@ -586,17 +507,13 @@ class ResetBudgetJob: start_time = time.time() users_to_reset: Optional[List[LiteLLM_UserTable]] = None try: - users_to_reset = await self.prisma_client.get_data( - table_name="user", query_type="find_all", reset_at=now - ) + users_to_reset = await self.prisma_client.get_data(table_name="user", query_type="find_all", reset_at=now) updated_users: List[LiteLLM_UserTable] = [] failed_users = [] if users_to_reset is not None and len(users_to_reset) > 0: for user in users_to_reset: try: - updated_user = await ResetBudgetJob._reset_budget_for_user( - user=user, current_time=now - ) + updated_user = await ResetBudgetJob._reset_budget_for_user(user=user, current_time=now) if updated_user is not None: updated_users.append(updated_user) else: @@ -608,27 +525,19 @@ class ResetBudgetJob: ) except Exception as e: failed_users.append({"user": user, "error": str(e)}) - verbose_proxy_logger.exception( - "Failed to reset budget for user: %s", user - ) + verbose_proxy_logger.exception("Failed to reset budget for user: %s", user) - verbose_proxy_logger.debug( - "Updated users %s", json.dumps(updated_users, indent=4, default=str) - ) + verbose_proxy_logger.debug("Updated users %s", json.dumps(updated_users, indent=4, default=str)) if updated_users: await self._write_user_reset_updates(updated_users=updated_users) for u in updated_users: user_id = getattr(u, "user_id", None) if user_id: - await self._invalidate_spend_counter( - f"spend:user:{user_id}" - ) + await self._invalidate_spend_counter(f"spend:user:{user_id}") end_time = time.time() if len(failed_users) > 0: # If any users failed to reset - raise Exception( - f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}" - ) + raise Exception(f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}") asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( @@ -639,13 +548,9 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, - "users_found": json.dumps( - users_to_reset, indent=4, default=str - ), + "users_found": json.dumps(users_to_reset, indent=4, default=str), "num_users_updated": len(updated_users), - "users_updated": json.dumps( - updated_users, indent=4, default=str - ), + "users_updated": json.dumps(updated_users, indent=4, default=str), "num_users_failed": len(failed_users), "users_failed": json.dumps(failed_users, indent=4, default=str), }, @@ -663,9 +568,7 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, - "users_found": json.dumps( - users_to_reset, indent=4, default=str - ), + "users_found": json.dumps(users_to_reset, indent=4, default=str), }, ) ) @@ -679,17 +582,13 @@ class ResetBudgetJob: start_time = time.time() teams_to_reset: Optional[List[LiteLLM_TeamTable]] = None try: - teams_to_reset = await self.prisma_client.get_data( - table_name="team", query_type="find_all", reset_at=now - ) + teams_to_reset = await self.prisma_client.get_data(table_name="team", query_type="find_all", reset_at=now) updated_teams: List[LiteLLM_TeamTable] = [] failed_teams = [] if teams_to_reset is not None and len(teams_to_reset) > 0: for team in teams_to_reset: try: - updated_team = await ResetBudgetJob._reset_budget_for_team( - team=team, current_time=now - ) + updated_team = await ResetBudgetJob._reset_budget_for_team(team=team, current_time=now) if updated_team is not None: updated_teams.append(updated_team) else: @@ -701,27 +600,19 @@ class ResetBudgetJob: ) except Exception as e: failed_teams.append({"team": team, "error": str(e)}) - verbose_proxy_logger.exception( - "Failed to reset budget for team: %s", team - ) + verbose_proxy_logger.exception("Failed to reset budget for team: %s", team) - verbose_proxy_logger.debug( - "Updated teams %s", json.dumps(updated_teams, indent=4, default=str) - ) + verbose_proxy_logger.debug("Updated teams %s", json.dumps(updated_teams, indent=4, default=str)) if updated_teams: await self._write_team_reset_updates(updated_teams=updated_teams) for t in updated_teams: team_id = getattr(t, "team_id", None) if team_id: - await self._invalidate_spend_counter( - f"spend:team:{team_id}" - ) + await self._invalidate_spend_counter(f"spend:team:{team_id}") end_time = time.time() if len(failed_teams) > 0: # If any teams failed to reset - raise Exception( - f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}" - ) + raise Exception(f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}") asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( @@ -732,13 +623,9 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, - "teams_found": json.dumps( - teams_to_reset, indent=4, default=str - ), + "teams_found": json.dumps(teams_to_reset, indent=4, default=str), "num_teams_updated": len(updated_teams), - "teams_updated": json.dumps( - updated_teams, indent=4, default=str - ), + "teams_updated": json.dumps(updated_teams, indent=4, default=str), "num_teams_failed": len(failed_teams), "teams_failed": json.dumps(failed_teams, indent=4, default=str), }, @@ -756,9 +643,7 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, - "teams_found": json.dumps( - teams_to_reset, indent=4, default=str - ), + "teams_found": json.dumps(teams_to_reset, indent=4, default=str), }, ) ) @@ -777,24 +662,16 @@ class ResetBudgetJob: reset_at_str = window.get("reset_at") if not reset_at_str: return False - reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace( - tzinfo=None - ) + reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None) if reset_at > now: return False spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 - ) + await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=0.0) except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to reset Redis counter %s: %s", counter_key, redis_err - ) - window["reset_at"] = get_budget_reset_time( - budget_duration=window["budget_duration"] - ).isoformat() + verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) + window["reset_at"] = get_budget_reset_time(budget_duration=window["budget_duration"]).isoformat() return True async def reset_budget_windows(self) -> None: @@ -816,8 +693,7 @@ class ResetBudgetJob: # --- Keys --- try: key_rows = await self.prisma_client.db.query_raw( - 'SELECT token, budget_limits FROM "LiteLLM_VerificationToken" ' - "WHERE budget_limits IS NOT NULL" + 'SELECT token, budget_limits FROM "LiteLLM_VerificationToken" WHERE budget_limits IS NOT NULL' ) for row in key_rows: raw = row["budget_limits"] @@ -826,12 +702,8 @@ class ResetBudgetJob: windows: list = raw if isinstance(raw, list) else json.loads(raw) changed = False for window in windows: - counter_key = ( - f"spend:key:{row['token']}:window:{window['budget_duration']}" - ) - if await ResetBudgetJob._reset_expired_window( - window, counter_key, spend_counter_cache, now - ): + counter_key = f"spend:key:{row['token']}:window:{window['budget_duration']}" + if await ResetBudgetJob._reset_expired_window(window, counter_key, spend_counter_cache, now): changed = True if changed: await VerificationTokenRepository(self.prisma_client).table.update( @@ -839,15 +711,12 @@ class ResetBudgetJob: data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) except Exception as e: - verbose_proxy_logger.exception( - "Failed to reset budget windows for keys: %s", e - ) + verbose_proxy_logger.exception("Failed to reset budget windows for keys: %s", e) # --- Teams --- try: team_rows = await self.prisma_client.db.query_raw( - 'SELECT team_id, budget_limits FROM "LiteLLM_TeamTable" ' - "WHERE budget_limits IS NOT NULL" + 'SELECT team_id, budget_limits FROM "LiteLLM_TeamTable" WHERE budget_limits IS NOT NULL' ) for row in team_rows: raw = row["budget_limits"] @@ -857,9 +726,7 @@ class ResetBudgetJob: changed = False for window in windows: counter_key = f"spend:team:{row['team_id']}:window:{window['budget_duration']}" - if await ResetBudgetJob._reset_expired_window( - window, counter_key, spend_counter_cache, now - ): + if await ResetBudgetJob._reset_expired_window(window, counter_key, spend_counter_cache, now): changed = True if changed: await TeamRepository(self.prisma_client).table.update( @@ -867,9 +734,7 @@ class ResetBudgetJob: data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) except Exception as e: - verbose_proxy_logger.exception( - "Failed to reset budget windows for teams: %s", e - ) + verbose_proxy_logger.exception("Failed to reset budget windows for teams: %s", e) @staticmethod async def _reset_budget_common( @@ -894,32 +759,20 @@ class ResetBudgetJob: get_budget_reset_time, ) - item.budget_reset_at = get_budget_reset_time( - budget_duration=item.budget_duration - ) + item.budget_reset_at = get_budget_reset_time(budget_duration=item.budget_duration) return item except Exception as e: - verbose_proxy_logger.exception( - "Error resetting budget for %s: %s. Item: %s", item_type, e, item - ) + verbose_proxy_logger.exception("Error resetting budget for %s: %s. Item: %s", item_type, e, item) raise e @staticmethod - async def _reset_budget_for_team( - team: LiteLLM_TeamTable, current_time: datetime - ) -> Optional[LiteLLM_TeamTable]: - await ResetBudgetJob._reset_budget_common( - item=team, current_time=current_time, item_type="team" - ) + async def _reset_budget_for_team(team: LiteLLM_TeamTable, current_time: datetime) -> Optional[LiteLLM_TeamTable]: + await ResetBudgetJob._reset_budget_common(item=team, current_time=current_time, item_type="team") return team @staticmethod - async def _reset_budget_for_user( - user: LiteLLM_UserTable, current_time: datetime - ) -> Optional[LiteLLM_UserTable]: - await ResetBudgetJob._reset_budget_common( - item=user, current_time=current_time, item_type="user" - ) + async def _reset_budget_for_user(user: LiteLLM_UserTable, current_time: datetime) -> Optional[LiteLLM_UserTable]: + await ResetBudgetJob._reset_budget_common(item=user, current_time=current_time, item_type="user") return user @staticmethod @@ -929,9 +782,7 @@ class ResetBudgetJob: try: enduser.spend = 0.0 except Exception as e: - verbose_proxy_logger.exception( - "Error resetting budget for enduser: %s. Item: %s", e, enduser - ) + verbose_proxy_logger.exception("Error resetting budget for enduser: %s. Item: %s", e, enduser) raise e return enduser @@ -945,13 +796,9 @@ class ResetBudgetJob: get_budget_reset_time, ) - budget.budget_reset_at = get_budget_reset_time( - budget_duration=budget.budget_duration - ) + budget.budget_reset_at = get_budget_reset_time(budget_duration=budget.budget_duration) except Exception as e: - verbose_proxy_logger.exception( - "Error resetting budget_reset_at for budget: %s. Item: %s", e, budget - ) + verbose_proxy_logger.exception("Error resetting budget_reset_at for budget: %s. Item: %s", e, budget) raise e return budget @@ -959,7 +806,5 @@ class ResetBudgetJob: async def _reset_budget_for_key( key: LiteLLM_VerificationToken, current_time: datetime ) -> Optional[LiteLLM_VerificationToken]: - await ResetBudgetJob._reset_budget_common( - item=key, current_time=current_time, item_type="key" - ) + await ResetBudgetJob._reset_budget_common(item=key, current_time=current_time, item_type="key") return key diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 8ef127ac706..09921a3ac1d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -67,9 +67,7 @@ class UserApiKeyCache(DualCache): ) -> Union[Any, Optional[BaseModel]]: if model_type is None and "model_type" in kwargs: model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) - cached = super().get_cache( - key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs - ) + cached = super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) if model_type is None: return cached if cached is None: @@ -77,8 +75,7 @@ class UserApiKeyCache(DualCache): decoded = CacheCodec.deserialize(cached, model_type=model_type) if decoded is None: verbose_proxy_logger.error( - "UserApiKeyCache.get_cache failed to deserialize cached value for " - "key=%r model_type=%s", + "UserApiKeyCache.get_cache failed to deserialize cached value for key=%r model_type=%s", key, getattr(model_type, "__name__", str(model_type)), ) @@ -125,8 +122,7 @@ class UserApiKeyCache(DualCache): decoded = CacheCodec.deserialize(cached, model_type=model_type) if decoded is None: verbose_proxy_logger.error( - "UserApiKeyCache.async_get_cache failed to deserialize cached value for " - "key=%r model_type=%s", + "UserApiKeyCache.async_get_cache failed to deserialize cached value for key=%r model_type=%s", key, getattr(model_type, "__name__", str(model_type)), ) @@ -136,16 +132,12 @@ class UserApiKeyCache(DualCache): def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) payload = CacheCodec.serialize(value, model_type=model_type) - return super().set_cache( - key=key, value=payload, local_only=local_only, **kwargs - ) + return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) payload = CacheCodec.serialize(value, model_type=model_type) - return await super().async_set_cache( - key=key, value=payload, local_only=local_only, **kwargs - ) + return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) async def async_set_cache_pipeline( # type: ignore[override] self, cache_list: list, local_only: bool = False, **kwargs @@ -154,13 +146,8 @@ class UserApiKeyCache(DualCache): Batch writes with the same Codec boundary as ``async_set_cache`` without ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. """ - normalized = [ - (key, CacheCodec.serialize(value, model_type=None)) - for key, value in cache_list - ] - return await super().async_set_cache_pipeline( - cache_list=normalized, local_only=local_only, **kwargs - ) + normalized = [(key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list] + return await super().async_set_cache_pipeline(cache_list=normalized, local_only=local_only, **kwargs) def get_management_object_ttl(cache: DualCache) -> float: diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index 381b0f815d2..8cde91e32b2 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -68,11 +68,7 @@ class ComplianceChecker: check_name="Guardrails applied", article="Art. 9", passed=has_guardrails, - detail=( - f"{len(self.guardrails)} guardrail(s) applied" - if has_guardrails - else "No guardrails applied" - ), + detail=(f"{len(self.guardrails)} guardrail(s) applied" if has_guardrails else "No guardrails applied"), ) def _check_art_5_content_screened(self) -> ComplianceCheckResult: @@ -112,11 +108,7 @@ class ComplianceChecker: check_name="Audit record complete", article="Art. 12", passed=audit_complete, - detail=( - "All required audit fields present" - if audit_complete - else f"Missing: {', '.join(missing)}" - ), + detail=("All required audit fields present" if audit_complete else f"Missing: {', '.join(missing)}"), ) # ── GDPR Helper Methods ────────────────────────────────────────────────── @@ -179,11 +171,7 @@ class ComplianceChecker: check_name="Audit record complete", article="Art. 30", passed=audit_complete, - detail=( - "All required audit fields present" - if audit_complete - else f"Missing: {', '.join(missing)}" - ), + detail=("All required audit fields present" if audit_complete else f"Missing: {', '.join(missing)}"), ) # ── Main Compliance Check Methods ──────────────────────────────────────── diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 52bd9d66b7d..bc871479356 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -245,10 +245,7 @@ async def _process_binary_request( content_type = "application/pdf" if not isinstance(content, bytes): - raise TypeError( - "aretrieve_container_file_content expected bytes, got " - f"{type(content).__name__}" - ) + raise TypeError(f"aretrieve_container_file_content expected bytes, got {type(content).__name__}") return Response( content=content, @@ -455,9 +452,7 @@ def register_container_file_endpoints(router: APIRouter) -> None: is_multipart = endpoint_config.get("is_multipart", False) # Create handler with correct signature for path params - handler = _create_handler_for_path_params( - path_params, route_type, returns_binary, is_multipart - ) + handler = _create_handler_for_path_params(path_params, route_type, returns_binary, is_multipart) # Register routes route_method = getattr(router, method) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 8118d53b9f6..583d8db7d30 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -47,15 +47,11 @@ def _allowed_container_ids_cache_key(owner_scopes: List[str]) -> str: return json.dumps(sorted(owner_scopes)) -def _container_model_object_id( - original_container_id: str, custom_llm_provider: str -) -> str: +def _container_model_object_id(original_container_id: str, custom_llm_provider: str) -> str: return f"{CONTAINER_OBJECT_PURPOSE}:{custom_llm_provider}:{original_container_id}" -def decode_container_id_for_ownership( - container_id: str, custom_llm_provider: str -) -> Tuple[str, str]: +def decode_container_id_for_ownership(container_id: str, custom_llm_provider: str) -> Tuple[str, str]: decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) original_container_id = decoded.get("response_id", container_id) decoded_provider = decoded.get("custom_llm_provider") @@ -79,9 +75,7 @@ async def get_container_forwarding_params( # ``model_id``. Recover it from the encoded ``unified_object_id`` # captured on the ownership row at create time — when the router # selected a specific deployment that ID embeds the model_id. - stored_id = await _get_stored_container_id( - original_container_id, custom_llm_provider - ) + stored_id = await _get_stored_container_id(original_container_id, custom_llm_provider) if stored_id and stored_id != container_id: stored_decoded = ResponsesAPIRequestUtils._decode_container_id(stored_id) stored_model_id = stored_decoded.get("model_id") @@ -140,15 +134,11 @@ async def record_container_owners_from_responses_response( custom_llm_provider: Optional[str] = None, ) -> None: """Track containers created implicitly by code interpreter in /v1/responses.""" - container_ids = ( - ResponsesAPIRequestUtils.collect_container_ids_from_responses_response(response) - ) + container_ids = ResponsesAPIRequestUtils.collect_container_ids_from_responses_response(response) if not container_ids: return - resolved_provider = ( - custom_llm_provider or _custom_llm_provider_from_responses_response(response) - ) + resolved_provider = custom_llm_provider or _custom_llm_provider_from_responses_response(response) for container_id in container_ids: try: @@ -163,8 +153,7 @@ async def record_container_owners_from_responses_response( # batch — other containers in the same response should still # get recorded so their follow-up file API calls don't 403. verbose_proxy_logger.exception( - "Failed to record container ownership from responses output " - "for container_id=%s: %s", + "Failed to record container ownership from responses output for container_id=%s: %s", container_id, e, ) @@ -177,9 +166,7 @@ async def record_container_owner( ) -> Any: container_id = _get_response_id(response) if container_id is None: - verbose_proxy_logger.warning( - "Skipping container ownership tracking because provider response has no id" - ) + verbose_proxy_logger.warning("Skipping container ownership tracking because provider response has no id") return response owner = get_primary_resource_owner_scope(user_api_key_dict) if owner is None: @@ -195,12 +182,8 @@ async def record_container_owner( detail="Unable to record container ownership: caller has no identity scope.", ) - original_container_id, resolved_provider = decode_container_id_for_ownership( - container_id, custom_llm_provider - ) - model_object_id = _container_model_object_id( - original_container_id, resolved_provider - ) + original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) + model_object_id = _container_model_object_id(original_container_id, resolved_provider) file_object = _dump_response(response) file_object["custom_llm_provider"] = resolved_provider file_object["provider_container_id"] = original_container_id @@ -209,9 +192,7 @@ async def record_container_owner( prisma_client = await _get_prisma_client() if prisma_client is None: - verbose_proxy_logger.warning( - "Skipping container ownership tracking because prisma_client is None" - ) + verbose_proxy_logger.warning("Skipping container ownership tracking because prisma_client is None") return response table = ManagedObjectRepository(prisma_client).table @@ -219,9 +200,7 @@ async def record_container_owner( if existing is not None: if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE: raise HTTPException(status_code=500, detail="Unable to track container") - if not user_can_access_resource_owner( - getattr(existing, "created_by", None), user_api_key_dict - ): + if not user_can_access_resource_owner(getattr(existing, "created_by", None), user_api_key_dict): raise HTTPException(status_code=403, detail="Forbidden") await table.update( where={"model_object_id": model_object_id}, @@ -251,18 +230,12 @@ async def record_container_owner( # tuples self-correct on the 60s TTL. caller_scopes = get_resource_owner_scopes(user_api_key_dict) if caller_scopes: - _ALLOWED_CONTAINER_IDS_CACHE.delete_cache( - _allowed_container_ids_cache_key(caller_scopes) - ) + _ALLOWED_CONTAINER_IDS_CACHE.delete_cache(_allowed_container_ids_cache_key(caller_scopes)) return response -async def _get_container_owner( - original_container_id: str, custom_llm_provider: str -) -> Optional[str]: - model_object_id = _container_model_object_id( - original_container_id, custom_llm_provider - ) +async def _get_container_owner(original_container_id: str, custom_llm_provider: str) -> Optional[str]: + model_object_id = _container_model_object_id(original_container_id, custom_llm_provider) cached = _CONTAINER_OWNER_CACHE.get_cache(model_object_id) if cached == _NEGATIVE_OWNER_SENTINEL: @@ -281,24 +254,16 @@ async def _get_container_owner( } ) owner = getattr(row, "created_by", None) if row is not None else None - _CONTAINER_OWNER_CACHE.set_cache( - model_object_id, owner if owner is not None else _NEGATIVE_OWNER_SENTINEL - ) + _CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner if owner is not None else _NEGATIVE_OWNER_SENTINEL) stored_id = getattr(row, "unified_object_id", None) if row is not None else None _CONTAINER_STORED_ID_CACHE.set_cache( model_object_id, - ( - stored_id - if isinstance(stored_id, str) and stored_id - else _NEGATIVE_STORED_ID_SENTINEL - ), + (stored_id if isinstance(stored_id, str) and stored_id else _NEGATIVE_STORED_ID_SENTINEL), ) return owner -async def _get_stored_container_id( - original_container_id: str, custom_llm_provider: str -) -> Optional[str]: +async def _get_stored_container_id(original_container_id: str, custom_llm_provider: str) -> Optional[str]: """Return the ``unified_object_id`` stored at create time, if any. Used by :func:`get_container_forwarding_params` to recover the @@ -306,9 +271,7 @@ async def _get_stored_container_id( value is the encoded form produced by ``encode_container_id_in_response`` when the router selected a specific deployment. """ - model_object_id = _container_model_object_id( - original_container_id, custom_llm_provider - ) + model_object_id = _container_model_object_id(original_container_id, custom_llm_provider) cached = _CONTAINER_STORED_ID_CACHE.get_cache(model_object_id) if cached == _NEGATIVE_STORED_ID_SENTINEL: @@ -329,11 +292,7 @@ async def _get_stored_container_id( stored_id = getattr(row, "unified_object_id", None) if row is not None else None _CONTAINER_STORED_ID_CACHE.set_cache( model_object_id, - ( - stored_id - if isinstance(stored_id, str) and stored_id - else _NEGATIVE_STORED_ID_SENTINEL - ), + (stored_id if isinstance(stored_id, str) and stored_id else _NEGATIVE_STORED_ID_SENTINEL), ) return stored_id if isinstance(stored_id, str) and stored_id else None @@ -343,9 +302,7 @@ async def assert_user_can_access_container( user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str, ) -> Tuple[str, str]: - original_container_id, resolved_provider = decode_container_id_for_ownership( - container_id, custom_llm_provider - ) + original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) if is_proxy_admin(user_api_key_dict): return original_container_id, resolved_provider @@ -370,9 +327,7 @@ def _get_container_list_data(response: Any) -> Optional[List[Any]]: return data if isinstance(data, list) else None -def _set_container_list_data( - response: Any, data: List[Any], removed_filtered_items: bool = False -) -> Any: +def _set_container_list_data(response: Any, data: List[Any], removed_filtered_items: bool = False) -> Any: if isinstance(response, dict): response["data"] = data if data: @@ -418,11 +373,7 @@ async def _get_allowed_container_ids( "created_by": {"in": owner_scopes}, } ) - allowed_ids = { - row.model_object_id - for row in rows - if getattr(row, "model_object_id", None) is not None - } + allowed_ids = {row.model_object_id for row in rows if getattr(row, "model_object_id", None) is not None} # ``InMemoryCache.get_cache`` attempts ``json.loads`` on the stored # value; passing a set would round-trip through that path # unnecessarily. Store as a list and rehydrate above. @@ -448,15 +399,8 @@ async def filter_container_list_response( container_id = _get_response_id(item) if container_id is None: continue - original_container_id, resolved_provider = decode_container_id_for_ownership( - container_id, custom_llm_provider - ) - if ( - _container_model_object_id(original_container_id, resolved_provider) - in allowed_container_ids - ): + original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) + if _container_model_object_id(original_container_id, resolved_provider) in allowed_container_ids: filtered.append(item) - return _set_container_list_data( - response, filtered, removed_filtered_items=len(filtered) != len(data) - ) + return _set_container_list_data(response, filtered, removed_filtered_items=len(filtered) != len(data)) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index a716857111b..e94016c555e 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -28,9 +28,7 @@ class CredentialHelperUtils: """Encrypt values in credential.credential_values and add to DB""" encrypted_credential_values = {} for key, value in (credential.credential_values or {}).items(): - encrypted_credential_values[key] = encrypt_value_helper( - value, new_encryption_key - ) + encrypted_credential_values[key] = encrypt_value_helper(value, new_encryption_key) # Return a new object to avoid mutating the caller's credential, which # is kept in memory and should remain unencrypted. @@ -75,9 +73,7 @@ async def create_credential( model = llm_router.get_deployment(credential.model_id) if model is None: raise HTTPException(status_code=404, detail="Model not found") - credential_values = llm_router.get_deployment_credentials( - credential.model_id - ) + credential_values = llm_router.get_deployment_credentials(credential.model_id) if credential_values is None: raise HTTPException(status_code=404, detail="Model not found") credential.credential_values = credential_values @@ -92,9 +88,7 @@ async def create_credential( credential_values=credential.credential_values, credential_info=credential.credential_info, ) - encrypted_credential = CredentialHelperUtils.encrypt_credential_values( - processed_credential - ) + encrypted_credential = CredentialHelperUtils.encrypt_credential_values(processed_credential) credentials_dict = encrypted_credential.model_dump() credentials_dict_jsonified = jsonify_object(credentials_dict) await CredentialsRepository(prisma_client).create( @@ -150,9 +144,7 @@ async def get_credentials( async def get_credential_by_name( request: Request, fastapi_response: Response, - credential_name: str = Path( - ..., description="The credential name, percent-decoded; may contain slashes" - ), + credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -230,9 +222,7 @@ async def get_credential_by_model( async def delete_credential( request: Request, fastapi_response: Response, - credential_name: str = Path( - ..., description="The credential name, percent-decoded; may contain slashes" - ), + credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -249,11 +239,7 @@ async def delete_credential( await CredentialsRepository(prisma_client).delete_by_name(credential_name) ## DELETE FROM LITELLM ## - litellm.credential_list = [ - cred - for cred in litellm.credential_list - if cred.credential_name != credential_name - ] + litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name] return {"success": True, "message": "Credential deleted successfully"} except Exception as e: return handle_exception_on_proxy(e) @@ -284,9 +270,7 @@ def update_db_credential( # update litellm params if encrypted_credential.credential_values: # Encrypt any sensitive values - encrypted_params = { - k: v for k, v in encrypted_credential.credential_values.items() - } + encrypted_params = {k: v for k, v in encrypted_credential.credential_values.items()} merged_credential.credential_values.update(encrypted_params) @@ -309,9 +293,7 @@ async def update_credential( request: Request, fastapi_response: Response, credential: CredentialItem, - credential_name: str = Path( - ..., description="The credential name, percent-decoded; may contain slashes" - ), + credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -361,11 +343,7 @@ async def update_credential( ) # Remove old entry if renamed, then use upsert_credentials to handle duplicates if new_name != credential_name: - litellm.credential_list = [ - c - for c in litellm.credential_list - if c.credential_name != credential_name - ] + litellm.credential_list = [c for c in litellm.credential_list if c.credential_name != credential_name] CredentialAccessor.upsert_credentials([updated_in_memory]) return {"success": True, "message": "Credential updated successfully"} diff --git a/litellm/proxy/custom_auth_auto.py b/litellm/proxy/custom_auth_auto.py index c8991520898..97e928b6546 100644 --- a/litellm/proxy/custom_auth_auto.py +++ b/litellm/proxy/custom_auth_auto.py @@ -11,9 +11,7 @@ from fastapi import Request from litellm.proxy._types import ProxyException, UserAPIKeyAuth -async def user_api_key_auth( - request: Request, api_key: str -) -> Union[UserAPIKeyAuth, str]: +async def user_api_key_auth(request: Request, api_key: str) -> Union[UserAPIKeyAuth, str]: try: if api_key.startswith("my-custom-key"): return "sk-P1zJMdsqCPNN54alZd_ETw" diff --git a/litellm/proxy/custom_sso.py b/litellm/proxy/custom_sso.py index 43fb9f97ce4..4f621a9659e 100644 --- a/litellm/proxy/custom_sso.py +++ b/litellm/proxy/custom_sso.py @@ -21,9 +21,7 @@ from litellm.proxy import proxy_server async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: try: if userIDPInfo.id is None: - raise ValueError( - f"No ID found for user. userIDPInfo.id is None {userIDPInfo}" - ) + raise ValueError(f"No ID found for user. userIDPInfo.id is None {userIDPInfo}") # Access extra fields from the IDP response (requires GENERIC_USER_EXTRA_ATTRIBUTES env var) # Example: Set GENERIC_USER_EXTRA_ATTRIBUTES="group,NTID,domain" to capture these fields @@ -32,9 +30,7 @@ async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: # check if user exists in litellm proxy DB if proxy_server.prisma_client is not None: - _user_info = await proxy_server.prisma_client.get_data( - user_id=userIDPInfo.id - ) + _user_info = await proxy_server.prisma_client.get_data(user_id=userIDPInfo.id) return SSOUserDefinedValues( models=[], diff --git a/litellm/proxy/db/base_client.py b/litellm/proxy/db/base_client.py index 07f0ecdc47d..6e5fff44c79 100644 --- a/litellm/proxy/db/base_client.py +++ b/litellm/proxy/db/base_client.py @@ -21,17 +21,13 @@ class CustomDB: """ pass - def update_data( - self, key: str, value: Any, table_name: Literal["user", "key", "config"] - ): + def update_data(self, key: str, value: Any, table_name: Literal["user", "key", "config"]): """ For cost tracking logic """ pass - def delete_data( - self, keys: List[str], table_name: Literal["user", "key", "config"] - ): + def delete_data(self, keys: List[str], table_name: Literal["user", "key", "config"]): """ For /key/delete endpoint s """ diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index d9e21fc5d2a..411faa7e4e0 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -154,9 +154,7 @@ async def create_missing_views(db: _db): verbose_logger.debug("MonthlyGlobalSpendPerKey Created!") try: - await db.query_raw( - """SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""" - ) + await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""") verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Exists!") except Exception as e: error_msg = str(e).lower() diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 57a062b509a..cf4c3e98f00 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -91,11 +91,7 @@ def _extract_cache_creation_tokens(usage_obj: dict) -> int: if explicit: return int(explicit) details = usage_obj.get("prompt_tokens_details") or {} - return int( - details.get("cache_write_tokens", 0) - or details.get("cache_creation_tokens", 0) - or 0 - ) + return int(details.get("cache_write_tokens", 0) or details.get("cache_creation_tokens", 0) or 0) class DBSpendUpdateWriter: @@ -279,13 +275,9 @@ class DBSpendUpdateWriter: # --- MCP tool calls --- sl_object = kwargs.get("standard_logging_object") if sl_object is not None: - mcp_metadata = (sl_object.get("metadata", {}) or {}).get( - "mcp_tool_call_metadata" - ) + mcp_metadata = (sl_object.get("metadata", {}) or {}).get("mcp_tool_call_metadata") if mcp_metadata and isinstance(mcp_metadata, dict): - tool_name = mcp_metadata.get( - "namespaced_tool_name" - ) or mcp_metadata.get("name") + tool_name = mcp_metadata.get("namespaced_tool_name") or mcp_metadata.get("name") mcp_server_name = mcp_metadata.get("mcp_server_name") if tool_name: _enqueue(tool_name, origin=mcp_server_name or "user_defined") @@ -304,9 +296,7 @@ class DBSpendUpdateWriter: # (Anthropic format: tools[].name, no "function" wrapper) --- passthrough_payload = kwargs.get("passthrough_logging_payload") or {} request_body = ( - passthrough_payload.get("request_body") - if isinstance(passthrough_payload, dict) - else None + passthrough_payload.get("request_body") if isinstance(passthrough_payload, dict) else None ) or {} for tool_def in request_body.get("tools") or []: if not isinstance(tool_def, dict): @@ -316,9 +306,7 @@ class DBSpendUpdateWriter: _enqueue(name) # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) --- - if completion_response is not None and hasattr( - completion_response, "choices" - ): + if completion_response is not None and hasattr(completion_response, "choices"): for choice in completion_response.choices or []: message = getattr(choice, "message", None) if message is None: @@ -334,9 +322,7 @@ class DBSpendUpdateWriter: if tool_name: _enqueue(tool_name) except Exception as e: - verbose_proxy_logger.debug( - "_enqueue_tool_registry_upsert error (non-blocking): %s", e - ) + verbose_proxy_logger.debug("_enqueue_tool_registry_upsert error (non-blocking): %s", e) async def _batch_database_updates( self, @@ -543,9 +529,7 @@ class DBSpendUpdateWriter: try: if prisma_client is not None: # update user_ids = [user_id] - if ( - litellm.max_budget > 0 - ): # track global proxy budget, if user set max budget + if litellm.max_budget > 0: # track global proxy budget, if user set max budget user_ids.append(litellm_proxy_budget_name) for _id in user_ids: @@ -623,8 +607,7 @@ class DBSpendUpdateWriter: ) except Exception as e: spend_log_error( - "Spend tracking - failed to enqueue team spend update. " - "team_id=%s, response_cost=%s - %s", + "Spend tracking - failed to enqueue team spend update. team_id=%s, response_cost=%s - %s", team_id, response_cost, str(e), @@ -654,8 +637,7 @@ class DBSpendUpdateWriter: ) except Exception as e: spend_log_error( - "Spend tracking - failed to enqueue org spend update. " - "org_id=%s, response_cost=%s - %s", + "Spend tracking - failed to enqueue org spend update. org_id=%s, response_cost=%s - %s", org_id, response_cost, str(e), @@ -682,8 +664,7 @@ class DBSpendUpdateWriter: ) except Exception as e: spend_log_error( - "Spend tracking - failed to enqueue agent spend update. " - "agent_id=%s, response_cost=%s - %s", + "Spend tracking - failed to enqueue agent spend update. agent_id=%s, response_cost=%s - %s", agent_id, response_cost, str(e), @@ -714,9 +695,7 @@ class DBSpendUpdateWriter: if isinstance(request_tags, str): tags = safe_json_loads(request_tags, default=[]) if not tags: - verbose_proxy_logger.debug( - f"Failed to parse request_tags JSON: {request_tags}" - ) + verbose_proxy_logger.debug(f"Failed to parse request_tags JSON: {request_tags}") return elif isinstance(request_tags, list): tags = request_tags @@ -735,8 +714,7 @@ class DBSpendUpdateWriter: ) except Exception as e: spend_log_error( - "Spend tracking - failed to enqueue tag spend update. " - "request_tags=%s, response_cost=%s - %s", + "Spend tracking - failed to enqueue tag spend update. request_tags=%s, response_cost=%s - %s", request_tags, response_cost, str(e), @@ -762,9 +740,7 @@ class DBSpendUpdateWriter: async with prisma_client._spend_log_transactions_lock: prisma_client.spend_log_transactions.append(payload) else: - verbose_proxy_logger.debug( - "prisma_client is None. Skipping writing spend logs to db." - ) + verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return prisma_client @@ -849,42 +825,14 @@ class DBSpendUpdateWriter: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d", - len( - db_spend_update_transactions.get("key_list_transactions") - or {} - ), - len( - db_spend_update_transactions.get("user_list_transactions") - or {} - ), - len( - db_spend_update_transactions.get("team_list_transactions") - or {} - ), - len( - db_spend_update_transactions.get("org_list_transactions") - or {} - ), - len( - db_spend_update_transactions.get( - "end_user_list_transactions" - ) - or {} - ), - len( - db_spend_update_transactions.get( - "team_member_list_transactions" - ) - or {} - ), - len( - db_spend_update_transactions.get("tag_list_transactions") - or {} - ), - len( - db_spend_update_transactions.get("agent_list_transactions") - or {} - ), + len(db_spend_update_transactions.get("key_list_transactions") or {}), + len(db_spend_update_transactions.get("user_list_transactions") or {}), + len(db_spend_update_transactions.get("team_list_transactions") or {}), + len(db_spend_update_transactions.get("org_list_transactions") or {}), + len(db_spend_update_transactions.get("end_user_list_transactions") or {}), + len(db_spend_update_transactions.get("team_member_list_transactions") or {}), + len(db_spend_update_transactions.get("tag_list_transactions") or {}), + len(db_spend_update_transactions.get("agent_list_transactions") or {}), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -958,7 +906,9 @@ class DBSpendUpdateWriter: # Aggregate all in memory spend updates (key, user, end_user, team, team_member, org) and commit to db ################## Spend Update Transactions ################## - db_spend_update_transactions = await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + db_spend_update_transactions = ( + await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, @@ -1085,7 +1035,9 @@ class DBSpendUpdateWriter: ): verbose_proxy_logger.debug("acquired lock for daily tag spend updates") try: - daily_tag_spend_update_transactions = await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + daily_tag_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + ) if daily_tag_spend_update_transactions: await DBSpendUpdateWriter.update_daily_tag_spend( @@ -1118,9 +1070,7 @@ class DBSpendUpdateWriter: if items: await batch_upsert_tools(prisma_client=prisma_client, items=items) except Exception as e: - verbose_proxy_logger.debug( - "_flush_tool_discovery_queue error (non-blocking): %s", e - ) + verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e) async def _commit_spend_updates_to_db( self, @@ -1140,35 +1090,24 @@ class DBSpendUpdateWriter: ### UPDATE USER TABLE ### user_list_transactions = db_spend_update_transactions["user_list_transactions"] - verbose_proxy_logger.debug( - "User Spend transactions: {}".format(user_list_transactions) - ) - if ( - user_list_transactions is not None - and len(user_list_transactions.keys()) > 0 - ): + verbose_proxy_logger.debug("User Spend transactions: {}".format(user_list_transactions)) + if user_list_transactions is not None and len(user_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: + async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction: async with transaction.batch_() as batcher: # Sort by ID for consistent lock ordering across pods to prevent deadlocks. # batch_() issues statements sequentially within the tx, so iteration # order = lock acquisition order. - for user_id, response_cost in sorted( - user_list_transactions.items() - ): + for user_id, response_cost in sorted(user_list_transactions.items()): batcher.litellm_usertable.update_many( where={"user_id": user_id}, data={"spend": {"increment": response_cost}}, ) break except DB_CONNECTION_ERROR_TYPES as e: - if ( - i >= n_retry_times - ): # If we've reached the maximum number of retries + if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, @@ -1182,16 +1121,9 @@ class DBSpendUpdateWriter: ) ### UPDATE END-USER TABLE ### - end_user_list_transactions = db_spend_update_transactions[ - "end_user_list_transactions" - ] - verbose_proxy_logger.debug( - "End-User Spend transactions: {}".format(end_user_list_transactions) - ) - if ( - end_user_list_transactions is not None - and len(end_user_list_transactions.keys()) > 0 - ): + end_user_list_transactions = db_spend_update_transactions["end_user_list_transactions"] + verbose_proxy_logger.debug("End-User Spend transactions: {}".format(end_user_list_transactions)) + if end_user_list_transactions is not None and len(end_user_list_transactions.keys()) > 0: await ProxyUpdateSpend.update_end_user_spend( n_retry_times=n_retry_times, prisma_client=prisma_client, @@ -1200,21 +1132,15 @@ class DBSpendUpdateWriter: ) ### UPDATE KEY TABLE ### key_list_transactions = db_spend_update_transactions["key_list_transactions"] - verbose_proxy_logger.debug( - "KEY Spend transactions: {}".format(key_list_transactions) - ) + verbose_proxy_logger.debug("KEY Spend transactions: {}".format(key_list_transactions)) if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: + async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction: async with transaction.batch_() as batcher: # Sort by token for consistent lock ordering across pods to prevent deadlocks. - for token, response_cost in sorted( - key_list_transactions.items() - ): + for token, response_cost in sorted(key_list_transactions.items()): batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, data={ @@ -1224,9 +1150,7 @@ class DBSpendUpdateWriter: ) break except DB_CONNECTION_ERROR_TYPES as e: - if ( - i >= n_retry_times - ): # If we've reached the maximum number of retries + if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, @@ -1241,28 +1165,17 @@ class DBSpendUpdateWriter: ### UPDATE TEAM TABLE ### team_list_transactions = db_spend_update_transactions["team_list_transactions"] - verbose_proxy_logger.debug( - "Team Spend transactions: {}".format(team_list_transactions) - ) - if ( - team_list_transactions is not None - and len(team_list_transactions.keys()) > 0 - ): + verbose_proxy_logger.debug("Team Spend transactions: {}".format(team_list_transactions)) + if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: + async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction: async with transaction.batch_() as batcher: # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. - for team_id, response_cost in sorted( - team_list_transactions.items() - ): + for team_id, response_cost in sorted(team_list_transactions.items()): verbose_proxy_logger.debug( - "Updating spend for team id={} by {}".format( - team_id, response_cost - ) + "Updating spend for team id={} by {}".format(team_id, response_cost) ) batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists where={"team_id": team_id}, @@ -1270,9 +1183,7 @@ class DBSpendUpdateWriter: ) break except DB_CONNECTION_ERROR_TYPES as e: - if ( - i >= n_retry_times - ): # If we've reached the maximum number of retries + if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, @@ -1286,18 +1197,9 @@ class DBSpendUpdateWriter: ) ### UPDATE TEAM Membership TABLE with spend ### - team_member_list_transactions = db_spend_update_transactions[ - "team_member_list_transactions" - ] - verbose_proxy_logger.debug( - "Team Membership Spend transactions: {}".format( - team_member_list_transactions - ) - ) - if ( - team_member_list_transactions is not None - and len(team_member_list_transactions.keys()) > 0 - ): + team_member_list_transactions = db_spend_update_transactions["team_member_list_transactions"] + verbose_proxy_logger.debug("Team Membership Spend transactions: {}".format(team_member_list_transactions)) + if team_member_list_transactions is not None and len(team_member_list_transactions.keys()) > 0: # Track which team memberships will be updated for cache invalidation team_memberships_to_invalidate: List[tuple[str, str]] = [] for key in team_member_list_transactions.keys(): @@ -1309,15 +1211,11 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: + async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction: async with transaction.batch_() as batcher: # Sort by composite key for consistent lock ordering across pods to prevent deadlocks. # Key format "team_id::::user_id::" makes the string sort equivalent to sorting by (team_id, user_id). - for key, response_cost in sorted( - team_member_list_transactions.items() - ): + for key, response_cost in sorted(team_member_list_transactions.items()): # key is "team_id::::user_id::" team_id = key.split("::")[1] user_id = key.split("::")[3] @@ -1332,9 +1230,7 @@ class DBSpendUpdateWriter: # Transaction succeeded, break out of retry loop break except DB_CONNECTION_ERROR_TYPES as e: - if ( - i >= n_retry_times - ): # If we've reached the maximum number of retries + if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, @@ -1350,9 +1246,7 @@ class DBSpendUpdateWriter: # Invalidate cache for updated team memberships # This ensures budget checks read fresh spend data from the database if team_memberships_to_invalidate and proxy_logging_obj is not None: - user_api_key_cache = proxy_logging_obj.call_details.get( - "user_api_key_cache" - ) + user_api_key_cache = proxy_logging_obj.call_details.get("user_api_key_cache") if user_api_key_cache is not None: for user_id, team_id in team_memberships_to_invalidate: cache_key = "team_membership:{}:{}".format(user_id, team_id) @@ -1363,30 +1257,22 @@ class DBSpendUpdateWriter: ### UPDATE ORG TABLE ### org_list_transactions = db_spend_update_transactions["org_list_transactions"] - verbose_proxy_logger.debug( - "Org Spend transactions: {}".format(org_list_transactions) - ) + verbose_proxy_logger.debug("Org Spend transactions: {}".format(org_list_transactions)) if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: + async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction: async with transaction.batch_() as batcher: # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. - for org_id, response_cost in sorted( - org_list_transactions.items() - ): + for org_id, response_cost in sorted(org_list_transactions.items()): batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists where={"organization_id": org_id}, data={"spend": {"increment": response_cost}}, ) break except DB_CONNECTION_ERROR_TYPES as e: - if ( - i >= n_retry_times - ): # If we've reached the maximum number of retries + if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, @@ -1419,9 +1305,7 @@ class DBSpendUpdateWriter: ) ### UPDATE AGENT TABLE ### - agent_list_transactions = db_spend_update_transactions[ - "agent_list_transactions" - ] + agent_list_transactions = db_spend_update_transactions["agent_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Agent", transactions=agent_list_transactions, @@ -1461,14 +1345,10 @@ class DBSpendUpdateWriter: for i in range(n_retry_times + 1): start_time = time.time() try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: + async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction: async with transaction.batch_() as batcher: # Sort by entity_id for consistent lock ordering across pods to prevent deadlocks. - for entity_id, response_cost in sorted( - transactions.items() - ): + for entity_id, response_cost in sorted(transactions.items()): verbose_proxy_logger.debug( f"Updating spend for {entity_name} {where_field}={entity_id} by {response_cost}" ) @@ -1651,16 +1531,10 @@ class DBSpendUpdateWriter: "date": transaction["date"], "api_key": transaction["api_key"], "model": transaction["model"], - "custom_llm_provider": transaction.get( - "custom_llm_provider" - ) - or "", - "mcp_namespaced_tool_name": transaction.get( - "mcp_namespaced_tool_name" - ) - or "", - "endpoint": transaction.get("endpoint") + "custom_llm_provider": transaction.get("custom_llm_provider") or "", + "mcp_namespaced_tool_name": transaction.get("mcp_namespaced_tool_name") or "", + "endpoint": transaction.get("endpoint") or "", } } @@ -1674,100 +1548,55 @@ class DBSpendUpdateWriter: "api_key": transaction["api_key"], "model": transaction.get("model"), "model_group": transaction.get("model_group"), - "mcp_namespaced_tool_name": transaction.get( - "mcp_namespaced_tool_name" - ) - or "", - "custom_llm_provider": transaction.get( - "custom_llm_provider" - ), + "mcp_namespaced_tool_name": transaction.get("mcp_namespaced_tool_name") or "", + "custom_llm_provider": transaction.get("custom_llm_provider"), "endpoint": transaction.get("endpoint") or "", "prompt_tokens": transaction["prompt_tokens"], - "completion_tokens": transaction[ - "completion_tokens" - ], + "completion_tokens": transaction["completion_tokens"], "spend": transaction["spend"], "api_requests": transaction["api_requests"], - "successful_requests": transaction[ - "successful_requests" - ], - "failed_requests": transaction[ - "failed_requests" - ], + "successful_requests": transaction["successful_requests"], + "failed_requests": transaction["failed_requests"], } # Add cache-related fields if they exist if "cache_read_input_tokens" in transaction: - common_data["cache_read_input_tokens"] = ( - transaction.get( - "cache_read_input_tokens", 0 - ) + common_data["cache_read_input_tokens"] = transaction.get( + "cache_read_input_tokens", 0 ) if "cache_creation_input_tokens" in transaction: - common_data["cache_creation_input_tokens"] = ( - transaction.get( - "cache_creation_input_tokens", 0 - ) + common_data["cache_creation_input_tokens"] = transaction.get( + "cache_creation_input_tokens", 0 ) - if ( - entity_type == "tag" - and "request_id" in transaction - ): - common_data["request_id"] = transaction.get( - "request_id" - ) + if entity_type == "tag" and "request_id" in transaction: + common_data["request_id"] = transaction.get("request_id") # Create update data structure update_data = { - "prompt_tokens": { - "increment": transaction["prompt_tokens"] - }, - "completion_tokens": { - "increment": transaction[ - "completion_tokens" - ] - }, + "prompt_tokens": {"increment": transaction["prompt_tokens"]}, + "completion_tokens": {"increment": transaction["completion_tokens"]}, "spend": {"increment": transaction["spend"]}, - "api_requests": { - "increment": transaction["api_requests"] - }, - "successful_requests": { - "increment": transaction[ - "successful_requests" - ] - }, - "failed_requests": { - "increment": transaction["failed_requests"] - }, + "api_requests": {"increment": transaction["api_requests"]}, + "successful_requests": {"increment": transaction["successful_requests"]}, + "failed_requests": {"increment": transaction["failed_requests"]}, } # Add cache-related fields to update if they exist if "cache_read_input_tokens" in transaction: update_data["cache_read_input_tokens"] = { - "increment": transaction.get( - "cache_read_input_tokens", 0 - ) + "increment": transaction.get("cache_read_input_tokens", 0) } if "cache_creation_input_tokens" in transaction: update_data["cache_creation_input_tokens"] = { - "increment": transaction.get( - "cache_creation_input_tokens", 0 - ) + "increment": transaction.get("cache_creation_input_tokens", 0) } - if ( - entity_type == "tag" - and "request_id" in transaction - ): - update_data["request_id"] = transaction.get( - "request_id" - ) + if entity_type == "tag" and "request_id" in transaction: + update_data["request_id"] = transaction.get("request_id") # Add endpoint to update_data so existing rows get their endpoint field updated - update_data["endpoint"] = ( - transaction.get("endpoint") or "" - ) + update_data["endpoint"] = transaction.get("endpoint") or "" table.upsert( where=where_clause, @@ -1821,9 +1650,7 @@ class DBSpendUpdateWriter: if "transactions_to_process" in locals(): for key in transactions_to_process.keys(): # type: ignore daily_spend_transactions.pop(key, None) - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) + _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) @staticmethod async def update_daily_user_spend( @@ -1955,9 +1782,7 @@ class DBSpendUpdateWriter: self, payload: Union[dict, SpendLogsPayload], prisma_client: PrismaClient, - type: Literal[ - "user", "team", "org", "request_tags", "end_user", "agent" - ] = "user", + type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", ) -> Optional[BaseDailySpendTransaction]: common_expected_keys = ["startTime", "api_key"] if type == "user": @@ -1988,9 +1813,7 @@ class DBSpendUpdateWriter: return None elif "mcp_namespaced_tool_name" in payload: pass - elif "model" in payload and ( - "custom_llm_provider" not in payload or "model_group" not in payload - ): + elif "model" in payload and ("custom_llm_provider" not in payload or "model_group" not in payload): verbose_proxy_logger.debug( "Missing custom_llm_provider or model_group in payload, skipping from daily_user_spend_transactions" ) @@ -2051,27 +1874,19 @@ class DBSpendUpdateWriter: If key exists, update the transaction with the new spend and usage """ if prisma_client is None: - verbose_proxy_logger.debug( - "prisma_client is None. Skipping writing spend logs to db." - ) + verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return - base_daily_transaction = ( - await self._common_add_spend_log_transaction_to_daily_transaction( - payload, prisma_client, "user" - ) + base_daily_transaction = await self._common_add_spend_log_transaction_to_daily_transaction( + payload, prisma_client, "user" ) if base_daily_transaction is None: return endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{payload['user']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" - daily_transaction = DailyUserSpendTransaction( - user_id=payload["user"], **base_daily_transaction - ) - await self.daily_spend_update_queue.add_update( - update={daily_transaction_key: daily_transaction} - ) + daily_transaction = DailyUserSpendTransaction(user_id=payload["user"], **base_daily_transaction) + await self.daily_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_team_transaction( self, @@ -2079,32 +1894,22 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient] = None, ) -> None: if prisma_client is None: - verbose_proxy_logger.debug( - "prisma_client is None. Skipping writing spend logs to db." - ) + verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return - base_daily_transaction = ( - await self._common_add_spend_log_transaction_to_daily_transaction( - payload, prisma_client, "team" - ) + base_daily_transaction = await self._common_add_spend_log_transaction_to_daily_transaction( + payload, prisma_client, "team" ) if base_daily_transaction is None: return if payload["team_id"] is None: - verbose_proxy_logger.debug( - "team_id is None for request. Skipping incrementing team spend." - ) + verbose_proxy_logger.debug("team_id is None for request. Skipping incrementing team spend.") return endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{payload['team_id']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" - daily_transaction = DailyTeamSpendTransaction( - team_id=payload["team_id"], **base_daily_transaction - ) - await self.daily_team_spend_update_queue.add_update( - update={daily_transaction_key: daily_transaction} - ) + daily_transaction = DailyTeamSpendTransaction(team_id=payload["team_id"], **base_daily_transaction) + await self.daily_team_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_org_transaction( self, @@ -2113,15 +1918,11 @@ class DBSpendUpdateWriter: org_id: Optional[str] = None, ) -> None: if prisma_client is None: - verbose_proxy_logger.debug( - "prisma_client is None. Skipping writing spend logs to db." - ) + verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return if org_id is None: - verbose_proxy_logger.debug( - "organization_id is None for request. Skipping incrementing organization spend." - ) + verbose_proxy_logger.debug("organization_id is None for request. Skipping incrementing organization spend.") return payload_with_org = cast( @@ -2132,22 +1933,16 @@ class DBSpendUpdateWriter: }, ) - base_daily_transaction = ( - await self._common_add_spend_log_transaction_to_daily_transaction( - payload_with_org, prisma_client, "org" - ) + base_daily_transaction = await self._common_add_spend_log_transaction_to_daily_transaction( + payload_with_org, prisma_client, "org" ) if base_daily_transaction is None: return endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}_{endpoint_str}" - daily_transaction = DailyOrganizationSpendTransaction( - organization_id=org_id, **base_daily_transaction - ) - await self.daily_org_spend_update_queue.add_update( - update={daily_transaction_key: daily_transaction} - ) + daily_transaction = DailyOrganizationSpendTransaction(organization_id=org_id, **base_daily_transaction) + await self.daily_org_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_end_user_transaction( self, @@ -2155,16 +1950,12 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient] = None, ) -> None: if prisma_client is None: - verbose_proxy_logger.debug( - "prisma_client is None. Skipping writing spend logs to db." - ) + verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return end_user_id = payload.get("end_user") if end_user_id is None or end_user_id == "": - verbose_proxy_logger.debug( - "end_user is None or empty for request. Skipping incrementing end user spend." - ) + verbose_proxy_logger.debug("end_user is None or empty for request. Skipping incrementing end user spend.") return payload_with_end_user_id = cast( @@ -2175,22 +1966,16 @@ class DBSpendUpdateWriter: }, ) - base_daily_transaction = ( - await self._common_add_spend_log_transaction_to_daily_transaction( - payload_with_end_user_id, prisma_client, "end_user" - ) + base_daily_transaction = await self._common_add_spend_log_transaction_to_daily_transaction( + payload_with_end_user_id, prisma_client, "end_user" ) if base_daily_transaction is None: return endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{end_user_id}_{base_daily_transaction['date']}_{payload_with_end_user_id['api_key']}_{payload_with_end_user_id['model']}_{payload_with_end_user_id['custom_llm_provider']}_{endpoint_str}" - daily_transaction = DailyEndUserSpendTransaction( - end_user_id=end_user_id, **base_daily_transaction - ) - await self.daily_end_user_spend_update_queue.add_update( - update={daily_transaction_key: daily_transaction} - ) + daily_transaction = DailyEndUserSpendTransaction(end_user_id=end_user_id, **base_daily_transaction) + await self.daily_end_user_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_agent_transaction( self, @@ -2198,9 +1983,7 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient] = None, ) -> None: if prisma_client is None: - verbose_proxy_logger.debug( - "prisma_client is None. Skipping writing spend logs to db." - ) + verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return if payload["agent_id"] is None: return @@ -2211,21 +1994,15 @@ class DBSpendUpdateWriter: "agent_id": payload["agent_id"], }, ) - base_daily_transaction = ( - await self._common_add_spend_log_transaction_to_daily_transaction( - payload_with_agent_id, prisma_client, "agent" - ) + base_daily_transaction = await self._common_add_spend_log_transaction_to_daily_transaction( + payload_with_agent_id, prisma_client, "agent" ) if base_daily_transaction is None: return endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}" - daily_transaction = DailyAgentSpendTransaction( - agent_id=payload["agent_id"], **base_daily_transaction - ) - await self.daily_agent_spend_update_queue.add_update( - update={daily_transaction_key: daily_transaction} - ) + daily_transaction = DailyAgentSpendTransaction(agent_id=payload["agent_id"], **base_daily_transaction) + await self.daily_agent_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_tag_transaction( self, @@ -2233,22 +2010,16 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient] = None, ) -> None: if prisma_client is None: - verbose_proxy_logger.debug( - "prisma_client is None. Skipping writing spend logs to db." - ) + verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return - base_daily_transaction = ( - await self._common_add_spend_log_transaction_to_daily_transaction( - payload, prisma_client, "request_tags" - ) + base_daily_transaction = await self._common_add_spend_log_transaction_to_daily_transaction( + payload, prisma_client, "request_tags" ) if base_daily_transaction is None: return if payload["request_tags"] is None: - verbose_proxy_logger.debug( - "request_tags is None for request. Skipping incrementing tag spend." - ) + verbose_proxy_logger.debug("request_tags is None for request. Skipping incrementing tag spend.") return request_tags = [] @@ -2265,6 +2036,4 @@ class DBSpendUpdateWriter: tag=tag, **base_daily_transaction, request_id=payload["request_id"] ) - await self.daily_tag_spend_update_queue.add_update( - update={daily_transaction_key: daily_transaction} - ) + await self.daily_tag_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) diff --git a/litellm/proxy/db/db_transaction_queue/base_update_queue.py b/litellm/proxy/db/db_transaction_queue/base_update_queue.py index 7f1a7474690..fb6010e7d21 100644 --- a/litellm/proxy/db/db_transaction_queue/base_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/base_update_queue.py @@ -8,9 +8,7 @@ from typing import Optional from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceLogging -service_logger_obj = ( - ServiceLogging() -) # used for tracking metrics for In memory buffer, redis buffer, pod lock manager +service_logger_obj = ServiceLogging() # used for tracking metrics for In memory buffer, redis buffer, pod lock manager from litellm.constants import ( LITELLM_ASYNCIO_QUEUE_MAXSIZE, MAX_IN_MEMORY_QUEUE_FLUSH_COUNT, @@ -38,9 +36,7 @@ class BaseUpdateQueue: """Enqueue an update.""" verbose_proxy_logger.debug("Adding update to queue: %s", update) await self.update_queue.put(update) - await self._emit_new_item_added_to_queue_event( - queue_size=self.update_queue.qsize() - ) + await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize()) async def flush_all_updates_from_in_memory_queue(self): """Get all updates from the queue.""" @@ -48,9 +44,7 @@ class BaseUpdateQueue: while not self.update_queue.empty(): # Circuit breaker to ensure we're not stuck dequeuing updates. Protect CPU utilization if len(updates) >= MAX_IN_MEMORY_QUEUE_FLUSH_COUNT: - verbose_proxy_logger.debug( - "Max in memory queue flush count reached, stopping flush" - ) + verbose_proxy_logger.debug("Max in memory queue flush count reached, stopping flush") break updates.append(await self.update_queue.get()) return updates diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index a72ffad7e9b..19f8f4a94ad 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -54,8 +54,8 @@ class DailySpendUpdateQueue(BaseUpdateQueue): def __init__(self): super().__init__() - self.update_queue: asyncio.Queue[Dict[str, BaseDailySpendTransaction]] = ( - asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) + self.update_queue: asyncio.Queue[Dict[str, BaseDailySpendTransaction]] = asyncio.Queue( + maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE ) async def add_update(self, update: Dict[str, BaseDailySpendTransaction]): @@ -73,12 +73,8 @@ class DailySpendUpdateQueue(BaseUpdateQueue): Combine all updates in the queue into a single update. This is used to reduce the size of the in-memory queue. """ - updates: List[ - Dict[str, BaseDailySpendTransaction] - ] = await self.flush_all_updates_from_in_memory_queue() - aggregated_updates = self.get_aggregated_daily_spend_update_transactions( - updates - ) + updates: List[Dict[str, BaseDailySpendTransaction]] = await self.flush_all_updates_from_in_memory_queue() + aggregated_updates = self.get_aggregated_daily_spend_update_transactions(updates) await self.update_queue.put(aggregated_updates) async def flush_and_get_aggregated_daily_spend_update_transactions( @@ -92,9 +88,7 @@ class DailySpendUpdateQueue(BaseUpdateQueue): len(updates), ) aggregated_daily_spend_update_transactions = ( - DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( - updates - ) + DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(updates) ) verbose_proxy_logger.debug( "Aggregated daily spend update transactions: %s", @@ -107,22 +101,16 @@ class DailySpendUpdateQueue(BaseUpdateQueue): updates: List[Dict[str, BaseDailySpendTransaction]], ) -> Dict[str, BaseDailySpendTransaction]: """Aggregate updates by daily_transaction_key.""" - aggregated_daily_spend_update_transactions: Dict[ - str, BaseDailySpendTransaction - ] = {} + aggregated_daily_spend_update_transactions: Dict[str, BaseDailySpendTransaction] = {} for _update in updates: for _key, payload in _update.items(): if _key in aggregated_daily_spend_update_transactions: daily_transaction = aggregated_daily_spend_update_transactions[_key] daily_transaction["spend"] += payload["spend"] daily_transaction["prompt_tokens"] += payload["prompt_tokens"] - daily_transaction["completion_tokens"] += payload[ - "completion_tokens" - ] + daily_transaction["completion_tokens"] += payload["completion_tokens"] daily_transaction["api_requests"] += payload["api_requests"] - daily_transaction["successful_requests"] += payload[ - "successful_requests" - ] + daily_transaction["successful_requests"] += payload["successful_requests"] daily_transaction["failed_requests"] += payload["failed_requests"] # Add optional metrics cache_read_input_tokens and cache_creation_input_tokens diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 2cbc0646567..e04c2ba9a4e 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -106,9 +106,7 @@ end ) return False except Exception as e: - verbose_proxy_logger.error( - f"Error acquiring Redis lock for {cronjob_id}: {e}" - ) + verbose_proxy_logger.error(f"Error acquiring Redis lock for {cronjob_id}: {e}") return False async def release_lock( @@ -150,9 +148,7 @@ end cronjob_id, ) except Exception as e: - verbose_proxy_logger.error( - f"Error releasing Redis lock for {cronjob_id}: {e}" - ) + verbose_proxy_logger.error(f"Error releasing Redis lock for {cronjob_id}: {e}") async def _compare_and_delete_lock(self, lock_key: str) -> int: """ @@ -165,15 +161,11 @@ end if callable(script_register): try: if self._release_lock_script is None: - self._release_lock_script = script_register( - self._COMPARE_AND_DELETE_LOCK_SCRIPT - ) + self._release_lock_script = script_register(self._COMPARE_AND_DELETE_LOCK_SCRIPT) # acquire_lock stores the pod_id via async_set_cache, which # JSON-encodes the value; compare against the same encoding so # the Lua equality check matches and the lock is released - result = await self._release_lock_script( - keys=[lock_key], args=[json.dumps(self.pod_id)] - ) + result = await self._release_lock_script(keys=[lock_key], args=[json.dumps(self.pod_id)]) return int(result or 0) except Exception: # Lua execution failed (e.g. Redis restart cleared loaded scripts, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 6cbfb37396c..15d8c6c5d1e 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -74,8 +74,8 @@ class RedisUpdateBuffer: """ from litellm.proxy.proxy_server import general_settings - _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( - general_settings.get("use_redis_transaction_buffer", False) + _use_redis_transaction_buffer: Optional[Union[bool, str]] = general_settings.get( + "use_redis_transaction_buffer", False ) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -109,8 +109,7 @@ class RedisUpdateBuffer: values=list_of_transactions, ) verbose_proxy_logger.debug( - "Spend tracking - pushed spend updates to Redis buffer. " - "redis_key=%s, buffer_size=%s", + "Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s", redis_key, current_redis_buffer_size, ) @@ -120,8 +119,7 @@ class RedisUpdateBuffer: ) except Exception as e: verbose_proxy_logger.error( - "Spend tracking - failed to push spend updates to Redis (redis_key=%s). " - "Error: %s", + "Spend tracking - failed to push spend updates to Redis (redis_key=%s). Error: %s", redis_key, str(e), ) @@ -180,25 +178,29 @@ class RedisUpdateBuffer: ``` """ if self.redis_cache is None: - verbose_proxy_logger.debug( - "redis_cache is None, skipping store_in_memory_spend_updates_in_redis" - ) + verbose_proxy_logger.debug("redis_cache is None, skipping store_in_memory_spend_updates_in_redis") return # Get all transactions db_spend_update_transactions = await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() - daily_spend_update_transactions = await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - daily_team_spend_update_transactions = await daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - daily_org_spend_update_transactions = await daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - daily_end_user_spend_update_transactions = await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - daily_agent_spend_update_transactions = await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + daily_spend_update_transactions = ( + await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_team_spend_update_transactions = ( + await daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_org_spend_update_transactions = ( + await daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_end_user_spend_update_transactions = ( + await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_agent_spend_update_transactions = ( + await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) - verbose_proxy_logger.debug( - "ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions - ) - verbose_proxy_logger.debug( - "ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions - ) + verbose_proxy_logger.debug("ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions) + verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions) # Build a list of rpush operations, skipping empty/None transaction sets _queue_configs: List[Tuple[Any, str, ServiceTypes]] = [ @@ -293,18 +295,10 @@ class RedisUpdateBuffer: async def _restore_spend_updates_to_in_memory_queues( db_spend_update_transactions: Optional[DBSpendUpdateTransactions], daily_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]], - daily_team_spend_update_transactions: Optional[ - Dict[str, BaseDailySpendTransaction] - ], - daily_org_spend_update_transactions: Optional[ - Dict[str, BaseDailySpendTransaction] - ], - daily_end_user_spend_update_transactions: Optional[ - Dict[str, BaseDailySpendTransaction] - ], - daily_agent_spend_update_transactions: Optional[ - Dict[str, BaseDailySpendTransaction] - ], + daily_team_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]], + daily_org_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]], + daily_end_user_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]], + daily_agent_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]], spend_update_queue: SpendUpdateQueue, daily_spend_update_queue: DailySpendUpdateQueue, daily_team_spend_update_queue: DailySpendUpdateQueue, @@ -320,9 +314,7 @@ class RedisUpdateBuffer: because the source queues were already drained before the rpush. """ if db_spend_update_transactions is not None: - entity_entries: List[ - Tuple[Litellm_EntityType, Optional[Dict[str, float]]] - ] = [ + entity_entries: List[Tuple[Litellm_EntityType, Optional[Dict[str, float]]]] = [ ( Litellm_EntityType.USER, db_spend_update_transactions.get("user_list_transactions"), @@ -368,9 +360,7 @@ class RedisUpdateBuffer: ) ) - daily_pairs: List[ - Tuple[Optional[Dict[str, BaseDailySpendTransaction]], DailySpendUpdateQueue] - ] = [ + daily_pairs: List[Tuple[Optional[Dict[str, BaseDailySpendTransaction]], DailySpendUpdateQueue]] = [ (daily_spend_update_transactions, daily_spend_update_queue), (daily_team_spend_update_transactions, daily_team_spend_update_queue), (daily_org_spend_update_transactions, daily_org_spend_update_queue), @@ -496,9 +486,7 @@ class RedisUpdateBuffer: return None, None, None, None, None, None lpop_list: List[RedisPipelineLpopOperation] = [ - RedisPipelineLpopOperation( - key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT - ), + RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), RedisPipelineLpopOperation( key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, @@ -541,18 +529,14 @@ class RedisUpdateBuffer: daily_results.append(None) else: list_of_daily = [json.loads(t) for t in raw_results[slot]] # type: ignore - aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( - list_of_daily - ) + aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) daily_results.append(aggregated) return ( db_spend, cast(Optional[Dict[str, DailyUserSpendTransaction]], daily_results[0]), cast(Optional[Dict[str, DailyTeamSpendTransaction]], daily_results[1]), - cast( - Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2] - ), + cast(Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2]), cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]), cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]), ) @@ -564,7 +548,9 @@ class RedisUpdateBuffer: """ Flush in-memory daily tag spend updates and append them to Redis. """ - daily_tag_spend_update_transactions = await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + daily_tag_spend_update_transactions = ( + await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) await self._store_transactions_in_redis( transactions=daily_tag_spend_update_transactions, redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, @@ -585,9 +571,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - list_of_daily_spend_update_transactions = [ - json.loads(transaction) for transaction in list_of_transactions - ] + list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( Dict[str, DailyUserSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -609,9 +593,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - list_of_daily_spend_update_transactions = [ - json.loads(transaction) for transaction in list_of_transactions - ] + list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( Dict[str, DailyTeamSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -633,9 +615,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - list_of_daily_spend_update_transactions = [ - json.loads(transaction) for transaction in list_of_transactions - ] + list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( Dict[str, DailyOrganizationSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -657,9 +637,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - list_of_daily_spend_update_transactions = [ - json.loads(transaction) for transaction in list_of_transactions - ] + list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( Dict[str, DailyEndUserSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -681,9 +659,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - list_of_daily_spend_update_transactions = [ - json.loads(transaction) for transaction in list_of_transactions - ] + list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( Dict[str, DailyAgentSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -705,9 +681,7 @@ class RedisUpdateBuffer: ) if list_of_transactions is None: return None - list_of_daily_spend_update_transactions = [ - json.loads(transaction) for transaction in list_of_transactions - ] + list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( Dict[str, DailyTagSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index a4c23937b98..93f45fd1198 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -44,17 +44,13 @@ class SpendLogCleanup: pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager self.pod_lock_manager = pod_lock_manager - verbose_proxy_logger.info( - f"SpendLogCleanup initialized with batch size: {self.batch_size}" - ) + verbose_proxy_logger.info(f"SpendLogCleanup initialized with batch size: {self.batch_size}") def _should_delete_spend_logs(self) -> bool: """ Determines if logs should be deleted based on the max retention period in settings. """ - retention_setting = self.general_settings.get( - "maximum_spend_logs_retention_period" - ) + retention_setting = self.general_settings.get("maximum_spend_logs_retention_period") verbose_proxy_logger.info(f"Checking retention setting: {retention_setting}") if retention_setting is None: @@ -69,9 +65,7 @@ class SpendLogCleanup: ) retention_setting = f"{retention_setting}d" self.retention_seconds = duration_in_seconds(retention_setting) - verbose_proxy_logger.info( - f"Retention period set to {self.retention_seconds} seconds" - ) + verbose_proxy_logger.info(f"Retention period set to {self.retention_seconds} seconds") return True except ValueError as e: verbose_proxy_logger.warning( @@ -79,9 +73,7 @@ class SpendLogCleanup: ) return False - async def _delete_old_logs( - self, prisma_client: PrismaClient, cutoff_date: datetime - ) -> int: + async def _delete_old_logs(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int: """ Helper method to delete old logs in batches. Returns the total number of logs deleted. @@ -91,9 +83,7 @@ class SpendLogCleanup: consecutive_failures = 0 while True: if run_count > SPEND_LOG_RUN_LOOPS: - verbose_proxy_logger.info( - "Max logs deleted - 1,00,000, rest of the logs will be deleted in next run" - ) + verbose_proxy_logger.info("Max logs deleted - 1,00,000, rest of the logs will be deleted in next run") break # Step 1: Find logs and delete them in one go without fetching to application # Delete in batches, limited by self.batch_size @@ -126,10 +116,7 @@ class SpendLogCleanup: type(batch_exc).__name__, batch_exc, ) - if ( - consecutive_failures - >= SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES - ): + if consecutive_failures >= SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES: verbose_proxy_logger.error( "Aborting spend log cleanup after %d consecutive batch " "failures; total deleted before abort: %d", @@ -155,9 +142,7 @@ class SpendLogCleanup: verbose_proxy_logger.info(f"Deleted {deleted_count} logs in this batch") if deleted_count == 0: - verbose_proxy_logger.info( - f"No more logs to delete. Total deleted: {total_deleted}" - ) + verbose_proxy_logger.info(f"No more logs to delete. Total deleted: {total_deleted}") break total_deleted += deleted_count @@ -182,9 +167,7 @@ class SpendLogCleanup: return if self.retention_seconds is None: - verbose_proxy_logger.error( - "Retention seconds is None, cannot proceed with cleanup" - ) + verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup") return # If we have a pod lock manager, try to acquire the lock @@ -203,20 +186,14 @@ class SpendLogCleanup: verbose_proxy_logger.info("Another pod is already running cleanup") return - cutoff_date = datetime.now(timezone.utc) - timedelta( - seconds=float(self.retention_seconds) - ) - verbose_proxy_logger.info( - f"Removing logs older than {cutoff_date.isoformat()}" - ) + cutoff_date = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds)) + verbose_proxy_logger.info(f"Removing logs older than {cutoff_date.isoformat()}") if self.general_settings.get( "use_spend_logs_partitioning", False ) and await self.partition_manager.is_partitioned(prisma_client): await self.partition_manager.ensure_partitions(prisma_client) - dropped = await self.partition_manager.drop_partitions_older_than( - prisma_client, cutoff_date - ) + dropped = await self.partition_manager.drop_partitions_older_than(prisma_client, cutoff_date) verbose_proxy_logger.info( "Dropped %d expired spend-log partitions: %s", len(dropped), @@ -227,9 +204,7 @@ class SpendLogCleanup: # or in a partition that spans the cutoff, so retention must # also delete those stragglers row-wise. total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) - verbose_proxy_logger.info( - f"Deleted {total_deleted} expired logs not covered by dropped partitions" - ) + verbose_proxy_logger.info(f"Deleted {total_deleted} expired logs not covered by dropped partitions") else: total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) verbose_proxy_logger.info(f"Deleted {total_deleted} logs") @@ -245,12 +220,6 @@ class SpendLogCleanup: return # Return after error handling finally: # Only release the lock if it was actually acquired - if ( - lock_acquired - and self.pod_lock_manager - and self.pod_lock_manager.redis_cache - ): - await self.pod_lock_manager.release_lock( - cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME - ) + if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: + await self.pod_lock_manager.release_lock(cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME) verbose_proxy_logger.info("Released cleanup lock") diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py index eee0f862b4e..932675a6ac3 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -59,9 +59,7 @@ def partition_name(start: date) -> str: return f"{SPEND_LOGS_TABLE}_p{start.strftime('%Y%m%d')}" -def upcoming_partitions( - today: date, interval: PartitionInterval, ahead: int -) -> List[Tuple[str, date, date]]: +def upcoming_partitions(today: date, interval: PartitionInterval, ahead: int) -> List[Tuple[str, date, date]]: """ Specs (name, lower_inclusive, upper_exclusive) for the current period plus the next `ahead` periods, so writes always have a partition to land in. @@ -93,9 +91,7 @@ def parse_partition_upper_bound(bound_expr: str) -> Optional[datetime]: return None -def select_partitions_to_drop( - partitions: List[Tuple[str, Optional[datetime]]], cutoff: datetime -) -> List[str]: +def select_partitions_to_drop(partitions: List[Tuple[str, Optional[datetime]]], cutoff: datetime) -> List[str]: """ Names of partitions whose entire range is older than `cutoff` (upper bound <= cutoff). `cutoff` and the bounds are UTC-naive. Partitions without a @@ -112,8 +108,7 @@ class SpendLogsPartitionManager: ): if interval not in VALID_PARTITION_INTERVALS: verbose_proxy_logger.warning( - "Invalid SPEND_LOG_PARTITION_INTERVAL %r, falling back to 'day'. " - "Supported values: %s", + "Invalid SPEND_LOG_PARTITION_INTERVAL %r, falling back to 'day'. Supported values: %s", interval, sorted(VALID_PARTITION_INTERVALS), ) @@ -163,14 +158,10 @@ class SpendLogsPartitionManager: ) ensured.append(name) except Exception as e: - verbose_proxy_logger.warning( - "Failed to ensure spend-log partition %s: %s", name, e - ) + verbose_proxy_logger.warning("Failed to ensure spend-log partition %s: %s", name, e) return ensured - async def _list_partitions( - self, prisma_client - ) -> List[Tuple[str, Optional[datetime]]]: + async def _list_partitions(self, prisma_client) -> List[Tuple[str, Optional[datetime]]]: rows = await prisma_client.db.query_raw( """ SELECT c.relname AS name, @@ -184,14 +175,9 @@ class SpendLogsPartitionManager: """, SPEND_LOGS_TABLE, ) - return [ - (row["name"], parse_partition_upper_bound(row.get("bound") or "")) - for row in rows - ] + return [(row["name"], parse_partition_upper_bound(row.get("bound") or "")) for row in rows] - async def drop_partitions_older_than( - self, prisma_client, cutoff: datetime - ) -> List[str]: + async def drop_partitions_older_than(self, prisma_client, cutoff: datetime) -> List[str]: """DROP every partition whose whole range is older than `cutoff`.""" cutoff_naive = cutoff.astimezone(timezone.utc).replace(tzinfo=None) partitions = await self._list_partitions(prisma_client) @@ -202,7 +188,5 @@ class SpendLogsPartitionManager: await prisma_client.db.execute_raw(f'DROP TABLE IF EXISTS "{name}"') dropped.append(name) except Exception as e: - verbose_proxy_logger.warning( - "Failed to drop spend-log partition %s: %s", name, e - ) + verbose_proxy_logger.warning("Failed to drop spend-log partition %s: %s", name, e) return dropped diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 727e8dc1d5a..0689fc00b02 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -22,9 +22,7 @@ class SpendUpdateQueue(BaseUpdateQueue): def __init__(self): super().__init__() - self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue( - maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE - ) + self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) async def flush_and_get_aggregated_db_spend_update_transactions( self, @@ -53,9 +51,7 @@ class SpendUpdateQueue(BaseUpdateQueue): async def aggregate_queue_updates(self): """Concatenate all updates in the queue to reduce the size of in-memory queue""" - updates: List[ - SpendUpdateQueueItem - ] = await self.flush_all_updates_from_in_memory_queue() + updates: List[SpendUpdateQueueItem] = await self.flush_all_updates_from_in_memory_queue() aggregated_updates = self._get_aggregated_spend_update_queue_item(updates) for update in aggregated_updates: await self.update_queue.put(update) @@ -127,9 +123,7 @@ class SpendUpdateQueue(BaseUpdateQueue): for _key, update in _in_memory_map.items(): aggregated_spend_updates.append(update) - verbose_proxy_logger.debug( - "Aggregated spend updates: %s", aggregated_spend_updates - ) + verbose_proxy_logger.debug("Aggregated spend updates: %s", aggregated_spend_updates) return aggregated_spend_updates def get_aggregated_db_spend_update_transactions( @@ -182,37 +176,21 @@ class SpendUpdateQueue(BaseUpdateQueue): # Type-safe access using if/elif statements if dict_key == "user_list_transactions": - transactions_dict = db_spend_update_transactions[ - "user_list_transactions" - ] + transactions_dict = db_spend_update_transactions["user_list_transactions"] elif dict_key == "end_user_list_transactions": - transactions_dict = db_spend_update_transactions[ - "end_user_list_transactions" - ] + transactions_dict = db_spend_update_transactions["end_user_list_transactions"] elif dict_key == "key_list_transactions": - transactions_dict = db_spend_update_transactions[ - "key_list_transactions" - ] + transactions_dict = db_spend_update_transactions["key_list_transactions"] elif dict_key == "team_list_transactions": - transactions_dict = db_spend_update_transactions[ - "team_list_transactions" - ] + transactions_dict = db_spend_update_transactions["team_list_transactions"] elif dict_key == "team_member_list_transactions": - transactions_dict = db_spend_update_transactions[ - "team_member_list_transactions" - ] + transactions_dict = db_spend_update_transactions["team_member_list_transactions"] elif dict_key == "org_list_transactions": - transactions_dict = db_spend_update_transactions[ - "org_list_transactions" - ] + transactions_dict = db_spend_update_transactions["org_list_transactions"] elif dict_key == "tag_list_transactions": - transactions_dict = db_spend_update_transactions[ - "tag_list_transactions" - ] + transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": - transactions_dict = db_spend_update_transactions[ - "agent_list_transactions" - ] + transactions_dict = db_spend_update_transactions["agent_list_transactions"] else: continue diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index 16a3ada40f2..5d23bcaa944 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -34,9 +34,7 @@ class ToolDiscoveryQueue: if not tool_name: return if tool_name in self._seen_tool_names: - verbose_proxy_logger.debug( - "ToolDiscoveryQueue: skipping already-seen tool %s", tool_name - ) + verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name) return self._seen_tool_names.add(tool_name) self._pending.append(item) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index ae2307658dd..4f220916a4a 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -96,46 +96,26 @@ class DatabaseURLSettings(BaseSettings): database_url: str | None = Field(default=None, validation_alias="DATABASE_URL") direct_url: str | None = Field(default=None, validation_alias="DIRECT_URL") database_host: str | None = Field(default=None, validation_alias="DATABASE_HOST") - database_port: str = Field( - default=_DEFAULT_PG_PORT, validation_alias="DATABASE_PORT" - ) + database_port: str = Field(default=_DEFAULT_PG_PORT, validation_alias="DATABASE_PORT") database_user: str | None = Field( default=None, validation_alias=AliasChoices("DATABASE_USER", "DATABASE_USERNAME"), ) database_name: str | None = Field(default=None, validation_alias="DATABASE_NAME") - database_schema: str | None = Field( - default=None, validation_alias="DATABASE_SCHEMA" - ) - database_password: str | None = Field( - default=None, validation_alias="DATABASE_PASSWORD" - ) + database_schema: str | None = Field(default=None, validation_alias="DATABASE_SCHEMA") + database_password: str | None = Field(default=None, validation_alias="DATABASE_PASSWORD") # Read replica - database_url_read_replica: str | None = Field( - default=None, validation_alias="DATABASE_URL_READ_REPLICA" - ) - database_host_read_replica: str | None = Field( - default=None, validation_alias="DATABASE_HOST_READ_REPLICA" - ) - database_port_read_replica: str | None = Field( - default=None, validation_alias="DATABASE_PORT_READ_REPLICA" - ) + database_url_read_replica: str | None = Field(default=None, validation_alias="DATABASE_URL_READ_REPLICA") + database_host_read_replica: str | None = Field(default=None, validation_alias="DATABASE_HOST_READ_REPLICA") + database_port_read_replica: str | None = Field(default=None, validation_alias="DATABASE_PORT_READ_REPLICA") database_user_read_replica: str | None = Field( default=None, - validation_alias=AliasChoices( - "DATABASE_USER_READ_REPLICA", "DATABASE_USERNAME_READ_REPLICA" - ), - ) - database_name_read_replica: str | None = Field( - default=None, validation_alias="DATABASE_NAME_READ_REPLICA" - ) - database_schema_read_replica: str | None = Field( - default=None, validation_alias="DATABASE_SCHEMA_READ_REPLICA" - ) - database_password_read_replica: str | None = Field( - default=None, validation_alias="DATABASE_PASSWORD_READ_REPLICA" + validation_alias=AliasChoices("DATABASE_USER_READ_REPLICA", "DATABASE_USERNAME_READ_REPLICA"), ) + database_name_read_replica: str | None = Field(default=None, validation_alias="DATABASE_NAME_READ_REPLICA") + database_schema_read_replica: str | None = Field(default=None, validation_alias="DATABASE_SCHEMA_READ_REPLICA") + database_password_read_replica: str | None = Field(default=None, validation_alias="DATABASE_PASSWORD_READ_REPLICA") @classmethod def from_env(cls) -> "DatabaseURLSettings": @@ -170,9 +150,7 @@ class DatabaseURLSettings(BaseSettings): name = cast(str, self.database_name) # IAM token is already URL-quoted by generate_iam_auth_token; # user/name embedded raw (parity with proxy_cli.py / IAMEndpoint). - token = rds_iam_token.generate_iam_auth_token( - db_host=host, db_port=self.database_port, db_user=user - ) + token = rds_iam_token.generate_iam_auth_token(db_host=host, db_port=self.database_port, db_user=user) url = f"postgresql://{user}:{token}@{host}:{self.database_port}/{name}" if self.database_schema: url += f"?schema={self.database_schema}" @@ -230,9 +208,7 @@ class DatabaseURLSettings(BaseSettings): ) user = cast(str, user) name = cast(str, name) - token = rds_iam_token.generate_iam_auth_token( - db_host=host, db_port=port, db_user=user - ) + token = rds_iam_token.generate_iam_auth_token(db_host=host, db_port=port, db_user=user) url = f"postgresql://{user}:{token}@{host}:{port}/{name}" if schema: url += f"?schema={schema}" diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 57ebb9678cb..6367c34341f 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -39,9 +39,7 @@ class DynamoDBWrapper(CustomDB): def set_env_vars_based_on_arn(self): if self.database_arguments.aws_role_name is None: return - verbose_proxy_logger.debug( - f"DynamoDB: setting env vars based on arn={self.database_arguments.aws_role_name}" - ) + verbose_proxy_logger.debug(f"DynamoDB: setting env vars based on arn={self.database_arguments.aws_role_name}") import os import boto3 @@ -65,9 +63,7 @@ class DynamoDBWrapper(CustomDB): aws_secret_access_key = assumed_role["Credentials"]["SecretAccessKey"] aws_session_token = assumed_role["Credentials"]["SessionToken"] - verbose_proxy_logger.debug( - f"Got STS assumed Role, aws_access_key_id={aws_access_key_id}" - ) + verbose_proxy_logger.debug(f"Got STS assumed Role, aws_access_key_id={aws_access_key_id}") # set these in the env so aiodynamo can use them os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index c500e727595..48066945131 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -305,9 +305,7 @@ async def call_with_db_reconnect_retry( ( lock_timeout_seconds if lock_timeout_seconds is not None - else getattr( - prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None - ) + else getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None) ), _DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS, ) @@ -332,8 +330,7 @@ async def call_with_db_reconnect_retry( ) except Exception as reconnect_exc: verbose_proxy_logger.warning( - "DB reconnect attempt raised; preserving original transport error. " - "reason=%s reconnect_error=%s", + "DB reconnect attempt raised; preserving original transport error. reason=%s reconnect_error=%s", reason, reconnect_exc, ) diff --git a/litellm/proxy/db/log_db_metrics.py b/litellm/proxy/db/log_db_metrics.py index 837e94f1a85..ba7d6d6fbee 100644 --- a/litellm/proxy/db/log_db_metrics.py +++ b/litellm/proxy/db/log_db_metrics.py @@ -69,9 +69,7 @@ def log_db_metrics(func): args is not None and len(args) > 1 and isinstance(args[1], dict) ): passed_kwargs = args[1] - parent_otel_span = _get_parent_otel_span_from_kwargs( - kwargs=passed_kwargs - ) + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=passed_kwargs) if parent_otel_span is not None: # No metadata dump: identity rides on Baggage, and the full # request metadata (auth blob, response headers, tokens) must diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index d133ddc9d1a..4042755f80d 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -260,9 +260,7 @@ class PrismaWrapper: return self.FALLBACK_REFRESH_INTERVAL_SECONDS # Calculate when we should refresh (expiration - buffer) - refresh_at = expiration_time - timedelta( - seconds=self.TOKEN_REFRESH_BUFFER_SECONDS - ) + refresh_at = expiration_time - timedelta(seconds=self.TOKEN_REFRESH_BUFFER_SECONDS) # How long until refresh time? now = datetime.utcnow() @@ -281,9 +279,7 @@ class PrismaWrapper: if expiration_time is None: # If we can't parse the token, assume it's expired to trigger refresh - verbose_proxy_logger.debug( - "Could not parse token expiration, treating as expired" - ) + verbose_proxy_logger.debug("Could not parse token expiration, treating as expired") return True return datetime.utcnow() > expiration_time @@ -303,9 +299,7 @@ class PrismaWrapper: if self._iam_endpoint is not None: endpoint = self._iam_endpoint - token = generate_iam_auth_token( - db_host=endpoint.host, db_port=endpoint.port, db_user=endpoint.user - ) + token = generate_iam_auth_token(db_host=endpoint.host, db_port=endpoint.port, db_user=endpoint.user) _db_url = endpoint.build_url(token) else: db_host = os.getenv("DATABASE_HOST") @@ -317,9 +311,7 @@ class PrismaWrapper: db_name = os.getenv("DATABASE_NAME") db_schema = os.getenv("DATABASE_SCHEMA") - token = generate_iam_auth_token( - db_host=db_host, db_port=db_port, db_user=db_user - ) + token = generate_iam_auth_token(db_host=db_host, db_port=db_port, db_user=db_user) _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" if db_schema: @@ -381,13 +373,9 @@ class PrismaWrapper: """ from prisma import Prisma # type: ignore - if ( - expected_generation is not None - and expected_generation != self._engine_generation - ): + if expected_generation is not None and expected_generation != self._engine_generation: verbose_proxy_logger.info( - "%sSkipping Prisma client recreate: engine already replaced " - "(generation %s != expected %s).", + "%sSkipping Prisma client recreate: engine already replaced (generation %s != expected %s).", self._log_prefix, self._engine_generation, expected_generation, @@ -437,9 +425,7 @@ class PrismaWrapper: Prisma client connection is established. """ if not self.iam_token_db_auth: - verbose_proxy_logger.debug( - "IAM token auth not enabled, skipping token refresh task" - ) + verbose_proxy_logger.debug("IAM token auth not enabled, skipping token refresh task") return if self._token_refresh_task is not None: @@ -467,9 +453,7 @@ class PrismaWrapper: except asyncio.CancelledError: pass self._token_refresh_task = None - verbose_proxy_logger.info( - "%sStopped RDS IAM token refresh background task", self._log_prefix - ) + verbose_proxy_logger.info("%sStopped RDS IAM token refresh background task", self._log_prefix) async def _token_refresh_loop(self) -> None: """ @@ -497,15 +481,11 @@ class PrismaWrapper: await asyncio.sleep(sleep_seconds) # Refresh the token - verbose_proxy_logger.info( - "%sProactively refreshing RDS IAM token...", self._log_prefix - ) + verbose_proxy_logger.info("%sProactively refreshing RDS IAM token...", self._log_prefix) await self._safe_refresh_token() except asyncio.CancelledError: - verbose_proxy_logger.info( - "%sRDS IAM token refresh loop cancelled", self._log_prefix - ) + verbose_proxy_logger.info("%sRDS IAM token refresh loop cancelled", self._log_prefix) break except Exception as e: verbose_proxy_logger.error( @@ -644,9 +624,7 @@ class PrismaManager: return dname @staticmethod - def setup_database( - use_migrate: bool = False, use_v2_resolver: bool = False - ) -> bool: + def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool: """ Set up the database using either prisma migrate or prisma db push @@ -669,9 +647,7 @@ class PrismaManager: try: from litellm_proxy_extras.utils import ProxyExtrasDBManager except ImportError as e: - verbose_proxy_logger.error( - f"\033[1;31mLiteLLM: Failed to import proxy extras. Got {e}\033[0m" - ) + verbose_proxy_logger.error(f"\033[1;31mLiteLLM: Failed to import proxy extras. Got {e}\033[0m") return False prisma_dir = PrismaManager._get_prisma_dir() @@ -699,14 +675,8 @@ class PrismaManager: time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt - retry_msg = ( - f" Retrying... ({attempts_left} attempts left)" - if attempts_left > 0 - else "" - ) - verbose_proxy_logger.warning( - f"The process failed to execute. Details: {e}.{retry_msg}" - ) + retry_msg = f" Retrying... ({attempts_left} attempts left)" if attempts_left > 0 else "" + verbose_proxy_logger.warning(f"The process failed to execute. Details: {e}.{retry_msg}") time.sleep(random.randrange(5, 15)) finally: os.chdir(original_dir) diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index d752c6c5718..7ae60121c6f 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -194,18 +194,12 @@ class RoutingPrismaWrapper: if self._reader.iam_token_db_auth: new_reader_url = self._reader.get_rds_iam_token() if not new_reader_url: - raise RuntimeError( - "Failed to generate fresh IAM token for read replica" - ) - await self._reader.recreate_prisma_client( - new_reader_url, http_client=http_client - ) + raise RuntimeError("Failed to generate fresh IAM token for read replica") + await self._reader.recreate_prisma_client(new_reader_url, http_client=http_client) return reader_url = os.getenv("DATABASE_URL_READ_REPLICA", "") if not reader_url: - raise RuntimeError( - "DATABASE_URL_READ_REPLICA not set; cannot recreate read replica client" - ) + raise RuntimeError("DATABASE_URL_READ_REPLICA not set; cannot recreate read replica client") await self._reader.recreate_prisma_client(reader_url, http_client=http_client) def __getattr__(self, name: str) -> Any: @@ -216,11 +210,7 @@ class RoutingPrismaWrapper: # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / # tx are callables and stay on the writer untouched. - if ( - not callable(writer_attr) - and hasattr(writer_attr, "find_many") - and hasattr(writer_attr, "create") - ): + if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"): try: reader_attr = getattr(self._reader, name) except AttributeError: diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 2226aeb4b0a..079cbd163dc 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -72,9 +72,7 @@ class SpendCounterReseed: return lock @staticmethod - async def from_db( - prisma_client: Optional["PrismaClient"], counter_key: str - ) -> Optional[float]: + async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> Optional[float]: """ Read the authoritative spend for a counter from the DB. @@ -93,9 +91,7 @@ class SpendCounterReseed: try: if counter_key.startswith("spend:key:"): token = counter_key[len("spend:key:") :] - row = await VerificationTokenRepository( - prisma_client - ).table.find_unique(where={"token": token}) + row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) elif counter_key.startswith("spend:team_member:"): suffix = counter_key[len("spend:team_member:") :] if ":" not in suffix: @@ -106,29 +102,21 @@ class SpendCounterReseed: ) elif counter_key.startswith("spend:team:"): team_id = counter_key[len("spend:team:") :] - row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) elif counter_key.startswith("spend:user:"): user_id = counter_key[len("spend:user:") :] - row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) elif counter_key.startswith("spend:end_user:"): return None elif counter_key.startswith("spend:tag:"): return None elif counter_key.startswith("spend:org:"): org_id = counter_key[len("spend:org:") :] - row = await OrganizationRepository(prisma_client).table.find_unique( - where={"organization_id": org_id} - ) + row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id}) else: return None except Exception: - verbose_proxy_logger.exception( - "SpendCounterReseed.from_db: failed for %s", counter_key - ) + verbose_proxy_logger.exception("SpendCounterReseed.from_db: failed for %s", counter_key) return None if row is None: return None @@ -171,9 +159,7 @@ class SpendCounterReseed: redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: - val = await spend_counter_cache.redis_cache.async_get_cache( - key=counter_key - ) + val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: return float(val) redis_clean_miss = True @@ -204,20 +190,14 @@ class SpendCounterReseed: if seeded: current_value = float(db_spend) else: - cached = await spend_counter_cache.redis_cache.async_get_cache( - key=counter_key - ) - current_value = ( - float(cached) if cached is not None else float(db_spend) - ) + cached = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) + current_value = float(cached) if cached is not None else float(db_spend) spend_counter_cache.in_memory_cache.set_cache( key=counter_key, value=current_value, ) else: - await spend_counter_cache.async_increment_cache( - key=counter_key, value=db_spend, refresh_ttl=True - ) + await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -269,16 +249,8 @@ class SpendCounterReseed: if not response: return 0.0 first_row = response[0] - sum_row = ( - first_row.get("_sum") - if isinstance(first_row, dict) - else getattr(first_row, "_sum", None) - ) - spend = ( - sum_row.get("spend") - if isinstance(sum_row, dict) - else getattr(sum_row, "spend", None) - ) + sum_row = first_row.get("_sum") if isinstance(first_row, dict) else getattr(first_row, "_sum", None) + spend = sum_row.get("spend") if isinstance(sum_row, dict) else getattr(sum_row, "spend", None) return float(spend or 0.0) @staticmethod @@ -295,9 +267,7 @@ class SpendCounterReseed: redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: - val = await spend_counter_cache.redis_cache.async_get_cache( - key=counter_key - ) + val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: return float(val) redis_clean_miss = True @@ -326,17 +296,11 @@ class SpendCounterReseed: if seeded: current_value = window_spend else: - current_cached_value = ( - await spend_counter_cache.redis_cache.async_get_cache( - key=counter_key - ) - ) + current_cached_value = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if current_cached_value is None: - current_value = ( - await spend_counter_cache.redis_cache.async_increment( - key=counter_key, - value=window_spend, - ) + current_value = await spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=window_spend, ) else: current_value = float(current_cached_value) @@ -345,9 +309,7 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache( - key=counter_key, value=window_spend - ) + await spend_counter_cache.async_increment_cache(key=counter_key, value=window_spend) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced_window: failed to warm counter %s", diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 77c06a465f4..80036e235f7 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -42,29 +42,19 @@ def _parse_tool_names_from_payload(payload: Dict[str, Any]) -> Set[str]: # Response: OpenAI-style tool_calls[].function.name or choices[0].message.tool_calls response_raw = payload.get("response") if response_raw: - response_obj = ( - safe_json_loads(response_raw, default=None) - if isinstance(response_raw, str) - else response_raw - ) + response_obj = safe_json_loads(response_raw, default=None) if isinstance(response_raw, str) else response_raw if isinstance(response_obj, dict): _add_tool_calls_to_set(response_obj.get("tool_calls"), tool_names) choices = response_obj.get("choices") if isinstance(choices, list) and choices: - msg = ( - choices[0].get("message") if isinstance(choices[0], dict) else None - ) + msg = choices[0].get("message") if isinstance(choices[0], dict) else None if isinstance(msg, dict): _add_tool_calls_to_set(msg.get("tool_calls"), tool_names) # Request body: tools[].function.name request_raw = payload.get("proxy_server_request") if request_raw: - request_obj = ( - safe_json_loads(request_raw, default=None) - if isinstance(request_raw, str) - else request_raw - ) + request_obj = safe_json_loads(request_raw, default=None) if isinstance(request_raw, str) else request_raw if isinstance(request_obj, dict): body = request_obj.get("body", request_obj) if isinstance(body, dict): @@ -147,6 +137,4 @@ async def process_spend_logs_tool_usage( skip_duplicates=True, ) except Exception as e: - verbose_proxy_logger.warning( - "Tool usage tracking (SpendLogToolIndex) failed (non-fatal): %s", e - ) + verbose_proxy_logger.warning("Tool usage tracking (SpendLogToolIndex) failed (non-fatal): %s", e) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 08bc8944b92..5f7de772a2c 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -121,13 +121,9 @@ async def batch_upsert_tools( }, }, ) - verbose_proxy_logger.debug( - "tool_registry_writer: upserted %d tool(s)", len(data) - ) + verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(data)) except Exception as e: - verbose_proxy_logger.error( - "tool_registry_writer batch_upsert_tools error: %s", e - ) + verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e) async def list_tools( @@ -204,9 +200,7 @@ async def update_tool_policy( ) return await get_tool(prisma_client, tool_name) except Exception as e: - verbose_proxy_logger.error( - "tool_registry_writer update_tool_policy error: %s", e - ) + verbose_proxy_logger.error("tool_registry_writer update_tool_policy error: %s", e) return None @@ -231,9 +225,7 @@ async def get_tools_by_names( for row in rows } except Exception as e: - verbose_proxy_logger.error( - "tool_registry_writer get_tools_by_names error: %s", e - ) + verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e) return {} @@ -286,9 +278,7 @@ async def list_overrides_for_tool( ) return out except Exception as e: - verbose_proxy_logger.error( - "tool_registry_writer list_overrides_for_tool error: %s", e - ) + verbose_proxy_logger.error("tool_registry_writer list_overrides_for_tool error: %s", e) return [] @@ -316,12 +306,10 @@ class ToolPolicyRegistry: reason="sync_tool_policy_from_db_tools_lookup_failure", ) self._tool_input_policies = { - row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted" - for row in tools + row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted" for row in tools } self._tool_output_policies = { - row.tool_name: getattr(row, "output_policy", "untrusted") or "untrusted" - for row in tools + row.tool_name: getattr(row, "output_policy", "untrusted") or "untrusted" for row in tools } perms = await call_with_db_reconnect_retry( @@ -343,9 +331,7 @@ class ToolPolicyRegistry: len(self._blocked_tools_by_op_id), ) except Exception as e: - verbose_proxy_logger.exception( - "ToolPolicyRegistry sync_tool_policy_from_db error: %s", e - ) + verbose_proxy_logger.exception("ToolPolicyRegistry sync_tool_policy_from_db error: %s", e) raise def get_input_policy(self, tool_name: str) -> str: @@ -414,9 +400,7 @@ async def add_tool_to_object_permission_blocked( ) return True except Exception as e: - verbose_proxy_logger.error( - "tool_registry_writer add_tool_to_object_permission_blocked error: %s", e - ) + verbose_proxy_logger.error("tool_registry_writer add_tool_to_object_permission_blocked error: %s", e) return False diff --git a/litellm/proxy/dd_span_tagger.py b/litellm/proxy/dd_span_tagger.py index 7bbfe50a01b..08b7d928d0e 100644 --- a/litellm/proxy/dd_span_tagger.py +++ b/litellm/proxy/dd_span_tagger.py @@ -48,9 +48,7 @@ class DDSpanTagger: """ try: if user_api_key_dict.key_alias: - set_active_span_tag( - "litellm.key_alias", str(user_api_key_dict.key_alias) - ) + set_active_span_tag("litellm.key_alias", str(user_api_key_dict.key_alias)) if user_api_key_dict.token: set_active_span_tag("litellm.key_hash", str(user_api_key_dict.token)) if requested_model: diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 0e7e67aa37f..ba2e0a39de3 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -11,9 +11,7 @@ router = APIRouter() @router.get("/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) -@router.get( - "/litellm/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints -) # if mounted at root path +@router.get("/litellm/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) # if mounted at root path async def get_ui_config(): from litellm.proxy.auth.auth_utils import _has_user_setup_sso from litellm.proxy.utils import get_proxy_base_url, get_server_root_path diff --git a/litellm/proxy/example_config_yaml/custom_callbacks.py b/litellm/proxy/example_config_yaml/custom_callbacks.py index 9e86f931537..cf62417c44b 100644 --- a/litellm/proxy/example_config_yaml/custom_callbacks.py +++ b/litellm/proxy/example_config_yaml/custom_callbacks.py @@ -4,9 +4,7 @@ import traceback # this file is to test litellm/proxy -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import inspect @@ -28,11 +26,7 @@ class MyCustomHandler(CustomLogger): print_verbose(f"{blue_color_code}Initialized LiteLLM custom logger") try: print_verbose("Logger Initialized with following methods:") - methods = [ - method - for method in dir(self) - if inspect.ismethod(getattr(self, method)) - ] + methods = [method for method in dir(self) if inspect.ismethod(getattr(self, method))] # Pretty print_verbose the methods for method in methods: diff --git a/litellm/proxy/example_config_yaml/custom_callbacks1.py b/litellm/proxy/example_config_yaml/custom_callbacks1.py index b261f5a83af..ae4e90beff2 100644 --- a/litellm/proxy/example_config_yaml/custom_callbacks1.py +++ b/litellm/proxy/example_config_yaml/custom_callbacks1.py @@ -8,9 +8,7 @@ from litellm.types.utils import CallTypesLiteral # This file includes the custom callbacks for LiteLLM Proxy # Once defined, these can be passed in proxy_config.yaml -class MyCustomHandler( - CustomLogger -): # https://docs.litellm.ai/docs/observability/custom_callback#callback-class +class MyCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/observability/custom_callback#callback-class # Class variables or attributes def __init__(self): pass diff --git a/litellm/proxy/example_config_yaml/custom_guardrail.py b/litellm/proxy/example_config_yaml/custom_guardrail.py index 84d404d1e65..a755390743e 100644 --- a/litellm/proxy/example_config_yaml/custom_guardrail.py +++ b/litellm/proxy/example_config_yaml/custom_guardrail.py @@ -23,9 +23,7 @@ class GuardrailForLBTestingA(CustomGuardrail): call_type: CallTypesLiteral, ) -> Optional[Union[Exception, str, dict]]: guardrail_lb_call_count["A"] += 1 - verbose_proxy_logger.info( - f"GuardrailForLBTestingA called. Total A calls: {guardrail_lb_call_count['A']}" - ) + verbose_proxy_logger.info(f"GuardrailForLBTestingA called. Total A calls: {guardrail_lb_call_count['A']}") return data @@ -40,9 +38,7 @@ class GuardrailForLBTestingB(CustomGuardrail): call_type: CallTypesLiteral, ) -> Optional[Union[Exception, str, dict]]: guardrail_lb_call_count["B"] += 1 - verbose_proxy_logger.info( - f"GuardrailForLBTestingB called. Total B calls: {guardrail_lb_call_count['B']}" - ) + verbose_proxy_logger.info(f"GuardrailForLBTestingB called. Total B calls: {guardrail_lb_call_count['B']}") return data @@ -79,9 +75,7 @@ class myCustomGuardrail(CustomGuardrail): _content = _content.replace("litellm", "********") message["content"] = _content - verbose_proxy_logger.debug( - "async_pre_call_hook: Message after masking %s", _messages - ) + verbose_proxy_logger.debug("async_pre_call_hook: Message after masking %s", _messages) return data diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index 17a7c09321a..ef2943df76d 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -51,9 +51,7 @@ def get_fine_tuning_provider_config( ): global fine_tuning_config if fine_tuning_config is None: - raise ValueError( - "fine_tuning_config is not set, set it on your config.yaml file." - ) + raise ValueError("fine_tuning_config is not set, set it on your config.yaml file.") for setting in fine_tuning_config: if setting.get("custom_llm_provider") == custom_llm_provider: return setting @@ -110,9 +108,7 @@ async def create_fine_tuning_job( data = fine_tuning_request.model_dump(exclude_none=True) try: if premium_user is not True: - raise ValueError( - f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") # Convert Pydantic model to dict verbose_proxy_logger.debug( @@ -146,14 +142,10 @@ async def create_fine_tuning_job( if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) - response = cast( - LiteLLMFineTuningJob, await llm_router.acreate_fine_tuning_job(**data) - ) + response = cast(LiteLLMFineTuningJob, await llm_router.acreate_fine_tuning_job(**data)) response.training_file = unified_file_id response._hidden_params["unified_file_id"] = unified_file_id ## ELSE, Route based on custom_llm_provider @@ -169,9 +161,7 @@ async def create_fine_tuning_job( response = await litellm.acreate_fine_tuning_job(**data) if response is None: - raise ValueError( - "Invalid request, No litellm managed file id or custom_llm_provider provided." - ) + raise ValueError("Invalid request, No litellm managed file id or custom_llm_provider provided.") ### CALL HOOKS ### - modify outgoing data _response = await proxy_logging_obj.post_call_success_hook( @@ -184,9 +174,7 @@ async def create_fine_tuning_job( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -212,9 +200,7 @@ async def create_fine_tuning_job( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.create_fine_tuning_job(): Exception occurred - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.create_fine_tuning_job(): Exception occurred - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -258,9 +244,7 @@ async def retrieve_fine_tuning_job( data: dict = {"fine_tuning_job_id": fine_tuning_job_id} try: if premium_user is not True: - raise ValueError( - f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") # Include original request and headers in the data base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -281,24 +265,18 @@ async def retrieve_fine_tuning_job( except Exception: request_body = {} - custom_llm_provider = ( - request_body.get("custom_llm_provider", None) or custom_llm_provider - ) + custom_llm_provider = request_body.get("custom_llm_provider", None) or custom_llm_provider ## CHECK IF MANAGED FILE ID unified_finetuning_job_id: Union[str, Literal[False]] = False response: Optional[LiteLLMFineTuningJob] = None if fine_tuning_job_id: - unified_finetuning_job_id = _is_base64_encoded_unified_file_id( - fine_tuning_job_id - ) + unified_finetuning_job_id = _is_base64_encoded_unified_file_id(fine_tuning_job_id) if unified_finetuning_job_id: if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) response = cast( LiteLLMFineTuningJob, @@ -306,14 +284,10 @@ async def retrieve_fine_tuning_job( **data, ), ) - response._hidden_params["unified_finetuning_job_id"] = ( - unified_finetuning_job_id - ) + response._hidden_params["unified_finetuning_job_id"] = unified_finetuning_job_id elif custom_llm_provider: # get configs for custom_llm_provider - llm_provider_config = get_fine_tuning_provider_config( - custom_llm_provider=custom_llm_provider - ) + llm_provider_config = get_fine_tuning_provider_config(custom_llm_provider=custom_llm_provider) if llm_provider_config is not None: data.update(llm_provider_config) @@ -339,9 +313,7 @@ async def retrieve_fine_tuning_job( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -368,9 +340,7 @@ async def retrieve_fine_tuning_job( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.retrieve_fine_tuning_job(): Exception occurred - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.retrieve_fine_tuning_job(): Exception occurred - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -420,9 +390,7 @@ async def list_fine_tuning_jobs( data: dict = {} try: if premium_user is not True: - raise ValueError( - f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") # Include original request and headers in the data base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -461,9 +429,7 @@ async def list_fine_tuning_jobs( return response elif custom_llm_provider: # get configs for custom_llm_provider - llm_provider_config = get_fine_tuning_provider_config( - custom_llm_provider=custom_llm_provider - ) + llm_provider_config = get_fine_tuning_provider_config(custom_llm_provider=custom_llm_provider) if llm_provider_config is not None: data.update(llm_provider_config) @@ -503,9 +469,7 @@ async def list_fine_tuning_jobs( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.list_fine_tuning_jobs(): Exception occurred - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.list_fine_tuning_jobs(): Exception occurred - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -549,9 +513,7 @@ async def cancel_fine_tuning_job( data: dict = {"fine_tuning_job_id": fine_tuning_job_id} try: if premium_user is not True: - raise ValueError( - f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") # Include original request and headers in the data base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -578,16 +540,12 @@ async def cancel_fine_tuning_job( unified_finetuning_job_id: Union[str, Literal[False]] = False response: Optional[LiteLLMFineTuningJob] = None if fine_tuning_job_id: - unified_finetuning_job_id = _is_base64_encoded_unified_file_id( - fine_tuning_job_id - ) + unified_finetuning_job_id = _is_base64_encoded_unified_file_id(fine_tuning_job_id) if unified_finetuning_job_id: if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) response = cast( LiteLLMFineTuningJob, @@ -595,14 +553,10 @@ async def cancel_fine_tuning_job( **data, ), ) - response._hidden_params["unified_finetuning_job_id"] = ( - unified_finetuning_job_id - ) + response._hidden_params["unified_finetuning_job_id"] = unified_finetuning_job_id else: # get configs for custom_llm_provider - llm_provider_config = get_fine_tuning_provider_config( - custom_llm_provider=custom_llm_provider - ) + llm_provider_config = get_fine_tuning_provider_config(custom_llm_provider=custom_llm_provider) if llm_provider_config is not None: data.update(llm_provider_config) @@ -628,9 +582,7 @@ async def cancel_fine_tuning_job( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -657,8 +609,6 @@ async def cancel_fine_tuning_job( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.cancel_fine_tuning_job(): Exception occurred - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.cancel_fine_tuning_job(): Exception occurred - {}".format(str(e)) ) raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/google_endpoints/agents_endpoints.py b/litellm/proxy/google_endpoints/agents_endpoints.py index 779284023a0..e3dd466b0b5 100644 --- a/litellm/proxy/google_endpoints/agents_endpoints.py +++ b/litellm/proxy/google_endpoints/agents_endpoints.py @@ -98,11 +98,7 @@ def _merge_query_params_into_data(data: dict, request: Request) -> dict: raw_template = query_params.get("litellm_params_template") if raw_template: try: - template = ( - json.loads(raw_template) - if isinstance(raw_template, str) - else raw_template - ) + template = json.loads(raw_template) if isinstance(raw_template, str) else raw_template except (json.JSONDecodeError, ValueError): template = {} if isinstance(template, dict): diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 6427835c250..37bcb3e6d6d 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -173,11 +173,7 @@ async def google_count_tokens(request: Request, model_name: str): from litellm.proxy._types import TokenCountRequest # Translate contents to openai format messages using the adapter - messages = ( - GoogleGenAIAdapter() - .translate_generate_content_to_completion(model_name, contents) - .get("messages", []) - ) + messages = GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", []) token_request = TokenCountRequest( model=model_name, diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 766ef0cf9f6..1d31e33c77c 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -23,9 +23,7 @@ from typing import Any, Callable, Dict, FrozenSet, Iterator, List # ``pre_call_hook`` directly with the sync name. Embedding, moderation, # audio, and transcription endpoints are deliberately excluded — text # guardrails on those paths are a separate scope. -TEXT_CONTENT_CALL_TYPES: FrozenSet[str] = frozenset( - {"completion", "acompletion", "aresponses"} -) +TEXT_CONTENT_CALL_TYPES: FrozenSet[str] = frozenset({"completion", "acompletion", "aresponses"}) def is_text_content_call_type(call_type: str) -> bool: @@ -61,9 +59,7 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]: if isinstance(input_value, str): return [{"role": "user", "content": input_value}] if isinstance(input_value, list): - if input_value and all( - isinstance(item, dict) and "role" in item for item in input_value - ): + if input_value and all(isinstance(item, dict) and "role" in item for item in input_value): return list(input_value) # Mixed lists (content-part dicts + bare strings) and pure # string/dict lists all become a single user message; the content @@ -141,9 +137,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visited if isinstance(input_value, list): # List of full messages: rewrite each message's content. - if input_value and all( - isinstance(item, dict) and "role" in item for item in input_value - ): + if input_value and all(isinstance(item, dict) and "role" in item for item in input_value): for item in input_value: if "content" in item: item["content"] = _rewrite_content(item["content"]) @@ -166,9 +160,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visited -def apply_redacted_messages_back( - data: Dict[str, Any], redacted_messages: List[Dict[str, Any]] -) -> None: +def apply_redacted_messages_back(data: Dict[str, Any], redacted_messages: List[Dict[str, Any]]) -> None: """Write redacted messages back to whichever field(s) the caller used. Mask/anonymize paths take a synthesised messages list (from @@ -201,9 +193,7 @@ def has_non_string_content(data: Dict[str, Any]) -> bool: messages = data.get("messages") if isinstance(messages, list): for message in messages: - if isinstance(message, dict) and not isinstance( - message.get("content"), str - ): + if isinstance(message, dict) and not isinstance(message.get("content"), str): if message.get("content") is not None: return True input_value = data.get("input") diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index c2175ce95e0..d3a5d649f17 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -184,9 +184,7 @@ async def list_guardrails_v2( is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN try: - guardrails = await GUARDRAIL_REGISTRY.get_all_guardrails_from_db( - prisma_client=prisma_client - ) + guardrails = await GUARDRAIL_REGISTRY.get_all_guardrails_from_db(prisma_client=prisma_client) excluded_guardrail_ids: set = set() if not is_admin: @@ -205,9 +203,7 @@ async def list_guardrails_v2( guardrail_configs: List[GuardrailInfoResponse] = [] seen_guardrail_ids: set = excluded_guardrail_ids.copy() for guardrail in guardrails: - litellm_params: Optional[Union[LitellmParams, dict]] = guardrail.get( - "litellm_params" - ) + litellm_params: Optional[Union[LitellmParams, dict]] = guardrail.get("litellm_params") litellm_params_dict = ( litellm_params.model_dump(exclude_none=True) if isinstance(litellm_params, LitellmParams) @@ -219,9 +215,7 @@ async def list_guardrails_v2( number_of_asterisks=4, ) masked_litellm_params = ( - BaseLitellmParams(**masked_litellm_params_dict) - if masked_litellm_params_dict - else None + BaseLitellmParams(**masked_litellm_params_dict) if masked_litellm_params_dict else None ) guardrail_configs.append( GuardrailInfoResponse( @@ -262,9 +256,7 @@ async def list_guardrails_v2( number_of_asterisks=4, ) masked_in_memory_litellm_params_typed = ( - BaseLitellmParams(**masked_in_memory_litellm_params) - if masked_in_memory_litellm_params - else None + BaseLitellmParams(**masked_in_memory_litellm_params) if masked_in_memory_litellm_params else None ) guardrail_configs.append( GuardrailInfoResponse( @@ -355,17 +347,13 @@ async def create_guardrail( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - result = await GUARDRAIL_REGISTRY.add_guardrail_to_db( - guardrail=request.guardrail, prisma_client=prisma_client - ) + result = await GUARDRAIL_REGISTRY.add_guardrail_to_db(guardrail=request.guardrail, prisma_client=prisma_client) guardrail_name = result.get("guardrail_name", "Unknown") guardrail_id = result.get("guardrail_id", "Unknown") try: - IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( - guardrail=cast(Guardrail, result), source="db" - ) + IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(guardrail=cast(Guardrail, result), source="db") verbose_proxy_logger.info( f"Immediate sync: Successfully initialized guardrail '{guardrail_name}' (ID: {guardrail_id})" ) @@ -373,13 +361,9 @@ async def create_guardrail( # Configuration error — roll back the DB write so the guardrail isn't orphaned if prisma_client is not None: try: - await GuardrailsRepository(prisma_client).table.delete( - where={"guardrail_id": guardrail_id} - ) + await GuardrailsRepository(prisma_client).table.delete(where={"guardrail_id": guardrail_id}) except Exception as rollback_err: - verbose_proxy_logger.warning( - f"Rollback failed for guardrail '{guardrail_id}': {rollback_err}" - ) + verbose_proxy_logger.warning(f"Rollback failed for guardrail '{guardrail_id}': {rollback_err}") raise HTTPException( status_code=400, detail=f"Guardrail configuration error: {init_error}", @@ -474,9 +458,7 @@ async def update_guardrail( ) if existing_guardrail is None: - raise HTTPException( - status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") result = await GUARDRAIL_REGISTRY.update_guardrail_in_db( guardrail_id=guardrail_id, @@ -550,9 +532,7 @@ async def delete_guardrail( ) if existing_guardrail is None: - raise HTTPException( - status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") result = await GUARDRAIL_REGISTRY.delete_guardrail_from_db( guardrail_id=guardrail_id, prisma_client=prisma_client @@ -588,9 +568,7 @@ class RegisterGuardrailRequest(BaseModel): """Request body for POST /guardrails/register. Follows Generic Guardrail API config.""" guardrail_name: str - litellm_params: Dict[ - str, Any - ] # guardrail, mode, api_base required; api_key, headers, etc. optional + litellm_params: Dict[str, Any] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: Optional[Dict[str, Any]] = None team_id: Optional[str] = None @@ -617,7 +595,9 @@ class GuardrailSubmissionItem(BaseModel): guardrail_name: str status: str # pending_review | active | rejected team_id: Optional[str] = None - team_guardrail: bool = False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails + team_guardrail: bool = ( + False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails + ) litellm_params: Optional[Dict[str, Any]] = None guardrail_info: Optional[Dict[str, Any]] = None submitted_by_user_id: Optional[str] = None @@ -714,9 +694,7 @@ async def register_guardrail( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - "Error checking guardrail name uniqueness: %s", e - ) + verbose_proxy_logger.exception("Error checking guardrail name uniqueness: %s", e) raise HTTPException(status_code=500, detail=str(e)) now = datetime.now(timezone.utc) @@ -724,9 +702,7 @@ async def register_guardrail( guardrail_info = dict(request.guardrail_info or {}) guardrail_info["submitted_by_user_id"] = user_api_key_dict.user_id guardrail_info["submitted_by_email"] = user_api_key_dict.user_email - guardrail_info["team_guardrail"] = ( - True # Mark as team submission for filtering/display - ) + guardrail_info["team_guardrail"] = True # Mark as team submission for filtering/display guardrail_info_str = safe_dumps(guardrail_info) try: @@ -796,9 +772,7 @@ def _row_to_submission_item(row: Any) -> GuardrailSubmissionItem: guardrail_info = _parse_json_field(row.guardrail_info) or {} team_guardrail = row.team_id is not None raw_params = _parse_json_field(row.litellm_params) or {} - masked_params = _get_masked_values( - raw_params, unmasked_length=4, number_of_asterisks=4 - ) + masked_params = _get_masked_values(raw_params, unmasked_length=4, number_of_asterisks=4) return GuardrailSubmissionItem( guardrail_id=row.guardrail_id, guardrail_name=row.guardrail_name, @@ -865,9 +839,7 @@ async def list_guardrail_submissions( # Non-admin with no team memberships: nothing visible. return ListGuardrailSubmissionsResponse( submissions=[], - summary=GuardrailSubmissionSummary( - total=0, pending_review=0, active=0, rejected=0 - ), + summary=GuardrailSubmissionSummary(total=0, pending_review=0, active=0, rejected=0), ) where_clause["team_id"] = {"in": visible_team_ids} @@ -879,12 +851,8 @@ async def list_guardrail_submissions( # Derive summary counts from the full result set total = len(all_team_rows) - pending_review = sum( - 1 for r in all_team_rows if (r.status or "active") == "pending_review" - ) - active_count = sum( - 1 for r in all_team_rows if (r.status or "active") == "active" - ) + pending_review = sum(1 for r in all_team_rows if (r.status or "active") == "pending_review") + active_count = sum(1 for r in all_team_rows if (r.status or "active") == "active") rejected = sum(1 for r in all_team_rows if (r.status or "active") == "rejected") # Apply filters to get the submissions list @@ -901,13 +869,9 @@ async def list_guardrail_submissions( if search_lower in (r.guardrail_name or "").lower() or ( isinstance(r.guardrail_info, dict) - and search_lower - in str((r.guardrail_info or {}).get("description", "")).lower() - ) - or ( - isinstance(r.guardrail_info, str) - and search_lower in r.guardrail_info.lower() + and search_lower in str((r.guardrail_info or {}).get("description", "")).lower() ) + or (isinstance(r.guardrail_info, str) and search_lower in r.guardrail_info.lower()) ] items = [_row_to_submission_item(r) for r in rows] @@ -943,13 +907,9 @@ async def get_guardrail_submission( is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN try: - row = await GuardrailsRepository(prisma_client).table.find_unique( - where={"guardrail_id": guardrail_id} - ) + row = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) if row is None: - raise HTTPException( - status_code=404, detail="Guardrail submission not found" - ) + raise HTTPException(status_code=404, detail="Guardrail submission not found") if not is_admin: visible_team_ids = await _get_user_team_ids(user_api_key_dict) if row.team_id is None or row.team_id not in visible_team_ids: @@ -984,13 +944,9 @@ async def approve_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await GuardrailsRepository(prisma_client).table.find_unique( - where={"guardrail_id": guardrail_id} - ) + row = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) if row is None: - raise HTTPException( - status_code=404, detail="Guardrail submission not found" - ) + raise HTTPException(status_code=404, detail="Guardrail submission not found") if row.status != "pending_review": raise HTTPException( status_code=400, @@ -1018,9 +974,7 @@ async def approve_guardrail_submission( "team_id": row.team_id, } try: - IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( - guardrail=cast(Guardrail, guardrail_dict), source="db" - ) + IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail(guardrail=cast(Guardrail, guardrail_dict), source="db") verbose_proxy_logger.info( "Approved guardrail %s (ID: %s) and initialized in memory", row.guardrail_name, @@ -1070,13 +1024,9 @@ async def reject_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await GuardrailsRepository(prisma_client).table.find_unique( - where={"guardrail_id": guardrail_id} - ) + row = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) if row is None: - raise HTTPException( - status_code=404, detail="Guardrail submission not found" - ) + raise HTTPException(status_code=404, detail="Guardrail submission not found") if row.status != "pending_review": raise HTTPException( status_code=400, @@ -1173,25 +1123,17 @@ async def patch_guardrail( ) if existing_guardrail is None: - raise HTTPException( - status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") # Create updated guardrail object guardrail_name = ( - request.guardrail_name - if request.guardrail_name is not None - else existing_guardrail.get("guardrail_name") + request.guardrail_name if request.guardrail_name is not None else existing_guardrail.get("guardrail_name") ) # Update litellm_params if default_on is provided or pii_entities_config is provided - litellm_params = LitellmParams( - **dict(existing_guardrail.get("litellm_params", {})) - ) + litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {}))) if request.litellm_params is not None: - requested_litellm_params = request.litellm_params.model_dump( - exclude_unset=True - ) + requested_litellm_params = request.litellm_params.model_dump(exclude_unset=True) litellm_params_dict = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) litellm_params = LitellmParams(**litellm_params_dict) @@ -1290,34 +1232,23 @@ async def get_guardrail_info(guardrail_id: str): raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = ( - GUARDRAIL_DEFINITION_LOCATION.DB - ) + guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.DB result = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db( guardrail_id=guardrail_id, prisma_client=prisma_client ) if result is None: - in_memory = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id( - guardrail_id=guardrail_id - ) + in_memory = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id(guardrail_id=guardrail_id) # Only return config-loaded entries here. A DB-backed entry that's # missing from the DB is stale (deleted on another pod, awaiting # reconciliation on this one) and must surface as 404. - if ( - in_memory is not None - and IN_MEMORY_GUARDRAIL_HANDLER.get_source(guardrail_id) == "config" - ): + if in_memory is not None and IN_MEMORY_GUARDRAIL_HANDLER.get_source(guardrail_id) == "config": result = in_memory guardrail_definition_location = GUARDRAIL_DEFINITION_LOCATION.CONFIG if result is None: - raise HTTPException( - status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") - litellm_params: Optional[Union[LitellmParams, dict]] = result.get( - "litellm_params" - ) + litellm_params: Optional[Union[LitellmParams, dict]] = result.get("litellm_params") result_litellm_params_dict = ( litellm_params.model_dump(exclude_none=True) if isinstance(litellm_params, LitellmParams) @@ -1328,11 +1259,7 @@ async def get_guardrail_info(guardrail_id: str): unmasked_length=4, number_of_asterisks=4, ) - masked_litellm_params = ( - BaseLitellmParams(**masked_litellm_params_dict) - if masked_litellm_params_dict - else None - ) + masked_litellm_params = BaseLitellmParams(**masked_litellm_params_dict) if masked_litellm_params_dict else None return GuardrailInfoResponse( guardrail_id=result.get("guardrail_id"), @@ -1450,9 +1377,7 @@ async def get_category_yaml(category_name: str): "file_type": file_type, } except Exception as e: - raise HTTPException( - status_code=500, detail=f"Error reading category file: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Error reading category file: {str(e)}") @router.get( @@ -1484,9 +1409,7 @@ async def get_major_airlines(): airlines = json.load(f) return {"airlines": airlines} except Exception as e: - raise HTTPException( - status_code=500, detail=f"Error reading major_airlines.json: {str(e)}" - ) from e + raise HTTPException(status_code=500, detail=f"Error reading major_airlines.json: {str(e)}") from e @router.post( @@ -1561,13 +1484,9 @@ async def validate_blocked_words_file(request: Dict[str, str]): if "action" not in word_data: errors.append(f"Entry {idx}: missing 'action' field") elif word_data["action"] not in ["BLOCK", "MASK"]: - errors.append( - f"Entry {idx}: action must be 'BLOCK' or 'MASK', got '{word_data['action']}'" - ) + errors.append(f"Entry {idx}: action must be 'BLOCK' or 'MASK', got '{word_data['action']}'") - if "description" in word_data and not isinstance( - word_data["description"], str - ): + if "description" in word_data and not isinstance(word_data["description"], str): errors.append(f"Entry {idx}: 'description' must be a string") if errors: @@ -1609,9 +1528,7 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str: return "dict" # Handle Literal types - if hasattr(field_annotation, "__origin__") and hasattr( - field_annotation, "__args__" - ): + if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"): # Check for Literal types (Python 3.8+) origin = field_annotation.__origin__ if hasattr(origin, "__name__") and origin.__name__ == "Literal": @@ -1704,8 +1621,7 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Check if the annotation is still a generic TypeVar (not specialized) if isinstance(field_annotation, TypeVar) or ( - hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is TypeVar + hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar ): return True @@ -1718,9 +1634,7 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Handle Optional[T] where T is still a TypeVar if hasattr(field_annotation, "__args__"): - non_none_args = [ - arg for arg in field_annotation.__args__ if arg is not type(None) - ] + non_none_args = [arg for arg in field_annotation.__args__ if arg is not type(None)] if non_none_args and isinstance(non_none_args[0], TypeVar): return True @@ -1824,9 +1738,7 @@ def _extract_fields_recursive( field_annotation = field.annotation # Skip optional_params if it's not meaningfully overridden - if _should_skip_optional_params( - field_name=field_name, field_annotation=field_annotation - ): + if _should_skip_optional_params(field_name=field_name, field_annotation=field_annotation): continue # Handle Optional types and get the actual type @@ -1848,9 +1760,7 @@ def _extract_fields_recursive( if is_basemodel_subclass: # Recursively get fields from the nested model - nested_fields = _extract_fields_recursive( - cast(Type[BaseModel], field_annotation), depth + 1 - ) + nested_fields = _extract_fields_recursive(cast(Type[BaseModel], field_annotation), depth + 1) fields[field_name] = { "description": description, "required": required, @@ -1938,9 +1848,7 @@ async def get_provider_specific_params(): lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel) tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel) - tool_permission_fields["ui_friendly_name"] = ( - ToolPermissionGuardrailConfigModel.ui_friendly_name() - ) + tool_permission_fields["ui_friendly_name"] = ToolPermissionGuardrailConfigModel.ui_friendly_name() # Return the provider-specific parameters provider_params = { @@ -2185,9 +2093,7 @@ async def test_custom_code_guardrail( ) -def _resolve_guardrail_input_type( - active_guardrail: CustomGuardrail, input_type: str -) -> Literal["request", "response"]: +def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: str) -> Literal["request", "response"]: """Return the effective input_type, auto-upgrading to 'response' for post_call guardrails.""" if input_type == "request": hook = getattr(active_guardrail, "event_hook", None) @@ -2196,16 +2102,12 @@ def _resolve_guardrail_input_type( return "response" if input_type == "response" else "request" -def _patch_logging_obj_for_guardrail( - litellm_logging_obj: Any, request: ApplyGuardrailRequest -) -> None: +def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None: """Configure the logging object so Langfuse/OTEL extract input and output correctly.""" litellm_logging_obj.call_type = "pass_through_endpoint" litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint" litellm_logging_obj.update_messages( - request.messages - if request.messages - else [{"role": "user", "content": request.text}] + request.messages if request.messages else [{"role": "user", "content": request.text}] ) @@ -2253,9 +2155,7 @@ async def _emit_guardrail_success_logs( cache_hit=False, ) except Exception: - verbose_proxy_logger.exception( - "apply_guardrail: async_success_handler failed" - ) + verbose_proxy_logger.exception("apply_guardrail: async_success_handler failed") try: thread_pool_executor.submit( litellm_logging_obj.success_handler, @@ -2265,9 +2165,7 @@ async def _emit_guardrail_success_logs( False, ) except Exception: - verbose_proxy_logger.exception( - "apply_guardrail: success_handler submit failed" - ) + verbose_proxy_logger.exception("apply_guardrail: success_handler submit failed") return response @@ -2308,10 +2206,8 @@ async def apply_guardrail( start_time = datetime.now(timezone.utc) try: - active_guardrail: Optional[CustomGuardrail] = ( - GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( - guardrail_name=request.guardrail_name - ) + active_guardrail: Optional[CustomGuardrail] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( + guardrail_name=request.guardrail_name ) if active_guardrail is None: raise HTTPException( @@ -2337,18 +2233,14 @@ async def apply_guardrail( _patch_logging_obj_for_guardrail(litellm_logging_obj, request) request_data: dict = {"messages": request.messages} if request.messages else {} - _input_type = _resolve_guardrail_input_type( - active_guardrail, request.input_type - ) + _input_type = _resolve_guardrail_input_type(active_guardrail, request.input_type) guardrailed_inputs = await active_guardrail.apply_guardrail( inputs={"texts": [request.text]}, request_data=request_data, input_type=_input_type, ) response_text = guardrailed_inputs.get("texts", []) - response = ApplyGuardrailResponse( - response_text=response_text[0] if response_text else request.text - ) + response = ApplyGuardrailResponse(response_text=response_text[0] if response_text else request.text) except Exception as e: if litellm_logging_obj is not None and not isinstance(e, HTTPException): try: @@ -2357,9 +2249,7 @@ async def apply_guardrail( traceback_exception=traceback.format_exc(), ) except Exception: - verbose_proxy_logger.exception( - "apply_guardrail: async_failure_handler failed" - ) + verbose_proxy_logger.exception("apply_guardrail: async_failure_handler failed") try: thread_pool_executor.submit( litellm_logging_obj.failure_handler, @@ -2367,9 +2257,7 @@ async def apply_guardrail( traceback.format_exc(), ) except Exception: - verbose_proxy_logger.exception( - "apply_guardrail: failure_handler submit failed" - ) + verbose_proxy_logger.exception("apply_guardrail: failure_handler submit failed") try: transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -2379,9 +2267,7 @@ async def apply_guardrail( if isinstance(transformed_exception, Exception): e = transformed_exception except Exception: - verbose_proxy_logger.exception( - "apply_guardrail: post_call_failure_hook failed" - ) + verbose_proxy_logger.exception("apply_guardrail: post_call_failure_hook failed") raise handle_exception_on_proxy(e) # Success logging outside except so a hook error never triggers failure handlers. diff --git a/litellm/proxy/guardrails/guardrail_helpers.py b/litellm/proxy/guardrails/guardrail_helpers.py index e9703114603..677ac66fcd0 100644 --- a/litellm/proxy/guardrails/guardrail_helpers.py +++ b/litellm/proxy/guardrails/guardrail_helpers.py @@ -7,9 +7,7 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy.proxy_server import LiteLLM_TeamTable, UserAPIKeyAuth from litellm.types.guardrails import * -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path def can_modify_guardrails(team_obj: Optional[LiteLLM_TeamTable]) -> bool: @@ -18,9 +16,7 @@ def can_modify_guardrails(team_obj: Optional[LiteLLM_TeamTable]) -> bool: team_metadata = team_obj.metadata or {} - if team_metadata.get("guardrails", None) is not None and isinstance( - team_metadata.get("guardrails"), Dict - ): + if team_metadata.get("guardrails", None) is not None and isinstance(team_metadata.get("guardrails"), Dict): if team_metadata.get("guardrails", {}).get("modify_guardrails", None) is False: return False @@ -56,16 +52,12 @@ async def should_proceed_based_on_metadata(data: dict, guardrail_name: str) -> b continue # lookup the guardrail in guardrail_name_config_map - guardrail_item: GuardrailItem = litellm.guardrail_name_config_map[ - _guardrail_name - ] + guardrail_item: GuardrailItem = litellm.guardrail_name_config_map[_guardrail_name] guardrail_callbacks = guardrail_item.callbacks requested_callback_names.extend(guardrail_callbacks) - verbose_proxy_logger.debug( - "requested_callback_names %s", requested_callback_names - ) + verbose_proxy_logger.debug("requested_callback_names %s", requested_callback_names) if guardrail_name in requested_callback_names: return True @@ -75,9 +67,7 @@ async def should_proceed_based_on_metadata(data: dict, guardrail_name: str) -> b return True -async def should_proceed_based_on_api_key( - user_api_key_dict: UserAPIKeyAuth, guardrail_name: str -) -> bool: +async def should_proceed_based_on_api_key(user_api_key_dict: UserAPIKeyAuth, guardrail_name: str) -> bool: """ checks if this guardrail should be applied to this call """ @@ -105,9 +95,7 @@ async def should_proceed_based_on_api_key( continue # lookup the guardrail in guardrail_name_config_map - guardrail_item: GuardrailItem = litellm.guardrail_name_config_map[ - _guardrail_name - ] + guardrail_item: GuardrailItem = litellm.guardrail_name_config_map[_guardrail_name] guardrail_callbacks = guardrail_item.callbacks if guardrail_name in guardrail_callbacks: diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index d70c8e4f310..e7d9406ae3b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -44,9 +44,7 @@ class AimGuardrailMissingSecrets(Exception): class AimGuardrail(CustomGuardrail): - def __init__( - self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs - ): + def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs): ssl_verify = kwargs.pop("ssl_verify", None) self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, @@ -59,12 +57,8 @@ class AimGuardrail(CustomGuardrail): "pass it as a parameter to the guardrail in the config file" ) raise AimGuardrailMissingSecrets(msg) - self.api_base = ( - api_base or os.environ.get("AIM_API_BASE") or "https://api.aim.security" - ) - self.ws_api_base = self.api_base.replace("http://", "ws://").replace( - "https://", "wss://" - ) + self.api_base = api_base or os.environ.get("AIM_API_BASE") or "https://api.aim.security" + self.ws_api_base = self.api_base.replace("http://", "ws://").replace("https://", "wss://") self.dlp_entities: list[dict] = [] self._max_dlp_entities = 100 super().__init__(**kwargs) @@ -77,9 +71,7 @@ class AimGuardrail(CustomGuardrail): call_type: CallTypesLiteral, ) -> Union[Exception, str, dict, None]: verbose_proxy_logger.debug("Inside AIM Pre-Call Hook") - return await self.call_aim_guardrail( - data, hook="pre_call", key_alias=user_api_key_dict.key_alias - ) + return await self.call_aim_guardrail(data, hook="pre_call", key_alias=user_api_key_dict.key_alias) async def async_moderation_hook( self, @@ -89,14 +81,10 @@ class AimGuardrail(CustomGuardrail): ) -> Union[Exception, str, dict, None]: verbose_proxy_logger.debug("Inside AIM Moderation Hook") - await self.call_aim_guardrail( - data, hook="moderation", key_alias=user_api_key_dict.key_alias - ) + await self.call_aim_guardrail(data, hook="moderation", key_alias=user_api_key_dict.key_alias) return data - async def call_aim_guardrail( - self, data: dict, hook: str, key_alias: Optional[str] - ) -> dict: + async def call_aim_guardrail(self, data: dict, hook: str, key_alias: Optional[str]) -> dict: user_email = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") call_id = data.get("litellm_call_id") headers = self._build_aim_headers( @@ -179,9 +167,7 @@ class AimGuardrail(CustomGuardrail): async def call_aim_guardrail_on_output( self, request_data: dict, output: str, hook: str, key_alias: Optional[str] ) -> Optional[dict]: - user_email = ( - request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") - ) + user_email = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") call_id = request_data.get("litellm_call_id") response = await self.async_handler.post( f"{self.api_base}/fw/v1/analyze", @@ -191,30 +177,21 @@ class AimGuardrail(CustomGuardrail): user_email=user_email, litellm_call_id=call_id, ), - json={ - "messages": build_inspection_messages(request_data) - + [{"role": "assistant", "content": output}] - }, + json={"messages": build_inspection_messages(request_data) + [{"role": "assistant", "content": output}]}, ) response.raise_for_status() res = response.json() required_action = res.get("required_action") action_type = required_action and required_action.get("action_type", None) if action_type and action_type == "block_action": - return self._handle_block_action_on_output( - res["analysis_result"], required_action - ) + return self._handle_block_action_on_output(res["analysis_result"], required_action) redacted_chat = res.get("redacted_chat", None) if action_type and action_type == "anonymize_action" and redacted_chat: - return { - "redacted_output": redacted_chat["all_redacted_messages"][-1]["content"] - } + return {"redacted_output": redacted_chat["all_redacted_messages"][-1]["content"]} return {"redacted_output": output} - def _handle_block_action_on_output( - self, analysis_result: Any, required_action: Any - ) -> dict | None: + def _handle_block_action_on_output(self, analysis_result: Any, required_action: Any) -> dict | None: detection_message = required_action.get("detection_message", None) verbose_proxy_logger.info( "Aim: detected: {detected}, enabled policies: {policies}".format( @@ -290,19 +267,13 @@ class AimGuardrail(CustomGuardrail): for choice, aim_output_guardrail_result in zip(choices_to_inspect, results): if isinstance(aim_output_guardrail_result, BaseException): raise aim_output_guardrail_result - if aim_output_guardrail_result and aim_output_guardrail_result.get( - "detection_message" - ): + if aim_output_guardrail_result and aim_output_guardrail_result.get("detection_message"): raise self._rejection( aim_output_guardrail_result.get("detection_message"), openai_code="content_policy_violation", ) - if aim_output_guardrail_result and aim_output_guardrail_result.get( - "redacted_output" - ): - choice.message.content = aim_output_guardrail_result.get( - "redacted_output" - ) + if aim_output_guardrail_result and aim_output_guardrail_result.get("redacted_output"): + choice.message.content = aim_output_guardrail_result.get("redacted_output") return response async def async_post_call_streaming_iterator_hook( @@ -311,9 +282,7 @@ class AimGuardrail(CustomGuardrail): response, request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: - user_email = ( - request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") - ) + user_email = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") call_id = request_data.get("litellm_call_id") async with connect( f"{self.ws_api_base}/fw/v1/analyze/stream", @@ -324,9 +293,7 @@ class AimGuardrail(CustomGuardrail): litellm_call_id=call_id, ), ) as websocket: - sender = asyncio.create_task( - self.forward_the_stream_to_aim(websocket, response) - ) + sender = asyncio.create_task(self.forward_the_stream_to_aim(websocket, response)) while True: result = json.loads(await websocket.recv()) if verified_chunk := result.get("verified_chunk"): @@ -339,9 +306,7 @@ class AimGuardrail(CustomGuardrail): from litellm.proxy.proxy_server import StreamingCallbackError raise StreamingCallbackError(blocking_message) - verbose_proxy_logger.error( - f"Unknown message received from AIM: {result}" - ) + verbose_proxy_logger.error(f"Unknown message received from AIM: {result}") return async def forward_the_stream_to_aim( diff --git a/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py index 1e3dd906b9f..9fe098f29e8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py @@ -16,9 +16,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" akto_api_key=getattr(litellm_params, "akto_api_key", None), akto_account_id=getattr(litellm_params, "akto_account_id", None), akto_vxlan_id=getattr(litellm_params, "akto_vxlan_id", None), - unreachable_fallback=getattr( - litellm_params, "unreachable_fallback", "fail_closed" - ), + unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"), guardrail_timeout=getattr(litellm_params, "guardrail_timeout", None), guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, diff --git a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py index ece311666c5..be9c9cb1be7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py +++ b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py @@ -77,27 +77,17 @@ class AktoGuardrail(CustomGuardrail): ) self.background_tasks: set = set() - self.akto_base_url = ( - akto_base_url or os.environ.get("AKTO_GUARDRAIL_API_BASE", "") - ).rstrip("/") + self.akto_base_url = (akto_base_url or os.environ.get("AKTO_GUARDRAIL_API_BASE", "")).rstrip("/") if not self.akto_base_url: - raise ValueError( - "akto_base_url is required. Set AKTO_GUARDRAIL_API_BASE or pass it in litellm_params." - ) + raise ValueError("akto_base_url is required. Set AKTO_GUARDRAIL_API_BASE or pass it in litellm_params.") self.akto_api_key = akto_api_key or os.environ.get("AKTO_API_KEY", "") if not self.akto_api_key: - raise ValueError( - "akto_api_key is required. Set AKTO_API_KEY or pass it in litellm_params." - ) + raise ValueError("akto_api_key is required. Set AKTO_API_KEY or pass it in litellm_params.") - self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( - unreachable_fallback - ) + self.unreachable_fallback: Literal["fail_closed", "fail_open"] = unreachable_fallback self.guardrail_timeout = guardrail_timeout or DEFAULT_GUARDRAIL_TIMEOUT - self.akto_account_id = akto_account_id or os.environ.get( - "AKTO_ACCOUNT_ID", "1000000" - ) + self.akto_account_id = akto_account_id or os.environ.get("AKTO_ACCOUNT_ID", "1000000") self.akto_vxlan_id = akto_vxlan_id or os.environ.get("AKTO_VXLAN_ID", "0") kwargs["supported_event_hooks"] = [ @@ -183,9 +173,7 @@ class AktoGuardrail(CustomGuardrail): body["model"] = request_data["model"] else: texts = inputs.get("texts", []) - body["messages"] = ( - [{"role": "user", "content": t} for t in texts] if texts else [] - ) + body["messages"] = [{"role": "user", "content": t} for t in texts] if texts else [] tools = inputs.get("tools") if tools: @@ -211,23 +199,15 @@ class AktoGuardrail(CustomGuardrail): texts = inputs.get("texts", []) if texts: - return { - "choices": [ - {"message": {"content": t, "role": "assistant"}} for t in texts - ] - } + return {"choices": [{"message": {"content": t, "role": "assistant"}} for t in texts]} return {} @staticmethod def build_tag_metadata(request_data: dict) -> Dict[str, str]: """Build tag/metadata dict with user_id and team_id for Akto tracking.""" tag: Dict[str, str] = {"gen-ai": "Gen AI"} - user_id = AktoGuardrail.resolve_metadata_value( - request_data, "user_api_key_user_id" - ) - team_id = AktoGuardrail.resolve_metadata_value( - request_data, "user_api_key_team_id" - ) + user_id = AktoGuardrail.resolve_metadata_value(request_data, "user_api_key_user_id") + team_id = AktoGuardrail.resolve_metadata_value(request_data, "user_api_key_team_id") if user_id: tag["user_id"] = user_id if team_id: @@ -256,23 +236,15 @@ class AktoGuardrail(CustomGuardrail): response_headers: Dict[str, str] = {} if include_response: response_body = self.build_response_body(inputs, request_data) - response_payload = json.dumps( - {"body": json.dumps(response_body)} - ) # Double-encoded + response_payload = json.dumps({"body": json.dumps(response_body)}) # Double-encoded response_headers = {"content-type": "application/json"} # Extract client IP from proxy headers ip = "" proxy_req = request_data.get("proxy_server_request", {}) - proxy_headers = ( - proxy_req.get("headers", {}) if isinstance(proxy_req, dict) else {} - ) + proxy_headers = proxy_req.get("headers", {}) if isinstance(proxy_req, dict) else {} if isinstance(proxy_headers, dict): - ip = ( - proxy_headers.get("x-forwarded-for") - or proxy_headers.get("x-real-ip") - or "" - ) + ip = proxy_headers.get("x-forwarded-for") or proxy_headers.get("x-real-ip") or "" if "," in ip: ip = ip.split(",")[0].strip() @@ -281,9 +253,7 @@ class AktoGuardrail(CustomGuardrail): "requestHeaders": json.dumps(request_headers), "responseHeaders": json.dumps(response_headers), "method": "POST", - "requestPayload": json.dumps( - {"body": json.dumps(request_body)} - ), # Double-encoded + "requestPayload": json.dumps({"body": json.dumps(request_body)}), # Double-encoded "responsePayload": response_payload, "ip": ip, "destIp": "127.0.0.1", @@ -423,9 +393,7 @@ class AktoGuardrail(CustomGuardrail): if input_type == "request": # Pre_call: awaited guardrail check (no ingestion) - payload = self.build_akto_payload( - inputs, request_data, include_response=False - ) + payload = self.build_akto_payload(inputs, request_data, include_response=False) try: response = await self.send_request( guardrails=True, @@ -451,9 +419,7 @@ class AktoGuardrail(CustomGuardrail): ) blocked_payload["responsePayload"] = json.dumps( { - "body": json.dumps( - {"x-blocked-by": "Akto Proxy", "reason": reason} - ), + "body": json.dumps({"x-blocked-by": "Akto Proxy", "reason": reason}), } ) blocked_payload["responseHeaders"] = json.dumps( @@ -476,9 +442,7 @@ class AktoGuardrail(CustomGuardrail): elif input_type == "response": # Post_call: fire-and-forget combined guardrail + ingest - payload = self.build_akto_payload( - inputs, request_data, include_response=True - ) + payload = self.build_akto_payload(inputs, request_data, include_response=True) task = asyncio.create_task( self.fire_and_forget_request( guardrails=True, diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py index 35f776e5bd1..ba9c8398152 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py @@ -8,9 +8,7 @@ import os import sys -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import json import sys from typing import TYPE_CHECKING, Any, List, Literal, Optional, Type @@ -39,12 +37,8 @@ if TYPE_CHECKING: class AporiaGuardrail(CustomGuardrail): - def __init__( - self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs - ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs): + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.aporia_api_key = api_key or os.environ["APORIO_API_KEY"] self.aporia_api_base = api_base or os.environ["APORIO_API_BASE"] super().__init__(**kwargs) @@ -67,9 +61,7 @@ class AporiaGuardrail(CustomGuardrail): return new_messages - async def prepare_aporia_request( - self, new_messages: List[dict], response_string: Optional[str] = None - ) -> dict: + async def prepare_aporia_request(self, new_messages: List[dict], response_string: Optional[str] = None) -> dict: data: dict[str, Any] = {} if new_messages is not None: data["messages"] = new_messages @@ -93,13 +85,9 @@ class AporiaGuardrail(CustomGuardrail): new_messages: List[dict], response_string: Optional[str] = None, ): - data = await self.prepare_aporia_request( - new_messages=new_messages, response_string=response_string - ) + data = await self.prepare_aporia_request(new_messages=new_messages, response_string=response_string) - data.update( - self.get_guardrail_dynamic_request_body_params(request_data=request_data) - ) + data.update(self.get_guardrail_dynamic_request_body_params(request_data=request_data)) _json_data = json.dumps(data) @@ -132,9 +120,7 @@ class AporiaGuardrail(CustomGuardrail): if response.status_code == 200: # check if the response was flagged _json_response = response.json() - action: str = _json_response.get( - "action" - ) # possible values are modify, passthrough, block, rephrase + action: str = _json_response.get("action") # possible values are modify, passthrough, block, rephrase if action == "block": raise HTTPException( status_code=400, @@ -170,9 +156,7 @@ class AporiaGuardrail(CustomGuardrail): new_messages=data.get("messages", []), ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) pass @@ -223,13 +207,9 @@ class AporiaGuardrail(CustomGuardrail): request_data=data, new_messages=new_messages, ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) else: - verbose_proxy_logger.warning( - "Aporia AI: not running guardrail. No messages in data" - ) + verbose_proxy_logger.warning("Aporia AI: not running guardrail. No messages in data") pass @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py index 449dd42ba59..243c4ad408b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py @@ -49,13 +49,9 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" }, ) else: - raise ValueError( - f"Azure Content Safety: {azure_guardrail} is not a valid guardrail" - ) + raise ValueError(f"Azure Content Safety: {azure_guardrail} is not a valid guardrail") - litellm.logging_callback_manager.add_litellm_callback( - azure_content_safety_guardrail - ) + litellm.logging_callback_manager.add_litellm_callback(azure_content_safety_guardrail) return azure_content_safety_guardrail diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 7d2dfce0711..b178efbda59 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -36,16 +36,12 @@ class AzureGuardrailBase: # (typically CustomGuardrail). super().__init__(**kwargs) - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_key = api_key self.api_base = api_base self.api_version: str = kwargs.get("api_version") or "2024-09-01" - async def _post_to_content_safety( - self, endpoint_path: str, request_body: Dict[str, Any] - ) -> Dict[str, Any]: + async def _post_to_content_safety(self, endpoint_path: str, request_body: Dict[str, Any]) -> Dict[str, Any]: """POST to an Azure Content Safety endpoint with standard auth headers. Args: @@ -62,18 +58,14 @@ class AzureGuardrailBase: "Content-Type": "application/json", } - verbose_proxy_logger.debug( - "Azure Content Safety request [%s]: %s", endpoint_path, request_body - ) + verbose_proxy_logger.debug("Azure Content Safety request [%s]: %s", endpoint_path, request_body) response = await self.async_handler.post( url=url, headers=headers, json=request_body, ) response_json: Dict[str, Any] = response.json() - verbose_proxy_logger.debug( - "Azure Content Safety response [%s]: %s", endpoint_path, response_json - ) + verbose_proxy_logger.debug("Azure Content Safety response [%s]: %s", endpoint_path, response_json) return response_json @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index 5f7e04cfb8b..788fe5b05c7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -63,13 +63,9 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai **kwargs, ) - verbose_proxy_logger.debug( - f"Initialized Azure Prompt Shield Guardrail: {guardrail_name}" - ) + verbose_proxy_logger.debug(f"Initialized Azure Prompt Shield Guardrail: {guardrail_name}") - async def async_make_request( - self, user_prompt: str - ) -> "AzurePromptShieldGuardrailResponse": + async def async_make_request(self, user_prompt: str) -> "AzurePromptShieldGuardrailResponse": """ Make a request to the Azure Prompt Shield API. @@ -84,19 +80,13 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai AzurePromptShieldGuardrailResponse, ) - chunks = self.split_text_by_words( - user_prompt, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH - ) + chunks = self.split_text_by_words(user_prompt, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH) last_response: Optional[AzurePromptShieldGuardrailResponse] = None for chunk in chunks: - request_body = AzurePromptShieldGuardrailRequestBody( - documents=[], userPrompt=chunk - ) - response_json = await self._post_to_content_safety( - "text:shieldPrompt", cast(dict, request_body) - ) + request_body = AzurePromptShieldGuardrailRequestBody(documents=[], userPrompt=chunk) + response_json = await self._post_to_content_safety("text:shieldPrompt", cast(dict, request_body)) last_response = cast(AzurePromptShieldGuardrailResponse, response_json) @@ -136,16 +126,12 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai ) new_messages: Optional[List[AllMessageValues]] = data.get("messages") if new_messages is None: - verbose_proxy_logger.warning( - "Azure Prompt Shield: not running guardrail. No messages in data" - ) + verbose_proxy_logger.warning("Azure Prompt Shield: not running guardrail. No messages in data") return data user_prompt = self.get_user_prompt(new_messages) if user_prompt: - verbose_proxy_logger.debug( - f"Azure Prompt Shield: User prompt: {user_prompt}" - ) + verbose_proxy_logger.debug(f"Azure Prompt Shield: User prompt: {user_prompt}") await self.async_make_request( user_prompt=user_prompt, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index e21be6ffdbe..c8553926559 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -73,21 +73,15 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr "SelfHarm", "Violence", ], - "blocklistNames": cast( - Optional[List[str]], kwargs.get("blocklistNames") or None - ), + "blocklistNames": cast(Optional[List[str]], kwargs.get("blocklistNames") or None), "haltOnBlocklistHit": kwargs.get("haltOnBlocklistHit") or False, "outputType": kwargs.get("outputType") or "FourSeverityLevels", } - self.severity_threshold = ( - int(severity_threshold) if severity_threshold else None - ) + self.severity_threshold = int(severity_threshold) if severity_threshold else None self.severity_threshold_by_category = severity_threshold_by_category - verbose_proxy_logger.info( - f"Initialized Azure Text Moderation Guardrail: {guardrail_name}" - ) + verbose_proxy_logger.info(f"Initialized Azure Text Moderation Guardrail: {guardrail_name}") @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: @@ -97,9 +91,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr return AzureContentSafetyTextModerationConfigModel - async def async_make_request( - self, text: str - ) -> "AzureTextModerationGuardrailResponse": + async def async_make_request(self, text: str) -> "AzureTextModerationGuardrailResponse": """ Make a request to the Azure Text Moderation API. @@ -123,9 +115,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr text=chunk, **self.optional_params_request_body, # type: ignore[misc] ) - response_json = await self._post_to_content_safety( - "text:analyze", cast(dict, request_body) - ) + response_json = await self._post_to_content_safety("text:analyze", cast(dict, request_body)) chunk_response = cast(AzureTextModerationGuardrailResponse, response_json) @@ -147,9 +137,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr assert last_response is not None return last_response - def check_severity_threshold( - self, response: "AzureTextModerationGuardrailResponse" - ) -> Literal[True]: + def check_severity_threshold(self, response: "AzureTextModerationGuardrailResponse") -> Literal[True]: """ - Check if threshold set by category - Check if general severity threshold set @@ -158,9 +146,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr if self.severity_threshold_by_category: for category in response["categoriesAnalysis"]: - severity_category_threshold_item = ( - self.severity_threshold_by_category.get(category["category"]) - ) + severity_category_threshold_item = self.severity_threshold_by_category.get(category["category"]) if ( severity_category_threshold_item is not None and category["severity"] >= severity_category_threshold_item @@ -170,9 +156,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr detail={ "error": "Azure Content Safety Guardrail: {} crossed severity {}, Got severity: {}".format( category["category"], - self.severity_threshold_by_category.get( - category["category"] - ), + self.severity_threshold_by_category.get(category["category"]), category["severity"], ) }, @@ -190,10 +174,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr ) }, ) - if ( - self.severity_threshold is None - and self.severity_threshold_by_category is None - ): + if self.severity_threshold is None and self.severity_threshold_by_category is None: for category in response["categoriesAnalysis"]: if category["severity"] >= self.default_severity_threshold: raise HTTPException( @@ -227,16 +208,12 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr ) new_messages: Optional[List[AllMessageValues]] = data.get("messages") if new_messages is None: - verbose_proxy_logger.warning( - "Azure Text Moderation: not running guardrail. No messages in data" - ) + verbose_proxy_logger.warning("Azure Text Moderation: not running guardrail. No messages in data") return data user_prompt = self.get_user_prompt(new_messages) if user_prompt: - verbose_proxy_logger.info( - f"Azure Text Moderation: User prompt: {user_prompt}" - ) + verbose_proxy_logger.info(f"Azure Text Moderation: User prompt: {user_prompt}") await self.async_make_request( text=user_prompt, ) @@ -264,9 +241,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr ) return response - async def async_post_call_streaming_hook( - self, user_api_key_dict: UserAPIKeyAuth, response: str - ) -> Any: + async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> Any: try: if response is not None and len(response) > 0: await self.async_make_request( @@ -285,9 +260,7 @@ def _message_content_to_text(content: Any) -> str: return content if isinstance(content, list): text_parts = [ - item.get("text") - for item in content - if isinstance(item, dict) and isinstance(item.get("text"), str) + item.get("text") for item in content if isinstance(item, dict) and isinstance(item.get("text"), str) ] return "\n".join(part for part in text_parts if part) return "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 19055f1083a..6e46f971dd8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -8,9 +8,7 @@ import os import sys -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import json import sys from typing import ( @@ -153,15 +151,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): disable_exception_on_block: Optional[bool] = False, **kwargs, ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.guardrailIdentifier = guardrailIdentifier self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" - self.experimental_use_latest_role_message_only = bool( - kwargs.get("experimental_use_latest_role_message_only") - ) + self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only")) # store kwargs as optional_params self.optional_params = kwargs @@ -190,9 +184,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.guardrailVersion, ) - def _create_bedrock_input_content_request( - self, messages: Optional[List[AllMessageValues]] - ) -> BedrockRequest: + def _create_bedrock_input_content_request(self, messages: Optional[List[AllMessageValues]]) -> BedrockRequest: """ Create a bedrock request for the input content - the LLM request. """ @@ -209,9 +201,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # exclusively when assembling the OUTPUT request, so a caller cannot use # a grounding_source/query tag to change how input-safety policies treat # their content (which would be an input-guardrail bypass). - bedrock_request_content.append( - BedrockContentItem(text=BedrockTextContent(text=block.text)) - ) + bedrock_request_content.append(BedrockContentItem(text=BedrockTextContent(text=block.text))) bedrock_request["content"] = bedrock_request_content return bedrock_request @@ -238,9 +228,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): has_grounding = len(bedrock_request_content) > 0 # Append the response (the content to guard) after any grounding blocks; assign # unconditionally so harvested grounding blocks survive a non-ModelResponse input. - bedrock_request_content.extend( - self._build_response_content_items(response, has_grounding=has_grounding) - ) + bedrock_request_content.extend(self._build_response_content_items(response, has_grounding=has_grounding)) bedrock_request["content"] = bedrock_request_content return bedrock_request @@ -283,18 +271,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ bedrock_request: BedrockRequest = BedrockRequest(source=source) if source == "INPUT": - bedrock_request = self._create_bedrock_input_content_request( - messages=messages - ) + bedrock_request = self._create_bedrock_input_content_request(messages=messages) elif source == "OUTPUT": - bedrock_request = self._create_bedrock_output_content_request( - response=response, messages=messages - ) + bedrock_request = self._create_bedrock_output_content_request(response=response, messages=messages) return bedrock_request - def get_content_items_for_message( - self, message: AllMessageValues - ) -> Optional[List[QualifiedTextBlock]]: + def get_content_items_for_message(self, message: AllMessageValues) -> Optional[List[QualifiedTextBlock]]: """ Flatten a message into text blocks, preserving any contextual-grounding qualifier carried by the content-block ``type`` (grounding_source / query). @@ -311,9 +293,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): for item in content: if isinstance(item, dict) and "text" in item: qualifier = _CONTENT_TYPE_TO_QUALIFIER.get(item.get("type", "")) - blocks.append( - QualifiedTextBlock(text=item["text"], qualifier=qualifier) - ) + blocks.append(QualifiedTextBlock(text=item["text"], qualifier=qualifier)) elif isinstance(item, str): blocks.append(QualifiedTextBlock(text=item, qualifier=None)) return blocks @@ -325,9 +305,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): text_content["qualifiers"] = [block.qualifier] return BedrockContentItem(text=text_content) - def _collect_grounding_blocks( - self, messages: Optional[List[AllMessageValues]] - ) -> List[QualifiedTextBlock]: + def _collect_grounding_blocks(self, messages: Optional[List[AllMessageValues]]) -> List[QualifiedTextBlock]: """Harvest grounding_source/query blocks from the request for an OUTPUT scan. ``grounding_source`` is honored only from app-authored roles (system / @@ -343,10 +321,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): for block in self.get_content_items_for_message(message=message) or []: if block.qualifier == "query": grounding.append(block) - elif ( - block.qualifier == "grounding_source" - and role in _GROUNDING_SOURCE_TRUSTED_ROLES - ): + elif block.qualifier == "grounding_source" and role in _GROUNDING_SOURCE_TRUSTED_ROLES: grounding.append(block) return grounding @@ -375,9 +350,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): target_indices=[latest_index], ) - def _find_latest_message_index( - self, messages: List[AllMessageValues], target_role: str - ) -> Optional[int]: + def _find_latest_message_index(self, messages: List[AllMessageValues], target_role: str) -> Optional[int]: for index in range(len(messages) - 1, -1, -1): if messages[index].get("role", None) == target_role: return index @@ -390,11 +363,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if isinstance(content, str): return 1 if isinstance(content, list): - return sum( - 1 - for item in content - if isinstance(item, dict) and item.get("text") is not None - ) + return sum(1 for item in content if isinstance(item, dict) and item.get("text") is not None) return 0 def _locate_message_texts_slice( @@ -441,9 +410,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): leaking tool/assistant content to the INPUT scan (https://github.com/BerriAI/litellm/issues/23476). """ - mock_messages: list[AllMessageValues] = [ - ChatCompletionUserMessage(role="user", content=text) for text in texts - ] + mock_messages: list[AllMessageValues] = [ChatCompletionUserMessage(role="user", content=text) for text in texts] if self.experimental_use_latest_role_message_only is not True: return ApplyGuardrailMessageSelection( @@ -470,22 +437,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # No role information available (e.g. raw-text callers like # /guardrails/apply_guardrail) — keep the legacy behavior of # scanning the latest text only. - filter_result = self._prepare_guardrail_messages_for_role( - messages=mock_messages - ) + filter_result = self._prepare_guardrail_messages_for_role(messages=mock_messages) return ApplyGuardrailMessageSelection( filtered_messages=filter_result.payload_messages or mock_messages, scanned_slice=None, scanned_role_subset=False, ) - latest_user_index = self._find_latest_message_index( - structured_messages, target_role="user" - ) + latest_user_index = self._find_latest_message_index(structured_messages, target_role="user") if latest_user_index is None: - verbose_proxy_logger.debug( - "Bedrock Guardrail: no user-role message in request, skipping INPUT scan" - ) + verbose_proxy_logger.debug("Bedrock Guardrail: no user-role message in request, skipping INPUT scan") return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True) selected_message = structured_messages[latest_user_index] @@ -556,9 +517,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): merged_messages = list(original_messages) if not merged_messages: merged_messages = list(updated_target_messages) - for replacement_index, updated_message in zip( - target_indices, updated_target_messages - ): + for replacement_index, updated_message in zip(target_indices, updated_target_messages): if replacement_index < len(merged_messages): merged_messages[replacement_index] = updated_message @@ -583,9 +542,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): aws_role_name = self.optional_params.get("aws_role_name", None) aws_session_name = self.optional_params.get("aws_session_name", None) aws_profile_name = self.optional_params.get("aws_profile_name", None) - aws_web_identity_token = self.optional_params.get( - "aws_web_identity_token", None - ) + aws_web_identity_token = self.optional_params.get("aws_web_identity_token", None) aws_sts_endpoint = self.optional_params.get("aws_sts_endpoint", None) ### SET REGION NAME ### @@ -619,15 +576,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - aws_bedrock_runtime_endpoint = self.optional_params.get( - "aws_bedrock_runtime_endpoint", None - ) + aws_bedrock_runtime_endpoint = self.optional_params.get("aws_bedrock_runtime_endpoint", None) _, proxy_endpoint_url = self.get_runtime_endpoint( api_base=None, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_region_name=aws_region_name, ) - proxy_endpoint_url = f"{proxy_endpoint_url}/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply" + proxy_endpoint_url = ( + f"{proxy_endpoint_url}/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply" + ) # api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply" encoded_data = json.dumps(data).encode("utf-8") @@ -641,9 +598,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): try: from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError( - "Missing boto3 to call bedrock. Run 'pip install boto3'." - ) + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") headers["Authorization"] = f"Bearer {aws_bearer_token}" request = AWSRequest( method="POST", @@ -656,9 +611,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError( - "Missing boto3 to call bedrock. Run 'pip install boto3'." - ) + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) request = AWSRequest( @@ -689,20 +642,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time = datetime.now() credentials, aws_region_name = self._load_credentials() bedrock_request_data: dict = dict( - self.convert_to_bedrock_format( - source=source, messages=messages, response=response - ) - ) - bedrock_guardrail_response: BedrockGuardrailResponse = ( - BedrockGuardrailResponse() + self.convert_to_bedrock_format(source=source, messages=messages, response=response) ) + bedrock_guardrail_response: BedrockGuardrailResponse = BedrockGuardrailResponse() api_key: Optional[str] = None if request_data: - dynamic_request_body_params = ( - self.get_guardrail_dynamic_request_body_params( - request_data=request_data - ) - ) + dynamic_request_body_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data) bedrock_request_data.update( { key: value @@ -733,11 +678,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if logging_event_type is not None: event_type = logging_event_type else: - event_type = ( - GuardrailEventHooks.pre_call - if source == "INPUT" - else GuardrailEventHooks.post_call - ) + event_type = GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call try: httpx_response = await self.async_handler.post( @@ -768,15 +709,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): duration=(datetime.now() - start_time).total_seconds(), event_type=event_type, ) - raise HTTPException( - status_code=status_code, detail=detail_message - ) from e + raise HTTPException(status_code=status_code, detail=detail_message) from e except HTTPException: raise # Endpoint down, timeout, or other HTTP/network errors - verbose_proxy_logger.error( - "Bedrock AI: failed to make guardrail request: %s", str(e) - ) + verbose_proxy_logger.error("Bedrock AI: failed to make guardrail request: %s", str(e)) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, @@ -801,9 +738,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): guardrail_provider=self.guardrail_provider, guardrail_json_response=_json_response, request_data=request_data or {}, - guardrail_status=self._get_bedrock_guardrail_response_status( - response=httpx_response - ), + guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), @@ -818,16 +753,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): redact_nested_match_and_regex_keys(_json_response), ) bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) - if self._should_raise_guardrail_blocked_exception( - bedrock_guardrail_response - ): - raise self._get_http_exception_for_blocked_guardrail( - bedrock_guardrail_response - ) + if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): + raise self._get_http_exception_for_blocked_guardrail(bedrock_guardrail_response) else: - status_code, detail_message = self._parse_bedrock_guardrail_error_response( - httpx_response - ) + status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) verbose_proxy_logger.error( "Bedrock AI: error in response. Status code: %s, response: %s", httpx_response.status_code, @@ -870,9 +799,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return "Exception" in payload.get("Output", {}).get("__type", "") - def _get_bedrock_guardrail_response_status( - self, response: httpx.Response - ) -> GuardrailStatus: + def _get_bedrock_guardrail_response_status(self, response: httpx.Response) -> GuardrailStatus: """ Get the status of the bedrock guardrail response. @@ -889,9 +816,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): try: _json_response = response.json() bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) - if self._should_raise_guardrail_blocked_exception( - bedrock_guardrail_response - ): + if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): return "guardrail_intervened" except Exception: pass @@ -899,9 +824,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return "success" return "guardrail_failed_to_respond" - def _parse_bedrock_guardrail_error_response( - self, response: httpx.Response - ) -> Tuple[int, str]: + def _parse_bedrock_guardrail_error_response(self, response: httpx.Response) -> Tuple[int, str]: """ Parse AWS Bedrock guardrail error response body to extract status code and message. @@ -927,9 +850,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return (status_code, err) return (status_code, message) - def _build_tracing_detail( - self, response: BedrockGuardrailResponse - ) -> GuardrailTracingDetail: + def _build_tracing_detail(self, response: BedrockGuardrailResponse) -> GuardrailTracingDetail: """ Build the tracing detail from the raw Bedrock response, before redaction, so downstream loggers (OTEL, Langfuse, ...) get the @@ -948,9 +869,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): tracing_detail["guardrail_action"] = bedrock_action return tracing_detail - def _extract_violation_category_names( - self, response: BedrockGuardrailResponse - ) -> List[str]: + def _extract_violation_category_names(self, response: BedrockGuardrailResponse) -> List[str]: """ Flatten the BLOCKED assessments into a list of human-readable category names suitable for queryable OTEL / standard-logging attributes. @@ -976,9 +895,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): names.append(label) return names - def _extract_blocked_assessments( - self, response: BedrockGuardrailResponse - ) -> List[dict]: + def _extract_blocked_assessments(self, response: BedrockGuardrailResponse) -> List[dict]: """ Walk the Bedrock guardrail response and emit a structured list of BLOCKED assessment entries describing exactly which policies fired. @@ -1024,9 +941,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if f.get("action") == "BLOCKED" ] if content_matches: - blocked.append( - {"policy": "contentPolicy", "matches": content_matches} - ) + blocked.append({"policy": "contentPolicy", "matches": content_matches}) # Word policy word_policy = assessment.get("wordPolicy") @@ -1118,18 +1033,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Get the HTTP exception for a blocked guardrail. """ bedrock_guardrail_output_text: str = "" - outputs: Optional[List[BedrockGuardrailOutput]] = ( - response.get("outputs", []) or [] - ) + outputs: Optional[List[BedrockGuardrailOutput]] = response.get("outputs", []) or [] if outputs: for output in outputs: if output.get("text"): bedrock_guardrail_output_text += output.get("text") or "" if self.disable_exception_on_block is True: - return GuardrailInterventionNormalStringError( - message=bedrock_guardrail_output_text - ) + return GuardrailInterventionNormalStringError(message=bedrock_guardrail_output_text) detail: Dict[str, Any] = { "error": "Violated guardrail policy", @@ -1146,9 +1057,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return HTTPException(status_code=400, detail=detail) - def _should_raise_guardrail_blocked_exception( - self, response: BedrockGuardrailResponse - ) -> bool: + def _should_raise_guardrail_blocked_exception(self, response: BedrockGuardrailResponse) -> bool: """ Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions. @@ -1244,9 +1153,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): data: dict, call_type: CallTypesLiteral, ) -> Union[Exception, str, dict, None]: - verbose_proxy_logger.debug( - "Inside Bedrock Pre-Call Hook for call_type: %s", call_type - ) + verbose_proxy_logger.debug("Inside Bedrock Pre-Call Hook for call_type: %s", call_type) from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -1263,26 +1170,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Handle None case if new_messages is None: - verbose_proxy_logger.debug( - "No messages found for call_type, skipping guardrail" - ) + verbose_proxy_logger.debug("No messages found for call_type, skipping guardrail") return data filter_result = self._prepare_guardrail_messages_for_role(messages=new_messages) filtered_messages = filter_result.payload_messages if not filtered_messages: - verbose_proxy_logger.debug( - "No user-role messages available for guardrail payload" - ) + verbose_proxy_logger.debug("No user-role messages available for guardrail payload") return data ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", @@ -1307,16 +1208,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): target_indices=filter_result.target_indices, ) if isinstance(bedrock_guardrail_response, str): - data["mock_response"] = self.create_guardrail_blocked_response( - response=bedrock_guardrail_response - ) + data["mock_response"] = self.create_guardrail_blocked_response(response=bedrock_guardrail_response) ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## ######################################################### - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data async def async_moderation_hook( @@ -1339,25 +1236,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if new_messages is None: - verbose_proxy_logger.warning( - "Bedrock AI: not running guardrail. No messages in data" - ) + verbose_proxy_logger.warning("Bedrock AI: not running guardrail. No messages in data") return filter_result = self._prepare_guardrail_messages_for_role(messages=new_messages) filtered_messages = filter_result.payload_messages if not filtered_messages: - verbose_proxy_logger.debug( - "Bedrock AI: not running guardrail. No user-role messages" - ) + verbose_proxy_logger.debug("Bedrock AI: not running guardrail. No user-role messages") return ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", @@ -1382,16 +1273,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): target_indices=filter_result.target_indices, ) if isinstance(bedrock_guardrail_response, str): - data["mock_response"] = self.create_guardrail_blocked_response( - response=bedrock_guardrail_response - ) + data["mock_response"] = self.create_guardrail_blocked_response(response=bedrock_guardrail_response) ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## ######################################################### - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -1406,19 +1293,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) from litellm.types.guardrails import GuardrailEventHooks - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return new_messages: Optional[List[AllMessageValues]] = data.get("messages") if new_messages is None: - verbose_proxy_logger.warning( - "Bedrock AI: not running guardrail. No messages in data" - ) + verbose_proxy_logger.warning("Bedrock AI: not running guardrail. No messages in data") return # Check if the ModelResponse has text content in its choices @@ -1427,16 +1307,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): has_text_content = False for choice in response.choices: if isinstance(choice, litellm.Choices): - if choice.message.content and isinstance( - choice.message.content, str - ): + if choice.message.content and isinstance(choice.message.content, str): has_text_content = True break if not has_text_content: - verbose_proxy_logger.warning( - "Bedrock AI: not running guardrail. No output text in response" - ) + verbose_proxy_logger.warning("Bedrock AI: not running guardrail. No output text in response") return ######################################################### @@ -1463,9 +1339,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ########## 2. Apply masking to response with output guardrail response ########## ######################################################### if isinstance(output_content_bedrock, str): - response = self.create_guardrail_blocked_response( - response=output_content_bedrock - ) + response = self.create_guardrail_blocked_response(response=output_content_bedrock) elif output_content_bedrock is not None: self._apply_masking_to_response( response=response, @@ -1475,9 +1349,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## ######################################################### - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) ########### HELPER FUNCTIONS for bedrock guardrails ############################ ############################################################################## @@ -1500,19 +1372,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if isinstance(bedrock_guardrail_response, str): return messages # Get masked texts from guardrail response - masked_texts = self._extract_masked_texts_from_response( - bedrock_guardrail_response - ) + masked_texts = self._extract_masked_texts_from_response(bedrock_guardrail_response) # If guardrail provided masked output, use it regardless of masking flags # because the guardrail has already determined this content needs anonymization if masked_texts: - verbose_proxy_logger.debug( - "Bedrock guardrail provided masked output, applying to messages" - ) - return self._apply_masking_to_messages( - messages=messages, masked_texts=masked_texts - ) + verbose_proxy_logger.debug("Bedrock guardrail provided masked output, applying to messages") + return self._apply_masking_to_messages(messages=messages, masked_texts=masked_texts) # If masking is enabled but no masked texts available, still try to apply # (this maintains backward compatibility for edge cases) @@ -1545,9 +1411,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async for chunk in response: all_chunks.append(chunk) - assembled_model_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = stream_chunk_builder( + assembled_model_response: Optional[Union[ModelResponse, TextCompletionResponse]] = stream_chunk_builder( chunks=all_chunks, ) if isinstance(assembled_model_response, ModelResponse): @@ -1558,9 +1422,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # pre_call / during_call. Bedrock will raise if the response # violates the guardrail policy. ################################################################### - output_guardrail_response: Optional[ - Union[BedrockGuardrailResponse, str] - ] = None + output_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = None try: output_guardrail_response = await self.make_bedrock_api_request( source="OUTPUT", @@ -1576,9 +1438,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ########## 2. Apply masking to response with output guardrail response ########## ######################################################################### if isinstance(output_guardrail_response, str): - assembled_model_response = self.create_guardrail_blocked_response( - response=output_guardrail_response - ) + assembled_model_response = self.create_guardrail_blocked_response(response=output_guardrail_response) elif output_guardrail_response is not None: self._apply_masking_to_response( response=assembled_model_response, @@ -1588,9 +1448,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################################### ########## 3. Return the (potentially masked) chunks ########## ######################################################################### - mock_response = MockResponseIterator( - model_response=assembled_model_response - ) + mock_response = MockResponseIterator(model_response=assembled_model_response) # Return the reconstructed stream async for chunk in mock_response: @@ -1599,9 +1457,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): for chunk in all_chunks: yield chunk - def _extract_masked_texts_from_response( - self, bedrock_guardrail_response: BedrockGuardrailResponse - ) -> List[str]: + def _extract_masked_texts_from_response(self, bedrock_guardrail_response: BedrockGuardrailResponse) -> List[str]: """ Extract all masked text outputs from the guardrail response. @@ -1612,9 +1468,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): List of masked text strings """ masked_output_text: List[str] = [] - masked_outputs: Optional[List[BedrockGuardrailOutput]] = ( - bedrock_guardrail_response.get("outputs", []) or [] - ) + masked_outputs: Optional[List[BedrockGuardrailOutput]] = bedrock_guardrail_response.get("outputs", []) or [] if not masked_outputs: verbose_proxy_logger.debug("No masked outputs found in guardrail response") return [] @@ -1712,31 +1566,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): bedrock_guardrail_response: Response from Bedrock guardrail containing masked content """ # Get masked texts from guardrail response - masked_texts = self._extract_masked_texts_from_response( - bedrock_guardrail_response - ) + masked_texts = self._extract_masked_texts_from_response(bedrock_guardrail_response) if not masked_texts: - verbose_proxy_logger.debug( - "No masked outputs found, skipping response masking" - ) + verbose_proxy_logger.debug("No masked outputs found, skipping response masking") return - verbose_proxy_logger.debug( - "Applying masking to response with %d masked texts", len(masked_texts) - ) + verbose_proxy_logger.debug("Applying masking to response with %d masked texts", len(masked_texts)) # Apply masking to ModelResponse if isinstance(response, litellm.ModelResponse): self._apply_masking_to_model_response(response, masked_texts) else: - verbose_proxy_logger.warning( - "Unsupported response type for masking: %s", type(response) - ) + verbose_proxy_logger.warning("Unsupported response type for masking: %s", type(response)) - def _apply_masking_to_model_response( - self, response: litellm.ModelResponse, masked_texts: List[str] - ) -> None: + def _apply_masking_to_model_response(self, response: litellm.ModelResponse, masked_texts: List[str]) -> None: """ Apply masked texts to a ModelResponse object. @@ -1753,27 +1597,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if masking_index < len(masked_texts): choice.message.content = masked_texts[masking_index] masking_index += 1 - verbose_proxy_logger.debug( - "Applied masking to choice message content" - ) + verbose_proxy_logger.debug("Applied masking to choice message content") elif isinstance(choice, StreamingChoices): # For streaming responses, modify delta content if choice.delta.content and isinstance(choice.delta.content, str): if masking_index < len(masked_texts): choice.delta.content = masked_texts[masking_index] masking_index += 1 - verbose_proxy_logger.debug( - "Applied masking to choice delta content" - ) + verbose_proxy_logger.debug("Applied masking to choice delta content") elif isinstance(choice, TextChoices): # For text completions if choice.text and isinstance(choice.text, str): if masking_index < len(masked_texts): choice.text = masked_texts[masking_index] masking_index += 1 - verbose_proxy_logger.debug( - "Applied masking to choice text content" - ) + verbose_proxy_logger.debug("Applied masking to choice text content") async def apply_guardrail( self, @@ -1804,9 +1642,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # dict.get("texts", []) would return None if the key exists with a None value. texts = inputs.get("texts") or [] try: - verbose_proxy_logger.debug( - f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)" - ) + verbose_proxy_logger.debug(f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)") masked_texts = [] @@ -1824,20 +1660,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Bedrock will throw an error if there is no text to process if filtered_messages: - _log_hook = ( - GuardrailEventHooks.pre_call - if input_type == "request" - else GuardrailEventHooks.post_call - ) + _log_hook = GuardrailEventHooks.pre_call if input_type == "request" else GuardrailEventHooks.post_call # Map the abstract input_type to the Bedrock source parameter. # "request" -> INPUT (scan user-supplied content) # "response" -> OUTPUT (scan model-generated content) # Bedrock guardrail policies are often configured differently # for Input vs Output (e.g. PII blocking only on Output), so # the source MUST match where the text originated. - bedrock_source: Literal["INPUT", "OUTPUT"] = ( - "OUTPUT" if input_type == "response" else "INPUT" - ) + bedrock_source: Literal["INPUT", "OUTPUT"] = "OUTPUT" if input_type == "response" else "INPUT" if bedrock_source == "OUTPUT": # Build a synthetic ModelResponse whose choices carry the # text(s) to scan, so _create_bedrock_output_content_request @@ -1897,9 +1727,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): scanned_role_subset=scanned_role_subset, ) - verbose_proxy_logger.debug( - "Bedrock Guardrail: Successfully applied guardrail" - ) + verbose_proxy_logger.debug("Bedrock Guardrail: Successfully applied guardrail") inputs["texts"] = masked_texts return inputs @@ -1913,7 +1741,5 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # from properly blocking the call. raise except Exception as e: - verbose_proxy_logger.error( - "Bedrock Guardrail: Failed to apply guardrail: %s", str(e) - ) + verbose_proxy_logger.error("Bedrock Guardrail: Failed to apply guardrail: %s", str(e)) raise Exception(f"Bedrock guardrail failed: {str(e)}") diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py index b8a2111c011..40ed634d39a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -57,9 +57,7 @@ def initialize_guardrail( _get_param(litellm_params, guardrail, "confidence_threshold", 0.5), ) ) - detect_execution_intent = bool( - _get_param(litellm_params, guardrail, "detect_execution_intent", True) - ) + detect_execution_intent = bool(_get_param(litellm_params, guardrail, "detect_execution_intent", True)) mode = _get_param(litellm_params, guardrail, "mode") event_hook = cast( Optional[Union[Literal["pre_call", "post_call", "during_call"], List[str]]], diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py index 49c3dc00cd3..ea66f416e15 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -51,9 +51,7 @@ LANGUAGE_ALIASES: Dict[str, str] = { } # Tags that indicate non-executable / plain text (lower confidence when block-all) -NON_EXECUTABLE_TAGS: frozenset = frozenset( - {"text", "plaintext", "plain", "markdown", "md", "output", "result"} -) +NON_EXECUTABLE_TAGS: frozenset = frozenset({"text", "plaintext", "plain", "markdown", "md", "output", "result"}) # Regex: fenced code block with optional language tag. Handles ```lang\n...\n``` # Content between fences; does not handle nested ``` inside body (documented edge case). @@ -340,22 +338,15 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): action: Literal["block", "mask"] = "block", confidence_threshold: float = 0.5, detect_execution_intent: bool = True, - event_hook: Optional[ - Union[Literal["pre_call", "post_call", "during_call"], List[str]] - ] = None, + event_hook: Optional[Union[Literal["pre_call", "post_call", "during_call"], List[str]]] = None, default_on: bool = False, **kwargs: Any, ) -> None: # Normalize to type expected by CustomGuardrail - _event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = ( - None - ) + _event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = None if event_hook is not None: if isinstance(event_hook, list): - _event_hook = [ - GuardrailEventHooks(h) if isinstance(h, str) else h - for h in event_hook - ] + _event_hook = [GuardrailEventHooks(h) if isinstance(h, str) else h for h in event_hook] else: _event_hook = GuardrailEventHooks(event_hook) super().__init__( @@ -387,9 +378,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): return BlockCodeExecutionGuardrailConfigModel - def _find_blocks( - self, text: str - ) -> List[Tuple[int, int, str, str, float, CodeBlockActionTaken]]: + def _find_blocks(self, text: str) -> List[Tuple[int, int, str, str, float, CodeBlockActionTaken]]: """ Find all fenced code blocks in text. Returns list of (start, end, language_tag, block_content, confidence, action_taken). @@ -401,9 +390,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): tag_in_list = not self.block_all and _normalize_language(tag) in [ _normalize_language(t) for t in (self.blocked_languages or []) ] - is_blocked = _is_blocked_language( - tag, self.blocked_languages, self.block_all - ) + is_blocked = _is_blocked_language(tag, self.blocked_languages, self.block_all) confidence = _confidence_for_block(tag, self.block_all, tag_in_list) if not is_blocked: action_taken: CodeBlockActionTaken = "allow" @@ -411,9 +398,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): action_taken = "block" else: action_taken = "log_only" - results.append( - (m.start(), m.end(), tag or "(none)", body, confidence, action_taken) - ) + results.append((m.start(), m.end(), tag or "(none)", body, confidence, action_taken)) return results def _scan_text( @@ -453,11 +438,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): blocks = self._find_blocks(text) # For requests, check execution intent; for responses, skip this check - has_execution_intent = ( - not is_response - and self.detect_execution_intent - and _has_execution_intent(text) - ) + has_execution_intent = not is_response and self.detect_execution_intent and _has_execution_intent(text) if not blocks: if has_execution_intent and self.action == "block": @@ -493,9 +474,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): "type": "code_block", "language": tag, "confidence": round(confidence, 2), - "action_taken": ( - "block" if effective_block else action_taken - ), + "action_taken": ("block" if effective_block else action_taken), }, ) ) @@ -513,9 +492,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): new_text = "".join(parts) return new_text, should_raise - def _raise_block_error( - self, language: str, is_output: bool, request_data: dict - ) -> None: + def _raise_block_error(self, language: str, is_output: bool, request_data: dict) -> None: if language == "execution_request": msg = "Content blocked: execution request detected" else: @@ -595,11 +572,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): } if max_confidence is not None: tracing_kw["confidence_score"] = max_confidence - event_type = ( - GuardrailEventHooks.pre_call - if input_type == "request" - else GuardrailEventHooks.post_call - ) + event_type = GuardrailEventHooks.pre_call if input_type == "request" else GuardrailEventHooks.post_call self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider="block_code_execution", guardrail_json_response=guardrail_response, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index d8e33e13b36..bf5a0a5f262 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -49,9 +49,7 @@ class CatoNetworksGuardrailMissingSecrets(Exception): class CatoNetworksGuardrail(CustomGuardrail): - def __init__( - self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs - ): + def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs): ssl_verify = kwargs.pop("ssl_verify", None) self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, @@ -64,24 +62,14 @@ class CatoNetworksGuardrail(CustomGuardrail): "pass it as a parameter to the guardrail in the config file" ) raise CatoNetworksGuardrailMissingSecrets(msg) - self.api_base = ( - api_base - or os.environ.get("CATO_API_BASE") - or "https://api.aisec.catonetworks.com" - ) + self.api_base = api_base or os.environ.get("CATO_API_BASE") or "https://api.aisec.catonetworks.com" self.api_base = self.api_base.rstrip("/") - self.ws_api_base = self.api_base.replace("http://", "ws://").replace( - "https://", "wss://" - ) - self._ws_connect_ssl_kwargs = self._build_ws_ssl_kwargs( - ssl_verify, self.ws_api_base - ) + self.ws_api_base = self.api_base.replace("http://", "ws://").replace("https://", "wss://") + self._ws_connect_ssl_kwargs = self._build_ws_ssl_kwargs(ssl_verify, self.ws_api_base) super().__init__(**kwargs) @staticmethod - def _build_ws_ssl_kwargs( - ssl_verify: Optional[Union[bool, str]], ws_api_base: str - ) -> dict: + def _build_ws_ssl_kwargs(ssl_verify: Optional[Union[bool, str]], ws_api_base: str) -> dict: """Resolve the ``ssl`` argument for ``websockets.connect``. Mirrors the ``ssl_verify`` handling applied to the HTTP handler so a custom Cato instance behind TLS honours the same verification settings for streaming.""" @@ -149,9 +137,7 @@ class CatoNetworksGuardrail(CustomGuardrail): for message in data.get("messages") or []: if isinstance(message, dict) and isinstance(message.get("content"), list): parts = build_inspection_messages({"messages": [message]}) - flattened.append( - {**message, "content": parts[0]["content"] if parts else ""} - ) + flattened.append({**message, "content": parts[0]["content"] if parts else ""}) else: flattened.append(message) for _field, messages in cls._extra_inspection_sources(data): @@ -165,11 +151,7 @@ class CatoNetworksGuardrail(CustomGuardrail): if isinstance(prompt, str): return [{"role": "user", "content": prompt}] if prompt else [] if isinstance(prompt, list): - return [ - {"role": "user", "content": part} - for part in prompt - if isinstance(part, str) and part - ] + return [{"role": "user", "content": part} for part in prompt if isinstance(part, str) and part] return [] @staticmethod @@ -227,15 +209,12 @@ class CatoNetworksGuardrail(CustomGuardrail): sources.append(("input", input_messages)) instructions = data.get("instructions") if isinstance(instructions, str) and instructions: - sources.append( - ("instructions", [{"role": "system", "content": instructions}]) - ) + sources.append(("instructions", [{"role": "system", "content": instructions}])) prompt_messages = cls._prompt_inspection_messages(data.get("prompt")) if prompt_messages: sources.append(("prompt", prompt_messages)) schema_strings = [ - {"role": "system", "content": container[key]} - for container, key in cls._iter_schema_string_refs(data) + {"role": "system", "content": container[key]} for container, key in cls._iter_schema_string_refs(data) ] if schema_strings: sources.append(("schema_strings", schema_strings)) @@ -298,8 +277,7 @@ class CatoNetworksGuardrail(CustomGuardrail): data["messages"] = [ ( {**original, "content": redacted_messages[idx]["content"]} - if idx < len(redacted_messages) - and redacted_messages[idx].get("content") is not None + if idx < len(redacted_messages) and redacted_messages[idx].get("content") is not None else original ) for idx, original in enumerate(original_messages) @@ -371,19 +349,14 @@ class CatoNetworksGuardrail(CustomGuardrail): user_email=user_email, litellm_call_id=call_id, ), - json={ - "messages": inspection_messages - + [{"role": "assistant", "content": output}] - }, + json={"messages": inspection_messages + [{"role": "assistant", "content": output}]}, ) response.raise_for_status() res = response.json() required_action = res.get("required_action") action_type = required_action and required_action.get("action_type", None) if action_type and action_type == "block_action": - self._handle_block_action_on_output( - res.get("analysis_result", {}), required_action - ) + self._handle_block_action_on_output(res.get("analysis_result", {}), required_action) redacted_chat = res.get("redacted_chat", None) if action_type and action_type == "anonymize_action" and redacted_chat: @@ -394,9 +367,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return {"redacted_output": redacted_output} return None - def _handle_block_action_on_output( - self, analysis_result: Any, required_action: Any - ) -> None: + def _handle_block_action_on_output(self, analysis_result: Any, required_action: Any) -> None: detection_message = required_action.get("detection_message", None) verbose_proxy_logger.info( "Cato: detected: {detected}, enabled policies: {policies}".format( @@ -492,9 +463,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return fragments @staticmethod - def _apply_responses_output_fragment( - container: Any, key: str, redacted: str - ) -> None: + def _apply_responses_output_fragment(container: Any, key: str, redacted: str) -> None: if isinstance(container, dict): container[key] = redacted else: @@ -533,22 +502,14 @@ class CatoNetworksGuardrail(CustomGuardrail): if not isinstance(choice, Choices): continue for target, text in self._output_fragments(choice.message): - redacted_output = await self._inspect_output_text( - data, text, user_api_key_dict, user_email - ) + redacted_output = await self._inspect_output_text(data, text, user_api_key_dict, user_email) if redacted_output is not None: - self._apply_output_fragment( - choice.message, target, redacted_output - ) + self._apply_output_fragment(choice.message, target, redacted_output) elif isinstance(response, ResponsesAPIResponse): for container, key, text in self._responses_output_fragments(response): - redacted_output = await self._inspect_output_text( - data, text, user_api_key_dict, user_email - ) + redacted_output = await self._inspect_output_text(data, text, user_api_key_dict, user_email) if redacted_output is not None: - self._apply_responses_output_fragment( - container, key, redacted_output - ) + self._apply_responses_output_fragment(container, key, redacted_output) return response async def async_post_call_streaming_iterator_hook( @@ -571,9 +532,7 @@ class CatoNetworksGuardrail(CustomGuardrail): ), **self._ws_connect_ssl_kwargs, ) as websocket: - sender = asyncio.create_task( - self.forward_the_stream_to_cato(websocket, response) - ) + sender = asyncio.create_task(self.forward_the_stream_to_cato(websocket, response)) try: while True: raw_message = await self._await_cato_message(websocket, sender) @@ -585,16 +544,12 @@ class CatoNetworksGuardrail(CustomGuardrail): return if blocking_message := result.get("blocking_message"): raise StreamingCallbackError(blocking_message) - verbose_proxy_logger.error( - f"Unknown message received from Cato: {result}" - ) + verbose_proxy_logger.error(f"Unknown message received from Cato: {result}") return finally: await self._cancel_background_task(sender) - async def _await_cato_message( - self, websocket: ClientConnection, sender: asyncio.Task - ) -> Any: + async def _await_cato_message(self, websocket: ClientConnection, sender: asyncio.Task) -> Any: """Wait for the next Cato message, surfacing a dead forwarding task instead of blocking.""" from litellm.proxy.proxy_server import StreamingCallbackError @@ -603,15 +558,11 @@ class CatoNetworksGuardrail(CustomGuardrail): await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED) if sender.done() and (sender_exc := sender.exception()) is not None: await self._cancel_background_task(recv_task) - raise StreamingCallbackError( - "Cato guardrail upstream stream failed" - ) from sender_exc + raise StreamingCallbackError("Cato guardrail upstream stream failed") from sender_exc try: return await recv_task except ConnectionClosed as exc: - raise StreamingCallbackError( - "Cato guardrail connection closed unexpectedly" - ) from exc + raise StreamingCallbackError("Cato guardrail connection closed unexpectedly") from exc async def forward_the_stream_to_cato( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/__init__.py index 774a0334072..4e191c3db52 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/__init__.py @@ -27,33 +27,15 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" guardrail_name=guardrail_name, api_key=litellm_params.api_key, api_base=litellm_params.api_base, - inspection_type=_get_optional_value( - litellm_params, optional_params, "inspection_type" - ), - inspect_path=_get_optional_value( - litellm_params, optional_params, "inspect_path" - ), - enabled_rules=_get_optional_value( - litellm_params, optional_params, "enabled_rules" - ), - integration_profile_id=_get_optional_value( - litellm_params, optional_params, "integration_profile_id" - ), - integration_profile_version=_get_optional_value( - litellm_params, optional_params, "integration_profile_version" - ), - integration_tenant_id=_get_optional_value( - litellm_params, optional_params, "integration_tenant_id" - ), - integration_type=_get_optional_value( - litellm_params, optional_params, "integration_type" - ), - on_flagged_action=_get_optional_value( - litellm_params, optional_params, "on_flagged_action" - ), - fallback_on_error=_get_optional_value( - litellm_params, optional_params, "fallback_on_error" - ), + inspection_type=_get_optional_value(litellm_params, optional_params, "inspection_type"), + inspect_path=_get_optional_value(litellm_params, optional_params, "inspect_path"), + enabled_rules=_get_optional_value(litellm_params, optional_params, "enabled_rules"), + integration_profile_id=_get_optional_value(litellm_params, optional_params, "integration_profile_id"), + integration_profile_version=_get_optional_value(litellm_params, optional_params, "integration_profile_version"), + integration_tenant_id=_get_optional_value(litellm_params, optional_params, "integration_tenant_id"), + integration_type=_get_optional_value(litellm_params, optional_params, "integration_type"), + on_flagged_action=_get_optional_value(litellm_params, optional_params, "on_flagged_action"), + fallback_on_error=_get_optional_value(litellm_params, optional_params, "fallback_on_error"), timeout=_get_optional_value(litellm_params, optional_params, "timeout"), event_hook=litellm_params.mode, default_on=litellm_params.default_on or False, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 79e81b6c4c4..561e6ce5f2b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -165,11 +165,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): ) self.api_key: str = resolved_api_key - self.api_base: str = ( - api_base - or os.environ.get("CISCO_AI_DEFENSE_API_BASE") - or CISCO_DEFAULT_API_BASE - ).rstrip("/") + self.api_base: str = (api_base or os.environ.get("CISCO_AI_DEFENSE_API_BASE") or CISCO_DEFAULT_API_BASE).rstrip( + "/" + ) self.inspection_type: str = self._resolve_choice( value=inspection_type, @@ -179,34 +177,21 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): setting_name="inspection_type", ) - inferred = self._infer_inspection_type_from_mode( - kwargs.get("event_hook"), self.inspection_type - ) + inferred = self._infer_inspection_type_from_mode(kwargs.get("event_hook"), self.inspection_type) if inferred != self.inspection_type: verbose_proxy_logger.info( - "Cisco AI Defense: inferred inspection_type=%s from " - "MCP-only event_hook configuration (was %s)", + "Cisco AI Defense: inferred inspection_type=%s from MCP-only event_hook configuration (was %s)", inferred, self.inspection_type, ) self.inspection_type = inferred if inspect_path: - self.inspect_path = ( - inspect_path if inspect_path.startswith("/") else f"/{inspect_path}" - ) + self.inspect_path = inspect_path if inspect_path.startswith("/") else f"/{inspect_path}" else: - self.inspect_path = ( - CISCO_MCP_INSPECT_PATH - if self.inspection_type == "mcp" - else CISCO_CHAT_INSPECT_PATH - ) + self.inspect_path = CISCO_MCP_INSPECT_PATH if self.inspection_type == "mcp" else CISCO_CHAT_INSPECT_PATH - self.enabled_rules = ( - [self._normalize_rule(rule) for rule in enabled_rules] - if enabled_rules - else None - ) + self.enabled_rules = [self._normalize_rule(rule) for rule in enabled_rules] if enabled_rules else None self.integration_profile_id = integration_profile_id self.integration_profile_version = integration_profile_version self.integration_tenant_id = integration_tenant_id @@ -233,18 +218,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): resolved_timeout = self._coerce_timeout(timeout) else: env_timeout = os.environ.get("CISCO_AI_DEFENSE_TIMEOUT") - resolved_timeout = ( - self._coerce_timeout(env_timeout) if env_timeout is not None else None - ) - self.timeout: float = ( - resolved_timeout - if resolved_timeout is not None - else DEFAULT_TIMEOUT_SECONDS - ) + resolved_timeout = self._coerce_timeout(env_timeout) if env_timeout is not None else None + self.timeout: float = resolved_timeout if resolved_timeout is not None else DEFAULT_TIMEOUT_SECONDS - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) # Register broadly; runtime filtering happens in ``_surface_matches``. supported_event_hooks = [ @@ -295,8 +272,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if candidate in allowed: return candidate verbose_proxy_logger.warning( - "Cisco AI Defense guardrail: invalid value '%s' for %s, falling " - "back to default '%s'. Allowed values: %s", + "Cisco AI Defense guardrail: invalid value '%s' for %s, falling back to default '%s'. Allowed values: %s", candidate, setting_name, default, @@ -310,8 +286,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): parsed = float(value) except (TypeError, ValueError): verbose_proxy_logger.warning( - "Cisco AI Defense guardrail: invalid timeout value '%s', " - "using default %ss", + "Cisco AI Defense guardrail: invalid timeout value '%s', using default %ss", value, DEFAULT_TIMEOUT_SECONDS, ) @@ -354,29 +329,23 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if not self._surface_matches(is_mcp): verbose_proxy_logger.debug( - "Cisco AI Defense guardrail: call_type=%s does not match " - "configured inspection_type=%s, skipping", + "Cisco AI Defense guardrail: call_type=%s does not match configured inspection_type=%s, skipping", call_type, self.inspection_type, ) return data - event_type = ( - GuardrailEventHooks.pre_mcp_call if is_mcp else GuardrailEventHooks.pre_call - ) + event_type = GuardrailEventHooks.pre_mcp_call if is_mcp else GuardrailEventHooks.pre_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data if is_mcp: - await self._inspect_mcp_request( - data=data, user_api_key_dict=user_api_key_dict - ) + await self._inspect_mcp_request(data=data, user_api_key_dict=user_api_key_dict) else: messages = self._extract_inspect_messages_from_request(data) if not messages: verbose_proxy_logger.debug( - "Cisco AI Defense guardrail: no scannable messages in " - "pre-call request, skipping" + "Cisco AI Defense guardrail: no scannable messages in pre-call request, skipping" ) return data await self._inspect_chat( @@ -385,9 +354,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): user_api_key_dict=user_api_key_dict, ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @log_guardrail_information @@ -411,18 +378,12 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if not self._surface_matches(is_mcp): return data - event_type = ( - GuardrailEventHooks.during_mcp_call - if is_mcp - else GuardrailEventHooks.during_call - ) + event_type = GuardrailEventHooks.during_mcp_call if is_mcp else GuardrailEventHooks.during_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data if is_mcp: - await self._inspect_mcp_request( - data=data, user_api_key_dict=user_api_key_dict - ) + await self._inspect_mcp_request(data=data, user_api_key_dict=user_api_key_dict) else: messages = self._extract_inspect_messages_from_request(data) if not messages: @@ -433,9 +394,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): user_api_key_dict=user_api_key_dict, ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @log_guardrail_information @@ -448,19 +407,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if self.inspection_type != "chat": return response - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return response response_messages = self._extract_response_messages(response) if not response_messages: verbose_proxy_logger.debug( - "Cisco AI Defense guardrail: no response content to scan, " - "skipping post-call analysis" + "Cisco AI Defense guardrail: no response content to scan, skipping post-call analysis" ) return response @@ -475,9 +428,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): response_obj=response, ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response async def async_post_call_streaming_iterator_hook( @@ -495,12 +446,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): yield chunk return - if ( - self.should_run_guardrail( - data=request_data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: async for chunk in response: yield chunk return @@ -526,8 +472,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if not isinstance(all_chunks[0], (ModelResponse, ModelResponseStream)): verbose_proxy_logger.warning( - "Cisco AI Defense guardrail (%s): unsupported streaming " - "chunk shape (%s) — failing closed.", + "Cisco AI Defense guardrail (%s): unsupported streaming chunk shape (%s) — failing closed.", self.guardrail_name, type(all_chunks[0]).__name__, ) @@ -551,13 +496,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): response_messages = self._extract_response_messages(assembled) original_stream_text = self._extract_streaming_chunk_scan_text(all_chunks) - assembled_text = " ".join( - m.get("content", "") for m in response_messages if isinstance(m, dict) - ) + assembled_text = " ".join(m.get("content", "") for m in response_messages if isinstance(m, dict)) if original_stream_text and original_stream_text not in assembled_text: - response_messages.append( - {"role": "assistant", "content": original_stream_text} - ) + response_messages.append({"role": "assistant", "content": original_stream_text}) if not response_messages: for chunk in all_chunks: yield chunk @@ -591,9 +532,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): exc, ) error_obj = { - "message": ( - "Cisco AI Defense streaming scan failed — response withheld." - ), + "message": ("Cisco AI Defense streaming scan failed — response withheld."), "type": "guardrail_scan_error", "code": 500, "guardrail": self.guardrail_name, @@ -601,9 +540,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): yield f"data: {json.dumps({'error': error_obj})}\n\n" return - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) if self._streaming_content_was_modified(all_chunks, assembled): mock_iterator = MockResponseIterator(model_response=assembled) @@ -613,9 +550,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): for chunk in all_chunks: yield chunk - def _build_block_payload( - self, context: _ScanContext, verdict: _CiscoVerdict - ) -> Dict[str, Any]: + def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> Dict[str, Any]: """Canonical block payload used across all four block paths. Same dict is the ``HTTPException.detail`` for chat / MCP request @@ -646,25 +581,17 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): payload, so this is a near-passthrough that just adds ``code`` / ``guardrail`` defaults for non-Cisco / unstructured details. """ - error_obj: Dict[str, Any] = ( - dict(exc.detail) - if isinstance(exc.detail, dict) - else {"message": str(exc.detail)} - ) + error_obj: Dict[str, Any] = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)} error_obj.setdefault("message", error_obj.get("error", "Guardrail block")) error_obj.setdefault("code", exc.status_code) error_obj.setdefault("guardrail", self.guardrail_name) return error_obj @classmethod - def _streaming_content_was_modified( - cls, original_chunks: List[Any], assembled: ModelResponse - ) -> bool: + def _streaming_content_was_modified(cls, original_chunks: List[Any], assembled: ModelResponse) -> bool: """Decide whether redact changed content or tool/function arguments.""" original_text = cls._extract_streaming_chunk_scan_text(original_chunks) - assembled_text = " ".join( - m.get("content", "") for m in cls._extract_response_messages(assembled) - ) + assembled_text = " ".join(m.get("content", "") for m in cls._extract_response_messages(assembled)) return original_text != assembled_text @classmethod @@ -772,9 +699,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): "action": verdict.action, "is_safe": verdict.is_safe, "severity": verdict.severity, - "classifications": ( - list(verdict.classifications) if verdict.classifications else [] - ), + "classifications": (list(verdict.classifications) if verdict.classifications else []), "rule_violations": sorted( { rule.get("rule_name") @@ -800,9 +725,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): fields[target_key] = value payload = {k: v for k, v in fields.items() if v not in (None, [], "")} - line = "CISCO_AI_DEFENSE_DECISION " + json.dumps( - payload, default=str, sort_keys=True, separators=(",", ":") - ) + line = "CISCO_AI_DEFENSE_DECISION " + json.dumps(payload, default=str, sort_keys=True, separators=(",", ":")) if verdict.action == _ACTION_ALLOW: verbose_proxy_logger.info(line) @@ -853,9 +776,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): payload = self._build_chat_payload(messages, request_data, user_api_key_dict) start_time = datetime.now() try: - inspect_response = await self._post_inspection( - url=url, payload=payload, surface="chat" - ) + inspect_response = await self._post_inspection(url=url, payload=payload, surface="chat") except HTTPException: # Re-raise; _post_inspection only raises CiscoAIDefenseGuardrailAPIError, # but be defensive in case downstream evolves. @@ -926,17 +847,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): except Exception: body_snippet = "" raise CiscoAIDefenseGuardrailAPIError( - f"Cisco AI Defense {surface} API returned HTTP {status_code}: " - f"{body_snippet}" + f"Cisco AI Defense {surface} API returned HTTP {status_code}: {body_snippet}" ) from exc except httpx.TimeoutException as exc: raise CiscoAIDefenseGuardrailAPIError( f"Cisco AI Defense {surface} API call timed out after {self.timeout}s" ) from exc except httpx.RequestError as exc: - raise CiscoAIDefenseGuardrailAPIError( - f"Cisco AI Defense {surface} API request failed: {exc}" - ) from exc + raise CiscoAIDefenseGuardrailAPIError(f"Cisco AI Defense {surface} API request failed: {exc}") from exc try: return response.json() @@ -1038,9 +956,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): normalized["classification"] = classification return normalized - raise ValueError( - f"Cisco AI Defense guardrail: invalid rule definition: {rule!r}" - ) + raise ValueError(f"Cisco AI Defense guardrail: invalid rule definition: {rule!r}") # ------------------------------------------------------------------ # Response processing @@ -1078,15 +994,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): jsonrpc_error = self._extract_jsonrpc_error(inspect_response) if jsonrpc_error is not None: verbose_proxy_logger.warning( - "Cisco AI Defense guardrail: API returned JSON-RPC error " - "envelope (code=%s message=%s)", + "Cisco AI Defense guardrail: API returned JSON-RPC error envelope (code=%s message=%s)", jsonrpc_error.get("code"), jsonrpc_error.get("message"), ) return self._handle_api_error( CiscoAIDefenseGuardrailAPIError( - f"AI Defense error code={jsonrpc_error.get('code')} " - f"message={jsonrpc_error.get('message')}" + f"AI Defense error code={jsonrpc_error.get('code')} message={jsonrpc_error.get('message')}" ), request_data=request_data, start_time=start_time, @@ -1103,11 +1017,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): # examples & SDK return `classifications` (plural). Accept both. classifications = ( verdict_dict.get("classifications") - or ( - [verdict_dict["classification"]] - if verdict_dict.get("classification") - else [] - ) + or ([verdict_dict["classification"]] if verdict_dict.get("classification") else []) or [] ) verdict = _CiscoVerdict( @@ -1140,9 +1050,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): ) else: logging_event_type = ( - GuardrailEventHooks.post_call - if context.direction == "output" - else GuardrailEventHooks.pre_call + GuardrailEventHooks.post_call if context.direction == "output" else GuardrailEventHooks.pre_call ) self.add_standard_logging_guardrail_information_to_request_data( @@ -1151,11 +1059,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): inspect_response, surface=context.surface, action=action ), request_data=request_data, - guardrail_status=( - "guardrail_intervened" - if action in (_ACTION_BLOCK, _ACTION_REDACT) - else "success" - ), + guardrail_status=("guardrail_intervened" if action in (_ACTION_BLOCK, _ACTION_REDACT) else "success"), start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, @@ -1171,9 +1075,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return inspect_response if action == _ACTION_REDACT: - redacted = self._apply_redaction( - request_data, response_obj, context, verdict - ) + redacted = self._apply_redaction(request_data, response_obj, context, verdict) if redacted: verbose_proxy_logger.info( "Cisco AI Defense guardrail (%s): redaction applied (event_id=%s)", @@ -1196,17 +1098,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): ) verbose_proxy_logger.info( - "Cisco AI Defense guardrail (%s): violation in monitor mode — " - "request allowed to proceed (event_id=%s)", + "Cisco AI Defense guardrail (%s): violation in monitor mode — request allowed to proceed (event_id=%s)", context.surface, verdict.event_id, ) return inspect_response @staticmethod - def _stash_verdict_on_request( - request_data: dict, context: _ScanContext, verdict: _CiscoVerdict - ) -> None: + def _stash_verdict_on_request(request_data: dict, context: _ScanContext, verdict: _CiscoVerdict) -> None: """Surface the Cisco verdict on the request metadata for observability.""" metadata_store = request_data.setdefault("metadata", {}) if not isinstance(metadata_store, dict): @@ -1221,9 +1120,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): metadata_store[f"{prefix}_severity"] = verdict.severity if verdict.rules: metadata_store[f"{prefix}_rules"] = [ - rule.get("rule_name") - for rule in verdict.rules - if isinstance(rule, dict) + rule.get("rule_name") for rule in verdict.rules if isinstance(rule, dict) ] if verdict.event_id: metadata_store[f"{prefix}_event_id"] = verdict.event_id @@ -1307,9 +1204,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return any(key in payload for key in cls._DECISION_FIELDS) @classmethod - def _unwrap_verdict_envelope( - cls, inspect_response: Dict[str, Any] - ) -> Dict[str, Any]: + def _unwrap_verdict_envelope(cls, inspect_response: Dict[str, Any]) -> Dict[str, Any]: """Return the dict that actually holds is_safe / action / rules. Cisco AI Defense returns the verdict at different nesting depths @@ -1461,25 +1356,17 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): ``on_flagged_action``). """ if context.surface == "mcp" and context.direction == "input": - return self._redact_mcp_input( - request_data, verdict.sanitized_text, verdict.sanitized_mcp_arguments - ) + return self._redact_mcp_input(request_data, verdict.sanitized_text, verdict.sanitized_mcp_arguments) if context.surface == "mcp" and context.direction == "output": if response_obj is None: return False if verdict.sanitized_text: - return self._set_mcp_tool_response_text( - response_obj, verdict.sanitized_text - ) + return self._set_mcp_tool_response_text(response_obj, verdict.sanitized_text) return False if context.surface == "chat" and context.direction == "input": - return self._redact_chat_input( - request_data, verdict.sanitized_text, verdict.sanitized_messages - ) + return self._redact_chat_input(request_data, verdict.sanitized_text, verdict.sanitized_messages) if context.surface == "chat" and context.direction == "output": - return self._redact_chat_output( - response_obj, verdict.sanitized_text, verdict.sanitized_messages - ) + return self._redact_chat_output(response_obj, verdict.sanitized_text, verdict.sanitized_messages) return False @staticmethod @@ -1507,9 +1394,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): ): if not isinstance(args_path, dict): continue - string_keys = [ - key for key, value in args_path.items() if isinstance(value, str) - ] + string_keys = [key for key, value in args_path.items() if isinstance(value, str)] if len(string_keys) != 1: continue args_path[string_keys[0]] = sanitized_text @@ -1543,9 +1428,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return instructions_redacted if sanitized_messages: if uses_input: - rewritten = self._sanitized_messages_to_responses_input( - sanitized_messages - ) + rewritten = self._sanitized_messages_to_responses_input(sanitized_messages) if rewritten is not None: request_data["input"] = rewritten return True @@ -1554,9 +1437,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return True if sanitized_text: if uses_input: - rewritten_input = self._rewrite_responses_input_text( - request_data.get("input"), sanitized_text - ) + rewritten_input = self._rewrite_responses_input_text(request_data.get("input"), sanitized_text) if rewritten_input is not None: request_data["input"] = rewritten_input return True @@ -1589,17 +1470,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if instruction_text: request_data["instructions"] = instruction_text return True - if sanitized_text and not any( - key in request_data for key in ("input", "messages", "prompt") - ): + if sanitized_text and not any(key in request_data for key in ("input", "messages", "prompt")): request_data["instructions"] = sanitized_text return True return False @classmethod - def _instruction_text_from_messages( - cls, messages: List[Dict[str, Any]] - ) -> Optional[str]: + def _instruction_text_from_messages(cls, messages: List[Dict[str, Any]]) -> Optional[str]: for message in messages: if not isinstance(message, dict): continue @@ -1610,18 +1487,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return None @classmethod - def _non_instruction_messages( - cls, messages: Optional[List[Dict[str, Any]]] - ) -> Optional[List[Dict[str, Any]]]: + def _non_instruction_messages(cls, messages: Optional[List[Dict[str, Any]]]) -> Optional[List[Dict[str, Any]]]: if messages is None: return None return [ message for message in messages - if not ( - isinstance(message, dict) - and cls._is_instruction_role(message.get("role")) - ) + if not (isinstance(message, dict) and cls._is_instruction_role(message.get("role"))) ] @staticmethod @@ -1661,15 +1533,11 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): choices = getattr(response_obj, "choices", None) if isinstance(choices, list): - return self._redact_model_response_choices( - choices, sanitized_text, sanitized_messages - ) + return self._redact_model_response_choices(choices, sanitized_text, sanitized_messages) output_items = getattr(response_obj, "output", None) if isinstance(output_items, list): - return self._redact_responses_api_output( - output_items, sanitized_text, sanitized_messages - ) + return self._redact_responses_api_output(output_items, sanitized_text, sanitized_messages) return False @@ -1689,9 +1557,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): replacement = next(msg_iter, None) replacement_text = sanitized_text or "[REDACTED]" if replacement is not None: - text = CiscoAIDefenseGuardrail._normalize_message_content( - replacement.get("content") - ) + text = CiscoAIDefenseGuardrail._normalize_message_content(replacement.get("content")) if text: replacement_text = text choice.message.content = text @@ -1700,9 +1566,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if getattr(choice.message, "content", None): choice.message.content = replacement_text applied = True - if CiscoAIDefenseGuardrail._redact_message_reasoning_fields( - choice.message, replacement_text - ): + if CiscoAIDefenseGuardrail._redact_message_reasoning_fields(choice.message, replacement_text): applied = True CiscoAIDefenseGuardrail._clear_tool_call_arguments(choice.message) return applied @@ -1715,9 +1579,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if getattr(msg, "content", None): msg.content = sanitized_text applied = True - if CiscoAIDefenseGuardrail._redact_message_reasoning_fields( - msg, sanitized_text - ): + if CiscoAIDefenseGuardrail._redact_message_reasoning_fields(msg, sanitized_text): applied = True CiscoAIDefenseGuardrail._clear_tool_call_arguments(msg) return applied @@ -1735,9 +1597,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): for message in sanitized_messages: if not isinstance(message, dict): continue - text = CiscoAIDefenseGuardrail._normalize_message_content( - message.get("content") - ) + text = CiscoAIDefenseGuardrail._normalize_message_content(message.get("content")) if text: replacement = text break @@ -1751,9 +1611,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return applied @classmethod - def _redact_message_reasoning_fields( - cls, message: object, replacement_text: str - ) -> bool: + def _redact_message_reasoning_fields(cls, message: object, replacement_text: str) -> bool: """Remove preserved reasoning fields and expose the sanitized text.""" if not cls._extract_message_reasoning_parts(message): return False @@ -1786,22 +1644,12 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _clear_tool_call_arguments(cls, message: object) -> None: """Clear tool-call / function-call arguments after Cisco redaction.""" - tool_calls = ( - message.get("tool_calls") - if isinstance(message, dict) - else getattr(message, "tool_calls", None) - ) + tool_calls = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) for tc in tool_calls or []: - fn = ( - tc.get("function") - if isinstance(tc, dict) - else getattr(tc, "function", None) - ) + fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) cls._clear_arguments_field(fn) function_call = ( - message.get("function_call") - if isinstance(message, dict) - else getattr(message, "function_call", None) + message.get("function_call") if isinstance(message, dict) else getattr(message, "function_call", None) ) cls._clear_arguments_field(function_call) @@ -1814,17 +1662,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): replacement_text: Optional[str] = sanitized_text if not replacement_text and sanitized_messages: replacement_text = " ".join( - self._normalize_message_content(m.get("content")) - for m in sanitized_messages - if isinstance(m, dict) + self._normalize_message_content(m.get("content")) for m in sanitized_messages if isinstance(m, dict) ).strip() if not replacement_text: return False applied = False for item in output_items: - content = getattr(item, "content", None) or ( - item.get("content") if isinstance(item, dict) else None - ) + content = getattr(item, "content", None) or (item.get("content") if isinstance(item, dict) else None) if isinstance(content, list): for part in content: if isinstance(part, dict): @@ -1839,11 +1683,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): applied = True except (AttributeError, TypeError, ValueError): continue - args = ( - item.get("arguments") - if isinstance(item, dict) - else getattr(item, "arguments", None) - ) + args = item.get("arguments") if isinstance(item, dict) else getattr(item, "arguments", None) if isinstance(args, str) and args: self._clear_arguments_field(item) applied = True @@ -1866,17 +1706,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): content = m.get("content") if isinstance(content, str): ptype = "output_text" if role == "assistant" else "input_text" - out.append( - {"role": role, "content": [{"type": ptype, "text": content}]} - ) + out.append({"role": role, "content": [{"type": ptype, "text": content}]}) elif isinstance(content, list): out.append({"role": role, "content": content}) return out or None @staticmethod - def _rewrite_responses_input_text( - original_input: object, sanitized_text: str - ) -> Optional[object]: + def _rewrite_responses_input_text(original_input: object, sanitized_text: str) -> Optional[object]: """Apply ``sanitized_text`` to a Responses API ``input`` value. Handles plain string, list of message items (rewrites the last @@ -1958,17 +1794,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): end_time = datetime.now() duration = (end_time - start_time).total_seconds() if surface == "mcp": - evt = ( - GuardrailEventHooks.during_mcp_call - if direction == "output" - else GuardrailEventHooks.pre_mcp_call - ) + evt = GuardrailEventHooks.during_mcp_call if direction == "output" else GuardrailEventHooks.pre_mcp_call else: - evt = ( - GuardrailEventHooks.post_call - if direction == "output" - else GuardrailEventHooks.pre_call - ) + evt = GuardrailEventHooks.post_call if direction == "output" else GuardrailEventHooks.pre_call self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self._PROVIDER_NAME, guardrail_json_response={ @@ -1986,8 +1814,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if self.fallback_on_error == "allow": verbose_proxy_logger.warning( - "Cisco AI Defense guardrail: API unavailable, proceeding " - "without scanning (fallback_on_error='allow')" + "Cisco AI Defense guardrail: API unavailable, proceeding without scanning (fallback_on_error='allow')" ) return { "is_safe": True, @@ -2000,8 +1827,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): detail={ "error": "Cisco AI Defense guardrail unavailable", "message": ( - "Cisco AI Defense scanning service is temporarily " - "unavailable and fallback_on_error='block'" + "Cisco AI Defense scanning service is temporarily unavailable and fallback_on_error='block'" ), "error_type": type(error).__name__, }, @@ -2017,9 +1843,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): # parts, ``output_text`` for assistant turns, ``summary_text`` / # ``reasoning_text`` for reasoning summaries that may appear in # conversation history). - _TEXT_PART_TYPES = frozenset( - {"text", "input_text", "output_text", "summary_text", "reasoning_text"} - ) + _TEXT_PART_TYPES = frozenset({"text", "input_text", "output_text", "summary_text", "reasoning_text"}) @staticmethod def _extract_inspect_messages_from_request( @@ -2028,9 +1852,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): """Build {role, content} messages for the Cisco AI Defense chat API.""" messages: List[Dict[str, str]] = [] - instructions_text = CiscoAIDefenseGuardrail._normalize_message_content( - data.get("instructions") - ) + instructions_text = CiscoAIDefenseGuardrail._normalize_message_content(data.get("instructions")) if instructions_text: messages.append({"role": "system", "content": instructions_text}) @@ -2042,14 +1864,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if not role: continue parts: List[str] = [] - text = CiscoAIDefenseGuardrail._normalize_message_content( - message.get("content") - ) + text = CiscoAIDefenseGuardrail._normalize_message_content(message.get("content")) if text: parts.append(text) - parts.extend( - CiscoAIDefenseGuardrail._extract_message_tool_argument_parts(message) - ) + parts.extend(CiscoAIDefenseGuardrail._extract_message_tool_argument_parts(message)) if parts: messages.append({"role": role, "content": " ".join(parts)}) @@ -2058,14 +1876,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): # message-shaped dicts (with role + nested content array), or # a flat list of content-part dicts. Flatten properly so the # scan sees every text segment, not just the top-level ones. - messages.extend( - CiscoAIDefenseGuardrail._flatten_responses_input(data.get("input")) - ) + messages.extend(CiscoAIDefenseGuardrail._flatten_responses_input(data.get("input"))) if not messages and data.get("prompt") is not None: - prompt_text = CiscoAIDefenseGuardrail._normalize_message_content( - data.get("prompt") - ) + prompt_text = CiscoAIDefenseGuardrail._normalize_message_content(data.get("prompt")) if prompt_text: messages.append({"role": "user", "content": prompt_text}) @@ -2163,24 +1977,18 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if not isinstance(part, dict): continue part_type = part.get("type") - if part_type in CiscoAIDefenseGuardrail._TEXT_PART_TYPES and part.get( - "text" - ): + if part_type in CiscoAIDefenseGuardrail._TEXT_PART_TYPES and part.get("text"): parts.append(str(part["text"])) continue nested = part.get("content") if nested is not None: - nested_text = CiscoAIDefenseGuardrail._normalize_message_content( - nested - ) + nested_text = CiscoAIDefenseGuardrail._normalize_message_content(nested) if nested_text: parts.append(nested_text) for key in ("arguments", "output"): value = part.get(key) if value: - parts.append( - CiscoAIDefenseGuardrail._normalize_message_content(value) - ) + parts.append(CiscoAIDefenseGuardrail._normalize_message_content(value)) return " ".join(parts) return str(content) @@ -2200,21 +2008,11 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): if not isinstance(choice, Choices): continue parts: List[str] = [] - content = CiscoAIDefenseGuardrail._normalize_message_content( - getattr(choice.message, "content", None) - ) + content = CiscoAIDefenseGuardrail._normalize_message_content(getattr(choice.message, "content", None)) if content: parts.append(content) - parts.extend( - CiscoAIDefenseGuardrail._extract_message_tool_argument_parts( - choice.message - ) - ) - parts.extend( - CiscoAIDefenseGuardrail._extract_message_reasoning_parts( - choice.message - ) - ) + parts.extend(CiscoAIDefenseGuardrail._extract_message_tool_argument_parts(choice.message)) + parts.extend(CiscoAIDefenseGuardrail._extract_message_reasoning_parts(choice.message)) if parts: result.append({"role": "assistant", "content": " ".join(parts)}) return result @@ -2233,17 +2031,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return [] output_parts: List[str] = [] for item in output_items: - get = ( - item.get - if isinstance(item, dict) - else (lambda k: getattr(item, k, None)) - ) + get = item.get if isinstance(item, dict) else (lambda k: getattr(item, k, None)) for part in get("content") or []: - pget = ( - part.get - if isinstance(part, dict) - else (lambda k: getattr(part, k, None)) - ) + pget = part.get if isinstance(part, dict) else (lambda k: getattr(part, k, None)) for key in ("text", "reasoning", "thinking"): value = pget(key) if isinstance(value, str) and value: @@ -2296,19 +2086,13 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _extract_message_tool_argument_parts(cls, message: object) -> List[str]: parts: List[str] = [] - tool_calls = ( - message.get("tool_calls") - if isinstance(message, dict) - else getattr(message, "tool_calls", None) - ) + tool_calls = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) for tool_call in tool_calls or []: args = cls._extract_tool_call_arguments(tool_call) if args: parts.append(args) function_call = ( - message.get("function_call") - if isinstance(message, dict) - else getattr(message, "function_call", None) + message.get("function_call") if isinstance(message, dict) else getattr(message, "function_call", None) ) if function_call is not None: args = cls._extract_function_call_arguments(function_call) @@ -2321,11 +2105,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): """Pull ``function.arguments`` off a tool_calls entry (dict or model).""" if tool_call is None: return None - function = ( - tool_call.get("function") - if isinstance(tool_call, dict) - else getattr(tool_call, "function", None) - ) + function = tool_call.get("function") if isinstance(tool_call, dict) else getattr(tool_call, "function", None) return CiscoAIDefenseGuardrail._extract_function_call_arguments(function) @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py index bb691c171db..2b53d71e8be 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py @@ -55,13 +55,9 @@ class _CiscoAIDefenseMcpMixin: _PROVIDER_NAME: str guardrail_name: Optional[str] - def should_run_guardrail( - self, data: dict, event_type: GuardrailEventHooks - ) -> bool: ... + def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool: ... - async def _post_inspection( - self, url: str, payload: Dict[str, Any], surface: str - ) -> Dict[str, Any]: ... + async def _post_inspection(self, url: str, payload: Dict[str, Any], surface: str) -> Dict[str, Any]: ... def _handle_api_error( self, @@ -131,18 +127,14 @@ class _CiscoAIDefenseMcpMixin: ) ): verbose_proxy_logger.debug( - "Cisco AI Defense guardrail (%s): no MCP mode configured " - "— skipping MCP response scan.", + "Cisco AI Defense guardrail (%s): no MCP mode configured — skipping MCP response scan.", self.guardrail_name, ) return None mcp_tool_response = self._extract_mcp_tool_call_response(response_obj) if mcp_tool_response is None: - verbose_proxy_logger.debug( - "Cisco AI Defense guardrail: no MCP tool response payload " - "to scan, skipping" - ) + verbose_proxy_logger.debug("Cisco AI Defense guardrail: no MCP tool response payload to scan, skipping") return None original_response = kwargs.get("original_response") @@ -150,22 +142,14 @@ class _CiscoAIDefenseMcpMixin: await self._inspect_mcp_response( request_data=request_data, response=mcp_tool_response, - redact_response_obj=( - original_response - if original_response is not None - else mcp_tool_response - ), + redact_response_obj=(original_response if original_response is not None else mcp_tool_response), ) except HTTPException as exc: - blocking_response = self._build_blocking_mcp_response( - detail=exc.detail, original_response_obj=response_obj - ) + blocking_response = self._build_blocking_mcp_response(detail=exc.detail, original_response_obj=response_obj) self._replace_mcp_tool_response(response_obj, blocking_response) if original_response is not None: self._replace_mcp_tool_response(original_response, blocking_response) - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) verbose_proxy_logger.warning( "Cisco AI Defense guardrail (%s): MCP response blocked — " "tool output replaced with synthesized violation message.", @@ -173,9 +157,7 @@ class _CiscoAIDefenseMcpMixin: ) return blocking_response - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) return None def _build_blocking_mcp_response( @@ -195,9 +177,7 @@ class _CiscoAIDefenseMcpMixin: else: payload = { "error": "Blocked by Cisco AI Defense Guardrail", - "message": ( - str(detail) if detail else "Blocked by Cisco AI Defense Guardrail" - ), + "message": (str(detail) if detail else "Blocked by Cisco AI Defense Guardrail"), "provider": self._PROVIDER_NAME, "guardrail": self.guardrail_name, "surface": "mcp", @@ -210,32 +190,22 @@ class _CiscoAIDefenseMcpMixin: hidden_params: Any = original_hidden else: response_cost = getattr(original_hidden, "response_cost", None) - hidden_params = ( - HiddenParams(response_cost=response_cost) - if response_cost is not None - else HiddenParams() - ) + hidden_params = HiddenParams(response_cost=response_cost) if response_cost is not None else HiddenParams() return MCPPostCallResponseObject( - mcp_tool_call_response=[ - TextContent(type="text", text=_json.dumps(payload)) - ], + mcp_tool_call_response=[TextContent(type="text", text=_json.dumps(payload))], hidden_params=hidden_params, ) @staticmethod - def _replace_mcp_tool_response( - response_obj: object, replacement_obj: object - ) -> bool: + def _replace_mcp_tool_response(response_obj: object, replacement_obj: object) -> bool: replacement = getattr(replacement_obj, "mcp_tool_call_response", None) if replacement is None: return False inner = getattr(response_obj, "mcp_tool_call_response", None) if inner is not None: - if _CiscoAIDefenseMcpMixin._replace_mcp_tool_response( - inner, replacement_obj - ): + if _CiscoAIDefenseMcpMixin._replace_mcp_tool_response(inner, replacement_obj): return True try: setattr(response_obj, "mcp_tool_call_response", replacement) @@ -246,9 +216,7 @@ class _CiscoAIDefenseMcpMixin: content = getattr(response_obj, "content", None) if isinstance(content, list): content[:] = replacement - structured_replacement = ( - _CiscoAIDefenseMcpMixin._replacement_structured_content(replacement) - ) + structured_replacement = _CiscoAIDefenseMcpMixin._replacement_structured_content(replacement) if hasattr(response_obj, "structuredContent"): try: setattr(response_obj, "structuredContent", structured_replacement) @@ -269,16 +237,12 @@ class _CiscoAIDefenseMcpMixin: result = response_obj.get("result") if isinstance(result, dict): result["content"] = replacement - result["structuredContent"] = ( - _CiscoAIDefenseMcpMixin._replacement_structured_content(replacement) - ) + result["structuredContent"] = _CiscoAIDefenseMcpMixin._replacement_structured_content(replacement) result["isError"] = True return True response_obj["result"] = { "content": replacement, - "structuredContent": _CiscoAIDefenseMcpMixin._replacement_structured_content( - replacement - ), + "structuredContent": _CiscoAIDefenseMcpMixin._replacement_structured_content(replacement), "isError": True, } return True @@ -292,11 +256,7 @@ class _CiscoAIDefenseMcpMixin: if not isinstance(replacement, list) or not replacement: return None first = replacement[0] - text = ( - first.get("text") - if isinstance(first, dict) - else getattr(first, "text", None) - ) + text = first.get("text") if isinstance(first, dict) else getattr(first, "text", None) return {"result": text} if isinstance(text, str) else None @staticmethod @@ -320,16 +280,11 @@ class _CiscoAIDefenseMcpMixin: url = f"{self.api_base}{self.inspect_path}" payload = self._build_mcp_request_payload(data=data) if payload is None: - verbose_proxy_logger.debug( - "Cisco AI Defense guardrail: could not build MCP request " - "payload, skipping" - ) + verbose_proxy_logger.debug("Cisco AI Defense guardrail: could not build MCP request payload, skipping") return {} start_time = datetime.now() try: - inspect_response = await self._post_inspection( - url=url, payload=payload, surface="mcp" - ) + inspect_response = await self._post_inspection(url=url, payload=payload, surface="mcp") except HTTPException: raise except Exception as exc: @@ -364,16 +319,11 @@ class _CiscoAIDefenseMcpMixin: response=response, ) if payload is None: - verbose_proxy_logger.debug( - "Cisco AI Defense guardrail: could not build MCP response " - "payload, skipping" - ) + verbose_proxy_logger.debug("Cisco AI Defense guardrail: could not build MCP response payload, skipping") return {} start_time = datetime.now() try: - inspect_response = await self._post_inspection( - url=url, payload=payload, surface="mcp" - ) + inspect_response = await self._post_inspection(url=url, payload=payload, surface="mcp") except HTTPException: raise except Exception as exc: @@ -392,9 +342,7 @@ class _CiscoAIDefenseMcpMixin: request_data=request_data, context=_ScanContext(surface="mcp", direction="output"), start_time=start_time, - response_obj=( - response if redact_response_obj is None else redact_response_obj - ), + response_obj=(response if redact_response_obj is None else redact_response_obj), ) def _build_mcp_request_payload( @@ -419,9 +367,7 @@ class _CiscoAIDefenseMcpMixin: "params": data.get("params") or {}, } - tool_name = ( - data.get("mcp_tool_name") or data.get("tool_name") or data.get("name") - ) + tool_name = data.get("mcp_tool_name") or data.get("tool_name") or data.get("name") if not tool_name: return None @@ -471,9 +417,7 @@ class _CiscoAIDefenseMcpMixin: def _hydrate_mcp_tool_context(request_data: Dict[str, Any]) -> None: metadata = request_data.get("mcp_tool_call_metadata") if metadata is None: - nested = request_data.get("metadata") or request_data.get( - "litellm_metadata" - ) + nested = request_data.get("metadata") or request_data.get("litellm_metadata") if isinstance(nested, dict): metadata = nested.get("mcp_tool_call_metadata") if not isinstance(metadata, dict): @@ -515,14 +459,11 @@ class _CiscoAIDefenseMcpMixin: return { "jsonrpc": "2.0", "id": response.get("id") or "litellm-mcp", - "result": _CiscoAIDefenseMcpMixin._build_mcp_result( - content=content, source=response - ), + "result": _CiscoAIDefenseMcpMixin._build_mcp_result(content=content, source=response), } if isinstance(response, list): if response and all( - isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) - for item in response + isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) for item in response ): response_fields = dict(response) inner_content = response_fields.get("content") @@ -554,9 +495,7 @@ class _CiscoAIDefenseMcpMixin: return { "jsonrpc": "2.0", "id": "litellm-mcp", - "result": _CiscoAIDefenseMcpMixin._build_mcp_result( - content=content, source=response - ), + "result": _CiscoAIDefenseMcpMixin._build_mcp_result(content=content, source=response), } return None @@ -565,15 +504,9 @@ class _CiscoAIDefenseMcpMixin: content: List[Any], source: object = None, ) -> Dict[str, Any]: - result: Dict[str, Any] = { - "content": [_serialize_mcp_content_item(item) for item in content] - } + result: Dict[str, Any] = {"content": [_serialize_mcp_content_item(item) for item in content]} for key in ("structuredContent", "isError"): - value = ( - source.get(key) - if isinstance(source, dict) - else getattr(source, key, None) - ) + value = source.get(key) if isinstance(source, dict) else getattr(source, key, None) if value is not None and (key != "isError" or isinstance(value, bool)): result[key] = value return result @@ -611,10 +544,7 @@ class _CiscoAIDefenseMcpMixin: if ( isinstance(response_obj, list) and response_obj - and all( - isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) - for item in response_obj - ) + and all(isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) for item in response_obj) ): for index, item in enumerate(response_obj): if item[0] == "structuredContent": @@ -628,9 +558,7 @@ class _CiscoAIDefenseMcpMixin: pass elif isinstance(response_obj, dict): result = response_obj.get("result") - target: Dict[Any, Any] = ( - result if isinstance(result, dict) else response_obj - ) + target: Dict[Any, Any] = result if isinstance(result, dict) else response_obj if "structuredContent" in target: target["structuredContent"] = replacement replaced = True @@ -650,8 +578,7 @@ class _CiscoAIDefenseMcpMixin: return content if isinstance(response_obj, list): if response_obj and all( - isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) - for item in response_obj + isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) for item in response_obj ): inner_content = dict(response_obj).get("content") if isinstance(inner_content, list): diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 248202b644c..84b0a3b8eba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -51,9 +51,7 @@ class _ImageUrlContentPart(BaseModel): image_url: _ImageUrl -_ContentPart = Annotated[ - Union[_TextContentPart, _ImageUrlContentPart], Field(discriminator="type") -] +_ContentPart = Annotated[Union[_TextContentPart, _ImageUrlContentPart], Field(discriminator="type")] class _Message(BaseModel): @@ -96,11 +94,7 @@ def _extract_text_from_content(content: object) -> str: if isinstance(content, str): return content if isinstance(content, list): - parts = [ - item.get("text", "") - for item in content - if isinstance(item, dict) and item.get("type") == "text" - ] + parts = [item.get("text", "") for item in content if isinstance(item, dict) and item.get("type") == "text"] return "\n".join(parts) return "" @@ -137,9 +131,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): api_base (Optional[str]): The CrowdStrike AIDR API base URL. Reads from CS_AIDR_BASE_URL env var if None. **kwargs: Additional arguments passed to the CustomGuardrail base class. """ - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_key = api_key or os.environ.get("CS_AIDR_TOKEN") if not self.api_key: @@ -159,9 +151,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): f"Initialized CrowdStrike AIDR Guardrail: name={guardrail_name}, api_base={self.api_base}" ) - async def _call_crowdstrike_aidr_guard( - self, payload: dict[str, Any], hook_name: str - ) -> dict[str, Any]: + async def _call_crowdstrike_aidr_guard(self, payload: dict[str, Any], hook_name: str) -> dict[str, Any]: """ Makes the API call to the CrowdStrike AIDR AI Guard endpoint. The function itself will raise an error if a response should be blocked, @@ -190,9 +180,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): f"CrowdStrike AIDR Guardrail ({hook_name}): Calling endpoint {endpoint} with payload: {payload}" ) - response = await self.async_handler.post( - url=endpoint, json=payload, headers=headers - ) + response = await self.async_handler.post(url=endpoint, json=payload, headers=headers) response.raise_for_status() result: dict[str, Any] = response.json() @@ -214,9 +202,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): return result - def _build_guard_input_for_request( - self, inputs: GenericGuardrailAPIInputs - ) -> Optional[_GuardInput]: + def _build_guard_input_for_request(self, inputs: GenericGuardrailAPIInputs) -> Optional[_GuardInput]: guard_input = _GuardInput(messages=[], tools=[]) structured_messages = inputs.get("structured_messages") texts = inputs.get("texts", []) @@ -227,17 +213,11 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): content = _normalize_content(message.get("content")) if content is None or len(content) == 0: content = "" - guard_input.messages.append( - _Message(role=message["role"], content=content) - ) + guard_input.messages.append(_Message(role=message["role"], content=content)) elif texts: - guard_input.messages = [ - _Message(role="user", content=text) for text in texts - ] + guard_input.messages = [_Message(role="user", content=text) for text in texts] else: - verbose_proxy_logger.warning( - "CrowdStrike AIDR Guardrail: No messages or texts provided for input request" - ) + verbose_proxy_logger.warning("CrowdStrike AIDR Guardrail: No messages or texts provided for input request") return None if tools: @@ -250,9 +230,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): ) -> Optional[_GuardInput]: output_texts: list[str] = inputs.get("texts", []) if len(output_texts) == 0: - verbose_proxy_logger.warning( - "CrowdStrike AIDR Guardrail: No text in output response." - ) + verbose_proxy_logger.warning("CrowdStrike AIDR Guardrail: No text in output response.") return None input_messages = request_data.get("messages", []) @@ -261,8 +239,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): messages=[ _Message(role=role, content=content) for (role, content) in ( - (message["role"], _normalize_content(message.get("content"))) - for message in input_messages + (message["role"], _normalize_content(message.get("content"))) for message in input_messages ) if content is not None and len(content) > 0 ] @@ -275,19 +252,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): num_assistant_messages: int, ) -> list[str]: transformed_messages = guard_output.get("messages", []) - tail = ( - transformed_messages[-num_assistant_messages:] - if num_assistant_messages > 0 - else [] - ) - return [ - ( - _extract_text_from_content(msg.get("content")) - if isinstance(msg, dict) - else "" - ) - for msg in tail - ] + tail = transformed_messages[-num_assistant_messages:] if num_assistant_messages > 0 else [] + return [(_extract_text_from_content(msg.get("content")) if isinstance(msg, dict) else "") for msg in tail] @log_guardrail_information @override @@ -298,9 +264,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - verbose_proxy_logger.debug( - f"CrowdStrike AIDR Guardrail: Applying guardrail to {input_type}" - ) + verbose_proxy_logger.debug(f"CrowdStrike AIDR Guardrail: Applying guardrail to {input_type}") # Extract inputs texts = inputs.get("texts", []) @@ -343,14 +307,10 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): extra_info["user_name"] = user_email ai_guard_payload["extra_info"] = extra_info - ai_guard_response = await self._call_crowdstrike_aidr_guard( - ai_guard_payload, hook_name - ) + ai_guard_response = await self._call_crowdstrike_aidr_guard(ai_guard_payload, hook_name) if "body" in request_data or "messages" in request_data: - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) result = ai_guard_response.get("result", {}) if not result.get("transformed"): @@ -363,18 +323,13 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): # for every message in guard_output. all_messages = guard_output.get("messages", []) transformed_texts = [ - _extract_text_from_content( - msg.get("content") if isinstance(msg, dict) else "" - ) - for msg in all_messages + _extract_text_from_content(msg.get("content") if isinstance(msg, dict) else "") for msg in all_messages ] else: # For responses, guard_input contained history + assistant messages # appended at the end. Extract only the assistant tail. num_assistant = len(texts) - transformed_texts = self._extract_transformed_texts( - guard_output, num_assistant - ) + transformed_texts = self._extract_transformed_texts(guard_output, num_assistant) result_inputs: GenericGuardrailAPIInputs = {"texts": transformed_texts} if tools: diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py index a956688fd43..3b031999f40 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -21,9 +21,7 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail( - litellm_params: "LitellmParams", guardrail: "Guardrail" -) -> CustomCodeGuardrail: +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> CustomCodeGuardrail: """ Initialize a custom code guardrail. @@ -43,9 +41,7 @@ def initialize_guardrail( # Get the custom code from litellm_params custom_code = getattr(litellm_params, "custom_code", None) if not custom_code: - raise ValueError( - "Custom code guardrail requires 'custom_code' in litellm_params" - ) + raise ValueError("Custom code guardrail requires 'custom_code' in litellm_params") custom_code_guardrail = CustomCodeGuardrail( guardrail_name=guardrail_name, diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index 22d4548aa99..9021f023156 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -158,9 +158,7 @@ class CustomCodeGuardrail(CustomGuardrail): apply_fn = exec_globals["apply_guardrail"] if not callable(apply_fn): - raise CustomCodeCompilationError( - "'apply_guardrail' must be a callable function" - ) + raise CustomCodeCompilationError("'apply_guardrail' must be a callable function") self._compiled_function = apply_fn @@ -176,9 +174,7 @@ class CustomCodeGuardrail(CustomGuardrail): try: self._do_compile() - verbose_proxy_logger.debug( - f"Custom code guardrail '{self.guardrail_name}' compiled successfully" - ) + verbose_proxy_logger.debug(f"Custom code guardrail '{self.guardrail_name}' compiled successfully") except SyntaxError as e: self._compile_error = f"Syntax error in custom code: {e}" @@ -225,9 +221,7 @@ class CustomCodeGuardrail(CustomGuardrail): """ if self._compiled_function is None: if self._compile_error: - raise CustomCodeExecutionError( - f"Custom code guardrail not compiled: {self._compile_error}" - ) + raise CustomCodeExecutionError(f"Custom code guardrail not compiled: {self._compile_error}") raise CustomCodeExecutionError("Custom code guardrail not compiled") try: @@ -258,9 +252,7 @@ class CustomCodeGuardrail(CustomGuardrail): # Pre-call block uses passthrough; must not wrap as execution error (500) raise except Exception as e: - verbose_proxy_logger.error( - f"Custom code guardrail '{self.guardrail_name}' execution error: {e}" - ) + verbose_proxy_logger.error(f"Custom code guardrail '{self.guardrail_name}' execution error: {e}") raise CustomCodeExecutionError( f"Custom code guardrail execution failed: {e}", details={ @@ -322,9 +314,7 @@ class CustomCodeGuardrail(CustomGuardrail): action = result.get("action", "allow") if action == "allow": - verbose_proxy_logger.debug( - f"Custom code guardrail '{self.guardrail_name}': Allowing {input_type}" - ) + verbose_proxy_logger.debug(f"Custom code guardrail '{self.guardrail_name}': Allowing {input_type}") return inputs elif action == "block": @@ -356,9 +346,7 @@ class CustomCodeGuardrail(CustomGuardrail): ) elif action == "modify": - verbose_proxy_logger.debug( - f"Custom code guardrail '{self.guardrail_name}': Modifying {input_type}" - ) + verbose_proxy_logger.debug(f"Custom code guardrail '{self.guardrail_name}': Modifying {input_type}") # Apply modifications modified_inputs = dict(inputs) @@ -376,8 +364,7 @@ class CustomCodeGuardrail(CustomGuardrail): else: verbose_proxy_logger.warning( - f"Custom code guardrail '{self.guardrail_name}': " - f"Unknown action '{action}'. Treating as allow." + f"Custom code guardrail '{self.guardrail_name}': Unknown action '{action}'. Treating as allow." ) return inputs @@ -404,9 +391,7 @@ class CustomCodeGuardrail(CustomGuardrail): try: self.custom_code = new_code self._do_compile() - verbose_proxy_logger.info( - f"Custom code guardrail '{self.guardrail_name}': Code updated successfully" - ) + verbose_proxy_logger.info(f"Custom code guardrail '{self.guardrail_name}': Code updated successfully") except SyntaxError as e: # Rollback on failure self.custom_code = old_code diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index de7690635d8..e60b900428c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -31,9 +31,7 @@ def allow() -> Dict[str, Any]: return {"action": "allow"} -def block( - reason: str, detection_info: Optional[Dict[str, Any]] = None -) -> Dict[str, Any]: +def block(reason: str, detection_info: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """ Block the request/response with a reason. @@ -228,9 +226,7 @@ def json_schema_valid(obj: Any, schema: Dict[str, Any]) -> bool: return False -def _basic_json_schema_validate( - obj: Any, schema: Dict[str, Any], max_depth: int = 50 -) -> bool: +def _basic_json_schema_validate(obj: Any, schema: Dict[str, Any], max_depth: int = 50) -> bool: """ Basic JSON schema validation without external library. Handles: type, required, properties @@ -287,9 +283,7 @@ def _basic_json_schema_validate( # Common URL pattern for extraction -_URL_PATTERN = re.compile( - r"https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[^\s]*", re.IGNORECASE -) +_URL_PATTERN = re.compile(r"https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[^\s]*", re.IGNORECASE) def extract_urls(text: str) -> List[str]: @@ -460,9 +454,7 @@ async def http_request( method = method.upper() allowed_methods = {"GET", "POST", "PUT", "DELETE", "PATCH"} if method not in allowed_methods: - return _http_error_response( - f"Invalid HTTP method: {method}. Allowed: {', '.join(allowed_methods)}" - ) + return _http_error_response(f"Invalid HTTP method: {method}. Allowed: {', '.join(allowed_methods)}") # Apply timeout limits if timeout is None: @@ -477,9 +469,7 @@ async def http_request( ) try: - response = await _execute_http_request( - client, method, url, headers, body, timeout - ) + response = await _execute_http_request(client, method, url, headers, body, timeout) return _http_success_response(response) except httpx.TimeoutException as e: @@ -510,21 +500,13 @@ async def _execute_http_request( if method == "GET": return await client.get(url=url, headers=headers) elif method == "POST": - return await client.post( - url=url, headers=headers, json=json_body, data=data_body, timeout=timeout - ) + return await client.post(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) elif method == "PUT": - return await client.put( - url=url, headers=headers, json=json_body, data=data_body, timeout=timeout - ) + return await client.put(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) elif method == "DELETE": - return await client.delete( - url=url, headers=headers, json=json_body, data=data_body, timeout=timeout - ) + return await client.delete(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) elif method == "PATCH": - return await client.patch( - url=url, headers=headers, json=json_body, data=data_body, timeout=timeout - ) + return await client.patch(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) else: raise ValueError(f"Unsupported HTTP method: {method}") @@ -570,9 +552,7 @@ async def http_post( Returns: Same as http_request """ - return await http_request( - url=url, method="POST", headers=headers, body=body, timeout=timeout - ) + return await http_request(url=url, method="POST", headers=headers, body=body, timeout=timeout) # ============================================================================= diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py index bb4a74efc82..51277936069 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py @@ -55,9 +55,7 @@ class myCustomGuardrail(CustomGuardrail): _content = _content.replace("litellm", "********") message["content"] = _content - verbose_proxy_logger.debug( - "async_pre_call_hook: Message after masking %s", _messages - ) + verbose_proxy_logger.debug("async_pre_call_hook: Message after masking %s", _messages) return data diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index 59381149809..b02f1030592 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -48,9 +48,7 @@ class DynamoAIGuardrails(CustomGuardrail): policy_ids: List[str] = [], **kwargs, ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) # Set API configuration self.api_key = api_key or os.getenv("DYNAMOAI_API_KEY") @@ -59,9 +57,7 @@ class DynamoAIGuardrails(CustomGuardrail): "DynamoAI API key is required. Set DYNAMOAI_API_KEY environment variable or pass api_key parameter." ) - self.api_base = api_base or os.getenv( - "DYNAMOAI_API_BASE", "https://api.dynamo.ai" - ) + self.api_base = api_base or os.getenv("DYNAMOAI_API_BASE", "https://api.dynamo.ai") self.api_url = f"{self.api_base}/v1/moderation/analyze/" # Model ID for tracking/logging purposes @@ -69,9 +65,7 @@ class DynamoAIGuardrails(CustomGuardrail): # Policy IDs - get from parameter, env var, or use empty list env_policy_ids = os.getenv("DYNAMOAI_POLICY_IDS", "") - self.policy_ids = policy_ids or ( - env_policy_ids.split(",") if env_policy_ids else [] - ) + self.policy_ids = policy_ids or (env_policy_ids.split(",") if env_policy_ids else []) self.guardrail_name = guardrail_name self.guardrail_provider = "dynamoai" @@ -184,9 +178,7 @@ class DynamoAIGuardrails(CustomGuardrail): raise - def _process_dynamoai_guardrails_response( - self, response: DynamoAIResponse - ) -> DynamoAIProcessedResult: + def _process_dynamoai_guardrails_response(self, response: DynamoAIResponse) -> DynamoAIProcessedResult: """ Process the response from the DynamoAI Guardrails API @@ -213,9 +205,7 @@ class DynamoAIGuardrails(CustomGuardrail): # Check for action in multiple places policy_action = ( - applied_policy.get("action") - or (policy_outputs.get("action") if policy_outputs else None) - or "NONE" + applied_policy.get("action") or (policy_outputs.get("action") if policy_outputs else None) or "NONE" ) # Only include policies with BLOCK action @@ -226,9 +216,7 @@ class DynamoAIGuardrails(CustomGuardrail): "action": policy_action, "method": policy_info.get("method"), "description": policy_info.get("description"), - "message": ( - policy_outputs.get("message") if policy_outputs else None - ), + "message": (policy_outputs.get("message") if policy_outputs else None), } return { @@ -236,9 +224,7 @@ class DynamoAIGuardrails(CustomGuardrail): "violation_details": violation_details, } - def _determine_guardrail_status( - self, response_json: DynamoAIResponse - ) -> GuardrailStatus: + def _determine_guardrail_status(self, response_json: DynamoAIResponse) -> GuardrailStatus: """ Determine the guardrail status based on DynamoAI API response. @@ -266,9 +252,7 @@ class DynamoAIGuardrails(CustomGuardrail): return "success" except Exception as e: - verbose_proxy_logger.error( - "Error determining DynamoAI guardrail status: %s", str(e) - ) + verbose_proxy_logger.error("Error determining DynamoAI guardrail status: %s", str(e)) return "guardrail_failed_to_respond" def _create_error_message(self, processed_result: DynamoAIProcessedResult) -> str: @@ -284,9 +268,7 @@ class DynamoAIGuardrails(CustomGuardrail): violations_detected = processed_result["violations_detected"] violation_details = processed_result["violation_details"] - error_message = ( - f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n" - ) + error_message = f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n" for policy_name in violations_detected: error_message += f"- {policy_name.upper()}:\n" @@ -338,9 +320,7 @@ class DynamoAIGuardrails(CustomGuardrail): event_type=GuardrailEventHooks.pre_call, ) - verbose_proxy_logger.debug( - "Guardrails async_pre_call_hook result=%s", result - ) + verbose_proxy_logger.debug("Guardrails async_pre_call_hook result=%s", result) # Process the guardrails response processed_result = self._process_dynamoai_guardrails_response(result) @@ -352,9 +332,7 @@ class DynamoAIGuardrails(CustomGuardrail): raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -387,9 +365,7 @@ class DynamoAIGuardrails(CustomGuardrail): event_type=GuardrailEventHooks.during_call, ) - verbose_proxy_logger.debug( - "Guardrails async_moderation_hook result=%s", result - ) + verbose_proxy_logger.debug("Guardrails async_moderation_hook result=%s", result) # Process the guardrails response processed_result = self._process_dynamoai_guardrails_response(result) @@ -401,9 +377,7 @@ class DynamoAIGuardrails(CustomGuardrail): raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -425,12 +399,7 @@ class DynamoAIGuardrails(CustomGuardrail): ) from litellm.types.guardrails import GuardrailEventHooks - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return verbose_proxy_logger.debug("async_post_call_success_hook response=%s", response) @@ -443,9 +412,7 @@ class DynamoAIGuardrails(CustomGuardrail): for choice in response.choices: if isinstance(choice, litellm.Choices): - if choice.message.content and isinstance( - choice.message.content, str - ): + if choice.message.content and isinstance(choice.message.content, str): has_text_content = True dynamoai_messages.append( { @@ -455,9 +422,7 @@ class DynamoAIGuardrails(CustomGuardrail): ) if not has_text_content: - verbose_proxy_logger.warning( - "DynamoAI: not running guardrail. No output text in response" - ) + verbose_proxy_logger.warning("DynamoAI: not running guardrail. No output text in response") return if dynamoai_messages: @@ -468,9 +433,7 @@ class DynamoAIGuardrails(CustomGuardrail): event_type=GuardrailEventHooks.post_call, ) - verbose_proxy_logger.debug( - "Guardrails async_post_call_success_hook result=%s", result - ) + verbose_proxy_logger.debug("Guardrails async_post_call_success_hook result=%s", result) # Process the guardrails response processed_result = self._process_dynamoai_guardrails_response(result) @@ -482,9 +445,7 @@ class DynamoAIGuardrails(CustomGuardrail): raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 63541a1e2f9..f9bb13ad64e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -59,9 +59,7 @@ class EnkryptAIGuardrails(CustomGuardrail): policy_name: Optional[str] = None, **kwargs, ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) # Set API configuration self.api_key = api_key or os.getenv("ENKRYPTAI_API_KEY") @@ -70,9 +68,7 @@ class EnkryptAIGuardrails(CustomGuardrail): "EnkryptAI API key is required. Set ENKRYPTAI_API_KEY environment variable or pass api_key parameter." ) - self.api_base = api_base or os.getenv( - "ENKRYPTAI_API_BASE", "https://api.enkryptai.com" - ) + self.api_base = api_base or os.getenv("ENKRYPTAI_API_BASE", "https://api.enkryptai.com") self.api_url = f"{self.api_base}/guardrails/policy/detect" # Policy name can be passed as parameter or use guardrail_name @@ -188,9 +184,7 @@ class EnkryptAIGuardrails(CustomGuardrail): raise - def _process_enkryptai_guardrails_response( - self, response: EnkryptAIResponse - ) -> EnkryptAIProcessedResult: + def _process_enkryptai_guardrails_response(self, response: EnkryptAIResponse) -> EnkryptAIProcessedResult: """ Process the response from the Enkrypt AI Guardrails API @@ -221,9 +215,7 @@ class EnkryptAIGuardrails(CustomGuardrail): return {"attacks_detected": detected_attacks, "attack_details": attack_details} - def _determine_guardrail_status( - self, response_json: EnkryptAIResponse - ) -> GuardrailStatus: + def _determine_guardrail_status(self, response_json: EnkryptAIResponse) -> GuardrailStatus: """ Determine the guardrail status based on EnkryptAI API response. @@ -237,9 +229,7 @@ class EnkryptAIGuardrails(CustomGuardrail): return "guardrail_failed_to_respond" # Process the response to check for violations - processed_result = self._process_enkryptai_guardrails_response( - response_json - ) + processed_result = self._process_enkryptai_guardrails_response(response_json) attacks_detected = processed_result["attacks_detected"] if attacks_detected: @@ -248,9 +238,7 @@ class EnkryptAIGuardrails(CustomGuardrail): return "success" except Exception as e: - verbose_proxy_logger.error( - "Error determining EnkryptAI guardrail status: %s", str(e) - ) + verbose_proxy_logger.error("Error determining EnkryptAI guardrail status: %s", str(e)) return "guardrail_failed_to_respond" def _create_error_message(self, processed_result: EnkryptAIProcessedResult) -> str: @@ -266,9 +254,7 @@ class EnkryptAIGuardrails(CustomGuardrail): attacks_detected = processed_result["attacks_detected"] attack_details = processed_result["attack_details"] - error_message = ( - f"Guardrail failed: {len(attacks_detected)} violation(s) detected\n\n" - ) + error_message = f"Guardrail failed: {len(attacks_detected)} violation(s) detected\n\n" for attack_type in attacks_detected: error_message += f"- {attack_type.upper()}:\n" @@ -281,18 +267,12 @@ class EnkryptAIGuardrails(CustomGuardrail): elif attack_type == "pii": error_message += f" PII Detected: {details.get('pii', {})}\n" elif attack_type == "toxicity": - toxic_types = [ - k - for k, v in details.items() - if isinstance(v, (int, float)) and v > 0.5 - ] + toxic_types = [k for k, v in details.items() if isinstance(v, (int, float)) and v > 0.5] error_message += f" Types: {', '.join(toxic_types)}\n" elif attack_type == "keyword_detected": error_message += f" Keywords: {details.get('detected_keywords', [])}\n" elif attack_type == "bias": - error_message += ( - f" Bias Detected: {details.get('bias_detected', False)}\n" - ) + error_message += f" Bias Detected: {details.get('bias_detected', False)}\n" else: error_message += f" Details: {details}\n" error_message += "\n" @@ -331,14 +311,10 @@ class EnkryptAIGuardrails(CustomGuardrail): request_data=data, ) - verbose_proxy_logger.debug( - "Guardrails async_pre_call_hook result: %s", result - ) + verbose_proxy_logger.debug("Guardrails async_pre_call_hook result: %s", result) # Process the guardrails response - processed_result = self._process_enkryptai_guardrails_response( - result - ) + processed_result = self._process_enkryptai_guardrails_response(result) attacks_detected = processed_result["attacks_detected"] # If any attacks are detected, raise an error @@ -347,9 +323,7 @@ class EnkryptAIGuardrails(CustomGuardrail): raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -383,14 +357,10 @@ class EnkryptAIGuardrails(CustomGuardrail): request_data=data, ) - verbose_proxy_logger.debug( - "Guardrails async_moderation_hook result: %s", result - ) + verbose_proxy_logger.debug("Guardrails async_moderation_hook result: %s", result) # Process the guardrails response - processed_result = self._process_enkryptai_guardrails_response( - result - ) + processed_result = self._process_enkryptai_guardrails_response(result) attacks_detected = processed_result["attacks_detected"] # If any attacks are detected, raise an error @@ -399,9 +369,7 @@ class EnkryptAIGuardrails(CustomGuardrail): raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -423,17 +391,10 @@ class EnkryptAIGuardrails(CustomGuardrail): ) from litellm.types.guardrails import GuardrailEventHooks - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return - verbose_proxy_logger.debug( - "async_post_call_success_hook response: %s", response - ) + verbose_proxy_logger.debug("async_post_call_success_hook response: %s", response) # Check if the ModelResponse has text content in its choices # to avoid sending empty content to EnkryptAI (e.g., during tool calls) @@ -441,39 +402,27 @@ class EnkryptAIGuardrails(CustomGuardrail): has_text_content = False for choice in response.choices: if isinstance(choice, litellm.Choices): - if choice.message.content and isinstance( - choice.message.content, str - ): + if choice.message.content and isinstance(choice.message.content, str): has_text_content = True break if not has_text_content: - verbose_proxy_logger.warning( - "EnkryptAI: not running guardrail. No output text in response" - ) + verbose_proxy_logger.warning("EnkryptAI: not running guardrail. No output text in response") return for choice in response.choices: if isinstance(choice, litellm.Choices): - verbose_proxy_logger.debug( - "async_post_call_success_hook choice: %s", choice - ) - if choice.message.content and isinstance( - choice.message.content, str - ): + verbose_proxy_logger.debug("async_post_call_success_hook choice: %s", choice) + if choice.message.content and isinstance(choice.message.content, str): result = await self._call_enkryptai_guardrails( prompt=choice.message.content, request_data=data, ) - verbose_proxy_logger.debug( - "Guardrails async_post_call_success_hook result: %s", result - ) + verbose_proxy_logger.debug("Guardrails async_post_call_success_hook result: %s", result) # Process the guardrails response - processed_result = self._process_enkryptai_guardrails_response( - result - ) + processed_result = self._process_enkryptai_guardrails_response(result) attacks_detected = processed_result["attacks_detected"] # If any attacks are detected, raise an error @@ -482,9 +431,7 @@ class EnkryptAIGuardrails(CustomGuardrail): raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @log_guardrail_information async def apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index 777c45bcd83..2386f80e819 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -15,12 +15,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" api_base=litellm_params.api_base, api_key=litellm_params.api_key, headers=getattr(litellm_params, "headers", None), - additional_provider_specific_params=getattr( - litellm_params, "additional_provider_specific_params", {} - ), - unreachable_fallback=getattr( - litellm_params, "unreachable_fallback", "fail_closed" - ), + additional_provider_specific_params=getattr(litellm_params, "additional_provider_specific_params", {}), + unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"), fail_on_error=getattr(litellm_params, "fail_on_error", True), extra_headers=getattr(litellm_params, "extra_headers", None), guardrail_name=guardrail.get("guardrail_name", ""), @@ -28,9 +24,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" default_on=litellm_params.default_on, ) - litellm.logging_callback_manager.add_litellm_callback( - _generic_guardrail_api_callback - ) + litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback) return _generic_guardrail_api_callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 8d3abe2bdc3..df80ea09de0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -128,31 +128,19 @@ def _extract_inbound_headers( # 3) Pre-call: headers stored in request metadata metadata_headers = (request_data.get("metadata") or {}).get("headers") if metadata_headers: - return _sanitize_inbound_headers( - metadata_headers, extra_allowlist=extra_allowlist - ) + return _sanitize_inbound_headers(metadata_headers, extra_allowlist=extra_allowlist) - litellm_metadata_headers = (request_data.get("litellm_metadata") or {}).get( - "headers" - ) + litellm_metadata_headers = (request_data.get("litellm_metadata") or {}).get("headers") if litellm_metadata_headers: - return _sanitize_inbound_headers( - litellm_metadata_headers, extra_allowlist=extra_allowlist - ) + return _sanitize_inbound_headers(litellm_metadata_headers, extra_allowlist=extra_allowlist) # 4) Post-call: headers not present on response; fallback to logging object if logging_obj and getattr(logging_obj, "model_call_details", None): try: details = logging_obj.model_call_details or {} - headers = ( - details.get("litellm_params", {}) - .get("metadata", {}) - .get("headers", None) - ) + headers = details.get("litellm_params", {}).get("metadata", {}).get("headers", None) if headers: - return _sanitize_inbound_headers( - headers, extra_allowlist=extra_allowlist - ) + return _sanitize_inbound_headers(headers, extra_allowlist=extra_allowlist) except Exception: pass @@ -192,9 +180,7 @@ class GenericGuardrailAPI(CustomGuardrail): extra_headers: Optional[list] = None, **kwargs, ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.headers = headers or {} self.extra_headers = extra_headers or [] @@ -217,13 +203,9 @@ class GenericGuardrailAPI(CustomGuardrail): else: self.api_base = base_url - self.additional_provider_specific_params = ( - additional_provider_specific_params or {} - ) + self.additional_provider_specific_params = additional_provider_specific_params or {} - self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( - unreachable_fallback - ) + self.unreachable_fallback: Literal["fail_closed", "fail_open"] = unreachable_fallback self.fail_on_error: bool = True if fail_on_error is None else fail_on_error @@ -237,13 +219,9 @@ class GenericGuardrailAPI(CustomGuardrail): super().__init__(**kwargs) - verbose_proxy_logger.debug( - "Generic Guardrail API initialized with api_base: %s", self.api_base - ) + verbose_proxy_logger.debug("Generic Guardrail API initialized with api_base: %s", self.api_base) - def _extract_user_api_key_metadata( - self, request_data: dict - ) -> GenericGuardrailAPIMetadata: + def _extract_user_api_key_metadata(self, request_data: dict) -> GenericGuardrailAPIMetadata: """ Extract user API key metadata from request_data. @@ -279,9 +257,7 @@ class GenericGuardrailAPI(CustomGuardrail): # handle user_api_key_token = user_api_key_hash if metadata_dict.get("user_api_key_token") is not None: - result_metadata["user_api_key_hash"] = metadata_dict.get( - "user_api_key_token" - ) + result_metadata["user_api_key_hash"] = metadata_dict.get("user_api_key_token") verbose_proxy_logger.debug( "Generic Guardrail API: Extracted user metadata: %s", @@ -299,9 +275,7 @@ class GenericGuardrailAPI(CustomGuardrail): error: Exception, http_status_code: Optional[int] = None, ) -> GenericGuardrailAPIInputs: - status_suffix = ( - f" http_status_code={http_status_code}" if http_status_code else "" - ) + status_suffix = f" http_status_code={http_status_code}" if http_status_code else "" verbose_proxy_logger.critical( "Generic Guardrail API error (fail-open). Proceeding without guardrail.%s " "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s", @@ -355,13 +329,9 @@ class GenericGuardrailAPI(CustomGuardrail): logging_obj: Optional["LiteLLMLoggingObj"], is_unreachable: bool = True, ) -> GenericGuardrailAPIInputs: - unreachable_fail_open = ( - is_unreachable and self.unreachable_fallback == "fail_open" - ) + unreachable_fail_open = is_unreachable and self.unreachable_fallback == "fail_open" if unreachable_fail_open or not self.fail_on_error: - http_status_code = getattr( - getattr(error, "response", None), "status_code", None - ) + http_status_code = getattr(getattr(error, "response", None), "status_code", None) return self._fail_open_passthrough( inputs=inputs, input_type=input_type, @@ -369,9 +339,7 @@ class GenericGuardrailAPI(CustomGuardrail): error=error, **({"http_status_code": http_status_code} if http_status_code else {}), ) - verbose_proxy_logger.error( - "Generic Guardrail API: failed to make request: %s", str(error) - ) + verbose_proxy_logger.error("Generic Guardrail API: failed to make request: %s", str(error)) raise Exception(f"Generic Guardrail API failed: {str(error)}") @log_guardrail_information @@ -428,11 +396,7 @@ class GenericGuardrailAPI(CustomGuardrail): # Extract user API key metadata user_metadata = self._extract_user_api_key_metadata(request_data) - extra_allowlist = ( - {h.lower() for h in self.extra_headers if isinstance(h, str)} - if self.extra_headers - else None - ) + extra_allowlist = {h.lower() for h in self.extra_headers if isinstance(h, str)} if self.extra_headers else None inbound_headers = _extract_inbound_headers( request_data=request_data, logging_obj=logging_obj, @@ -449,11 +413,7 @@ class GenericGuardrailAPI(CustomGuardrail): request_headers=inbound_headers, litellm_version=litellm_version, images=images, - tools=( - [GuardrailToolParam.model_validate(t) for t in tools] - if tools - else None - ), + tools=([GuardrailToolParam.model_validate(t) for t in tools] if tools else None), structured_messages=structured_messages, tool_calls=tool_calls, additional_provider_specific_params=additional_params, @@ -474,21 +434,15 @@ class GenericGuardrailAPI(CustomGuardrail): response.raise_for_status() response_json = response.json() - verbose_proxy_logger.debug( - "Generic Guardrail API response: %s", response_json - ) + verbose_proxy_logger.debug("Generic Guardrail API response: %s", response_json) guardrail_response = GenericGuardrailAPIResponse.from_dict(response_json) # Handle the response if guardrail_response.action == "BLOCKED": # Block the request - error_message = ( - guardrail_response.blocked_reason or "Content violates policy" - ) - verbose_proxy_logger.warning( - "Generic Guardrail API blocked request: %s", error_message - ) + error_message = guardrail_response.blocked_reason or "Content violates policy" + verbose_proxy_logger.warning("Generic Guardrail API blocked request: %s", error_message) raise GuardrailRaisedException( guardrail_name=GUARDRAIL_NAME, message=error_message, @@ -505,9 +459,7 @@ class GenericGuardrailAPI(CustomGuardrail): except GuardrailRaisedException: raise except Timeout as e: - return self._handle_guardrail_request_error( - e, inputs, input_type, logging_obj - ) + return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj) except httpx.HTTPStatusError as e: status_code = getattr(getattr(e, "response", None), "status_code", None) is_unreachable = status_code in (502, 503, 504) @@ -515,10 +467,6 @@ class GenericGuardrailAPI(CustomGuardrail): e, inputs, input_type, logging_obj, is_unreachable=is_unreachable ) except httpx.RequestError as e: - return self._handle_guardrail_request_error( - e, inputs, input_type, logging_obj - ) + return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj) except Exception as e: - return self._handle_guardrail_request_error( - e, inputs, input_type, logging_obj, is_unreachable=False - ) + return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj, is_unreachable=False) diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py index c6dee3f841d..603b9a7db26 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py @@ -14,9 +14,7 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail( - litellm_params: "LitellmParams", guardrail: "Guardrail" -) -> GraySwanGuardrail: +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> GraySwanGuardrail: import litellm guardrail_name = guardrail.get("guardrail_name") @@ -29,29 +27,16 @@ def initialize_guardrail( guardrail_name=guardrail_name, api_key=litellm_params.api_key, api_base=litellm_params.api_base, - on_flagged_action=_get_config_value( - litellm_params, optional_params, "on_flagged_action" - ), - violation_threshold=_get_config_value( - litellm_params, optional_params, "violation_threshold" - ), - reasoning_mode=_get_config_value( - litellm_params, optional_params, "reasoning_mode" - ), + on_flagged_action=_get_config_value(litellm_params, optional_params, "on_flagged_action"), + violation_threshold=_get_config_value(litellm_params, optional_params, "violation_threshold"), + reasoning_mode=_get_config_value(litellm_params, optional_params, "reasoning_mode"), categories=_get_config_value(litellm_params, optional_params, "categories"), policy_id=_get_config_value(litellm_params, optional_params, "policy_id"), - streaming_end_of_stream_only=_get_config_value( - litellm_params, optional_params, "streaming_end_of_stream_only" - ) + streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only") or False, - streaming_sampling_rate=_get_config_value( - litellm_params, optional_params, "streaming_sampling_rate" - ) - or 5, + streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate") or 5, fail_open=_get_config_value(litellm_params, optional_params, "fail_open"), - guardrail_timeout=_get_config_value( - litellm_params, optional_params, "guardrail_timeout" - ), + guardrail_timeout=_get_config_value(litellm_params, optional_params, "guardrail_timeout"), event_hook=litellm_params.mode, default_on=litellm_params.default_on, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 8955bffc125..a14d2fc8608 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -77,9 +77,7 @@ class GraySwanGuardrail(CustomGuardrail): guardrail_timeout: Optional[float] = 30.0, **kwargs: Any, ) -> None: - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) api_key_value = api_key or os.getenv("GRAYSWAN_API_KEY") if not api_key_value: @@ -109,9 +107,7 @@ class GraySwanGuardrail(CustomGuardrail): self.categories = categories self.policy_id = policy_id self.fail_open = True if fail_open is None else bool(fail_open) - self.guardrail_timeout = ( - 30.0 if guardrail_timeout is None else float(guardrail_timeout) - ) + self.guardrail_timeout = 30.0 if guardrail_timeout is None else float(guardrail_timeout) # Streaming configuration self.streaming_end_of_stream_only = streaming_end_of_stream_only @@ -212,13 +208,9 @@ class GraySwanGuardrail(CustomGuardrail): messages = [{"role": role, "content": text} for text in texts] # Get dynamic params from request metadata - dynamic_body = ( - self.get_guardrail_dynamic_request_body_params(request_data) or {} - ) + dynamic_body = self.get_guardrail_dynamic_request_body_params(request_data) or {} if dynamic_body: - verbose_proxy_logger.debug( - "Gray Swan Guardrail: dynamic extra_body=%s", safe_dumps(dynamic_body) - ) + verbose_proxy_logger.debug("Gray Swan Guardrail: dynamic extra_body=%s", safe_dumps(dynamic_body)) # Prepare and send payload payload = self._prepare_payload(messages, dynamic_body, request_data) @@ -240,9 +232,7 @@ class GraySwanGuardrail(CustomGuardrail): if self._is_grayswan_exception(exc): raise end_time = time.time() - status_code = getattr(exc, "status_code", None) or getattr( - exc, "exception_status_code", None - ) + status_code = getattr(exc, "status_code", None) or getattr(exc, "exception_status_code", None) self._log_guardrail_failure( exc=exc, request_data=request_data or {}, @@ -363,12 +353,8 @@ class GraySwanGuardrail(CustomGuardrail): elif self.on_flagged_action == "passthrough": # For passthrough mode, we need to handle violations detections = [detection_info] - violation_message = self._format_violation_message( - detections, is_output=not is_input - ) - verbose_proxy_logger.info( - "Gray Swan Guardrail: Passthrough mode - handling violation" - ) + violation_message = self._format_violation_message(detections, is_output=not is_input) + verbose_proxy_logger.info("Gray Swan Guardrail: Passthrough mode - handling violation") # If hook_type is provided and in pre/during call, raise exception if hook_type in [ @@ -410,14 +396,10 @@ class GraySwanGuardrail(CustomGuardrail): ) response.raise_for_status() result = response.json() - verbose_proxy_logger.debug( - "Gray Swan Guardrail: monitor response %s", safe_dumps(result) - ) + verbose_proxy_logger.debug("Gray Swan Guardrail: monitor response %s", safe_dumps(result)) return result except Exception as exc: - status_code = getattr(exc, "status_code", None) or getattr( - exc, "exception_status_code", None - ) + status_code = getattr(exc, "status_code", None) or getattr(exc, "exception_status_code", None) raise GraySwanGuardrailAPIError(str(exc), status_code=status_code) from exc def _process_response_internal( @@ -485,15 +467,11 @@ class GraySwanGuardrail(CustomGuardrail): }, ) elif self.on_flagged_action == "monitor": - verbose_proxy_logger.info( - "Gray Swan Guardrail: Monitoring mode - allowing flagged content" - ) + verbose_proxy_logger.info("Gray Swan Guardrail: Monitoring mode - allowing flagged content") return inputs elif self.on_flagged_action == "passthrough": # Replace content with violation message - violation_message = self._format_violation_message( - detection_info, is_output=is_output - ) + violation_message = self._format_violation_message(detection_info, is_output=is_output) verbose_proxy_logger.info( "Gray Swan Guardrail: Passthrough mode - replacing content with violation message" ) @@ -549,17 +527,13 @@ class GraySwanGuardrail(CustomGuardrail): if isinstance(litellm_metadata, dict) and litellm_metadata: cleaned_litellm_metadata = dict(litellm_metadata) # cleaned_litellm_metadata.pop("user_api_key_auth", None) - sanitized = safe_json_loads( - safe_dumps(cleaned_litellm_metadata), default={} - ) + sanitized = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) if isinstance(sanitized, dict) and sanitized: payload["litellm_metadata"] = sanitized return payload - def _format_violation_message( - self, detection_info: Any, is_output: bool = False - ) -> str: + def _format_violation_message(self, detection_info: Any, is_output: bool = False) -> str: """ Format detection info into a user-friendly violation message. @@ -577,9 +551,7 @@ class GraySwanGuardrail(CustomGuardrail): detection_info = detection_info[0] # Extract fields from detection_info dict - detection_dict: dict = ( - detection_info if isinstance(detection_info, dict) else {} - ) + detection_dict: dict = detection_info if isinstance(detection_info, dict) else {} violation_score = detection_dict.get("violation_score", 0.0) violated_rules = detection_dict.get("violated_rules", []) mutation = detection_dict.get("mutation", False) @@ -595,14 +567,10 @@ class GraySwanGuardrail(CustomGuardrail): if violated_rules: formatted_rules = self._format_violated_rules(violated_rules) if formatted_rules: - message_parts.append( - f"It was violating the rule(s): {formatted_rules}." - ) + message_parts.append(f"It was violating the rule(s): {formatted_rules}.") if mutation: - message_parts.append( - "Mutation effort to make the harmful intention disguised was DETECTED." - ) + message_parts.append("Mutation effort to make the harmful intention disguised was DETECTED.") if ipi: message_parts.append("Indirect Prompt Injection was DETECTED.") diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py index 454e06fdbc5..ed3d184c148 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py @@ -23,9 +23,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" event_hook=litellm_params.mode, default_on=litellm_params.default_on, guard_name=litellm_params.guard_name, - guardrails_ai_api_input_format=getattr( - litellm_params, "guardrails_ai_api_input_format", "llmOutput" - ), + guardrails_ai_api_input_format=getattr(litellm_params, "guardrails_ai_api_input_format", "llmOutput"), ) litellm.logging_callback_manager.add_litellm_callback(_guardrails_ai_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index 2d04517de2b..71c426a3367 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -70,9 +70,7 @@ class GuardrailsAI(CustomGuardrail): "GuardrailsAIException - Please pass the Guardrails AI guard name via 'litellm_params::guard_name'" ) # store kwargs as optional_params - self.guardrails_ai_api_base = ( - api_base or os.getenv("GUARDRAILS_AI_API_BASE") or "http://0.0.0.0:8000" - ) + self.guardrails_ai_api_base = api_base or os.getenv("GUARDRAILS_AI_API_BASE") or "http://0.0.0.0:8000" self.guardrails_ai_guard_name = guard_name self.optional_params = kwargs self.guardrails_ai_api_input_format = guardrails_ai_api_input_format @@ -83,9 +81,7 @@ class GuardrailsAI(CustomGuardrail): ] super().__init__(supported_event_hooks=supported_event_hooks, **kwargs) - async def make_guardrails_ai_api_request( - self, llm_output: str, request_data: dict - ) -> GuardrailsAIResponse: + async def make_guardrails_ai_api_request(self, llm_output: str, request_data: dict) -> GuardrailsAIResponse: from httpx import URL data = { @@ -94,11 +90,7 @@ class GuardrailsAI(CustomGuardrail): } _json_data = json.dumps(data) response = await litellm.module_level_aclient.post( - url=str( - URL(self.guardrails_ai_api_base).join( - f"guards/{self.guardrails_ai_guard_name}/validate" - ) - ), + url=str(URL(self.guardrails_ai_api_base).join(f"guards/{self.guardrails_ai_guard_name}/validate")), data=_json_data, headers={ "Content-Type": "application/json", @@ -116,9 +108,7 @@ class GuardrailsAI(CustomGuardrail): ) return _json_response - async def make_guardrails_ai_api_request_pre_call_request( - self, text_input: str, request_data: dict - ) -> str: + async def make_guardrails_ai_api_request_pre_call_request(self, text_input: str, request_data: dict) -> str: from httpx import URL # This branch of code does not work with current version of GuardrailsAI API (as of July 2025), and it is unclear if it ever worked. @@ -137,11 +127,7 @@ class GuardrailsAI(CustomGuardrail): } _json_data = json.dumps(data) response = await litellm.module_level_aclient.post( - url=str( - URL(self.guardrails_ai_api_base).join( - f"guards/{self.guardrails_ai_guard_name}/validate" - ) - ), + url=str(URL(self.guardrails_ai_api_base).join(f"guards/{self.guardrails_ai_guard_name}/validate")), data=_json_data, headers={ "Content-Type": "application/json", @@ -182,12 +168,8 @@ class GuardrailsAI(CustomGuardrail): text_input=text, request_data=data ) else: - _result = await self.make_guardrails_ai_api_request( - llm_output=text, request_data=data - ) - updated_text = ( - _result.get("validatedOutput") or _result.get("rawLlmOutput") or text - ) + _result = await self.make_guardrails_ai_api_request(llm_output=text, request_data=data) + updated_text = _result.get("validatedOutput") or _result.get("rawLlmOutput") or text data["messages"] = set_last_user_message(data["messages"], updated_text) return data @@ -214,9 +196,7 @@ class GuardrailsAI(CustomGuardrail): ]: # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm return await self.process_input(data=data, call_type=call_type) - async def async_logging_hook( - self, kwargs: dict, result: Any, call_type: str - ) -> Tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: if call_type == "acompletion" or call_type == "completion": kwargs = await self.process_input(data=kwargs, call_type=call_type) @@ -247,13 +227,9 @@ class GuardrailsAI(CustomGuardrail): response_str: str = get_content_from_model_response(response) if response_str is not None and len(response_str) > 0: - await self.make_guardrails_ai_api_request( - llm_output=response_str, request_data=data - ) + await self.make_guardrails_ai_api_request(llm_output=response_str, request_data=data) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py index cf1b3bc97e2..9b7934b7705 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py @@ -24,9 +24,7 @@ def _coerce_event_hook( return GuardrailEventHooks(mode) -def initialize_guardrail( - litellm_params: LitellmParams, guardrail: Guardrail -) -> HeadroomGuardrail: +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> HeadroomGuardrail: import litellm _callback = HeadroomGuardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 9d63e5d144a..2228ccf3997 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -42,15 +42,10 @@ class HeadroomGuardrail(CustomGuardrail): api_key: str | None = None, model: str | None = None, guardrail_name: str | None = None, - event_hook: GuardrailEventHooks - | list[GuardrailEventHooks] - | Mode - | None = None, + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, default_on: bool = False, ): - self.headroom_api_base = ( - api_base or get_secret_str("HEADROOM_API_BASE") or "" - ).rstrip("/") + self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/") if not self.headroom_api_base: raise ValueError( "Headroom guardrail requires an API base URL. " @@ -180,9 +175,7 @@ class HeadroomGuardrail(CustomGuardrail): return inputs if self._should_bypass(request_data): - verbose_proxy_logger.debug( - "Headroom: %s header set; skipping compression", BYPASS_HEADER - ) + verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs structured_messages = inputs.get("structured_messages") diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/__init__.py index cd71d55991e..53881efa346 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/__init__.py @@ -13,9 +13,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" api_id = litellm_params.api_id if hasattr(litellm_params, "api_id") else None auth_url = litellm_params.auth_url if hasattr(litellm_params, "auth_url") else None - version: int | None = ( - litellm_params.version if hasattr(litellm_params, "version") else None - ) + version: int | None = litellm_params.version if hasattr(litellm_params, "version") else None _hiddenlayer_callback: HiddenlayerGuardrail | HiddenlayerGuardrailV2 if not version or version < 2: diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index cabb659fa08..287f108b070 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -72,32 +72,18 @@ class HiddenlayerGuardrail(CustomGuardrail): **kwargs: Any, ) -> None: self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") - self.hiddenlayer_client_secret = api_key or os.getenv( - "HIDDENLAYER_CLIENT_SECRET" - ) - self.api_base = ( - api_base - or os.getenv("HIDDENLAYER_API_BASE") - or "https://api.hiddenlayer.ai" - ) + self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET") + self.api_base = api_base or os.getenv("HIDDENLAYER_API_BASE") or "https://api.hiddenlayer.ai" self.jwt_token = None - auth_url = ( - auth_url - or os.getenv("HIDDENLAYER_AUTH_URL") - or "https://auth.hiddenlayer.ai" - ) + auth_url = auth_url or os.getenv("HIDDENLAYER_AUTH_URL") or "https://auth.hiddenlayer.ai" if is_saas(self.api_base): if not self.hiddenlayer_client_id: - raise RuntimeError( - "`api_id` cannot be None when using the SaaS version of HiddenLayer." - ) + raise RuntimeError("`api_id` cannot be None when using the SaaS version of HiddenLayer.") if not self.hiddenlayer_client_secret: - raise RuntimeError( - "`api_key` cannot be None when using the SaaS version of HiddenLayer." - ) + raise RuntimeError("`api_key` cannot be None when using the SaaS version of HiddenLayer.") self.jwt_token = _get_jwt( auth_url=auth_url, @@ -110,9 +96,7 @@ class HiddenlayerGuardrail(CustomGuardrail): api_key=self.hiddenlayer_client_secret, ) - self._http_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) super().__init__(**kwargs) @log_guardrail_information @@ -129,9 +113,7 @@ class HiddenlayerGuardrail(CustomGuardrail): # I.e request can specify gpt-4o-mini but the response from the server will be # gpt-4o-mini-2025-11-01. We need the model to be consistent so that inferences # will be grouped correctly on the Hiddenlayer side - model_name = ( - logging_obj.model if logging_obj and logging_obj.model else "unknown" - ) + model_name = logging_obj.model if logging_obj and logging_obj.model else "unknown" hl_request_metadata = {"model": model_name} # We need the hiddenlayer project id and requester id on both the input and output @@ -141,15 +123,9 @@ class HiddenlayerGuardrail(CustomGuardrail): # from the logger object on the response from the model. headers = request_data.get("proxy_server_request", {}).get("headers", {}) if not headers and logging_obj and logging_obj.model_call_details: - headers = ( - logging_obj.model_call_details.get("litellm_params", {}) - .get("metadata", {}) - .get("headers", {}) - ) + headers = logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {}) - hl_request_metadata["requester_id"] = ( - headers.get("hl-requester-id") or "LiteLLM" - ) + hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM" project_id = headers.get("hl-project-id") if scan_params := inputs.get("structured_messages"): @@ -179,9 +155,7 @@ class HiddenlayerGuardrail(CustomGuardrail): if result.get("evaluation", {}).get("action") == HiddenlayerAction.BLOCK: detected_reasons = [ - entry.get("name", "unknown") - for entry in result.get("analysis", []) - if entry.get("detected") + entry.get("name", "unknown") for entry in result.get("analysis", []) if entry.get("detected") ] threat_level = result.get("evaluation", {}).get("threat_level") raise HTTPException( @@ -200,9 +174,7 @@ class HiddenlayerGuardrail(CustomGuardrail): last_content = modified_data["input"]["messages"][-1]["content"] if isinstance(last_content, list): texts = [ - item["text"] - for item in last_content - if isinstance(item, dict) and item.get("type") == "text" + item["text"] for item in last_content if isinstance(item, dict) and item.get("type") == "text" ] inputs["texts"] = texts if texts else [""] else: @@ -213,9 +185,7 @@ class HiddenlayerGuardrail(CustomGuardrail): last_content = modified_data["output"]["messages"][-1]["content"] if isinstance(last_content, list): texts = [ - item["text"] - for item in last_content - if isinstance(item, dict) and item.get("type") == "text" + item["text"] for item in last_content if isinstance(item, dict) and item.get("type") == "text" ] inputs["texts"] = texts if texts else [""] else: @@ -306,32 +276,18 @@ class HiddenlayerGuardrailV2(CustomGuardrail): **kwargs: Any, ) -> None: self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID") - self.hiddenlayer_client_secret = api_key or os.getenv( - "HIDDENLAYER_CLIENT_SECRET" - ) - self.api_base = ( - api_base - or os.getenv("HIDDENLAYER_API_BASE") - or "https://api.hiddenlayer.ai" - ) + self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET") + self.api_base = api_base or os.getenv("HIDDENLAYER_API_BASE") or "https://api.hiddenlayer.ai" self.jwt_token = None - auth_url = ( - auth_url - or os.getenv("HIDDENLAYER_AUTH_URL") - or "https://auth.hiddenlayer.ai" - ) + auth_url = auth_url or os.getenv("HIDDENLAYER_AUTH_URL") or "https://auth.hiddenlayer.ai" if is_saas(self.api_base): if not self.hiddenlayer_client_id: - raise RuntimeError( - "`api_id` cannot be None when using the SaaS version of HiddenLayer." - ) + raise RuntimeError("`api_id` cannot be None when using the SaaS version of HiddenLayer.") if not self.hiddenlayer_client_secret: - raise RuntimeError( - "`api_key` cannot be None when using the SaaS version of HiddenLayer." - ) + raise RuntimeError("`api_key` cannot be None when using the SaaS version of HiddenLayer.") self.jwt_token = _get_jwt( auth_url=auth_url, @@ -344,9 +300,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): api_key=self.hiddenlayer_client_secret, ) - self._http_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) super().__init__(**kwargs) @log_guardrail_information @@ -366,11 +320,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): # from the logger object on the response from the model. headers = request_data.get("proxy_server_request", {}).get("headers", {}) if not headers and logging_obj and logging_obj.model_call_details: - headers = ( - logging_obj.model_call_details.get("litellm_params", {}) - .get("metadata", {}) - .get("headers", {}) - ) + headers = logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {}) # put our roundtrip id in the header to the model so we get it on the way back from the model if "hl-roundtrip-id" not in headers: @@ -379,9 +329,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): proxy_req["headers"]["hl-roundtrip-id"] = str(uuid4()) headers["hl-roundtrip-id"] = proxy_req["headers"]["hl-roundtrip-id"] - hl_headers = { - h.lower(): v for h, v in headers.items() if h.lower().startswith("hl-") - } + hl_headers = {h.lower(): v for h, v in headers.items() if h.lower().startswith("hl-")} if "hl-requester-id" not in hl_headers: hl_headers["hl-requester-id"] = "LiteLLM" @@ -401,9 +349,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): "index": 0, "message": { "role": "assistant", - "content": ( - inputs["texts"][0] if inputs.get("texts") else "" - ), + "content": (inputs["texts"][0] if inputs.get("texts") else ""), }, "finish_reason": "stop", } @@ -434,9 +380,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): content = message.get("content", "") if isinstance(content, list): text_parts = [ - item["text"] - for item in content - if isinstance(item, dict) and item.get("type") == "text" + item["text"] for item in content if isinstance(item, dict) and item.get("type") == "text" ] if text_parts: new_texts.append(" ".join(text_parts)) @@ -446,9 +390,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): inputs["texts"] = new_texts elif input_type == "response" and inputs.get("texts"): - inputs["texts"] = [ - output.get("choices", [{}])[-1].get("message", {}).get("content", "") - ] + inputs["texts"] = [output.get("choices", [{}])[-1].get("message", {}).get("content", "")] elif input_type == "response" and inputs.get("tool_calls"): inputs["tool_calls"] = output diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py index 2f22e4c33d6..6c8ea8a7064 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py @@ -26,9 +26,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" verify_ssl = getattr(litellm_params, "verify_ssl", True) # Get optional params - optional_params = getattr( - litellm_params, "optional_params", IBMDetectorOptionalParams() - ) + optional_params = getattr(litellm_params, "optional_params", IBMDetectorOptionalParams()) detector_params = getattr(optional_params, "detector_params", {}) extra_headers = getattr(optional_params, "extra_headers", {}) score_threshold = getattr(optional_params, "score_threshold", None) diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 36d70a37c7c..27ba9f3467c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -60,15 +60,11 @@ class IBMGuardrailDetector(CustomGuardrail): self.base_url = base_url if not self.base_url: - raise ValueError( - "IBM Guardrails base_url is required. Pass base_url parameter." - ) + raise ValueError("IBM Guardrails base_url is required. Pass base_url parameter.") self.detector_id = detector_id if not self.detector_id: - raise ValueError( - "IBM Guardrails detector_id is required. Pass detector_id parameter." - ) + raise ValueError("IBM Guardrails detector_id is required. Pass detector_id parameter.") self.is_detector_server = is_detector_server self.detector_params = detector_params or {} @@ -157,15 +153,12 @@ class IBMGuardrailDetector(CustomGuardrail): # Add guardrail information to request trace if request_data: - guardrail_status = self._determine_guardrail_status_detector_server( - response_json - ) + guardrail_status = self._determine_guardrail_status_detector_server(response_json) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={ "detections": [ - [detection for detection in message_detections] - for message_detections in response_json + [detection for detection in message_detections] for message_detections in response_json ] }, request_data=request_data, @@ -251,9 +244,7 @@ class IBMGuardrailDetector(CustomGuardrail): # Add guardrail information to request trace if request_data: - guardrail_status = self._determine_guardrail_status_orchestrator( - response_json - ) + guardrail_status = self._determine_guardrail_status_orchestrator(response_json) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response=dict(response_json), @@ -288,9 +279,7 @@ class IBMGuardrailDetector(CustomGuardrail): raise - def _filter_detections_by_threshold( - self, detections: List[IBMDetectorDetection] - ) -> List[IBMDetectorDetection]: + def _filter_detections_by_threshold(self, detections: List[IBMDetectorDetection]) -> List[IBMDetectorDetection]: """ Filter detections based on score threshold. @@ -303,11 +292,7 @@ class IBMGuardrailDetector(CustomGuardrail): if self.score_threshold is None: return detections - return [ - detection - for detection in detections - if detection.get("score", 0.0) >= self.score_threshold - ] + return [detection for detection in detections if detection.get("score", 0.0) >= self.score_threshold] def _determine_guardrail_status_detector_server( self, response_json: List[List[IBMDetectorDetection]] @@ -340,9 +325,7 @@ class IBMGuardrailDetector(CustomGuardrail): return "success" except Exception as e: - verbose_proxy_logger.error( - "Error determining IBM Detector Server guardrail status: %s", str(e) - ) + verbose_proxy_logger.error("Error determining IBM Detector Server guardrail status: %s", str(e)) return "guardrail_failed_to_respond" def _determine_guardrail_status_orchestrator( @@ -370,14 +353,10 @@ class IBMGuardrailDetector(CustomGuardrail): return "success" except Exception as e: - verbose_proxy_logger.error( - "Error determining IBM Orchestrator guardrail status: %s", str(e) - ) + verbose_proxy_logger.error("Error determining IBM Orchestrator guardrail status: %s", str(e)) return "guardrail_failed_to_respond" - def _create_error_message_detector_server( - self, detections_list: List[List[IBMDetectorDetection]] - ) -> str: + def _create_error_message_detector_server(self, detections_list: List[List[IBMDetectorDetection]]) -> str: """ Create a detailed error message from detector server response. @@ -391,9 +370,7 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = "IBM Guardrail Detector failed:\n\n" for idx, message_detections in enumerate(detections_list): - filtered_detections = self._filter_detections_by_threshold( - message_detections - ) + filtered_detections = self._filter_detections_by_threshold(message_detections) if filtered_detections: error_message += f"Message {idx + 1}:\n" total_detections += len(filtered_detections) @@ -402,22 +379,15 @@ class IBMGuardrailDetector(CustomGuardrail): detection_type = detection.get("detection_type", "unknown") score = detection.get("score", 0.0) text = detection.get("text", "") - error_message += ( - f" - {detection_type.upper()} (score: {score:.3f})\n" - ) + error_message += f" - {detection_type.upper()} (score: {score:.3f})\n" error_message += f" Text: '{text}'\n" error_message += "\n" - error_message = ( - f"IBM Guardrail Detector failed: {total_detections} violation(s) detected\n\n" - + error_message - ) + error_message = f"IBM Guardrail Detector failed: {total_detections} violation(s) detected\n\n" + error_message return error_message.strip() - def _create_error_message_orchestrator( - self, detections: List[IBMDetectorDetection] - ) -> str: + def _create_error_message_orchestrator(self, detections: List[IBMDetectorDetection]) -> str: """ Create a detailed error message from orchestrator response. @@ -475,9 +445,7 @@ class IBMGuardrailDetector(CustomGuardrail): event_type=GuardrailEventHooks.pre_call, ) - verbose_proxy_logger.debug( - "IBM Detector Server async_pre_call_hook result: %s", result - ) + verbose_proxy_logger.debug("IBM Detector Server async_pre_call_hook result: %s", result) # Check if any detections were found has_violations = False @@ -507,15 +475,11 @@ class IBMGuardrailDetector(CustomGuardrail): filtered = self._filter_detections_by_threshold(orchestrator_result) if filtered and self.block_on_detection: - error_message = self._create_error_message_orchestrator( - orchestrator_result - ) + error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -550,9 +514,7 @@ class IBMGuardrailDetector(CustomGuardrail): event_type=GuardrailEventHooks.during_call, ) - verbose_proxy_logger.debug( - "IBM Detector Server async_moderation_hook result: %s", result - ) + verbose_proxy_logger.debug("IBM Detector Server async_moderation_hook result: %s", result) # Check if any detections were found has_violations = False @@ -582,15 +544,11 @@ class IBMGuardrailDetector(CustomGuardrail): filtered = self._filter_detections_by_threshold(orchestrator_result) if filtered and self.block_on_detection: - error_message = self._create_error_message_orchestrator( - orchestrator_result - ) + error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -612,17 +570,10 @@ class IBMGuardrailDetector(CustomGuardrail): ) from litellm.types.guardrails import GuardrailEventHooks - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return - verbose_proxy_logger.debug( - "async_post_call_success_hook response: %s", response - ) + verbose_proxy_logger.debug("async_post_call_success_hook response: %s", response) # Check if the ModelResponse has text content in its choices # to avoid sending empty content to IBM Detector (e.g., during tool calls) @@ -630,9 +581,7 @@ class IBMGuardrailDetector(CustomGuardrail): has_text_content = False for choice in response.choices: if isinstance(choice, litellm.Choices): - if choice.message.content and isinstance( - choice.message.content, str - ): + if choice.message.content and isinstance(choice.message.content, str): has_text_content = True break @@ -645,12 +594,8 @@ class IBMGuardrailDetector(CustomGuardrail): contents_to_check: List[str] = [] for choice in response.choices: if isinstance(choice, litellm.Choices): - verbose_proxy_logger.debug( - "async_post_call_success_hook choice: %s", choice - ) - if choice.message.content and isinstance( - choice.message.content, str - ): + verbose_proxy_logger.debug("async_post_call_success_hook choice: %s", choice) + if choice.message.content and isinstance(choice.message.content, str): contents_to_check.append(choice.message.content) if contents_to_check: @@ -670,17 +615,13 @@ class IBMGuardrailDetector(CustomGuardrail): # Check if any detections were found has_violations = False for message_detections in result: - filtered = self._filter_detections_by_threshold( - message_detections - ) + filtered = self._filter_detections_by_threshold(message_detections) if filtered: has_violations = True break if has_violations and self.block_on_detection: - error_message = self._create_error_message_detector_server( - result - ) + error_message = self._create_error_message_detector_server(result) raise ValueError(error_message) else: @@ -697,19 +638,13 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result, ) - filtered = self._filter_detections_by_threshold( - orchestrator_result - ) + filtered = self._filter_detections_by_threshold(orchestrator_result) if filtered and self.block_on_detection: - error_message = self._create_error_message_orchestrator( - orchestrator_result - ) + error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index 953275acf14..4575504feb6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -54,15 +54,9 @@ class JavelinGuardrail(CustomGuardrail): application: Optional[str] = None, """ - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.javelin_api_key = api_key or get_secret_str("JAVELIN_API_KEY") - self.api_base = ( - api_base - or get_secret_str("JAVELIN_API_BASE") - or "https://api-dev.javelin.live" - ) + self.api_base = api_base or get_secret_str("JAVELIN_API_BASE") or "https://api-dev.javelin.live" self.api_version = api_version self.guardrail_name = guardrail_name self.javelin_guard_name = javelin_guard_name or guardrail_name @@ -103,9 +97,7 @@ class JavelinGuardrail(CustomGuardrail): exception_str = "" try: - verbose_proxy_logger.debug( - "Javelin Guardrail: Calling Javelin guard API with request: %s", request - ) + verbose_proxy_logger.debug("Javelin Guardrail: Calling Javelin guard API with request: %s", request) url = f"{self.api_base}/{self.api_version}/guardrail/{self.javelin_guard_name}/apply" verbose_proxy_logger.debug("Javelin Guardrail: Calling URL: %s", url) response = await self.async_handler.post( @@ -113,9 +105,7 @@ class JavelinGuardrail(CustomGuardrail): headers=headers, json=dict(request), ) - verbose_proxy_logger.debug( - "Javelin Guardrail: Javelin guard API response: %s", response.json() - ) + verbose_proxy_logger.debug("Javelin Guardrail: Javelin guard API response: %s", response.json()) response_data = response.json() # Ensure the response has the required assessments field if "assessments" not in response_data: @@ -184,9 +174,7 @@ class JavelinGuardrail(CustomGuardrail): event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug( - "Javelin Guardrail: not running guardrail. Guardrail is disabled." - ) + verbose_proxy_logger.debug("Javelin Guardrail: not running guardrail. Guardrail is disabled.") return data if "messages" not in data: @@ -198,11 +186,7 @@ class JavelinGuardrail(CustomGuardrail): clean_metadata = {} if self.metadata: - clean_metadata = { - k: v - for k, v in self.metadata.items() - if k != "standard_logging_guardrail_information" - } + clean_metadata = {k: v for k, v in self.metadata.items() if k != "standard_logging_guardrail_information"} javelin_guard_request = JavelinGuardRequest( input=JavelinGuardInput(text=text), @@ -219,14 +203,10 @@ class JavelinGuardrail(CustomGuardrail): should_reject = False # Debug: Log the full Javelin response - verbose_proxy_logger.debug( - "Javelin Guardrail: Full Javelin response: %s", javelin_response - ) + verbose_proxy_logger.debug("Javelin Guardrail: Full Javelin response: %s", javelin_response) for assessment in assessments: - verbose_proxy_logger.debug( - "Javelin Guardrail: Processing assessment: %s", assessment - ) + verbose_proxy_logger.debug("Javelin Guardrail: Processing assessment: %s", assessment) for assessment_type, assessment_data in assessment.items(): verbose_proxy_logger.debug( "Javelin Guardrail: Processing assessment_type: %s, data: %s", @@ -278,9 +258,7 @@ class JavelinGuardrail(CustomGuardrail): }, ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index 72b9b7dc3c1..3804d1cb93f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -8,9 +8,7 @@ import os import sys -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import json import sys from typing import Dict, List, Literal, Optional, Union @@ -56,15 +54,11 @@ class lakeraAI_Moderation(CustomGuardrail): api_key: Optional[str] = None, **kwargs, ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.lakera_api_key = api_key or os.environ.get("LAKERA_API_KEY") or "" self.moderation_check = moderation_check self.category_thresholds = category_thresholds - self.api_base = ( - api_base or get_secret("LAKERA_API_BASE") or "https://api.lakera.ai" - ) + self.api_base = api_base or get_secret("LAKERA_API_BASE") or "https://api.lakera.ai" super().__init__(**kwargs) #### CALL HOOKS - proxy only #### @@ -79,15 +73,9 @@ class lakeraAI_Moderation(CustomGuardrail): if self.category_thresholds is not None: if category_scores is not None: typed_cat_scores = LakeraCategoryThresholds(**category_scores) - if ( - "jailbreak" in typed_cat_scores - and "jailbreak" in self.category_thresholds - ): + if "jailbreak" in typed_cat_scores and "jailbreak" in self.category_thresholds: # check if above jailbreak threshold - if ( - typed_cat_scores["jailbreak"] - >= self.category_thresholds["jailbreak"] - ): + if typed_cat_scores["jailbreak"] >= self.category_thresholds["jailbreak"]: raise HTTPException( status_code=400, detail={ @@ -95,14 +83,8 @@ class lakeraAI_Moderation(CustomGuardrail): "lakera_ai_response": response, }, ) - if ( - "prompt_injection" in typed_cat_scores - and "prompt_injection" in self.category_thresholds - ): - if ( - typed_cat_scores["prompt_injection"] - >= self.category_thresholds["prompt_injection"] - ): + if "prompt_injection" in typed_cat_scores and "prompt_injection" in self.category_thresholds: + if typed_cat_scores["prompt_injection"] >= self.category_thresholds["prompt_injection"]: raise HTTPException( status_code=400, detail={ @@ -150,9 +132,7 @@ class lakeraAI_Moderation(CustomGuardrail): text = "" _json_data: str = "" if "messages" in data and isinstance(data["messages"], list): - prompt_injection_obj: Optional[GuardrailItem] = ( - litellm.guardrail_name_config_map.get("prompt_injection") - ) + prompt_injection_obj: Optional[GuardrailItem] = litellm.guardrail_name_config_map.get("prompt_injection") if prompt_injection_obj is not None: enabled_roles = prompt_injection_obj.enabled_roles else: @@ -168,9 +148,7 @@ class lakeraAI_Moderation(CustomGuardrail): stringified_roles.append(role.value) elif isinstance(role, str): stringified_roles.append(role) - lakera_input_dict: Dict = { - role: None for role in INPUT_POSITIONING_MAP.keys() - } + lakera_input_dict: Dict = {role: None for role in INPUT_POSITIONING_MAP.keys()} system_message = None tool_call_messages: List = [] for message in data["messages"]: @@ -212,15 +190,11 @@ class lakeraAI_Moderation(CustomGuardrail): lakera_input = [ v - for k, v in sorted( - lakera_input_dict.items(), key=lambda x: INPUT_POSITIONING_MAP[x[0]] - ) + for k, v in sorted(lakera_input_dict.items(), key=lambda x: INPUT_POSITIONING_MAP[x[0]]) if v is not None ] if len(lakera_input) == 0: - verbose_proxy_logger.debug( - "Skipping lakera prompt injection, no roles with messages found" - ) + verbose_proxy_logger.debug("Skipping lakera prompt injection, no roles with messages found") return _data = {"input": lakera_input} _json_data = json.dumps( @@ -327,17 +301,10 @@ class lakeraAI_Moderation(CustomGuardrail): else: # v2 guardrails implementation - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is not True: return None - return await self._check( - data=data, user_api_key_dict=user_api_key_dict, call_type=call_type - ) + return await self._check(data=data, user_api_key_dict=user_api_key_dict, call_type=call_type) @log_guardrail_information async def async_moderation_hook( @@ -366,6 +333,4 @@ class lakeraAI_Moderation(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return - return await self._check( - data=data, user_api_key_dict=user_api_key_dict, call_type=call_type - ) + return await self._check(data=data, user_api_key_dict=user_api_key_dict, call_type=call_type) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index c6af4d3c428..e79a3e7b3d8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -59,14 +59,10 @@ class LakeraAIGuardrail(CustomGuardrail): dev_info: Optional[bool] = True, on_flagged: Optional[str] = "block", Action to take when content is flagged: "block" or "monitor" """ - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.lakera_api_key = api_key or os.environ.get("LAKERA_API_KEY") or "" self.project_id = project_id - self.api_base = ( - api_base or get_secret_str("LAKERA_API_BASE") or "https://api.lakera.ai" - ) + self.api_base = api_base or get_secret_str("LAKERA_API_BASE") or "https://api.lakera.ai" self.payload: Optional[bool] = payload self.breakdown: Optional[bool] = breakdown self.metadata: Optional[Dict] = metadata @@ -106,9 +102,7 @@ class LakeraAIGuardrail(CustomGuardrail): headers={"Authorization": f"Bearer {self.lakera_api_key}"}, json=request, ) - verbose_proxy_logger.debug( - "Lakera AI v2 guard response: %s", response.json() - ) + verbose_proxy_logger.debug("Lakera AI v2 guard response: %s", response.json()) lakera_response = LakeraAIResponse(**response.json()) return lakera_response, masked_entity_count except Exception as e: @@ -121,9 +115,7 @@ class LakeraAIGuardrail(CustomGuardrail): #################################################### guardrail_json_response: Union[Exception, str, dict, List[dict]] = {} if status == "success": - copy_lakera_response_dict = ( - dict(copy.deepcopy(lakera_response)) if lakera_response else {} - ) + copy_lakera_response_dict = dict(copy.deepcopy(lakera_response)) if lakera_response else {} # payload contains PII, we don't want to log it copy_lakera_response_dict.pop("payload") guardrail_json_response = copy_lakera_response_dict @@ -214,17 +206,13 @@ class LakeraAIGuardrail(CustomGuardrail): event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug( - "Lakera AI: not running guardrail. Guardrail is disabled." - ) + verbose_proxy_logger.debug("Lakera AI: not running guardrail. Guardrail is disabled.") return data # Covers multimodal list content + Responses-API input. new_messages = build_inspection_messages(data) if not new_messages: - verbose_proxy_logger.warning( - "Lakera AI: not running guardrail. No inspectable text in data" - ) + verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data") return data # Mask-in-place uses offsets returned by Lakera and can only @@ -248,10 +236,7 @@ class LakeraAIGuardrail(CustomGuardrail): ######################################################### if lakera_guardrail_response.get("flagged") is True: # If only PII violations exist, mask the PII (string input only). - if ( - self._is_only_pii_violation(lakera_guardrail_response) - and not is_multimodal_input - ): + if self._is_only_pii_violation(lakera_guardrail_response) and not is_multimodal_input: redacted_messages = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, @@ -261,9 +246,7 @@ class LakeraAIGuardrail(CustomGuardrail): # backend reads ``input``; writing only to ``messages`` # would let unredacted PII reach the LLM for /v1/responses. apply_redacted_messages_back(data, list(redacted_messages)) # type: ignore[arg-type] - verbose_proxy_logger.debug( - "Lakera AI: Masked PII in messages instead of blocking request" - ) + verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request") else: # Check on_flagged setting if self.on_flagged == "monitor": @@ -275,16 +258,12 @@ class LakeraAIGuardrail(CustomGuardrail): # Either non-PII violations, or PII on multimodal input # (which cannot be masked in place without dropping # image/audio parts) — raise the standard block error. - raise self._get_http_exception_for_blocked_guardrail( - lakera_guardrail_response - ) + raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## ######################################################### - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -304,9 +283,7 @@ class LakeraAIGuardrail(CustomGuardrail): new_messages = build_inspection_messages(data) if not new_messages: - verbose_proxy_logger.warning( - "Lakera AI: not running guardrail. No inspectable text in data" - ) + verbose_proxy_logger.warning("Lakera AI: not running guardrail. No inspectable text in data") return # See ``async_pre_call_hook`` — multimodal input degrades to @@ -326,10 +303,7 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 2. Handle flagged content ########## ######################################################### if lakera_guardrail_response.get("flagged") is True: - if ( - self._is_only_pii_violation(lakera_guardrail_response) - and not is_multimodal_input - ): + if self._is_only_pii_violation(lakera_guardrail_response) and not is_multimodal_input: redacted_messages = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, @@ -339,25 +313,19 @@ class LakeraAIGuardrail(CustomGuardrail): # backend reads ``input``; writing only to ``messages`` # would let unredacted PII reach the LLM for /v1/responses. apply_redacted_messages_back(data, list(redacted_messages)) # type: ignore[arg-type] - verbose_proxy_logger.debug( - "Lakera AI: Masked PII in messages instead of blocking request" - ) + verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request") else: if self.on_flagged == "monitor": verbose_proxy_logger.warning( "Lakera Guardrail: Monitoring mode - violation detected but allowing request" ) elif self.on_flagged == "block": - raise self._get_http_exception_for_blocked_guardrail( - lakera_guardrail_response - ) + raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## ######################################################### - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -422,34 +390,22 @@ class LakeraAIGuardrail(CustomGuardrail): for idx, msg in enumerate(assistant_messages): if idx < len(choice_indices): choice_idx = choice_indices[idx] - response_dict["choices"][choice_idx]["message"]["content"] = ( - msg.get("content", "") - ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + response_dict["choices"][choice_idx]["message"]["content"] = msg.get("content", "") + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return ModelResponse(**response_dict) if self.on_flagged == "monitor": - verbose_proxy_logger.warning( - "Lakera Guardrail: Post-call violation detected in monitor mode" - ) + verbose_proxy_logger.warning("Lakera Guardrail: Post-call violation detected in monitor mode") # Allow response to proceed elif self.on_flagged == "block": - raise self._get_http_exception_for_blocked_guardrail( - lakera_guardrail_response - ) + raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) # Record applied guardrail - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response - def _is_only_pii_violation( - self, lakera_response: Optional[LakeraAIResponse] - ) -> bool: + def _is_only_pii_violation(self, lakera_response: Optional[LakeraAIResponse]) -> bool: """ Returns True if there are only PII violations in the response. """ @@ -472,9 +428,7 @@ class LakeraAIGuardrail(CustomGuardrail): # Return True only if there are violations and they are all PII return has_violations - def _get_http_exception_for_blocked_guardrail( - self, lakera_response: Optional[LakeraAIResponse] - ) -> HTTPException: + def _get_http_exception_for_blocked_guardrail(self, lakera_response: Optional[LakeraAIResponse]) -> HTTPException: """ Get the HTTP exception for a blocked guardrail, similar to Bedrock's implementation. """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index eea378e43bf..36fbd73c5bd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -103,14 +103,10 @@ class LassoGuardrail(CustomGuardrail): mask: Optional[bool] = False, **kwargs, ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.lasso_api_key = lasso_api_key or api_key or os.environ.get("LASSO_API_KEY") self.user_id = user_id or os.environ.get("LASSO_USER_ID") - self.conversation_id = conversation_id or os.environ.get( - "LASSO_CONVERSATION_ID" - ) + self.conversation_id = conversation_id or os.environ.get("LASSO_CONVERSATION_ID") self.mask = mask or False if self.lasso_api_key is None: @@ -119,11 +115,7 @@ class LassoGuardrail(CustomGuardrail): "pass it as a parameter to the guardrail in the config file" ) - self.api_base = ( - api_base - or os.getenv("LASSO_API_BASE") - or "https://server.lasso.security/gateway/v3" - ) + self.api_base = api_base or os.getenv("LASSO_API_BASE") or "https://server.lasso.security/gateway/v3" verbose_proxy_logger.debug( f"Lasso guardrail initialized: {kwargs.get('guardrail_name', 'unknown')}, " @@ -213,9 +205,7 @@ class LassoGuardrail(CustomGuardrail): # The conversation_id is being stored in the cache so it can be used by the post_call hook self._get_or_generate_conversation_id(data, global_cache) - return await self._run_lasso_guardrail( - data, global_cache, message_type="PROMPT" - ) + return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT") @log_guardrail_information async def async_moderation_hook( @@ -269,9 +259,7 @@ class LassoGuardrail(CustomGuardrail): continue msg = choice.message if msg.content: - response_messages.append( - {"role": "assistant", "content": msg.content} - ) + response_messages.append({"role": "assistant", "content": msg.content}) for call in getattr(msg, "tool_calls", None) or []: call_id, name, input_data = self._extract_tool_call_fields(call) if not call_id or not name: @@ -298,50 +286,31 @@ class LassoGuardrail(CustomGuardrail): # Handle masking for post-call if self.mask: headers = self._prepare_headers(response_data, global_cache) - payload = self._prepare_payload( - response_messages, response_data, global_cache, "COMPLETION" - ) + payload = self._prepare_payload(response_messages, response_data, global_cache, "COMPLETION") api_url = f"{self.api_base}/classifix" try: - lasso_response = await self._call_lasso_api( - headers=headers, payload=payload, api_url=api_url - ) + lasso_response = await self._call_lasso_api(headers=headers, payload=payload, api_url=api_url) self._process_lasso_response(lasso_response) # Apply masking to the actual response if masked content is available masked_messages = lasso_response.get("messages") - if ( - lasso_response.get("violations_detected") - and masked_messages - ): - self._apply_masking_to_model_response( - response, masked_messages - ) - verbose_proxy_logger.debug( - "Applied Lasso masking to model response" - ) + if lasso_response.get("violations_detected") and masked_messages: + self._apply_masking_to_model_response(response, masked_messages) + verbose_proxy_logger.debug("Applied Lasso masking to model response") except Exception as e: if isinstance(e, HTTPException): raise e - verbose_proxy_logger.error( - f"Error in post-call Lasso masking: {str(e)}" - ) - raise LassoGuardrailAPIError( - f"Failed to apply post-call masking: {str(e)}" - ) + verbose_proxy_logger.error(f"Error in post-call Lasso masking: {str(e)}") + raise LassoGuardrailAPIError(f"Failed to apply post-call masking: {str(e)}") else: # Use the same data for conversation_id consistency (no cache access needed) - await self._run_lasso_guardrail( - response_data, cache=global_cache, message_type="COMPLETION" - ) + await self._run_lasso_guardrail(response_data, cache=global_cache, message_type="COMPLETION") verbose_proxy_logger.debug("Post-call Lasso validation completed") else: verbose_proxy_logger.warning("No response messages found to validate") else: - verbose_proxy_logger.warning( - f"Unexpected response type for post-call hook: {type(response)}" - ) + verbose_proxy_logger.warning(f"Unexpected response type for post-call hook: {type(response)}") return response @@ -429,11 +398,7 @@ class LassoGuardrail(CustomGuardrail): HTTPException: If blocking violations are detected """ raw_messages: List[Dict[str, Any]] = data.get("messages") or [] - messages: List[Dict[str, Any]] = ( - self._expand_messages_for_classification(raw_messages) - if raw_messages - else [] - ) + messages: List[Dict[str, Any]] = self._expand_messages_for_classification(raw_messages) if raw_messages else [] messages_count = len(messages) if data.get("input") is not None: # Responses-API payloads carry text in data["input"]. Inspect it @@ -449,9 +414,7 @@ class LassoGuardrail(CustomGuardrail): # classify endpoint (which still raises on BLOCK actions) and # leave the original payload intact. if self.mask and not has_non_string_content(data): - return await self._handle_masking( - data, cache, message_type, messages, messages_count - ) + return await self._handle_masking(data, cache, message_type, messages, messages_count) return await self._handle_classification(data, cache, message_type, messages) async def _handle_classification( @@ -490,9 +453,7 @@ class LassoGuardrail(CustomGuardrail): headers = self._prepare_headers(data, cache) payload = self._prepare_payload(messages, data, cache, message_type) api_url = f"{self.api_base}/classifix" - response = await self._call_lasso_api( - headers=headers, payload=payload, api_url=api_url - ) + response = await self._call_lasso_api(headers=headers, payload=payload, api_url=api_url) self._process_lasso_response(response) # Apply masking to messages if violations detected and masked messages are available. @@ -503,17 +464,11 @@ class LassoGuardrail(CustomGuardrail): masked_for_messages = masked[:messages_count] masked_for_input = masked[messages_count:] if data.get("messages"): - data["messages"] = self._map_masked_messages_back( - data["messages"], masked_for_messages - ) + data["messages"] = self._map_masked_messages_back(data["messages"], masked_for_messages) # Also update data["input"] for Responses-API payloads so the # unredacted text doesn't leak through that field. if isinstance(data.get("input"), str): - text_parts = [ - msg["content"] - for msg in masked_for_input - if isinstance(msg.get("content"), str) - ] + text_parts = [msg["content"] for msg in masked_for_input if isinstance(msg.get("content"), str)] if text_parts: data["input"] = "\n".join(text_parts) self._log_masking_applied(message_type, dict(response)) @@ -562,10 +517,7 @@ class LassoGuardrail(CustomGuardrail): 1 for m in original_messages if m.get("role") != "tool" - and ( - (isinstance(m.get("content"), str) and m.get("content")) - or isinstance(m.get("content"), list) - ) + and ((isinstance(m.get("content"), str) and m.get("content")) or isinstance(m.get("content"), list)) ) apply_text_cursor = original_text_count == len(masked_text) if not apply_text_cursor and masked_text: @@ -595,9 +547,7 @@ class LassoGuardrail(CustomGuardrail): msg["content"] = masked_text[text_cursor] text_cursor += 1 if role == "assistant" and orig_msg.get("tool_calls"): - msg["tool_calls"] = self._update_tool_calls_from_masked( - orig_msg["tool_calls"], masked_tool_use - ) + msg["tool_calls"] = self._update_tool_calls_from_masked(orig_msg["tool_calls"], masked_tool_use) elif isinstance(content, list): # Multimodal list content was flattened to a text string before @@ -607,14 +557,10 @@ class LassoGuardrail(CustomGuardrail): msg["content"] = masked_text[text_cursor] text_cursor += 1 if role == "assistant" and orig_msg.get("tool_calls"): - msg["tool_calls"] = self._update_tool_calls_from_masked( - orig_msg["tool_calls"], masked_tool_use - ) + msg["tool_calls"] = self._update_tool_calls_from_masked(orig_msg["tool_calls"], masked_tool_use) elif role == "assistant" and not content and orig_msg.get("tool_calls"): - msg["tool_calls"] = self._update_tool_calls_from_masked( - orig_msg["tool_calls"], masked_tool_use - ) + msg["tool_calls"] = self._update_tool_calls_from_masked(orig_msg["tool_calls"], masked_tool_use) result.append(msg) @@ -673,14 +619,10 @@ class LassoGuardrail(CustomGuardrail): elif error.response.status_code == 429: raise LassoGuardrailAPIError("Lasso API rate limit exceeded") else: - raise LassoGuardrailAPIError( - f"API error: {error.response.status_code}" - ) + raise LassoGuardrailAPIError(f"API error: {error.response.status_code}") # Generic error handling - raise LassoGuardrailAPIError( - f"Failed to verify request safety with Lasso API: {str(error)}" - ) + raise LassoGuardrailAPIError(f"Failed to verify request safety with Lasso API: {str(error)}") def _log_masking_applied( self, @@ -700,9 +642,7 @@ class LassoGuardrail(CustomGuardrail): }, ) - def _expand_messages_for_classification( - self, messages: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: + def _expand_messages_for_classification(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Convert raw OpenAI-format messages to Lasso API format with content blocks. @@ -718,9 +658,7 @@ class LassoGuardrail(CustomGuardrail): if role == "tool": tool_call_id = msg.get("tool_call_id") if not tool_call_id: - verbose_proxy_logger.warning( - "Skipping tool message without tool_call_id" - ) + verbose_proxy_logger.warning("Skipping tool message without tool_call_id") continue # Flatten multimodal list content to text so Lasso's # tool_result.content field receives a string. @@ -728,9 +666,7 @@ class LassoGuardrail(CustomGuardrail): text_parts = [ part["text"] for part in content - if isinstance(part, dict) - and part.get("type") == "text" - and part.get("text") + if isinstance(part, dict) and part.get("type") == "text" and part.get("text") ] tool_result_content = "\n".join(text_parts) else: @@ -752,9 +688,7 @@ class LassoGuardrail(CustomGuardrail): text_parts = [ part["text"] for part in content - if isinstance(part, dict) - and part.get("type") == "text" - and part.get("text") + if isinstance(part, dict) and part.get("type") == "text" and part.get("text") ] if text_parts: expanded.append({"role": role, "content": "\n".join(text_parts)}) @@ -872,9 +806,7 @@ class LassoGuardrail(CustomGuardrail): ) -> LassoResponse: """Call the Lasso API and return the response.""" url = api_url or f"{self.api_base}/classify" - verbose_proxy_logger.debug( - f"Calling Lasso API with messageType: {payload.get('messageType')}" - ) + verbose_proxy_logger.debug(f"Calling Lasso API with messageType: {payload.get('messageType')}") response = await self.async_handler.post( url=url, headers=headers, @@ -912,9 +844,7 @@ class LassoGuardrail(CustomGuardrail): """ if response and response.get("violations_detected") is True: violated_deputies = self._parse_violated_deputies(response) - verbose_proxy_logger.warning( - f"Lasso guardrail detected violations: {violated_deputies}" - ) + verbose_proxy_logger.warning(f"Lasso guardrail detected violations: {violated_deputies}") # Check if any findings have "BLOCK" action blocking_violations = self._check_for_blocking_actions(response) @@ -1001,11 +931,7 @@ class LassoGuardrail(CustomGuardrail): masked_text.append(content) # Count text-bearing choices to verify 1:1 mapping with masked texts. - original_text_count = sum( - 1 - for c in model_response.choices - if hasattr(c, "message") and c.message.content - ) + original_text_count = sum(1 for c in model_response.choices if hasattr(c, "message") and c.message.content) apply_text = original_text_count == len(masked_text) if not apply_text and masked_text: verbose_proxy_logger.warning( @@ -1025,9 +951,7 @@ class LassoGuardrail(CustomGuardrail): if msg.content and apply_text and text_cursor < len(masked_text): msg.content = masked_text[text_cursor] text_cursor += 1 - verbose_proxy_logger.debug( - f"Applied masked text content to choice {text_cursor}" - ) + verbose_proxy_logger.debug(f"Applied masked text content to choice {text_cursor}") for call in getattr(msg, "tool_calls", None) or []: call_id = self._get_field(call, "id") @@ -1042,9 +966,7 @@ class LassoGuardrail(CustomGuardrail): func = getattr(call, "function", None) if func: func.arguments = json.dumps(masked_input) - verbose_proxy_logger.debug( - f"Applied masked tool_call arguments for call_id={call_id}" - ) + verbose_proxy_logger.debug(f"Applied masked tool_call arguments for call_id={call_id}") @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py index 8eb49602647..8f1cee2672a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py @@ -44,16 +44,10 @@ def initialize_guardrail( severity_threshold=getattr(litellm_params, "severity_threshold", "medium"), llm_router=llm_router, image_model=getattr(litellm_params, "image_model", None), - competitor_intent_config=getattr( - litellm_params, "competitor_intent_config", None - ), - end_session_after_n_fails=getattr( - litellm_params, "end_session_after_n_fails", None - ), + competitor_intent_config=getattr(litellm_params, "competitor_intent_config", None), + end_session_after_n_fails=getattr(litellm_params, "end_session_after_n_fails", None), on_violation=getattr(litellm_params, "on_violation", None), - realtime_violation_message=getattr( - litellm_params, "realtime_violation_message", None - ), + realtime_violation_message=getattr(litellm_params, "realtime_violation_message", None), ) litellm.logging_callback_manager.add_litellm_callback(content_filter_guardrail) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py index 9ab5b9c1d5b..04d8459b821 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py @@ -106,7 +106,9 @@ AIRLINE_COMPARISON_SIGNALS = [ # Explicit markers: strong override when present. AIRLINE_EXPLICIT_COMPETITOR_MARKER = r"\b(airways?|airline|carrier)\b" -AIRLINE_EXPLICIT_OTHER_MEANING_MARKER = r"\b(fly|travel|going|visit|layover|stopover|transit)\b.{0,12}\b(to|in|via|from)\b.{0,8}\b" +AIRLINE_EXPLICIT_OTHER_MEANING_MARKER = ( + r"\b(fly|travel|going|visit|layover|stopover|transit)\b.{0,12}\b(to|in|via|from)\b.{0,8}\b" +) _MAJOR_AIRLINES_PATH = Path(__file__).resolve().parent / "major_airlines.json" @@ -159,25 +161,17 @@ class AirlineCompetitorIntentChecker(BaseCompetitorIntentChecker): if not merged.get("explicit_competitor_marker"): merged["explicit_competitor_marker"] = AIRLINE_EXPLICIT_COMPETITOR_MARKER if not merged.get("explicit_other_meaning_marker"): - merged["explicit_other_meaning_marker"] = ( - AIRLINE_EXPLICIT_OTHER_MEANING_MARKER - ) + merged["explicit_other_meaning_marker"] = AIRLINE_EXPLICIT_OTHER_MEANING_MARKER if not merged.get("domain_words"): merged["domain_words"] = ["airline", "airlines", "carrier"] if not merged.get("competitors"): - merged["competitors"] = _load_competitors_excluding_brand( - merged.get("brand_self") or [] - ) + merged["competitors"] = _load_competitors_excluding_brand(merged.get("brand_self") or []) super().__init__(merged) self._other_meaning_signals = list(merged.get("other_meaning_signals") or []) self._competitor_signals = list(merged.get("competitor_signals") or []) self._other_meaning_anchors = list(merged.get("other_meaning_anchors") or []) - self._explicit_competitor_marker = _compile_marker( - merged.get("explicit_competitor_marker") - ) - self._explicit_other_meaning_marker = _compile_marker( - merged.get("explicit_other_meaning_marker") - ) + self._explicit_competitor_marker = _compile_marker(merged.get("explicit_competitor_marker")) + self._explicit_other_meaning_marker = _compile_marker(merged.get("explicit_other_meaning_marker")) def _classify_ambiguous(self, text: str, token: str) -> Tuple[str, float]: """Other meaning vs competitor using airline signals and explicit markers.""" @@ -188,10 +182,7 @@ class AirlineCompetitorIntentChecker(BaseCompetitorIntentChecker): and _word_boundary_match(text_lower, token.lower()) ): return "COMPETITOR", 0.85 - if ( - self._explicit_other_meaning_marker - and self._explicit_other_meaning_marker.search(text_lower) - ): + if self._explicit_other_meaning_marker and self._explicit_other_meaning_marker.search(text_lower): return "OTHER_MEANING", 0.85 # Operational-only: baggage/lounge/check-in/refund with no comparison → product query has_comparison = _count_signals(text_lower, AIRLINE_COMPARISON_SIGNALS) > 0 diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/base.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/base.py index 4ebff5fb3c1..e41f69784c6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/base.py @@ -16,9 +16,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor ZERO_WIDTH = re.compile(r"[\u200b-\u200d\u2060\ufeff]") LEET = {"@": "a", "4": "a", "0": "o", "3": "e", "1": "i", "5": "s", "7": "t"} -OTHER_MEANING_DEFAULT_THRESHOLD = ( - 0.65 # Below this → treat as non-competitor (safe default). -) +OTHER_MEANING_DEFAULT_THRESHOLD = 0.65 # Below this → treat as non-competitor (safe default). def normalize(text: str) -> str: @@ -66,12 +64,8 @@ class BaseCompetitorIntentChecker: """ def __init__(self, config: Dict[str, Any]) -> None: - self.brand_self: List[str] = [ - s.lower().strip() for s in (config.get("brand_self") or []) if s - ] - competitors: List[str] = [ - s.lower().strip() for s in (config.get("competitors") or []) if s - ] + self.brand_self: List[str] = [s.lower().strip() for s in (config.get("brand_self") or []) if s] + competitors: List[str] = [s.lower().strip() for s in (config.get("competitors") or []) if s] aliases_map: Dict[str, List[str]] = config.get("competitor_aliases") or {} self.competitor_canonical: Dict[str, str] = {} self._competitor_tokens: Set[str] = set() @@ -84,9 +78,7 @@ class BaseCompetitorIntentChecker: self._competitor_tokens.add(a) self.competitor_canonical[a] = c - other: List[str] = [ - s.lower().strip() for s in (config.get("locations") or []) if s - ] + other: List[str] = [s.lower().strip() for s in (config.get("locations") or []) if s] self._other_meaning_tokens: Set[str] = set(other) self._ambiguous: Set[str] = self._competitor_tokens & self._other_meaning_tokens @@ -94,12 +86,8 @@ class BaseCompetitorIntentChecker: self.threshold_high = float(config.get("threshold_high", 0.70)) self.threshold_medium = float(config.get("threshold_medium", 0.45)) self.threshold_low = float(config.get("threshold_low", 0.30)) - self.reframe_message_template: Optional[str] = config.get( - "reframe_message_template" - ) - self.refuse_message_template: Optional[str] = config.get( - "refuse_message_template" - ) + self.reframe_message_template: Optional[str] = config.get("reframe_message_template") + self.refuse_message_template: Optional[str] = config.get("refuse_message_template") self._comparison_words: List[str] = list( config.get("comparison_words") or [ @@ -114,9 +102,7 @@ class BaseCompetitorIntentChecker: "ranked", ] ) - self._domain_words: List[str] = [ - s.lower().strip() for s in (config.get("domain_words") or []) if s - ] + self._domain_words: List[str] = [s.lower().strip() for s in (config.get("domain_words") or []) if s] def _classify_ambiguous(self, text: str, token: str) -> Tuple[str, float]: """ @@ -154,19 +140,13 @@ class BaseCompetitorIntentChecker: for b in self.brand_self: if _word_boundary_match(normalized, b): entities["brand_self"].append(b) - evidence.append( - {"type": "entity", "key": "brand_self", "value": b, "match": b} - ) + evidence.append({"type": "entity", "key": "brand_self", "value": b, "match": b}) matches = self._find_matches(text) if not matches: - has_comparison = any( - re.search(r"\b" + re.escape(w) + r"\b", normalized) - for w in self._comparison_words - ) + has_comparison = any(re.search(r"\b" + re.escape(w) + r"\b", normalized) for w in self._comparison_words) has_domain = self._domain_words and any( - re.search(r"\b" + re.escape(w) + r"\b", normalized) - for w in self._domain_words + re.search(r"\b" + re.escape(w) + r"\b", normalized) for w in self._domain_words ) if has_comparison and has_domain: evidence.append( @@ -201,9 +181,7 @@ class BaseCompetitorIntentChecker: for token, canonical, _ in matches: label, conf = self._classify_ambiguous(normalized, token) if label == "OTHER_MEANING": - evidence.append( - {"type": "signal", "key": "other_meaning", "match": token} - ) + evidence.append({"type": "signal", "key": "other_meaning", "match": token}) continue if label == "COMPETITOR": competitor_resolved.append(canonical) @@ -238,14 +216,9 @@ class BaseCompetitorIntentChecker: "evidence": evidence, } - has_comparison = any( - re.search(r"\b" + re.escape(w) + r"\b", normalized) - for w in self._comparison_words - ) + has_comparison = any(re.search(r"\b" + re.escape(w) + r"\b", normalized) for w in self._comparison_words) if has_comparison: - evidence.append( - {"type": "signal", "key": "comparison", "match": "comparison language"} - ) + evidence.append({"type": "signal", "key": "comparison", "match": "comparison language"}) confidence = 0.75 if has_comparison else 0.55 if confidence >= self.threshold_high: intent = "competitor_comparison" @@ -256,9 +229,7 @@ class BaseCompetitorIntentChecker: else: intent = "other" - resolved_action_hint: CompetitorActionHint = cast( - CompetitorActionHint, self.policy.get(intent, "allow") - ) + resolved_action_hint: CompetitorActionHint = cast(CompetitorActionHint, self.policy.get(intent, "allow")) if intent == "log_only": resolved_action_hint = "log_only" if intent == "other": @@ -268,8 +239,7 @@ class BaseCompetitorIntentChecker: "intent": cast(CompetitorIntentType, intent), "confidence": round(confidence, 2), "entities": entities, - "signals": ["competitor_resolved"] - + (["comparison"] if has_comparison else []), + "signals": ["competitor_resolved"] + (["comparison"] if has_comparison else []), "action_hint": resolved_action_hint, "evidence": evidence, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index cc944fb46dc..ede36b23216 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -113,25 +113,17 @@ class CategoryConfig: self.keywords = keywords self.exceptions = [e.lower() for e in exceptions] # New fields for conditional child safety logic - self.identifier_words = ( - [w.lower() for w in identifier_words] if identifier_words else [] - ) + self.identifier_words = [w.lower() for w in identifier_words] if identifier_words else [] self.always_block_keywords = always_block_keywords or [] self.inherit_from = inherit_from - self.additional_block_words = ( - [w.lower() for w in additional_block_words] - if additional_block_words - else [] - ) + self.additional_block_words = [w.lower() for w in additional_block_words] if additional_block_words else [] # Phrase patterns: regex patterns for catching paraphrases self.phrase_patterns: List[Tuple[str, Pattern]] = [] for p in phrase_patterns or []: try: self.phrase_patterns.append((p, re.compile(p, re.IGNORECASE))) except re.error: - verbose_proxy_logger.warning( - f"Invalid phrase pattern in {category_name}: {p}" - ) + verbose_proxy_logger.warning(f"Invalid phrase pattern in {category_name}: {p}") class ContentFilterGuardrail(CustomGuardrail): @@ -158,9 +150,7 @@ class ContentFilterGuardrail(CustomGuardrail): patterns: Optional[List[ContentFilterPattern]] = None, blocked_words: Optional[List[BlockedWord]] = None, blocked_words_file: Optional[str] = None, - event_hook: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode] - ] = None, + event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = None, default_on: bool = False, pattern_redaction_format: Optional[str] = None, keyword_redaction_tag: Optional[str] = None, @@ -203,9 +193,7 @@ class ContentFilterGuardrail(CustomGuardrail): self.guardrail_provider = "litellm_content_filter" self.config_guardrail_id = guardrail_id self.config_policy_template = policy_template - self.pattern_redaction_format = ( - pattern_redaction_format or self.PATTERN_REDACTION_FORMAT - ) + self.pattern_redaction_format = pattern_redaction_format or self.PATTERN_REDACTION_FORMAT self.keyword_redaction_tag = keyword_redaction_tag or self.KEYWORD_REDACTION_STR self.severity_threshold = severity_threshold self.llm_router = llm_router @@ -216,9 +204,7 @@ class ContentFilterGuardrail(CustomGuardrail): str, Tuple[str, str, ContentFilterAction] ] = {} # keyword -> (category, severity, action) # Always-block keywords are checked after exceptions (exceptions take precedence) - self.always_block_category_keywords: Dict[ - str, Tuple[str, str, ContentFilterAction] - ] = {} + self.always_block_category_keywords: Dict[str, Tuple[str, str, ContentFilterAction]] = {} # Store conditional categories (identifier_words + block_words) self.conditional_categories: Dict[ str, Dict[str, Any] @@ -272,25 +258,16 @@ class ContentFilterGuardrail(CustomGuardrail): f"and {len(self.blocked_words)} blocked words" ) verbose_proxy_logger.debug( - f"Loaded {len(self.loaded_categories)} categories with " - f"{len(self.category_keywords)} keywords" + f"Loaded {len(self.loaded_categories)} categories with {len(self.category_keywords)} keywords" ) - def _init_competitor_intent_checker( - self, competitor_intent_config: Dict[str, Any] - ) -> None: + def _init_competitor_intent_checker(self, competitor_intent_config: Dict[str, Any]) -> None: try: - competitor_intent_type = competitor_intent_config.get( - "competitor_intent_type", "airline" - ) + competitor_intent_type = competitor_intent_config.get("competitor_intent_type", "airline") if competitor_intent_type == "generic": - self._competitor_intent_checker = BaseCompetitorIntentChecker( - competitor_intent_config - ) + self._competitor_intent_checker = BaseCompetitorIntentChecker(competitor_intent_config) else: - self._competitor_intent_checker = AirlineCompetitorIntentChecker( - competitor_intent_config - ) + self._competitor_intent_checker = AirlineCompetitorIntentChecker(competitor_intent_config) verbose_proxy_logger.debug( "ContentFilterGuardrail: competitor intent checker enabled (%s)", competitor_intent_type, @@ -336,13 +313,10 @@ class ContentFilterGuardrail(CustomGuardrail): common = os.path.commonpath([resolved, allowed]) except ValueError: # commonpath() raises ValueError on Windows when paths span different drives - raise ValueError( - f"Category file path '{path}' is outside the allowed categories directory" - ) + raise ValueError(f"Category file path '{path}' is outside the allowed categories directory") if common != allowed: raise ValueError( - f"Category file path '{path}' is outside the allowed " - f"categories directory '{categories_dir}'" + f"Category file path '{path}' is outside the allowed categories directory '{categories_dir}'" ) def _resolve_category_file_path(self, file_path: str) -> str: @@ -380,10 +354,7 @@ class ContentFilterGuardrail(CustomGuardrail): and ``LITELLM_CONTENT_FILTER_ALLOW_EXTERNAL_PATHS`` is not set. """ module_dir = os.path.dirname(__file__) - allow_external = ( - os.environ.get("LITELLM_CONTENT_FILTER_ALLOW_EXTERNAL_PATHS", "").lower() - == "true" - ) + allow_external = os.environ.get("LITELLM_CONTENT_FILTER_ALLOW_EXTERNAL_PATHS", "").lower() == "true" if os.path.isabs(file_path) or os.path.exists(file_path): if not allow_external: @@ -417,9 +388,7 @@ class ContentFilterGuardrail(CustomGuardrail): # path anyway to reject traversal attempts (e.g. "../../../../etc/passwd") # regardless of CWD or whether the target file exists. if not allow_external: - self._assert_within_categories_dir( - os.path.join(module_dir, file_path), module_dir - ) + self._assert_within_categories_dir(os.path.join(module_dir, file_path), module_dir) return file_path def _load_categories(self, categories: List[ContentFilterCategoryConfig]) -> None: @@ -439,30 +408,23 @@ class ContentFilterGuardrail(CustomGuardrail): for cat_config in categories: category_name = cat_config.get("category") if not category_name or not isinstance(category_name, str): - verbose_proxy_logger.warning( - "Category name missing or invalid in config, skipping" - ) + verbose_proxy_logger.warning("Category name missing or invalid in config, skipping") continue # Prevent path traversal via category_name (e.g. "../../etc/passwd") if not re.match(r"^[a-zA-Z0-9_\-]+$", category_name): - verbose_proxy_logger.warning( - f"Category name '{category_name}' contains invalid characters, skipping" - ) + verbose_proxy_logger.warning(f"Category name '{category_name}' contains invalid characters, skipping") continue enabled = cat_config.get("enabled", True) action = cat_config.get("action") severity_threshold = ( - cat_config.get("severity_threshold", self.severity_threshold) - or self.severity_threshold + cat_config.get("severity_threshold", self.severity_threshold) or self.severity_threshold ) custom_file = cat_config.get("category_file") if not enabled: - verbose_proxy_logger.debug( - f"Category {category_name} is disabled, skipping" - ) + verbose_proxy_logger.debug(f"Category {category_name} is disabled, skipping") continue # Load category file (custom or default) @@ -470,9 +432,7 @@ class ContentFilterGuardrail(CustomGuardrail): try: category_file_path = self._resolve_category_file_path(custom_file) except ValueError as e: - verbose_proxy_logger.warning( - f"Category {category_name}: invalid category_file path, skipping. {e}" - ) + verbose_proxy_logger.warning(f"Category {category_name}: invalid category_file path, skipping. {e}") continue else: # Try .yaml first, then .json (e.g. harm_toxic_abuse.json) @@ -486,9 +446,7 @@ class ContentFilterGuardrail(CustomGuardrail): category_file_path = yaml_path # will trigger "not found" below if not os.path.exists(category_file_path): - verbose_proxy_logger.warning( - f"Category file not found: {category_file_path}, skipping" - ) + verbose_proxy_logger.warning(f"Category file not found: {category_file_path}, skipping") continue try: @@ -496,14 +454,11 @@ class ContentFilterGuardrail(CustomGuardrail): self.loaded_categories[category_name] = category_config_obj # Use action from config, or default from category file - category_action = ContentFilterAction( - action if action else category_config_obj.default_action - ) + category_action = ContentFilterAction(action if action else category_config_obj.default_action) # Handle conditional categories (with identifier_words + block words) if category_config_obj.identifier_words and ( - category_config_obj.inherit_from - or category_config_obj.additional_block_words + category_config_obj.inherit_from or category_config_obj.additional_block_words ): self._load_conditional_category( category_name, @@ -545,9 +500,7 @@ class ContentFilterGuardrail(CustomGuardrail): f"conditional: {bool(category_config_obj.identifier_words)}" ) except Exception as e: - verbose_proxy_logger.error( - f"Error loading category {category_name}: {e}" - ) + verbose_proxy_logger.error(f"Error loading category {category_name}: {e}") def _load_conditional_category( self, @@ -590,9 +543,7 @@ class ContentFilterGuardrail(CustomGuardrail): verbose_proxy_logger.warning( f"Category {category_name}: inherit_from '{inherit_from}' file not found at {categories_dir}" ) - verbose_proxy_logger.debug( - f"Tried paths: {inherit_yaml_path}, {inherit_json_path}" - ) + verbose_proxy_logger.debug(f"Tried paths: {inherit_yaml_path}, {inherit_json_path}") if inherit_file_path: # Load the inherited category @@ -627,9 +578,7 @@ class ContentFilterGuardrail(CustomGuardrail): f"{len(block_words)} block words" ) if inherit_from and category_config_obj.additional_block_words: - inherited_count = len(block_words) - len( - category_config_obj.additional_block_words - ) + inherited_count = len(block_words) - len(category_config_obj.additional_block_words) log_msg += ( f" ({len(category_config_obj.additional_block_words)} additional + " f"{inherited_count} from {inherit_from})" @@ -641,9 +590,7 @@ class ContentFilterGuardrail(CustomGuardrail): verbose_proxy_logger.info(log_msg) except Exception as e: - verbose_proxy_logger.error( - f"Error loading conditional category for {category_name}: {e}" - ) + verbose_proxy_logger.error(f"Error loading conditional category for {category_name}: {e}") def _load_category_file(self, file_path: str) -> CategoryConfig: """ @@ -703,9 +650,7 @@ class ContentFilterGuardrail(CustomGuardrail): continue match_str = item.get("match") or "" raw_severity = item.get("severity", 2) - severity = severity_map.get( - raw_severity if isinstance(raw_severity, int) else 2, "medium" - ) + severity = severity_map.get(raw_severity if isinstance(raw_severity, int) else 2, "medium") for phrase in match_str.split("|"): phrase = phrase.strip().lower() if not phrase or phrase in seen: @@ -759,9 +704,7 @@ class ContentFilterGuardrail(CustomGuardrail): keyword_regex: Optional[Pattern] = None if extra_config.get("keyword_pattern"): - keyword_regex = re.compile( - extra_config["keyword_pattern"], re.IGNORECASE - ) + keyword_regex = re.compile(extra_config["keyword_pattern"], re.IGNORECASE) self.compiled_patterns.append( { @@ -772,9 +715,7 @@ class ContentFilterGuardrail(CustomGuardrail): "allow_word_numbers": bool(extra_config.get("allow_word_numbers")), } ) - verbose_proxy_logger.debug( - f"Added pattern: {pattern_name} with action {pattern_config.action}" - ) + verbose_proxy_logger.debug(f"Added pattern: {pattern_name} with action {pattern_config.action}") except Exception as e: verbose_proxy_logger.error(f"Error adding pattern {pattern_config}: {e}") raise @@ -799,19 +740,11 @@ class ContentFilterGuardrail(CustomGuardrail): data = yaml.safe_load(f) if not isinstance(data, dict) or "blocked_words" not in data: - raise ValueError( - "Invalid format: file must contain 'blocked_words' key with list of words" - ) + raise ValueError("Invalid format: file must contain 'blocked_words' key with list of words") for word_data in data["blocked_words"]: - if ( - not isinstance(word_data, dict) - or "keyword" not in word_data - or "action" not in word_data - ): - verbose_proxy_logger.warning( - f"Skipping invalid word entry: {word_data}" - ) + if not isinstance(word_data, dict) or "keyword" not in word_data or "action" not in word_data: + verbose_proxy_logger.warning(f"Skipping invalid word entry: {word_data}") continue keyword = word_data["keyword"].lower() @@ -820,17 +753,13 @@ class ContentFilterGuardrail(CustomGuardrail): self.blocked_words[keyword] = (action, description) - verbose_proxy_logger.info( - f"Loaded {len(data['blocked_words'])} blocked words from {file_path}" - ) + verbose_proxy_logger.info(f"Loaded {len(data['blocked_words'])} blocked words from {file_path}") except FileNotFoundError: raise FileNotFoundError(f"Blocked words file not found: {file_path}") except Exception as e: raise Exception(f"Error loading blocked words file {file_path}: {str(e)}") - def _find_pattern_spans( - self, text: str, pattern_entry: Dict[str, Any] - ) -> List[Tuple[int, int]]: + def _find_pattern_spans(self, text: str, pattern_entry: Dict[str, Any]) -> List[Tuple[int, int]]: """Return all match spans for a pattern, applying contextual rules if required.""" regex: Pattern = pattern_entry["regex"] @@ -919,9 +848,7 @@ class ContentFilterGuardrail(CustomGuardrail): merged.append((start, end)) return merged - def _mask_spans( - self, text: str, spans: List[Tuple[int, int]], redaction: str - ) -> str: + def _mask_spans(self, text: str, spans: List[Tuple[int, int]], redaction: str) -> str: """Apply masking for the provided spans using the given redaction tag.""" if not spans: @@ -952,9 +879,7 @@ class ContentFilterGuardrail(CustomGuardrail): return "".join(digits) if digits else None - def _check_patterns( - self, text: str - ) -> Optional[Tuple[str, str, ContentFilterAction]]: + def _check_patterns(self, text: str) -> Optional[Tuple[str, str, ContentFilterAction]]: """ Check text against all compiled regex patterns. @@ -971,9 +896,7 @@ class ContentFilterGuardrail(CustomGuardrail): matched_text = text[start:end] pattern_name = pattern_entry["pattern_name"] action = pattern_entry["action"] - verbose_proxy_logger.debug( - f"Pattern '{pattern_name}' matched: {matched_text[:20]}..." - ) + verbose_proxy_logger.debug(f"Pattern '{pattern_name}' matched: {matched_text[:20]}...") return (matched_text, pattern_name, action) return None @@ -1104,9 +1027,7 @@ class ContentFilterGuardrail(CustomGuardrail): for pattern_str, pattern in config.phrase_patterns: if pattern.search(text): - verbose_proxy_logger.warning( - f"Phrase pattern match in {category_name}: '{pattern_str}'" - ) + verbose_proxy_logger.warning(f"Phrase pattern match in {category_name}: '{pattern_str}'") return ( f"phrase: {pattern_str}", category_name, @@ -1134,9 +1055,7 @@ class ContentFilterGuardrail(CustomGuardrail): # Check exceptions first — they take precedence over always-block keywords too. for exception in exceptions: if exception in text_lower: - verbose_proxy_logger.debug( - f"Exception phrase '{exception}' found, skipping category keyword check" - ) + verbose_proxy_logger.debug(f"Exception phrase '{exception}' found, skipping category keyword check") return None # Always-block keywords are checked after exceptions. @@ -1152,9 +1071,7 @@ class ContentFilterGuardrail(CustomGuardrail): keyword_pattern = r"\b" + keyword_pattern_str + r"\b" keyword_found = bool(re.search(keyword_pattern, text_lower)) if keyword_found: - verbose_proxy_logger.debug( - f"Always-block keyword '{keyword}' found in category '{category}'" - ) + verbose_proxy_logger.debug(f"Always-block keyword '{keyword}' found in category '{category}'") return (keyword, category, severity, action) # Check category keywords @@ -1198,9 +1115,7 @@ class ContentFilterGuardrail(CustomGuardrail): return (keyword, category, severity, action) return None - def _check_blocked_words( - self, text: str - ) -> Optional[Tuple[str, ContentFilterAction, Optional[str]]]: + def _check_blocked_words(self, text: str) -> Optional[Tuple[str, ContentFilterAction, Optional[str]]]: """ Check text for blocked keywords. @@ -1232,9 +1147,7 @@ class ContentFilterGuardrail(CustomGuardrail): text_lower = text.lower() for keyword, (action, description) in self.blocked_words.items(): if keyword in text_lower: - verbose_proxy_logger.debug( - f"Blocked word '{keyword}' found with action {action}" - ) + verbose_proxy_logger.debug(f"Blocked word '{keyword}' found with action {action}") return (keyword, action, description) return None @@ -1259,8 +1172,7 @@ class ContentFilterGuardrail(CustomGuardrail): if action == ContentFilterAction.BLOCK: error_msg = ( - f"Content blocked: {category_name} conditional match '{matched_phrase}' detected " - f"(severity: {severity})" + f"Content blocked: {category_name} conditional match '{matched_phrase}' detected (severity: {severity})" ) verbose_proxy_logger.warning(error_msg) raise HTTPException( @@ -1298,10 +1210,7 @@ class ContentFilterGuardrail(CustomGuardrail): detections.append(category_detection) if action == ContentFilterAction.BLOCK: - error_msg = ( - f"Content blocked: {category_name} category keyword '{keyword}' detected " - f"(severity: {severity})" - ) + error_msg = f"Content blocked: {category_name} category keyword '{keyword}' detected (severity: {severity})" verbose_proxy_logger.warning(error_msg) raise HTTPException( status_code=400, @@ -1351,9 +1260,7 @@ class ContentFilterGuardrail(CustomGuardrail): detail={"error": error_msg, "pattern": pattern_name}, ) elif action == ContentFilterAction.MASK: - redaction_tag = self.pattern_redaction_format.format( - pattern_name=pattern_name.upper() - ) + redaction_tag = self.pattern_redaction_format.format(pattern_name=pattern_name.upper()) text = self._mask_spans(text, spans, redaction_tag) verbose_proxy_logger.info(f"Masked all {pattern_name} matches in content") @@ -1368,9 +1275,7 @@ class ContentFilterGuardrail(CustomGuardrail): detections: Optional[List[ContentFilterDetection]], ) -> str: """Handle blocked word match detection and action.""" - verbose_proxy_logger.debug( - f"Blocked word '{keyword}' found with action {action}" - ) + verbose_proxy_logger.debug(f"Blocked word '{keyword}' found with action {action}") if detections is not None: blocked_word_detection: BlockedWordDetection = { @@ -1406,9 +1311,7 @@ class ContentFilterGuardrail(CustomGuardrail): return text - def _filter_single_text( - self, text: str, detections: Optional[List[ContentFilterDetection]] = None - ) -> str: + def _filter_single_text(self, text: str, detections: Optional[List[ContentFilterDetection]] = None) -> str: """ Apply all content filtering checks to a single text. @@ -1436,25 +1339,19 @@ class ContentFilterGuardrail(CustomGuardrail): conditional_match = self._check_conditional_categories(text, all_exceptions) if conditional_match: matched_phrase, category_name, severity, action = conditional_match - self._handle_conditional_match( - matched_phrase, category_name, severity, action, detections - ) + self._handle_conditional_match(matched_phrase, category_name, severity, action, detections) # Check phrase patterns (regex-based paraphrase detection) phrase_match = self._check_phrase_patterns(text, all_exceptions) if phrase_match: matched_phrase, category_name, severity, action = phrase_match - self._handle_conditional_match( - matched_phrase, category_name, severity, action, detections - ) + self._handle_conditional_match(matched_phrase, category_name, severity, action, detections) # Check category keywords category_keyword_match = self._check_category_keywords(text, all_exceptions) if category_keyword_match: keyword, category_name, severity, action = category_keyword_match - text = self._handle_category_keyword_match( - keyword, category_name, severity, action, text, detections - ) + text = self._handle_category_keyword_match(keyword, category_name, severity, action, text, detections) # Check regex patterns - process ALL patterns, not just first match for pattern_entry in self.compiled_patterns: @@ -1462,18 +1359,14 @@ class ContentFilterGuardrail(CustomGuardrail): if spans: pattern_name = pattern_entry["pattern_name"] action = pattern_entry["action"] - text = self._handle_pattern_match( - pattern_name, action, text, spans, detections - ) + text = self._handle_pattern_match(pattern_name, action, text, spans, detections) # Check blocked words - iterate through ALL blocked words text_lower = text.lower() for keyword, (action, description) in self.blocked_words.items(): keyword_pattern_str = self._keyword_to_regex_pattern(keyword) if re.search(keyword_pattern_str, text_lower): - text = self._handle_blocked_word_match( - keyword, action, description, text, detections - ) + text = self._handle_blocked_word_match(keyword, action, description, text, detections) text_lower = text.lower() # Update after masking return text @@ -1500,14 +1393,10 @@ class ContentFilterGuardrail(CustomGuardrail): Returns: Text with sensitive content masked """ - redaction_tag = self.pattern_redaction_format.format( - pattern_name=pattern_name.upper() - ) + redaction_tag = self.pattern_redaction_format.format(pattern_name=pattern_name.upper()) return redaction_tag - async def _process_images( - self, images: List[str], detections: List[ContentFilterDetection] - ) -> None: + async def _process_images(self, images: List[str], detections: List[ContentFilterDetection]) -> None: """ Process images by describing them and applying content filtering. @@ -1562,15 +1451,11 @@ class ContentFilterGuardrail(CustomGuardrail): # e.detail can be a string or dict if isinstance(e.detail, dict) and "error" in e.detail: detail_dict = cast(Dict[str, Any], e.detail) - detail_dict["error"] = ( - detail_dict["error"] + " (Image description): " + description - ) + detail_dict["error"] = detail_dict["error"] + " (Image description): " + description elif isinstance(e.detail, str): e.detail = e.detail + " (Image description): " + description else: - e.detail = ( - "Content blocked: Image description detected" + description - ) + e.detail = "Content blocked: Image description detected" + description raise e def _count_masked_entities( @@ -1593,24 +1478,16 @@ class ContentFilterGuardrail(CustomGuardrail): if detection_type == "pattern": pattern_detection = cast(PatternDetection, detection) pattern_name = pattern_detection["pattern_name"] - masked_entity_count[pattern_name] = ( - masked_entity_count.get(pattern_name, 0) + 1 - ) + masked_entity_count[pattern_name] = masked_entity_count.get(pattern_name, 0) + 1 elif detection_type == "blocked_word": entity_type = "blocked_word" - masked_entity_count[entity_type] = ( - masked_entity_count.get(entity_type, 0) + 1 - ) + masked_entity_count[entity_type] = masked_entity_count.get(entity_type, 0) + 1 elif detection_type == "category_keyword": category_detection = cast(CategoryKeywordDetection, detection) category = category_detection["category"] - masked_entity_count[category] = ( - masked_entity_count.get(category, 0) + 1 - ) + masked_entity_count[category] = masked_entity_count.get(category, 0) + 1 - def _build_match_details( - self, detections: List[ContentFilterDetection] - ) -> List[dict]: + def _build_match_details(self, detections: List[ContentFilterDetection]) -> List[dict]: """Build match_details list from content filter detections.""" match_details: List[dict] = [] for detection in detections: @@ -1618,14 +1495,10 @@ class ContentFilterGuardrail(CustomGuardrail): detail: dict = {"type": detection["type"], "action_taken": action_taken} if detection["type"] == "pattern": detail["detection_method"] = "regex" - detail["snippet"] = cast(PatternDetection, detection).get( - "pattern_name", "" - ) + detail["snippet"] = cast(PatternDetection, detection).get("pattern_name", "") elif detection["type"] == "blocked_word": detail["detection_method"] = "keyword" - detail["snippet"] = cast(BlockedWordDetection, detection).get( - "keyword", "" - ) + detail["snippet"] = cast(BlockedWordDetection, detection).get("keyword", "") elif detection["type"] == "category_keyword": detail["detection_method"] = "keyword" cat_det = cast(CategoryKeywordDetection, detection) @@ -1663,10 +1536,7 @@ class ContentFilterGuardrail(CustomGuardrail): """Get comma-separated policy template names from loaded categories.""" if not self.loaded_categories: return None - names = [ - cat.description or cat.category_name - for cat in self.loaded_categories.values() - ] + names = [cat.description or cat.category_name for cat in self.loaded_categories.values()] return ", ".join(names) if names else None def _compute_risk_score( @@ -1735,9 +1605,7 @@ class ContentFilterGuardrail(CustomGuardrail): self._competitor_intent_checker, "refuse_message_template", None ): msg = self._competitor_intent_checker.refuse_message_template or msg - verbose_proxy_logger.warning( - "ContentFilterGuardrail: competitor intent refuse - %s", intent_val - ) + verbose_proxy_logger.warning("ContentFilterGuardrail: competitor intent refuse - %s", intent_val) raise HTTPException( status_code=400, detail={ @@ -1755,9 +1623,7 @@ class ContentFilterGuardrail(CustomGuardrail): self._competitor_intent_checker, "reframe_message_template", None ): msg = self._competitor_intent_checker.reframe_message_template or msg - verbose_proxy_logger.info( - "ContentFilterGuardrail: competitor intent reframe - %s", intent_val - ) + verbose_proxy_logger.info("ContentFilterGuardrail: competitor intent reframe - %s", intent_val) self.raise_passthrough_exception( violation_message=msg, request_data=request_data, @@ -1791,31 +1657,18 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str: Exception string if guardrail failed """ # Convert TypedDict detections to regular dicts for JSON serialization - guardrail_json_response: Union[Exception, str, dict, List[dict]] = [ - dict(detection) for detection in detections - ] + guardrail_json_response: Union[Exception, str, dict, List[dict]] = [dict(detection) for detection in detections] if status != "success": - guardrail_json_response = ( - exception_str - if exception_str - else [dict(detection) for detection in detections] - ) + guardrail_json_response = exception_str if exception_str else [dict(detection) for detection in detections] # Competitor intent: add confidence and classification to tracing if present tracing_kw: Dict[str, Any] = { "guardrail_id": self.config_guardrail_id or self.guardrail_name, - "policy_template": self.config_policy_template - or self._get_policy_templates(), - "detection_method": ( - self._get_detection_methods(detections) if detections else None - ), - "match_details": ( - self._build_match_details(detections) if detections else None - ), + "policy_template": self.config_policy_template or self._get_policy_templates(), + "detection_method": (self._get_detection_methods(detections) if detections else None), + "match_details": (self._build_match_details(detections) if detections else None), "patterns_checked": self._get_patterns_checked_count(), - "risk_score": self._compute_risk_score( - detections, masked_entity_count, status - ), + "risk_score": self._compute_risk_score(detections, masked_entity_count, status), } for d in detections: if isinstance(d, dict) and d.get("type") == "competitor_intent": @@ -1874,9 +1727,7 @@ class ContentFilterGuardrail(CustomGuardrail): await self._process_images(images, detections) # Process texts - verbose_proxy_logger.debug( - f"ContentFilterGuardrail: Applying guardrail to {len(texts)} text(s)" - ) + verbose_proxy_logger.debug(f"ContentFilterGuardrail: Applying guardrail to {len(texts)} text(s)") processed_texts = [] for text in texts: @@ -1884,15 +1735,11 @@ class ContentFilterGuardrail(CustomGuardrail): if self._competitor_intent_checker and text: intent_result = self._competitor_intent_checker.run(text) if intent_result.get("intent", "other") != "other": - self._apply_competitor_intent_policy( - intent_result, request_data, detections - ) + self._apply_competitor_intent_policy(intent_result, request_data, detections) filtered_text = self._filter_single_text(text, detections=detections) processed_texts.append(filtered_text) - verbose_proxy_logger.debug( - "ContentFilterGuardrail: Guardrail applied successfully" - ) + verbose_proxy_logger.debug("ContentFilterGuardrail: Guardrail applied successfully") inputs["texts"] = processed_texts # Count masked entities by type @@ -1965,8 +1812,7 @@ class ContentFilterGuardrail(CustomGuardrail): is_final = bool(getattr(choice, "finish_reason", None)) if isinstance(content, str) and content: accumulated_text_by_choice[choice_index] = ( - accumulated_text_by_choice.get(choice_index, "") - + content + accumulated_text_by_choice.get(choice_index, "") + content ) elif not is_final: continue @@ -1985,23 +1831,15 @@ class ContentFilterGuardrail(CustomGuardrail): # matches are guaranteed to be re-found. Keeping # only each choice's latest scan avoids duplicate # detections in the final log row. - masked_text = self._filter_single_text( - text_to_scan, detections=choice_detections - ) + masked_text = self._filter_single_text(text_to_scan, detections=choice_detections) if is_final and masked_text.endswith(" "): masked_text = masked_text[:-1] - latest_detections_by_choice[choice_index] = ( - choice_detections - ) + latest_detections_by_choice[choice_index] = choice_detections except HTTPException: - latest_detections_by_choice[choice_index] = ( - choice_detections - ) + latest_detections_by_choice[choice_index] = choice_detections raise except Exception as e: - verbose_proxy_logger.error( - f"ContentFilterGuardrail: Error in masking: {e}" - ) + verbose_proxy_logger.error(f"ContentFilterGuardrail: Error in masking: {e}") masked_text = text_to_scan # Fallback to current text # Determine how much can be safely yielded @@ -2010,17 +1848,11 @@ class ContentFilterGuardrail(CustomGuardrail): else: safe_to_yield_len = max(0, len(masked_text) - buffer_size) - yielded_masked_text_len = yielded_masked_text_len_by_choice.get( - choice_index, 0 - ) + yielded_masked_text_len = yielded_masked_text_len_by_choice.get(choice_index, 0) if safe_to_yield_len > yielded_masked_text_len: - new_masked_content = masked_text[ - yielded_masked_text_len:safe_to_yield_len - ] + new_masked_content = masked_text[yielded_masked_text_len:safe_to_yield_len] choice.delta.content = new_masked_content - yielded_masked_text_len_by_choice[choice_index] = ( - safe_to_yield_len - ) + yielded_masked_text_len_by_choice[choice_index] = safe_to_yield_len else: # Hold content by yielding empty content on this choice # while preserving chunk metadata and other choices. @@ -2033,8 +1865,7 @@ class ContentFilterGuardrail(CustomGuardrail): # Any remaining content (should have been handled by is_final, but just in case) if any( - yielded_masked_text_len_by_choice.get(choice_index, 0) - < len(accumulated_text) + yielded_masked_text_len_by_choice.get(choice_index, 0) < len(accumulated_text) for choice_index, accumulated_text in accumulated_text_by_choice.items() ): # We already reached the end of the generator diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py index aedc6acc810..3e20ada1cfa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py @@ -186,21 +186,15 @@ def _confusion_matrix(checker, cases: List[dict], label: str): tn += 1 elif expected == "BLOCK" and actual == "ALLOW": fn += 1 - wrong.append( - f" FN (score={score:.3f}): {case['sentence']!r:60s} — {case['test']}" - ) + wrong.append(f" FN (score={score:.3f}): {case['sentence']!r:60s} — {case['test']}") elif expected == "ALLOW" and actual == "BLOCK": fp += 1 - wrong.append( - f" FP (score={score:.3f}): {case['sentence']!r:60s} — {case['test']}" - ) + wrong.append(f" FP (score={score:.3f}): {case['sentence']!r:60s} — {case['test']}") total = tp + tn + fp + fn precision = tp / (tp + fp) if (tp + fp) > 0 else 0 recall = tp / (tp + fn) if (tp + fn) > 0 else 0 - f1 = ( - 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 - ) + f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 accuracy = (tp + tn) / total if total > 0 else 0 # Latency stats @@ -593,6 +587,4 @@ class TestInvestmentLlmJudgeClaude: return _load_jsonl("block_investment.jsonl") def test_confusion_matrix(self, blocker, cases): - _confusion_matrix( - blocker, cases, "Block Investment — LLM Judge (claude-haiku-4.5)" - ) + _confusion_matrix(blocker, cases, "Block Investment — LLM Judge (claude-haiku-4.5)") diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py index 27e554a1025..4879808d2eb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py @@ -36,8 +36,7 @@ for pattern_data in _PATTERNS_DATA["patterns"]: # Build lookup dictionaries from JSON PREBUILT_PATTERNS: Dict[str, str] = { - pattern_data["name"]: pattern_data["pattern"] - for pattern_data in _PATTERNS_DATA["patterns"] + pattern_data["name"]: pattern_data["pattern"] for pattern_data in _PATTERNS_DATA["patterns"] } @@ -53,11 +52,7 @@ KNOWN_PATTERN_KEYS = { PATTERN_EXTRA_CONFIG: Dict[str, Dict[str, Any]] = {} for pattern_data in _PATTERNS_DATA["patterns"]: - extra_config = { - key: value - for key, value in pattern_data.items() - if key not in KNOWN_PATTERN_KEYS - } + extra_config = {key: value for key, value in pattern_data.items() if key not in KNOWN_PATTERN_KEYS} PATTERN_EXTRA_CONFIG[pattern_data["name"]] = extra_config @@ -76,10 +71,7 @@ def get_compiled_pattern(pattern_name: str) -> Pattern: """ if pattern_name not in PREBUILT_PATTERNS: available_patterns = ", ".join(PREBUILT_PATTERNS.keys()) - raise ValueError( - f"Unknown pattern name: '{pattern_name}'. " - f"Available patterns: {available_patterns}" - ) + raise ValueError(f"Unknown pattern name: '{pattern_name}'. Available patterns: {available_patterns}") return re.compile(PREBUILT_PATTERNS[pattern_name], re.IGNORECASE) @@ -105,15 +97,13 @@ for pattern_data in _PATTERNS_DATA["patterns"]: # Build display names mapping from JSON PATTERN_DISPLAY_NAMES: Dict[str, str] = { - pattern_data["name"]: pattern_data["display_name"] - for pattern_data in _PATTERNS_DATA["patterns"] + pattern_data["name"]: pattern_data["display_name"] for pattern_data in _PATTERNS_DATA["patterns"] } # Build descriptions mapping from JSON PATTERN_DESCRIPTIONS: Dict[str, str] = { - pattern_data["name"]: pattern_data["description"] - for pattern_data in _PATTERNS_DATA["patterns"] + pattern_data["name"]: pattern_data["description"] for pattern_data in _PATTERNS_DATA["patterns"] } @@ -172,18 +162,14 @@ def get_available_content_categories() -> List[Dict[str, str]]: "name": category_data["category_name"], "display_name": display_name, "description": category_data.get("description", ""), - "default_action": category_data.get( - "default_action", "BLOCK" - ), + "default_action": category_data.get("default_action", "BLOCK"), } ) except Exception as e: # Skip files that can't be loaded but log the error for debugging from litellm._logging import verbose_proxy_logger - verbose_proxy_logger.warning( - f"Failed to load category file {filename}: {str(e)}" - ) + verbose_proxy_logger.warning(f"Failed to load category file {filename}: {str(e)}") continue elif filename.endswith(".json"): # JSON category files (e.g. harm_toxic_abuse.json) - no YAML header, use filename @@ -191,9 +177,7 @@ def get_available_content_categories() -> List[Dict[str, str]]: try: if category_name == "harm_toxic_abuse": display_name = "Harmful Toxic Abuse" - description = ( - "Detects harmful, toxic, or abusive language and content" - ) + description = "Detects harmful, toxic, or abusive language and content" else: display_name = category_name.replace("_", " ").title() description = f"Content category: {display_name}" diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index bc234c0a6f8..5445425a1d1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -68,8 +68,7 @@ def _build_judge_prompt( response_text: str, ) -> str: criteria_block = "\n".join( - f"- {c.get('name', '')} (weight {c.get('weight', 0)}%): {c.get('description', '')}" - for c in criteria + f"- {c.get('name', '')} (weight {c.get('weight', 0)}%): {c.get('description', '')}" for c in criteria ) conversation = "\n".join( f"{m.get('role', 'user').upper()}: {_extract_text_from_content(m.get('content', ''))}" @@ -93,27 +92,16 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): criteria: List[Dict[str, Any]], overall_threshold: float = 80.0, on_failure: Literal["block", "log"] = "block", - event_hook: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks]] - ] = None, + event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = None, default_on: bool = False, **kwargs: Any, ) -> None: - _event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = ( - None - ) + _event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = None if event_hook is not None: if isinstance(event_hook, list): - _event_hook = [ - GuardrailEventHooks(h) if isinstance(h, str) else h - for h in event_hook - ] + _event_hook = [GuardrailEventHooks(h) if isinstance(h, str) else h for h in event_hook] else: - _event_hook = ( - GuardrailEventHooks(event_hook) - if isinstance(event_hook, str) - else event_hook - ) + _event_hook = GuardrailEventHooks(event_hook) if isinstance(event_hook, str) else event_hook super().__init__( guardrail_name=guardrail_name, @@ -174,20 +162,14 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): try: judge_result = await self._run_judge(messages, response_text) except Exception as judge_err: - verbose_logger.warning( - f"llm_as_a_judge guardrail: judge call failed, failing open. Error: {judge_err}" - ) + verbose_logger.warning(f"llm_as_a_judge guardrail: judge call failed, failing open. Error: {judge_err}") status = "guardrail_failed_to_respond" return inputs try: - overall_score = max( - 0.0, min(100.0, float(judge_result.get("overall_score", 100))) - ) + overall_score = max(0.0, min(100.0, float(judge_result.get("overall_score", 100)))) except (TypeError, ValueError): - verbose_logger.warning( - "llm_as_a_judge: invalid overall_score from judge, failing open" - ) + verbose_logger.warning("llm_as_a_judge: invalid overall_score from judge, failing open") return inputs passed = overall_score >= self.overall_threshold @@ -251,9 +233,7 @@ def initialize_guardrail( judge_model = _get_litellm_param(litellm_params, guardrail, "judge_model") if not judge_model: - raise ValueError( - "llm_as_a_judge guardrail requires judge_model in litellm_params" - ) + raise ValueError("llm_as_a_judge guardrail requires judge_model in litellm_params") criteria = _get_litellm_param(litellm_params, guardrail, "criteria") or [] if not criteria: @@ -261,19 +241,13 @@ def initialize_guardrail( weight_total = sum(float(c.get("weight", 0)) for c in criteria) if abs(weight_total - 100) > 0.5: - raise ValueError( - f"llm_as_a_judge criterion weights must sum to 100 (got {weight_total})" - ) + raise ValueError(f"llm_as_a_judge criterion weights must sum to 100 (got {weight_total})") on_failure = _get_litellm_param(litellm_params, guardrail, "on_failure", "block") if on_failure not in _VALID_ON_FAILURE: - raise ValueError( - f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure}'" - ) + raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure}'") - overall_threshold = float( - _get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0) - ) + overall_threshold = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0)) mode = _get_litellm_param(litellm_params, guardrail, "mode") event_hook: Optional[GuardrailEventHooks] = None @@ -287,9 +261,7 @@ def initialize_guardrail( overall_threshold=overall_threshold, on_failure=on_failure, event_hook=event_hook, - default_on=bool( - _get_litellm_param(litellm_params, guardrail, "default_on", False) - ), + default_on=bool(_get_litellm_param(litellm_params, guardrail, "default_on", False)), ) litellm.logging_callback_manager.add_litellm_callback(instance) return instance diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py index 5060fade8cd..237364f9714 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py @@ -14,9 +14,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" # Default to always-on. Only disable if the user explicitly sets default_on: false. # We check the raw guardrail dict because LitellmParams normalizes None → False, # making it impossible to distinguish "not set" from "explicitly false" via litellm_params. - _raw_default_on = ( - cast(Dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") - ) + _raw_default_on = cast(Dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") _default_on = False if _raw_default_on is False else True _callback = MCPEndUserPermissionGuardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py index a485ff26463..524d087cfb7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py @@ -79,26 +79,18 @@ class MCPEndUserPermissionGuardrail(CustomGuardrail): if not tools: return inputs - allowed_mcp_servers = ( - await self._get_allowed_mcp_servers_from_object_permission( - object_permission - ) - ) + allowed_mcp_servers = await self._get_allowed_mcp_servers_from_object_permission(object_permission) if allowed_mcp_servers is None: return inputs # No restrictions → pass through unchanged - verbose_proxy_logger.debug( - f"MCP guardrail: end user restricted to MCP servers: {allowed_mcp_servers}" - ) + verbose_proxy_logger.debug(f"MCP guardrail: end user restricted to MCP servers: {allowed_mcp_servers}") filtered_tools = [] removed_tools = [] for tool in tools: tool_name = self._get_tool_name_from_definition(tool) - server_name = ( - self._extract_mcp_server_name(tool_name) if tool_name else None - ) + server_name = self._extract_mcp_server_name(tool_name) if tool_name else None if server_name is None: # Not an MCP tool (no prefix) or unrecognised format → keep @@ -134,24 +126,18 @@ class MCPEndUserPermissionGuardrail(CustomGuardrail): Uses get_end_user_object (same path as auth) so no extra DB round-trip when the cache is warm. """ - end_user_id = MCPEndUserPermissionGuardrail._get_end_user_id_from_request_data( - request_data - ) + end_user_id = MCPEndUserPermissionGuardrail._get_end_user_id_from_request_data(request_data) if not end_user_id: return None - end_user_object = await MCPEndUserPermissionGuardrail._fetch_end_user_object( - end_user_id - ) - return ( - end_user_object.object_permission if end_user_object is not None else None - ) + end_user_object = await MCPEndUserPermissionGuardrail._fetch_end_user_object(end_user_id) + return end_user_object.object_permission if end_user_object is not None else None @staticmethod def _get_end_user_id_from_request_data(request_data: dict) -> Optional[str]: - return request_data.get("user_api_key_end_user_id") or request_data.get( - "litellm_metadata", {} - ).get("user_api_key_end_user_id") + return request_data.get("user_api_key_end_user_id") or request_data.get("litellm_metadata", {}).get( + "user_api_key_end_user_id" + ) @staticmethod async def _fetch_end_user_object(end_user_id: str): # type: ignore[return] @@ -179,9 +165,7 @@ class MCPEndUserPermissionGuardrail(CustomGuardrail): route="/mcp", ) except Exception as e: - verbose_proxy_logger.warning( - f"MCP guardrail: failed to fetch end_user_object for '{end_user_id}': {e}" - ) + verbose_proxy_logger.warning(f"MCP guardrail: failed to fetch end_user_object for '{end_user_id}': {e}") return None # ------------------------------------------------------------------ @@ -210,11 +194,7 @@ class MCPEndUserPermissionGuardrail(CustomGuardrail): MCPRequestHandler, ) - access_group_servers = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - mcp_access_groups - ) - ) + access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups) return list(set(direct_mcp_servers + access_group_servers)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py index abea9014a11..434a136dc91 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py @@ -10,9 +10,7 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail( - litellm_params: "LitellmParams", guardrail: "Guardrail" -) -> MCPJWTSigner: +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> MCPJWTSigner: import litellm guardrail_name = guardrail.get("guardrail_name") diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 0f299f4c5f7..46b1afb5db7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -108,9 +108,7 @@ def _load_private_key_from_env(env_var: str) -> RSAPrivateKey: """Load an RSA private key from an env var (PEM string or file:// path).""" key_material = os.environ.get(env_var, "") if not key_material: - raise ValueError( - f"MCPJWTSigner: environment variable '{env_var}' is set but empty." - ) + raise ValueError(f"MCPJWTSigner: environment variable '{env_var}' is set but empty.") if key_material.startswith("file://"): path = key_material[len("file://") :] with open(path, "rb") as f: @@ -131,11 +129,7 @@ def _generate_rsa_key_pair() -> RSAPrivateKey: def _int_to_base64url(n: int) -> str: """Encode an integer as a base64url string (no padding).""" byte_length = (n.bit_length() + 7) // 8 - return ( - base64.urlsafe_b64encode(n.to_bytes(byte_length, byteorder="big")) - .rstrip(b"=") - .decode("ascii") - ) + return base64.urlsafe_b64encode(n.to_bytes(byte_length, byteorder="big")).rstrip(b"=").decode("ascii") def _compute_kid(public_key: Any) -> str: @@ -253,9 +247,7 @@ class MCPJWTSigner(CustomGuardrail): if key_material: self._private_key = _load_private_key_from_env(self.SIGNING_KEY_ENV) self._persistent_key: bool = True - verbose_proxy_logger.info( - "MCPJWTSigner: loaded RSA key from env var %s", self.SIGNING_KEY_ENV - ) + verbose_proxy_logger.info("MCPJWTSigner: loaded RSA key from env var %s", self.SIGNING_KEY_ENV) else: self._private_key = _generate_rsa_key_pair() self._persistent_key = False @@ -269,23 +261,14 @@ class MCPJWTSigner(CustomGuardrail): # --- Core config --- self.issuer: str = ( - issuer - or os.environ.get("MCP_JWT_ISSUER") - or os.environ.get("LITELLM_EXTERNAL_URL") - or "litellm" - ) - self.audience: str = ( - audience or os.environ.get("MCP_JWT_AUDIENCE") or self.DEFAULT_AUDIENCE + issuer or os.environ.get("MCP_JWT_ISSUER") or os.environ.get("LITELLM_EXTERNAL_URL") or "litellm" ) + self.audience: str = audience or os.environ.get("MCP_JWT_AUDIENCE") or self.DEFAULT_AUDIENCE resolved_ttl = int( - ttl_seconds - if ttl_seconds is not None - else os.environ.get("MCP_JWT_TTL_SECONDS", str(self.DEFAULT_TTL)) + ttl_seconds if ttl_seconds is not None else os.environ.get("MCP_JWT_TTL_SECONDS", str(self.DEFAULT_TTL)) ) if resolved_ttl <= 0: - raise ValueError( - f"MCPJWTSigner: ttl_seconds must be > 0, got {resolved_ttl}" - ) + raise ValueError(f"MCPJWTSigner: ttl_seconds must be > 0, got {resolved_ttl}") self.ttl_seconds: int = resolved_ttl # --- FR-5: Verify + re-sign --- @@ -311,9 +294,7 @@ class MCPJWTSigner(CustomGuardrail): # --- FR-14: Two-token model --- self.channel_token_audience: Optional[str] = channel_token_audience - self.channel_token_ttl: int = ( - channel_token_ttl if channel_token_ttl is not None else self.ttl_seconds - ) + self.channel_token_ttl: int = channel_token_ttl if channel_token_ttl is not None else self.ttl_seconds # --- FR-15: Incoming claim validation --- self.required_claims: List[str] = required_claims or [] @@ -336,8 +317,7 @@ class MCPJWTSigner(CustomGuardrail): _mcp_jwt_signer_instance = self verbose_proxy_logger.info( - "MCPJWTSigner initialized: issuer=%s audience=%s ttl=%ds kid=%s " - "verify=%s channel_token=%s debug=%s", + "MCPJWTSigner initialized: issuer=%s audience=%s ttl=%ds kid=%s verify=%s channel_token=%s debug=%s", self.issuer, self.audience, self.ttl_seconds, @@ -395,12 +375,8 @@ class MCPJWTSigner(CustomGuardrail): malformed response doesn't permanently disable JWT verification. """ now = time.time() - cache_expired = ( - now - self._oidc_discovery_fetched_at - ) >= self._OIDC_DISCOVERY_TTL - if ( - self._oidc_discovery_doc is None or cache_expired - ) and self.access_token_discovery_uri: + cache_expired = (now - self._oidc_discovery_fetched_at) >= self._OIDC_DISCOVERY_TTL + if (self._oidc_discovery_doc is None or cache_expired) and self.access_token_discovery_uri: doc = await _fetch_oidc_discovery(self.access_token_discovery_uri) if "jwks_uri" in doc: self._oidc_discovery_doc = doc @@ -473,9 +449,7 @@ class MCPJWTSigner(CustomGuardrail): if self.verify_issuer: decode_kwargs["issuer"] = self.verify_issuer - payload: Dict[str, Any] = jwt.decode( - raw_token, signing_jwk.key, **decode_kwargs - ) + payload: Dict[str, Any] = jwt.decode(raw_token, signing_jwk.key, **decode_kwargs) return payload async def _introspect_opaque_token(self, token: str) -> Dict[str, Any]: @@ -585,18 +559,14 @@ class MCPJWTSigner(CustomGuardrail): value = str(tid) if tid else None else: - verbose_proxy_logger.warning( - "MCPJWTSigner: unknown end_user_claim_source %r — skipping", source - ) + verbose_proxy_logger.warning("MCPJWTSigner: unknown end_user_claim_source %r — skipping", source) continue if value: return value # Final fallback for service accounts with no user identity - token = getattr(user_api_key_dict, "token", None) or getattr( - user_api_key_dict, "api_key", None - ) + token = getattr(user_api_key_dict, "token", None) or getattr(user_api_key_dict, "api_key", None) if token: return "apikey:" + hashlib.sha256(str(token).encode()).hexdigest()[:16] return "litellm-proxy" @@ -626,9 +596,7 @@ class MCPJWTSigner(CustomGuardrail): if self.allowed_scopes is not None: return " ".join(self.allowed_scopes) - tool_name = ( - re.sub(r"[^a-zA-Z0-9_\-]", "_", raw_tool_name) if raw_tool_name else "" - ) + tool_name = re.sub(r"[^a-zA-Z0-9_\-]", "_", raw_tool_name) if raw_tool_name else "" if tool_name: scopes = ["mcp:tools/call", f"mcp:tools/{tool_name}:call"] elif call_type == "call_mcp_tool": @@ -825,18 +793,12 @@ class MCPJWTSigner(CustomGuardrail): "Proceeding without incoming token verification." ) except Exception as exc: - verbose_proxy_logger.error( - "MCPJWTSigner: incoming token verification failed: %s", exc - ) + verbose_proxy_logger.error("MCPJWTSigner: incoming token verification failed: %s", exc) from fastapi import HTTPException raise HTTPException( status_code=401, - detail={ - "error": ( - f"MCPJWTSigner: incoming token verification failed: {exc}" - ) - }, + detail={"error": (f"MCPJWTSigner: incoming token verification failed: {exc}")}, ) elif not raw_token and self.access_token_discovery_uri: verbose_proxy_logger.debug( @@ -856,9 +818,7 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ # Build outbound access token # ------------------------------------------------------------------ - claims = self._build_claims( - user_api_key_dict, hook_data, jwt_claims, call_type=call_type - ) + claims = self._build_claims(user_api_key_dict, hook_data, jwt_claims, call_type=call_type) signed_token = jwt.encode( claims, @@ -892,15 +852,12 @@ class MCPJWTSigner(CustomGuardrail): # FR-9: Debug header # ------------------------------------------------------------------ if self.debug_headers: - new_headers["x-litellm-mcp-debug"] = self._build_debug_header( - claims, self._kid - ) + new_headers["x-litellm-mcp-debug"] = self._build_debug_header(claims, self._kid) hook_data["extra_headers"] = new_headers verbose_proxy_logger.debug( - "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d " - "verified=%s channel=%s call_type=%s", + "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d verified=%s channel=%s call_type=%s", claims.get("sub"), claims.get("act", {}).get("sub"), hook_data.get("mcp_tool_name"), @@ -942,19 +899,13 @@ async def inject_mcp_jwt_headers_for_upstream( "incoming_bearer_token": incoming_bearer_token, "extra_headers": merged, } - call_type: CallTypesLiteral = ( - "list_mcp_tools" if for_list_tools else "call_mcp_tool" - ) + call_type: CallTypesLiteral = "list_mcp_tools" if for_list_tools else "call_mcp_tool" try: from litellm.proxy.proxy_server import ( # noqa: PLC0415 proxy_logging_obj as _proxy_logging, ) - shared_cache = ( - _proxy_logging.internal_usage_cache.dual_cache - if _proxy_logging is not None - else DualCache() - ) + shared_cache = _proxy_logging.internal_usage_cache.dual_cache if _proxy_logging is not None else DualCache() except Exception: shared_cache = DualCache() diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 794edf092dd..3aeed1a25bb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -40,12 +40,7 @@ class MCPSecurityGuardrail(CustomGuardrail): data: dict, call_type: str, ) -> Optional[Union[Exception, str, dict]]: - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is not True: return data unregistered = self._find_unregistered_mcp_servers(data) @@ -98,9 +93,7 @@ class MCPSecurityGuardrail(CustomGuardrail): if not tools or not isinstance(tools, list): return set() - requested_servers = MCPSecurityGuardrail._extract_mcp_server_names_from_tools( - tools - ) + requested_servers = MCPSecurityGuardrail._extract_mcp_server_names_from_tools(tools) if not requested_servers: return set() diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py index 9ead2a63b60..fded89aa9b8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/__init__.py @@ -16,9 +16,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" # client_secret can be passed via the standard api_key field or as # a dedicated client_secret parameter. - client_secret = litellm_params.api_key or getattr( - litellm_params, "client_secret", None - ) + client_secret = litellm_params.api_key or getattr(litellm_params, "client_secret", None) if not tenant_id: raise ValueError("Microsoft Purview: tenant_id is required") @@ -36,9 +34,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" tenant_id=str(tenant_id), client_id=str(client_id), client_secret=str(client_secret), - purview_app_name=str( - getattr(litellm_params, "purview_app_name", None) or "LiteLLM" - ), + purview_app_name=str(getattr(litellm_params, "purview_app_name", None) or "LiteLLM"), user_id_field=str(getattr(litellm_params, "user_id_field", None) or "user_id"), event_hook=litellm_params.mode, default_on=litellm_params.default_on, diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index a7ed1d40913..830056db0f4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -18,9 +18,7 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues GRAPH_API_BASE = "https://graph.microsoft.com/v1.0" -TOKEN_ENDPOINT_TEMPLATE = ( - "https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token" -) +TOKEN_ENDPOINT_TEMPLATE = "https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token" GRAPH_SCOPE = "https://graph.microsoft.com/.default" # Protection scope cache TTL in seconds (1 hour, per Microsoft recommendation). @@ -49,9 +47,7 @@ class PurviewGuardrailBase: # (typically CustomGuardrail). super().__init__(**kwargs) - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.tenant_id = tenant_id self.client_id = client_id self.client_secret = client_secret @@ -63,9 +59,7 @@ class PurviewGuardrailBase: # Protection scope cache: user_id -> (etag, scope_response, fetched_at) # Capped at 1000 entries (LRU eviction) to avoid unbounded growth. - self._scope_cache: OrderedDict[str, Tuple[str, Dict[str, Any], float]] = ( - OrderedDict() - ) + self._scope_cache: OrderedDict[str, Tuple[str, Dict[str, Any], float]] = OrderedDict() self._scope_cache_maxsize = 1000 # Use a threading.Lock (not asyncio.Lock) because this lock is acquired # from both the proxy's main asyncio event loop and from short-lived @@ -113,9 +107,7 @@ class PurviewGuardrailBase: # token was actually received, not when the request started. with self._cache_lock: self._token_cache = (access_token, time.time() + expires_in) - verbose_proxy_logger.debug( - "Purview: acquired new OAuth2 token (expires_in=%ds)", expires_in - ) + verbose_proxy_logger.debug("Purview: acquired new OAuth2 token (expires_in=%ds)", expires_in) return access_token # ------------------------------------------------------------------ @@ -142,9 +134,7 @@ class PurviewGuardrailBase: headers.update(extra_headers) verbose_proxy_logger.debug("Purview Graph POST %s", url) - response = await self.async_handler.post( - url=url, headers=headers, json=json_body - ) + response = await self.async_handler.post(url=url, headers=headers, json=json_body) response.raise_for_status() response_json: Dict[str, Any] = response.json() response_headers = dict(response.headers) @@ -155,9 +145,7 @@ class PurviewGuardrailBase: # Protection scopes # ------------------------------------------------------------------ - async def _compute_protection_scopes( - self, user_id: str - ) -> Tuple[str, Dict[str, Any]]: + async def _compute_protection_scopes(self, user_id: str) -> Tuple[str, Dict[str, Any]]: """Call protectionScopes/compute and cache with ETag. Returns: @@ -172,10 +160,7 @@ class PurviewGuardrailBase: self._scope_cache.move_to_end(user_id) return cached[0], cached[1] - url = ( - f"{GRAPH_API_BASE}/users/{encoded_user_id}" - "/dataSecurityAndGovernance/protectionScopes/compute" - ) + url = f"{GRAPH_API_BASE}/users/{encoded_user_id}/dataSecurityAndGovernance/protectionScopes/compute" body: Dict[str, Any] = { "activities": "uploadText,downloadText", "locations": [ @@ -225,10 +210,7 @@ class PurviewGuardrailBase: correlation_id: Optional conversation/thread ID. """ encoded_user_id = self._encode_graph_user_id(user_id) - url = ( - f"{GRAPH_API_BASE}/users/{encoded_user_id}" - "/dataSecurityAndGovernance/processContent" - ) + url = f"{GRAPH_API_BASE}/users/{encoded_user_id}/dataSecurityAndGovernance/processContent" body: Dict[str, Any] = { "contentToProcess": { "contentEntries": [ @@ -279,9 +261,7 @@ class PurviewGuardrailBase: # User ID resolution # ------------------------------------------------------------------ - def _resolve_user_id( - self, data: Dict[str, Any], user_api_key_dict: Any - ) -> Optional[str]: + def _resolve_user_id(self, data: Dict[str, Any], user_api_key_dict: Any) -> Optional[str]: """Resolve the Entra user object ID from request data or auth context. Returns the strongest available identity walking down four sources, in @@ -324,9 +304,7 @@ class PurviewGuardrailBase: md = litellm_params.get("metadata") return md if isinstance(md, dict) else {} - def _resolve_trusted_user_id( - self, data: Dict[str, Any], user_api_key_dict: Any - ) -> Optional[str]: + def _resolve_trusted_user_id(self, data: Dict[str, Any], user_api_key_dict: Any) -> Optional[str]: """Resolve user ID from API-key/JWT-bound identity for blocking DLP. Uses only ``UserAPIKeyAuth.user_id`` (bound on the LiteLLM key or JWT). @@ -347,9 +325,7 @@ class PurviewGuardrailBase: return None - def _resolve_user_id_from_logging_kwargs( - self, kwargs: Dict[str, Any] - ) -> Optional[str]: + def _resolve_user_id_from_logging_kwargs(self, kwargs: Dict[str, Any]) -> Optional[str]: """Trusted-identity-only resolver for logging-only hooks. Uses only the proxy-injected ``user_api_key_user_id`` (populated from @@ -423,9 +399,7 @@ class PurviewGuardrailBase: joined = "\n".join(s.strip() for s in prompt if isinstance(s, str)) return joined.strip() or None if all(isinstance(x, int) for x in prompt): - verbose_proxy_logger.debug( - "Purview DLP: completions prompt is token ids only; skipping text scan" - ) + verbose_proxy_logger.debug("Purview DLP: completions prompt is token ids only; skipping text scan") return None str_parts = [x for x in prompt if isinstance(x, str)] if str_parts: @@ -445,33 +419,19 @@ class PurviewGuardrailBase: args: List[str] = [] # tool_calls: [{"function": {"arguments": "..."}}] - tool_calls = ( - message.get("tool_calls") - if isinstance(message, dict) - else getattr(message, "tool_calls", None) - ) + tool_calls = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) if tool_calls: for tc in tool_calls: - fn = ( - tc.get("function") - if isinstance(tc, dict) - else getattr(tc, "function", None) - ) + fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) if fn is None: continue - arguments = ( - fn.get("arguments") - if isinstance(fn, dict) - else getattr(fn, "arguments", None) - ) + arguments = fn.get("arguments") if isinstance(fn, dict) else getattr(fn, "arguments", None) if isinstance(arguments, str) and arguments.strip(): args.append(arguments) # Legacy function_call: {"arguments": "..."} function_call = ( - message.get("function_call") - if isinstance(message, dict) - else getattr(message, "function_call", None) + message.get("function_call") if isinstance(message, dict) else getattr(message, "function_call", None) ) if function_call is not None: arguments = ( @@ -484,9 +444,7 @@ class PurviewGuardrailBase: return args - def get_prompt_text_for_dlp( - self, messages: List["AllMessageValues"] - ) -> Optional[str]: + def get_prompt_text_for_dlp(self, messages: List["AllMessageValues"]) -> Optional[str]: """Concatenate text from every chat message (all roles) for pre-call DLP. Evaluates the same payload the model receives, not only the trailing user diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py index ee0bac64d4f..c6471cfcf51 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py @@ -155,9 +155,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): status = "guardrail_failed_to_respond" if block_on_violation: upstream_status = exc.response.status_code - client_status = ( - 502 if upstream_status in (401, 403) else upstream_status - ) + client_status = 502 if upstream_status in (401, 403) else upstream_status headers: Optional[Dict[str, str]] = None retry_after = exc.response.headers.get("retry-after") if retry_after: @@ -269,20 +267,14 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): msg = chat_choice.message if msg is None: continue - raw = ( - msg.get("content") - if isinstance(msg, dict) - else getattr(msg, "content", None) - ) + raw = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None) if isinstance(raw, str) and raw.strip(): parts.append(raw) # Include tool-call arguments returned by the model parts.extend(self._extract_tool_call_args_from_message(msg)) return parts - def _assemble_responses_api_from_chunks( - self, chunks: List[Any] - ) -> Tuple[bool, Optional[ResponsesAPIResponse]]: + def _assemble_responses_api_from_chunks(self, chunks: List[Any]) -> Tuple[bool, Optional[ResponsesAPIResponse]]: """Extract the final ``ResponsesAPIResponse`` from a buffered Responses API stream. Returns a ``(is_responses_api_stream, assembled)`` tuple so the caller @@ -304,9 +296,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): final = candidate return looks_like_responses_api, final - def _responses_api_input_to_str( - self, data: Dict[str, Any], raise_on_failure: bool = False - ) -> Optional[str]: + def _responses_api_input_to_str(self, data: Dict[str, Any], raise_on_failure: bool = False) -> Optional[str]: """Extract DLP-scannable text from a Responses API request ``input`` field. ``input`` may be a plain string or a list of input items (messages). In @@ -563,15 +553,12 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): status_code=400, detail={ "error": ( - "Microsoft Purview DLP: Unable to assemble streamed " - "response for scanning; blocking response." + "Microsoft Purview DLP: Unable to assemble streamed response for scanning; blocking response." ), }, ) - if isinstance( - assembled_response, (TextCompletionResponse, ResponsesAPIResponse) - ): + if isinstance(assembled_response, (TextCompletionResponse, ResponsesAPIResponse)): parts = self._completion_response_text_parts(assembled_response) if parts: combined = "\n\n---\n\n".join(parts) @@ -613,9 +600,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): # Logging-only hook — audit without blocking # ------------------------------------------------------------------ - def logging_hook( - self, kwargs: dict, result: Any, call_type: str - ) -> Tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: """Fire-and-forget async audit logging; returns original (kwargs, result) immediately. In the proxy's async success path, litellm independently calls both @@ -638,22 +623,16 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): # the deferral is observable if the framework ever stops dispatching # async_logging_hook on a given code path (otherwise audit silently # drops). - verbose_proxy_logger.debug( - "Purview audit: deferring to async_logging_hook (running event loop detected)" - ) + verbose_proxy_logger.debug("Purview audit: deferring to async_logging_hook (running event loop detected)") return kwargs, result except RuntimeError: pass async def _log_safe() -> None: try: - await self.async_logging_hook( - kwargs=kwargs, result=result, call_type=call_type - ) + await self.async_logging_hook(kwargs=kwargs, result=result, call_type=call_type) except Exception as exc: - verbose_proxy_logger.error( - "Purview audit background logging error: %s", exc - ) + verbose_proxy_logger.error("Purview audit background logging error: %s", exc) def _run_in_new_loop() -> None: new_loop = asyncio.new_event_loop() @@ -669,9 +648,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): return kwargs, result - async def async_logging_hook( - self, kwargs: dict, result: Any, call_type: str - ) -> Tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: """Send both prompt and response to Purview for audit logging. Errors are logged but never raised — this mode is non-blocking. @@ -701,9 +678,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): else: messages = kwargs.get("messages") if messages: - prompt_text = self.get_prompt_text_for_dlp( - cast(List[Any], messages) - ) + prompt_text = self.get_prompt_text_for_dlp(cast(List[Any], messages)) if prompt_text: await self._check_content( diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 7334e82549b..19b6fa77911 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -73,9 +73,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): VertexBase.__init__(self) # Then set our attributes (this ensures project_id is not overwritten) - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.template_id = template_id self.project_id = project_id self.location = location or "us-central1" @@ -98,18 +96,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return self.api_endpoint return f"https://modelarmor.{self.location}.rep.googleapis.com" - def _create_sanitize_request( - self, content: str, source: Literal["user_prompt", "model_response"] - ) -> dict: + def _create_sanitize_request(self, content: str, source: Literal["user_prompt", "model_response"]) -> dict: """Create request body for Model Armor API with correct camelCase field names.""" if source == "user_prompt": return {"userPromptData": {"text": content}} else: return {"modelResponseData": {"text": content}} - def _extract_content_from_response( - self, response: Union[Any, ModelResponse] - ) -> str: + def _extract_content_from_response(self, response: Union[Any, ModelResponse]) -> str: """ Extract text content from model response. @@ -125,9 +119,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # For non-ModelResponse types (e.g., TTS, images), return empty string # These response types are not text-based and shouldn't be processed by text guardrails - verbose_proxy_logger.debug( - "Model Armor: Skipping non-ModelResponse type: %s", type(response).__name__ - ) + verbose_proxy_logger.debug("Model Armor: Skipping non-ModelResponse type: %s", type(response).__name__) return "" async def make_model_armor_request( @@ -166,9 +158,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): elif content is not None: body = self._create_sanitize_request(content, source) else: - raise ValueError( - "Either content or file_bytes and file_type must be provided." - ) + raise ValueError("Either content or file_bytes and file_type must be provided.") # Set headers headers = { @@ -214,9 +204,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return await json_response return json_response - def sanitize_file_prompt( - self, file_bytes: bytes, file_type: str, source: str = "user_prompt" - ) -> dict: + def sanitize_file_prompt(self, file_bytes: bytes, file_type: str, source: str = "user_prompt") -> dict: """ Helper to build the request body for file prompt sanitization for Model Armor. file_type should be one of: PLAINTEXT_UTF8, PDF, WORD_DOCUMENT, EXCEL_DOCUMENT, POWERPOINT_DOCUMENT, TXT, CSV @@ -226,21 +214,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): base64_data = base64.b64encode(file_bytes).decode("utf-8") if source == "user_prompt": - return { - "userPromptData": { - "byteItem": {"byteDataType": file_type, "byteData": base64_data} - } - } + return {"userPromptData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}} else: - return { - "modelResponseData": { - "byteItem": {"byteDataType": file_type, "byteData": base64_data} - } - } + return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}} - def _should_block_content( - self, armor_response: dict, allow_sanitization: bool = False - ) -> bool: + def _should_block_content(self, armor_response: dict, allow_sanitization: bool = False) -> bool: """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult.""" sanitization_result = armor_response.get("sanitizationResult", {}) filter_results = sanitization_result.get("filterResults", {}) @@ -256,20 +234,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before if filt.get("raiFilterResult", {}).get("matchState") == "MATCH_FOUND": return True - if ( - filt.get("piAndJailbreakFilterResult", {}).get("matchState") - == "MATCH_FOUND" - ): + if filt.get("piAndJailbreakFilterResult", {}).get("matchState") == "MATCH_FOUND": return True - if ( - filt.get("maliciousUriFilterResult", {}).get("matchState") - == "MATCH_FOUND" - ): + if filt.get("maliciousUriFilterResult", {}).get("matchState") == "MATCH_FOUND": return True - if ( - filt.get("csamFilterFilterResult", {}).get("matchState") - == "MATCH_FOUND" - ): + if filt.get("csamFilterFilterResult", {}).get("matchState") == "MATCH_FOUND": return True if filt.get("virusScanFilterResult", {}).get("matchState") == "MATCH_FOUND": return True @@ -339,16 +308,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): This prevents circular references in logging. """ # Retrieve the Model Armor response & status stored on the per-request `metadata` object. - metadata = ( - request_data.get("metadata", {}) if isinstance(request_data, dict) else {} - ) + metadata = request_data.get("metadata", {}) if isinstance(request_data, dict) else {} guardrail_response = metadata.get("_model_armor_response", {}) # Determine status – default to "success" but prefer the explicit value if present. - guardrail_status: GuardrailStatus = metadata.get( - "_model_armor_status", "success" - ) # type: ignore + guardrail_status: GuardrailStatus = metadata.get("_model_armor_status", "success") # type: ignore self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, @@ -382,9 +347,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): messages = data.get("messages") if not messages: - verbose_proxy_logger.warning( - "Model Armor: not running guardrail. No messages in data" - ) + verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data") return data # Extract content from messages using helper from common_utils @@ -410,31 +373,23 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # This ensures each request logs its own Model Armor response instead of a potentially stale value # overwritten by another coroutine. if isinstance(data, dict): - metadata = data.setdefault( - "metadata", {} - ) # ensures metadata exists and is unique per request + metadata = data.setdefault("metadata", {}) # ensures metadata exists and is unique per request metadata["_model_armor_response"] = armor_response # Pre-compute guardrail status for downstream logging. A blocked response will eventually raise # an HTTPException, however in scenarios where the caller decides to ignore the exception (e.g. # fail_on_error=False) we still want the correct status reflected. metadata["_model_armor_status"] = ( "blocked" - if self._should_block_content( - armor_response, allow_sanitization=self.mask_request_content - ) + if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content) else "success" ) # Add guardrail to applied_guardrails BEFORE potential blocking # This ensures guardrail is recorded even when it blocks the request - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) # Check if content should be blocked - if self._should_block_content( - armor_response, allow_sanitization=self.mask_request_content - ): + if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content): raise HTTPException( status_code=400, detail={ @@ -452,16 +407,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): set_last_user_message, ) - data["messages"] = set_last_user_message( - messages, sanitized_content - ) + data["messages"] = set_last_user_message(messages, sanitized_content) except HTTPException: raise except Exception as e: - verbose_proxy_logger.error( - "Model Armor pre-call error: %s", str(e), exc_info=True - ) + verbose_proxy_logger.error("Model Armor pre-call error: %s", str(e), exc_info=True) # Depending on configuration, either fail or continue if self.optional_params.get("fail_on_error", True): raise @@ -488,9 +439,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): messages = data.get("messages") if not messages: - verbose_proxy_logger.warning( - "Model Armor: not running guardrail. No messages in data" - ) + verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data") return data # Extract content from messages @@ -516,22 +465,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): metadata["_model_armor_response"] = armor_response metadata["_model_armor_status"] = ( "blocked" - if self._should_block_content( - armor_response, allow_sanitization=self.mask_request_content - ) + if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content) else "success" ) # Add guardrail to applied_guardrails BEFORE potential blocking # This ensures guardrail is recorded even when it blocks the request - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) # Check if content should be blocked - if self._should_block_content( - armor_response, allow_sanitization=self.mask_request_content - ): + if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content): raise HTTPException( status_code=400, detail={ @@ -548,16 +491,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): set_last_user_message, ) - data["messages"] = set_last_user_message( - messages, sanitized_content - ) + data["messages"] = set_last_user_message(messages, sanitized_content) except HTTPException: raise except Exception as e: - verbose_proxy_logger.error( - "Model Armor moderation error: %s", str(e), exc_info=True - ) + verbose_proxy_logger.error("Model Armor moderation error: %s", str(e), exc_info=True) if self.optional_params.get("fail_on_error", True): raise @@ -576,20 +515,13 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): add_guardrail_to_applied_guardrails_header, ) - if ( - self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return # Extract content from response content = self._extract_content_from_response(response) if not content: - verbose_proxy_logger.debug( - "Model Armor: No text content to process in response, skipping guardrail" - ) + verbose_proxy_logger.debug("Model Armor: No text content to process in response, skipping guardrail") return # Make Model Armor request @@ -613,15 +545,13 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else "success" ), } - standard_logging_guardrail_information = ( - StandardLoggingGuardrailInformation( - guardrail_name=self.guardrail_name, - guardrail_provider="model_armor", - guardrail_mode=GuardrailEventHooks.post_call, - guardrail_response=model_armor_logged_object, - guardrail_status="success", - start_time=data.get("start_time"), - ) + standard_logging_guardrail_information = StandardLoggingGuardrailInformation( + guardrail_name=self.guardrail_name, + guardrail_provider="model_armor", + guardrail_mode=GuardrailEventHooks.post_call, + guardrail_response=model_armor_logged_object, + guardrail_status="success", + start_time=data.get("start_time"), ) add_guardrail_response_to_standard_logging_object( litellm_logging_obj=data.get("litellm_logging_obj"), @@ -630,14 +560,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Add guardrail to applied_guardrails BEFORE potential blocking # This ensures guardrail is recorded even when it blocks the request - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) # Check if content should be blocked - if self._should_block_content( - armor_response, allow_sanitization=self.mask_response_content - ): + if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content): raise HTTPException( status_code=400, detail={ @@ -660,9 +586,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error( - "Model Armor post-call error: %s", str(e), exc_info=True - ) + verbose_proxy_logger.error("Model Armor post-call error: %s", str(e), exc_info=True) if self.optional_params.get("fail_on_error", True): raise @@ -705,9 +629,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): metadata = request_data.setdefault("metadata", {}) metadata["_model_armor_response"] = armor_response metadata["_model_armor_status"] = ( - "blocked" - if self._should_block_content(armor_response) - else "success" + "blocked" if self._should_block_content(armor_response) else "success" ) # Add guardrail to applied_guardrails BEFORE potential blocking @@ -741,9 +663,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): choice.message.content = sanitized_content # Return sanitized stream - mock_response = MockResponseIterator( - model_response=assembled_response - ) + mock_response = MockResponseIterator(model_response=assembled_response) async for chunk in mock_response: yield chunk return @@ -752,11 +672,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Yield error as SSE event so create_response() detects it and # returns a proper JSON error response with the correct status code. # (Raising from a generator hits create_response's generic except → 500.) - detail = ( - e.detail - if isinstance(e.detail, dict) - else {"message": str(e.detail)} - ) + detail = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} error_value = detail.get("error", detail) if isinstance(error_value, dict): error_obj = dict(error_value) @@ -766,15 +682,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): yield f"data: {json.dumps({'error': error_obj})}\n\n" # type: ignore[misc] return except Exception as e: - verbose_proxy_logger.error( - "Model Armor streaming error: %s", str(e), exc_info=True - ) + verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True) if self.optional_params.get("fail_on_error", True): raise else: - verbose_proxy_logger.debug( - "Model Armor: No text content in streaming response, skipping guardrail" - ) + verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail") # Return original chunks if no sanitization needed for chunk in all_chunks: diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py index 1eea74a1e68..f0e6c6677a3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py @@ -14,9 +14,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" if isinstance(use_v2, str): use_v2 = use_v2.lower() == "true" if use_v2: - return initialize_guardrail_v2( - litellm_params=litellm_params, guardrail=guardrail - ) + return initialize_guardrail_v2(litellm_params=litellm_params, guardrail=guardrail) import litellm diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index aadfdf66531..7e8a22a66a8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -125,35 +125,25 @@ class NomaGuardrail(CustomGuardrail): ) _LEGACY_NOMA_DEPRECATION_WARNED = True - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self._responses_transform_handler = LiteLLMResponsesTransformationHandler() self.api_key = api_key or os.environ.get("NOMA_API_KEY") - self.api_base = api_base or os.environ.get( - "NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE - ) + self.api_base = api_base or os.environ.get("NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE) self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID") self.default_application_id = "litellm" if monitor_mode is None: - self.monitor_mode = ( - os.environ.get("NOMA_MONITOR_MODE", "false").lower() == "true" - ) + self.monitor_mode = os.environ.get("NOMA_MONITOR_MODE", "false").lower() == "true" else: self.monitor_mode = monitor_mode if block_failures is None: - self.block_failures = ( - os.environ.get("NOMA_BLOCK_FAILURES", "true").lower() == "true" - ) + self.block_failures = os.environ.get("NOMA_BLOCK_FAILURES", "true").lower() == "true" else: self.block_failures = block_failures if anonymize_input is None: - self.anonymize_input = ( - os.environ.get("NOMA_ANONYMIZE_INPUT", "false").lower() == "true" - ) + self.anonymize_input = os.environ.get("NOMA_ANONYMIZE_INPUT", "false").lower() == "true" else: self.anonymize_input = anonymize_input @@ -167,9 +157,7 @@ class NomaGuardrail(CustomGuardrail): try: asyncio.create_task(coro) except Exception as e: - verbose_proxy_logger.error( - f"Failed to create background Noma task: {str(e)}" - ) + verbose_proxy_logger.error(f"Failed to create background Noma task: {str(e)}") async def _process_user_message_check( self, @@ -185,10 +173,8 @@ class NomaGuardrail(CustomGuardrail): if not messages: return None - input_items, instructions = ( - self._responses_transform_handler.convert_chat_completion_messages_to_responses_api( # type: ignore[arg-type] - messages - ) + input_items, instructions = self._responses_transform_handler.convert_chat_completion_messages_to_responses_api( # type: ignore[arg-type] + messages ) if instructions: @@ -232,22 +218,16 @@ class NomaGuardrail(CustomGuardrail): ) if self.monitor_mode: - await self._handle_verdict_background( - USER_ROLE, json.dumps(input_items), response_json - ) + await self._handle_verdict_background(USER_ROLE, json.dumps(input_items), response_json) return json.dumps(input_items) # Check if we should anonymize content if self._should_anonymize(response_json, USER_ROLE): - anonymized_content = self._extract_anonymized_content( - response_json, USER_ROLE - ) + anonymized_content = self._extract_anonymized_content(response_json, USER_ROLE) if anonymized_content: # Replace the user message content with anonymized version self._replace_user_message_content(request_data, anonymized_content) - verbose_proxy_logger.debug( - f"Noma guardrail anonymized user message: {anonymized_content}" - ) + verbose_proxy_logger.debug(f"Noma guardrail anonymized user message: {anonymized_content}") return anonymized_content await self._check_verdict(USER_ROLE, json.dumps(input_items), response_json) @@ -314,22 +294,16 @@ class NomaGuardrail(CustomGuardrail): ) if self.monitor_mode: - await self._handle_verdict_background( - ASSISTANT_ROLE, json.dumps(content), response_json - ) + await self._handle_verdict_background(ASSISTANT_ROLE, json.dumps(content), response_json) return content # Check if we should anonymize content if self._should_anonymize(response_json, ASSISTANT_ROLE): - anonymized_content = self._extract_anonymized_content( - response_json, ASSISTANT_ROLE - ) + anonymized_content = self._extract_anonymized_content(response_json, ASSISTANT_ROLE) if anonymized_content: # Replace the LLM response content with anonymized version self._replace_llm_response_content(response, anonymized_content) - verbose_proxy_logger.debug( - f"Noma guardrail anonymized LLM response: {anonymized_content}" - ) + verbose_proxy_logger.debug(f"Noma guardrail anonymized LLM response: {anonymized_content}") return anonymized_content await self._check_verdict(ASSISTANT_ROLE, content, response_json) @@ -368,9 +342,7 @@ class NomaGuardrail(CustomGuardrail): return "guardrail_failed_to_respond" except Exception as e: - verbose_proxy_logger.error( - f"Error determining NOMA guardrail status: {str(e)}" - ) + verbose_proxy_logger.error(f"Error determining NOMA guardrail status: {str(e)}") return "guardrail_failed_to_respond" def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool: @@ -412,9 +384,7 @@ class NomaGuardrail(CustomGuardrail): # Return True only if sensitive data was detected AND no other detectors have result=true return sensitive_data_detected and len(failed_detectors) == 0 - def _extract_anonymized_content( - self, response_json: dict, message_type: MessageRole - ) -> Optional[str]: + def _extract_anonymized_content(self, response_json: dict, message_type: MessageRole) -> Optional[str]: """ Extract anonymized content from Noma API response. @@ -433,11 +403,7 @@ class NomaGuardrail(CustomGuardrail): # Find the scan result matching the message type (role) for result_item in scan_result: if result_item.get("role") == message_type: - return ( - result_item.get("results", {}) - .get("anonymizedContent", {}) - .get("anonymized", "") - ) + return result_item.get("results", {}).get("anonymizedContent", {}).get("anonymized", "") return None @@ -479,9 +445,7 @@ class NomaGuardrail(CustomGuardrail): for result_item in scan_result: if result_item.get("role") == message_type: - return self._should_only_sensitive_data_failed( - result_item.get("results", {}) - ) + return self._should_only_sensitive_data_failed(result_item.get("results", {})) return False @@ -500,9 +464,7 @@ class NomaGuardrail(CustomGuardrail): return result_obj.get("result") is True - def _replace_user_message_content( - self, request_data: dict, anonymized_content: str - ): + def _replace_user_message_content(self, request_data: dict, anonymized_content: str): """ Replace the user message content in request data with anonymized version. @@ -520,9 +482,7 @@ class NomaGuardrail(CustomGuardrail): messages[i]["content"] = anonymized_content break - def _replace_llm_response_content( - self, response: LLMResponse, anonymized_content: str - ): + def _replace_llm_response_content(self, response: LLMResponse, anonymized_content: str): """ Replace the LLM response content with anonymized version. @@ -547,9 +507,7 @@ class NomaGuardrail(CustomGuardrail): try: await self._process_user_message_check(request_data, user_auth) except Exception as e: - verbose_proxy_logger.error( - f"Noma background user message check failed: {str(e)}" - ) + verbose_proxy_logger.error(f"Noma background user message check failed: {str(e)}") async def _check_llm_response_background( self, @@ -561,9 +519,7 @@ class NomaGuardrail(CustomGuardrail): try: await self._process_llm_response_check(request_data, response, user_auth) except Exception as e: - verbose_proxy_logger.error( - f"Noma background response check failed: {str(e)}" - ) + verbose_proxy_logger.error(f"Noma background response check failed: {str(e)}") async def _handle_verdict_background( self, @@ -585,9 +541,7 @@ class NomaGuardrail(CustomGuardrail): msg = f"Noma guardrail allowed {type} message: {message}" verbose_proxy_logger.info(msg) except Exception as e: - verbose_proxy_logger.error( - f"Noma background verdict handling failed: {str(e)}" - ) + verbose_proxy_logger.error(f"Noma background verdict handling failed: {str(e)}") async def async_pre_call_hook( self, @@ -608,19 +562,13 @@ class NomaGuardrail(CustomGuardrail): # In monitor mode, run Noma check in background and return immediately if self.monitor_mode: try: - self._create_background_noma_check( - self._check_user_message_background(data, user_api_key_dict) - ) + self._create_background_noma_check(self._check_user_message_background(data, user_api_key_dict)) except Exception as e: - verbose_proxy_logger.error( - f"Failed to start background Noma pre-call check: {str(e)}" - ) + verbose_proxy_logger.error(f"Failed to start background Noma pre-call check: {str(e)}") return data try: - return await self._check_user_message( - data, user_api_key_dict, GuardrailEventHooks.pre_call - ) + return await self._check_user_message(data, user_api_key_dict, GuardrailEventHooks.pre_call) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -662,19 +610,13 @@ class NomaGuardrail(CustomGuardrail): # In monitor mode, run Noma check in background and return immediately if self.monitor_mode: try: - self._create_background_noma_check( - self._check_user_message_background(data, user_api_key_dict) - ) + self._create_background_noma_check(self._check_user_message_background(data, user_api_key_dict)) except Exception as e: - verbose_proxy_logger.error( - f"Failed to start background Noma moderation check: {str(e)}" - ) + verbose_proxy_logger.error(f"Failed to start background Noma moderation check: {str(e)}") return data try: - return await self._check_user_message( - data, user_api_key_dict, GuardrailEventHooks.during_call - ) + return await self._check_user_message(data, user_api_key_dict, GuardrailEventHooks.during_call) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -714,20 +656,14 @@ class NomaGuardrail(CustomGuardrail): if self.monitor_mode: try: self._create_background_noma_check( - self._check_llm_response_background( - data, response, user_api_key_dict - ) + self._check_llm_response_background(data, response, user_api_key_dict) ) except Exception as e: - verbose_proxy_logger.error( - f"Failed to start background Noma post-call check: {str(e)}" - ) + verbose_proxy_logger.error(f"Failed to start background Noma post-call check: {str(e)}") return response try: - return await self._check_llm_response( - data, response, user_api_key_dict, GuardrailEventHooks.post_call - ) + return await self._check_llm_response(data, response, user_api_key_dict, GuardrailEventHooks.post_call) except NomaBlockedMessage: # Blocked requests were already logged in _process_llm_response_check with "blocked" status raise @@ -759,9 +695,7 @@ class NomaGuardrail(CustomGuardrail): event_type: Optional[GuardrailEventHooks] = None, ) -> Union[Exception, str, dict, None]: """Check user message for policy violations""" - user_message = await self._process_user_message_check( - request_data, user_auth, event_type - ) + user_message = await self._process_user_message_check(request_data, user_auth, event_type) if not user_message: return request_data @@ -775,9 +709,7 @@ class NomaGuardrail(CustomGuardrail): event_type: Optional[GuardrailEventHooks] = None, ) -> Any: """Check LLM response for policy violations""" - content = await self._process_llm_response_check( - request_data, response, user_auth, event_type - ) + content = await self._process_llm_response_check(request_data, response, user_auth, event_type) if not content: return response @@ -796,9 +728,7 @@ class NomaGuardrail(CustomGuardrail): **({"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}), **({"X-Noma-Request-ID": call_id} if call_id else {}), } - endpoint = urljoin( - self.api_base or "https://api.noma.security/", NomaGuardrail._AIDR_ENDPOINT - ) + endpoint = urljoin(self.api_base or "https://api.noma.security/", NomaGuardrail._AIDR_ENDPOINT) response = await self.async_handler.post( endpoint, @@ -807,20 +737,12 @@ class NomaGuardrail(CustomGuardrail): **payload, "x-noma-context": { "applicationId": extra_data.get("application_id") - or request_data.get("metadata", {}) - .get("headers", {}) - .get("x-noma-application-id") + or request_data.get("metadata", {}).get("headers", {}).get("x-noma-application-id") or self.application_id or user_auth.key_alias or self.default_application_id, - "ipAddress": request_data.get("metadata", {}).get( - "requester_ip_address", None - ), - "userId": ( - user_auth.user_email - if user_auth.user_email - else user_auth.user_id - ), + "ipAddress": request_data.get("metadata", {}).get("requester_ip_address", None), + "userId": (user_auth.user_email if user_auth.user_email else user_auth.user_id), "sessionId": call_id, "requestId": llm_request_id, }, @@ -883,9 +805,9 @@ class NomaGuardrail(CustomGuardrail): if not all_chunks: return - assembled_model_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = stream_chunk_builder(chunks=all_chunks) + assembled_model_response: Optional[Union[ModelResponse, TextCompletionResponse]] = stream_chunk_builder( + chunks=all_chunks + ) if isinstance(assembled_model_response, ModelResponse): try: @@ -900,9 +822,7 @@ class NomaGuardrail(CustomGuardrail): except Exception as e: if self.block_failures: raise - verbose_proxy_logger.error( - f"Noma streaming post-call hook failed: {str(e)}" - ) + verbose_proxy_logger.error(f"Noma streaming post-call hook failed: {str(e)}") for chunk in all_chunks: yield chunk return diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 6aeaac949a9..1cf3dcd9ac4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -53,33 +53,23 @@ class NomaV2Guardrail(CustomGuardrail): block_failures: Optional[bool] = None, **kwargs: Any, ) -> None: - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_key = api_key or os.environ.get("NOMA_API_KEY") - self.api_base = ( - api_base or os.environ.get("NOMA_API_BASE") or _DEFAULT_API_BASE - ).rstrip("/") + self.api_base = (api_base or os.environ.get("NOMA_API_BASE") or _DEFAULT_API_BASE).rstrip("/") self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID") if monitor_mode is None: - self.monitor_mode = ( - os.environ.get("NOMA_MONITOR_MODE", "false").lower() == "true" - ) + self.monitor_mode = os.environ.get("NOMA_MONITOR_MODE", "false").lower() == "true" else: self.monitor_mode = monitor_mode if block_failures is None: - self.block_failures = ( - os.environ.get("NOMA_BLOCK_FAILURES", "true").lower() == "true" - ) + self.block_failures = os.environ.get("NOMA_BLOCK_FAILURES", "true").lower() == "true" else: self.block_failures = block_failures if self._requires_api_key(api_base=self.api_base) and not self.api_key: - raise ValueError( - "Noma v2 guardrail requires api_key when using Noma SaaS endpoint" - ) + raise ValueError("Noma v2 guardrail requires api_key when using Noma SaaS endpoint") if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ @@ -140,9 +130,7 @@ class NomaV2Guardrail(CustomGuardrail): ) -> dict: payload_request_data = self._sanitize_payload_for_transport(request_data) if logging_obj is not None: - payload_request_data["litellm_logging_obj"] = getattr( - logging_obj, "model_call_details", None - ) + payload_request_data["litellm_logging_obj"] = getattr(logging_obj, "model_call_details", None) payload: dict[str, Any] = { "inputs": inputs, @@ -278,9 +266,7 @@ class NomaV2Guardrail(CustomGuardrail): if application_id is None: application_id = self._get_non_empty_str( request_data.get("litellm_metadata", {}).get("user_api_key_alias") - ) or self._get_non_empty_str( - request_data.get("metadata", {}).get("user_api_key_alias") - ) + ) or self._get_non_empty_str(request_data.get("metadata", {}).get("user_api_key_alias")) try: payload = self._build_scan_payload( @@ -308,17 +294,13 @@ class NomaV2Guardrail(CustomGuardrail): action=action, ) - guardrail_status = ( - "success" if action == _Action.NONE else "guardrail_intervened" - ) + guardrail_status = "success" if action == _Action.NONE else "guardrail_intervened" return processed_inputs except NomaBlockedMessage as e: guardrail_status = "guardrail_intervened" guardrail_json_response = ( - response_json - if isinstance(response_json, dict) - else getattr(e, "detail", {"error": "blocked"}) + response_json if isinstance(response_json, dict) else getattr(e, "detail", {"error": "blocked"}) ) raise except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index a07f5371355..b411d0fb9eb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -77,9 +77,7 @@ class OnyxGuardrail(CustomGuardrail): detection_message = "Unknown violation" if "violated_rules" in result: detection_message = ", ".join(result["violated_rules"]) - verbose_proxy_logger.warning( - f"Request blocked by Onyx Guard. Violations: {detection_message}." - ) + verbose_proxy_logger.warning(f"Request blocked by Onyx Guard. Violations: {detection_message}.") raise HTTPException( status_code=400, detail=f"Request blocked by Onyx Guard. Violations: {detection_message}.", @@ -94,9 +92,7 @@ class OnyxGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - conversation_id = ( - logging_obj.litellm_call_id if logging_obj else str(uuid.uuid4()) - ) + conversation_id = logging_obj.litellm_call_id if logging_obj else str(uuid.uuid4()) verbose_proxy_logger.info( "Running Onyx Guard apply_guardrail hook", diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py index e1d9a7ce505..bf19200d8c8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py @@ -29,9 +29,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" "streaming_end_of_stream_only": _get_config_value( litellm_params, optional_params, "streaming_end_of_stream_only" ), - "streaming_sampling_rate": _get_config_value( - litellm_params, optional_params, "streaming_sampling_rate" - ), + "streaming_sampling_rate": _get_config_value(litellm_params, optional_params, "streaming_sampling_rate"), }, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index b3b8fbdb2a5..093ac693d5e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -58,9 +58,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): guardrail_name: str, api_key: Optional[str] = None, api_base: Optional[str] = None, - model: Optional[ - Literal["omni-moderation-latest", "text-moderation-latest"] - ] = None, + model: Optional[Literal["omni-moderation-latest", "text-moderation-latest"]] = None, streaming_end_of_stream_only: Optional[bool] = None, streaming_sampling_rate: Optional[int] = None, **kwargs, @@ -80,27 +78,19 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): **kwargs, ) - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) # Store configuration self.api_key = api_key or self._get_api_key() self.api_base = api_base or "https://api.openai.com/v1" - self.model: Literal["omni-moderation-latest", "text-moderation-latest"] = ( - model or "omni-moderation-latest" - ) + self.model: Literal["omni-moderation-latest", "text-moderation-latest"] = model or "omni-moderation-latest" # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook # via getattr(guardrail_to_apply, "streaming_*", default). self.streaming_end_of_stream_only: bool = ( - False - if streaming_end_of_stream_only is None - else streaming_end_of_stream_only - ) - self.streaming_sampling_rate: int = ( - 5 if streaming_sampling_rate is None else streaming_sampling_rate + False if streaming_end_of_stream_only is None else streaming_end_of_stream_only ) + self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate if not self.api_key: raise ValueError( @@ -142,9 +132,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): json=request_body, ) - verbose_proxy_logger.debug( - "OpenAI Moderation guard response: %s", response.json() - ) + verbose_proxy_logger.debug("OpenAI Moderation guard response: %s", response.json()) if response.status_code != 200: raise HTTPException( @@ -159,9 +147,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): return OpenAIModerationResponse(**response.json()) - def _check_moderation_result( - self, moderation_response: "OpenAIModerationResponse" - ) -> None: + def _check_moderation_result(self, moderation_response: "OpenAIModerationResponse") -> None: """ Check if the moderation response indicates harmful content and raise exception if needed. """ @@ -235,9 +221,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): text_to_moderate = "\n".join(texts) if not text_to_moderate: - verbose_proxy_logger.debug( - "OpenAI Moderation: No text content to moderate in inputs" - ) + verbose_proxy_logger.debug("OpenAI Moderation: No text content to moderate in inputs") return inputs # Make moderation request @@ -309,9 +293,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): instead of the stringified exception. """ guardrail_status: GuardrailStatus = ( - "guardrail_intervened" - if self._is_guardrail_intervention(e) - else "guardrail_failed_to_respond" + "guardrail_intervened" if self._is_guardrail_intervention(e) else "guardrail_failed_to_respond" ) if isinstance(request_data, dict): @@ -321,9 +303,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): metadata = {} # Use the stashed moderation response if available, fall back to exception - guardrail_response: Union[dict, Exception, str] = metadata.pop( - "_openai_moderation_response", e - ) + guardrail_response: Union[dict, Exception, str] = metadata.pop("_openai_moderation_response", e) self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 2ebbeb31c0b..7986d4294a4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -75,19 +75,11 @@ class OvalixGuardrail(CustomGuardrail): post_checkpoint_id: Optional[str] = None, **kwargs: Any, ): - self._tracker_api_base = tracker_api_base or os.environ.get( - "OVALIX_TRACKER_API_BASE" - ) - self._tracker_api_key = tracker_api_key or os.environ.get( - "OVALIX_TRACKER_API_KEY" - ) + self._tracker_api_base = tracker_api_base or os.environ.get("OVALIX_TRACKER_API_BASE") + self._tracker_api_key = tracker_api_key or os.environ.get("OVALIX_TRACKER_API_KEY") self._application_id = application_id or os.environ.get("OVALIX_APPLICATION_ID") - self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get( - "OVALIX_PRE_CHECKPOINT_ID" - ) - self._post_checkpoint_id = post_checkpoint_id or os.environ.get( - "OVALIX_POST_CHECKPOINT_ID" - ) + self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get("OVALIX_PRE_CHECKPOINT_ID") + self._post_checkpoint_id = post_checkpoint_id or os.environ.get("OVALIX_POST_CHECKPOINT_ID") if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [] @@ -102,9 +94,7 @@ class OvalixGuardrail(CustomGuardrail): encoding="utf-8", ) - self._async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) super().__init__(**kwargs) verbose_proxy_logger.debug( @@ -115,58 +105,32 @@ class OvalixGuardrail(CustomGuardrail): self._post_checkpoint_id, ) - def _validate_config( - self, supported_event_hooks: List[GuardrailEventHooks] - ) -> None: + def _validate_config(self, supported_event_hooks: List[GuardrailEventHooks]) -> None: """Ensure required secrets and checkpoint IDs are set; auto-add hooks when IDs are present.""" errors: List[str] = [] if not self._tracker_api_base: - errors.append( - "Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base" - ) + errors.append("Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base") if not self._tracker_api_key: - errors.append( - "Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key" - ) + errors.append("Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key") if not self._application_id: - errors.append( - "Application ID, set OVALIX_APPLICATION_ID or pass application_id" - ) - if ( - not self._pre_checkpoint_id - and GuardrailEventHooks.pre_call in supported_event_hooks - ): - errors.append( - "Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id" - ) - if ( - not self._post_checkpoint_id - and GuardrailEventHooks.post_call in supported_event_hooks - ): - errors.append( - "Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id" - ) + errors.append("Application ID, set OVALIX_APPLICATION_ID or pass application_id") + if not self._pre_checkpoint_id and GuardrailEventHooks.pre_call in supported_event_hooks: + errors.append("Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id") + if not self._post_checkpoint_id and GuardrailEventHooks.post_call in supported_event_hooks: + errors.append("Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id") if not self._pre_checkpoint_id and not self._post_checkpoint_id: errors.append( "Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id" ) if errors: - raise OvalixGuardrailMissingSecrets( - "Missing Ovalix guardrail configuration errors: " + ". ".join(errors) - ) + raise OvalixGuardrailMissingSecrets("Missing Ovalix guardrail configuration errors: " + ". ".join(errors)) # auto-add hooks when checkpoint IDs are present - if ( - self._pre_checkpoint_id - and GuardrailEventHooks.pre_call not in supported_event_hooks - ): + if self._pre_checkpoint_id and GuardrailEventHooks.pre_call not in supported_event_hooks: supported_event_hooks.append(GuardrailEventHooks.pre_call) - if ( - self._post_checkpoint_id - and GuardrailEventHooks.post_call not in supported_event_hooks - ): + if self._post_checkpoint_id and GuardrailEventHooks.post_call not in supported_event_hooks: supported_event_hooks.append(GuardrailEventHooks.post_call) def _get_actor(self, data: dict) -> str: @@ -276,13 +240,9 @@ class OvalixGuardrail(CustomGuardrail): is_first_response = True for llm_response in reversed(texts): try: - resp = await self._call_checkpoint( - llm_response, checkpoint_id, actor, session_id - ) + resp = await self._call_checkpoint(llm_response, checkpoint_id, actor, session_id) except Exception as e: - verbose_proxy_logger.exception( - "Ovalix apply_guardrail checkpoint call failed: %s", e - ) + verbose_proxy_logger.exception("Ovalix apply_guardrail checkpoint call failed: %s", e) raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=f"Ovalix guardrail error: {e!s}", @@ -290,18 +250,13 @@ class OvalixGuardrail(CustomGuardrail): ) from e action_type = (resp.get("action_type") or "").lower() - blocking_message = ( - self._get_trackers_corrected_message(resp) - or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE - ) + blocking_message = self._get_trackers_corrected_message(resp) or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE if action_type == BLOCKED_ACTION_TYPE and is_first_response: self._block_current_message(blocking_message) elif action_type == BLOCKED_ACTION_TYPE: post_guardrail_texts.insert(0, blocking_message) else: - corrected_text = ( - self._get_trackers_corrected_message(resp) or llm_response - ) + corrected_text = self._get_trackers_corrected_message(resp) or llm_response post_guardrail_texts.insert(0, corrected_text) is_first_response = False return post_guardrail_texts diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 2974febe022..3d3c5403993 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -77,9 +77,7 @@ class PangeaHandler(CustomGuardrail): api_base (Optional[str]): The Pangea API base URL. Reads from PANGEA_API_BASE env var or uses default if None. **kwargs: Additional arguments passed to the CustomGuardrail base class. """ - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_key = api_key or os.environ.get("PANGEA_API_KEY") if not self.api_key: raise PangeaGuardrailMissingSecrets( @@ -87,11 +85,7 @@ class PangeaHandler(CustomGuardrail): ) # Default Pangea base URL if not provided - self.api_base = ( - api_base - or os.environ.get("PANGEA_API_BASE") - or "https://ai-guard.aws.us.pangea.cloud" - ) + self.api_base = api_base or os.environ.get("PANGEA_API_BASE") or "https://ai-guard.aws.us.pangea.cloud" self.pangea_input_recipe = pangea_input_recipe self.pangea_output_recipe = pangea_output_recipe @@ -110,9 +104,7 @@ class PangeaHandler(CustomGuardrail): f"Initialized Pangea Guardrail: name={guardrail_name}, recipe={pangea_input_recipe}, api_base={self.api_base}" ) - async def _call_pangea_ai_guard( - self, api: str, payload: dict, hook_name: str - ) -> dict: + async def _call_pangea_ai_guard(self, api: str, payload: dict, hook_name: str) -> dict: """ Makes the API call to the Pangea AI Guard endpoint. The function itself will raise an error in the case that a response @@ -143,17 +135,13 @@ class PangeaHandler(CustomGuardrail): f"Pangea Guardrail ({hook_name}): Calling endpoint {endpoint} with payload: {payload}" ) - response = await self.async_handler.post( - url=endpoint, json=payload, headers=headers - ) + response = await self.async_handler.post(url=endpoint, json=payload, headers=headers) response.raise_for_status() result = response.json() if result.get("result", {}).get("blocked"): - verbose_proxy_logger.warning( - f"Pangea Guardrail ({hook_name}): Request blocked. Response: {result}" - ) + verbose_proxy_logger.warning(f"Pangea Guardrail ({hook_name}): Request blocked. Response: {result}") raise HTTPException( status_code=400, # Bad Request, indicating violation detail={ @@ -190,12 +178,8 @@ class PangeaHandler(CustomGuardrail): if self.pangea_input_recipe: ai_guard_payload["recipe"] = self.pangea_input_recipe - ai_guard_response = await self._call_pangea_ai_guard( - "v1beta/guard", ai_guard_payload, "async_pre_call_hook" - ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + ai_guard_response = await self._call_pangea_ai_guard("v1beta/guard", ai_guard_payload, "async_pre_call_hook") + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) if not ai_guard_response.get("result", {}).get("transformed"): return @@ -223,9 +207,7 @@ class PangeaHandler(CustomGuardrail): return data try: - return await self._async_pre_call_hook( - user_api_key_dict, cache, data, call_type - ) + return await self._async_pre_call_hook(user_api_key_dict, cache, data, call_type) except HTTPException: raise except Exception as e: @@ -282,12 +264,8 @@ class PangeaHandler(CustomGuardrail): if self.pangea_output_recipe: ai_guard_payload["recipe"] = self.pangea_output_recipe - ai_guard_response = await self._call_pangea_ai_guard( - "v1beta/guard", ai_guard_payload, "async_pre_call_hook" - ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + ai_guard_response = await self._call_pangea_ai_guard("v1beta/guard", ai_guard_payload, "async_pre_call_hook") + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) if not ai_guard_response.get("result", {}).get("transformed"): return @@ -319,9 +297,7 @@ class PangeaHandler(CustomGuardrail): ) return data try: - return await self._async_post_call_success_hook( - data, user_api_key_dict, response - ) + return await self._async_post_call_success_hook(data, user_api_key_dict, response) except HTTPException: raise except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 04d68edeb80..c522ffad35d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -111,9 +111,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Store configuration with env var fallbacks self.api_key = api_key or os.getenv("PANW_PRISMA_AIRS_API_KEY") self.api_base = ( - api_base - or os.getenv("PANW_PRISMA_AIRS_API_BASE") - or "https://service.api.aisecurity.paloaltonetworks.com" + api_base or os.getenv("PANW_PRISMA_AIRS_API_BASE") or "https://service.api.aisecurity.paloaltonetworks.com" ) self.profile_name = profile_name @@ -204,16 +202,12 @@ class PanwPrismaAirsHandler(CustomGuardrail): return "" - def _extract_text_from_content_list( - self, content_list: List[Dict[str, Any]] - ) -> str: + def _extract_text_from_content_list(self, content_list: List[Dict[str, Any]]) -> str: """Extract text from content list format.""" text_parts = [ part.get("text", "") for part in content_list - if isinstance(part, dict) - and part.get("type") == "text" - and part.get("text") + if isinstance(part, dict) and part.get("type") == "text" and part.get("text") ] return " ".join(text_parts) if text_parts else "" @@ -234,31 +228,19 @@ class PanwPrismaAirsHandler(CustomGuardrail): text_parts.append(str(choice.message.content)) # Extract tool call arguments - if ( - hasattr(choice.message, "tool_calls") - and choice.message.tool_calls - ): + if hasattr(choice.message, "tool_calls") and choice.message.tool_calls: for tool_call in choice.message.tool_calls: - if hasattr(tool_call, "function") and hasattr( - tool_call.function, "arguments" - ): + if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"): text_parts.append(str(tool_call.function.arguments)) # Extract function call arguments (legacy) - if ( - hasattr(choice.message, "function_call") - and choice.message.function_call - ): + if hasattr(choice.message, "function_call") and choice.message.function_call: if hasattr(choice.message.function_call, "arguments"): - text_parts.append( - str(choice.message.function_call.arguments) - ) + text_parts.append(str(choice.message.function_call.arguments)) return " ".join(text_parts) if text_parts else "" except (AttributeError, IndexError) as e: - verbose_proxy_logger.error( - f"PANW Prisma AIRS: Error extracting response text: {str(e)}" - ) + verbose_proxy_logger.error(f"PANW Prisma AIRS: Error extracting response text: {str(e)}") return "" async def _call_panw_api( @@ -300,9 +282,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): panw_metadata = { "app_user": ( - (metadata.get("app_user") or metadata.get("user") or "litellm_user") - if metadata - else "litellm_user" + (metadata.get("app_user") or metadata.get("user") or "litellm_user") if metadata else "litellm_user" ), "ai_model": metadata.get("model", "unknown") if metadata else "unknown", "app_name": app_name_value, @@ -363,15 +343,12 @@ class PanwPrismaAirsHandler(CustomGuardrail): headers = { "Content-Type": "application/json", "Accept": "application/json", - "x-pan-token": self.api_key - or "", # api_key validated in __init__, never None + "x-pan-token": self.api_key or "", # api_key validated in __init__, never None } try: # Use LiteLLM's async HTTP client - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) # Bypass wrapper to access follow_redirects parameter response = await async_client.client.post( # type: ignore[attr-defined] @@ -387,18 +364,14 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Validate response format if "action" not in result: - verbose_proxy_logger.error( - f"PANW Prisma AIRS: Invalid API response format: {result}" - ) + verbose_proxy_logger.error(f"PANW Prisma AIRS: Invalid API response format: {result}") return {"action": "block", "category": "api_error"} # Check for profile-related errors from PANW API if result.get("action") == "block" and "error" in result: error_msg = str(result.get("error", "")).lower() if "profile" in error_msg and ( - "not found" in error_msg - or "required" in error_msg - or "invalid" in error_msg + "not found" in error_msg or "required" in error_msg or "invalid" in error_msg ): verbose_proxy_logger.error( f"PANW Prisma AIRS: Profile configuration error. " @@ -424,9 +397,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): if status == 400: diag_parts = ["PANW Prisma AIRS: HTTP 400 from AIRS API."] if tool_event is not None: - diag_parts.append( - f"tool_event.metadata={tool_event.get('metadata')}" - ) + diag_parts.append(f"tool_event.metadata={tool_event.get('metadata')}") has_input = "input" in tool_event input_len = len(tool_event["input"]) if has_input else 0 diag_parts.append(f"input present={has_input}, len={input_len}") @@ -454,9 +425,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): } elif status == 429 or status >= 500: # Transient: rate-limit and server errors — safe to fail-open - verbose_proxy_logger.error( - f"PANW Prisma AIRS: API error (HTTP {status}): {error_body[:500]}" - ) + verbose_proxy_logger.error(f"PANW Prisma AIRS: API error (HTTP {status}): {error_body[:500]}") return { "action": "block", "category": f"http_{status}_error", @@ -465,9 +434,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): else: # Permanent 4xx client errors (400, 404, etc.) — must not bypass scanning if status != 400: # 400 already logged with diagnostics above - verbose_proxy_logger.error( - f"PANW Prisma AIRS: API error (HTTP {status}): {error_body[:500]}" - ) + verbose_proxy_logger.error(f"PANW Prisma AIRS: API error (HTTP {status}): {error_body[:500]}") return { "action": "block", "category": f"http_{status}_error", @@ -483,9 +450,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): } except httpx.RequestError as e: - verbose_proxy_logger.error( - f"PANW Prisma AIRS: Network/request error: {str(e)}" - ) + verbose_proxy_logger.error(f"PANW Prisma AIRS: Network/request error: {str(e)}") return { "action": "block", "category": "network_error", @@ -519,9 +484,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): or getattr(server, "server_id", None) or "unknown" ) - return global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.get( - mcp_tool_name, "unknown" - ) + return global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.get(mcp_tool_name, "unknown") except ImportError: return "unknown" except Exception: @@ -531,9 +494,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) return "unknown" - def _get_masked_text( - self, scan_result: Dict[str, Any], is_response: bool = False - ) -> Optional[str]: + def _get_masked_text(self, scan_result: Dict[str, Any], is_response: bool = False) -> Optional[str]: """Extract masked text from PANW scan result.""" masked_key = "response_masked_data" if is_response else "prompt_masked_data" masked_data = scan_result.get(masked_key) @@ -614,13 +575,9 @@ class PanwPrismaAirsHandler(CustomGuardrail): "PANW Prisma AIRS: MCP request blocked but masked instead (mask_request_content=True)" ) else: - verbose_proxy_logger.info( - "PANW Prisma AIRS: MCP request allowed with PII masking applied" - ) + verbose_proxy_logger.info("PANW Prisma AIRS: MCP request allowed with PII masking applied") - def _apply_masking_to_messages( - self, messages: List[Dict[str, Any]], masked_text: str - ) -> List[Dict[str, Any]]: + def _apply_masking_to_messages(self, messages: List[Dict[str, Any]], masked_text: str) -> List[Dict[str, Any]]: """Apply masked text to the last user message.""" if not messages: return messages @@ -633,18 +590,14 @@ class PanwPrismaAirsHandler(CustomGuardrail): if isinstance(content, str): new_message["content"] = masked_text elif isinstance(content, list): - new_message["content"] = self._mask_content_list( - content, masked_text - ) + new_message["content"] = self._mask_content_list(content, masked_text) idx = len(messages) - i - 1 return messages[:idx] + [new_message] + messages[idx + 1 :] return messages - def _apply_masking_to_response( - self, response: ModelResponse, masked_text: str - ) -> None: + def _apply_masking_to_response(self, response: ModelResponse, masked_text: str) -> None: """ Apply masked text to all content in response in-place. Handles message content, tool calls, and function calls across all choices. @@ -668,22 +621,15 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Mask tool call arguments if hasattr(choice.message, "tool_calls") and choice.message.tool_calls: for tool_call in choice.message.tool_calls: - if hasattr(tool_call, "function") and hasattr( - tool_call.function, "arguments" - ): + if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"): tool_call.function.arguments = masked_text # Mask function call arguments (legacy) - if ( - hasattr(choice.message, "function_call") - and choice.message.function_call - ): + if hasattr(choice.message, "function_call") and choice.message.function_call: if hasattr(choice.message.function_call, "arguments"): choice.message.function_call.arguments = masked_text - def _build_error_detail( - self, scan_result: Dict[str, Any], is_response: bool = False - ) -> Dict[str, Any]: + def _build_error_detail(self, scan_result: Dict[str, Any], is_response: bool = False) -> Dict[str, Any]: """Build enhanced error detail with scan information.""" action_type = "Response" if is_response else "Prompt" code_suffix = "_response_blocked" if is_response else "_blocked" @@ -766,16 +712,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): if is_config else "Security scan failed - request blocked for safety" ), - "type": ( - "guardrail_config_error" - if is_config - else "guardrail_scan_error" - ), - "code": ( - "panw_prisma_airs_config_error" - if is_config - else "panw_prisma_airs_scan_failed" - ), + "type": ("guardrail_config_error" if is_config else "guardrail_scan_error"), + "code": ("panw_prisma_airs_config_error" if is_config else "panw_prisma_airs_scan_failed"), "guardrail": self.guardrail_name, "category": category, } @@ -888,10 +826,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): if isinstance(delta, str): parts.append(delta) # Defense-in-depth: handle dict chat.completion.chunk format - elif ( - isinstance(chunk, dict) - and chunk.get("object") == "chat.completion.chunk" - ): + elif isinstance(chunk, dict) and chunk.get("object") == "chat.completion.chunk": for choice in chunk.get("choices") or []: if isinstance(choice, dict): delta = choice.get("delta") or {} @@ -907,9 +842,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): parts.append(text) return "".join(parts) - async def _scan_raw_streaming_text( - self, text: str, request_data: dict, start_time: datetime - ) -> None: + async def _scan_raw_streaming_text(self, text: str, request_data: dict, start_time: datetime) -> None: """Scan text from non-ModelResponse streaming chunks. Raises HTTPException(400) on block. Note: response masking is not supported on raw streaming paths @@ -964,9 +897,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): duration=(end_time - start_time).total_seconds(), event_type=GuardrailEventHooks.post_call, ) - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) def _check_and_mark_scanned(self, data: dict, scan_type: str) -> bool: """ @@ -984,8 +915,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): call_id = str(uuid.uuid4()) data["litellm_call_id"] = call_id verbose_proxy_logger.warning( - "PANW Prisma AIRS: litellm_call_id missing from request data, " - "synthesized %s for %s scan deduplication", + "PANW Prisma AIRS: litellm_call_id missing from request data, synthesized %s for %s scan deduplication", call_id, scan_type, ) @@ -994,9 +924,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): litellm_metadata = data.setdefault("litellm_metadata", {}) if litellm_metadata.get(scan_key): - verbose_proxy_logger.debug( - f"PANW Prisma AIRS: Skipping duplicate {scan_type}-call scan" - ) + verbose_proxy_logger.debug(f"PANW Prisma AIRS: Skipping duplicate {scan_type}-call scan") return True # Already scanned litellm_metadata[scan_key] = True @@ -1093,11 +1021,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): guardrail_provider=self._PROVIDER_NAME, guardrail_json_response=scan_result, request_data=data, - guardrail_status=( - "success" - if scan_result.get("action") == "allow" - else "guardrail_intervened" - ), + guardrail_status=("success" if scan_result.get("action") == "allow" else "guardrail_intervened"), start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), @@ -1112,44 +1036,30 @@ class PanwPrismaAirsHandler(CustomGuardrail): if action == "allow": if masked_text: if messages: - data["messages"] = self._apply_masking_to_messages( - messages, masked_text - ) + data["messages"] = self._apply_masking_to_messages(messages, masked_text) elif "prompt" in data: data["prompt"] = masked_text - verbose_proxy_logger.info( - f"PANW Prisma AIRS: Prompt allowed with masking (Category: {category})" - ) + verbose_proxy_logger.info(f"PANW Prisma AIRS: Prompt allowed with masking (Category: {category})") else: - verbose_proxy_logger.info( - f"PANW Prisma AIRS: Prompt allowed (Category: {category})" - ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + verbose_proxy_logger.info(f"PANW Prisma AIRS: Prompt allowed (Category: {category})") + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return None # Action is "block" - check if we should mask instead of blocking if masked_text and self.mask_request_content: if messages: - data["messages"] = self._apply_masking_to_messages( - messages, masked_text - ) + data["messages"] = self._apply_masking_to_messages(messages, masked_text) elif "prompt" in data: data["prompt"] = masked_text verbose_proxy_logger.warning( "PANW Prisma AIRS: Prompt blocked but masked instead (mask_request_content=True)" ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return None # Block the request error_detail = self._build_error_detail(scan_result, is_response=False) - verbose_proxy_logger.warning( - f"PANW Prisma AIRS: {error_detail['error']['message']}" - ) + verbose_proxy_logger.warning(f"PANW Prisma AIRS: {error_detail['error']['message']}") raise HTTPException(status_code=400, detail=error_detail) except HTTPException: @@ -1202,9 +1112,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): response_text = self._extract_response_text(response) if not response_text: - verbose_proxy_logger.warning( - "PANW Prisma AIRS: No response content found to scan" - ) + verbose_proxy_logger.warning("PANW Prisma AIRS: No response content found to scan") return response # Prepare metadata - include user's metadata for profile override @@ -1233,11 +1141,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): guardrail_provider=self._PROVIDER_NAME, guardrail_json_response=scan_result, request_data=data, - guardrail_status=( - "success" - if scan_result.get("action") == "allow" - else "guardrail_intervened" - ), + guardrail_status=("success" if scan_result.get("action") == "allow" else "guardrail_intervened"), start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), @@ -1252,16 +1156,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): if action == "allow": if masked_text: self._apply_masking_to_response(response, masked_text) - verbose_proxy_logger.info( - f"PANW Prisma AIRS: Response allowed with masking (Category: {category})" - ) + verbose_proxy_logger.info(f"PANW Prisma AIRS: Response allowed with masking (Category: {category})") else: - verbose_proxy_logger.info( - f"PANW Prisma AIRS: Response allowed (Category: {category})" - ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + verbose_proxy_logger.info(f"PANW Prisma AIRS: Response allowed (Category: {category})") + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response # Action is "block" - check if we should mask instead of blocking @@ -1270,16 +1168,12 @@ class PanwPrismaAirsHandler(CustomGuardrail): verbose_proxy_logger.warning( "PANW Prisma AIRS: Response blocked but masked instead (mask_response_content=True)" ) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response # Block the response error_detail = self._build_error_detail(scan_result, is_response=True) - verbose_proxy_logger.warning( - f"PANW Prisma AIRS: {error_detail['error']['message']}" - ) + verbose_proxy_logger.warning(f"PANW Prisma AIRS: {error_detail['error']['message']}") raise HTTPException(status_code=400, detail=error_detail) except HTTPException: @@ -1312,9 +1206,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): response_text = self._extract_response_text(assembled_model_response) if not response_text or not response_text.strip(): - verbose_proxy_logger.info( - "PANW Prisma AIRS: No content to scan in streaming response" - ) + verbose_proxy_logger.info("PANW Prisma AIRS: No content to scan in streaming response") return ( content_was_modified, assembled_model_response, @@ -1349,9 +1241,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): f"PANW Prisma AIRS: Streaming response allowed with masking (Category: {category})" ) else: - verbose_proxy_logger.info( - f"PANW Prisma AIRS: Streaming response allowed (Category: {category})" - ) + verbose_proxy_logger.info(f"PANW Prisma AIRS: Streaming response allowed (Category: {category})") elif masked_text and self.mask_response_content: self._apply_masking_to_response(assembled_model_response, masked_text) content_was_modified = True @@ -1360,9 +1250,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) else: error_detail = self._build_error_detail(scan_result, is_response=True) - verbose_proxy_logger.warning( - f"PANW Prisma AIRS: {error_detail['error']['message']}" - ) + verbose_proxy_logger.warning(f"PANW Prisma AIRS: {error_detail['error']['message']}") raise HTTPException(status_code=400, detail=error_detail) return content_was_modified, assembled_model_response, scan_result @@ -1382,9 +1270,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Check if guardrail should run for this request - if not self.should_run_guardrail( - data=request_data, event_type=GuardrailEventHooks.post_call - ): + if not self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call): async for chunk in response: yield chunk return @@ -1416,9 +1302,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): return # Handle /v1/responses streaming: chunks are Pydantic events (not ModelResponse/ModelResponseStream) - if all_chunks and not isinstance( - all_chunks[0], (ModelResponse, ModelResponseStream) - ): + if all_chunks and not isinstance(all_chunks[0], (ModelResponse, ModelResponseStream)): text = self._extract_text_from_streaming_events(all_chunks) await self._scan_raw_streaming_text(text, request_data, start_time) for chunk in all_chunks: @@ -1434,9 +1318,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): content_was_modified, assembled_model_response, scan_result, - ) = await self._scan_and_process_streaming_response( - assembled_model_response, request_data, start_time - ) + ) = await self._scan_and_process_streaming_response(assembled_model_response, request_data, start_time) if scan_result.get("_is_transient") or scan_result.get("_always_block"): self._handle_api_error_with_logging( @@ -1458,11 +1340,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): guardrail_provider=self._PROVIDER_NAME, guardrail_json_response=scan_result, request_data=request_data, - guardrail_status=( - "success" - if scan_result.get("action") == "allow" - else "guardrail_intervened" - ), + guardrail_status=("success" if scan_result.get("action") == "allow" else "guardrail_intervened"), start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), @@ -1477,9 +1355,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Only use MockResponseIterator if content was modified # Otherwise, yield original chunks to preserve streaming behavior if content_was_modified: - mock_response = MockResponseIterator( - model_response=assembled_model_response - ) + mock_response = MockResponseIterator(model_response=assembled_model_response) async for chunk in mock_response: yield chunk else: @@ -1494,9 +1370,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Yield error as SSE event so create_response() detects it and # returns a proper JSON error response with the correct status code. # (Raising from a generator hits create_response's generic except → 500.) - detail = ( - e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} - ) + detail = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} error_obj: Dict[str, Any] = dict(detail.get("error", detail)) # type: ignore[arg-type] error_obj["code"] = e.status_code yield f"data: {json.dumps({'error': error_obj})}\n\n" @@ -1536,9 +1410,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): tool_name: Optional[str] = None args_text: Optional[str] = None - if hasattr(tool_call, "function") and hasattr( - tool_call.function, "arguments" - ): + if hasattr(tool_call, "function") and hasattr(tool_call.function, "arguments"): args_text = tool_call.function.arguments tool_name = getattr(tool_call.function, "name", None) elif isinstance(tool_call, dict): @@ -1567,11 +1439,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) if scan_result.get("_is_transient") or scan_result.get("_always_block"): - event_type = ( - GuardrailEventHooks.post_call - if is_response - else GuardrailEventHooks.pre_call - ) + event_type = GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call self._handle_api_error_with_logging( scan_result=scan_result, data=request_data, @@ -1591,14 +1459,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if masked_text: self._set_tool_call_arguments(tool_call, masked_text) elif masked_text and ( - (is_response and self.mask_response_content) - or (not is_response and self.mask_request_content) + (is_response and self.mask_response_content) or (not is_response and self.mask_request_content) ): self._set_tool_call_arguments(tool_call, masked_text) else: - error_detail = self._build_error_detail( - scan_result, is_response=is_response - ) + error_detail = self._build_error_detail(scan_result, is_response=is_response) raise HTTPException(status_code=400, detail=error_detail) @staticmethod @@ -1606,9 +1471,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): """Set masked text on a tool call's function arguments, handling both object and dict forms.""" if hasattr(tool_call, "function"): tool_call.function.arguments = masked_text - elif isinstance(tool_call, dict) and isinstance( - tool_call.get("function"), dict - ): + elif isinstance(tool_call, dict) and isinstance(tool_call.get("function"), dict): tool_call["function"]["arguments"] = masked_text @staticmethod @@ -1786,11 +1649,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): call_id = getattr(logging_obj, "litellm_call_id", None) if not call_id: # Use MCP name fallback: mcp_tool_name (canonical) or name (/mcp-rest path) - _mcp_tool = str( - request_data.get("mcp_tool_name") - or self._mcp_name_fallback(request_data) - or "" - ).strip() + _mcp_tool = str(request_data.get("mcp_tool_name") or self._mcp_name_fallback(request_data) or "").strip() if input_type == "request" and logging_obj is None and _mcp_tool: # Synthesize a tool-prefixed call_id for AIRS grouping. # Slug: lowercase, non-alphanum → "-", truncate to 40 chars. @@ -1818,8 +1677,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): call_id = str(uuid.uuid4()) request_data["litellm_call_id"] = call_id verbose_proxy_logger.warning( - "PANW Prisma AIRS: litellm_call_id missing, synthesized %s " - "(input_type=%s)", + "PANW Prisma AIRS: litellm_call_id missing, synthesized %s (input_type=%s)", call_id, input_type, ) @@ -1834,9 +1692,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Enrich request_data with metadata from logging_obj (post-call metadata loss). # Merge: logging_obj provides the base, request_data keys win on conflict. if logging_obj: - _lp = (getattr(logging_obj, "model_call_details", {}) or {}).get( - "litellm_params", {} - ) or {} + _lp = (getattr(logging_obj, "model_call_details", {}) or {}).get("litellm_params", {}) or {} _orig_meta = _lp.get("metadata") or {} if _orig_meta: existing_meta = request_data.get("metadata") @@ -1860,17 +1716,13 @@ class PanwPrismaAirsHandler(CustomGuardrail): if self._use_latest_user_only(request_data, logging_obj): original_messages = request_data.get("messages") if original_messages: - scannable_indices = self._get_latest_user_text_indices( - texts, original_messages - ) + scannable_indices = self._get_latest_user_text_indices(texts, original_messages) # Fall through to existing role filtering if: # - not Anthropic, OR flag explicitly False, OR # - no original messages, OR # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: - scannable_indices = self._get_scannable_text_indices( - texts, structured_messages - ) + scannable_indices = self._get_scannable_text_indices(texts, structured_messages) for i, text in enumerate(texts): if not text or not text.strip(): @@ -1891,11 +1743,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # Handle API errors (transient/config) if scan_result.get("_is_transient") or scan_result.get("_always_block"): - event_type = ( - GuardrailEventHooks.post_call - if is_response - else GuardrailEventHooks.pre_call - ) + event_type = GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call self._handle_api_error_with_logging( scan_result=scan_result, data=request_data, @@ -1913,14 +1761,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if action == "allow": new_texts.append(masked_text if masked_text else text) elif masked_text and ( - (is_response and self.mask_response_content) - or (not is_response and self.mask_request_content) + (is_response and self.mask_response_content) or (not is_response and self.mask_request_content) ): new_texts.append(masked_text) else: - error_detail = self._build_error_detail( - scan_result, is_response=is_response - ) + error_detail = self._build_error_detail(scan_result, is_response=is_response) raise HTTPException(status_code=400, detail=error_detail) # Scan tool call arguments — same masking policy as texts. @@ -1943,17 +1788,13 @@ class PanwPrismaAirsHandler(CustomGuardrail): # We send a tool_event so AIRS can apply tool-aware policies. # REST MCP path sets "name"/"arguments"; canonical keys are # "mcp_tool_name"/"mcp_arguments". Check canonical first, then fallback. - mcp_tool_name = request_data.get("mcp_tool_name") or self._mcp_name_fallback( - request_data - ) + mcp_tool_name = request_data.get("mcp_tool_name") or self._mcp_name_fallback(request_data) if mcp_tool_name and input_type == "request": mcp_tool_event: Dict[str, Any] = { "metadata": { "ecosystem": "mcp", "method": "tools/call", - "server_name": self._get_mcp_server_name( - request_data, mcp_tool_name - ), + "server_name": self._get_mcp_server_name(request_data, mcp_tool_name), "tool_invoked": mcp_tool_name, }, } @@ -1974,9 +1815,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): call_id=call_id, ) - if mcp_scan_result.get("_is_transient") or mcp_scan_result.get( - "_always_block" - ): + if mcp_scan_result.get("_is_transient") or mcp_scan_result.get("_always_block"): self._handle_api_error_with_logging( scan_result=mcp_scan_result, data=request_data, @@ -2001,15 +1840,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): elif masked_text and self.mask_request_content: self._apply_mcp_masking(request_data, mcp_arguments, masked_text) else: - error_detail = self._build_error_detail( - mcp_scan_result, is_response=False - ) + error_detail = self._build_error_detail(mcp_scan_result, is_response=False) raise HTTPException(status_code=400, detail=error_detail) inputs["texts"] = new_texts - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) return inputs @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py index 5ef6f32ead5..405edb482d6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py @@ -33,18 +33,10 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" event_hook=litellm_params.mode, default_on=litellm_params.default_on, async_mode=_get_config_value(litellm_params, optional_params, "async_mode"), - persist_session=_get_config_value( - litellm_params, optional_params, "persist_session" - ), - include_scanners=_get_config_value( - litellm_params, optional_params, "include_scanners" - ), - include_evidence=_get_config_value( - litellm_params, optional_params, "include_evidence" - ), - fallback_on_error=_get_config_value( - litellm_params, optional_params, "fallback_on_error" - ), + persist_session=_get_config_value(litellm_params, optional_params, "persist_session"), + include_scanners=_get_config_value(litellm_params, optional_params, "include_scanners"), + include_evidence=_get_config_value(litellm_params, optional_params, "include_evidence"), + fallback_on_error=_get_config_value(litellm_params, optional_params, "fallback_on_error"), timeout=_get_config_value(litellm_params, optional_params, "timeout"), ) litellm.logging_callback_manager.add_litellm_callback(_pillar_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index c9c73053a0e..f976839787b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -50,9 +50,7 @@ def _encode_json_for_header(data: Any) -> str: return quote(json_payload, safe="") -def _truncate_evidence_payload( - evidence: Any, max_bytes: int = MAX_PILLAR_HEADER_VALUE_BYTES -) -> Tuple[Any, str, bool]: +def _truncate_evidence_payload(evidence: Any, max_bytes: int = MAX_PILLAR_HEADER_VALUE_BYTES) -> Tuple[Any, str, bool]: """ Truncate evidence payload so the encoded header value stays within max_bytes. @@ -91,17 +89,9 @@ def _truncate_evidence_payload( if evidence_text: step = max(1, len(evidence_text) // 2) while len(encoded.encode("utf-8")) > max_bytes and evidence_text: - evidence_text = ( - evidence_text[:-step] - if len(evidence_text) > step - else evidence_text[:-1] - ) + evidence_text = evidence_text[:-step] if len(evidence_text) > step else evidence_text[:-1] step = max(1, step // 2) - truncated_text = ( - f"{evidence_text}...[truncated]" - if evidence_text - else "[truncated]" - ) + truncated_text = f"{evidence_text}...[truncated]" if evidence_text else "[truncated]" working_entry["evidence"] = truncated_text working_entry["evidence_truncated"] = True encoded = _encode_json_for_header(truncated) @@ -125,9 +115,7 @@ def build_pillar_response_headers(metadata_store: Dict[str, Any]) -> Dict[str, s headers["x-pillar-flagged"] = str(metadata_store["pillar_flagged"]).lower() if "pillar_scanners" in metadata_store: - headers["x-pillar-scanners"] = _encode_json_for_header( - metadata_store["pillar_scanners"] - ) + headers["x-pillar-scanners"] = _encode_json_for_header(metadata_store["pillar_scanners"]) if "pillar_evidence" in metadata_store: truncated_evidence, encoded_value, truncated_flag = _truncate_evidence_payload( @@ -139,9 +127,7 @@ def build_pillar_response_headers(metadata_store: Dict[str, Any]) -> Dict[str, s headers["x-pillar-evidence"] = encoded_value if "pillar_session_id_response" in metadata_store: - headers["x-pillar-session-id"] = quote( - str(metadata_store["pillar_session_id_response"]), safe="" - ) + headers["x-pillar-session-id"] = quote(str(metadata_store["pillar_session_id_response"]), safe="") if headers: metadata_store["pillar_response_headers"] = headers @@ -177,9 +163,7 @@ class PillarGuardrail(CustomGuardrail): SUPPORTED_FALLBACK_ACTIONS = ["allow", "block"] DEFAULT_FALLBACK_ACTION = "allow" BASE_API_URL = "https://api.pillar.security" - DEFAULT_TIMEOUT = ( - 5.0 # 5 seconds - fast failure detection with graceful degradation - ) + DEFAULT_TIMEOUT = 5.0 # 5 seconds - fast failure detection with graceful degradation def __init__( self, @@ -211,9 +195,7 @@ class PillarGuardrail(CustomGuardrail): LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always automatically passed as X-LiteLLM-* headers to enable application/user tracking. """ - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_key = api_key or os.environ.get("PILLAR_API_KEY") if self.api_key is None: @@ -231,14 +213,10 @@ class PillarGuardrail(CustomGuardrail): self.on_flagged_action = action else: if action: - verbose_proxy_logger.warning( - f"Invalid action '{action}', using default" - ) + verbose_proxy_logger.warning(f"Invalid action '{action}', using default") self.on_flagged_action = self.DEFAULT_ON_FLAGGED_ACTION - verbose_proxy_logger.debug( - f"Pillar Guardrail: Initialized with on_flagged_action: {self.on_flagged_action}" - ) + verbose_proxy_logger.debug(f"Pillar Guardrail: Initialized with on_flagged_action: {self.on_flagged_action}") self.async_mode = self._resolve_bool_config( provided_value=async_mode, @@ -276,18 +254,14 @@ class PillarGuardrail(CustomGuardrail): ) self.fallback_on_error = self.DEFAULT_FALLBACK_ACTION - verbose_proxy_logger.debug( - f"Pillar Guardrail: Initialized with fallback_on_error: {self.fallback_on_error}" - ) + verbose_proxy_logger.debug(f"Pillar Guardrail: Initialized with fallback_on_error: {self.fallback_on_error}") # Set timeout with graceful fallback on invalid configuration if timeout is not None: self.timeout = timeout else: try: - self.timeout = float( - os.environ.get("PILLAR_TIMEOUT", str(self.DEFAULT_TIMEOUT)) - ) + self.timeout = float(os.environ.get("PILLAR_TIMEOUT", str(self.DEFAULT_TIMEOUT))) except (ValueError, TypeError): verbose_proxy_logger.warning( f"Pillar Guardrail: Invalid PILLAR_TIMEOUT value '{os.environ.get('PILLAR_TIMEOUT')}', " @@ -350,18 +324,14 @@ class PillarGuardrail(CustomGuardrail): """ event_type = GuardrailEventHooks.pre_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug( - f"Pillar Guardrail: Pre-call scanning disabled for {self.guardrail_name}" - ) + verbose_proxy_logger.debug(f"Pillar Guardrail: Pre-call scanning disabled for {self.guardrail_name}") return data verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook") result = await self.run_pillar_guardrail(data, user_api_key_dict) # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result @@ -397,18 +367,14 @@ class PillarGuardrail(CustomGuardrail): """ event_type = GuardrailEventHooks.during_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug( - f"Pillar Guardrail: During-call scanning disabled for {self.guardrail_name}" - ) + verbose_proxy_logger.debug(f"Pillar Guardrail: During-call scanning disabled for {self.guardrail_name}") return data verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook") result = await self.run_pillar_guardrail(data, user_api_key_dict) # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result @@ -435,9 +401,7 @@ class PillarGuardrail(CustomGuardrail): """ event_type = GuardrailEventHooks.post_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug( - f"Pillar Guardrail: Post-call scanning disabled for {self.guardrail_name}" - ) + verbose_proxy_logger.debug(f"Pillar Guardrail: Post-call scanning disabled for {self.guardrail_name}") return response verbose_proxy_logger.debug("Pillar Guardrail: Post-call hook") @@ -445,15 +409,11 @@ class PillarGuardrail(CustomGuardrail): # Extract response messages in the format Pillar expects response_dict = response.model_dump() if hasattr(response, "model_dump") else {} # type: ignore[union-attr] response_messages = [ - choice.get("message") - for choice in response_dict.get("choices", []) - if choice.get("message") + choice.get("message") for choice in response_dict.get("choices", []) if choice.get("message") ] if not response_messages: - verbose_proxy_logger.debug( - "Pillar Guardrail: No response content to scan, skipping post-call analysis" - ) + verbose_proxy_logger.debug("Pillar Guardrail: No response content to scan, skipping post-call analysis") return response # Create complete conversation: original messages + response messages @@ -464,9 +424,7 @@ class PillarGuardrail(CustomGuardrail): await self.run_pillar_guardrail(post_call_data, user_api_key_dict) # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response @@ -474,9 +432,7 @@ class PillarGuardrail(CustomGuardrail): # CORE LOGIC METHOD # ========================================================================= - async def run_pillar_guardrail( - self, data: dict, user_api_key_dict: UserAPIKeyAuth - ) -> dict: + async def run_pillar_guardrail(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> dict: """ Core method to run the Pillar guardrail scan. @@ -492,9 +448,7 @@ class PillarGuardrail(CustomGuardrail): """ # Check if messages are present if not data.get("messages"): - verbose_proxy_logger.debug( - "Pillar Guardrail: No messages detected, bypassing security scan" - ) + verbose_proxy_logger.debug("Pillar Guardrail: No messages detected, bypassing security scan") return data try: @@ -516,9 +470,7 @@ class PillarGuardrail(CustomGuardrail): raise e # Handle API communication errors based on fallback_on_error setting - verbose_proxy_logger.error( - f"Pillar Guardrail: API communication failed - {str(e)}" - ) + verbose_proxy_logger.error(f"Pillar Guardrail: API communication failed - {str(e)}") return self._handle_api_error(e, data) @@ -602,9 +554,7 @@ class PillarGuardrail(CustomGuardrail): return headers - def _set_bool_header( - self, headers: Dict[str, str], header_name: str, value: Optional[bool] - ) -> None: + def _set_bool_header(self, headers: Dict[str, str], header_name: str, value: Optional[bool]) -> None: """Apply a boolean value as a lowercase string HTTP header when provided.""" if value is None: @@ -745,9 +695,7 @@ class PillarGuardrail(CustomGuardrail): ) return payload - async def _call_pillar_api( - self, headers: Dict[str, str], payload: Dict[str, Any] - ) -> Dict[str, Any]: + async def _call_pillar_api(self, headers: Dict[str, str], payload: Dict[str, Any]) -> Dict[str, Any]: """ Call the Pillar API and return the response. @@ -772,14 +720,10 @@ class PillarGuardrail(CustomGuardrail): flagged = res.get("flagged") session_id = res.get("session_id") - verbose_proxy_logger.debug( - f"Pillar Guardrail: Analysis complete - flagged={flagged}, session={session_id}" - ) + verbose_proxy_logger.debug(f"Pillar Guardrail: Analysis complete - flagged={flagged}, session={session_id}") return res - def _process_pillar_response( - self, pillar_response: Dict[str, Any], original_data: dict - ) -> None: + def _process_pillar_response(self, pillar_response: Dict[str, Any], original_data: dict) -> None: """ Process the Pillar API response and handle detections based on configuration. @@ -796,25 +740,19 @@ class PillarGuardrail(CustomGuardrail): flagged = pillar_response.get("flagged", False) metadata_field = get_metadata_variable_name_from_kwargs(original_data) - if metadata_field not in original_data or not isinstance( - original_data.get(metadata_field), dict - ): + if metadata_field not in original_data or not isinstance(original_data.get(metadata_field), dict): original_data[metadata_field] = {} metadata_store = original_data[metadata_field] # Backwards compatibility - ensure metadata alias exists when different key used if metadata_field != "metadata": - if "metadata" not in original_data or not isinstance( - original_data.get("metadata"), dict - ): + if "metadata" not in original_data or not isinstance(original_data.get("metadata"), dict): original_data["metadata"] = metadata_store # Store session_id from Pillar response for potential reuse pillar_session_id = pillar_response.get("session_id") if pillar_session_id: - verbose_proxy_logger.debug( - f"Pillar Guardrail: Received session_id from server: {pillar_session_id}" - ) + verbose_proxy_logger.debug(f"Pillar Guardrail: Received session_id from server: {pillar_session_id}") # Store in request metadata for use in subsequent hooks if "pillar_session_id" not in metadata_store: metadata_store["pillar_session_id"] = pillar_session_id @@ -832,9 +770,7 @@ class PillarGuardrail(CustomGuardrail): if self.on_flagged_action == "block": self._raise_pillar_detection_exception(pillar_response) elif self.on_flagged_action == "mask": - verbose_proxy_logger.info( - "Pillar Guardrail: Masking mode - masking flagged content" - ) + verbose_proxy_logger.info("Pillar Guardrail: Masking mode - masking flagged content") masked_messages = pillar_response.get("masked_session_messages", []) if masked_messages: original_data["messages"] = masked_messages @@ -843,15 +779,11 @@ class PillarGuardrail(CustomGuardrail): "Pillar Guardrail: Masking requested but no masked_session_messages in response" ) elif self.on_flagged_action == "monitor": - verbose_proxy_logger.info( - "Pillar Guardrail: Monitoring mode - allowing flagged content to proceed" - ) + verbose_proxy_logger.info("Pillar Guardrail: Monitoring mode - allowing flagged content to proceed") build_pillar_response_headers(metadata_store) - def _raise_pillar_detection_exception( - self, pillar_response: Dict[str, Any] - ) -> None: + def _raise_pillar_detection_exception(self, pillar_response: Dict[str, Any]) -> None: """ Raise an HTTPException for Pillar security detections. @@ -877,9 +809,7 @@ class PillarGuardrail(CustomGuardrail): "pillar_response": pillar_response_dict, } - verbose_proxy_logger.warning( - "Pillar Guardrail: Request blocked - Security threats detected" - ) + verbose_proxy_logger.warning("Pillar Guardrail: Request blocked - Security threats detected") raise HTTPException(status_code=400, detail=error_detail) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 1367f95c5c8..95876a55eab 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -78,13 +78,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): apply_to_output: bool = False, presidio_ad_hoc_recognizers: Optional[str] = None, logging_only: Optional[bool] = None, - pii_entities_config: Optional[ - Dict[Union[PiiEntityType, str], PiiAction] - ] = None, + pii_entities_config: Optional[Dict[Union[PiiEntityType, str], PiiAction]] = None, presidio_language: Optional[str] = None, - presidio_score_thresholds: Optional[ - Dict[Union[PiiEntityType, str], float] - ] = None, + presidio_score_thresholds: Optional[Dict[Union[PiiEntityType, str], float]] = None, presidio_entities_deny_list: Optional[List[Union[PiiEntityType, str]]] = None, **kwargs, ): @@ -104,22 +100,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if (self.output_parse_pii or self.apply_to_output) and not logging_only: current_hook = self.event_hook if isinstance(current_hook, str) and current_hook != "post_call": - self.event_hook = cast( - List[GuardrailEventHooks], [current_hook, "post_call"] - ) + self.event_hook = cast(List[GuardrailEventHooks], [current_hook, "post_call"]) elif isinstance(current_hook, list) and "post_call" not in current_hook: - self.event_hook = cast( - List[GuardrailEventHooks], current_hook + ["post_call"] - ) - self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = ( - pii_entities_config or {} - ) - self.presidio_score_thresholds: Dict[Union[PiiEntityType, str], float] = ( - presidio_score_thresholds or {} - ) - self.presidio_entities_deny_list: List[Union[PiiEntityType, str]] = ( - presidio_entities_deny_list or [] - ) + self.event_hook = cast(List[GuardrailEventHooks], current_hook + ["post_call"]) + self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = pii_entities_config or {} + self.presidio_score_thresholds: Dict[Union[PiiEntityType, str], float] = presidio_score_thresholds or {} + self.presidio_entities_deny_list: List[Union[PiiEntityType, str]] = presidio_entities_deny_list or [] self.presidio_language = presidio_language or "en" # Shared HTTP session to prevent memory leaks (issue #14540) self._http_session: Optional[aiohttp.ClientSession] = None @@ -145,13 +131,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except FileNotFoundError: raise Exception(f"File not found. file_path={ad_hoc_recognizers}") except json.JSONDecodeError as e: - raise Exception( - f"Error decoding JSON file: {str(e)}, file_path={ad_hoc_recognizers}" - ) + raise Exception(f"Error decoding JSON file: {str(e)}, file_path={ad_hoc_recognizers}") except Exception as e: - raise Exception( - f"An error occurred: {str(e)}, file_path={ad_hoc_recognizers}" - ) + raise Exception(f"An error occurred: {str(e)}, file_path={ad_hoc_recognizers}") self.validate_environment( presidio_analyzer_api_base=presidio_analyzer_api_base, presidio_anonymizer_api_base=presidio_anonymizer_api_base, @@ -162,12 +144,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): presidio_analyzer_api_base: Optional[str] = None, presidio_anonymizer_api_base: Optional[str] = None, ): - self.presidio_analyzer_api_base: Optional[str] = ( - presidio_analyzer_api_base or get_secret("PRESIDIO_ANALYZER_API_BASE", None) + self.presidio_analyzer_api_base: Optional[str] = presidio_analyzer_api_base or get_secret( + "PRESIDIO_ANALYZER_API_BASE", None ) # type: ignore - self.presidio_anonymizer_api_base: Optional[str] = ( - presidio_anonymizer_api_base - or litellm.get_secret("PRESIDIO_ANONYMIZER_API_BASE", None) + self.presidio_anonymizer_api_base: Optional[str] = presidio_anonymizer_api_base or litellm.get_secret( + "PRESIDIO_ANONYMIZER_API_BASE", None ) # type: ignore if self.presidio_analyzer_api_base is None: @@ -179,9 +160,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): or self.presidio_analyzer_api_base.startswith("https://") ): # add http:// if unset, assume communicating over private network - e.g. render - self.presidio_analyzer_api_base = ( - "http://" + self.presidio_analyzer_api_base - ) + self.presidio_analyzer_api_base = "http://" + self.presidio_analyzer_api_base if self.presidio_anonymizer_api_base is None: raise Exception("Missing `PRESIDIO_ANONYMIZER_API_BASE` from environment") @@ -192,9 +171,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): or self.presidio_anonymizer_api_base.startswith("https://") ): # add http:// if unset, assume communicating over private network - e.g. render - self.presidio_anonymizer_api_base = ( - "http://" + self.presidio_anonymizer_api_base - ) + self.presidio_anonymizer_api_base = "http://" + self.presidio_anonymizer_api_base @asynccontextmanager async def _get_session_iterator( @@ -221,10 +198,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Background thread/loop -> use loop-bound session cache # This avoids "attached to a different loop" or "no running event loop" errors # when accessing the shared session created in the main loop - if ( - current_loop not in self._loop_sessions - or self._loop_sessions[current_loop].closed - ): + if current_loop not in self._loop_sessions or self._loop_sessions[current_loop].closed: self._loop_sessions[current_loop] = aiohttp.ClientSession() yield self._loop_sessions[current_loop] @@ -247,9 +221,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """Return True if pii_entities_config has any BLOCK action (fail-closed on analyzer errors).""" if not self.pii_entities_config: return False - return any( - action == PiiAction.BLOCK for action in self.pii_entities_config.values() - ) + return any(action == PiiAction.BLOCK for action in self.pii_entities_config.values()) def _get_presidio_analyze_request_payload( self, @@ -284,9 +256,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): analyze_payload["language"] = presidio_config.language casted_analyze_payload: dict = cast(dict, analyze_payload) - casted_analyze_payload.update( - self.get_guardrail_dynamic_request_body_params(request_data=request_data) - ) + casted_analyze_payload.update(self.get_guardrail_dynamic_request_body_params(request_data=request_data)) return cast(PresidioAnalyzeRequest, casted_analyze_payload) async def analyze_text( @@ -302,9 +272,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Skip empty or whitespace-only text to avoid Presidio errors # Common in tool/function calling where assistant content is empty if not text or len(text.strip()) == 0: - verbose_proxy_logger.debug( - "Skipping Presidio analysis for empty/whitespace-only text" - ) + verbose_proxy_logger.debug("Skipping Presidio analysis for empty/whitespace-only text") return [] if self.mock_redacted_text is not None: @@ -315,12 +283,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Make the request to /analyze analyze_url = f"{self.presidio_analyzer_api_base}analyze" - analyze_payload: PresidioAnalyzeRequest = ( - self._get_presidio_analyze_request_payload( - text=text, - presidio_config=presidio_config, - request_data=request_data, - ) + analyze_payload: PresidioAnalyzeRequest = self._get_presidio_analyze_request_payload( + text=text, + presidio_config=presidio_config, + request_data=request_data, ) verbose_proxy_logger.debug( @@ -332,20 +298,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): def _fail_on_invalid_response( reason: str, ) -> List[PresidioAnalyzeResponseItem]: - should_fail_closed = ( - bool(self.pii_entities_config) - or self.output_parse_pii - or self.apply_to_output - ) + should_fail_closed = bool(self.pii_entities_config) or self.output_parse_pii or self.apply_to_output if should_fail_closed: raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=f"Presidio analyzer returned invalid response; cannot verify PII when PII protection is configured: {reason}", should_wrap_with_default_message=False, ) - verbose_proxy_logger.warning( - "Presidio analyzer %s, returning empty list", reason - ) + verbose_proxy_logger.warning("Presidio analyzer %s, returning empty list", reason) return [] async with session.post( @@ -380,9 +340,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if isinstance(analyze_results, dict): if "error" in analyze_results: - return _fail_on_invalid_response( - f"error: {analyze_results.get('error')}" - ) + return _fail_on_invalid_response(f"error: {analyze_results.get('error')}") # If it's a dict but not an error, try to process it as a single item verbose_proxy_logger.debug( "Presidio returned dict (not list), attempting to process as single item" @@ -390,9 +348,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): try: return [PresidioAnalyzeResponseItem(**analyze_results)] except Exception as e: - return _fail_on_invalid_response( - f"failed to parse dict response: {e}" - ) + return _fail_on_invalid_response(f"failed to parse dict response: {e}") # Handle unexpected types (str, None, etc.) - e.g. from malformed/error if not isinstance(analyze_results, list): @@ -445,9 +401,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) as response: if response.status >= 400: error_body = await response.text() - raise Exception( - f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}" - ) + raise Exception(f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}") content_type = getattr( response, "content_type", @@ -472,9 +426,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for item in redacted_text.get("items", []): entity_type = item.get("entity_type", None) if entity_type is not None: - masked_entity_count[entity_type] = ( - masked_entity_count.get(entity_type, 0) + 1 - ) + masked_entity_count[entity_type] = masked_entity_count.get(entity_type, 0) + 1 return redacted_text["text"] def _finalize_presidio_anonymize_numbered_tokens( @@ -524,9 +476,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): replacement = f"{replacement}_{seq}" pii_tokens[replacement] = text[start:end] new_text = new_text[:start] + replacement + new_text[end:] - masked_entity_count[entity_type] = ( - masked_entity_count.get(entity_type, 0) + 1 - ) + masked_entity_count[entity_type] = masked_entity_count.get(entity_type, 0) + 1 return new_text async def anonymize_text( @@ -552,9 +502,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): verbose_proxy_logger.debug("redacted_text: %s", redacted_text) if not output_parse_pii: - return self._finalize_presidio_anonymize_simple( - redacted_text, masked_entity_count - ) + return self._finalize_presidio_anonymize_simple(redacted_text, masked_entity_count) return self._finalize_presidio_anonymize_numbered_tokens( text, analyze_results, request_data, masked_entity_count @@ -563,14 +511,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # Sanitize exception to avoid leaking the original text (which may # contain API keys or other secrets) in error responses. error_str = str(e) - if ( - "Invalid anonymizer response" in error_str - or "Presidio anonymizer returned" in error_str - ): + if "Invalid anonymizer response" in error_str or "Presidio anonymizer returned" in error_str: raise - raise Exception( - f"Presidio PII anonymization failed: {type(e).__name__}" - ) from e + raise Exception(f"Presidio PII anonymization failed: {type(e).__name__}") from e def filter_analyze_results_by_score( self, analyze_results: Union[List[PresidioAnalyzeResponseItem], Dict] @@ -586,16 +529,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return analyze_results filtered_results: List[PresidioAnalyzeResponseItem] = [] - deny_list_strings = [ - getattr(x, "value", str(x)) for x in self.presidio_entities_deny_list - ] + deny_list_strings = [getattr(x, "value", str(x)) for x in self.presidio_entities_deny_list] for item in analyze_results: entity_type = item.get("entity_type") str_entity_type = str( - getattr(entity_type, "value", entity_type) - if entity_type is not None - else entity_type + getattr(entity_type, "value", entity_type) if entity_type is not None else entity_type ) if entity_type and str_entity_type in deny_list_strings: continue @@ -635,10 +574,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if entity_type: # Check if entity_type is in config (supports both enum and string) - if ( - entity_type in self.pii_entities_config - and self.pii_entities_config[entity_type] == PiiAction.BLOCK - ): + if entity_type in self.pii_entities_config and self.pii_entities_config[entity_type] == PiiAction.BLOCK: raise BlockedPiiEntityError( entity_type=entity_type, guardrail_name=self.guardrail_name, @@ -673,16 +609,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): verbose_proxy_logger.debug("analyze_results: %s", analyze_results) # Apply score threshold filtering if configured - analyze_results = self.filter_analyze_results_by_score( - analyze_results=analyze_results - ) + analyze_results = self.filter_analyze_results_by_score(analyze_results=analyze_results) #################################################### # Blocked Entities check #################################################### - self.raise_exception_if_blocked_entities_detected( - analyze_results=analyze_results - ) + self.raise_exception_if_blocked_entities_detected(analyze_results=analyze_results) # Then anonymize the text using the analysis results anonymized_text = await self.anonymize_text( @@ -757,9 +689,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if messages is None: return data tasks = [] - task_mappings: List[ - Tuple[int, Optional[int]] - ] = [] # Track (message_index, content_index) for each task + task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each task for msg_idx, m in enumerate(messages): content = m.get("content", None) @@ -774,9 +704,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data=data, ) ) - task_mappings.append( - (msg_idx, None) - ) # None indicates string content + task_mappings.append((msg_idx, None)) # None indicates string content elif isinstance(content, list): for content_idx, c in enumerate(content): text_str = c.get("text", None) @@ -803,23 +731,17 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if content is None: continue if isinstance(content, str) and content_idx_optional is None: - messages[msg_idx]["content"] = ( - r # replace content with redacted string - ) + messages[msg_idx]["content"] = r # replace content with redacted string elif isinstance(content, list) and content_idx_optional is not None: messages[msg_idx]["content"][content_idx_optional]["text"] = r - verbose_proxy_logger.debug( - f"Presidio PII Masking: Redacted pii message: {data['messages']}" - ) + verbose_proxy_logger.debug(f"Presidio PII Masking: Redacted pii message: {data['messages']}") data["messages"] = messages return data except Exception as e: raise e - def logging_hook( - self, kwargs: dict, result: Any, call_type: str - ) -> Tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: from concurrent.futures import ThreadPoolExecutor def run_in_new_loop(): @@ -828,9 +750,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): try: asyncio.set_event_loop(new_loop) return new_loop.run_until_complete( - self.async_logging_hook( - kwargs=kwargs, result=result, call_type=call_type - ) + self.async_logging_hook(kwargs=kwargs, result=result, call_type=call_type) ) finally: new_loop.close() @@ -849,20 +769,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # No running event loop, we can safely run in this thread return run_in_new_loop() - async def async_logging_hook( - self, kwargs: dict, result: Any, call_type: str - ) -> Tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: """ Masks the input before logging to langfuse, datadog, etc. """ - if ( - call_type == "completion" or call_type == "acompletion" - ): # /chat/completions requests + if call_type == "completion" or call_type == "acompletion": # /chat/completions requests messages: Optional[List] = kwargs.get("messages", None) tasks = [] - task_mappings: List[ - Tuple[int, Optional[int]] - ] = [] # Track (message_index, content_index) for each task + task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each task if messages is None: return kwargs, result @@ -882,9 +796,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data=kwargs, ) ) # need to pass separately b/c presidio has context window limits - task_mappings.append( - (msg_idx, None) - ) # None indicates string content + task_mappings.append((msg_idx, None)) # None indicates string content elif isinstance(content, list): for content_idx, c in enumerate(content): text_str = c.get("text", None) @@ -911,15 +823,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if content is None: continue if isinstance(content, str) and content_idx_optional is None: - messages[msg_idx]["content"] = ( - r # replace content with redacted string - ) + messages[msg_idx]["content"] = r # replace content with redacted string elif isinstance(content, list) and content_idx_optional is not None: messages[msg_idx]["content"][content_idx_optional]["text"] = r - verbose_proxy_logger.debug( - f"Presidio PII Masking: Redacted pii message: {messages}" - ) + verbose_proxy_logger.debug(f"Presidio PII Masking: Redacted pii message: {messages}") kwargs["messages"] = messages return kwargs, result @@ -942,9 +850,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return await self._process_anthropic_response_for_pii( response=cast(dict, response), request_data=data, mode="mask" ) - return await self._mask_output_response( - response=response, request_data=data - ) + return await self._mask_output_response(response=response, request_data=data) if self.output_parse_pii is False and litellm.output_parse_pii is False: return response @@ -1011,12 +917,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): metadata = (request_data.get("metadata") or {}) if request_data else {} pii_tokens = metadata.get("pii_tokens", {}) if not pii_tokens and mode == "unmask": - verbose_proxy_logger.debug( - "No pii_tokens in metadata for Anthropic response unmask" - ) - presidio_config = self.get_presidio_settings_from_request_data( - request_data or {} - ) + verbose_proxy_logger.debug("No pii_tokens in metadata for Anthropic response unmask") + presidio_config = self.get_presidio_settings_from_request_data(request_data or {}) content = response.get("content") if not isinstance(content, list): @@ -1053,12 +955,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): metadata = (request_data.get("metadata") or {}) if request_data else {} pii_tokens = metadata.get("pii_tokens", {}) if not pii_tokens and mode == "unmask": - verbose_proxy_logger.debug( - "No pii_tokens found in request_data['metadata'] — nothing to unmask" - ) - presidio_config = self.get_presidio_settings_from_request_data( - request_data or {} - ) + verbose_proxy_logger.debug("No pii_tokens found in request_data['metadata'] — nothing to unmask") + presidio_config = self.get_presidio_settings_from_request_data(request_data or {}) for choice in response.choices: message = getattr(choice, "message", None) @@ -1103,9 +1001,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): args = function.arguments if isinstance(args, str): if mode == "unmask": - function.arguments = self._unmask_pii_text( - args, pii_tokens - ) + function.arguments = self._unmask_pii_text(args, pii_tokens) elif mode == "mask": function.arguments = await self.check_pii( text=args, @@ -1120,9 +1016,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): args = function_call.arguments if isinstance(args, str): if mode == "unmask": - function_call.arguments = self._unmask_pii_text( - args, pii_tokens - ) + function_call.arguments = self._unmask_pii_text(args, pii_tokens) elif mode == "mask": function_call.arguments = await self.check_pii( text=args, @@ -1208,9 +1102,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return - assembled_model_response = stream_chunk_builder( - chunks=all_chunks, messages=request_data.get("messages") - ) + assembled_model_response = stream_chunk_builder(chunks=all_chunks, messages=request_data.get("messages")) if not isinstance(assembled_model_response, ModelResponse): for chunk in all_chunks: @@ -1223,9 +1115,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): mode="mask", ) - mock_response_stream = convert_model_response_to_streaming( - assembled_model_response - ) + mock_response_stream = convert_model_response_to_streaming(assembled_model_response) yield mock_response_stream except Exception as e: @@ -1254,9 +1144,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): and delta.get("type") == "text_delta" and isinstance(delta.get("text"), str) ): - unmasked = _OPTIONAL_PresidioPIIMasking._unmask_pii_text( - delta["text"], pii_tokens - ) + unmasked = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(delta["text"], pii_tokens) if unmasked != delta["text"]: event["delta"]["text"] = unmasked line = "data: " + json.dumps(event, ensure_ascii=False) @@ -1266,9 +1154,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return "\n".join(result_lines).encode("utf-8") - def _unmask_responses_api_completed_chunk( - self, chunk: Any, pii_tokens: Dict[str, str] - ) -> None: + def _unmask_responses_api_completed_chunk(self, chunk: Any, pii_tokens: Dict[str, str]) -> None: """ Unmask PII tokens in-place for a ``response.completed`` Responses API event. @@ -1287,15 +1173,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for content_block in content: if isinstance(content_block, dict): if isinstance(content_block.get("text"), str): - content_block["text"] = self._unmask_pii_text( - content_block["text"], pii_tokens - ) - elif hasattr(content_block, "text") and isinstance( - content_block.text, str - ): - content_block.text = self._unmask_pii_text( - content_block.text, pii_tokens - ) + content_block["text"] = self._unmask_pii_text(content_block["text"], pii_tokens) + elif hasattr(content_block, "text") and isinstance(content_block.text, str): + content_block.text = self._unmask_pii_text(content_block.text, pii_tokens) async def _stream_pii_unmasking( self, @@ -1356,9 +1236,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk return - self._preserve_usage_from_last_chunk( - assembled_model_response, remaining_chunks - ) + self._preserve_usage_from_last_chunk(assembled_model_response, remaining_chunks) await self._process_response_for_pii( response=assembled_model_response, @@ -1366,9 +1244,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): mode="unmask", ) - mock_response_stream = convert_model_response_to_streaming( - assembled_model_response - ) + mock_response_stream = convert_model_response_to_streaming(assembled_model_response) yield mock_response_stream except Exception as e: @@ -1390,18 +1266,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): The base class declares ModelResponseStream only. """ if self.apply_to_output: - async for chunk in self._stream_apply_output_masking( - response, request_data - ): + async for chunk in self._stream_apply_output_masking(response, request_data): yield chunk return metadata = (request_data.get("metadata") or {}) if request_data else {} pii_tokens = metadata.get("pii_tokens", {}) if not pii_tokens and request_data: - verbose_proxy_logger.debug( - "No pii_tokens in request_data['metadata'] for streaming unmask path" - ) + verbose_proxy_logger.debug("No pii_tokens in request_data['metadata'] for streaming unmask path") if not (self.output_parse_pii and pii_tokens): async for chunk in response: yield chunk @@ -1421,9 +1293,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if last_chunk_usage: setattr(assembled_model_response, "usage", last_chunk_usage) - def get_presidio_settings_from_request_data( - self, data: dict - ) -> Optional[PresidioPerRequestConfig]: + def get_presidio_settings_from_request_data(self, data: dict) -> Optional[PresidioPerRequestConfig]: if "metadata" in data: _metadata = data.get("metadata", None) if _metadata is None: @@ -1489,6 +1359,4 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if litellm_params.presidio_score_thresholds: self.presidio_score_thresholds = litellm_params.presidio_score_thresholds if litellm_params.presidio_entities_deny_list: - self.presidio_entities_deny_list = ( - litellm_params.presidio_entities_deny_list - ) + self.presidio_entities_deny_list = litellm_params.presidio_entities_deny_list diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index b3e761869b0..23f349e0b18 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -35,23 +35,17 @@ class PromptSecurityGuardrail(CustomGuardrail): check_tool_results: Optional[bool] = None, **kwargs, ): - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.api_key = api_key or os.environ.get("PROMPT_SECURITY_API_KEY") self.api_base = api_base or os.environ.get("PROMPT_SECURITY_API_BASE") self.user = user or os.environ.get("PROMPT_SECURITY_USER") - self.system_prompt = system_prompt or os.environ.get( - "PROMPT_SECURITY_SYSTEM_PROMPT" - ) + self.system_prompt = system_prompt or os.environ.get("PROMPT_SECURITY_SYSTEM_PROMPT") # Configure whether to check tool/function results for indirect prompt injection # Default: False (Filter out tool/function messages) # True: Transform to "other" role and send to API if check_tool_results is None: - check_tool_results_env = os.environ.get( - "PROMPT_SECURITY_CHECK_TOOL_RESULTS", "false" - ).lower() + check_tool_results_env = os.environ.get("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "false").lower() self.check_tool_results = check_tool_results_env in ("true", "1", "yes") else: self.check_tool_results = check_tool_results @@ -158,9 +152,7 @@ class PromptSecurityGuardrail(CustomGuardrail): messages = [{"role": "user", "content": text} for text in texts] # Process any embedded files/images in messages - messages = await self.process_message_files( - messages, user_api_key_alias=user_api_key_alias - ) + messages = await self.process_message_files(messages, user_api_key_alias=user_api_key_alias) # Also process standalone images from inputs if images: @@ -170,9 +162,7 @@ class PromptSecurityGuardrail(CustomGuardrail): filtered_messages = self.filter_messages_by_role(messages) if not filtered_messages: - verbose_proxy_logger.debug( - "Prompt Security Guardrail: No messages to check after filtering" - ) + verbose_proxy_logger.debug("Prompt Security Guardrail: No messages to check after filtering") return inputs # Call Prompt Security API @@ -214,8 +204,7 @@ class PromptSecurityGuardrail(CustomGuardrail): if action == "block": raise HTTPException( status_code=400, - detail="Blocked by Prompt Security, Violations: " - + ", ".join(violations), + detail="Blocked by Prompt Security, Violations: " + ", ".join(violations), ) elif action == "modify": # Extract modified texts from modified_messages @@ -277,8 +266,7 @@ class PromptSecurityGuardrail(CustomGuardrail): if action == "block": raise HTTPException( status_code=400, - detail="Blocked by Prompt Security, Violations: " - + ", ".join(violations), + detail="Blocked by Prompt Security, Violations: " + ", ".join(violations), ) elif action == "modify": modified_text = result.get("modified_text") @@ -304,9 +292,7 @@ class PromptSecurityGuardrail(CustomGuardrail): texts.append(text) return texts - async def _process_standalone_images( - self, images: List[str], user_api_key_alias: Optional[str] - ) -> None: + async def _process_standalone_images(self, images: List[str], user_api_key_alias: Optional[str]) -> None: """Process standalone images from inputs (data URLs).""" for image_url in images: if image_url.startswith("data:"): @@ -390,13 +376,9 @@ class PromptSecurityGuardrail(CustomGuardrail): ) if not job_id: - raise HTTPException( - status_code=500, detail="Failed to get jobId from Prompt Security" - ) + raise HTTPException(status_code=500, detail="Failed to get jobId from Prompt Security") - verbose_proxy_logger.debug( - "Prompt Security Guardrail: File sanitization started with jobId=%s", job_id - ) + verbose_proxy_logger.debug("Prompt Security Guardrail: File sanitization started with jobId=%s", job_id) # Step 2: Poll for results for attempt in range(self.max_poll_attempts): @@ -443,22 +425,14 @@ class PromptSecurityGuardrail(CustomGuardrail): ) continue else: - raise HTTPException( - status_code=500, detail=f"Unexpected sanitization status: {status}" - ) + raise HTTPException(status_code=500, detail=f"Unexpected sanitization status: {status}") raise HTTPException(status_code=408, detail="File sanitization timeout") - async def _process_image_url_item( - self, item: dict, user_api_key_alias: Optional[str] - ) -> dict: + async def _process_image_url_item(self, item: dict, user_api_key_alias: Optional[str]) -> dict: """Process and sanitize image_url items.""" image_url_data = item.get("image_url", {}) - url = ( - image_url_data.get("url", "") - if isinstance(image_url_data, dict) - else image_url_data - ) + url = image_url_data.get("url", "") if isinstance(image_url_data, dict) else image_url_data if not url.startswith("data:"): return item @@ -485,30 +459,22 @@ class PromptSecurityGuardrail(CustomGuardrail): if action == "modify": sanitized_content = sanitization_result.get("content", "") if sanitized_content: - sanitized_encoded = base64.b64encode( - sanitized_content.encode() - ).decode() + sanitized_encoded = base64.b64encode(sanitized_content.encode()).decode() sanitized_url = f"{header},{sanitized_encoded}" if isinstance(image_url_data, dict): image_url_data["url"] = sanitized_url else: item["image_url"] = sanitized_url - verbose_proxy_logger.info( - "File content modified by Prompt Security" - ) + verbose_proxy_logger.info("File content modified by Prompt Security") return item except HTTPException: raise except Exception as e: verbose_proxy_logger.error(f"Error sanitizing image file: {str(e)}") - raise HTTPException( - status_code=500, detail=f"File sanitization failed: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"File sanitization failed: {str(e)}") - async def _process_document_item( - self, item: dict, user_api_key_alias: Optional[str] - ) -> dict: + async def _process_document_item(self, item: dict, user_api_key_alias: Optional[str]) -> dict: """Process and sanitize document/file items.""" doc_data = item.get("document") or item.get("file") or item @@ -531,9 +497,7 @@ class PromptSecurityGuardrail(CustomGuardrail): else: file_data = base64.b64decode(doc_content) mime_type = ( - doc_data.get("mime_type", "application/pdf") - if isinstance(doc_data, dict) - else "application/pdf" + doc_data.get("mime_type", "application/pdf") if isinstance(doc_data, dict) else "application/pdf" ) if "pdf" in mime_type: @@ -564,9 +528,7 @@ class PromptSecurityGuardrail(CustomGuardrail): sanitized_content = sanitization_result.get("content", "") if sanitized_content: sanitized_encoded = base64.b64encode( - sanitized_content - if isinstance(sanitized_content, bytes) - else sanitized_content.encode() + sanitized_content if isinstance(sanitized_content, bytes) else sanitized_content.encode() ).decode() if url.startswith("data:") and header: @@ -576,22 +538,16 @@ class PromptSecurityGuardrail(CustomGuardrail): elif isinstance(doc_data, dict): doc_data["data"] = sanitized_encoded - verbose_proxy_logger.info( - "Document content modified by Prompt Security" - ) + verbose_proxy_logger.info("Document content modified by Prompt Security") return item except HTTPException: raise except Exception as e: verbose_proxy_logger.error(f"Error sanitizing document: {str(e)}") - raise HTTPException( - status_code=500, detail=f"Document sanitization failed: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Document sanitization failed: {str(e)}") - async def process_message_files( - self, messages: list, user_api_key_alias: Optional[str] = None - ) -> list: + async def process_message_files(self, messages: list, user_api_key_alias: Optional[str] = None) -> list: """Process messages and sanitize any file content (images, documents, PDFs, etc.).""" processed_messages = [] @@ -607,13 +563,9 @@ class PromptSecurityGuardrail(CustomGuardrail): if isinstance(item, dict): item_type = item.get("type") if item_type == "image_url": - item = await self._process_image_url_item( - item, user_api_key_alias - ) + item = await self._process_image_url_item(item, user_api_key_alias) elif item_type in ["document", "file"]: - item = await self._process_document_item( - item, user_api_key_alias - ) + item = await self._process_document_item(item, user_api_key_alias) processed_content.append(item) @@ -645,11 +597,7 @@ class PromptSecurityGuardrail(CustomGuardrail): if self.check_tool_results: transformed_message = { "role": "other", - **{ - key: value - for key, value in message.items() - if key != "role" - }, + **{key: value for key, value in message.items() if key != "role"}, } filtered_messages.append(transformed_message) transformed_count += 1 @@ -688,10 +636,7 @@ class PromptSecurityGuardrail(CustomGuardrail): @staticmethod def _redact_headers(headers: dict) -> dict: - return { - name: ("REDACTED" if name.lower() == "app-id" else value) - for name, value in headers.items() - } + return {name: ("REDACTED" if name.lower() == "app-id" else value) for name, value in headers.items()} def _log_api_request( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index d9c4ecb61ae..7012b0d8d5d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -65,9 +65,7 @@ class PromptGuardGuardrail(CustomGuardrail): "the guardrail config." ) - self.api_base = ( - api_base or os.environ.get("PROMPTGUARD_API_BASE") or _DEFAULT_API_BASE - ).rstrip("/") + self.api_base = (api_base or os.environ.get("PROMPTGUARD_API_BASE") or _DEFAULT_API_BASE).rstrip("/") if block_on_error is None: env = os.environ.get("PROMPTGUARD_BLOCK_ON_ERROR", "true") @@ -175,12 +173,7 @@ class PromptGuardGuardrail(CustomGuardrail): confidence = result.get("confidence", 0.0) raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message=( - f"Blocked by PromptGuard: " - f"{threat_type} " - f"(confidence={confidence}, " - f"event_id={event_id})" - ), + message=(f"Blocked by PromptGuard: {threat_type} (confidence={confidence}, event_id={event_id})"), ) if decision == "redact": diff --git a/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py b/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py index a1bab6dbac9..6d1c14a934b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py @@ -26,9 +26,7 @@ class QostodianNexus(GenericGuardrailAPI): api_base: Optional[str] = None, **kwargs, ): - api_base = api_base or os.environ.get( - "QOSTODIAN_NEXUS_API_BASE", "http://nexus:8800" - ) + api_base = api_base or os.environ.get("QOSTODIAN_NEXUS_API_BASE", "http://nexus:8800") kwargs["guardrail_name"] = kwargs.get("guardrail_name", GUARDRAIL_NAME) @@ -41,9 +39,7 @@ class QostodianNexus(GenericGuardrailAPI): ] existing = kwargs.get("extra_headers") or [] - kwargs["extra_headers"] = nexus_headers + [ - h for h in existing if h not in nexus_headers - ] + kwargs["extra_headers"] = nexus_headers + [h for h in existing if h not in nexus_headers] super().__init__( api_base=api_base, diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py index b9f7aed4a26..8c29cfcd309 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py @@ -19,12 +19,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" hallucinations_check=getattr(litellm_params, "hallucinations_check", None), grounding_check=getattr(litellm_params, "grounding_check", None), pii_check=getattr(litellm_params, "pii_check", None), - content_moderation_check=getattr( - litellm_params, "content_moderation_check", None - ), - tool_selection_quality_check=getattr( - litellm_params, "tool_selection_quality_check", None - ), + content_moderation_check=getattr(litellm_params, "content_moderation_check", None), + tool_selection_quality_check=getattr(litellm_params, "tool_selection_quality_check", None), assertions=getattr(litellm_params, "assertions", None), on_flagged=getattr(litellm_params, "on_flagged", "block"), guardrail_name=guardrail.get("guardrail_name", ""), diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index 6486da7f714..54f47cfbfa1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -62,11 +62,7 @@ class QualifireGuardrail(CustomGuardrail): assertions: Custom assertions to validate against the output on_flagged: Action when content is flagged: "block" or "monitor" """ - self.qualifire_api_key = ( - api_key - or get_secret_str("QUALIFIRE_API_KEY") - or os.environ.get("QUALIFIRE_API_KEY") - ) + self.qualifire_api_key = api_key or get_secret_str("QUALIFIRE_API_KEY") or os.environ.get("QUALIFIRE_API_KEY") self.qualifire_api_base = ( api_base or get_secret_str("QUALIFIRE_BASE_URL") @@ -88,9 +84,7 @@ class QualifireGuardrail(CustomGuardrail): self.prompt_injections = True # Initialize async HTTP client for direct API calls - self.async_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.GuardrailCallback - ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) super().__init__(**kwargs) @@ -108,9 +102,7 @@ class QualifireGuardrail(CustomGuardrail): ] ) - def _convert_messages_to_api_format( - self, messages: List[AllMessageValues] - ) -> List[Dict[str, Any]]: + def _convert_messages_to_api_format(self, messages: List[AllMessageValues]) -> List[Dict[str, Any]]: """ Convert LiteLLM messages to Qualifire API format. Supports tool calls for tool_selection_quality_check. @@ -176,9 +168,7 @@ class QualifireGuardrail(CustomGuardrail): return api_messages - def _convert_tools_to_api_format( - self, tools: Optional[List[Any]] - ) -> Optional[List[Dict[str, Any]]]: + def _convert_tools_to_api_format(self, tools: Optional[List[Any]]) -> Optional[List[Dict[str, Any]]]: """ Convert OpenAI-format tools to Qualifire API format. @@ -422,9 +412,7 @@ class QualifireGuardrail(CustomGuardrail): HTTPException: If content is blocked """ # Get dynamic params from request body (allows runtime overrides) - dynamic_params = self.get_guardrail_dynamic_request_body_params( - request_data=request_data - ) + dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data) # Extract messages from structured_messages or request_data messages: Optional[List[AllMessageValues]] = inputs.get("structured_messages") @@ -451,9 +439,7 @@ class QualifireGuardrail(CustomGuardrail): if texts: messages = [{"role": "user", "content": texts[-1] if texts else ""}] # type: ignore else: - verbose_proxy_logger.debug( - "Qualifire Guardrail: No messages or texts found, skipping" - ) + verbose_proxy_logger.debug("Qualifire Guardrail: No messages or texts found, skipping") return inputs # Get available tools from request_data for tool_selection_quality_check diff --git a/litellm/proxy/guardrails/guardrail_hooks/repelloai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/repelloai/__init__.py index 93c5221f111..9dc060ac9ae 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/repelloai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/repelloai/__init__.py @@ -22,9 +22,7 @@ def _event_hook_from_mode( return GuardrailEventHooks(mode) -def initialize_guardrail( - litellm_params: "LitellmParams", guardrail: "Guardrail" -) -> RepelloAIGuardrail: +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RepelloAIGuardrail: import litellm _repelloai_callback = RepelloAIGuardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py index a796fedbe27..f8971ceb405 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py @@ -107,9 +107,7 @@ class RepelloAIGuardrail(CustomGuardrail): for item in items: if isinstance(item, str) and item: texts.append(item) - remaining: list[object] = [ - v for k, v in current.items() if k not in _SCHEMA_EXTRACTED_KEYS - ] + remaining: list[object] = [v for k, v in current.items() if k not in _SCHEMA_EXTRACTED_KEYS] stack.extend(reversed(remaining)) elif _is_object_list(current): stack.extend(reversed(current)) @@ -142,17 +140,10 @@ class RepelloAIGuardrail(CustomGuardrail): asset_id: str | None = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", guardrail_name: str | None = None, - event_hook: ( - GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None - ) = None, + event_hook: (GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None) = None, default_on: bool = False, ): - self.repelloai_api_key = ( - api_key - or get_secret_str("ARGUS_API_KEY") - or get_secret_str("REPELLOAI_API_KEY") - or "" - ) + self.repelloai_api_key = api_key or get_secret_str("ARGUS_API_KEY") or get_secret_str("REPELLOAI_API_KEY") or "" if not self.repelloai_api_key: raise RepelloAIGuardrailMissingSecrets( "Couldn't get Repello API key. Set `ARGUS_API_KEY` in the environment " @@ -166,11 +157,7 @@ class RepelloAIGuardrail(CustomGuardrail): "dashboard and set `asset_id` on the guardrail in the config file." ) - self.api_base = ( - api_base - or get_secret_str("REPELLOAI_API_BASE") - or DEFAULT_REPELLOAI_API_BASE - ) + self.api_base = api_base or get_secret_str("REPELLOAI_API_BASE") or DEFAULT_REPELLOAI_API_BASE self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" ) @@ -214,9 +201,7 @@ class RepelloAIGuardrail(CustomGuardrail): self._raise_for_config_error(response) response.raise_for_status() try: - repelloai_response = TypeAdapter( - RepelloAIAnalyzeResponse - ).validate_json(response.text) + repelloai_response = TypeAdapter(RepelloAIAnalyzeResponse).validate_json(response.text) except ValidationError as e: raise HTTPException( status_code=500, @@ -225,17 +210,13 @@ class RepelloAIGuardrail(CustomGuardrail): "status_code": response.status_code, }, ) from e - verbose_proxy_logger.debug( - "RepelloAI Argus response: %s", repelloai_response - ) + verbose_proxy_logger.debug("RepelloAI Argus response: %s", repelloai_response) if self._verdict_blocks(repelloai_response): status = "guardrail_intervened" return repelloai_response except HTTPException as e: status = "guardrail_failed_to_respond" - guardrail_json_response = ( - str(e.detail) if not isinstance(e.detail, (dict, list)) else e.detail - ) # type: ignore[assignment] + guardrail_json_response = str(e.detail) if not isinstance(e.detail, (dict, list)) else e.detail # type: ignore[assignment] raise except HTTPError as e: status = "guardrail_failed_to_respond" @@ -244,9 +225,7 @@ class RepelloAIGuardrail(CustomGuardrail): except Exception as e: status = "guardrail_failed_to_respond" guardrail_json_response = str(e) - raise HTTPException( - status_code=500, detail={"error": "RepelloAI Argus guardrail failed"} - ) from e + raise HTTPException(status_code=500, detail={"error": "RepelloAI Argus guardrail failed"}) from e finally: end_time = datetime.now() if repelloai_response is not None: @@ -273,9 +252,7 @@ class RepelloAIGuardrail(CustomGuardrail): }, ) - def _verdict_blocks( - self, repelloai_response: RepelloAIAnalyzeResponse | None - ) -> bool: + def _verdict_blocks(self, repelloai_response: RepelloAIAnalyzeResponse | None) -> bool: if repelloai_response is None: return False verdict = repelloai_response.get("verdict") @@ -298,9 +275,7 @@ class RepelloAIGuardrail(CustomGuardrail): ) return None - def _raise_if_blocked( - self, repelloai_response: RepelloAIAnalyzeResponse | None - ) -> None: + def _raise_if_blocked(self, repelloai_response: RepelloAIAnalyzeResponse | None) -> None: if repelloai_response is None: return if self._verdict_blocks(repelloai_response): @@ -311,9 +286,7 @@ class RepelloAIGuardrail(CustomGuardrail): self._log_flagged_verdict(repelloai_response) @classmethod - def _format_blocked_detail( - cls, repelloai_response: RepelloAIAnalyzeResponse - ) -> str: + def _format_blocked_detail(cls, repelloai_response: RepelloAIAnalyzeResponse) -> str: policies = repelloai_response.get("policies_violated") if not isinstance(policies, list) or not policies: return "Blocked by RepelloAI Argus guardrail." @@ -348,11 +321,7 @@ class RepelloAIGuardrail(CustomGuardrail): @staticmethod def _extract_prompt_message_text(data: dict[str, object]) -> list[str]: messages = build_inspection_messages(data) - return [ - content - for message in messages - if isinstance(content := message.get("content"), str) and content - ] + return [content for message in messages if isinstance(content := message.get("content"), str) and content] @staticmethod def _extract_input_text_parts(content: object) -> list[str]: @@ -420,9 +389,7 @@ class RepelloAIGuardrail(CustomGuardrail): text = self._extract_prompt_text(data) if not text: - verbose_proxy_logger.warning( - "RepelloAI Argus: no inspectable prompt text in data - skipping." - ) + verbose_proxy_logger.warning("RepelloAI Argus: no inspectable prompt text in data - skipping.") return data repelloai_response = await self._call_analyze( @@ -433,9 +400,7 @@ class RepelloAIGuardrail(CustomGuardrail): ) self._raise_if_blocked(repelloai_response) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data async def async_post_call_success_hook( @@ -457,9 +422,7 @@ class RepelloAIGuardrail(CustomGuardrail): text = self._extract_response_text(response) if not text: - verbose_proxy_logger.warning( - "RepelloAI Argus: no inspectable response text - skipping." - ) + verbose_proxy_logger.warning("RepelloAI Argus: no inspectable response text - skipping.") return response repelloai_response = await self._call_analyze( @@ -470,9 +433,7 @@ class RepelloAIGuardrail(CustomGuardrail): ) self._raise_if_blocked(repelloai_response) - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response async def async_post_call_streaming_iterator_hook( @@ -501,11 +462,7 @@ class RepelloAIGuardrail(CustomGuardrail): assembled = litellm_main.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] chunks=chunks ) - text = ( - self._extract_response_text(assembled) - if isinstance(assembled, ModelResponse) - else None - ) + text = self._extract_response_text(assembled) if isinstance(assembled, ModelResponse) else None if text: repelloai_response = await self._call_analyze( text=text, @@ -519,9 +476,7 @@ class RepelloAIGuardrail(CustomGuardrail): from litellm.proxy.proxy_server import StreamingCallbackError raise StreamingCallbackError("Blocked by RepelloAI Argus guardrail") - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) else: verbose_proxy_logger.warning( "RepelloAI Argus: no inspectable text in streamed response; skipping scan. " @@ -553,9 +508,7 @@ class RepelloAIGuardrail(CustomGuardrail): return RepelloAIGuardrail._extract_responses_api_text(response_dict) @classmethod - def _extract_chat_completion_text( - cls, response_dict: dict[str, object] - ) -> str | None: + def _extract_chat_completion_text(cls, response_dict: dict[str, object]) -> str | None: choices = response_dict.get("choices") if not _is_object_list(choices): return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py index ab347130a30..4ad29bbeae8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py @@ -9,9 +9,7 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail( - litellm_params: "LitellmParams", guardrail: "Guardrail" -) -> RubrikLogger: +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RubrikLogger: import litellm rubrik_callback = RubrikLogger( diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/__init__.py index 124eefa7b5e..6c7eaf3ac24 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/__init__.py @@ -43,15 +43,13 @@ def initialize_guardrail( if llm_router is None: raise ValueError( - "SemanticGuard requires llm_router for embeddings. " - "Configure a model_list with an embedding model." + "SemanticGuard requires llm_router for embeddings. Configure a model_list with an embedding model." ) semantic_guardrail = SemanticGuardrail( guardrail_name=guardrail_name, llm_router=llm_router, - embedding_model=getattr(litellm_params, "embedding_model", None) - or DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL, + embedding_model=getattr(litellm_params, "embedding_model", None) or DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL, similarity_threshold=getattr(litellm_params, "similarity_threshold", None) or DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD, route_templates=getattr(litellm_params, "route_templates", None), diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py index 10a50c39e35..ad05e7656c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py @@ -52,18 +52,14 @@ class SemanticGuardRouteLoader: def load_custom_routes_file(file_path: str) -> List[Dict[str, Any]]: """Load custom routes from a YAML file.""" if not os.path.exists(file_path): - raise ValueError( - f"SemanticGuard: custom routes file not found: {file_path}" - ) + raise ValueError(f"SemanticGuard: custom routes file not found: {file_path}") with open(file_path, "r") as f: data = yaml.safe_load(f) if isinstance(data, list): return data if isinstance(data, dict): return [data] - raise ValueError( - f"SemanticGuard: invalid custom routes file format in {file_path}" - ) + raise ValueError(f"SemanticGuard: invalid custom routes file format in {file_path}") @classmethod def build_routes( diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index 21c01da029b..5840480f0da 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -52,9 +52,7 @@ class SemanticGuardrail(CustomGuardrail): custom_routes_file: Optional[str] = None, custom_routes: Optional[List[Dict[str, Any]]] = None, on_flagged_action: str = "block", - event_hook: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode] - ] = None, + event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = None, default_on: bool = False, **kwargs, ): @@ -83,18 +81,13 @@ class SemanticGuardrail(CustomGuardrail): ) if not routes: - raise ValueError( - "SemanticGuardrail: no routes configured. " - "Provide route_templates or custom_routes." - ) + raise ValueError("SemanticGuardrail: no routes configured. Provide route_templates or custom_routes.") - self.semantic_router: "SemanticRouter" = ( - SemanticGuardRouteLoader.build_semantic_router( - routes=routes, - litellm_router=llm_router, - embedding_model=embedding_model, - global_threshold=similarity_threshold, - ) + self.semantic_router: "SemanticRouter" = SemanticGuardRouteLoader.build_semantic_router( + routes=routes, + litellm_router=llm_router, + embedding_model=embedding_model, + global_threshold=similarity_threshold, ) self.route_count = len(routes) @@ -112,9 +105,7 @@ class SemanticGuardrail(CustomGuardrail): call_type: str, ): """Check user messages against semantic routes before LLM call.""" - messages = self.get_guardrails_messages_for_call_type( - call_type=CallTypes(call_type), data=data - ) + messages = self.get_guardrails_messages_for_call_type(call_type=CallTypes(call_type), data=data) if not messages: return None @@ -179,10 +170,7 @@ def _extract_user_text(messages: List) -> str: if isinstance(content, str): return content if isinstance(content, list): - return " ".join( - block.get("text", "") if isinstance(block, dict) else str(block) - for block in content - ) + return " ".join(block.get("text", "") if isinstance(block, dict) else str(block) for block in content) return "" @@ -204,9 +192,7 @@ def _content_to_text(content: Any) -> str: return content if isinstance(content, list): text_parts = [ - block.get("text") - for block in content - if isinstance(block, dict) and isinstance(block.get("text"), str) + block.get("text") for block in content if isinstance(block, dict) and isinstance(block.get("text"), str) ] return " ".join(part for part in text_parts if part) return "" @@ -220,10 +206,7 @@ def _handle_match( data: dict, ) -> None: """Block or passthrough based on config.""" - violation_msg = ( - f"Request blocked by semantic guardrail '{guardrail.guardrail_name}'. " - f"Matched route: {route_name}" - ) + violation_msg = f"Request blocked by semantic guardrail '{guardrail.guardrail_name}'. Matched route: {route_name}" detection_info = { "route_name": route_name, @@ -232,8 +215,7 @@ def _handle_match( } verbose_logger.warning( - f"SemanticGuard match: route={route_name}, score={similarity_score}, " - f"action={guardrail.on_flagged_action}" + f"SemanticGuard match: route={route_name}, score={similarity_score}, action={guardrail.on_flagged_action}" ) if guardrail.on_flagged_action == "passthrough": diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index b0932015ab3..2171be235e5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -63,15 +63,9 @@ class ToolPermissionGuardrail(CustomGuardrail): self._load_rules(rules) # Normalize to lowercase for case-insensitive handling - self.default_action = ( - default_action.lower() - if isinstance(default_action, str) - else default_action - ) + self.default_action = default_action.lower() if isinstance(default_action, str) else default_action self.on_disallowed_action = ( - on_disallowed_action.lower() - if isinstance(on_disallowed_action, str) - else on_disallowed_action + on_disallowed_action.lower() if isinstance(on_disallowed_action, str) else on_disallowed_action ) verbose_proxy_logger.debug( @@ -94,11 +88,7 @@ class ToolPermissionGuardrail(CustomGuardrail): compiled_patterns: Dict[str, Dict[str, re.Pattern]] = {} for rule_item in rules or []: - rule = ( - rule_item - if isinstance(rule_item, ToolPermissionRule) - else ToolPermissionRule(**rule_item) - ) + rule = rule_item if isinstance(rule_item, ToolPermissionRule) else ToolPermissionRule(**rule_item) target_patterns: Dict[str, Optional[re.Pattern]] = { "tool_name": None, @@ -108,25 +98,19 @@ class ToolPermissionGuardrail(CustomGuardrail): try: target_patterns["tool_name"] = re.compile(rule.tool_name) except re.error as exc: - raise ValueError( - f"Invalid regex for tool_name in rule '{rule.id}': {exc}" - ) from exc + raise ValueError(f"Invalid regex for tool_name in rule '{rule.id}': {exc}") from exc if rule.tool_type is not None: try: target_patterns["tool_type"] = re.compile(rule.tool_type) except re.error as exc: - raise ValueError( - f"Invalid regex for tool_type in rule '{rule.id}': {exc}" - ) from exc + raise ValueError(f"Invalid regex for tool_type in rule '{rule.id}': {exc}") from exc rule_patterns: Dict[str, re.Pattern] = {} for path, pattern in (rule.allowed_param_patterns or {}).items(): try: rule_patterns[path] = re.compile(pattern) except re.error as exc: - raise ValueError( - f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}" - ) from exc + raise ValueError(f"Invalid regex in allowed_param_patterns for rule '{rule.id}': {exc}") from exc parsed_rules.append(rule) compiled_targets[rule.id] = target_patterns @@ -140,9 +124,7 @@ class ToolPermissionGuardrail(CustomGuardrail): self._compiled_rule_targets = compiled_targets self._compiled_rule_patterns = compiled_patterns - def update_in_memory_litellm_params( - self, litellm_params: Union[LitellmParams, dict] - ) -> None: + def update_in_memory_litellm_params(self, litellm_params: Union[LitellmParams, dict]) -> None: """Apply updated params in place, rebuilding the compiled rule state. The base implementation only ``setattr``s raw fields, which would leave @@ -197,9 +179,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return ToolPermissionGuardrailConfigModel - def _matches_regex( - self, pattern: Optional[re.Pattern], value: Optional[str] - ) -> bool: + def _matches_regex(self, pattern: Optional[re.Pattern], value: Optional[str]) -> bool: if pattern is None: return True if value is None: @@ -220,12 +200,8 @@ class ToolPermissionGuardrail(CustomGuardrail): name_required = rule.tool_name is not None type_required = rule.tool_type is not None - name_matched = ( - self._matches_regex(name_pattern, tool_name) if name_required else True - ) - type_matched = ( - self._matches_regex(type_pattern, tool_type) if type_required else True - ) + name_matched = self._matches_regex(name_pattern, tool_name) if name_required else True + type_matched = self._matches_regex(type_pattern, tool_type) if type_required else True overall_match = name_matched and type_matched should_check_params = name_required and name_matched @@ -247,9 +223,7 @@ class ToolPermissionGuardrail(CustomGuardrail): Returns: Tuple of (is_allowed, rule_id, message) """ - verbose_proxy_logger.debug( - f"Checking permission for tool: {tool_name or tool_type}" - ) + verbose_proxy_logger.debug(f"Checking permission for tool: {tool_name or tool_type}") # Check each rule in order for rule in self.rules: @@ -261,7 +235,9 @@ class ToolPermissionGuardrail(CustomGuardrail): if matches: is_allowed = rule.decision == "allow" tool_identifier = tool_name or tool_type or "unknown_tool" - default_message = f"Tool '{tool_identifier}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" + default_message = ( + f"Tool '{tool_identifier}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" + ) message = self.render_violation_message( default=default_message, context={ @@ -468,9 +444,7 @@ class ToolPermissionGuardrail(CustomGuardrail): function={"name": function_name, "arguments": arguments}, ) - def _extract_tool_calls_from_response( - self, response: ModelResponse - ) -> List[ChatCompletionMessageToolCall]: + def _extract_tool_calls_from_response(self, response: ModelResponse) -> List[ChatCompletionMessageToolCall]: """ Extract tool_calls from all choices in a model response. @@ -514,9 +488,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return tool_choice if self._get_mapping_value(tool_choice, "type") != "function": return None - return self._get_mapping_value( - self._get_mapping_value(tool_choice, "function"), "name" - ) + return self._get_mapping_value(self._get_mapping_value(tool_choice, "function"), "name") def _get_named_function_call(self, data: dict) -> Optional[str]: function_call = data.get("function_call") @@ -563,9 +535,7 @@ class ToolPermissionGuardrail(CustomGuardrail): if not denied_tool_names: return data - verbose_proxy_logger.info( - f"Blocking {len(denied_tool_names)} unauthorized tool uses" - ) + verbose_proxy_logger.info(f"Blocking {len(denied_tool_names)} unauthorized tool uses") # Create a mapping of tool_use_id to error result error_tool_names = set() @@ -585,9 +555,7 @@ class ToolPermissionGuardrail(CustomGuardrail): functions = data.get("functions") if functions is not None: data["functions"] = [ - function - for function in functions - if self._get_legacy_function_name(function) not in error_tool_names + function for function in functions if self._get_legacy_function_name(function) not in error_tool_names ] named_tool_choice = self._get_named_tool_choice(data) @@ -617,9 +585,7 @@ class ToolPermissionGuardrail(CustomGuardrail): if error.rule_id: error_message += f" (Rule: {error.rule_id})" - return ToolResult( - tool_use_id=tool_call.id, content=error_message, is_error=True - ) + return ToolResult(tool_use_id=tool_call.id, content=error_message, is_error=True) def _modify_response_with_permission_errors( self, @@ -636,9 +602,7 @@ class ToolPermissionGuardrail(CustomGuardrail): if not denied_tools: return - verbose_proxy_logger.info( - f"Blocking {len(denied_tools)} unauthorized tool uses" - ) + verbose_proxy_logger.info(f"Blocking {len(denied_tools)} unauthorized tool uses") # Create a mapping of tool_use_id to error result error_results = {} @@ -661,9 +625,7 @@ class ToolPermissionGuardrail(CustomGuardrail): else: filtered_tool_calls.append(tool_call) - choice.message.tool_calls = ( - filtered_tool_calls if filtered_tool_calls else None - ) + choice.message.tool_calls = filtered_tool_calls if filtered_tool_calls else None legacy_tool_call = self._legacy_function_call_to_tool_call( getattr(choice.message, "function_call", None), choice_index @@ -678,9 +640,7 @@ class ToolPermissionGuardrail(CustomGuardrail): if error_messages: existing_content = choice.message.content if existing_content: - choice.message.content = ( - existing_content + "\n\n" + "\n".join(error_messages) - ) + choice.message.content = existing_content + "\n\n" + "\n".join(error_messages) else: choice.message.content = "\n".join(error_messages) @@ -730,13 +690,9 @@ class ToolPermissionGuardrail(CustomGuardrail): if denied_tool_names: data = self._modify_request_with_permission_errors(data, denied_tool_names) - verbose_proxy_logger.debug( - "Tool Permission Guardrail Pre-Call Hook: All tools allowed" - ) + verbose_proxy_logger.debug("Tool Permission Guardrail Pre-Call Hook: All tools allowed") - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @log_guardrail_information @@ -757,16 +713,10 @@ class ToolPermissionGuardrail(CustomGuardrail): if not isinstance(response, ModelResponse): return response - verbose_proxy_logger.debug( - "Tool Permission Guardrail Post-Call Hook: Checking response" - ) + verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: Checking response") - if not self.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ): - verbose_proxy_logger.debug( - "Tool Permission Guardrail: Skipping check (not enabled)" - ) + if not self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call): + verbose_proxy_logger.debug("Tool Permission Guardrail: Skipping check (not enabled)") return response # Extract tool_calls from the response @@ -776,9 +726,7 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail: No tool uses found") return response - verbose_proxy_logger.debug( - f"Tool Permission Guardrail: Found {len(tool_calls)} tool calls" - ) + verbose_proxy_logger.debug(f"Tool Permission Guardrail: Found {len(tool_calls)} tool calls") # Check permissions for each tool use denied_tools = [] @@ -811,13 +759,9 @@ class ToolPermissionGuardrail(CustomGuardrail): if denied_tools: self._modify_response_with_permission_errors(response, denied_tools) else: - verbose_proxy_logger.debug( - "Tool Permission Guardrail Post-Call Hook: All tools allowed" - ) + verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed") - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response async def async_post_call_streaming_iterator_hook( @@ -845,45 +789,31 @@ class ToolPermissionGuardrail(CustomGuardrail): async for chunk in response: all_chunks.append(chunk) - assembled_model_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = stream_chunk_builder( + assembled_model_response: Optional[Union[ModelResponse, TextCompletionResponse]] = stream_chunk_builder( chunks=all_chunks, ) if isinstance(assembled_model_response, ModelResponse): verbose_proxy_logger.debug("Tool Permission Guardrail: Checking response") # Extract tool_calls from the response - tool_calls = self._extract_tool_calls_from_response( - assembled_model_response - ) + tool_calls = self._extract_tool_calls_from_response(assembled_model_response) if not tool_calls: - verbose_proxy_logger.debug( - "Tool Permission Guardrail: No tool uses found" - ) - mock_response = MockResponseIterator( - model_response=assembled_model_response - ) + verbose_proxy_logger.debug("Tool Permission Guardrail: No tool uses found") + mock_response = MockResponseIterator(model_response=assembled_model_response) async for chunk in mock_response: yield chunk return - verbose_proxy_logger.debug( - f"Tool Permission Guardrail: Found {len(tool_calls)} tool calls" - ) + verbose_proxy_logger.debug(f"Tool Permission Guardrail: Found {len(tool_calls)} tool calls") # Check permissions for each tool use denied_tools = [] for tool_call in tool_calls: - is_allowed, rule_id, message = self._get_permission_for_tool_call( - tool_call - ) + is_allowed, rule_id, message = self._get_permission_for_tool_call(tool_call) if not is_allowed and message is not None: - verbose_proxy_logger.warning( - f"Tool Permission Guardrail: {message}" - ) + verbose_proxy_logger.warning(f"Tool Permission Guardrail: {message}") if self.on_disallowed_action == "block": raise GuardrailRaisedException( @@ -906,17 +836,11 @@ class ToolPermissionGuardrail(CustomGuardrail): ) if denied_tools: - self._modify_response_with_permission_errors( - assembled_model_response, denied_tools - ) + self._modify_response_with_permission_errors(assembled_model_response, denied_tools) else: - verbose_proxy_logger.debug( - "Tool Permission Guardrail Post-Call Hook: All tools allowed" - ) + verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed") - mock_response = MockResponseIterator( - model_response=assembled_model_response - ) + mock_response = MockResponseIterator(model_response=assembled_model_response) # Return the reconstructed stream async for chunk in mock_response: yield chunk diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index 12510d051d7..7b9e88fb6e5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -131,9 +131,7 @@ class ToolPolicyGuardrail(CustomGuardrail): tool_names = [ t["function"]["name"] for t in tools - if isinstance(t, dict) - and isinstance(t.get("function"), dict) - and t["function"].get("name") + if isinstance(t, dict) and isinstance(t.get("function"), dict) and t["function"].get("name") ] if not tool_names: route = _get_request_route_from_data(request_data) @@ -189,9 +187,7 @@ class ToolPolicyGuardrail(CustomGuardrail): # For each tool with input_policy=trusted, check if conversation # contains output from tools with output_policy=untrusted if input_type == "response": - trusted_input_tools = [ - name for name in tool_names if policy_map.get(name) == "trusted" - ] + trusted_input_tools = [name for name in tool_names if policy_map.get(name) == "trusted"] if trusted_input_tools: messages = request_data.get("messages") or [] tc_id_to_name = _resolve_tool_names_from_messages(messages) @@ -201,9 +197,7 @@ class ToolPolicyGuardrail(CustomGuardrail): if msg.get("role") != "tool": continue tool_call_id = msg.get("tool_call_id") - source_tool = ( - tc_id_to_name.get(tool_call_id, "") if tool_call_id else "" - ) + source_tool = tc_id_to_name.get(tool_call_id, "") if tool_call_id else "" if not source_tool: continue if registry.get_output_policy(source_tool) == "untrusted": diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 09fff71062b..79c29670d93 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -29,9 +29,7 @@ A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME = "unified_llm_guardrails" -def _get_a2a_request_id( - responses_so_far: List[Any], request_data: dict -) -> Optional[str]: +def _get_a2a_request_id(responses_so_far: List[Any], request_data: dict) -> Optional[str]: """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting.""" for item in responses_so_far: if isinstance(item, dict) and "id" in item: @@ -59,9 +57,7 @@ def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> N BaseTranslation, ) - user_metadata = BaseTranslation.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) + user_metadata = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) if user_metadata: data["litellm_metadata"] = user_metadata @@ -109,10 +105,7 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.pre_mcp_call - if ( - guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) - is not True - ): + if guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True: verbose_proxy_logger.debug( "UnifiedLLMGuardrails: Pre-call scanning disabled for %s", guardrail_to_apply.guardrail_name, @@ -120,9 +113,7 @@ class UnifiedLLMGuardrails(CustomLogger): return data if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = ( - load_guardrail_translation_mappings() - ) + endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() try: if CallTypes(call_type) not in endpoint_guardrail_translation_mappings: @@ -130,9 +121,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ValueError: return data # handle unmapped call types - endpoint_translation = endpoint_guardrail_translation_mappings[ - CallTypes(call_type) - ]() + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() _ensure_litellm_metadata(data, user_api_key_dict) @@ -143,9 +132,7 @@ class UnifiedLLMGuardrails(CustomLogger): ) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=guardrail_to_apply.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_to_apply.guardrail_name) return data async def async_moderation_hook( @@ -169,10 +156,7 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.during_mcp_call - if ( - guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) - is not True - ): + if guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True: verbose_proxy_logger.debug( "UnifiedLLMGuardrails: Pre-call scanning disabled for %s", guardrail_to_apply.guardrail_name, @@ -180,18 +164,11 @@ class UnifiedLLMGuardrails(CustomLogger): return data if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = ( - load_guardrail_translation_mappings() - ) - if ( - call_type is not None - and CallTypes(call_type) not in endpoint_guardrail_translation_mappings - ): + endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: return data - endpoint_translation = endpoint_guardrail_translation_mappings[ - CallTypes(call_type) - ]() + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() _ensure_litellm_metadata(data, user_api_key_dict) @@ -225,17 +202,10 @@ class UnifiedLLMGuardrails(CustomLogger): if guardrail_to_apply is None: return - if ( - guardrail_to_apply.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): + if guardrail_to_apply.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True: return - verbose_proxy_logger.debug( - "async_post_call_success_hook response: %s", response - ) + verbose_proxy_logger.debug("async_post_call_success_hook response: %s", response) call_type: Optional[CallTypesLiteral] = None if user_api_key_dict.request_route is not None: @@ -249,9 +219,7 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is None: litellm_logging_obj = data.get("litellm_logging_obj") logging_call_type = ( - getattr(litellm_logging_obj, "call_type", None) - if litellm_logging_obj is not None - else None + getattr(litellm_logging_obj, "call_type", None) if litellm_logging_obj is not None else None ) if logging_call_type in ( CallTypes.pass_through.value, @@ -263,16 +231,12 @@ class UnifiedLLMGuardrails(CustomLogger): return response if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = ( - load_guardrail_translation_mappings() - ) + endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() if CallTypes(call_type) not in endpoint_guardrail_translation_mappings: return response - endpoint_translation = endpoint_guardrail_translation_mappings[ - CallTypes(call_type) - ]() + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() response = await endpoint_translation.process_output_response( response=response, # type: ignore @@ -282,9 +246,7 @@ class UnifiedLLMGuardrails(CustomLogger): request_data=data, ) # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=guardrail_to_apply.guardrail_name - ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_to_apply.guardrail_name) return response @@ -309,9 +271,7 @@ class UnifiedLLMGuardrails(CustomLogger): global endpoint_guardrail_translation_mappings - guardrail_to_apply: CustomGuardrail = request_data.pop( - "guardrail_to_apply", None - ) + guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None) # Get streaming configuration from guardrail or optional_params sampling_rate = 5 @@ -319,30 +279,18 @@ class UnifiedLLMGuardrails(CustomLogger): if guardrail_to_apply is not None: # Check direct attributes on guardrail first - sampling_rate = getattr( - guardrail_to_apply, "streaming_sampling_rate", sampling_rate - ) - end_of_stream_only = getattr( - guardrail_to_apply, "streaming_end_of_stream_only", end_of_stream_only - ) + sampling_rate = getattr(guardrail_to_apply, "streaming_sampling_rate", sampling_rate) + end_of_stream_only = getattr(guardrail_to_apply, "streaming_end_of_stream_only", end_of_stream_only) # Also check guardrail_config dict if present guardrail_config = getattr(guardrail_to_apply, "guardrail_config", {}) if isinstance(guardrail_config, dict): - sampling_rate = guardrail_config.get( - "streaming_sampling_rate", sampling_rate - ) - end_of_stream_only = guardrail_config.get( - "streaming_end_of_stream_only", end_of_stream_only - ) + sampling_rate = guardrail_config.get("streaming_sampling_rate", sampling_rate) + end_of_stream_only = guardrail_config.get("streaming_end_of_stream_only", end_of_stream_only) # Also check optional_params as fallback - sampling_rate = self.optional_params.get( - "streaming_sampling_rate", sampling_rate - ) - end_of_stream_only = self.optional_params.get( - "streaming_end_of_stream_only", end_of_stream_only - ) + sampling_rate = self.optional_params.get("streaming_sampling_rate", sampling_rate) + end_of_stream_only = self.optional_params.get("streaming_end_of_stream_only", end_of_stream_only) if guardrail_to_apply is None: async for item in response: @@ -350,12 +298,7 @@ class UnifiedLLMGuardrails(CustomLogger): return event_type: GuardrailEventHooks = GuardrailEventHooks.post_call - if ( - guardrail_to_apply.should_run_guardrail( - data=request_data, event_type=event_type - ) - is not True - ): + if guardrail_to_apply.should_run_guardrail(data=request_data, event_type=event_type) is not True: verbose_proxy_logger.debug( "UnifiedLLMGuardrails: Post-call streaming scanning disabled for %s", guardrail_to_apply.guardrail_name, @@ -366,9 +309,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Initialize translation mappings if needed if endpoint_guardrail_translation_mappings is None: - endpoint_guardrail_translation_mappings = ( - load_guardrail_translation_mappings() - ) + endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() # Infer call type from first chunk call_type = None @@ -389,10 +330,7 @@ class UnifiedLLMGuardrails(CustomLogger): call_type = _infer_call_type(call_type=None, completion_response=item) # type: ignore # If call type not supported, just pass through all chunks - if ( - call_type is None - or CallTypes(call_type) not in endpoint_guardrail_translation_mappings - ): + if call_type is None or CallTypes(call_type) not in endpoint_guardrail_translation_mappings: yield item async for remaining_item in response: yield remaining_item @@ -420,9 +358,7 @@ class UnifiedLLMGuardrails(CustomLogger): # string, permanently losing this chunk's content. original_item = copy.deepcopy(item) - endpoint_translation = endpoint_guardrail_translation_mappings[ - CallTypes(call_type) - ]() + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() try: await endpoint_translation.process_output_streaming_response( @@ -437,11 +373,7 @@ class UnifiedLLMGuardrails(CustomLogger): # For A2A (NDJSON), yield an in-stream JSON-RPC error so the client sees it. if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: request_id = _get_a2a_request_id(responses_so_far, request_data) - detail = ( - e.detail - if isinstance(e.detail, dict) - else {"message": str(e.detail)} - ) + detail = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} error_chunk = ( json.dumps( { @@ -453,11 +385,7 @@ class UnifiedLLMGuardrails(CustomLogger): "error", detail.get("message", str(e.detail)), ), - "data": { - k: v - for k, v in detail.items() - if k not in ("error", "message") - }, + "data": {k: v for k, v in detail.items() if k not in ("error", "message")}, }, } ) @@ -471,19 +399,14 @@ class UnifiedLLMGuardrails(CustomLogger): yield item # Stream has ended - do final processing with all collected chunks - if ( - call_type is not None - and CallTypes(call_type) in endpoint_guardrail_translation_mappings - ): + if call_type is not None and CallTypes(call_type) in endpoint_guardrail_translation_mappings: verbose_proxy_logger.debug( "Processing final streaming response with all %s chunks for guardrail %s", len(responses_so_far), guardrail_to_apply.guardrail_name, ) - endpoint_translation = endpoint_guardrail_translation_mappings[ - CallTypes(call_type) - ]() + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() try: await endpoint_translation.process_output_streaming_response( @@ -496,11 +419,7 @@ class UnifiedLLMGuardrails(CustomLogger): except HTTPException as e: if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: request_id = _get_a2a_request_id(responses_so_far, request_data) - detail = ( - e.detail - if isinstance(e.detail, dict) - else {"message": str(e.detail)} - ) + detail = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} error_chunk = ( json.dumps( { @@ -508,14 +427,8 @@ class UnifiedLLMGuardrails(CustomLogger): "id": request_id, "error": { "code": -32603, - "message": detail.get( - "error", detail.get("message", str(e.detail)) - ), - "data": { - k: v - for k, v in detail.items() - if k not in ("error", "message") - }, + "message": detail.get("error", detail.get("message", str(e.detail))), + "data": {k: v for k, v in detail.items() if k not in ("error", "message")}, }, } ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py index 337cb9a9f29..9a8893734e1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py @@ -106,14 +106,10 @@ class VigilGuardGuardrail(CustomGuardrail): self.api_key = resolved_key fallback = (unreachable_fallback or "fail_closed").lower() - self.unreachable_fallback: _FallbackMode = ( - "fail_open" if fallback == "fail_open" else "fail_closed" - ) + self.unreachable_fallback: _FallbackMode = "fail_open" if fallback == "fail_open" else "fail_closed" self.timeout: httpx.Timeout = ( - _DEFAULT_VIGIL_TIMEOUT - if timeout is None - else httpx.Timeout(timeout, connect=min(timeout, 5.0)) + _DEFAULT_VIGIL_TIMEOUT if timeout is None else httpx.Timeout(timeout, connect=min(timeout, 5.0)) ) self.async_handler: _AsyncPostHandler = async_handler or get_async_httpx_client( @@ -146,11 +142,7 @@ class VigilGuardGuardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: texts = inputs.get("texts") or [] has_text = any(isinstance(text, str) and text.strip() for text in texts) - tool_call_args = ( - self._tool_call_arguments(inputs.get("tool_calls")) - if input_type == "response" - else [] - ) + tool_call_args = self._tool_call_arguments(inputs.get("tool_calls")) if input_type == "response" else [] if not has_text and not tool_call_args: return inputs @@ -164,9 +156,7 @@ class VigilGuardGuardrail(CustomGuardrail): continue try: - analysis = await self._analyze( - text=text, source=source, metadata=metadata - ) + analysis = await self._analyze(text=text, source=source, metadata=metadata) except ( httpx.HTTPError, LiteLLMTimeout, @@ -184,8 +174,7 @@ class VigilGuardGuardrail(CustomGuardrail): decision = analysis.get("decision") if isinstance(analysis, dict) else None if decision not in _VALID_DECISIONS: verbose_proxy_logger.error( - "Vigil Guard unrecognized decision for guardrail_name=%s " - "source=%s: %r", + "Vigil Guard unrecognized decision for guardrail_name=%s source=%s: %r", self.guardrail_name, source, decision, @@ -217,24 +206,19 @@ class VigilGuardGuardrail(CustomGuardrail): result_tool_calls = inputs.get("tool_calls") for tc_index, arguments in tool_call_args: try: - analysis = await self._analyze( - text=arguments, source=source, metadata=metadata - ) + analysis = await self._analyze(text=arguments, source=source, metadata=metadata) except ( httpx.HTTPError, LiteLLMTimeout, JSONDecodeError, OSError, ) as exc: - return self._handle_backend_failure( - exc, inputs, source, result_texts, result_tool_calls - ) + return self._handle_backend_failure(exc, inputs, source, result_texts, result_tool_calls) decision = analysis.get("decision") if isinstance(analysis, dict) else None if decision not in _VALID_DECISIONS: verbose_proxy_logger.error( - "Vigil Guard unrecognized decision for guardrail_name=%s " - "source=%s: %r", + "Vigil Guard unrecognized decision for guardrail_name=%s source=%s: %r", self.guardrail_name, source, decision, @@ -281,8 +265,7 @@ class VigilGuardGuardrail(CustomGuardrail): ) return self._build_output(inputs, final_texts, final_tool_calls) verbose_proxy_logger.error( - "Vigil Guard backend failure with fail_closed; blocking request. " - "guardrail_name=%s source=%s error=%s", + "Vigil Guard backend failure with fail_closed; blocking request. guardrail_name=%s source=%s error=%s", self.guardrail_name, source, str(exc), @@ -321,20 +304,14 @@ class VigilGuardGuardrail(CustomGuardrail): pairs: List[Tuple[int, str]] = [] if isinstance(tool_calls, list): for index, tool_call in enumerate(tool_calls): - function = ( - tool_call.get("function") if isinstance(tool_call, dict) else None - ) - arguments = ( - function.get("arguments") if isinstance(function, dict) else None - ) + function = tool_call.get("function") if isinstance(tool_call, dict) else None + arguments = function.get("arguments") if isinstance(function, dict) else None if isinstance(arguments, str) and arguments.strip(): pairs.append((index, arguments)) return pairs @staticmethod - def _set_tool_call_arguments( - tool_calls: Any, index: int, arguments: str - ) -> List[Any]: + def _set_tool_call_arguments(tool_calls: Any, index: int, arguments: str) -> List[Any]: updated = list(tool_calls) tool_call = dict(updated[index]) function = dict(tool_call.get("function") or {}) @@ -343,9 +320,7 @@ class VigilGuardGuardrail(CustomGuardrail): updated[index] = tool_call return updated - async def _analyze( - self, text: str, source: str, metadata: Dict[str, Any] - ) -> Dict[str, Any]: + async def _analyze(self, text: str, source: str, metadata: Dict[str, Any]) -> Dict[str, Any]: payload = { "text": text, "source": source, @@ -360,9 +335,7 @@ class VigilGuardGuardrail(CustomGuardrail): response = await self._post_with_retry(endpoint, headers, payload) return response.json() - async def _post_with_retry( - self, endpoint: str, headers: Dict[str, str], payload: Dict[str, Any] - ) -> httpx.Response: + async def _post_with_retry(self, endpoint: str, headers: Dict[str, str], payload: Dict[str, Any]) -> httpx.Response: for attempt in range(2): try: response = await self.async_handler.post( @@ -419,9 +392,7 @@ class VigilGuardGuardrail(CustomGuardrail): return value return original - def _collect_metadata( - self, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"] - ) -> Dict[str, Any]: + def _collect_metadata(self, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"]) -> Dict[str, Any]: sources: List[dict] = [] if isinstance(request_data, dict): sources.append(request_data) @@ -466,9 +437,7 @@ class VigilGuardGuardrail(CustomGuardrail): return None @staticmethod - def _extract_call_id( - request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"] - ) -> Optional[str]: + def _extract_call_id(request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"]) -> Optional[str]: if logging_obj is not None: call_id = getattr(logging_obj, "litellm_call_id", None) if isinstance(call_id, str) and call_id: diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index 5acfc5403ce..36c22753404 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -98,9 +98,7 @@ class XecGuardGuardrail(CustomGuardrail): "the guardrail config." ) - self.api_base = ( - api_base or os.environ.get("XECGUARD_API_BASE") or _DEFAULT_API_BASE - ).rstrip("/") + self.api_base = (api_base or os.environ.get("XECGUARD_API_BASE") or _DEFAULT_API_BASE).rstrip("/") self.xecguard_model = xecguard_model or _DEFAULT_MODEL self.policy_names = policy_names @@ -115,9 +113,7 @@ class XecGuardGuardrail(CustomGuardrail): else: self.block_on_error = block_on_error - self.grounding_strictness = ( - grounding_strictness or _DEFAULT_GROUNDING_STRICTNESS - ) + self.grounding_strictness = grounding_strictness or _DEFAULT_GROUNDING_STRICTNESS self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, @@ -179,16 +175,11 @@ class XecGuardGuardrail(CustomGuardrail): messages=messages, documents=documents, ) - if ( - grounding_result is not None - and grounding_result.get("decision") == "UNSAFE" - ): + if grounding_result is not None and grounding_result.get("decision") == "UNSAFE": raise HTTPException( status_code=400, detail={ - "error": self._format_grounding_block_message( - grounding_result - ), + "error": self._format_grounding_block_message(grounding_result), "guardrail_name": self.guardrail_name or "xecguard", "xecguard_response": grounding_result, }, @@ -212,11 +203,8 @@ class XecGuardGuardrail(CustomGuardrail): isinstance(kwargs, dict) and "litellm_params" in kwargs and "metadata" in kwargs["litellm_params"] - and "standard_logging_guardrail_information" - in kwargs["litellm_params"]["metadata"] - and kwargs["litellm_params"]["metadata"][ - "standard_logging_guardrail_information" - ] + and "standard_logging_guardrail_information" in kwargs["litellm_params"]["metadata"] + and kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] ): return kwargs, result @@ -252,9 +240,7 @@ class XecGuardGuardrail(CustomGuardrail): return kwargs, result guardrail_status: GuardrailStatus = ( - "guardrail_intervened" - if scan_result.get("decision") == "UNSAFE" - else "success" + "guardrail_intervened" if scan_result.get("decision") == "UNSAFE" else "success" ) end_time = datetime.now() kwargs["standard_logging_object"]["guardrail_information"] = { @@ -295,11 +281,7 @@ class XecGuardGuardrail(CustomGuardrail): asyncio.set_event_loop(loop) if loop.is_running(): return kwargs, result - loop.run_until_complete( - self.async_logging_hook( - kwargs=kwargs, result=result, call_type=call_type - ) - ) + loop.run_until_complete(self.async_logging_hook(kwargs=kwargs, result=result, call_type=call_type)) except Exception as exc: verbose_proxy_logger.debug( "XecGuard sync logging_hook swallowed exception: %s", @@ -321,9 +303,7 @@ class XecGuardGuardrail(CustomGuardrail): "model": self.xecguard_model, "scan_type": scan_type, "messages": messages, - "policy_names": ( - self.policy_names if self.policy_names else _DEFAULT_POLICIES - ), + "policy_names": (self.policy_names if self.policy_names else _DEFAULT_POLICIES), } return await self._post( path=_SCAN_ENDPOINT, @@ -381,9 +361,7 @@ class XecGuardGuardrail(CustomGuardrail): raise HTTPException( status_code=400, detail={ - "error": ( - f"XecGuard API unreachable (block_on_error=True): {exc}" - ), + "error": (f"XecGuard API unreachable (block_on_error=True): {exc}"), "guardrail_name": self.guardrail_name or "xecguard", }, ) from exc @@ -407,9 +385,7 @@ class XecGuardGuardrail(CustomGuardrail): the request data is incomplete. """ raw_messages = request_data.get("messages") or [] - messages: List[dict] = [ - self._normalize_message(m) for m in raw_messages if isinstance(m, dict) - ] + messages: List[dict] = [self._normalize_message(m) for m in raw_messages if isinstance(m, dict)] if input_type == "request": if not messages: @@ -422,9 +398,7 @@ class XecGuardGuardrail(CustomGuardrail): return messages # input_type == "response" - assistant_text = self._extract_assistant_text_from_response( - request_data.get("response") - ) + assistant_text = self._extract_assistant_text_from_response(request_data.get("response")) if assistant_text is None: return [] messages.append({"role": "assistant", "content": assistant_text}) @@ -512,9 +486,7 @@ class XecGuardGuardrail(CustomGuardrail): parts = [ item.get("text") for item in content - if isinstance(item, dict) - and item.get("type") == "text" - and isinstance(item.get("text"), str) + if isinstance(item, dict) and item.get("type") == "text" and isinstance(item.get("text"), str) ] joined = "\n".join(p for p in parts if p) return joined or None diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index e0e59bfef3b..9c660620582 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -42,11 +42,7 @@ class ZscalerAIGuard(CustomGuardrail): "ZSCALER_AI_GUARD_URL", "https://api.us1.zseclipse.net/v1/detection/execute-policy", ) - self.policy_id = ( - policy_id - if policy_id is not None - else int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1)) - ) + self.policy_id = policy_id if policy_id is not None else int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1)) self.api_key = api_key or os.getenv("ZSCALER_AI_GUARD_API_KEY") self.send_user_api_key_alias = ( send_user_api_key_alias @@ -56,14 +52,12 @@ class ZscalerAIGuard(CustomGuardrail): self.send_user_api_key_user_id = ( send_user_api_key_user_id if send_user_api_key_user_id is not None - else os.getenv("SEND_USER_API_KEY_USER_ID", "False").lower() - in ("true", "1") + else os.getenv("SEND_USER_API_KEY_USER_ID", "False").lower() in ("true", "1") ) self.send_user_api_key_team_id = ( send_user_api_key_team_id if send_user_api_key_team_id is not None - else os.getenv("SEND_USER_API_KEY_TEAM_ID", "False").lower() - in ("true", "1") + else os.getenv("SEND_USER_API_KEY_TEAM_ID", "False").lower() in ("true", "1") ) verbose_proxy_logger.debug( @@ -77,9 +71,7 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.debug("ZscalerAIGuard Initializing ...") @staticmethod - def _resolve_metadata_value( - request_data: Optional[dict], key: str - ) -> Optional[str]: + def _resolve_metadata_value(request_data: Optional[dict], key: str) -> Optional[str]: """ Resolve metadata value from request_data, checking both metadata locations. @@ -161,9 +153,7 @@ class ZscalerAIGuard(CustomGuardrail): user_api_key_metadata.get("zguard_policy_id") if "zguard_policy_id" in user_api_key_metadata else ( - team_metadata.get("zguard_policy_id") - if "zguard_policy_id" in team_metadata - else self.policy_id + team_metadata.get("zguard_policy_id") if "zguard_policy_id" in team_metadata else self.policy_id ) ) ) @@ -171,19 +161,14 @@ class ZscalerAIGuard(CustomGuardrail): kwargs = {} if self.send_user_api_key_alias: - kwargs["user_api_key_alias"] = ( - self._resolve_metadata_value(request_data, "user_api_key_alias") - or "N/A" - ) + kwargs["user_api_key_alias"] = self._resolve_metadata_value(request_data, "user_api_key_alias") or "N/A" if self.send_user_api_key_team_id: kwargs["user_api_key_team_id"] = ( - self._resolve_metadata_value(request_data, "user_api_key_team_id") - or "N/A" + self._resolve_metadata_value(request_data, "user_api_key_team_id") or "N/A" ) if self.send_user_api_key_user_id: kwargs["user_api_key_user_id"] = ( - self._resolve_metadata_value(request_data, "user_api_key_user_id") - or "N/A" + self._resolve_metadata_value(request_data, "user_api_key_user_id") or "N/A" ) verbose_proxy_logger.debug(f"inside apply_guardrail kwargs: {kwargs}") @@ -201,20 +186,13 @@ class ZscalerAIGuard(CustomGuardrail): content=concatenated_text, **kwargs, ) - verbose_proxy_logger.debug( - f"response from zscaler ai guards: {zscaler_ai_guard_result}" - ) - if ( - zscaler_ai_guard_result - and zscaler_ai_guard_result.get("action") == "BLOCK" - ): + verbose_proxy_logger.debug(f"response from zscaler ai guards: {zscaler_ai_guard_result}") + if zscaler_ai_guard_result and zscaler_ai_guard_result.get("action") == "BLOCK": blocking_info = zscaler_ai_guard_result.get("zscaler_ai_guard_response") error_message = f"Content blocked by Zscaler AI Guard: {self.extract_blocking_info(blocking_info)}" raise Exception(error_message) except Exception as e: - verbose_proxy_logger.error( - "ZscalerAIGuard: Failed to apply guardrail: %s", str(e) - ) + verbose_proxy_logger.error("ZscalerAIGuard: Failed to apply guardrail: %s", str(e)) raise e verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.") @@ -256,9 +234,7 @@ class ZscalerAIGuard(CustomGuardrail): if self.send_user_api_key_alias: verbose_proxy_logger.debug(f"kwargs: {kwargs}") user_api_key_alias = kwargs.get("user_api_key_alias", "N/A") - verbose_proxy_logger.debug( - f"kwargs user_api_key_alias: {user_api_key_alias}" - ) + verbose_proxy_logger.debug(f"kwargs user_api_key_alias: {user_api_key_alias}") extra_headers.update({"user-api-key-alias": user_api_key_alias}) if self.send_user_api_key_team_id: @@ -273,9 +249,7 @@ class ZscalerAIGuard(CustomGuardrail): return extra_headers async def _send_request(self, url, headers, data): - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) response = await async_client.post( f"{url}", @@ -289,12 +263,8 @@ class ZscalerAIGuard(CustomGuardrail): def _handle_response(self, response, direction): # Raise exceptions on critical errors to stop the request if response.status_code == 429: # Rate limit - verbose_proxy_logger.error( - "Zscaler AI Guard rate limit reached. Blocking request." - ) - user_facing_error = self._create_user_facing_error( - "Rate limit reached. status_code: 429" - ) + verbose_proxy_logger.error("Zscaler AI Guard rate limit reached. Blocking request.") + user_facing_error = self._create_user_facing_error("Rate limit reached. status_code: 429") # This exception will be caught by the proxy and returned to the user raise HTTPException(status_code=500, detail=user_facing_error) @@ -302,9 +272,7 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.error( f"Zscaler AI Guard service is unavailable (Status: {response.status_code}). Blocking request." ) - user_facing_error = self._create_user_facing_error( - f"Service is unavailable (HTTP {response.status_code})" - ) + user_facing_error = self._create_user_facing_error(f"Service is unavailable (HTTP {response.status_code})") raise HTTPException(status_code=500, detail=user_facing_error) if response.status_code == 200: @@ -341,23 +309,15 @@ class ZscalerAIGuard(CustomGuardrail): raise HTTPException(status_code=500, detail=user_facing_error) else: errorMsg = json_response.get("errorMsg", None) - verbose_proxy_logger.error( - f"statusCode in response: {statusCode_in_response}, errorMsg: {errorMsg}" - ) + verbose_proxy_logger.error(f"statusCode in response: {statusCode_in_response}, errorMsg: {errorMsg}") user_facing_error = self._create_user_facing_error( f"statusCode in response: {statusCode_in_response}, errorMsg: {errorMsg}" ) raise HTTPException(status_code=500, detail=user_facing_error) else: - verbose_proxy_logger.error( - f"Zscaler AI Guard status_code - {response.status_code}" - ) - user_facing_error = self._create_user_facing_error( - f"Response status code: {response.status_code}" - ) - raise HTTPException( - status_code=response.status_code, detail=user_facing_error - ) + verbose_proxy_logger.error(f"Zscaler AI Guard status_code - {response.status_code}") + user_facing_error = self._create_user_facing_error(f"Response status code: {response.status_code}") + raise HTTPException(status_code=response.status_code, detail=user_facing_error) async def make_zscaler_ai_guard_api_call( self, zscaler_ai_guard_url, api_key, policy_id, direction, content, **kwargs @@ -378,9 +338,7 @@ class ZscalerAIGuard(CustomGuardrail): if policy_id is not None and policy_id >= 1: data["policyId"] = policy_id try: - response = await self._send_request( - zscaler_ai_guard_url, extra_headers, data - ) + response = await self._send_request(zscaler_ai_guard_url, extra_headers, data) return self._handle_response(response, direction) except Exception as e: verbose_proxy_logger.error(f"{e}. Blocking request.") diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 9af43950837..e8abf66a6f7 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -130,10 +130,7 @@ def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail) _ENTERPRISE_SecretDetection, ) except ImportError: - raise Exception( - "Trying to use Secret Detection" - + CommonProxyErrors.missing_enterprise_package.value - ) + raise Exception("Trying to use Secret Detection" + CommonProxyErrors.missing_enterprise_package.value) _secret_detection_object = _ENTERPRISE_SecretDetection( detect_secrets_config=litellm_params.detect_secrets_config, @@ -204,12 +201,9 @@ def initialize_panw_prisma_airs(litellm_params, guardrail): raise ValueError("PANW Prisma AIRS: profile_name is required") _panw_callback = PanwPrismaAirsHandler( - guardrail_name=guardrail.get( - "guardrail_name", "panw_prisma_airs" - ), # Use .get() with default + guardrail_name=guardrail.get("guardrail_name", "panw_prisma_airs"), # Use .get() with default api_key=litellm_params.api_key, - api_base=litellm_params.api_base - or "https://service.api.aisecurity.paloaltonetworks.com/v1/scan/sync/request", + api_base=litellm_params.api_base or "https://service.api.aisecurity.paloaltonetworks.com/v1/scan/sync/request", profile_name=litellm_params.profile_name, default_on=litellm_params.default_on, mask_on_block=getattr(litellm_params, "mask_on_block", False), diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 734330a2167..8962073fe7a 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -114,9 +114,7 @@ def get_guardrail_initializer_from_hooks(): # For directories with just initialize_guardrail, use the directory name as the key initialize_fn = getattr(module, "initialize_guardrail") discovered_initializers[item] = initialize_fn - verbose_proxy_logger.debug( - f"Found initialize_guardrail function in {module_path}" - ) + verbose_proxy_logger.debug(f"Found initialize_guardrail function in {module_path}") except ImportError as e: verbose_proxy_logger.error(f"Could not import {module_path}: {e}") @@ -222,17 +220,11 @@ class GuardrailRegistry: ########################################################### ########### In memory management helpers for guardrails ########### ############################################################ - def get_initialized_guardrail_callback( - self, guardrail_name: str - ) -> Optional[CustomGuardrail]: + def get_initialized_guardrail_callback(self, guardrail_name: str) -> Optional[CustomGuardrail]: """ Returns the initialized guardrail callback for a given guardrail name """ - active_guardrails = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=CustomGuardrail - ) - ) + active_guardrails = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) for active_guardrail in active_guardrails: if isinstance(active_guardrail, CustomGuardrail): if active_guardrail.guardrail_name == guardrail_name: @@ -242,9 +234,7 @@ class GuardrailRegistry: ########################################################### ########### DB management helpers for guardrails ########### ############################################################ - async def add_guardrail_to_db( - self, guardrail: Guardrail, prisma_client: PrismaClient - ): + async def add_guardrail_to_db(self, guardrail: Guardrail, prisma_client: PrismaClient): """ Add a guardrail to the database """ @@ -255,9 +245,7 @@ class GuardrailRegistry: if hasattr(litellm_params_obj, "model_dump"): litellm_params_dict = litellm_params_obj.model_dump() else: - litellm_params_dict = ( - dict(litellm_params_obj) if litellm_params_obj else {} - ) + litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} litellm_params: str = safe_dumps(litellm_params_dict) guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) @@ -280,25 +268,19 @@ class GuardrailRegistry: except Exception as e: raise Exception(f"Error adding guardrail to DB: {str(e)}") - async def delete_guardrail_from_db( - self, guardrail_id: str, prisma_client: PrismaClient - ): + async def delete_guardrail_from_db(self, guardrail_id: str, prisma_client: PrismaClient): """ Delete a guardrail from the database """ try: # Delete from DB - await GuardrailsRepository(prisma_client).table.delete( - where={"guardrail_id": guardrail_id} - ) + await GuardrailsRepository(prisma_client).table.delete(where={"guardrail_id": guardrail_id}) return {"message": f"Guardrail {guardrail_id} deleted successfully"} except Exception as e: raise Exception(f"Error deleting guardrail from DB: {str(e)}") - async def update_guardrail_in_db( - self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient - ): + async def update_guardrail_in_db(self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient): """ Update a guardrail in the database """ @@ -309,9 +291,7 @@ class GuardrailRegistry: if hasattr(litellm_params_obj, "model_dump"): litellm_params_dict = litellm_params_obj.model_dump() else: - litellm_params_dict = ( - dict(litellm_params_obj) if litellm_params_obj else {} - ) + litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} litellm_params: str = safe_dumps(litellm_params_dict) guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) @@ -340,9 +320,7 @@ class GuardrailRegistry: Only rows with status == "active" are returned (pending_review and rejected are excluded). """ try: - guardrails_from_db = await GuardrailsRepository( - prisma_client - ).table.find_many( + guardrails_from_db = await GuardrailsRepository(prisma_client).table.find_many( where={"status": "active"}, order={"created_at": "desc"}, ) @@ -355,9 +333,7 @@ class GuardrailRegistry: except Exception as e: raise Exception(f"Error getting guardrails from DB: {str(e)}") - async def get_guardrail_by_id_from_db( - self, guardrail_id: str, prisma_client: PrismaClient - ) -> Optional[Guardrail]: + async def get_guardrail_by_id_from_db(self, guardrail_id: str, prisma_client: PrismaClient) -> Optional[Guardrail]: """ Get a guardrail by its ID from the database """ @@ -430,9 +406,7 @@ class InMemoryGuardrailHandler: guardrail_id = guardrail.get("guardrail_id") or str(uuid.uuid4()) guardrail["guardrail_id"] = guardrail_id if guardrail_id in self.IN_MEMORY_GUARDRAILS: - verbose_proxy_logger.debug( - "guardrail_id already exists in IN_MEMORY_GUARDRAILS" - ) + verbose_proxy_logger.debug("guardrail_id already exists in IN_MEMORY_GUARDRAILS") # Honor the caller's source even on the early-return path so a # racing polling tick or a hot-reload of config can correct an # entry's provenance. @@ -448,21 +422,14 @@ class InMemoryGuardrailHandler: else: litellm_params = litellm_params_data - if ( - "category_thresholds" in litellm_params_data - and litellm_params_data["category_thresholds"] - ): - lakera_category_thresholds = LakeraCategoryThresholds( - **litellm_params_data["category_thresholds"] - ) + if "category_thresholds" in litellm_params_data and litellm_params_data["category_thresholds"]: + lakera_category_thresholds = LakeraCategoryThresholds(**litellm_params_data["category_thresholds"]) litellm_params.category_thresholds = lakera_category_thresholds if litellm_params.api_key and litellm_params.api_key.startswith("os.environ/"): litellm_params.api_key = str(get_secret(litellm_params.api_key)) - if litellm_params.api_base and litellm_params.api_base.startswith( - "os.environ/" - ): + if litellm_params.api_base and litellm_params.api_base.startswith("os.environ/"): litellm_params.api_base = str(get_secret(litellm_params.api_base)) guardrail_type = litellm_params.guardrail @@ -533,18 +500,14 @@ class InMemoryGuardrailHandler: This initializes it by adding it to the litellm callback manager """ if not config_file_path: - raise Exception( - "GuardrailsAIException - Please pass the config_file_path to initialize_guardrails_v2" - ) + raise Exception("GuardrailsAIException - Please pass the config_file_path to initialize_guardrails_v2") verbose_proxy_logger.debug( "Initializing custom guardrail: %s", guardrail_type, ) - _guardrail_class = get_instance_fn( - guardrail_type, config_file_path=config_file_path - ) + _guardrail_class = get_instance_fn(guardrail_type, config_file_path=config_file_path) mode = litellm_params.mode if mode is None: @@ -591,16 +554,10 @@ class InMemoryGuardrailHandler: self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail self._sources[guardrail_id] = source - custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.get( - guardrail_id - ) + custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.get(guardrail_id) if custom_guardrail_callback: - updated_litellm_params = cast( - LitellmParams, guardrail.get("litellm_params", {}) - ) - custom_guardrail_callback.update_in_memory_litellm_params( - litellm_params=updated_litellm_params - ) + updated_litellm_params = cast(LitellmParams, guardrail.get("litellm_params", {})) + custom_guardrail_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ @@ -615,15 +572,11 @@ class InMemoryGuardrailHandler: self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None) self._sources.pop(guardrail_id, None) - custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop( - guardrail_id, None - ) + custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop(guardrail_id, None) if custom_guardrail_callback is None: return - litellm.logging_callback_manager.remove_callback_from_all_lists( - custom_guardrail_callback - ) + litellm.logging_callback_manager.remove_callback_from_all_lists(custom_guardrail_callback) def list_in_memory_guardrails(self) -> List[Guardrail]: """ @@ -658,8 +611,7 @@ class InMemoryGuardrailHandler: ] for guardrail_id in stale_ids: verbose_proxy_logger.info( - "Reconcile: removing stale DB-backed guardrail '%s' from memory " - "(deleted in DB by another pod)", + "Reconcile: removing stale DB-backed guardrail '%s' from memory (deleted in DB by another pod)", guardrail_id, ) self.delete_in_memory_guardrail(guardrail_id) @@ -693,9 +645,7 @@ class InMemoryGuardrailHandler: return params return params - def _has_guardrail_params_changed( - self, guardrail_id: str, new_guardrail: Guardrail - ) -> bool: + def _has_guardrail_params_changed(self, guardrail_id: str, new_guardrail: Guardrail) -> bool: """ Check if guardrail params or name have changed compared to in-memory version. Returns True if params/name changed or guardrail doesn't exist in memory. @@ -709,12 +659,8 @@ class InMemoryGuardrailHandler: return True # Compare litellm_params - existing_dict = self._normalize_litellm_params_for_comparison( - existing.get("litellm_params") - ) - new_dict = self._normalize_litellm_params_for_comparison( - new_guardrail.get("litellm_params") - ) + existing_dict = self._normalize_litellm_params_for_comparison(existing.get("litellm_params")) + new_dict = self._normalize_litellm_params_for_comparison(new_guardrail.get("litellm_params")) # Compare and identify specific differences changed_fields = {} @@ -730,9 +676,7 @@ class InMemoryGuardrailHandler: # Log differences if any found if changed_fields: - verbose_proxy_logger.debug( - f"Guardrail params changed. Differences: {changed_fields}" - ) + verbose_proxy_logger.debug(f"Guardrail params changed. Differences: {changed_fields}") # Return True if any fields changed return len(changed_fields) > 0 @@ -749,9 +693,7 @@ class InMemoryGuardrailHandler: """ guardrail_id = guardrail.get("guardrail_id") if not guardrail_id: - verbose_proxy_logger.error( - "Cannot reinitialize guardrail without guardrail_id" - ) + verbose_proxy_logger.error("Cannot reinitialize guardrail without guardrail_id") return None # Remove from memory if exists (also removes from callbacks) @@ -759,9 +701,7 @@ class InMemoryGuardrailHandler: self.delete_in_memory_guardrail(guardrail_id) # Initialize fresh (will add new callback to litellm.callbacks) - return self.initialize_guardrail( - guardrail=guardrail, config_file_path=config_file_path, source=source - ) + return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source) def sync_guardrail_from_db( self, guardrail: Guardrail, config_file_path: Optional[str] = None diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 9e45f18232c..34961afb8ff 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -53,9 +53,7 @@ def _populate_router_guardrail_list(guardrail_list: List[Guardrail]) -> None: from litellm.types.router import GuardrailTypedDict if llm_router is None: - verbose_proxy_logger.debug( - "Router not initialized yet, skipping guardrail_list population" - ) + verbose_proxy_logger.debug("Router not initialized yet, skipping guardrail_list population") return router_guardrail_list: List[GuardrailTypedDict] = [] @@ -68,16 +66,10 @@ def _populate_router_guardrail_list(guardrail_list: List[Guardrail]) -> None: # Get the callback instance from the registry callback = None if guardrail_id: - callback = IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.get( - guardrail_id - ) + callback = IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.get(guardrail_id) # Build litellm_params dict for the router - params_dict = ( - litellm_params.model_dump() - if hasattr(litellm_params, "model_dump") - else dict(litellm_params) - ) + params_dict = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) router_guardrail: GuardrailTypedDict = GuardrailTypedDict( guardrail_name=guardrail_name or "", @@ -94,9 +86,7 @@ def _populate_router_guardrail_list(guardrail_list: List[Guardrail]) -> None: router_guardrail_list.append(router_guardrail) llm_router.guardrail_list = router_guardrail_list - verbose_proxy_logger.debug( - f"Populated router guardrail_list with {len(router_guardrail_list)} guardrails" - ) + verbose_proxy_logger.debug(f"Populated router guardrail_list with {len(router_guardrail_list)} guardrails") ### LEGACY IMPLEMENTATION ### @@ -137,9 +127,7 @@ def initialize_guardrails( if guardrail.logging_only is True: if callback == "presidio": - callback_specific_params["presidio"] = { - "logging_only": True - } # type: ignore + callback_specific_params["presidio"] = {"logging_only": True} # type: ignore default_on_callbacks_list = list(default_on_callbacks) if len(default_on_callbacks_list) > 0: @@ -153,7 +141,5 @@ def initialize_guardrails( return litellm.guardrail_name_config_map except Exception as e: - verbose_proxy_logger.exception( - "error initializing guardrails {}".format(str(e)) - ) + verbose_proxy_logger.exception("error initializing guardrails {}".format(str(e))) raise e diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py index fb1c0d72ee7..02f11e3c20a 100644 --- a/litellm/proxy/guardrails/tool_name_extraction.py +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -40,12 +40,8 @@ def _extract_mcp_tool_names(data: dict) -> List[str]: def _register_standalone_extractors() -> None: if STANDALONE_EXTRACTORS: return - STANDALONE_EXTRACTORS[CallTypes.generate_content.value] = ( - _extract_generate_content_tool_names - ) - STANDALONE_EXTRACTORS[CallTypes.agenerate_content.value] = ( - _extract_generate_content_tool_names - ) + STANDALONE_EXTRACTORS[CallTypes.generate_content.value] = _extract_generate_content_tool_names + STANDALONE_EXTRACTORS[CallTypes.agenerate_content.value] = _extract_generate_content_tool_names STANDALONE_EXTRACTORS[CallTypes.call_mcp_tool.value] = _extract_mcp_tool_names diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index d8457cf9c86..e03bdbb95d2 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -123,10 +123,7 @@ def _prev_fail_rates(metrics_prev: Any, id_attr: str) -> Dict[str, float]: prev_agg_raw[gid] = {"req": 0, "blocked": 0} prev_agg_raw[gid]["req"] += r prev_agg_raw[gid]["blocked"] += b - return { - gid: (100.0 * v["blocked"] / v["req"]) if v["req"] else 0.0 - for gid, v in prev_agg_raw.items() - } + return {gid: (100.0 * v["blocked"] / v["req"]) if v["req"] else 0.0 for gid, v in prev_agg_raw.items()} def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: @@ -137,20 +134,13 @@ def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]: chart_by_date[d] = {"passed": 0, "blocked": 0} chart_by_date[d]["passed"] += int(m.passed_count or 0) chart_by_date[d]["blocked"] += int(m.blocked_count or 0) - return [ - {"date": d, "passed": v["passed"], "blocked": v["blocked"]} - for d, v in sorted(chart_by_date.items()) - ] + return [{"date": d, "passed": v["passed"], "blocked": v["blocked"]} for d, v in sorted(chart_by_date.items())] def _get_guardrail_attrs(g: Any) -> tuple[Any, str]: """Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict.""" - gid = getattr(g, "guardrail_id", None) or ( - g.get("guardrail_id") if isinstance(g, dict) else None - ) - name = getattr(g, "guardrail_name", None) or ( - g.get("guardrail_name") if isinstance(g, dict) else None - ) + gid = getattr(g, "guardrail_id", None) or (g.get("guardrail_id") if isinstance(g, dict) else None) + name = getattr(g, "guardrail_name", None) or (g.get("guardrail_name") if isinstance(g, dict) else None) return gid, (name or gid or "") @@ -173,13 +163,9 @@ def _guardrail_overview_rows( break req, blocked = a["requests"], a["blocked"] fail_rate = (100.0 * blocked / req) if req else 0.0 - litellm_params = ( - (g.litellm_params or {}) if isinstance(g.litellm_params, dict) else {} - ) + litellm_params = (g.litellm_params or {}) if isinstance(g.litellm_params, dict) else {} provider = str(litellm_params.get("guardrail", "Unknown")) - guardrail_info = ( - (g.guardrail_info or {}) if isinstance(g.guardrail_info, dict) else {} - ) + guardrail_info = (g.guardrail_info or {}) if isinstance(g.guardrail_info, dict) else {} gtype = str(guardrail_info.get("type", "Guardrail")) prev_fail = 0.0 for k in lookup_keys: @@ -270,9 +256,7 @@ async def guardrails_usage_overview( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return UsageOverviewResponse( - rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0 - ) + return UsageOverviewResponse(rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0) now = datetime.now(timezone.utc) end = end_date or now.strftime("%Y-%m-%d") @@ -288,23 +272,17 @@ async def guardrails_usage_overview( ) # Previous period for trend - start_prev = ( - datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7) - ).strftime("%Y-%m-%d") - metrics_prev = await DailyGuardrailMetricsRepository( - prisma_client - ).table.find_many(where={"date": {"gte": start_prev, "lt": start}}) + start_prev = (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d") + metrics_prev = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( + where={"date": {"gte": start_prev, "lt": start}} + ) agg = _aggregate_daily_metrics(metrics, "guardrail_id") prev_agg = _prev_fail_rates(metrics_prev, "guardrail_id") chart = _chart_from_metrics(metrics) total_requests = sum(a["requests"] for a in agg.values()) total_blocked = sum(a["blocked"] for a in agg.values()) - pass_rate = ( - (100.0 * (total_requests - total_blocked) / total_requests) - if total_requests - else 100.0 - ) + pass_rate = (100.0 * (total_requests - total_blocked) / total_requests) if total_requests else 100.0 rows = _guardrail_overview_rows(guardrails, agg, prev_agg) return UsageOverviewResponse( rows=rows, @@ -343,9 +321,7 @@ async def guardrails_usage_detail( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") - guardrail = await GuardrailsRepository(prisma_client).table.find_unique( - where={"guardrail_id": guardrail_id} - ) + guardrail = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id}) if not guardrail: from fastapi import HTTPException @@ -427,9 +403,7 @@ def _build_usage_logs_where( ) -> Dict[str, Any]: where: Dict[str, Any] = {} if guardrail_ids: - where["guardrail_id"] = ( - {"in": guardrail_ids} if len(guardrail_ids) > 1 else guardrail_ids[0] - ) + where["guardrail_id"] = {"in": guardrail_ids} if len(guardrail_ids) > 1 else guardrail_ids[0] if policy_id: where["policy_id"] = policy_id if start_date or end_date: @@ -448,9 +422,7 @@ def _build_usage_logs_where( return where -def _usage_log_entry_from_row( - r: Any, sl: Any, action_filter: Optional[str] -) -> Optional[UsageLogEntry]: +def _usage_log_entry_from_row(r: Any, sl: Any, action_filter: Optional[str]) -> Optional[UsageLogEntry]: meta = sl.metadata if isinstance(meta, str): try: @@ -476,9 +448,7 @@ def _usage_log_entry_from_row( duration = entry_for_guardrail.get("duration") if duration is not None: latency_val = round(float(duration) * 1000, 0) - score_val = entry_for_guardrail.get( - "confidence_score" - ) or entry_for_guardrail.get("risk_score") + score_val = entry_for_guardrail.get("confidence_score") or entry_for_guardrail.get("risk_score") if score_val is not None: score_val = round(float(score_val), 2) resp = entry_for_guardrail.get("guardrail_response") @@ -488,11 +458,7 @@ def _usage_log_entry_from_row( reason_val = str(resp)[:500] if action_filter and action_val != action_filter: return None - ts = ( - sl.startTime.isoformat() - if hasattr(sl.startTime, "isoformat") - else str(sl.startTime) - ) + ts = sl.startTime.isoformat() if hasattr(sl.startTime, "isoformat") else str(sl.startTime) return UsageLogEntry( id=r.request_id, timestamp=ts, @@ -590,28 +556,18 @@ async def guardrails_usage_logs( if logical_name and logical_name not in effective_guardrail_ids: effective_guardrail_ids.append(logical_name) - where = _build_usage_logs_where( - effective_guardrail_ids or None, policy_id, start_date, end_date - ) - index_rows = await SpendLogGuardrailIndexRepository( - prisma_client - ).table.find_many( + where = _build_usage_logs_where(effective_guardrail_ids or None, policy_id, start_date, end_date) + index_rows = await SpendLogGuardrailIndexRepository(prisma_client).table.find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, take=page_size + 1, ) - total = await SpendLogGuardrailIndexRepository(prisma_client).table.count( - where=where - ) + total = await SpendLogGuardrailIndexRepository(prisma_client).table.count(where=where) request_ids = [r.request_id for r in index_rows[:page_size]] if not request_ids: - return UsageLogsResponse( - logs=[], total=total, page=page, page_size=page_size - ) - spend_logs = await SpendLogsRepository(prisma_client).table.find_many( - where={"request_id": {"in": request_ids}} - ) + return UsageLogsResponse(logs=[], total=total, page=page, page_size=page_size) + spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}}) log_by_id = {s.request_id: s for s in spend_logs} logs_out: List[UsageLogEntry] = [] for r in index_rows[:page_size]: @@ -621,9 +577,7 @@ async def guardrails_usage_logs( entry = _usage_log_entry_from_row(r, sl, action) if entry is not None: logs_out.append(entry) - return UsageLogsResponse( - logs=logs_out, total=total, page=page, page_size=page_size - ) + return UsageLogsResponse(logs=logs_out, total=total, page=page, page_size=page_size) except Exception as e: from litellm.proxy.utils import handle_exception_on_proxy @@ -648,9 +602,7 @@ async def policies_usage_overview( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return UsageOverviewResponse( - rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0 - ) + return UsageOverviewResponse(rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0) now = datetime.now(timezone.utc) end = end_date or now.strftime("%Y-%m-%d") @@ -661,14 +613,10 @@ async def policies_usage_overview( metrics = await DailyPolicyMetricsRepository(prisma_client).table.find_many( where={"date": {"gte": start, "lte": end}} ) - metrics_prev = await DailyPolicyMetricsRepository( - prisma_client - ).table.find_many( + metrics_prev = await DailyPolicyMetricsRepository(prisma_client).table.find_many( where={ "date": { - "gte": ( - datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7) - ).strftime("%Y-%m-%d"), + "gte": (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d"), "lt": start, } } @@ -678,11 +626,7 @@ async def policies_usage_overview( chart = _chart_from_metrics(metrics) total_requests = sum(a["requests"] for a in agg.values()) total_blocked = sum(a["blocked"] for a in agg.values()) - pass_rate = ( - (100.0 * (total_requests - total_blocked) / total_requests) - if total_requests - else 100.0 - ) + pass_rate = (100.0 * (total_requests - total_blocked) / total_requests) if total_requests else 100.0 rows = _policy_overview_rows(policies, agg, prev_agg) return UsageOverviewResponse( rows=rows, diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index c55c47ca774..eb1979074d4 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -40,9 +40,7 @@ def _parse_guardrail_info_from_payload(payload: Dict[str, Any]) -> List[Dict[str return [] if not isinstance(meta, dict): return [] - info = meta.get("guardrail_information") or meta.get( - "standard_logging_guardrail_information" - ) + info = meta.get("guardrail_information") or meta.get("standard_logging_guardrail_information") if not isinstance(info, list): return [] return info @@ -89,9 +87,7 @@ async def process_spend_logs_guardrail_usage( date_key = _date_str(start_time) for entry in _parse_guardrail_info_from_payload(payload): - guardrail_id = ( - entry.get("guardrail_id") or entry.get("guardrail_name") or "" - ) + guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" if not guardrail_id: continue key = (guardrail_id, date_key) @@ -141,9 +137,7 @@ async def process_spend_logs_guardrail_usage( skip_duplicates=True, ) except Exception as e: - verbose_proxy_logger.debug( - "Guardrail usage tracking: index create_many skipped: %s", e - ) + verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e) # Upsert daily guardrail metrics (counts only; latency/score dropped) for (guardrail_id, date_key), agg in daily_guardrail.items(): @@ -175,6 +169,4 @@ async def process_spend_logs_guardrail_usage( }, ) except Exception as e: - verbose_proxy_logger.warning( - "Guardrail usage tracking failed (non-fatal): %s", e - ) + verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index ed91645eb43..13055775f66 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -43,14 +43,10 @@ MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"] # endpoints that reject unknown fields with 400 "Unknown parameter: # 'max_tokens'". Allow-list so new modes are safe by default. # Per-deployment override: `model_info.health_check_supports_max_tokens`. -_MAX_TOKEN_SUPPORT_MODES: frozenset[str] = frozenset( - {"chat", "completion", "responses"} -) +_MAX_TOKEN_SUPPORT_MODES: frozenset[str] = frozenset({"chat", "completion", "responses"}) -def _resolve_health_check_mode( - model_info: Mapping[str, object], litellm_params: Mapping[str, object] -) -> str | None: +def _resolve_health_check_mode(model_info: Mapping[str, object], litellm_params: Mapping[str, object]) -> str | None: """ Effective mode for a deployment's health-check probe. @@ -72,9 +68,7 @@ def _resolve_health_check_mode( return None -def _should_inject_health_check_max_tokens( - model_info: Mapping[str, object], mode: str | None -) -> bool: +def _should_inject_health_check_max_tokens(model_info: Mapping[str, object], mode: str | None) -> bool: """ Whether the health-check probe should include `max_tokens`. @@ -90,9 +84,7 @@ def _should_inject_health_check_max_tokens( # Health-check modes that forward `reasoning_effort` to the provider (chat-style calls). -_HEALTH_CHECK_MODES_SUPPORTING_REASONING_EFFORT = frozenset( - (None, "chat", "completion") -) +_HEALTH_CHECK_MODES_SUPPORTING_REASONING_EFFORT = frozenset((None, "chat", "completion")) def _get_process_rss_mb() -> Optional[float]: @@ -211,9 +203,7 @@ async def _run_model_health_check(model: dict): ) -async def _run_health_checks_with_bounded_concurrency( - models: list, concurrency_limit: int -) -> tuple[list, int]: +async def _run_health_checks_with_bounded_concurrency(models: list, concurrency_limit: int) -> tuple[list, int]: """ Run health checks with at most `concurrency_limit` active tasks. Preserves result ordering to match `models`. @@ -274,13 +264,9 @@ async def _perform_health_check( peak_in_flight = 0 if isinstance(max_concurrency, int) and max_concurrency > 0: dispatch_mode = "bounded" - results, peak_in_flight = await _run_health_checks_with_bounded_concurrency( - model_list, max_concurrency - ) + results, peak_in_flight = await _run_health_checks_with_bounded_concurrency(model_list, max_concurrency) else: - tasks = [ - asyncio.create_task(_run_model_health_check(model)) for model in model_list - ] + tasks = [asyncio.create_task(_run_model_health_check(model)) for model in model_list] peak_in_flight = len(tasks) results = await asyncio.gather(*tasks, return_exceptions=True) @@ -330,9 +316,7 @@ async def _perform_health_check( cleaned["model_id"] = _model_id if isinstance(is_healthy, Exception): exceptions_by_model_id[_model_id] = is_healthy - cleaned["exception_status"] = getattr( - is_healthy, "status_code", 500 - ) + cleaned["exception_status"] = getattr(is_healthy, "status_code", 500) unhealthy_endpoints.append(cleaned) return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id @@ -389,9 +373,7 @@ def _health_check_deployment_is_wildcard(litellm_params: dict) -> bool: return "*" in _deployment_model_string_for_health_check(litellm_params) -def _resolve_health_check_max_tokens( - model_info: dict, litellm_params: dict -) -> Optional[int]: +def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> Optional[int]: """ Pick max_tokens for the health check request. @@ -418,9 +400,7 @@ def _resolve_health_check_max_tokens( except Exception: is_reasoning = False tokens_reasoning = model_info.get("health_check_max_tokens_reasoning", None) - tokens_non_reasoning = model_info.get( - "health_check_max_tokens_non_reasoning", None - ) + tokens_non_reasoning = model_info.get("health_check_max_tokens_non_reasoning", None) if tokens_reasoning is not None or tokens_non_reasoning is not None: if is_reasoning and tokens_reasoning is not None: return int(tokens_reasoning) @@ -438,9 +418,7 @@ def _resolve_health_check_max_tokens( return None -def _update_litellm_params_for_health_check( - model_info: dict, litellm_params: dict -) -> dict: +def _update_litellm_params_for_health_check(model_info: dict, litellm_params: dict) -> dict: """ Update the litellm params for health check. @@ -463,9 +441,7 @@ def _update_litellm_params_for_health_check( model_info, mode, # any-ok: untyped router config dict ): - _resolved_max_tokens = _resolve_health_check_max_tokens( - model_info, litellm_params - ) + _resolved_max_tokens = _resolve_health_check_max_tokens(model_info, litellm_params) if _resolved_max_tokens is not None: litellm_params["max_tokens"] = _resolved_max_tokens @@ -556,9 +532,7 @@ async def perform_health_check( if not model_list: if cli_model: - model_list = [ - {"model_name": cli_model, "litellm_params": {"model": cli_model}} - ] + model_list = [{"model_name": cli_model, "litellm_params": {"model": cli_model}}] else: if instrumentation_enabled: logger.debug( @@ -573,26 +547,18 @@ async def perform_health_check( # Filter by model_id first so a single deployment is checked when id is specified if model_id is not None: - _by_id = [ - x for x in model_list if (x.get("model_info") or {}).get("id") == model_id - ] + _by_id = [x for x in model_list if (x.get("model_info") or {}).get("id") == model_id] if _by_id: model_list = _by_id elif model is not None: - _new_model_list = [ - x for x in model_list if x["litellm_params"]["model"] == model - ] + _new_model_list = [x for x in model_list if x["litellm_params"]["model"] == model] if _new_model_list == []: _new_model_list = [x for x in model_list if x["model_name"] == model] model_list = _new_model_list if health_check_skip_disabled_background_models: model_list = [ - x - for x in model_list - if not (x.get("model_info") or {}).get( - "disable_background_health_check", False - ) + x for x in model_list if not (x.get("model_info") or {}).get("disable_background_health_check", False) ] if not model_list: if instrumentation_enabled: diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index 5c5f8929a34..eb6d27ca20b 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -76,13 +76,9 @@ class SharedHealthCheckManager: ) if acquired: - verbose_proxy_logger.info( - "Pod %s acquired health check lock", self.pod_id - ) + verbose_proxy_logger.info("Pod %s acquired health check lock", self.pod_id) else: - verbose_proxy_logger.debug( - "Pod %s failed to acquire health check lock", self.pod_id - ) + verbose_proxy_logger.debug("Pod %s failed to acquire health check lock", self.pod_id) return bool(acquired) except Exception as e: @@ -100,9 +96,7 @@ class SharedHealthCheckManager: current_owner = await self.redis_cache.async_get_cache(lock_key) if current_owner == self.pod_id: await self.redis_cache.async_delete_cache(lock_key) - verbose_proxy_logger.info( - "Pod %s released health check lock", self.pod_id - ) + verbose_proxy_logger.info("Pod %s released health check lock", self.pod_id) except Exception as e: verbose_proxy_logger.error("Error releasing health check lock: %s", str(e)) @@ -141,9 +135,7 @@ class SharedHealthCheckManager: return cached_results except Exception as e: - verbose_proxy_logger.error( - "Error getting cached health check results: %s", str(e) - ) + verbose_proxy_logger.error("Error getting cached health check results: %s", str(e)) return None async def cache_health_check_results( @@ -246,9 +238,7 @@ class SharedHealthCheckManager: ) # Cache the results - await self.cache_health_check_results( - healthy_endpoints, unhealthy_endpoints - ) + await self.cache_health_check_results(healthy_endpoints, unhealthy_endpoints) return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id @@ -269,9 +259,7 @@ class SharedHealthCheckManager: # Lock not acquired — poll for cached results until the lock # holder finishes or the lock expires, rather than falling back # to a redundant local health check after only 2 seconds. - verbose_proxy_logger.debug( - "Pod %s waiting for other pod to complete health check", self.pod_id - ) + verbose_proxy_logger.debug("Pod %s waiting for other pod to complete health check", self.pod_id) poll_interval = 5 # seconds between cache checks max_wait = self.lock_ttl # wait at most as long as the lock can live @@ -338,9 +326,7 @@ class SharedHealthCheckManager: current_owner = await self.redis_cache.async_get_cache(lock_key) return current_owner is not None and current_owner != self.pod_id except Exception as e: - verbose_proxy_logger.error( - "Error checking health check lock status: %s", str(e) - ) + verbose_proxy_logger.error("Error checking health check lock status: %s", str(e)) return False async def get_health_check_status(self) -> Dict[str, Any]: @@ -369,9 +355,7 @@ class SharedHealthCheckManager: cached_results = await self.get_cached_health_check_results() status["cache_available"] = cached_results is not None if cached_results: - status["cache_age_seconds"] = time.time() - cached_results.get( - "timestamp", 0 - ) + status["cache_age_seconds"] = time.time() - cached_results.get("timestamp", 0) status["last_checked_by"] = cached_results.get("checked_by") except Exception as e: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 928d00109fa..c20991b8d43 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -72,9 +72,7 @@ def _reject_os_environ_references(params: dict) -> None: if isinstance(value, str) and value.startswith("os.environ/"): raise HTTPException( status_code=400, - detail={ - "error": "Environment variable references are not permitted in request parameters." - }, + detail={"error": "Environment variable references are not permitted in request parameters."}, ) if isinstance(value, (dict, list)) and id(value) not in seen: seen.add(id(value)) @@ -102,13 +100,8 @@ def get_callback_identifier(callback): if hasattr(callback, "callback_name") and callback.callback_name: return callback.callback_name if hasattr(callback, "__class__"): - callback_strs = CustomLoggerRegistry.get_all_callback_strs_from_class_type( - callback.__class__ - ) - if ( - hasattr(callback, "callback_name") - and callback.callback_name in callback_strs - ): + callback_strs = CustomLoggerRegistry.get_all_callback_strs_from_class_type(callback.__class__) + if hasattr(callback, "callback_name") and callback.callback_name in callback_strs: return callback.callback_name if callback_strs: return callback_strs[0] @@ -188,9 +181,7 @@ async def health_services_endpoint( ) if service is None: - raise HTTPException( - status_code=400, detail={"error": "Service must be specified."} - ) + raise HTTPException(status_code=400, detail={"error": "Service must be specified."}) if service not in [ "slack_budget_alerts", @@ -215,9 +206,7 @@ async def health_services_endpoint( ]: raise HTTPException( status_code=400, - detail={ - "error": f"Service must be in list. Service={service} not in {services}" - }, + detail={"error": f"Service must be in list. Service={service} not in {services}"}, ) service_in_success_callbacks = False @@ -256,11 +245,7 @@ async def health_services_endpoint( response = await datadog_logger.async_health_check() return { "status": response["status"], - "message": ( - response["error_message"] - if response["status"] == "unhealthy" - else "Datadog is healthy" - ), + "message": (response["error_message"] if response["status"] == "unhealthy" else "Datadog is healthy"), } elif service == "datadog_metrics": from litellm.integrations.datadog.datadog_metrics import ( @@ -270,21 +255,15 @@ async def health_services_endpoint( get_custom_logger_compatible_class, ) - datadog_metrics_logger = get_custom_logger_compatible_class( - "datadog_metrics" - ) + datadog_metrics_logger = get_custom_logger_compatible_class("datadog_metrics") if datadog_metrics_logger is None: - datadog_metrics_logger = DatadogMetricsLogger( - start_periodic_flush=False - ) + datadog_metrics_logger = DatadogMetricsLogger(start_periodic_flush=False) assert isinstance(datadog_metrics_logger, DatadogMetricsLogger) response = await datadog_metrics_logger.async_health_check() return { "status": response["status"], "message": ( - response["error_message"] - if response["status"] == "unhealthy" - else "Datadog Metrics is healthy" + response["error_message"] if response["status"] == "unhealthy" else "Datadog Metrics is healthy" ), } elif service == "arize": @@ -294,11 +273,7 @@ async def health_services_endpoint( response = await arize_logger.async_health_check() return { "status": response["status"], - "message": ( - response["error_message"] - if response["status"] == "unhealthy" - else "Arize is healthy" - ), + "message": (response["error_message"] if response["status"] == "unhealthy" else "Arize is healthy"), } elif service == "galileo": from litellm.integrations.galileo import GalileoObserve @@ -307,11 +282,7 @@ async def health_services_endpoint( response = await galileo_logger.async_health_check() return { "status": response["status"], - "message": ( - response["error_message"] - if response["status"] == "unhealthy" - else "Galileo is healthy" - ), + "message": (response["error_message"] if response["status"] == "unhealthy" else "Galileo is healthy"), } elif service == "langfuse": from litellm.integrations.langfuse.langfuse import LangFuseLogger @@ -332,9 +303,7 @@ async def health_services_endpoint( if not _is_proxy_admin(user_api_key_dict): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Only proxy admins can trigger the New Relic test event." - }, + detail={"error": "Only proxy admins can trigger the New Relic test event."}, ) from litellm.integrations.newrelic.newrelic import NewRelicLogger @@ -377,19 +346,12 @@ async def health_services_endpoint( if "slack" in general_settings.get("alerting", []): # test_message = f"""\n🚨 `ProjectedLimitExceededError` 💸\n\n`Key Alias:` litellm-ui-test-alert \n`Expected Day of Error`: 28th March \n`Current Spend`: $100.00 \n`Projected Spend at end of month`: $1000.00 \n`Soft Limit`: $700""" # check if user has opted into unique_alert_webhooks - if ( - proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url - is not None - ): - for ( - alert_type - ) in proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url: + if proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url is not None: + for alert_type in proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url: # only test alert if it's in active alert types if ( - proxy_logging_obj.slack_alerting_instance.alert_types - is not None - and alert_type - not in proxy_logging_obj.slack_alerting_instance.alert_types + proxy_logging_obj.slack_alerting_instance.alert_types is not None + and alert_type not in proxy_logging_obj.slack_alerting_instance.alert_types ): continue @@ -422,16 +384,10 @@ async def health_services_endpoint( ) if prisma_client is not None: - asyncio.create_task( - proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report() - ) - asyncio.create_task( - proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report() - ) + asyncio.create_task(proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report()) + asyncio.create_task(proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report()) - alert_types = ( - proxy_logging_obj.slack_alerting_instance.alert_types or [] - ) + alert_types = proxy_logging_obj.slack_alerting_instance.alert_types or [] alert_types = list(alert_types) return { "status": "success", @@ -442,9 +398,7 @@ async def health_services_endpoint( raise HTTPException( status_code=422, detail={ - "error": '"{}" not in proxy config: general_settings. Unable to test this.'.format( - service - ) + "error": '"{}" not in proxy config: general_settings. Unable to test this.'.format(service) }, ) if service == "email": @@ -473,9 +427,7 @@ async def health_services_endpoint( except Exception as e: verbose_proxy_logger.error( - "litellm.proxy.proxy_server.health_services_endpoint(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.health_services_endpoint(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): @@ -555,9 +507,7 @@ async def _save_health_check_to_db( checked_by=user_id, ) except Exception as db_error: - verbose_proxy_logger.warning( - f"Failed to save health check to database for model {model_name}: {db_error}" - ) + verbose_proxy_logger.warning(f"Failed to save health check to database for model {model_name}: {db_error}") # Continue execution - don't let database save failure break health checks @@ -690,9 +640,7 @@ async def _save_health_check_results_if_changed( if last_check.checked_at: from datetime import datetime, timezone - time_since_last_check = ( - datetime.now(timezone.utc) - last_check.checked_at - ).total_seconds() + time_since_last_check = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds() # Only skip if status unchanged AND checked recently (within 1 hour) # This ensures we still get periodic updates even if status is stable if time_since_last_check < 3600: # 1 hour threshold @@ -763,9 +711,7 @@ async def _save_background_health_checks_to_db( checked_by, ) except Exception as db_error: - verbose_proxy_logger.warning( - f"Failed to save background health checks to database: {db_error}" - ) + verbose_proxy_logger.warning(f"Failed to save background health checks to database: {db_error}") # Continue execution - don't let database save failure break health checks @@ -809,20 +755,11 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict: for key in ("healthy_endpoints", "unhealthy_endpoints"): eps = out.get(key) if isinstance(eps, list): - out[key] = [ - ( - {k: v for k, v in ep.items() if k not in drop} - if isinstance(ep, dict) - else ep - ) - for ep in eps - ] + out[key] = [({k: v for k, v in ep.items() if k not in drop} if isinstance(ep, dict) else ep) for ep in eps] return out -def _resolve_targeted_model_ids( - model_list: list, model: Optional[str], model_id: Optional[str] -) -> Optional[set]: +def _resolve_targeted_model_ids(model_list: list, model: Optional[str], model_id: Optional[str]) -> Optional[set]: """ Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of deployment IDs the response should be scoped to. @@ -859,9 +796,7 @@ def _resolve_targeted_model_ids( return target_ids -def _filter_health_check_results_by_model_ids( - results: dict, allowed_model_ids: set -) -> dict: +def _filter_health_check_results_by_model_ids(results: dict, allowed_model_ids: set) -> dict: """ Restrict a cached background health-check result dict to endpoints whose model_id is in ``allowed_model_ids``. @@ -874,15 +809,9 @@ def _filter_health_check_results_by_model_ids( downstream transform (e.g. _strip_admin_only_fields_from_health_result) cannot accidentally mutate the shared ``health_check_results`` cache. """ - healthy = [ - dict(ep) - for ep in (results.get("healthy_endpoints") or []) - if ep.get("model_id") in allowed_model_ids - ] + healthy = [dict(ep) for ep in (results.get("healthy_endpoints") or []) if ep.get("model_id") in allowed_model_ids] unhealthy = [ - dict(ep) - for ep in (results.get("unhealthy_endpoints") or []) - if ep.get("model_id") in allowed_model_ids + dict(ep) for ep in (results.get("unhealthy_endpoints") or []) if ep.get("model_id") in allowed_model_ids ] return { "healthy_endpoints": healthy, @@ -956,9 +885,7 @@ def _health_endpoint_resolve_target_model_name( try: deployment = llm_router.get_deployment(model_id=model_id) except Exception as e: - verbose_proxy_logger.error( - f"Error getting deployment for model_id {model_id}: {e}" - ) + verbose_proxy_logger.error(f"Error getting deployment for model_id {model_id}: {e}") raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={"error": f"Model with ID {model_id} not found"}, @@ -975,12 +902,8 @@ def _health_endpoint_resolve_target_model_name( async def health_endpoint( response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - model: Optional[str] = fastapi.Query( - None, description="Specify the model name (optional)" - ), - model_id: Optional[str] = fastapi.Query( - None, description="Specify the model ID (optional)" - ), + model: Optional[str] = fastapi.Query(None, description="Specify the model name (optional)"), + model_id: Optional[str] = fastapi.Query(None, description="Specify the model ID (optional)"), ): """ 🚨 USE `/health/liveliness` to health check the proxy 🚨 @@ -1019,9 +942,7 @@ async def health_endpoint( _hc_filter = health_check_filter_kwargs_from_general_settings(general_settings) start_time = time.time() - target_model = _health_endpoint_resolve_target_model_name( - model, model_id, llm_router - ) + target_model = _health_endpoint_resolve_target_model_name(model, model_id, llm_router) is_admin = _is_proxy_admin(user_api_key_dict) model_specific_request = bool(model or model_id) @@ -1040,9 +961,7 @@ async def health_endpoint( response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE if is_admin: return result - response.headers["Litellm-Health-Field-Notice"] = ( - "api_base and api_version are admin-only on this endpoint" - ) + response.headers["Litellm-Health-Field-Notice"] = "api_base and api_version are admin-only on this endpoint" return _strip_admin_only_fields_from_health_result(result) try: @@ -1084,20 +1003,14 @@ async def health_endpoint( # stays in the list, matching nothing; denied rather than # unrestricted, mirroring _resolve_key_models_for_auth_check. accessible_models = list(user_api_key_dict.models) - if ( - SpecialModelNames.all_team_models.value in accessible_models - and user_api_key_dict.team_id is not None - ): + if SpecialModelNames.all_team_models.value in accessible_models and user_api_key_dict.team_id is not None: accessible_models = list(user_api_key_dict.team_models) restrict_to_allowed_models = ( - len(accessible_models) > 0 - and SpecialModelNames.all_proxy_models.value not in accessible_models + len(accessible_models) > 0 and SpecialModelNames.all_proxy_models.value not in accessible_models ) if restrict_to_allowed_models: allowed_models = set(accessible_models) - _llm_model_list = [ - m for m in _llm_model_list if m.get("model_name") in allowed_models - ] + _llm_model_list = [m for m in _llm_model_list if m.get("model_name") in allowed_models] if use_background_health_checks: # The cached background result covers every model. When the # caller targets a specific model/model_id we have to narrow the @@ -1115,12 +1028,8 @@ async def health_endpoint( # _llm_model_list is already scoped to the caller's allowed # model_names above, so targeted_ids is implicitly the # intersection of "targeted" and "allowed." - filter_ids = ( - targeted_ids if targeted_ids is not None else allowed_model_ids - ) - filtered = _filter_health_check_results_by_model_ids( - health_check_results, filter_ids - ) + filter_ids = targeted_ids if targeted_ids is not None else allowed_model_ids + filtered = _filter_health_check_results_by_model_ids(health_check_results, filter_ids) if targeted_ids is None and not allowed_model_ids: # Caller has accessible model_names but none of the # matching deployments expose a model_info.id, so the @@ -1146,11 +1055,7 @@ async def health_endpoint( # Admin caller targeting a specific model: filter the cache # so the response (and the targeted-503 check) reflects only # that deployment, not the global aggregate. - return _post_process( - _filter_health_check_results_by_model_ids( - health_check_results, targeted_ids - ) - ) + return _post_process(_filter_health_check_results_by_model_ids(health_check_results, targeted_ids)) return _post_process(health_check_results) else: router_result = await _perform_health_check_and_save( @@ -1168,28 +1073,18 @@ async def health_endpoint( return _post_process(router_result) except Exception as e: verbose_proxy_logger.error( - "litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) raise e -@router.get( - "/health/history", tags=["health"], dependencies=[Depends(user_api_key_auth)] -) +@router.get("/health/history", tags=["health"], dependencies=[Depends(user_api_key_auth)]) async def health_check_history_endpoint( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - model: Optional[str] = fastapi.Query( - None, description="Filter by specific model name" - ), - status_filter: Optional[str] = fastapi.Query( - None, description="Filter by status (healthy/unhealthy)" - ), - limit: int = fastapi.Query( - 100, description="Number of records to return", ge=1, le=1000 - ), + model: Optional[str] = fastapi.Query(None, description="Filter by specific model name"), + status_filter: Optional[str] = fastapi.Query(None, description="Filter by status (healthy/unhealthy)"), + limit: int = fastapi.Query(100, description="Number of records to return", ge=1, le=1000), offset: int = fastapi.Query(0, description="Number of records to skip", ge=0), ): """ @@ -1224,9 +1119,7 @@ async def health_check_history_endpoint( ) -@router.get( - "/health/latest", tags=["health"], dependencies=[Depends(user_api_key_auth)] -) +@router.get("/health/latest", tags=["health"], dependencies=[Depends(user_api_key_auth)]) async def latest_health_checks_endpoint( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -1242,9 +1135,7 @@ async def latest_health_checks_endpoint( # Convert to dict format for JSON response using helper function checks_data = { - ( - check.model_id if check.model_id else check.model_name - ): _convert_health_check_to_dict(check) + (check.model_id if check.model_id else check.model_name): _convert_health_check_to_dict(check) for check in latest_checks } @@ -1260,9 +1151,7 @@ async def latest_health_checks_endpoint( ) -@router.get( - "/health/shared-status", tags=["health"], dependencies=[Depends(user_api_key_auth)] -) +@router.get("/health/shared-status", tags=["health"], dependencies=[Depends(user_api_key_auth)]) async def shared_health_check_status_endpoint( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -1301,18 +1190,14 @@ async def shared_health_check_status_endpoint( verbose_proxy_logger.error(f"Error getting shared health check status: {e}") raise HTTPException( status_code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={ - "error": f"Failed to retrieve shared health check status: {str(e)}" - }, + detail={"error": f"Failed to retrieve shared health check status: {str(e)}"}, ) def _read_license_data() -> Optional[Dict[str, Any]]: from litellm.proxy.proxy_server import _license_check, premium_user_data - license_data: Optional[EnterpriseLicenseData] = ( - premium_user_data or _license_check.airgapped_license_data - ) + license_data: Optional[EnterpriseLicenseData] = premium_user_data or _license_check.airgapped_license_data if ( license_data is None @@ -1396,9 +1281,7 @@ async def _db_health_readiness_check(): try: time_diff = datetime.now() - db_health_cache["last_updated"] - if db_health_cache["status"] == "connected" and time_diff < timedelta( - seconds=15 - ): + if db_health_cache["status"] == "connected" and time_diff < timedelta(seconds=15): return db_health_cache if prisma_client is None: @@ -1412,25 +1295,17 @@ async def _db_health_readiness_check(): db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} if PrismaDBExceptionHandler.is_database_transport_error(e): try: - verbose_proxy_logger.warning( - "_db_health_readiness_check: health_check failed, attempting reconnect" - ) - await prisma_client.attempt_db_reconnect( - reason="health_readiness_check" - ) + verbose_proxy_logger.warning("_db_health_readiness_check: health_check failed, attempting reconnect") + await prisma_client.attempt_db_reconnect(reason="health_readiness_check") await prisma_client.health_check() - verbose_proxy_logger.info( - "_db_health_readiness_check: reconnect succeeded" - ) + verbose_proxy_logger.info("_db_health_readiness_check: reconnect succeeded") db_health_cache = { "status": "connected", "last_updated": datetime.now(), } return db_health_cache except Exception: - verbose_proxy_logger.error( - "_db_health_readiness_check: reconnect failed" - ) + verbose_proxy_logger.error("_db_health_readiness_check: reconnect failed") return db_health_cache @@ -1541,9 +1416,7 @@ async def _get_health_readiness_details( try: # this was returning a JSON of the values in some of the callbacks # all we need is the callback name, hence we do str(callback) - success_callback_names = [ - callback_name(x) for x in litellm.success_callback - ] + success_callback_names = [callback_name(x) for x in litellm.success_callback] except AttributeError: # don't let this block the /health/readiness response, if we can't convert to str -> return litellm.success_callback success_callback_names = litellm.success_callback @@ -1948,14 +1821,10 @@ async def test_model_connection( try: deployment_by_id = None if request_model_id: - deployment_by_id = llm_router.get_deployment( - model_id=request_model_id - ) + deployment_by_id = llm_router.get_deployment(model_id=request_model_id) if deployment_by_id is not None: - config_litellm_params = deployment_by_id.litellm_params.model_dump( - exclude_none=True - ) + config_litellm_params = deployment_by_id.litellm_params.model_dump(exclude_none=True) elif model_name: # Fall back to model_name lookup for callers (e.g. the # "Add Model" wizard, or curl) that don't supply an id. @@ -1968,10 +1837,7 @@ async def test_model_connection( all_deployments = llm_router.get_model_list(model_name=None) if all_deployments: for deployment in all_deployments: - if ( - deployment.get("litellm_params", {}).get("model") - == model_name - ): + if deployment.get("litellm_params", {}).get("model") == model_name: deployments = [deployment] break @@ -1979,13 +1845,10 @@ async def test_model_connection( # Use the first deployment's litellm_params as base # config. These already have resolved environment # variables from proxy config. - config_litellm_params = dict( - deployments[0].get("litellm_params", {}) - ) + config_litellm_params = dict(deployments[0].get("litellm_params", {})) except Exception as e: verbose_proxy_logger.debug( - f"Could not find model {model_name} in router: {e}. " - "Proceeding with request params only." + f"Could not find model {model_name} in router: {e}. Proceeding with request params only." ) # Merge: config params (from proxy config) as base, request params override @@ -2021,9 +1884,7 @@ async def test_model_connection( ) # Clean the result for display - cleaned_result = _clean_endpoint_data( - {**litellm_params, **result}, details=True - ) + cleaned_result = _clean_endpoint_data({**litellm_params, **result}, details=True) return { "status": "error" if "error" in result else "success", diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 0db661fb508..d729e339a51 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -47,9 +47,7 @@ def get_proxy_hook( Factory method to get a proxy hook instance by name """ if hook_name not in PROXY_HOOKS: - raise ValueError( - f"Unknown hook: {hook_name}. Available hooks: {list(PROXY_HOOKS.keys())}" - ) + raise ValueError(f"Unknown hook: {hook_name}. Available hooks: {list(PROXY_HOOKS.keys())}") return PROXY_HOOKS[hook_name] diff --git a/litellm/proxy/hooks/azure_content_safety.py b/litellm/proxy/hooks/azure_content_safety.py index ddcd1540597..c41effb4783 100644 --- a/litellm/proxy/hooks/azure_content_safety.py +++ b/litellm/proxy/hooks/azure_content_safety.py @@ -42,9 +42,7 @@ class _PROXY_AzureContentSafety( self.thresholds = self._configure_thresholds(thresholds) - self.client = ContentSafetyClient( - self.endpoint, AzureKeyCredential(self.api_key) - ) + self.client = ContentSafetyClient(self.endpoint, AzureKeyCredential(self.api_key)) def _configure_thresholds(self, thresholds=None): default_thresholds = { @@ -66,9 +64,7 @@ class _PROXY_AzureContentSafety( def _compute_result(self, response): result = {} - category_severity = { - item.category: item.severity for item in response.categories_analysis - } + category_severity = {item.category: item.severity for item in response.categories_analysis} for category in self.text_category: severity = category_severity.get(category) if severity is not None: @@ -92,9 +88,7 @@ class _PROXY_AzureContentSafety( try: response = await self.client.analyze_text(request) except self.azure_http_error: - verbose_proxy_logger.debug( - "Error in Azure Content-Safety: %s", traceback.format_exc() - ) + verbose_proxy_logger.debug("Error in Azure Content-Safety: %s", traceback.format_exc()) verbose_proxy_logger.debug(traceback.format_exc()) raise diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 91c604e6204..3477af36285 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -257,18 +257,11 @@ class _PROXY_BatchRateLimiter(CustomLogger): if general_settings.get("disable_batch_input_file_rate_limiting") is True: return True, None - skip_providers = ( - general_settings.get("skip_batch_input_file_rate_limiting_for_providers") - or [] - ) + skip_providers = general_settings.get("skip_batch_input_file_rate_limiting_for_providers") or [] if skip_providers: - batch_provider = self._resolve_batch_provider( - self._get_batch_routing_model(data) - ) + batch_provider = self._resolve_batch_provider(self._get_batch_routing_model(data)) if batch_provider and batch_provider in skip_providers: - verbose_proxy_logger.debug( - f"Skipping batch input file processing for provider={batch_provider}" - ) + verbose_proxy_logger.debug(f"Skipping batch input file processing for provider={batch_provider}") return True, None descriptors = self._create_batch_rate_limit_descriptors( @@ -276,16 +269,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): data=data, ) if not self._has_applicable_batch_rate_limits(descriptors): - verbose_proxy_logger.debug( - "Skipping batch input file processing: no rate limits configured" - ) + verbose_proxy_logger.debug("Skipping batch input file processing: no rate limits configured") return True, None return False, descriptors - def _warn_if_unsupported_model_skip_configured( - self, general_settings: Dict - ) -> None: + def _warn_if_unsupported_model_skip_configured(self, general_settings: Dict) -> None: """Warn once that ``skip_batch_input_file_rate_limiting_for_models`` is a no-op. A per-model skip is intentionally not honored because the model a batch @@ -406,25 +395,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Find the descriptor for this status descriptor_index = next( - ( - i - for i, d in enumerate(descriptors) - if d.get("key") == status.get("descriptor_key") - ), + (i for i, d in enumerate(descriptors) if d.get("key") == status.get("descriptor_key")), 0, ) descriptor: RateLimitDescriptor = ( - descriptors[descriptor_index] - if descriptors - else {"key": "", "value": "", "rate_limit": None} + descriptors[descriptor_index] if descriptors else {"key": "", "value": "", "rate_limit": None} ) now = datetime.now().timestamp() window_size = self.parallel_request_limiter.window_size reset_time = now + window_size - reset_time_formatted = datetime.fromtimestamp(reset_time).strftime( - "%Y-%m-%d %H:%M:%S UTC" - ) + reset_time_formatted = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC") remaining_display = max(0, status["limit_remaining"]) current_limit = status["current_limit"] @@ -444,9 +425,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"Limit resets at: {reset_time_formatted}" ) - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - requested_model - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(requested_model) raise ProxyRateLimitError( detail=detail, headers={ @@ -488,16 +467,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): "requests": batch_usage.request_count, "tokens": batch_usage.total_tokens, } - increments: List[Dict[Literal["requests", "tokens"], int]] = [ - increment for _ in descriptors - ] + increments: List[Dict[Literal["requests", "tokens"], int]] = [increment for _ in descriptors] - rate_limit_response = ( - await self.parallel_request_limiter.atomic_check_and_increment_by_n( - descriptors=descriptors, - increments=increments, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) + rate_limit_response = await self.parallel_request_limiter.atomic_check_and_increment_by_n( + descriptors=descriptors, + increments=increments, + parent_otel_span=user_api_key_dict.parent_otel_span, ) if rate_limit_response["overall_code"] == "OVER_LIMIT": @@ -542,23 +517,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): is_managed_file = _is_base64_encoded_unified_file_id(file_id) # For managed files the unified file id encodes the proxy model # alias(es) the file was uploaded for; auth validates against those. - target_model_names = ( - get_models_from_unified_file_id(is_managed_file) - if is_managed_file - else [] - ) + target_model_names = get_models_from_unified_file_id(is_managed_file) if is_managed_file else [] if is_managed_file and user_api_key_dict is not None: file_content = await self._fetch_managed_file_content( file_id=file_id, user_api_key_dict=user_api_key_dict, ) else: - provider_file_id, fetch_kwargs = ( - self._resolve_batch_input_file_fetch_params( - file_id=file_id, - custom_llm_provider=custom_llm_provider, - data=data or {}, - ) + provider_file_id, fetch_kwargs = self._resolve_batch_input_file_fetch_params( + file_id=file_id, + custom_llm_provider=custom_llm_provider, + data=data or {}, ) # For non-managed files, use the standard litellm.afile_content file_content = await litellm.afile_content( @@ -570,8 +539,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_content_bytes = getattr(file_content, "content", None) if not isinstance(file_content_bytes, bytes): raise ValueError( - f"Expected bytes content from file retrieval for {file_id}, " - f"got {type(file_content_bytes)}" + f"Expected bytes content from file retrieval for {file_id}, got {type(file_content_bytes)}" ) # Single streaming pass over the JSONL lines, accounting each row @@ -637,9 +605,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) raise except Exception as e: - verbose_proxy_logger.error( - f"Error counting input file usage for {file_id}: {str(e)}" - ) + verbose_proxy_logger.error(f"Error counting input file usage for {file_id}: {str(e)}") raise async def _enforce_batch_file_model_access( @@ -696,10 +662,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): raise HTTPException( status_code=403, detail={ - "error": ( - "Batch input file model access could not be " - "validated against the current team." - ) + "error": ("Batch input file model access could not be validated against the current team.") }, ) from e @@ -785,15 +748,11 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Get the managed files hook if proxy_logging_obj is None: - raise ValueError( - "proxy_logging_obj not available. Cannot access managed files hook." - ) + raise ValueError("proxy_logging_obj not available. Cannot access managed files hook.") managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: - raise ValueError( - "Managed files hook not found. Cannot access managed file." - ) + raise ValueError("Managed files hook not found. Cannot access managed file.") if not isinstance(managed_files_obj, BaseFileEndpoints): raise ValueError("Managed files hook is not a BaseFileEndpoints instance.") @@ -845,23 +804,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) return data - verbose_proxy_logger.debug( - "Batch rate limiter: Handling batch creation rate limiting" - ) + verbose_proxy_logger.debug("Batch rate limiter: Handling batch creation rate limiting") try: # Extract input_file_id from data input_file_id = data.get("input_file_id") if not input_file_id: - verbose_proxy_logger.debug( - "No input_file_id in batch request, skipping rate limiting" - ) + verbose_proxy_logger.debug("No input_file_id in batch request, skipping rate limiting") return data - should_skip, batch_rate_limit_descriptors = ( - self._should_skip_batch_input_file_processing( - data=data, user_api_key_dict=user_api_key_dict - ) + should_skip, batch_rate_limit_descriptors = self._should_skip_batch_input_file_processing( + data=data, user_api_key_dict=user_api_key_dict ) if should_skip: return data @@ -870,9 +823,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): custom_llm_provider = data.get("custom_llm_provider", "openai") # Count tokens and requests from input file - verbose_proxy_logger.debug( - f"Counting tokens from batch input file: {input_file_id}" - ) + verbose_proxy_logger.debug(f"Counting tokens from batch input file: {input_file_id}") batch_usage = await self.count_input_file_usage( file_id=input_file_id, custom_llm_provider=custom_llm_provider, @@ -881,8 +832,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) verbose_proxy_logger.debug( - f"Batch input file usage - Tokens: {batch_usage.total_tokens}, " - f"Requests: {batch_usage.request_count}" + f"Batch input file usage - Tokens: {batch_usage.total_tokens}, Requests: {batch_usage.request_count}" ) # Store batch usage in data for later reference @@ -898,17 +848,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): descriptors=batch_rate_limit_descriptors, ) - verbose_proxy_logger.debug( - "Batch rate limit check passed, counters incremented" - ) + verbose_proxy_logger.debug("Batch rate limit check passed, counters incremented") return data except HTTPException: # Re-raise HTTP exceptions (rate limit exceeded) raise except Exception as e: - verbose_proxy_logger.error( - f"Error in batch rate limiting: {str(e)}", exc_info=True - ) + verbose_proxy_logger.error(f"Error in batch rate limiting: {str(e)}", exc_info=True) # Don't block the request if rate limiting fails return data diff --git a/litellm/proxy/hooks/batch_redis_get.py b/litellm/proxy/hooks/batch_redis_get.py index f734b19681d..8b11e185fb9 100644 --- a/litellm/proxy/hooks/batch_redis_get.py +++ b/litellm/proxy/hooks/batch_redis_get.py @@ -25,9 +25,7 @@ class _PROXY_BatchRedisRequests(CustomLogger): self.async_get_cache ) # map the litellm 'get_cache' function to our custom function - def print_verbose( - self, print_statement, debug_level: Literal["INFO", "DEBUG"] = "DEBUG" - ): + def print_verbose(self, print_statement, debug_level: Literal["INFO", "DEBUG"] = "DEBUG"): if debug_level == "DEBUG": verbose_proxy_logger.debug(print_statement) elif debug_level == "INFO": @@ -66,41 +64,29 @@ class _PROXY_BatchRedisRequests(CustomLogger): - Check if `litellm.Cache` is redis - Get the relevant values """ - if litellm.cache.type is not None and isinstance( - litellm.cache.cache, RedisCache - ): + if litellm.cache.type is not None and isinstance(litellm.cache.cache, RedisCache): # Initialize an empty list to store the keys keys = [] self.print_verbose(f"cache_key_name: {cache_key_name}") # Use the SCAN iterator to fetch keys matching the pattern - keys = await litellm.cache.cache.async_scan_iter( - pattern=cache_key_name, count=100 - ) + keys = await litellm.cache.cache.async_scan_iter(pattern=cache_key_name, count=100) # If you need the truly "last" based on time or another criteria, # ensure your key naming or storage strategy allows this determination # Here you would sort or filter the keys as needed based on your strategy self.print_verbose(f"redis keys: {keys}") if len(keys) > 0: - key_value_dict = ( - await litellm.cache.cache.async_batch_get_cache( - key_list=keys - ) - ) + key_value_dict = await litellm.cache.cache.async_batch_get_cache(key_list=keys) ## Add to cache if len(key_value_dict.items()) > 0: - await cache.in_memory_cache.async_set_cache_pipeline( - cache_list=list(key_value_dict.items()), ttl=60 - ) + await cache.in_memory_cache.async_set_cache_pipeline(cache_list=list(key_value_dict.items()), ttl=60) ## Set cache namespace if it's a miss data["metadata"]["redis_namespace"] = cache_key_name except HTTPException as e: raise e except Exception as e: verbose_proxy_logger.error( - "litellm.proxy.hooks.batch_redis_get.py::async_pre_call_hook(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.hooks.batch_redis_get.py::async_pre_call_hook(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) @@ -122,28 +108,14 @@ class _PROXY_BatchRedisRequests(CustomLogger): *args, **kwargs ) # returns ":" - we pass redis_namespace in async_pre_call_hook. Done to avoid rewriting the async_set_cache logic - if ( - cache_key is not None - and self.in_memory_cache is not None - and litellm.cache is not None - ): + if cache_key is not None and self.in_memory_cache is not None and litellm.cache is not None: cache_control_args = kwargs.get("cache", {}) - max_age = cache_control_args.get( - "s-max-age", cache_control_args.get("s-maxage", float("inf")) - ) - cached_result = self.in_memory_cache.get_cache( - cache_key, *args, **kwargs - ) + max_age = cache_control_args.get("s-max-age", cache_control_args.get("s-maxage", float("inf"))) + cached_result = self.in_memory_cache.get_cache(cache_key, *args, **kwargs) if cached_result is None: - cached_result = await litellm.cache.cache.async_get_cache( - cache_key, *args, **kwargs - ) + cached_result = await litellm.cache.cache.async_get_cache(cache_key, *args, **kwargs) if cached_result is not None: - await self.in_memory_cache.async_set_cache( - cache_key, cached_result, ttl=60 - ) - return litellm.cache._get_cache_logic( - cached_result=cached_result, max_age=max_age - ) + await self.in_memory_cache.async_set_cache(cache_key, cached_result, ttl=60) + return litellm.cache._get_cache_logic(cached_result=cached_result, max_age=max_age) except Exception: return None diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index 7edc5f4698c..f2e31b77761 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -63,9 +63,7 @@ class DynamicRateLimiterCache: current_minute = dt.strftime("%H-%M") key_name = "{}:{}".format(current_minute, model) - await self.cache.async_set_cache_sadd( - key=key_name, value=value, ttl=self.ttl - ) + await self.cache.async_set_cache_sadd(key=key_name, value=value, ttl=self.ttl) except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.hooks.dynamic_rate_limiter.py::async_set_cache_sadd(): Exception occured - {}".format( @@ -85,9 +83,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): async def check_available_usage( self, model: str, priority: Optional[str] = None - ) -> Tuple[ - Optional[int], Optional[int], Optional[int], Optional[int], Optional[int] - ]: + ) -> Tuple[Optional[int], Optional[int], Optional[int], Optional[int], Optional[int]]: """ For a given model, get its available tpm @@ -105,15 +101,10 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): """ try: # Get model info first for conversion - model_group_info: Optional[ModelGroupInfo] = ( - self.llm_router.get_model_group_info(model_group=model) - ) + model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(model_group=model) weight: float = 1 - if ( - litellm.priority_reservation is None - or priority not in litellm.priority_reservation - ): + if litellm.priority_reservation is None or priority not in litellm.priority_reservation: verbose_proxy_logger.error( "Priority Reservation not set. priority={}, but litellm.priority_reservation is {}.".format( priority, litellm.priority_reservation @@ -128,9 +119,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): value = litellm.priority_reservation[priority] weight = convert_priority_to_percent(value, model_group_info) - active_projects = await self.internal_usage_cache.async_get_cache( - model=model - ) + active_projects = await self.internal_usage_cache.async_get_cache(model=model) ( current_model_tpm, current_model_rpm, @@ -206,23 +195,17 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): - Raise RateLimitError if no tpm/rpm available """ if "model" in data: - key_priority: Optional[str] = user_api_key_dict.metadata.get( - "priority", None - ) + key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None) ( available_tpm, available_rpm, model_tpm, model_rpm, active_projects, - ) = await self.check_available_usage( - model=data["model"], priority=key_priority - ) + ) = await self.check_available_usage(model=data["model"], priority=key_priority) ### CHECK TPM ### if available_tpm is not None and available_tpm == 0: - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - data.get("model") - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model")) raise ProxyRateLimitError( detail={ "error": "Key={} over available TPM={}. Model TPM={}, Active keys={}".format( @@ -238,9 +221,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) ### CHECK RPM ### elif available_rpm is not None and available_rpm == 0: - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - data.get("model") - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model")) raise ProxyRateLimitError( detail={ "error": "Key={} over available RPM={}. Model RPM={}, Active keys={}".format( @@ -264,34 +245,22 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) return None - async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response - ): + async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): try: if isinstance(response, ModelResponse): - model_info = self.llm_router.get_model_info( - id=response._hidden_params["model_id"] - ) - assert model_info is not None, ( - "Model info for model with id={} is None".format( - response._hidden_params["model_id"] - ) - ) - key_priority: Optional[str] = user_api_key_dict.metadata.get( - "priority", None + model_info = self.llm_router.get_model_info(id=response._hidden_params["model_id"]) + assert model_info is not None, "Model info for model with id={} is None".format( + response._hidden_params["model_id"] ) + key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None) ( available_tpm, available_rpm, model_tpm, model_rpm, active_projects, - ) = await self.check_available_usage( - model=model_info["model_name"], priority=key_priority - ) - response._hidden_params[ - "additional_headers" - ] = { # Add additional response headers - easier debugging + ) = await self.check_available_usage(model=model_info["model_name"], priority=key_priority) + response._hidden_params["additional_headers"] = { # Add additional response headers - easier debugging "x-litellm-model_group": model_info["model_name"], "x-ratelimit-remaining-litellm-project-tokens": available_tpm, "x-ratelimit-remaining-litellm-project-requests": available_rpm, diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 493afe6105a..6e4a6fe1a51 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -77,9 +77,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): time_provider: Optional[Callable[[], datetime]] = None, ): self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache) - self.v3_limiter = _PROXY_MaxParallelRequestsHandler_v3( - self.internal_usage_cache, time_provider=time_provider - ) + self.v3_limiter = _PROXY_MaxParallelRequestsHandler_v3(self.internal_usage_cache, time_provider=time_provider) def update_variables(self, llm_router: Router): self.llm_router = llm_router @@ -113,18 +111,11 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ttl=local_cache_ttl, ) - def _get_priority_weight( - self, priority: Optional[str], model_info: Optional[ModelGroupInfo] = None - ) -> float: + def _get_priority_weight(self, priority: Optional[str], model_info: Optional[ModelGroupInfo] = None) -> float: """Get the weight for a given priority from litellm.priority_reservation""" weight: float = _get_priority_settings().default_priority - if ( - litellm.priority_reservation is None - or priority not in litellm.priority_reservation - ): - verbose_proxy_logger.debug( - "Priority Reservation not set for the given priority." - ) + if litellm.priority_reservation is None or priority not in litellm.priority_reservation: + verbose_proxy_logger.debug("Priority Reservation not set for the given priority.") elif priority is not None and litellm.priority_reservation is not None: if os.getenv("LITELLM_LICENSE", None) is None: verbose_proxy_logger.error( @@ -135,9 +126,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): weight = convert_priority_to_percent(value, model_info) return weight - def _get_priority_from_user_api_key_dict( - self, user_api_key_dict: UserAPIKeyAuth - ) -> Optional[str]: + def _get_priority_from_user_api_key_dict(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: """ Get priority from user_api_key_dict. @@ -161,9 +150,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return priority - def _normalize_priority_weights( - self, model_info: ModelGroupInfo - ) -> Dict[str, float]: + def _normalize_priority_weights(self, model_info: ModelGroupInfo) -> Dict[str, float]: """ Normalize priority weights if they sum to > 1.0 @@ -182,9 +169,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if total_weight > 1.0: normalized = {k: v / total_weight for k, v in weights.items()} - verbose_proxy_logger.debug( - f"Normalized over-allocated priorities: {weights} -> {normalized}" - ) + verbose_proxy_logger.debug(f"Normalized over-allocated priorities: {weights} -> {normalized}") return normalized return weights @@ -220,9 +205,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if has_explicit_priority and priority is not None: # Explicit priority: get its specific allocation - priority_weight = normalized_weights.get( - priority, self._get_priority_weight(priority, model_info) - ) + priority_weight = normalized_weights.get(priority, self._get_priority_weight(priority, model_info)) # Use unique key per priority level priority_key = f"{model}:{priority}" else: @@ -260,9 +243,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) # Query Redis directly for current counter value (skip local cache for consistency) - counter_value = await self._get_saturation_value_from_cache( - counter_key=counter_key - ) + counter_value = await self._get_saturation_value_from_cache(counter_key=counter_key) if counter_value is not None: current_requests = int(counter_value) @@ -270,8 +251,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): max_saturation = max(max_saturation, rpm_saturation) verbose_proxy_logger.debug( - f"Model {model} RPM: {current_requests}/{model_group_info.rpm} " - f"({rpm_saturation:.1%})" + f"Model {model} RPM: {current_requests}/{model_group_info.rpm} ({rpm_saturation:.1%})" ) # Query TPM saturation @@ -282,9 +262,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): rate_limit_type="tokens", ) - counter_value = await self._get_saturation_value_from_cache( - counter_key=counter_key - ) + counter_value = await self._get_saturation_value_from_cache(counter_key=counter_key) if counter_value is not None: current_tokens = float(counter_value) @@ -292,20 +270,15 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): max_saturation = max(max_saturation, tpm_saturation) verbose_proxy_logger.debug( - f"Model {model} TPM: {current_tokens}/{model_group_info.tpm} " - f"({tpm_saturation:.1%})" + f"Model {model} TPM: {current_tokens}/{model_group_info.tpm} ({tpm_saturation:.1%})" ) - verbose_proxy_logger.debug( - f"Model {model} overall saturation: {max_saturation:.1%}" - ) + verbose_proxy_logger.debug(f"Model {model} overall saturation: {max_saturation:.1%}") return max_saturation except Exception as e: - verbose_proxy_logger.error( - f"Error checking saturation for {model}: {str(e)}" - ) + verbose_proxy_logger.error(f"Error checking saturation for {model}: {str(e)}") # Fail open: assume not saturated on error return 0.0 @@ -329,9 +302,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return descriptors # Get model group info - model_group_info: Optional[ModelGroupInfo] = ( - self.llm_router.get_model_group_info(model_group=model) - ) + model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(model_group=model) if model_group_info is None: return descriptors @@ -389,16 +360,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): key="model_saturation_check", value=model, rate_limit={ - "requests_per_unit": ( - model_group_info.rpm * high_limit_multiplier - if model_group_info.rpm - else None - ), - "tokens_per_unit": ( - model_group_info.tpm * high_limit_multiplier - if model_group_info.tpm - else None - ), + "requests_per_unit": (model_group_info.rpm * high_limit_multiplier if model_group_info.rpm else None), + "tokens_per_unit": (model_group_info.tpm * high_limit_multiplier if model_group_info.tpm else None), "window_size": self.v3_limiter.window_size, }, ) @@ -489,9 +452,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) - verbose_proxy_logger.debug( - f"Atomic check+increment response: {json.dumps(atomic_response, indent=2)}" - ) + verbose_proxy_logger.debug(f"Atomic check+increment response: {json.dumps(atomic_response, indent=2)}") if atomic_response["overall_code"] == "OVER_LIMIT": resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model) @@ -514,16 +475,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): "rate_limit_type": str(status["rate_limit_type"]), "x-litellm-priority": priority or "default", }, - rate_limit_type=map_v3_rate_limit_type( - status["rate_limit_type"] - ), + rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]), model=resolved_model, llm_provider=llm_provider, ) if descriptor_key == "priority_model": verbose_proxy_logger.debug( - f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " - f"priority: {priority}" + f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, priority: {priority}" ) raise ProxyRateLimitError( detail={ @@ -542,9 +500,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): "x-litellm-priority": priority or "default", "x-litellm-saturation": f"{saturation:.2%}", }, - rate_limit_type=map_v3_rate_limit_type( - status["rate_limit_type"] - ), + rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]), model=resolved_model, llm_provider=llm_provider, ) @@ -566,16 +522,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): raise ProxyRateLimitError( detail={ "error": "Rate limit exceeded", - "descriptor_key": ( - offending["descriptor_key"] if offending else "unknown" - ), - "rate_limit_type": ( - str(offending["rate_limit_type"]) if offending else "unknown" - ), + "descriptor_key": (offending["descriptor_key"] if offending else "unknown"), + "rate_limit_type": (str(offending["rate_limit_type"]) if offending else "unknown"), }, - rate_limit_type=map_v3_rate_limit_type( - offending["rate_limit_type"] if offending else None - ), + rate_limit_type=map_v3_rate_limit_type(offending["rate_limit_type"] if offending else None), headers={ "retry-after": str(self.v3_limiter.window_size), "x-litellm-priority": priority or "default", @@ -602,8 +552,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) data["litellm_proxy_rate_limit_response"] = { "overall_code": atomic_response["overall_code"], - "statuses": atomic_response["statuses"] - + priority_tracking_response["statuses"], + "statuses": atomic_response["statuses"] + priority_tracking_response["statuses"], } else: data["litellm_proxy_rate_limit_response"] = atomic_response @@ -653,18 +602,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return None model = data["model"] - priority = self._get_priority_from_user_api_key_dict( - user_api_key_dict=user_api_key_dict - ) + priority = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) # Get model configuration - model_group_info: Optional[ModelGroupInfo] = ( - self.llm_router.get_model_group_info(model_group=model) - ) + model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(model_group=model) if model_group_info is None: - verbose_proxy_logger.debug( - f"No model group info for {model}, allowing request" - ) + verbose_proxy_logger.debug(f"No model group info for {model}, allowing request") return None try: @@ -695,17 +638,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): except HTTPException: raise except Exception as e: - verbose_proxy_logger.error( - f"Error in dynamic rate limiter: {str(e)}, allowing request" - ) + verbose_proxy_logger.error(f"Error in dynamic rate limiter: {str(e)}, allowing request") # Fail open on unexpected errors return None return None - async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response - ): + async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ Post-call hook to add rate limit headers to response. Leverages v3 limiter's post-call hook functionality. @@ -718,17 +657,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # Add additional priority-specific headers if isinstance(response, ModelResponse): - priority = self._get_priority_from_user_api_key_dict( - user_api_key_dict=user_api_key_dict - ) + priority = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) # Get existing additional headers - additional_headers = ( - getattr(response, "_hidden_params", {}).get( - "additional_headers", {} - ) - or {} - ) + additional_headers = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} # Add priority information additional_headers["x-litellm-priority"] = priority or "default" @@ -742,9 +674,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return response except Exception as e: - verbose_proxy_logger.exception( - f"Error in dynamic rate limiter v3 post-call hook: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error in dynamic rate limiter v3 post-call hook: {str(e)}") return response async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -765,9 +695,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): from litellm.types.utils import Usage try: - verbose_proxy_logger.debug( - "INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING" - ) + verbose_proxy_logger.debug("INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING") litellm_parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) @@ -782,9 +710,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # Get priority from user_api_key_auth_metadata in standard_logging_metadata # This is where user_api_key_dict.metadata is stored during pre-call - user_api_key_auth_metadata = ( - standard_logging_metadata.get("user_api_key_auth_metadata") or {} - ) + user_api_key_auth_metadata = standard_logging_metadata.get("user_api_key_auth_metadata") or {} key_priority: Optional[str] = user_api_key_auth_metadata.get("priority") # Get total tokens from response @@ -856,15 +782,11 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # Only log 'priority' if it's known safe; otherwise, redact. SAFE_PRIORITIES = {"low", "medium", "high", "default"} - logged_priority = ( - key_priority if key_priority in SAFE_PRIORITIES else "REDACTED" - ) + logged_priority = key_priority if key_priority in SAFE_PRIORITIES else "REDACTED" verbose_proxy_logger.debug( f"[Dynamic Rate Limiter] Incremented tokens by {total_tokens} for " f"model={model_group}, priority={logged_priority}" ) except Exception as e: - verbose_proxy_logger.exception( - f"Error in dynamic rate limiter success event: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error in dynamic rate limiter success event: {str(e)}") diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index ee6c063e8e7..ebac86c022d 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -48,9 +48,7 @@ class KeyManagementEventHooks: # Send email notification - non-blocking, independent operation if data.send_invite_email is True: try: - await KeyManagementEventHooks._send_key_created_email( - response.model_dump(exclude_none=True) - ) + await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True)) except Exception as e: verbose_proxy_logger.warning(f"Failed to send key created email: {e}") @@ -85,9 +83,7 @@ class KeyManagementEventHooks: team_id=data.team_id, ) except Exception as e: - verbose_proxy_logger.warning( - f"Failed to store virtual key in secret manager: {e}" - ) + verbose_proxy_logger.warning(f"Failed to store virtual key in secret manager: {e}") @staticmethod async def async_key_updated_hook( @@ -153,13 +149,8 @@ class KeyManagementEventHooks: # Store the generated key in the secret manager - non-blocking, independent operation if data is not None and response.token_id is not None: try: - initial_secret_name = ( - existing_key_row.key_alias - or f"virtual-key-{existing_key_row.token}" - ) - new_secret_name = ( - response.key_alias or data.key_alias or initial_secret_name - ) + initial_secret_name = existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}" + new_secret_name = response.key_alias or data.key_alias or initial_secret_name verbose_proxy_logger.info( "Updating secret in secret manager: secret_name=%s", new_secret_name, @@ -176,9 +167,7 @@ class KeyManagementEventHooks: new_secret_name, ) except Exception as e: - verbose_proxy_logger.warning( - f"Failed to rotate virtual key in secret manager: {e}" - ) + verbose_proxy_logger.warning(f"Failed to rotate virtual key in secret manager: {e}") # Send key rotated email if configured - non-blocking, independent operation try: @@ -206,9 +195,7 @@ class KeyManagementEventHooks: object_id=existing_key_row.token, action="rotated", updated_values=response.model_dump_json(exclude_none=True), - before_value=existing_key_row.model_dump_json( - exclude_none=True - ), + before_value=existing_key_row.model_dump_json(exclude_none=True), ) ) ) @@ -262,15 +249,11 @@ class KeyManagementEventHooks: ) ) # delete the keys from the secret manager - await KeyManagementEventHooks._delete_virtual_keys_from_secret_manager( - keys_being_deleted=keys_being_deleted - ) + await KeyManagementEventHooks._delete_virtual_keys_from_secret_manager(keys_being_deleted=keys_being_deleted) pass @staticmethod - async def _store_virtual_key_in_secret_manager( - secret_name: str, secret_token: str, team_id: Optional[str] = None - ): + async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str, team_id: Optional[str] = None): """ Store a virtual key in the secret manager @@ -287,20 +270,14 @@ class KeyManagementEventHooks: # store the key in the secret manager if isinstance(litellm.secret_manager_client, BaseSecretManager): tags = getattr(litellm._key_management_settings, "tags", None) - description = getattr( - litellm._key_management_settings, "description", None - ) - optional_params = await KeyManagementEventHooks._get_secret_manager_optional_params( - team_id - ) + description = getattr(litellm._key_management_settings, "description", None) + optional_params = await KeyManagementEventHooks._get_secret_manager_optional_params(team_id) verbose_proxy_logger.debug( f"Creating secret with {secret_name} and tags={tags} and description={description}" ) await litellm.secret_manager_client.async_write_secret( - secret_name=KeyManagementEventHooks._get_secret_name( - secret_name - ), + secret_name=KeyManagementEventHooks._get_secret_name(secret_name), description=description, secret_value=secret_token, tags=tags, @@ -331,25 +308,17 @@ class KeyManagementEventHooks: # store the key in the secret manager if isinstance(litellm.secret_manager_client, BaseSecretManager): - optional_params = await KeyManagementEventHooks._get_secret_manager_optional_params( - team_id - ) + optional_params = await KeyManagementEventHooks._get_secret_manager_optional_params(team_id) await litellm.secret_manager_client.async_rotate_secret( - current_secret_name=KeyManagementEventHooks._get_secret_name( - current_secret_name - ), - new_secret_name=KeyManagementEventHooks._get_secret_name( - new_secret_name - ), + current_secret_name=KeyManagementEventHooks._get_secret_name(current_secret_name), + new_secret_name=KeyManagementEventHooks._get_secret_name(new_secret_name), new_secret_value=new_secret_value, optional_params=optional_params, ) @staticmethod def _get_secret_name(secret_name: str) -> str: - if litellm._key_management_settings.prefix_for_stored_virtual_keys.endswith( - "/" - ): + if litellm._key_management_settings.prefix_for_stored_virtual_keys.endswith("/"): return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{secret_name}" else: return f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{secret_name}" @@ -378,14 +347,10 @@ class KeyManagementEventHooks: if team_id not in team_settings_cache: team_settings_cache[ team_id - ] = await KeyManagementEventHooks._get_secret_manager_optional_params( - team_id - ) + ] = await KeyManagementEventHooks._get_secret_manager_optional_params(team_id) optional_params = team_settings_cache[team_id] await litellm.secret_manager_client.async_delete_secret( - secret_name=KeyManagementEventHooks._get_secret_name( - key.key_alias - ), + secret_name=KeyManagementEventHooks._get_secret_name(key.key_alias), optional_params=optional_params, ) else: @@ -420,9 +385,7 @@ class KeyManagementEventHooks: user_api_key_cache=user_api_key_cache, ) except Exception as exc: # pragma: no cover - defensive logging - verbose_proxy_logger.debug( - f"Unable to load team metadata for team_id={team_id}: {exc}" - ) + verbose_proxy_logger.debug(f"Unable to load team metadata for team_id={team_id}: {exc}") return None metadata = getattr(team_obj, "metadata", None) @@ -455,10 +418,8 @@ class KeyManagementEventHooks: BaseEmailLogger, ) - initialized_email_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=BaseEmailLogger - ) + initialized_email_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger ) if len(initialized_email_loggers) > 0: return True @@ -483,9 +444,7 @@ class KeyManagementEventHooks: """ # Early exit if email is not enabled if not KeyManagementEventHooks._is_email_sending_enabled(): - verbose_proxy_logger.debug( - "Email sending not enabled, skipping key created email" - ) + verbose_proxy_logger.debug("Email sending not enabled, skipping key created email") return from litellm.proxy.proxy_server import general_settings, proxy_logging_obj @@ -501,10 +460,8 @@ class KeyManagementEventHooks: SendKeyCreatedEmailEvent, ) - initialized_email_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=BaseEmailLogger - ) + initialized_email_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger ) if len(initialized_email_loggers) > 0: event = SendKeyCreatedEmailEvent( @@ -553,9 +510,7 @@ class KeyManagementEventHooks: ) @staticmethod - async def _send_key_rotated_email( - response: dict, existing_key_alias: Optional[str] - ): + async def _send_key_rotated_email(response: dict, existing_key_alias: Optional[str]): """ Send key rotated email if email sending is enabled. @@ -564,9 +519,7 @@ class KeyManagementEventHooks: """ # Early exit if email is not enabled if not KeyManagementEventHooks._is_email_sending_enabled(): - verbose_proxy_logger.debug( - "Email sending not enabled, skipping key rotated email" - ) + verbose_proxy_logger.debug("Email sending not enabled, skipping key rotated email") return try: @@ -575,9 +528,7 @@ class KeyManagementEventHooks: ) except ImportError: # Enterprise package not installed - v0 doesn't support key rotated email - verbose_proxy_logger.debug( - "Enterprise package not installed, skipping key rotated email" - ) + verbose_proxy_logger.debug("Enterprise package not installed, skipping key rotated email") return try: @@ -585,9 +536,7 @@ class KeyManagementEventHooks: SendKeyRotatedEmailEvent, ) except ImportError: - verbose_proxy_logger.debug( - "Enterprise types not available, skipping key rotated email" - ) + verbose_proxy_logger.debug("Enterprise types not available, skipping key rotated email") return event = SendKeyRotatedEmailEvent( @@ -606,10 +555,8 @@ class KeyManagementEventHooks: ########################## # v2 integration for emails ########################## - initialized_email_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=BaseEmailLogger - ) + initialized_email_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger ) if len(initialized_email_loggers) > 0: for email_logger in initialized_email_loggers: diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 77ed3493a0c..12370ea1536 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -96,9 +96,7 @@ class SkillsInjectionHook(CustomLogger): if not skills or not isinstance(skills, list): return data - verbose_proxy_logger.debug( - f"SkillsInjectionHook: Processing {len(skills)} skills" - ) + verbose_proxy_logger.debug(f"SkillsInjectionHook: Processing {len(skills)} skills") litellm_skills: List[LiteLLM_SkillsTable] = [] anthropic_skills: List[Dict[str, Any]] = [] @@ -118,9 +116,7 @@ class SkillsInjectionHook(CustomLogger): if db_skill: litellm_skills.append(db_skill) else: - verbose_proxy_logger.warning( - f"SkillsInjectionHook: Skill '{skill_id}' not found in LiteLLM DB" - ) + verbose_proxy_logger.warning(f"SkillsInjectionHook: Skill '{skill_id}' not found in LiteLLM DB") else: # Native Anthropic skill - pass through anthropic_skills.append(skill) @@ -250,9 +246,7 @@ class SkillsInjectionHook(CustomLogger): # Inject skill content into system prompt if skill_contents: - data = self.prompt_handler.inject_skill_content_to_messages( - data, skill_contents - ) + data = self.prompt_handler.inject_skill_content_to_messages(data, skill_contents) # Add litellm_code_execution tool if we have skill files if all_skill_files: @@ -301,9 +295,7 @@ class SkillsInjectionHook(CustomLogger): user_api_key_dict=user_api_key_dict, ) except Exception as e: - verbose_proxy_logger.warning( - f"SkillsInjectionHook: Error fetching skill {skill_id}: {e}" - ) + verbose_proxy_logger.warning(f"SkillsInjectionHook: Error fetching skill {skill_id}: {e}") return None def _is_anthropic_model(self, model: str) -> bool: @@ -354,24 +346,20 @@ class SkillsInjectionHook(CustomLogger): litellm_metadata = request_data.get("litellm_metadata") or {} metadata = request_data.get("metadata") or {} - code_exec_enabled = litellm_metadata.get( + code_exec_enabled = litellm_metadata.get("_litellm_code_execution_enabled") or metadata.get( "_litellm_code_execution_enabled" - ) or metadata.get("_litellm_code_execution_enabled") + ) if not code_exec_enabled: return None # Get skill files - skill_files_by_id = litellm_metadata.get("_skill_files") or metadata.get( - "_skill_files", {} - ) + skill_files_by_id = litellm_metadata.get("_skill_files") or metadata.get("_skill_files", {}) all_skill_files: Dict[str, bytes] = {} for files_dict in skill_files_by_id.values(): all_skill_files.update(files_dict) if not all_skill_files: - verbose_proxy_logger.warning( - "SkillsInjectionHook: No skill files found, cannot execute code" - ) + verbose_proxy_logger.warning("SkillsInjectionHook: No skill files found, cannot execute code") return None # Check for tool calls - handle both Anthropic and OpenAI formats @@ -384,19 +372,14 @@ class SkillsInjectionHook(CustomLogger): for tc in tool_calls: tool_name = tc.get("name", "") # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx) - if ( - tool_name == LiteLLMInternalTools.CODE_EXECUTION.value - or tool_name.startswith(LITELLM_SKILL_ID_PREFIX) - ): + if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith(LITELLM_SKILL_ID_PREFIX): has_executable_tool = True break if not has_executable_tool: return None - verbose_proxy_logger.debug( - "SkillsInjectionHook: Detected tool call, starting execution loop" - ) + verbose_proxy_logger.debug("SkillsInjectionHook: Detected tool call, starting execution loop") # Start the agentic loop return await self._execute_code_loop_messages_api( @@ -427,10 +410,7 @@ class SkillsInjectionHook(CustomLogger): "input": block.get("input", {}), } ) - elif ( - hasattr(block, "type") - and getattr(block, "type", None) == "tool_use" - ): + elif hasattr(block, "type") and getattr(block, "type", None) == "tool_use": tool_calls.append( { "id": getattr(block, "id", None), @@ -448,11 +428,7 @@ class SkillsInjectionHook(CustomLogger): { "id": tc.id, "name": tc.function.name, - "input": ( - json.loads(tc.function.arguments) - if tc.function.arguments - else {} - ), + "input": (json.loads(tc.function.arguments) if tc.function.arguments else {}), } ) @@ -479,9 +455,7 @@ class SkillsInjectionHook(CustomLogger): # Ensure response is not None if response is None: - verbose_proxy_logger.error( - "SkillsInjectionHook: Response is None, cannot execute code loop" - ) + verbose_proxy_logger.error("SkillsInjectionHook: Response is None, cannot execute code loop") return None model = data.get("model", "") @@ -541,9 +515,7 @@ class SkillsInjectionHook(CustomLogger): # Execute if it's litellm_code_execution OR a skill tool if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: code = tool_input.get("code", "") - result = await self._execute_code( - code, skill_files, executor, generated_files - ) + result = await self._execute_code(code, skill_files, executor, generated_files) elif tool_name.startswith(LITELLM_SKILL_ID_PREFIX): # Skill tool - execute the skill's code result = await self._execute_skill_tool( @@ -564,9 +536,7 @@ class SkillsInjectionHook(CustomLogger): messages.append({"role": "user", "content": tool_results}) # Make next LLM call - verbose_proxy_logger.debug( - f"SkillsInjectionHook: Making LLM call iteration {iteration + 2}" - ) + verbose_proxy_logger.debug(f"SkillsInjectionHook: Making LLM call iteration {iteration + 2}") try: current_response = await litellm.anthropic.acreate( model=model, @@ -575,17 +545,13 @@ class SkillsInjectionHook(CustomLogger): max_tokens=max_tokens, ) if current_response is None: - verbose_proxy_logger.error( - "SkillsInjectionHook: LLM call returned None" - ) + verbose_proxy_logger.error("SkillsInjectionHook: LLM call returned None") return self._attach_files_to_response(response, generated_files) except Exception as e: verbose_proxy_logger.error(f"SkillsInjectionHook: LLM call failed: {e}") return self._attach_files_to_response(response, generated_files) - verbose_proxy_logger.warning( - f"SkillsInjectionHook: Max iterations ({self.max_iterations}) reached" - ) + verbose_proxy_logger.warning(f"SkillsInjectionHook: Max iterations ({self.max_iterations}) reached") return self._attach_files_to_response(current_response, generated_files) async def _execute_code( @@ -597,9 +563,7 @@ class SkillsInjectionHook(CustomLogger): ) -> str: """Execute code in sandbox and return result string.""" try: - verbose_proxy_logger.debug( - f"SkillsInjectionHook: Executing code ({len(code)} chars)" - ) + verbose_proxy_logger.debug(f"SkillsInjectionHook: Executing code ({len(code)} chars)") exec_result = executor.execute(code=code, skill_files=skill_files) @@ -636,20 +600,12 @@ class SkillsInjectionHook(CustomLogger): """Execute a skill tool by generating and running code based on skill content.""" # Generate code based on available skill modules # Look for Python modules in the skill - python_modules = [ - p - for p in skill_files.keys() - if p.endswith(".py") and not p.endswith("__init__.py") - ] + python_modules = [p for p in skill_files.keys() if p.endswith(".py") and not p.endswith("__init__.py")] # Try to find the main builder/creator module main_module = None for mod in python_modules: - if ( - "builder" in mod.lower() - or "creator" in mod.lower() - or "generator" in mod.lower() - ): + if "builder" in mod.lower() or "creator" in mod.lower() or "generator" in mod.lower(): main_module = mod break @@ -805,9 +761,7 @@ print('No executable skill module found') ) # Make next LLM call using the messages API - verbose_proxy_logger.debug( - f"SkillsInjectionHook: Making LLM call iteration {iteration + 2}" - ) + verbose_proxy_logger.debug(f"SkillsInjectionHook: Making LLM call iteration {iteration + 2}") current_response = await litellm.anthropic.acreate( model=model, messages=messages, @@ -816,9 +770,7 @@ print('No executable skill module found') ) # Max iterations reached - verbose_proxy_logger.warning( - f"SkillsInjectionHook: Max iterations ({self.max_iterations}) reached" - ) + verbose_proxy_logger.warning(f"SkillsInjectionHook: Max iterations ({self.max_iterations}) reached") return self._attach_files_to_response(current_response, generated_files) async def _execute_code_tool( @@ -833,9 +785,7 @@ print('No executable skill module found') args = json.loads(tool_call.function.arguments) code = args.get("code", "") - verbose_proxy_logger.debug( - f"SkillsInjectionHook: Executing code ({len(code)} chars)" - ) + verbose_proxy_logger.debug(f"SkillsInjectionHook: Executing code ({len(code)} chars)") exec_result = executor.execute( code=code, @@ -861,8 +811,7 @@ print('No executable skill module found') tool_result += f"\n- {f['name']} ({len(file_content)} bytes)" verbose_proxy_logger.debug( - f"SkillsInjectionHook: Generated file {f['name']} " - f"({len(file_content)} bytes)" + f"SkillsInjectionHook: Generated file {f['name']} ({len(file_content)} bytes)" ) if exec_result.get("error"): @@ -871,9 +820,7 @@ print('No executable skill module found') return tool_result except Exception as e: - verbose_proxy_logger.error( - f"SkillsInjectionHook: Code execution failed: {e}" - ) + verbose_proxy_logger.error(f"SkillsInjectionHook: Code execution failed: {e}") return f"Code execution failed: {str(e)}" def _attach_files_to_response( @@ -893,9 +840,7 @@ print('No executable skill module found') # Handle dict response (Anthropic/messages API format) if isinstance(response, dict): response["_litellm_generated_files"] = generated_files - verbose_proxy_logger.debug( - f"SkillsInjectionHook: Attached {len(generated_files)} files to dict response" - ) + verbose_proxy_logger.debug(f"SkillsInjectionHook: Attached {len(generated_files)} files to dict response") return response # Handle object response (OpenAI format) @@ -910,9 +855,7 @@ print('No executable skill module found') response.model_extra = {} response.model_extra["_litellm_generated_files"] = generated_files - verbose_proxy_logger.debug( - f"SkillsInjectionHook: Attached {len(generated_files)} files to response" - ) + verbose_proxy_logger.debug(f"SkillsInjectionHook: Attached {len(generated_files)} files to response") return response diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 769348a0b88..1983675c5f3 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -45,9 +45,7 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): ) user_counter_key = f"spend:user:{user_id}" - if user_counter_key in get_reserved_counter_keys( - user_api_key_dict.budget_reservation - ): + if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation): return from litellm.proxy.proxy_server import get_current_spend @@ -66,9 +64,7 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): # CHECK IF REQUEST ALLOWED if curr_spend >= max_budget: - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - data.get("model") if data else None - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) raise ProxyRateLimitError( detail="Max budget limit reached.", rate_limit_type=RateLimitType.BUDGET, diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 20bfeb3a6d5..53de9e0c5a3 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -75,10 +75,8 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): ) if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.increment_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - MAX_BUDGET_SESSION_INCREMENT_SCRIPT - ) + self.increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + MAX_BUDGET_SESSION_INCREMENT_SCRIPT ) else: self.increment_script = None @@ -113,9 +111,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): ) if current_spend >= max_budget: - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - data.get("model") if data else None - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) raise ProxyRateLimitError( detail=( f"Session budget exceeded for session {session_id}. " @@ -189,9 +185,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return None - def _get_max_budget_per_session( - self, user_api_key_dict: UserAPIKeyAuth - ) -> Optional[float]: + def _get_max_budget_per_session(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[float]: """Extract max_budget_per_session from agent litellm_params.""" agent_id = user_api_key_dict.agent_id if agent_id is None: @@ -216,16 +210,13 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): """Read current accumulated spend for a session.""" if self.internal_usage_cache.dual_cache.redis_cache is not None: try: - result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache( - key=cache_key - ) + result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key) if result is not None: return float(result) return 0.0 except Exception as e: verbose_proxy_logger.warning( - "MaxBudgetPerSessionHandler: Redis GET failed, " - "falling back to in-memory: %s", + "MaxBudgetPerSessionHandler: Redis GET failed, falling back to in-memory: %s", str(e), ) @@ -249,8 +240,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return float(result) except Exception as e: verbose_proxy_logger.warning( - "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, " - "falling back to in-memory: %s", + "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, falling back to in-memory: %s", str(e), ) diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 525214ff6be..093351c7d6d 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -70,16 +70,12 @@ class _PROXY_MaxIterationsHandler(CustomLogger): def __init__(self, internal_usage_cache: InternalUsageCache): self.internal_usage_cache = internal_usage_cache - self.ttl = int( - os.getenv("LITELLM_MAX_ITERATIONS_TTL", DEFAULT_MAX_ITERATIONS_TTL) - ) + self.ttl = int(os.getenv("LITELLM_MAX_ITERATIONS_TTL", DEFAULT_MAX_ITERATIONS_TTL)) # Register Lua script with Redis if available (same pattern as v3 limiter) if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.increment_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - MAX_ITERATIONS_INCREMENT_SCRIPT - ) + self.increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + MAX_ITERATIONS_INCREMENT_SCRIPT ) else: self.increment_script = None @@ -117,9 +113,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger): current_count = await self._increment_and_get(cache_key) if current_count > max_iterations: - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - data.get("model") if data else None - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) raise ProxyRateLimitError( detail=( f"Max iterations exceeded for session {session_id}. " diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 9888baf897e..92b74848188 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -102,19 +102,13 @@ class SemanticToolFilterHook(CustomLogger): openai_tools_as_dicts.append(tool_dict) elif hasattr(tool, "dict"): tool_dict = tool.dict(exclude_none=True) - verbose_proxy_logger.debug( - f"Converted Pydantic tool (v1) to dict: {type(tool).__name__} -> dict" - ) + verbose_proxy_logger.debug(f"Converted Pydantic tool (v1) to dict: {type(tool).__name__} -> dict") openai_tools_as_dicts.append(tool_dict) elif isinstance(tool, dict): - verbose_proxy_logger.debug( - f"Tool is already a dict with keys: {list(tool.keys())}" - ) + verbose_proxy_logger.debug(f"Tool is already a dict with keys: {list(tool.keys())}") openai_tools_as_dicts.append(tool) else: - verbose_proxy_logger.warning( - f"Tool is unknown type: {type(tool)}, passing as-is" - ) + verbose_proxy_logger.warning(f"Tool is unknown type: {type(tool)}, passing as-is") openai_tools_as_dicts.append(tool) verbose_proxy_logger.debug( @@ -132,17 +126,9 @@ class SemanticToolFilterHook(CustomLogger): 2. Responses API function tools are always native. 3. Everything else is looked up by name in the MCP registry. """ - if ( - isinstance(tool, dict) - and tool.get("type") == "function" - and isinstance(tool.get("function"), dict) - ): + if isinstance(tool, dict) and tool.get("type") == "function" and isinstance(tool.get("function"), dict): return False - if ( - isinstance(tool, dict) - and tool.get("type") == "function" - and isinstance(tool.get("name"), str) - ): + if isinstance(tool, dict) and tool.get("type") == "function" and isinstance(tool.get("name"), str): return False name, _ = self.filter._extract_tool_info(tool) return bool(name) and name in self.filter._tool_map @@ -183,8 +169,7 @@ class SemanticToolFilterHook(CustomLogger): ) else: verbose_proxy_logger.info( - f"Semantic tool filter: all {len(native_tools)} tools " - f"are native, no MCP filtering applied" + f"Semantic tool filter: all {len(native_tools)} tools are native, no MCP filtering applied" ) async def async_pre_call_hook( @@ -201,9 +186,7 @@ class SemanticToolFilterHook(CustomLogger): tools list to only include semantically relevant tools. """ if call_type not in ("completion", "acompletion", "aresponses"): - verbose_proxy_logger.debug( - f"Skipping semantic filter for call_type={call_type}" - ) + verbose_proxy_logger.debug(f"Skipping semantic filter for call_type={call_type}") return None tools = data.get("tools") @@ -215,16 +198,10 @@ class SemanticToolFilterHook(CustomLogger): # filter_tools/_extract_tool_info cannot name-match, so we skip # semantic filtering and return early. if self._should_expand_mcp_tools(tools): - verbose_proxy_logger.debug( - "Detected litellm_proxy MCP references, expanding before semantic filtering" - ) + verbose_proxy_logger.debug("Detected litellm_proxy MCP references, expanding before semantic filtering") try: - native_tools_before_expand = [ - t - for t in tools - if not (isinstance(t, dict) and t.get("type") == "mcp") - ] + native_tools_before_expand = [t for t in tools if not (isinstance(t, dict) and t.get("type") == "mcp")] expanded_tools = await self._expand_mcp_tools(tools, user_api_key_dict) @@ -232,13 +209,10 @@ class SemanticToolFilterHook(CustomLogger): if native_tools_before_expand: data["tools"] = native_tools_before_expand verbose_proxy_logger.warning( - "No MCP tools expanded, preserving " - f"{len(native_tools_before_expand)} native tools" + f"No MCP tools expanded, preserving {len(native_tools_before_expand)} native tools" ) return data - verbose_proxy_logger.warning( - "No tools expanded from MCP references" - ) + verbose_proxy_logger.warning("No tools expanded from MCP references") return None data["tools"] = native_tools_before_expand + expanded_tools @@ -250,18 +224,14 @@ class SemanticToolFilterHook(CustomLogger): return data except Exception as e: - verbose_proxy_logger.error( - f"Failed to expand MCP references: {e}", exc_info=True - ) + verbose_proxy_logger.error(f"Failed to expand MCP references: {e}", exc_info=True) return None messages = data.get("messages", []) if not messages: messages = data.get("input", []) if not messages: - verbose_proxy_logger.debug( - "No messages in request, skipping semantic filter" - ) + verbose_proxy_logger.debug("No messages in request, skipping semantic filter") return None if not self.filter.enabled: @@ -271,9 +241,7 @@ class SemanticToolFilterHook(CustomLogger): try: user_query = self.filter.extract_user_query(messages) if not user_query: - verbose_proxy_logger.debug( - "No user query found, skipping semantic filter" - ) + verbose_proxy_logger.debug("No user query found, skipping semantic filter") return None native_tools: list[object] = [] @@ -334,9 +302,7 @@ class SemanticToolFilterHook(CustomLogger): return data except Exception as e: - verbose_proxy_logger.warning( - f"Semantic tool filter hook failed: {e}. Proceeding with all tools." - ) + verbose_proxy_logger.warning(f"Semantic tool filter hook failed: {e}. Proceeding with all tools.") return None async def async_post_call_response_headers_hook( @@ -363,10 +329,7 @@ class SemanticToolFilterHook(CustomLogger): tool_names_csv = metadata.get("litellm_semantic_filter_tools", "") if tool_names_csv: if len(tool_names_csv) > MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: - tool_names_csv = ( - tool_names_csv[: MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH - 3] - + "..." - ) + tool_names_csv = tool_names_csv[: MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH - 3] + "..." headers["x-litellm-semantic-filter-tools"] = tool_names_csv @@ -379,11 +342,7 @@ class SemanticToolFilterHook(CustomLogger): tool_names = [] for tool in tools: - name = ( - tool.get("name", "") - if isinstance(tool, dict) - else getattr(tool, "name", "") - ) + name = tool.get("name", "") if isinstance(tool, dict) else getattr(tool, "name", "") if name: tool_names.append(name) @@ -413,19 +372,13 @@ class SemanticToolFilterHook(CustomLogger): return None if llm_router is None: - verbose_proxy_logger.warning( - "Cannot initialize semantic filter: llm_router is None" - ) + verbose_proxy_logger.warning("Cannot initialize semantic filter: llm_router is None") return None try: - embedding_model = config.get( - "embedding_model", DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL - ) + embedding_model = config.get("embedding_model", DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL) top_k = config.get("top_k", DEFAULT_MCP_SEMANTIC_FILTER_TOP_K) - similarity_threshold = config.get( - "similarity_threshold", DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD - ) + similarity_threshold = config.get("similarity_threshold", DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD) semantic_filter = SemanticMCPToolFilter( embedding_model=embedding_model, @@ -450,12 +403,9 @@ class SemanticToolFilterHook(CustomLogger): except ImportError as e: verbose_proxy_logger.warning( - f"semantic-router not installed. Install with: " - f"pip install 'litellm[semantic-router]'. Error: {e}" + f"semantic-router not installed. Install with: pip install 'litellm[semantic-router]'. Error: {e}" ) return None except Exception as e: - verbose_proxy_logger.exception( - f"Failed to initialize MCP semantic tool filter: {e}" - ) + verbose_proxy_logger.exception(f"Failed to initialize MCP semantic tool filter: {e}") return None diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 3c96067da87..5ebb5819c20 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -57,16 +57,11 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model=model, internal_model_max_budget=internal_model_max_budget ) if _current_model_budget_info is None: - verbose_proxy_logger.debug( - f"Model {model} not found in internal_model_max_budget" - ) + verbose_proxy_logger.debug(f"Model {model} not found in internal_model_max_budget") return True # check if current model is within budget - if ( - _current_model_budget_info.max_budget - and _current_model_budget_info.max_budget > 0 - ): + if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0: _current_spend = await self._get_virtual_key_spend_for_model( user_api_key_hash=user_api_key_dict.token, model=model, @@ -112,16 +107,11 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model=model, internal_model_max_budget=internal_model_max_budget ) if _current_model_budget_info is None: - verbose_proxy_logger.debug( - f"Model {model} not found in end_user_model_max_budget" - ) + verbose_proxy_logger.debug(f"Model {model} not found in end_user_model_max_budget") return True # check if current model is within budget - if ( - _current_model_budget_info.max_budget - and _current_model_budget_info.max_budget > 0 - ): + if _current_model_budget_info.max_budget and _current_model_budget_info.max_budget > 0: _current_spend = await self._get_end_user_spend_for_model( end_user_id=end_user_id, model=model, @@ -147,7 +137,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): key_budget_config: BudgetConfig, ) -> Optional[float]: # 1. model: directly look up `model` - end_user_model_spend_cache_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" + end_user_model_spend_cache_key = ( + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" + ) _current_spend = await self.dual_cache.async_get_cache( key=end_user_model_spend_cache_key, ) @@ -175,7 +167,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): """ # 1. model: directly look up `model` - virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}" + virtual_key_model_spend_cache_key = ( + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}" + ) _current_spend = await self.dual_cache.async_get_cache( key=virtual_key_model_spend_cache_key, ) @@ -198,9 +192,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): 1. Check if `model` is in `internal_model_max_budget` 2. If not, check if `model` without custom llm provider is in `internal_model_max_budget` """ - return internal_model_max_budget.get( - model, None - ) or internal_model_max_budget.get( + return internal_model_max_budget.get(model, None) or internal_model_max_budget.get( self._get_model_without_custom_llm_provider(model), None ) @@ -226,9 +218,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d """ verbose_proxy_logger.debug("in RouterBudgetLimiting.async_log_success_event") - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: verbose_proxy_logger.debug( "Skipping _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: standard_logging_payload is None" @@ -237,18 +227,12 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): _litellm_params: dict = kwargs.get("litellm_params", {}) or {} _metadata: dict = _litellm_params.get("metadata", {}) or {} - user_api_key_model_max_budget: Optional[dict] = _metadata.get( - "user_api_key_model_max_budget", None - ) + user_api_key_model_max_budget: Optional[dict] = _metadata.get("user_api_key_model_max_budget", None) user_api_key_end_user_model_max_budget: Optional[dict] = _metadata.get( "user_api_key_end_user_model_max_budget", None ) - if ( - user_api_key_model_max_budget is None - or len(user_api_key_model_max_budget) == 0 - ) and ( - user_api_key_end_user_model_max_budget is None - or len(user_api_key_end_user_model_max_budget) == 0 + if (user_api_key_model_max_budget is None or len(user_api_key_model_max_budget) == 0) and ( + user_api_key_end_user_model_max_budget is None or len(user_api_key_end_user_model_max_budget) == 0 ): verbose_proxy_logger.debug( "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget and user_api_key_end_user_model_max_budget are None or empty." @@ -263,15 +247,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): # Falling back to the deployment-level "model" field preserves # behaviour for non-proxy or non-router deployments where model_group # is None. - model = standard_logging_payload.get( - "model_group" - ) or standard_logging_payload.get("model") - virtual_key = standard_logging_payload.get("metadata", {}).get( - "user_api_key_hash" - ) - end_user_id = standard_logging_payload.get( - "end_user" - ) or standard_logging_payload.get("metadata", {}).get( + model = standard_logging_payload.get("model_group") or standard_logging_payload.get("model") + virtual_key = standard_logging_payload.get("metadata", {}).get("user_api_key_hash") + end_user_id = standard_logging_payload.get("end_user") or standard_logging_payload.get("metadata", {}).get( "user_api_key_end_user_id" ) @@ -290,7 +268,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model=model, internal_model_max_budget=internal_model_max_budget ) if key_budget_config is not None and key_budget_config.budget_duration: - virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" + virtual_spend_key = ( + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" + ) virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}" await self._increment_spend_for_key( budget_config=key_budget_config, @@ -311,7 +291,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model=model, internal_model_max_budget=internal_model_max_budget ) if key_budget_config is not None and key_budget_config.budget_duration: - end_user_spend_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" + end_user_spend_key = ( + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" + ) end_user_start_time_key = f"end_user_budget_start_time:{end_user_id}" await self._increment_spend_for_key( budget_config=key_budget_config, @@ -325,7 +307,5 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): verbose_proxy_logger.debug( "current state of in memory cache %s", - json.dumps( - self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str - ), + json.dumps(self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str), ) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 37da5671b64..1ed76d5b1e3 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -68,9 +68,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): rate_limit_type: Literal["key", "model_per_key", "user", "customer", "team"], values_to_update_in_cache: List[Tuple[Any, Any]], ) -> dict: - verbose_proxy_logger.info( - f"Current Usage of {rate_limit_type} in this minute: {current}" - ) + verbose_proxy_logger.info(f"Current Usage of {rate_limit_type} in this minute: {current}") if current is None: if max_parallel_requests == 0 or tpm_limit == 0 or rpm_limit == 0: # base case — at least one dimension is set to 0 (effectively @@ -120,9 +118,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): else: triggered_type = RateLimitType.REQUESTS requested_model = data.get("model") if data else None - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - requested_model - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(requested_model) raise ProxyRateLimitError( detail=f"LiteLLM Rate Limit Handler for rate limit type = {rate_limit_type}. {CommonProxyErrors.max_parallel_request_limit_reached.value}. current rpm: {current['current_rpm']}, rpm limit: {rpm_limit}, current tpm: {current['current_tpm']}, tpm limit: {tpm_limit}, current max_parallel_requests: {current['current_requests']}, max_parallel_requests: {max_parallel_requests}", headers={"retry-after": str(self.time_to_next_minute())}, @@ -186,9 +182,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): error_message = "Max parallel request limit reached" if additional_details is not None: error_message = error_message + " " + additional_details - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - requested_model - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(requested_model) raise ProxyRateLimitError( detail=error_message, headers={"retry-after": str(self.time_to_next_minute())}, @@ -253,9 +247,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): max_parallel_requests = sys.maxsize if data is None: data = {} - global_max_parallel_requests = data.get("metadata", {}).get( - "global_max_parallel_requests", None - ) + global_max_parallel_requests = data.get("metadata", {}).get("global_max_parallel_requests", None) tpm_limit = getattr(user_api_key_dict, "tpm_limit", sys.maxsize) if tpm_limit is None: tpm_limit = sys.maxsize @@ -306,15 +298,9 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): cache_objects: CacheObject = await self.get_all_cache_objects( current_global_requests=( - "global_max_parallel_requests" - if global_max_parallel_requests is not None - else None - ), - request_count_api_key=( - f"{api_key}::{precise_minute}::request_count" - if api_key is not None - else None + "global_max_parallel_requests" if global_max_parallel_requests is not None else None ), + request_count_api_key=(f"{api_key}::{precise_minute}::request_count" if api_key is not None else None), request_count_api_key_model=( f"{api_key}::{_model}::{precise_minute}::request_count" if api_key is not None and _model is not None @@ -356,16 +342,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # Check if request under RPM/TPM per model for a given API Key _model = data.get("model", None) - _tpm_limit_for_key_model = get_key_model_tpm_limit( - user_api_key_dict, model_name=_model - ) - _rpm_limit_for_key_model = get_key_model_rpm_limit( - user_api_key_dict, model_name=_model - ) + _tpm_limit_for_key_model = get_key_model_tpm_limit(user_api_key_dict, model_name=_model) + _rpm_limit_for_key_model = get_key_model_rpm_limit(user_api_key_dict, model_name=_model) if _tpm_limit_for_key_model is not None or _rpm_limit_for_key_model is not None: - request_count_api_key = ( - f"{api_key}::{_model}::{precise_minute}::request_count" - ) + request_count_api_key = f"{api_key}::{_model}::{precise_minute}::request_count" tpm_limit_for_model = None rpm_limit_for_model = None @@ -464,12 +444,8 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # End-User Rate Limits # Only enforce if user passed `user` to /chat, /completions, /embeddings if user_api_key_dict.end_user_id: - end_user_tpm_limit = getattr( - user_api_key_dict, "end_user_tpm_limit", sys.maxsize - ) - end_user_rpm_limit = getattr( - user_api_key_dict, "end_user_rpm_limit", sys.maxsize - ) + end_user_tpm_limit = getattr(user_api_key_dict, "end_user_tpm_limit", sys.maxsize) + end_user_rpm_limit = getattr(user_api_key_dict, "end_user_rpm_limit", sys.maxsize) if end_user_tpm_limit is None: end_user_tpm_limit = sys.maxsize @@ -477,9 +453,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): end_user_rpm_limit = sys.maxsize # now do the same tpm/rpm checks - request_count_api_key = ( - f"{user_api_key_dict.end_user_id}::{precise_minute}::request_count" - ) + request_count_api_key = f"{user_api_key_dict.end_user_id}::{precise_minute}::request_count" # print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}") await self.check_key_in_limits( @@ -511,9 +485,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): get_model_group_from_litellm_kwargs, ) - litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs( - kwargs=kwargs - ) + litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(kwargs=kwargs) try: self.print_verbose("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") @@ -521,24 +493,15 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "global_max_parallel_requests", None ) user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"] - user_api_key_user_id = kwargs["litellm_params"]["metadata"].get( - "user_api_key_user_id", None - ) - user_api_key_team_id = kwargs["litellm_params"]["metadata"].get( - "user_api_key_team_id", None - ) + user_api_key_user_id = kwargs["litellm_params"]["metadata"].get("user_api_key_user_id", None) + user_api_key_team_id = kwargs["litellm_params"]["metadata"].get("user_api_key_team_id", None) user_api_key_model_max_budget = kwargs["litellm_params"]["metadata"].get( "user_api_key_model_max_budget", None ) user_api_key_end_user_id = kwargs.get("user") - user_api_key_metadata = ( - kwargs["litellm_params"]["metadata"].get("user_api_key_metadata", {}) - or {} - ) - user_api_key_team_metadata = kwargs["litellm_params"]["metadata"].get( - "user_api_key_team_metadata", None - ) + user_api_key_metadata = kwargs["litellm_params"]["metadata"].get("user_api_key_metadata", {}) or {} + user_api_key_team_metadata = kwargs["litellm_params"]["metadata"].get("user_api_key_team_metadata", None) user_api_key_dict = UserAPIKeyAuth( api_key=user_api_key, metadata=user_api_key_metadata, @@ -568,9 +531,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): total_tokens = 0 - if isinstance( - response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse) - ): + if isinstance(response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse)): total_tokens = response_obj.usage.total_tokens # type: ignore # ------------ @@ -580,9 +541,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): values_to_update_in_cache = [] if user_api_key is not None: - request_count_api_key = ( - f"{user_api_key}::{precise_minute}::request_count" - ) + request_count_api_key = f"{user_api_key}::{precise_minute}::request_count" current = await self.internal_usage_cache.async_get_cache( key=request_count_api_key, @@ -599,9 +558,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "current_rpm": current["current_rpm"], } - self.print_verbose( - f"updated_value in success call: {new_val}, precise_minute: {precise_minute}" - ) + self.print_verbose(f"updated_value in success call: {new_val}, precise_minute: {precise_minute}") values_to_update_in_cache.append((request_count_api_key, new_val)) # ------------ @@ -609,14 +566,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # ------------ model_group = get_model_group_from_litellm_kwargs(kwargs) _success_tpm_limit = ( - get_key_model_tpm_limit(user_api_key_dict, model_name=model_group) - if model_group is not None - else None + get_key_model_tpm_limit(user_api_key_dict, model_name=model_group) if model_group is not None else None ) _success_rpm_limit = ( - get_key_model_rpm_limit(user_api_key_dict, model_name=model_group) - if model_group is not None - else None + get_key_model_rpm_limit(user_api_key_dict, model_name=model_group) if model_group is not None else None ) if ( user_api_key is not None @@ -629,9 +582,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): or _success_rpm_limit is not None ) ): - request_count_api_key = ( - f"{user_api_key}::{model_group}::{precise_minute}::request_count" - ) + request_count_api_key = f"{user_api_key}::{model_group}::{precise_minute}::request_count" current = await self.internal_usage_cache.async_get_cache( key=request_count_api_key, @@ -648,9 +599,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "current_rpm": current["current_rpm"], } - self.print_verbose( - f"updated_value in success call: {new_val}, precise_minute: {precise_minute}" - ) + self.print_verbose(f"updated_value in success call: {new_val}, precise_minute: {precise_minute}") values_to_update_in_cache.append((request_count_api_key, new_val)) # ------------ @@ -665,9 +614,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ): total_tokens = response_obj.usage.total_tokens # type: ignore - request_count_api_key = ( - f"{user_api_key_user_id}::{precise_minute}::request_count" - ) + request_count_api_key = f"{user_api_key_user_id}::{precise_minute}::request_count" current = await self.internal_usage_cache.async_get_cache( key=request_count_api_key, @@ -684,9 +631,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "current_rpm": current["current_rpm"], } - self.print_verbose( - f"updated_value in success call: {new_val}, precise_minute: {precise_minute}" - ) + self.print_verbose(f"updated_value in success call: {new_val}, precise_minute: {precise_minute}") values_to_update_in_cache.append((request_count_api_key, new_val)) # ------------ @@ -701,9 +646,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ): total_tokens = response_obj.usage.total_tokens # type: ignore - request_count_api_key = ( - f"{user_api_key_team_id}::{precise_minute}::request_count" - ) + request_count_api_key = f"{user_api_key_team_id}::{precise_minute}::request_count" current = await self.internal_usage_cache.async_get_cache( key=request_count_api_key, @@ -720,9 +663,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "current_rpm": current["current_rpm"], } - self.print_verbose( - f"updated_value in success call: {new_val}, precise_minute: {precise_minute}" - ) + self.print_verbose(f"updated_value in success call: {new_val}, precise_minute: {precise_minute}") values_to_update_in_cache.append((request_count_api_key, new_val)) # ------------ @@ -737,9 +678,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ): total_tokens = response_obj.usage.total_tokens # type: ignore - request_count_api_key = ( - f"{user_api_key_end_user_id}::{precise_minute}::request_count" - ) + request_count_api_key = f"{user_api_key_end_user_id}::{precise_minute}::request_count" current = await self.internal_usage_cache.async_get_cache( key=request_count_api_key, @@ -756,9 +695,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "current_rpm": current["current_rpm"], } - self.print_verbose( - f"updated_value in success call: {new_val}, precise_minute: {precise_minute}" - ) + self.print_verbose(f"updated_value in success call: {new_val}, precise_minute: {precise_minute}") values_to_update_in_cache.append((request_count_api_key, new_val)) await self.internal_usage_cache.async_batch_set_cache( @@ -772,22 +709,16 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: self.print_verbose("Inside Max Parallel Request Failure Hook") - litellm_parent_otel_span: Union[Span, None] = ( - _get_parent_otel_span_from_kwargs(kwargs=kwargs) - ) + litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(kwargs=kwargs) _metadata = kwargs["litellm_params"].get("metadata", {}) or {} - global_max_parallel_requests = _metadata.get( - "global_max_parallel_requests", None - ) + global_max_parallel_requests = _metadata.get("global_max_parallel_requests", None) user_api_key = _metadata.get("user_api_key", None) self.print_verbose(f"user_api_key: [set={user_api_key is not None}]") if user_api_key is None: return ## decrement call count if call failed - if CommonProxyErrors.max_parallel_request_limit_reached.value in str( - kwargs["exception"] - ): + if CommonProxyErrors.max_parallel_request_limit_reached.value in str(kwargs["exception"]): pass # ignore failed calls due to max limit being reached else: # ------------ @@ -817,9 +748,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): current_minute = datetime.now().strftime("%M") precise_minute = f"{current_date}-{current_hour}-{current_minute}" - request_count_api_key = ( - f"{user_api_key}::{precise_minute}::request_count" - ) + request_count_api_key = f"{user_api_key}::{precise_minute}::request_count" # ------------ # Update usage @@ -847,11 +776,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): litellm_parent_otel_span=litellm_parent_otel_span, ) # save in cache for up to 1 min. except Exception as e: - verbose_proxy_logger.exception( - "Inside Parallel Request Limiter: An exception occurred - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("Inside Parallel Request Limiter: An exception occurred - {}".format(str(e))) async def get_internal_user_object( self, @@ -884,14 +809,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): return _user_id_rate_limits.model_dump() except Exception as e: - verbose_proxy_logger.debug( - "Parallel Request Limiter: Error getting user object", str(e) - ) + verbose_proxy_logger.debug("Parallel Request Limiter: Error getting user object", str(e)) return None - async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response - ): + async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ Retrieve the key's remaining rate limits. """ @@ -901,9 +822,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): current_minute = datetime.now().strftime("%M") precise_minute = f"{current_date}-{current_hour}-{current_minute}" request_count_api_key = f"{api_key}::{precise_minute}::request_count" - current: Optional[ - CurrentItemRateLimit - ] = await self.internal_usage_cache.async_get_cache( + current: Optional[CurrentItemRateLimit] = await self.internal_usage_cache.async_get_cache( key=request_count_api_key, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) @@ -914,38 +833,28 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key_tpm_limit: Optional[int] = None if current is not None: if user_api_key_dict.rpm_limit is not None: - key_remaining_rpm_limit = ( - user_api_key_dict.rpm_limit - current["current_rpm"] - ) + key_remaining_rpm_limit = user_api_key_dict.rpm_limit - current["current_rpm"] key_rpm_limit = user_api_key_dict.rpm_limit if user_api_key_dict.tpm_limit is not None: - key_remaining_tpm_limit = ( - user_api_key_dict.tpm_limit - current["current_tpm"] - ) + key_remaining_tpm_limit = user_api_key_dict.tpm_limit - current["current_tpm"] key_tpm_limit = user_api_key_dict.tpm_limit if hasattr(response, "_hidden_params"): _hidden_params = getattr(response, "_hidden_params") else: _hidden_params = None - if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) - ): + if _hidden_params is not None and (isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict)): if isinstance(_hidden_params, BaseModel): _hidden_params = _hidden_params.model_dump() _additional_headers = _hidden_params.get("additional_headers", {}) or {} if key_remaining_rpm_limit is not None: - _additional_headers["x-ratelimit-remaining-requests"] = ( - key_remaining_rpm_limit - ) + _additional_headers["x-ratelimit-remaining-requests"] = key_remaining_rpm_limit if key_rpm_limit is not None: _additional_headers["x-ratelimit-limit-requests"] = key_rpm_limit if key_remaining_tpm_limit is not None: - _additional_headers["x-ratelimit-remaining-tokens"] = ( - key_remaining_tpm_limit - ) + _additional_headers["x-ratelimit-remaining-tokens"] = key_remaining_tpm_limit if key_tpm_limit is not None: _additional_headers["x-ratelimit-limit-tokens"] = key_tpm_limit @@ -955,6 +864,4 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): {**_hidden_params, "additional_headers": _additional_headers}, ) - return await super().async_post_call_success_hook( - data, user_api_key_dict, response - ) + return await super().async_post_call_success_hook(data, user_api_key_dict, response) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5ca1b2caccf..8522294d120 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -291,20 +291,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.internal_usage_cache = internal_usage_cache self._time_provider = time_provider or datetime.now if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.batch_rate_limiter_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - BATCH_RATE_LIMITER_SCRIPT - ) + self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + BATCH_RATE_LIMITER_SCRIPT ) - self.token_increment_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - TOKEN_INCREMENT_SCRIPT - ) + self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + TOKEN_INCREMENT_SCRIPT ) self.check_and_increment_by_n_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - CHECK_AND_INCREMENT_BY_N_SCRIPT - ) + self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT) ) else: self.batch_rate_limiter_script = None @@ -317,9 +311,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # behavior) instead of reserving an estimated budget upfront, shedding # the extra per-request Redis Lua round-trip and the global-lock # in-memory fallback that the reservation path incurs. - self.tpm_reservation_enabled = ( - os.getenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "true").lower() == "true" - ) + self.tpm_reservation_enabled = os.getenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "true").lower() == "true" # Batch rate limiter (lazy loaded) self._batch_rate_limiter: Optional[Any] = None @@ -353,9 +345,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parallel_request_limiter=self, ) except Exception as e: - verbose_proxy_logger.debug( - f"Could not load batch rate limiter: {str(e)}" - ) + verbose_proxy_logger.debug(f"Could not load batch rate limiter: {str(e)}") return self._batch_rate_limiter def _get_current_time(self) -> datetime: @@ -414,13 +404,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): case _: total_chars = 0 - estimated_input_tokens = ( - max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 - ) + estimated_input_tokens = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 - explicit_max_tokens = data.get("max_tokens") or data.get( - "max_completion_tokens" - ) + explicit_max_tokens = data.get("max_tokens") or data.get("max_completion_tokens") match (explicit_max_tokens, input_text): case (mt, _) if mt is not None: @@ -442,9 +428,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # the smallest TPM limit this request will be charged against, # so a small per-tenant TPM cap can't be tripped by the floor # alone. - output_floor = self._no_max_tokens_output_floor( - min_configured_tpm_limit - ) + output_floor = self._no_max_tokens_output_floor(min_configured_tpm_limit) max_tokens_estimate = max(estimated_input_tokens, output_floor) total_estimated = estimated_input_tokens + max_tokens_estimate @@ -466,11 +450,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ from litellm.caching.redis_cluster_cache import RedisClusterCache - return ( - self.internal_usage_cache.dual_cache.redis_cache is not None - and isinstance( - self.internal_usage_cache.dual_cache.redis_cache, RedisClusterCache - ) + return self.internal_usage_cache.dual_cache.redis_cache is not None and isinstance( + self.internal_usage_cache.dual_cache.redis_cache, RedisClusterCache ) async def in_memory_cache_sliding_window( @@ -524,9 +505,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): litellm_parent_otel_span=None, local_only=True, ) - new_counter_value = ( - int(current_counter) if current_counter is not None else 0 - ) + increment_value + new_counter_value = (int(current_counter) if current_counter is not None else 0) + increment_value await self.internal_usage_cache.async_set_cache( key=counter_key, value=new_counter_value, @@ -570,16 +549,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): counter_key = keys_to_fetch[i + 1] counter_value = cache_values[i + 1] requests_limit = key_metadata[window_key]["requests_limit"] - max_parallel_requests_limit = key_metadata[window_key][ - "max_parallel_requests_limit" - ] + max_parallel_requests_limit = key_metadata[window_key]["max_parallel_requests_limit"] tokens_limit = key_metadata[window_key]["tokens_limit"] # Determine which limit to use for current_limit and limit_remaining current_limit: Optional[int] = None - rate_limit_type: Optional[ - Literal["requests", "tokens", "max_parallel_requests"] - ] = None + rate_limit_type: Optional[Literal["requests", "tokens", "max_parallel_requests"]] = None if counter_key.endswith(":requests"): current_limit = requests_limit rate_limit_type = "requests" @@ -598,11 +573,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): item_code = "OVER_LIMIT" # Only compute limit_remaining if current_limit is not None - limit_remaining = ( - current_limit - int(counter_value) - if counter_value is not None - else current_limit - ) + limit_remaining = current_limit - int(counter_value) if counter_value is not None else current_limit statuses.append( { @@ -696,9 +667,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) all_cache_values.extend(group_cache_values) except Exception as e: - verbose_proxy_logger.warning( - f"Redis Lua script failed for hash tag {hash_tag}: {str(e)}" - ) + verbose_proxy_logger.warning(f"Redis Lua script failed for hash tag {hash_tag}: {str(e)}") # Fallback to in-memory cache for this group group_cache_values = await self.in_memory_cache_sliding_window( keys=group_keys, @@ -755,15 +724,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): rate_limit_set = False if requests_limit is not None: - rpm_key = self.create_rate_limit_keys( - descriptor_key, descriptor_value, "requests" - ) + rpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "requests") keys_to_fetch.extend([window_key, rpm_key]) rate_limit_set = True if tokens_limit is not None: - tpm_key = self.create_rate_limit_keys( - descriptor_key, descriptor_value, "tokens" - ) + tpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "tokens") keys_to_fetch.extend([window_key, tpm_key]) rate_limit_set = True if max_parallel_requests_limit is not None: @@ -777,14 +742,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): continue key_metadata[window_key] = { - "requests_limit": ( - int(requests_limit) if requests_limit is not None else None - ), + "requests_limit": (int(requests_limit) if requests_limit is not None else None), "tokens_limit": int(tokens_limit) if tokens_limit is not None else None, "max_parallel_requests_limit": ( - int(max_parallel_requests_limit) - if max_parallel_requests_limit is not None - else None + int(max_parallel_requests_limit) if max_parallel_requests_limit is not None else None ), "window_size": int(window_size), "descriptor_key": descriptor_key, @@ -798,9 +759,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if cache_values is not None: - rate_limit_response = self.is_cache_list_over_limit( - keys_to_fetch, cache_values, key_metadata - ) + rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata) if rate_limit_response["overall_code"] == "OVER_LIMIT": return rate_limit_response @@ -854,9 +813,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): window_size=self.window_size, ) - rate_limit_response = self.is_cache_list_over_limit( - keys_to_fetch, cache_values, key_metadata - ) + rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata) return rate_limit_response async def atomic_check_and_increment_by_n( @@ -892,10 +849,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): counter, mirroring `should_rate_limit`'s shape. """ if len(descriptors) != len(increments): - raise ValueError( - "atomic_check_and_increment_by_n: descriptors and increments " - "must have the same length" - ) + raise ValueError("atomic_check_and_increment_by_n: descriptors and increments must have the same length") # Build per-descriptor (keys, args, meta) groups. All keys within a # group share the descriptor's {key:value} hash tag, so a single Lua @@ -923,9 +877,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=parent_otel_span, ) - flat_meta: List[Dict[str, Any]] = [ - m for _keys, _args, group_meta in descriptor_groups for m in group_meta - ] + flat_meta: List[Dict[str, Any]] = [m for _keys, _args, group_meta in descriptor_groups for m in group_meta] async with self._check_and_increment_lock: return await self._atomic_check_and_increment_in_memory( per_counter_meta=flat_meta, @@ -954,9 +906,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): meta: List[Dict[str, Any]] = [] for rate_limit_type in ("requests", "tokens"): - rlt: Literal["requests", "tokens"] = cast( - Literal["requests", "tokens"], rate_limit_type - ) + rlt: Literal["requests", "tokens"] = cast(Literal["requests", "tokens"], rate_limit_type) if rlt == "requests": limit_value = rate_limit.get("requests_per_unit") inc_amount = int(increment_amounts.get("requests", 0) or 0) @@ -965,9 +915,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): inc_amount = int(increment_amounts.get("tokens", 0) or 0) if limit_value is None or inc_amount <= 0: continue - counter_key = self.create_rate_limit_keys( - descriptor_key, descriptor_value, rlt - ) + counter_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, rlt) # Counter-key TTL and window_size are conceptually distinct # ("how long the counter Redis key lives" vs "how long the # sliding window is"). Kept as separate values so a future @@ -977,9 +925,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): keys.extend([window_key, counter_key]) # 4-tuple matches the Lua ARGV layout: # [limit, increment, ttl_seconds, window_size_seconds]. - args.extend( - [int(limit_value), inc_amount, ttl_seconds, window_size_seconds] - ) + args.extend([int(limit_value), inc_amount, ttl_seconds, window_size_seconds]) meta.append( { "descriptor_key": descriptor_key, @@ -1028,9 +974,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): f"{self.window_size}s)." ) await self._refund_applied_descriptor_groups(applied) - flat_meta: List[Dict[str, Any]] = [ - m for _k, _a, group_meta in descriptor_groups for m in group_meta - ] + flat_meta: List[Dict[str, Any]] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta] async with self._check_and_increment_lock: return await self._atomic_check_and_increment_in_memory( per_counter_meta=flat_meta, @@ -1069,8 +1013,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: verbose_proxy_logger.warning( - f"Failed to refund {entry['counter_key']} on " - f"cross-descriptor rollback: {e}" + f"Failed to refund {entry['counter_key']} on cross-descriptor rollback: {e}" ) def _build_atomic_response( @@ -1150,9 +1093,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): litellm_parent_otel_span=parent_otel_span, local_only=True, ) - window_expired = ( - window_start is None or (now_int - int(window_start)) >= window_size - ) + window_expired = window_start is None or (now_int - int(window_start)) >= window_size current_counter = ( 0 if window_expired @@ -1172,26 +1113,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): RateLimitStatus( code="OVER_LIMIT", current_limit=meta["current_limit"], - limit_remaining=max( - 0, meta["current_limit"] - current_counter - ), + limit_remaining=max(0, meta["current_limit"] - current_counter), rate_limit_type=meta["rate_limit_type"], descriptor_key=meta["descriptor_key"], ) ], ) - descriptor_state.append( - {"window_expired": window_expired, "current": current_counter} - ) + descriptor_state.append({"window_expired": window_expired, "current": current_counter}) # Pass 2: apply increments. statuses: List[RateLimitStatus] = [] for meta, state in zip(per_counter_meta, descriptor_state): - new_counter = ( - meta["increment"] - if state["window_expired"] - else state["current"] + meta["increment"] - ) + new_counter = meta["increment"] if state["window_expired"] else state["current"] + meta["increment"] if state["window_expired"]: await self.internal_usage_cache.async_set_cache( key=meta["window_key"], @@ -1235,9 +1168,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): shared primitive. """ tpm_descriptors: List[RateLimitDescriptor] = [ - d - for d in descriptors - if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None + d for d in descriptors if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None ] if not tpm_descriptors: return RateLimitResponse(overall_code="OK", statuses=[]) @@ -1258,8 +1189,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Global org rate limits if user_api_key_dict.org_id is not None and ( - user_api_key_dict.organization_rpm_limit is not None - or user_api_key_dict.organization_tpm_limit is not None + user_api_key_dict.organization_rpm_limit is not None or user_api_key_dict.organization_tpm_limit is not None ): descriptors.append( RateLimitDescriptor( @@ -1275,26 +1205,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Model specific org rate limits if ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "organization_metadata", "model_rpm_limit" - ) + get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_rpm_limit") is not None - or get_model_rate_limit_from_metadata( - user_api_key_dict, "organization_metadata", "model_tpm_limit" - ) + or get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_tpm_limit") is not None ): _tpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "organization_metadata", "model_tpm_limit" - ) - or {} + get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_tpm_limit") or {} ) _rpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "organization_metadata", "model_rpm_limit" - ) - or {} + get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_rpm_limit") or {} ) should_check_rate_limit = False @@ -1307,13 +1227,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): model_specific_tpm_limit = None model_specific_rpm_limit = None if requested_model in _tpm_limit_for_team_model: - model_specific_tpm_limit = _tpm_limit_for_team_model[ - requested_model - ] + model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model] if requested_model in _rpm_limit_for_team_model: - model_specific_rpm_limit = _rpm_limit_for_team_model[ - requested_model - ] + model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model] descriptors.append( RateLimitDescriptor( key="model_per_organization", @@ -1350,12 +1266,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not requested_model: return - _tpm_limit_for_key_model = get_key_model_tpm_limit( - user_api_key_dict, model_name=requested_model - ) - _rpm_limit_for_key_model = get_key_model_rpm_limit( - user_api_key_dict, model_name=requested_model - ) + _tpm_limit_for_key_model = get_key_model_tpm_limit(user_api_key_dict, model_name=requested_model) + _rpm_limit_for_key_model = get_key_model_rpm_limit(user_api_key_dict, model_name=requested_model) if _tpm_limit_for_key_model is None and _rpm_limit_for_key_model is None: return @@ -1365,20 +1277,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Check if model has any rate limits configured should_check_rate_limit = ( - requested_model in _tpm_limit_for_key_model - or requested_model in _rpm_limit_for_key_model + requested_model in _tpm_limit_for_key_model or requested_model in _rpm_limit_for_key_model ) if not should_check_rate_limit: return # Get model-specific limits - model_specific_tpm_limit: Optional[int] = _tpm_limit_for_key_model.get( - requested_model - ) - model_specific_rpm_limit: Optional[int] = _rpm_limit_for_key_model.get( - requested_model - ) + model_specific_tpm_limit: Optional[int] = _tpm_limit_for_key_model.get(requested_model) + model_specific_rpm_limit: Optional[int] = _rpm_limit_for_key_model.get(requested_model) descriptors.append( RateLimitDescriptor( @@ -1537,9 +1444,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return global_agent_registry.get_agent_by_id(agent_id=agent_id) - def _get_resolved_agent_id( - self, user_api_key_dict: UserAPIKeyAuth, data: dict - ) -> Optional[str]: + def _get_resolved_agent_id(self, user_api_key_dict: UserAPIKeyAuth, data: dict) -> Optional[str]: """ Resolve the agent_id from either the API key or request metadata. Key-level agent_id takes precedence over metadata/header-supplied agent_id. @@ -1667,8 +1572,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # User rate limits if user_api_key_dict.user_id and ( - user_api_key_dict.user_rpm_limit is not None - or user_api_key_dict.user_tpm_limit is not None + user_api_key_dict.user_rpm_limit is not None or user_api_key_dict.user_tpm_limit is not None ): descriptors.append( RateLimitDescriptor( @@ -1684,8 +1588,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Team rate limits if user_api_key_dict.team_id and ( - user_api_key_dict.team_rpm_limit is not None - or user_api_key_dict.team_tpm_limit is not None + user_api_key_dict.team_rpm_limit is not None or user_api_key_dict.team_tpm_limit is not None ): descriptors.append( RateLimitDescriptor( @@ -1701,12 +1604,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Team Member rate limits if user_api_key_dict.user_id and ( - user_api_key_dict.team_member_rpm_limit is not None - or user_api_key_dict.team_member_tpm_limit is not None + user_api_key_dict.team_member_rpm_limit is not None or user_api_key_dict.team_member_tpm_limit is not None ): - team_member_value = ( - f"{user_api_key_dict.team_id}:{user_api_key_dict.user_id}" - ) + team_member_value = f"{user_api_key_dict.team_id}:{user_api_key_dict.user_id}" descriptors.append( RateLimitDescriptor( key="team_member", @@ -1721,8 +1621,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # End user rate limits if user_api_key_dict.end_user_id and ( - user_api_key_dict.end_user_rpm_limit is not None - or user_api_key_dict.end_user_tpm_limit is not None + user_api_key_dict.end_user_rpm_limit is not None or user_api_key_dict.end_user_tpm_limit is not None ): descriptors.append( RateLimitDescriptor( @@ -1763,12 +1662,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): get_team_model_rpm_limit(user_api_key_dict) is not None or get_team_model_tpm_limit(user_api_key_dict) is not None ): - _tpm_limit_for_team_model = ( - get_team_model_tpm_limit(user_api_key_dict) or {} - ) - _rpm_limit_for_team_model = ( - get_team_model_rpm_limit(user_api_key_dict) or {} - ) + _tpm_limit_for_team_model = get_team_model_tpm_limit(user_api_key_dict) or {} + _rpm_limit_for_team_model = get_team_model_rpm_limit(user_api_key_dict) or {} should_check_rate_limit = False if requested_model in _tpm_limit_for_team_model: should_check_rate_limit = True @@ -1779,13 +1674,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): model_specific_tpm_limit = None model_specific_rpm_limit = None if requested_model in _tpm_limit_for_team_model: - model_specific_tpm_limit = _tpm_limit_for_team_model[ - requested_model - ] + model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model] if requested_model in _rpm_limit_for_team_model: - model_specific_rpm_limit = _rpm_limit_for_team_model[ - requested_model - ] + model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model] descriptors.append( RateLimitDescriptor( key="model_per_team", @@ -1861,9 +1752,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return False except Exception as e: - verbose_proxy_logger.debug( - f"Error checking model failure status: {str(e)}, defaulting to enforce limits" - ) + verbose_proxy_logger.debug(f"Error checking model failure status: {str(e)}, defaulting to enforce limits") # Fail safe: enforce limits if we can't check return True @@ -1882,39 +1771,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) -> None: """Add team model rate limit descriptor from team_metadata if applicable.""" if ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "team_metadata", "model_rpm_limit" - ) - is not None - or get_model_rate_limit_from_metadata( - user_api_key_dict, "team_metadata", "model_tpm_limit" - ) - is not None + get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") is not None + or get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") is not None ): _tpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "team_metadata", "model_tpm_limit" - ) - or {} + get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {} ) _rpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "team_metadata", "model_rpm_limit" - ) - or {} + get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") or {} ) should_check_rate_limit = ( - requested_model in _tpm_limit_for_team_model - or requested_model in _rpm_limit_for_team_model + requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model ) if should_check_rate_limit and requested_model is not None: - model_specific_tpm_limit = _tpm_limit_for_team_model.get( - requested_model - ) - model_specific_rpm_limit = _rpm_limit_for_team_model.get( - requested_model - ) + model_specific_tpm_limit = _tpm_limit_for_team_model.get(requested_model) + model_specific_rpm_limit = _rpm_limit_for_team_model.get(requested_model) descriptors.append( RateLimitDescriptor( key="model_per_team", @@ -1935,39 +1807,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) -> None: """Add project model rate limit descriptor from project_metadata if applicable.""" if ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "project_metadata", "model_rpm_limit" - ) - is not None - or get_model_rate_limit_from_metadata( - user_api_key_dict, "project_metadata", "model_tpm_limit" - ) - is not None + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_rpm_limit") is not None + or get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_tpm_limit") is not None ): _tpm_limit_for_project_model = ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "project_metadata", "model_tpm_limit" - ) - or {} + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_tpm_limit") or {} ) _rpm_limit_for_project_model = ( - get_model_rate_limit_from_metadata( - user_api_key_dict, "project_metadata", "model_rpm_limit" - ) - or {} + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_rpm_limit") or {} ) should_check_rate_limit = ( - requested_model in _tpm_limit_for_project_model - or requested_model in _rpm_limit_for_project_model + requested_model in _tpm_limit_for_project_model or requested_model in _rpm_limit_for_project_model ) if should_check_rate_limit and requested_model is not None: - model_specific_tpm_limit = _tpm_limit_for_project_model.get( - requested_model - ) - model_specific_rpm_limit = _rpm_limit_for_project_model.get( - requested_model - ) + model_specific_tpm_limit = _tpm_limit_for_project_model.get(requested_model) + model_specific_rpm_limit = _rpm_limit_for_project_model.get(requested_model) descriptors.append( RateLimitDescriptor( key="model_per_project", @@ -1994,17 +1849,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): (desc for desc in descriptors if desc["key"] == descriptor_key), None, ) - descriptor_value = ( - matching_descriptor["value"] - if matching_descriptor is not None - else "unknown" - ) + descriptor_value = matching_descriptor["value"] if matching_descriptor is not None else "unknown" now = self._get_current_time().timestamp() reset_time = now + self.window_size - reset_time_formatted = datetime.fromtimestamp(reset_time).strftime( - "%Y-%m-%d %H:%M:%S UTC" - ) + reset_time_formatted = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC") remaining_display = max(0, status["limit_remaining"]) rate_limit_type = status["rate_limit_type"] @@ -2017,9 +1866,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): f"Limit resets at: {reset_time_formatted}" ) - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( - requested_model - ) + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(requested_model) raise ProxyRateLimitError( detail=detail, headers={ @@ -2056,9 +1903,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Check if the call type has a specific rate limiter # eg. for Batch APIs we need to use the batch rate limiter to read the input file and count the tokens and requests ######################################################### - call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type( - call_type=call_type - ) + call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type(call_type=call_type) if call_type_specific_rate_limiter: return await call_type_specific_rate_limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -2113,11 +1958,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # Org Level Rate Limits - descriptors.extend( - self.create_organization_rate_limit_descriptor( - user_api_key_dict, requested_model - ) - ) + descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model)) # Only check rate limits if we have descriptors with actual limits if descriptors: # First pass: RPM and max_parallel_requests sliding-window check. @@ -2172,20 +2013,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # spend past the limit before post-call reconciliation runs. # Skip when the request already sets max_tokens or has no # generation budget at all (embeddings). - capped_floor = self._no_max_tokens_output_floor( - min_configured_tpm_limit - ) + capped_floor = self._no_max_tokens_output_floor(min_configured_tpm_limit) baseline_floor = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION has_explicit_max_tokens = ( - data.get("max_tokens") is not None - or data.get("max_completion_tokens") is not None + data.get("max_tokens") is not None or data.get("max_completion_tokens") is not None ) is_embedding = data.get("input") is not None - if ( - capped_floor < baseline_floor - and not has_explicit_max_tokens - and not is_embedding - ): + if capped_floor < baseline_floor and not has_explicit_max_tokens and not is_embedding: data["max_tokens"] = capped_floor # Floor at 1 token so contentless requests (/responses, @@ -2229,8 +2063,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_scopes: List[Tuple[str, str]] = [ (d["key"], d["value"]) for d in descriptors - if (d.get("rate_limit") or {}).get("tokens_per_unit") - is not None + if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None ] self._stash_reservation_in_data( data=data, @@ -2246,15 +2079,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # silently drops all token headers. stored_response = data.get("litellm_proxy_rate_limit_response") if isinstance(stored_response, dict): - stored_response.setdefault("statuses", []).extend( - tpm_response["statuses"] - ) + stored_response.setdefault("statuses", []).extend(tpm_response["statuses"]) elif tpm_response["statuses"]: data["litellm_proxy_rate_limit_response"] = tpm_response - verbose_proxy_logger.debug( - f"TPM tokens reserved: {estimated_tokens} for model {requested_model}" - ) + verbose_proxy_logger.debug(f"TPM tokens reserved: {estimated_tokens} for model {requested_model}") # Defense-in-depth: scrub any stash key that escaped onto data # top-level (stale cache hit, router pass, test fixture) before the @@ -2329,14 +2158,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get cached tokens to exclude from input/total if rate_limit_type in ("input", "total"): - if ( - hasattr(usage, "prompt_tokens_details") - and usage.prompt_tokens_details is not None - ): - cached_tokens = ( - getattr(usage.prompt_tokens_details, "cached_tokens", 0) - or 0 - ) + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: + cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 elif isinstance(usage, dict): # Responses API usage comes as a dict @@ -2375,9 +2198,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for _hash_tag, group_keys in key_groups.items(): # Get operations for this hash tag group - group_operations = [ - op for op in pipeline_operations if op["key"] in group_keys - ] + group_operations = [op for op in pipeline_operations if op["key"] in group_keys] keys = [] args = [] @@ -2412,9 +2233,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Check if script is available if self.token_increment_script is None: - verbose_proxy_logger.debug( - "TTL preservation script not available, using regular pipeline" - ) + verbose_proxy_logger.debug("TTL preservation script not available, using regular pipeline") await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=pipeline_operations, litellm_parent_otel_span=parent_otel_span, @@ -2429,9 +2248,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.warning( - f"TTL preservation failed, falling back to regular pipeline: {str(e)}" - ) + verbose_proxy_logger.warning(f"TTL preservation failed, falling back to regular pipeline: {str(e)}") # Fallback to regular pipeline on error await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=pipeline_operations, @@ -2441,9 +2258,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def get_rate_limit_type(self) -> Literal["output", "input", "total"]: from litellm.proxy.proxy_server import general_settings - specified_rate_limit_type = general_settings.get( - "token_rate_limit_type", "total" - ) + specified_rate_limit_type = general_settings.get("token_rate_limit_type", "total") if specified_rate_limit_type not in [ "output", "input", @@ -2479,21 +2294,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ``reserved_scopes`` is serialized as a list of [key, value] pairs so it round-trips through JSON-based metadata transports. """ - scopes_payload: Optional[List[List[str]]] = ( - [[k, v] for k, v in reserved_scopes] if reserved_scopes else None - ) + scopes_payload: Optional[List[List[str]]] = [[k, v] for k, v in reserved_scopes] if reserved_scopes else None - cls._stash_value_in_metadata_channels( - data=data, key=TPM_RESERVED_TOKENS_KEY, value=estimated_tokens - ) + cls._stash_value_in_metadata_channels(data=data, key=TPM_RESERVED_TOKENS_KEY, value=estimated_tokens) if reserved_model: - cls._stash_value_in_metadata_channels( - data=data, key=TPM_RESERVED_MODEL_KEY, value=reserved_model - ) + cls._stash_value_in_metadata_channels(data=data, key=TPM_RESERVED_MODEL_KEY, value=reserved_model) if scopes_payload is not None: - cls._stash_value_in_metadata_channels( - data=data, key=TPM_RESERVED_SCOPES_KEY, value=scopes_payload - ) + cls._stash_value_in_metadata_channels(data=data, key=TPM_RESERVED_SCOPES_KEY, value=scopes_payload) @staticmethod def _lookup_stashed_value( @@ -2529,9 +2336,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): kwargs: Any, standard_logging_metadata: Optional[Dict[str, Any]] = None, ) -> int: - candidate = cls._lookup_stashed_value( - kwargs, standard_logging_metadata, TPM_RESERVED_TOKENS_KEY - ) + candidate = cls._lookup_stashed_value(kwargs, standard_logging_metadata, TPM_RESERVED_TOKENS_KEY) try: return int(candidate or 0) except (TypeError, ValueError): @@ -2549,9 +2354,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): of whether the router later set a different ``model_group`` in ``litellm_params.metadata``. """ - candidate = cls._lookup_stashed_value( - kwargs, standard_logging_metadata, TPM_RESERVED_MODEL_KEY - ) + candidate = cls._lookup_stashed_value(kwargs, standard_logging_metadata, TPM_RESERVED_MODEL_KEY) return candidate if isinstance(candidate, str) and candidate else None @classmethod @@ -2567,9 +2370,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): already carry +reserved on the counter) and the full actual to unreserved ones (which were never charged). """ - candidate = cls._lookup_stashed_value( - kwargs, standard_logging_metadata, TPM_RESERVED_SCOPES_KEY - ) + candidate = cls._lookup_stashed_value(kwargs, standard_logging_metadata, TPM_RESERVED_SCOPES_KEY) if not isinstance(candidate, list): return set() scopes: Set[Tuple[str, str]] = set() @@ -2590,11 +2391,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): standard_logging_metadata: Optional[Dict[str, Any]] = None, ) -> bool: """True if a prior callback already refunded this request's reservation.""" - return bool( - cls._lookup_stashed_value( - kwargs, standard_logging_metadata, TPM_RESERVATION_RELEASED_KEY - ) - ) + return bool(cls._lookup_stashed_value(kwargs, standard_logging_metadata, TPM_RESERVATION_RELEASED_KEY)) @staticmethod def _mark_reservation_released(data: Any) -> None: @@ -2639,19 +2436,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): user_api_key = standard_logging_metadata.get("user_api_key_hash") user_api_key_user_id = standard_logging_metadata.get("user_api_key_user_id") user_api_key_team_id = standard_logging_metadata.get("user_api_key_team_id") - user_api_key_organization_id = standard_logging_metadata.get( - "user_api_key_org_id" - ) - user_api_key_project_id = standard_logging_metadata.get( - "user_api_key_project_id" - ) + user_api_key_organization_id = standard_logging_metadata.get("user_api_key_org_id") + user_api_key_project_id = standard_logging_metadata.get("user_api_key_project_id") user_api_key_end_user_id = ( kwargs.get("user") if isinstance(kwargs, dict) else None ) or standard_logging_metadata.get("user_api_key_end_user_id") agent_id = standard_logging_metadata.get("agent_id") - session_id = standard_logging_metadata.get( - "session_id" - ) or standard_logging_metadata.get("trace_id") + session_id = standard_logging_metadata.get("session_id") or standard_logging_metadata.get("trace_id") targets: List[Tuple[str, str]] = [] if user_api_key: @@ -2661,9 +2452,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if user_api_key_team_id: targets.append(("team", user_api_key_team_id)) if user_api_key_team_id and user_api_key_user_id: - targets.append( - ("team_member", f"{user_api_key_team_id}:{user_api_key_user_id}") - ) + targets.append(("team_member", f"{user_api_key_team_id}:{user_api_key_user_id}")) if user_api_key_end_user_id: targets.append(("end_user", user_api_key_end_user_id)) if user_api_key_organization_id: @@ -2672,9 +2461,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if user_api_key: targets.append(("model_per_key", f"{user_api_key}:{model_group}")) if user_api_key_team_id: - targets.append( - ("model_per_team", f"{user_api_key_team_id}:{model_group}") - ) + targets.append(("model_per_team", f"{user_api_key_team_id}:{model_group}")) if user_api_key_organization_id: targets.append( ( @@ -2762,9 +2549,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ), ): _usage = getattr(response_obj, "usage", None) - total_tokens = self._get_total_tokens_from_usage( - usage=_usage, rate_limit_type=rate_limit_type - ) + total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type) reserved_tokens = self._get_reserved_tokens_from_kwargs( kwargs=kwargs, @@ -2843,13 +2628,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): rate_limit_type = self.get_rate_limit_type() - litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs( - kwargs - ) + litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(kwargs) try: - verbose_proxy_logger.debug( - "INSIDE parallel request limiter ASYNC SUCCESS LOGGING" - ) + verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") pipeline_operations = self._build_success_event_pipeline_operations( kwargs=kwargs, @@ -2864,9 +2645,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.exception( - f"Error in rate limit success event: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error in rate limit success event: {str(e)}") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ @@ -2880,9 +2659,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) try: - litellm_parent_otel_span: Union[Span, None] = ( - _get_parent_otel_span_from_kwargs(kwargs) - ) + litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(kwargs) standard_logging_object = kwargs.get("standard_logging_object") or {} standard_logging_metadata = standard_logging_object.get("metadata") or {} user_api_key = standard_logging_metadata.get("user_api_key_hash") @@ -2919,9 +2696,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) if reserved_tokens > 0: - verbose_proxy_logger.debug( - f"Releasing reserved TPM tokens on failure: {reserved_tokens}" - ) + verbose_proxy_logger.debug(f"Releasing reserved TPM tokens on failure: {reserved_tokens}") # Refund only against the scopes the reservation actually # charged. _build_reservation_aware_tpm_ops with # actual_tokens=0 emits -reserved on reserved scopes and 0 @@ -2943,22 +2718,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if pipeline_operations: - await ( - self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, + litellm_parent_otel_span=litellm_parent_otel_span, ) if reserved_tokens > 0: self._mark_reservation_released(kwargs) except Exception as e: - verbose_proxy_logger.exception( - f"Error in rate limit failure event: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error in rate limit failure event: {str(e)}") - async def async_release_max_parallel_requests_on_disconnect( - self, user_api_key_dict: UserAPIKeyAuth - ) -> None: + async def async_release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None: """ Release the api-key ``max_parallel_requests`` slot that ``async_pre_call_hook`` reserved, for a request that ended without @@ -2971,10 +2740,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): runs, so without this the counter leaks one slot per cancelled stream until the key wedges at its limit. """ - if ( - not user_api_key_dict.api_key - or user_api_key_dict.max_parallel_requests is None - ): + if not user_api_key_dict.api_key or user_api_key_dict.max_parallel_requests is None: return await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( @@ -2996,9 +2762,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): litellm_parent_otel_span=None, ) - async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response - ): + async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ Post-call hook to update rate limit headers in the response. """ @@ -3018,25 +2782,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): _hidden_params = None if _hidden_params is not None and ( - isinstance(_hidden_params, BaseModel) - or isinstance(_hidden_params, dict) + isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict) ): if isinstance(_hidden_params, BaseModel): _hidden_params = _hidden_params.model_dump() - _additional_headers = ( - _hidden_params.get("additional_headers", {}) or {} - ) + _additional_headers = _hidden_params.get("additional_headers", {}) or {} # Add rate limit headers for status in litellm_proxy_rate_limit_response["statuses"]: prefix = f"x-ratelimit-{status['descriptor_key']}" - _additional_headers[ - f"{prefix}-remaining-{status['rate_limit_type']}" - ] = status["limit_remaining"] - _additional_headers[ - f"{prefix}-limit-{status['rate_limit_type']}" - ] = status["current_limit"] + _additional_headers[f"{prefix}-remaining-{status['rate_limit_type']}"] = status[ + "limit_remaining" + ] + _additional_headers[f"{prefix}-limit-{status['rate_limit_type']}"] = status["current_limit"] setattr( response, @@ -3045,9 +2804,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) except Exception as e: - verbose_proxy_logger.exception( - f"Error in rate limit post-call hook: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error in rate limit post-call hook: {str(e)}") async def async_post_call_failure_hook( self, @@ -3082,9 +2839,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): standard_logging_metadata=None, key=RATE_LIMIT_DESCRIPTORS_KEY, ) - descriptors: List[RateLimitDescriptor] = ( - stashed if isinstance(stashed, list) else [] - ) + descriptors: List[RateLimitDescriptor] = stashed if isinstance(stashed, list) else [] ops: List[RedisPipelineIncrementOperation] = [] for descriptor in descriptors: rate_limit = descriptor.get("rate_limit") or {} @@ -3102,19 +2857,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) if ops: - verbose_proxy_logger.debug( - f"Releasing reserved TPM tokens on proxy-level " - f"rejection: {reserved_tokens}" - ) - await ( - self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=ops, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - ) + verbose_proxy_logger.debug(f"Releasing reserved TPM tokens on proxy-level rejection: {reserved_tokens}") + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=ops, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) self._mark_reservation_released(request_data) except Exception as e: - verbose_proxy_logger.exception( - f"Error releasing TPM reservation on post-call failure: {e}" - ) + verbose_proxy_logger.exception(f"Error releasing TPM reservation on post-call failure: {e}") return None diff --git a/litellm/proxy/hooks/prompt_injection_detection.py b/litellm/proxy/hooks/prompt_injection_detection.py index bfe85edf4cb..390074bf005 100644 --- a/litellm/proxy/hooks/prompt_injection_detection.py +++ b/litellm/proxy/hooks/prompt_injection_detection.py @@ -77,10 +77,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): def update_environment(self, router: Optional[Router] = None): self.llm_router = router - if ( - self.prompt_injection_params is not None - and self.prompt_injection_params.llm_api_check is True - ): + if self.prompt_injection_params is not None and self.prompt_injection_params.llm_api_check is True: if self.llm_router is None: raise Exception( "PromptInjectionDetection: Model List not set. Required for Prompt Injection detection." @@ -91,8 +88,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): ) if ( self.prompt_injection_params.llm_api_name is None - or self.prompt_injection_params.llm_api_name - not in self.llm_router.model_names + or self.prompt_injection_params.llm_api_name not in self.llm_router.model_names ): raise Exception( "PromptInjectionDetection: Invalid LLM API Name. LLM API Name must be a 'model_name' in 'model_list'." @@ -104,9 +100,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): for adj in self.adjectives: for prep in self.prepositions: phrase = " ".join(filter(None, [verb, adj, prep])).strip() - if ( - len(phrase.split()) > 2 - ): # additional check to ensure more than 2 words + if len(phrase.split()) > 2: # additional check to ensure more than 2 words combinations.append(phrase.lower()) return combinations @@ -171,30 +165,22 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): if self.prompt_injection_params is not None: # 1. check if heuristics check turned on if self.prompt_injection_params.heuristics_check is True: - is_prompt_attack = self.check_user_input_similarity( - user_input=formatted_prompt - ) + is_prompt_attack = self.check_user_input_similarity(user_input=formatted_prompt) if is_prompt_attack is True: raise HTTPException( status_code=400, - detail={ - "error": "Rejected message. This is a prompt injection attack." - }, + detail={"error": "Rejected message. This is a prompt injection attack."}, ) # 2. check if vector db similarity check turned on [TODO] Not Implemented yet if self.prompt_injection_params.vector_db_check is True: pass else: - is_prompt_attack = self.check_user_input_similarity( - user_input=formatted_prompt - ) + is_prompt_attack = self.check_user_input_similarity(user_input=formatted_prompt) if is_prompt_attack is True: raise HTTPException( status_code=400, - detail={ - "error": "Rejected message. This is a prompt injection attack." - }, + detail={"error": "Rejected message. This is a prompt injection attack."}, ) return data @@ -229,9 +215,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): "audio_transcription", ], ) -> Optional[bool]: - self.print_verbose( - f"IN ASYNC MODERATION HOOK - self.prompt_injection_params = {self.prompt_injection_params}" - ) + self.print_verbose(f"IN ASYNC MODERATION HOOK - self.prompt_injection_params = {self.prompt_injection_params}") if self.prompt_injection_params is None: return None @@ -264,24 +248,15 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger): ) self.print_verbose(f"Received LLM Moderation response: {response}") - self.print_verbose( - f"llm_api_fail_call_string: {self.prompt_injection_params.llm_api_fail_call_string}" - ) - if isinstance(response, litellm.ModelResponse) and isinstance( - response.choices[0], litellm.Choices - ): - if ( - self.prompt_injection_params.llm_api_fail_call_string - in response.choices[0].message.content - ): # type: ignore + self.print_verbose(f"llm_api_fail_call_string: {self.prompt_injection_params.llm_api_fail_call_string}") + if isinstance(response, litellm.ModelResponse) and isinstance(response.choices[0], litellm.Choices): + if self.prompt_injection_params.llm_api_fail_call_string in response.choices[0].message.content: # type: ignore is_prompt_attack = True if is_prompt_attack is True: raise HTTPException( status_code=400, - detail={ - "error": "Rejected message. This is a prompt injection attack." - }, + detail={"error": "Rejected message. This is a prompt injection attack."}, ) return is_prompt_attack diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index e3506c78096..9c09231cd9f 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -36,9 +36,7 @@ from litellm.utils import get_end_user_id_for_cost_tracking class _ProxyDBLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - await self._PROXY_track_cost_callback( - kwargs, response_obj, start_time, end_time - ) + await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) async def async_post_call_failure_hook( self, @@ -48,17 +46,11 @@ class _ProxyDBLogger(CustomLogger): traceback_str: Optional[str] = None, ): try: - await _release_budget_reservation( - budget_reservation=user_api_key_dict.budget_reservation - ) + await _release_budget_reservation(budget_reservation=user_api_key_dict.budget_reservation) except Exception: - verbose_proxy_logger.exception( - "Failed to release budget reservation during failure handling" - ) + verbose_proxy_logger.exception("Failed to release budget reservation during failure handling") try: - await _invalidate_budget_reservation_counters( - budget_reservation=user_api_key_dict.budget_reservation - ) + await _invalidate_budget_reservation_counters(budget_reservation=user_api_key_dict.budget_reservation) if user_api_key_dict.budget_reservation is not None: user_api_key_dict.budget_reservation["finalized"] = True except Exception: @@ -70,17 +62,14 @@ class _ProxyDBLogger(CustomLogger): if _ProxyDBLogger._should_track_errors_in_db() is False: return elif request_route is not None and not ( - RouteChecks.is_llm_api_route(route=request_route) - or RouteChecks.is_info_route(route=request_route) + RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route) ): return from litellm.proxy.proxy_server import proxy_logging_obj _metadata = dict( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["status"] = "failure" @@ -122,17 +111,13 @@ class _ProxyDBLogger(CustomLogger): existing_metadata["tags"] = existing_litellm_metadata.get("tags") request_data["litellm_params"]["proxy_server_request"] = ( - request_data.get("proxy_server_request") - or existing_litellm_params.get("proxy_server_request") - or {} + request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} ) request_data["litellm_params"]["metadata"] = existing_metadata # Preserve model name and custom_llm_provider if "model" not in request_data: - request_data["model"] = existing_litellm_params.get( - "model" - ) or request_data.get("model", "") + request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "") if "custom_llm_provider" not in request_data: request_data["custom_llm_provider"] = existing_litellm_params.get( "custom_llm_provider" @@ -146,13 +131,11 @@ class _ProxyDBLogger(CustomLogger): _litellm_logging_obj = request_data.get("litellm_logging_obj") if _litellm_logging_obj is not None: if not request_data.get("standard_logging_object"): - request_data["standard_logging_object"] = getattr( - _litellm_logging_obj, "model_call_details", {} - ).get("standard_logging_object") - if request_data.get("litellm_trace_id") is None: - request_data["litellm_trace_id"] = getattr( - _litellm_logging_obj, "litellm_trace_id", None + request_data["standard_logging_object"] = getattr(_litellm_logging_obj, "model_call_details", {}).get( + "standard_logging_object" ) + if request_data.get("litellm_trace_id") is None: + request_data["litellm_trace_id"] = getattr(_litellm_logging_obj, "litellm_trace_id", None) # Use the actual request start time from the logging object so that # failed requests record the real duration instead of 0. @@ -169,9 +152,7 @@ class _ProxyDBLogger(CustomLogger): # real partial spend to this failure row instead of zero. recovered_response_cost = 0.0 if isinstance(request_data.get("combined_usage_object"), litellm.Usage): - recovered_response_cost = max( - float(request_data.get("response_cost") or 0.0), 0.0 - ) + recovered_response_cost = max(float(request_data.get("response_cost") or 0.0), 0.0) await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, @@ -190,9 +171,7 @@ class _ProxyDBLogger(CustomLogger): async def _PROXY_track_cost_callback( self, kwargs, # kwargs to completion - completion_response: Optional[ - Union[litellm.ModelResponse, Any] - ], # response from completion + completion_response: Optional[Union[litellm.ModelResponse, Any]], # response from completion start_time=None, end_time=None, # start/end time for completion ): @@ -211,21 +190,15 @@ class _ProxyDBLogger(CustomLogger): litellm_params = kwargs.get("litellm_params", {}) or {} end_user_id = get_end_user_id_for_cost_tracking(litellm_params) metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) - budget_reservation = _get_budget_reservation_from_metadata( - metadata=metadata - ) + budget_reservation = _get_budget_reservation_from_metadata(metadata=metadata) user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) - sl_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + sl_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) response_cost = ( - sl_object.get("response_cost", None) - if sl_object is not None - else kwargs.get("response_cost", None) + sl_object.get("response_cost", None) if sl_object is not None else kwargs.get("response_cost", None) ) tags = _get_request_tags_for_cost_tracking( sl_object=sl_object, @@ -236,9 +209,7 @@ class _ProxyDBLogger(CustomLogger): user_api_key = metadata.get("user_api_key", None) if kwargs.get("cache_hit", False) is True: response_cost = 0.0 - verbose_proxy_logger.debug( - f"Cache Hit: response_cost {response_cost}, for user_id {user_id}" - ) + verbose_proxy_logger.debug(f"Cache Hit: response_cost {response_cost}, for user_id {user_id}") verbose_proxy_logger.debug( f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" @@ -281,19 +252,15 @@ class _ProxyDBLogger(CustomLogger): ) ) - await ( - proxy_logging_obj.slack_alerting_instance.customer_spend_alert( - token=user_api_key, - key_alias=key_alias, - end_user_id=end_user_id, - response_cost=response_cost, - max_budget=end_user_max_budget, - ) + await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( + token=user_api_key, + key_alias=key_alias, + end_user_id=end_user_id, + response_cost=response_cost, + max_budget=end_user_max_budget, ) elif budget_reservation is not None: - await _release_budget_reservation( - budget_reservation=budget_reservation - ) + await _release_budget_reservation(budget_reservation=budget_reservation) else: await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. @@ -301,9 +268,7 @@ class _ProxyDBLogger(CustomLogger): # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls. if sl_object is None and ( - not kwargs.get("model") - or kwargs.get("call_type") - in ("_aresponses_websocket", "_arealtime") + not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime") ): verbose_proxy_logger.warning( "Cost tracking - skipping, no standard_logging_object for call_type=%s", @@ -311,8 +276,7 @@ class _ProxyDBLogger(CustomLogger): ) return if kwargs.get("stream") is not True or ( - kwargs.get("stream") is True - and "complete_streaming_response" in kwargs + kwargs.get("stream") is True and "complete_streaming_response" in kwargs ): if sl_object is not None: cost_tracking_failure_debug_info: Union[dict, str] = ( @@ -320,9 +284,7 @@ class _ProxyDBLogger(CustomLogger): or "response_cost_failure_debug_info is None in standard_logging_object" ) else: - cost_tracking_failure_debug_info = ( - "standard_logging_object not found" - ) + cost_tracking_failure_debug_info = "standard_logging_object not found" model = kwargs.get("model") raise Exception( f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" @@ -331,9 +293,7 @@ class _ProxyDBLogger(CustomLogger): error_msg = f"Error in tracking cost callback - {str(e)}\n Traceback:{traceback.format_exc()}" model = kwargs.get("model", "") metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) - litellm_metadata = kwargs.get("litellm_params", {}).get( - "litellm_metadata", {} - ) + litellm_metadata = kwargs.get("litellm_params", {}).get("litellm_metadata", {}) old_metadata = kwargs.get("litellm_params", {}).get("metadata", {}) call_type = kwargs.get("call_type", "") error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n chosen_metadata: {metadata}\n litellm_metadata: {litellm_metadata}\n old_metadata: {old_metadata}\n call_type: {call_type}\n" @@ -442,12 +402,7 @@ def _should_track_cost_callback( if ProxyUpdateSpend.disable_spend_updates() is True: return False - if ( - user_api_key is not None - or user_id is not None - or team_id is not None - or end_user_id is not None - ): + if user_api_key is not None or user_id is not None or team_id is not None or end_user_id is not None: return True return False @@ -516,13 +471,9 @@ async def _update_database_and_spend_counters( try: await _release_budget_reservation(budget_reservation=budget_reservation) except Exception: - verbose_proxy_logger.exception( - "Failed to release budget reservation after database update failed" - ) + verbose_proxy_logger.exception("Failed to release budget reservation after database update failed") try: - await _invalidate_budget_reservation_counters( - budget_reservation=budget_reservation - ) + await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after release failed" @@ -543,9 +494,7 @@ async def _update_database_and_spend_counters( except Exception: if budget_reservation is not None: try: - await _invalidate_budget_reservation_counters( - budget_reservation=budget_reservation - ) + await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after spend counter update failed" diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 3a23347f351..a7b05c57dac 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -50,27 +50,17 @@ class ResponsesIDSecurity(CustomLogger): if call_type == "aresponses": # check 'previous_response_id' if present in the data previous_response_id = data.get("previous_response_id") - if previous_response_id and self._is_encrypted_response_id( - previous_response_id - ): - original_response_id, user_id, team_id = self._decrypt_response_id( - previous_response_id - ) - self.check_user_access_to_response_id( - user_id, team_id, user_api_key_dict - ) + if previous_response_id and self._is_encrypted_response_id(previous_response_id): + original_response_id, user_id, team_id = self._decrypt_response_id(previous_response_id) + self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) data["previous_response_id"] = original_response_id elif call_type in {"aget_responses", "adelete_responses", "acancel_responses"}: response_id = data.get("response_id") if response_id and self._is_encrypted_response_id(response_id): - original_response_id, user_id, team_id = self._decrypt_response_id( - response_id - ) + original_response_id, user_id, team_id = self._decrypt_response_id(response_id) - self.check_user_access_to_response_id( - user_id, team_id, user_api_key_dict - ) + self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) data["response_id"] = original_response_id return data @@ -118,9 +108,7 @@ class ResponsesIDSecurity(CustomLogger): return False remaining_string = split_result[1] - decrypted_value = decrypt_value_helper( - value=remaining_string, key="response_id", return_original_value=True - ) + decrypted_value = decrypt_value_helper(value=remaining_string, key="response_id", return_original_value=True) if decrypted_value is None: return False @@ -129,9 +117,7 @@ class ResponsesIDSecurity(CustomLogger): return True return False - def _decrypt_response_id( - self, response_id: str - ) -> Tuple[str, Optional[str], Optional[str]]: + def _decrypt_response_id(self, response_id: str) -> Tuple[str, Optional[str], Optional[str]]: """ Returns: - original_response_id: the original response id @@ -143,9 +129,7 @@ class ResponsesIDSecurity(CustomLogger): return response_id, None, None remaining_string = split_result[1] - decrypted_value = decrypt_value_helper( - value=remaining_string, key="response_id", return_original_value=True - ) + decrypted_value = decrypt_value_helper(value=remaining_string, key="response_id", return_original_value=True) if decrypted_value is None: return response_id, None, None @@ -207,11 +191,7 @@ class ResponsesIDSecurity(CustomLogger): response_id = getattr(response, "id", None) response_obj = getattr(response, "response", None) - if ( - response_id - and isinstance(response_id, str) - and response_id.startswith("resp_") - ): + if response_id and isinstance(response_id, str) and response_id.startswith("resp_"): # Check request-scoped cache first (for streaming consistency) if request_cache is not None and response_id in request_cache: setattr(response, "id", request_cache[response_id]) @@ -222,9 +202,7 @@ class ResponsesIDSecurity(CustomLogger): user_api_key_dict.team_id or "", ) - encoded_user_id_and_response_id = encrypt_value_helper( - value=encrypted_response_id - ) + encoded_user_id_and_response_id = encrypt_value_helper(value=encrypted_response_id) encrypted_id = f"resp_{encoded_user_id_and_response_id}" if request_cache is not None: request_cache[response_id] = encrypted_id @@ -240,9 +218,7 @@ class ResponsesIDSecurity(CustomLogger): user_api_key_dict.user_id or "", user_api_key_dict.team_id or "", ) - encoded_user_id_and_response_id = encrypt_value_helper( - value=encrypted_response_id - ) + encoded_user_id_and_response_id = encrypt_value_helper(value=encrypted_response_id) encrypted_id = f"resp_{encoded_user_id_and_response_id}" if request_cache is not None: request_cache[response_obj.id] = encrypted_id @@ -269,9 +245,7 @@ class ResponsesIDSecurity(CustomLogger): if isinstance(response, ResponsesAPIResponse): response = cast( ResponsesAPIResponse, - self._encrypt_response_id( - response, user_api_key_dict, request_cache=None - ), + self._encrypt_response_id(response, user_api_key_dict, request_cache=None), ) return response @@ -290,7 +264,5 @@ class ResponsesIDSecurity(CustomLogger): == "/v1/responses" # only encrypt the response id for the responses api and not general_settings.get("disable_responses_id_security", False) ): - chunk = self._encrypt_response_id( - chunk, user_api_key_dict, request_encryption_cache - ) + chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) yield chunk diff --git a/litellm/proxy/hooks/sensitive_data_routing.py b/litellm/proxy/hooks/sensitive_data_routing.py index 0a907b1d71c..b4f44b5e41a 100644 --- a/litellm/proxy/hooks/sensitive_data_routing.py +++ b/litellm/proxy/hooks/sensitive_data_routing.py @@ -77,24 +77,16 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): ] return "|".join(principal) if principal else "default" - async def _get_routed_model( - self, session_id: str, user_api_key_dict: Optional[UserAPIKeyAuth] - ) -> Optional[str]: + async def _get_routed_model(self, session_id: str, user_api_key_dict: Optional[UserAPIKeyAuth]) -> Optional[str]: """Get the model this session should be routed to, if any.""" - cache_key = self._make_cache_key( - session_id, self._resolve_tenant(user_api_key_dict) - ) + cache_key = self._make_cache_key(session_id, self._resolve_tenant(user_api_key_dict)) if self.internal_usage_cache.dual_cache.redis_cache is not None: try: - result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache( - key=cache_key - ) + result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key) if result is not None: routed_model = str(result) - remaining_ttl = await self.internal_usage_cache.dual_cache.redis_cache.async_get_ttl( - key=cache_key - ) + remaining_ttl = await self.internal_usage_cache.dual_cache.redis_cache.async_get_ttl(key=cache_key) await self.internal_usage_cache.async_set_cache( key=cache_key, value=routed_model, @@ -132,9 +124,7 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger): route the session to a specific model. The override is scoped to the requesting principal so sessions from different tenants cannot collide. """ - cache_key = self._make_cache_key( - session_id, self._resolve_tenant(user_api_key_dict) - ) + cache_key = self._make_cache_key(session_id, self._resolve_tenant(user_api_key_dict)) verbose_proxy_logger.info( "SensitiveDataRoutingHandler: Setting session routing session_id=%s model=%s guardrail=%s ttl=%s", diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 8122c5e68c6..6122f0594e8 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -62,9 +62,7 @@ class UserManagementEventHooks: where={"user_id": response.user_id} ) - user_row_litellm_typed = LiteLLM_UserTable( - **user_row.model_dump(exclude_none=True) - ) + user_row_litellm_typed = LiteLLM_UserTable(**user_row.model_dump(exclude_none=True)) asyncio.create_task( UserManagementEventHooks.create_internal_user_audit_log( user_id=user_row_litellm_typed.user_id, @@ -73,15 +71,11 @@ class UserManagementEventHooks: user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, before_value=None, - after_value=user_row_litellm_typed.model_dump_json( - exclude_none=True - ), + after_value=user_row_litellm_typed.model_dump_json(exclude_none=True), ) ) except Exception as e: - verbose_proxy_logger.warning( - "Unable to create audit log for user on `/user/new` - {}".format(str(e)) - ) + verbose_proxy_logger.warning("Unable to create audit log for user on `/user/new` - {}".format(str(e))) pass @staticmethod @@ -117,16 +111,13 @@ class UserManagementEventHooks: use_enterprise_email_hooks = True except ImportError: verbose_proxy_logger.warning( - "Defaulting to using Legacy Email Hooks." - + CommonProxyErrors.missing_enterprise_package.value + "Defaulting to using Legacy Email Hooks." + CommonProxyErrors.missing_enterprise_package.value ) use_enterprise_email_hooks = False if use_enterprise_email_hooks and (data.send_invite_email is True): - initialized_email_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=BaseEmailLogger # type: ignore - ) + initialized_email_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=BaseEmailLogger # type: ignore ) if len(initialized_email_loggers) > 0: for email_logger in initialized_email_loggers: diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index c217116e45f..8178cad9038 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -140,9 +140,7 @@ async def image_generation( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### CALL HOOKS ### - modify outgoing data (guardrails, otel, etc.) @@ -189,9 +187,7 @@ async def image_generation( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.error( - "litellm.proxy.proxy_server.image_generation(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.image_generation(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): @@ -252,13 +248,9 @@ async def image_edit_api( ``` """ if image is not None and image_array is not None: - raise HTTPException( - status_code=422, detail="Cannot specify both 'image' and 'image[]'" - ) + raise HTTPException(status_code=422, detail="Cannot specify both 'image' and 'image[]'") if mask is not None and mask_array is not None: - raise HTTPException( - status_code=422, detail="Cannot specify both 'mask' and 'mask[]'" - ) + raise HTTPException(status_code=422, detail="Cannot specify both 'mask' and 'mask[]'") if image is None and image_array is not None: image = image_array if mask is None and mask_array is not None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 36881765596..0ffb0337545 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -33,9 +33,7 @@ from litellm.proxy.common_utils.callback_utils import ( from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers # Cache special headers as a frozenset for O(1) lookup performance -_SPECIAL_HEADERS_CACHE = frozenset( - v.value.lower() for v in SpecialHeaders._member_map_.values() -) +_SPECIAL_HEADERS_CACHE = frozenset(v.value.lower() for v in SpecialHeaders._member_map_.values()) # Matches any header of the form x--session-id (case-insensitive). # Excludes the two explicit litellm headers which are handled with higher priority. @@ -189,9 +187,7 @@ _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset( ) _CLIENT_MOCK_CONTROL_FIELDS = frozenset({"mock_response", "mock_tool_calls"}) _ALLOW_CLIENT_MOCK_RESPONSE_METADATA_KEY = "allow_client_mock_response" -_ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY = ( - "allow_client_message_redaction_opt_out" -) +_ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY = "allow_client_message_redaction_opt_out" # Per-request pricing parameters mutate cost-tracking output and (via # ``litellm.completion`` → ``register_model``) the process-wide @@ -199,9 +195,7 @@ _ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY = ( # not to user-supplied request bodies, so the proxy strips them before they # reach the call path. Built from the Pydantic model so newly-added pricing # fields are covered automatically. -_CLIENT_PRICING_CONTROL_FIELDS = frozenset( - CustomPricingLiteLLMParams.model_fields.keys() -) +_CLIENT_PRICING_CONTROL_FIELDS = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) # ``model_info`` carries the same pricing fields when read by # ``use_custom_pricing_for_model``; strip from metadata for the same reason. _CLIENT_PRICING_METADATA_FIELDS = frozenset({"model_info"}) @@ -256,10 +250,7 @@ def _strip_untrusted_request_header_controls( return for header_name in list(headers.keys()): - if ( - isinstance(header_name, str) - and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS - ): + if isinstance(header_name, str) and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: if allow_client_message_redaction_opt_out: continue headers.pop(header_name, None) @@ -278,10 +269,7 @@ def _key_or_team_metadata_flag_is_true( metadata_key: str, ) -> bool: for admin_metadata in (user_api_key_dict.metadata, user_api_key_dict.team_metadata): - if ( - isinstance(admin_metadata, dict) - and admin_metadata.get(metadata_key) is True - ): + if isinstance(admin_metadata, dict) and admin_metadata.get(metadata_key) is True: return True return False @@ -422,9 +410,7 @@ def is_claude_code_user_agent(user_agent: str) -> bool: return user_agent.startswith("claude-cli/") -def should_auto_drop_params_for_claude_code( - user_agent: str, data: dict, proxy_config: ProxyConfig -) -> bool: +def should_auto_drop_params_for_claude_code(user_agent: str, data: dict, proxy_config: ProxyConfig) -> bool: """drop_params defaults to on for Claude Code so its Anthropic-specific params (e.g. thinking) don't fail requests routed to non-Anthropic providers. An explicit drop_params from the caller or in the operator's @@ -434,12 +420,8 @@ def should_auto_drop_params_for_claude_code( if "drop_params" in data: return False config = getattr(proxy_config, "config", None) - litellm_settings = ( - config.get("litellm_settings") if isinstance(config, dict) else None - ) - return not ( - isinstance(litellm_settings, dict) and "drop_params" in litellm_settings - ) + litellm_settings = config.get("litellm_settings") if isinstance(config, dict) else None + return not (isinstance(litellm_settings, dict) and "drop_params" in litellm_settings) def safe_add_api_version_from_query_params(data: dict, request: Request): @@ -451,9 +433,7 @@ def safe_add_api_version_from_query_params(data: dict, request: Request): except KeyError: pass except Exception as e: - verbose_logger.exception( - "error checking api version in query params: %s", str(e) - ) + verbose_logger.exception("error checking api version in query params: %s", str(e)) def convert_key_logging_metadata_to_callback( @@ -500,9 +480,7 @@ def convert_key_logging_metadata_to_callback( return team_callback_settings_obj -def _get_validated_callback_metadata( - item: dict, *, source: str -) -> Optional[AddTeamCallback]: +def _get_validated_callback_metadata(item: dict, *, source: str) -> Optional[AddTeamCallback]: try: return AddTeamCallback(**item) except (PydanticValidationError, ValueError) as e: @@ -521,19 +499,13 @@ class KeyAndTeamLoggingSettings: @staticmethod def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): - if ( - user_api_key_dict.metadata is not None - and "logging" in user_api_key_dict.metadata - ): + if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata: return decrypt_callback_vars(user_api_key_dict.metadata).get("logging") return None @staticmethod def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): - if ( - user_api_key_dict.team_metadata is not None - and "logging" in user_api_key_dict.team_metadata - ): + if user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata: return decrypt_callback_vars(user_api_key_dict.team_metadata).get("logging") return None @@ -542,11 +514,11 @@ def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - key_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + key_dynamic_logging_settings: Optional[dict] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( + user_api_key_dict ) - team_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) + team_dynamic_logging_settings: Optional[dict] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( + user_api_key_dict ) ######################################################################################### # Key-based callbacks @@ -575,10 +547,7 @@ def _get_dynamic_logging_metadata( ######################################################################################### # Deprecated format - maintained for backwards compatibility ######################################################################################### - elif ( - user_api_key_dict.team_metadata is not None - and "callback_settings" in user_api_key_dict.team_metadata - ): + elif user_api_key_dict.team_metadata is not None and "callback_settings" in user_api_key_dict.team_metadata: """ callback_settings = { { @@ -591,17 +560,13 @@ def _get_dynamic_logging_metadata( team_metadata = decrypt_callback_vars(user_api_key_dict.team_metadata) callback_settings = team_metadata.get("callback_settings", None) or {} callback_settings_obj = TeamCallbackMetadata(**callback_settings) - verbose_proxy_logger.debug( - "Team callback settings activated: %s", callback_settings_obj - ) + verbose_proxy_logger.debug("Team callback settings activated: %s", callback_settings_obj) ######################################################################################### # Enter here when configured on the config.yaml file. ######################################################################################### elif user_api_key_dict.team_id is not None: - callback_settings_obj = ( - LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config( - team_id=user_api_key_dict.team_id, proxy_config=proxy_config - ) + callback_settings_obj = LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config( + team_id=user_api_key_dict.team_id, proxy_config=proxy_config ) return callback_settings_obj @@ -628,29 +593,21 @@ def clean_headers( from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key clean_headers = {} - litellm_key_lower = ( - litellm_key_header_name.lower() if litellm_key_header_name is not None else None - ) + litellm_key_lower = litellm_key_header_name.lower() if litellm_key_header_name is not None else None for header, value in headers.items(): header_lower = header.lower() if header_lower == "authorization" and is_anthropic_oauth_key(value): - if ( - authenticated_with_header is None - or authenticated_with_header.lower() != "authorization" - ): + if authenticated_with_header is None or authenticated_with_header.lower() != "authorization": clean_headers[header] = value continue # Special handling for x-api-key: forward it based on authenticated_with_header elif header_lower == "x-api-key": if forward_llm_provider_auth_headers and ( - authenticated_with_header is None - or authenticated_with_header.lower() != "x-api-key" + authenticated_with_header is None or authenticated_with_header.lower() != "x-api-key" ): clean_headers[header] = value - elif ( - forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE - ): + elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: if litellm_key_lower and header_lower == litellm_key_lower: continue if header_lower == "authorization": @@ -768,23 +725,17 @@ class LiteLLMProxyRequestSetup: user_header_mapping = general_settings.get("user_header_mappings") if not user_header_mapping: return user_api_key_dict - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - user_header_mapping - ) + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(user_header_mapping) if not header_name: return user_api_key_dict - header_value = LiteLLMProxyRequestSetup._get_case_insensitive_header( - headers, header_name - ) + header_value = LiteLLMProxyRequestSetup._get_case_insensitive_header(headers, header_name) if header_value: user_api_key_dict.user_id = header_value return user_api_key_dict return user_api_key_dict @staticmethod - def get_user_from_headers( - headers: dict, general_settings: Optional[Dict] = None - ) -> Optional[str]: + def get_user_from_headers(headers: dict, general_settings: Optional[Dict] = None) -> Optional[str]: """ Get the user from the specified header if `general_settings.user_header_name` is set. """ @@ -796,29 +747,20 @@ class LiteLLMProxyRequestSetup: return None if not isinstance(header_name, str): - raise TypeError( - f"Expected user_header_name to be a str but got {type(header_name)}" - ) + raise TypeError(f"Expected user_header_name to be a str but got {type(header_name)}") - user = LiteLLMProxyRequestSetup._get_case_insensitive_header( - headers, header_name - ) + user = LiteLLMProxyRequestSetup._get_case_insensitive_header(headers, header_name) if user is not None: verbose_logger.info(f'found user "{user}" in header "{header_name}"') return user @staticmethod - def get_openai_org_id_from_headers( - headers: dict, general_settings: Optional[Dict] = None - ) -> Optional[str]: + def get_openai_org_id_from_headers(headers: dict, general_settings: Optional[Dict] = None) -> Optional[str]: """ Get the OpenAI Org ID from the headers. """ - if ( - general_settings is not None - and general_settings.get("forward_openai_org_id") is not True - ): + if general_settings is not None and general_settings.get("forward_openai_org_id") is not True: return None for header, value in headers.items(): if header.lower() == "openai-organization": @@ -827,9 +769,7 @@ class LiteLLMProxyRequestSetup: return None @staticmethod - def add_headers_to_llm_call( - headers: dict, user_api_key_dict: UserAPIKeyAuth - ) -> dict: + def add_headers_to_llm_call(headers: dict, user_api_key_dict: UserAPIKeyAuth) -> dict: """ Add headers to the LLM call @@ -840,10 +780,8 @@ class LiteLLMProxyRequestSetup: returned_headers = LiteLLMProxyRequestSetup._get_forwardable_headers(headers) if litellm.add_user_information_to_llm_headers is True: - litellm_logging_metadata_headers = ( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + litellm_logging_metadata_headers = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict ) for k, v in litellm_logging_metadata_headers.items(): if v is None: @@ -861,9 +799,7 @@ class LiteLLMProxyRequestSetup: return returned_headers @staticmethod - def add_headers_to_llm_call_by_model_group( - data: dict, headers: dict, user_api_key_dict: UserAPIKeyAuth - ) -> dict: + def add_headers_to_llm_call_by_model_group(data: dict, headers: dict, user_api_key_dict: UserAPIKeyAuth) -> dict: """ Add headers to the LLM call by model group """ @@ -875,8 +811,7 @@ class LiteLLMProxyRequestSetup: if ( data_model is not None and litellm.model_group_settings is not None - and litellm.model_group_settings.forward_client_headers_to_llm_api - is not None + and litellm.model_group_settings.forward_client_headers_to_llm_api is not None and _check_model_access_helper( model=data_model, llm_router=llm_router, @@ -885,9 +820,7 @@ class LiteLLMProxyRequestSetup: team_id=user_api_key_dict.team_id, ) # handles aliases, wildcards, etc. ): - _headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call( - headers, user_api_key_dict - ) + _headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call(headers, user_api_key_dict) if _headers != {}: data["headers"] = _headers return data @@ -896,11 +829,7 @@ class LiteLLMProxyRequestSetup: def get_internal_user_header_from_mapping(user_header_mapping) -> Optional[str]: if not user_header_mapping: return None - items = ( - user_header_mapping - if isinstance(user_header_mapping, list) - else [user_header_mapping] - ) + items = user_header_mapping if isinstance(user_header_mapping, list) else [user_header_mapping] for item in items: if not isinstance(item, dict): continue @@ -926,18 +855,11 @@ class LiteLLMProxyRequestSetup: """ data = LitellmDataForBackendLLMCall() - if ( - general_settings - and general_settings.get("forward_client_headers_to_llm_api") is True - ): - _headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call( - headers, user_api_key_dict - ) + if general_settings and general_settings.get("forward_client_headers_to_llm_api") is True: + _headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call(headers, user_api_key_dict) if _headers != {}: data["headers"] = _headers - _organization = LiteLLMProxyRequestSetup.get_openai_org_id_from_headers( - headers, general_settings - ) + _organization = LiteLLMProxyRequestSetup.get_openai_org_id_from_headers(headers, general_settings) if _organization is not None: data["organization"] = _organization @@ -945,9 +867,7 @@ class LiteLLMProxyRequestSetup: if timeout is not None: data["timeout"] = timeout - stream_timeout = LiteLLMProxyRequestSetup._get_stream_timeout_from_request( - headers - ) + stream_timeout = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers) if stream_timeout is not None: data["stream_timeout"] = stream_timeout @@ -971,11 +891,7 @@ class LiteLLMProxyRequestSetup: from litellm.proxy._types import LitellmMetadataFromRequestHeaders metadata_from_headers = LitellmMetadataFromRequestHeaders() - spend_logs_metadata = ( - LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers( - headers - ) - ) + spend_logs_metadata = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) if spend_logs_metadata is not None: metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata @@ -989,18 +905,14 @@ class LiteLLMProxyRequestSetup: if agent_id_from_header: metadata_from_headers["agent_id"] = agent_id_from_header - verbose_proxy_logger.debug( - f"Extracted agent_id from header: {agent_id_from_header}" - ) + verbose_proxy_logger.debug(f"Extracted agent_id from header: {agent_id_from_header}") if chain_id: metadata_from_headers["trace_id"] = chain_id metadata_from_headers["session_id"] = chain_id data["litellm_session_id"] = chain_id data["litellm_trace_id"] = chain_id - verbose_proxy_logger.debug( - f"Extracted chain_id from header (trace-id/session-id): {chain_id}" - ) + verbose_proxy_logger.debug(f"Extracted chain_id from header (trace-id/session-id): {chain_id}") if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update(metadata_from_headers) @@ -1026,9 +938,7 @@ class LiteLLMProxyRequestSetup: user_api_key_user_email=user_api_key_dict.user_email, user_api_key_request_route=user_api_key_dict.request_route, user_api_key_budget_reset_at=( - user_api_key_dict.budget_reset_at.isoformat() - if user_api_key_dict.budget_reset_at - else None + user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None ), user_api_key_auth_metadata=user_api_key_dict.metadata, ) @@ -1043,15 +953,11 @@ class LiteLLMProxyRequestSetup: """ Adds the `UserAPIKeyAuth` object to the request metadata. """ - user_api_key_logged_metadata = ( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + user_api_key_logged_metadata = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict ) data[_metadata_variable_name].update(user_api_key_logged_metadata) - data[_metadata_variable_name]["user_api_key"] = ( - user_api_key_dict.api_key - ) # this is just the hashed token + data[_metadata_variable_name]["user_api_key"] = user_api_key_dict.api_key # this is just the hashed token # Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none _key_agent_id = getattr(user_api_key_dict, "agent_id", None) @@ -1063,9 +969,7 @@ class LiteLLMProxyRequestSetup: user_api_key_dict, "end_user_max_budget", None ) if user_api_key_dict.budget_reservation is not None: - data[_metadata_variable_name]["user_api_key_budget_reservation"] = ( - user_api_key_dict.budget_reservation - ) + data[_metadata_variable_name]["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation # Add the full UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict return data @@ -1091,22 +995,15 @@ class LiteLLMProxyRequestSetup: # ignore any special fields added_metadata = {} for k, v in management_endpoint_metadata.items(): - if k not in ( - LiteLLM_ManagementEndpoint_MetadataFields_Premium - + LiteLLM_ManagementEndpoint_MetadataFields - ): + if k not in (LiteLLM_ManagementEndpoint_MetadataFields_Premium + LiteLLM_ManagementEndpoint_MetadataFields): added_metadata[k] = v if data[_metadata_variable_name].get("user_api_key_auth_metadata") is None: data[_metadata_variable_name]["user_api_key_auth_metadata"] = {} - data[_metadata_variable_name]["user_api_key_auth_metadata"].update( - added_metadata - ) + data[_metadata_variable_name]["user_api_key_auth_metadata"].update(added_metadata) return data @staticmethod - def add_key_level_controls( - key_metadata: Optional[dict], data: dict, _metadata_variable_name: str - ): + def add_key_level_controls(key_metadata: Optional[dict], data: dict, _metadata_variable_name: str): if key_metadata is None: return data if "cache" in key_metadata: @@ -1118,21 +1015,13 @@ class LiteLLMProxyRequestSetup: ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = ( - LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=key_metadata["tags"], - ) + data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=key_metadata["tags"], ) - if "disable_global_guardrails" in key_metadata and isinstance( - key_metadata["disable_global_guardrails"], bool - ): - data[_metadata_variable_name]["disable_global_guardrails"] = key_metadata[ - "disable_global_guardrails" - ] - if "spend_logs_metadata" in key_metadata and isinstance( - key_metadata["spend_logs_metadata"], dict - ): + if "disable_global_guardrails" in key_metadata and isinstance(key_metadata["disable_global_guardrails"], bool): + data[_metadata_variable_name]["disable_global_guardrails"] = key_metadata["disable_global_guardrails"] + if "spend_logs_metadata" in key_metadata and isinstance(key_metadata["spend_logs_metadata"], dict): if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance( data[_metadata_variable_name]["spend_logs_metadata"], dict ): @@ -1140,18 +1029,12 @@ class LiteLLMProxyRequestSetup: if ( key not in data[_metadata_variable_name]["spend_logs_metadata"] ): # don't override k-v pair sent by request (user request) - data[_metadata_variable_name]["spend_logs_metadata"][key] = ( - value - ) + data[_metadata_variable_name]["spend_logs_metadata"][key] = value else: - data[_metadata_variable_name]["spend_logs_metadata"] = key_metadata[ - "spend_logs_metadata" - ] + data[_metadata_variable_name]["spend_logs_metadata"] = key_metadata["spend_logs_metadata"] ## KEY-LEVEL DISABLE FALLBACKS - if "disable_fallbacks" in key_metadata and isinstance( - key_metadata["disable_fallbacks"], bool - ): + if "disable_fallbacks" in key_metadata and isinstance(key_metadata["disable_fallbacks"], bool): data["disable_fallbacks"] = key_metadata["disable_fallbacks"] ## KEY-LEVEL METADATA @@ -1203,11 +1086,7 @@ class LiteLLMProxyRequestSetup: callback_vars_dict.pop("success_callback", None) callback_vars_dict.pop("failure_callback", None) callback_vars_dict = { - key: ( - litellm.utils.get_secret(value, default_value=value) or value - if isinstance(value, str) - else value - ) + key: (litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else value) for key, value in callback_vars_dict.items() } @@ -1320,9 +1199,7 @@ class LiteLLMProxyRequestSetup: return if isinstance(raw_header_tags, str): - header_tags: List[str] = [ - t.strip() for t in raw_header_tags.split(",") if t.strip() - ] + header_tags: List[str] = [t.strip() for t in raw_header_tags.split(",") if t.strip()] elif isinstance(raw_header_tags, list): header_tags = [t for t in raw_header_tags if isinstance(t, str) and t] else: @@ -1384,12 +1261,8 @@ async def add_litellm_data_to_request( # Strip internal-only keys from user input before the proxy sets its own. # These keys are injected by the proxy itself below — user-supplied values # must not be trusted. - _allow_client_mock_response = _key_or_team_allows_client_mock_response( - user_api_key_dict - ) - _allow_client_message_redaction_opt_out = ( - _key_or_team_allows_client_message_redaction_opt_out(user_api_key_dict) - ) + _allow_client_mock_response = _key_or_team_allows_client_mock_response(user_api_key_dict) + _allow_client_message_redaction_opt_out = _key_or_team_allows_client_message_redaction_opt_out(user_api_key_dict) for _internal_key in _UNTRUSTED_ROOT_CONTROL_FIELDS: if _allow_client_mock_response and _internal_key in _CLIENT_MOCK_CONTROL_FIELDS: continue @@ -1406,9 +1279,7 @@ async def add_litellm_data_to_request( forward_llm_auth = False if general_settings: - forward_llm_auth = general_settings.get( - "forward_llm_provider_auth_headers", False - ) + forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) if not forward_llm_auth: forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False) # Determine which header was used for authentication @@ -1427,9 +1298,7 @@ async def add_litellm_data_to_request( _headers: Dict[str, str] = clean_headers( request.headers, litellm_key_header_name=( - general_settings.get("litellm_key_header_name") - if general_settings is not None - else None + general_settings.get("litellm_key_header_name") if general_settings is not None else None ), forward_llm_provider_auth_headers=forward_llm_auth, authenticated_with_header=authenticated_with_header, @@ -1491,9 +1360,7 @@ async def add_litellm_data_to_request( ) # Expose request headers under the metadata field for guardrails (fixes #17477) - if _metadata_variable_name in data and isinstance( - data[_metadata_variable_name], dict - ): + if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name]["headers"] = _headers # check for forwardable headers @@ -1583,15 +1450,10 @@ async def add_litellm_data_to_request( if isinstance(_user_meta, dict): _strip_untrusted_request_header_controls( _user_meta.get("headers"), - allow_client_message_redaction_opt_out=( - _allow_client_message_redaction_opt_out - ), + allow_client_message_redaction_opt_out=(_allow_client_message_redaction_opt_out), ) for _k in [ - k - for k in _user_meta - if k.startswith("user_api_key_") - or k in _UNTRUSTED_METADATA_CONTROL_FIELDS + k for k in _user_meta if k.startswith("user_api_key_") or k in _UNTRUSTED_METADATA_CONTROL_FIELDS ]: _user_meta.pop(_k, None) @@ -1626,9 +1488,7 @@ async def add_litellm_data_to_request( # them — from leaking into requester_metadata where guardrails and audit # paths may read from it. if "metadata" in data and isinstance(data["metadata"], dict): - data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy( - data["metadata"] - ) + data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"]) # Merge litellm_metadata into the metadata variable (preserving existing # values). Runs after the user_api_key_* / _pipeline_managed_guardrails @@ -1648,8 +1508,8 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["litellm_api_version"] = version if general_settings is not None: - data[_metadata_variable_name]["global_max_parallel_requests"] = ( - general_settings.get("global_max_parallel_requests", None) + data[_metadata_variable_name]["global_max_parallel_requests"] = general_settings.get( + "global_max_parallel_requests", None ) ### KEY-LEVEL Controls @@ -1666,21 +1526,13 @@ async def add_litellm_data_to_request( request_tags=data[_metadata_variable_name].get("tags"), tags_to_add=team_metadata["tags"], ) - if "disable_global_guardrails" in team_metadata and isinstance( - team_metadata["disable_global_guardrails"], bool - ): - data[_metadata_variable_name]["disable_global_guardrails"] = team_metadata[ - "disable_global_guardrails" - ] + if "disable_global_guardrails" in team_metadata and isinstance(team_metadata["disable_global_guardrails"], bool): + data[_metadata_variable_name]["disable_global_guardrails"] = team_metadata["disable_global_guardrails"] if "opted_out_global_guardrails" in team_metadata and isinstance( team_metadata["opted_out_global_guardrails"], list ): - data[_metadata_variable_name]["opted_out_global_guardrails"] = team_metadata[ - "opted_out_global_guardrails" - ] - if "spend_logs_metadata" in team_metadata and isinstance( - team_metadata["spend_logs_metadata"], dict - ): + data[_metadata_variable_name]["opted_out_global_guardrails"] = team_metadata["opted_out_global_guardrails"] + if "spend_logs_metadata" in team_metadata and isinstance(team_metadata["spend_logs_metadata"], dict): if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance( data[_metadata_variable_name]["spend_logs_metadata"], dict ): @@ -1690,9 +1542,7 @@ async def add_litellm_data_to_request( ): # don't override k-v pair sent by request (user request) data[_metadata_variable_name]["spend_logs_metadata"][key] = value else: - data[_metadata_variable_name]["spend_logs_metadata"] = team_metadata[ - "spend_logs_metadata" - ] + data[_metadata_variable_name]["spend_logs_metadata"] = team_metadata["spend_logs_metadata"] ## PROJECT-LEVEL TAGS project_metadata = user_api_key_dict.project_metadata or {} @@ -1703,50 +1553,32 @@ async def add_litellm_data_to_request( ) ## TEAM-LEVEL METADATA - data = ( - LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( - data=data, - management_endpoint_metadata=team_metadata, - _metadata_variable_name=_metadata_variable_name, - ) + data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( + data=data, + management_endpoint_metadata=team_metadata, + _metadata_variable_name=_metadata_variable_name, ) # Team spend, budget - used by prometheus.py - data[_metadata_variable_name]["user_api_key_team_max_budget"] = ( - user_api_key_dict.team_max_budget - ) - data[_metadata_variable_name]["user_api_key_team_spend"] = ( - user_api_key_dict.team_spend - ) - data[_metadata_variable_name]["user_api_key_request_route"] = ( - user_api_key_dict.request_route - ) + data[_metadata_variable_name]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget + data[_metadata_variable_name]["user_api_key_team_spend"] = user_api_key_dict.team_spend + data[_metadata_variable_name]["user_api_key_request_route"] = user_api_key_dict.request_route # API Key spend, budget - used by prometheus.py data[_metadata_variable_name]["user_api_key_spend"] = user_api_key_dict.spend - data[_metadata_variable_name]["user_api_key_max_budget"] = ( - user_api_key_dict.max_budget - ) - data[_metadata_variable_name]["user_api_key_model_max_budget"] = ( - user_api_key_dict.model_max_budget - ) + data[_metadata_variable_name]["user_api_key_max_budget"] = user_api_key_dict.max_budget + data[_metadata_variable_name]["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget data[_metadata_variable_name]["user_api_key_end_user_model_max_budget"] = ( user_api_key_dict.end_user_model_max_budget ) # User spend, budget - used by prometheus.py # Follow same pattern as team and API key budgets - data[_metadata_variable_name]["user_api_key_user_spend"] = ( - user_api_key_dict.user_spend - ) - data[_metadata_variable_name]["user_api_key_user_max_budget"] = ( - user_api_key_dict.user_max_budget - ) + data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend + data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata - data[_metadata_variable_name]["user_api_key_team_metadata"] = ( - user_api_key_dict.team_metadata - ) + data[_metadata_variable_name]["user_api_key_team_metadata"] = user_api_key_dict.team_metadata data[_metadata_variable_name]["user_api_key_object_permission_id"] = getattr( user_api_key_dict, "object_permission_id", None ) @@ -1758,15 +1590,11 @@ async def add_litellm_data_to_request( # Carry the proxy-receive instant via metadata (like `endpoint`) so the # OTel layer can compute pre-request latency, including on the failure # path after the logging object is popped. - data[_metadata_variable_name]["litellm_received_at"] = getattr( - request.state, "litellm_received_at", None - ) + data[_metadata_variable_name]["litellm_received_at"] = getattr(request.state, "litellm_received_at", None) # OTEL Controls / Tracing # Add the OTEL Parent Trace before sending it LiteLLM - data[_metadata_variable_name]["litellm_parent_otel_span"] = ( - user_api_key_dict.parent_otel_span - ) + data[_metadata_variable_name]["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _add_otel_traceparent_to_data(data, request=request) ### END-USER SPECIFIC PARAMS ### @@ -1799,11 +1627,7 @@ async def add_litellm_data_to_request( # Add User-Agent user_agent = "" - if ( - request is not None - and hasattr(request, "headers") - and "user-agent" in request.headers - ): + if request is not None and hasattr(request, "headers") and "user-agent" in request.headers: user_agent = request.headers["user-agent"] data[_metadata_variable_name]["user_agent"] = user_agent @@ -1838,10 +1662,7 @@ async def add_litellm_data_to_request( data[k] = v # Add disabled callbacks from key metadata - if ( - user_api_key_dict.metadata - and "litellm_disabled_callbacks" in user_api_key_dict.metadata - ): + if user_api_key_dict.metadata and "litellm_disabled_callbacks" in user_api_key_dict.metadata: disabled_callbacks = user_api_key_dict.metadata["litellm_disabled_callbacks"] if disabled_callbacks and isinstance(disabled_callbacks, list): data["litellm_disabled_callbacks"] = disabled_callbacks @@ -1868,9 +1689,7 @@ async def add_litellm_data_to_request( user_api_key_dict=user_api_key_dict, ) - verbose_proxy_logger.debug( - "[PROXY] returned data from litellm_pre_call_utils: %s", data - ) + verbose_proxy_logger.debug("[PROXY] returned data from litellm_pre_call_utils: %s", data) # Team/Project credential overrides from model_config # Placed after the debug log to avoid leaking credential secrets in logs @@ -1927,11 +1746,7 @@ def _update_model_if_team_alias_exists( and are resolved via map_team_model in route_llm_request. """ _model = data.get("model") - if ( - _model - and user_api_key_dict.team_model_aliases - and _model in user_api_key_dict.team_model_aliases - ): + if _model and user_api_key_dict.team_model_aliases and _model in user_api_key_dict.team_model_aliases: from litellm.proxy.proxy_server import llm_router # Skip alias rewrite if this model resolves to team-specific deployments @@ -1943,15 +1758,11 @@ def _update_model_if_team_alias_exists( # Cached at module level to avoid hot-path secret lookups on every request. global _ENABLE_TEAM_STALE_ALIAS_BYPASS if _ENABLE_TEAM_STALE_ALIAS_BYPASS is None: - _ENABLE_TEAM_STALE_ALIAS_BYPASS = get_secret_bool( - "LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", False - ) + _ENABLE_TEAM_STALE_ALIAS_BYPASS = get_secret_bool("LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", False) enable_stale_alias_bypass = _ENABLE_TEAM_STALE_ALIAS_BYPASS # Check if the alias points to a team-scoped UUID name # (format: "model_name_{team_id}_{uuid}") - is_stale_team_alias = aliased_target.startswith( - f"model_name_{user_api_key_dict.team_id}_" - ) + is_stale_team_alias = aliased_target.startswith(f"model_name_{user_api_key_dict.team_id}_") if is_stale_team_alias and llm_router: # This is a stale alias from pre-PR deployments. # Check if current team deployments exist for the public name. @@ -1963,10 +1774,7 @@ def _update_model_if_team_alias_exists( warning_key = f"{user_api_key_dict.team_id}:{_model}:{aliased_target}" if warning_key not in _STALE_TEAM_ALIAS_WARNING_KEYS: _STALE_TEAM_ALIAS_WARNING_KEYS[warning_key] = None - while ( - len(_STALE_TEAM_ALIAS_WARNING_KEYS) - > _MAX_STALE_ALIAS_WARNING_KEYS - ): + while len(_STALE_TEAM_ALIAS_WARNING_KEYS) > _MAX_STALE_ALIAS_WARNING_KEYS: _STALE_TEAM_ALIAS_WARNING_KEYS.popitem(last=False) verbose_proxy_logger.warning( "Stale team model alias detected for model='%s', team_id='%s'. " @@ -2115,9 +1923,7 @@ def _resolve_provider_from_deployment( for name in candidates: try: - deployment = llm_router.get_deployment_by_model_group_name( - model_group_name=name - ) + deployment = llm_router.get_deployment_by_model_group_name(model_group_name=name) except Exception: deployment = None if deployment is None: @@ -2174,33 +1980,26 @@ def _resolve_credential_from_model_config( for name in model_names_to_try: model_entry = model_config.get(name) if model_entry: - credential_name = _extract_credential_from_entry( - model_entry, provider=provider - ) + credential_name = _extract_credential_from_entry(model_entry, provider=provider) if credential_name: return credential_name _safe_name = str(name).replace("\n", "").replace("\r", "") verbose_proxy_logger.debug( - "model_config entry '%s' found but has no litellm_credentials, " - "trying next candidate", + "model_config entry '%s' found but has no litellm_credentials, trying next candidate", _safe_name, ) # Default check default_entry = model_config.get("defaultconfig") if default_entry: - credential_name = _extract_credential_from_entry( - default_entry, provider=provider - ) + credential_name = _extract_credential_from_entry(default_entry, provider=provider) if credential_name: return credential_name return None -def _extract_credential_from_entry( - entry: dict, provider: Optional[str] = None -) -> Optional[str]: +def _extract_credential_from_entry(entry: dict, provider: Optional[str] = None) -> Optional[str]: """ Extract litellm_credentials from a model_config entry. @@ -2229,9 +2028,7 @@ def _extract_credential_from_entry( return None -def _get_enforced_params( - general_settings: Optional[dict], user_api_key_dict: UserAPIKeyAuth -) -> Optional[list]: +def _get_enforced_params(general_settings: Optional[dict], user_api_key_dict: UserAPIKeyAuth) -> Optional[list]: enforced_params: Optional[list] = None if general_settings is not None: enforced_params = general_settings.get("enforced_params") @@ -2331,28 +2128,19 @@ def _add_guardrails_from_key_or_team_metadata( # Add key-level guardrails first if key_metadata and "guardrails" in key_metadata: - if ( - isinstance(key_metadata["guardrails"], list) - and len(key_metadata["guardrails"]) > 0 - ): + if isinstance(key_metadata["guardrails"], list) and len(key_metadata["guardrails"]) > 0: _premium_user_check() combined_guardrails.update(key_metadata["guardrails"]) # Add team-level guardrails (set automatically handles duplicates) if team_metadata and "guardrails" in team_metadata: - if ( - isinstance(team_metadata["guardrails"], list) - and len(team_metadata["guardrails"]) > 0 - ): + if isinstance(team_metadata["guardrails"], list) and len(team_metadata["guardrails"]) > 0: _premium_user_check() combined_guardrails.update(team_metadata["guardrails"]) # Add project-level guardrails (set automatically handles duplicates) if project_metadata and "guardrails" in project_metadata: - if ( - isinstance(project_metadata["guardrails"], list) - and len(project_metadata["guardrails"]) > 0 - ): + if isinstance(project_metadata["guardrails"], list) and len(project_metadata["guardrails"]) > 0: _premium_user_check() combined_guardrails.update(project_metadata["guardrails"]) @@ -2394,44 +2182,31 @@ def _add_guardrails_from_policies_in_metadata( # Add key-level policies first if key_metadata and "policies" in key_metadata: - if ( - isinstance(key_metadata["policies"], list) - and len(key_metadata["policies"]) > 0 - ): + if isinstance(key_metadata["policies"], list) and len(key_metadata["policies"]) > 0: _premium_user_check() policy_names.update(key_metadata["policies"]) # Add team-level policies if team_metadata and "policies" in team_metadata: - if ( - isinstance(team_metadata["policies"], list) - and len(team_metadata["policies"]) > 0 - ): + if isinstance(team_metadata["policies"], list) and len(team_metadata["policies"]) > 0: _premium_user_check() policy_names.update(team_metadata["policies"]) # Add project-level policies if project_metadata and "policies" in project_metadata: - if ( - isinstance(project_metadata["policies"], list) - and len(project_metadata["policies"]) > 0 - ): + if isinstance(project_metadata["policies"], list) and len(project_metadata["policies"]) > 0: _premium_user_check() policy_names.update(project_metadata["policies"]) if not policy_names: return - verbose_proxy_logger.debug( - f"Policy engine: resolving guardrails from key/team policies: {policy_names}" - ) + verbose_proxy_logger.debug(f"Policy engine: resolving guardrails from key/team policies: {policy_names}") # Check if policy registry is initialized registry = get_policy_registry() if not registry.is_initialized(): - verbose_proxy_logger.debug( - "Policy engine not initialized, skipping policy resolution from metadata" - ) + verbose_proxy_logger.debug("Policy engine not initialized, skipping policy resolution from metadata") return # Build context for policy resolution (model from request data) @@ -2454,9 +2229,7 @@ def _add_guardrails_from_policies_in_metadata( f"Policy engine: resolved guardrails from policy '{policy_name}': {resolved_policy.guardrails}" ) else: - verbose_proxy_logger.warning( - f"Policy engine: policy '{policy_name}' not found in registry" - ) + verbose_proxy_logger.warning(f"Policy engine: policy '{policy_name}' not found in registry") if not resolved_guardrails: return @@ -2502,23 +2275,13 @@ async def move_guardrails_to_metadata( team_metadata = user_api_key_dict.team_metadata project_metadata = user_api_key_dict.project_metadata or {} - has_key_config = key_metadata and ( - "guardrails" in key_metadata or "policies" in key_metadata - ) - has_team_config = team_metadata and ( - "guardrails" in team_metadata or "policies" in team_metadata - ) - has_project_config = project_metadata and ( - "guardrails" in project_metadata or "policies" in project_metadata - ) - has_request_config = ( - "guardrails" in data or "guardrail_config" in data or "policies" in data - ) + has_key_config = key_metadata and ("guardrails" in key_metadata or "policies" in key_metadata) + has_team_config = team_metadata and ("guardrails" in team_metadata or "policies" in team_metadata) + has_project_config = project_metadata and ("guardrails" in project_metadata or "policies" in project_metadata) + has_request_config = "guardrails" in data or "guardrail_config" in data or "policies" in data # Only check policy engine if no local config (avoid import + registry lookup) - if not ( - has_key_config or has_team_config or has_project_config or has_request_config - ): + if not (has_key_config or has_team_config or has_project_config or has_request_config): from litellm.proxy.policy_engine.policy_registry import get_policy_registry if not get_policy_registry().is_initialized(): @@ -2574,13 +2337,9 @@ async def move_guardrails_to_metadata( if "guardrail_config" in data[_metadata_variable_name] and isinstance( data[_metadata_variable_name]["guardrail_config"], dict ): - data[_metadata_variable_name]["guardrail_config"].update( - request_body_guardrail_config - ) + data[_metadata_variable_name]["guardrail_config"].update(request_body_guardrail_config) else: - data[_metadata_variable_name]["guardrail_config"] = ( - request_body_guardrail_config - ) + data[_metadata_variable_name]["guardrail_config"] = request_body_guardrail_config def _is_policy_version_id(s: str) -> bool: @@ -2621,23 +2380,17 @@ def _match_and_track_policies( # Get matching policies via attachments (with match reasons for attribution) attachment_registry = get_attachment_registry() - matches_with_reasons = attachment_registry.get_attached_policies_with_reasons( - context - ) + matches_with_reasons = attachment_registry.get_attached_policies_with_reasons(context) matching_policy_names = [m["policy_name"] for m in matches_with_reasons] policy_reasons = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} - verbose_proxy_logger.debug( - f"Policy engine: matched policies via attachments: {matching_policy_names}" - ) + verbose_proxy_logger.debug(f"Policy engine: matched policies via attachments: {matching_policy_names}") # Combine attachment-based policies with dynamic request body policies all_policy_names = set(matching_policy_names) if request_body_policies and isinstance(request_body_policies, list): all_policy_names.update(request_body_policies) - verbose_proxy_logger.debug( - f"Policy engine: added dynamic policies from request body: {request_body_policies}" - ) + verbose_proxy_logger.debug(f"Policy engine: added dynamic policies from request body: {request_body_policies}") if not all_policy_names: return [], {} @@ -2649,22 +2402,14 @@ def _match_and_track_policies( policies=policies_override, ) - verbose_proxy_logger.debug( - f"Policy engine: applied policies (conditions matched): {applied_policy_names}" - ) + verbose_proxy_logger.debug(f"Policy engine: applied policies (conditions matched): {applied_policy_names}") # Track applied policies in metadata for response headers for policy_name in applied_policy_names: - add_policy_to_applied_policies_header( - request_data=data, policy_name=policy_name - ) + add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) # Track policy attribution sources for x-litellm-policy-sources header - applied_reasons = { - name: policy_reasons[name] - for name in applied_policy_names - if name in policy_reasons - } + applied_reasons = {name: policy_reasons[name] for name in applied_policy_names if name in policy_reasons} add_policy_sources_to_metadata(request_data=data, policy_sources=applied_reasons) return applied_policy_names, policy_reasons @@ -2688,9 +2433,7 @@ def _apply_resolved_guardrails_to_metadata( policy_names=policy_names, ) - verbose_proxy_logger.debug( - f"Policy engine: resolved guardrails: {resolved_guardrails}" - ) + verbose_proxy_logger.debug(f"Policy engine: resolved guardrails: {resolved_guardrails}") # Resolve pipelines from matching policies pipelines = PolicyResolver.resolve_pipelines_for_context( @@ -2706,16 +2449,11 @@ def _apply_resolved_guardrails_to_metadata( # Track pipeline-managed guardrails to exclude from independent execution pipeline_managed_guardrails: set = set() if pipelines: - pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails( - pipelines - ) + pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails(pipelines) data[metadata_variable_name]["_guardrail_pipelines"] = pipelines - data[metadata_variable_name]["_pipeline_managed_guardrails"] = ( - pipeline_managed_guardrails - ) + data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails verbose_proxy_logger.debug( - f"Policy engine: resolved {len(pipelines)} pipeline(s), " - f"managed guardrails: {pipeline_managed_guardrails}" + f"Policy engine: resolved {len(pipelines)} pipeline(s), managed guardrails: {pipeline_managed_guardrails}" ) if not resolved_guardrails and not pipelines: @@ -2732,9 +2470,7 @@ def _apply_resolved_guardrails_to_metadata( combined -= pipeline_managed_guardrails data[metadata_variable_name]["guardrails"] = list(combined) - verbose_proxy_logger.debug( - f"Policy engine: added guardrails to request metadata: {list(combined)}" - ) + verbose_proxy_logger.debug(f"Policy engine: added guardrails to request metadata: {list(combined)}") async def add_guardrails_from_policy_engine( @@ -2774,9 +2510,7 @@ async def add_guardrails_from_policy_engine( f"policy_count={len(registry.get_all_policies())}" ) if not registry.is_initialized(): - verbose_proxy_logger.debug( - "Policy engine not initialized, skipping policy matching" - ) + verbose_proxy_logger.debug("Policy engine not initialized, skipping policy matching") return # Extract tags and build context @@ -2818,13 +2552,9 @@ async def add_guardrails_from_policy_engine( pname, policy = result merged_policies[pname] = policy fetched_policy_names.append(pname) - verbose_proxy_logger.debug( - f"Policy engine: loaded version by ID policy_{policy_id} -> {pname}" - ) + verbose_proxy_logger.debug(f"Policy engine: loaded version by ID policy_{policy_id} -> {pname}") else: - verbose_proxy_logger.debug( - f"Policy engine: policy version {policy_id} not found in cache, skipping" - ) + verbose_proxy_logger.debug(f"Policy engine: policy version {policy_id} not found in cache, skipping") # Build request body list: names + policy names from fetched versions request_body_policies = request_body_names + fetched_policy_names diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py index a96a5431294..89f70f9ea5f 100644 --- a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -34,9 +34,7 @@ from litellm.types.proxy.callback_logs_endpoints import ( # Routes the Python proxy exposes for the Rust data-plane gateway to call into # (logging today; auth/budgets later). Namespaced under /v1/rust_control_plane so # they're clearly distinct from the proxy's own control-plane/management routes. -rust_control_plane_router = APIRouter( - prefix="/v1/rust_control_plane", tags=["rust control plane"] -) +rust_control_plane_router = APIRouter(prefix="/v1/rust_control_plane", tags=["rust control plane"]) class CallbackLogsReplayer: @@ -71,9 +69,7 @@ class CallbackLogsReplayer: model = payload.get("model") or "" call_type = payload.get("call_type") or "acompletion" start_time = CallbackLogsReplayer._epoch_to_datetime(payload.get("startTime")) - call_id = ( - payload.get("litellm_call_id") or payload.get("id") or str(uuid.uuid4()) - ) + call_id = payload.get("litellm_call_id") or payload.get("id") or str(uuid.uuid4()) logging_obj = LiteLLMLogging( model=model, @@ -157,9 +153,7 @@ class CallbackLogsReplayer: end_time=end_time, ) - async def replay_batch( - self, records: list[CallbackLogRecord] - ) -> CallbackLogsResponse: + async def replay_batch(self, records: list[CallbackLogRecord]) -> CallbackLogsResponse: """Replay a batch; a single bad record never sinks the rest. Each failure is reported back with its batch index so the caller can retry/triage it.""" processed = 0 @@ -180,9 +174,7 @@ class CallbackLogsReplayer: processed, len(failures), ) - return CallbackLogsResponse( - processed=processed, failed=len(failures), failures=failures - ) + return CallbackLogsResponse(processed=processed, failed=len(failures), failures=failures) @rust_control_plane_router.post( diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 65f7ffc9081..6ae97bfce85 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -111,69 +111,47 @@ async def _invalidate_cache_access_group(access_group_id: str) -> None: # --------------------------------------------------------------------------- -async def _sync_add_access_group_to_teams( - tx, team_ids: List[str], access_group_id: str -) -> None: +async def _sync_add_access_group_to_teams(tx, team_ids: List[str], access_group_id: str) -> None: """Add access_group_id to each team's access_group_ids (idempotent).""" for team_id in team_ids: team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) if team is not None and access_group_id not in (team.access_group_ids or []): await tx.litellm_teamtable.update( where={"team_id": team_id}, - data={ - "access_group_ids": list(team.access_group_ids or []) - + [access_group_id] - }, + data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]}, ) -async def _sync_remove_access_group_from_teams( - tx, team_ids: List[str], access_group_id: str -) -> None: +async def _sync_remove_access_group_from_teams(tx, team_ids: List[str], access_group_id: str) -> None: """Remove access_group_id from each team's access_group_ids (idempotent).""" for team_id in team_ids: team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) if team is not None and access_group_id in (team.access_group_ids or []): await tx.litellm_teamtable.update( where={"team_id": team_id}, - data={ - "access_group_ids": [ - ag for ag in team.access_group_ids if ag != access_group_id - ] - }, + data={"access_group_ids": [ag for ag in team.access_group_ids if ag != access_group_id]}, ) -async def _sync_add_access_group_to_keys( - tx, key_tokens: List[str], access_group_id: str -) -> None: +async def _sync_add_access_group_to_keys(tx, key_tokens: List[str], access_group_id: str) -> None: """Add access_group_id to each key's access_group_ids (idempotent).""" for token in key_tokens: key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) if key is not None and access_group_id not in (key.access_group_ids or []): await tx.litellm_verificationtoken.update( where={"token": token}, - data={ - "access_group_ids": list(key.access_group_ids or []) - + [access_group_id] - }, + data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]}, ) -async def _sync_remove_access_group_from_keys( - tx, key_tokens: List[str], access_group_id: str -) -> None: +async def _sync_remove_access_group_from_keys(tx, key_tokens: List[str], access_group_id: str) -> None: """Remove access_group_id from each key's access_group_ids (idempotent).""" for token in key_tokens: key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) if key is not None and access_group_id in (key.access_group_ids or []): await tx.litellm_verificationtoken.update( where={"token": token}, - data={ - "access_group_ids": [ - ag for ag in key.access_group_ids if ag != access_group_id - ] - }, + data={"access_group_ids": [ag for ag in key.access_group_ids if ag != access_group_id]}, ) @@ -201,9 +179,7 @@ async def _patch_team_caches_add_access_group( if cached_team.access_group_ids is None: cached_team.access_group_ids = [access_group_id] elif access_group_id not in cached_team.access_group_ids: - cached_team.access_group_ids = list(cached_team.access_group_ids) + [ - access_group_id - ] + cached_team.access_group_ids = list(cached_team.access_group_ids) + [access_group_id] else: continue await _cache_team_object( @@ -229,9 +205,7 @@ async def _patch_team_caches_remove_access_group( parent_otel_span=None, ) if cached_team is not None and cached_team.access_group_ids: - cached_team.access_group_ids = [ - ag for ag in cached_team.access_group_ids if ag != access_group_id - ] + cached_team.access_group_ids = [ag for ag in cached_team.access_group_ids if ag != access_group_id] await _cache_team_object( team_id=team_id, team_table=cached_team, @@ -257,9 +231,7 @@ async def _patch_key_caches_add_access_group( if cached_key.access_group_ids is None: cached_key.access_group_ids = [access_group_id] elif access_group_id not in cached_key.access_group_ids: - cached_key.access_group_ids = list(cached_key.access_group_ids) + [ - access_group_id - ] + cached_key.access_group_ids = list(cached_key.access_group_ids) + [access_group_id] else: continue await _cache_key_object( @@ -283,9 +255,7 @@ async def _patch_key_caches_remove_access_group( model_type=UserAPIKeyAuth, ) if cached_key is not None and cached_key.access_group_ids: - cached_key.access_group_ids = [ - ag for ag in cached_key.access_group_ids if ag != access_group_id - ] + cached_key.access_group_ids = [ag for ag in cached_key.access_group_ids if ag != access_group_id] await _cache_key_object( hashed_token=token, user_api_key_obj=cached_key, @@ -309,9 +279,7 @@ async def create_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupResponse: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw( - CommonProxyErrors.db_not_connected_error.value - ) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) try: async with prisma_client.db.tx() as tx: @@ -339,12 +307,8 @@ async def create_access_group( ) # Sync team and key tables to reference the new access group - await _sync_add_access_group_to_teams( - tx, data.assigned_team_ids or [], record.access_group_id - ) - await _sync_add_access_group_to_keys( - tx, data.assigned_key_ids or [], record.access_group_id - ) + await _sync_add_access_group_to_teams(tx, data.assigned_team_ids or [], record.access_group_id) + await _sync_add_access_group_to_keys(tx, data.assigned_key_ids or [], record.access_group_id) except HTTPException: raise except Exception as e: @@ -383,13 +347,9 @@ async def list_access_groups( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> List[AccessGroupResponse]: _require_admin_view(user_api_key_dict) - prisma_client = get_prisma_client_or_throw( - CommonProxyErrors.db_not_connected_error.value - ) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - records = await AccessGroupRepository(prisma_client).table.find_many( - order={"created_at": "desc"} - ) + records = await AccessGroupRepository(prisma_client).table.find_many(order={"created_at": "desc"}) return [_record_to_response(r) for r in records] @@ -402,13 +362,9 @@ async def get_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupResponse: _require_admin_view(user_api_key_dict) - prisma_client = get_prisma_client_or_throw( - CommonProxyErrors.db_not_connected_error.value - ) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - record = await AccessGroupRepository(prisma_client).table.find_unique( - where={"access_group_id": access_group_id} - ) + record = await AccessGroupRepository(prisma_client).table.find_unique(where={"access_group_id": access_group_id}) if record is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -427,9 +383,7 @@ async def update_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupResponse: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw( - CommonProxyErrors.db_not_connected_error.value - ) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) update_fields = data.model_dump(exclude_unset=True) update_data: dict = {"updated_by": user_api_key_dict.user_id} @@ -459,9 +413,7 @@ async def update_access_group( async with prisma_client.db.tx() as tx: # Read inside the transaction so delta computation is consistent with the write, # avoiding a TOCTOU race where a concurrent update could make deltas stale. - existing = await tx.litellm_accessgrouptable.find_unique( - where={"access_group_id": access_group_id} - ) + existing = await tx.litellm_accessgrouptable.find_unique(where={"access_group_id": access_group_id}) if existing is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -471,14 +423,10 @@ async def update_access_group( old_team_ids: Set[str] = set(existing.assigned_team_ids or []) old_key_ids: Set[str] = set(existing.assigned_key_ids or []) new_team_ids: Set[str] = ( - set(update_fields["assigned_team_ids"] or []) - if "assigned_team_ids" in update_fields - else old_team_ids + set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids ) new_key_ids: Set[str] = ( - set(update_fields["assigned_key_ids"] or []) - if "assigned_key_ids" in update_fields - else old_key_ids + set(update_fields["assigned_key_ids"] or []) if "assigned_key_ids" in update_fields else old_key_ids ) teams_to_add = list(new_team_ids - old_team_ids) @@ -492,13 +440,9 @@ async def update_access_group( ) await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id) - await _sync_remove_access_group_from_teams( - tx, teams_to_remove, access_group_id - ) + await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id) await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id) - await _sync_remove_access_group_from_keys( - tx, keys_to_remove, access_group_id - ) + await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id) except HTTPException: raise except Exception as e: @@ -513,18 +457,12 @@ async def update_access_group( from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache await _cache_access_group_record(record) - await _patch_team_caches_add_access_group( - teams_to_add, access_group_id, user_api_key_cache, proxy_logging_obj - ) + await _patch_team_caches_add_access_group(teams_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) await _patch_team_caches_remove_access_group( teams_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj ) - await _patch_key_caches_add_access_group( - keys_to_add, access_group_id, user_api_key_cache, proxy_logging_obj - ) - await _patch_key_caches_remove_access_group( - keys_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj - ) + await _patch_key_caches_add_access_group(keys_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_key_caches_remove_access_group(keys_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj) return _record_to_response(record) @@ -538,18 +476,14 @@ async def delete_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> None: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw( - CommonProxyErrors.db_not_connected_error.value - ) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) try: affected_team_ids: List[str] = [] affected_key_tokens: List[str] = [] async with prisma_client.db.tx() as tx: - existing = await tx.litellm_accessgrouptable.find_unique( - where={"access_group_id": access_group_id} - ) + existing = await tx.litellm_accessgrouptable.find_unique(where={"access_group_id": access_group_id}) if existing is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -561,9 +495,9 @@ async def delete_access_group( teams_with_group = await tx.litellm_teamtable.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) - all_affected_team_ids: Set[str] = { - team.team_id for team in teams_with_group - } | set(existing.assigned_team_ids or []) + all_affected_team_ids: Set[str] = {team.team_id for team in teams_with_group} | set( + existing.assigned_team_ids or [] + ) affected_team_ids = list(all_affected_team_ids) # Union of: keys that have this access_group_id in their own access_group_ids @@ -571,54 +505,32 @@ async def delete_access_group( keys_with_group = await tx.litellm_verificationtoken.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) - all_affected_key_tokens: Set[str] = { - key.token for key in keys_with_group - } | set(existing.assigned_key_ids or []) + all_affected_key_tokens: Set[str] = {key.token for key in keys_with_group} | set( + existing.assigned_key_ids or [] + ) affected_key_tokens = list(all_affected_key_tokens) # Update teams returned by find_many directly — we already have their data. for team in teams_with_group: await tx.litellm_teamtable.update( where={"team_id": team.team_id}, - data={ - "access_group_ids": [ - ag - for ag in (team.access_group_ids or []) - if ag != access_group_id - ] - }, + data={"access_group_ids": [ag for ag in (team.access_group_ids or []) if ag != access_group_id]}, ) # Use _sync_remove only for out-of-sync teams not found by the hasSome query. - out_of_sync_team_ids = set(existing.assigned_team_ids or []) - { - t.team_id for t in teams_with_group - } - await _sync_remove_access_group_from_teams( - tx, list(out_of_sync_team_ids), access_group_id - ) + out_of_sync_team_ids = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group} + await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id) # Update keys returned by find_many directly — we already have their data. for key in keys_with_group: await tx.litellm_verificationtoken.update( where={"token": key.token}, - data={ - "access_group_ids": [ - ag - for ag in (key.access_group_ids or []) - if ag != access_group_id - ] - }, + data={"access_group_ids": [ag for ag in (key.access_group_ids or []) if ag != access_group_id]}, ) # Use _sync_remove only for out-of-sync keys not found by the hasSome query. - out_of_sync_key_tokens = set(existing.assigned_key_ids or []) - { - k.token for k in keys_with_group - } - await _sync_remove_access_group_from_keys( - tx, list(out_of_sync_key_tokens), access_group_id - ) + out_of_sync_key_tokens = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group} + await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id) - await tx.litellm_accessgrouptable.delete( - where={"access_group_id": access_group_id} - ) + await tx.litellm_accessgrouptable.delete(where={"access_group_id": access_group_id}) from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache @@ -643,9 +555,7 @@ async def delete_access_group( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=CommonProxyErrors.db_not_connected_error.value, ) - if "P2025" in str(e) or ( - "record" in str(e).lower() and "not found" in str(e).lower() - ): + if "P2025" in str(e) or ("record" in str(e).lower() and "not found" in str(e).lower()): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Access group '{access_group_id}' not found", diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index e35ec2933d0..6b70a9064df 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -60,23 +60,15 @@ async def new_budget( ) # Validate budget values are not negative - if budget_obj.max_budget is not None and ( - not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0 - ): + if budget_obj.max_budget is not None and (not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}"}, ) - if budget_obj.soft_budget is not None and ( - not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0 - ): + if budget_obj.soft_budget is not None and (not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}" - }, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, ) # Validate model_max_budget if present @@ -92,9 +84,7 @@ async def new_budget( # if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None: - budget_obj.budget_reset_at = get_budget_reset_time( - budget_duration=budget_obj.budget_duration - ) + budget_obj.budget_reset_at = get_budget_reset_time(budget_duration=budget_obj.budget_duration) budget_obj_json = budget_obj.model_dump(exclude_none=True) budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries @@ -111,9 +101,7 @@ async def new_budget( raise raise HTTPException( status_code=400, - detail={ - "error": f"Budget with id '{budget_obj.budget_id}' already exists." - }, + detail={"error": f"Budget with id '{budget_obj.budget_id}' already exists."}, ) return response @@ -153,23 +141,15 @@ async def update_budget( raise HTTPException(status_code=400, detail={"error": "budget_id is required"}) # Validate budget values are not negative - if budget_obj.max_budget is not None and ( - not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0 - ): + if budget_obj.max_budget is not None and (not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}"}, ) - if budget_obj.soft_budget is not None and ( - not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0 - ): + if budget_obj.soft_budget is not None and (not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}" - }, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, ) # Validate model_max_budget if present in update @@ -185,13 +165,8 @@ async def update_budget( # recompute budget_reset_at when the duration changes, unless the caller pinned a reset time explicitly recomputed_reset_at = ( - { - "budget_reset_at": get_budget_reset_time( - budget_duration=budget_obj.budget_duration - ) - } - if budget_obj.budget_duration is not None - and "budget_reset_at" not in budget_obj.model_fields_set + {"budget_reset_at": get_budget_reset_time(budget_duration=budget_obj.budget_duration)} + if budget_obj.budget_duration is not None and "budget_reset_at" not in budget_obj.model_fields_set else {} ) @@ -227,9 +202,7 @@ async def info_budget(data: BudgetRequest): if len(data.budgets) == 0: raise HTTPException( status_code=400, - detail={ - "error": f"Specify list of budget id's to query. Passed in={data.budgets}" - }, + detail={"error": f"Specify list of budget id's to query. Passed in={data.budgets}"}, ) response = await BudgetRepository(prisma_client).table.find_many( where={"budget_id": {"in": data.budgets}}, @@ -275,9 +248,7 @@ async def budget_settings( ) ## get budget item from db - db_budget_row = await BudgetRepository(prisma_client).table.find_first( - where={"budget_id": budget_id} - ) + db_budget_row = await BudgetRepository(prisma_client).table.find_first(where={"budget_id": budget_id}) if db_budget_row is not None: db_budget_row_dict = db_budget_row.model_dump(exclude_none=True) @@ -380,8 +351,6 @@ async def delete_budget( }, ) - response = await BudgetRepository(prisma_client).table.delete( - where={"budget_id": data.id} - ) + response = await BudgetRepository(prisma_client).table.delete(where={"budget_id": data.id}) return response diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index b6ddf2d8e07..6aa3dbbf902 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -68,9 +68,7 @@ def _log_audit_task_exception(task: "asyncio.Task[None]") -> None: return exc = task.exception() if exc is not None: - verbose_proxy_logger.warning( - "Failed to write cache-settings audit log: %s", exc - ) + verbose_proxy_logger.warning("Failed to write cache-settings audit log: %s", exc) async def _emit_cache_settings_audit_log( @@ -102,19 +100,13 @@ async def _emit_cache_settings_audit_log( request_data=LiteLLM_AuditLogs( id=str(uuid.uuid4()), updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by - or user_api_key_dict.user_id - or litellm_proxy_admin_name, + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, changed_by_api_key=user_api_key_dict.api_key, table_name=LitellmTableNames.CACHE_CONFIG_TABLE_NAME, object_id="cache_config", action=action, - updated_values=json.dumps( - {"settings": _redact_settings(after_settings)}, default=str - ), - before_value=json.dumps( - {"settings": _redact_settings(before_settings)}, default=str - ), + updated_values=json.dumps({"settings": _redact_settings(after_settings)}, default=str), + before_value=json.dumps({"settings": _redact_settings(before_settings)}, default=str), ) ) ) @@ -163,9 +155,7 @@ class CacheSettingsManager: try: cache_config = await call_with_db_reconnect_retry( prisma_client, - lambda: CacheConfigRepository(prisma_client).table.find_unique( - where={"id": "cache_config"} - ), + lambda: CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}), reason="init_cache_settings_in_db_lookup_failure", ) if cache_config is not None and cache_config.cache_settings: @@ -177,26 +167,17 @@ class CacheSettingsManager: cache_settings_dict = cache_settings_json # Decrypt cache settings - decrypted_settings = proxy_config._decrypt_db_variables( - variables_dict=cache_settings_dict - ) + decrypted_settings = proxy_config._decrypt_db_variables(variables_dict=cache_settings_dict) # Remove redis_type if present (UI-only field, not a Cache parameter) # We derive it for UI in get_cache_settings endpoint - cache_params = { - k: v for k, v in decrypted_settings.items() if k != "redis_type" - } + cache_params = {k: v for k, v in decrypted_settings.items() if k != "redis_type"} # Check if cache params have changed - if ( - CacheSettingsManager._last_cache_params is not None - and CacheSettingsManager._cache_params_equal( - CacheSettingsManager._last_cache_params, cache_params - ) + if CacheSettingsManager._last_cache_params is not None and CacheSettingsManager._cache_params_equal( + CacheSettingsManager._last_cache_params, cache_params ): - verbose_proxy_logger.debug( - "Cache settings unchanged, skipping reinitialization" - ) + verbose_proxy_logger.debug("Cache settings unchanged, skipping reinitialization") return # Initialize cache only if params changed or cache not initialized @@ -226,29 +207,19 @@ class CacheSettingsManager: class CacheSettingsResponse(BaseModel): - fields: List[CacheSettingsField] = Field( - description="List of all configurable cache settings with metadata" - ) - current_values: Dict[str, Any] = Field( - description="Current values of cache settings" - ) - redis_type_descriptions: Dict[str, str] = Field( - description="Descriptions for each Redis type option" - ) + fields: List[CacheSettingsField] = Field(description="List of all configurable cache settings with metadata") + current_values: Dict[str, Any] = Field(description="Current values of cache settings") + redis_type_descriptions: Dict[str, str] = Field(description="Descriptions for each Redis type option") class CacheTestRequest(BaseModel): - cache_settings: Dict[str, Any] = Field( - description="Cache settings to test connection with" - ) + cache_settings: Dict[str, Any] = Field(description="Cache settings to test connection with") class CacheTestResponse(BaseModel): status: str = Field(description="Connection status: 'success' or 'failed'") message: str = Field(description="Connection result message") - error: Optional[str] = Field( - default=None, description="Error message if connection failed" - ) + error: Optional[str] = Field(default=None, description="Error message if connection failed") class CacheSettingsUpdateRequest(BaseModel): @@ -280,9 +251,7 @@ async def get_cache_settings( # Try to get cache settings from database current_values = {} if prisma_client is not None: - cache_config = await CacheConfigRepository(prisma_client).table.find_unique( - where={"id": "cache_config"} - ) + cache_config = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) if cache_config is not None and cache_config.cache_settings: # Decrypt cache settings cache_settings_json = cache_config.cache_settings @@ -292,9 +261,7 @@ async def get_cache_settings( cache_settings_dict = cache_settings_json # Decrypt environment variables - decrypted_settings = proxy_config._decrypt_db_variables( - variables_dict=cache_settings_dict - ) + decrypted_settings = proxy_config._decrypt_db_variables(variables_dict=cache_settings_dict) # Derive redis_type for UI based on settings # UI uses redis_type to show/hide fields, backend only stores 'type' @@ -308,9 +275,7 @@ async def get_cache_settings( # Mask credential fields so the GET response never carries # plaintext Redis / Sentinel passwords off the server. - current_values = mask_sensitive_keys( - decrypted_settings, _CACHE_SENSITIVE_FIELDS - ) + current_values = mask_sensitive_keys(decrypted_settings, _CACHE_SENSITIVE_FIELDS) # Update field values with current values for field in cache_fields: @@ -324,9 +289,7 @@ async def get_cache_settings( ) except Exception as e: verbose_proxy_logger.error(f"Error fetching cache settings: {str(e)}") - raise HTTPException( - status_code=500, detail=f"Error fetching cache settings: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Error fetching cache settings: {str(e)}") @router.post( @@ -349,9 +312,7 @@ async def test_cache_connection( try: cache_settings = request.cache_settings.copy() - verbose_proxy_logger.debug( - "Testing cache connection with settings: %s", cache_settings - ) + verbose_proxy_logger.debug("Testing cache connection with settings: %s", cache_settings) # Only support Redis for now if cache_settings.get("type") != "redis": @@ -413,9 +374,7 @@ async def update_cache_settings( if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) try: @@ -423,9 +382,7 @@ async def update_cache_settings( # Snapshot the prior settings (key set only — values get redacted in # the audit row) so the audit-log entry shows which fields changed. - existing_row = await CacheConfigRepository(prisma_client).table.find_unique( - where={"id": "cache_config"} - ) + existing_row = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) before_settings: Optional[Dict[str, Any]] = None if existing_row is not None and existing_row.cache_settings: try: @@ -435,9 +392,7 @@ async def update_cache_settings( action: AUDIT_ACTIONS = "updated" if existing_row is not None else "created" # Encrypt sensitive fields (keep redis_type for storage) - encrypted_settings = proxy_config._encrypt_env_variables( - environment_variables=cache_settings - ) + encrypted_settings = proxy_config._encrypt_env_variables(environment_variables=cache_settings) # Save to database await CacheConfigRepository(prisma_client).table.upsert( @@ -455,14 +410,10 @@ async def update_cache_settings( # Reinitialize cache with new settings # Decrypt for initialization - decrypted_settings = proxy_config._decrypt_db_variables( - variables_dict=encrypted_settings - ) + decrypted_settings = proxy_config._decrypt_db_variables(variables_dict=encrypted_settings) # Remove redis_type if present (UI-only field, not a Cache parameter) - cache_params = { - k: v for k, v in decrypted_settings.items() if k != "redis_type" - } + cache_params = {k: v for k, v in decrypted_settings.items() if k != "redis_type"} # Initialize cache (frontend sends type="redis", not redis_type) proxy_config._init_cache(cache_params=cache_params) @@ -492,6 +443,4 @@ async def update_cache_settings( } except Exception as e: verbose_proxy_logger.error(f"Error updating cache settings: {str(e)}") - raise HTTPException( - status_code=500, detail=f"Error updating cache settings: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Error updating cache settings: {str(e)}") diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index ba41852f6d8..60cb3ccd30d 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -48,9 +48,7 @@ def update_metrics(existing_metrics: SpendMetrics, record: Any) -> SpendMetrics: existing_metrics.completion_tokens += completion_tokens existing_metrics.total_tokens += prompt_tokens + completion_tokens existing_metrics.cache_read_input_tokens += record.cache_read_input_tokens or 0 - existing_metrics.cache_creation_input_tokens += ( - record.cache_creation_input_tokens or 0 - ) + existing_metrics.cache_creation_input_tokens += record.cache_creation_input_tokens or 0 existing_metrics.api_requests += record.api_requests or 0 existing_metrics.successful_requests += record.successful_requests or 0 existing_metrics.failed_requests += record.failed_requests or 0 @@ -62,9 +60,7 @@ def _is_user_agent_tag(tag: Optional[str]) -> bool: if not tag: return False normalized_tag = tag.strip().lower() - return normalized_tag.startswith("user-agent:") or normalized_tag.startswith( - "user agent:" - ) + return normalized_tag.startswith("user-agent:") or normalized_tag.startswith("user agent:") def compute_tag_metadata_totals(records: List[Any]) -> SpendMetrics: @@ -108,33 +104,21 @@ def update_breakdown_metrics( if record.model and record.model not in breakdown.models: breakdown.models[record.model] = MetricWithMetadata( metrics=SpendMetrics(), - metadata=model_metadata.get( - record.model, {} - ), # Add any model-specific metadata here + metadata=model_metadata.get(record.model, {}), # Add any model-specific metadata here ) if record.model: - breakdown.models[record.model].metrics = update_metrics( - breakdown.models[record.model].metrics, record - ) + breakdown.models[record.model].metrics = update_metrics(breakdown.models[record.model].metrics, record) # Update API key breakdown for this model if record.api_key not in breakdown.models[record.model].api_key_breakdown: - breakdown.models[record.model].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), - ), - ) + breakdown.models[record.model].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), ) - breakdown.models[record.model].api_key_breakdown[ - record.api_key - ].metrics = update_metrics( + breakdown.models[record.model].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.models[record.model].api_key_breakdown[record.api_key].metrics, record, ) @@ -151,29 +135,16 @@ def update_breakdown_metrics( ) # Update API key breakdown for this model - if ( - record.api_key - not in breakdown.model_groups[record.model_group].api_key_breakdown - ): - breakdown.model_groups[record.model_group].api_key_breakdown[ - record.api_key - ] = KeyMetricWithMetadata( + if record.api_key not in breakdown.model_groups[record.model_group].api_key_breakdown: + breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), ), ) - breakdown.model_groups[record.model_group].api_key_breakdown[ - record.api_key - ].metrics = update_metrics( - breakdown.model_groups[record.model_group] - .api_key_breakdown[record.api_key] - .metrics, + breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key].metrics = update_metrics( + breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key].metrics, record, ) @@ -188,32 +159,21 @@ def update_breakdown_metrics( ) # Update API key breakdown for this MCP server - if ( - record.api_key - not in breakdown.mcp_servers[ - record.mcp_namespaced_tool_name - ].api_key_breakdown - ): - breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[ - record.api_key - ] = KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None + if record.api_key not in breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown: + breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = ( + KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), - ), + ) ) breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[ record.api_key ].metrics = update_metrics( - breakdown.mcp_servers[record.mcp_namespaced_tool_name] - .api_key_breakdown[record.api_key] - .metrics, + breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key].metrics, record, ) @@ -222,32 +182,20 @@ def update_breakdown_metrics( if provider not in breakdown.providers: breakdown.providers[provider] = MetricWithMetadata( metrics=SpendMetrics(), - metadata=provider_metadata.get( - provider, {} - ), # Add any provider-specific metadata here + metadata=provider_metadata.get(provider, {}), # Add any provider-specific metadata here ) - breakdown.providers[provider].metrics = update_metrics( - breakdown.providers[provider].metrics, record - ) + breakdown.providers[provider].metrics = update_metrics(breakdown.providers[provider].metrics, record) # Update API key breakdown for this provider if record.api_key not in breakdown.providers[provider].api_key_breakdown: - breakdown.providers[provider].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), - ), - ) + breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), ) - breakdown.providers[provider].api_key_breakdown[ - record.api_key - ].metrics = update_metrics( + breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, record, ) @@ -265,25 +213,15 @@ def update_breakdown_metrics( # Update API key breakdown for this endpoint if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown: - breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), - ), - ) + breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), ) - breakdown.endpoints[record.endpoint].api_key_breakdown[ - record.api_key - ].metrics = update_metrics( - breakdown.endpoints[record.endpoint] - .api_key_breakdown[record.api_key] - .metrics, + breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics( + breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics, record, ) @@ -292,53 +230,33 @@ def update_breakdown_metrics( breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), ), # Add any api_key-specific metadata here ) - breakdown.api_keys[record.api_key].metrics = update_metrics( - breakdown.api_keys[record.api_key].metrics, record - ) + breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record) # Update entity-specific metrics if entity_id_field is provided if entity_id_field: entity_value = getattr(record, entity_id_field, None) - entity_value = ( - entity_value if entity_value else "Unassigned" - ) # allow for null entity_id_field + entity_value = entity_value if entity_value else "Unassigned" # allow for null entity_id_field if entity_value not in breakdown.entities: breakdown.entities[entity_value] = MetricWithMetadata( metrics=SpendMetrics(), - metadata=( - entity_metadata_field.get(entity_value, {}) - if entity_metadata_field - else {} - ), + metadata=(entity_metadata_field.get(entity_value, {}) if entity_metadata_field else {}), ) - breakdown.entities[entity_value].metrics = update_metrics( - breakdown.entities[entity_value].metrics, record - ) + breakdown.entities[entity_value].metrics = update_metrics(breakdown.entities[entity_value].metrics, record) # Update API key breakdown for this entity if record.api_key not in breakdown.entities[entity_value].api_key_breakdown: - breakdown.entities[entity_value].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), - ), - ) + breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), ) - breakdown.entities[entity_value].api_key_breakdown[ - record.api_key - ].metrics = update_metrics( + breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, record, ) @@ -358,17 +276,13 @@ async def get_api_key_metadata( key_records = await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(api_keys)}} ) - result = { - k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records - } + result = {k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records} # For any keys not found in the active table, check the deleted keys table missing_keys = api_keys - set(result.keys()) if missing_keys: try: - deleted_key_records = await DeletedVerificationTokenRepository( - prisma_client - ).table.find_many( + deleted_key_records = await DeletedVerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(missing_keys)}}, order={"deleted_at": "desc"}, ) @@ -428,9 +342,7 @@ def _build_where_conditions( ) -> Dict[str, Any]: """Build prisma where clause for daily activity queries.""" # Adjust dates for timezone if provided - adjusted_start, adjusted_end = _adjust_dates_for_timezone( - start_date, end_date, timezone_offset_minutes - ) + adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) where_conditions: Dict[str, Any] = { "date": { @@ -489,9 +401,7 @@ def _build_aggregated_sql_query( if pg_table is None: raise ValueError(f"Unknown table name: {table_name}") - adjusted_start, adjusted_end = _adjust_dates_for_timezone( - start_date, end_date, timezone_offset_minutes - ) + adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) sql_conditions: List[str] = [] sql_params: List[Any] = [] @@ -610,9 +520,7 @@ def _aggregate_spend_records_sync( "breakdown": BreakdownMetrics(), } - grouped_data[date_str]["metrics"] = update_metrics( - grouped_data[date_str]["metrics"], record - ) + grouped_data[date_str]["metrics"] = update_metrics(grouped_data[date_str]["metrics"], record) grouped_data[date_str]["breakdown"] = update_breakdown_metrics( grouped_data[date_str]["breakdown"], @@ -710,9 +618,7 @@ def _record_to_spend_metrics(record: Any) -> SpendMetrics: ) -def _key_metadata( - api_key_metadata: Dict[str, Dict[str, Any]], api_key: str -) -> KeyMetadata: +def _key_metadata(api_key_metadata: Dict[str, Dict[str, Any]], api_key: str) -> KeyMetadata: meta = api_key_metadata.get(api_key, {}) return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id")) @@ -739,9 +645,7 @@ def _aggregate_grouping_sets_records_sync( grouped_data[date_str] = bucket return bucket - def assign_metric_with_metadata( - target: Dict[str, MetricWithMetadata], key: str, metrics: SpendMetrics - ) -> None: + def assign_metric_with_metadata(target: Dict[str, MetricWithMetadata], key: str, metrics: SpendMetrics) -> None: existing = target.get(key) if existing is None: target[key] = MetricWithMetadata(metrics=metrics, metadata={}) @@ -787,14 +691,10 @@ def _aggregate_grouping_sets_records_sync( assign_metric_with_metadata(breakdown.models, record.model, metrics) elif level == _GROUP_DATE_MODEL_API_KEY: if record.model and record.api_key: - assign_api_key_breakdown( - breakdown.models, record.model, record.api_key, metrics - ) + assign_api_key_breakdown(breakdown.models, record.model, record.api_key, metrics) elif level == _GROUP_DATE_MODEL_GROUP: if record.model_group: - assign_metric_with_metadata( - breakdown.model_groups, record.model_group, metrics - ) + assign_metric_with_metadata(breakdown.model_groups, record.model_group, metrics) elif level == _GROUP_DATE_MODEL_GROUP_API_KEY: if record.model_group and record.api_key: assign_api_key_breakdown( @@ -809,14 +709,10 @@ def _aggregate_grouping_sets_records_sync( elif level == _GROUP_DATE_PROVIDER_API_KEY: if record.api_key: provider = record.custom_llm_provider or "unknown" - assign_api_key_breakdown( - breakdown.providers, provider, record.api_key, metrics - ) + assign_api_key_breakdown(breakdown.providers, provider, record.api_key, metrics) elif level == _GROUP_DATE_MCP: if record.mcp_namespaced_tool_name: - assign_metric_with_metadata( - breakdown.mcp_servers, record.mcp_namespaced_tool_name, metrics - ) + assign_metric_with_metadata(breakdown.mcp_servers, record.mcp_namespaced_tool_name, metrics) elif level == _GROUP_DATE_MCP_API_KEY: if record.mcp_namespaced_tool_name and record.api_key: assign_api_key_breakdown( @@ -827,14 +723,10 @@ def _aggregate_grouping_sets_records_sync( ) elif level == _GROUP_DATE_ENDPOINT: if record.endpoint: - assign_metric_with_metadata( - breakdown.endpoints, record.endpoint, metrics - ) + assign_metric_with_metadata(breakdown.endpoints, record.endpoint, metrics) elif level == _GROUP_DATE_ENDPOINT_API_KEY: if record.endpoint and record.api_key: - assign_api_key_breakdown( - breakdown.endpoints, record.endpoint, record.api_key, metrics - ) + assign_api_key_breakdown(breakdown.endpoints, record.endpoint, record.api_key, metrics) results = [ DailySpendData( @@ -883,9 +775,7 @@ async def get_daily_activity( exclude_entity_ids: Optional[List[str]] = None, metadata_metrics_func: Optional[Callable[[List[Any]], SpendMetrics]] = None, timezone_offset_minutes: Optional[int] = None, - resolve_entity_metadata: Optional[ - Callable[[list[Any]], Awaitable[dict[str, dict]]] - ] = None, + resolve_entity_metadata: Optional[Callable[[list[Any]], Awaitable[dict[str, dict]]]] = None, ) -> SpendAnalyticsPaginatedResponse: """Common function to get daily activity for any entity type. @@ -920,9 +810,7 @@ async def get_daily_activity( ) # Get total count for pagination - total_count = await getattr(prisma_client.db, table_name).count( - where=where_conditions - ) + total_count = await getattr(prisma_client.db, table_name).count(where=where_conditions) # Fetch paginated results. # ``date`` alone is not a unique sort key -- a busy tenant has many @@ -1058,12 +946,8 @@ async def get_daily_activity_aggregated( total_api_requests=aggregated["totals"].api_requests, total_successful_requests=aggregated["totals"].successful_requests, total_failed_requests=aggregated["totals"].failed_requests, - total_cache_read_input_tokens=aggregated[ - "totals" - ].cache_read_input_tokens, - total_cache_creation_input_tokens=aggregated[ - "totals" - ].cache_creation_input_tokens, + total_cache_read_input_tokens=aggregated["totals"].cache_read_input_tokens, + total_cache_creation_input_tokens=aggregated["totals"].cache_creation_input_tokens, page=1, total_pages=1, has_more=False, @@ -1071,9 +955,7 @@ async def get_daily_activity_aggregated( ) except Exception as e: - verbose_proxy_logger.exception( - f"Error fetching aggregated daily activity: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error fetching aggregated daily activity: {str(e)}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Failed to fetch analytics: {str(e)}"}, diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index bc2da33672f..8162babef40 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -69,8 +69,7 @@ def require_caller_user_id_for_non_admin( status_code=status.HTTP_403_FORBIDDEN, detail={ "error": ( - "Service-account keys cannot query user analytics. " - "Use a user-bound key, or call as a proxy admin." + "Service-account keys cannot query user analytics. Use a user-bound key, or call as a proxy admin." ) }, ) @@ -94,35 +93,25 @@ def _check_passthrough_routes_caller_permission( if getattr(data, "allowed_passthrough_routes", None): raise HTTPException( status_code=403, - detail={ - "error": f"Only proxy admins can set `allowed_passthrough_routes` on a {entity}." - }, + detail={"error": f"Only proxy admins can set `allowed_passthrough_routes` on a {entity}."}, ) metadata = getattr(data, "metadata", None) if isinstance(metadata, dict) and metadata.get("allowed_passthrough_routes"): raise HTTPException( status_code=403, - detail={ - "error": f"Only proxy admins can set `metadata.allowed_passthrough_routes` on a {entity}." - }, + detail={"error": f"Only proxy admins can set `metadata.allowed_passthrough_routes` on a {entity}."}, ) -def _is_user_team_admin( - user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable -) -> bool: +def _is_user_team_admin(user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable) -> bool: for member in team_obj.members_with_roles: - if ( - member.user_id is not None and member.user_id == user_api_key_dict.user_id - ) and member.role == "admin": + if (member.user_id is not None and member.user_id == user_api_key_dict.user_id) and member.role == "admin": return True return False -async def _is_user_org_admin_for_team( - user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable -) -> bool: +async def _is_user_org_admin_for_team(user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable) -> bool: """ Check if user is an org admin for the team's organization. @@ -151,10 +140,7 @@ async def _is_user_org_admin_for_team( return False for m in caller_user.organization_memberships or []: - if ( - m.organization_id == team_obj.organization_id - and m.user_role == LitellmUserRoles.ORG_ADMIN.value - ): + if m.organization_id == team_obj.organization_id and m.user_role == LitellmUserRoles.ORG_ADMIN.value: return True return False @@ -227,22 +213,16 @@ async def _user_has_admin_privileges( # Check if user is team admin for any team if user_obj.teams is not None and len(user_obj.teams) > 0: # Get all teams user is in - teams = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": user_obj.teams}} - ) + teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_obj.teams}}) for team in teams: team_obj = LiteLLM_TeamTable(**team.model_dump()) - if _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): return True except Exception as e: # If there's an error checking, default to False for security - verbose_proxy_logger.debug( - f"Error checking admin privileges for user {user_api_key_dict.user_id}: {e}" - ) + verbose_proxy_logger.debug(f"Error checking admin privileges for user {user_api_key_dict.user_id}: {e}") return False return False @@ -274,9 +254,7 @@ def _org_admin_can_invite_user( return False if target_user_obj.organization_memberships is None: return False - target_org_ids = { - m.organization_id for m in target_user_obj.organization_memberships - } + target_org_ids = {m.organization_id for m in target_user_obj.organization_memberships} return bool(admin_org_ids & target_org_ids) @@ -304,9 +282,7 @@ async def _team_admin_can_invite_user( if not target_user_obj.teams or len(target_user_obj.teams) == 0: return False - teams = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": admin_user_obj.teams}} - ) + teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": admin_user_obj.teams}}) admin_team_ids = [ team.team_id for team in teams @@ -390,9 +366,7 @@ async def admin_can_invite_user( return False except Exception as e: - verbose_proxy_logger.debug( - f"Error checking invite permission for user {user_api_key_dict.user_id}: {e}" - ) + verbose_proxy_logger.debug(f"Error checking invite permission for user {user_api_key_dict.user_id}: {e}") return False @@ -446,10 +420,7 @@ def _is_set_budget_value(value: Any) -> bool: def _has_meaningful_budget_limit(budget_values: Dict[str, Any]) -> bool: """A budget is meaningful if at least one limit is actually set; an empty list (no model restriction) and None both count as unset.""" - return any( - _is_set_budget_value(budget_values.get(field)) - for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS - ) + return any(_is_set_budget_value(budget_values.get(field)) for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS) async def _upsert_budget_and_membership( @@ -483,9 +454,7 @@ async def _upsert_budget_and_membership( if "budget_duration" in write_data: duration = write_data["budget_duration"] write_data["budget_reset_at"] = ( - get_budget_reset_time(budget_duration=duration) - if duration is not None - else None + get_budget_reset_time(budget_duration=duration) if duration is not None else None ) is_shared_default = ( @@ -501,9 +470,7 @@ async def _upsert_budget_and_membership( ) if existing_budget_id is not None and not is_shared_default: - existing_budget = await tx.litellm_budgettable.find_unique( - where={"budget_id": existing_budget_id} - ) + existing_budget = await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) merged = existing_budget.model_dump() if existing_budget is not None else {} merged.update(write_data) if not _has_meaningful_budget_limit(merged): @@ -521,9 +488,7 @@ async def _upsert_budget_and_membership( } if is_shared_default: - default_budget_row = await tx.litellm_budgettable.find_unique( - where={"budget_id": existing_budget_id} - ) + default_budget_row = await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) if default_budget_row is not None: default_budget_dict = default_budget_row.model_dump() for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: @@ -534,9 +499,7 @@ async def _upsert_budget_and_membership( create_data.update(write_data) if create_data.get("budget_duration") is not None: - create_data["budget_reset_at"] = get_budget_reset_time( - budget_duration=create_data["budget_duration"] - ) + create_data["budget_reset_at"] = get_budget_reset_time(budget_duration=create_data["budget_duration"]) else: create_data.pop("budget_reset_at", None) diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 97cb5eeddc4..d13644c8ae5 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -61,9 +61,7 @@ def _log_audit_task_exception(task: "asyncio.Task[None]") -> None: return exc = task.exception() if exc is not None: - verbose_proxy_logger.warning( - "Failed to write hashicorp-vault config audit log: %s", exc - ) + verbose_proxy_logger.warning("Failed to write hashicorp-vault config audit log: %s", exc) async def _emit_hashicorp_vault_audit_log( @@ -96,19 +94,13 @@ async def _emit_hashicorp_vault_audit_log( request_data=LiteLLM_AuditLogs( id=str(uuid.uuid4()), updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by - or user_api_key_dict.user_id - or litellm_proxy_admin_name, + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, changed_by_api_key=user_api_key_dict.api_key, table_name=LitellmTableNames.CONFIG_OVERRIDES_TABLE_NAME, object_id="hashicorp_vault", action=action, - updated_values=json.dumps( - {"config": _redact_config(after_config)}, default=str - ), - before_value=json.dumps( - {"config": _redact_config(before_config)}, default=str - ), + updated_values=json.dumps({"config": _redact_config(after_config)}, default=str), + before_value=json.dumps({"config": _redact_config(before_config)}, default=str), ) ) ) @@ -143,9 +135,7 @@ _sensitive_masker = SensitiveDataMasker() # --- Shared helpers --- -def _mask_sensitive_fields( - data: Dict[str, Any], sensitive_fields: Set[str] -) -> Dict[str, Any]: +def _mask_sensitive_fields(data: Dict[str, Any], sensitive_fields: Set[str]) -> Dict[str, Any]: """Mask sensitive fields for API responses. Non-sensitive fields are left as-is.""" masked = {} for key, value in data.items(): @@ -280,12 +270,8 @@ async def update_hashicorp_vault_config( # Validate that the config has enough fields to initialize has_vault_addr = bool(config_data.get("vault_addr")) has_token_auth = bool(config_data.get("vault_token")) - has_approle_auth = bool( - config_data.get("approle_role_id") and config_data.get("approle_secret_id") - ) - has_tls_cert_auth = bool( - config_data.get("client_cert") and config_data.get("client_key") - ) + has_approle_auth = bool(config_data.get("approle_role_id") and config_data.get("approle_secret_id")) + has_tls_cert_auth = bool(config_data.get("client_cert") and config_data.get("client_key")) if not has_vault_addr: raise HTTPException( @@ -311,9 +297,7 @@ async def update_hashicorp_vault_config( proxy_config.initialize_secret_manager(key_management_system="hashicorp_vault") except Exception as e: _set_env_vars(previous_env) - verbose_proxy_logger.exception( - "Error reinitializing Hashicorp Vault secret manager: %s", str(e) - ) + verbose_proxy_logger.exception("Error reinitializing Hashicorp Vault secret manager: %s", str(e)) raise HTTPException( status_code=500, detail=f"Failed to initialize secret manager: {e}", @@ -455,23 +439,17 @@ async def delete_hashicorp_vault_config( before_config: Optional[Dict[str, Any]] = None if existing_record is not None and existing_record.config_value is not None: try: - before_config = proxy_config._decrypt_db_variables( - _parse_config_value(existing_record.config_value) - ) + before_config = proxy_config._decrypt_db_variables(_parse_config_value(existing_record.config_value)) except Exception: before_config = None # Delete DB record if it exists — ignore if not found deleted = False try: - await ConfigOverridesRepository(prisma_client).table.delete( - where={"config_type": "hashicorp_vault"} - ) + await ConfigOverridesRepository(prisma_client).table.delete(where={"config_type": "hashicorp_vault"}) deleted = True except RecordNotFoundError: - verbose_proxy_logger.debug( - "No existing Hashicorp Vault config record to delete" - ) + verbose_proxy_logger.debug("No existing Hashicorp Vault config record to delete") _clear_hashicorp_vault_state(proxy_config) @@ -532,9 +510,7 @@ async def test_hashicorp_vault_connection( # Step 2: Verify the token is valid via token/lookup-self try: - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.SecretManager - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SecretManager) lookup_url = f"{client.vault_addr}/v1/auth/token/lookup-self" if client.vault_namespace: headers["X-Vault-Namespace"] = client.vault_namespace diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index b5ae8f93be6..cd2c5704778 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -59,50 +59,32 @@ def _resolve_model_for_cost_lookup(model: str) -> Tuple[str, Optional[str]]: model_info = first_deployment.get("model_info", {}) # Check base_model first (needed for Azure custom deployment names) - base_model = model_info.get("base_model") or litellm_params.get( - "base_model" - ) + base_model = model_info.get("base_model") or litellm_params.get("base_model") if base_model: - verbose_proxy_logger.debug( - f"Resolved model '{model}' to base_model '{base_model}' from router" - ) + verbose_proxy_logger.debug(f"Resolved model '{model}' to base_model '{base_model}' from router") custom_llm_provider = litellm_params.get("custom_llm_provider") return ( str(base_model), - ( - str(custom_llm_provider) - if custom_llm_provider is not None - else None - ), + (str(custom_llm_provider) if custom_llm_provider is not None else None), ) resolved_model = litellm_params.get("model") if resolved_model: - verbose_proxy_logger.debug( - f"Resolved model '{model}' to '{resolved_model}' from router" - ) + verbose_proxy_logger.debug(f"Resolved model '{model}' to '{resolved_model}' from router") custom_llm_provider = litellm_params.get("custom_llm_provider") return ( str(resolved_model), - ( - str(custom_llm_provider) - if custom_llm_provider is not None - else None - ), + (str(custom_llm_provider) if custom_llm_provider is not None else None), ) except Exception as e: - verbose_proxy_logger.debug( - f"Could not resolve model '{model}' from router: {e}" - ) + verbose_proxy_logger.debug(f"Could not resolve model '{model}' from router: {e}") # Return original model if not resolved return model, custom_llm_provider -def _calculate_period_costs( - num_requests, cost_per_request, input_cost, output_cost, margin_cost -): +def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost): """ Calculate costs for a given number of requests. @@ -192,9 +174,7 @@ async def update_cost_discount_config( if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) # Validate that all providers are valid LiteLLM providers @@ -214,9 +194,7 @@ async def update_cost_discount_config( # Validate discount values are between 0 and 1 for provider, discount in cost_discount_config.items(): if not isinstance(discount, (int, float)): - raise HTTPException( - status_code=400, detail=f"Discount for {provider} must be a number" - ) + raise HTTPException(status_code=400, detail=f"Discount for {provider} must be a number") if not (0 <= discount <= 1): raise HTTPException( status_code=400, @@ -240,9 +218,7 @@ async def update_cost_discount_config( # Update in-memory litellm.cost_discount_config litellm.cost_discount_config = cost_discount_config - verbose_proxy_logger.info( - f"Updated cost_discount_config: {cost_discount_config}" - ) + verbose_proxy_logger.info(f"Updated cost_discount_config: {cost_discount_config}") return { "message": "Cost discount configuration updated successfully", @@ -336,9 +312,7 @@ async def update_cost_margin_config( if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) # Validate that all providers are valid LiteLLM providers (except "global") @@ -478,9 +452,7 @@ async def estimate_cost( # Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4') resolved_model, resolved_provider = _resolve_model_for_cost_lookup(request.model) - verbose_proxy_logger.debug( - f"Cost estimate: request.model='{request.model}' resolved to '{resolved_model}'" - ) + verbose_proxy_logger.debug(f"Cost estimate: request.model='{request.model}' resolved to '{resolved_model}'") # Create a mock response with usage for completion_cost mock_response = ModelResponse( @@ -523,9 +495,7 @@ async def estimate_cost( input_cost = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0 output_cost = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0 - margin_cost = ( - cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 - ) + margin_cost = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 # Get model info for per-token pricing display try: diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index ecc338c2052..7c8a9b88191 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -133,9 +133,7 @@ async def unblock_user(data: BlockUsers): ): raise HTTPException( status_code=400, - detail={ - "error": "Blocked user check was never set. This call has no effect." - }, + detail={"error": "Blocked user check was never set. This call has no effect."}, ) if isinstance(litellm.blocked_user_list, list): @@ -144,9 +142,7 @@ async def unblock_user(data: BlockUsers): else: raise HTTPException( status_code=500, - detail={ - "error": "`blocked_user_list` must be set as a list. Filepaths can't be updated." - }, + detail={"error": "`blocked_user_list` must be set as a list. Filepaths can't be updated."}, ) return {"blocked_users": litellm.blocked_user_list} @@ -169,10 +165,7 @@ def new_budget_request(data: NewCustomerRequest) -> Optional[BudgetNewRequest]: if budget_kv_pairs: budget_request = BudgetNewRequest(**budget_kv_pairs) - if ( - budget_request.budget_reset_at is None - and budget_request.budget_duration is not None - ): + if budget_request.budget_reset_at is None and budget_request.budget_duration is not None: budget_request.budget_reset_at = datetime.utcnow() + timedelta( seconds=duration_in_seconds(duration=budget_request.budget_duration) ) @@ -199,9 +192,7 @@ async def _handle_customer_object_permission_update( """ if "object_permission" in non_default_values: existing_object_permission_id = ( - end_user_table_data_typed.object_permission_id - if end_user_table_data_typed is not None - else None + end_user_table_data_typed.object_permission_id if end_user_table_data_typed is not None else None ) object_permission_id = await handle_update_object_permission_common( data_json=non_default_values, @@ -342,10 +333,8 @@ async def new_end_user( budget_record = await BudgetRepository(prisma_client).table.create( data={ **_new_budget.model_dump(exclude_unset=True), - "created_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, # type: ignore - "updated_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, # type: ignore + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, } ) except Exception as e: @@ -424,9 +413,7 @@ async def new_end_user( dependencies=[Depends(user_api_key_auth)], ) async def end_user_info( - end_user_id: str = fastapi.Query( - description="End User ID in the request parameters" - ), + end_user_id: str = fastapi.Query(description="End User ID in the request parameters"), ): """ Get information about an end-user. An `end_user` is a customer (external user) of the proxy. @@ -564,14 +551,10 @@ async def update_end_user( # get non default values for key non_default_values = {} for k, v in data_json.items(): - if ( - v is not None - and v - not in ( - [], - {}, - 0, - ) + if v is not None and v not in ( + [], + {}, + 0, ): # models default to [], spend defaults to 0, we should not reset these values non_default_values[k] = v @@ -588,9 +571,7 @@ async def update_end_user( param="user_id", ) - end_user_table_data_typed = LiteLLM_EndUserTable( - **end_user_table_data.model_dump() - ) + end_user_table_data_typed = LiteLLM_EndUserTable(**end_user_table_data.model_dump()) ## Get budget table data ## end_user_budget_table = end_user_table_data_typed.litellm_budget_table @@ -620,27 +601,19 @@ async def update_end_user( if budget_table_data: if end_user_budget_table is None: ## Create new budget ## - budget_table_data_record = await BudgetRepository( - prisma_client - ).table.create( + budget_table_data_record = await BudgetRepository(prisma_client).table.create( data={ **budget_table_data, - "created_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, }, include={"end_users": True}, ) - update_end_user_table_data["budget_id"] = ( - budget_table_data_record.budget_id - ) + update_end_user_table_data["budget_id"] = budget_table_data_record.budget_id else: ## Update existing budget ## - budget_table_data_record = await BudgetRepository( - prisma_client - ).table.update( + budget_table_data_record = await BudgetRepository(prisma_client).table.update( where={"budget_id": end_user_budget_table.budget_id}, data=budget_table_data, ) @@ -665,12 +638,8 @@ async def update_end_user( include={"litellm_budget_table": True, "object_permission": True}, # type: ignore ) if response is None: - raise ValueError( - f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}" - ) - verbose_proxy_logger.debug( - f"received response from updating prisma client. response={response}" - ) + raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}") + verbose_proxy_logger.debug(f"received response from updating prisma client. response={response}") # Convert to dict and clean up recursive fields response_dict = response.model_dump() @@ -693,9 +662,7 @@ async def update_end_user( except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.update_end_user(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.update_end_user(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -740,25 +707,17 @@ async def delete_end_user( raise Exception("Not connected to DB!") verbose_proxy_logger.debug("/customer/delete: Received data = %s", data) - if ( - data.user_ids is not None - and isinstance(data.user_ids, list) - and len(data.user_ids) > 0 - ): + if data.user_ids is not None and isinstance(data.user_ids, list) and len(data.user_ids) > 0: # First check if all users exist existing_users = await EndUserRepository(prisma_client).table.find_many( where={"user_id": {"in": data.user_ids}} ) existing_user_ids = {user.user_id for user in existing_users} - missing_user_ids = [ - user_id for user_id in data.user_ids if user_id not in existing_user_ids - ] + missing_user_ids = [user_id for user_id in data.user_ids if user_id not in existing_user_ids] if missing_user_ids: raise ProxyException( - message="End User Id(s)={} do not exist in db".format( - ", ".join(missing_user_ids) - ), + message="End User Id(s)={} do not exist in db".format(", ".join(missing_user_ids)), type="not_found", code=404, param="user_ids", @@ -768,13 +727,10 @@ async def delete_end_user( response = await EndUserRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) - verbose_proxy_logger.debug( - f"received response from updating prisma client. response={response}" - ) + verbose_proxy_logger.debug(f"received response from updating prisma client. response={response}") return { "deleted_customers": response, - "message": "Successfully deleted customers with ids: " - + str(data.user_ids), + "message": "Successfully deleted customers with ids: " + str(data.user_ids), } else: raise ValueError(f"user_id is required, passed user_id = {data.user_ids}") @@ -782,9 +738,7 @@ async def delete_end_user( # update based on remaining passed in values except Exception as e: verbose_proxy_logger.error( - "litellm.proxy.proxy_server.delete_end_user(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.delete_end_user(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -824,11 +778,7 @@ async def list_end_user( ): raise HTTPException( status_code=401, - detail={ - "error": "Admin-only endpoint. Your user role={}".format( - user_api_key_dict.user_role - ) - }, + detail={"error": "Admin-only endpoint. Your user role={}".format(user_api_key_dict.user_role)}, ) if prisma_client is None: @@ -898,11 +848,7 @@ async def get_customer_daily_activity( ): raise HTTPException( status_code=401, - detail={ - "error": "Admin-only endpoint. Your user role={}".format( - user_api_key_dict.user_role - ) - }, + detail={"error": "Admin-only endpoint. Your user role={}".format(user_api_key_dict.user_role)}, ) from litellm.proxy.proxy_server import prisma_client @@ -917,17 +863,13 @@ async def get_customer_daily_activity( end_user_ids_list = end_user_ids.split(",") if end_user_ids else None exclude_end_user_ids_list: Optional[List[str]] = None if exclude_end_user_ids: - exclude_end_user_ids_list = ( - exclude_end_user_ids.split(",") if exclude_end_user_ids else None - ) + exclude_end_user_ids_list = exclude_end_user_ids.split(",") if exclude_end_user_ids else None # Fetch organization aliases for metadata where_condition = {} if end_user_ids_list: where_condition["user_id"] = {"in": list(end_user_ids_list)} - end_user_aliases = await EndUserRepository(prisma_client).table.find_many( - where=where_condition - ) + end_user_aliases = await EndUserRepository(prisma_client).table.find_many(where=where_condition) end_user_alias_metadata = {e.user_id: {"alias": e.alias} for e in end_user_aliases} # Query daily activity for organizations diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index 1333122c87a..dc594923166 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -96,9 +96,7 @@ async def create_fallback( ) # Validate that all fallback models exist in the router - invalid_fallback_models = [ - m for m in data.fallback_models if m not in model_names - ] + invalid_fallback_models = [m for m in data.fallback_models if m not in model_names] if invalid_fallback_models: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -136,9 +134,7 @@ async def create_fallback( fallback_key = "content_policy_fallbacks" # Get existing fallbacks - existing_fallbacks: List[Dict[str, List[str]]] = router_settings.get( - fallback_key, [] - ) + existing_fallbacks: List[Dict[str, List[str]]] = router_settings.get(fallback_key, []) # Update or add the fallback configuration fallback_updated = False @@ -226,16 +222,12 @@ async def get_fallback( ) # Get fallbacks using the existing utility function - fallback_models = get_all_fallbacks( - model=model, llm_router=llm_router, fallback_type=fallback_type - ) + fallback_models = get_all_fallbacks(model=model, llm_router=llm_router, fallback_type=fallback_type) if not fallback_models: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail={ - "error": f"No {fallback_type} fallbacks configured for model '{model}'" - }, + detail={"error": f"No {fallback_type} fallbacks configured for model '{model}'"}, ) return FallbackGetResponse( @@ -311,9 +303,7 @@ async def delete_fallback( fallback_key = "content_policy_fallbacks" # Get existing fallbacks - existing_fallbacks: List[Dict[str, List[str]]] = router_settings.get( - fallback_key, [] - ) + existing_fallbacks: List[Dict[str, List[str]]] = router_settings.get(fallback_key, []) # Find and remove the fallback configuration fallback_found = False @@ -327,9 +317,7 @@ async def delete_fallback( if not fallback_found: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail={ - "error": f"No {fallback_type} fallbacks configured for model '{model}'" - }, + detail={"error": f"No {fallback_type} fallbacks configured for model '{model}'"}, ) # Update router settings diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 99e276cfbb4..77e2f354bd4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -97,41 +97,25 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d auto_create_key = data_json.pop("auto_create_key", True) if auto_create_key is False: - data_json["table_name"] = ( - "user" # only create a user, don't create key if 'auto_create_key' set to False - ) + data_json["table_name"] = "user" # only create a user, don't create key if 'auto_create_key' set to False if litellm.default_internal_user_params and ( - data.user_role != LitellmUserRoles.PROXY_ADMIN.value - and data.user_role != LitellmUserRoles.PROXY_ADMIN + data.user_role != LitellmUserRoles.PROXY_ADMIN.value and data.user_role != LitellmUserRoles.PROXY_ADMIN ): for key, value in litellm.default_internal_user_params.items(): if key == "available_teams": continue elif key not in data_json or data_json[key] is None: data_json[key] = value - elif ( - key == "models" - and isinstance(data_json[key], list) - and len(data_json[key]) == 0 - ): + elif key == "models" and isinstance(data_json[key], list) and len(data_json[key]) == 0: data_json[key] = value ## INTERNAL USER ROLE ONLY DEFAULT PARAMS ## - if ( - data.user_role is not None - and data.user_role == LitellmUserRoles.INTERNAL_USER.value - ): - if ( - litellm.max_internal_user_budget is not None - and data_json.get("max_budget") is None - ): + if data.user_role is not None and data.user_role == LitellmUserRoles.INTERNAL_USER.value: + if litellm.max_internal_user_budget is not None and data_json.get("max_budget") is None: data_json["max_budget"] = litellm.max_internal_user_budget - if ( - litellm.internal_user_budget_duration is not None - and data_json.get("budget_duration") is None - ): + if litellm.internal_user_budget_duration is not None and data_json.get("budget_duration") is None: data_json["budget_duration"] = litellm.internal_user_budget_duration data_json.pop("teams", None) # handled separately @@ -169,24 +153,18 @@ async def _check_duplicate_user_field( if case_insensitive: where_clause[field_name]["mode"] = "insensitive" - existing_user = await UserRepository(prisma_client).table.find_first( - where=where_clause - ) + existing_user = await UserRepository(prisma_client).table.find_first(where=where_clause) if existing_user is not None: existing_value = getattr(existing_user, field_name, value) error_label = label or field_name raise HTTPException( status_code=409, - detail={ - "error": f"User with {error_label} {existing_value} already exists" - }, + detail={"error": f"User with {error_label} {existing_value} already exists"}, ) -async def _check_duplicate_user_email( - user_email: Optional[str], prisma_client: Any -) -> None: +async def _check_duplicate_user_email(user_email: Optional[str], prisma_client: Any) -> None: """ Helper function to check if a user email already exists in the database. """ @@ -270,9 +248,7 @@ async def _add_user_to_team( user_api_key_dict=user_api_key_dict, ) except HTTPException as e: - if e.status_code == 400 and ( - "already exists" in str(e) or "doesn't exist" in str(e) - ): + if e.status_code == 400 and ("already exists" in str(e) or "doesn't exist" in str(e)): verbose_proxy_logger.debug( "litellm.proxy.management_endpoints.internal_user_endpoints.new_user(): User already exists in team - {}".format( str(e) @@ -291,10 +267,7 @@ async def _add_user_to_team( str(e) ) ) - elif ( - isinstance(e, ProxyException) - and ProxyErrorTypes.team_member_already_in_team in e.type - ): + elif isinstance(e, ProxyException) and ProxyErrorTypes.team_member_already_in_team in e.type: verbose_proxy_logger.debug( "litellm.proxy.management_endpoints.internal_user_endpoints.new_user(): User already exists in team - {}".format( str(e) @@ -435,9 +408,7 @@ async def new_user( from litellm.proxy.proxy_server import _license_check, prisma_client if prisma_client is None: - raise HTTPException( - status_code=400, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) if prisma_client is None: raise HTTPException( @@ -460,8 +431,7 @@ async def new_user( # Check if user_api_key_dict is actually a UserAPIKeyAuth instance (not a Depends object) # This can happen when the function is called directly in tests if ( - data.user_role - in [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] + data.user_role in [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] and isinstance(user_api_key_dict, UserAPIKeyAuth) and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN ): @@ -476,9 +446,7 @@ async def new_user( teams = data.teams if teams is None: teams = check_if_default_team_set() - organization_ids = cast( - Optional[List[str]], data_json.pop("organizations", None) - ) + organization_ids = cast(Optional[List[str]], data_json.pop("organizations", None)) response = await generate_key_helper_fn(request_type="user", **data_json) # Admin UI Logic @@ -539,9 +507,7 @@ async def new_user( return new_user_response except Exception as e: - verbose_proxy_logger.exception( - "/user/new: Exception occured - {}".format(str(e)) - ) + verbose_proxy_logger.exception("/user/new: Exception occured - {}".format(str(e))) raise handle_exception_on_proxy(e) @@ -626,18 +592,14 @@ def get_user_id_from_request(request: Request) -> Optional[str]: return user_id -def _normalize_user_info_user_id( - request: Request, user_id: Optional[str] -) -> Optional[str]: +def _normalize_user_info_user_id(request: Request, user_id: Optional[str]) -> Optional[str]: """Normalize URL-decoded user_id while preserving '+' characters.""" if user_id is not None and " " in user_id: return get_user_id_from_request(request=request) return user_id -def _enforce_user_info_access( - user_id: Optional[str], user_api_key_dict: UserAPIKeyAuth -) -> None: +def _enforce_user_info_access(user_id: Optional[str], user_api_key_dict: UserAPIKeyAuth) -> None: """Re-validate that the caller may read the resolved ``user_id`` after URL-decoding has been finalized. @@ -663,8 +625,7 @@ def _enforce_user_info_access( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=( - f"key not allowed to access this user's info. user_id={user_id}, " - f"key's user_id={user_api_key_dict.user_id}" + f"key not allowed to access this user's info. user_id={user_id}, key's user_id={user_api_key_dict.user_id}" ), ) @@ -703,9 +664,7 @@ async def _get_user_info_teams( query_type="find_all", ) elif user_api_key_dict.user_id is not None and user_id is None: - caller_user_info = await prisma_client.get_data( - user_id=user_api_key_dict.user_id - ) + caller_user_info = await prisma_client.get_data(user_id=user_api_key_dict.user_id) caller_team_ids = getattr(caller_user_info, "teams", None) if caller_team_ids: teams_2 = await prisma_client.get_data( @@ -749,14 +708,10 @@ def _build_user_info_response( returned_keys = _process_keys_for_user_info(keys=keys, all_teams=teams_1) team_list.sort(key=lambda x: getattr(x, "team_alias", "") or "") - _user_info = ( - user_info.model_dump() if isinstance(user_info, BaseModel) else user_info - ) + _user_info = user_info.model_dump() if isinstance(user_info, BaseModel) else user_info if isinstance(_user_info, dict): _user_info.pop("password", None) - _user_info["metadata"] = _redact_scim_enterprise_metadata( - _user_info.get("metadata") - ) + _user_info["metadata"] = _redact_scim_enterprise_metadata(_user_info.get("metadata")) return UserInfoResponse( user_id=user_id, @@ -775,9 +730,7 @@ def _build_user_info_response( @management_endpoint_wrapper async def user_info( request: Request, - user_id: Optional[str] = fastapi.Query( - default=None, description="User ID in the request parameters" - ), + user_id: Optional[str] = fastapi.Query(default=None, description="User ID in the request parameters"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -803,13 +756,8 @@ async def user_info( raise Exception( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - if ( - user_id is None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - ): - return await _get_user_info_for_proxy_admin( - user_api_key_dict=user_api_key_dict - ) + if user_id is None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return await _get_user_info_for_proxy_admin(user_api_key_dict=user_api_key_dict) elif user_id is None: user_id = user_api_key_dict.user_id ## GET USER ROW ## @@ -848,11 +796,7 @@ async def user_info( return response_data except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.user_info(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.user_info(): Exception occured - {}".format(str(e))) raise handle_exception_on_proxy(e) @@ -880,9 +824,7 @@ async def _check_user_info_v2_access( # Helper: fetch the target user row (reused across branches) async def _fetch_target_user(): - return await UserRepository(prisma_client).table.find_unique( - where={"user_id": target_user_id} - ) + return await UserRepository(prisma_client).table.find_unique(where={"user_id": target_user_id}) # Rule 1: Proxy admins — fetch and return the target row directly if _user_has_admin_view(user_api_key_dict): @@ -905,14 +847,10 @@ async def _check_user_info_v2_access( return None # Get all teams the caller belongs to - teams = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": caller_user.teams}} - ) + teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": caller_user.teams}}) for team in teams: team_obj = LiteLLM_TeamTable(**team.model_dump()) - if _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): # Check if target user is in this team if team.team_id in (target_user.teams or []): return target_user @@ -929,9 +867,7 @@ async def _check_user_info_v2_access( @management_endpoint_wrapper async def user_info_v2( request: Request, - user_id: Optional[str] = fastapi.Query( - default=None, description="User ID in the request parameters" - ), + user_id: Optional[str] = fastapi.Query(default=None, description="User ID in the request parameters"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1009,9 +945,7 @@ async def user_info_v2( ) except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.user_info_v2(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.user_info_v2(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -1065,9 +999,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): admin_user_info = await prisma_client.get_data(user_id=admin_user_id) if admin_user_info is not None: admin_user_info = ( - admin_user_info.model_dump() - if isinstance(admin_user_info, BaseModel) - else admin_user_info + admin_user_info.model_dump() if isinstance(admin_user_info, BaseModel) else admin_user_info ) if isinstance(admin_user_info, dict): admin_user_info.pop("password", None) @@ -1109,14 +1041,8 @@ def _process_keys_for_user_info( if _key.get("team_id") == UI_SESSION_TOKEN_TEAM_ID: continue - if ( - "team_id" in _key - and _key["team_id"] is not None - and _key["team_id"] != "litellm-dashboard" - ): - team_info = get_team_from_list( - team_list=all_teams, team_id=_key["team_id"] - ) + if "team_id" in _key and _key["team_id"] is not None and _key["team_id"] != "litellm-dashboard": + team_info = get_team_from_list(team_list=all_teams, team_id=_key["team_id"]) if team_info is not None: team_alias = getattr(team_info, "team_alias", None) _key["team_alias"] = team_alias @@ -1166,13 +1092,9 @@ def _update_internal_user_params( ): # applies internal user limits, if user role updated non_default_values["max_budget"] = litellm.max_internal_user_budget - if ( - "budget_duration" not in non_default_values - ): # applies internal user limits, if user role updated + if "budget_duration" not in non_default_values: # applies internal user limits, if user role updated if is_internal_user and litellm.internal_user_budget_duration is not None: - non_default_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + non_default_values["budget_duration"] = litellm.internal_user_budget_duration from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time non_default_values["budget_reset_at"] = get_budget_reset_time( @@ -1194,13 +1116,9 @@ async def _schedule_user_update_audit_log( if prisma_client is None: return try: - updated_user_row = await UserRepository(prisma_client).table.find_first( - where={"user_id": response["user_id"]} - ) + updated_user_row = await UserRepository(prisma_client).table.find_first(where={"user_id": response["user_id"]}) if updated_user_row: - user_row_typed = LiteLLM_UserTable( - **updated_user_row.model_dump(exclude_none=True) - ) + user_row_typed = LiteLLM_UserTable(**updated_user_row.model_dump(exclude_none=True)) asyncio.create_task( UserManagementEventHooks.create_internal_user_audit_log( user_id=user_row_typed.user_id, @@ -1208,18 +1126,12 @@ async def _schedule_user_update_audit_log( litellm_changed_by=litellm_changed_by or user_api_key_dict.user_id, user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, - before_value=( - existing_user_row.model_dump_json(exclude_none=True) - if existing_user_row - else None - ), + before_value=(existing_user_row.model_dump_json(exclude_none=True) if existing_user_row else None), after_value=user_row_typed.model_dump_json(exclude_none=True), ) ) except Exception as audit_error: - verbose_proxy_logger.warning( - f"Failed to create audit log for user {response.get('user_id')}: {audit_error}" - ) + verbose_proxy_logger.warning(f"Failed to create audit log for user {response.get('user_id')}: {audit_error}") def _check_user_update_authz( @@ -1228,19 +1140,12 @@ def _check_user_update_authz( existing_user_row: Optional[BaseModel], ) -> None: """Authorization checks for /user/update — raises HTTPException on failure.""" - if ( - user_request.user_role is not None - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - ): - raise HTTPException( - status_code=403, detail="Only proxy admins can modify user roles." - ) + if user_request.user_role is not None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + raise HTTPException(status_code=403, detail="Only proxy admins can modify user roles.") if existing_user_row is not None: typed_row = LiteLLM_UserTable(**existing_user_row.model_dump(exclude_none=True)) - if not can_user_call_user_update( - user_api_key_dict=user_api_key_dict, user_info=typed_row - ): + if not can_user_call_user_update(user_api_key_dict=user_api_key_dict, user_info=typed_row): raise HTTPException( status_code=403, detail={ @@ -1271,9 +1176,7 @@ async def _invalidate_user_spend_counter_if_changed( if non_default_values.get("spend") is not None: from litellm.proxy.proxy_server import _invalidate_spend_counter - await _invalidate_spend_counter( - counter_key=f"spend:user:{non_default_values['user_id']}" - ) + await _invalidate_spend_counter(counter_key=f"spend:user:{non_default_values['user_id']}") async def _update_single_user_helper( @@ -1296,9 +1199,7 @@ async def _update_single_user_helper( raise ValueError("Either user_id or user_email must be provided") data_json: dict = user_request.model_dump(exclude_unset=True) - non_default_values = _update_internal_user_params( - data_json=data_json, data=user_request - ) + non_default_values = _update_internal_user_params(data_json=data_json, data=user_request) _hash_password_in_dict(non_default_values) existing_user_row: Optional[BaseModel] = None @@ -1314,26 +1215,17 @@ async def _update_single_user_helper( _check_user_update_authz(user_request, user_api_key_dict, existing_user_row) if existing_user_row is not None: - existing_user_row = LiteLLM_UserTable( - **existing_user_row.model_dump(exclude_none=True) - ) + existing_user_row = LiteLLM_UserTable(**existing_user_row.model_dump(exclude_none=True)) # Prevent budget self-escalation (GHSA-wvg4-6222-3q4r): non-admin callers # must not be able to raise their own budget/spend fields. # can_user_call_user_update() already restricts non-admins to self-updates, # so this guard only fires for self-escalation attempts. _target_user_id = user_request.user_id or ( - getattr(existing_user_row, "user_id", None) - if existing_user_row is not None - else None + getattr(existing_user_row, "user_id", None) if existing_user_row is not None else None ) - _is_self_update = ( - _target_user_id is not None and user_api_key_dict.user_id == _target_user_id - ) - if ( - _is_self_update - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - ): + _is_self_update = _target_user_id is not None and user_api_key_dict.user_id == _target_user_id + if _is_self_update and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: _protected_fields = ("max_budget", "soft_budget", "spend") for _field in _protected_fields: if _field in non_default_values: @@ -1345,9 +1237,7 @@ async def _update_single_user_helper( ) existing_metadata = ( - cast(Dict, getattr(existing_user_row, "metadata", {}) or {}) - if existing_user_row is not None - else {} + cast(Dict, getattr(existing_user_row, "metadata", {}) or {}) if existing_user_row is not None else {} ) non_default_values = prepare_metadata_fields( @@ -1377,11 +1267,7 @@ async def _update_single_user_helper( query_type="find_all", ) - if ( - existing_user_rows - and isinstance(existing_user_rows, list) - and len(existing_user_rows) > 0 - ): + if existing_user_rows and isinstance(existing_user_rows, list) and len(existing_user_rows) > 0: for existing_user in existing_user_rows: non_default_values["user_id"] = existing_user.user_id response = await prisma_client.update_data( @@ -1394,9 +1280,7 @@ async def _update_single_user_helper( # Create new user if not found non_default_values["user_id"] = str(uuid.uuid4()) non_default_values["user_email"] = user_request.user_email - response = await prisma_client.insert_data( - data=non_default_values, table_name="user" - ) + response = await prisma_client.insert_data(data=non_default_values, table_name="user") if response is not None: await _schedule_user_update_audit_log( @@ -1503,9 +1387,7 @@ async def user_update( return response except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.user_update(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.user_update(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): @@ -1546,11 +1428,7 @@ async def bulk_update_processed_users( # Record success results.append( UserUpdateResult( - user_id=( - response.get("user_id") - if response - else user_request.user_id - ), + user_id=(response.get("user_id") if response else user_request.user_id), user_email=user_request.user_email, success=True, updated_user=response, @@ -1664,26 +1542,17 @@ async def bulk_user_update( ) # Only proxy admins can modify user_role in bulk updates - _bulk_role = ( - getattr(data.user_updates, "user_role", None) if data.user_updates else None - ) + _bulk_role = getattr(data.user_updates, "user_role", None) if data.user_updates else None if _bulk_role is None and data.users: - _bulk_role = next( - (u.user_role for u in data.users if u.user_role is not None), None - ) - if ( - _bulk_role is not None - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - ): + _bulk_role = next((u.user_role for u in data.users if u.user_role is not None), None) + if _bulk_role is not None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, detail="Only proxy admins can modify user roles.", ) # Determine the list of users to update - users_to_update: Union[ - List[UpdateUserRequest], List[UpdateUserRequestNoUserIDorEmail] - ] = [] + users_to_update: Union[List[UpdateUserRequest], List[UpdateUserRequestNoUserIDorEmail]] = [] if data.all_users and data.user_updates: # Only proxy admins can update all users at once @@ -1693,9 +1562,7 @@ async def bulk_user_update( detail="Only proxy admins can update all users at once.", ) # Optimized path for updating all users directly in database - all_users_in_db = await UserRepository(prisma_client).table.find_many( - order={"created_at": "desc"} - ) + all_users_in_db = await UserRepository(prisma_client).table.find_many(order={"created_at": "desc"}) if not all_users_in_db: raise HTTPException( @@ -1715,9 +1582,7 @@ async def bulk_user_update( # Apply update transformations (reuse existing logic) data_json: dict = data.user_updates.model_dump(exclude_unset=True) - non_default_values = _update_internal_user_params( - data_json=data_json, data=data.user_updates - ) + non_default_values = _update_internal_user_params(data_json=data_json, data=data.user_updates) # Remove user identification fields since we're updating by user_id non_default_values.pop("user_id", None) @@ -1752,8 +1617,7 @@ async def bulk_user_update( UserManagementEventHooks.create_internal_user_audit_log( user_id=user_api_key_dict.user_id or "", action="updated", - litellm_changed_by=litellm_changed_by - or user_api_key_dict.user_id, + litellm_changed_by=litellm_changed_by or user_api_key_dict.user_id, user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, before_value=f"Updated {len(all_users_in_db)} users", @@ -1761,9 +1625,7 @@ async def bulk_user_update( ) ) except Exception as audit_error: - verbose_proxy_logger.warning( - f"Failed to create bulk audit log: {audit_error}" - ) + verbose_proxy_logger.warning(f"Failed to create bulk audit log: {audit_error}") except Exception as e: verbose_proxy_logger.exception(f"Failed to perform bulk update: {e}") @@ -1851,9 +1713,7 @@ async def get_user_key_counts( return result -def _validate_sort_params( - sort_by: Optional[str], sort_order: str -) -> Optional[Dict[str, str]]: +def _validate_sort_params(sort_by: Optional[str], sort_order: str) -> Optional[Dict[str, str]]: order_by: Dict[str, str] = {} if sort_by is None: @@ -1870,9 +1730,7 @@ def _validate_sort_params( if sort_by not in valid_columns: raise HTTPException( status_code=400, - detail={ - "error": f"Invalid sort column. Must be one of: {', '.join(valid_columns)}" - }, + detail={"error": f"Invalid sort column. Must be one of: {', '.join(valid_columns)}"}, ) # Validate sort_order @@ -1907,9 +1765,7 @@ async def _authorize_user_list_request( if user_api_key_dict.user_id is None: raise HTTPException( status_code=403, - detail={ - "error": "Only proxy admins and organization admins can list users." - }, + detail={"error": "Only proxy admins and organization admins can list users."}, ) try: caller_user = await get_user_object( @@ -1922,16 +1778,12 @@ async def _authorize_user_list_request( except ValueError: raise HTTPException( status_code=403, - detail={ - "error": "Only proxy admins and organization admins can list users." - }, + detail={"error": "Only proxy admins and organization admins can list users."}, ) if caller_user is None: raise HTTPException( status_code=403, - detail={ - "error": "Only proxy admins and organization admins can list users." - }, + detail={"error": "Only proxy admins and organization admins can list users."}, ) allowed_org_ids = [ @@ -1942,23 +1794,17 @@ async def _authorize_user_list_request( if not allowed_org_ids: raise HTTPException( status_code=403, - detail={ - "error": "Only proxy admins and organization admins can list users." - }, + detail={"error": "Only proxy admins and organization admins can list users."}, ) # If client also sent organization_ids, intersect with allowed orgs if organization_ids: - requested = set( - oid.strip() for oid in organization_ids.split(",") if oid.strip() - ) + requested = set(oid.strip() for oid in organization_ids.split(",") if oid.strip()) intersection = list(requested & set(allowed_org_ids)) if not intersection: raise HTTPException( status_code=403, - detail={ - "error": "You do not have org_admin access to the requested organization(s)." - }, + detail={"error": "You do not have org_admin access to the requested organization(s)."}, ) allowed_org_ids = intersection @@ -1972,32 +1818,18 @@ async def _authorize_user_list_request( response_model=UserListResponse, ) async def get_users( - role: Optional[str] = fastapi.Query( - default=None, description="Filter users by role" - ), - user_ids: Optional[str] = fastapi.Query( - default=None, description="Get list of users by user_ids" - ), - sso_user_ids: Optional[str] = fastapi.Query( - default=None, description="Get list of users by sso_user_id" - ), - user_email: Optional[str] = fastapi.Query( - default=None, description="Filter users by partial email match" - ), - team: Optional[str] = fastapi.Query( - default=None, description="Filter users by team id" - ), + role: Optional[str] = fastapi.Query(default=None, description="Filter users by role"), + user_ids: Optional[str] = fastapi.Query(default=None, description="Get list of users by user_ids"), + sso_user_ids: Optional[str] = fastapi.Query(default=None, description="Get list of users by sso_user_id"), + user_email: Optional[str] = fastapi.Query(default=None, description="Filter users by partial email match"), + team: Optional[str] = fastapi.Query(default=None, description="Filter users by team id"), page: int = fastapi.Query(default=1, ge=1, description="Page number"), - page_size: int = fastapi.Query( - default=25, ge=1, le=100, description="Number of items per page" - ), + page_size: int = fastapi.Query(default=25, ge=1, le=100, description="Number of items per page"), sort_by: Optional[str] = fastapi.Query( default=None, description="Column to sort by (e.g. 'user_id', 'user_email', 'created_at', 'spend')", ), - sort_order: str = fastapi.Query( - default="asc", description="Sort order ('asc' or 'desc')" - ), + sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), organization_ids: Optional[str] = fastapi.Query( default=None, description="Filter users by organization membership. Comma-separated list of org IDs.", @@ -2091,13 +1923,9 @@ async def get_users( } if organization_ids: - org_id_list = [ - oid.strip() for oid in organization_ids.split(",") if oid.strip() - ] + org_id_list = [oid.strip() for oid in organization_ids.split(",") if oid.strip()] if org_id_list: - where_conditions["organization_memberships"] = { - "some": {"organization_id": {"in": org_id_list}} - } + where_conditions["organization_memberships"] = {"some": {"organization_id": {"in": org_id_list}}} ## Filter any none fastapi.Query params - e.g. where_conditions: {'user_email': {'contains': Query(None), 'mode': 'insensitive'}, 'teams': {'has': Query(None)}} where_conditions = {k: v for k, v in where_conditions.items() if v is not None} @@ -2105,30 +1933,22 @@ async def get_users( # Build order_by conditions order_by: Optional[Dict[str, str]] = ( - _validate_sort_params(sort_by, sort_order) - if sort_by is not None and isinstance(sort_by, str) - else None + _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) users = await UserRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, - order=( - order_by if order_by else {"created_at": "desc"} - ), # Default to created_at desc if no sort specified + order=(order_by if order_by else {"created_at": "desc"}), # Default to created_at desc if no sort specified ) # Get total count of user rows - total_count = await UserRepository(prisma_client).table.count( - where=where_conditions - ) + total_count = await UserRepository(prisma_client).table.count(where=where_conditions) # Get key count for each user if users is not None: - user_key_counts = await get_user_key_counts( - prisma_client, [user.user_id for user in users] - ) + user_key_counts = await get_user_key_counts(prisma_client, [user.user_id for user in users]) else: user_key_counts = {} @@ -2142,14 +1962,8 @@ async def get_users( if users is not None: for user in users: user_dump = user.model_dump() - user_dump["metadata"] = _redact_scim_enterprise_metadata( - user_dump.get("metadata") - ) - user_list.append( - LiteLLM_UserTableWithKeyCount( - **user_dump, key_count=user_key_counts.get(user.user_id, 0) - ) - ) + user_dump["metadata"] = _redact_scim_enterprise_metadata(user_dump.get("metadata")) + user_list.append(LiteLLM_UserTableWithKeyCount(**user_dump, key_count=user_key_counts.get(user.user_id, 0))) else: user_list = [] @@ -2218,9 +2032,7 @@ async def delete_user( # cross-check data.user_ids against the caller's scope, so without this # loop an org-admin of org-A could delete users in org-B by supplying # {"user_ids": [victim_in_org_B], "organization_id": "org-A"}. - caller_is_proxy_admin = ( - user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - ) + caller_is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value caller_admin_org_ids: set = set() if not caller_is_proxy_admin: caller_memberships = ( @@ -2233,24 +2045,20 @@ async def delete_user( if user_api_key_dict.user_id else [] ) - caller_admin_org_ids = { - m.organization_id for m in caller_memberships if m.organization_id - } + caller_admin_org_ids = {m.organization_id for m in caller_memberships if m.organization_id} if not caller_admin_org_ids: raise HTTPException( status_code=403, - detail={ - "error": "Only PROXY_ADMIN or ORG_ADMIN users may delete users." - }, + detail={"error": "Only PROXY_ADMIN or ORG_ADMIN users may delete users."}, ) # Batch-fetch target memberships once before the per-user loop. Avoids # an N+1 DB call when delete_user is called with a large user_ids list. target_org_ids_by_user: Dict[str, set] = {} if not caller_is_proxy_admin: - all_target_memberships = await OrganizationMembershipRepository( - prisma_client - ).table.find_many(where={"user_id": {"in": data.user_ids}}) + all_target_memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + where={"user_id": {"in": data.user_ids}} + ) for m in all_target_memberships: if not m.organization_id: continue @@ -2258,9 +2066,7 @@ async def delete_user( # check that all teams passed exist for user_id in data.user_ids: - user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) if user_row is None: raise HTTPException( @@ -2312,9 +2118,7 @@ async def delete_user( ) ## CLEANUP MEMBERS_WITH_ROLES - fetch_all_teams = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": user_row.teams}} - ) + fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}}) teams_to_update = [] for team in fetch_all_teams: is_member_in_team, new_team_members = _cleanup_members_with_roles( @@ -2326,9 +2130,7 @@ async def delete_user( ), ) if is_member_in_team: - _db_new_team_members: List[dict] = [ - m.model_dump() for m in new_team_members - ] + _db_new_team_members: List[dict] = [m.model_dump() for m in new_team_members] team.members_with_roles = json.dumps(_db_new_team_members) teams_to_update.append(team) @@ -2342,9 +2144,7 @@ async def delete_user( # End of Audit logging ## DELETE ASSOCIATED KEYS - await VerificationTokenRepository(prisma_client).table.delete_many( - where={"user_id": {"in": data.user_ids}} - ) + await VerificationTokenRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) ## DELETE ASSOCIATED INVITATION LINKS await InvitationLinkRepository(prisma_client).table.delete_many( @@ -2358,19 +2158,13 @@ async def delete_user( ) ## DELETE ASSOCIATED ORGANIZATION MEMBERSHIPS - await OrganizationMembershipRepository(prisma_client).table.delete_many( - where={"user_id": {"in": data.user_ids}} - ) + await OrganizationMembershipRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) ## DELETE ASSOCIATED TEAM MEMBERSHIPS - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"user_id": {"in": data.user_ids}} - ) + await TeamMembershipRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) ## DELETE USERS - deleted_users = await UserRepository(prisma_client).table.delete_many( - where={"user_id": {"in": data.user_ids}} - ) + deleted_users = await UserRepository(prisma_client).table.delete_many(where={"user_id": {"in": data.user_ids}}) return deleted_users @@ -2398,18 +2192,14 @@ async def add_internal_user_to_organization( try: # Check if organization_id exists - organization_row = await OrganizationRepository( - prisma_client - ).table.find_unique(where={"organization_id": organization_id}) + organization_row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": organization_id} + ) if organization_row is None: - raise Exception( - f"Organization not found, passed organization_id={organization_id}" - ) + raise Exception(f"Organization not found, passed organization_id={organization_id}") # Create a new organization membership entry - new_membership = await OrganizationMembershipRepository( - prisma_client - ).table.create( + new_membership = await OrganizationMembershipRepository(prisma_client).table.create( data={ "user_id": user_id, "organization_id": organization_id, @@ -2465,9 +2255,7 @@ async def _resolve_org_filter_for_user_search( # This allows team admins who are org members to search users in their org. member_org_ids: List[str] = [] if caller_user is not None: - member_org_ids = [ - m.organization_id for m in (caller_user.organization_memberships or []) - ] + member_org_ids = [m.organization_id for m in (caller_user.organization_memberships or [])] if member_org_ids: return member_org_ids @@ -2511,17 +2299,13 @@ async def _resolve_team_org_filter( except HTTPException: raise HTTPException( status_code=403, - detail={ - "error": f"scope_user_search_to_org is enabled but team '{team_id}' was not found." - }, + detail={"error": f"scope_user_search_to_org is enabled but team '{team_id}' was not found."}, ) if not _is_user_team_admin(user_api_key_dict, team_obj): raise HTTPException( status_code=403, - detail={ - "error": "scope_user_search_to_org is enabled. You must be an admin of this team to search users." - }, + detail={"error": "scope_user_search_to_org is enabled. You must be an admin of this team to search users."}, ) if team_obj.organization_id: @@ -2545,22 +2329,14 @@ async def _resolve_team_org_filter( }, ) async def ui_view_users( - user_id: Optional[str] = fastapi.Query( - default=None, description="User ID in the request parameters" - ), - user_email: Optional[str] = fastapi.Query( - default=None, description="User email in the request parameters" - ), + user_id: Optional[str] = fastapi.Query(default=None, description="User ID in the request parameters"), + user_email: Optional[str] = fastapi.Query(default=None, description="User email in the request parameters"), team_id: Optional[str] = fastapi.Query( default=None, description="Team ID — used when a team admin searches for users to add to their team", ), - page: int = fastapi.Query( - default=1, description="Page number for pagination", ge=1 - ), - page_size: int = fastapi.Query( - default=50, description="Number of items per page", ge=1, le=100 - ), + page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), + page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=100), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -2614,14 +2390,10 @@ async def ui_view_users( # Apply org filter when scope_user_search_to_org is ON and caller is not proxy admin if org_filter_ids is not None: - where_conditions["organization_memberships"] = { - "some": {"organization_id": {"in": org_filter_ids}} - } + where_conditions["organization_memberships"] = {"some": {"organization_id": {"in": org_filter_ids}}} # Query users with pagination and filters - users: Optional[List[BaseModel]] = await UserRepository( - prisma_client - ).table.find_many( + users: Optional[List[BaseModel]] = await UserRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -2643,23 +2415,14 @@ async def ui_view_users( # Using shared metric helper implementations from common_daily_activity -async def _resolve_user_email_metadata( - prisma_client: "PrismaClient", records: list[Any] -) -> dict[str, dict]: +async def _resolve_user_email_metadata(prisma_client: "PrismaClient", records: list[Any]) -> dict[str, dict]: """Map each user_id on the page to its email/alias so the Usage dashboard can label the 'Spend Per User' chart with the email instead of the raw UUID.""" - user_ids = { - record.user_id for record in records if getattr(record, "user_id", None) - } + user_ids = {record.user_id for record in records if getattr(record, "user_id", None)} if not user_ids: return {} - users = await UserRepository(prisma_client).table.find_many( - where={"user_id": {"in": list(user_ids)}} - ) - return { - user.user_id: {"user_email": user.user_email, "user_alias": user.user_alias} - for user in users - } + users = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": list(user_ids)}}) + return {user.user_id: {"user_email": user.user_email, "user_alias": user.user_alias} for user in users} @router.get( @@ -2690,12 +2453,8 @@ async def get_user_daily_activity( default=None, description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", ), - page: int = fastapi.Query( - default=1, description="Page number for pagination", ge=1 - ), - page_size: int = fastapi.Query( - default=50, description="Items per page", ge=1, le=1000 - ), + page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), + page_size: int = fastapi.Query(default=50, description="Items per page", ge=1, le=1000), timezone: Optional[int] = fastapi.Query( default=None, description="Timezone offset in minutes from UTC (e.g., 480 for PST). " @@ -2745,9 +2504,7 @@ async def get_user_daily_activity( if user_id != caller_user_id: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Non-admin users can only view their own spend data." - }, + detail={"error": "Non-admin users can only view their own spend data."}, ) entity_id = user_id @@ -2764,17 +2521,13 @@ async def get_user_daily_activity( page=page, page_size=page_size, timezone_offset_minutes=timezone, - resolve_entity_metadata=lambda records: _resolve_user_email_metadata( - prisma_client, records - ), + resolve_entity_metadata=lambda records: _resolve_user_email_metadata(prisma_client, records), ) except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - "/spend/daily/analytics: Exception occured - {}".format(str(e)) - ) + verbose_proxy_logger.exception("/spend/daily/analytics: Exception occured - {}".format(str(e))) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Failed to fetch analytics: {str(e)}"}, @@ -2846,9 +2599,7 @@ async def get_user_daily_activity_aggregated( if user_id != caller_user_id: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Non-admin users can only view their own spend data." - }, + detail={"error": "Non-admin users can only view their own spend data."}, ) entity_id = user_id @@ -2868,9 +2619,7 @@ async def get_user_daily_activity_aggregated( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - "/user/daily/activity/aggregated: Exception occured - {}".format(str(e)) - ) + verbose_proxy_logger.exception("/user/daily/activity/aggregated: Exception occured - {}".format(str(e))) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Failed to fetch analytics: {str(e)}"}, diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index a5a364c3679..e1496ae9bf1 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -43,9 +43,7 @@ async def create_jwt_key_mapping( from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=403, detail="Only proxy admins can create JWT key mappings" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can create JWT key mappings") if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") @@ -62,9 +60,7 @@ async def create_jwt_key_mapping( if data.description is not None: create_data["description"] = data.description - new_mapping = await JWTKeyMappingRepository(prisma_client).table.create( - data=create_data - ) + new_mapping = await JWTKeyMappingRepository(prisma_client).table.create(data=create_data) # Invalidate cache cache_key = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}" @@ -100,9 +96,7 @@ async def update_jwt_key_mapping( from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=403, detail="Only proxy admins can update JWT key mappings" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can update JWT key mappings") if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") @@ -114,9 +108,7 @@ async def update_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( - where={"id": data.id} - ) + old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -158,18 +150,14 @@ async def delete_jwt_key_mapping( from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=403, detail="Only proxy admins can delete JWT key mappings" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can delete JWT key mappings") if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") try: # Get old mapping for cache invalidation - old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( - where={"id": data.id} - ) + old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id}) if old_mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") @@ -198,9 +186,7 @@ async def list_jwt_key_mappings( # Admin Viewer follows the read-parity rule. if not _user_has_admin_view(user_api_key_dict): - raise HTTPException( - status_code=403, detail="Only proxy admins can list JWT key mappings" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can list JWT key mappings") if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") @@ -238,23 +224,17 @@ async def info_jwt_key_mapping( # Admin Viewer follows the read-parity rule. if not _user_has_admin_view(user_api_key_dict): - raise HTTPException( - status_code=403, detail="Only proxy admins can get JWT key mapping info" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can get JWT key mapping info") if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") try: - mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( - where={"id": id} - ) + mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": id}) if mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") return _to_response(mapping) except HTTPException: raise except Exception: - raise HTTPException( - status_code=500, detail="Failed to get JWT key mapping info." - ) + raise HTTPException(status_code=500, detail="Failed to get JWT key mapping info.") diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 67f9103d4b7..eed7d869a5d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -133,14 +133,10 @@ async def _check_custom_key_allowed(custom_key_value: Optional[str]) -> None: ui_settings = await get_ui_settings_cached() if ui_settings.get("disable_custom_api_keys", False) is True: - verbose_proxy_logger.warning( - "Custom API key rejected: disable_custom_api_keys is enabled" - ) + verbose_proxy_logger.warning("Custom API key rejected: disable_custom_api_keys is enabled") raise HTTPException( status_code=403, - detail={ - "error": "Custom API key values are disabled by your administrator. Keys must be auto-generated." - }, + detail={"error": "Custom API key values are disabled by your administrator. Keys must be auto-generated."}, ) @@ -148,9 +144,7 @@ def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]): return data.team_id is not None -def _get_user_in_team( - team_table: LiteLLM_TeamTableCachedObj, user_id: Optional[str] -) -> Optional[Member]: +def _get_user_in_team(team_table: LiteLLM_TeamTableCachedObj, user_id: Optional[str]) -> Optional[Member]: if user_id is None: return None for member in team_table.members_with_roles: @@ -223,10 +217,7 @@ def _is_allowed_to_make_key_request( Relevant issue: https://github.com/BerriAI/litellm/issues/7336 """ ## BASE CASE - PROXY ADMIN - if ( - user_api_key_dict.user_role is not None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - ): + if user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return True if user_id is not None: @@ -237,10 +228,7 @@ def _is_allowed_to_make_key_request( ) if team_id is not None: - if ( - user_api_key_dict.team_id is not None - and user_api_key_dict.team_id == UI_TEAM_ID - ): + if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == UI_TEAM_ID: return True # handle https://github.com/BerriAI/litellm/issues/7482 return True @@ -254,9 +242,7 @@ def _team_key_operation_team_member_check( route: KeyManagementRoutes, ): if assigned_user_id is not None: - key_assigned_user_in_team = _get_user_in_team( - team_table=team_table, user_id=assigned_user_id - ) + key_assigned_user_in_team = _get_user_in_team(team_table=team_table, user_id=assigned_user_id) if key_assigned_user_in_team is None: raise HTTPException( @@ -264,13 +250,10 @@ def _team_key_operation_team_member_check( detail=f"User={assigned_user_id} not assigned to team={team_table.team_id}", ) - team_member_object = _get_user_in_team( - team_table=team_table, user_id=user_api_key_dict.user_id - ) + team_member_object = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id) is_admin = ( - user_api_key_dict.user_role is not None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) if is_admin: @@ -282,8 +265,7 @@ def _team_key_operation_team_member_check( ) elif ( "allowed_team_member_roles" in team_key_generation - and team_member_object.role - not in team_key_generation["allowed_team_member_roles"] + and team_member_object.role not in team_key_generation["allowed_team_member_roles"] ): raise HTTPException( status_code=400, @@ -298,9 +280,7 @@ def _team_key_operation_team_member_check( return True -def _key_generation_required_param_check( - data: GenerateKeyRequest, required_params: Optional[List[str]] -): +def _key_generation_required_param_check(data: GenerateKeyRequest, required_params: Optional[List[str]]): if required_params is None: return True @@ -322,10 +302,7 @@ def _team_key_generation_check( ): if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return True - if ( - litellm.key_generation_settings is not None - and "team_key_generation" in litellm.key_generation_settings - ): + if litellm.key_generation_settings is not None and "team_key_generation" in litellm.key_generation_settings: _team_key_generation = litellm.key_generation_settings["team_key_generation"] else: _team_key_generation = TeamUIKeyGenerationConfig( @@ -359,10 +336,7 @@ def _personal_key_membership_check( user_api_key_dict: UserAPIKeyAuth, personal_key_generation: Optional[PersonalUIKeyGenerationConfig], ): - if ( - personal_key_generation is None - or "allowed_user_roles" not in personal_key_generation - ): + if personal_key_generation is None or "allowed_user_roles" not in personal_key_generation: return True if user_api_key_dict.user_role not in personal_key_generation["allowed_user_roles"]: @@ -374,18 +348,14 @@ def _personal_key_membership_check( return True -def _personal_key_generation_check( - user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest -): +def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest): if ( litellm.key_generation_settings is None or litellm.key_generation_settings.get("personal_key_generation") is None ): return True - _personal_key_generation = litellm.key_generation_settings[ - "personal_key_generation" - ] # type: ignore + _personal_key_generation = litellm.key_generation_settings["personal_key_generation"] # type: ignore _personal_key_membership_check( user_api_key_dict, @@ -413,8 +383,7 @@ def key_generation_check( ## check if key is for team or individual is_team_key = _is_team_key(data=data) _is_admin = ( - user_api_key_dict.user_role is not None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) if is_team_key: if team_table is None and litellm.key_generation_settings is not None: @@ -437,9 +406,7 @@ def key_generation_check( route=route, ) else: - return _personal_key_generation_check( - user_api_key_dict=user_api_key_dict, data=data - ) + return _personal_key_generation_check(user_api_key_dict=user_api_key_dict, data=data) def common_key_access_checks( @@ -567,9 +534,7 @@ def _check_allowed_routes_caller_permission( return if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return - if allow_safe_presets and all( - r in _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS for r in allowed_routes - ): + if allow_safe_presets and all(r in _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS for r in allowed_routes): return raise HTTPException( status_code=403, @@ -613,9 +578,7 @@ async def validate_team_id_used_in_service_account_request( return True -_BUDGET_NUMERIC_KEYS = frozenset( - ["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"] -) +_BUDGET_NUMERIC_KEYS = frozenset(["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"]) def _enforce_upperbound_key_params( @@ -638,9 +601,7 @@ def _enforce_upperbound_key_params( if not math.isfinite(value): raise HTTPException( status_code=400, - detail={ - "error": f"{key} must be a finite number. Received: {value}" - }, + detail={"error": f"{key} must be a finite number. Received: {value}"}, ) if litellm.upperbound_key_generate_params is None: @@ -702,11 +663,7 @@ async def _common_key_generation_helper( premium_user=premium_user, ) - if ( - data.metadata is not None - and data.metadata.get("service_account_id") is not None - and data.team_id is None - ): + if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None: await validate_team_id_used_in_service_account_request( team_id=data.team_id, prisma_client=prisma_client, @@ -744,10 +701,7 @@ async def _common_key_generation_helper( # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. - is_ui_session_team_key = ( - user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID - and _requested_team_id is not None - ) + is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None # Session tokens (lite login) carry max_budget=None to avoid a per-session # LLM spend cap, but that None must not be read as "unlimited delegation # authority". A personal key (no team) has no team-budget enforcement at @@ -771,11 +725,7 @@ async def _common_key_generation_helper( delegation_ceiling = ( user_api_key_dict.max_budget if user_api_key_dict.max_budget is not None - else ( - team_table.max_budget - if user_api_key_dict.is_session_token and team_table is not None - else None - ) + else (team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None) ) if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -884,9 +834,7 @@ async def _common_key_generation_helper( from litellm.proxy.proxy_server import premium_user if premium_user is not True and data_json["tags"] is not None: - raise ValueError( - f"Only premium users can add tags to keys. {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"Only premium users can add tags to keys. {CommonProxyErrors.not_premium_user.value}") _metadata = data_json.get("metadata") if not _metadata: @@ -901,8 +849,7 @@ async def _common_key_generation_helper( object_permission=data_json.get("object_permission"), team_obj=team_table, prisma_client=prisma_client, - is_proxy_admin=user_api_key_dict.user_role - == LitellmUserRoles.PROXY_ADMIN.value, + is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value, ) if normalized_object_permission is not None: data_json["object_permission"] = normalized_object_permission @@ -928,16 +875,10 @@ async def _common_key_generation_helper( # Validate user-provided key format if data.key is not None and not data.key.startswith("sk-"): - _masked = ( - "{}****{}".format(data.key[:4], data.key[-4:]) - if len(data.key) > 8 - else "****" - ) + _masked = "{}****{}".format(data.key[:4], data.key[-4:]) if len(data.key) > 8 else "****" raise HTTPException( status_code=400, - detail={ - "error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {_masked}" - }, + detail={"error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {_masked}"}, ) # check org key limits - done here to handle inheriting org id from team @@ -980,13 +921,9 @@ async def _common_key_generation_helper( prisma_client=prisma_client, ) - response = await generate_key_helper_fn( - request_type="key", **data_json, table_name="key" - ) + response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key") - response["soft_budget"] = ( - data.soft_budget - ) # include the user-input soft budget in the response + response["soft_budget"] = data.soft_budget # include the user-input soft budget in the response response = GenerateKeyResponse(**response) @@ -1033,22 +970,14 @@ def _check_key_model_specific_limits( for key in keys: if key.metadata.get("model_rpm_limit", None) is not None: for model, rpm_limit in key.metadata.get("model_rpm_limit", {}).items(): - model_specific_rpm_limit[model] = ( - model_specific_rpm_limit.get(model, 0) + rpm_limit - ) + model_specific_rpm_limit[model] = model_specific_rpm_limit.get(model, 0) + rpm_limit if key.metadata.get("model_tpm_limit", None) is not None: for model, tpm_limit in key.metadata.get("model_tpm_limit", {}).items(): - model_specific_tpm_limit[model] = ( - model_specific_tpm_limit.get(model, 0) + tpm_limit - ) + model_specific_tpm_limit[model] = model_specific_tpm_limit.get(model, 0) + tpm_limit if model_rpm_limit is not None: for model, rpm_limit in model_rpm_limit.items(): - if ( - entity_rpm_limit is not None - and model_specific_rpm_limit.get(model, 0) + rpm_limit - > entity_rpm_limit - ): + if entity_rpm_limit is not None and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_rpm_limit: raise HTTPException( status_code=400, detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}", @@ -1057,8 +986,7 @@ def _check_key_model_specific_limits( entity_model_specific_rpm_limit = entity_model_rpm_limit_dict.get(model) if ( entity_model_specific_rpm_limit - and model_specific_rpm_limit.get(model, 0) + rpm_limit - > entity_model_specific_rpm_limit + and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_model_specific_rpm_limit ): raise HTTPException( status_code=400, @@ -1067,11 +995,7 @@ def _check_key_model_specific_limits( if model_tpm_limit is not None: for model, tpm_limit in model_tpm_limit.items(): - if ( - entity_tpm_limit is not None - and model_specific_tpm_limit.get(model, 0) + tpm_limit - > entity_tpm_limit - ): + if entity_tpm_limit is not None and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_tpm_limit: raise HTTPException( status_code=400, detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}", @@ -1080,8 +1004,7 @@ def _check_key_model_specific_limits( entity_model_specific_tpm_limit = entity_model_tpm_limit_dict.get(model) if ( entity_model_specific_tpm_limit - and model_specific_tpm_limit.get(model, 0) + tpm_limit - > entity_model_specific_tpm_limit + and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_model_specific_tpm_limit ): raise HTTPException( status_code=400, @@ -1179,10 +1102,7 @@ async def _check_team_key_limits( Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" """ - if ( - data.tpm_limit_type != "guaranteed_throughput" - and data.rpm_limit_type != "guaranteed_throughput" - ): + if data.tpm_limit_type != "guaranteed_throughput" and data.rpm_limit_type != "guaranteed_throughput": return # get all team keys # calculate allocated tpm/rpm limit @@ -1246,15 +1166,9 @@ async def _check_project_key_limits( # Validate key max_budget <= project max_budget project_max_budget = None if project_obj.litellm_budget_table is not None: - project_max_budget = getattr( - project_obj.litellm_budget_table, "max_budget", None - ) + project_max_budget = getattr(project_obj.litellm_budget_table, "max_budget", None) - if ( - data.max_budget is not None - and project_max_budget is not None - and data.max_budget > project_max_budget - ): + if data.max_budget is not None and project_max_budget is not None and data.max_budget > project_max_budget: raise HTTPException( status_code=400, detail={ @@ -1342,13 +1256,9 @@ async def _validate_caller_can_assign_key_org( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) - memberships = ( - getattr(user_row, "organization_memberships", None) if user_row else None - ) + memberships = getattr(user_row, "organization_memberships", None) if user_row else None member_org_ids = { - membership.organization_id - for membership in (memberships or []) - if membership.organization_id is not None + membership.organization_id for membership in (memberships or []) if membership.organization_id is not None } if organization_id not in member_org_ids: raise HTTPException( @@ -1375,10 +1285,7 @@ async def _check_org_key_limits( data.metadata.get("tpm_limit_type", None) if data.metadata else None ) - if ( - tpm_limit_type != "guaranteed_throughput" - and rpm_limit_type != "guaranteed_throughput" - ): + if tpm_limit_type != "guaranteed_throughput" and rpm_limit_type != "guaranteed_throughput": return # get all organization keys # calculate allocated tpm/rpm limit @@ -1511,23 +1418,15 @@ async def generate_key_fn( # Validate budget values are not negative and are finite numbers # (GHSA-2rv4-xv66-fpjg): float('nan') passes `< 0` because nan < 0 is False. - if data.max_budget is not None and ( - not math.isfinite(data.max_budget) or data.max_budget < 0 - ): + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) - if data.soft_budget is not None and ( - not math.isfinite(data.soft_budget) or data.soft_budget < 0 - ): + if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}" - }, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) if user_custom_key_generate is not None: @@ -1538,9 +1437,7 @@ async def generate_key_fn( decision = result.get("decision", True) message = result.get("message", "Authentication Failed - Custom Auth Rule") if not decision: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, detail=message - ) + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) _check_allowed_routes_caller_permission( allowed_routes=data.allowed_routes, @@ -1575,9 +1472,7 @@ async def generate_key_fn( check_db_only=True, ) except Exception as e: - verbose_proxy_logger.debug( - f"Error getting team object in `/key/generate`: {e}" - ) + verbose_proxy_logger.debug(f"Error getting team object in `/key/generate`: {e}") # For non-admin callers, team must exist (LIT-1884) if not _is_proxy_admin: raise HTTPException( @@ -1617,9 +1512,7 @@ async def generate_key_fn( except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -1752,9 +1645,7 @@ async def generate_service_account_key_fn( check_db_only=True, ) except Exception as e: - verbose_proxy_logger.debug( - f"Error getting team object in `/key/generate`: {e}" - ) + verbose_proxy_logger.debug(f"Error getting team object in `/key/generate`: {e}") team_table = None if team_table is not None: @@ -1781,9 +1672,7 @@ async def generate_service_account_key_fn( ) -def prepare_metadata_fields( - data: BaseModel, non_default_values: dict, existing_metadata: dict -) -> dict: +def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_metadata: dict) -> dict: """ Check LiteLLM_ManagementEndpoint_MetadataFields (proxy/_types.py) for fields that are allowed to be updated """ @@ -1799,8 +1688,7 @@ def prepare_metadata_fields( if existing_value is None: continue if casted_metadata is None or ( - reserved_field in casted_metadata - and casted_metadata[reserved_field] != existing_value + reserved_field in casted_metadata and casted_metadata[reserved_field] != existing_value ): raise HTTPException( status_code=400, @@ -1826,9 +1714,7 @@ def prepare_metadata_fields( except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - {}".format(str(e)) ) non_default_values["metadata"] = encrypt_callback_vars(casted_metadata) @@ -1863,10 +1749,7 @@ async def prepare_key_update_data( value=getattr(data, field), ) for k, v in data_json.items(): - if ( - k in LiteLLM_ManagementEndpoint_MetadataFields - or k in LiteLLM_ManagementEndpoint_MetadataFields_Premium - ): + if k in LiteLLM_ManagementEndpoint_MetadataFields or k in LiteLLM_ManagementEndpoint_MetadataFields_Premium: continue non_default_values[k] = v @@ -1900,9 +1783,7 @@ async def prepare_key_update_data( initialized_windows = [] for window in raw_windows: w = window if isinstance(window, dict) else window.model_dump() - w["reset_at"] = get_budget_reset_time( - budget_duration=w["budget_duration"] - ).isoformat() + w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat() initialized_windows.append(w) non_default_values["budget_limits"] = json.dumps(initialized_windows) else: @@ -1923,13 +1804,8 @@ async def prepare_key_update_data( validate_model_max_budget(non_default_values["model_max_budget"]) # Serialize router_settings to JSON if present - if ( - "router_settings" in non_default_values - and non_default_values["router_settings"] is not None - ): - non_default_values["router_settings"] = safe_dumps( - non_default_values["router_settings"] - ) + if "router_settings" in non_default_values and non_default_values["router_settings"] is not None: + non_default_values["router_settings"] = safe_dumps(non_default_values["router_settings"]) non_default_values = prepare_metadata_fields( data=data, non_default_values=non_default_values, existing_metadata=_metadata @@ -1957,16 +1833,12 @@ async def _handle_update_object_permission( # Add the object_permission_id to data_json if one was created/updated if object_permission_id is not None: data_json["object_permission_id"] = object_permission_id - verbose_proxy_logger.debug( - f"updated object_permission_id: {object_permission_id}" - ) + verbose_proxy_logger.debug(f"updated object_permission_id: {object_permission_id}") return data_json -def is_different_team( - data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken -) -> bool: +def is_different_team(data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken) -> bool: if data.team_id is None: return False if existing_key_row.team_id is None: @@ -1987,9 +1859,7 @@ def _validate_max_budget(max_budget: Optional[float]) -> None: if max_budget is not None and (not math.isfinite(max_budget) or max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {max_budget}"}, ) @@ -2017,9 +1887,7 @@ async def _get_and_validate_existing_key( hashed_token = _hash_token_if_needed(token=token) - existing_key_row = await VerificationTokenRepository( - prisma_client - ).table.find_unique(where={"token": hashed_token}) + existing_key_row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) if existing_key_row is None: raise ProxyException( @@ -2077,14 +1945,12 @@ async def _process_single_key_update( # Check team member permissions if prisma_client is not None: - await ( - TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, - existing_key_row=existing_key_row, - user_api_key_cache=user_api_key_cache, - ) + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=existing_key_row, + user_api_key_cache=user_api_key_cache, ) # Custom key update hook @@ -2140,9 +2006,7 @@ async def _process_single_key_update( ) # Prepare update data - non_default_values = await prepare_key_update_data( - data=update_key_request, existing_key_row=existing_key_row - ) + non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row) # Update key in database if prisma_client is None: @@ -2323,9 +2187,7 @@ async def _validate_update_key_data( # non-budget change means the caller was authorized — skip the redundant # _check_key_admin_access that would otherwise require team/org admin status. _key_is_team_key = getattr(existing_key_row, "team_id", None) is not None - can_skip_admin_check = ( - caller_is_creator or _key_is_team_key - ) and not _is_budget_change + can_skip_admin_check = (caller_is_creator or _key_is_team_key) and not _is_budget_change if (not _is_proxy_admin) and prisma_client is not None and not can_skip_admin_check: hashed_key = existing_key_row.token await _check_key_admin_access( @@ -2333,9 +2195,7 @@ async def _validate_update_key_data( hashed_token=hashed_key, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - route=( - "/key/update (max_budget/spend)" if _is_budget_change else "/key/update" - ), + route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"), ) # Check team limits if key has a team_id (from request or existing key) @@ -2372,12 +2232,8 @@ async def _validate_update_key_data( ) # Validate key against project limits if project_id is being set - _project_id_to_check = getattr(data, "project_id", None) or getattr( - existing_key_row, "project_id", None - ) - if _project_id_to_check is not None and ( - data.models is not None or data.max_budget is not None - ): + _project_id_to_check = getattr(data, "project_id", None) or getattr(existing_key_row, "project_id", None) + if _project_id_to_check is not None and (data.models is not None or data.max_budget is not None): await _check_project_key_limits( project_id=_project_id_to_check, data=data, @@ -2392,11 +2248,7 @@ async def _validate_update_key_data( # IDOR. The check mirrors the membership rule already used on the # `/key/list` filter path in `validate_key_list_check`. _existing_org_id = getattr(existing_key_row, "organization_id", None) - if ( - data.organization_id is not None - and data.organization_id != _existing_org_id - and not _is_proxy_admin - ): + if data.organization_id is not None and data.organization_id != _existing_org_id and not _is_proxy_admin: await _validate_caller_can_assign_key_org( user_api_key_dict=user_api_key_dict, organization_id=data.organization_id, @@ -2461,14 +2313,10 @@ async def _validate_update_key_data( is_proxy_admin=_is_proxy_admin, ) if normalized_object_permission is not None: - data.object_permission = LiteLLM_ObjectPermissionBase( - **normalized_object_permission - ) + data.object_permission = LiteLLM_ObjectPermissionBase(**normalized_object_permission) -@router.post( - "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def update_key_fn( request: Request, @@ -2558,14 +2406,10 @@ async def update_key_fn( try: # Validate budget values are not negative and are finite numbers - if data.max_budget is not None and ( - not math.isfinite(data.max_budget) or data.max_budget < 0 - ): + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) data_json: dict = data.model_dump(exclude_unset=True) @@ -2596,15 +2440,11 @@ async def update_key_fn( decision = result.get("decision", True) message = result.get("message", "Authentication Failed - Custom Auth Rule") if not decision: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, detail=message - ) + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) - non_default_values = await prepare_key_update_data( - data=data, existing_key_row=existing_key_row - ) + non_default_values = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) # Only validate key_alias format if it's actually being changed new_key_alias = non_default_values.get("key_alias", None) @@ -2642,14 +2482,10 @@ async def update_key_fn( from litellm.proxy.proxy_server import spend_counter_cache counter_key = f"spend:key:{_hash_token_if_needed(key)}" - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=data.spend, ttl=60 - ) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=data.spend, ttl=60) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=data.spend, ttl=60 - ) + await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=data.spend, ttl=60) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to update spend counter %s in Redis after key spend update: %s. " @@ -2675,9 +2511,7 @@ async def update_key_fn( # update based on remaining passed in values except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.update_key_fn(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.update_key_fn(): Exception occured - {}".format(str(e)) ) if isinstance(e, HTTPException): raise ProxyException( @@ -2782,9 +2616,7 @@ async def bulk_update_keys( if len(data.keys) > MAX_BATCH_SIZE: raise HTTPException( status_code=400, - detail={ - "error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.keys)} keys." - }, + detail={"error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.keys)} keys."}, ) successful_updates: List[SuccessfulKeyUpdate] = [] @@ -2818,9 +2650,7 @@ async def bulk_update_keys( ) except Exception as e: - verbose_proxy_logger.exception( - f"Failed to update key {key_update_item.key}: {e}" - ) + verbose_proxy_logger.exception(f"Failed to update key {key_update_item.key}: {e}") if isinstance(e, HTTPException): error_detail = e.detail @@ -2952,9 +2782,7 @@ async def bulk_update_team_keys( # `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT` # excludes NULLs, so explicitly OR `false` with `null` to include them. now = datetime.now(timezone.utc) - existing_keys = await VerificationTokenRepository( - prisma_client - ).table.find_many( + existing_keys = await VerificationTokenRepository(prisma_client).table.find_many( where={ "team_id": data.team_id, "AND": [ @@ -2977,9 +2805,7 @@ async def bulk_update_team_keys( if data.key_ids is None or len(data.key_ids) == 0: raise HTTPException( status_code=400, - detail={ - "error": "key_ids must be provided when all_keys_in_team is False" - }, + detail={"error": "key_ids must be provided when all_keys_in_team is False"}, ) # Dedupe by hashed form — duplicates collapse to one update. requested_tokens = [] @@ -2992,9 +2818,7 @@ async def bulk_update_team_keys( seen_hashes.add(h) requested_tokens.append(k) hashed_key_ids.append(h) - existing_keys = await VerificationTokenRepository( - prisma_client - ).table.find_many( + existing_keys = await VerificationTokenRepository(prisma_client).table.find_many( where={"team_id": data.team_id, "token": {"in": hashed_key_ids}} ) @@ -3009,21 +2833,17 @@ async def bulk_update_team_keys( models=[], ) ) - await ( - TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, - existing_key_row=auth_anchor, - user_api_key_cache=user_api_key_cache, - ) + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=auth_anchor, + user_api_key_cache=user_api_key_cache, ) # Block metadata.allowed_passthrough_routes for non-admins — the runtime # route checker reads it from key/team metadata to grant passthrough. - _check_passthrough_routes_caller_permission( - data=data.update_fields, user_api_key_dict=user_api_key_dict - ) + _check_passthrough_routes_caller_permission(data=data.update_fields, user_api_key_dict=user_api_key_dict) if not requested_tokens: raise HTTPException( @@ -3064,15 +2884,11 @@ async def bulk_update_team_keys( existing_key_row=existing_by_token[db_token], ) - successful_updates.append( - SuccessfulKeyUpdate(key=token, key_info=updated_key_info) - ) + successful_updates.append(SuccessfulKeyUpdate(key=token, key_info=updated_key_info)) except Exception as e: # Log the hashed prefix — `token` may be a raw sk-... and ERROR logs persist. - verbose_proxy_logger.exception( - f"Failed to update key {db_token[:12]}... in team {data.team_id}: {e}" - ) + verbose_proxy_logger.exception(f"Failed to update key {db_token[:12]}... in team {data.team_id}: {e}") failed_updates.append( _build_failed_team_key_update( token=token, @@ -3129,9 +2945,7 @@ async def validate_key_team_change( ) # Check if the key's user_id is a member of the team - member_object = _get_user_in_team( - team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id - ) + member_object = _get_user_in_team(team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id) if key.user_id is not None: if not member_object: raise HTTPException( @@ -3161,9 +2975,7 @@ async def validate_key_team_change( ) -@router.post( - "/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def delete_key_fn( data: KeyRequest, @@ -3207,9 +3019,7 @@ async def delete_key_fn( litellm_changed_by = None ## only allow user to delete keys they own - verbose_proxy_logger.debug( - f"user_api_key_dict.user_role: {user_api_key_dict.user_role}" - ) + verbose_proxy_logger.debug(f"user_api_key_dict.user_role: {user_api_key_dict.user_role}") num_keys_to_be_deleted = 0 deleted_keys = [] @@ -3271,9 +3081,7 @@ async def delete_key_fn( return {"deleted_keys": deleted_keys} except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.delete_key_fn(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.delete_key_fn(): Exception occured - {}".format(str(e)) ) raise handle_exception_on_proxy(e) @@ -3285,8 +3093,7 @@ async def _get_model_max_budget_current_spend( user_api_key_cache: UserApiKeyCache, ) -> float: virtual_key_model_spend_cache_key = ( - f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:" - f"{api_key_hash}:{model}:{budget_config.budget_duration}" + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}" ) current_spend: float | None = await user_api_key_cache.async_get_cache( key=virtual_key_model_spend_cache_key, @@ -3380,9 +3187,7 @@ async def info_key_fn_v2( # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query = list(data.keys) if data.keys else [] if data.key_aliases: - alias_rows = await VerificationTokenRepository( - prisma_client - ).table.find_many( + alias_rows = await VerificationTokenRepository(prisma_client).table.find_many( where={"key_alias": {"in": data.key_aliases}}, include={"litellm_budget_table": True}, ) @@ -3392,9 +3197,7 @@ async def info_key_fn_v2( if not tokens_to_query: return {"key": data.keys, "info": []} - key_info = await prisma_client.get_data( - token=tokens_to_query, table_name="key", query_type="find_all" - ) + key_info = await prisma_client.get_data(token=tokens_to_query, table_name="key", query_type="find_all") if not key_info: return {"key": data.keys, "info": []} @@ -3430,14 +3233,10 @@ async def info_key_fn_v2( raise handle_exception_on_proxy(e) -@router.get( - "/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) +@router.get("/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def info_key_fn( - key: Optional[str] = fastapi.Query( - default=None, description="Key in the request parameters" - ), + key: Optional[str] = fastapi.Query(default=None, description="Key in the request parameters"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -3538,9 +3337,7 @@ def _check_model_access_group( return True for model in models: - if llm_router._is_model_access_group_for_wildcard_route( - model_access_group=model - ): + if llm_router._is_model_access_group_for_wildcard_route(model_access_group=model): if not premium_user: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -3555,9 +3352,7 @@ def _check_model_access_group( async def generate_key_helper_fn( - request_type: Literal[ - "user", "key" - ], # identifies if this request is from /user/new or /key/generate + request_type: Literal["user", "key"], # identifies if this request is from /user/new or /key/generate duration: Optional[str] = None, models: list = [], aliases: dict = {}, @@ -3566,9 +3361,7 @@ async def generate_key_helper_fn( key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key key_budget_duration: Optional[str] = None, budget_id: Optional[float] = None, # budget id <-> LiteLLM_BudgetTable - soft_budget: Optional[ - float - ] = None, # soft_budget is used to set soft Budgets Per user + soft_budget: Optional[float] = None, # soft_budget is used to set soft Budgets Per user max_budget: Optional[float] = None, # max_budget is used to Budget Per user blocked: Optional[bool] = None, budget_duration: Optional[str] = None, # max_budget is used to Budget Per user @@ -3607,9 +3400,7 @@ async def generate_key_helper_fn( updated_by: Optional[str] = None, allowed_routes: Optional[list] = None, sso_user_id: Optional[str] = None, - object_permission_id: Optional[ - str - ] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable + object_permission_id: Optional[str] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable object_permission: Optional[LiteLLM_ObjectPermissionBase] = None, auto_rotate: Optional[bool] = None, rotation_interval: Optional[str] = None, @@ -3620,9 +3411,7 @@ async def generate_key_helper_fn( from litellm.proxy.proxy_server import premium_user, prisma_client if prisma_client is None: - raise Exception( - "Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys " - ) + raise Exception("Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys ") if token is None: if key is not None: @@ -3653,18 +3442,14 @@ async def generate_key_helper_fn( initialized_windows = [] for window in budget_limits: w = dict(window) if not isinstance(window, dict) else {**window} - w["reset_at"] = get_budget_reset_time( - budget_duration=w["budget_duration"] - ).isoformat() + w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat() initialized_windows.append(w) budget_limits_json = json.dumps(initialized_windows) aliases_json = json.dumps(aliases) config_json = json.dumps(config) permissions_json = json.dumps(permissions) - router_settings_json = ( - safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) - ) + router_settings_json = safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) # Add model_rpm_limit and model_tpm_limit to metadata if model_rpm_limit is not None: @@ -3775,19 +3560,12 @@ async def generate_key_helper_fn( if isinstance(saved_token["metadata"], str): saved_token["metadata"] = json.loads(saved_token["metadata"]) if isinstance(saved_token["permissions"], str): - if ( - "get_spend_routes" in saved_token["permissions"] - and premium_user is not True - ): - raise ValueError( - "get_spend_routes permission is only available for LiteLLM Enterprise users" - ) + if "get_spend_routes" in saved_token["permissions"] and premium_user is not True: + raise ValueError("get_spend_routes permission is only available for LiteLLM Enterprise users") saved_token["permissions"] = json.loads(saved_token["permissions"]) if isinstance(saved_token["model_max_budget"], str): - saved_token["model_max_budget"] = json.loads( - saved_token["model_max_budget"] - ) + saved_token["model_max_budget"] = json.loads(saved_token["model_max_budget"]) router_settings = cast(Optional[dict], saved_token.get("router_settings")) if router_settings is not None and isinstance(router_settings, str): try: @@ -3796,19 +3574,13 @@ async def generate_key_helper_fn( # If it's not valid JSON/YAML, keep as is or set to empty dict saved_token["router_settings"] = {} - if saved_token.get("expires", None) is not None and isinstance( - saved_token["expires"], datetime - ): + if saved_token.get("expires", None) is not None and isinstance(saved_token["expires"], datetime): saved_token["expires"] = saved_token["expires"].isoformat() if prisma_client is not None: - if ( - table_name is None or table_name == "user" - ): # do not auto-create users for `/key/generate` + if table_name is None or table_name == "user": # do not auto-create users for `/key/generate` ## CREATE USER (If necessary) if query_type == "insert_data": - user_row = await prisma_client.insert_data( - data=user_data, table_name="user" - ) + user_row = await prisma_client.insert_data(data=user_data, table_name="user") if user_row is None: raise Exception("Failed to create user") @@ -3829,22 +3601,16 @@ async def generate_key_helper_fn( ## CREATE KEY verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data) - create_key_response = await prisma_client.insert_data( - data=key_data, table_name="key" - ) + create_key_response = await prisma_client.insert_data(data=key_data, table_name="key") key_data["token_id"] = getattr(create_key_response, "token", None) - key_data["litellm_budget_table"] = getattr( - create_key_response, "litellm_budget_table", None - ) + key_data["litellm_budget_table"] = getattr(create_key_response, "litellm_budget_table", None) key_data["created_at"] = getattr(create_key_response, "created_at", None) key_data["updated_at"] = getattr(create_key_response, "updated_at", None) # Deserialize router_settings from JSON string to dict for response router_settings_value = key_data.get("router_settings") - if router_settings_value is not None and isinstance( - router_settings_value, str - ): + if router_settings_value is not None and isinstance(router_settings_value, str): try: key_data["router_settings"] = yaml.safe_load(router_settings_value) except yaml.YAMLError: @@ -3852,9 +3618,7 @@ async def generate_key_helper_fn( key_data["router_settings"] = {} except Exception as e: verbose_proxy_logger.error( - "litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): @@ -3889,13 +3653,8 @@ async def _team_key_deletion_check( user_api_key_cache=user_api_key_cache, check_db_only=True, ) - if ( - litellm.key_generation_settings is not None - and "team_key_generation" in litellm.key_generation_settings - ): - _team_key_generation = litellm.key_generation_settings[ - "team_key_generation" - ] + if litellm.key_generation_settings is not None and "team_key_generation" in litellm.key_generation_settings: + _team_key_generation = litellm.key_generation_settings["team_key_generation"] else: _team_key_generation = TeamUIKeyGenerationConfig( allowed_team_member_roles=["admin", "user"], @@ -3912,9 +3671,7 @@ async def _team_key_deletion_check( else: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail={ - "error": f"Team not found in db, and user not proxy admin. Team id = {key_info.team_id}" - }, + detail={"error": f"Team not found in db, and user not proxy admin. Team id = {key_info.team_id}"}, ) return False @@ -3976,10 +3733,7 @@ async def can_modify_verification_token( return True # Check if the key belongs to the user (they own it) - if ( - key_info.user_id is not None - and key_info.user_id == user_api_key_dict.user_id - ): + if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: return True # Not team admin and doesn't own the key @@ -4023,11 +3777,9 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[ - LiteLLM_VerificationToken - ] = await VerificationTokenRepository(prisma_client).table.find_many( - where={"token": {"in": tokens}} - ) + _keys_being_deleted: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"token": {"in": tokens}}) if len(_keys_being_deleted) == 0: raise HTTPException( @@ -4050,9 +3802,7 @@ async def delete_verification_tokens( else: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "You are not authorized to delete this key" - }, + detail={"error": "You are not authorized to delete this key"}, ) await _persist_deleted_verification_tokens( keys=authorized_keys, @@ -4064,28 +3814,19 @@ async def delete_verification_tokens( if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: deleted_tokens = await prisma_client.delete_data(tokens=tokens) if deleted_tokens is not None and len(deleted_tokens) != len(tokens): - failed_tokens = [ - token for token in tokens if token not in deleted_tokens - ] + failed_tokens = [token for token in tokens if token not in deleted_tokens] else: - deletion_tasks = [ - prisma_client.delete_data(tokens=[key.token]) - for key in authorized_keys - ] + deletion_tasks = [prisma_client.delete_data(tokens=[key.token]) for key in authorized_keys] await asyncio.gather(*deletion_tasks) deleted_tokens = [key.token for key in authorized_keys] if len(deleted_tokens) != len(tokens): - failed_tokens = [ - token for token in tokens if token not in deleted_tokens - ] + failed_tokens = [token for token in tokens if token not in deleted_tokens] else: raise Exception("DB not connected. prisma_client is None") except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.delete_verification_tokens(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.delete_verification_tokens(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) raise e @@ -4162,9 +3903,7 @@ async def _save_deleted_verification_token_records( """Save deleted verification token records to the database.""" if not records: return - await DeletedVerificationTokenRepository(prisma_client).table.create_many( - data=records - ) + await DeletedVerificationTokenRepository(prisma_client).table.create_many(data=records) async def _persist_deleted_verification_tokens( @@ -4192,9 +3931,9 @@ async def delete_key_aliases( user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: - _keys_being_deleted = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"key_alias": {"in": key_aliases}}) + _keys_being_deleted = await VerificationTokenRepository(prisma_client).table.find_many( + where={"key_alias": {"in": key_aliases}} + ) tokens = [key.token for key in _keys_being_deleted] return await delete_verification_tokens( @@ -4235,9 +3974,7 @@ async def _rotate_master_key( # 2. process model table if models: decrypted_models = proxy_config.decrypt_model_list_from_db(new_models=models) - verbose_proxy_logger.debug( - "ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models) - ) + verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models)) new_models = [] for model in decrypted_models: new_model = await _add_model_to_db( @@ -4295,9 +4032,7 @@ async def _rotate_master_key( new_master_key=new_master_key, ) except Exception as e: - verbose_proxy_logger.warning( - "Failed to rotate MCP server credentials: %s", str(e) - ) + verbose_proxy_logger.warning("Failed to rotate MCP server credentials: %s", str(e)) # 4b. process MCP user-scoped credentials table (BYOK + OAuth2 tokens) try: @@ -4306,9 +4041,7 @@ async def _rotate_master_key( new_master_key=new_master_key, ) except Exception as e: - verbose_proxy_logger.warning( - "Failed to rotate MCP user credentials: %s", str(e) - ) + verbose_proxy_logger.warning("Failed to rotate MCP user credentials: %s", str(e)) # 4c. process MCP per-user environment variables table try: @@ -4352,14 +4085,10 @@ async def _rotate_master_key( }, ) except Exception as e: - verbose_proxy_logger.error( - f"Failed to re-encrypt credential {cred.credential_name}: {str(e)}" - ) + verbose_proxy_logger.error(f"Failed to re-encrypt credential {cred.credential_name}: {str(e)}") # Continue with next credential instead of failing entire rotation continue - verbose_proxy_logger.debug( - f"Successfully re-encrypted {len(credentials)} credentials with new master key" - ) + verbose_proxy_logger.debug(f"Successfully re-encrypted {len(credentials)} credentials with new master key") async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: @@ -4396,9 +4125,7 @@ async def _insert_deprecated_key( new_token_hash: Hash of the new replacement key grace_period: Duration string (e.g. "24h", "2d") or None/empty for immediate revoke """ - grace_period_value = grace_period or os.getenv( - "LITELLM_KEY_ROTATION_GRACE_PERIOD", "" - ) + grace_period_value = grace_period or os.getenv("LITELLM_KEY_ROTATION_GRACE_PERIOD", "") if not grace_period_value: return @@ -4489,9 +4216,7 @@ async def _execute_virtual_key_regeneration( if data is not None: # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) - non_default_values = await prepare_key_update_data( - data=data, existing_key_row=key_in_db - ) + non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db) # Only validate key_alias format if it's actually being changed new_key_alias = non_default_values.get("key_alias") if new_key_alias != key_in_db.key_alias: @@ -4713,14 +4438,12 @@ async def regenerate_key_fn( ) # check if user has permission to regenerate key - await ( - TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_REGENERATE, - prisma_client=prisma_client, - existing_key_row=_key_in_db, - user_api_key_cache=user_api_key_cache, - ) + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_REGENERATE, + prisma_client=prisma_client, + existing_key_row=_key_in_db, + user_api_key_cache=user_api_key_cache, ) # check if user has ownership permission to regenerate key @@ -4818,9 +4541,7 @@ async def _check_proxy_or_team_admin_for_key( ) -def _validate_reset_spend_value( - reset_to: Any, key_in_db: LiteLLM_VerificationToken -) -> float: +def _validate_reset_spend_value(reset_to: Any, key_in_db: LiteLLM_VerificationToken) -> float: if not isinstance(reset_to, (int, float)): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -4839,9 +4560,7 @@ def _validate_reset_spend_value( if reset_to > current_spend: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": f"reset_to ({reset_to}) must be <= current spend ({current_spend})" - }, + detail={"error": f"reset_to ({reset_to}) must be <= current spend ({current_spend})"}, ) max_budget = key_in_db.max_budget @@ -4937,14 +4656,10 @@ async def reset_key_spend_fn( from litellm.proxy.proxy_server import spend_counter_cache _counter_key = f"spend:key:{hashed_api_key}" - spend_counter_cache.in_memory_cache.set_cache( - key=_counter_key, value=reset_to, ttl=60 - ) + spend_counter_cache.in_memory_cache.set_cache(key=_counter_key, value=reset_to, ttl=60) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_cache( - key=_counter_key, value=reset_to, ttl=60 - ) + await spend_counter_cache.redis_cache.async_set_cache(key=_counter_key, value=reset_to, ttl=60) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to update spend counter %s in Redis: %s. " @@ -4989,9 +4704,7 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[BaseModel] = await UserRepository( - prisma_client - ).table.find_unique( + complete_user_info_db_obj: Optional[BaseModel] = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) @@ -5026,14 +4739,9 @@ async def validate_key_list_check( ) if organization_id: - if ( - complete_user_info.organization_memberships is None - or organization_id - not in [ - membership.organization_id - for membership in complete_user_info.organization_memberships - ] - ): + if complete_user_info.organization_memberships is None or organization_id not in [ + membership.organization_id for membership in complete_user_info.organization_memberships + ]: raise ProxyException( message="You are not authorized to check this organization's keys", type=ProxyErrorTypes.bad_request_error, @@ -5043,9 +4751,7 @@ async def validate_key_list_check( if key_hash: try: - key_info = await VerificationTokenRepository( - prisma_client - ).table.find_unique( + key_info = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": key_hash}, ) except Exception: @@ -5078,9 +4784,9 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[List[BaseModel]] = await TeamRepository( - prisma_client - ).table.find_many(where={"team_id": {"in": complete_user_info.teams}}) + teams: Optional[List[BaseModel]] = await TeamRepository(prisma_client).table.find_many( + where={"team_id": {"in": complete_user_info.teams}} + ) if teams is None: return [] @@ -5093,9 +4799,7 @@ def _get_admin_team_ids_from_objects( ) -> List[str]: """Filter team objects to those where the user is an admin.""" return [ - team.team_id - for team in team_objects - if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team) + team.team_id for team in team_objects if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team) ] @@ -5175,36 +4879,24 @@ async def list_keys( description="Filter keys by user ID. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.", ), team_id: Optional[str] = Query(None, description="Filter keys by team ID"), - organization_id: Optional[str] = Query( - None, description="Filter keys by organization ID" - ), + organization_id: Optional[str] = Query(None, description="Filter keys by organization ID"), key_hash: Optional[str] = Query(None, description="Filter keys by key hash"), key_alias: Optional[str] = Query( None, description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.", ), return_full_object: bool = Query(False, description="Return full key object"), - include_team_keys: bool = Query( - False, description="Include all keys for teams that user is an admin of." - ), - include_created_by_keys: bool = Query( - False, description="Include keys created by the user" - ), + include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."), + include_created_by_keys: bool = Query(False, description="Include keys created by the user"), sort_by: Optional[str] = Query( default=None, description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), - expand: Optional[List[str]] = Query( - None, description="Expand related objects (e.g. 'user')" - ), - status: Optional[str] = Query( - None, description="Filter by status (e.g. 'deleted')" - ), + expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), + status: Optional[str] = Query(None, description="Filter by status (e.g. 'deleted')"), project_id: Optional[str] = Query(None, description="Filter keys by project ID"), - access_group_id: Optional[str] = Query( - None, description="Filter keys by access group ID" - ), + access_group_id: Optional[str] = Query(None, description="Filter keys by access group ID"), substring_matching: bool = Query( False, description="If true (proxy admins only), match user_id/key_alias as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id/key_alias filter must never return another user's keys.", @@ -5241,9 +4933,7 @@ async def list_keys( if status is not None and status != "deleted": raise HTTPException( status_code=400, - detail={ - "error": "Invalid status value. Currently only 'deleted' is supported." - }, + detail={"error": "Invalid status value. Currently only 'deleted' is supported."}, ) complete_user_info = await validate_key_list_check( @@ -5279,11 +4969,9 @@ async def list_keys( # Non-admin members with /key/list permission get full team-key # visibility for that team — matching the UI contract that # granting this permission lets them see all keys within the team. - list_permission_team_ids = ( - _get_team_ids_with_key_list_permission_from_objects( - user_api_key_dict=user_api_key_dict, - team_objects=team_objects, - ) + list_permission_team_ids = _get_team_ids_with_key_list_permission_from_objects( + user_api_key_dict=user_api_key_dict, + team_objects=team_objects, ) if list_permission_team_ids: admin_team_ids = list({*admin_team_ids, *list_permission_team_ids}) @@ -5340,9 +5028,7 @@ async def list_keys( message=getattr(e, "detail", f"error({str(e)})"), type=ProxyErrorTypes.internal_server_error, param=getattr(e, "param", "None"), - code=getattr( - e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR - ), + code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR), ) elif isinstance(e, ProxyException): raise e @@ -5370,16 +5056,12 @@ async def _apply_non_admin_alias_scope( # Look up the user's teams from the user table user_teams: List[str] = [] if user_api_key_dict.user_id: - user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_api_key_dict.user_id} - ) + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_api_key_dict.user_id}) if user_row is not None: user_teams = getattr(user_row, "teams", []) or [] if user_teams: - team_placeholders = ", ".join( - f"${len(query_params) + i + 1}" for i in range(len(user_teams)) - ) + team_placeholders = ", ".join(f"${len(query_params) + i + 1}" for i in range(len(user_teams))) query_params.extend(user_teams) scope_conditions.append(f"team_id IN ({team_placeholders})") @@ -5400,12 +5082,8 @@ async def key_aliases( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), page: int = Query(1, ge=1, description="Page number"), size: int = Query(50, ge=1, le=100, description="Page size"), - search: Optional[str] = Query( - None, description="Search key aliases (case-insensitive partial match)" - ), - team_id: Optional[str] = Query( - None, description="Filter aliases to keys belonging to this team" - ), + search: Optional[str] = Query(None, description="Search key aliases (case-insensitive partial match)"), + team_id: Optional[str] = Query(None, description="Filter aliases to keys belonging to this team"), ) -> Dict[str, Any]: """ Lists key aliases with pagination and optional search. @@ -5450,9 +5128,7 @@ async def key_aliases( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, ] if not is_proxy_admin: - await _apply_non_admin_alias_scope( - user_api_key_dict, prisma_client, query_params, where_parts - ) + await _apply_non_admin_alias_scope(user_api_key_dict, prisma_client, query_params, where_parts) if search: query_params.append(f"%{search}%") @@ -5479,9 +5155,7 @@ async def key_aliases( f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) - aliases: List[str] = [ - row["key_alias"] for row in alias_rows if row.get("key_alias") - ] + aliases: List[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages = -(-total_count // size) if total_count > 0 else 0 verbose_proxy_logger.debug( @@ -5516,9 +5190,7 @@ async def key_aliases( ) -def _validate_sort_params( - sort_by: Optional[str], sort_order: str -) -> Optional[Dict[str, str]]: +def _validate_sort_params(sort_by: Optional[str], sort_order: str) -> Optional[Dict[str, str]]: order_by: Dict[str, str] = {} if sort_by is None: @@ -5535,9 +5207,7 @@ def _validate_sort_params( if sort_by not in valid_columns: raise HTTPException( status_code=400, - detail={ - "error": f"Invalid sort column. Must be one of: {', '.join(valid_columns)}" - }, + detail={"error": f"Invalid sort column. Must be one of: {', '.join(valid_columns)}"}, ) # Validate sort_order @@ -5632,9 +5302,7 @@ def _build_key_filter_conditions( ) else: # User is not a member of any team, only show non-team created_by keys - or_conditions.append( - {"AND": [{"created_by": user_id}, {"team_id": None}]} - ) + or_conditions.append({"AND": [{"created_by": user_id}, {"team_id": None}]}) else: # No team membership info provided (backward compatibility for # direct _list_key_helper callers like Prometheus) @@ -5647,9 +5315,7 @@ def _build_key_filter_conditions( # Add condition for member team service accounts (members only see keys with user_id=NULL) if member_team_ids: # Exclude teams where user is already admin (those are covered above with full visibility) - member_only_team_ids = [ - tid for tid in member_team_ids if tid not in (admin_team_ids or []) - ] + member_only_team_ids = [tid for tid in member_team_ids if tid not in (admin_team_ids or [])] if member_only_team_ids: or_conditions.append( { @@ -5690,9 +5356,7 @@ async def _list_key_helper( key_hash: Optional[str], exclude_team_id: Optional[str] = None, return_full_object: bool = False, - admin_team_ids: Optional[ - List[str] - ] = None, # New parameter for teams where user is admin + admin_team_ids: Optional[List[str]] = None, # New parameter for teams where user is admin member_team_ids: Optional[ List[str] ] = None, # Team IDs where user is a member (any role) - for service account visibility @@ -5748,9 +5412,7 @@ async def _list_key_helper( verbose_proxy_logger.debug(f"Pagination: skip={skip}, take={size}") order_by: Optional[Dict[str, str]] = ( - _validate_sort_params(sort_by, sort_order) - if sort_by is not None and isinstance(sort_by, str) - else None + _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) # Determine which table to query based on status @@ -5791,9 +5453,7 @@ async def _list_key_helper( # Get total count of keys if use_deleted_table: - total_count = await DeletedVerificationTokenRepository( - prisma_client - ).table.count( + total_count = await DeletedVerificationTokenRepository(prisma_client).table.count( where=where # type: ignore ) else: @@ -5813,9 +5473,7 @@ async def _list_key_helper( created_by_ids = [key.created_by for key in keys if key.created_by] all_ids = list(set(user_ids + created_by_ids)) # Remove duplicates if all_ids: - users = await UserRepository(prisma_client).table.find_many( - where={"user_id": {"in": all_ids}} - ) + users = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": all_ids}}) user_map = {user.user_id: user for user in users} # Prepare response @@ -5871,9 +5529,7 @@ def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]: return { "OR": [ {"team_id": None}, # Include records where team_id is null - { - "team_id": {"not": UI_SESSION_TOKEN_TEAM_ID} - }, # Include records where team_id != UI_SESSION_TOKEN_TEAM_ID + {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, # Include records where team_id != UI_SESSION_TOKEN_TEAM_ID ] } @@ -5900,9 +5556,7 @@ async def _check_key_admin_access( return # Look up the target key to find its team - target_key_row = await VerificationTokenRepository(prisma_client).table.find_unique( - where={"token": hashed_token} - ) + target_key_row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) if target_key_row is None: raise HTTPException( status_code=404, @@ -5918,13 +5572,9 @@ async def _check_key_admin_access( check_db_only=True, ) if team_obj is not None: - if _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): return - if await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj): return raise HTTPException( @@ -5936,9 +5586,7 @@ async def _check_key_admin_access( ) -@router.post( - "/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def block_key( data: BlockKeyRequest, @@ -6004,9 +5652,7 @@ async def block_key( ) # Check if the key exists before trying to block it - existing_record = await VerificationTokenRepository( - prisma_client - ).table.find_unique(where={"token": hashed_token}) + existing_record = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6051,9 +5697,7 @@ async def block_key( return record -@router.post( - "/key/unblock", tags=["key management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/key/unblock", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def unblock_key( data: BlockKeyRequest, @@ -6119,9 +5763,7 @@ async def unblock_key( ) # Check if the key exists before trying to unblock it - existing_record = await VerificationTokenRepository( - prisma_client - ).table.find_unique(where={"token": hashed_token}) + existing_record = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -6331,9 +5973,7 @@ async def test_key_logging( request=request, ) data["mock_response"] = "test response" - await litellm.acompletion( - **data - ) # make mock completion call to trigger key based callbacks + await litellm.acompletion(**data) # make mock completion call to trigger key based callbacks except Exception as e: return LoggingCallbackStatus( callbacks=logging_callbacks, @@ -6341,9 +5981,7 @@ async def test_key_logging( details=f"Logging test failed: {str(e)}", ) - await asyncio.sleep( - 2 - ) # wait for callbacks to run, callbacks use batching so wait for the flush event + await asyncio.sleep(2) # wait for callbacks to run, callbacks use batching so wait for the flush event # Check if any logger exceptions were triggered log_contents = log_capture_string.getvalue() @@ -6416,9 +6054,7 @@ async def _enforce_unique_key_alias( # Exclude the current key from the uniqueness check where_clause["NOT"] = {"token": existing_key_token} - existing_key = await VerificationTokenRepository( - prisma_client - ).table.find_first(where=where_clause) + existing_key = await VerificationTokenRepository(prisma_client).table.find_first(where=where_clause) if existing_key is not None: raise ProxyException( message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", @@ -6451,11 +6087,7 @@ def validate_model_max_budget(model_max_budget: Optional[Dict]) -> None: assert isinstance(_model, str) # Normalize to dict (Pydantic may already parse nested values as BudgetConfig) - _info = ( - _budget_info.model_dump() - if hasattr(_budget_info, "model_dump") - else dict(_budget_info) - ) + _info = _budget_info.model_dump() if hasattr(_budget_info, "model_dump") else dict(_budget_info) # /CRUD endpoints can pass budget_limit as a string, so we need to convert it to a float if "budget_limit" in _info: _info["budget_limit"] = float(_info["budget_limit"]) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 4ab308990b5..ab9d04a4eb4 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -75,9 +75,7 @@ TEMPORARY_MCP_SERVER_TTL_SECONDS = 300 TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX = "litellm:mcp:temporary_server" -def does_mcp_server_exist( - mcp_server_records: Iterable[Any], mcp_server_id: str -) -> bool: +def does_mcp_server_exist(mcp_server_records: Iterable[Any], mcp_server_id: str) -> bool: """ Check if the mcp server with the given id exists in the iterable of mcp servers. @@ -205,13 +203,9 @@ if MCP_AVAILABLE: if validation_result.is_valid: continue - error_messages_text = ( - f"Invalid MCP tool prefix '{value}' provided via {field_name}" - ) + error_messages_text = f"Invalid MCP tool prefix '{value}' provided via {field_name}" if validation_result.warnings: - error_messages_text = ( - error_messages_text + "\n" + "\n".join(validation_result.warnings) - ) + error_messages_text = error_messages_text + "\n" + "\n".join(validation_result.warnings) raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": error_messages_text}, @@ -229,9 +223,7 @@ if MCP_AVAILABLE: general_settings as proxy_general_settings, ) - required_fields: Optional[List[str]] = proxy_general_settings.get( - "mcp_required_fields" - ) + required_fields: Optional[List[str]] = proxy_general_settings.get("mcp_required_fields") if not required_fields: return @@ -291,9 +283,7 @@ if MCP_AVAILABLE: return server.server_name return server.server_id - def _build_mcp_registry_entry_for_server( - server: MCPServer, base_url: str - ) -> Dict[str, Any]: + def _build_mcp_registry_entry_for_server(server: MCPServer, base_url: str) -> Dict[str, Any]: server_name = _build_mcp_registry_server_name(server) title = server_name description = server_name @@ -339,11 +329,7 @@ if MCP_AVAILABLE: return now = datetime.utcnow() - expired_ids = [ - server_id - for server_id, entry in _temporary_mcp_servers.items() - if entry.expires_at <= now - ] + expired_ids = [server_id for server_id, entry in _temporary_mcp_servers.items() if entry.expires_at <= now] for server_id in expired_ids: _temporary_mcp_servers.pop(server_id, None) @@ -357,9 +343,7 @@ if MCP_AVAILABLE: ) return server - async def _cache_temporary_mcp_server_in_redis( - server: MCPServer, ttl_seconds: int - ) -> None: + async def _cache_temporary_mcp_server_in_redis(server: MCPServer, ttl_seconds: int) -> None: """ Best-effort write-through to Redis so temporary MCP OAuth sessions are shared across proxy instances. Keep local in-memory cache as fallback. @@ -375,15 +359,11 @@ if MCP_AVAILABLE: try: encrypted_payload = encrypt_value_helper(payload_json) except Exception as e: - verbose_proxy_logger.debug( - f"Failed to encrypt temporary MCP server payload for Redis cache: {str(e)}" - ) + verbose_proxy_logger.debug(f"Failed to encrypt temporary MCP server payload for Redis cache: {str(e)}") return if not isinstance(encrypted_payload, str): - verbose_proxy_logger.debug( - "Encrypted temporary MCP payload is not a string; skipping Redis cache write" - ) + verbose_proxy_logger.debug("Encrypted temporary MCP payload is not a string; skipping Redis cache write") return try: @@ -393,9 +373,7 @@ if MCP_AVAILABLE: ttl=max(1, ttl_seconds), ) except Exception as e: - verbose_proxy_logger.debug( - f"Failed to write temporary MCP server to Redis cache: {str(e)}" - ) + verbose_proxy_logger.debug(f"Failed to write temporary MCP server to Redis cache: {str(e)}") async def _get_temporary_mcp_server_from_redis( server_id: str, @@ -417,9 +395,7 @@ if MCP_AVAILABLE: key=f"{TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX}:{server_id}" ) except Exception as e: - verbose_proxy_logger.debug( - f"Failed reading temporary MCP server from Redis cache: {str(e)}" - ) + verbose_proxy_logger.debug(f"Failed reading temporary MCP server from Redis cache: {str(e)}") return None if not isinstance(cached_server, str): @@ -438,9 +414,7 @@ if MCP_AVAILABLE: try: loaded = json.loads(decrypted_json) except Exception as e: - verbose_proxy_logger.debug( - f"Invalid decrypted temporary MCP payload in Redis cache: {str(e)}" - ) + verbose_proxy_logger.debug(f"Invalid decrypted temporary MCP payload in Redis cache: {str(e)}") return None if not isinstance(loaded, dict): return None @@ -449,9 +423,7 @@ if MCP_AVAILABLE: try: return MCPServer(**payload_dict) except Exception as e: - verbose_proxy_logger.debug( - f"Invalid temporary MCP server payload in Redis cache: {str(e)}" - ) + verbose_proxy_logger.debug(f"Invalid temporary MCP server payload in Redis cache: {str(e)}") return None async def get_cached_temporary_mcp_server( @@ -618,9 +590,7 @@ if MCP_AVAILABLE: if not payload.server_id or payload.credentials: return payload - existing_server = global_mcp_server_manager.get_mcp_server_by_id( - payload.server_id - ) + existing_server = global_mcp_server_manager.get_mcp_server_by_id(payload.server_id) if existing_server is None: return payload @@ -635,17 +605,11 @@ if MCP_AVAILABLE: inherited_credentials["scopes"] = existing_server.scopes # AWS SigV4 fields if existing_server.aws_access_key_id: - inherited_credentials["aws_access_key_id"] = ( - existing_server.aws_access_key_id - ) + inherited_credentials["aws_access_key_id"] = existing_server.aws_access_key_id if existing_server.aws_secret_access_key: - inherited_credentials["aws_secret_access_key"] = ( - existing_server.aws_secret_access_key - ) + inherited_credentials["aws_secret_access_key"] = existing_server.aws_secret_access_key if existing_server.aws_session_token: - inherited_credentials["aws_session_token"] = ( - existing_server.aws_session_token - ) + inherited_credentials["aws_session_token"] = existing_server.aws_session_token if existing_server.aws_region_name: inherited_credentials["aws_region_name"] = existing_server.aws_region_name if existing_server.aws_service_name: @@ -767,10 +731,7 @@ if MCP_AVAILABLE: try: mcp_servers = await MCPServerRepository(prisma_client).table.find_many() for server in mcp_servers: - if ( - hasattr(server, "mcp_access_groups") - and server.mcp_access_groups - ): + if hasattr(server, "mcp_access_groups") and server.mcp_access_groups: access_groups.update(server.mcp_access_groups) except Exception as e: verbose_proxy_logger.debug(f"Error getting MCP access groups: {e}") @@ -814,9 +775,7 @@ if MCP_AVAILABLE: registry_servers.append({"server": _build_builtin_registry_entry(base_url)}) # Centralized IP-based filtering: external callers only see public servers - registered_servers = list( - global_mcp_server_manager.get_filtered_registry(client_ip).values() - ) + registered_servers = list(global_mcp_server_manager.get_filtered_registry(client_ip).values()) registered_servers.sort(key=_build_mcp_registry_server_name) @@ -876,9 +835,7 @@ if MCP_AVAILABLE: for server_id in all_allowed_ids: server = global_mcp_server_manager.get_mcp_server_by_id(server_id) if server is not None: - mcp_server_table = global_mcp_server_manager._build_mcp_server_table( - server - ) + mcp_server_table = global_mcp_server_manager._build_mcp_server_table(server) servers.append(mcp_server_table) return _redact_mcp_credentials_list(servers) @@ -895,17 +852,12 @@ if MCP_AVAILABLE: aligned with the cards actually rendered: an admin in view_all mode sees every server even when their key carries no per-server MCP grant. """ - if ( - _get_user_mcp_management_mode() == "view_all" - and not _is_restricted_virtual_key_request(user_api_key_dict) - ): + if _get_user_mcp_management_mode() == "view_all" and not _is_restricted_virtual_key_request(user_api_key_dict): return await global_mcp_server_manager.get_all_mcp_servers_unfiltered() aggregated: Dict[str, LiteLLM_MCPServerTable] = {} for auth_context in await build_effective_auth_contexts(user_api_key_dict): - for server in await global_mcp_server_manager.get_all_allowed_mcp_servers( - user_api_key_auth=auth_context - ): + for server in await global_mcp_server_manager.get_all_allowed_mcp_servers(user_api_key_auth=auth_context): aggregated.setdefault(server.server_id, server) return list(aggregated.values()) @@ -937,9 +889,7 @@ if MCP_AVAILABLE: """ # If team_id is provided, return team-scoped servers + allow_all_keys servers - is_restricted_virtual_key = _is_restricted_virtual_key_request( - user_api_key_dict - ) + is_restricted_virtual_key = _is_restricted_virtual_key_request(user_api_key_dict) if team_id is not None and isinstance(team_id, str) and team_id.strip(): # Restricted virtual keys must not use the team_id filter to # bypass their own access limitations. @@ -976,9 +926,7 @@ if MCP_AVAILABLE: detail="You do not have permission to view MCP servers for this team.", ) - redacted_mcp_servers = await _get_team_scoped_mcp_server_list( - sanitized_team_id - ) + redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id) else: servers = await _resolve_accessible_mcp_servers(user_api_key_dict) redacted_mcp_servers = _redact_mcp_credentials_list(servers) @@ -996,15 +944,9 @@ if MCP_AVAILABLE: user_id = user_api_key_dict.user_id or "" if user_id and _byok_prisma_client is not None: - byok_server_ids = [ - s.server_id - for s in redacted_mcp_servers - if getattr(s, "is_byok", False) - ] + byok_server_ids = [s.server_id for s in redacted_mcp_servers if getattr(s, "is_byok", False)] if byok_server_ids: - cred_rows = await MCPUserCredentialsRepository( - _byok_prisma_client - ).table.find_many( + cred_rows = await MCPUserCredentialsRepository(_byok_prisma_client).table.find_many( where={"user_id": user_id, "server_id": {"in": byok_server_ids}} ) cred_set = {r.server_id for r in cred_rows} @@ -1057,19 +999,12 @@ if MCP_AVAILABLE: user_mcp_management_mode = _get_user_mcp_management_mode() if user_mcp_management_mode == "view_all": - servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_unfiltered( - server_ids=server_ids - ) - return [ - {"server_id": server.server_id, "status": server.status} - for server in servers - ] + servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_unfiltered(server_ids=server_ids) + return [{"server_id": server.server_id, "status": server.status} for server in servers] auth_contexts = await build_effective_auth_contexts(user_api_key_dict) - server_status_map: Dict[ - str, Optional[Literal["healthy", "unhealthy", "unknown"]] - ] = {} + server_status_map: Dict[str, Optional[Literal["healthy", "unhealthy", "unknown"]]] = {} for auth_context in auth_contexts: servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( user_api_key_auth=auth_context, @@ -1079,10 +1014,7 @@ if MCP_AVAILABLE: if server.server_id not in server_status_map: server_status_map[server.server_id] = server.status - return [ - {"server_id": server_id, "status": status} - for server_id, status in server_status_map.items() - ] + return [{"server_id": server_id, "status": status} for server_id, status in server_status_map.items()] @router.post( "/server/register", @@ -1112,9 +1044,7 @@ if MCP_AVAILABLE: if not user_api_key_dict.team_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "Registration requires an API key associated with a team. Use a team-scoped key." - }, + detail={"error": "Registration requires an API key associated with a team. Use a team-scoped key."}, ) # stdio servers spawn a local subprocess on the proxy host with the @@ -1135,9 +1065,7 @@ if MCP_AVAILABLE: }, ) - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") validate_and_normalize_mcp_server_payload(payload) _validate_mcp_required_fields(payload) @@ -1180,14 +1108,10 @@ if MCP_AVAILABLE: ): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Admin access required to view MCP server submissions." - }, + detail={"error": "Admin access required to view MCP server submissions."}, ) - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") submissions = await get_mcp_submissions(prisma_client) if not _user_is_full_admin(user_api_key_dict): @@ -1212,14 +1136,10 @@ if MCP_AVAILABLE: if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Admin access required to approve MCP server submissions." - }, + detail={"error": "Admin access required to approve MCP server submissions."}, ) - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") existing = await get_mcp_server(prisma_client, server_id) if existing is None: @@ -1260,14 +1180,10 @@ if MCP_AVAILABLE: if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Admin access required to reject MCP server submissions." - }, + detail={"error": "Admin access required to reject MCP server submissions."}, ) - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") existing = await get_mcp_server(prisma_client, server_id) if existing is None: @@ -1313,9 +1229,7 @@ if MCP_AVAILABLE: --header 'Authorization: Bearer your_api_key_here' ``` """ - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") # check to see if server exists (DB first, then registry for config-based servers) mcp_server = await get_mcp_server(prisma_client, server_id) @@ -1327,22 +1241,15 @@ if MCP_AVAILABLE: client_ip = IPAddressUtils.get_mcp_client_ip(request) registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) - if ( - registry_server is not None - and not global_mcp_server_manager._is_server_accessible_from_ip( - registry_server, client_ip - ) + if registry_server is not None and not global_mcp_server_manager._is_server_accessible_from_ip( + registry_server, client_ip ): registry_server = None if registry_server is None: # Try lookup by server_name or alias (client may use display name in URL) - registry_server = global_mcp_server_manager.get_mcp_server_by_name( - server_id, client_ip=client_ip - ) + registry_server = global_mcp_server_manager.get_mcp_server_by_name(server_id, client_ip=client_ip) if registry_server is not None: - mcp_server = global_mcp_server_manager._build_mcp_server_table( - registry_server - ) + mcp_server = global_mcp_server_manager._build_mcp_server_table(registry_server) if mcp_server is None: raise HTTPException( @@ -1352,25 +1259,17 @@ if MCP_AVAILABLE: # Implement authz restriction from requested user is_admin_view = _user_has_admin_view(user_api_key_dict) - is_restricted_virtual_key = _is_restricted_virtual_key_request( - user_api_key_dict - ) + is_restricted_virtual_key = _is_restricted_virtual_key_request(user_api_key_dict) if not is_admin_view: # Perform authz check BEFORE any health check (avoid side-effects for # unauthorized callers). if from_db: - mcp_server_records = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) + mcp_server_records = await get_all_mcp_servers_for_user(prisma_client, user_api_key_dict) exists = does_mcp_server_exist(mcp_server_records, server_id) else: # Registry/config server: use same access logic as list endpoint - allowed_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_dict - ) - ) + allowed_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_dict) exists = mcp_server.server_id in allowed_server_ids if not exists: @@ -1390,19 +1289,13 @@ if MCP_AVAILABLE: # Perform health check on the server using server manager try: - health_result = await global_mcp_server_manager.health_check_server( - server_id - ) + health_result = await global_mcp_server_manager.health_check_server(server_id) # Update the server object with health check results - mcp_server.status = ( - health_result.status if health_result.status else "unknown" - ) + mcp_server.status = health_result.status if health_result.status else "unknown" mcp_server.last_health_check = health_result.last_health_check mcp_server.health_check_error = health_result.health_check_error except Exception as e: - verbose_proxy_logger.debug( - f"Error performing health check on server {server_id}: {e}" - ) + verbose_proxy_logger.debug(f"Error performing health check on server {server_id}: {e}") mcp_server.status = "unknown" mcp_server.last_health_check = datetime.now() mcp_server.health_check_error = str(e) @@ -1435,9 +1328,7 @@ if MCP_AVAILABLE: """ Allow users to add a new external mcp server. """ - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") # Validate and normalize payload fields validate_and_normalize_mcp_server_payload(payload) @@ -1458,9 +1349,7 @@ if MCP_AVAILABLE: ): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": f"MCP Server with id {payload.server_id} is special and cannot be used." - }, + detail={"error": f"MCP Server with id {payload.server_id} is special and cannot be used."}, ) if payload.server_id is not None: @@ -1469,9 +1358,7 @@ if MCP_AVAILABLE: if mcp_server is not None: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another." - }, + detail={"error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."}, ) # TODO: audit log for create @@ -1507,8 +1394,7 @@ if MCP_AVAILABLE: await global_mcp_server_manager.reload_servers_from_database() except Exception as e: verbose_proxy_logger.exception( - f"MCP server {new_mcp_server.server_id} created but in-memory " - f"registry refresh failed: {str(e)}" + f"MCP server {new_mcp_server.server_id} created but in-memory registry refresh failed: {str(e)}" ) return _redact_mcp_credentials(new_mcp_server) @@ -1555,11 +1441,9 @@ if MCP_AVAILABLE: ) try: - temporary_server = ( - await global_mcp_server_manager.build_mcp_server_from_table( - temp_record, - credentials_are_encrypted=False, - ) + temporary_server = await global_mcp_server_manager.build_mcp_server_from_table( + temp_record, + credentials_are_encrypted=False, ) _cache_temporary_mcp_server( temporary_server, @@ -1570,9 +1454,7 @@ if MCP_AVAILABLE: ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, ) except Exception as e: - verbose_proxy_logger.exception( - f"Error caching temporary mcp server: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error caching temporary mcp server: {str(e)}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Error caching temporary mcp server: {str(e)}"}, @@ -1668,9 +1550,7 @@ if MCP_AVAILABLE: return UserAPIKeyAuth() request_data = await _read_request_body(request=request) - request_data = populate_request_with_path_params( - request_data=request_data, request=request - ) + request_data = populate_request_with_path_params(request_data=request_data, request=request) return await _user_api_key_auth_builder( request=request, @@ -1697,9 +1577,7 @@ if MCP_AVAILABLE: client_ip = IPAddressUtils.get_mcp_client_ip(request) if request else None server = global_mcp_server_manager.get_mcp_server_by_id( server_id - ) or global_mcp_server_manager.get_mcp_server_by_name( - server_id, client_ip=client_ip - ) + ) or global_mcp_server_manager.get_mcp_server_by_name(server_id, client_ip=client_ip) if server is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1719,11 +1597,7 @@ if MCP_AVAILABLE: ) allowed_server_ids: Set[str] = set() for auth_context in await build_effective_auth_contexts(user_api_key_dict): - allowed_server_ids.update( - await global_mcp_server_manager.get_allowed_mcp_servers( - auth_context - ) - ) + allowed_server_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(auth_context)) if server.server_id not in allowed_server_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -1748,9 +1622,7 @@ if MCP_AVAILABLE: response_type: Optional[str] = None, scope: Optional[str] = None, ): - mcp_server = await _get_cached_temporary_mcp_server_or_404( - server_id, user_api_key_dict, request=request - ) + mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) # Use the server's stored client_id when the caller doesn't supply one resolved_client_id = mcp_server.client_id or client_id or "" if not resolved_client_id: @@ -1794,9 +1666,7 @@ if MCP_AVAILABLE: refresh_token: Optional[str] = Form(None), scope: Optional[str] = Form(None), ): - mcp_server = await _get_cached_temporary_mcp_server_or_404( - server_id, user_api_key_dict, request=request - ) + mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) resolved_client_id = mcp_server.client_id or client_id or "" if not resolved_client_id: raise HTTPException( @@ -1832,9 +1702,7 @@ if MCP_AVAILABLE: server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - mcp_server = await _get_cached_temporary_mcp_server_or_404( - server_id, user_api_key_dict, request=request - ) + mcp_server = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) request_data = await _read_request_body(request=request) data: dict = {**request_data} @@ -1927,12 +1795,8 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Store a BYOK credential for the calling user.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) - mcp_server = await _authorize_and_fetch_mcp_server( - prisma_client, user_api_key_dict, server_id - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + mcp_server = await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) if not getattr(mcp_server, "is_byok", False): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -1945,9 +1809,7 @@ if MCP_AVAILABLE: detail={"error": "User ID not found in token"}, ) if payload.save: - await store_user_credential( - prisma_client, user_id, server_id, payload.credential - ) + await store_user_credential(prisma_client, user_id, server_id, payload.credential) from litellm.proxy._experimental.mcp_server.server import ( _invalidate_byok_cred_cache, ) @@ -1969,9 +1831,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Remove the calling user's BYOK credential.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( @@ -2004,12 +1864,8 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Persist the OAuth2 access token obtained by the calling user.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) - await _authorize_and_fetch_mcp_server( - prisma_client, user_api_key_dict, server_id - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( @@ -2049,9 +1905,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Revoke/delete the user's OAuth2 credential.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( @@ -2061,9 +1915,7 @@ if MCP_AVAILABLE: # Only delete if the stored credential is actually an OAuth2 token. # This prevents accidentally deleting a BYOK credential if one exists # for the same (user_id, server_id) pair. - cred_to_delete = await get_user_oauth_credential( - prisma_client, user_id, server_id - ) + cred_to_delete = await get_user_oauth_credential(prisma_client, user_id, server_id) if cred_to_delete is not None: try: await delete_user_credential(prisma_client, user_id, server_id) @@ -2087,9 +1939,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return credential status (has_credential, expiry) without exposing the token.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( @@ -2098,9 +1948,7 @@ if MCP_AVAILABLE: ) cred = await get_user_oauth_credential(prisma_client, user_id, server_id) if cred is None: - return MCPOAuthUserCredentialStatus( - server_id=server_id, has_credential=False, is_expired=False - ) + return MCPOAuthUserCredentialStatus(server_id=server_id, has_credential=False, is_expired=False) expires_at: Optional[str] = cred.get("expires_at") is_expired = False if expires_at: @@ -2128,9 +1976,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return all servers the calling user has connected via OAuth2.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( @@ -2142,10 +1988,7 @@ if MCP_AVAILABLE: return [] # Fetch server metadata for display names — single batch query instead of N+1. server_ids = [c["server_id"] for c in oauth_creds] - servers = { - srv.server_id: srv - for srv in await get_mcp_servers(prisma_client, server_ids) - } + servers = {srv.server_id: srv for srv in await get_mcp_servers(prisma_client, server_ids)} items: List[MCPUserCredentialListItem] = [] for cred in oauth_creds: sid = cred["server_id"] @@ -2184,9 +2027,7 @@ if MCP_AVAILABLE: if server is None: registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) if registry_server is not None: - server = global_mcp_server_manager._build_mcp_server_table( - registry_server - ) + server = global_mcp_server_manager._build_mcp_server_table(registry_server) if _user_has_admin_view(user_api_key_dict): if server is None: @@ -2198,9 +2039,7 @@ if MCP_AVAILABLE: allowed_server_ids: set[str] = set() for auth_context in await build_effective_auth_contexts(user_api_key_dict): - allowed_server_ids.update( - await global_mcp_server_manager.get_allowed_mcp_servers(auth_context) - ) + allowed_server_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(auth_context)) if server is None or server.server_id not in allowed_server_ids: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -2224,9 +2063,7 @@ if MCP_AVAILABLE: each value ``is_set`` and never echoes the decrypted secret back, so a leaked token can't be used to exfiltrate the raw upstream credential. """ - global_values, user_specs = parse_admin_env_vars( - getattr(server, "env_vars", None) - ) + global_values, user_specs = parse_admin_env_vars(getattr(server, "env_vars", None)) # An empty-valued global is not a usable fallback, so it must not mark a # referenced per-user var as covered, matching the empty-global filter in # _resolve_static_headers_with_env_vars. Otherwise this endpoint reports no @@ -2245,9 +2082,7 @@ if MCP_AVAILABLE: static_headers = {} referenced = collect_env_var_references(strings=static_headers.values()) user_var_names = {spec["name"] for spec in user_specs} - blocking = { - name for name in (referenced & user_var_names) if name not in global_values - } + blocking = {name for name in (referenced & user_var_names) if name not in global_values} required: List[MCPUserEnvVarSpec] = [] missing_count = 0 @@ -2287,18 +2122,14 @@ if MCP_AVAILABLE: server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> MCPUserEnvVarsStatus: - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) - server = await _authorize_and_fetch_mcp_server( - prisma_client, user_api_key_dict, server_id - ) + server = await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) stored = await get_user_env_vars(prisma_client, user_id, server_id) return _compute_user_env_var_status(server=server, stored_values=stored) @@ -2320,18 +2151,14 @@ if MCP_AVAILABLE: payload: MCPUserEnvVarsRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> MCPUserEnvVarsStatus: - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) - server = await _authorize_and_fetch_mcp_server( - prisma_client, user_api_key_dict, server_id - ) + server = await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) # Only known per-user var names declared by the admin are accepted — # never persist arbitrary keys the user invents. Submitted values are # merged over the existing set so a user updating one credential does @@ -2339,12 +2166,8 @@ if MCP_AVAILABLE: # back); an omitted/empty field keeps its stored value. _, user_specs = parse_admin_env_vars(getattr(server, "env_vars", None)) allowed_names = {spec["name"] for spec in user_specs} - updates = { - k: v for k, v in payload.values.items() if k in allowed_names and v != "" - } - merged = await merge_user_env_vars( - prisma_client, user_id, server_id, updates, allowed_names - ) + updates = {k: v for k, v in payload.values.items() if k in allowed_names and v != ""} + merged = await merge_user_env_vars(prisma_client, user_id, server_id, updates, allowed_names) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( invalidate_user_env_vars_cache, ) @@ -2363,18 +2186,14 @@ if MCP_AVAILABLE: server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> MCPUserEnvVarsStatus: - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) - server = await _authorize_and_fetch_mcp_server( - prisma_client, user_api_key_dict, server_id - ) + server = await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) await delete_user_env_vars(prisma_client, user_id, server_id) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( invalidate_user_env_vars_cache, @@ -2394,9 +2213,7 @@ if MCP_AVAILABLE: async def list_mcp_user_env_var_status( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> List[MCPUserEnvVarsStatus]: - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") user_id = user_api_key_dict.user_id or "" if not user_id: return [] @@ -2408,9 +2225,7 @@ if MCP_AVAILABLE: statuses: List[MCPUserEnvVarsStatus] = [] for server in accessible: stored = stored_bulk.get(server.server_id, {}) - status_obj = _compute_user_env_var_status( - server=server, stored_values=stored - ) + status_obj = _compute_user_env_var_status(server=server, stored_values=stored) if status_obj.required: statuses.append(status_obj) return statuses @@ -2472,9 +2287,7 @@ if MCP_AVAILABLE: if mcp_server_record_updated is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail={ - "error": f"MCP Server not found, passed server_id={payload.server_id}" - }, + detail={"error": f"MCP Server not found, passed server_id={payload.server_id}"}, ) await global_mcp_server_manager.update_server(mcp_server_record_updated) @@ -2525,9 +2338,7 @@ if MCP_AVAILABLE: litellm.public_mcp_servers = [] for server_id in request.mcp_server_ids: - server = global_mcp_server_manager.get_mcp_server_by_id( - server_id=server_id - ) + server = global_mcp_server_manager.get_mcp_server_by_id(server_id=server_id) if server is None: raise HTTPException( status_code=404, @@ -2540,9 +2351,7 @@ if MCP_AVAILABLE: if "litellm_settings" not in config or config["litellm_settings"] is None: config["litellm_settings"] = {} - config["litellm_settings"]["public_mcp_servers"] = ( - litellm.public_mcp_servers - ) + config["litellm_settings"]["public_mcp_servers"] = litellm.public_mcp_servers # Save the updated config await proxy_config.save_config(new_config=config) @@ -2580,9 +2389,7 @@ if MCP_AVAILABLE: with open(_MCP_REGISTRY_PATH, "r") as f: data: Dict[str, Any] = json.load(f) except Exception as e: - verbose_proxy_logger.warning( - f"Failed to load MCP registry from {_MCP_REGISTRY_PATH}: {e}" - ) + verbose_proxy_logger.warning(f"Failed to load MCP registry from {_MCP_REGISTRY_PATH}: {e}") data = {"servers": []} _mcp_registry_cache = data return data @@ -2593,9 +2400,7 @@ if MCP_AVAILABLE: dependencies=[Depends(user_api_key_auth)], ) async def discover_mcp_servers( - query: Optional[str] = Query( - None, description="Search filter for server names and descriptions" - ), + query: Optional[str] = Query(None, description="Search filter for server names and descriptions"), category: Optional[str] = Query(None, description="Filter by category"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -2675,9 +2480,7 @@ if MCP_AVAILABLE: try: return _load_openapi_registry() except Exception as e: - verbose_proxy_logger.warning( - f"Failed to load OpenAPI registry from {_OPENAPI_REGISTRY_PATH}: {e}" - ) + verbose_proxy_logger.warning(f"Failed to load OpenAPI registry from {_OPENAPI_REGISTRY_PATH}: {e}") return {"apis": []} # --------------------------------------------------------------------------- @@ -2708,9 +2511,7 @@ if MCP_AVAILABLE: litellm_changed_by: Optional[str] = Header(None), ): """Create a named toolset — a curated selection of {server_id, tool_name} pairs.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -2729,9 +2530,7 @@ if MCP_AVAILABLE: except UniqueViolationError: raise HTTPException( status_code=status.HTTP_409_CONFLICT, - detail={ - "error": f"A toolset named '{payload.toolset_name}' already exists." - }, + detail={"error": f"A toolset named '{payload.toolset_name}' already exists."}, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -2749,9 +2548,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return toolsets the calling key is allowed to access.""" - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") is_admin = _user_has_admin_view(user_api_key_dict) op = user_api_key_dict.object_permission # mcp_toolsets=None or [] both mean "not restricted by toolsets". @@ -2774,9 +2571,7 @@ if MCP_AVAILABLE: toolset_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") # Non-admin keys may only fetch toolsets they've been explicitly granted. if not _user_has_admin_view(user_api_key_dict): op = user_api_key_dict.object_permission @@ -2804,9 +2599,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header(None), ): - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -2842,9 +2635,7 @@ if MCP_AVAILABLE: global_mcp_server_manager, ) - global_mcp_server_manager.invalidate_toolset_cache( - getattr(payload, "toolset_id", None) - ) + global_mcp_server_manager.invalidate_toolset_cache(getattr(payload, "toolset_id", None)) return result @router.delete( @@ -2858,9 +2649,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: Optional[str] = Header(None), ): - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index a8551f6333a..1ca29eee89d 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -48,9 +48,7 @@ def validate_models_exist(model_names: List[str], llm_router) -> Tuple[bool, Lis return (len(missing) == 0, missing) -def add_access_group_to_deployment( - model_info: Dict[str, Any], access_group: str -) -> Tuple[Dict[str, Any], bool]: +def add_access_group_to_deployment(model_info: Dict[str, Any], access_group: str) -> Tuple[Dict[str, Any], bool]: """ Add an access group to a deployment's model_info. @@ -96,13 +94,9 @@ async def update_deployments_with_access_group( verbose_proxy_logger.debug(f"Updating deployments for model_name: {model_name}") # Get all deployments with this model_name - deployments = await ModelRepository(prisma_client).table.find_many( - where={"model_name": model_name} - ) + deployments = await ModelRepository(prisma_client).table.find_many(where={"model_name": model_name}) - verbose_proxy_logger.debug( - f"Found {len(deployments)} deployments for model_name: {model_name}" - ) + verbose_proxy_logger.debug(f"Found {len(deployments)} deployments for model_name: {model_name}") # If no deployments found, this is a config model (not in DB) if len(deployments) == 0: @@ -153,15 +147,11 @@ async def update_specific_deployments_with_access_group( models_updated = 0 for model_id in model_ids: verbose_proxy_logger.debug(f"Updating specific deployment model_id: {model_id}") - deployment = await ModelRepository(prisma_client).table.find_unique( - where={"model_id": model_id} - ) + deployment = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}) if deployment is None: raise HTTPException( status_code=400, - detail={ - "error": f"Deployment with model_id '{model_id}' not found in Database." - }, + detail={"error": f"Deployment with model_id '{model_id}' not found in Database."}, ) model_info = deployment.model_info or {} updated_model_info, was_modified = add_access_group_to_deployment( @@ -174,15 +164,11 @@ async def update_specific_deployments_with_access_group( data={"model_info": json.dumps(updated_model_info)}, ) models_updated += 1 - verbose_proxy_logger.debug( - f"Updated deployment {model_id} with access group: {access_group}" - ) + verbose_proxy_logger.debug(f"Updated deployment {model_id} with access group: {access_group}") return models_updated -def remove_access_group_from_deployment( - model_info: Dict[str, Any], access_group: str -) -> Tuple[Dict[str, Any], bool]: +def remove_access_group_from_deployment(model_info: Dict[str, Any], access_group: str) -> Tuple[Dict[str, Any], bool]: """ Remove an access group from a deployment's model_info. @@ -291,9 +277,7 @@ async def create_model_group( prisma_client, ) - verbose_proxy_logger.debug( - f"Creating access group: {data.access_group} with models: {data.model_names}" - ) + verbose_proxy_logger.debug(f"Creating access group: {data.access_group} with models: {data.model_names}") # Validation: Check if access_group is provided if not data.access_group or not data.access_group.strip(): @@ -309,9 +293,7 @@ async def create_model_group( if not has_model_names and not has_model_ids: raise HTTPException( status_code=400, - detail={ - "error": "Either model_names or model_ids must be provided and non-empty" - }, + detail={"error": "Either model_names or model_ids must be provided and non-empty"}, ) # If model_ids is provided, use it (more precise targeting) @@ -340,9 +322,7 @@ async def create_model_group( try: # Check if access group already exists - existing_access_groups = await get_all_access_groups_from_db( - prisma_client=prisma_client - ) + existing_access_groups = await get_all_access_groups_from_db(prisma_client=prisma_client) if data.access_group in existing_access_groups: raise HTTPException( @@ -384,9 +364,7 @@ async def create_model_group( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - f"Error creating access group '{data.access_group}': {str(e)}" - ) + verbose_proxy_logger.exception(f"Error creating access group '{data.access_group}': {str(e)}") raise HTTPException( status_code=500, detail={"error": f"Failed to create access group: {str(e)}"}, @@ -425,9 +403,7 @@ async def list_access_groups( ) try: - access_groups_map = await get_all_access_groups_from_db( - prisma_client=prisma_client - ) + access_groups_map = await get_all_access_groups_from_db(prisma_client=prisma_client) # Sort by access group name access_groups_list = sorted( @@ -482,9 +458,7 @@ async def get_access_group_info( ) try: - access_groups_map = await get_all_access_groups_from_db( - prisma_client=prisma_client - ) + access_groups_map = await get_all_access_groups_from_db(prisma_client=prisma_client) if access_group not in access_groups_map: raise HTTPException( @@ -497,9 +471,7 @@ async def get_access_group_info( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - f"Error getting access group info for '{access_group}': {str(e)}" - ) + verbose_proxy_logger.exception(f"Error getting access group info for '{access_group}': {str(e)}") raise HTTPException( status_code=500, detail={"error": f"Failed to get access group info: {str(e)}"}, @@ -553,9 +525,7 @@ async def update_access_group( detail={"error": "Database not connected."}, ) - verbose_proxy_logger.debug( - f"Updating access group: {access_group} with models: {data.model_names}" - ) + verbose_proxy_logger.debug(f"Updating access group: {access_group} with models: {data.model_names}") # Validation: Check that at least one of model_names or model_ids is provided has_model_names = data.model_names and len(data.model_names) > 0 @@ -564,18 +534,14 @@ async def update_access_group( if not has_model_names and not has_model_ids: raise HTTPException( status_code=400, - detail={ - "error": "Either model_names or model_ids must be provided and non-empty" - }, + detail={"error": "Either model_names or model_ids must be provided and non-empty"}, ) use_model_ids = has_model_ids # Validation: Check if access group exists try: - access_groups_map = await get_all_access_groups_from_db( - prisma_client=prisma_client - ) + access_groups_map = await get_all_access_groups_from_db(prisma_client=prisma_client) if access_group not in access_groups_map: raise HTTPException( status_code=404, @@ -654,9 +620,7 @@ async def update_access_group( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - f"Error updating access group '{access_group}': {str(e)}" - ) + verbose_proxy_logger.exception(f"Error updating access group '{access_group}': {str(e)}") raise HTTPException( status_code=500, detail={"error": f"Failed to update access group: {str(e)}"}, @@ -705,9 +669,7 @@ async def delete_access_group( # Validation: Check if access group exists try: - access_groups_map = await get_all_access_groups_from_db( - prisma_client=prisma_client - ) + access_groups_map = await get_all_access_groups_from_db(prisma_client=prisma_client) if access_group not in access_groups_map: raise HTTPException( status_code=404, @@ -757,9 +719,7 @@ async def delete_access_group( except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - f"Error deleting access group '{access_group}': {str(e)}" - ) + verbose_proxy_logger.exception(f"Error deleting access group '{access_group}': {str(e)}") raise HTTPException( status_code=500, detail={"error": f"Failed to delete access group: {str(e)}"}, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index df4d46b9098..afe71084a45 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -78,21 +78,15 @@ async def update_team(*args, **kwargs): class UpdatePublicModelGroupsRequest(BaseModel): """Request model for updating public model groups""" - model_groups: List[str] = Field( - description="List of model group names to make public" - ) + model_groups: List[str] = Field(description="List of model group names to make public") model_config = ConfigDict(extra="forbid") -async def get_db_model( - model_id: str, prisma_client: PrismaClient -) -> Optional[Deployment]: +async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Optional[Deployment]: db_model = cast( Optional[BaseModel], - await ModelRepository(prisma_client).table.find_unique( - where={"model_id": model_id} - ), + await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}), ) if not db_model: @@ -102,9 +96,7 @@ async def get_db_model( return deployment_pydantic_obj -def update_db_model( - db_model: Deployment, updated_patch: updateDeployment -) -> PrismaCompatibleUpdateDBModel: +def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: merged_deployment_dict = DeploymentTypedDict( model_name=db_model.model_name, litellm_params=LiteLLMParamsTypedDict( @@ -120,10 +112,7 @@ def update_db_model( if updated_patch.litellm_params: # Encrypt any sensitive values encrypted_params = { - k: encrypt_value_helper(v) - for k, v in updated_patch.litellm_params.model_dump( - exclude_none=True - ).items() + k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() } merged_deployment_dict["litellm_params"].update(encrypted_params) # type: ignore @@ -132,9 +121,7 @@ def update_db_model( if updated_patch.model_info: if "model_info" not in merged_deployment_dict: merged_deployment_dict["model_info"] = {} - merged_deployment_dict["model_info"].update( - updated_patch.model_info.model_dump(exclude_none=True) - ) + merged_deployment_dict["model_info"].update(updated_patch.model_info.model_dump(exclude_none=True)) # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI # passes through (which today re-sends the OLD pricing on every save) cannot @@ -147,18 +134,12 @@ def update_db_model( # clear propagates to both blobs. if updated_patch.litellm_params: for field in updated_patch.litellm_params.model_fields_set: - if ( - field in SPECIAL_MODEL_INFO_PARAMS - and getattr(updated_patch.litellm_params, field) is None - ): + if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: merged_deployment_dict["litellm_params"].pop(field, None) # type: ignore merged_deployment_dict.get("model_info", {}).pop(field, None) if updated_patch.model_info: for field in updated_patch.model_info.model_fields_set: - if ( - field in SPECIAL_MODEL_INFO_PARAMS - and getattr(updated_patch.model_info, field) is None - ): + if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: merged_deployment_dict["model_info"].pop(field, None) # type: ignore merged_deployment_dict.get("litellm_params", {}).pop(field, None) # type: ignore @@ -166,14 +147,10 @@ def update_db_model( prisma_compatible_model_dict = PrismaCompatibleUpdateDBModel() if "model_name" in merged_deployment_dict: - prisma_compatible_model_dict["model_name"] = merged_deployment_dict[ - "model_name" - ] + prisma_compatible_model_dict["model_name"] = merged_deployment_dict["model_name"] if "litellm_params" in merged_deployment_dict: - prisma_compatible_model_dict["litellm_params"] = json.dumps( - merged_deployment_dict["litellm_params"] - ) + prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"]) if "model_info" in merged_deployment_dict: model_info = merged_deployment_dict["model_info"] @@ -268,10 +245,7 @@ async def patch_model( # Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins # passed the auth check above for team-scoped models, but they must not # be able to unblock (or block) a model their proxy admin has paused. - if ( - patch_data.blocked is not None - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN - ): + if patch_data.blocked is not None and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise ProxyException( message="Only proxy admins can change a model's blocked flag.", type=ProxyErrorTypes.auth_error.value, @@ -288,9 +262,7 @@ async def patch_model( ) # Add metadata about update - update_data["updated_by"] = ( - user_api_key_dict.user_id or litellm_proxy_admin_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()) # Perform partial update @@ -375,10 +347,7 @@ async def _set_model_blocked_status( ) if db_model is None: - if ( - llm_router - and llm_router.get_deployment(model_id=data.model_id) is not None - ): + if llm_router and llm_router.get_deployment(model_id=data.model_id) is not None: raise ProxyException( message="Cannot edit config-based model. Store model in DB via /model/new first.", type=ProxyErrorTypes.validation_error.value, @@ -411,9 +380,7 @@ async def _set_model_blocked_status( table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, before_value=db_model.model_dump_json(exclude_none=True), after_value=( - updated_model.model_dump_json(exclude_none=True) - if isinstance(updated_model, BaseModel) - else None + updated_model.model_dump_json(exclude_none=True) if isinstance(updated_model, BaseModel) else None ), litellm_changed_by=litellm_changed_by, litellm_proxy_admin_name=litellm_proxy_admin_name, @@ -511,16 +478,12 @@ async def _add_model_to_db( _litellm_params_dict = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name = model_params.litellm_params.model for k, v in _litellm_params_dict.items(): - encrypted_value = encrypt_value_helper( - value=v, new_encryption_key=new_encryption_key - ) + encrypted_value = encrypt_value_helper(value=v, new_encryption_key=new_encryption_key) model_params.litellm_params[k] = encrypted_value _data: dict = { "model_id": model_params.model_info.id, "model_name": model_params.model_name, - "litellm_params": model_params.litellm_params.model_dump_json( - exclude_none=True - ), # type: ignore + "litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), # type: ignore "model_info": model_params.model_info.model_dump_json( # type: ignore exclude_none=True ), @@ -692,11 +655,7 @@ def _get_public_model_name( return name.startswith(f"model_name_{team_id}_") incoming = patch_data.model_name - if ( - incoming - and not _is_internal_shape(incoming) - and incoming != db_model.model_name - ): + if incoming and not _is_internal_shape(incoming) and incoming != db_model.model_name: return incoming if db_model.model_info and db_model.model_info.team_public_model_name: @@ -784,14 +743,10 @@ async def delete_team_models( deleted_model_ids: List[str] = [] async with prisma_client.db.tx() as tx: for team_id in team_ids: - rows = await _get_team_deployments( - team_id, prisma_client, table=tx.litellm_proxymodeltable - ) + rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] if model_ids: - await tx.litellm_proxymodeltable.delete_many( - where={"model_id": {"in": model_ids}} - ) + await tx.litellm_proxymodeltable.delete_many(where={"model_id": {"in": model_ids}}) deleted_model_ids.extend(model_ids) if llm_router is not None: @@ -853,37 +808,23 @@ async def _remove_unbacked_team_models( public_model_name=model_params.model_name, prisma_client=prisma_client, ) - names_to_remove = { - alias - for alias_team_id, alias in removed_model_aliases - if alias_team_id == team_id - } + names_to_remove = {alias for alias_team_id, alias in removed_model_aliases if alias_team_id == team_id} if model_params.model_info.team_public_model_name is not None: names_to_remove.add(model_params.model_info.team_public_model_name) if names_to_remove: - names_to_remove -= await _get_team_public_model_names( - team_id=team_id, prisma_client=prisma_client - ) + names_to_remove -= await _get_team_public_model_names(team_id=team_id, prisma_client=prisma_client) if not names_to_remove: return - existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) + existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) if existing_team_row is None: return updated_team_row = await prisma_client.db.litellm_teamtable.update( where={"team_id": team_id}, - data={ - "models": [ - model - for model in existing_team_row.models - if model not in names_to_remove - ] - }, + data={"models": [model for model in existing_team_row.models if model not in names_to_remove]}, include={"object_permission": True}, # type: ignore ) await _refresh_cached_team( @@ -925,9 +866,7 @@ async def _update_existing_team_model_assignment( return value if isinstance(value, str) else None return None - old_public_name = ( - db_model.model_info.team_public_model_name if db_model.model_info else None - ) + old_public_name = 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 @@ -944,8 +883,7 @@ async def _update_existing_team_model_assignment( other_deployments_with_old_name = [ d for d in team_deployments - if d.model_name != db_model.model_name - and _get_team_public_model_name(d.model_info) == old_public_name + if d.model_name != db_model.model_name and _get_team_public_model_name(d.model_info) == old_public_name ] # Add new name first, then delete old name to prevent access loss on partial failure @@ -1003,14 +941,9 @@ class ModelManagementAuthChecks: status_code=403, detail={"error": CommonProxyErrors.not_premium_user.value}, ) - if ( - user_api_key_dict.user_role - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - ): + if user_api_key_dict.user_role and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return True - elif team_obj is None or not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + elif team_obj is None or not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): raise HTTPException( status_code=403, detail={ @@ -1043,11 +976,7 @@ class ModelManagementAuthChecks: if _existing_team_row is None: raise HTTPException( status_code=400, - detail={ - "error": "Team id={} does not exist in db".format( - model_params.model_info.team_id - ) - }, + detail={"error": "Team id={} does not exist in db".format(model_params.model_info.team_id)}, ) existing_team_row = LiteLLM_TeamTable(**_existing_team_row.model_dump()) @@ -1068,10 +997,7 @@ class ModelManagementAuthChecks: allow_missing_team: bool = False, ) -> Literal[True]: ## Check team model auth - if ( - model_params.model_info is not None - and model_params.model_info.team_id is not None - ): + if model_params.model_info is not None and model_params.model_info.team_id is not None: team_obj_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1084,17 +1010,11 @@ class ModelManagementAuthChecks: return True raise HTTPException( status_code=403, - detail={ - "error": "Only a proxy admin can delete a model whose team has been deleted." - }, + detail={"error": "Only a proxy admin can delete a model whose team has been deleted."}, ) raise HTTPException( status_code=400, - detail={ - "error": "Team id={} does not exist in db".format( - model_params.model_info.team_id - ) - }, + detail={"error": "Team id={} does not exist in db".format(model_params.model_info.team_id)}, ) team_obj = LiteLLM_TeamTable(**team_obj_row.model_dump()) @@ -1158,9 +1078,7 @@ async def delete_model( }, ) - model_in_db = await ModelRepository(prisma_client).table.find_unique( - where={"model_id": model_info.id} - ) + model_in_db = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_info.id}) if model_in_db is None: raise HTTPException( status_code=400, @@ -1183,9 +1101,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result = await ModelRepository(prisma_client).table.delete( - where={"model_id": model_info.id} - ) + result = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1223,15 +1139,11 @@ async def delete_model( else: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) except Exception as e: - verbose_proxy_logger.exception( - f"Failed to delete model. Due to error - {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to delete model. Due to error - {str(e)}") if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "detail", f"Authentication Error({str(e)})"), @@ -1261,9 +1173,7 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases = await ModelTableRepository(prisma_client).table.find_many( - include={"team": True} - ) + team_model_aliases = await ModelTableRepository(prisma_client).table.find_many(include={"team": True}) tasks = [] removed_model_aliases = [] for team_model_alias in team_model_aliases: @@ -1271,9 +1181,7 @@ async def delete_team_model_alias( id = team_model_alias.id if public_model_name in model_aliases.values(): - key = list(model_aliases.keys())[ - list(model_aliases.values()).index(public_model_name) - ] + key = list(model_aliases.keys())[list(model_aliases.values()).index(public_model_name)] if team_model_alias.team is not None: removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] @@ -1381,9 +1289,7 @@ async def add_new_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) - await proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) + await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) # don't let failed slack alert block the /model/new response _alerting = general_settings.get("alerting", []) or [] if "slack" in _alerting: @@ -1399,17 +1305,13 @@ async def add_new_model( else: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) if model_response is None: raise HTTPException( status_code=500, - detail={ - "error": "Failed to add model to db. Check your server logs for more details." - }, + detail={"error": "Failed to add model to db. Check your server logs for more details."}, ) ## CREATE AUDIT LOG ## @@ -1421,9 +1323,7 @@ async def add_new_model( table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, before_value=None, after_value=( - model_response.model_dump_json(exclude_none=True) - if isinstance(model_response, BaseModel) - else None + model_response.model_dump_json(exclude_none=True) if isinstance(model_response, BaseModel) else None ), litellm_changed_by=user_api_key_dict.user_id, litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, @@ -1434,9 +1334,7 @@ async def add_new_model( except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.add_new_model(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.add_new_model(): Exception occured - {}".format(str(e)) ) if isinstance(e, HTTPException): raise ProxyException( @@ -1497,20 +1395,13 @@ async def update_model( if _model_id is None: raise Exception("model_info.id not provided") - _existing_litellm_params = await ModelRepository( - prisma_client - ).table.find_unique(where={"model_id": _model_id}) + _existing_litellm_params = await ModelRepository(prisma_client).table.find_unique(where={"model_id": _model_id}) if _existing_litellm_params is None: - if ( - llm_router is not None - and llm_router.get_deployment(model_id=_model_id) is not None - ): + if llm_router is not None and llm_router.get_deployment(model_id=_model_id) is not None: raise HTTPException( status_code=400, - detail={ - "error": "Can't edit model. Model in config. Store model in db via `/model/new`. to edit." - }, + detail={"error": "Can't edit model. Model in config. Store model in db via `/model/new`. to edit."}, ) else: raise Exception("model not found") @@ -1525,16 +1416,12 @@ async def update_model( # update DB if store_model_in_db is True: - _existing_litellm_params_dict = dict( - _existing_litellm_params.litellm_params - ) + _existing_litellm_params_dict = dict(_existing_litellm_params.litellm_params) if model_params.litellm_params is None: raise Exception("litellm_params not provided") - _new_litellm_params_dict = model_params.litellm_params.dict( - exclude_none=True - ) + _new_litellm_params_dict = model_params.litellm_params.dict(exclude_none=True) ### ENCRYPT PARAMS ### for k, v in _new_litellm_params_dict.items(): @@ -1548,10 +1435,7 @@ async def update_model( for key, value in _mp.items(): if value is not None: merged_dictionary[key] = value - elif ( - key in _existing_litellm_params_dict - and _existing_litellm_params_dict[key] is not None - ): + elif key in _existing_litellm_params_dict and _existing_litellm_params_dict[key] is not None: merged_dictionary[key] = _existing_litellm_params_dict[key] else: pass @@ -1593,9 +1477,7 @@ async def update_model( return model_response except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.update_model(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.update_model(): Exception occured - {}".format(str(e)) ) if isinstance(e, HTTPException): raise ProxyException( @@ -1660,9 +1542,7 @@ async def update_public_model_groups( if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) # Load existing config first (this may overwrite in-memory litellm settings @@ -1810,9 +1690,7 @@ async def clear_cache(): ) if llm_router is None or prisma_client is None: - verbose_proxy_logger.debug( - "llm_router or prisma_client is None, skipping cache clear" - ) + verbose_proxy_logger.debug("llm_router or prisma_client is None, skipping cache clear") return try: @@ -1841,14 +1719,10 @@ async def clear_cache(): llm_router.auto_routers.clear() # Reload only DB models - await proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) + await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) verbose_proxy_logger.debug( f"Cleared {len(db_model_ids)} DB models, preserved {len(config_models)} config models" ) except Exception as e: - verbose_proxy_logger.exception( - f"Failed to clear cache and reload models. Due to error - {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to clear cache and reload models. Due to error - {str(e)}") diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index a45382b54d0..138a55d9227 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -92,10 +92,7 @@ async def _verify_org_access( ) for m in caller_user.organization_memberships or []: - if ( - m.organization_id == organization_id - and m.user_role == LitellmUserRoles.ORG_ADMIN.value - ): + if m.organization_id == organization_id and m.user_role == LitellmUserRoles.ORG_ADMIN.value: return raise HTTPException( @@ -214,40 +211,25 @@ async def new_organization( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if ( - user_api_key_dict.user_role is None - or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN - ): + if user_api_key_dict.user_role is None or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=401, - detail={ - "error": f"Only admins can create orgs. Your role is = {user_api_key_dict.user_role}" - }, + detail={"error": f"Only admins can create orgs. Your role is = {user_api_key_dict.user_role}"}, ) if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) # Validate budget values are not negative - if data.max_budget is not None and ( - not math.isfinite(data.max_budget) or data.max_budget < 0 - ): + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) - if data.soft_budget is not None and ( - not math.isfinite(data.soft_budget) or data.soft_budget < 0 - ): + if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}" - }, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) user_object_correct_type: Optional[LiteLLM_UserTable] = None @@ -307,9 +289,7 @@ async def new_organization( ) for m in data.models: - await can_user_call_model( - m, llm_router=llm_router, user_object=user_object_correct_type - ) + await can_user_call_model(m, llm_router=llm_router, user_object=user_object_correct_type) organization_row = LiteLLM_OrganizationTable( **data.json(exclude_none=True), @@ -326,12 +306,8 @@ async def new_organization( value=getattr(data, field), ) - new_organization_row = prisma_client.jsonify_object( - organization_row.json(exclude_none=True) - ) - verbose_proxy_logger.info( - f"new_organization_row: {json.dumps(new_organization_row, indent=2)}" - ) + new_organization_row = prisma_client.jsonify_object(organization_row.json(exclude_none=True)) + verbose_proxy_logger.info(f"new_organization_row: {json.dumps(new_organization_row, indent=2)}") response = await OrganizationRepository(prisma_client).table.create( data={ **new_organization_row, # type: ignore @@ -375,20 +351,14 @@ async def get_organization_daily_activity( org_ids_list = organization_ids.split(",") if organization_ids else None exclude_org_ids_list: Optional[List[str]] = None if exclude_organization_ids: - exclude_org_ids_list = ( - exclude_organization_ids.split(",") if exclude_organization_ids else None - ) + exclude_org_ids_list = exclude_organization_ids.split(",") if exclude_organization_ids else None # Restrict non-proxy-admins to only organizations where they are org_admin if not _user_has_admin_view(user_api_key_dict): - memberships = await OrganizationMembershipRepository( - prisma_client - ).table.find_many(where={"user_id": user_api_key_dict.user_id}) - admin_org_ids = [ - m.organization_id - for m in memberships - if m.user_role == LitellmUserRoles.ORG_ADMIN.value - ] + memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + where={"user_id": user_api_key_dict.user_id} + ) + admin_org_ids = [m.organization_id for m in memberships if m.user_role == LitellmUserRoles.ORG_ADMIN.value] if org_ids_list is None: # Default to orgs where user is org_admin org_ids_list = admin_org_ids @@ -398,24 +368,15 @@ async def get_organization_daily_activity( if org_id not in admin_org_ids: raise HTTPException( status_code=403, - detail={ - "error": "User is not org_admin for Organization= {}.".format( - org_id - ) - }, + detail={"error": "User is not org_admin for Organization= {}.".format(org_id)}, ) # Fetch organization aliases for metadata where_condition = {} if org_ids_list: where_condition["organization_id"] = {"in": list(org_ids_list)} - org_aliases = await OrganizationRepository(prisma_client).table.find_many( - where=where_condition - ) - org_alias_metadata = { - o.organization_id: {"organization_alias": o.organization_alias} - for o in org_aliases - } + org_aliases = await OrganizationRepository(prisma_client).table.find_many(where=where_condition) + org_alias_metadata = {o.organization_id: {"organization_alias": o.organization_alias} for o in org_aliases} # Query daily activity for organizations return await get_daily_activity( @@ -448,9 +409,7 @@ async def _set_object_permission( return None if data.object_permission is not None: - created_object_permission = await ObjectPermissionRepository( - prisma_client - ).table.create( + created_object_permission = await ObjectPermissionRepository(prisma_client).table.create( data=data.object_permission.model_dump(exclude_none=True), ) del data.object_permission @@ -489,31 +448,21 @@ async def update_organization( # Transform UI payload to expected format raw_data = await request.json() - raw_data_with_flat_budget_fields = ( - handle_nested_budget_structure_in_organization_update_request(raw_data) - ) + raw_data_with_flat_budget_fields = handle_nested_budget_structure_in_organization_update_request(raw_data) # Create validated data model data = LiteLLM_OrganizationTableUpdate(**raw_data_with_flat_budget_fields) # Validate budget values are not negative - if data.max_budget is not None and ( - not math.isfinite(data.max_budget) or data.max_budget < 0 - ): + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) - if data.soft_budget is not None and ( - not math.isfinite(data.soft_budget) or data.soft_budget < 0 - ): + if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}" - }, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) if data.updated_by is None: @@ -534,30 +483,22 @@ async def update_organization( prisma_client=prisma_client, ) - existing_organization_row = await OrganizationRepository( - prisma_client - ).table.find_unique( + existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": data.organization_id}, ) if existing_organization_row is None: - raise ValueError( - f"Organization not found for organization_id={data.organization_id}" - ) + raise ValueError(f"Organization not found for organization_id={data.organization_id}") updated_organization_row_json = data.model_dump(exclude_none=True) # Merge metadata from existing organization with updated metadata if updated_organization_row_json.get("metadata") is not None: existing_metadata = existing_organization_row.metadata or {} updated_metadata = updated_organization_row_json.get("metadata", {}) - merged_metadata = _update_dictionary( - existing_dict=existing_metadata.copy(), new_dict=updated_metadata - ) + merged_metadata = _update_dictionary(existing_dict=existing_metadata.copy(), new_dict=updated_metadata) updated_organization_row_json["metadata"] = merged_metadata - updated_organization_row = prisma_client.jsonify_object( - updated_organization_row_json - ) + updated_organization_row = prisma_client.jsonify_object(updated_organization_row_json) if data.object_permission is not None: updated_organization_row = await handle_update_object_permission( data_json=updated_organization_row, @@ -566,16 +507,12 @@ async def update_organization( # Handle budget updates if budget fields are provided budget_fields = { - k: v - for k, v in data.model_dump().items() - if k in LiteLLM_BudgetTable.model_fields.keys() and v is not None + k: v for k, v in data.model_dump().items() if k in LiteLLM_BudgetTable.model_fields.keys() and v is not None } if budget_fields and existing_organization_row.budget_id: await update_budget( - budget_obj=BudgetNewRequest( - budget_id=existing_organization_row.budget_id, **budget_fields - ), + budget_obj=BudgetNewRequest(budget_id=existing_organization_row.budget_id, **budget_fields), user_api_key_dict=user_api_key_dict, ) @@ -653,17 +590,13 @@ async def delete_organization( deleted_orgs = [] for organization_id in data.organization_ids: # delete all teams in the organization - await TeamRepository(prisma_client).table.delete_many( - where={"organization_id": organization_id} - ) + await TeamRepository(prisma_client).table.delete_many(where={"organization_id": organization_id}) # delete all members in the organization await OrganizationMembershipRepository(prisma_client).table.delete_many( where={"organization_id": organization_id} ) # delete all keys in the organization - await VerificationTokenRepository(prisma_client).table.delete_many( - where={"organization_id": organization_id} - ) + await VerificationTokenRepository(prisma_client).table.delete_many(where={"organization_id": organization_id}) # delete the organization deleted_org = await OrganizationRepository(prisma_client).table.delete( where={"organization_id": organization_id}, @@ -747,12 +680,10 @@ async def list_organization( ) # if internal user - get orgs they are a member of (with optional filters) else: - org_memberships = await OrganizationMembershipRepository( - prisma_client - ).table.find_many(where={"user_id": user_api_key_dict.user_id}) - membership_org_ids = [ - membership.organization_id for membership in org_memberships - ] + org_memberships = await OrganizationMembershipRepository(prisma_client).table.find_many( + where={"user_id": user_api_key_dict.user_id} + ) + membership_org_ids = [membership.organization_id for membership in org_memberships] # Combine membership filter with provided filters if membership_org_ids: @@ -763,9 +694,7 @@ async def list_organization( response = [] else: where_conditions["organization_id"] = org_id - response = await OrganizationRepository( - prisma_client - ).table.find_many( + response = await OrganizationRepository(prisma_client).table.find_many( where=where_conditions, include={ "litellm_budget_table": True, @@ -816,9 +745,9 @@ async def info_organization( prisma_client=prisma_client, ) - response: Optional[ - LiteLLM_OrganizationTableWithMembers - ] = await OrganizationRepository(prisma_client).table.find_unique( + response: Optional[LiteLLM_OrganizationTableWithMembers] = await OrganizationRepository( + prisma_client + ).table.find_unique( where={"organization_id": organization_id}, include={ "litellm_budget_table": True, @@ -835,9 +764,7 @@ async def info_organization( if response is None: raise HTTPException(status_code=404, detail={"error": "Organization not found"}) - response_pydantic_obj = LiteLLM_OrganizationTableWithMembers( - **response.model_dump() - ) + response_pydantic_obj = LiteLLM_OrganizationTableWithMembers(**response.model_dump()) return response_pydantic_obj @@ -862,9 +789,7 @@ async def deprecated_info_organization( if len(data.organizations) == 0: raise HTTPException( status_code=400, - detail={ - "error": f"Specify list of organization id's to query. Passed in={data.organizations}" - }, + detail={"error": f"Specify list of organization id's to query. Passed in={data.organizations}"}, ) # Verify caller has access to each requested organization @@ -952,9 +877,9 @@ async def organization_member_add( ) # Check if organization exists - existing_organization_row = await OrganizationRepository( - prisma_client - ).table.find_unique(where={"organization_id": data.organization_id}) + existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": data.organization_id} + ) if existing_organization_row is None: raise HTTPException( status_code=404, @@ -1009,17 +934,15 @@ async def organization_member_add( ) -async def find_member_if_email( - user_email: str, prisma_client: PrismaClient -) -> LiteLLM_UserTable: +async def find_member_if_email(user_email: str, prisma_client: PrismaClient) -> LiteLLM_UserTable: """ Find a member if the user_email is in LiteLLM_UserTable """ try: - existing_user_email_row: BaseModel = await UserRepository( - prisma_client - ).table.find_unique(where={"user_email": user_email}) + existing_user_email_row: BaseModel = await UserRepository(prisma_client).table.find_unique( + where={"user_email": user_email} + ) except Exception: raise HTTPException( status_code=400, @@ -1027,9 +950,7 @@ async def find_member_if_email( "error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead." }, ) - existing_user_email_row_pydantic = LiteLLM_UserTable( - **existing_user_email_row.model_dump() - ) + existing_user_email_row_pydantic = LiteLLM_UserTable(**existing_user_email_row.model_dump()) return existing_user_email_row_pydantic @@ -1067,9 +988,9 @@ async def organization_member_update( ) # Check if organization exists - existing_organization_row = await OrganizationRepository( - prisma_client - ).table.find_unique(where={"organization_id": data.organization_id}) + existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": data.organization_id} + ) if existing_organization_row is None: raise HTTPException( status_code=400, @@ -1080,15 +1001,11 @@ async def organization_member_update( # Check if member exists in organization if data.user_email is not None and data.user_id is None: - existing_user_email_row = await find_member_if_email( - data.user_email, prisma_client - ) + existing_user_email_row = await find_member_if_email(data.user_email, prisma_client) data.user_id = existing_user_email_row.user_id try: - existing_organization_membership = await OrganizationMembershipRepository( - prisma_client - ).table.find_unique( + existing_organization_membership = await OrganizationMembershipRepository(prisma_client).table.find_unique( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1106,21 +1023,15 @@ async def organization_member_update( if existing_organization_membership is None: raise HTTPException( status_code=404, - detail={ - "error": f"Member not found in organization for user_id={data.user_id}" - }, + detail={"error": f"Member not found in organization for user_id={data.user_id}"}, ) # Reject attempts to change the role of a global PROXY_ADMIN via # org-scoped operations. An org-admin of any org could otherwise # alter a PROXY_ADMIN user's per-org role, which has downstream # effects on admin UI filtering and scope derivation. - target_user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": data.user_id} - ) - if target_user_row is not None and getattr( - target_user_row, "user_role", None - ) in ( + target_user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": data.user_id}) + if target_user_row is not None and getattr(target_user_row, "user_role", None) in ( LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, ): @@ -1129,8 +1040,7 @@ async def organization_member_update( status_code=403, detail={ "error": ( - "Only PROXY_ADMIN may modify the organization " - "role of a user who is a global PROXY_ADMIN." + "Only PROXY_ADMIN may modify the organization role of a user who is a global PROXY_ADMIN." ) }, ) @@ -1150,18 +1060,12 @@ async def organization_member_update( # if budget_id is None, create a new budget budget_id = existing_organization_membership.budget_id or str(uuid.uuid4()) if existing_organization_membership.budget_id is None: - new_budget_obj = BudgetNewRequest( - budget_id=budget_id, max_budget=data.max_budget_in_organization - ) - await new_budget( - budget_obj=new_budget_obj, user_api_key_dict=user_api_key_dict - ) + new_budget_obj = BudgetNewRequest(budget_id=budget_id, max_budget=data.max_budget_in_organization) + await new_budget(budget_obj=new_budget_obj, user_api_key_dict=user_api_key_dict) else: # update budget table with new max_budget await update_budget( - budget_obj=BudgetNewRequest( - budget_id=budget_id, max_budget=data.max_budget_in_organization - ), + budget_obj=BudgetNewRequest(budget_id=budget_id, max_budget=data.max_budget_in_organization), user_api_key_dict=user_api_key_dict, ) @@ -1175,9 +1079,9 @@ async def organization_member_update( }, data={"budget_id": budget_id}, ) - final_organization_membership: Optional[ - BaseModel - ] = await OrganizationMembershipRepository(prisma_client).table.find_unique( + final_organization_membership: Optional[BaseModel] = await OrganizationMembershipRepository( + prisma_client + ).table.find_unique( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1190,9 +1094,7 @@ async def organization_member_update( if final_organization_membership is None: raise HTTPException( status_code=400, - detail={ - "error": f"Member not found in organization={data.organization_id} for user_id={data.user_id}" - }, + detail={"error": f"Member not found in organization={data.organization_id} for user_id={data.user_id}"}, ) final_organization_membership_pydantic = LiteLLM_OrganizationMembershipTable( @@ -1235,14 +1137,10 @@ async def organization_member_delete( ) if data.user_email is not None and data.user_id is None: - existing_user_email_row = await find_member_if_email( - data.user_email, prisma_client - ) + existing_user_email_row = await find_member_if_email(data.user_email, prisma_client) data.user_id = existing_user_email_row.user_id - member_to_delete = await OrganizationMembershipRepository( - prisma_client - ).table.delete( + member_to_delete = await OrganizationMembershipRepository(prisma_client).table.delete( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1276,15 +1174,15 @@ async def add_member_to_organization( existing_user_email_row = None ## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable if member.user_id is not None: - existing_user_id_row = await UserRepository( - prisma_client - ).table.find_unique(where={"user_id": member.user_id}) + existing_user_id_row = await UserRepository(prisma_client).table.find_unique( + where={"user_id": member.user_id} + ) if existing_user_id_row is None and member.user_email is not None: try: - existing_user_email_row = await UserRepository( - prisma_client - ).table.find_unique(where={"user_email": member.user_email}) + existing_user_email_row = await UserRepository(prisma_client).table.find_unique( + where={"user_email": member.user_email} + ) except Exception as e: raise ValueError( f"Potential NON-Existent or Duplicate user email in DB: Error finding a unique instance of user_email={member.user_email} in LiteLLM_UserTable.: {e}" @@ -1299,17 +1197,13 @@ async def add_member_to_organization( user_email=member.user_email, ) - _returned_user = await prisma_client.insert_data( - data=new_user_defaults, table_name="user" - ) # type: ignore + _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore if _returned_user is not None: user_object = LiteLLM_UserTable(**_returned_user.model_dump()) elif existing_user_email_row is not None and len(existing_user_email_row) > 1: raise HTTPException( status_code=400, - detail={ - "error": "Multiple users with this email found in db. Please use 'user_id' instead." - }, + detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."}, ) elif existing_user_email_row is not None: user_object = LiteLLM_UserTable(**existing_user_email_row.model_dump()) @@ -1318,9 +1212,7 @@ async def add_member_to_organization( else: raise HTTPException( status_code=404, - detail={ - "error": f"User not found for user_id={member.user_id} and user_email={member.user_email}" - }, + detail={"error": f"User not found for user_id={member.user_id} and user_email={member.user_email}"}, ) if user_object is None: @@ -1329,21 +1221,15 @@ async def add_member_to_organization( ) # Add user to organization - _organization_membership = await OrganizationMembershipRepository( - prisma_client - ).table.create( + _organization_membership = await OrganizationMembershipRepository(prisma_client).table.create( data={ "organization_id": organization_id, "user_id": user_object.user_id, "user_role": member.role, } ) - organization_membership = LiteLLM_OrganizationMembershipTable( - **_organization_membership.model_dump() - ) + organization_membership = LiteLLM_OrganizationMembershipTable(**_organization_membership.model_dump()) return user_object, organization_membership except Exception as e: - raise ValueError( - f"Error adding member={member} to organization={organization_id}: {e}" - ) + raise ValueError(f"Error adding member={member} to organization={organization_id}: {e}") diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 0ae11435122..00ca52bb081 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -87,9 +87,7 @@ class AiPolicySuggester: valid_ids = {t["id"] for t in templates} result["selected_templates"] = [ - s - for s in result.get("selected_templates", []) - if s.get("template_id") in valid_ids + s for s in result.get("selected_templates", []) if s.get("template_id") in valid_ids ] return result diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 57578d98b75..f50dcb3054f 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -164,9 +164,7 @@ async def apply_policies( current_inputs = cast(GenericGuardrailAPIInputs, dict(inputs)) for guardrail_name in sorted(guardrail_name_set): - callback = guardrail_registry.get_initialized_guardrail_callback( - guardrail_name=guardrail_name - ) + callback = guardrail_registry.get_initialized_guardrail_callback(guardrail_name=guardrail_name) if callback is None: verbose_proxy_logger.debug( "apply_policies: guardrail '%s' not found, skipping", @@ -207,9 +205,7 @@ async def apply_policies( return {"inputs": current_inputs, "guardrail_errors": guardrail_errors} -def _chat_body_from_inputs( - inputs: GenericGuardrailAPIInputs, agent_id: str, request_data: dict -) -> dict: +def _chat_body_from_inputs(inputs: GenericGuardrailAPIInputs, agent_id: str, request_data: dict) -> dict: """Build a chat completion request body from guardrail inputs and agent_id.""" messages: List[dict] structured = inputs.get("structured_messages") @@ -259,19 +255,13 @@ def _request_with_json_body(body: dict) -> Request: class TestPoliciesAndGuardrailsRequest(BaseModel): """Request body for POST /utils/test_policies_and_guardrails.""" - policy_names: Optional[List[str]] = Field( - default=None, description="Policy names to resolve guardrails from" - ) - guardrail_names: Optional[List[str]] = Field( - default=None, description="Guardrail names to apply directly" - ) + policy_names: Optional[List[str]] = Field(default=None, description="Policy names to resolve guardrails from") + guardrail_names: Optional[List[str]] = Field(default=None, description="Guardrail names to apply directly") inputs_list: List[GenericGuardrailAPIInputs] = Field( default=[], description="List of GenericGuardrailAPIInputs; each item processed separately (for batch compliance testing).", ) - request_data: dict = Field( - default_factory=dict, description="Request context (model, user_id, etc.)" - ) + request_data: dict = Field(default_factory=dict, description="Request context (model, user_id, etc.)") input_type: Literal["request", "response"] = Field( default="request", description="Whether inputs are request or response" ) @@ -356,9 +346,7 @@ async def test_policies_and_guardrails( guardrail_name="response_rejection", ) try: - model_response = ModelResponse.model_validate( - item["agent_response"] - ) + model_response = ModelResponse.model_validate(item["agent_response"]) handler = OpenAIChatCompletionsHandler() await handler.process_output_response( response=model_response, @@ -443,9 +431,7 @@ async def validate_policy( from litellm.proxy.policy_engine.policy_validator import PolicyValidator from litellm.proxy.proxy_server import prisma_client - verbose_proxy_logger.debug( - f"Validating policy configuration with {len(data.policies)} policies" - ) + verbose_proxy_logger.debug(f"Validating policy configuration with {len(data.policies)} policies") validator = PolicyValidator(prisma_client=prisma_client) @@ -534,9 +520,7 @@ async def get_policy_info( detail=f"Policy '{policy_name}' not found", ) - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name=policy_name, policies=registry.get_all_policies() - ) + resolved = PolicyResolver.resolve_policy_guardrails(policy_name=policy_name, policies=registry.get_all_policies()) return PolicyInfoResponse( policy_name=policy_name, @@ -605,9 +589,7 @@ async def test_policy_matching( matching_policy_names = PolicyMatcher.get_matching_policies(context=context) # Resolve guardrails - resolved_guardrails = PolicyResolver.resolve_guardrails_for_context( - context=context, policies=policies - ) + resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(context=context, policies=policies) return PolicyTestResponse( context=context, @@ -616,9 +598,7 @@ async def test_policy_matching( ) -POLICY_TEMPLATES_GITHUB_URL = ( - "https://raw.githubusercontent.com/BerriAI/litellm/main/policy_templates.json" -) +POLICY_TEMPLATES_GITHUB_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/policy_templates.json" def _load_policy_templates_from_local_backup() -> list: @@ -673,9 +653,7 @@ async def get_policy_templates( if response.status_code == 200: return response.json() except Exception as e: - verbose_proxy_logger.debug( - "Failed to fetch policy templates from GitHub, using local backup: %s", e - ) + verbose_proxy_logger.debug("Failed to fetch policy templates from GitHub, using local backup: %s", e) return _load_policy_templates_from_local_backup() @@ -704,15 +682,11 @@ def _validate_enrichment_request(data: EnrichTemplateRequest) -> tuple[dict, dic templates = _load_policy_templates_from_local_backup() template = next((t for t in templates if t.get("id") == data.template_id), None) if template is None: - raise HTTPException( - status_code=404, detail=f"Template '{data.template_id}' not found" - ) + raise HTTPException(status_code=404, detail=f"Template '{data.template_id}' not found") llm_enrichment = template.get("llm_enrichment") if llm_enrichment is None: - raise HTTPException( - status_code=400, detail="Template does not support LLM enrichment" - ) + raise HTTPException(status_code=400, detail="Template does not support LLM enrichment") # Validate competitors list size if provided if data.competitors and len(data.competitors) > MAX_COMPETITOR_NAMES: @@ -754,9 +728,7 @@ async def enrich_policy_template( if data.competitors: competitors = data.competitors else: - prompt = llm_enrichment["prompt"].replace( - "{{" + llm_enrichment["parameter"] + "}}", brand_name - ) + prompt = llm_enrichment["prompt"].replace("{{" + llm_enrichment["parameter"] + "}}", brand_name) competitors = await _discover_competitors_via_llm(prompt, model=model) variations_map = await _generate_competitor_variations(competitors, model=model) @@ -822,11 +794,7 @@ async def _stream_llm_competitor_names( while "\n" in buffer: line, buffer = buffer.split("\n", 1) name = _clean_competitor_line(line) - if ( - name - and name.lower() not in existing_lower - and count < MAX_COMPETITOR_NAMES - ): + if name and name.lower() not in existing_lower and count < MAX_COMPETITOR_NAMES: existing_lower.add(name.lower()) count += 1 yield name, False @@ -851,13 +819,9 @@ async def _stream_competitor_events( for comp in competitors: yield f"data: {json.dumps({'type': 'competitor', 'name': comp})}\n\n" - refinement_prompt = _build_refinement_prompt( - data.instruction, competitors, brand_name - ) + refinement_prompt = _build_refinement_prompt(data.instruction, competitors, brand_name) try: - async for name, _ in _stream_llm_competitor_names( - refinement_prompt, model, competitors - ): + async for name, _ in _stream_llm_competitor_names(refinement_prompt, model, competitors): if name: competitors.append(name) yield f"data: {json.dumps({'type': 'competitor', 'name': name})}\n\n" @@ -871,9 +835,7 @@ async def _stream_competitor_events( yield f"data: {json.dumps({'type': 'competitor', 'name': comp})}\n\n" else: # Initial discovery mode - prompt = llm_enrichment["prompt"].replace( - "{{" + llm_enrichment["parameter"] + "}}", brand_name - ) + prompt = llm_enrichment["prompt"].replace("{{" + llm_enrichment["parameter"] + "}}", brand_name) try: async for name, _ in _stream_llm_competitor_names(prompt, model, []): if name: @@ -932,9 +894,7 @@ def _clean_competitor_line(line: str) -> Optional[str]: return name if name and len(name) > 1 else None -async def _generate_competitor_variations( - competitors: list, model: str = DEFAULT_COMPETITOR_DISCOVERY_MODEL -) -> dict: +async def _generate_competitor_variations(competitors: list, model: str = DEFAULT_COMPETITOR_DISCOVERY_MODEL) -> dict: """Generate common misspellings, abbreviations, and alternate names for each competitor.""" if not competitors: return {} @@ -983,18 +943,14 @@ def _parse_variations_response(raw: str, competitors: list) -> dict[str, list[st if canonical is None: continue variations = [ - v.strip() - for v in variations_str.split(",") - if v.strip() and v.strip().lower() != canonical.lower() + v.strip() for v in variations_str.split(",") if v.strip() and v.strip().lower() != canonical.lower() ] variations_map[canonical] = variations return variations_map -async def _discover_competitors_via_llm( - prompt: str, model: str = DEFAULT_COMPETITOR_DISCOVERY_MODEL -) -> list: +async def _discover_competitors_via_llm(prompt: str, model: str = DEFAULT_COMPETITOR_DISCOVERY_MODEL) -> list: """Call an onboarded LLM to discover competitor names.""" try: from litellm.proxy.proxy_server import llm_router @@ -1007,11 +963,7 @@ async def _discover_competitors_via_llm( temperature=COMPETITOR_LLM_TEMPERATURE, ) raw = response.choices[0].message.content or "" # type: ignore - competitors = [ - name - for line in raw.strip().split("\n") - if (name := _clean_competitor_line(line)) is not None - ] + competitors = [name for line in raw.strip().split("\n") if (name := _clean_competitor_line(line)) is not None] return competitors[:MAX_COMPETITOR_NAMES] except Exception as e: verbose_proxy_logger.error("LLM competitor discovery failed: %s", e) @@ -1038,9 +990,7 @@ def _build_competitor_guardrail_definitions( output_blocked = _build_name_blocked_words(competitors, all_names) recommendation_blocked = _build_recommendation_blocked_words(competitors, all_names) - comparison_blocked = _build_comparison_blocked_words( - competitors, all_names, brand_name - ) + comparison_blocked = _build_comparison_blocked_words(competitors, all_names, brand_name) blocked_words_map = { "competitor-output-blocker": output_blocked, @@ -1064,25 +1014,17 @@ def _build_competitor_guardrail_definitions( return enriched -def _build_name_blocked_words( - competitors: list[str], all_names: dict[str, list[str]] -) -> list[dict]: +def _build_name_blocked_words(competitors: list[str], all_names: dict[str, list[str]]) -> list[dict]: """Build blocked word entries for direct competitor name mentions.""" result = [] for comp in competitors: for name in all_names[comp]: - desc = ( - f"Competitor: {comp}" - if name == comp - else f"Competitor variation ({comp}): {name}" - ) + desc = f"Competitor: {comp}" if name == comp else f"Competitor variation ({comp}): {name}" result.append({"keyword": name, "action": "BLOCK", "description": desc}) return result -def _build_recommendation_blocked_words( - competitors: list[str], all_names: dict[str, list[str]] -) -> list[dict]: +def _build_recommendation_blocked_words(competitors: list[str], all_names: dict[str, list[str]]) -> list[dict]: """Build blocked word entries for competitor recommendations.""" result = [] for comp in competitors: @@ -1176,9 +1118,7 @@ class GuardrailTestResultEntry(TypedDict): class TestPolicyTemplateRequest(BaseModel): - guardrail_definitions: List[dict] = Field( - description="All guardrailDefinitions from the policy template" - ) + guardrail_definitions: List[dict] = Field(description="All guardrailDefinitions from the policy template") text: str = Field(description="Test input text to run guardrails against") @@ -1263,9 +1203,7 @@ async def _test_guardrail_definitions( request_data={}, input_type="request", ) - output_text = ( - output.get("texts", [text])[0] if output.get("texts") else text - ) + output_text = output.get("texts", [text])[0] if output.get("texts") else text if output_text != text: action = "masked" diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 557eb608136..d46ad41eee3 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -27,24 +27,16 @@ router = APIRouter() class RouterSettingsResponse(BaseModel): - fields: List[RouterSettingsField] = Field( - description="List of all configurable router settings with metadata" - ) - current_values: Dict[str, Any] = Field( - description="Current values of router settings" - ) - routing_strategy_descriptions: Dict[str, str] = Field( - description="Descriptions for each routing strategy option" - ) + fields: List[RouterSettingsField] = Field(description="List of all configurable router settings with metadata") + current_values: Dict[str, Any] = Field(description="Current values of router settings") + routing_strategy_descriptions: Dict[str, str] = Field(description="Descriptions for each routing strategy option") class RouterFieldsResponse(BaseModel): fields: List[RouterSettingsField] = Field( description="List of all configurable router settings with metadata (without field values)" ) - routing_strategy_descriptions: Dict[str, str] = Field( - description="Descriptions for each routing strategy option" - ) + routing_strategy_descriptions: Dict[str, str] = Field(description="Descriptions for each routing strategy option") def _get_routing_strategies_from_router_class() -> List[str]: @@ -90,9 +82,7 @@ async def get_router_settings( available_routing_strategies = _get_routing_strategies_from_router_class() # Get router settings fields from types file - router_fields = [ - field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS - ] + router_fields = [field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS] # Populate routing_strategy field with available options and descriptions for field in router_fields: @@ -108,9 +98,7 @@ async def get_router_settings( if llm_router is not None: # Router exposes routing groups as private `_routing_groups`; the # generic `hasattr` loop below would miss them. - current_values["routing_groups"] = [ - group.model_dump() for group in llm_router._routing_groups.values() - ] + current_values["routing_groups"] = [group.model_dump() for group in llm_router._routing_groups.values()] for field in router_fields: if field.field_name == "routing_groups": continue @@ -163,9 +151,7 @@ async def get_router_fields( available_routing_strategies = _get_routing_strategies_from_router_class() # Get router settings fields from types file - router_fields = [ - field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS - ] + router_fields = [field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS] # Populate routing_strategy field with available options for field in router_fields: diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index 866f5baf3a4..cc3f18f593d 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -23,16 +23,12 @@ class ScimTransformations: from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail={"error": "No database connected"} - ) + raise HTTPException(status_code=500, detail={"error": "No database connected"}) # Get user's teams/groups groups = [] for team_id in user.teams or []: - team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if team: team_alias = getattr(team, "team_alias", team.team_id) groups.append(SCIMUserGroup(value=team.team_id, display=team_alias)) @@ -53,9 +49,7 @@ class ScimTransformations: schemas = ["urn:ietf:params:scim:schemas:core:2.0:User"] enterprise_user = None if metadata.get(SCIM_ENTERPRISE_METADATA_KEY): - enterprise_user = SCIMEnterpriseUser.model_validate( - metadata[SCIM_ENTERPRISE_METADATA_KEY] - ) + enterprise_user = SCIMEnterpriseUser.model_validate(metadata[SCIM_ENTERPRISE_METADATA_KEY]) schemas.append(SCIM_ENTERPRISE_USER_SCHEMA) return SCIMUser( @@ -96,9 +90,7 @@ class ScimTransformations: """ metadata = user.metadata or {} if "scim_metadata" in metadata: - scim_metadata: LiteLLM_UserScimMetadata = LiteLLM_UserScimMetadata( - **metadata["scim_metadata"] - ) + scim_metadata: LiteLLM_UserScimMetadata = LiteLLM_UserScimMetadata(**metadata["scim_metadata"]) if scim_metadata.familyName and len(scim_metadata.familyName) > 0: return scim_metadata.familyName @@ -113,9 +105,7 @@ class ScimTransformations: """ metadata = user.metadata or {} if "scim_metadata" in metadata: - scim_metadata: LiteLLM_UserScimMetadata = LiteLLM_UserScimMetadata( - **metadata["scim_metadata"] - ) + scim_metadata: LiteLLM_UserScimMetadata = LiteLLM_UserScimMetadata(**metadata["scim_metadata"]) if scim_metadata.givenName and len(scim_metadata.givenName) > 0: return scim_metadata.givenName @@ -130,9 +120,7 @@ class ScimTransformations: from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail={"error": "No database connected"} - ) + raise HTTPException(status_code=500, detail={"error": "No database connected"}) if isinstance(team, dict): team = LiteLLM_TeamTable(**team) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index d5da0372a8f..808b80cd1ed 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -107,17 +107,11 @@ class UserProvisionerHelpers: "user_alias": new_user_request.user_alias, "teams": new_user_request.teams, "metadata": safe_dumps(new_user_request.metadata), - **( - {"user_role": new_user_request.user_role} - if admin_group is not None - else {} - ), + **({"user_role": new_user_request.user_role} if admin_group is not None else {}), }, ) - return await ScimTransformations.transform_litellm_user_to_scim_user( - updated_user - ) + return await ScimTransformations.transform_litellm_user_to_scim_user(updated_user) class ScimUserData(TypedDict): @@ -162,14 +156,10 @@ async def _check_user_exists(user_id: str): """Check if user exists and return user, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) if not user: - raise HTTPException( - status_code=404, detail={"error": f"User not found with ID: {user_id}"} - ) + raise HTTPException(status_code=404, detail={"error": f"User not found with ID: {user_id}"}) return user @@ -178,14 +168,10 @@ async def _check_team_exists(team_id: str): """Check if team exists and return team, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if not team: - raise HTTPException( - status_code=404, detail={"error": f"Group not found with ID: {team_id}"} - ) + raise HTTPException(status_code=404, detail={"error": f"Group not found with ID: {team_id}"}) return team @@ -234,9 +220,7 @@ def _build_scim_metadata( metadata["scim_active"] = active if enterprise is not None: - metadata[SCIM_ENTERPRISE_METADATA_KEY] = enterprise.model_dump( - by_alias=True, exclude_none=True - ) + metadata[SCIM_ENTERPRISE_METADATA_KEY] = enterprise.model_dump(by_alias=True, exclude_none=True) return metadata @@ -259,9 +243,7 @@ async def _get_scim_upsert_user_setting() -> bool: # Default to True if not set (backward compatibility) return bool(scim_upsert_user) except Exception as e: - verbose_proxy_logger.warning( - f"Error reading scim_upsert_user setting, defaulting to True: {e}" - ) + verbose_proxy_logger.warning(f"Error reading scim_upsert_user setting, defaulting to True: {e}") # Default to True for backward compatibility return True @@ -297,9 +279,7 @@ async def _get_scim_admin_group() -> Optional[str]: litellm_settings = config.get("litellm_settings", {}) or {} return litellm_settings.get("scim_admin_group") or None except Exception as e: - verbose_proxy_logger.warning( - f"Error reading scim_admin_group setting, defaulting to None: {e}" - ) + verbose_proxy_logger.warning(f"Error reading scim_admin_group setting, defaulting to None: {e}") return None @@ -323,20 +303,13 @@ def _resolve_scim_user_role( return default_role -async def _scim_groups_from_team_ids( - prisma_client: Any, team_ids: list[str] -) -> list[SCIMUserGroup]: +async def _scim_groups_from_team_ids(prisma_client: Any, team_ids: list[str]) -> list[SCIMUserGroup]: """ Build SCIMUserGroup objects from team ids, populating display from each team's alias so admin-group matching by display name works the same way it does on PUT (where SCIM groups carry display names natively). """ - teams = [ - await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) - for team_id in team_ids - ] + teams = [await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) for team_id in team_ids] return [ SCIMUserGroup( value=team_id, @@ -346,9 +319,7 @@ async def _scim_groups_from_team_ids( ] -async def _recompute_scim_member_roles( - prisma_client: Any, user_ids: Iterable[str] -) -> None: +async def _recompute_scim_member_roles(prisma_client: Any, user_ids: Iterable[str]) -> None: """ Recompute and persist each user's global proxy role from their resulting team membership. No-op unless scim_admin_group is configured, so a SCIM group write @@ -361,9 +332,7 @@ async def _recompute_scim_member_roles( default_role = _default_scim_user_role() for user_id in user_ids: - user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) if user is None: continue resolved_role = _resolve_scim_user_role( @@ -411,9 +380,7 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe ) # Check if user exists - user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) if user: existing_member_ids.append(user_id) @@ -453,9 +420,7 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: members: List[SCIMMember] = [] for member_id in member_ids: - user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": member_id} - ) + user = await UserRepository(prisma_client).table.find_unique(where={"user_id": member_id}) if user: display_name = user.user_email or user.user_id members.append(SCIMMember(value=user.user_id, display=display_name)) @@ -463,9 +428,7 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: return members -async def _handle_team_membership_changes( - user_id: str, existing_teams: List[str], new_teams: List[str] -) -> None: +async def _handle_team_membership_changes(user_id: str, existing_teams: List[str], new_teams: List[str]) -> None: """Handle adding/removing user from teams based on changes.""" existing_teams_set = set(existing_teams) new_teams_set = set(new_teams) @@ -486,9 +449,7 @@ SCIM_BLOCKED_METADATA_KEY = "scim_blocked" def _key_was_scim_blocked(metadata: Any) -> bool: """True if a verification token carries the SCIM-block marker in metadata.""" - return ( - isinstance(metadata, dict) and metadata.get(SCIM_BLOCKED_METADATA_KEY) is True - ) + return isinstance(metadata, dict) and metadata.get(SCIM_BLOCKED_METADATA_KEY) is True async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: @@ -527,17 +488,11 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: return 0 for key_row in affected_keys: - current_metadata: Dict[str, Any] = ( - dict(key_row.metadata) if isinstance(key_row.metadata, dict) else {} - ) + current_metadata: Dict[str, Any] = dict(key_row.metadata) if isinstance(key_row.metadata, dict) else {} if blocked: new_metadata = {**current_metadata, SCIM_BLOCKED_METADATA_KEY: True} else: - new_metadata = { - k: v - for k, v in current_metadata.items() - if k != SCIM_BLOCKED_METADATA_KEY - } + new_metadata = {k: v for k, v in current_metadata.items() if k != SCIM_BLOCKED_METADATA_KEY} await VerificationTokenRepository(prisma_client).table.update( where={"token": key_row.token}, data={"blocked": blocked, "metadata": safe_dumps(new_metadata)}, @@ -575,12 +530,8 @@ async def _delete_rows_referencing_user(prisma_client: Any, *, user_id: str) -> ] } ) - await OrganizationMembershipRepository(prisma_client).table.delete_many( - where={"user_id": user_id} - ) - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"user_id": user_id} - ) + await OrganizationMembershipRepository(prisma_client).table.delete_many(where={"user_id": user_id}) + await TeamMembershipRepository(prisma_client).table.delete_many(where={"user_id": user_id}) def _scim_active_value(metadata: Optional[Dict[str, Any]]) -> Optional[bool]: @@ -593,9 +544,7 @@ def _scim_active_value(metadata: Optional[Dict[str, Any]]) -> Optional[bool]: return bool(value) -async def _create_user_if_not_exists( - user_id: str, created_via: str = "scim_group" -) -> Optional[NewUserResponse]: +async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_group") -> Optional[NewUserResponse]: """ Helper function to create a user if they don't exist. @@ -1040,9 +989,7 @@ async def get_users( where_conditions["user_email"] = filter_value # Get users from database - users: List[LiteLLM_UserTable] = await UserRepository( - prisma_client - ).table.find_many( + users: List[LiteLLM_UserTable] = await UserRepository(prisma_client).table.find_many( where=where_conditions, skip=(startIndex - 1), take=count, @@ -1050,16 +997,12 @@ async def get_users( ) # Get total count for pagination - total_count = await UserRepository(prisma_client).table.count( - where=where_conditions - ) + total_count = await UserRepository(prisma_client).table.count(where=where_conditions) # Convert to SCIM format scim_users: List[SCIMUser] = [] for user in users: - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user=user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user=user) scim_users.append(scim_user) return SCIMListResponse( @@ -1118,15 +1061,11 @@ async def create_user( # Check if user already exists if user.userName: - existing_user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user.userName} - ) + existing_user = await UserRepository(prisma_client).table.find_unique(where={"user_id": user.userName}) if existing_user: raise HTTPException( status_code=409, - detail={ - "error": f"User already exists with username: {user.userName}" - }, + detail={"error": f"User already exists with username: {user.userName}"}, ) # Create user in database @@ -1139,9 +1078,7 @@ async def create_user( default_role = _default_scim_user_role() admin_group = await _get_scim_admin_group() - resolved_role = _resolve_scim_user_role( - user.groups or [], admin_group, default_role - ) + resolved_role = _resolve_scim_user_role(user.groups or [], admin_group, default_role) new_user_request = NewUserRequest( user_id=user_id, @@ -1167,13 +1104,9 @@ async def create_user( data=new_user_request, ) - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user=created_user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user=created_user) return scim_user - except ( - HTTPException - ) as e: # allow exceptions like SCIMUserAlreadyExists to be raised + except HTTPException as e: # allow exceptions like SCIMUserAlreadyExists to be raised raise e except Exception as e: raise handle_exception_on_proxy(e) @@ -1212,9 +1145,7 @@ async def update_user( # SCIM active state when omitted — otherwise a vanilla PUT to a # deactivated user would silently re-enable them and unblock their keys. client_set_active = "active" in user.model_fields_set - scim_active_for_metadata = ( - user_data["active"] if client_set_active else prev_active - ) + scim_active_for_metadata = user_data["active"] if client_set_active else prev_active metadata = _build_scim_metadata( user_data["given_name"], @@ -1250,15 +1181,11 @@ async def update_user( if client_set_active: new_active = _scim_active_value(metadata) - if new_active is not None and new_active != ( - True if prev_active is None else prev_active - ): + if new_active is not None and new_active != (True if prev_active is None else prev_active): await _set_user_keys_blocked(user_id=user_id, blocked=not new_active) # Convert back to SCIM format - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - updated_user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(updated_user) return scim_user @@ -1286,9 +1213,7 @@ async def delete_user( teams = [] if existing_user.teams: for team_id in existing_user.teams: - team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if team: teams.append(team) @@ -1330,9 +1255,7 @@ def _extract_group_values(value: Any) -> List[str]: return group_values -def _handle_displayname_update( - op_type: str, value: Any, update_data: Dict[str, Any] -) -> None: +def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: """Handle displayname updates.""" if op_type == "remove": update_data["user_alias"] = None @@ -1340,9 +1263,7 @@ def _handle_displayname_update( update_data["user_alias"] = str(value) -def _handle_externalid_update( - op_type: str, value: Any, update_data: Dict[str, Any] -) -> None: +def _handle_externalid_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None: """Handle externalid updates.""" if op_type == "remove": update_data["sso_user_id"] = None @@ -1363,9 +1284,7 @@ def _handle_active_update(op_type: str, value: Any, metadata: Dict[str, Any]) -> metadata["scim_active"] = bool_val -def _handle_name_update( - path: str, op_type: str, value: Any, scim_metadata: Dict[str, Any] -) -> None: +def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict[str, Any]) -> None: """Handle name field updates (givenName, familyName).""" if path == "name.givenname": if op_type == "remove": @@ -1379,9 +1298,7 @@ def _handle_name_update( scim_metadata["familyName"] = str(value) -def _handle_group_operations( - op_type: str, value: Any, teams_set: Set[str] -) -> Optional[Set[str]]: +def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str]) -> Optional[Set[str]]: """Handle group/team membership operations.""" group_values = _extract_group_values(value) if op_type == "replace": @@ -1394,9 +1311,7 @@ def _handle_group_operations( return None -def _handle_generic_metadata( - path: str, op_type: str, value: Any, metadata: Dict[str, Any] -) -> None: +def _handle_generic_metadata(path: str, op_type: str, value: Any, metadata: Dict[str, Any]) -> None: """Handle generic metadata operations for unknown paths.""" if op_type == "remove": metadata.pop(path, None) @@ -1482,20 +1397,14 @@ async def patch_team_membership( team_id=_team_id, member=Member(user_id=user_id, role="user"), ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) except ProxyException as e: # Handle duplicate membership gracefully - this is idempotent if e.type == ProxyErrorTypes.team_member_already_in_team: - verbose_proxy_logger.debug( - f"User {user_id} is already in team {_team_id}, skipping add" - ) + verbose_proxy_logger.debug(f"User {user_id} is already in team {_team_id}, skipping add") else: - verbose_proxy_logger.exception( - f"Error adding user to team {_team_id}: {e}" - ) + verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") except Exception as e: verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}") @@ -1503,14 +1412,10 @@ async def patch_team_membership( try: await team_member_delete( data=TeamMemberDeleteRequest(team_id=_team_id, user_id=user_id), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) except Exception as e: - verbose_proxy_logger.exception( - f"Error removing user from team {_team_id}: {e}" - ) + verbose_proxy_logger.exception(f"Error removing user from team {_team_id}: {e}") return True @@ -1575,14 +1480,10 @@ async def patch_user( data=update_data, ) - if new_active is not None and new_active != ( - True if prev_active is None else prev_active - ): + if new_active is not None and new_active != (True if prev_active is None else prev_active): await _set_user_keys_blocked(user_id=user_id, blocked=not new_active) - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - updated_user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(updated_user) return scim_user @@ -1630,9 +1531,7 @@ async def get_groups( ) # Get total count for pagination - total_count = await TeamRepository(prisma_client).table.count( - where=where_conditions - ) + total_count = await TeamRepository(prisma_client).table.count(where=where_conditions) # Convert to SCIM format scim_groups = [] @@ -1685,9 +1584,7 @@ async def get_group( try: team = await _check_team_exists(group_id) - scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( - team - ) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(team) verbose_proxy_logger.debug(f"SCIM GET GROUP response: {scim_group}") return scim_group @@ -1718,9 +1615,7 @@ async def create_group( team_id = group.id or group.externalId or str(uuid.uuid4()) # Check if team already exists - existing_team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if existing_team: raise HTTPException( @@ -1730,10 +1625,7 @@ async def create_group( # Extract and validate group members (all users must exist) member_result = await _extract_group_member_ids(group) - members_with_roles = [ - Member(user_id=member_id, role="user") - for member_id in member_result.all_member_ids - ] + members_with_roles = [Member(user_id=member_id, role="user") for member_id in member_result.all_member_ids] # Create team in database created_team = await new_team( @@ -1748,9 +1640,7 @@ async def create_group( await _recompute_scim_member_roles(prisma_client, member_result.all_member_ids) - scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( - created_team - ) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(created_team) return scim_group except Exception as e: raise handle_exception_on_proxy(e) @@ -1780,12 +1670,8 @@ async def update_group( # Extract and validate group members (all users must exist) member_result = await _extract_group_member_ids(group) - verbose_proxy_logger.debug( - f"SCIM PUT GROUP all_member_ids: {member_result.all_member_ids}" - ) - verbose_proxy_logger.debug( - f"SCIM PUT GROUP created_users: {len(member_result.created_users)}" - ) + verbose_proxy_logger.debug(f"SCIM PUT GROUP all_member_ids: {member_result.all_member_ids}") + verbose_proxy_logger.debug(f"SCIM PUT GROUP created_users: {len(member_result.created_users)}") # Prepare update data existing_metadata = existing_team.metadata if existing_team.metadata else {} @@ -1820,17 +1706,11 @@ async def update_group( alias_changed = existing_team.team_alias != group.displayName await _recompute_scim_member_roles( prisma_client, - ( - current_members | final_members - if alias_changed - else current_members ^ final_members - ), + (current_members | final_members if alias_changed else current_members ^ final_members), ) # Convert to SCIM format and return - scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( - updated_team - ) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(updated_team) return scim_group except Exception as e: @@ -1857,9 +1737,7 @@ async def delete_group( # For each member, remove this team from their teams list for member_id in member_ids: - user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": member_id} - ) + user = await UserRepository(prisma_client).table.find_unique(where={"user_id": member_id}) if user: current_teams = user.teams or [] if group_id in current_teams: @@ -1924,9 +1802,7 @@ async def _process_group_patch_operations( detail={"error": "Invalid member: user ID cannot be empty."}, ) - user = await UserRepository(prisma_client).table.find_unique( - where={"user_id": member_id} - ) + user = await UserRepository(prisma_client).table.find_unique(where={"user_id": member_id}) if user: valid_members.append(member_id) else: @@ -1989,9 +1865,7 @@ async def _apply_group_patch_updates( return updated_team -async def _handle_group_membership_changes( - group_id: str, current_members: Set[str], final_members: Set[str] -): +async def _handle_group_membership_changes(group_id: str, current_members: Set[str], final_members: Set[str]): """Handle adding/removing members from the group.""" members_to_add = final_members - current_members members_to_remove = current_members - final_members @@ -2039,29 +1913,21 @@ async def patch_group( existing_team = await _check_team_exists(group_id) # Process patch operations - update_data, final_members = await _process_group_patch_operations( - patch_ops, existing_team, prisma_client - ) + update_data, final_members = await _process_group_patch_operations(patch_ops, existing_team, prisma_client) # Track current members BEFORE update for comparison current_members = set(await _get_team_member_user_ids_from_team(existing_team)) # Apply updates to the database - updated_team = await _apply_group_patch_updates( - group_id, update_data, final_members, prisma_client - ) + updated_team = await _apply_group_patch_updates(group_id, update_data, final_members, prisma_client) # Refresh team data from database to get the latest state after concurrent updates # This prevents race conditions when multiple PATCH requests come in simultaneously - refreshed_team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": group_id} - ) + refreshed_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) if refreshed_team: # Re-read current members from refreshed team to account for concurrent updates refreshed_current_members = set( - await _get_team_member_user_ids_from_team( - LiteLLM_TeamTable(**refreshed_team.model_dump()) - ) + await _get_team_member_user_ids_from_team(LiteLLM_TeamTable(**refreshed_team.model_dump())) ) # Use the refreshed members for comparison current_members = refreshed_current_members @@ -2076,17 +1942,11 @@ async def patch_group( alias_changed = new_alias != existing_team.team_alias await _recompute_scim_member_roles( prisma_client, - ( - current_members | final_members - if alias_changed - else current_members ^ final_members - ), + (current_members | final_members if alias_changed else current_members ^ final_members), ) # Refresh team one more time to get final state after membership changes - final_team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": group_id} - ) + final_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": group_id}) if final_team: updated_team = final_team diff --git a/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py index 04e44c623d1..00d021efd80 100644 --- a/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py +++ b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py @@ -56,30 +56,18 @@ class CustomMicrosoftSSO(MicrosoftSSO): Override to support custom endpoints via environment variables. Falls back to default Microsoft endpoints if not set. """ - custom_authorization_endpoint = os.getenv( - "MICROSOFT_AUTHORIZATION_ENDPOINT", None - ) + custom_authorization_endpoint = os.getenv("MICROSOFT_AUTHORIZATION_ENDPOINT", None) custom_token_endpoint = os.getenv("MICROSOFT_TOKEN_ENDPOINT", None) custom_userinfo_endpoint = os.getenv("MICROSOFT_USERINFO_ENDPOINT", None) # Use custom endpoints if set, otherwise use defaults authorization_endpoint = ( - custom_authorization_endpoint - or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/authorize" - ) - token_endpoint = ( - custom_token_endpoint - or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/token" - ) - userinfo_endpoint = ( - custom_userinfo_endpoint or f"https://graph.microsoft.com/{self.version}/me" + custom_authorization_endpoint or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/authorize" ) + token_endpoint = custom_token_endpoint or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/token" + userinfo_endpoint = custom_userinfo_endpoint or f"https://graph.microsoft.com/{self.version}/me" - if ( - custom_authorization_endpoint - or custom_token_endpoint - or custom_userinfo_endpoint - ): + if custom_authorization_endpoint or custom_token_endpoint or custom_userinfo_endpoint: verbose_proxy_logger.debug( f"Using custom Microsoft SSO endpoints - " f"authorization: {authorization_endpoint}, " diff --git a/litellm/proxy/management_endpoints/sso_helper_utils.py b/litellm/proxy/management_endpoints/sso_helper_utils.py index 7b296a6646f..fc10ab11658 100644 --- a/litellm/proxy/management_endpoints/sso_helper_utils.py +++ b/litellm/proxy/management_endpoints/sso_helper_utils.py @@ -19,9 +19,6 @@ def has_admin_ui_access(user_role: str) -> bool: bool: True if user is 'proxy_admin' or 'proxy_admin_view_only', False otherwise. """ - if ( - user_role != LitellmUserRoles.PROXY_ADMIN.value - and user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value - ): + if user_role != LitellmUserRoles.PROXY_ADMIN.value and user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: return False return True diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index f0bb8bdb5ff..3cf933ee84c 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -68,11 +68,7 @@ async def _get_internal_user_api_keys( where={"user_id": user_id}, select={"token": True}, ) - user_api_keys.update( - key_record.token - for key_record in key_records - if getattr(key_record, "token", None) - ) + user_api_keys.update(key_record.token for key_record in key_records if getattr(key_record, "token", None)) return sorted(user_api_keys) @@ -82,9 +78,7 @@ async def _get_tag_list_scope( user_api_key_dict: UserAPIKeyAuth, ) -> Optional[Dict[str, dict]]: user_role = user_api_key_dict.user_role - if user_api_key_has_admin_view(user_api_key_dict) or ( - user_role is None or not user_role.is_internal_user_role - ): + if user_api_key_has_admin_view(user_api_key_dict) or (user_role is None or not user_role.is_internal_user_role): return None scoped_api_keys = await _get_internal_user_api_keys( @@ -100,9 +94,7 @@ async def _get_tag_daily_activity_api_key_filter( requested_api_key: Optional[str], ) -> Optional[Union[str, List[str]]]: user_role = user_api_key_dict.user_role - if user_api_key_has_admin_view(user_api_key_dict) or ( - user_role is None or not user_role.is_internal_user_role - ): + if user_api_key_has_admin_view(user_api_key_dict) or (user_role is None or not user_role.is_internal_user_role): return requested_api_key scoped_api_keys = await _get_internal_user_api_keys( @@ -117,18 +109,14 @@ async def _get_tag_daily_activity_api_key_filter( async def _get_model_names(prisma_client, model_ids: list) -> Dict[str, str]: """Helper function to get model names from model IDs""" try: - models = await ModelRepository(prisma_client).table.find_many( - where={"model_id": {"in": model_ids}} - ) + models = await ModelRepository(prisma_client).table.find_many(where={"model_id": {"in": model_ids}}) return {model.model_id: model.model_name for model in models} except Exception as e: verbose_proxy_logger.error(f"Error getting model names: {str(e)}") return {} -async def get_deployments_by_model( - model: str, llm_router: "Router" -) -> List["Deployment"]: +async def get_deployments_by_model(model: str, llm_router: "Router") -> List["Deployment"]: """ Get all deployments by model """ @@ -188,22 +176,14 @@ async def new_tag( ) if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) if llm_router is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.no_llm_router.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.no_llm_router.value) try: # Check if tag already exists - existing_tag = await TagRepository(prisma_client).table.find_unique( - where={"tag_name": tag.name} - ) + existing_tag = await TagRepository(prisma_client).table.find_unique(where={"tag_name": tag.name}) if existing_tag is not None: - raise HTTPException( - status_code=400, detail=f"Tag {tag.name} already exists" - ) + raise HTTPException(status_code=400, detail=f"Tag {tag.name} already exists") # Handle budget creation/assignment using common helper budget_id = await handle_budget_for_entity( @@ -275,9 +255,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): try: # Get current model from database to preserve encrypted fields - db_model = await ModelRepository(prisma_client).table.find_unique( - where={"model_id": deployment.model_info.id} - ) + db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": deployment.model_info.id}) if db_model is None: raise HTTPException( @@ -343,9 +321,7 @@ async def update_tag( try: # Check if tag exists - existing_tag = await TagRepository(prisma_client).table.find_unique( - where={"tag_name": tag.name} - ) + existing_tag = await TagRepository(prisma_client).table.find_unique(where={"tag_name": tag.name}) if existing_tag is None: raise HTTPException(status_code=404, detail=f"Tag {tag.name} not found") @@ -431,9 +407,7 @@ async def info_tag( found_tag_names = {tag.tag_name for tag in tag_records} missing_tags = [name for name in data.names if name not in found_tag_names] if missing_tags: - raise HTTPException( - status_code=404, detail=f"Tags not found: {missing_tags}" - ) + raise HTTPException(status_code=404, detail=f"Tags not found: {missing_tags}") # Build response requested_tags = {} @@ -457,10 +431,7 @@ async def info_tag( } # Add budget info if available - if ( - hasattr(tag_record, "litellm_budget_table") - and tag_record.litellm_budget_table - ): + if hasattr(tag_record, "litellm_budget_table") and tag_record.litellm_budget_table: tag_dict["litellm_budget_table"] = tag_record.litellm_budget_table requested_tags[tag_record.tag_name] = tag_dict @@ -470,9 +441,7 @@ async def info_tag( raise HTTPException(status_code=500, detail=str(e)) -def _validate_tag_list_date_range( - start_date: Optional[str], end_date: Optional[str] -) -> None: +def _validate_tag_list_date_range(start_date: Optional[str], end_date: Optional[str]) -> None: """Require both dates together, and enforce YYYY-MM-DD format with start <= end.""" if (start_date is None) != (end_date is None): raise HTTPException( @@ -554,9 +523,7 @@ async def list_tags( if tag_scope is not None and not used_tag_names: return [] - stored_tag_where = ( - {"tag_name": {"in": used_tag_names}} if tag_scope is not None else None - ) + stored_tag_where = {"tag_name": {"in": used_tag_names}} if tag_scope is not None else None ## QUERY STORED TAGS ## tag_records = await TagRepository(prisma_client).table.find_many( @@ -587,10 +554,7 @@ async def list_tags( } # Add budget info if available - if ( - hasattr(tag_record, "litellm_budget_table") - and tag_record.litellm_budget_table - ): + if hasattr(tag_record, "litellm_budget_table") and tag_record.litellm_budget_table: tag_dict["litellm_budget_table"] = tag_record.litellm_budget_table list_of_tags.append(tag_dict) @@ -634,9 +598,7 @@ async def delete_tag( try: # Check if tag exists - existing_tag = await TagRepository(prisma_client).table.find_unique( - where={"tag_name": data.name} - ) + existing_tag = await TagRepository(prisma_client).table.find_unique(where={"tag_name": data.name}) if existing_tag is None: raise HTTPException(status_code=404, detail=f"Tag {data.name} not found") diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 94c84387be3..025b7c4210e 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -57,16 +57,10 @@ def _redact_callback_secrets(metadata: Any) -> Any: if isinstance(logging_entries, list): for entry in logging_entries: if isinstance(entry, dict) and isinstance(entry.get("callback_vars"), dict): - entry["callback_vars"] = { - k: _CALLBACK_VARS_REDACTED for k in entry["callback_vars"] - } + entry["callback_vars"] = {k: _CALLBACK_VARS_REDACTED for k in entry["callback_vars"]} callback_settings = redacted.get("callback_settings") - if isinstance(callback_settings, dict) and isinstance( - callback_settings.get("callback_vars"), dict - ): - callback_settings["callback_vars"] = { - k: _CALLBACK_VARS_REDACTED for k in callback_settings["callback_vars"] - } + if isinstance(callback_settings, dict) and isinstance(callback_settings.get("callback_vars"), dict): + callback_settings["callback_vars"] = {k: _CALLBACK_VARS_REDACTED for k in callback_settings["callback_vars"]} return redacted @@ -120,9 +114,7 @@ async def _emit_team_callback_audit_log( request_data=LiteLLM_AuditLogs( id=str(uuid.uuid4()), updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by - or user_api_key_dict.user_id - or litellm_proxy_admin_name, + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, changed_by_api_key=user_api_key_dict.api_key, table_name=LitellmTableNames.TEAM_TABLE_NAME, object_id=team_id, @@ -200,15 +192,11 @@ async def add_team_callbacks( ) # Check if team_id exists already - _existing_team = await prisma_client.get_data( - team_id=team_id, table_name="team", query_type="find_unique" - ) + _existing_team = await prisma_client.get_data(team_id=team_id, table_name="team", query_type="find_unique") if _existing_team is None: raise HTTPException( status_code=400, - detail={ - "error": f"Team id = {team_id} does not exist. Please use a different team id." - }, + detail={"error": f"Team id = {team_id} does not exist. Please use a different team id."}, ) # IDOR guard: only proxy admins / org admins / team admins of THIS @@ -222,12 +210,8 @@ async def add_team_callbacks( # store team callback settings in metadata team_metadata = _existing_team.metadata - team_callback_settings: List[dict] = team_metadata.get( - "logging" - ) # will be dict of type AddTeamCallback - if team_callback_settings is None or not isinstance( - team_callback_settings, list - ): + team_callback_settings: List[dict] = team_metadata.get("logging") # will be dict of type AddTeamCallback + if team_callback_settings is None or not isinstance(team_callback_settings, list): team_callback_settings = [] ## check if it already exists, for the same callback event @@ -274,9 +258,7 @@ async def add_team_callbacks( raise e except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.add_team_callbacks(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.add_team_callbacks(): Exception occured - {}".format(str(e)) ) raise ProxyException( message="Internal Server Error, " + str(e), @@ -322,9 +304,7 @@ async def disable_team_logging( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Check if team exists - _existing_team = await prisma_client.get_data( - team_id=team_id, table_name="team", query_type="find_unique" - ) + _existing_team = await prisma_client.get_data(team_id=team_id, table_name="team", query_type="find_unique") if _existing_team is None: raise HTTPException( status_code=404, @@ -363,9 +343,7 @@ async def disable_team_logging( if updated_team is None: raise HTTPException( status_code=404, - detail={ - "error": f"Team id = {team_id} does not exist. Error updating team logging" - }, + detail={"error": f"Team id = {team_id} does not exist. Error updating team logging"}, ) # Disabling a team's logging callbacks is itself a logging-control @@ -397,9 +375,7 @@ async def disable_team_logging( except ProxyException: raise except Exception as e: - verbose_proxy_logger.error( - f"litellm.proxy.proxy_server.disable_team_logging(): Exception occurred - {str(e)}" - ) + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.disable_team_logging(): Exception occurred - {str(e)}") verbose_proxy_logger.debug(traceback.format_exc()) raise ProxyException( message="Internal Server Error, " + str(e), @@ -451,9 +427,7 @@ async def get_team_callbacks( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Check if team_id exists - _existing_team = await prisma_client.get_data( - team_id=team_id, table_name="team", query_type="find_unique" - ) + _existing_team = await prisma_client.get_data(team_id=team_id, table_name="team", query_type="find_unique") if _existing_team is None: raise HTTPException( status_code=404, @@ -494,9 +468,7 @@ async def get_team_callbacks( raise except Exception as e: verbose_proxy_logger.error( - "litellm.proxy.proxy_server.get_team_callbacks(): Exception occurred - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.get_team_callbacks(): Exception occurred - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ac7d92aaa42..fb835535af8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -193,9 +193,7 @@ async def _verify_team_access( if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): return - if await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj): return raise HTTPException( @@ -253,9 +251,7 @@ class TeamMemberBudgetHandler: ) if data.team_alias is not None: - budget_id = ( - f"team-{data.team_alias.replace(' ', '-')}-budget-{uuid.uuid4().hex}" - ) + budget_id = f"team-{data.team_alias.replace(' ', '-')}-budget-{uuid.uuid4().hex}" else: budget_id = f"team-budget-{uuid.uuid4().hex}" @@ -282,9 +278,7 @@ class TeamMemberBudgetHandler: # Add team_member_budget_id as metadata field to team table if new_team_data_json.get("metadata") is None: new_team_data_json["metadata"] = {} - new_team_data_json["metadata"]["team_member_budget_id"] = ( - team_member_budget_table.budget_id - ) + new_team_data_json["metadata"]["team_member_budget_id"] = team_member_budget_table.budget_id # Remove team member fields from new_team_data_json TeamMemberBudgetHandler._clean_team_member_fields(new_team_data_json) @@ -416,9 +410,7 @@ class TeamMemberBudgetHandler: return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = await TeamMembershipRepository( - prisma_client - ).table.find_many(where={"team_id": team_id}) + existing_memberships = await TeamMembershipRepository(prisma_client).table.find_many(where={"team_id": team_id}) existing_user_ids = {m.user_id for m in existing_memberships} # Identify members with no existing membership row. @@ -548,22 +540,14 @@ def _check_team_model_specific_limits( for team in teams: if team.metadata and team.metadata.get("model_rpm_limit", None) is not None: for model, rpm_limit in team.metadata.get("model_rpm_limit", {}).items(): - model_specific_rpm_limit[model] = ( - model_specific_rpm_limit.get(model, 0) + rpm_limit - ) + model_specific_rpm_limit[model] = model_specific_rpm_limit.get(model, 0) + rpm_limit if team.metadata and team.metadata.get("model_tpm_limit", None) is not None: for model, tpm_limit in team.metadata.get("model_tpm_limit", {}).items(): - model_specific_tpm_limit[model] = ( - model_specific_tpm_limit.get(model, 0) + tpm_limit - ) + model_specific_tpm_limit[model] = model_specific_tpm_limit.get(model, 0) + tpm_limit if model_rpm_limit is not None: for model, rpm_limit in model_rpm_limit.items(): - if ( - entity_rpm_limit is not None - and model_specific_rpm_limit.get(model, 0) + rpm_limit - > entity_rpm_limit - ): + if entity_rpm_limit is not None and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_rpm_limit: raise HTTPException( status_code=400, detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Team RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}", @@ -572,8 +556,7 @@ def _check_team_model_specific_limits( entity_model_specific_rpm_limit = entity_model_rpm_limit_dict.get(model) if ( entity_model_specific_rpm_limit - and model_specific_rpm_limit.get(model, 0) + rpm_limit - > entity_model_specific_rpm_limit + and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_model_specific_rpm_limit ): raise HTTPException( status_code=400, @@ -582,11 +565,7 @@ def _check_team_model_specific_limits( if model_tpm_limit is not None: for model, tpm_limit in model_tpm_limit.items(): - if ( - entity_tpm_limit is not None - and model_specific_tpm_limit.get(model, 0) + tpm_limit - > entity_tpm_limit - ): + if entity_tpm_limit is not None and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_tpm_limit: raise HTTPException( status_code=400, detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Team TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}", @@ -595,8 +574,7 @@ def _check_team_model_specific_limits( entity_model_specific_tpm_limit = entity_model_tpm_limit_dict.get(model) if ( entity_model_specific_tpm_limit - and model_specific_tpm_limit.get(model, 0) + tpm_limit - > entity_model_specific_tpm_limit + and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_model_specific_tpm_limit ): raise HTTPException( status_code=400, @@ -616,12 +594,8 @@ def _check_team_rpm_tpm_limits( Raises an error if we're overallocating. """ if teams is not None and len(teams) > 0: - allocated_tpm = sum( - team.tpm_limit for team in teams if team.tpm_limit is not None - ) - allocated_rpm = sum( - team.rpm_limit for team in teams if team.rpm_limit is not None - ) + allocated_tpm = sum(team.tpm_limit for team in teams if team.tpm_limit is not None) + allocated_rpm = sum(team.rpm_limit for team in teams if team.rpm_limit is not None) else: allocated_tpm = 0 allocated_rpm = 0 @@ -781,10 +755,7 @@ async def _check_org_team_limits( data.metadata.get("tpm_limit_type", None) if data.metadata else None ) - if ( - tpm_limit_type != "guaranteed_throughput" - and rpm_limit_type != "guaranteed_throughput" - ): + if tpm_limit_type != "guaranteed_throughput" and rpm_limit_type != "guaranteed_throughput": return # get all organization teams # calculate allocated tpm/rpm limit @@ -838,11 +809,7 @@ async def _check_user_team_limits( user_id_upsert=False, ) - if ( - user_obj is not None - and user_obj.max_budget is not None - and data.max_budget > user_obj.max_budget - ): + if user_obj is not None and user_obj.max_budget is not None and data.max_budget > user_obj.max_budget: raise HTTPException( status_code=400, detail={ @@ -906,9 +873,7 @@ def _check_team_budget_update_authority( if existing_team_max_budget is None: return - budget_explicitly_set = "max_budget" in ( - getattr(data, "model_fields_set", None) or set() - ) + budget_explicitly_set = "max_budget" in (getattr(data, "model_fields_set", None) or set()) if budget_explicitly_set and data.max_budget is None: raise HTTPException( status_code=403, @@ -1038,14 +1003,10 @@ async def new_team( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Validate budget values are not negative - if data.max_budget is not None and ( - not math.isfinite(data.max_budget) or data.max_budget < 0 - ): + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) if data.team_member_budget is not None and ( not math.isfinite(data.team_member_budget) or data.team_member_budget < 0 @@ -1056,14 +1017,10 @@ async def new_team( "error": f"team_member_budget must be a non-negative finite number. Received: {data.team_member_budget}" }, ) - if data.soft_budget is not None and ( - not math.isfinite(data.soft_budget) or data.soft_budget < 0 - ): + if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}" - }, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) if data.soft_budget is not None: @@ -1079,9 +1036,7 @@ async def new_team( # Check if license is over limit total_teams = await TeamRepository(prisma_client).table.count() - if total_teams and _license_check.is_team_count_over_limit( - team_count=total_teams - ): + if total_teams and _license_check.is_team_count_over_limit(team_count=total_teams): raise HTTPException( status_code=403, detail="License is over limit. Please contact support@berri.ai to upgrade your license.", @@ -1097,9 +1052,7 @@ async def new_team( if _existing_team_id is not None: raise HTTPException( status_code=400, - detail={ - "error": f"Team id = {data.team_id} already exists. Please use a different team id." - }, + detail={"error": f"Team id = {data.team_id} already exists. Please use a different team id."}, ) # check org key limits - done here to handle inheriting org id from team @@ -1148,8 +1101,7 @@ async def new_team( data.max_budget = default_budget if ( - user_api_key_dict.user_role is None - or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + user_api_key_dict.user_role is None or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN ): # don't restrict proxy admin # Only validate user budget/models/tpm/rpm for standalone teams (not org-scoped) # For org-scoped teams, validation is done by _check_org_team_limits() @@ -1168,13 +1120,9 @@ async def new_team( creating_user_in_list = True if creating_user_in_list is False: - data.members_with_roles.append( - Member(role="admin", user_id=user_api_key_dict.user_id) - ) + data.members_with_roles.append(Member(role="admin", user_id=user_api_key_dict.user_id)) - _check_passthrough_routes_caller_permission( - data, user_api_key_dict, entity="team" - ) + _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") ## ADD TO MODEL TABLE _model_id = None @@ -1255,9 +1203,7 @@ async def new_team( initialized_windows = [] for window in complete_team_data.budget_limits: w = window if isinstance(window, dict) else window.model_dump() - w["reset_at"] = get_budget_reset_time( - budget_duration=w["budget_duration"] - ).isoformat() + w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat() initialized_windows.append(w) complete_team_data.budget_limits = initialized_windows # type: ignore[assignment] @@ -1272,20 +1218,14 @@ async def new_team( # Serialize router_settings to JSON (matching key creation pattern) router_settings_value = getattr(data, "router_settings", None) router_settings_json = ( - safe_dumps(router_settings_value) - if router_settings_value is not None - else safe_dumps({}) + safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({}) ) complete_team_data_dict["router_settings"] = router_settings_json if complete_team_data_dict.get("metadata") is not None: - complete_team_data_dict["metadata"] = encrypt_callback_vars( - complete_team_data_dict["metadata"] - ) + complete_team_data_dict["metadata"] = encrypt_callback_vars(complete_team_data_dict["metadata"]) - complete_team_data_dict = prisma_client.jsonify_team_object( - db_data=complete_team_data_dict - ) + complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) team_row: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.create( data=complete_team_data_dict, @@ -1436,9 +1376,7 @@ async def _auto_add_team_members_to_organization( are never explicitly added to organizations. This silently upserts missing members rather than blocking the team move. """ - org_member_ids = ( - {m.user_id for m in organization.members} if organization.members else set() - ) + org_member_ids = {m.user_id for m in organization.members} if organization.members else set() for member in team.members_with_roles: if member.user_id is None: continue @@ -1488,9 +1426,7 @@ async def fetch_and_validate_organization( HTTPException: If llm_router is None, organization not found, or validation fails """ if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) organization_row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": organization_id}, @@ -1500,15 +1436,10 @@ async def fetch_and_validate_organization( if organization_row is None: raise HTTPException( status_code=404, - detail={ - "error": f"Organization not found, passed organization_id={organization_id}" - }, + detail={"error": f"Organization not found, passed organization_id={organization_id}"}, ) - is_proxy_admin = ( - user_api_key_dict is not None - and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - ) + is_proxy_admin = user_api_key_dict is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN organization = LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()) validate_team_org_change( team=LiteLLM_TeamTable(**existing_team_row.model_dump()), @@ -1590,13 +1521,9 @@ def validate_team_org_change( # thereby injecting members into that org without org admin approval. if not is_proxy_admin: team_members = [m.user_id for m in team.members_with_roles] - org_members = ( - [m.user_id for m in organization.members] if organization.members else [] - ) + org_members = [m.user_id for m in organization.members] if organization.members else [] not_in_org = [ - m - for m in team_members - if m not in org_members and m != SpecialProxyStrings.default_user_id.value + m for m in team_members if m not in org_members and m != SpecialProxyStrings.default_user_id.value ] if len(not_in_org) > 0: raise HTTPException( @@ -1634,9 +1561,7 @@ def validate_team_org_change( return True -@router.post( - "/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def update_team( data: UpdateTeamRequest, @@ -1727,20 +1652,14 @@ async def update_team( ) if data.team_id is None: - raise HTTPException( - status_code=400, detail={"error": "No team id passed in"} - ) + raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) verbose_proxy_logger.debug("/team/update - %s", data) # Validate budget values are not negative - if data.max_budget is not None and ( - not math.isfinite(data.max_budget) or data.max_budget < 0 - ): + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}" - }, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, ) if data.team_member_budget is not None and ( not math.isfinite(data.team_member_budget) or data.team_member_budget < 0 @@ -1751,19 +1670,13 @@ async def update_team( "error": f"team_member_budget must be a non-negative finite number. Received: {data.team_member_budget}" }, ) - if data.soft_budget is not None and ( - not math.isfinite(data.soft_budget) or data.soft_budget < 0 - ): + if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): raise HTTPException( status_code=400, - detail={ - "error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}" - }, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": data.team_id} - ) + existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) if existing_team_row is None: raise HTTPException( @@ -1777,16 +1690,10 @@ async def update_team( user_api_key_dict=user_api_key_dict, ) - _check_passthrough_routes_caller_permission( - data, user_api_key_dict, entity="team" - ) + _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") if data.soft_budget is not None: - max_budget_to_check = ( - data.max_budget - if data.max_budget is not None - else existing_team_row.max_budget - ) + max_budget_to_check = data.max_budget if data.max_budget is not None else existing_team_row.max_budget if max_budget_to_check is not None: if data.soft_budget >= max_budget_to_check: raise HTTPException( @@ -1798,14 +1705,8 @@ async def update_team( if data.max_budget is not None: existing_soft_budget = getattr(existing_team_row, "soft_budget", None) - soft_budget_to_check = ( - data.soft_budget - if data.soft_budget is not None - else existing_soft_budget - ) - if soft_budget_to_check is not None and isinstance( - soft_budget_to_check, (int, float) - ): + soft_budget_to_check = data.soft_budget if data.soft_budget is not None else existing_soft_budget + if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): if data.max_budget <= soft_budget_to_check: raise HTTPException( status_code=400, @@ -1814,9 +1715,7 @@ async def update_team( }, ) - if ( - data.organization_id is not None and len(data.organization_id) > 0 - ): # allow unsetting the organization_id + if data.organization_id is not None and len(data.organization_id) > 0: # allow unsetting the organization_id # If the caller is relocating the team to a different org, they # must also be PROXY_ADMIN or an org-admin of the DESTINATION org. # _verify_team_access above only checked the team's CURRENT org, @@ -1830,9 +1729,7 @@ async def update_team( ): # Is the caller org_admin of the destination org? caller_memberships = ( - await OrganizationMembershipRepository( - prisma_client - ).table.find_many( + await OrganizationMembershipRepository(prisma_client).table.find_many( where={ "user_id": user_api_key_dict.user_id, "organization_id": data.organization_id, @@ -1867,15 +1764,9 @@ async def update_team( # check org team limits - if updating team that belongs to an org org_id_to_check = ( - data.organization_id - if data.organization_id is not None - else existing_team_row.organization_id + data.organization_id if data.organization_id is not None else existing_team_row.organization_id ) - if ( - org_id_to_check is not None - and isinstance(org_id_to_check, str) - and prisma_client is not None - ): + if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None: org_table = await get_org_object( org_id=org_id_to_check, user_api_key_cache=user_api_key_cache, @@ -1902,9 +1793,7 @@ async def update_team( # Drop server-owned metadata keys from caller input so they can only # be written by the same code path that creates the underlying rows. if isinstance(updated_kv.get("metadata"), dict): - TeamMemberBudgetHandler.strip_system_managed_metadata_keys( - updated_kv["metadata"] - ) + TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) @@ -1920,14 +1809,11 @@ async def update_team( if field in updated_kv } - if ( - _team_member_fields_in_request - and TeamMemberBudgetHandler.should_create_budget( - team_member_budget=data.team_member_budget, - team_member_rpm_limit=data.team_member_rpm_limit, - team_member_tpm_limit=data.team_member_tpm_limit, - team_member_budget_duration=data.team_member_budget_duration, - ) + if _team_member_fields_in_request and TeamMemberBudgetHandler.should_create_budget( + team_member_budget=data.team_member_budget, + team_member_rpm_limit=data.team_member_rpm_limit, + team_member_tpm_limit=data.team_member_tpm_limit, + team_member_budget_duration=data.team_member_budget_duration, ): updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=existing_team_row, @@ -1940,9 +1826,7 @@ async def update_team( ) # Backfill team_memberships for members who joined before the # budget was configured — they won't have a membership row yet. - _backfill_budget_id = (updated_kv.get("metadata") or {}).get( - "team_member_budget_id" - ) + _backfill_budget_id = (updated_kv.get("metadata") or {}).get("team_member_budget_id") if _backfill_budget_id and existing_team_row.members_with_roles: await TeamMemberBudgetHandler.backfill_team_member_budget_entries( team_id=data.team_id, @@ -1986,16 +1870,11 @@ async def update_team( updated_kv["model_id"] = _model_id # Serialize router_settings to JSON if present (matching key update pattern) - if ( - "router_settings" in updated_kv - and updated_kv["router_settings"] is not None - ): + if "router_settings" in updated_kv and updated_kv["router_settings"] is not None: updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = await TeamRepository( - prisma_client - ).table.update( + team_row: Optional[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data=updated_kv, # `object_permission` is included so `_refresh_cached_team` @@ -2013,9 +1892,7 @@ async def update_team( detail={"error": "Team doesn't exist. Got={}".format(team_row)}, ) - verbose_proxy_logger.info( - "Successfully updated team - %s, info", team_row.team_id - ) + verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) await _refresh_cached_team( team_row=team_row, user_api_key_cache=user_api_key_cache, @@ -2054,16 +1931,12 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: initialized_windows = [] for window in data.budget_limits: w = window if isinstance(window, dict) else window.model_dump() - w["reset_at"] = get_budget_reset_time( - budget_duration=w["budget_duration"] - ).isoformat() + w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat() initialized_windows.append(w) updated_kv["budget_limits"] = json.dumps(initialized_windows) -async def handle_update_object_permission( - data_json: dict, existing_team_row: LiteLLM_TeamTable -) -> dict: +async def handle_update_object_permission(data_json: dict, existing_team_row: LiteLLM_TeamTable) -> dict: """ Handle the update of object permission for a team. @@ -2082,9 +1955,7 @@ async def handle_update_object_permission( # Add the object_permission_id to data_json if one was created/updated if object_permission_id is not None: data_json["object_permission_id"] = object_permission_id - verbose_proxy_logger.debug( - f"updated object_permission_id: {object_permission_id}" - ) + verbose_proxy_logger.debug(f"updated object_permission_id: {object_permission_id}") return data_json @@ -2095,9 +1966,7 @@ def _check_team_member_admin_add( ): if isinstance(member, Member) and member.role == "admin": if premium_user is not True: - raise ValueError( - f"Assigning team admins is a premium feature. {CommonProxyErrors.not_premium_user.value}" - ) + raise ValueError(f"Assigning team admins is a premium feature. {CommonProxyErrors.not_premium_user.value}") elif isinstance(member, List): for m in member: if m.role == "admin": @@ -2119,9 +1988,7 @@ def team_call_validation_checks( raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) if data.member is None: - raise HTTPException( - status_code=400, detail={"error": "No member/members passed in"} - ) + raise HTTPException(status_code=400, detail={"error": "No member/members passed in"}) try: _check_team_member_admin_add( @@ -2198,18 +2065,11 @@ async def _validate_team_member_add_permissions( the request matches the caller's own ``user_id`` and is being added with ``role="user"``. """ - if ( - getattr(user_api_key_dict, "user_role", None) - == LitellmUserRoles.PROXY_ADMIN.value - ): + if getattr(user_api_key_dict, "user_role", None) == LitellmUserRoles.PROXY_ADMIN.value: return - if _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ): + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=complete_team_data): return - if await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ): + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=complete_team_data): return if not _is_available_team( @@ -2230,11 +2090,7 @@ async def _validate_team_member_add_permissions( # budget and model controls stay admin-only. Reject them here so a # self-joining non-admin cannot set their own cap, reset window, or model # scope via the bypass. - if ( - data.max_budget_in_team is not None - or data.budget_duration is not None - or data.allowed_models is not None - ): + if data.max_budget_in_team is not None or data.budget_duration is not None or data.allowed_models is not None: raise HTTPException( status_code=403, detail={ @@ -2287,16 +2143,12 @@ async def _process_team_members( updated_team_memberships: List[LiteLLM_TeamMembership] = [] default_team_budget_id = ( - complete_team_data.metadata.get("team_member_budget_id") - if complete_team_data.metadata is not None - else None + complete_team_data.metadata.get("team_member_budget_id") if complete_team_data.metadata is not None else None ) # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models member_allowed_models = data.allowed_models - team_default_member_models = getattr( - complete_team_data, "default_team_member_models", None - ) + team_default_member_models = getattr(complete_team_data, "default_team_member_models", None) if member_allowed_models is None and team_default_member_models: member_allowed_models = team_default_member_models @@ -2367,21 +2219,14 @@ async def _update_team_members_list( # get user id if new_member.user_id is None and new_member.user_email is not None: for user in updated_users: - if ( - user.user_email is not None - and user.user_email == new_member.user_email - ): + if user.user_email is not None and user.user_email == new_member.user_email: new_member.user_id = user.user_id # Check if member already exists in team before adding member_already_exists = False for existing_member in complete_team_data.members_with_roles: - if ( - new_member.user_id is not None - and existing_member.user_id == new_member.user_id - ) or ( - new_member.user_email is not None - and existing_member.user_email == new_member.user_email + if (new_member.user_id is not None and existing_member.user_id == new_member.user_id) or ( + new_member.user_email is not None and existing_member.user_email == new_member.user_email ): member_already_exists = True break @@ -2399,11 +2244,8 @@ async def _update_team_members_list( # Check if member already exists in team before adding member_already_exists = False for existing_member in complete_team_data.members_with_roles: - if ( - nm.user_id is not None and existing_member.user_id == nm.user_id - ) or ( - nm.user_email is not None - and existing_member.user_email == nm.user_email + if (nm.user_id is not None and existing_member.user_id == nm.user_id) or ( + nm.user_email is not None and existing_member.user_email == nm.user_email ): member_already_exists = True break @@ -2458,9 +2300,7 @@ def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: try: prometheus_logger.set_team_members_metric(team) except Exception as e: - verbose_proxy_logger.debug( - "Prometheus: failed to emit team members metric: %s", str(e) - ) + verbose_proxy_logger.debug("Prometheus: failed to emit team members metric: %s", str(e)) async def _validate_and_populate_member_user_info( @@ -2495,9 +2335,7 @@ async def _validate_and_populate_member_user_info( query_type="find_all", ) - if users_by_email is None or ( - isinstance(users_by_email, list) and len(users_by_email) == 0 - ): + if users_by_email is None or (isinstance(users_by_email, list) and len(users_by_email) == 0): # User doesn't exist yet - this is fine, will be created later return member @@ -2541,11 +2379,7 @@ async def _validate_and_populate_member_user_info( query_type="find_all", ) - if ( - users_by_email - and isinstance(users_by_email, list) - and len(users_by_email) > 1 - ): + if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1: raise HTTPException( status_code=400, detail={ @@ -2559,9 +2393,7 @@ async def _validate_and_populate_member_user_info( # Case 3: Only user_id provided - populate user_email from DB if user exists if member.user_id is not None and member.user_email is None: - user_by_id = await UserRepository(prisma_client).table.find_unique( - where={"user_id": member.user_id} - ) + user_by_id = await UserRepository(prisma_client).table.find_unique(where={"user_id": member.user_id}) if user_by_id is None: # User doesn't exist yet - allow it to pass with user_email as None @@ -2634,9 +2466,7 @@ async def team_member_add( if existing_team_row is None: raise HTTPException( status_code=404, - detail={ - "error": f"Team not found for team_id={getattr(data, 'team_id', None)}" - }, + detail={"error": f"Team not found for team_id={getattr(data, 'team_id', None)}"}, ) complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) @@ -2680,9 +2510,7 @@ async def team_member_add( # Check if updated_team is None if updated_team is None: - raise HTTPException( - status_code=404, detail={"error": f"Team with id {data.team_id} not found"} - ) + raise HTTPException(status_code=404, detail={"error": f"Team with id {data.team_id} not found"}) _emit_team_members_metric(complete_team_data) @@ -2701,18 +2529,10 @@ def _cleanup_members_with_roles( is_member_in_team = False new_team_members: List[Member] = [] for m in existing_team_row.members_with_roles: - if ( - data.user_id is not None - and m.user_id is not None - and data.user_id == m.user_id - ): + if data.user_id is not None and m.user_id is not None and data.user_id == m.user_id: is_member_in_team = True continue - elif ( - data.user_email is not None - and m.user_email is not None - and data.user_email == m.user_email - ): + elif data.user_email is not None and m.user_email is not None and data.user_email == m.user_email: is_member_in_team = True continue new_team_members.append(m) @@ -2762,9 +2582,7 @@ async def team_member_delete( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": data.team_id} - ) + _existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) if _existing_team_row is None: raise HTTPException( @@ -2777,12 +2595,8 @@ async def team_member_delete( if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=existing_team_row - ) - and not await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=existing_team_row - ) + and not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=existing_team_row) + and not await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=existing_team_row) ): raise HTTPException( status_code=403, @@ -2826,9 +2640,7 @@ async def team_member_delete( where=key_val # type: ignore ) - if existing_user_rows is not None and ( - isinstance(existing_user_rows, list) and len(existing_user_rows) > 0 - ): + if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): for existing_user in existing_user_rows: team_list = [] if data.team_id in existing_user.teams: @@ -2862,9 +2674,9 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -2974,9 +2786,7 @@ async def team_member_update( _validate_budget_duration(data.budget_duration) - _existing_team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": data.team_id} - ) + _existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) if _existing_team_row is None: raise HTTPException( @@ -2989,12 +2799,8 @@ async def team_member_update( if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=existing_team_row - ) - and not await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=existing_team_row - ) + and not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=existing_team_row) + and not await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=existing_team_row) ): raise HTTPException( status_code=403, @@ -3027,9 +2833,7 @@ async def team_member_update( if received_user_id is None: raise HTTPException( status_code=400, - detail={ - "error": "User id doesn't exist in team table. Data={}".format(data) - }, + detail={"error": "User id doesn't exist in team table. Data={}".format(data)}, ) ## find the relevant team membership identified_budget_id: Optional[str] = None @@ -3209,9 +3013,7 @@ async def bulk_team_member_add( }, ) # get all users from the database - all_users_in_db = await UserRepository(prisma_client).table.find_many( - order={"created_at": "desc"} - ) + all_users_in_db = await UserRepository(prisma_client).table.find_many(order={"created_at": "desc"}) data.members = [ Member( user_id=user.user_id, @@ -3282,9 +3084,7 @@ async def bulk_team_member_add( ) -@router.post( - "/team/delete", tags=["team management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/team/delete", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def delete_team( data: DeleteTeamRequest, @@ -3329,9 +3129,9 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[BaseModel] = await TeamRepository( - prisma_client - ).table.find_unique(where={"team_id": team_id}) + team_row_base: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": team_id} + ) if team_row_base is None: raise Exception except Exception: @@ -3398,9 +3198,9 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"team_id": {"in": data.team_ids}}) + keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + where={"team_id": {"in": data.team_ids}} + ) if keys_to_delete: await _persist_deleted_verification_tokens( @@ -3446,9 +3246,7 @@ async def delete_team( await asyncio.gather(*tasks) ## DELETE TEAMS - deleted_teams = await prisma_client.delete_data( - team_id_list=data.team_ids, table_name="team" - ) + deleted_teams = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team") return deleted_teams @@ -3526,18 +3324,14 @@ async def _persist_deleted_team_records( ) -async def validate_membership( - user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable -): +async def validate_membership(user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable): if ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value ): return - if ( - user_api_key_dict.team_id == team_table.team_id - ): # allow team keys to check their info + if user_api_key_dict.team_id == team_table.team_id: # allow team keys to check their info return # Handle case where user_id is None (e.g., team key accessing different team) @@ -3566,9 +3360,7 @@ async def validate_membership( return # Check if user is an org admin for the team's organization - if await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=team_table - ): + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_table): return raise HTTPException( @@ -3616,18 +3408,12 @@ async def _resolve_team_access_group_resources(_team_info: Any) -> None: _team_info.access_group_agent_ids = list(agent_ids) -@router.get( - "/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)] -) +@router.get("/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def team_info( http_request: Request, - team_id: str = fastapi.Query( - default=None, description="Team ID in the request parameters" - ), - key_limit: int | None = fastapi.Query( - default=None, description="Limit the number of keys returned", gt=0 - ), + team_id: str = fastapi.Query(default=None, description="Team ID in the request parameters"), + key_limit: int | None = fastapi.Query(default=None, description="Limit the number of keys returned", gt=0), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -3659,9 +3445,7 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = await TeamRepository( - prisma_client - ).table.find_unique( + team_info: Optional[BaseModel] = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, include={"object_permission": True}, ) @@ -3706,9 +3490,7 @@ async def team_info( key.pop("token", None) ## GET ALL MEMBERSHIPS ## - returned_tm = await get_all_team_memberships( - prisma_client, [team_id], user_id=None - ) + returned_tm = await get_all_team_memberships(prisma_client, [team_id], user_id=None) if isinstance(team_info, dict): _team_info = TeamInfoResponseObjectTeamTable(**team_info) @@ -3719,9 +3501,7 @@ async def team_info( ## GET TEAM BUDGET (if exists) ## team_member_budget_id = ( - _team_info.metadata.get("team_member_budget_id") - if _team_info.metadata is not None - else None + _team_info.metadata.get("team_member_budget_id") if _team_info.metadata is not None else None ) if team_member_budget_id is not None: _team_info = await _add_team_member_budget_table( @@ -3806,9 +3586,7 @@ async def team_member_me( # Team keys / service-account keys without a user_id can't resolve "me". raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "API key has no associated user_id; cannot resolve 'me' for team membership." - }, + detail={"error": "API key has no associated user_id; cannot resolve 'me' for team membership."}, ) team_table = await get_team_object( @@ -3823,9 +3601,7 @@ async def team_member_me( # Match by user_id when present, else fall back to email — members # added by email may have user_id=None on the stored entry. if (m.user_id is not None and m.user_id == caller_user_id) or ( - m.user_email is not None - and caller_user_email is not None - and m.user_email == caller_user_email + m.user_email is not None and caller_user_email is not None and m.user_email == caller_user_email ): member_role = m.role break @@ -3836,9 +3612,7 @@ async def team_member_me( # actual members of the team. raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail={ - "error": f"User user_id={caller_user_id} is not a member of team_id={team_id}." - }, + detail={"error": f"User user_id={caller_user_id} is not a member of team_id={team_id}."}, ) membership = await get_team_membership( @@ -3884,9 +3658,7 @@ async def team_member_me( ) -@router.post( - "/team/block", tags=["team management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/team/block", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def block_team( data: BlockTeamRequest, @@ -3920,9 +3692,7 @@ async def block_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": data.team_id} - ) + existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) if existing_team is None: raise HTTPException( status_code=404, @@ -3943,9 +3713,7 @@ async def block_team( return record -@router.post( - "/team/unblock", tags=["team management"], dependencies=[Depends(user_api_key_auth)] -) +@router.post("/team/unblock", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def unblock_team( data: BlockTeamRequest, @@ -3973,9 +3741,7 @@ async def unblock_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": data.team_id} - ) + existing_team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) if existing_team is None: raise HTTPException( status_code=404, @@ -4022,9 +3788,7 @@ async def list_available_teams( return [] # filter out teams that the user is already a member of - user_info = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_api_key_dict.user_id} - ) + user_info = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_api_key_dict.user_id}) if user_info is None: raise HTTPException( status_code=404, @@ -4032,17 +3796,11 @@ async def list_available_teams( ) user_info_correct_type = LiteLLM_UserTable(**user_info.model_dump()) - available_teams = [ - team for team in available_teams if team not in user_info_correct_type.teams - ] + available_teams = [team for team in available_teams if team not in user_info_correct_type.teams] - available_teams_db = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": available_teams}} - ) + available_teams_db = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": available_teams}}) - available_teams_correct_type = [ - LiteLLM_TeamTable(**team.model_dump()) for team in available_teams_db - ] + available_teams_correct_type = [LiteLLM_TeamTable(**team.model_dump()) for team in available_teams_db] return available_teams_correct_type @@ -4076,8 +3834,7 @@ async def _get_org_admin_org_ids( org_ids = [ m.organization_id for m in (caller_user.organization_memberships or []) - if m.user_role == LitellmUserRoles.ORG_ADMIN.value - and m.organization_id is not None + if m.user_role == LitellmUserRoles.ORG_ADMIN.value and m.organization_id is not None ] return org_ids if org_ids else None @@ -4202,9 +3959,7 @@ def _convert_teams_to_response_models( keys_count_by_team: Optional[Dict[str, int]] = None, ) -> List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]: """Convert raw Prisma team rows to response models.""" - team_list: List[ - Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable] - ] = [] + team_list: List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] = [] counts = keys_count_by_team or {} for team in teams: try: @@ -4241,9 +3996,7 @@ async def _get_keys_count_by_team( bounded by page_size and uses the existing @@index([team_id]), so this is one DB round-trip per page. Returns an empty map when the page has no teams. """ - page_team_ids = [ - getattr(t, "team_id", None) for t in teams if getattr(t, "team_id", None) - ] + page_team_ids = [getattr(t, "team_id", None) for t in teams if getattr(t, "team_id", None)] if not page_team_ids: return {} @@ -4252,11 +4005,7 @@ async def _get_keys_count_by_team( where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, ) - return { - row["team_id"]: row.get("_count", {}).get("team_id", 0) - for row in grouped - if row.get("team_id") - } + return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")} async def _enforce_list_team_v2_access( @@ -4314,9 +4063,7 @@ async def _enforce_list_team_v2_access( ) else: # Not an org admin — fall back to standard route check - if not allowed_route_check_inside_route( - user_api_key_dict=user_api_key_dict, requested_user_id=user_id - ): + if not allowed_route_check_inside_route(user_api_key_dict=user_api_key_dict, requested_user_id=user_id): raise HTTPException( status_code=401, detail={ @@ -4359,22 +4106,14 @@ async def list_team_v2( default=None, description="Combined search: matches teams whose 'team_id' equals the value OR whose 'team_alias' contains it (case-insensitive).", ), - page: int = fastapi.Query( - default=1, description="Page number for pagination", ge=1 - ), - page_size: int = fastapi.Query( - default=10, description="Number of teams per page", ge=1, le=100 - ), + page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), + page_size: int = fastapi.Query(default=10, description="Number of teams per page", ge=1, le=100), sort_by: Optional[str] = fastapi.Query( default=None, description="Column to sort by (e.g. 'team_id', 'team_alias', 'created_at')", ), - sort_order: str = fastapi.Query( - default="asc", description="Sort order ('asc' or 'desc')" - ), - status: Optional[str] = fastapi.Query( - default=None, description="Filter by status (e.g. 'deleted')" - ), + sort_order: str = fastapi.Query(default="asc", description="Sort order ('asc' or 'desc')"), + status: Optional[str] = fastapi.Query(default=None, description="Filter by status (e.g. 'deleted')"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -4425,9 +4164,7 @@ async def list_team_v2( if status is not None and status != "deleted": raise HTTPException( status_code=400, - detail={ - "error": "Invalid status value. Currently only 'deleted' is supported." - }, + detail={"error": "Invalid status value. Currently only 'deleted' is supported."}, ) use_deleted_table = status == "deleted" @@ -4476,9 +4213,7 @@ async def list_team_v2( order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await DeletedTeamRepository(prisma_client).table.count( - where=where_conditions - ) + total_count = await DeletedTeamRepository(prisma_client).table.count(where=where_conditions) else: teams = await TeamRepository(prisma_client).table.find_many( where=where_conditions, @@ -4487,9 +4222,7 @@ async def list_team_v2( order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await TeamRepository(prisma_client).table.count( - where=where_conditions - ) + total_count = await TeamRepository(prisma_client).table.count(where=where_conditions) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division @@ -4501,21 +4234,13 @@ async def list_team_v2( keys_count_by_team = await _get_keys_count_by_team(prisma_client, teams) # Convert Prisma models to response models with members_count and keys_count - team_list = _convert_teams_to_response_models( - teams, use_deleted_table, keys_count_by_team=keys_count_by_team - ) + team_list = _convert_teams_to_response_models(teams, use_deleted_table, keys_count_by_team=keys_count_by_team) # Resolve resources inherited from access groups (single batch query) if not use_deleted_table: - team_items_with_ag = [ - t for t in team_list if isinstance(t, TeamListItem) and t.access_group_ids - ] + team_items_with_ag = [t for t in team_list if isinstance(t, TeamListItem) and t.access_group_ids] if team_items_with_ag: - all_ag_ids = [ - ag_id - for t in team_items_with_ag - for ag_id in (t.access_group_ids or []) - ] + all_ag_ids = [ag_id for t in team_items_with_ag for ag_id in (t.access_group_ids or [])] ag_lookup = await _batch_resolve_access_group_resources(all_ag_ids) for team_item in team_items_with_ag: models, mcp_ids, agent_ids = set(), set(), set() @@ -4557,9 +4282,7 @@ async def _authorize_and_filter_teams( if not is_proxy_admin: is_own_query = ( - user_id is not None - and user_api_key_dict.user_id is not None - and user_api_key_dict.user_id == user_id + user_id is not None and user_api_key_dict.user_id is not None and user_api_key_dict.user_id == user_id ) # Check if user is an org admin (even for own queries, so they see org teams) @@ -4575,8 +4298,7 @@ async def _authorize_and_filter_teams( allowed_org_ids = [ m.organization_id for m in (caller_user.organization_memberships or []) - if m.user_role == LitellmUserRoles.ORG_ADMIN.value - and m.organization_id is not None + if m.user_role == LitellmUserRoles.ORG_ADMIN.value and m.organization_id is not None ] if not allowed_org_ids: allowed_org_ids = None @@ -4603,32 +4325,22 @@ async def _authorize_and_filter_teams( return [ team for team in org_teams - if team.members_with_roles - and any(m.get("user_id") == user_id for m in team.members_with_roles) + if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles) ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response = await TeamRepository(prisma_client).table.find_many( - include={"litellm_model_table": True} - ) + response = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}) return [ team for team in response - if team.members_with_roles - and any(m.get("user_id") == user_id for m in team.members_with_roles) + if team.members_with_roles and any(m.get("user_id") == user_id for m in team.members_with_roles) ] else: # Proxy admin: all teams - return list( - await TeamRepository(prisma_client).table.find_many( - include={"litellm_model_table": True} - ) - ) + return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) -@router.get( - "/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)] -) +@router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def list_team( http_request: Request, @@ -4669,9 +4381,7 @@ async def list_team( ) _team_ids = [team.team_id for team in filtered_response] - returned_tm = await get_all_team_memberships( - prisma_client, _team_ids, user_id=user_id - ) + returned_tm = await get_all_team_memberships(prisma_client, _team_ids, user_id=user_id) returned_responses: List[TeamListResponseObject] = [] for team in filtered_response: @@ -4681,9 +4391,7 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await VerificationTokenRepository(prisma_client).table.find_many( - where={"team_id": team.team_id} - ) + keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id}) try: returned_responses.append( @@ -4704,15 +4412,9 @@ async def list_team( if organization_id is not None: if organization_id == SpecialManagementEndpointEnums.DEFAULT_ORGANIZATION.value: - returned_responses = [ - team for team in returned_responses if team.organization_id is None - ] + returned_responses = [team for team in returned_responses if team.organization_id is None] else: - returned_responses = [ - team - for team in returned_responses - if team.organization_id == organization_id - ] + returned_responses = [team for team in returned_responses if team.organization_id == organization_id] return returned_responses @@ -4747,9 +4449,7 @@ async def get_paginated_teams( ) return teams, total_count except Exception as e: - verbose_proxy_logger.exception( - f"[Non-Blocking] Error getting paginated teams: {e}" - ) + verbose_proxy_logger.exception(f"[Non-Blocking] Error getting paginated teams: {e}") return [], 0 @@ -4763,18 +4463,10 @@ async def get_paginated_teams( }, ) async def ui_view_teams( - team_id: Optional[str] = fastapi.Query( - default=None, description="Team ID in the request parameters" - ), - team_alias: Optional[str] = fastapi.Query( - default=None, description="Team alias in the request parameters" - ), - page: int = fastapi.Query( - default=1, description="Page number for pagination", ge=1 - ), - page_size: int = fastapi.Query( - default=50, description="Number of items per page", ge=1, le=100 - ), + team_id: Optional[str] = fastapi.Query(default=None, description="Team ID in the request parameters"), + team_alias: Optional[str] = fastapi.Query(default=None, description="Team alias in the request parameters"), + page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), + page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=100), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -4831,16 +4523,12 @@ async def ui_view_teams( raise HTTPException(status_code=500, detail=f"Error searching teams: {str(e)}") -def add_new_models_to_team( - team_obj: LiteLLM_TeamTable, new_models: List[str] -) -> List[str]: +def add_new_models_to_team(team_obj: LiteLLM_TeamTable, new_models: List[str]) -> List[str]: """ Add new models to a team's allowed model list. """ current_models = team_obj.models - if ( - current_models is not None and len(current_models) == 0 - ): # implies all model access + if current_models is not None and len(current_models) == 0: # implies all model access current_models = [SpecialModelNames.all_proxy_models.value] else: current_models = team_obj.models @@ -4887,9 +4575,7 @@ async def team_model_add( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": data.team_id} - ) + team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) if team_row is None: raise HTTPException( @@ -4902,12 +4588,8 @@ async def team_model_add( # Authorization check - only proxy admin, team admin, or org admin can add models if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ) - and not await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ) + and not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj) + and not await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj) ): raise HTTPException( status_code=403, @@ -4993,9 +4675,7 @@ async def team_model_delete( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": data.team_id} - ) + team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) if team_row is None: raise HTTPException( @@ -5008,12 +4688,8 @@ async def team_model_delete( # Authorization check - only proxy admin, team admin, or org admin can remove models if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ) - and not await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ) + and not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj) + and not await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj) ): raise HTTPException( status_code=403, @@ -5049,9 +4725,7 @@ async def team_model_delete( ) @management_endpoint_wrapper async def team_member_permissions( - team_id: str = fastapi.Query( - default=None, description="Team ID in the request parameters" - ), + team_id: str = fastapi.Query(default=None, description="Team ID in the request parameters"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> GetTeamMemberPermissionsResponse: """ @@ -5084,12 +4758,8 @@ async def team_member_permissions( if ( hasattr(user_api_key_dict, "user_role") and not _user_has_admin_view(user_api_key_dict) - and not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ) - and not await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ) + and not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=complete_team_data) + and not await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=complete_team_data) and not _is_available_team( team_id=complete_team_data.team_id, user_api_key_dict=user_api_key_dict, @@ -5106,9 +4776,7 @@ async def team_member_permissions( ) if existing_team_row.team_member_permissions is None: - existing_team_row.team_member_permissions = ( - TeamMemberPermissionChecks.default_team_member_permissions() - ) + existing_team_row.team_member_permissions = TeamMemberPermissionChecks.default_team_member_permissions() return GetTeamMemberPermissionsResponse( team_id=team_id, @@ -5157,12 +4825,8 @@ async def update_team_member_permissions( if ( hasattr(user_api_key_dict, "user_role") and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ) - and not await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ) + and not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=complete_team_data) + and not await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=complete_team_data) ): raise HTTPException( status_code=403, @@ -5233,13 +4897,9 @@ async def bulk_update_team_member_permissions( permissions_to_add = set(data.permissions) if data.team_ids: - teams_updated = await _append_permissions_to_specific_teams( - prisma_client, data.team_ids, permissions_to_add - ) + teams_updated = await _append_permissions_to_specific_teams(prisma_client, data.team_ids, permissions_to_add) else: - teams_updated = await _append_permissions_to_all_teams( - prisma_client, permissions_to_add - ) + teams_updated = await _append_permissions_to_all_teams(prisma_client, permissions_to_add) return { "message": "Team permissions updated successfully", @@ -5248,18 +4908,14 @@ async def bulk_update_team_member_permissions( } -async def _compute_and_batch_updates( - prisma_client, teams, permissions_to_add: set -) -> int: +async def _compute_and_batch_updates(prisma_client, teams, permissions_to_add: set) -> int: """Compute merged permissions and batch-write updates. Returns count of teams updated.""" updates = [] for team in teams: existing = set(team.team_member_permissions or []) if permissions_to_add <= existing: continue - merged = sorted( - existing | permissions_to_add - ) # normalise to alphabetical order + merged = sorted(existing | permissions_to_add) # normalise to alphabetical order updates.append((team.team_id, merged)) if updates: @@ -5274,9 +4930,7 @@ async def _compute_and_batch_updates( return len(updates) -async def _append_permissions_to_specific_teams( - prisma_client, team_ids: List[str], permissions_to_add: set -) -> int: +async def _append_permissions_to_specific_teams(prisma_client, team_ids: List[str], permissions_to_add: set) -> int: """Fetch specific teams by ID and append permissions.""" teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": team_ids}}, @@ -5293,9 +4947,7 @@ async def _append_permissions_to_specific_teams( return await _compute_and_batch_updates(prisma_client, teams, permissions_to_add) -async def _append_permissions_to_all_teams( - prisma_client, permissions_to_add: set -) -> int: +async def _append_permissions_to_all_teams(prisma_client, permissions_to_add: set) -> int: """Paginated read + batched write across all teams.""" teams_updated = 0 cursor = None @@ -5315,9 +4967,7 @@ async def _append_permissions_to_all_teams( if not teams: break - teams_updated += await _compute_and_batch_updates( - prisma_client, teams, permissions_to_add - ) + teams_updated += await _compute_and_batch_updates(prisma_client, teams, permissions_to_add) cursor = teams[-1].team_id @@ -5375,9 +5025,7 @@ async def get_team_daily_activity( exclude_team_ids_list: Optional[List[str]] = None if exclude_team_ids: - exclude_team_ids_list = ( - exclude_team_ids.split(",") if exclude_team_ids else None - ) + exclude_team_ids_list = exclude_team_ids.split(",") if exclude_team_ids else None if not _user_has_admin_view(user_api_key_dict): user_info = await get_user_object( @@ -5392,9 +5040,7 @@ async def get_team_daily_activity( if user_info is None: raise HTTPException( status_code=404, - detail={ - "error": "User= {} not found".format(user_api_key_dict.user_id) - }, + detail={"error": "User= {} not found".format(user_api_key_dict.user_id)}, ) if team_ids_list is None: @@ -5416,12 +5062,8 @@ async def get_team_daily_activity( where_condition = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases = await TeamRepository(prisma_client).table.find_many( - where=where_condition - ) - team_alias_metadata = { - t.team_id: {"team_alias": t.team_alias} for t in team_aliases - } + team_aliases = await TeamRepository(prisma_client).table.find_many(where=where_condition) + team_alias_metadata = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. @@ -5438,9 +5080,7 @@ async def get_team_daily_activity( has_full_team_view = True for team_alias in team_aliases: team_obj = LiteLLM_TeamTable(**team_alias.model_dump()) - is_admin = _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ) + is_admin = _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj) has_perm = _team_member_has_permission( user_api_key_dict=user_api_key_dict, team_obj=team_obj, @@ -5453,9 +5093,9 @@ async def get_team_daily_activity( # If user does not have full team view, filter by their API keys if not has_full_team_view: # Get all API keys for this user - user_keys = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"user_id": user_api_key_dict.user_id}) + user_keys = await VerificationTokenRepository(prisma_client).table.find_many( + where={"user_id": user_api_key_dict.user_id} + ) user_api_keys = [key.token for key in user_keys if key.token] # If user has no API keys, return empty result if not user_api_keys: diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index a9b57db8a6f..9d71761f115 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -114,14 +114,10 @@ async def list_tools( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - tools = await db_list_tools( - prisma_client=prisma_client, input_policy=input_policy - ) + tools = await db_list_tools(prisma_client=prisma_client, input_policy=input_policy) return ToolListResponse(tools=tools, total=len(tools)) except Exception as e: verbose_proxy_logger.exception("Error listing tools: %s", e) @@ -146,17 +142,13 @@ async def get_tool_detail( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name) if tool is None: raise HTTPException(status_code=404, detail=f"Tool '{tool_name}' not found") - overrides = await list_overrides_for_tool( - prisma_client=prisma_client, tool_name=tool_name - ) + overrides = await list_overrides_for_tool(prisma_client=prisma_client, tool_name=tool_name) return ToolDetailResponse(tool=tool, overrides=overrides) except HTTPException: raise @@ -235,9 +227,7 @@ async def get_tool_usage_logs( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: where: dict = {"tool_name": tool_name} @@ -246,16 +236,16 @@ async def get_tool_usage_logs( end_time_filter: Optional[datetime] = None if start_date: try: - start_time_filter = datetime.strptime( - start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S" - ).replace(tzinfo=timezone.utc) + start_time_filter = datetime.strptime(start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S").replace( + tzinfo=timezone.utc + ) except ValueError: pass if end_date: try: - end_time_filter = datetime.strptime( - end_date + "T23:59:59", "%Y-%m-%dT%H:%M:%S" - ).replace(tzinfo=timezone.utc) + end_time_filter = datetime.strptime(end_date + "T23:59:59", "%Y-%m-%dT%H:%M:%S").replace( + tzinfo=timezone.utc + ) except ValueError: pass if start_time_filter is not None or end_time_filter is not None: @@ -265,9 +255,7 @@ async def get_tool_usage_logs( if end_time_filter is not None: where["start_time"]["lte"] = end_time_filter - total = await SpendLogToolIndexRepository(prisma_client).table.count( - where=where - ) + total = await SpendLogToolIndexRepository(prisma_client).table.count(where=where) index_rows = await SpendLogToolIndexRepository(prisma_client).table.find_many( where=where, order={"start_time": "desc"}, @@ -276,13 +264,9 @@ async def get_tool_usage_logs( ) request_ids = [r.request_id for r in index_rows] if not request_ids: - return ToolUsageLogsResponse( - logs=[], total=total, page=page, page_size=page_size - ) + return ToolUsageLogsResponse(logs=[], total=total, page=page, page_size=page_size) - spend_logs = await SpendLogsRepository(prisma_client).table.find_many( - where={"request_id": {"in": request_ids}} - ) + spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}}) log_by_id = {s.request_id: s for s in spend_logs} logs_out: List[ToolUsageLogEntry] = [] @@ -290,11 +274,7 @@ async def get_tool_usage_logs( sl = log_by_id.get(r.request_id) if not sl: continue - ts = ( - sl.startTime.isoformat() - if hasattr(sl.startTime, "isoformat") - else str(sl.startTime) - ) + ts = sl.startTime.isoformat() if hasattr(sl.startTime, "isoformat") else str(sl.startTime) logs_out.append( ToolUsageLogEntry( id=sl.request_id, @@ -306,9 +286,7 @@ async def get_tool_usage_logs( ) ) - return ToolUsageLogsResponse( - logs=logs_out, total=total, page=page, page_size=page_size - ) + return ToolUsageLogsResponse(logs=logs_out, total=total, page=page, page_size=page_size) except HTTPException: raise except Exception as e: @@ -333,9 +311,7 @@ async def get_tool( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name) @@ -359,9 +335,7 @@ async def _resolve_key_hash_to_object_permission_id( hashed = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) if not hashed: return None - row = await VerificationTokenRepository(prisma_client).table.find_unique( - where={"token": hashed} - ) + row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) if row is None: return None op_id = getattr(row, "object_permission_id", None) @@ -376,12 +350,8 @@ async def _resolve_key_hash_to_object_permission_id( data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete( - where={"object_permission_id": new_id} - ) - row = await VerificationTokenRepository(prisma_client).table.find_unique( - where={"token": hashed} - ) + await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) + row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) return getattr(row, "object_permission_id", None) if row else None return new_id @@ -412,9 +382,7 @@ async def _resolve_team_id_to_object_permission_id( data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete( - where={"object_permission_id": new_id} - ) + await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id_clean}, select={"object_permission_id": True}, @@ -454,9 +422,7 @@ async def update_tool_policy( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: if data.team_id is not None or data.key_hash is not None: @@ -466,13 +432,9 @@ async def update_tool_policy( detail="Provide either team_id or key_hash, not both", ) if data.key_hash is not None: - op_id = await _resolve_key_hash_to_object_permission_id( - prisma_client, data.key_hash - ) + op_id = await _resolve_key_hash_to_object_permission_id(prisma_client, data.key_hash) else: - op_id = await _resolve_team_id_to_object_permission_id( - prisma_client, data.team_id or "" - ) + op_id = await _resolve_team_id_to_object_permission_id(prisma_client, data.team_id or "") if op_id is None: raise HTTPException( status_code=404, @@ -549,12 +511,8 @@ async def update_tool_policy( ) async def delete_tool_policy_override( tool_name: str, - team_id: Optional[str] = Query( - None, description="Team ID of the override to remove" - ), - key_hash: Optional[str] = Query( - None, description="Key hash of the override to remove" - ), + team_id: Optional[str] = Query(None, description="Team ID of the override to remove"), + key_hash: Optional[str] = Query(None, description="Key hash of the override to remove"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -568,9 +526,7 @@ async def delete_tool_policy_override( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) if team_id is None and key_hash is None: raise HTTPException( status_code=400, @@ -583,13 +539,9 @@ async def delete_tool_policy_override( ) try: if key_hash is not None: - op_id = await _resolve_key_hash_to_object_permission_id( - prisma_client, key_hash - ) + op_id = await _resolve_key_hash_to_object_permission_id(prisma_client, key_hash) else: - op_id = await _resolve_team_id_to_object_permission_id( - prisma_client, team_id or "" - ) + op_id = await _resolve_team_id_to_object_permission_id(prisma_client, team_id or "") if op_id is None: raise HTTPException( status_code=404, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index dc128e88d51..73ec56a82e6 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -138,16 +138,12 @@ router = APIRouter() # response convertors see the same fields in the PKCE path as in the non-PKCE path. _OAUTH_TOKEN_FIELDS = frozenset({"access_token", "id_token", "refresh_token"}) _CLI_SSO_FLOW_CACHE_KEY_PREFIX = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:flow" -_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX = ( - f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:start_rate_limit" -) +_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX = f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:start_rate_limit" _CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS = 60 _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30 _CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" _CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$") -_CLI_SSO_USER_CODE_RE = re.compile( - rf"^[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}-[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}$" -) +_CLI_SSO_USER_CODE_RE = re.compile(rf"^[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}-[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}$") _CLI_SSO_SCALAR_TYPES = (str, int, float, bool) _CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$") _CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset( @@ -186,9 +182,7 @@ def _is_valid_cli_sso_login_id(login_id: Optional[str]) -> bool: def _is_valid_cli_sso_user_code(user_code: str | None) -> bool: - return isinstance(user_code, str) and bool( - _CLI_SSO_USER_CODE_RE.fullmatch(user_code) - ) + return isinstance(user_code, str) and bool(_CLI_SSO_USER_CODE_RE.fullmatch(user_code)) def _cli_sso_verification_uri_complete_enabled() -> bool: @@ -220,15 +214,8 @@ def _cli_sso_start_response_body( } -def _get_cli_sso_start_rate_limit_cache_key( - request: Request, use_x_forwarded_for: Optional[bool] = False -) -> str: - client_ip = ( - _get_request_ip_address( - request=request, use_x_forwarded_for=use_x_forwarded_for - ) - or "unknown" - ) +def _get_cli_sso_start_rate_limit_cache_key(request: Request, use_x_forwarded_for: Optional[bool] = False) -> str: + client_ip = _get_request_ip_address(request=request, use_x_forwarded_for=use_x_forwarded_for) or "unknown" client_ip_hash = _hash_cli_sso_secret(client_ip) return f"{_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX}:{client_ip_hash}" @@ -274,9 +261,7 @@ def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool: expected_poll_secret_hash = flow.get("poll_secret_hash") - if not isinstance(expected_poll_secret_hash, str) or not isinstance( - poll_secret, str - ): + if not isinstance(expected_poll_secret_hash, str) or not isinstance(poll_secret, str): return False supplied_poll_secret_hash = _hash_cli_sso_secret(poll_secret) return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash) @@ -361,9 +346,7 @@ def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any: return current -def _extract_sso_claim_value( - result: Union[CustomOpenID, OpenID, dict], claim_path: str -) -> Any: +def _extract_sso_claim_value(result: Union[CustomOpenID, OpenID, dict], claim_path: str) -> Any: extra_fields = getattr(result, "extra_fields", None) if isinstance(extra_fields, dict): if claim_path in extra_fields: @@ -379,9 +362,7 @@ def _extract_sso_claim_value( return _get_nested_claim_value(result_dict, claim_path) -def _set_nested_metadata_value( - metadata: Dict[str, Any], key_path: str, value: Any -) -> None: +def _set_nested_metadata_value(metadata: Dict[str, Any], key_path: str, value: Any) -> None: placeholder = "\x00" parts = key_path.replace("\\.", placeholder).split(".") parts = [p.replace(placeholder, ".") for p in parts] @@ -428,18 +409,14 @@ def build_cli_sso_attribution_metadata( metadata: Dict[str, Any] = {} for source_claim, dest_key in claim_map: if not _is_safe_cli_sso_metadata_dest_key(dest_key): - verbose_proxy_logger.debug( - f"Skipping unsafe CLI SSO metadata destination key: {dest_key}" - ) + verbose_proxy_logger.debug(f"Skipping unsafe CLI SSO metadata destination key: {dest_key}") continue raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim) if not _is_safe_cli_sso_scalar_claim_value(raw_value): continue - _set_nested_metadata_value( - metadata=metadata, key_path=dest_key, value=raw_value - ) + _set_nested_metadata_value(metadata=metadata, key_path=dest_key, value=raw_value) return metadata @@ -454,9 +431,7 @@ def _merge_cli_sso_attribution_metadata( are merged iteratively so attribution claims do not clobber unrelated keys under the same parent. """ - pending: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [ - (existing_metadata, attribution_metadata) - ] + pending: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [(existing_metadata, attribution_metadata)] while pending: target, source = pending.pop() for key, value in source.items(): @@ -479,9 +454,7 @@ async def _persist_cli_sso_user_metadata( return try: - user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) existing_metadata: Dict[str, Any] = {} if user_row is not None: row_metadata = user_row.metadata @@ -501,9 +474,7 @@ async def _persist_cli_sso_user_metadata( f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}" ) except Exception as e: - verbose_proxy_logger.error( - f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}" - ) + verbose_proxy_logger.error(f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}") def _cli_poll_attribution_metadata_from_session( @@ -522,9 +493,7 @@ def _render_cli_sso_verification_page( ) -> str: escaped_verify_url = escape(verify_url, quote=True) escaped_browser_complete_token = escape(browser_complete_token, quote=True) - user_code_value_attr = ( - f' value="{escape(prefill_user_code, quote=True)}"' if prefill_user_code else "" - ) + user_code_value_attr = f' value="{escape(prefill_user_code, quote=True)}"' if prefill_user_code else "" instructions = ( "Confirm the verification code below to finish this login." if prefill_user_code @@ -603,9 +572,7 @@ async def cli_sso_start(request: Request): _check_cli_sso_start_rate_limit( request=request, cache=user_api_key_cache, - use_x_forwarded_for=bool( - (general_settings or {}).get("use_x_forwarded_for", False) - ), + use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), ) login_id = f"cli-{secrets.token_urlsafe(24)}" @@ -623,9 +590,7 @@ async def cli_sso_start(request: Request): verification_uri_complete: str | None = ( ( - get_custom_url( - request_base_url=str(request.base_url), route="sso/key/generate" - ) + get_custom_url(request_base_url=str(request.base_url), route="sso/key/generate") + "?" + urlencode( { @@ -646,9 +611,7 @@ async def cli_sso_start(request: Request): ) -@router.post( - "/sso/cli/complete/{login_id}", tags=["experimental"], include_in_schema=False -) +@router.post("/sso/cli/complete/{login_id}", tags=["experimental"], include_in_schema=False) async def cli_sso_complete(request: Request, login_id: str): from fastapi.responses import HTMLResponse @@ -664,15 +627,9 @@ async def cli_sso_complete(request: Request, login_id: str): body = (await request.body()).decode("utf-8") form_values = parse_qs(body) supplied_user_code = (form_values.get("user_code") or [""])[0] - supplied_browser_complete_token = ( - form_values.get("browser_complete_token") or [""] - )[0] - supplied_user_code_hash = _hash_cli_sso_secret( - _normalize_cli_sso_user_code(supplied_user_code) - ) - supplied_browser_complete_token_hash = _hash_cli_sso_secret( - supplied_browser_complete_token - ) + supplied_browser_complete_token = (form_values.get("browser_complete_token") or [""])[0] + supplied_user_code_hash = _hash_cli_sso_secret(_normalize_cli_sso_user_code(supplied_user_code)) + supplied_browser_complete_token_hash = _hash_cli_sso_secret(supplied_browser_complete_token) expected_user_code_hash = flow.get("user_code_hash") if not isinstance(expected_user_code_hash, str) or not secrets.compare_digest( @@ -681,9 +638,7 @@ async def cli_sso_complete(request: Request, login_id: str): raise HTTPException(status_code=400, detail="Invalid verification code") expected_browser_complete_token_hash = flow.get("browser_complete_token_hash") - if not isinstance( - expected_browser_complete_token_hash, str - ) or not secrets.compare_digest( + if not isinstance(expected_browser_complete_token_hash, str) or not secrets.compare_digest( supplied_browser_complete_token_hash, expected_browser_complete_token_hash ): raise HTTPException(status_code=400, detail="Invalid verification code") @@ -754,9 +709,7 @@ def determine_role_from_groups( for role in role_hierarchy: if role in role_mappings.roles: role_groups = role_mappings.roles[role] - if isinstance(role_groups, list) and user_groups_set.intersection( - set(role_groups) - ): + if isinstance(role_groups, list) and user_groups_set.intersection(set(role_groups)): verbose_proxy_logger.debug( f"User groups {user_groups} matched role '{role.value}' via groups: {role_groups}" ) @@ -799,9 +752,7 @@ def process_sso_jwt_access_token( import jwt try: - access_token_payload = jwt.decode( - access_token_str, options={"verify_signature": False} - ) + access_token_payload = jwt.decode(access_token_str, options={"verify_signature": False}) except jwt.exceptions.DecodeError: verbose_proxy_logger.debug( "Access token is not a valid JWT (possibly an opaque token), skipping JWT-based extraction" @@ -813,41 +764,29 @@ def process_sso_jwt_access_token( if isinstance(result, dict): result_team_ids: Optional[List[str]] = result.get("team_ids", []) if not result_team_ids: - team_ids = sso_jwt_handler.get_team_ids_from_jwt( - access_token_payload - ) + team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) result["team_ids"] = team_ids else: result_team_ids = getattr(result, "team_ids", []) if result else [] if not result_team_ids: - team_ids = sso_jwt_handler.get_team_ids_from_jwt( - access_token_payload - ) + team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) setattr(result, "team_ids", team_ids) # Extract user role from access token if not already set from UserInfo - existing_role = ( - result.get("user_role") - if isinstance(result, dict) - else getattr(result, "user_role", None) - ) + existing_role = result.get("user_role") if isinstance(result, dict) else getattr(result, "user_role", None) if existing_role is None: user_role: Optional[LitellmUserRoles] = None # Try role_mappings first (group-based role determination) if role_mappings is not None and role_mappings.roles: group_claim = role_mappings.group_claim - user_groups_raw: Any = get_nested_value( - access_token_payload, group_claim - ) + user_groups_raw: Any = get_nested_value(access_token_payload, group_claim) user_groups: List[str] = [] if isinstance(user_groups_raw, list): user_groups = [str(g) for g in user_groups_raw] elif isinstance(user_groups_raw, str): - user_groups = [ - g.strip() for g in user_groups_raw.split(",") if g.strip() - ] + user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] elif user_groups_raw is not None: user_groups = [str(user_groups_raw)] @@ -861,12 +800,8 @@ def process_sso_jwt_access_token( # Fallback: try GENERIC_USER_ROLE_ATTRIBUTE on the access token payload if user_role is None: - generic_user_role_attribute_name = os.getenv( - "GENERIC_USER_ROLE_ATTRIBUTE", "role" - ) - user_role_from_token = get_nested_value( - access_token_payload, generic_user_role_attribute_name - ) + generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") + user_role_from_token = get_nested_value(access_token_payload, generic_user_role_attribute_name) if user_role_from_token is not None: user_role = get_litellm_user_role(user_role_from_token) verbose_proxy_logger.debug( @@ -878,9 +813,7 @@ def process_sso_jwt_access_token( result["user_role"] = user_role else: setattr(result, "user_role", user_role) - verbose_proxy_logger.debug( - f"Set user_role='{user_role}' from JWT access token" - ) + verbose_proxy_logger.debug(f"Set user_role='{user_role}' from JWT access token") return access_token_payload @@ -921,11 +854,7 @@ async def google_login( return admin_ui_disabled() ####### Check if user is a Enterprise / Premium User ####### - if ( - microsoft_client_id is not None - or google_client_id is not None - or generic_client_id is not None - ): + if microsoft_client_id is not None or google_client_id is not None or generic_client_id is not None: if premium_user is not True: # Check if under 'free SSO user' limit if prisma_client is not None: @@ -1032,26 +961,14 @@ def generic_response_convertor( role_mappings: Optional["RoleMappings"] = None, team_mappings: Optional["TeamMappings"] = None, ) -> CustomOpenID: - generic_user_id_attribute_name = os.getenv( - "GENERIC_USER_ID_ATTRIBUTE", "preferred_username" - ) - generic_user_display_name_attribute_name = os.getenv( - "GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub" - ) - generic_user_email_attribute_name = os.getenv( - "GENERIC_USER_EMAIL_ATTRIBUTE", "email" - ) + generic_user_id_attribute_name = os.getenv("GENERIC_USER_ID_ATTRIBUTE", "preferred_username") + generic_user_display_name_attribute_name = os.getenv("GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub") + generic_user_email_attribute_name = os.getenv("GENERIC_USER_EMAIL_ATTRIBUTE", "email") - generic_user_first_name_attribute_name = os.getenv( - "GENERIC_USER_FIRST_NAME_ATTRIBUTE", "first_name" - ) - generic_user_last_name_attribute_name = os.getenv( - "GENERIC_USER_LAST_NAME_ATTRIBUTE", "last_name" - ) + generic_user_first_name_attribute_name = os.getenv("GENERIC_USER_FIRST_NAME_ATTRIBUTE", "first_name") + generic_user_last_name_attribute_name = os.getenv("GENERIC_USER_LAST_NAME_ATTRIBUTE", "last_name") - generic_provider_attribute_name = os.getenv( - "GENERIC_USER_PROVIDER_ATTRIBUTE", "provider" - ) + generic_provider_attribute_name = os.getenv("GENERIC_USER_PROVIDER_ATTRIBUTE", "provider") generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") @@ -1118,9 +1035,7 @@ def generic_response_convertor( # Fallback to existing logic if role_mappings not used if user_role is None: - user_role_from_sso = get_nested_value( - response, generic_user_role_attribute_name - ) + user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name) if user_role_from_sso is not None: role = get_litellm_user_role(user_role_from_sso) if role is not None: @@ -1139,12 +1054,8 @@ def generic_response_convertor( return CustomOpenID( id=get_nested_value(response, generic_user_id_attribute_name), - display_name=get_nested_value( - response, generic_user_display_name_attribute_name - ), - email=normalize_email( - get_nested_value(response, generic_user_email_attribute_name) - ), + display_name=get_nested_value(response, generic_user_display_name_attribute_name), + email=normalize_email(get_nested_value(response, generic_user_email_attribute_name)), first_name=get_nested_value(response, generic_user_first_name_attribute_name), last_name=get_nested_value(response, generic_user_last_name_attribute_name), provider=get_nested_value(response, generic_provider_attribute_name), @@ -1163,9 +1074,7 @@ def _setup_generic_sso_env_vars( generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) - generic_include_client_id = ( - os.getenv("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" - ) + generic_include_client_id = os.getenv("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" # Validate required environment variables if generic_client_secret is None: @@ -1200,9 +1109,7 @@ def _setup_generic_sso_env_vars( verbose_proxy_logger.debug( f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}" ) - verbose_proxy_logger.debug( - f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n" - ) + verbose_proxy_logger.debug(f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n") return ( generic_client_secret, @@ -1220,13 +1127,9 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: try: from litellm.proxy.utils import get_prisma_client_or_throw - prisma_client = get_prisma_client_or_throw( - "Prisma client is None, connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( - where={"id": "sso_config"} - ) + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict = dict(sso_db_record.sso_settings) @@ -1258,13 +1161,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: try: from litellm.proxy.utils import get_prisma_client_or_throw - prisma_client = get_prisma_client_or_throw( - "Prisma client is None, connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( - where={"id": "sso_config"} - ) + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict = dict(sso_db_record.sso_settings) @@ -1279,31 +1178,21 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: role_mappings = role_mappings_data if role_mappings: - verbose_proxy_logger.debug( - f"Loaded role_mappings for provider '{role_mappings.provider}'" - ) + verbose_proxy_logger.debug(f"Loaded role_mappings for provider '{role_mappings.provider}'") except Exception as e: verbose_proxy_logger.debug( f"Could not load role_mappings from database: {e}. Continuing with existing role logic." ) generic_role_mappings = os.getenv("GENERIC_ROLE_MAPPINGS_ROLES", None) - generic_role_mappings_group_claim = os.getenv( - "GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None - ) - generic_role_mappings_default_role = os.getenv( - "GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None - ) + generic_role_mappings_group_claim = os.getenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None) + generic_role_mappings_default_role = os.getenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None) if generic_role_mappings is not None: - verbose_proxy_logger.debug( - "Found role_mappings for generic provider in environment variables" - ) + verbose_proxy_logger.debug("Found role_mappings for generic provider in environment variables") import ast try: - generic_user_role_mappings_data: Dict[LitellmUserRoles, List[str]] = ( - ast.literal_eval(generic_role_mappings) - ) + generic_user_role_mappings_data: Dict[LitellmUserRoles, List[str]] = ast.literal_eval(generic_role_mappings) if isinstance(generic_user_role_mappings_data, dict): from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings @@ -1353,13 +1242,10 @@ def _handle_generic_sso_error( # 1. The error mentions PKCE/code verifier, AND # 2. PKCE is not currently configured (GENERIC_CLIENT_USE_PKCE != true) pkce_configured = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true" - if not pkce_configured and ( - "PKCE" in error_message or "code verifier" in error_message.lower() - ): - is_okta = ( - generic_authorization_endpoint - and "okta" in generic_authorization_endpoint.lower() - ) or (generic_token_endpoint and "okta" in generic_token_endpoint.lower()) + if not pkce_configured and ("PKCE" in error_message or "code verifier" in error_message.lower()): + is_okta = (generic_authorization_endpoint and "okta" in generic_authorization_endpoint.lower()) or ( + generic_token_endpoint and "okta" in generic_token_endpoint.lower() + ) provider_name = "Okta" if is_okta else "Your OAuth provider" detailed_message = ( @@ -1401,14 +1287,10 @@ def _handle_generic_sso_error( async def get_generic_sso_response( request: Request, jwt_handler: JWTHandler, - sso_jwt_handler: Optional[ - JWTHandler - ], # sso specific jwt handler - used for restricted sso group access control + sso_jwt_handler: Optional[JWTHandler], # sso specific jwt handler - used for restricted sso group access control generic_client_id: str, redirect_url: str, -) -> Tuple[ - Union[OpenID, dict], Optional[dict], Optional[dict] -]: # (result, received_response, access_token_payload) +) -> Tuple[Union[OpenID, dict], Optional[dict], Optional[dict]]: # (result, received_response, access_token_payload) # make generic sso provider from fastapi_sso.sso.base import DiscoveryDocument from fastapi_sso.sso.generic import create_provider @@ -1460,17 +1342,13 @@ async def get_generic_sso_response( verbose_proxy_logger.debug("calling generic_sso.verify_and_process") additional_generic_sso_headers_dict = _parse_generic_sso_headers() - code_verifier: Optional[str] = ( - None # assigned inside try; initialized for type tracking - ) + code_verifier: Optional[str] = None # assigned inside try; initialized for type tracking access_token_payload: Optional[dict] = None # decoded JWT access token claims try: - token_exchange_params = ( - await SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=request, - generic_include_client_id=generic_include_client_id, - ) + token_exchange_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=request, + generic_include_client_id=generic_include_client_id, ) # Extract code_verifier (and the cache key for deferred deletion) before calling fastapi-sso @@ -1493,16 +1371,9 @@ async def get_generic_sso_response( # code (Login-CSRF / token theft). url_state = request.query_params.get("state") cookie_state = request.cookies.get("litellm_oauth_state") - if ( - not url_state - or not cookie_state - or not secrets.compare_digest(url_state, cookie_state) - ): + if not url_state or not cookie_state or not secrets.compare_digest(url_state, cookie_state): raise ProxyException( - message=( - "Invalid OAuth state parameter — does not match " - "the browser-bound state cookie." - ), + message=("Invalid OAuth state parameter — does not match the browser-bound state cookie."), type=ProxyErrorTypes.auth_error, param="state", code=status.HTTP_400_BAD_REQUEST, @@ -1557,11 +1428,7 @@ async def get_generic_sso_response( # must not be exposed to callers. # Assign directly rather than relying on nonlocal mutation so that Pyright # can track that received_response is non-None from this point on. - received_response = { - k: v - for k, v in combined_response.items() - if k not in _OAUTH_TOKEN_FIELDS - } + received_response = {k: v for k, v in combined_response.items() if k not in _OAUTH_TOKEN_FIELDS} # In the PKCE path verify_and_process is skipped, so generic_sso.access_token # is never set. Read the token directly from the exchange response instead so # process_sso_jwt_access_token can extract JWT-embedded roles/teams. @@ -1608,14 +1475,10 @@ async def create_team_member_add_task(team_id, user_info): user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) except Exception as e: - verbose_proxy_logger.debug( - f"[Non-Blocking] Error trying to add sso user to db: {e}" - ) + verbose_proxy_logger.debug(f"[Non-Blocking] Error trying to add sso user to db: {e}") -async def add_missing_team_member( - user_info: Union[NewUserResponse, LiteLLM_UserTable], sso_teams: List[str] -): +async def add_missing_team_member(user_info: Union[NewUserResponse, LiteLLM_UserTable], sso_teams: List[str]): """ - Get missing teams (diff b/w user_info.team_ids and sso_teams) - Add missing user to missing teams @@ -1625,26 +1488,19 @@ async def add_missing_team_member( missing_teams = set(sso_teams) - set(user_teams) missing_teams_list = list(missing_teams) tasks = [] - tasks = [ - create_team_member_add_task(team_id, user_info) - for team_id in missing_teams_list - ] + tasks = [create_team_member_add_task(team_id, user_info) for team_id in missing_teams_list] try: await asyncio.gather(*tasks) except Exception as e: - verbose_proxy_logger.debug( - f"[Non-Blocking] Error trying to add sso user to db: {e}" - ) + verbose_proxy_logger.debug(f"[Non-Blocking] Error trying to add sso user to db: {e}") def get_disabled_non_admin_personal_key_creation(): key_generation_settings = litellm.key_generation_settings if key_generation_settings is None: return False - personal_key_generation = ( - key_generation_settings.get("personal_key_generation") or {} - ) + personal_key_generation = key_generation_settings.get("personal_key_generation") or {} allowed_user_roles = personal_key_generation.get("allowed_user_roles") or [] return bool("proxy_admin" in allowed_user_roles) @@ -1697,9 +1553,7 @@ async def get_user_info_from_db( potential_user_ids.append(_id) user_email = normalize_email( - getattr(result, "email", None) - if not isinstance(result, dict) - else result.get("email", None) + getattr(result, "email", None) if not isinstance(result, dict) else result.get("email", None) ) user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]] = None @@ -1735,9 +1589,7 @@ async def get_user_info_from_db( return user_info except Exception as e: - verbose_proxy_logger.exception( - f"[Non-Blocking] Error trying to add sso user to db: {e}" - ) + verbose_proxy_logger.exception(f"[Non-Blocking] Error trying to add sso user to db: {e}") return None @@ -1778,16 +1630,12 @@ def _build_sso_user_update_data( sso_role = getattr(result, "user_role", None) if sso_role is not None: # Convert enum to string if needed - sso_role_str = ( - sso_role.value if isinstance(sso_role, LitellmUserRoles) else sso_role - ) + sso_role_str = sso_role.value if isinstance(sso_role, LitellmUserRoles) else sso_role # Only include if it's a valid LiteLLM role if _should_use_role_from_sso_response(sso_role_str): update_data["user_role"] = sso_role_str - verbose_proxy_logger.info( - f"Updating user {user_id} role from SSO: {sso_role_str}" - ) + verbose_proxy_logger.info(f"Updating user {user_id} role from SSO: {sso_role_str}") return update_data @@ -1819,9 +1667,7 @@ async def _sync_user_role_from_jwt_role_map( if mapped_role is None: return - verbose_proxy_logger.info( - f"SSO jwt_litellm_role_map matched role: {mapped_role.value}" - ) + verbose_proxy_logger.info(f"SSO jwt_litellm_role_map matched role: {mapped_role.value}") # Update user_defined_values so downstream code uses the mapped role if user_defined_values is not None: @@ -1857,18 +1703,12 @@ def apply_user_info_values_to_sso_user_defined_values( if _should_use_role_from_sso_response(sso_role): # SSO provided a valid role, keep it and log that we're using it - verbose_proxy_logger.info( - f"Using SSO role: {sso_role} (DB role was: {db_role})" - ) + verbose_proxy_logger.info(f"Using SSO role: {sso_role} (DB role was: {db_role})") else: # SSO didn't provide a valid role, fall back to DB role or default if user_info is None or user_info.user_role is None: - user_defined_values["user_role"] = ( - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - ) - verbose_proxy_logger.debug( - "No SSO or DB role found, using default: INTERNAL_USER_VIEW_ONLY" - ) + user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + verbose_proxy_logger.debug("No SSO or DB role found, using default: INTERNAL_USER_VIEW_ONLY") else: user_defined_values["user_role"] = user_info.user_role verbose_proxy_logger.debug(f"Using DB role: {user_info.user_role}") @@ -1880,9 +1720,7 @@ def apply_user_info_values_to_sso_user_defined_values( return user_defined_values -async def check_and_update_if_proxy_admin_id( - user_role: str, user_id: str, prisma_client: Optional[PrismaClient] -): +async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prisma_client: Optional[PrismaClient]): """ - Check if user role in DB is admin - If not, update user role in DB to admin role @@ -1921,9 +1759,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): ) if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) sso_jwt_handler: Optional[JWTHandler] = None ui_access_mode = general_settings.get("ui_access_mode", None) @@ -1933,9 +1769,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, litellm_jwtauth=LiteLLM_JWTAuth( - team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get( - "sso_group_jwt_field", None - ), + team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get("sso_group_jwt_field", None), ), leeway=0, ) @@ -1953,9 +1787,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): param="master_key", code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) - redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( - request=request, sso_callback_route="sso/callback" - ) + redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(request=request, sso_callback_route="sso/callback") verbose_proxy_logger.info(f"Redirecting to {redirect_url}") result = None @@ -2052,9 +1884,7 @@ async def _fetch_cli_sso_team_details( team_details: List[Dict[str, Any]] = [] try: if teams: - prisma_teams = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": teams}} - ) + prisma_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}}) for team_row in prisma_teams: team_dict = team_row.model_dump() team_details.append( @@ -2064,9 +1894,7 @@ async def _fetch_cli_sso_team_details( } ) except Exception as e: - verbose_proxy_logger.error( - f"Error fetching team details for CLI SSO session: {e}" - ) + verbose_proxy_logger.error(f"Error fetching team details for CLI SSO session: {e}") return team_details @@ -2097,21 +1925,15 @@ async def _complete_cli_sso_callback_session( alternate_user_id=user_id, ) if user_info is None: - raise HTTPException( - status_code=500, detail="Failed to retrieve user information from SSO" - ) + raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") if not user_info.user_id: - raise HTTPException( - status_code=500, detail="Failed to retrieve user information from SSO" - ) + raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") teams: List[str] = [] if hasattr(user_info, "teams") and user_info.teams: teams = user_info.teams if isinstance(user_info.teams, list) else [] - team_details = await _fetch_cli_sso_team_details( - prisma_client=prisma_client, teams=teams - ) + team_details = await _fetch_cli_sso_team_details(prisma_client=prisma_client, teams=teams) attribution_metadata = build_cli_sso_attribution_metadata(result=result) if attribution_metadata: await _persist_cli_sso_user_metadata( @@ -2171,9 +1993,7 @@ async def cli_sso_callback( flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) if result is None: raise HTTPException( @@ -2185,11 +2005,9 @@ async def cli_sso_callback( result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result) try: - parsed_openid_result = ( - SSOAuthenticationHandler._get_user_email_and_id_from_result( - result=result_non_none, - generic_client_id=os.getenv("GENERIC_CLIENT_ID", None), - ) + parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result( + result=result_non_none, + generic_client_id=os.getenv("GENERIC_CLIENT_ID", None), ) verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}") user_defined_values = await _build_cli_sso_user_defined_values( @@ -2221,9 +2039,7 @@ async def cli_sso_callback( raise except Exception as e: verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to process CLI SSO: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Failed to process CLI SSO: {str(e)}") @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) @@ -2252,9 +2068,7 @@ async def cli_poll_key( try: flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache) - if not _verify_cli_sso_poll_secret( - flow=flow, poll_secret=x_litellm_cli_poll_secret - ): + if not _verify_cli_sso_poll_secret(flow=flow, poll_secret=x_litellm_cli_poll_secret): raise HTTPException(status_code=403, detail="Invalid CLI polling secret") if not flow.get("sso_complete") or not flow.get("user_code_verified"): @@ -2276,18 +2090,14 @@ async def cli_poll_key( # clients we return rich team details (id + alias); older clients # can continue to rely on the simple "teams" list. if team_id is None and len(user_teams) > 1: - verbose_proxy_logger.info( - f"Returning teams list for user {user_id} to select from: {user_teams}" - ) + verbose_proxy_logger.info(f"Returning teams list for user {user_id} to select from: {user_teams}") # Best-effort construction of team_details if it wasn't # already cached for some reason. team_details_response: Optional[List[Dict[str, Any]]] = None if isinstance(user_team_details, list) and user_team_details: team_details_response = user_team_details elif user_teams: - team_details_response = [ - {"team_id": t, "team_alias": None} for t in user_teams - ] + team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams] poll_response: Dict[str, Any] = { "status": "ready", "user_id": user_id, @@ -2295,9 +2105,7 @@ async def cli_poll_key( "team_details": team_details_response, "requires_team_selection": True, } - attribution_metadata = _cli_poll_attribution_metadata_from_session( - session_data - ) + attribution_metadata = _cli_poll_attribution_metadata_from_session(session_data) if attribution_metadata: poll_response["attribution_metadata"] = attribution_metadata return poll_response @@ -2316,11 +2124,7 @@ async def cli_poll_key( team_alias = None if team_id and isinstance(user_team_details, list): team_alias = next( - ( - team.get("team_alias") - for team in user_team_details - if team.get("team_id") == team_id - ), + (team.get("team_alias") for team in user_team_details if team.get("team_id") == team_id), None, ) @@ -2354,8 +2158,7 @@ async def cli_poll_key( session_max_budget = ( litellm.max_ui_session_budget - if user_budget is None - and (team_id is None or (team_budget_resolved and team_budget is None)) + if user_budget is None and (team_id is None or (team_budget_resolved and team_budget is None)) else None ) @@ -2369,9 +2172,7 @@ async def cli_poll_key( # Delete cache entry (single-use) user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) - verbose_proxy_logger.info( - f"CLI JWT generated for user: {user_id}, team: {team_id}" - ) + verbose_proxy_logger.info(f"CLI JWT generated for user: {user_id}, team: {team_id}") poll_response = { "status": "ready", "key": jwt_token, @@ -2382,9 +2183,7 @@ async def cli_poll_key( # present nicer information if needed. "team_details": user_team_details, } - attribution_metadata = _cli_poll_attribution_metadata_from_session( - session_data - ) + attribution_metadata = _cli_poll_attribution_metadata_from_session(session_data) if attribution_metadata: poll_response["attribution_metadata"] = attribution_metadata return poll_response @@ -2395,9 +2194,7 @@ async def cli_poll_key( raise except Exception as e: verbose_proxy_logger.error(f"Error polling for CLI JWT: {e}") - raise HTTPException( - status_code=500, detail=f"Error checking session status: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Error checking session status: {str(e)}") async def insert_sso_user( @@ -2414,9 +2211,7 @@ async def insert_sso_user( Returns: Tuple[str, str]: User ID and User Role """ - verbose_proxy_logger.debug( - f"Inserting SSO user into DB. User values: {user_defined_values}" - ) + verbose_proxy_logger.debug(f"Inserting SSO user into DB. User values: {user_defined_values}") if result_openid is None: raise ValueError("result_openid is None") if isinstance(result_openid, dict): @@ -2436,9 +2231,7 @@ async def insert_sso_user( preserved_role = sso_role user_defined_values.update(litellm.default_internal_user_params) # type: ignore user_defined_values["user_role"] = preserved_role # Restore preserved role - verbose_proxy_logger.debug( - f"Preserved SSO-extracted role '{preserved_role}'" - ) + verbose_proxy_logger.debug(f"Preserved SSO-extracted role '{preserved_role}'") else: # SSO didn't provide a valid role, apply all defaults including role user_defined_values.update(litellm.default_internal_user_params) # type: ignore @@ -2448,9 +2241,7 @@ async def insert_sso_user( if user_defined_values.get("max_budget") is None: user_defined_values["max_budget"] = litellm.max_internal_user_budget if user_defined_values.get("budget_duration") is None: - user_defined_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + user_defined_values["budget_duration"] = litellm.internal_user_budget_duration if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -2466,9 +2257,7 @@ async def insert_sso_user( ) if result_openid and hasattr(result_openid, "provider"): - new_user_request.metadata = { - "auth_provider": getattr(result_openid, "provider") - } + new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")} response = await new_user( data=new_user_request, @@ -2492,8 +2281,7 @@ async def get_ui_settings(request: Request): _api_doc_base_url = os.getenv("LITELLM_UI_API_DOC_BASE_URL", None) _is_sso_enabled = _has_user_setup_sso() disable_expensive_db_queries = ( - proxy_state.get_proxy_state_variable("spend_logs_row_count") - > MAX_SPENDLOG_ROWS_TO_QUERY + proxy_state.get_proxy_state_variable("spend_logs_row_count") > MAX_SPENDLOG_ROWS_TO_QUERY ) default_team_disabled = general_settings.get("default_team_disabled", False) if "PROXY_DEFAULT_TEAM_DISABLED" in os.environ: @@ -2506,9 +2294,7 @@ async def get_ui_settings(request: Request): "LITELLM_UI_API_DOC_BASE_URL": _api_doc_base_url, "DEFAULT_TEAM_DISABLED": default_team_disabled, "SSO_ENABLED": _is_sso_enabled, - "NUM_SPEND_LOGS_ROWS": proxy_state.get_proxy_state_variable( - "spend_logs_row_count" - ), + "NUM_SPEND_LOGS_ROWS": proxy_state.get_proxy_state_variable("spend_logs_row_count"), "DISABLE_EXPENSIVE_DB_QUERIES": disable_expensive_db_queries, } @@ -2562,9 +2348,7 @@ async def sso_readiness(): elif configured_provider == "generic": generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) - generic_authorization_endpoint = os.getenv( - "GENERIC_AUTHORIZATION_ENDPOINT", None - ) + generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) if generic_client_secret is None: @@ -2704,12 +2488,8 @@ class SSOAuthenticationHandler: from fastapi_sso.sso.generic import create_provider generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) - generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split( - " " - ) - generic_authorization_endpoint = os.getenv( - "GENERIC_AUTHORIZATION_ENDPOINT", None - ) + generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") + generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) if generic_client_secret is None: @@ -2809,9 +2589,7 @@ class SSOAuthenticationHandler: state_only_params[key] = value # Get the redirect response from fastapi-sso with only state param - redirect_response = await generic_sso.get_login_redirect( - **state_only_params - ) # type: ignore + redirect_response = await generic_sso.get_login_redirect(**state_only_params) # type: ignore # If PKCE is enabled, add PKCE parameters to the redirect URL if code_verifier and "state" in redirect_params: @@ -2831,9 +2609,7 @@ class SSOAuthenticationHandler: value={"code_verifier": code_verifier}, ttl=600, ) - verbose_proxy_logger.debug( - "PKCE code_verifier stored in cache (TTL: 600s)" - ) + verbose_proxy_logger.debug("PKCE code_verifier stored in cache (TTL: 600s)") # Add PKCE parameters to the authorization URL if pkce_params: @@ -2945,11 +2721,7 @@ class SSOAuthenticationHandler: microsoft_client_id: Optional[str] = None, generic_client_id: Optional[str] = None, ) -> bool: - if ( - google_client_id is not None - or microsoft_client_id is not None - or generic_client_id is not None - ): + if google_client_id is not None or microsoft_client_id is not None or generic_client_id is not None: return True return False @@ -2998,13 +2770,9 @@ class SSOAuthenticationHandler: user_id=user_id, ) - await UserRepository(prisma_client).table.update_many( - where={"user_id": user_id}, data=update_data - ) + await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data) else: - verbose_proxy_logger.info( - "user not in DB, inserting user into LiteLLM DB" - ) + verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") # user not in DB, insert User into LiteLLM DB user_info = await insert_sso_user( result_openid=result, @@ -3012,9 +2780,7 @@ class SSOAuthenticationHandler: ) return user_info except Exception as e: - verbose_proxy_logger.exception( - f"Error upserting SSO user into LiteLLM DB: {e}" - ) + verbose_proxy_logger.exception(f"Error upserting SSO user into LiteLLM DB: {e}") return user_info @staticmethod @@ -3029,9 +2795,7 @@ class SSOAuthenticationHandler: The `team_ids` field is populated by litellm after processing the SSO response """ if user_info is None: - verbose_proxy_logger.debug( - "User not found in LiteLLM DB, skipping team member addition" - ) + verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition") return sso_teams = getattr(result, "team_ids", []) await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) @@ -3053,9 +2817,7 @@ class SSOAuthenticationHandler: - if result.team_ids is a list, return True if the restricted_sso_group is in the list, otherwise return False """ - ui_access_mode = cast( - Optional[Union[Dict, str]], general_settings.get("ui_access_mode") - ) + ui_access_mode = cast(Optional[Union[Dict, str]], general_settings.get("ui_access_mode")) if ui_access_mode is None: return True @@ -3100,16 +2862,12 @@ class SSOAuthenticationHandler: code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) try: - team_obj = await TeamRepository(prisma_client).table.find_first( - where={"team_id": litellm_team_id} - ) + team_obj = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id}) verbose_proxy_logger.debug(f"Team object: {team_obj}") # only create a new team if it doesn't exist if team_obj: - verbose_proxy_logger.debug( - f"Team already exists: {litellm_team_id} - {litellm_team_name}" - ) + verbose_proxy_logger.debug(f"Team already exists: {litellm_team_id} - {litellm_team_name}") return team_request: NewTeamRequest = NewTeamRequest( @@ -3198,9 +2956,7 @@ class SSOAuthenticationHandler: Gets the user email and id from the OpenID result after validating the email domain """ user_email: Optional[str] = normalize_email(getattr(result, "email", None)) - user_id: Optional[str] = ( - getattr(result, "id", None) if result is not None else None - ) + user_id: Optional[str] = getattr(result, "id", None) if result is not None else None user_role: Optional[str] = None if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None: @@ -3221,32 +2977,20 @@ class SSOAuthenticationHandler: _user_role = getattr(result, "user_role", None) if _user_role is not None: # Convert enum to string if needed - user_role = ( - _user_role.value - if isinstance(_user_role, LitellmUserRoles) - else _user_role - ) - verbose_proxy_logger.debug( - f"Extracted user_role from SSO result: {user_role}" - ) + user_role = _user_role.value if isinstance(_user_role, LitellmUserRoles) else _user_role + verbose_proxy_logger.debug(f"Extracted user_role from SSO result: {user_role}") # generic client id - override with custom attribute name if specified if generic_client_id is not None and result is not None: - generic_user_role_attribute_name = os.getenv( - "GENERIC_USER_ROLE_ATTRIBUTE", "role" - ) + generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") user_id = getattr(result, "id", None) user_email = normalize_email(getattr(result, "email", None)) if user_role is None: - _role_from_attr = getattr( - result, generic_user_role_attribute_name, None - ) # type: ignore + _role_from_attr = getattr(result, generic_user_role_attribute_name, None) # type: ignore if _role_from_attr is not None: # Convert enum to string if needed user_role = ( - _role_from_attr.value - if isinstance(_role_from_attr, LitellmUserRoles) - else _role_from_attr + _role_from_attr.value if isinstance(_role_from_attr, LitellmUserRoles) else _role_from_attr ) if user_id is None and result is not None: @@ -3289,15 +3033,11 @@ class SSOAuthenticationHandler: from litellm.proxy.utils import get_prisma_client_or_throw from litellm.types.proxy.ui_sso import ReturnedUITokenObject - prisma_client = get_prisma_client_or_throw( - "Prisma client is None, connect a database to your proxy" - ) + prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") # User is Authe'd in - generate key for the UI to access Proxy - parsed_openid_result = ( - SSOAuthenticationHandler._get_user_email_and_id_from_result( - result=result, generic_client_id=generic_client_id - ) + parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result( + result=result, generic_client_id=generic_client_id ) user_email = parsed_openid_result.get("user_email") user_id = parsed_openid_result.get("user_id") @@ -3376,9 +3116,7 @@ class SSOAuthenticationHandler: "Unable to map user identity to known values. 'user_defined_values' is None. File an issue - https://github.com/BerriAI/litellm/issues" ) - verbose_proxy_logger.info( - f"user_defined_values for creating ui key: {user_defined_values}" - ) + verbose_proxy_logger.info(f"user_defined_values for creating ui key: {user_defined_values}") default_ui_key_values.update(user_defined_values) default_ui_key_values["request_type"] = "key" @@ -3390,18 +3128,13 @@ class SSOAuthenticationHandler: key = response["token"] # type: ignore user_id = response["user_id"] # type: ignore - user_role = ( - user_defined_values["user_role"] - or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - ) + user_role = user_defined_values["user_role"] or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value if user_id and isinstance(user_id, str): user_role = await check_and_update_if_proxy_admin_id( user_role=user_role, user_id=user_id, prisma_client=prisma_client ) - verbose_proxy_logger.debug( - f"user_role: {user_role}; ui_access_mode: {ui_access_mode}" - ) + verbose_proxy_logger.debug(f"user_role: {user_role}; ui_access_mode: {ui_access_mode}") ## CHECK IF ROLE ALLOWED TO USE PROXY ## is_admin_only_access = check_is_admin_only_access(ui_access_mode or {}) if is_admin_only_access: @@ -3414,19 +3147,12 @@ class SSOAuthenticationHandler: }, ) - disabled_non_admin_personal_key_creation = ( - get_disabled_non_admin_personal_key_creation() - ) - litellm_dashboard_ui = get_custom_url( - request_base_url=str(request.base_url), route="ui/" - ) + disabled_non_admin_personal_key_creation = get_disabled_non_admin_personal_key_creation() + litellm_dashboard_ui = get_custom_url(request_base_url=str(request.base_url), route="ui/") if get_secret_bool("EXPERIMENTAL_UI_LOGIN"): _user_info: Optional[LiteLLM_UserTable] = None - if ( - user_defined_values is not None - and user_defined_values["user_id"] is not None - ): + if user_defined_values is not None and user_defined_values["user_id"] is not None: _user_info = LiteLLM_UserTable( user_id=user_defined_values["user_id"], user_role=user_defined_values["user_role"] or user_role, @@ -3436,14 +3162,10 @@ class SSOAuthenticationHandler: if _user_info is None: raise HTTPException( status_code=401, - detail={ - "error": "User Information is required for experimental UI login" - }, + detail={"error": "User Information is required for experimental UI login"}, ) - key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - _user_info - ) + key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(_user_info) returned_ui_token_object = ReturnedUITokenObject( user_id=cast(str, user_id), @@ -3452,9 +3174,7 @@ class SSOAuthenticationHandler: user_role=user_role or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, login_method="sso", premium_user=premium_user, - auth_header_name=general_settings.get( - "litellm_key_header_name", "Authorization" - ), + auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"), disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, server_root_path=get_server_root_path(), ) @@ -3468,28 +3188,18 @@ class SSOAuthenticationHandler: # Control-plane cross-origin: store JWT behind a single-use opaque # code (60s TTL) so the token never appears in browser history / logs. # The control plane redeems it via POST /v3/login/exchange. - if return_to is not None and SSOAuthenticationHandler._validate_return_to( - return_to - ): + if return_to is not None and SSOAuthenticationHandler._validate_return_to(return_to): code = secrets.token_urlsafe(32) cache_key = f"login_code:{code}" cache_value = {"token": jwt_token, "redirect_url": return_to} if redis_usage_cache is not None: - await redis_usage_cache.async_set_cache( - key=cache_key, value=cache_value, ttl=60 - ) + await redis_usage_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) else: - await user_api_key_cache.async_set_cache( - key=cache_key, value=cache_value, ttl=60 - ) + await user_api_key_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) separator = "&" if "?" in return_to else "?" - redirect_url = ( - return_to + separator + urlencode({"login": "success", "code": code}) - ) - verbose_proxy_logger.info( - "Cross-origin SSO: redirecting to control plane with login code" - ) + redirect_url = return_to + separator + urlencode({"login": "success", "code": code}) + verbose_proxy_logger.info("Cross-origin SSO: redirecting to control plane with login code") redirect_response = RedirectResponse(url=redirect_url, status_code=303) redirect_response.delete_cookie("litellm_cp_return_to") return redirect_response @@ -3564,15 +3274,12 @@ class SSOAuthenticationHandler: state, ) else: - verbose_proxy_logger.debug( - "PKCE code_verifier retrieved from cache" - ) + verbose_proxy_logger.debug("PKCE code_verifier retrieved from cache") elif isinstance(cached_data, str): # Handle legacy format (plain string) for backward compatibility code_verifier = cached_data verbose_proxy_logger.warning( - "Retrieved code_verifier in legacy plain-string format. " - "Future storage will use dict format." + "Retrieved code_verifier in legacy plain-string format. Future storage will use dict format." ) else: # Defer the detailed ERROR log to the strict-mode branch below @@ -3615,12 +3322,8 @@ class SSOAuthenticationHandler: In strict mode (PKCE_STRICT_CACHE_MISS=true) raises ProxyException. Otherwise logs a warning and returns (token exchange proceeds without verifier). """ - active_cache = ( - redis_usage_cache if redis_usage_cache is not None else user_api_key_cache - ) - strict_cache_miss = ( - os.getenv("PKCE_STRICT_CACHE_MISS", "false").lower() == "true" - ) + active_cache = redis_usage_cache if redis_usage_cache is not None else user_api_key_cache + strict_cache_miss = os.getenv("PKCE_STRICT_CACHE_MISS", "false").lower() == "true" if strict_cache_miss: if empty_value_in_dict: await SSOAuthenticationHandler._delete_pkce_verifier(cache_key) @@ -3664,8 +3367,7 @@ class SSOAuthenticationHandler: "Configure Redis so all proxy instances share the PKCE verifier." ) verbose_proxy_logger.error( - "PKCE is enabled but no verifier found in cache for state '%s'. " - "%s Cache type: %s.", + "PKCE is enabled but no verifier found in cache for state '%s'. %s Cache type: %s.", state, cause, type(active_cache).__name__, @@ -3723,17 +3425,11 @@ class SSOAuthenticationHandler: """ # Generate a cryptographically random code_verifier (43 characters) # Using 32 random bytes which becomes 43 characters when base64-url-encoded - code_verifier = ( - base64.urlsafe_b64encode(secrets.token_bytes(32)) - .decode("utf-8") - .rstrip("=") - ) + code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("utf-8").rstrip("=") # Generate code_challenge using S256 method (SHA256) code_challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest() - code_challenge = ( - base64.urlsafe_b64encode(code_challenge_bytes).decode("utf-8").rstrip("=") - ) + code_challenge = base64.urlsafe_b64encode(code_challenge_bytes).decode("utf-8").rstrip("=") return code_verifier, code_challenge @@ -3788,9 +3484,7 @@ class SSOAuthenticationHandler: "token endpoint returned HTTP 200 but no access_token " f"(response keys: {sorted(token_response.keys())})" ) - verbose_proxy_logger.error( - "Token response missing or null access_token. detail=%s", detail - ) + verbose_proxy_logger.error("Token response missing or null access_token. detail=%s", detail) raise ProxyException( message=f"Token exchange failed: {detail}", type=ProxyErrorTypes.auth_error, @@ -3845,9 +3539,7 @@ class SSOAuthenticationHandler: if not include_client_id: # Use Basic Auth only when a secret is available; public PKCE clients omit it. if client_secret: - credentials = base64.b64encode( - f"{client_id}:{client_secret}".encode() - ).decode() + credentials = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() request_headers["Authorization"] = f"Basic {credentials}" else: token_data["client_id"] = client_id @@ -3856,9 +3548,7 @@ class SSOAuthenticationHandler: if client_secret: token_data["client_secret"] = client_secret - http_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.SSO_HANDLER - ) + http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SSO_HANDLER) try: response = await http_client.post( url=token_endpoint, @@ -3948,9 +3638,7 @@ class SSOAuthenticationHandler: if userinfo_endpoint: try: - client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.SSO_HANDLER - ) + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SSO_HANDLER) resp = await client.get( url=userinfo_endpoint, headers={ @@ -3985,9 +3673,7 @@ class SSOAuthenticationHandler: resp.text[:500], ) except Exception as e: - verbose_proxy_logger.warning( - "Userinfo endpoint error: %s, falling back to id_token", e - ) + verbose_proxy_logger.warning("Userinfo endpoint error: %s, falling back to id_token", e) # Only fall back to id_token when the userinfo request failed (None). # Empty dict ({}) and JSON null are both treated as failure (set to None above) since @@ -4001,9 +3687,7 @@ class SSOAuthenticationHandler: # jwt.decode returned an empty dict (payload-free JWT or provider bug). # Treat this the same as a missing userinfo — the session would have no # identity claims, which is equivalent to a broken session. - verbose_proxy_logger.warning( - "id_token decoded to an empty payload — treating as failure." - ) + verbose_proxy_logger.warning("id_token decoded to an empty payload — treating as failure.") userinfo = None except Exception as decode_err: verbose_proxy_logger.error("Failed to decode id_token: %s", decode_err) @@ -4031,7 +3715,9 @@ class SSOAuthenticationHandler: "and id_token decoded to an empty payload — no identity claims available" ) else: - detail = "no userinfo endpoint is configured (GENERIC_USERINFO_ENDPOINT) and no id_token was present" + detail = ( + "no userinfo endpoint is configured (GENERIC_USERINFO_ENDPOINT) and no id_token was present" + ) raise ProxyException( message=f"SSO user info unavailable: {detail}.", type=ProxyErrorTypes.auth_error, @@ -4107,9 +3793,7 @@ class MicrosoftSSOHandler: ) # Extract app roles from the id_token JWT - app_roles = MicrosoftSSOHandler.get_app_roles_from_id_token( - id_token=microsoft_sso.id_token - ) + app_roles = MicrosoftSSOHandler.get_app_roles_from_id_token(id_token=microsoft_sso.id_token) verbose_proxy_logger.debug(f"Extracted app roles from id_token: {app_roles}") # Combine groups and app roles @@ -4120,20 +3804,14 @@ class MicrosoftSSOHandler: role = get_litellm_user_role(role_str) if role is not None: user_role = role - verbose_proxy_logger.debug( - f"Found valid LitellmUserRoles '{role.value}' in app_roles" - ) + verbose_proxy_logger.debug(f"Found valid LitellmUserRoles '{role.value}' in app_roles") break - verbose_proxy_logger.debug( - f"Combined team_ids (groups + app roles): {user_team_ids}" - ) + verbose_proxy_logger.debug(f"Combined team_ids (groups + app roles): {user_team_ids}") # if user is trying to get the raw sso response for debugging, return the raw sso response if return_raw_sso_response: - original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = ( - user_team_ids - ) + original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = user_team_ids original_msft_result["app_roles"] = app_roles return original_msft_result or {} @@ -4153,9 +3831,7 @@ class MicrosoftSSOHandler: response = response or {} verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") openid_response = CustomOpenID( - email=normalize_email( - response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail") - ), + email=normalize_email(response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail")), display_name=response.get(MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE), provider="microsoft", id=response.get(MICROSOFT_USER_ID_ATTRIBUTE), @@ -4197,14 +3873,10 @@ class MicrosoftSSOHandler: roles = decoded_token.get("app_roles", []) or decoded_token.get("roles", []) if roles and isinstance(roles, list): - verbose_proxy_logger.debug( - f"Found {len(roles)} app role(s) in id_token: {roles}" - ) + verbose_proxy_logger.debug(f"Found {len(roles)} app role(s) in id_token: {roles}") return roles else: - verbose_proxy_logger.debug( - "No app roles found in id_token or roles claim is not a list" - ) + verbose_proxy_logger.debug("No app roles found in id_token or roles claim is not a list") return [] except Exception as e: @@ -4225,9 +3897,7 @@ class MicrosoftSSOHandler: List[str]: List of group IDs the user belongs to """ try: - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.SSO_HANDLER - ) + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SSO_HANDLER) # Handle MSFT Enterprise Application Groups service_principal_id = os.getenv("MICROSOFT_SERVICE_PRINCIPAL_ID", None) @@ -4242,9 +3912,7 @@ class MicrosoftSSOHandler: async_client=async_client, access_token=access_token, ) - verbose_proxy_logger.debug( - f"Service principal group IDs: {service_principal_group_ids}" - ) + verbose_proxy_logger.debug(f"Service principal group IDs: {service_principal_group_ids}") if len(service_principal_group_ids) > 0: await MicrosoftSSOHandler.create_litellm_teams_from_service_principal_team_ids( service_principal_teams=service_principal_teams, @@ -4252,44 +3920,30 @@ class MicrosoftSSOHandler: # Fetch user membership from Microsoft Graph API all_group_ids = [] - next_link: Optional[str] = ( - MicrosoftSSOHandler.graph_api_user_groups_endpoint - ) + next_link: Optional[str] = MicrosoftSSOHandler.graph_api_user_groups_endpoint auth_headers = {"Authorization": f"Bearer {access_token}"} page_count = 0 - while ( - next_link is not None - and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES - ): + while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: group_ids, next_link = await MicrosoftSSOHandler.fetch_and_parse_groups( url=next_link, headers=auth_headers, async_client=async_client ) all_group_ids.extend(group_ids) page_count += 1 - if ( - next_link is not None - and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES - ): + if next_link is not None and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: verbose_proxy_logger.warning( f"Reached maximum page limit of {MicrosoftSSOHandler.MAX_GRAPH_API_PAGES}. Some groups may not be included." ) # If service_principal_group_ids is not empty, only return group_ids that are in both all_group_ids and service_principal_group_ids if service_principal_group_ids and len(service_principal_group_ids) > 0: - all_group_ids = [ - group_id - for group_id in all_group_ids - if group_id in service_principal_group_ids - ] + all_group_ids = [group_id for group_id in all_group_ids if group_id in service_principal_group_ids] return all_group_ids except Exception as e: - verbose_proxy_logger.error( - f"Error getting user groups from Microsoft Graph API: {e}" - ) + verbose_proxy_logger.error(f"Error getting user groups from Microsoft Graph API: {e}") return [] @staticmethod @@ -4299,12 +3953,8 @@ class MicrosoftSSOHandler: """Helper function to fetch and parse group data from a URL""" response = await async_client.get(url, headers=headers) response_json = response.json() - response_typed = await MicrosoftSSOHandler._cast_graph_api_response_dict( - response=response_json - ) - group_ids = MicrosoftSSOHandler._get_group_ids_from_graph_api_response( - response=response_typed - ) + response_typed = await MicrosoftSSOHandler._cast_graph_api_response_dict(response=response_json) + group_ids = MicrosoftSSOHandler._get_group_ids_from_graph_api_response(response=response_typed) return group_ids, response_typed.get("odata_nextLink") @staticmethod @@ -4365,9 +4015,7 @@ class MicrosoftSSOHandler: response = await async_client.get(url, headers=headers) response_json = response.json() - verbose_proxy_logger.debug( - f"Response from service principal app role assigned to: {response_json}" - ) + verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}") group_ids: List[str] = [] service_principal_teams: List[MicrosoftServicePrincipalTeam] = [] @@ -4394,14 +4042,10 @@ class MicrosoftSSOHandler: When a user sets a `SERVICE_PRINCIPAL_ID` in the env, litellm will fetch groups under that service principal and create Litellm Teams from them """ - verbose_proxy_logger.debug( - f"Creating Litellm Teams from Service Principal Teams: {service_principal_teams}" - ) + verbose_proxy_logger.debug(f"Creating Litellm Teams from Service Principal Teams: {service_principal_teams}") for service_principal_team in service_principal_teams: litellm_team_id: Optional[str] = service_principal_team.get("principalId") - litellm_team_name: Optional[str] = service_principal_team.get( - "principalDisplayName" - ) + litellm_team_name: Optional[str] = service_principal_team.get("principalDisplayName") if not litellm_team_id: verbose_proxy_logger.debug( f"Skipping team creation for {litellm_team_name} because it has no principalId" @@ -4476,11 +4120,7 @@ async def debug_sso_login(request: Request): generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) ####### Check if user is a Enterprise / Premium User ####### - if ( - microsoft_client_id is not None - or google_client_id is not None - or generic_client_id is not None - ): + if microsoft_client_id is not None or google_client_id is not None or generic_client_id is not None: if premium_user is not True: raise ProxyException( message="You must be a LiteLLM Enterprise user to use SSO. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", @@ -4539,9 +4179,7 @@ async def debug_sso_callback(request: Request): prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, litellm_jwtauth=LiteLLM_JWTAuth( - team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get( - "sso_group_jwt_field", None - ), + team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get("sso_group_jwt_field", None), ), leeway=0, ) @@ -4615,16 +4253,8 @@ async def debug_sso_callback(request: Request): # Defense-in-depth: ensure no bearer tokens leak into the rendered HTML even if # a non-conforming IdP places them in its userinfo response. - safe_raw_claims = { - k: v - for k, v in (received_response or {}).items() - if k not in _OAUTH_TOKEN_FIELDS - } - safe_access_token_claims = { - k: v - for k, v in (access_token_payload or {}).items() - if k not in _OAUTH_TOKEN_FIELDS - } + safe_raw_claims = {k: v for k, v in (received_response or {}).items() if k not in _OAUTH_TOKEN_FIELDS} + safe_access_token_claims = {k: v for k, v in (access_token_payload or {}).items() if k not in _OAUTH_TOKEN_FIELDS} sso_payload = { "parsed_by_proxy": filtered_result, @@ -4633,9 +4263,7 @@ async def debug_sso_callback(request: Request): } # Replace the placeholder in the template with the actual data - sso_payload_json = json.dumps(sso_payload, indent=2, default=str).replace( - " str: """Build role-appropriate system prompt with today's date.""" tool_desc = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE - return ( - f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}Today's date: {date.today().isoformat()}" - ) + return f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}Today's date: {date.today().isoformat()}" # keep a public reference for test assertions @@ -258,9 +253,7 @@ async def _query_activity( ) -async def _fetch_usage_data( - start_date: str, end_date: str, user_id: Optional[str] = None -) -> Dict[str, Any]: +async def _fetch_usage_data(start_date: str, end_date: str, user_id: Optional[str] = None) -> Dict[str, Any]: resp = await _query_activity( TABLE_DAILY_USER_SPEND, ENTITY_FIELD_USER, @@ -272,9 +265,7 @@ async def _fetch_usage_data( return resp.model_dump(mode="json") -async def _fetch_team_usage_data( - start_date: str, end_date: str, team_ids: Optional[str] = None -) -> Dict[str, Any]: +async def _fetch_team_usage_data(start_date: str, end_date: str, team_ids: Optional[str] = None) -> Dict[str, Any]: resp = await _query_activity( TABLE_DAILY_TEAM_SPEND, ENTITY_FIELD_TEAM, @@ -285,9 +276,7 @@ async def _fetch_team_usage_data( return resp.model_dump(mode="json") -async def _fetch_tag_usage_data( - start_date: str, end_date: str, tags: Optional[str] = None -) -> Dict[str, Any]: +async def _fetch_tag_usage_data(start_date: str, end_date: str, tags: Optional[str] = None) -> Dict[str, Any]: resp = await _query_activity( TABLE_DAILY_TAG_SPEND, ENTITY_FIELD_TAG, @@ -324,12 +313,7 @@ def _ranked_lines( limit: int, ) -> List[str]: """Sort by spend descending, format each entry, and truncate.""" - return [ - fmt(name, vals) - for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[ - :limit - ] - ] + return [fmt(name, vals) for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[:limit]] def _summarise_usage_data(data: Dict[str, Any]) -> str: @@ -344,16 +328,12 @@ def _summarise_usage_data(data: Dict[str, Any]) -> str: f"Total Tokens: {meta.get('total_tokens', 0)}" ) - models = _accumulate_breakdown( - results, "models", ["spend", "api_requests", "total_tokens"] - ) + models = _accumulate_breakdown(results, "models", ["spend", "api_requests", "total_tokens"]) providers = _accumulate_breakdown(results, "providers", ["spend", "api_requests"]) model_lines = _ranked_lines( models, - lambda n, d: ( - f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)" - ), + lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)", TOP_N_MODELS, ) provider_lines = _ranked_lines( @@ -389,8 +369,7 @@ def _summarise_entity_data(data: Dict[str, Any], entity_label: str) -> str: for eid, d in sorted(totals.items(), key=lambda x: -x[1]["spend"]): label = d["alias"] if d["alias"] != eid else eid lines.append( - f"- {label} (ID: {eid}): ${d['spend']:.4f} | " - f"{int(d['requests'])} reqs | {int(d['tokens'])} tokens" + f"- {label} (ID: {eid}): ${d['spend']:.4f} | {int(d['requests'])} reqs | {int(d['tokens'])} tokens" ) return "\n".join(lines) @@ -505,23 +484,17 @@ async def _process_tool_call( yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "running"})) try: - tool_result = await _execute_tool_call( - handler, fn_name, fn_args, user_id, is_admin - ) + tool_result = await _execute_tool_call(handler, fn_name, fn_args, user_id, is_admin) yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "complete"})) except Exception as e: verbose_proxy_logger.error("Tool %s failed: %s", fn_name, e) tool_result = f"Error fetching {handler['label']}. Please try again." yield _sse(cast(SSEToolCallEvent, {**tool_event_base, "status": "error"})) - chat_messages.append( - {"role": "tool", "tool_call_id": tc.id, "content": tool_result} - ) + chat_messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_result}) -async def _stream_final_response( - model: str, chat_messages: List[Dict[str, Any]] -) -> AsyncIterator[str]: +async def _stream_final_response(model: str, chat_messages: List[Dict[str, Any]]) -> AsyncIterator[str]: """Stream the final LLM response after tool results are appended.""" yield _sse({"type": "status", "message": "Analyzing results..."}) @@ -545,9 +518,7 @@ async def stream_usage_ai_chat( ) -> AsyncIterator[str]: """Stream SSE events: status → tool_call → chunk → done.""" resolved_model = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL - truncated = ( - messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages - ) + truncated = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages chat_messages: List[Dict[str, Any]] = [ {"role": "system", "content": _build_system_prompt(is_admin)}, *truncated, diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py index d0df80fed0d..26515c749fb 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -22,9 +22,7 @@ class ChatMessage(BaseModel): class UsageAIChatRequest(BaseModel): - messages: List[ChatMessage] = Field( - ..., description="Chat messages (user/assistant history)" - ) + messages: List[ChatMessage] = Field(..., description="Chat messages (user/assistant history)") model: Optional[str] = Field(default=None, description="Model to use for AI chat") diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index 661487577c3..af65eb6c3d8 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -40,9 +40,7 @@ class TagActiveUsersResponse(BaseModel): tag: str active_users: int date: str # The specific date or period identifier - period_start: Optional[str] = ( - None # For WAU/MAU, this will be the start of the period - ) + period_start: Optional[str] = None # For WAU/MAU, this will be the start of the period period_end: Optional[str] = None # For WAU/MAU, this will be the end of the period @@ -198,9 +196,7 @@ async def get_daily_active_users( # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace( - hour=0, minute=0, second=0, microsecond=0 - ) + timedelta(days=1) + end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") # Calculate date range (last MAX_DAYS days) @@ -208,9 +204,7 @@ async def get_daily_active_users( start_date = start_dt.strftime("%Y-%m-%d") # Build SQL query with optional tag filter(s) - where_clause = ( - "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" - ) + where_clause = "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" params = [start_date, end_date] # Handle multiple tag filters (takes precedence over single tag filter) @@ -240,9 +234,7 @@ async def get_daily_active_users( db_response = await prisma_client.db.query_raw(sql_query, *params) results = [ - TagActiveUsersResponse( - tag=row["tag"], active_users=row["active_users"], date=row["date"] - ) + TagActiveUsersResponse(tag=row["tag"], active_users=row["active_users"], date=row["date"]) for row in db_response ] @@ -301,22 +293,16 @@ async def get_weekly_active_users( # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace( - hour=0, minute=0, second=0, microsecond=0 - ) + timedelta(days=1) + end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") # Calculate date range for all weeks (49 days total) # Start from 48 days before end_date to cover exactly MAX_WEEKS complete weeks - start_dt = end_dt - timedelta( - days=(MAX_WEEKS * 7 - 1) - ) # MAX_WEEKS weeks * 7 days - 1 + start_dt = end_dt - timedelta(days=(MAX_WEEKS * 7 - 1)) # MAX_WEEKS weeks * 7 days - 1 start_date = start_dt.strftime("%Y-%m-%d") # Build SQL query with optional tag filter(s) - where_clause = ( - "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" - ) + where_clause = "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" params = [start_date, end_date] # Handle multiple tag filters (takes precedence over single tag filter) @@ -366,9 +352,7 @@ async def get_weekly_active_users( TagActiveUsersResponse( tag=row["tag"], active_users=row["active_users"], - date=row[ - "date" - ], # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc. + date=row["date"], # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc. period_start=row["period_start"], period_end=row["period_end"], ) @@ -430,22 +414,16 @@ async def get_monthly_active_users( # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace( - hour=0, minute=0, second=0, microsecond=0 - ) + timedelta(days=1) + end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") # Calculate date range for all months (210 days total) # Start from 209 days before end_date to cover exactly MAX_MONTHS complete months - start_dt = end_dt - timedelta( - days=(MAX_MONTHS * 30 - 1) - ) # MAX_MONTHS months * 30 days - 1 + start_dt = end_dt - timedelta(days=(MAX_MONTHS * 30 - 1)) # MAX_MONTHS months * 30 days - 1 start_date = start_dt.strftime("%Y-%m-%d") # Build SQL query with optional tag filter(s) - where_clause = ( - "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" - ) + where_clause = "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" params = [start_date, end_date] # Handle multiple tag filters (takes precedence over single tag filter) @@ -662,9 +640,7 @@ async def get_per_user_analytics( # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace( - hour=0, minute=0, second=0, microsecond=0 - ) + timedelta(days=1) + end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") # Calculate date range (last 30 days) @@ -681,9 +657,7 @@ async def get_per_user_analytics( where_clause["tag"] = {"contains": tag_filter} # Get all tag records in the date range with optional tag filtering - tag_records = await DailyTagSpendRepository(prisma_client).table.find_many( - where=where_clause - ) + tag_records = await DailyTagSpendRepository(prisma_client).table.find_many(where=where_clause) # Get unique api_keys api_keys = set(record.api_key for record in tag_records if record.api_key) @@ -698,25 +672,19 @@ async def get_per_user_analytics( ) # Lookup user_id for each api_key - api_key_records = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"token": {"in": list(api_keys)}}) + api_key_records = await VerificationTokenRepository(prisma_client).table.find_many( + where={"token": {"in": list(api_keys)}} + ) # Create mapping from api_key to user_id - api_key_to_user_id = { - record.token: record.user_id for record in api_key_records if record.user_id - } + api_key_to_user_id = {record.token: record.user_id for record in api_key_records if record.user_id} # Get user emails for the user_ids user_ids = list(set(api_key_to_user_id.values())) - user_records = await UserRepository(prisma_client).table.find_many( - where={"user_id": {"in": user_ids}} - ) + user_records = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": user_ids}}) # Create mapping from user_id to user_email - user_id_to_email = { - record.user_id: record.user_email for record in user_records - } + user_id_to_email = {record.user_id: record.user_email for record in user_records} # Aggregate metrics by user user_metrics: Dict[str, PerUserMetrics] = {} @@ -738,17 +706,13 @@ async def get_per_user_analytics( user_metrics[user_id].user_agent = tag # Aggregate metrics - user_metrics[user_id].successful_requests += ( - record.successful_requests or 0 - ) + user_metrics[user_id].successful_requests += record.successful_requests or 0 user_metrics[user_id].failed_requests += record.failed_requests or 0 user_metrics[user_id].total_requests += record.api_requests or 0 # Calculate total_tokens from prompt_tokens + completion_tokens prompt_tokens = record.prompt_tokens or 0 completion_tokens = record.completion_tokens or 0 - user_metrics[user_id].total_tokens += int( - prompt_tokens + completion_tokens - ) + user_metrics[user_id].total_tokens += int(prompt_tokens + completion_tokens) user_metrics[user_id].spend += record.spend or 0.0 # Convert to list and sort by successful requests (descending) diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 57cc0dc6745..b2488d20127 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -121,9 +121,7 @@ async def _require_run( user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> Any: """Return the run or raise 404. For non-admin callers, also enforce key ownership.""" - run = await WorkflowRunRepository(prisma_client).table.find_unique( - where={"run_id": run_id} - ) + run = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id}) if run is None: raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") if user_api_key_dict is not None and not _is_admin(user_api_key_dict): @@ -155,9 +153,7 @@ async def create_workflow_run( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: create_data: Dict[str, Any] = { @@ -193,9 +189,7 @@ async def list_workflow_runs( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) where: Dict[str, Any] = {} if workflow_type: @@ -235,9 +229,7 @@ async def get_workflow_run( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: run = await WorkflowRunRepository(prisma_client).table.find_unique( @@ -272,9 +264,7 @@ async def update_workflow_run( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) update: Dict[str, Any] = {} if data.status is not None: @@ -324,9 +314,7 @@ async def append_workflow_event( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) await _require_run(prisma_client, run_id, user_api_key_dict) @@ -370,9 +358,7 @@ async def append_workflow_event( verbose_proxy_logger.exception("Error appending workflow event: %s", e) raise HTTPException(status_code=500, detail=str(e)) - raise HTTPException( - status_code=500, detail="Failed to append event" - ) # pragma: no cover + raise HTTPException(status_code=500, detail="Failed to append event") # pragma: no cover @router.get( @@ -389,9 +375,7 @@ async def list_workflow_events( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) await _require_run(prisma_client, run_id, user_api_key_dict) @@ -424,9 +408,7 @@ async def append_workflow_message( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) await _require_run(prisma_client, run_id, user_api_key_dict) @@ -441,9 +423,7 @@ async def append_workflow_message( } if data.session_id is not None: msg_data["session_id"] = data.session_id - msg = await WorkflowMessageRepository(prisma_client).table.create( - data=msg_data - ) + msg = await WorkflowMessageRepository(prisma_client).table.create(data=msg_data) return msg except Exception as e: @@ -462,9 +442,7 @@ async def append_workflow_message( verbose_proxy_logger.exception("Error appending workflow message: %s", e) raise HTTPException(status_code=500, detail=str(e)) - raise HTTPException( - status_code=500, detail="Failed to append message" - ) # pragma: no cover + raise HTTPException(status_code=500, detail="Failed to append message") # pragma: no cover @router.get( @@ -481,9 +459,7 @@ async def list_workflow_messages( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) await _require_run(prisma_client, run_id, user_api_key_dict) diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index 33599c3c622..c184ce6bca5 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -57,15 +57,10 @@ def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]: return _audit_log_callback_cache[name] instance: Optional[CustomLogger] - if ( - name == "s3_v2" - and getattr(litellm, "s3_audit_callback_params", None) is not None - ): + if name == "s3_v2" and getattr(litellm, "s3_audit_callback_params", None) is not None: from litellm.integrations.s3_v2 import S3Logger as S3V2Logger - instance = S3V2Logger( - s3_callback_params_override=litellm.s3_audit_callback_params - ) + instance = S3V2Logger(s3_callback_params_override=litellm.s3_audit_callback_params) else: from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, @@ -121,9 +116,7 @@ def _audit_log_task_done_callback(task: asyncio.Task) -> None: except asyncio.CancelledError: return if exc is not None: - verbose_proxy_logger.error( - "Audit log callback task failed: %s", exc, exc_info=exc - ) + verbose_proxy_logger.error("Audit log callback task failed: %s", exc, exc_info=exc) async def _dispatch_audit_log_to_callbacks( @@ -137,24 +130,18 @@ async def _dispatch_audit_log_to_callbacks( for callback in litellm.audit_log_callbacks: try: - resolved: Optional[CustomLogger] = ( - callback if isinstance(callback, CustomLogger) else None - ) + resolved: Optional[CustomLogger] = callback if isinstance(callback, CustomLogger) else None if isinstance(callback, str): resolved = _resolve_audit_log_callback(callback) if resolved is None: - verbose_proxy_logger.warning( - "Could not resolve audit log callback: %s", callback - ) + verbose_proxy_logger.warning("Could not resolve audit log callback: %s", callback) continue if isinstance(resolved, CustomLogger): task = asyncio.create_task(resolved.async_log_audit_log_event(payload)) task.add_done_callback(_audit_log_task_done_callback) except Exception as e: - verbose_proxy_logger.error( - "Failed dispatching audit log to callback: %s", e - ) + verbose_proxy_logger.error("Failed dispatching audit log to callback: %s", e) async def create_object_audit_log( @@ -180,9 +167,7 @@ async def create_object_audit_log( """ from litellm.secret_managers.main import get_secret_bool - _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool( - "LITELLM_STORE_AUDIT_LOGS" - ) + _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool("LITELLM_STORE_AUDIT_LOGS") if _store_audit_logs is not True: return @@ -214,9 +199,7 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): """ from litellm.secret_managers.main import get_secret_bool - _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool( - "LITELLM_STORE_AUDIT_LOGS" - ) + _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool("LITELLM_STORE_AUDIT_LOGS") if _store_audit_logs is not True: return @@ -237,9 +220,7 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): await _dispatch_audit_log_to_callbacks(request_data) if prisma_client is None: - verbose_proxy_logger.error( - "prisma_client is None, cannot write audit log to DB" - ) + verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB") return _request_data = request_data.model_dump(exclude_none=True) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 0a61fc5d3dd..9d5f716033f 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -51,9 +51,7 @@ async def attach_object_permission_to_dict( object_permission_id = data_dict.get("object_permission_id") if object_permission_id: - object_permission = await ObjectPermissionRepository( - prisma_client - ).table.find_unique( + object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id}, ) if object_permission: @@ -104,22 +102,16 @@ async def handle_update_object_permission_common( return None # Lookup existing object permission ID and update that entry - object_permission_id_to_use: str = existing_object_permission_id or str( - uuid.uuid4() - ) + object_permission_id_to_use: str = existing_object_permission_id or str(uuid.uuid4()) existing_object_permissions_dict: Dict = {} - existing_object_permission = await ObjectPermissionRepository( - prisma_client - ).table.find_unique( + existing_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id_to_use}, ) # Update the object permission if existing_object_permission is not None: - existing_object_permissions_dict = existing_object_permission.model_dump( - exclude_unset=True, exclude_none=True - ) + existing_object_permissions_dict = existing_object_permission.model_dump(exclude_unset=True, exclude_none=True) # Handle string JSON object permission if isinstance(new_object_permission, str): @@ -140,9 +132,7 @@ async def handle_update_object_permission_common( ######################################################### # Commit the update to the LiteLLM_ObjectPermissionTable ######################################################### - created_object_permission_row = await ObjectPermissionRepository( - prisma_client - ).table.upsert( + created_object_permission_row = await ObjectPermissionRepository(prisma_client).table.upsert( where={"object_permission_id": object_permission_id_to_use}, data={ "create": existing_object_permissions_dict, @@ -150,9 +140,7 @@ async def handle_update_object_permission_common( }, ) - verbose_proxy_logger.debug( - f"created_object_permission_row: {created_object_permission_row}" - ) + verbose_proxy_logger.debug(f"created_object_permission_row: {created_object_permission_row}") return created_object_permission_row.object_permission_id @@ -174,21 +162,13 @@ async def _set_object_permission( return data_json # Clean data: exclude None values and object_permission_id - clean_data = { - k: v - for k, v in permission_data.items() - if v is not None and k != "object_permission_id" - } + clean_data = {k: v for k, v in permission_data.items() if v is not None and k != "object_permission_id"} # Serialize mcp_tool_permissions to JSON string for GraphQL compatibility if "mcp_tool_permissions" in clean_data: - clean_data["mcp_tool_permissions"] = safe_dumps( - clean_data["mcp_tool_permissions"] - ) + clean_data["mcp_tool_permissions"] = safe_dumps(clean_data["mcp_tool_permissions"]) - created_permission = await ObjectPermissionRepository(prisma_client).table.create( - data=clean_data - ) + created_permission = await ObjectPermissionRepository(prisma_client).table.create(data=clean_data) data_json["object_permission_id"] = created_permission.object_permission_id data_json.pop("object_permission") @@ -270,9 +250,7 @@ async def _resolve_mcp_server_identifiers_to_ids( if not server_id: continue for identifier in identifiers: - if identifier == registry_key or _mcp_server_identifier_matches( - server, identifier - ): + if identifier == registry_key or _mcp_server_identifier_matches(server, identifier): resolved[identifier].add(server_id) return resolved @@ -312,8 +290,7 @@ def _rewrite_object_permission_mcp_tool_permissions( normalized_tool_permissions[server_id].extend(tools) object_permission["mcp_tool_permissions"] = { - server_id: _dedupe_preserving_order(tools) - for server_id, tools in normalized_tool_permissions.items() + server_id: _dedupe_preserving_order(tools) for server_id, tools in normalized_tool_permissions.items() } @@ -337,11 +314,7 @@ def _rewrite_object_permission_mcp_identifiers( def _flatten_resolved_mcp_server_ids( identifier_to_server_ids: Dict[str, Set[str]], ) -> Set[str]: - return { - server_id - for server_ids in identifier_to_server_ids.values() - for server_id in server_ids - } + return {server_id for server_ids in identifier_to_server_ids.values() for server_id in server_ids} async def _resolve_team_allowed_mcp_servers( @@ -361,9 +334,7 @@ async def _resolve_team_allowed_mcp_servers( ) direct_servers: List[str] = team_object_permission.mcp_servers or [] - access_group_servers: List[ - str - ] = await MCPRequestHandler._get_mcp_servers_from_access_groups( + access_group_servers: List[str] = await MCPRequestHandler._get_mcp_servers_from_access_groups( team_object_permission.mcp_access_groups or [] ) raw_tool_perms = team_object_permission.mcp_tool_permissions or {} @@ -375,9 +346,7 @@ async def _resolve_team_allowed_mcp_servers( identifiers=raw_servers, prisma_client=prisma_client, ) - unresolved_servers = { - server_id for server_id in raw_servers if not resolved_servers.get(server_id) - } + unresolved_servers = {server_id for server_id in raw_servers if not resolved_servers.get(server_id)} return _flatten_resolved_mcp_server_ids(resolved_servers) | unresolved_servers @@ -515,9 +484,7 @@ async def validate_key_mcp_servers_against_team( prisma_client=prisma_client, ) stale_identifiers = { - identifier - for identifier in requested_servers - if not identifier_to_server_ids.get(identifier) + identifier for identifier in requested_servers if not identifier_to_server_ids.get(identifier) } if stale_identifiers: verbose_proxy_logger.warning( @@ -528,9 +495,7 @@ async def validate_key_mcp_servers_against_team( object_permission=object_permission, identifier_to_server_ids=identifier_to_server_ids, ) - active_requested_servers = _flatten_resolved_mcp_server_ids( - identifier_to_server_ids - ) + active_requested_servers = _flatten_resolved_mcp_server_ids(identifier_to_server_ids) allowed_servers = all_allowed_servers if teamless_admin_assignment: diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index 2272a37488f..1353b9ed651 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -38,13 +38,8 @@ class TeamMemberPermissionChecks: - If team has no permissions set (None), fall back to DEFAULT_TEAM_MEMBER_PERMISSIONS. """ - if team_table.team_member_permissions is not None and isinstance( - team_table.team_member_permissions, list - ): - permissions = { - KeyManagementRoutes(permission) - for permission in team_table.team_member_permissions - } + if team_table.team_member_permissions is not None and isinstance(team_table.team_member_permissions, list): + permissions = {KeyManagementRoutes(permission) for permission in team_table.team_member_permissions} # Always include baseline permissions permissions.update(BASELINE_TEAM_MEMBER_PERMISSIONS) return list(permissions) @@ -93,17 +88,13 @@ class TeamMemberPermissionChecks: ) # 4. Extract `Member` object from `team_table` - key_assigned_user_in_team = _get_user_in_team( - team_table=team_table, user_id=user_api_key_dict.user_id - ) + key_assigned_user_in_team = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id) # 5. Check if the team member has permissions for the endpoint - has_permission = ( - TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint( - team_member_object=key_assigned_user_in_team, - team_table=team_table, - route=route, - ) + has_permission = TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint( + team_member_object=key_assigned_user_in_team, + team_table=team_table, + route=route, ) if not has_permission: raise ProxyException( @@ -130,21 +121,13 @@ class TeamMemberPermissionChecks: if team_member_object.role == "admin": return True - _team_member_permissions = ( - TeamMemberPermissionChecks.get_permissions_for_team_member( - team_member_object=team_member_object, - team_table=team_table, - ) - ) - team_member_permissions = ( - TeamMemberPermissionChecks._get_list_of_route_enum_as_str( - _team_member_permissions - ) + _team_member_permissions = TeamMemberPermissionChecks.get_permissions_for_team_member( + team_member_object=team_member_object, + team_table=team_table, ) + team_member_permissions = TeamMemberPermissionChecks._get_list_of_route_enum_as_str(_team_member_permissions) - if not RouteChecks.check_route_access( - route=route, allowed_routes=team_member_permissions - ): + if not RouteChecks.check_route_access(route=route, allowed_routes=team_member_permissions): raise ProxyException( message=f"Team member does not have permissions for endpoint: {route}. You only have access to the following endpoints: {team_member_permissions} for team {team_table.team_id}. To create keys for this team, please ask your proxy admin to check the team member permission settings and update the settings to allow team member users to create keys.", type=ProxyErrorTypes.team_member_permission_error, @@ -188,9 +171,7 @@ class TeamMemberPermissionChecks: if team_table is None: return - team_member_object = _get_user_in_team( - team_table=team_table, user_id=user_api_key_dict.user_id - ) + team_member_object = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id) # Team admins always bypass (consistent with other member-permission checks). if team_member_object is not None and team_member_object.role == "admin": @@ -242,9 +223,7 @@ class TeamMemberPermissionChecks: ) # 4. Extract `Member` object from `team_table` - team_member_object = _get_user_in_team( - team_table=team_table, user_id=user_api_key_dict.user_id - ) + team_member_object = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id) return team_member_object is not None @staticmethod diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 0c5439e26e2..11a99caebf5 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -38,17 +38,13 @@ from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.repositories.user_repository import UserRepository -def get_new_internal_user_defaults( - user_id: str, user_email: Optional[str] = None -) -> dict: +def get_new_internal_user_defaults(user_id: str, user_email: Optional[str] = None) -> dict: user_info = litellm.default_internal_user_params or {} returned_dict: SSOUserDefinedValues = { "models": user_info.get("models") or [], "max_budget": user_info.get("max_budget", litellm.max_internal_user_budget), - "budget_duration": user_info.get( - "budget_duration", litellm.internal_user_budget_duration - ), + "budget_duration": user_info.get("budget_duration", litellm.internal_user_budget_duration), "user_email": user_email or user_info.get("user_email", None), "user_id": user_id, "user_role": "internal_user", @@ -94,9 +90,7 @@ async def handle_budget_for_entity( budget_params = LiteLLM_BudgetTable.model_fields.keys() # Extract budget fields from data - _json_data = ( - data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data - ) + _json_data = data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data _budget_data = {k: v for k, v in _json_data.items() if k in budget_params} # Check if budget_id is explicitly provided in the data @@ -110,9 +104,7 @@ async def handle_budget_for_entity( elif _budget_data: # Create a new budget with the provided fields budget_row = LiteLLM_BudgetTable(**_budget_data) - new_budget_data = prisma_client.jsonify_object( - budget_row.model_dump(exclude_none=True) - ) + new_budget_data = prisma_client.jsonify_object(budget_row.model_dump(exclude_none=True)) _budget = await BudgetRepository(prisma_client).table.create( data={ @@ -132,9 +124,7 @@ async def handle_budget_for_entity( # If budget fields are provided, update the existing budget if _budget_data: await update_budget( - budget_obj=BudgetNewRequest( - budget_id=existing_budget_id, **_budget_data - ), + budget_obj=BudgetNewRequest(budget_id=existing_budget_id, **_budget_data), user_api_key_dict=user_api_key_dict, ) @@ -209,9 +199,7 @@ async def _clone_team_default_budget_for_member( # Start the member's budget window at clone time, not the pool's reset # timestamp — otherwise a member joining mid-cycle inherits a stale reset. if cloned_data.get("budget_duration"): - cloned_data["budget_reset_at"] = get_budget_reset_time( - cloned_data["budget_duration"] - ) + cloned_data["budget_reset_at"] = get_budget_reset_time(cloned_data["budget_duration"]) new_budget = await BudgetRepository(prisma_client).table.create(data=cloned_data) return new_budget.budget_id @@ -259,9 +247,7 @@ async def _resolve_member_budget_id( budget_data["allowed_models"] = allowed_models if budget_duration is not None: budget_data["budget_duration"] = budget_duration - budget_data["budget_reset_at"] = get_budget_reset_time( - budget_duration=budget_duration - ) + budget_data["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration) response = await BudgetRepository(prisma_client).table.create(data=budget_data) return response.budget_id @@ -300,9 +286,7 @@ async def add_new_member( if _returned_user is not None: returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) elif new_member.user_email is not None: - new_user_defaults = get_new_internal_user_defaults( - user_id=str(uuid.uuid4()), user_email=new_member.user_email - ) + new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email) ## user email is not unique acc. to prisma schema -> future improvement ### for now: check if it exists in db, if not - insert it existing_user_row: Optional[list] = await prisma_client.get_data( @@ -310,13 +294,9 @@ async def add_new_member( table_name="user", query_type="find_all", ) - if existing_user_row is None or ( - isinstance(existing_user_row, list) and len(existing_user_row) == 0 - ): + if existing_user_row is None or (isinstance(existing_user_row, list) and len(existing_user_row) == 0): new_user_defaults["teams"] = [team_id] - _returned_user = await prisma_client.insert_data( - data=new_user_defaults, table_name="user" - ) # type: ignore + _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore if _returned_user is not None: returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) @@ -331,9 +311,7 @@ async def add_new_member( elif len(existing_user_row) > 1: raise HTTPException( status_code=400, - detail={ - "error": "Multiple users with this email found in db. Please use 'user_id' instead." - }, + detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."}, ) _budget_id = await _resolve_member_budget_id( @@ -347,9 +325,7 @@ async def add_new_member( ) if _budget_id and returned_user is not None and returned_user.user_id is not None: - _returned_team_membership = await TeamMembershipRepository( - prisma_client - ).table.create( + _returned_team_membership = await TeamMembershipRepository(prisma_client).table.create( data={ "team_id": team_id, "user_id": returned_user.user_id, @@ -358,9 +334,7 @@ async def add_new_member( include={"litellm_budget_table": True}, ) - returned_team_membership = LiteLLM_TeamMembership( - **_returned_team_membership.model_dump() - ) + returned_team_membership = LiteLLM_TeamMembership(**_returned_team_membership.model_dump()) if returned_user is None: raise Exception("Unable to update user table with membership information!") @@ -457,10 +431,7 @@ async def send_management_endpoint_alert( } # Check if alerting is enabled - if ( - proxy_logging_obj is not None - and proxy_logging_obj.slack_alerting_instance is not None - ): + if proxy_logging_obj is not None and proxy_logging_obj.slack_alerting_instance is not None: # Virtual Key Events if function_name in management_function_to_event_name: _event_name: AlertType = management_function_to_event_name[function_name] @@ -474,12 +445,10 @@ async def send_management_endpoint_alert( # replace all "_" with " " and capitalize event_name = _event_name.replace("_", " ").title() - await ( - proxy_logging_obj.slack_alerting_instance.send_virtual_key_event_slack( - key_event=key_event, - event_name=event_name, - alert_type=_event_name, - ) + await proxy_logging_obj.slack_alerting_instance.send_virtual_key_event_slack( + key_event=key_event, + event_name=event_name, + alert_type=_event_name, ) @@ -500,11 +469,7 @@ def _redact_record_env_vars(record: Any) -> Any: object that is also returned to the caller. Records without an ``env_vars`` list are returned unchanged. """ - env_vars = ( - record.get("env_vars") - if isinstance(record, dict) - else getattr(record, "env_vars", None) - ) + env_vars = record.get("env_vars") if isinstance(record, dict) else getattr(record, "env_vars", None) if not isinstance(env_vars, list): return record redacted = [_redacted_env_var(entry) for entry in env_vars] @@ -526,9 +491,7 @@ def _redact_env_var_values(response: dict) -> None: scrubbed. Names, scopes, and descriptions are kept so traces stay useful. """ if isinstance(response.get("env_vars"), list): - response["env_vars"] = [ - _redacted_env_var(entry) for entry in response["env_vars"] - ] + response["env_vars"] = [_redacted_env_var(entry) for entry in response["env_vars"]] items = response.get("items") if isinstance(items, list): @@ -636,9 +599,7 @@ def management_endpoint_wrapper(func): result = await func(*args, **kwargs) end_time = datetime.now() try: - user_api_key_dict: UserAPIKeyAuth = ( - kwargs.get("user_api_key_dict") or UserAPIKeyAuth() - ) + user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") or UserAPIKeyAuth() await send_management_endpoint_alert( request_kwargs=kwargs, @@ -670,9 +631,7 @@ def management_endpoint_wrapper(func): except Exception as e: end_time = datetime.now() - user_api_key_dict: UserAPIKeyAuth = ( - kwargs.get("user_api_key_dict") or UserAPIKeyAuth() - ) + user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict") or UserAPIKeyAuth() parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None) if parent_otel_span is not None: try: diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 6f1ca3196fe..1d9704ba619 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -98,15 +98,11 @@ def _require_prisma(): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) return prisma_client -def _internal_error( - log_message: str, exc: Exception, default_detail: str -) -> HTTPException: +def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HTTPException: """ Build a 500 HTTPException with a generic, caller-safe `detail` while logging the actual exception server-side. Avoids leaking internal Prisma / @@ -116,9 +112,7 @@ def _internal_error( return HTTPException(status_code=500, detail=default_detail) -async def _assert_write_access( - prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth -) -> None: +async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth) -> None: """ Enforce ownership for mutations (PUT/DELETE). @@ -158,9 +152,7 @@ async def _assert_write_access( ) -async def _is_team_admin_for( - prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str -) -> bool: +async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool: """ True if the caller is a team admin of `team_id`, or an org admin for the team's organization. Mirrors the auth pattern used by team-management @@ -175,13 +167,9 @@ async def _is_team_admin_for( ) try: - team_obj = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + team_obj = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) except Exception as e: - verbose_proxy_logger.exception( - "Error loading team for write-auth check (team_id=%s): %s", team_id, e - ) + verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e) return False if team_obj is None: return False @@ -194,14 +182,10 @@ async def _is_team_admin_for( # initialized. In tests / non-proxy contexts that import path may fail — # treat any error as "not an org admin" rather than crashing the request. try: - if await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=team_obj - ): + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj): return True except Exception as e: - verbose_proxy_logger.debug( - "Org-admin check skipped during write-auth (team_id=%s): %s", team_id, e - ) + verbose_proxy_logger.debug("Org-admin check skipped during write-auth (team_id=%s): %s", team_id, e) return False @@ -217,9 +201,7 @@ def _is_unique_violation(exc: Exception) -> bool: if code == "P2002": return True msg = str(exc) - return ( - "P2002" in msg or "Unique" in msg or "unique" in msg or "UniqueViolation" in msg - ) + return "P2002" in msg or "Unique" in msg or "unique" in msg or "UniqueViolation" in msg def _resolve_scope( @@ -239,16 +221,8 @@ def _resolve_scope( a PROXY_ADMIN who is explicitly stamping a global/shared row. """ if _is_admin(user_api_key_dict): - user_id = ( - requested_user_id - if requested_user_id is not None - else user_api_key_dict.user_id - ) - team_id = ( - requested_team_id - if requested_team_id is not None - else user_api_key_dict.team_id - ) + user_id = requested_user_id if requested_user_id is not None else user_api_key_dict.user_id + team_id = requested_team_id if requested_team_id is not None else user_api_key_dict.team_id return user_id, team_id if requested_user_id is not None and requested_user_id != user_api_key_dict.user_id: @@ -374,27 +348,19 @@ async def list_memory( take=page_size, ) except Exception as e: - raise _internal_error( - "Error listing memory: %s", e, "Internal error listing memory entries." - ) + raise _internal_error("Error listing memory: %s", e, "Internal error listing memory entries.") return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total) -async def _find_memory_for_caller( - prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth -) -> Any: +async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth) -> Any: """Look up a memory row by key, scoped to the caller's visibility.""" key_filter: dict = {"key": key} vis = _visibility_filter(user_api_key_dict) where: dict = key_filter if vis is None else {"AND": [key_filter, vis]} - rows = await MemoryRepository(prisma_client).table.find_many( - where=where, take=1, order={"updated_at": "desc"} - ) + rows = await MemoryRepository(prisma_client).table.find_many(where=where, take=1, order={"updated_at": "desc"}) if not rows: - raise HTTPException( - status_code=404, detail=f"Memory with key '{key}' not found" - ) + raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found") return rows[0] @@ -488,9 +454,7 @@ async def upsert_memory( detail="Cannot create a new memory via PUT without a 'value'.", ) # PUT-create must honor admin scope override, matching POST semantics. - user_id, team_id = _resolve_scope( - user_api_key_dict, body.user_id, body.team_id - ) + user_id, team_id = _resolve_scope(user_api_key_dict, body.user_id, body.team_id) # Omit `metadata` when None so the column defaults to SQL NULL; # otherwise JSON-encode for Prisma — same pattern as # `create_memory` above. @@ -505,9 +469,7 @@ async def upsert_memory( if body.metadata is not None: create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) try: - row = await MemoryRepository(prisma_client).table.create( - data=create_data - ) + row = await MemoryRepository(prisma_client).table.create(data=create_data) except Exception as e: # Race: a concurrent PUT/POST created the row after our check. # Re-read and fall back to an update so the PUT stays idempotent @@ -523,9 +485,7 @@ async def upsert_memory( detail=f"Memory with key '{key}' already exists.", ) # Same write-authorization check as the non-race path. - await _assert_write_access( - prisma_client, existing_after_race, user_api_key_dict - ) + await _assert_write_access(prisma_client, existing_after_race, user_api_key_dict) row = await MemoryRepository(prisma_client).table.update( where={"memory_id": existing_after_race.memory_id}, data=data, @@ -533,9 +493,7 @@ async def upsert_memory( except HTTPException: raise except Exception as e: - raise _internal_error( - "Error upserting memory: %s", e, "Internal error updating memory entry." - ) + raise _internal_error("Error upserting memory: %s", e, "Internal error updating memory entry.") return _row_to_model(row) @@ -556,12 +514,8 @@ async def delete_memory( # Visibility != write authority — see the upsert handler for the rationale. await _assert_write_access(prisma_client, row, user_api_key_dict) try: - await MemoryRepository(prisma_client).table.delete( - where={"memory_id": row.memory_id} - ) + await MemoryRepository(prisma_client).table.delete(where={"memory_id": row.memory_id}) except Exception as e: - raise _internal_error( - "Error deleting memory: %s", e, "Internal error deleting memory entry." - ) + raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.") return MemoryDeleteResponse(key=key, deleted=True) diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index 7eb8ae83cb4..d9cf435235c 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -57,23 +57,13 @@ class PrometheusAuthMiddleware: await user_api_key_auth( request=request, api_key=request.headers.get(_AUTHORIZATION_HEADER) or "", - azure_api_key_header=request.headers.get( - SpecialHeaders.azure_authorization.value - ) - or "", - anthropic_api_key_header=request.headers.get( - SpecialHeaders.anthropic_authorization.value - ), + azure_api_key_header=request.headers.get(SpecialHeaders.azure_authorization.value) or "", + anthropic_api_key_header=request.headers.get(SpecialHeaders.anthropic_authorization.value), google_ai_studio_api_key_header=request.headers.get( SpecialHeaders.google_ai_studio_authorization.value ), - azure_apim_header=request.headers.get( - SpecialHeaders.azure_apim_authorization.value - ) - or "", - custom_litellm_key_header=request.headers.get( - SpecialHeaders.custom_litellm_api_key.value - ), + azure_apim_header=request.headers.get(SpecialHeaders.azure_apim_authorization.value) or "", + custom_litellm_key_header=request.headers.get(SpecialHeaders.custom_litellm_api_key.value), ) except Exception as e: # Send 401 response directly via ASGI protocol diff --git a/litellm/proxy/middleware/request_size_limit_middleware.py b/litellm/proxy/middleware/request_size_limit_middleware.py index 78a38e3572e..0015224279a 100644 --- a/litellm/proxy/middleware/request_size_limit_middleware.py +++ b/litellm/proxy/middleware/request_size_limit_middleware.py @@ -43,9 +43,7 @@ class RequestSizeLimitMiddleware: content_length = _get_content_length(scope=scope) if content_length is not None and content_length > max_request_size_bytes: - await _send_request_too_large( - send=send, max_request_size_mb=max_request_size_mb - ) + await _send_request_too_large(send=send, max_request_size_mb=max_request_size_mb) return received_body_bytes = 0 @@ -75,9 +73,7 @@ class RequestSizeLimitMiddleware: except RequestEntityTooLarge: if response_started: raise - await _send_request_too_large( - send=send, max_request_size_mb=max_request_size_mb - ) + await _send_request_too_large(send=send, max_request_size_mb=max_request_size_mb) def _mb_to_bytes(max_request_size_mb: Optional[Union[int, float]]) -> Optional[int]: diff --git a/litellm/proxy/middleware/security_headers_middleware.py b/litellm/proxy/middleware/security_headers_middleware.py index a090c8f027f..eebe9c9fe34 100644 --- a/litellm/proxy/middleware/security_headers_middleware.py +++ b/litellm/proxy/middleware/security_headers_middleware.py @@ -41,11 +41,7 @@ class SecurityHeadersMiddleware: async def send_with_security_headers(message: Message) -> None: if message["type"] == "http.response.start": headers = MutableHeaders(scope=message) - applied = ( - (*STATIC_SECURITY_HEADERS, HSTS_HEADER) - if _hsts_enabled() - else STATIC_SECURITY_HEADERS - ) + applied = (*STATIC_SECURITY_HEADERS, HSTS_HEADER) if _hsts_enabled() else STATIC_SECURITY_HEADERS for name, value in applied: headers.setdefault(name, value) await send(message) diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index e32fee6afc5..91699ad829a 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -64,12 +64,8 @@ async def _parse_multipart_form(request: Request) -> Dict[str, Any]: # request.form() may return either a FastAPI or Starlette UploadFile # depending on middleware; check both via isinstance (FastAPI's UploadFile # is a subclass of Starlette's) and fall back to duck-type check. - if uploaded_file is None or ( - not isinstance(uploaded_file, UploadFile) and not hasattr(uploaded_file, "read") - ): - raise ValueError( - "Multipart OCR request must include a 'file' field with the document to process" - ) + if uploaded_file is None or (not isinstance(uploaded_file, UploadFile) and not hasattr(uploaded_file, "read")): + raise ValueError("Multipart OCR request must include a 'file' field with the document to process") uploaded_file = cast(UploadFile, uploaded_file) @@ -144,8 +140,7 @@ async def _parse_ocr_request(request: Request) -> Dict[str, Any]: # Check if form data is available. if getattr(request, "_form", None) is not None: verbose_proxy_logger.debug( - "OCR request body is empty but form data is available from middleware — " - "processing as multipart form." + "OCR request body is empty but form data is available from middleware — processing as multipart form." ) return await _parse_multipart_form(request) diff --git a/litellm/proxy/openai_evals_endpoints/endpoints.py b/litellm/proxy/openai_evals_endpoints/endpoints.py index 98d409adf5f..565ba607a43 100644 --- a/litellm/proxy/openai_evals_endpoints/endpoints.py +++ b/litellm/proxy/openai_evals_endpoints/endpoints.py @@ -78,11 +78,7 @@ async def create_eval( # Extract model for routing (header > query > body) # When using extra_body={"model": "..."}, the OpenAI SDK merges it into the body - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -183,11 +179,7 @@ async def list_evals( data["order_by"] = order_by # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -275,11 +267,7 @@ async def get_eval( data["eval_id"] = eval_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -369,11 +357,7 @@ async def update_eval( data["eval_id"] = eval_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -461,11 +445,7 @@ async def delete_eval( data["eval_id"] = eval_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -553,11 +533,7 @@ async def cancel_eval( data["eval_id"] = eval_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -660,11 +636,7 @@ async def create_run( request.headers.get("x-litellm-model") or request.query_params.get("model") or data.get("model") - or ( - data.get("completion", {}).get("model") - if isinstance(data.get("completion"), dict) - else None - ) + or (data.get("completion", {}).get("model") if isinstance(data.get("completion"), dict) else None) ) if model: data["model"] = model @@ -934,11 +906,7 @@ async def cancel_run( data["run_id"] = run_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model @@ -1027,11 +995,7 @@ async def delete_run( data["run_id"] = run_id # Extract model for routing (header > query > body) - model = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") if model: data["model"] = model diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index bb3033e2a6c..efd1d6b6cee 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -91,9 +91,7 @@ def get_batch_id_from_unified_batch_id(file_id: str) -> str: return re.split(r"[;,]", batch_id, maxsplit=1)[0] -def encode_file_id_with_model( - file_id: str, model: str, id_type: Literal["file", "batch"] = "file" -) -> str: +def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", "batch"] = "file") -> str: """ Encode a file/batch ID with model routing information. @@ -141,9 +139,7 @@ def encode_batch_response_ids(response, model: str) -> None: """Encode all IDs in a batch response with model routing info (in-place).""" if not response or not hasattr(response, "id") or not response.id: return - response.id = encode_file_id_with_model( - file_id=response.id, model=model, id_type="batch" - ) + response.id = encode_file_id_with_model(file_id=response.id, model=model, id_type="batch") for attr in ("output_file_id", "error_file_id", "input_file_id"): if hasattr(response, attr) and getattr(response, attr): setattr( @@ -253,11 +249,7 @@ def extract_model_from_sources( model_from_id = decode_model_from_file_id(file_id) # Check other sources for model parameter - model_from_param = ( - data.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + model_from_param = data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") return model_from_id, model_from_param @@ -294,9 +286,7 @@ def get_credentials_for_model( if credentials is None: raise HTTPException( status_code=400, - detail={ - "error": f"Model '{model_id}' not found in model_list. Please check your config.yaml." - }, + detail={"error": f"Model '{model_id}' not found in model_list. Please check your config.yaml."}, ) return credentials @@ -328,13 +318,8 @@ def get_team_provider_credentials( return None def _provider_credentials(model_id: str) -> Optional[dict]: - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=model_id - ) - if ( - credentials is not None - and credentials.get("custom_llm_provider") == custom_llm_provider - ): + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider: return credentials return None @@ -406,9 +391,7 @@ def prepare_data_with_credentials( """ data.update(credentials) if include_internal_credentials: - data["_litellm_internal_model_credentials"] = MappingProxyType( - dict(credentials) - ) + data["_litellm_internal_model_credentials"] = MappingProxyType(dict(credentials)) data.pop("custom_llm_provider", None) if file_id is not None: @@ -555,9 +538,7 @@ def detect_content_type_from_filename(filename: str) -> str: return "application/octet-stream" -def normalize_mime_type_for_provider( - mime_type: str, provider: Optional[str] = None -) -> str: +def normalize_mime_type_for_provider(mime_type: str, provider: Optional[str] = None) -> str: """ Normalize MIME type for specific provider requirements. @@ -676,9 +657,7 @@ class FileCreationParams: self.target_storage = "default" # Strip whitespace from model names - self.target_model_names = [ - name.strip() for name in self.target_model_names if name.strip() - ] + self.target_model_names = [name.strip() for name in self.target_model_names if name.strip()] async def extract_file_creation_params( @@ -748,9 +727,7 @@ def _extract_target_model_names_simple( # Parse comma-separated string into list if isinstance(target_model_names_form, str): - return [ - name.strip() for name in target_model_names_form.split(",") if name.strip() - ] + return [name.strip() for name in target_model_names_form.split(",") if name.strip()] elif isinstance(target_model_names_form, list): return [str(name).strip() for name in target_model_names_form if name] @@ -758,9 +735,7 @@ def _extract_target_model_names_simple( def _is_target_model_names_key(key: str) -> bool: - return key == "target_model_names" or ( - key.startswith("target_model_names[") and key.endswith("]") - ) + return key == "target_model_names" or (key.startswith("target_model_names[") and key.endswith("]")) async def _extract_target_model_names_from_form(request: "Request") -> List[str]: @@ -837,11 +812,7 @@ def _extract_model_param(request: "Request", request_body: dict) -> Optional[str 2. Query parameter (?model=) 3. Header (x-litellm-model) """ - return ( - request_body.get("model") - or request.query_params.get("model") - or request.headers.get("x-litellm-model") - ) + return request_body.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") # ============================================================================ @@ -916,9 +887,7 @@ async def ensure_batch_response_managed_file_ids( model_name = hidden_params.get("model_name") unified_file_id = hidden_params.get("unified_file_id") if not model_name and isinstance(unified_file_id, str): - decoded_unified_file_id = ( - _is_base64_encoded_unified_file_id(unified_file_id) or unified_file_id - ) + decoded_unified_file_id = _is_base64_encoded_unified_file_id(unified_file_id) or unified_file_id target_model_names = get_models_from_unified_file_id(decoded_unified_file_id) if target_model_names: model_name = ",".join(target_model_names) @@ -946,20 +915,15 @@ async def ensure_batch_response_managed_file_ids( await managed_files_obj.store_unified_file_id( file_id=new_unified_file_id, file_object=None, - litellm_parent_otel_span=getattr( - user_api_key_dict, "parent_otel_span", None - ), + litellm_parent_otel_span=getattr(user_api_key_dict, "parent_otel_span", None), model_mappings={model_id: raw_file_id}, user_api_key_dict=user_api_key_dict, ) setattr(response, file_attr, new_unified_file_id) - verbose_proxy_logger.debug( - f"Converted batch {file_attr} {raw_file_id!r} to managed ID before DB write" - ) + verbose_proxy_logger.debug(f"Converted batch {file_attr} {raw_file_id!r} to managed ID before DB write") except Exception as e: verbose_proxy_logger.warning( - f"Failed to convert batch {file_attr}={raw_file_id!r} to managed ID " - f"before DB write: {e}" + f"Failed to convert batch {file_attr}={raw_file_id!r} to managed ID before DB write: {e}" ) @@ -1015,16 +979,12 @@ async def get_batch_from_database( # The stored batch object has the raw provider input_file_id. Resolve to unified ID. await resolve_input_file_id_to_unified(response, prisma_client) - verbose_proxy_logger.debug( - f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}" - ) + verbose_proxy_logger.debug(f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}") return db_batch_object, response except Exception as e: - verbose_proxy_logger.warning( - f"Failed to retrieve batch from ManagedObjectTable: {e}, falling back to provider" - ) + verbose_proxy_logger.warning(f"Failed to retrieve batch from ManagedObjectTable: {e}, falling back to provider") return None, None @@ -1083,9 +1043,7 @@ async def update_batch_in_database( f"Updating batch {batch_id} status from {db_batch_object.status} to {response.status}" ) else: - verbose_proxy_logger.info( - f"Updating batch {batch_id} status to {response.status} after {operation}" - ) + verbose_proxy_logger.info(f"Updating batch {batch_id} status to {response.status} after {operation}") # Normalize status for database storage db_status = response.status if response.status != "completed" else "complete" @@ -1116,13 +1074,8 @@ async def update_batch_in_database( # If the batch_processed column doesn't exist (old schema), # retry without it so the status update still succeeds. err_str = str(col_err).lower() - if ( - "batch_processed" in err_str - and update_data.get("batch_processed") is not None - ): - verbose_proxy_logger.warning( - f"batch_processed column not found, retrying update without it: {col_err}" - ) + if "batch_processed" in err_str and update_data.get("batch_processed") is not None: + verbose_proxy_logger.warning(f"batch_processed column not found, retrying update without it: {col_err}") update_data.pop("batch_processed", None) await ManagedObjectRepository(prisma_client).table.update( where={"unified_object_id": batch_id}, @@ -1131,6 +1084,4 @@ async def update_batch_in_database( else: raise except Exception as e: - verbose_proxy_logger.error( - f"Failed to update batch status in ManagedObjectTable: {e}" - ) + verbose_proxy_logger.error(f"Failed to update batch status in ManagedObjectTable: {e}") diff --git a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py index 543e345deb6..5546edba49f 100644 --- a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py +++ b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py @@ -107,9 +107,7 @@ class FileContentStreamingHandler: FileContentStreamingResult, await litellm.afile_content( **{ - "custom_llm_provider": cast( - FileContentProvider, custom_llm_provider - ), + "custom_llm_provider": cast(FileContentProvider, custom_llm_provider), "file_id": file_id, "stream": True, **data, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d8f8726b6b6..34ca43c603f 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -127,9 +127,7 @@ async def _deprecated_loadbalanced_create_file( if llm_router is None: raise HTTPException( status_code=500, - detail={ - "error": "LLM Router not initialized. Ensure models added to proxy." - }, + detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) response = await llm_router.acreate_file(model=router_model, **_create_file_request) @@ -210,9 +208,7 @@ async def route_create_file( original_id = response.id encoded_id = encode_file_id_with_model(file_id=original_id, model=model) response.id = encoded_id - verbose_proxy_logger.debug( - f"Encoded file ID: {original_id} -> {encoded_id} (model: {model})" - ) + verbose_proxy_logger.debug(f"Encoded file ID: {original_id} -> {encoded_id} (model: {model})") return response @@ -250,11 +246,7 @@ async def route_create_file( user_api_key_dict=user_api_key_dict, ) # EXISTING: Deprecated loadbalancing approach (for backwards compatibility when not using managed files) - elif ( - litellm.enable_loadbalancing_on_batch_endpoints is True - and is_router_model - and router_model is not None - ): + elif litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model and router_model is not None: response = await _deprecated_loadbalanced_create_file( llm_router=llm_router, router_model=router_model, @@ -262,17 +254,13 @@ async def route_create_file( ) else: # get configs for custom_llm_provider - llm_provider_config = get_files_provider_config( - custom_llm_provider=custom_llm_provider - ) + llm_provider_config = get_files_provider_config(custom_llm_provider=custom_llm_provider) if llm_provider_config is not None: # add llm_provider_config to data _create_file_request.update(llm_provider_config) _create_file_request.pop("custom_llm_provider", None) # type: ignore # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch - response = await litellm.acreate_file( - **_create_file_request, custom_llm_provider=custom_llm_provider - ) # type: ignore + response = await litellm.acreate_file(**_create_file_request, custom_llm_provider=custom_llm_provider) # type: ignore return response @@ -361,9 +349,7 @@ async def create_file( target_model_names_list = file_params.target_model_names model_param = file_params.model - validate_managed_files_requirement( - target_model_names=target_model_names_list, model=model_param - ) + validate_managed_files_requirement(target_model_names=target_model_names_list, model=model_param) # Prepare the data for forwarding @@ -385,10 +371,8 @@ async def create_file( expires_after: Optional[FileExpiresAfter] = None form_data_raw = await request.form() form_data_dict: Dict[str, Any] = dict(form_data_raw) - extracted_litellm_metadata: Optional[Dict[str, Any]] = ( - extract_nested_form_metadata( - form_data=form_data_dict, prefix="litellm_metadata[" - ) + extracted_litellm_metadata: Optional[Dict[str, Any]] = extract_nested_form_metadata( + form_data=form_data_dict, prefix="litellm_metadata[" ) expires_after_anchor = form_data_raw.get("expires_after[anchor]") expires_after_seconds_str = form_data_raw.get("expires_after[seconds]") @@ -473,18 +457,13 @@ async def create_file( json_obj = get_first_json_object(file_source) if json_obj: router_model = get_model_from_json_obj(json_object=json_obj) - is_router_model = is_known_model( - model=router_model, llm_router=llm_router - ) + is_router_model = is_known_model(model=router_model, llm_router=llm_router) # Apply team-level file expiry enforcement team_metadata = user_api_key_dict.team_metadata or {} enforced_file_expiry = team_metadata.get("enforced_file_expires_after") if enforced_file_expiry is not None: - if ( - "anchor" not in enforced_file_expiry - or "seconds" not in enforced_file_expiry - ): + if "anchor" not in enforced_file_expiry or "seconds" not in enforced_file_expiry: raise HTTPException( status_code=500, detail={ @@ -533,9 +512,7 @@ async def create_file( ) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ## POST CALL HOOKS ### @@ -567,9 +544,7 @@ async def create_file( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.create_file(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.create_file(): Exception occured - {}".format(str(e)) ) if isinstance(e, HTTPException): raise ProxyException( @@ -683,9 +658,7 @@ async def get_file_content( ) # Check if file is stored in a storage backend (check DB) - if hasattr(managed_files_obj, "prisma_client") and getattr( - managed_files_obj, "prisma_client", None - ): + if hasattr(managed_files_obj, "prisma_client") and getattr(managed_files_obj, "prisma_client", None): prisma_client = getattr(managed_files_obj, "prisma_client") db_file = await ManagedFileRepository(prisma_client).table.find_first( where={"unified_file_id": file_id} @@ -813,11 +786,7 @@ async def get_file_content( verbose_proxy_logger.debug( f"Retrieved file content using model: {model_used}" - + ( - f", file_id: {file_id} -> {original_file_id}" - if original_file_id - else "" - ) + + (f", file_id: {file_id} -> {original_file_id}" if original_file_id else "") ) else: # Fallback to default behavior (uses env variables or provider-based routing) @@ -831,9 +800,7 @@ async def get_file_content( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -854,9 +821,7 @@ async def get_file_content( ) httpx_response: Optional[httpx.Response] = getattr(response, "response", None) if httpx_response is None: - raise ValueError( - f"Invalid response - response.response is None - got {response}" - ) + raise ValueError(f"Invalid response - response.response is None - got {response}") return Response( content=httpx_response.content, @@ -869,9 +834,7 @@ async def get_file_content( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.retrieve_file_content(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.retrieve_file_content(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): @@ -986,12 +949,7 @@ async def get_file( response = await litellm.afile_retrieve(**data) # type: ignore # Keep the encoded ID in response if it was originally encoded - if ( - original_file_id - and response - and hasattr(response, "id") - and response.id - ): + if original_file_id and response and hasattr(response, "id") and response.id: response.id = file_id verbose_proxy_logger.debug( @@ -1033,9 +991,7 @@ async def get_file( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -1060,11 +1016,7 @@ async def get_file( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.retrieve_file(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_file(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -1239,9 +1191,7 @@ async def delete_file( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -1267,9 +1217,7 @@ async def delete_file( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.delete_file(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.delete_file(): Exception occured - {}".format(str(e)) ) if isinstance(e, HTTPException): raise ProxyException( @@ -1436,9 +1384,7 @@ async def list_files( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -1463,11 +1409,7 @@ async def list_files( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.list_files(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.list_files(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index 9adeeb995ac..5125931b28e 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -81,19 +81,15 @@ class StorageBackendFileService: file_naming_strategy="uuid", ) - verbose_proxy_logger.debug( - f"Storage backend upload complete: backend={target_storage}, url={storage_url}" - ) + verbose_proxy_logger.debug(f"Storage backend upload complete: backend={target_storage}, url={storage_url}") # Create file object with storage metadata - file_object = ( - StorageBackendFileService._create_file_object_with_storage_metadata( - file_content=file_content, - filename=filename, - purpose=purpose, - target_storage=target_storage, - storage_url=storage_url, - ) + file_object = StorageBackendFileService._create_file_object_with_storage_metadata( + file_content=file_content, + filename=filename, + purpose=purpose, + target_storage=target_storage, + storage_url=storage_url, ) # Store in managed files if target_model_names provided @@ -143,10 +139,7 @@ class StorageBackendFileService: ) # Store storage metadata in hidden params - if ( - not hasattr(file_object, "_hidden_params") - or file_object._hidden_params is None - ): + if not hasattr(file_object, "_hidden_params") or file_object._hidden_params is None: file_object._hidden_params = {} file_object._hidden_params.update( { @@ -174,19 +167,15 @@ class StorageBackendFileService: Returns: str: Base64-encoded unified file ID """ - unified_file_id_str = ( - SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - file_type, - str(uuid_module.uuid4()), - ",".join(target_model_names), - file_id, - None, - ) + unified_file_id_str = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + file_type, + str(uuid_module.uuid4()), + ",".join(target_model_names), + file_id, + None, ) - base64_unified_file_id = ( - base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=") - ) + base64_unified_file_id = base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=") return base64_unified_file_id @@ -213,12 +202,8 @@ class StorageBackendFileService: user_api_key_dict: User API key authentication data """ managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") - if not managed_files_obj or not isinstance( - managed_files_obj, BaseFileEndpoints - ): - verbose_proxy_logger.warning( - "Managed files hook not available, skipping managed files storage" - ) + if not managed_files_obj or not isinstance(managed_files_obj, BaseFileEndpoints): + verbose_proxy_logger.warning("Managed files hook not available, skipping managed files storage") return managed_files_obj = cast(Any, managed_files_obj) diff --git a/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py b/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py index 5fab1d504be..6456bc71594 100644 --- a/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py +++ b/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py @@ -38,9 +38,7 @@ class JsonPathExtractor: else: extracted_values.append(str(value)) except Exception as e: - verbose_proxy_logger.debug( - "Failed to extract field %s: %s", expr, str(e) - ) + verbose_proxy_logger.debug("Failed to extract field %s: %s", expr, str(e)) return "\n".join(extracted_values) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 5e688702f13..7e573de261b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -75,9 +75,7 @@ def create_request_copy(request: Request): } -def is_passthrough_request_using_router_model( - request_body: dict, llm_router: Optional[litellm.Router] -) -> bool: +def is_passthrough_request_using_router_model(request_body: dict, llm_router: Optional[litellm.Router]) -> bool: """ Returns True if the model is in the llm_router model names """ @@ -113,16 +111,12 @@ async def llm_passthrough_factory_proxy_route( model=None, ) if provider_config is None: - raise HTTPException( - status_code=404, detail=f"Provider {custom_llm_provider} not found" - ) + raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} not found") base_target_url = provider_config.get_api_base() if base_target_url is None: - raise HTTPException( - status_code=404, detail=f"Provider {custom_llm_provider} api base not found" - ) + raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} api base not found") encoded_endpoint = httpx.URL(endpoint).path @@ -134,9 +128,7 @@ async def llm_passthrough_factory_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) # Add or update query parameters @@ -197,17 +189,11 @@ async def gemini_proxy_route( [Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio) """ ## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY - google_ai_studio_api_key = request.query_params.get("key") or request.headers.get( - "x-goog-api-key" - ) + google_ai_studio_api_key = request.query_params.get("key") or request.headers.get("x-goog-api-key") - user_api_key_dict = await user_api_key_auth( - request=request, api_key=f"Bearer {google_ai_studio_api_key}" - ) + user_api_key_dict = await user_api_key_auth(request=request, api_key=f"Bearer {google_ai_studio_api_key}") - base_target_url = ( - os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" - ) + base_target_url = os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction @@ -218,9 +204,7 @@ async def gemini_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) # Add or update query parameters @@ -283,9 +267,7 @@ async def cohere_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) # Add or update query parameters @@ -332,9 +314,7 @@ async def vllm_proxy_route( from litellm.proxy.proxy_server import llm_router request_body = await get_request_body(request) - is_router_model = is_passthrough_request_using_router_model( - request_body, llm_router - ) + is_router_model = is_passthrough_request_using_router_model(request_body, llm_router) is_streaming_request = is_passthrough_request_streaming(request_body) if is_router_model and llm_router: result = cast( @@ -349,11 +329,7 @@ async def vllm_proxy_route( content=None, data=None, files=None, - json=( - request_body - if request.headers.get("content-type") == "application/json" - else None - ), + json=(request_body if request.headers.get("content-type") == "application/json" else None), params=None, headers=None, cookies=None, @@ -414,9 +390,7 @@ async def mistral_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) # Add or update query parameters @@ -459,9 +433,7 @@ async def milvus_proxy_route( Enable using Milvus `/vectors` endpoint as a pass-through endpoint. """ - provider_config = ProviderConfigManager.get_provider_vector_stores_config( - provider=LlmProviders.MILVUS - ) + provider_config = ProviderConfigManager.get_provider_vector_stores_config(provider=LlmProviders.MILVUS) if not provider_config: raise HTTPException( status_code=500, @@ -507,11 +479,7 @@ async def milvus_proxy_route( # get the vector store name from index registry index_object = ( - ( - litellm.vector_store_index_registry.get_vector_store_index_by_name( - vector_store_index_name=collection_name - ) - ) + (litellm.vector_store_index_registry.get_vector_store_index_by_name(vector_store_index_name=collection_name)) if litellm.vector_store_index_registry is not None else None ) @@ -536,9 +504,7 @@ async def milvus_proxy_route( user_api_key_dict=user_api_key_dict, ) litellm_params = vector_store.get("litellm_params") or {} - auth_credentials = provider_config.get_auth_credentials( - litellm_params=litellm_params - ) + auth_credentials = provider_config.get_auth_credentials(litellm_params=litellm_params) extra_headers = auth_credentials.get("headers") or {} @@ -549,9 +515,7 @@ async def milvus_proxy_route( ) if base_target_url is None: - raise Exception( - f"api_base not found in vector store configuration for {vector_store_name}" - ) + raise Exception(f"api_base not found in vector store configuration for {vector_store_name}") encoded_endpoint = httpx.URL(endpoint).path @@ -563,9 +527,7 @@ async def milvus_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) ## CREATE PASS-THROUGH endpoint_func = create_pass_through_route( @@ -608,11 +570,7 @@ async def anthropic_proxy_route( """ [Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion) """ - base_target_url = ( - os.getenv("ANTHROPIC_API_BASE") - or os.getenv("ANTHROPIC_BASE_URL") - or "https://api.anthropic.com" - ) + base_target_url = os.getenv("ANTHROPIC_API_BASE") or os.getenv("ANTHROPIC_BASE_URL") or "https://api.anthropic.com" encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction @@ -623,9 +581,7 @@ async def anthropic_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) # Add or update query parameters @@ -861,9 +817,7 @@ async def handle_bedrock_count_tokens( # Extract model from request body model = request_body.get("model") if not model: - raise HTTPException( - status_code=400, detail={"error": "Model is required in request body"} - ) + raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) # Get model parameters from router litellm_params = {"user_api_key_dict": user_api_key_dict} @@ -899,18 +853,14 @@ async def handle_bedrock_count_tokens( except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI - verbose_proxy_logger.error( - f"BedrockError in handle_bedrock_count_tokens: {str(e)}" - ) + verbose_proxy_logger.error(f"BedrockError in handle_bedrock_count_tokens: {str(e)}") raise HTTPException(status_code=e.status_code, detail={"error": e.message}) except HTTPException: # Re-raise HTTP exceptions as-is raise except Exception as e: verbose_proxy_logger.error(f"Error in handle_bedrock_count_tokens: {str(e)}") - raise HTTPException( - status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"} - ) + raise HTTPException(status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"}) async def bedrock_llm_proxy_route( @@ -968,9 +918,7 @@ async def bedrock_llm_proxy_route( ) # Check if this is a router model (from config.yaml) - is_router_model = is_passthrough_request_using_router_model( - request_body={"model": model}, llm_router=llm_router - ) + is_router_model = is_passthrough_request_using_router_model(request_body={"model": model}, llm_router=llm_router) # If router model, use dedicated router passthrough handler # This uses the same common processing path as non-router models @@ -995,9 +943,7 @@ async def bedrock_llm_proxy_route( ) # Fall back to existing implementation for direct Bedrock models - verbose_proxy_logger.debug( - f"Bedrock passthrough: Using direct Bedrock model '{model}' for endpoint '{endpoint}'" - ) + verbose_proxy_logger.debug(f"Bedrock passthrough: Using direct Bedrock model '{model}' for endpoint '{endpoint}'") data: Dict[str, Any] = {} base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) @@ -1062,9 +1008,7 @@ async def bedrock_proxy_route( aws_region_name = litellm.utils.get_secret(secret_name="AWS_REGION_NAME") if _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents - base_target_url = ( - f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" - ) + base_target_url = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" else: return await bedrock_llm_proxy_route( endpoint=endpoint, @@ -1082,9 +1026,7 @@ async def bedrock_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) # Add or update query parameters @@ -1099,9 +1041,7 @@ async def bedrock_proxy_route( data = await request.json() except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request = AWSRequest( - method="POST", url=str(updated_url), data=json.dumps(data), headers=headers - ) + _request = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) sigv4.add_auth(_request) prepped = _request.prepare() @@ -1158,9 +1098,7 @@ def _resolve_vertex_model_from_router( return encoded_endpoint, endpoint, vertex_project, vertex_location try: - deployment = llm_router.get_available_deployment_for_pass_through( - model=model_id - ) + deployment = llm_router.get_available_deployment_for_pass_through(model=model_id) if not deployment: return encoded_endpoint, endpoint, vertex_project, vertex_location @@ -1181,9 +1119,7 @@ def _resolve_vertex_model_from_router( # get_llm_provider returns (model, custom_llm_provider, dynamic_api_key, api_base) # For "vertex_ai/gemini-2.0-flash-exp" it returns: # model="gemini-2.0-flash-exp", custom_llm_provider="vertex_ai" - actual_model, custom_llm_provider, _, _ = get_llm_provider( - model=model_from_config - ) + actual_model, custom_llm_provider, _, _ = get_llm_provider(model=model_from_config) # Log only non-sensitive information (model names and provider), never API keys or secrets. safe_actual_model = actual_model @@ -1208,9 +1144,7 @@ def _resolve_vertex_model_from_router( endpoint = endpoint.replace(model_id, actual_model) except Exception as e: - verbose_proxy_logger.debug( - f"Error resolving vertex model from router for model {model_id}: {e}" - ) + verbose_proxy_logger.debug(f"Error resolving vertex model from router for model {model_id}: {e}") return encoded_endpoint, endpoint, vertex_project, vertex_location @@ -1249,14 +1183,8 @@ async def assemblyai_proxy_route( [Docs](https://api.assemblyai.com) """ # Set base URL based on the route - assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url( - url=str(request.url) - ) - base_target_url = ( - AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region( - region=assembly_region - ) - ) + assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(url=str(request.url)) + base_target_url = AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region(region=assembly_region) encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction if not encoded_endpoint.startswith("/"): @@ -1266,9 +1194,7 @@ async def assemblyai_proxy_route( # prefix that the operator configured on base_target_url. base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) # Add or update query parameters @@ -1338,11 +1264,7 @@ async def azure_proxy_route( ) # check if vector store index is_vector_store_index = ( - ( - litellm.vector_store_index_registry.is_vector_store_index( - vector_store_index_name=part - ) - ) + (litellm.vector_store_index_registry.is_vector_store_index(vector_store_index_name=part)) if litellm.vector_store_index_registry is not None else False ) @@ -1360,11 +1282,7 @@ async def azure_proxy_route( content=None, data=None, files=None, - json=( - request_body - if request.headers.get("content-type") == "application/json" - else None - ), + json=(request_body if request.headers.get("content-type") == "application/json" else None), params=None, headers=None, cookies=None, @@ -1406,10 +1324,8 @@ async def azure_proxy_route( ) elif is_vector_store_index: # get the api key from the provider config - provider_config = ( - ProviderConfigManager.get_provider_vector_stores_config( - provider=litellm.LlmProviders.AZURE_AI - ) + provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders.AZURE_AI ) if provider_config is None: raise Exception("Provider config not found for Azure AI") @@ -1425,11 +1341,7 @@ async def azure_proxy_route( ) # get the vector store name from index registry index_object = ( - ( - litellm.vector_store_index_registry.get_vector_store_index_by_name( - vector_store_index_name=part - ) - ) + (litellm.vector_store_index_registry.get_vector_store_index_by_name(vector_store_index_name=part)) if litellm.vector_store_index_registry is not None else None ) @@ -1448,9 +1360,7 @@ async def azure_proxy_route( user_api_key_dict=user_api_key_dict, ) litellm_params = vector_store.get("litellm_params") or {} - auth_credentials = provider_config.get_auth_credentials( - litellm_params=litellm_params - ) + auth_credentials = provider_config.get_auth_credentials(litellm_params=litellm_params) extra_headers = auth_credentials.get("headers") or {} @@ -1470,18 +1380,14 @@ async def azure_proxy_route( base_target_url = get_secret_str(secret_name="AZURE_API_BASE") if base_target_url is None: - raise Exception( - "Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure." - ) + raise Exception("Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure.") # Add or update query parameters azure_api_key = passthrough_endpoint_router.get_credentials( custom_llm_provider=litellm.LlmProviders.AZURE.value, region_name=None, ) if azure_api_key is None: - raise Exception( - "Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure." - ) + raise Exception("Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure.") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( endpoint=endpoint, @@ -1505,9 +1411,7 @@ class BaseVertexAIPassThroughHandler(ABC): @staticmethod @abstractmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: pass @@ -1517,9 +1421,7 @@ class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler): return "https://discoveryengine.googleapis.com/" @staticmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: return base_target_url @@ -1529,9 +1431,7 @@ class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler): return get_vertex_base_url(vertex_location) @staticmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: return get_vertex_base_url(vertex_location) @@ -1603,15 +1503,11 @@ def _override_vertex_params_from_router_credentials( if router_credentials is None: return vertex_project, vertex_location - verbose_proxy_logger.debug( - "Using vector store credentials to override vertex project and location" - ) + verbose_proxy_logger.debug("Using vector store credentials to override vertex project and location") litellm_params = router_credentials.get("litellm_params", {}) if not litellm_params: - verbose_proxy_logger.warning( - "Vector store credentials found but litellm_params is empty" - ) + verbose_proxy_logger.warning("Vector store credentials found but litellm_params is empty") return vertex_project, vertex_location # Extract vertex_project and vertex_location from litellm_params @@ -1626,9 +1522,7 @@ def _override_vertex_params_from_router_credentials( ) vertex_project = vector_store_project else: - verbose_proxy_logger.warning( - "Vector store credentials found but missing vertex_project in litellm_params" - ) + verbose_proxy_logger.warning("Vector store credentials found but missing vertex_project in litellm_params") if vector_store_location: verbose_proxy_logger.debug( @@ -1638,9 +1532,7 @@ def _override_vertex_params_from_router_credentials( ) vertex_location = vector_store_location else: - verbose_proxy_logger.warning( - "Vector store credentials found but missing vertex_location in litellm_params" - ) + verbose_proxy_logger.warning("Vector store credentials found but missing vertex_location in litellm_params") return vertex_project, vertex_location @@ -1678,14 +1570,10 @@ async def _prepare_vertex_auth_headers( headers_passed_through = False # Use headers from the incoming request if no vertex credentials are found - if ( - vertex_credentials is None or vertex_credentials.vertex_project is None - ) and router_credentials is None: + if (vertex_credentials is None or vertex_credentials.vertex_project is None) and router_credentials is None: headers = _safe_get_request_headers(request).copy() headers_passed_through = True - verbose_proxy_logger.debug( - "default_vertex_config not set, incoming request headers %s", headers - ) + verbose_proxy_logger.debug("default_vertex_config not set, incoming request headers %s", headers) headers.pop("content-length", None) headers.pop("host", None) else: @@ -1826,9 +1714,7 @@ async def _base_vertex_proxy_route( location=vertex_location, ) - base_target_url = get_vertex_pass_through_handler.get_default_base_target_url( - vertex_location - ) + base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) # Prepare authentication headers ( @@ -1923,21 +1809,15 @@ async def vertex_discovery_proxy_route( if vector_store_id_match: vector_store_id = vector_store_id_match.group(1) - verbose_proxy_logger.debug( - "Extracted vector store ID from endpoint: %s", vector_store_id - ) + verbose_proxy_logger.debug("Extracted vector store ID from endpoint: %s", vector_store_id) # Retrieve LiteLLM-managed vector store credentials if the datastore id # is registered with LiteLLM. Unknown datastore ids keep the existing # direct Vertex pass-through behavior. - vector_store_credentials = await get_litellm_managed_vector_store( - vector_store_id=vector_store_id - ) + vector_store_credentials = await get_litellm_managed_vector_store(vector_store_id=vector_store_id) if vector_store_credentials: - verbose_proxy_logger.debug( - "Found vector store credentials for ID: %s", vector_store_id - ) + verbose_proxy_logger.debug("Found vector store credentials for ID: %s", vector_store_id) else: verbose_proxy_logger.debug( "Vector store ID %s found in endpoint but no credentials found in registry", @@ -2033,9 +1913,7 @@ async def openai_proxy_route( region_name=None, ) if openai_api_key is None: - raise Exception( - "Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI." - ) + raise Exception("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( endpoint=endpoint, @@ -2107,17 +1985,12 @@ class BaseOpenAIPassThroughHandler: """ Appends the OpenAI-Beta header to the headers if the request is an OpenAI Assistants API request """ - if ( - RouteChecks._is_assistants_api_request(request) is True - and "OpenAI-Beta" not in headers - ): + if RouteChecks._is_assistants_api_request(request) is True and "OpenAI-Beta" not in headers: headers["OpenAI-Beta"] = "assistants=v2" return headers @staticmethod - def _assemble_headers( - api_key: Optional[str], request: Request, extra_headers: Optional[dict] = None - ) -> dict: + def _assemble_headers(api_key: Optional[str], request: Request, extra_headers: Optional[dict] = None) -> dict: base_headers = {} if api_key is not None: base_headers = { @@ -2132,31 +2005,20 @@ class BaseOpenAIPassThroughHandler: ) @staticmethod - def _join_url_paths( - base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders - ) -> str: + def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders) -> str: """ Properly joins a base URL with a path, preserving any existing path in the base URL. """ # Combine paths via the shared helper so any '..' in the path cannot # climb above the configured base path. joined_path_str = str( - base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, path - ) - ) + base_url.copy_with(path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, path)) ) # Apply OpenAI-specific path handling for both branches - if ( - custom_llm_provider == litellm.LlmProviders.OPENAI - and "/v1/" not in joined_path_str - ): + if custom_llm_provider == litellm.LlmProviders.OPENAI and "/v1/" not in joined_path_str: # Insert v1 after api.openai.com for OpenAI requests - joined_path_str = joined_path_str.replace( - "api.openai.com/", "api.openai.com/v1/" - ) + joined_path_str = joined_path_str.replace("api.openai.com/", "api.openai.com/v1/") return joined_path_str @@ -2205,10 +2067,7 @@ async def cursor_proxy_route( if cursor_api_key is None: for credential in litellm.credential_list: - if ( - credential.credential_info - and credential.credential_info.get("custom_llm_provider") == "cursor" - ): + if credential.credential_info and credential.credential_info.get("custom_llm_provider") == "cursor": cursor_api_key = credential.credential_values.get("api_key") credential_api_base = credential.credential_values.get("api_base") if credential_api_base: @@ -2228,9 +2087,7 @@ async def cursor_proxy_route( base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with( - path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path( - base_url, encoded_endpoint - ) + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) auth_value = base64.b64encode(f"{cursor_api_key}:".encode("utf-8")).decode("ascii") @@ -2307,9 +2164,7 @@ async def vertex_ai_live_websocket_passthrough( ) try: - resolved_location = resolved_location or ( - vertex_llm_base.get_default_vertex_location() - ) + resolved_location = resolved_location or (vertex_llm_base.get_default_vertex_location()) if model: resolved_location = vertex_llm_base.get_vertex_region( vertex_region=resolved_location, @@ -2325,9 +2180,7 @@ async def vertex_ai_live_websocket_passthrough( custom_llm_provider="vertex_ai_beta", ) except Exception as e: - verbose_proxy_logger.exception( - "Failed to prepare Vertex AI credentials for live passthrough" - ) + verbose_proxy_logger.exception("Failed to prepare Vertex AI credentials for live passthrough") # Log the authentication failure using proxy_logging_obj if proxy_logging_obj and user_api_key_dict: await proxy_logging_obj.post_call_failure_hook( @@ -2341,9 +2194,7 @@ async def vertex_ai_live_websocket_passthrough( host_location = resolved_location or vertex_llm_base.get_default_vertex_location() host = get_vertex_base_url(host_location).removeprefix("https://").rstrip("/") - service_url = ( - f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" - ) + service_url = f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" upstream_headers = { "Authorization": f"Bearer {access_token}", @@ -2461,9 +2312,7 @@ async def watsonx_proxy_route( ) if provider_config is None: - raise HTTPException( - status_code=404, detail="Watsonx passthrough config not found" - ) + raise HTTPException(status_code=404, detail="Watsonx passthrough config not found") # Get complete URL with version parameter complete_url, _ = provider_config.get_complete_url( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 8986166ba92..50e90699194 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -115,9 +115,7 @@ class AnthropicPassthroughLoggingHandler: def _resolve_costing_model(model: str, logging_obj: LiteLLMLoggingObj) -> str: if model and model != "unknown": return model - litellm_params = (getattr(logging_obj, "model_call_details", {}) or {}).get( - "litellm_params", {} - ) or {} + litellm_params = (getattr(logging_obj, "model_call_details", {}) or {}).get("litellm_params", {}) or {} deployment_model = litellm_params.get("model") if deployment_model and deployment_model != "unknown": return deployment_model @@ -205,13 +203,10 @@ class AnthropicPassthroughLoggingHandler: if not output_text: return try: - recovered_output_tokens = litellm.token_counter( - model=model, text=output_text, count_response_tokens=True - ) + recovered_output_tokens = litellm.token_counter(model=model, text=output_text, count_response_tokens=True) except Exception: verbose_proxy_logger.warning( - "Could not re-tokenize interrupted stream output; " - "keeping placeholder completion token count." + "Could not re-tokenize interrupted stream output; keeping placeholder completion token count." ) return if recovered_output_tokens <= (usage.completion_tokens or 0): @@ -243,18 +238,12 @@ class AnthropicPassthroughLoggingHandler: # perform_redaction scrubs this field only when stream is True, so setting # it on a non-streaming response would bypass message redaction. if logging_obj.model_call_details.get("stream") is True: - logging_obj.model_call_details["complete_streaming_response"] = ( - litellm_model_response - ) + logging_obj.model_call_details["complete_streaming_response"] = litellm_model_response try: # Get custom_llm_provider from logging object if available (e.g., azure_ai for Azure Anthropic) - custom_llm_provider = logging_obj.model_call_details.get( - "custom_llm_provider" - ) + custom_llm_provider = logging_obj.model_call_details.get("custom_llm_provider") - model = AnthropicPassthroughLoggingHandler._resolve_costing_model( - model, logging_obj - ) + model = AnthropicPassthroughLoggingHandler._resolve_costing_model(model, logging_obj) # Prepend custom_llm_provider to model if not already present model_for_cost = model @@ -263,11 +252,7 @@ class AnthropicPassthroughLoggingHandler: router_model_id = logging_obj.get_router_model_id() custom_pricing = use_custom_pricing_for_model( - litellm_params=( - logging_obj.litellm_params - if hasattr(logging_obj, "litellm_params") - else None - ) + litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None) ) response_cost = litellm.completion_cost( @@ -292,9 +277,7 @@ class AnthropicPassthroughLoggingHandler: ) if user: kwargs.setdefault("litellm_params", {}) - kwargs["litellm_params"].update( - {"proxy_server_request": {"body": {"user": user}}} - ) + kwargs["litellm_params"].update({"proxy_server_request": {"body": {"user": user}}}) # pretty print standard logging object verbose_proxy_logger.debug( @@ -307,14 +290,10 @@ class AnthropicPassthroughLoggingHandler: litellm_model_response.model = model logging_obj.model_call_details["model"] = model if not logging_obj.model_call_details.get("custom_llm_provider"): - logging_obj.model_call_details["custom_llm_provider"] = ( - litellm.LlmProviders.ANTHROPIC.value - ) + logging_obj.model_call_details["custom_llm_provider"] = litellm.LlmProviders.ANTHROPIC.value return kwargs except Exception as e: - verbose_proxy_logger.exception( - "Error creating Anthropic response logging payload: %s", e - ) + verbose_proxy_logger.exception("Error creating Anthropic response logging payload: %s", e) return kwargs @staticmethod @@ -346,21 +325,15 @@ class AnthropicPassthroughLoggingHandler: model = cast(str, litellm_logging_obj.model_call_details.get("model")) if not model or model == "unknown": - chunk_model = ( - AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks( - all_chunks - ) - ) + chunk_model = AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks(all_chunks) if chunk_model: model = chunk_model try: - complete_streaming_response = ( - AnthropicPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=all_chunks, - litellm_logging_obj=litellm_logging_obj, - model=model, - ) + complete_streaming_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, ) except Exception as e: # stream_chunk_builder re-raises assembly failures (as litellm.APIError) @@ -464,9 +437,7 @@ class AnthropicPassthroughLoggingHandler: event-loop CPU under concurrent streaming; collapsing the homogeneous text run removes O(num_output_tokens) of it. """ - collapsed = AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks( - all_chunks - ) + collapsed = AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) if collapsed is not None: return AnthropicPassthroughLoggingHandler._build_complete_streaming_response_legacy( all_chunks=collapsed, @@ -565,10 +536,7 @@ class AnthropicPassthroughLoggingHandler: elif etype == "content_block_delta": delta = data.get("delta") or {} dtype = delta.get("type") - if ( - dtype - in AnthropicPassthroughLoggingHandler._FAST_PATH_DISALLOWED_DELTA_TYPES - ): + if dtype in AnthropicPassthroughLoggingHandler._FAST_PATH_DISALLOWED_DELTA_TYPES: return None if dtype != "text_delta": return None @@ -582,11 +550,7 @@ class AnthropicPassthroughLoggingHandler: # disagrees with the current pending buffer, the stream is # interleaved -- fall back to legacy rather than risk merging # text from different blocks under a single index. - if ( - pending_text - and pending_index is not None - and cur_index != pending_index - ): + if pending_text and pending_index is not None and cur_index != pending_index: return None saw_any_text_delta = True pending_index = cur_index @@ -622,9 +586,7 @@ class AnthropicPassthroughLoggingHandler: - Converts generic chunks to litellm chunks (OpenAI format) - Builds complete response from litellm chunks """ - verbose_proxy_logger.debug( - "Building complete streaming response from %d chunks", len(all_chunks) - ) + verbose_proxy_logger.debug("Building complete streaming response from %d chunks", len(all_chunks)) anthropic_model_response_iterator = AnthropicModelResponseIterator( streaming_response=None, sync_stream=False, @@ -634,11 +596,7 @@ class AnthropicPassthroughLoggingHandler: # Process each chunk - a chunk may contain multiple SSE events for _chunk_str in all_chunks: # Split chunk into individual SSE events - individual_events = ( - AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events( - _chunk_str - ) - ) + individual_events = AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(_chunk_str) # Process each individual event for event_str in individual_events: @@ -646,9 +604,7 @@ class AnthropicPassthroughLoggingHandler: # Skip OpenAI-style [DONE] sentinels some Anthropic-compatible # providers emit. Match the whole SSE line so a valid chunk whose # text payload happens to contain "[DONE]" is not dropped. - if any( - line.strip() == "data: [DONE]" for line in event_str.split("\n") - ): + if any(line.strip() == "data: [DONE]" for line in event_str.split("\n")): continue transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk( chunk=event_str @@ -671,9 +627,7 @@ class AnthropicPassthroughLoggingHandler: chunks=all_openai_chunks, logging_obj=litellm_logging_obj, ) - verbose_proxy_logger.debug( - "Complete streaming response built: %s", complete_streaming_response - ) + verbose_proxy_logger.debug("Complete streaming response built: %s", complete_streaming_response) return complete_streaming_response @staticmethod @@ -718,11 +672,7 @@ class AnthropicPassthroughLoggingHandler: found_usage = False resolved_model = model for _chunk_str in all_chunks: - for ( - event_str - ) in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events( - _chunk_str - ): + for event_str in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(_chunk_str): data = AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) if not data: continue @@ -734,9 +684,7 @@ class AnthropicPassthroughLoggingHandler: usage = message.get("usage") or {} input_tokens = usage.get("input_tokens") or input_tokens cache_read = usage.get("cache_read_input_tokens") or cache_read - cache_creation = ( - usage.get("cache_creation_input_tokens") or cache_creation - ) + cache_creation = usage.get("cache_creation_input_tokens") or cache_creation _cc = usage.get("cache_creation") if isinstance(_cc, dict): cache_creation_5m = _cc.get("ephemeral_5m_input_tokens") @@ -794,16 +742,12 @@ class AnthropicPassthroughLoggingHandler: usage_object["server_tool_use"] = _server_tool_use if inference_geo is not None: usage_object["inference_geo"] = inference_geo - usage_obj = AnthropicConfig().calculate_usage( - usage_object=usage_object, reasoning_content=None - ) + usage_obj = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None) return ModelResponse( model=resolved_model, choices=[ Choices( - finish_reason=( - map_finish_reason(stop_reason) if stop_reason else "stop" - ), + finish_reason=(map_finish_reason(stop_reason) if stop_reason else "stop"), index=0, message=Message(role="assistant", content=""), ) @@ -840,13 +784,11 @@ class AnthropicPassthroughLoggingHandler: if httpx_response.status_code == 200 and "id" in _json_response: # Transform Anthropic response to LiteLLM batch format anthropic_batches_config = AnthropicBatchesConfig() - litellm_batch_response = ( - anthropic_batches_config.transform_retrieve_batch_response( - model=None, - raw_response=httpx_response, - logging_obj=logging_obj, - litellm_params={}, - ) + litellm_batch_response = anthropic_batches_config.transform_retrieve_batch_response( + model=None, + raw_response=httpx_response, + logging_obj=logging_obj, + litellm_params={}, ) # Set status to "validating" for newly created batches so polling mechanism picks them up # The polling mechanism only looks for status="validating" jobs @@ -877,28 +819,16 @@ class AnthropicPassthroughLoggingHandler: # Create unified object ID for tracking # Format: base64(litellm_proxy;model_id:{};llm_batch_id:{}) # For Anthropic passthrough, prefix model with "anthropic/" so router can determine provider - actual_model_id = ( - AnthropicPassthroughLoggingHandler.get_actual_model_id_from_router( - model_name - ) - ) + actual_model_id = AnthropicPassthroughLoggingHandler.get_actual_model_id_from_router(model_name) # If model not in router, use "anthropic/{model_name}" format so router can determine provider - if actual_model_id == model_name and not actual_model_id.startswith( - "anthropic/" - ): + if actual_model_id == model_name and not actual_model_id.startswith("anthropic/"): actual_model_id = f"anthropic/{model_name}" - unified_id_string = ( - SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( - actual_model_id, batch_id - ) - ) - unified_object_id = ( - base64.urlsafe_b64encode(unified_id_string.encode()) - .decode() - .rstrip("=") + unified_id_string = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( + actual_model_id, batch_id ) + unified_object_id = base64.urlsafe_b64encode(unified_id_string.encode()).decode().rstrip("=") # Store the managed object for cost tracking # This will be picked up by check_batch_cost polling mechanism @@ -1042,20 +972,14 @@ class AnthropicPassthroughLoggingHandler: from litellm.proxy.proxy_server import proxy_logging_obj managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_hook is not None and hasattr( - managed_files_hook, "store_unified_object_id" - ): + if managed_files_hook is not None and hasattr(managed_files_hook, "store_unified_object_id"): # Create a mock user API key dict for the managed object storage from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - _request_metadata = (kwargs.get("litellm_params", {}) or {}).get( - "metadata", {} - ) or {} + _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} user_api_key_dict = UserAPIKeyAuth( - user_id=_request_metadata.get( - "user_api_key_user_id", "default-user" - ), + user_id=_request_metadata.get("user_api_key_user_id", "default-user"), api_key="", team_id=_request_metadata.get("user_api_key_team_id"), team_alias=None, @@ -1100,9 +1024,7 @@ class AnthropicPassthroughLoggingHandler: ) except Exception as e: - verbose_proxy_logger.error( - f"Error storing Anthropic batch managed object: {e}" - ) + verbose_proxy_logger.error(f"Error storing Anthropic batch managed object: {e}") @staticmethod def get_actual_model_id_from_router(model_name: str) -> str: @@ -1115,20 +1037,14 @@ class AnthropicPassthroughLoggingHandler: if model_ids and len(model_ids) > 0: # Use the first model ID found actual_model_id = model_ids[0] - verbose_proxy_logger.info( - f"Found model ID in router: {actual_model_id}" - ) + verbose_proxy_logger.info(f"Found model ID in router: {actual_model_id}") return actual_model_id else: # Fallback to model name actual_model_id = model_name - verbose_proxy_logger.warning( - f"Model not found in router, using model name: {actual_model_id}" - ) + verbose_proxy_logger.warning(f"Model not found in router, using model name: {actual_model_id}") return actual_model_id else: # Fallback if router is not available - verbose_proxy_logger.warning( - f"Router not available, using model name: {model_name}" - ) + verbose_proxy_logger.warning(f"Router not available, using model name: {model_name}") return model_name diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py index 6277f6b4a75..7d7bc889120 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py @@ -99,18 +99,14 @@ class AssemblyAIPassthroughLoggingHandler: from ..pass_through_endpoints import pass_through_endpoint_logging model = response_body.get("speech_model", "") - verbose_proxy_logger.debug( - "response body %s", json.dumps(response_body, indent=4) - ) + verbose_proxy_logger.debug("response body %s", json.dumps(response_body, indent=4)) kwargs["model"] = model kwargs["custom_llm_provider"] = "assemblyai" response_cost: Optional[float] = None transcript_id = response_body.get("id") if transcript_id is None: - raise ValueError( - "Transcript ID is required to log the cost of the transcription" - ) + raise ValueError("Transcript ID is required to log the cost of the transcription") transcript_response = self._poll_assembly_for_transcript_response( transcript_id=transcript_id, url_route=url_route ) @@ -145,9 +141,7 @@ class AssemblyAIPassthroughLoggingHandler: ) # pretty print standard logging object - verbose_proxy_logger.debug( - "standard_logging_object= %s", json.dumps(standard_logging_object, indent=4) - ) + verbose_proxy_logger.debug("standard_logging_object= %s", json.dumps(standard_logging_object, indent=4)) logging_obj.model_call_details["model"] = model logging_obj.model_call_details["custom_llm_provider"] = "assemblyai" logging_obj.model_call_details["response_cost"] = response_cost @@ -155,9 +149,7 @@ class AssemblyAIPassthroughLoggingHandler: asyncio.run( pass_through_endpoint_logging._handle_logging( logging_obj=logging_obj, - standard_logging_response_object=self._get_response_to_log( - transcript_response - ), + standard_logging_response_object=self._get_response_to_log(transcript_response), result=result, start_time=start_time, end_time=end_time, @@ -168,9 +160,7 @@ class AssemblyAIPassthroughLoggingHandler: pass - def _get_response_to_log( - self, transcript_response: Optional[AssemblyAITranscriptResponse] - ) -> dict: + def _get_response_to_log(self, transcript_response: Optional[AssemblyAITranscriptResponse]) -> dict: if transcript_response is None: return {} return dict(transcript_response) @@ -193,24 +183,15 @@ class AssemblyAIPassthroughLoggingHandler: passthrough_endpoint_router, ) - _base_url = ( - self.assembly_ai_eu_base_url - if request_region == "eu" - else self.assembly_ai_base_url - ) + _base_url = self.assembly_ai_eu_base_url if request_region == "eu" else self.assembly_ai_base_url _api_key = passthrough_endpoint_router.get_credentials( custom_llm_provider="assemblyai", region_name=request_region, ) if _api_key is None: raise ValueError("AssemblyAI API key not found") - if ( - any(c in transcript_id for c in ("/", "\\", "#", "?")) - or ".." in transcript_id - ): - raise ValueError( - f"Invalid transcript_id {transcript_id!r}: contains disallowed characters" - ) + if any(c in transcript_id for c in ("/", "\\", "#", "?")) or ".." in transcript_id: + raise ValueError(f"Invalid transcript_id {transcript_id!r}: contains disallowed characters") safe_transcript_id = urllib.parse.quote(transcript_id, safe="") try: url = f"{_base_url}/v2/transcript/{safe_transcript_id}" @@ -237,21 +218,14 @@ class AssemblyAIPassthroughLoggingHandler: """ Poll the status of the transcript until it is completed or timeout (30 minutes) """ - for _ in range( - self.max_polling_attempts - ): # 180 attempts * 10s = 30 minutes max + for _ in range(self.max_polling_attempts): # 180 attempts * 10s = 30 minutes max transcript = self._get_assembly_transcript( - request_region=AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url( - url=url_route - ), + request_region=AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(url=url_route), transcript_id=transcript_id, ) if transcript is None: return None - if ( - transcript.get("status") == "completed" - or transcript.get("status") == "error" - ): + if transcript.get("status") == "completed" or transcript.get("status") == "error": return AssemblyAITranscriptResponse(**transcript) time.sleep(self.polling_interval) return None @@ -267,10 +241,8 @@ class AssemblyAIPassthroughLoggingHandler: _audio_duration = transcript_response.get("audio_duration") if _audio_duration is None: return None - _cost_per_second = ( - AssemblyAIPassthroughLoggingHandler.get_cost_per_second_for_assembly_model( - speech_model=speech_model - ) + _cost_per_second = AssemblyAIPassthroughLoggingHandler.get_cost_per_second_for_assembly_model( + speech_model=speech_model ) if _cost_per_second is None: return None diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py index a7ec2f0d368..980c05fa412 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py @@ -128,9 +128,7 @@ class BasePassthroughLoggingHandler(ABC): ) if user: kwargs.setdefault("litellm_params", {}) - kwargs["litellm_params"].update( - {"proxy_server_request": {"body": {"user": user}}} - ) + kwargs["litellm_params"].update({"proxy_server_request": {"body": {"user": user}}}) # Make standard logging object for Anthropic standard_logging_object = get_standard_logging_object_payload( @@ -155,9 +153,7 @@ class BasePassthroughLoggingHandler(ABC): logging_obj.model_call_details["model"] = model return kwargs except Exception as e: - verbose_proxy_logger.exception( - "Error creating LLM passthrough response logging payload: %s", e - ) + verbose_proxy_logger.exception("Error creating LLM passthrough response logging payload: %s", e) return kwargs @abstractmethod diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py index 0875f1d5508..70b09e101fc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py @@ -56,14 +56,8 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler): all_openai_chunks = [] for _chunk_str in all_chunks: try: - generic_chunk = ( - cohere_model_response_iterator.convert_str_chunk_to_generic_chunk( - chunk=_chunk_str - ) - ) - litellm_chunk = litellm_custom_stream_wrapper.chunk_creator( - chunk=generic_chunk - ) + generic_chunk = cohere_model_response_iterator.convert_str_chunk_to_generic_chunk(chunk=_chunk_str) + litellm_chunk = litellm_custom_stream_wrapper.chunk_creator(chunk=generic_chunk) if litellm_chunk is not None: all_openai_chunks.append(litellm_chunk) except (StopIteration, StopAsyncIteration): @@ -129,18 +123,16 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler): kwargs["custom_llm_provider"] = "cohere" # Extract user information for tracking - passthrough_logging_payload: Optional[ - PassthroughStandardLoggingPayload - ] = kwargs.get("passthrough_logging_payload") + passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = kwargs.get( + "passthrough_logging_payload" + ) if passthrough_logging_payload: user = handler_instance._get_user_from_metadata( passthrough_logging_payload=passthrough_logging_payload, ) if user: kwargs.setdefault("litellm_params", {}) - kwargs["litellm_params"].update( - {"proxy_server_request": {"body": {"user": user}}} - ) + kwargs["litellm_params"].update({"proxy_server_request": {"body": {"user": user}}}) # Create standard logging object if litellm_model_response is not None: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py index e7696e5a18a..63907e9638b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py @@ -108,9 +108,7 @@ class CursorPassthroughLoggingHandler: standard_logging_object = get_standard_logging_object_payload( kwargs=kwargs, - init_response_obj=StandardPassThroughResponseObject( - response=response_summary - ), + init_response_obj=StandardPassThroughResponseObject(response=response_summary), start_time=start_time, end_time=end_time, logging_obj=logging_obj, @@ -129,9 +127,7 @@ class CursorPassthroughLoggingHandler: "kwargs": kwargs, } except Exception as e: - verbose_proxy_logger.exception( - "Error in Cursor passthrough logging handler: %s", e - ) + verbose_proxy_logger.exception("Error in Cursor passthrough logging handler: %s", e) return { "result": StandardPassThroughResponseObject(response=result), "kwargs": kwargs, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index b05cb70f756..716dff21efc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -44,14 +44,12 @@ class GeminiPassthroughLoggingHandler: model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) gemini_video_config = GeminiVideoConfig() - litellm_video_response = ( - gemini_video_config.transform_video_create_response( - model=model, - raw_response=httpx_response, - logging_obj=logging_obj, - custom_llm_provider="gemini", - request_data=request_body, - ) + litellm_video_response = gemini_video_config.transform_video_create_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + custom_llm_provider="gemini", + request_data=request_body, ) logging_obj.model = model logging_obj.model_call_details["model"] = model @@ -84,21 +82,17 @@ class GeminiPassthroughLoggingHandler: # Use Gemini config for transformation instance_of_gemini_llm = litellm.GoogleAIStudioGeminiConfig() - litellm_model_response: ModelResponse = ( - instance_of_gemini_llm.transform_response( - model=model, - messages=[ - {"role": "user", "content": "no-message-pass-through-endpoint"} - ], - raw_response=httpx_response, - model_response=litellm.ModelResponse(), - logging_obj=logging_obj, - optional_params={}, - litellm_params={}, - api_key="", - request_data={}, - encoding=litellm.encoding, - ) + litellm_model_response: ModelResponse = instance_of_gemini_llm.transform_response( + model=model, + messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data={}, + encoding=litellm.encoding, ) kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content( litellm_model_response=litellm_model_response, @@ -140,16 +134,12 @@ class GeminiPassthroughLoggingHandler: - Logs in litellm callbacks """ kwargs: Dict[str, Any] = {} - model = model or GeminiPassthroughLoggingHandler.extract_model_from_url( - url_route - ) - complete_streaming_response = ( - GeminiPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=all_chunks, - litellm_logging_obj=litellm_logging_obj, - model=model, - url_route=url_route, - ) + model = model or GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) + complete_streaming_response = GeminiPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + url_route=url_route, ) if complete_streaming_response is None: @@ -204,9 +194,7 @@ class GeminiPassthroughLoggingHandler: continue all_openai_chunks.append(parsed_chunk) - complete_streaming_response = litellm.stream_chunk_builder( - chunks=all_openai_chunks - ) + complete_streaming_response = litellm.stream_chunk_builder(chunks=all_openai_chunks) return complete_streaming_response diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 1d76165c30c..f71bc167fb9 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -59,9 +59,7 @@ def _hostname_matches(hostname: str, suffixes: tuple) -> bool: Uses suffix matching (not a bare substring test) so look-alikes such as `cognitiveservices.azure.com.attacker.example` are not accepted. """ - return any( - hostname == suffix or hostname.endswith("." + suffix) for suffix in suffixes - ) + return any(hostname == suffix or hostname.endswith("." + suffix) for suffix in suffixes) def _is_openai_compatible_host(hostname: Optional[str]) -> bool: @@ -74,9 +72,7 @@ def _is_openai_compatible_host(hostname: Optional[str]) -> bool: """ if not hostname: return False - return _hostname_matches(hostname, _OPENAI_HOSTNAMES) or _hostname_matches( - hostname, _AZURE_OPENAI_HOSTNAMES - ) + return _hostname_matches(hostname, _OPENAI_HOSTNAMES) or _hostname_matches(hostname, _AZURE_OPENAI_HOSTNAMES) def _is_openai_compatible_url(url_route: Optional[str]) -> bool: @@ -119,10 +115,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url = urlparse(url_route) - return ( - _is_openai_compatible_host(parsed_url.hostname) - and "/v1/chat/completions" in parsed_url.path - ) + return _is_openai_compatible_host(parsed_url.hostname) and "/v1/chat/completions" in parsed_url.path @staticmethod def is_openai_image_generation_route(url_route: str) -> bool: @@ -130,10 +123,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url = urlparse(url_route) - return ( - _is_openai_compatible_host(parsed_url.hostname) - and "/v1/images/generations" in parsed_url.path - ) + return _is_openai_compatible_host(parsed_url.hostname) and "/v1/images/generations" in parsed_url.path @staticmethod def is_openai_image_editing_route(url_route: str) -> bool: @@ -141,10 +131,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url = urlparse(url_route) - return ( - _is_openai_compatible_host(parsed_url.hostname) - and "/v1/images/edits" in parsed_url.path - ) + return _is_openai_compatible_host(parsed_url.hostname) and "/v1/images/edits" in parsed_url.path @staticmethod def is_openai_responses_route(url_route: str) -> bool: @@ -197,9 +184,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return cost except Exception as e: - verbose_proxy_logger.warning( - f"Error calculating image generation cost: {str(e)}" - ) + verbose_proxy_logger.warning(f"Error calculating image generation cost: {str(e)}") return 0.0 @staticmethod @@ -233,9 +218,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return cost except Exception as e: - verbose_proxy_logger.warning( - f"Error calculating image editing cost: {str(e)}" - ) + verbose_proxy_logger.warning(f"Error calculating image editing cost: {str(e)}") return 0.0 @staticmethod @@ -291,25 +274,12 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, image editing, and responses API. """ # Check if this is a supported endpoint for cost tracking - is_chat_completions = ( - OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) - ) - is_image_generation = ( - OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) - ) - is_image_editing = ( - OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) - ) - is_responses = OpenAIPassthroughLoggingHandler.is_openai_responses_route( - url_route - ) + is_chat_completions = OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) + is_image_generation = OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) + is_image_editing = OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) + is_responses = OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) - if not ( - is_chat_completions - or is_image_generation - or is_image_editing - or is_responses - ): + if not (is_chat_completions or is_image_generation or is_image_editing or is_responses): # For unsupported endpoints, return None to let the system fall back to generic behavior return { "result": None, @@ -319,9 +289,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Extract model from request or response model = request_body.get("model", response_body.get("model", "")) if not model: - verbose_proxy_logger.warning( - "No model found in request or response for OpenAI passthrough cost tracking" - ) + verbose_proxy_logger.warning("No model found in request or response for OpenAI passthrough cost tracking") base_handler = OpenAIPassthroughLoggingHandler() return base_handler.passthrough_chat_handler( httpx_response=httpx_response, @@ -365,8 +333,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): api_key="", request_data=request_body, encoding=litellm.encoding, - json_mode=request_body.get("response_format", {}).get("type") - == "json_object", + json_mode=request_body.get("response_format", {}).get("type") == "json_object", litellm_params=existing_litellm_params, ) @@ -378,18 +345,14 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) elif is_image_generation: # Handle image generation cost calculation - response_cost = ( - OpenAIPassthroughLoggingHandler._calculate_image_generation_cost( - model=model, - response_body=response_body, - request_body=request_body, - ) + response_cost = OpenAIPassthroughLoggingHandler._calculate_image_generation_cost( + model=model, + response_body=response_body, + request_body=request_body, ) # Mark call type for downstream image-aware logic/metrics try: - logging_obj.call_type = ( - PassthroughCallTypes.passthrough_image_generation.value - ) + logging_obj.call_type = PassthroughCallTypes.passthrough_image_generation.value except Exception: pass # Create a simple response object for logging @@ -403,18 +366,14 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): litellm_model_response._hidden_params["response_cost"] = response_cost elif is_image_editing: # Handle image editing cost calculation - response_cost = ( - OpenAIPassthroughLoggingHandler._calculate_image_editing_cost( - model=model, - response_body=response_body, - request_body=request_body, - ) + response_cost = OpenAIPassthroughLoggingHandler._calculate_image_editing_cost( + model=model, + response_body=response_body, + request_body=request_body, ) # Mark call type for downstream image-aware logic/metrics try: - logging_obj.call_type = ( - PassthroughCallTypes.passthrough_image_generation.value - ) + logging_obj.call_type = PassthroughCallTypes.passthrough_image_generation.value except Exception: pass # Create a simple response object for logging @@ -447,17 +406,17 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): kwargs["custom_llm_provider"] = custom_llm_provider # Extract user information for tracking - passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( - kwargs.get("passthrough_logging_payload") + passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = kwargs.get( + "passthrough_logging_payload" ) if passthrough_logging_payload: user = handler_instance._get_user_from_metadata( passthrough_logging_payload=passthrough_logging_payload, ) if user: - kwargs["litellm_params"].setdefault( - "proxy_server_request", {} - ).setdefault("body", {})["user"] = user + kwargs["litellm_params"].setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = ( + user + ) # Create standard logging object if litellm_model_response is not None: @@ -492,9 +451,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): } except Exception as e: - verbose_proxy_logger.error( - f"Error in OpenAI passthrough cost tracking: {str(e)}" - ) + verbose_proxy_logger.error(f"Error in OpenAI passthrough cost tracking: {str(e)}") # Fall back to base handler without cost tracking base_handler = OpenAIPassthroughLoggingHandler() return base_handler.passthrough_chat_handler( @@ -541,17 +498,11 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) # Convert string chunk to dict - stripped_json_chunk = ( - BaseModelResponseIterator._string_to_dict_parser( - str_line=chunk_str - ) - ) + stripped_json_chunk = BaseModelResponseIterator._string_to_dict_parser(str_line=chunk_str) if stripped_json_chunk: # Parse the chunk using OpenAI's chunk parser - transformed_chunk = openai_iterator.chunk_parser( - chunk=stripped_json_chunk - ) + transformed_chunk = openai_iterator.chunk_parser(chunk=stripped_json_chunk) if transformed_chunk is not None: all_openai_chunks.append(transformed_chunk) @@ -560,22 +511,16 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): continue if not all_openai_chunks: - verbose_proxy_logger.warning( - "No valid chunks found in streaming response" - ) + verbose_proxy_logger.warning("No valid chunks found in streaming response") return None # Build complete response from chunks - complete_streaming_response = litellm.stream_chunk_builder( - chunks=all_openai_chunks - ) + complete_streaming_response = litellm.stream_chunk_builder(chunks=all_openai_chunks) return complete_streaming_response except Exception as e: - verbose_proxy_logger.error( - f"Error building complete streaming response: {str(e)}" - ) + verbose_proxy_logger.error(f"Error building complete streaming response: {str(e)}") return None @staticmethod @@ -606,17 +551,13 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) if complete_response is None: - verbose_proxy_logger.warning( - "Failed to build complete response from OpenAI streaming chunks" - ) + verbose_proxy_logger.warning("Failed to build complete response from OpenAI streaming chunks") return { "result": None, "kwargs": {}, } - custom_llm_provider = litellm_logging_obj.model_call_details.get( - "custom_llm_provider", "openai" - ) + custom_llm_provider = litellm_logging_obj.model_call_details.get("custom_llm_provider", "openai") # Calculate cost using LiteLLM's cost calculator response_cost = litellm.completion_cost( completion_response=complete_response, @@ -625,9 +566,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) # Preserve existing litellm_params to maintain metadata tags - existing_litellm_params = ( - litellm_logging_obj.model_call_details.get("litellm_params", {}) or {} - ) + existing_litellm_params = litellm_logging_obj.model_call_details.get("litellm_params", {}) or {} # Prepare kwargs for logging kwargs = { @@ -639,18 +578,16 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Extract user information for tracking passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( - litellm_logging_obj.model_call_details.get( - "passthrough_logging_payload" - ) + litellm_logging_obj.model_call_details.get("passthrough_logging_payload") ) if passthrough_logging_payload: user = handler_instance._get_user_from_metadata( passthrough_logging_payload=passthrough_logging_payload, ) if user: - kwargs["litellm_params"].setdefault( - "proxy_server_request", {} - ).setdefault("body", {})["user"] = user + kwargs["litellm_params"].setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = ( + user + ) # Create standard logging object get_standard_logging_object_payload( @@ -664,9 +601,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Update logging object with cost information litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details["custom_llm_provider"] = ( - custom_llm_provider - ) + litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider litellm_logging_obj.model_call_details["response_cost"] = response_cost verbose_proxy_logger.debug( @@ -679,9 +614,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): } except Exception as e: - verbose_proxy_logger.error( - f"Error in OpenAI streaming passthrough cost tracking: {str(e)}" - ) + verbose_proxy_logger.error(f"Error in OpenAI streaming passthrough cost tracking: {str(e)}") return { "result": None, "kwargs": {}, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 04fe74bbf25..158b629ad27 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -115,13 +115,9 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): # Convert aggregated modality totals back to details format for modality, totals in modality_totals.items(): if totals["prompt"] > 0: - aggregated["promptTokensDetails"].append( - {"modality": modality, "tokenCount": totals["prompt"]} - ) + aggregated["promptTokensDetails"].append({"modality": modality, "tokenCount": totals["prompt"]}) if totals["candidate"] > 0: - aggregated["candidatesTokensDetails"].append( - {"modality": modality, "tokenCount": totals["candidate"]} - ) + aggregated["candidatesTokensDetails"].append({"modality": modality, "tokenCount": totals["candidate"]}) # Add any additional fields from the first usage metadata first_usage = all_usage_metadata[0] @@ -150,19 +146,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ try: # Get model pricing information - model_info = get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - verbose_proxy_logger.debug( - f"Vertex AI Live API model info for '{model}': {model_info}" - ) + verbose_proxy_logger.debug(f"Vertex AI Live API model info for '{model}': {model_info}") # Check if pricing info is available if not model_info or not model_info.get("input_cost_per_token"): - verbose_proxy_logger.error( - f"No pricing info found for {model} in local model pricing database" - ) + verbose_proxy_logger.error(f"No pricing info found for {model} in local model pricing database") return 0.0 total_cost = 0.0 @@ -180,9 +170,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): # Handle modality-specific costs if present prompt_tokens_details = usage_metadata.get("promptTokensDetails", []) - candidates_tokens_details = usage_metadata.get( - "candidatesTokensDetails", [] - ) + candidates_tokens_details = usage_metadata.get("candidatesTokensDetails", []) # Process prompt tokens by modality for detail in prompt_tokens_details: @@ -190,15 +178,11 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): token_count = detail.get("tokenCount", 0) if modality == "AUDIO": - audio_cost_per_token = model_info.get( - "input_cost_per_audio_token", 0.0 - ) + audio_cost_per_token = model_info.get("input_cost_per_audio_token", 0.0) total_cost += token_count * audio_cost_per_token elif modality == "VIDEO": # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get( - "input_cost_per_video_per_second", 0.0 - ) + video_cost_per_token = model_info.get("input_cost_per_video_per_second", 0.0) total_cost += token_count * video_cost_per_token # TEXT tokens are already handled above @@ -208,22 +192,16 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): token_count = detail.get("tokenCount", 0) if modality == "AUDIO": - audio_cost_per_token = model_info.get( - "output_cost_per_audio_token", 0.0 - ) + audio_cost_per_token = model_info.get("output_cost_per_audio_token", 0.0) total_cost += token_count * audio_cost_per_token elif modality == "VIDEO": # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get( - "output_cost_per_video_per_second", 0.0 - ) + video_cost_per_token = model_info.get("output_cost_per_video_per_second", 0.0) total_cost += token_count * video_cost_per_token # TEXT tokens are already handled above # Handle web search costs if present - tool_use_prompt_token_count = usage_metadata.get( - "toolUsePromptTokenCount", 0 - ) + tool_use_prompt_token_count = usage_metadata.get("toolUsePromptTokenCount", 0) if tool_use_prompt_token_count > 0: # Web search typically has a fixed cost per request web_search_cost = model_info.get("web_search_cost_per_request", 0.0) @@ -243,9 +221,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): return total_cost except Exception as e: - verbose_proxy_logger.error( - f"Error calculating Vertex AI Live API cost: {e}" - ) + verbose_proxy_logger.error(f"Error calculating Vertex AI Live API cost: {e}") return 0.0 @staticmethod @@ -326,19 +302,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): # Extract model from request body or kwargs model = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09") custom_llm_provider = kwargs.get("custom_llm_provider", "vertex_ai") - verbose_proxy_logger.debug( - f"Vertex AI Live API model: {model}, custom_llm_provider: {custom_llm_provider}" - ) + verbose_proxy_logger.debug(f"Vertex AI Live API model: {model}, custom_llm_provider: {custom_llm_provider}") # Extract usage metadata from WebSocket messages - usage_metadata = self._extract_usage_metadata_from_websocket_messages( - websocket_messages - ) + usage_metadata = self._extract_usage_metadata_from_websocket_messages(websocket_messages) if not usage_metadata: - verbose_proxy_logger.warning( - "No usage metadata found in Vertex AI Live API WebSocket messages" - ) + verbose_proxy_logger.warning("No usage metadata found in Vertex AI Live API WebSocket messages") return { "result": None, "kwargs": kwargs, @@ -376,11 +346,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): import re allowed_pattern = re.compile(r"^[A-Za-z0-9._\-:]+$") - safe_model = ( - model - if isinstance(model, str) and allowed_pattern.match(model) - else "[REDACTED]" - ) + safe_model = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" verbose_proxy_logger.debug( f"Vertex AI Live API passthrough cost tracking - " f"Model: {safe_model}, Cost: ${response_cost:.6f}, " @@ -394,9 +360,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): } except Exception as e: - verbose_proxy_logger.error( - f"Error in Vertex AI Live API passthrough handler: {e}" - ) + verbose_proxy_logger.error(f"Error in Vertex AI Live API passthrough handler: {e}") return { "result": None, "kwargs": kwargs, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 538267e9c84..c61d48eda8c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -56,14 +56,12 @@ class VertexPassthroughLoggingHandler: model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) vertex_video_config = VertexAIVideoConfig() - litellm_video_response = ( - vertex_video_config.transform_video_create_response( - model=model, - raw_response=httpx_response, - logging_obj=logging_obj, - custom_llm_provider="vertex_ai", - request_data=request_body, - ) + litellm_video_response = vertex_video_config.transform_video_create_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + request_data=request_body, ) logging_obj.model = model @@ -97,21 +95,17 @@ class VertexPassthroughLoggingHandler: model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) instance_of_vertex_llm = litellm.VertexGeminiConfig() - litellm_model_response: ModelResponse = ( - instance_of_vertex_llm.transform_response( - model=model, - messages=[ - {"role": "user", "content": "no-message-pass-through-endpoint"} - ], - raw_response=httpx_response, - model_response=litellm.ModelResponse(), - logging_obj=logging_obj, - optional_params={}, - litellm_params={}, - api_key="", - request_data={}, - encoding=litellm.encoding, - ) + litellm_model_response: ModelResponse = instance_of_vertex_llm.transform_response( + model=model, + messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data={}, + encoding=litellm.encoding, ) kwargs = VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content( litellm_model_response=litellm_model_response, @@ -120,9 +114,7 @@ class VertexPassthroughLoggingHandler: start_time=start_time, end_time=end_time, logging_obj=logging_obj, - custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url( - url_route - ), + custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route), ) return { @@ -164,24 +156,22 @@ class VertexPassthroughLoggingHandler: model=model, vertex_publisher_or_api_spec=vertex_publisher_or_api_spec, ) - litellm_prediction_response = ( - vertex_ai_partner_model_config.transform_response( - model=model, - raw_response=httpx_response, - model_response=litellm_prediction_response, - logging_obj=logging_obj, - request_data={}, - encoding=litellm.encoding, - optional_params={}, - litellm_params={}, - api_key="", - messages=[ - { - "role": "user", - "content": "no-message-pass-through-endpoint", - } - ], - ) + litellm_prediction_response = vertex_ai_partner_model_config.transform_response( + model=model, + raw_response=httpx_response, + model_response=litellm_prediction_response, + logging_obj=logging_obj, + request_data={}, + encoding=litellm.encoding, + optional_params={}, + litellm_params={}, + api_key="", + messages=[ + { + "role": "user", + "content": "no-message-pass-through-endpoint", + } + ], ) kwargs = VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content( @@ -199,11 +189,9 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } elif "search" in url_route: - litellm_vs_response = ( - vertex_search_api_config.transform_search_vector_store_response( - response=httpx_response, - litellm_logging_obj=logging_obj, - ) + litellm_vs_response = vertex_search_api_config.transform_search_vector_store_response( + response=httpx_response, + litellm_logging_obj=logging_obj, ) response_cost = litellm.completion_cost( completion_response=litellm_vs_response, @@ -219,9 +207,7 @@ class VertexPassthroughLoggingHandler: kwargs["response_cost"] = response_cost kwargs["model"] = "vertex_ai/search_api" logging_obj.model_call_details.setdefault("litellm_params", {}) - logging_obj.model_call_details["litellm_params"]["base_model"] = ( - "vertex_ai/search_api" - ) + logging_obj.model_call_details["litellm_params"]["base_model"] = "vertex_ai/search_api" logging_obj.model_call_details["response_cost"] = response_cost return { @@ -267,45 +253,35 @@ class VertexPassthroughLoggingHandler: _json_response = httpx_response.json() - litellm_prediction_response: Union[ - ModelResponse, EmbeddingResponse, ImageResponse - ] = ModelResponse() + litellm_prediction_response: Union[ModelResponse, EmbeddingResponse, ImageResponse] = ModelResponse() if vertex_image_generation_class.is_image_generation_response(_json_response): - litellm_prediction_response = ( - vertex_image_generation_class.process_image_generation_response( - _json_response, - model_response=litellm.ImageResponse(), - model=model, - ) + litellm_prediction_response = vertex_image_generation_class.process_image_generation_response( + _json_response, + model_response=litellm.ImageResponse(), + model=model, ) - logging_obj.call_type = ( - PassthroughCallTypes.passthrough_image_generation.value - ) + logging_obj.call_type = PassthroughCallTypes.passthrough_image_generation.value elif VertexPassthroughLoggingHandler._is_multimodal_embedding_response( json_response=_json_response, ): # Use multimodal embedding transformation vertex_multimodal_config = VertexAIMultimodalEmbeddingConfig() - litellm_prediction_response = ( - vertex_multimodal_config.transform_embedding_response( - model=model, - raw_response=httpx_response, - model_response=litellm.EmbeddingResponse(), - logging_obj=logging_obj, - api_key="", - request_data={}, - optional_params={}, - litellm_params={}, - ) + litellm_prediction_response = vertex_multimodal_config.transform_embedding_response( + model=model, + raw_response=httpx_response, + model_response=litellm.EmbeddingResponse(), + logging_obj=logging_obj, + api_key="", + request_data={}, + optional_params={}, + litellm_params={}, ) else: - litellm_prediction_response = ( - litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( - response=_json_response, - model=model, - model_response=litellm.EmbeddingResponse(), - ) + litellm_prediction_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( + response=_json_response, + model=model, + model_response=litellm.EmbeddingResponse(), ) if isinstance(litellm_prediction_response, litellm.EmbeddingResponse): litellm_prediction_response.model = model @@ -383,9 +359,7 @@ class VertexPassthroughLoggingHandler: response_json=response_json, ) - custom_llm_provider = ( - VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route) - ) + custom_llm_provider = VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route) litellm_embedding_response.model = model logging_obj.model = model @@ -429,16 +403,12 @@ class VertexPassthroughLoggingHandler: - Logs in litellm callbacks """ kwargs: Dict[str, Any] = {} - model = model or VertexPassthroughLoggingHandler.extract_model_from_url( - url_route - ) - complete_streaming_response = ( - VertexPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=all_chunks, - litellm_logging_obj=litellm_logging_obj, - model=model, - url_route=url_route, - ) + model = model or VertexPassthroughLoggingHandler.extract_model_from_url(url_route) + complete_streaming_response = VertexPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + url_route=url_route, ) if complete_streaming_response is None: @@ -457,9 +427,7 @@ class VertexPassthroughLoggingHandler: start_time=start_time, end_time=end_time, logging_obj=litellm_logging_obj, - custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url( - url_route - ), + custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route), ) return { @@ -509,9 +477,7 @@ class VertexPassthroughLoggingHandler: continue all_openai_chunks.append(parsed_chunk) - complete_streaming_response = litellm.stream_chunk_builder( - chunks=all_openai_chunks - ) + complete_streaming_response = litellm.stream_chunk_builder(chunks=all_openai_chunks) return complete_streaming_response @@ -571,9 +537,7 @@ class VertexPassthroughLoggingHandler: @staticmethod def _get_custom_llm_provider_from_url(url: str) -> str: parsed_url = urlparse(url) - if parsed_url.hostname and parsed_url.hostname.endswith( - "generativelanguage.googleapis.com" - ): + if parsed_url.hostname and parsed_url.hostname.endswith("generativelanguage.googleapis.com"): return litellm.LlmProviders.GEMINI.value return litellm.LlmProviders.VERTEX_AI.value @@ -672,34 +636,24 @@ class VertexPassthroughLoggingHandler: # Only handle successful batch job creation (POST requests) if httpx_response.status_code == 200 and "name" in _json_response: # Transform Vertex AI response to LiteLLM batch format - litellm_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( - response=_json_response + litellm_batch_response = ( + VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( + response=_json_response + ) ) # Extract batch ID and model from the response - batch_id = VertexAIBatchTransformation._get_batch_id_from_vertex_ai_batch_response( - _json_response - ) + batch_id = VertexAIBatchTransformation._get_batch_id_from_vertex_ai_batch_response(_json_response) model_name = _json_response.get("model", "unknown") # Create unified object ID for tracking # Format: base64(litellm_proxy;model_id:{};llm_batch_id:{}) - actual_model_id = ( - VertexPassthroughLoggingHandler.get_actual_model_id_from_router( - model_name - ) - ) + actual_model_id = VertexPassthroughLoggingHandler.get_actual_model_id_from_router(model_name) - unified_id_string = ( - SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( - actual_model_id, batch_id - ) - ) - unified_object_id = ( - base64.urlsafe_b64encode(unified_id_string.encode()) - .decode() - .rstrip("=") + unified_id_string = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( + actual_model_id, batch_id ) + unified_object_id = base64.urlsafe_b64encode(unified_id_string.encode()).decode().rstrip("=") # Store the managed object for cost tracking # This will be picked up by check_batch_cost polling mechanism @@ -843,20 +797,14 @@ class VertexPassthroughLoggingHandler: from litellm.proxy.proxy_server import proxy_logging_obj managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_hook is not None and hasattr( - managed_files_hook, "store_unified_object_id" - ): + if managed_files_hook is not None and hasattr(managed_files_hook, "store_unified_object_id"): # Create a mock user API key dict for the managed object storage from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - _request_metadata = (kwargs.get("litellm_params", {}) or {}).get( - "metadata", {} - ) or {} + _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} user_api_key_dict = UserAPIKeyAuth( - user_id=_request_metadata.get( - "user_api_key_user_id", "default-user" - ), + user_id=_request_metadata.get("user_api_key_user_id", "default-user"), api_key="", team_id=_request_metadata.get("user_api_key_team_id"), team_alias=None, @@ -909,36 +857,22 @@ class VertexPassthroughLoggingHandler: if llm_router is not None: # Try to find the model in the router by the extracted model name - extracted_model_name = ( - VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path( - model_name - ) - ) + extracted_model_name = VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path(model_name) # Use the existing get_model_ids method from router model_ids = llm_router.get_model_ids(model_name=extracted_model_name) if model_ids and len(model_ids) > 0: # Use the first model ID found actual_model_id = model_ids[0] - verbose_proxy_logger.info( - f"Found model ID in router: {actual_model_id}" - ) + verbose_proxy_logger.info(f"Found model ID in router: {actual_model_id}") return actual_model_id else: # Fallback to constructed model name actual_model_id = extracted_model_name - verbose_proxy_logger.warning( - f"Model not found in router, using constructed name: {actual_model_id}" - ) + verbose_proxy_logger.warning(f"Model not found in router, using constructed name: {actual_model_id}") return actual_model_id else: # Fallback if router is not available - extracted_model_name = ( - VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path( - model_name - ) - ) - verbose_proxy_logger.warning( - f"Router not available, using constructed model name: {extracted_model_name}" - ) + extracted_model_name = VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path(model_name) + verbose_proxy_logger.warning(f"Router not available, using constructed model name: {extracted_model_name}") return extracted_model_name diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 9c0fbe30fc3..f6970c2a287 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -234,9 +234,7 @@ def _managed_id_matches_provider(unified_id: str, provider: str) -> bool: # /openai/v1/files -> /v1/files # /azure/openai/files -> /files (_canonical_path then prepends /v1/) # /azure_ai/openai/files -> /files -_PASSTHROUGH_PREFIX_RE = re.compile( - r"^/(?:azure(?:_ai)?/)?openai(?:_passthrough)?(?=/|$)" -) +_PASSTHROUGH_PREFIX_RE = re.compile(r"^/(?:azure(?:_ai)?/)?openai(?:_passthrough)?(?=/|$)") def _canonical_path(route: str) -> str: @@ -295,10 +293,7 @@ async def _resolve_one( if payload.provider != provider: raise HTTPException( status_code=404, - detail=( - f"Managed ID was minted for provider '{payload.provider}', " - f"not '{provider}'." - ), + detail=(f"Managed ID was minted for provider '{payload.provider}', not '{provider}'."), ) row_created_by: Optional[str] = None @@ -407,18 +402,13 @@ async def _guard_raw_provider_id( where={"flat_model_file_ids": {"has": raw_id}}, ) except Exception: - verbose_proxy_logger.debug( - "managed_id_rewriter: raw file-id guard lookup failed", exc_info=True - ) + verbose_proxy_logger.debug("managed_id_rewriter: raw file-id guard lookup failed", exc_info=True) return provider_rows = [ - row - for row in (candidates or []) - if _managed_id_matches_provider(row.unified_file_id, provider) + row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider) ] if provider_rows and not any( - can_access_resource(user_api_key_dict, row.created_by, row.team_id) - for row in provider_rows + can_access_resource(user_api_key_dict, row.created_by, row.team_id) for row in provider_rows ): raise HTTPException(status_code=404, detail="Managed resource not found.") return @@ -433,13 +423,9 @@ async def _guard_raw_provider_id( where={"model_object_id": f"passthrough:{provider}:{raw_id}"} ) except Exception: - verbose_proxy_logger.debug( - "managed_id_rewriter: raw object-id guard lookup failed", exc_info=True - ) + verbose_proxy_logger.debug("managed_id_rewriter: raw object-id guard lookup failed", exc_info=True) return - if existing is not None and not can_access_resource( - user_api_key_dict, existing.created_by, existing.team_id - ): + if existing is not None and not can_access_resource(user_api_key_dict, existing.created_by, existing.team_id): raise HTTPException(status_code=404, detail="Managed resource not found.") @@ -448,9 +434,7 @@ async def _guard_raw_provider_id( # --------------------------------------------------------------------------- -def _build_managed_file_object( - snapshot: Optional[Dict[str, Any]], managed_id: str -) -> Optional[OpenAIFileObject]: +def _build_managed_file_object(snapshot: Optional[Dict[str, Any]], managed_id: str) -> Optional[OpenAIFileObject]: """Build an ``OpenAIFileObject`` (with the managed ID swapped in) from an upstream file response so the DB-served list returns the same metadata as a direct file GET. Returns ``None`` when no usable snapshot is available, in @@ -461,8 +445,7 @@ def _build_managed_file_object( return OpenAIFileObject(**{**snapshot, "id": managed_id}) except Exception: verbose_proxy_logger.debug( - "managed_id_rewriter: file object snapshot incomplete; " - "storing file row without list metadata", + "managed_id_rewriter: file object snapshot incomplete; storing file row without list metadata", exc_info=True, ) return None @@ -502,20 +485,12 @@ async def _mint_or_reuse_file( ) except Exception: candidates = [] - verbose_proxy_logger.debug( - "managed_id_rewriter: file dedup lookup failed", exc_info=True - ) + verbose_proxy_logger.debug("managed_id_rewriter: file dedup lookup failed", exc_info=True) provider_rows = [ - row - for row in (candidates or []) - if _managed_id_matches_provider(row.unified_file_id, provider) + row for row in (candidates or []) if _managed_id_matches_provider(row.unified_file_id, provider) ] owned_row = next( - ( - row - for row in provider_rows - if can_access_resource(user_api_key_dict, row.created_by, row.team_id) - ), + (row for row in provider_rows if can_access_resource(user_api_key_dict, row.created_by, row.team_id)), None, ) if owned_row is not None: @@ -537,8 +512,7 @@ async def _mint_or_reuse_file( # different owner already holds (two upstream accounts under one # provider name); the file is the caller's, so leave it unmanaged. verbose_proxy_logger.debug( - "managed_id_rewriter: file dedup hit different owner on create; " - "leaving raw id unmanaged for prefix=%s", + "managed_id_rewriter: file dedup hit different owner on create; leaving raw id unmanaged for prefix=%s", raw_id.split("-", 1)[0], ) return raw_id @@ -553,15 +527,11 @@ async def _mint_or_reuse_file( try: await managed_files_hook.store_unified_file_id( file_id=managed_id, - file_object=_build_managed_file_object( - file_object_snapshot, managed_id - ), + file_object=_build_managed_file_object(file_object_snapshot, managed_id), litellm_parent_otel_span=None, model_mappings={ _passthrough_sentinel_model_id(provider): raw_id, - _PASSTHROUGH_PROVIDER_MARKER_KEY: _passthrough_provider_marker( - provider - ), + _PASSTHROUGH_PROVIDER_MARKER_KEY: _passthrough_provider_marker(provider), }, user_api_key_dict=user_api_key_dict, ) @@ -570,8 +540,7 @@ async def _mint_or_reuse_file( # back to the raw id (as when no persistence is available) to keep the # caller's freshly-created resource reachable rather than orphaned. verbose_proxy_logger.warning( - "managed_id_rewriter: could not persist file row; " - "leaving raw id unmanaged", + "managed_id_rewriter: could not persist file row; leaving raw id unmanaged", exc_info=True, ) return raw_id @@ -603,9 +572,7 @@ async def _mint_or_reuse_object( async def _reuse_existing(existing: Any, refresh_snapshot: bool) -> str: """Resolve an already-persisted namespaced row: enforce the access check, optionally refresh the snapshot, and return its managed ID.""" - if not can_access_resource( - user_api_key_dict, existing.created_by, existing.team_id - ): + if not can_access_resource(user_api_key_dict, existing.created_by, existing.team_id): if not is_create_route: # Retrieve / cancel / delete: the caller supplied a raw ID whose # managed row belongs to someone else. A raw ID only reaches the @@ -655,9 +622,7 @@ async def _mint_or_reuse_object( where={"model_object_id": namespaced_model_object_id} ) except Exception: - verbose_proxy_logger.debug( - "managed_id_rewriter: object dedup lookup failed", exc_info=True - ) + verbose_proxy_logger.debug("managed_id_rewriter: object dedup lookup failed", exc_info=True) existing = None if existing is not None: @@ -705,8 +670,7 @@ async def _mint_or_reuse_object( # back to the raw id (as when no persistence is available) to keep the # caller's freshly-created resource reachable rather than orphaned. verbose_proxy_logger.warning( - "managed_id_rewriter: could not persist object row; " - "leaving raw id unmanaged", + "managed_id_rewriter: could not persist object row; leaving raw id unmanaged", exc_info=True, ) return raw_id @@ -891,13 +855,9 @@ async def _build_list_where_with_cursor( if resource_kind == "files" else ManagedObjectRepository(prisma_client).table ) - cursor_field = ( - "unified_file_id" if resource_kind == "files" else "unified_object_id" - ) + cursor_field = "unified_file_id" if resource_kind == "files" else "unified_object_id" try: - cursor_row = await cursor_table.find_first( - where={**owner_filter, cursor_field: cursor_id} - ) + cursor_row = await cursor_table.find_first(where={**owner_filter, cursor_field: cursor_id}) if cursor_row is not None: if after_id: op = "lt" @@ -947,9 +907,7 @@ async def _fetch_list_rows( take=fetch_limit, ) except Exception: - verbose_proxy_logger.warning( - "managed_id_rewriter: list DB query failed", exc_info=True - ) + verbose_proxy_logger.warning("managed_id_rewriter: list DB query failed", exc_info=True) return None @@ -977,15 +935,11 @@ async def _fetch_provider_scoped_list_rows( """ scoped_where = dict(where) if resource_kind == "files": - scoped_where["flat_model_file_ids"] = { - "has": _passthrough_provider_marker(provider) - } + scoped_where["flat_model_file_ids"] = {"has": _passthrough_provider_marker(provider)} else: scoped_where["model_object_id"] = {"startswith": f"passthrough:{provider}:"} - rows = await _fetch_list_rows( - prisma_client, resource_kind, scoped_where, fetch_order, fetch_limit - ) + rows = await _fetch_list_rows(prisma_client, resource_kind, scoped_where, fetch_order, fetch_limit) if rows is None: return [], False @@ -1020,9 +974,7 @@ def _serialize_batch_list_item(row: Any) -> Dict[str, Any]: return item -def _list_boundary_ids( - rows: List[Any], resource_kind: str -) -> Tuple[Optional[str], Optional[str]]: +def _list_boundary_ids(rows: List[Any], resource_kind: str) -> Tuple[Optional[str], Optional[str]]: if not rows: return None, None id_attr = "unified_file_id" if resource_kind == "files" else "unified_object_id" @@ -1061,9 +1013,7 @@ async def list_passthrough_ids_from_db( owner_filter = build_owner_filter(user_api_key_dict) if owner_filter is None: - verbose_proxy_logger.warning( - "managed_id_rewriter: list denied — caller has no user_id or team_id" - ) + verbose_proxy_logger.warning("managed_id_rewriter: list denied — caller has no user_id or team_id") return _empty_list_response() raw_limit, fetch_limit = _parse_list_limit(query_params) @@ -1134,14 +1084,10 @@ async def rewrite_path_ids( new_segments.append(quote(raw, safe="-_.~")) changed = True else: - await _guard_raw_provider_id( - decoded_seg, provider, user_api_key_dict, prisma_client, budget - ) + await _guard_raw_provider_id(decoded_seg, provider, user_api_key_dict, prisma_client, budget) new_segments.append(seg) if changed: - verbose_proxy_logger.debug( - "managed_id_rewriter: path ids rewritten provider=%s", provider - ) + verbose_proxy_logger.debug("managed_id_rewriter: path ids rewritten provider=%s", provider) return "/".join(new_segments) if changed else path @@ -1164,14 +1110,10 @@ async def rewrite_query_ids( for key, val in list(mutated.items()): if isinstance(val, str): if is_managed(val): - mutated[key] = await _resolve_one( - val, provider, user_api_key_dict, prisma_client, managed_files_hook - ) + mutated[key] = await _resolve_one(val, provider, user_api_key_dict, prisma_client, managed_files_hook) rewritten_keys.append(key) else: - await _guard_raw_provider_id( - val, provider, user_api_key_dict, prisma_client, budget - ) + await _guard_raw_provider_id(val, provider, user_api_key_dict, prisma_client, budget) if rewritten_keys: verbose_proxy_logger.debug( "managed_id_rewriter: query ids rewritten provider=%s keys=%s", @@ -1221,18 +1163,12 @@ async def rewrite_body_ids( return node elif isinstance(node, str): if is_managed(node): - return await _resolve_one( - node, provider, user_api_key_dict, prisma_client, managed_files_hook - ) - await _guard_raw_provider_id( - node, provider, user_api_key_dict, prisma_client, budget - ) + return await _resolve_one(node, provider, user_api_key_dict, prisma_client, managed_files_hook) + await _guard_raw_provider_id(node, provider, user_api_key_dict, prisma_client, budget) return node return node rewritten = await _walk(body, 0) if rewritten is not body: - verbose_proxy_logger.debug( - "managed_id_rewriter: body ids rewritten provider=%s", provider - ) + verbose_proxy_logger.debug("managed_id_rewriter: body ids rewritten provider=%s", provider) return rewritten diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cea671e7510..391540f385e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -81,9 +81,7 @@ router = APIRouter() pass_through_endpoint_logging = PassThroughEndpointLogging() # Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: Dict[ - str, Dict[str, Union[str, bool, List[str], Dict[str, Any]]] -] = {} +_registered_pass_through_routes: Dict[str, Dict[str, Union[str, bool, List[str], Dict[str, Any]]]] = {} def get_response_body(response: httpx.Response) -> Optional[dict]: @@ -113,13 +111,9 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona # langfuse requires b64 encoded headers - we construct that here _langfuse_public_key = custom_headers["LANGFUSE_PUBLIC_KEY"] _langfuse_secret_key = custom_headers["LANGFUSE_SECRET_KEY"] - if isinstance( - _langfuse_public_key, str - ) and _langfuse_public_key.startswith("os.environ/"): + if isinstance(_langfuse_public_key, str) and _langfuse_public_key.startswith("os.environ/"): _langfuse_public_key = get_secret_str(_langfuse_public_key) - if isinstance( - _langfuse_secret_key, str - ) and _langfuse_secret_key.startswith("os.environ/"): + if isinstance(_langfuse_secret_key, str) and _langfuse_secret_key.startswith("os.environ/"): _langfuse_secret_key = get_secret_str(_langfuse_secret_key) headers["Authorization"] = "Basic " + b64encode( f"{_langfuse_public_key}:{_langfuse_secret_key}".encode("utf-8") @@ -128,9 +122,7 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona # for all other headers headers[key] = value if isinstance(value, str) and "os.environ/" in value: - verbose_proxy_logger.debug( - "pass through endpoint - looking up 'os.environ/' variable" - ) + verbose_proxy_logger.debug("pass through endpoint - looking up 'os.environ/' variable") # get string section that is os.environ/ start_index = value.find("os.environ/") _variable_name = value[start_index:] @@ -230,9 +222,7 @@ async def chat_completion_pass_through_endpoint( # skip router if user passed their key if "api_key" in data: llm_response = asyncio.create_task(litellm.aadapter_completion(**data)) - elif ( - llm_router is not None and data["model"] in router_model_names - ): # model in router model list + elif llm_router is not None and data["model"] in router_model_names: # model in router model list llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) elif ( llm_router is not None @@ -240,34 +230,24 @@ async def chat_completion_pass_through_endpoint( and data["model"] in llm_router.model_group_alias ): # model set in model_group_alias llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) - elif llm_router is not None and llm_router.has_model_id( - data["model"] - ): # model in router model list + elif llm_router is not None and llm_router.has_model_id(data["model"]): # model in router model list llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) elif ( llm_router is not None and data["model"] not in router_model_names - and ( - llm_router.default_deployment is not None - or len(llm_router.pattern_router.patterns) > 0 - ) + and (llm_router.default_deployment is not None or len(llm_router.pattern_router.patterns) > 0) ): # check for wildcard routes or default deployment before checking deployment_names llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) elif ( llm_router is not None and data["model"] in llm_router.deployment_names ): # model in router deployments, calling a specific deployment on the router (lowest priority) - llm_response = asyncio.create_task( - llm_router.aadapter_completion(**data, specific_deployment=True) - ) + llm_response = asyncio.create_task(llm_router.aadapter_completion(**data, specific_deployment=True)) elif user_model is not None: # `litellm --model ` llm_response = asyncio.create_task(litellm.aadapter_completion(**data)) else: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "completion: Invalid model name passed in model=" - + data.get("model", "") - }, + detail={"error": "completion: Invalid model name passed in model=" + data.get("model", "")}, ) # Await the llm_response task @@ -281,9 +261,7 @@ async def chat_completion_pass_through_endpoint( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) verbose_proxy_logger.debug("final response: %s", response) @@ -305,11 +283,7 @@ async def chat_completion_pass_through_endpoint( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.completion(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - {}".format(str(e))) error_msg = f"{str(e)}" raise ProxyException( message=getattr(e, "message", error_msg), @@ -341,19 +315,13 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "keep-alive", } - return_headers = { - key: value - for key, value in headers.items() - if key.lower() not in excluded_headers - } + return_headers = {key: value for key, value in headers.items() if key.lower() not in excluded_headers} if litellm_call_id: return_headers["x-litellm-call-id"] = litellm_call_id if custom_headers: # Ensure custom headers don't override actual upstream response headers or let framework defaults (like content-length: 0) interfere. sanitized_custom_headers = { - key: value - for key, value in custom_headers.items() - if key.lower() not in excluded_headers + key: value for key, value in custom_headers.items() if key.lower() not in excluded_headers } return_headers.update(sanitized_custom_headers) @@ -432,10 +400,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): headers=headers, params=requested_query_params, ) - elif ( - HttpPassThroughEndpointHelpers.is_multipart(request) is True - and forward_multipart - ): + elif HttpPassThroughEndpointHelpers.is_multipart(request) is True and forward_multipart: # Forward multipart via make_multipart_http_request even when _parsed_body is # non-empty (pass_through_request always injects litellm_logging_obj, etc.). # forward_multipart is False when custom_body was supplied (JSON body despite @@ -494,9 +459,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): files = [ ( field_name, - await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - upload_file=field_value - ), + await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(upload_file=field_value), ) for field_name, field_value in form_items if isinstance(field_value, (StarletteUploadFile, UploadFile)) @@ -509,9 +472,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): ) field_order = { field_name: index - for index, field_name in enumerate( - dict.fromkeys(field_name for field_name, _ in non_file_items) - ) + for index, field_name in enumerate(dict.fromkeys(field_name for field_name, _ in non_file_items)) } form_data_dict = { field_name: [value for _, value in group] @@ -569,9 +530,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): litellm_params_in_body[k] = _parsed_body.pop(k, None) _metadata = dict( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) ) litellm_metadata = litellm_params_in_body.pop("litellm_metadata", None) @@ -608,16 +567,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "passthrough_logging_payload": passthrough_logging_payload, } - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload return kwargs @staticmethod - def construct_target_url_with_subpath( - base_target: str, subpath: str, include_subpath: Optional[bool] - ) -> str: + def construct_target_url_with_subpath(base_target: str, subpath: str, include_subpath: Optional[bool]) -> str: """ Helper function to construct the full target URL with subpath handling. @@ -695,9 +650,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): return stream -def _carry_guardrail_logging_info( - request_data: dict, guardrail_data: Optional[dict] -) -> None: +def _carry_guardrail_logging_info(request_data: dict, guardrail_data: Optional[dict]) -> None: """Copy guardrail logging entries from ``guardrail_data`` onto ``request_data``. Post-call guardrails run against a throwaway ``hook_data`` dict (its @@ -814,9 +767,7 @@ async def pass_through_request( ).encode("ascii") ) - endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type( - str(url) - ) + endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type(str(url)) # SigV4-signed callers (e.g. Bedrock) attach the exact bytes that were # signed via request.state; we must send those instead of re-encoding the @@ -829,16 +780,12 @@ async def pass_through_request( if _request_state is not None else None ) - if state_raw_body is not None and not isinstance( - state_raw_body, (str, bytes, bytearray) - ): + if state_raw_body is not None and not isinstance(state_raw_body, (str, bytes, bytearray)): state_raw_body = None # Skip body parsing for multipart requests - make_multipart_http_request will handle it # But if custom_body is provided (e.g., JSON parsed despite multipart content-type), use it - is_multipart = ( - HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body - ) + is_multipart = HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body if custom_body: _parsed_body = custom_body @@ -869,16 +816,12 @@ async def pass_through_request( if "metadata" not in _parsed_body: _parsed_body["metadata"] = {} _parsed_body["metadata"]["guardrails"] = guardrails_to_run - verbose_proxy_logger.debug( - f"Added guardrails to passthrough request metadata: {guardrails_to_run}" - ) + verbose_proxy_logger.debug(f"Added guardrails to passthrough request metadata: {guardrails_to_run}") ## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it # Surface the requested model (when the body carries one) so logging/spans # read e.g. ``chat gpt-4o`` instead of ``chat unknown``. - passthrough_model = ( - _parsed_body.get("model") if isinstance(_parsed_body, dict) else None - ) or "unknown" + passthrough_model = (_parsed_body.get("model") if isinstance(_parsed_body, dict) else None) or "unknown" start_time = datetime.now() logging_obj = Logging( model=passthrough_model, @@ -928,9 +871,7 @@ async def pass_through_request( # Store custom_llm_provider in kwargs and logging object if provided if custom_llm_provider: logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider - logging_obj.model_call_details["litellm_params"] = kwargs.get( - "litellm_params", {} - ) + logging_obj.model_call_details["litellm_params"] = kwargs.get("litellm_params", {}) # done for supporting 'parallel_request_limiter.py' with pass-through endpoints logging_obj.update_environment_variables( @@ -943,9 +884,7 @@ async def pass_through_request( logging_obj.model_call_details["litellm_call_id"] = litellm_call_id # combine url with query params for logging - requested_query_params: Optional[dict] = query_params or dict( - request.query_params - ) + requested_query_params: Optional[dict] = query_params or dict(request.query_params) ## PASSTHROUGH MANAGED ID RESOLUTION (INPUT) ## # Resolve managed IDs in path, query params, and body back to raw @@ -956,22 +895,15 @@ async def pass_through_request( general_settings as proxy_general_settings, ) - _managed_id_provider = resolve_passthrough_managed_id_provider( - custom_llm_provider - ) + _managed_id_provider = resolve_passthrough_managed_id_provider(custom_llm_provider) - if ( - proxy_general_settings.get("passthrough_managed_object_ids", False) - and _managed_id_provider is not None - ): + if proxy_general_settings.get("passthrough_managed_object_ids", False) and _managed_id_provider is not None: verbose_proxy_logger.debug( "pass_through_endpoint: managed-id input rewrite enabled for route=%s method=%s", request.url.path, request.method, ) - _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook( - "managed_files" - ) + _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook("managed_files") if _passthrough_managed_hook is not None: from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( rewrite_body_ids, @@ -1045,9 +977,7 @@ async def pass_through_request( from litellm.proxy.proxy_server import prisma_client as _list_prisma if ( - is_passthrough_list_route( - _managed_id_provider, request.method, get_request_route(request) - ) + is_passthrough_list_route(_managed_id_provider, request.method, get_request_route(request)) and _list_prisma is not None ): _list_result = await list_passthrough_ids_from_db( @@ -1071,9 +1001,7 @@ async def pass_through_request( requested_query_params_str = None if requested_query_params: - requested_query_params_str = "&".join( - f"{k}={v}" for k, v in requested_query_params.items() - ) + requested_query_params_str = "&".join(f"{k}={v}" for k, v in requested_query_params.items()) logging_url = str(url) if requested_query_params_str: @@ -1091,11 +1019,9 @@ async def pass_through_request( "headers": headers, }, ) - stream = ( - HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( - parsed_body=_parsed_body or {}, - stream=stream, - ) + stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( + parsed_body=_parsed_body or {}, + stream=stream, ) if stream: @@ -1103,23 +1029,19 @@ async def pass_through_request( logging_obj.model_call_details["stream"] = True if is_multipart: - response = ( - await HttpPassThroughEndpointHelpers.make_multipart_http_request( - request=request, - async_client=async_client, - url=url, - headers=headers, - requested_query_params=requested_query_params, - stream=True, - ) + response = await HttpPassThroughEndpointHelpers.make_multipart_http_request( + request=request, + async_client=async_client, + url=url, + headers=headers, + requested_query_params=requested_query_params, + stream=True, ) else: # SigV4-signed callers (Bedrock) supply the exact pre-signed bytes; # otherwise httpx encodes the parsed JSON dict as before. body_kwargs: Dict[str, Any] = ( - {"content": state_raw_body} - if state_raw_body is not None - else {"json": _parsed_body} + {"content": state_raw_body} if state_raw_body is not None else {"json": _parsed_body} ) req = async_client.build_request( request.method, @@ -1134,9 +1056,7 @@ async def pass_through_request( try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) + raise HTTPException(status_code=e.response.status_code, detail=await e.response.aread()) # Call response headers hook for streaming pass-through _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( @@ -1177,16 +1097,14 @@ async def pass_through_request( content=state_raw_body, ) else: - response = ( - await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( - request=request, - async_client=async_client, - url=url, - headers=headers, - requested_query_params=requested_query_params, - _parsed_body=_parsed_body, - forward_multipart=is_multipart, - ) + response = await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( + request=request, + async_client=async_client, + url=url, + headers=headers, + requested_query_params=requested_query_params, + _parsed_body=_parsed_body, + forward_multipart=is_multipart, ) verbose_proxy_logger.debug("response.headers= %s", response.headers) @@ -1197,9 +1115,7 @@ async def pass_through_request( try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) + raise HTTPException(status_code=e.response.status_code, detail=await e.response.aread()) # Call response headers hook for detected streaming pass-through _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( @@ -1232,9 +1148,7 @@ async def pass_through_request( try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=e.response.text - ) + raise HTTPException(status_code=e.response.status_code, detail=e.response.text) if response.status_code >= 300: raise HTTPException(status_code=response.status_code, detail=response.text) @@ -1293,9 +1207,7 @@ async def pass_through_request( request.method, response.status_code, ) - _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook( - "managed_files" - ) + _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook("managed_files") if _passthrough_managed_hook is not None: from litellm.proxy.auth.auth_utils import get_request_route from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( @@ -1431,9 +1343,7 @@ async def pass_through_request( ) else: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format(str(e)) ) ######################################################### @@ -1448,11 +1358,7 @@ async def pass_through_request( if logging_obj is not None: request_payload["litellm_logging_obj"] = logging_obj - if ( - "model" not in request_payload - and _parsed_body - and isinstance(_parsed_body, dict) - ): + if "model" not in request_payload and _parsed_body and isinstance(_parsed_body, dict): request_payload["model"] = _parsed_body.get("model", "") if "custom_llm_provider" not in request_payload and custom_llm_provider: request_payload["custom_llm_provider"] = custom_llm_provider @@ -1643,25 +1549,18 @@ def create_pass_through_route( stream, ) = await _parse_request_data_by_content_type(request) - if not InitPassThroughEndpointHelpers.is_registered_pass_through_route( - route=path - ): + if not InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=path): raise HTTPException( status_code=404, detail=f"Pass-through endpoint {endpoint} not found. This could have been deleted or not yet added to the proxy.", ) - passthrough_params = ( - InitPassThroughEndpointHelpers.get_registered_pass_through_route( - route=path, method=request.method - ) + passthrough_params = InitPassThroughEndpointHelpers.get_registered_pass_through_route( + route=path, method=request.method ) if ( passthrough_params is None - and InitPassThroughEndpointHelpers.get_registered_pass_through_route( - route=path - ) - is not None + and InitPassThroughEndpointHelpers.get_registered_pass_through_route(route=path) is not None ): raise HTTPException( status_code=status.HTTP_405_METHOD_NOT_ALLOWED, @@ -1683,40 +1582,26 @@ def create_pass_through_route( # Extract and cast parameters with proper types param_target = target_params.get("target") or target param_custom_headers = target_params.get("custom_headers", custom_headers) - param_forward_headers = target_params.get( - "forward_headers", _forward_headers - ) - param_merge_query_params = target_params.get( - "merge_query_params", _merge_query_params - ) - param_cost_per_request = target_params.get( - "cost_per_request", cost_per_request - ) + param_forward_headers = target_params.get("forward_headers", _forward_headers) + param_merge_query_params = target_params.get("merge_query_params", _merge_query_params) + param_cost_per_request = target_params.get("cost_per_request", cost_per_request) param_guardrails = target_params.get("guardrails", None) param_default_query_params = target_params.get("default_query_params", None) param_timeout = target_params.get("timeout", timeout) # Construct the full target URL with subpath if needed - full_target = ( - HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( - base_target=cast(str, param_target), - subpath=subpath, - include_subpath=include_subpath, - ) + full_target = HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( + base_target=cast(str, param_target), + subpath=subpath, + include_subpath=include_subpath, ) # Ensure custom_headers is a dict. Botocore returns a HeadersDict # for SigV4-prepared requests, which is a Mapping but not a dict. - headers_dict = ( - dict(param_custom_headers) - if isinstance(param_custom_headers, Mapping) - else {} - ) + headers_dict = dict(param_custom_headers) if isinstance(param_custom_headers, Mapping) else {} # Ensure query_params and custom_body are dicts or None - final_query_params = ( - query_params_data if isinstance(query_params_data, dict) else {} - ) + final_query_params = query_params_data if isinstance(query_params_data, dict) else {} if query_params: final_query_params.update(query_params) # Programmatic callers set LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY on @@ -1741,9 +1626,7 @@ def create_pass_through_route( forward_headers=cast(Optional[bool], param_forward_headers), merge_query_params=cast(Optional[bool], param_merge_query_params), query_params=final_query_params, - default_query_params=cast( - Optional[dict], param_default_query_params - ), + default_query_params=cast(Optional[dict], param_default_query_params), stream=is_streaming_request or stream, custom_body=final_custom_body, cost_per_request=cast(Optional[float], param_cost_per_request), @@ -1844,16 +1727,12 @@ async def websocket_passthrough_request( websocket_messages: list[dict[str, Any]] = [] litellm_call_id = str(uuid.uuid4()) - verbose_proxy_logger.info( - f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}" - ) + verbose_proxy_logger.info(f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}") # Only accept the WebSocket if requested (for generic usage) if accept_websocket: await websocket.accept() - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): WebSocket connection accepted" - ) + verbose_proxy_logger.debug(f"WebSocket passthrough ({endpoint}): WebSocket connection accepted") # Prepare headers for the upstream connection upstream_headers = custom_headers.copy() @@ -1892,9 +1771,7 @@ async def websocket_passthrough_request( # Create a dummy request object for WebSocket connections to maintain compatibility # with the existing _init_kwargs_for_pass_through_endpoint function class DummyRequest: - def __init__( - self, url: str, method: str = "WEBSOCKET", headers: Optional[dict] = None - ): + def __init__(self, url: str, method: str = "WEBSOCKET", headers: Optional[dict] = None): self.url = url self.method = method self.headers = headers or {} @@ -1948,9 +1825,7 @@ async def websocket_passthrough_request( ) try: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}" - ) + verbose_proxy_logger.debug(f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}") async with connect( target, additional_headers=upstream_headers, @@ -1980,36 +1855,20 @@ async def websocket_passthrough_request( ) try: client_message = json.loads(text_data) - if ( - isinstance(client_message, dict) - and "setup" in client_message - ): + if isinstance(client_message, dict) and "setup" in client_message: setup_data = client_message["setup"] verbose_proxy_logger.debug( f"WebSocket passthrough ({endpoint}): Found setup data in client message: {setup_data}" ) - if ( - isinstance(setup_data, dict) - and "model" in setup_data - ): - extracted_model = ( - _extract_model_from_vertex_ai_setup( - setup_data - ) - ) + if isinstance(setup_data, dict) and "model" in setup_data: + extracted_model = _extract_model_from_vertex_ai_setup(setup_data) if extracted_model: kwargs["model"] = extracted_model - kwargs["custom_llm_provider"] = ( - "vertex_ai-language-models" - ) + kwargs["custom_llm_provider"] = "vertex_ai-language-models" # Update logging object with correct model logging_obj.model = extracted_model - logging_obj.model_call_details[ - "model" - ] = extracted_model - logging_obj.model_call_details[ - "custom_llm_provider" - ] = "vertex_ai" + logging_obj.model_call_details["model"] = extracted_model + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" verbose_proxy_logger.info( f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from client setup message" ) @@ -2058,18 +1917,14 @@ async def websocket_passthrough_request( verbose_proxy_logger.debug( f"WebSocket passthrough ({endpoint}): Processing server setup response for model extraction" ) - extracted_model = _extract_model_from_vertex_ai_setup( - setup_response - ) + extracted_model = _extract_model_from_vertex_ai_setup(setup_response) if extracted_model: kwargs["model"] = extracted_model kwargs["custom_llm_provider"] = "vertex_ai_language_models" # Update logging object with correct model logging_obj.model = extracted_model logging_obj.model_call_details["model"] = extracted_model - logging_obj.model_call_details["custom_llm_provider"] = ( - "vertex_ai_language_models" - ) + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models" verbose_proxy_logger.debug( f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response" ) @@ -2105,19 +1960,13 @@ async def websocket_passthrough_request( pass except (ConnectionClosedOK, ConnectionClosedError) as e: - verbose_proxy_logger.debug( - f"Upstream WebSocket connection closed: {e}" - ) + verbose_proxy_logger.debug(f"Upstream WebSocket connection closed: {e}") pass except asyncio.CancelledError: - verbose_proxy_logger.debug( - "asyncio.CancelledError in forward_upstream_to_client" - ) + verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client") raise except Exception as e: - verbose_proxy_logger.debug( - f"Exception in forward_upstream_to_client: {e}" - ) + verbose_proxy_logger.debug(f"Exception in forward_upstream_to_client: {e}") verbose_proxy_logger.exception( f"WebSocket passthrough ({endpoint}): error forwarding upstream message" ) @@ -2129,9 +1978,7 @@ async def websocket_passthrough_request( asyncio.create_task(forward_upstream_to_client()), ] - done, pending = await asyncio.wait( - tasks, return_when=asyncio.FIRST_COMPLETED - ) + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) # Cancel remaining tasks for task in pending: @@ -2212,9 +2059,7 @@ async def websocket_passthrough_request( ) except InvalidStatus as exc: - verbose_proxy_logger.exception( - f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection" - ) + verbose_proxy_logger.exception(f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection") # Prepare request payload for logging request_payload = {} @@ -2240,9 +2085,7 @@ async def websocket_passthrough_request( reason="Upstream connection rejected", ) except Exception as e: - verbose_proxy_logger.exception( - f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket" - ) + verbose_proxy_logger.exception(f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket") # Prepare request payload for logging request_payload = {} @@ -2554,9 +2397,7 @@ class InitPassThroughEndpointHelpers: """Remove all routes for a specific endpoint ID from the registry and clean up corresponding entries from LiteLLMRoutes.openai_routes.""" keys_to_remove = [ - key - for key, value in _registered_pass_through_routes.items() - if value["endpoint_id"] == endpoint_id + key for key, value in _registered_pass_through_routes.items() if value["endpoint_id"] == endpoint_id ] for key in keys_to_remove: route_info = _registered_pass_through_routes[key] @@ -2570,9 +2411,7 @@ class InitPassThroughEndpointHelpers: if wildcard_path in openai_routes: openai_routes.remove(wildcard_path) del _registered_pass_through_routes[key] - verbose_proxy_logger.debug( - "Removed pass-through route from registry: %s", key - ) + verbose_proxy_logger.debug("Removed pass-through route from registry: %s", key) @staticmethod def clear_all_pass_through_routes(): @@ -2617,9 +2456,7 @@ class InitPassThroughEndpointHelpers: if normalized_route.startswith(mapped_route): return True - comparison_route = InitPassThroughEndpointHelpers._route_for_registry_lookup( - route - ) + comparison_route = InitPassThroughEndpointHelpers._route_for_registry_lookup(route) # Fast path: check if any registered route key contains this path # Keys are in format: "{endpoint_id}:exact:{path}:{methods}" or "{endpoint_id}:subpath:{path}:{methods}" @@ -2633,22 +2470,15 @@ class InitPassThroughEndpointHelpers: if route_type == "exact" and comparison_route == registered_path: return True elif route_type == "subpath": - if ( - comparison_route == registered_path - or comparison_route.startswith(registered_path + "/") - ): + if comparison_route == registered_path or comparison_route.startswith(registered_path + "/"): return True return False @staticmethod - def get_registered_pass_through_route( - route: str, method: Optional[str] = None - ) -> Optional[Dict[str, Any]]: + def get_registered_pass_through_route(route: str, method: Optional[str] = None) -> Optional[Dict[str, Any]]: """Get passthrough params for a given route and optionally filter by HTTP method""" - comparison_route = InitPassThroughEndpointHelpers._route_for_registry_lookup( - route - ) + comparison_route = InitPassThroughEndpointHelpers._route_for_registry_lookup(route) for key in _registered_pass_through_routes.keys(): parts = key.split(":", 3) # Split into [endpoint_id, type, path, methods?] if len(parts) >= 3: @@ -2659,9 +2489,7 @@ class InitPassThroughEndpointHelpers: # but keep supporting test fixtures / older registry entries that # only encoded methods in the route key. methods_entry = _registered_pass_through_routes[key].get("methods", []) - route_methods: List[str] = ( - methods_entry if isinstance(methods_entry, list) else [] - ) + route_methods: List[str] = methods_entry if isinstance(methods_entry, list) else [] if not route_methods and len(parts) == 4: route_methods = parts[3].split(",") @@ -2670,10 +2498,7 @@ class InitPassThroughEndpointHelpers: if route_type == "exact" and comparison_route == registered_path: path_matches = True elif route_type == "subpath": - if ( - comparison_route == registered_path - or comparison_route.startswith(registered_path + "/") - ): + if comparison_route == registered_path or comparison_route.startswith(registered_path + "/"): path_matches = True # If path matches and method filter is provided, check if method is allowed @@ -2714,9 +2539,7 @@ async def _register_pass_through_endpoint( if path is None: raise ValueError("Path is required for pass-through endpoint") - custom_headers = await set_env_variables_in_header( - custom_headers=endpoint_data.get("headers") - ) + custom_headers = await set_env_variables_in_header(custom_headers=endpoint_data.get("headers")) forward_headers = endpoint_data.get("forward_headers") merge_query_params = endpoint_data.get("merge_query_params") default_query_params = endpoint_data.get("default_query_params") @@ -2741,9 +2564,7 @@ async def _register_pass_through_endpoint( cost_per_request = endpoint_data.get("cost_per_request") timeout = endpoint_data.get("timeout") - verbose_proxy_logger.debug( - "Initializing pass through endpoint: %s (ID: %s)", path, endpoint_id - ) + verbose_proxy_logger.debug("Initializing pass through endpoint: %s (ID: %s)", path, endpoint_id) InitPassThroughEndpointHelpers.add_exact_path_route( app=app, path=path, @@ -2790,9 +2611,7 @@ async def _register_pass_through_endpoint( ) visited_endpoints.add(f"{endpoint_id}:subpath:{path}:{methods_str}") - verbose_proxy_logger.debug( - "Added new pass through endpoint: %s (ID: %s)", path, endpoint_id - ) + verbose_proxy_logger.debug("Added new pass through endpoint: %s (ID: %s)", path, endpoint_id) async def initialize_pass_through_endpoints( @@ -2841,9 +2660,7 @@ async def initialize_pass_through_endpoints( # get a list of all registered pass-through endpoints # mark the ones that are visited in the list # remove the ones that are not visited from the list - registered_pass_through_endpoints = ( - InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() - ) + registered_pass_through_endpoints = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() visited_endpoints: set[str] = set() @@ -2979,10 +2796,7 @@ async def _filter_endpoints_by_team_allowed_routes( # retrieve team metadata team_metadata = team.metadata - if ( - team_metadata is not None - and team_metadata.get("allowed_passthrough_routes") is not None - ): + if team_metadata is not None and team_metadata.get("allowed_passthrough_routes") is not None: ## FILTER pass_through_endpoints by allowed_passthrough_routes pass_through_endpoints = [ endpoint @@ -3097,11 +2911,7 @@ async def update_pass_through_endpoints( # Find the index for updating the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = ( - PassThroughGenericEndpoint(**endpoint) - if isinstance(endpoint, dict) - else endpoint - ) + _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -3109,18 +2919,14 @@ async def update_pass_through_endpoints( if endpoint_index is None: raise HTTPException( status_code=404, - detail={ - "error": f"Could not find index for endpoint with ID '{endpoint_id}'" - }, + detail={"error": f"Could not find index for endpoint with ID '{endpoint_id}'"}, ) # Only merge fields the caller explicitly sent so omitted fields keep their # stored value. Without exclude_unset, defaults like auth=True would overwrite # an existing auth=false entry on any unrelated edit. # Exclude is_from_config as it's a response-only field (computed at read time) - update_data = data.model_dump( - exclude_unset=True, exclude_none=True, exclude={"is_from_config"} - ) + update_data = data.model_dump(exclude_unset=True, exclude_none=True, exclude={"is_from_config"}) # Start with existing endpoint data endpoint_dict = found_endpoint.model_dump() @@ -3151,9 +2957,7 @@ async def update_pass_through_endpoints( config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) # Re-register the route with updated headers _custom_headers: Optional[dict] = updated_endpoint.headers or {} @@ -3194,9 +2998,7 @@ async def update_pass_through_endpoints( timeout=updated_endpoint.timeout, ) - return PassThroughEndpointResponse( - endpoints=[updated_endpoint] if updated_endpoint else [] - ) + return PassThroughEndpointResponse(endpoints=[updated_endpoint] if updated_endpoint else []) @router.post( @@ -3224,9 +3026,7 @@ async def create_pass_through_endpoints( field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict ) except Exception: - response = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Auto-generate ID if not provided # Exclude is_from_config as it's a response-only field (computed at read time) @@ -3245,9 +3045,7 @@ async def create_pass_through_endpoints( field_value=response.field_value, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) # Return the created endpoint with the generated ID created_endpoint = PassThroughGenericEndpoint(**data_dict) @@ -3320,9 +3118,7 @@ async def delete_pass_through_endpoints( field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict ) except Exception: - response = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Update field by removing endpoint pass_through_endpoint_data: Optional[List] = response.field_value @@ -3338,21 +3134,13 @@ async def delete_pass_through_endpoints( if found_endpoint is None: raise HTTPException( status_code=400, - detail={ - "error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format( - endpoint_id - ) - }, + detail={"error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format(endpoint_id)}, ) # Find the index for deleting from the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = ( - PassThroughGenericEndpoint(**endpoint) - if isinstance(endpoint, dict) - else endpoint - ) + _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -3360,9 +3148,7 @@ async def delete_pass_through_endpoints( if endpoint_index is None: raise HTTPException( status_code=400, - detail={ - "error": f"Could not find index for endpoint with ID '{endpoint_id}'" - }, + detail={"error": f"Could not find index for endpoint with ID '{endpoint_id}'"}, ) # Remove the endpoint @@ -3378,9 +3164,7 @@ async def delete_pass_through_endpoints( field_value=pass_through_endpoint_data, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) return PassThroughEndpointResponse(endpoints=[response_obj]) @@ -3418,6 +3202,4 @@ async def initialize_pass_through_endpoints_in_db(): Gets all pass-through endpoints from db and initializes them in the proxy server. """ pass_through_endpoints = await _get_pass_through_endpoints_from_db() - await initialize_pass_through_endpoints( - pass_through_endpoints=pass_through_endpoints - ) + await initialize_pass_through_endpoints(pass_through_endpoints=pass_through_endpoints) diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index a32659e45bd..7f62a822dee 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -17,9 +17,7 @@ class PassthroughEndpointRouter: def __init__(self): self.credentials: Dict[str, str] = {} - self.deployment_key_to_vertex_credentials: Dict[ - str, VertexPassThroughCredentials - ] = {} + self.deployment_key_to_vertex_credentials: Dict[str, VertexPassThroughCredentials] = {} self.default_vertex_config: Optional[VertexPassThroughCredentials] = None def set_pass_through_credentials( @@ -38,9 +36,7 @@ class PassthroughEndpointRouter: """ credential_name = self._get_credential_name_for_provider( custom_llm_provider=custom_llm_provider, - region_name=self._get_region_name_from_api_base( - api_base=api_base, custom_llm_provider=custom_llm_provider - ), + region_name=self._get_region_name_from_api_base(api_base=api_base, custom_llm_provider=custom_llm_provider), ) if api_key is None: raise ValueError("api_key is required for setting pass-through credentials") @@ -55,20 +51,14 @@ class PassthroughEndpointRouter: custom_llm_provider=custom_llm_provider, region_name=region_name, ) - verbose_router_logger.debug( - f"Pass-through llm endpoints router, looking for credentials for {credential_name}" - ) + verbose_router_logger.debug(f"Pass-through llm endpoints router, looking for credentials for {credential_name}") if credential_name in self.credentials: verbose_router_logger.debug(f"Found credentials for {credential_name}") return self.credentials[credential_name] else: - verbose_router_logger.debug( - f"No credentials found for {credential_name}, looking for env variable" - ) - _env_variable_name = ( - self._get_default_env_variable_name_passthrough_endpoint( - custom_llm_provider=custom_llm_provider, - ) + verbose_router_logger.debug(f"No credentials found for {credential_name}, looking for env variable") + _env_variable_name = self._get_default_env_variable_name_passthrough_endpoint( + custom_llm_provider=custom_llm_provider, ) return get_secret_str(_env_variable_name) @@ -125,22 +115,16 @@ class PassthroughEndpointRouter: location=location, ) if deployment_key is None: - verbose_router_logger.debug( - "No deployment key found for project-id, location" - ) + verbose_router_logger.debug("No deployment key found for project-id, location") return vertex_pass_through_credentials = VertexPassThroughCredentials( vertex_project=project_id, vertex_location=location, vertex_credentials=vertex_credentials, ) - self.deployment_key_to_vertex_credentials[deployment_key] = ( - vertex_pass_through_credentials - ) + self.deployment_key_to_vertex_credentials[deployment_key] = vertex_pass_through_credentials - def _get_deployment_key( - self, project_id: Optional[str], location: Optional[str] - ) -> Optional[str]: + def _get_deployment_key(self, project_id: Optional[str], location: Optional[str]) -> Optional[str]: """ Get the deployment key for the given project-id, location """ @@ -148,9 +132,7 @@ class PassthroughEndpointRouter: return None return f"{project_id}-{location}" - def get_vector_store_credentials( - self, vector_store_id: str - ) -> Optional[LiteLLM_ManagedVectorStore]: + def get_vector_store_credentials(self, vector_store_id: str) -> Optional[LiteLLM_ManagedVectorStore]: """ Get the vector store credentials for the given vector store id """ diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index 5683491fedc..662c921bd4b 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -179,17 +179,11 @@ class PassthroughGuardrailHandler: HTTPException if a guardrail blocks the request """ if not PassthroughGuardrailHandler.is_enabled(guardrails_config): - verbose_proxy_logger.debug( - "Passthrough guardrails not enabled, skipping guardrail execution" - ) + verbose_proxy_logger.debug("Passthrough guardrails not enabled, skipping guardrail execution") return request_data - guardrail_names = PassthroughGuardrailHandler.get_guardrail_names( - guardrails_config - ) - verbose_proxy_logger.debug( - "Executing passthrough guardrails: %s", guardrail_names - ) + guardrail_names = PassthroughGuardrailHandler.get_guardrail_names(guardrails_config) + verbose_proxy_logger.debug("Executing passthrough guardrails: %s", guardrail_names) # Add to request metadata so guardrails know which to run from litellm.proxy.pass_through_endpoints.passthrough_context import ( @@ -200,9 +194,7 @@ class PassthroughGuardrailHandler: request_data["metadata"] = {} # Set guardrails in metadata using dict format for compatibility - request_data["metadata"]["guardrails"] = { - name: True for name in guardrail_names - } + request_data["metadata"]["guardrails"] = {name: True for name in guardrail_names} # Store passthrough guardrails config in request-scoped context set_passthrough_guardrails_config(guardrails_config) @@ -240,20 +232,14 @@ class PassthroughGuardrailHandler: ) # Normalize config to dict format (handles both list and dict) - normalized_config = PassthroughGuardrailHandler.normalize_config( - passthrough_guardrails_config - ) + normalized_config = PassthroughGuardrailHandler.normalize_config(passthrough_guardrails_config) if normalized_config is None: - verbose_proxy_logger.debug( - "Passthrough guardrails not configured, skipping guardrail collection" - ) + verbose_proxy_logger.debug("Passthrough guardrails not configured, skipping guardrail collection") return None if len(normalized_config) == 0: - verbose_proxy_logger.debug( - "Passthrough guardrails config is empty, skipping" - ) + verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None # Passthrough is enabled - collect guardrails @@ -262,9 +248,7 @@ class PassthroughGuardrailHandler: # Add passthrough-specific guardrails for guardrail_name in normalized_config.keys(): guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug( - "Added passthrough-specific guardrail: %s", guardrail_name - ) + verbose_proxy_logger.debug("Added passthrough-specific guardrail: %s", guardrail_name) # Add org/team/key level guardrails using shared helper temp_data: Dict[str, Any] = {"metadata": {}} @@ -280,9 +264,7 @@ class PassthroughGuardrailHandler: for guardrail_name in inherited_guardrails: if guardrail_name not in guardrails_to_run: guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug( - "Added inherited guardrail (key/team level): %s", guardrail_name - ) + verbose_proxy_logger.debug("Added inherited guardrail (key/team level): %s", guardrail_name) verbose_proxy_logger.debug( "Collected guardrails for passthrough endpoint: %s", @@ -319,9 +301,7 @@ class PassthroughGuardrailHandler: if passthrough_config is None: return None - settings = PassthroughGuardrailHandler.get_settings( - passthrough_config, guardrail_name - ) + settings = PassthroughGuardrailHandler.get_settings(passthrough_config, guardrail_name) if settings is None: return None diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 3a5728fd66f..1bdd57507e3 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -105,9 +105,7 @@ class PassThroughStreamingHandler: ) ) except Exception as e: - verbose_proxy_logger.error( - f"Error scheduling chunk_processor logging: {str(e)}" - ) + verbose_proxy_logger.error(f"Error scheduling chunk_processor logging: {str(e)}") @staticmethod async def _route_streaming_logging_to_handler( @@ -157,9 +155,7 @@ class PassThroughStreamingHandler: **kwargs, ) except Exception as e: - verbose_proxy_logger.error( - f"Error in _route_streaming_logging_to_handler: {str(e)}" - ) + verbose_proxy_logger.error(f"Error in _route_streaming_logging_to_handler: {str(e)}") @staticmethod def _build_passthrough_logging_result( @@ -180,27 +176,23 @@ class PassThroughStreamingHandler: be unit-tested in isolation. Still invoked synchronously on the event loop; an off-loop dispatch is a future change, not part of this PR. """ - all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines( - raw_bytes - ) - standard_logging_response_object: Optional[ - PassThroughEndpointLoggingResultValues - ] = None + all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) + standard_logging_response_object: Optional[PassThroughEndpointLoggingResultValues] = None kwargs: dict = {} if endpoint_type == EndpointType.ANTHROPIC: - anthropic_passthrough_logging_handler_result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body, - endpoint_type=endpoint_type, - start_time=start_time, - all_chunks=all_chunks, - end_time=end_time, - ) - standard_logging_response_object = ( - anthropic_passthrough_logging_handler_result["result"] + anthropic_passthrough_logging_handler_result = ( + AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + ) ) + standard_logging_response_object = anthropic_passthrough_logging_handler_result["result"] kwargs = anthropic_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.VERTEX_AI: vertex_passthrough_logging_handler_result = ( @@ -216,9 +208,7 @@ class PassThroughStreamingHandler: model=model, ) ) - standard_logging_response_object = ( - vertex_passthrough_logging_handler_result["result"] - ) + standard_logging_response_object = vertex_passthrough_logging_handler_result["result"] kwargs = vertex_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.OPENAI: openai_passthrough_logging_handler_result = ( @@ -233,9 +223,7 @@ class PassThroughStreamingHandler: end_time=end_time, ) ) - standard_logging_response_object = ( - openai_passthrough_logging_handler_result["result"] - ) + standard_logging_response_object = openai_passthrough_logging_handler_result["result"] kwargs = openai_passthrough_logging_handler_result["kwargs"] if standard_logging_response_object is None: diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 46043d10a06..ee651a15afe 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -54,9 +54,7 @@ class PassThroughEndpointLogging: # Cohere self.TRACKED_COHERE_ROUTES = ["/v2/chat", "/v1/embed"] - self.assemblyai_passthrough_logging_handler = ( - AssemblyAIPassthroughLoggingHandler() - ) + self.assemblyai_passthrough_logging_handler = AssemblyAIPassthroughLoggingHandler() # Langfuse self.TRACKED_LANGFUSE_ROUTES = ["/langfuse/"] @@ -102,11 +100,7 @@ class PassThroughEndpointLogging: # async-only loggers (e.g. the proxy spend logger) firing regardless of # how the call-type classification evolves. await logging_obj.dispatch_success_handlers( - result=( - json.dumps(result) - if isinstance(result, dict) - else standard_logging_response_object - ), + result=(json.dumps(result) if isinstance(result, dict) else standard_logging_response_object), start_time=start_time, end_time=end_time, cache_hit=False, @@ -135,41 +129,33 @@ class PassThroughEndpointLogging: standard_logging_response_object: Optional[Any] = None if self.is_gemini_route(url_route, custom_llm_provider): - gemini_passthrough_logging_handler_result = ( - GeminiPassthroughLoggingHandler.gemini_passthrough_handler( - httpx_response=httpx_response, - response_body=response_body or {}, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - gemini_passthrough_logging_handler_result["result"] + gemini_passthrough_logging_handler_result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) + standard_logging_response_object = gemini_passthrough_logging_handler_result["result"] kwargs = gemini_passthrough_logging_handler_result["kwargs"] elif self.is_vertex_route(url_route): - vertex_passthrough_logging_handler_result = ( - VertexPassthroughLoggingHandler.vertex_passthrough_handler( - httpx_response=httpx_response, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - vertex_passthrough_logging_handler_result["result"] + vertex_passthrough_logging_handler_result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=httpx_response, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) + standard_logging_response_object = vertex_passthrough_logging_handler_result["result"] kwargs = vertex_passthrough_logging_handler_result["kwargs"] elif self.is_anthropic_route(url_route): anthropic_passthrough_logging_handler_result = ( @@ -187,73 +173,57 @@ class PassThroughEndpointLogging: ) ) - standard_logging_response_object = ( - anthropic_passthrough_logging_handler_result["result"] - ) + standard_logging_response_object = anthropic_passthrough_logging_handler_result["result"] kwargs = anthropic_passthrough_logging_handler_result["kwargs"] elif self.is_cohere_route(url_route): - cohere_passthrough_logging_handler_result = ( - cohere_passthrough_logging_handler.cohere_passthrough_handler( - httpx_response=httpx_response, - response_body=response_body or {}, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - cohere_passthrough_logging_handler_result["result"] + cohere_passthrough_logging_handler_result = cohere_passthrough_logging_handler.cohere_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) + standard_logging_response_object = cohere_passthrough_logging_handler_result["result"] kwargs = cohere_passthrough_logging_handler_result["kwargs"] - elif self.is_openai_route(url_route) and self._is_supported_openai_endpoint( - url_route - ): + elif self.is_openai_route(url_route) and self._is_supported_openai_endpoint(url_route): from .llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) - openai_passthrough_logging_handler_result = ( - OpenAIPassthroughLoggingHandler.openai_passthrough_handler( - httpx_response=httpx_response, - response_body=response_body or {}, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - openai_passthrough_logging_handler_result["result"] + openai_passthrough_logging_handler_result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) + standard_logging_response_object = openai_passthrough_logging_handler_result["result"] kwargs = openai_passthrough_logging_handler_result["kwargs"] elif self.is_cursor_route(url_route, custom_llm_provider): - cursor_passthrough_logging_handler_result = ( - CursorPassthroughLoggingHandler.cursor_passthrough_handler( - httpx_response=httpx_response, - response_body=response_body or {}, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - cursor_passthrough_logging_handler_result["result"] + cursor_passthrough_logging_handler_result = CursorPassthroughLoggingHandler.cursor_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, ) + standard_logging_response_object = cursor_passthrough_logging_handler_result["result"] kwargs = cursor_passthrough_logging_handler_result["kwargs"] elif self.is_vertex_ai_live_route(url_route): from .llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import ( @@ -263,27 +233,21 @@ class PassThroughEndpointLogging: vertex_ai_live_handler = VertexAILivePassthroughLoggingHandler() # For WebSocket responses, response_body should be a list of messages - websocket_messages: list[dict[str, Any]] = ( - response_body if isinstance(response_body, list) else [] - ) + websocket_messages: list[dict[str, Any]] = response_body if isinstance(response_body, list) else [] - vertex_ai_live_handler_result = ( - vertex_ai_live_handler.vertex_ai_live_passthrough_handler( - websocket_messages=websocket_messages, - logging_obj=logging_obj, - url_route=url_route, - start_time=start_time, - end_time=end_time, - request_body=request_body, - **kwargs, - ) + vertex_ai_live_handler_result = vertex_ai_live_handler.vertex_ai_live_passthrough_handler( + websocket_messages=websocket_messages, + logging_obj=logging_obj, + url_route=url_route, + start_time=start_time, + end_time=end_time, + request_body=request_body, + **kwargs, ) standard_logging_response_object = vertex_ai_live_handler_result["result"] kwargs = vertex_ai_live_handler_result["kwargs"] - return_dict["standard_logging_response_object"] = ( - standard_logging_response_object - ) + return_dict["standard_logging_response_object"] = standard_logging_response_object return_dict["kwargs"] = kwargs return return_dict @@ -303,19 +267,10 @@ class PassThroughEndpointLogging: custom_llm_provider: Optional[str] = None, **kwargs, ): - standard_logging_response_object: Optional[ - PassThroughEndpointLoggingResultValues - ] = None - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + standard_logging_response_object: Optional[PassThroughEndpointLoggingResultValues] = None + logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload if self.is_assemblyai_route(url_route): - if ( - AssemblyAIPassthroughLoggingHandler._should_log_request( - httpx_response.request.method - ) - is not True - ): + if AssemblyAIPassthroughLoggingHandler._should_log_request(httpx_response.request.method) is not True: return self.assemblyai_passthrough_logging_handler.assemblyai_passthrough_logging_handler( httpx_response=httpx_response, @@ -333,31 +288,25 @@ class PassThroughEndpointLogging: # Don't log langfuse pass-through requests return else: - normalized_llm_passthrough_logging_payload = ( - self.normalize_llm_passthrough_logging_payload( - httpx_response=httpx_response, - response_body=response_body, - request_body=request_body, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - ) - standard_logging_response_object = ( - normalized_llm_passthrough_logging_payload[ - "standard_logging_response_object" - ] + normalized_llm_passthrough_logging_payload = self.normalize_llm_passthrough_logging_payload( + httpx_response=httpx_response, + response_body=response_body, + request_body=request_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + custom_llm_provider=custom_llm_provider, + **kwargs, ) + standard_logging_response_object = normalized_llm_passthrough_logging_payload[ + "standard_logging_response_object" + ] kwargs = normalized_llm_passthrough_logging_payload["kwargs"] if standard_logging_response_object is None: - standard_logging_response_object = StandardPassThroughResponseObject( - response=httpx_response.text - ) + standard_logging_response_object = StandardPassThroughResponseObject(response=httpx_response.text) kwargs = self._set_cost_per_request( logging_obj=logging_obj, @@ -417,9 +366,7 @@ class PassThroughEndpointLogging: return True return False - def is_cursor_route( - self, url_route: str, custom_llm_provider: Optional[str] = None - ): + def is_cursor_route(self, url_route: str, custom_llm_provider: Optional[str] = None): """Check if the URL route is a Cursor Cloud Agents API route.""" if custom_llm_provider == "cursor": return True @@ -448,9 +395,7 @@ class PassThroughEndpointLogging: return _is_openai_compatible_url(url_route) - def is_gemini_route( - self, url_route: str, custom_llm_provider: Optional[str] = None - ): + def is_gemini_route(self, url_route: str, custom_llm_provider: Optional[str] = None): """Check if the URL route is a Gemini API route.""" for route in self.TRACKED_GEMINI_ROUTES: if route in url_route and custom_llm_provider == "gemini": @@ -474,9 +419,7 @@ class PassThroughEndpointLogging: return ( OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) - or OpenAIPassthroughLoggingHandler.is_openai_image_generation_route( - url_route - ) + or OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) or OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) or OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) ) @@ -497,11 +440,7 @@ class PassThroughEndpointLogging: # Check if cost per request is set ######################################################### if passthrough_logging_payload.get("cost_per_request") is not None: - kwargs["response_cost"] = passthrough_logging_payload.get( - "cost_per_request" - ) - logging_obj.model_call_details["response_cost"] = ( - passthrough_logging_payload.get("cost_per_request") - ) + kwargs["response_cost"] = passthrough_logging_payload.get("cost_per_request") + logging_obj.model_call_details["response_cost"] = passthrough_logging_payload.get("cost_per_request") return kwargs diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py index 6a94f78fbe7..b6bd594096f 100644 --- a/litellm/proxy/plugin_routes.py +++ b/litellm/proxy/plugin_routes.py @@ -80,11 +80,7 @@ def _request_strip_headers() -> frozenset[str]: x-litellm-api-key, and any configured custom key header — so a plugin can never be handed the caller's live litellm key (confused-deputy escalation). """ - return ( - _HOP_BY_HOP_STRIP - | SpecialHeaders.litellm_credential_header_names() - | _configured_key_header_names() - ) + return _HOP_BY_HOP_STRIP | SpecialHeaders.litellm_credential_header_names() | _configured_key_header_names() # Headers to strip from plugin RESPONSES before returning to the browser. @@ -142,9 +138,7 @@ def _plugin_fernet(plugin_name: str) -> Fernet: _CLAIM_TTL_SECONDS = 30 # identity claims expire after 30 s -def issue_plugin_session_claim( - plugin_name: str, user_id: str | None, user_role: str | None -) -> str: +def issue_plugin_session_claim(plugin_name: str, user_id: str | None, user_role: str | None) -> str: """Issue a short-lived, audience-scoped identity claim for the plugin. The claim contains {user_id, user_role, plugin, exp}. Crucially it @@ -167,9 +161,7 @@ def verify_plugin_session_claim(plugin_name: str, ciphertext: str) -> dict: the claim is expired. Returns the decoded claim dict on success. """ try: - raw = _plugin_fernet(plugin_name).decrypt( - ciphertext.encode(), ttl=_CLAIM_TTL_SECONDS - ) + raw = _plugin_fernet(plugin_name).decrypt(ciphertext.encode(), ttl=_CLAIM_TTL_SECONDS) claim = json.loads(raw) except (InvalidToken, Exception) as exc: raise ValueError("Invalid, tampered, or expired plugin session claim") from exc @@ -192,9 +184,7 @@ def register_plugins_from_config(general_settings: dict[str, object]) -> None: """ raw = general_settings.get("plugins") entries: list[object] = raw if isinstance(raw, list) else [] - new_registry = { - p.name: p for p in (PluginConfig.model_validate(entry) for entry in entries) - } + new_registry = {p.name: p for p in (PluginConfig.model_validate(entry) for entry in entries)} _plugin_registry.clear() _plugin_registry.update(new_registry) @@ -245,14 +235,10 @@ async def plugin_auth_token( detail="LITELLM_SALT_KEY is not configured; plugin iframe auth unavailable.", ) if plugin_name not in _plugin_registry: - raise HTTPException( - status_code=404, detail=f"Plugin '{plugin_name}' is not registered." - ) + raise HTTPException(status_code=404, detail=f"Plugin '{plugin_name}' is not registered.") user_id = getattr(user_api_key_dict, "user_id", None) user_role = getattr(user_api_key_dict, "user_role", None) - return { - "session_claim": issue_plugin_session_claim(plugin_name, user_id, user_role) - } + return {"session_claim": issue_plugin_session_claim(plugin_name, user_id, user_role)} @router.api_route( @@ -299,9 +285,7 @@ async def plugin_proxy( # Strip caller credentials and hop-by-hop headers from forwarded request strip = _request_strip_headers() - forward_headers = { - k: v for k, v in request.headers.items() if k.lower() not in strip - } + forward_headers = {k: v for k, v in request.headers.items() if k.lower() not in strip} # Inject plugin's own credential as upstream auth (if configured) plugin_key = plugin.plugin_key @@ -318,9 +302,7 @@ async def plugin_proxy( if user_role: forward_headers["x-litellm-user-role"] = str(user_role) - handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.PassThroughEndpoint - ) + handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.PassThroughEndpoint) try: req = handler.client.build_request( method=request.method, diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index fb1e2652e8a..8ef509810ba 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -57,9 +57,7 @@ class AttachmentRegistry: try: attachment = self._parse_attachment(attachment_data) self._attachments.append(attachment) - verbose_proxy_logger.debug( - f"Loaded attachment for policy: {attachment.policy}" - ) + verbose_proxy_logger.debug(f"Loaded attachment for policy: {attachment.policy}") except Exception as e: verbose_proxy_logger.error(f"Error loading attachment: {str(e)}") raise ValueError(f"Invalid attachment: {str(e)}") from e @@ -96,13 +94,9 @@ class AttachmentRegistry: Returns: List of policy names that are attached to matching scopes """ - return [ - r["policy_name"] for r in self.get_attached_policies_with_reasons(context) - ] + return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)] - def get_attached_policies_with_reasons( - self, context: PolicyMatchContext - ) -> List[Dict[str, Any]]: + def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> List[Dict[str, Any]]: """ Get list of policy names and match reasons for the given context. @@ -135,9 +129,7 @@ class AttachmentRegistry: return results @staticmethod - def _describe_match_reason( - attachment: PolicyAttachment, context: PolicyMatchContext - ) -> str: + def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: """Describe why an attachment matched the context.""" from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher @@ -146,11 +138,7 @@ class AttachmentRegistry: reasons = [] if attachment.tags and context.tags: - matching_tags = [ - t - for t in context.tags - if PolicyMatcher.matches_pattern(t, attachment.tags) - ] + matching_tags = [t for t in context.tags if PolicyMatcher.matches_pattern(t, attachment.tags)] if matching_tags: reasons.append(f"tag:{matching_tags[0]}") if attachment.teams and context.team_alias: @@ -238,9 +226,7 @@ class AttachmentRegistry: self._attachments = [a for a in self._attachments if a.policy != policy_name] removed_count = original_count - len(self._attachments) if removed_count > 0: - verbose_proxy_logger.debug( - f"Removed {removed_count} attachment(s) for policy: {policy_name}" - ) + verbose_proxy_logger.debug(f"Removed {removed_count} attachment(s) for policy: {policy_name}") return removed_count def remove_attachment_by_id(self, attachment_id: str) -> bool: @@ -279,9 +265,7 @@ class AttachmentRegistry: PolicyAttachmentDBResponse with the created attachment """ try: - created_attachment = await PolicyAttachmentRepository( - prisma_client - ).table.create( + created_attachment = await PolicyAttachmentRepository(prisma_client).table.create( data={ "policy_name": attachment_request.policy_name, "scope": attachment_request.scope, @@ -341,17 +325,15 @@ class AttachmentRegistry: """ try: # Get attachment before deleting - attachment = await PolicyAttachmentRepository( - prisma_client - ).table.find_unique(where={"attachment_id": attachment_id}) + attachment = await PolicyAttachmentRepository(prisma_client).table.find_unique( + where={"attachment_id": attachment_id} + ) if attachment is None: raise Exception(f"Attachment with ID {attachment_id} not found") # Delete from DB - await PolicyAttachmentRepository(prisma_client).table.delete( - where={"attachment_id": attachment_id} - ) + await PolicyAttachmentRepository(prisma_client).table.delete(where={"attachment_id": attachment_id}) # Note: In-memory attachments don't have IDs, so we need to sync from DB # to properly update in-memory state @@ -378,9 +360,9 @@ class AttachmentRegistry: PolicyAttachmentDBResponse if found, None otherwise """ try: - attachment = await PolicyAttachmentRepository( - prisma_client - ).table.find_unique(where={"attachment_id": attachment_id}) + attachment = await PolicyAttachmentRepository(prisma_client).table.find_unique( + where={"attachment_id": attachment_id} + ) if attachment is None: return None @@ -416,9 +398,7 @@ class AttachmentRegistry: List of PolicyAttachmentDBResponse objects """ try: - attachments = await PolicyAttachmentRepository( - prisma_client - ).table.find_many( + attachments = await PolicyAttachmentRepository(prisma_client).table.find_many( order={"created_at": "desc"}, ) @@ -462,23 +442,15 @@ class AttachmentRegistry: attachment = PolicyAttachment( policy=attachment_response.policy_name, scope=attachment_response.scope, - teams=( - attachment_response.teams if attachment_response.teams else None - ), + teams=(attachment_response.teams if attachment_response.teams else None), keys=attachment_response.keys if attachment_response.keys else None, - models=( - attachment_response.models - if attachment_response.models - else None - ), + models=(attachment_response.models if attachment_response.models else None), tags=attachment_response.tags if attachment_response.tags else None, ) self._attachments.append(attachment) self._initialized = True - verbose_proxy_logger.info( - f"Synced {len(attachments)} attachments from DB to in-memory registry" - ) + verbose_proxy_logger.info(f"Synced {len(attachments)} attachments from DB to in-memory registry") except Exception as e: verbose_proxy_logger.exception(f"Error syncing attachments from DB: {e}") raise Exception(f"Error syncing attachments from DB: {str(e)}") diff --git a/litellm/proxy/policy_engine/condition_evaluator.py b/litellm/proxy/policy_engine/condition_evaluator.py index 1f1dea15a1d..de27b6d6d23 100644 --- a/litellm/proxy/policy_engine/condition_evaluator.py +++ b/litellm/proxy/policy_engine/condition_evaluator.py @@ -49,9 +49,7 @@ class ConditionEvaluator: condition=condition.model, model=context.model, ): - verbose_proxy_logger.debug( - f"Condition failed: model={context.model} did not match {condition.model}" - ) + verbose_proxy_logger.debug(f"Condition failed: model={context.model} did not match {condition.model}") return False return True @@ -76,10 +74,7 @@ class ConditionEvaluator: # Handle list of values if isinstance(condition, list): - return any( - ConditionEvaluator._matches_pattern(pattern, model) - for pattern in condition - ) + return any(ConditionEvaluator._matches_pattern(pattern, model) for pattern in condition) # Single value - check as pattern return ConditionEvaluator._matches_pattern(condition, model) diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index d5529f5bfe1..a3b0d1a6cdb 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -44,12 +44,8 @@ def _print_policies_on_startup( condition = policy_data.get("condition") description = policy_data.get("description") - guardrails_add = ( - guardrails.get("add", []) if isinstance(guardrails, dict) else [] - ) - guardrails_remove = ( - guardrails.get("remove", []) if isinstance(guardrails, dict) else [] - ) + guardrails_add = guardrails.get("add", []) if isinstance(guardrails, dict) else [] + guardrails_remove = guardrails.get("remove", []) if isinstance(guardrails, dict) else [] inherit_str = f" (inherits: {inherit})" if inherit else "" print( # noqa: T201 @@ -62,9 +58,7 @@ def _print_policies_on_startup( if guardrails_remove: print(f" guardrails.remove: {guardrails_remove}") # noqa: T201 if condition: - model_condition = ( - condition.get("model") if isinstance(condition, dict) else None - ) + model_condition = condition.get("model") if isinstance(condition, dict) else None if model_condition: print(f" condition.model: {model_condition}") # noqa: T201 @@ -152,33 +146,26 @@ async def init_policies( if validation_result.errors: for error in validation_result.errors: verbose_proxy_logger.error( - f"Policy validation error in '{error.policy_name}': " - f"[{error.error_type}] {error.message}" + f"Policy validation error in '{error.policy_name}': [{error.error_type}] {error.message}" ) if validation_result.warnings: for warning in validation_result.warnings: verbose_proxy_logger.warning( - f"Policy validation warning in '{warning.policy_name}': " - f"[{warning.error_type}] {warning.message}" + f"Policy validation warning in '{warning.policy_name}': [{warning.error_type}] {warning.message}" ) # Fail if there are errors and fail_on_error is True if not validation_result.valid and fail_on_error: - error_messages = [ - f"[{e.policy_name}] {e.message}" for e in validation_result.errors - ] + error_messages = [f"[{e.policy_name}] {e.message}" for e in validation_result.errors] raise ValueError( - f"Policy validation failed with {len(validation_result.errors)} error(s):\n" - + "\n".join(error_messages) + f"Policy validation failed with {len(validation_result.errors)} error(s):\n" + "\n".join(error_messages) ) # Load policies into registry (even with warnings) try: policy_registry.load_policies(policies_config) - verbose_proxy_logger.info( - f"Successfully loaded {len(policies_config)} policies" - ) + verbose_proxy_logger.info(f"Successfully loaded {len(policies_config)} policies") except Exception as e: verbose_proxy_logger.error(f"Failed to load policies: {str(e)}") raise @@ -187,9 +174,7 @@ async def init_policies( if policy_attachments_config: try: attachment_registry.load_attachments(policy_attachments_config) - verbose_proxy_logger.info( - f"Successfully loaded {len(policy_attachments_config)} policy attachments" - ) + verbose_proxy_logger.info(f"Successfully loaded {len(policy_attachments_config)} policy attachments") except Exception as e: verbose_proxy_logger.error(f"Failed to load policy attachments: {str(e)}") raise @@ -264,9 +249,7 @@ def get_policies_summary() -> Dict[str, Any]: "description": policy.description if policy else None, "guardrails_add": policy.guardrails.get_add() if policy else [], "guardrails_remove": policy.guardrails.get_remove() if policy else [], - "condition": ( - policy.condition.model_dump() if policy and policy.condition else None - ), + "condition": (policy.condition.model_dump() if policy and policy.condition else None), "resolved_guardrails": resolved_policy.guardrails, "inheritance_chain": resolved_policy.inheritance_chain, } diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 1fde9109e19..1f507bb4c54 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -92,8 +92,7 @@ class PipelineExecutor: step_results.append(step_result) verbose_proxy_logger.debug( - f"Pipeline '{policy_name}' step {i}: guardrail={step.guardrail}, " - f"outcome={outcome}, action={action}" + f"Pipeline '{policy_name}' step {i}: guardrail={step.guardrail}, outcome={outcome}, action={action}" ) # Forward modified data to next step if pass_data is True @@ -120,8 +119,7 @@ class PipelineExecutor: return PipelineExecutionResult( terminal_action="modify_response", step_results=step_results, - modify_response_message=step.modify_response_message - or error_detail, + modify_response_message=step.modify_response_message or error_detail, ) # action == "next" → continue to next step @@ -160,9 +158,7 @@ class PipelineExecutor: """ callback = PipelineExecutor.find_guardrail_callback(step.guardrail) if callback is None: - verbose_proxy_logger.warning( - f"Pipeline: guardrail '{step.guardrail}' not found in callbacks" - ) + verbose_proxy_logger.warning(f"Pipeline: guardrail '{step.guardrail}' not found in callbacks") return ("error", None, f"Guardrail '{step.guardrail}' not found", None) try: @@ -209,9 +205,7 @@ class PipelineExecutor: error_msg = _extract_error_message(e) return ("fail", None, error_msg, e) else: - verbose_proxy_logger.error( - f"Pipeline: unexpected error from guardrail '{step.guardrail}': {e}" - ) + verbose_proxy_logger.error(f"Pipeline: unexpected error from guardrail '{step.guardrail}': {e}") return ("error", None, str(e), e) @staticmethod diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 3a25e249a4e..4a94fc3d43c 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -88,9 +88,7 @@ async def list_policies(version_status: Optional[str] = None): raise HTTPException(status_code=500, detail="Database not connected") try: - policies = await get_policy_registry().get_all_policies_from_db( - prisma_client, version_status=version_status - ) + policies = await get_policy_registry().get_all_policies_from_db(prisma_client, version_status=version_status) return PolicyListDBResponse(policies=policies, total_count=len(policies)) except Exception as e: verbose_proxy_logger.exception(f"Error listing policies: {e}") @@ -260,11 +258,7 @@ async def update_policy_version_status( raise except Exception as e: verbose_proxy_logger.exception(f"Error updating version status: {e}") - if ( - "invalid status" in str(e).lower() - or "only draft" in str(e).lower() - or "cannot promote" in str(e).lower() - ): + if "invalid status" in str(e).lower() or "only draft" in str(e).lower() or "cannot promote" in str(e).lower(): raise HTTPException(status_code=400, detail=str(e)) if "not found" in str(e).lower(): raise HTTPException(status_code=404, detail=str(e)) @@ -358,9 +352,7 @@ async def get_policy(policy_id: str): prisma_client=prisma_client, ) if result is None: - raise HTTPException( - status_code=404, detail=f"Policy with ID {policy_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Policy with ID {policy_id} not found") return result except HTTPException: raise @@ -406,9 +398,7 @@ async def update_policy( prisma_client=prisma_client, ) if existing is None: - raise HTTPException( - status_code=404, detail=f"Policy with ID {policy_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Policy with ID {policy_id} not found") if getattr(existing, "version_status", "production") != "draft": raise HTTPException( status_code=400, @@ -464,9 +454,7 @@ async def delete_policy(policy_id: str): prisma_client=prisma_client, ) if existing is None: - raise HTTPException( - status_code=404, detail=f"Policy with ID {policy_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Policy with ID {policy_id} not found") result = await get_policy_registry().delete_policy_from_db( policy_id=policy_id, @@ -520,9 +508,7 @@ async def get_resolved_guardrails(policy_id: str): prisma_client=prisma_client, ) if policy is None: - raise HTTPException( - status_code=404, detail=f"Policy with ID {policy_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Policy with ID {policy_id} not found") # Resolve guardrails resolved = await get_policy_registry().resolve_guardrails_from_db( @@ -653,12 +639,8 @@ async def list_policy_attachments(): raise HTTPException(status_code=500, detail="Database not connected") try: - attachments = await get_attachment_registry().get_all_attachments_from_db( - prisma_client - ) - return PolicyAttachmentListResponse( - attachments=attachments, total_count=len(attachments) - ) + attachments = await get_attachment_registry().get_all_attachments_from_db(prisma_client) + return PolicyAttachmentListResponse(attachments=attachments, total_count=len(attachments)) except Exception as e: verbose_proxy_logger.exception(f"Error listing policy attachments: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -720,9 +702,7 @@ async def create_policy_attachment( try: # Verify the policy has a production version (attachments resolve against production) - policies = await get_policy_registry().get_all_policies_from_db( - prisma_client, version_status="production" - ) + policies = await get_policy_registry().get_all_policies_from_db(prisma_client, version_status="production") policy_names = {p.policy_name for p in policies} if request.policy_name not in policy_names: raise HTTPException( diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index b2788e6355b..acaa9629d83 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -45,9 +45,7 @@ class PolicyMatcher: for pattern in patterns: # Use existing wildcard pattern matching helper - if RouteChecks._route_matches_wildcard_pattern( - route=value, pattern=pattern - ): + if RouteChecks._route_matches_wildcard_pattern(route=value, pattern=pattern): return True return False @@ -88,9 +86,7 @@ class PolicyMatcher: if not context.tags: return False # Match if ANY context tag matches ANY scope tag pattern - if not any( - PolicyMatcher.matches_pattern(tag, scope_tags) for tag in context.tags - ): + if not any(PolicyMatcher.matches_pattern(tag, scope_tags) for tag in context.tags): return False return True @@ -114,9 +110,7 @@ class PolicyMatcher: registry = get_attachment_registry() if not registry.is_initialized(): - verbose_proxy_logger.debug( - "AttachmentRegistry not initialized, returning empty list" - ) + verbose_proxy_logger.debug("AttachmentRegistry not initialized, returning empty list") return [] return registry.get_attached_policies(context) @@ -172,9 +166,7 @@ class PolicyMatcher: if policy is None: continue # Policy matches if it has no condition OR condition evaluates to True - if policy.condition is None or ConditionEvaluator.evaluate( - policy.condition, context - ): + if policy.condition is None or ConditionEvaluator.evaluate(policy.condition, context): matching_policies.append(policy_name) return matching_policies diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index d6265516269..0dec93251f8 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -92,9 +92,7 @@ class PolicyRegistry: self._policies[policy_name] = policy verbose_proxy_logger.debug(f"Loaded policy: {policy_name}") except Exception as e: - verbose_proxy_logger.error( - f"Error loading policy '{policy_name}': {str(e)}" - ) + verbose_proxy_logger.error(f"Error loading policy '{policy_name}': {str(e)}") raise ValueError(f"Invalid policy '{policy_name}': {str(e)}") from e self._initialized = True @@ -120,9 +118,7 @@ class PolicyRegistry: ) else: # Handle legacy format where guardrails might be a list - guardrails = PolicyGuardrails( - add=guardrails_data if guardrails_data else None - ) + guardrails = PolicyGuardrails(add=guardrails_data if guardrails_data else None) # Parse condition (simple model-based condition) condition = None @@ -150,10 +146,7 @@ class PolicyRegistry: return None steps_data = pipeline_data.get("steps", []) - steps = [ - PipelineStep(**step_data) if isinstance(step_data, dict) else step_data - for step_data in steps_data - ] + steps = [PipelineStep(**step_data) if isinstance(step_data, dict) else step_data for step_data in steps_data] return GuardrailPipeline( mode=pipeline_data.get("mode", "pre_call"), @@ -296,9 +289,7 @@ class PolicyRegistry: validated_pipeline = GuardrailPipeline(**policy_request.pipeline) data["pipeline"] = json.dumps(validated_pipeline.model_dump()) - created_policy = await PolicyRepository(prisma_client).table.create( - data=data - ) + created_policy = await PolicyRepository(prisma_client).table.create(data=data) # Also add to in-memory registry policy = self._parse_policy( @@ -310,11 +301,7 @@ class PolicyRegistry: "add": policy_request.guardrails_add, "remove": policy_request.guardrails_remove, }, - "condition": ( - policy_request.condition.model_dump() - if policy_request.condition - else None - ), + "condition": (policy_request.condition.model_dump() if policy_request.condition else None), "pipeline": policy_request.pipeline, }, ) @@ -348,16 +335,12 @@ class PolicyRegistry: Exception: If policy is not in draft status (only drafts are editable). """ try: - existing = await PolicyRepository(prisma_client).table.find_unique( - where={"policy_id": policy_id} - ) + existing = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) if existing is None: raise Exception(f"Policy with ID {policy_id} not found") version_status = getattr(existing, "version_status", "production") if version_status != "draft": - raise Exception( - f"Only draft versions can be updated. This policy has status '{version_status}'." - ) + raise Exception(f"Only draft versions can be updated. This policy has status '{version_status}'.") # Build update data - only include fields that are set update_data: Dict[str, Any] = { @@ -376,9 +359,7 @@ class PolicyRegistry: if policy_request.guardrails_remove is not None: update_data["guardrails_remove"] = policy_request.guardrails_remove if policy_request.condition is not None: - update_data["condition"] = json.dumps( - policy_request.condition.model_dump() - ) + update_data["condition"] = json.dumps(policy_request.condition.model_dump()) if policy_request.pipeline is not None: validated_pipeline = GuardrailPipeline(**policy_request.pipeline) update_data["pipeline"] = json.dumps(validated_pipeline.model_dump()) @@ -414,9 +395,7 @@ class PolicyRegistry: Dict with "message" and optional "warning" if production was deleted. """ try: - policy = await PolicyRepository(prisma_client).table.find_unique( - where={"policy_id": policy_id} - ) + policy = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) if policy is None: raise Exception(f"Policy with ID {policy_id} not found") @@ -425,13 +404,9 @@ class PolicyRegistry: policy_name = policy.policy_name # Delete from DB - await PolicyRepository(prisma_client).table.delete( - where={"policy_id": policy_id} - ) + await PolicyRepository(prisma_client).table.delete(where={"policy_id": policy_id}) - result: Dict[str, Any] = { - "message": f"Policy {policy_id} deleted successfully" - } + result: Dict[str, Any] = {"message": f"Policy {policy_id} deleted successfully"} # Remove from in-memory registry only if this was the production version if version_status == "production": @@ -462,9 +437,7 @@ class PolicyRegistry: PolicyDBResponse if found, None otherwise """ try: - policy = await PolicyRepository(prisma_client).table.find_unique( - where={"policy_id": policy_id} - ) + policy = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) if policy is None: return None @@ -474,9 +447,7 @@ class PolicyRegistry: verbose_proxy_logger.exception(f"Error getting policy from DB: {e}") raise Exception(f"Error getting policy from DB: {str(e)}") - def get_policy_by_id_for_request( - self, policy_id: str - ) -> Optional[Tuple[str, Policy]]: + def get_policy_by_id_for_request(self, policy_id: str) -> Optional[Tuple[str, Policy]]: """ Return a policy version by ID from in-memory cache (no DB access). @@ -535,9 +506,7 @@ class PolicyRegistry: """ try: self._policies = {} - production = await self.get_all_policies_from_db( - prisma_client, version_status="production" - ) + production = await self.get_all_policies_from_db(prisma_client, version_status="production") for policy_response in production: policy = self._parse_policy( policy_response.policy_name, @@ -605,9 +574,7 @@ class PolicyRegistry: try: # Load only production versions so inheritance resolves against production - policies = await self.get_all_policies_from_db( - prisma_client, version_status="production" - ) + policies = await self.get_all_policies_from_db(prisma_client, version_status="production") # Build a temporary in-memory map for resolution temp_policies = {} @@ -691,15 +658,11 @@ class PolicyRegistry: """ try: if source_policy_id is not None: - source = await PolicyRepository(prisma_client).table.find_unique( - where={"policy_id": source_policy_id} - ) + source = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": source_policy_id}) if source is None: raise Exception(f"Source policy {source_policy_id} not found") if source.policy_name != policy_name: - raise Exception( - f"Source policy name '{source.policy_name}' does not match '{policy_name}'" - ) + raise Exception(f"Source policy name '{source.policy_name}' does not match '{policy_name}'") else: # Find current production version for this policy_name prod = await PolicyRepository(prisma_client).table.find_first( @@ -709,9 +672,7 @@ class PolicyRegistry: } ) if prod is None: - raise Exception( - f"No production version found for policy '{policy_name}'" - ) + raise Exception(f"No production version found for policy '{policy_name}'") source = prod # Next version number @@ -748,16 +709,10 @@ class PolicyRegistry: # Prisma expects Json fields as JSON strings on create (same as add_policy_to_db) if source.condition is not None: data["condition"] = ( - json.dumps(source.condition) - if isinstance(source.condition, dict) - else source.condition + json.dumps(source.condition) if isinstance(source.condition, dict) else source.condition ) if source.pipeline is not None: - data["pipeline"] = ( - json.dumps(source.pipeline) - if isinstance(source.pipeline, dict) - else source.pipeline - ) + data["pipeline"] = json.dumps(source.pipeline) if isinstance(source.pipeline, dict) else source.pipeline created = await PolicyRepository(prisma_client).table.create(data=data) return _row_to_policy_db_response(created) @@ -791,13 +746,9 @@ class PolicyRegistry: """ try: if new_status not in ("published", "production"): - raise Exception( - f"Invalid status '{new_status}'. Use 'published' or 'production'." - ) + raise Exception(f"Invalid status '{new_status}'. Use 'published' or 'production'.") - row = await PolicyRepository(prisma_client).table.find_unique( - where={"policy_id": policy_id} - ) + row = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id}) if row is None: raise Exception(f"Policy with ID {policy_id} not found") @@ -807,9 +758,7 @@ class PolicyRegistry: if new_status == "published": if current != "draft": - raise Exception( - f"Only draft versions can be published. Current status: '{current}'." - ) + raise Exception(f"Only draft versions can be published. Current status: '{current}'.") updated = await PolicyRepository(prisma_client).table.update( where={"policy_id": policy_id}, data={ @@ -828,9 +777,7 @@ class PolicyRegistry: ) # Plan: "draft -> production" NOT allowed if current == "draft": - raise Exception( - "Cannot promote draft directly to production. Publish the version first." - ) + raise Exception("Cannot promote draft directly to production. Publish the version first.") # Demote current production to published await PolicyRepository(prisma_client).table.update_many( @@ -896,12 +843,8 @@ class PolicyRegistry: PolicyVersionCompareResponse with both versions and field_diffs """ try: - a = await PolicyRepository(prisma_client).table.find_unique( - where={"policy_id": policy_id_a} - ) - b = await PolicyRepository(prisma_client).table.find_unique( - where={"policy_id": policy_id_b} - ) + a = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id_a}) + b = await PolicyRepository(prisma_client).table.find_unique(where={"policy_id": policy_id_b}) if a is None: raise Exception(f"Policy {policy_id_a} not found") if b is None: @@ -951,13 +894,9 @@ class PolicyRegistry: Dict with success message """ try: - await PolicyRepository(prisma_client).table.delete_many( - where={"policy_name": policy_name} - ) + await PolicyRepository(prisma_client).table.delete_many(where={"policy_name": policy_name}) self.remove_policy(policy_name) - return { - "message": f"All versions of policy '{policy_name}' deleted successfully" - } + return {"message": f"All versions of policy '{policy_name}' deleted successfully"} except Exception as e: verbose_proxy_logger.exception(f"Error deleting all versions: {e}") raise Exception(f"Error deleting all versions: {str(e)}") diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 84dcbcfd746..e5c2693cf21 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -97,9 +97,7 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple: unnamed_count = 0 for key in keys: key_alias = key.key_alias or "" - key_tags = _get_tags_from_metadata( - key.metadata, getattr(key, "metadata_json", None) - ) + key_tags = _get_tags_from_metadata(key.metadata, getattr(key, "metadata_json", None)) if key_tags and any( RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in key_tags @@ -153,8 +151,7 @@ async def _find_affected_by_team_patterns( for team in all_teams: team_alias = team.team_alias or "" if team_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) - for pat in team_patterns + RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns ): if team_alias not in existing_teams: new_teams.append(team_alias) @@ -179,9 +176,7 @@ async def _find_affected_by_team_patterns( return new_teams, new_keys, unnamed_keys_count -async def _find_affected_keys_by_alias( - prisma_client: object, key_patterns: list, existing_keys: list -) -> list: +async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list, existing_keys: list) -> list: """Find keys whose alias matches the given patterns.""" affected: list = [] @@ -194,8 +189,7 @@ async def _find_affected_keys_by_alias( for key in keys: key_alias = key.key_alias or "" if key_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) - for pat in key_patterns + RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns ): if key_alias not in existing_keys: affected.append(key_alias) @@ -260,9 +254,7 @@ async def resolve_policies_for_context( ) # Get matching policies with reasons - match_results = get_attachment_registry().get_attached_policies_with_reasons( - context=context - ) + match_results = get_attachment_registry().get_attached_policies_with_reasons(context=context) if not match_results: return PolicyResolveResponse( diff --git a/litellm/proxy/policy_engine/policy_resolver.py b/litellm/proxy/policy_engine/policy_resolver.py index 2c0a5334b05..f5520f4b4b4 100644 --- a/litellm/proxy/policy_engine/policy_resolver.py +++ b/litellm/proxy/policy_engine/policy_resolver.py @@ -49,9 +49,7 @@ class PolicyResolver: visited = set() if policy_name in visited: - verbose_proxy_logger.warning( - f"Circular inheritance detected for policy '{policy_name}'" - ) + verbose_proxy_logger.warning(f"Circular inheritance detected for policy '{policy_name}'") return [] policy = policies.get(policy_name) @@ -92,9 +90,7 @@ class PolicyResolver: """ from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator - inheritance_chain = PolicyResolver.resolve_inheritance_chain( - policy_name=policy_name, policies=policies - ) + inheritance_chain = PolicyResolver.resolve_inheritance_chain(policy_name=policy_name, policies=policies) # Start with empty set of guardrails guardrails: Set[str] = set() @@ -164,9 +160,7 @@ class PolicyResolver: # Use provided policy names or get matching policies via attachments matching_policy_names = ( - policy_names - if policy_names is not None - else PolicyMatcher.get_matching_policies(context=context) + policy_names if policy_names is not None else PolicyMatcher.get_matching_policies(context=context) ) if not matching_policy_names: @@ -186,9 +180,7 @@ class PolicyResolver: context=context, ) all_guardrails.update(resolved.guardrails) - verbose_proxy_logger.debug( - f"Policy '{policy_name}' contributes guardrails: {resolved.guardrails}" - ) + verbose_proxy_logger.debug(f"Policy '{policy_name}' contributes guardrails: {resolved.guardrails}") result = list(all_guardrails) verbose_proxy_logger.debug(f"Final guardrails for context: {result}") @@ -226,9 +218,7 @@ class PolicyResolver: policies = registry.get_all_policies() matching_policy_names = ( - policy_names - if policy_names is not None - else PolicyMatcher.get_matching_policies(context=context) + policy_names if policy_names is not None else PolicyMatcher.get_matching_policies(context=context) ) if not matching_policy_names: return [] @@ -241,8 +231,7 @@ class PolicyResolver: if policy.pipeline is not None: pipelines.append((policy_name, policy.pipeline)) verbose_proxy_logger.debug( - f"Policy '{policy_name}' has pipeline with " - f"{len(policy.pipeline.steps)} steps" + f"Policy '{policy_name}' has pipeline with {len(policy.pipeline.steps)} steps" ) return pipelines diff --git a/litellm/proxy/policy_engine/policy_validator.py b/litellm/proxy/policy_engine/policy_validator.py index 46796fbae28..626bbbc1ce5 100644 --- a/litellm/proxy/policy_engine/policy_validator.py +++ b/litellm/proxy/policy_engine/policy_validator.py @@ -74,15 +74,9 @@ class PolicyValidator: ) guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails() - return { - g.get("guardrail_name", "") - for g in guardrails - if g.get("guardrail_name") - } + return {g.get("guardrail_name", "") for g in guardrails if g.get("guardrail_name")} except Exception as e: - verbose_proxy_logger.warning( - f"Could not get guardrails from registry: {str(e)}" - ) + verbose_proxy_logger.warning(f"Could not get guardrails from registry: {str(e)}") return set() async def check_team_alias_exists(self, team_alias: str) -> bool: @@ -104,9 +98,7 @@ class PolicyValidator: ) return team is not None except Exception as e: - verbose_proxy_logger.warning( - f"Could not check team alias '{team_alias}': {str(e)}" - ) + verbose_proxy_logger.warning(f"Could not check team alias '{team_alias}': {str(e)}") return True # Assume valid on error async def check_key_alias_exists(self, key_alias: str) -> bool: @@ -123,16 +115,12 @@ class PolicyValidator: return True # Can't validate without DB, assume valid try: - key = await VerificationTokenRepository( - self.prisma_client - ).table.find_first( + key = await VerificationTokenRepository(self.prisma_client).table.find_first( where={"key_alias": key_alias}, ) return key is not None except Exception as e: - verbose_proxy_logger.warning( - f"Could not check key alias '{key_alias}': {str(e)}" - ) + verbose_proxy_logger.warning(f"Could not check key alias '{key_alias}': {str(e)}") return True # Assume valid on error def check_model_exists(self, model: str) -> bool: @@ -155,11 +143,7 @@ class PolicyValidator: # Check if model matches any pattern via pattern router if hasattr(self.llm_router, "pattern_router"): - pattern_deployments = ( - self.llm_router.pattern_router.get_deployments_by_pattern( - model=model - ) - ) + pattern_deployments = self.llm_router.pattern_router.get_deployments_by_pattern(model=model) if pattern_deployments: return True @@ -238,11 +222,7 @@ class PolicyValidator: else: # Recursively check parent with decremented depth visited.add(policy_name) - errors.extend( - self._validate_inheritance_chain( - policy.inherit, policies, visited, max_depth - 1 - ) - ) + errors.extend(self._validate_inheritance_chain(policy.inherit, policies, visited, max_depth - 1)) return errors @@ -303,9 +283,7 @@ class PolicyValidator: errors.extend(pipeline_errors) # Validate inheritance - inheritance_errors = self._validate_inheritance_chain( - policy_name=policy_name, policies=policies - ) + inheritance_errors = self._validate_inheritance_chain(policy_name=policy_name, policies=policies) errors.extend(inheritance_errors) return PolicyValidationResponse( @@ -336,8 +314,7 @@ class PolicyValidator: policy_name=policy_name, error_type=PolicyValidationErrorType.INVALID_GUARDRAIL, message=( - f"Pipeline step {i} guardrail '{step.guardrail}' " - f"is not in the policy's guardrails.add list" + f"Pipeline step {i} guardrail '{step.guardrail}' is not in the policy's guardrails.add list" ), field="pipeline.steps", value=step.guardrail, @@ -350,10 +327,7 @@ class PolicyValidator: PolicyValidationError( policy_name=policy_name, error_type=PolicyValidationErrorType.INVALID_GUARDRAIL, - message=( - f"Pipeline step {i} guardrail '{step.guardrail}' " - f"not found in guardrail registry" - ), + message=(f"Pipeline step {i} guardrail '{step.guardrail}' not found in guardrail registry"), field="pipeline.steps", value=step.guardrail, ) diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 251d1e56287..0f6944fac79 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -5,9 +5,7 @@ import os import subprocess import sys -sys.path.insert( - 0, os.path.abspath("./") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("./")) # Adds the parent directory to the system path from litellm._logging import verbose_proxy_logger from litellm.proxy.proxy_cli import run_server @@ -23,6 +21,4 @@ exit_code = result.returncode if exit_code != 0: verbose_proxy_logger.info(f"'prisma generate' failed with exit code {exit_code}.") - verbose_proxy_logger.error( - f"'prisma generate' stderr: {result.stderr}" - ) # Log stderr + verbose_proxy_logger.error(f"'prisma generate' stderr: {result.stderr}") # Log stderr diff --git a/litellm/proxy/prometheus_cleanup.py b/litellm/proxy/prometheus_cleanup.py index 6353588532a..2a22b1c5fae 100644 --- a/litellm/proxy/prometheus_cleanup.py +++ b/litellm/proxy/prometheus_cleanup.py @@ -21,13 +21,9 @@ def wipe_directory(directory: str) -> None: os.remove(filepath) deleted += 1 except OSError as e: - verbose_proxy_logger.warning( - f"Failed to delete stale prometheus file {filepath}: {e}" - ) + verbose_proxy_logger.warning(f"Failed to delete stale prometheus file {filepath}: {e}") if deleted: - verbose_proxy_logger.info( - f"Prometheus cleanup: wiped {deleted} stale .db files from {directory}" - ) + verbose_proxy_logger.info(f"Prometheus cleanup: wiped {deleted} stale .db files from {directory}") def mark_worker_exit(worker_pid: int) -> None: @@ -38,10 +34,6 @@ def mark_worker_exit(worker_pid: int) -> None: from prometheus_client import multiprocess multiprocess.mark_process_dead(worker_pid) - verbose_proxy_logger.info( - f"Prometheus cleanup: marked worker {worker_pid} as dead" - ) + verbose_proxy_logger.info(f"Prometheus cleanup: marked worker {worker_pid} as dead") except Exception as e: - verbose_proxy_logger.warning( - f"Failed to mark prometheus worker {worker_pid} as dead: {e}" - ) + verbose_proxy_logger.warning(f"Failed to mark prometheus worker {worker_pid} as dead: {e}") diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index c0d6794108a..91843e28283 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -186,18 +186,14 @@ def get_latest_prompt_versions(prompts: List[PromptSpec]) -> List[PromptSpec]: if base_id not in latest_prompts: latest_prompts[base_id] = prompt else: - existing_version = get_version_number( - prompt_id=latest_prompts[base_id].prompt_id - ) + existing_version = get_version_number(prompt_id=latest_prompts[base_id].prompt_id) if version > existing_version: latest_prompts[base_id] = prompt return list(latest_prompts.values()) -async def get_next_version_for_prompt( - prisma_client, prompt_id: str, environment: str = "development" -) -> int: +async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment: str = "development") -> int: """ Get the next version number for a prompt in a specific environment. @@ -430,9 +426,7 @@ async def get_prompt_versions( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): - raise HTTPException( - status_code=403, detail="Only proxy admins can view prompt versions" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can view prompt versions") base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) @@ -486,16 +480,12 @@ async def get_prompt_versions( versioned_prompts.sort(key=lambda p: p.version or 1, reverse=True) if not versioned_prompts: - raise HTTPException( - status_code=404, detail=f"No versions found for prompt ID {base_prompt_id}" - ) + raise HTTPException(status_code=404, detail=f"No versions found for prompt ID {base_prompt_id}") return ListPromptsResponse(prompts=versioned_prompts) -def _get_prompt_template( - prompt_spec: PromptSpec, base_prompt_id: str -) -> Optional[PromptTemplateBase]: +def _get_prompt_template(prompt_spec: PromptSpec, base_prompt_id: str) -> Optional[PromptTemplateBase]: """Resolve the raw prompt template from dotprompt content or the in-memory registry.""" from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY @@ -514,9 +504,7 @@ def _get_prompt_template( metadata=parsed.get("metadata"), ) else: - prompt_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id( - prompt_spec.prompt_id - ) + prompt_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(prompt_spec.prompt_id) if prompt_callback is not None: integration_name = prompt_callback.integration_name if integration_name == "dotprompt": @@ -525,9 +513,7 @@ def _get_prompt_template( ) if isinstance(prompt_callback, DotpromptManager): - template = ( - prompt_callback.prompt_manager.get_all_prompts_as_json() - ) + template = prompt_callback.prompt_manager.get_all_prompts_as_json() if template is not None and len(template) == 1: template_id = list(template.keys())[0] return PromptTemplateBase( @@ -592,9 +578,7 @@ async def get_prompt_info( ## CHECK IF USER HAS ACCESS TO PROMPT prompts: Optional[List[str]] = None if user_api_key_dict.metadata is not None: - prompts = cast( - Optional[List[str]], user_api_key_dict.metadata.get("prompts", None) - ) + prompts = cast(Optional[List[str]], user_api_key_dict.metadata.get("prompts", None)) if prompts is not None and prompt_id not in prompts: raise HTTPException(status_code=400, detail=f"Prompt {prompt_id} not found") if user_api_key_dict.user_role is not None and ( @@ -617,17 +601,13 @@ async def get_prompt_info( where={"prompt_id": base_prompt_id}, distinct=["environment"], ) - all_environments = sorted( - set(row.environment for row in all_prompt_rows if row.environment) - ) + all_environments = sorted(set(row.environment for row in all_prompt_rows if row.environment)) # If environment is specified, find the version in that environment from DB # If prompt_id has a version suffix (e.g., "testprompt.v2"), fetch that specific version # Otherwise fetch the latest version in that environment prompt_spec = None - requested_version = ( - get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None - ) + requested_version = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None if environment and prisma_client is not None: where_clause: Dict[str, Any] = { "prompt_id": base_prompt_id, @@ -656,8 +636,7 @@ async def get_prompt_info( if prompt_spec is None: raise HTTPException( status_code=400, - detail=f"Prompt {prompt_id} not found" - + (f" in environment {environment}" if environment else ""), + detail=f"Prompt {prompt_id} not found" + (f" in environment {environment}" if environment else ""), ) # Extract version number from the prompt_id @@ -728,14 +707,10 @@ async def create_prompt( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): - raise HTTPException( - status_code=403, detail="Only proxy admins can create prompts" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can create prompts") if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Extract environment from request @@ -772,9 +747,7 @@ async def create_prompt( prompt_spec = create_versioned_prompt_spec(db_prompt=prompt_db_entry) # Initialize the prompt - initialized_prompt = IN_MEMORY_PROMPT_REGISTRY.initialize_prompt( - prompt=prompt_spec, config_file_path=None - ) + initialized_prompt = IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec, config_file_path=None) if initialized_prompt is None: raise HTTPException(status_code=500, detail="Failed to initialize prompt") @@ -828,14 +801,10 @@ async def update_prompt( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): - raise HTTPException( - status_code=403, detail="Only proxy admins can update prompts" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can update prompts") if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Strip version suffix from prompt_id if present (e.g., "jack_success.v1" -> "jack_success") @@ -849,9 +818,7 @@ async def update_prompt( ) # Check if any version of this prompt exists (in any environment) - existing_prompts = await PromptRepository(prisma_client).table.find_many( - where={"prompt_id": base_prompt_id} - ) + existing_prompts = await PromptRepository(prisma_client).table.find_many(where={"prompt_id": base_prompt_id}) if not existing_prompts: raise HTTPException( @@ -861,10 +828,7 @@ async def update_prompt( # Check if it's a config prompt existing_in_memory = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - if ( - existing_in_memory - and existing_in_memory.prompt_info.prompt_type == "config" - ): + if existing_in_memory and existing_in_memory.prompt_info.prompt_type == "config": raise HTTPException( status_code=400, detail="Cannot update config prompts.", @@ -897,9 +861,7 @@ async def update_prompt( prompt_spec = create_versioned_prompt_spec(db_prompt=prompt_db_entry) # Initialize the new version - initialized_prompt = IN_MEMORY_PROMPT_REGISTRY.initialize_prompt( - prompt=prompt_spec, config_file_path=None - ) + initialized_prompt = IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec, config_file_path=None) if initialized_prompt is None: raise HTTPException(status_code=500, detail="Failed to update prompt") @@ -949,14 +911,10 @@ async def delete_prompt( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): - raise HTTPException( - status_code=403, detail="Only proxy admins can delete prompts" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can delete prompts") if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Try to get prompt directly first @@ -968,16 +926,12 @@ async def delete_prompt( prompt_id=prompt_id, all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, ) - existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id( - latest_prompt_id - ) + existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) # Use the resolved prompt_id for deletion prompt_id = latest_prompt_id if existing_prompt is None: - raise HTTPException( - status_code=404, detail=f"Prompt with ID {prompt_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Prompt with ID {prompt_id} not found") if existing_prompt.prompt_info.prompt_type == "config": raise HTTPException( @@ -1001,8 +955,7 @@ async def delete_prompt( prompts_to_delete = [ pid for pid, prompt in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.items() - if get_base_prompt_id(prompt_id=pid) == base_prompt_id - and prompt.environment == environment + if get_base_prompt_id(prompt_id=pid) == base_prompt_id and prompt.environment == environment ] for pid in prompts_to_delete: del IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[pid] @@ -1021,17 +974,13 @@ async def delete_prompt( raise HTTPException(status_code=500, detail=str(e)) -def _reload_prompt_in_registry( - registry: Any, versioned_id: str, updated_prompt_spec: PromptSpec -) -> PromptSpec: +def _reload_prompt_in_registry(registry: Any, versioned_id: str, updated_prompt_spec: PromptSpec) -> PromptSpec: """Remove stale entry and re-initialize the prompt in the in-memory registry.""" if versioned_id in registry.IN_MEMORY_PROMPTS: del registry.IN_MEMORY_PROMPTS[versioned_id] if versioned_id in registry.prompt_id_to_custom_prompt: del registry.prompt_id_to_custom_prompt[versioned_id] - initialized = registry.initialize_prompt( - prompt=updated_prompt_spec, config_file_path=None - ) + initialized = registry.initialize_prompt(prompt=updated_prompt_spec, config_file_path=None) if initialized is None: raise HTTPException(status_code=500, detail="Failed to patch prompt") return initialized @@ -1079,24 +1028,16 @@ async def patch_prompt( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): - raise HTTPException( - status_code=403, detail="Only proxy admins can patch prompts" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can patch prompts") if prisma_client is None: - raise HTTPException( - status_code=500, detail=CommonProxyErrors.db_not_connected_error.value - ) + raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: # Resolve the target row: find the latest version in the given environment base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) env = environment or "development" - requested_version = ( - get_version_number(prompt_id=prompt_id) - if prompt_id != base_prompt_id - else None - ) + requested_version = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None # Build query to find the exact row by composite unique key find_where: Dict[str, Any] = { @@ -1140,16 +1081,10 @@ async def patch_prompt( # Update fields if provided updated_litellm_params = ( - request.litellm_params - if request.litellm_params is not None - else current_litellm_params + request.litellm_params if request.litellm_params is not None else current_litellm_params ) - updated_prompt_info = ( - request.prompt_info - if request.prompt_info is not None - else current_prompt_info - ) + updated_prompt_info = request.prompt_info if request.prompt_info is not None else current_prompt_info # Ensure we have valid litellm_params if updated_litellm_params is None: @@ -1169,13 +1104,9 @@ async def patch_prompt( data=update_data, ) - updated_prompt_spec = create_versioned_prompt_spec( - db_prompt=updated_prompt_db_entry - ) + updated_prompt_spec = create_versioned_prompt_spec(db_prompt=updated_prompt_db_entry) - return _reload_prompt_in_registry( - IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec - ) + return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec) except HTTPException as e: raise e @@ -1241,44 +1172,30 @@ async def test_prompt( try: # Parse the dotprompt content and create PromptTemplate prompt_manager = PromptManager() - frontmatter, template_content = prompt_manager._parse_frontmatter( - content=request.dotprompt_content - ) + frontmatter, template_content = prompt_manager._parse_frontmatter(content=request.dotprompt_content) # Create PromptTemplate to leverage existing parameter extraction logic - template = PromptTemplate( - content=template_content, metadata=frontmatter, template_id="test_prompt" - ) + template = PromptTemplate(content=template_content, metadata=frontmatter, template_id="test_prompt") # Extract model from template if not template.model: - raise HTTPException( - status_code=400, detail="Model is required in dotprompt metadata" - ) + raise HTTPException(status_code=400, detail="Model is required in dotprompt metadata") # Always render the template to extract system messages and other metadata variables = request.prompt_variables or {} - rendered_content = prompt_manager.jinja_env.from_string( - template_content - ).render(**variables) + rendered_content = prompt_manager.jinja_env.from_string(template_content).render(**variables) # Convert rendered content to messages using DotpromptManager's method dotprompt_manager = DotpromptManager() - rendered_messages = dotprompt_manager._convert_to_messages( - rendered_content=rendered_content - ) + rendered_messages = dotprompt_manager._convert_to_messages(rendered_content=rendered_content) if not rendered_messages: - raise HTTPException( - status_code=400, detail="No messages found in rendered prompt" - ) + raise HTTPException(status_code=400, detail="No messages found in rendered prompt") # If conversation history is provided, use it but preserve system messages if request.conversation_history: # Extract system messages from rendered prompt - system_messages = [ - msg for msg in rendered_messages if msg.get("role") == "system" - ] + system_messages = [msg for msg in rendered_messages if msg.get("role") == "system"] # Use conversation history for user/assistant messages messages = system_messages + request.conversation_history else: @@ -1387,9 +1304,7 @@ async def convert_prompt_file_to_json( } except Exception as e: - raise HTTPException( - status_code=500, detail=f"Error converting prompt file: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Error converting prompt file: {str(e)}") finally: # Clean up temp file diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index ae5ce177853..e4d6fb30a1d 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -51,9 +51,7 @@ def get_prompt_initializer_from_integrations(): module_path = f"litellm.integrations.{item}" try: # Import the module - verbose_proxy_logger.debug( - f"Discovering prompt integrations in: {module_path}" - ) + verbose_proxy_logger.debug(f"Discovering prompt integrations in: {module_path}") module = importlib.import_module(module_path) @@ -97,9 +95,7 @@ class InMemoryPromptRegistry: Prompt id to Prompt object mapping """ - self.prompt_id_to_custom_prompt: Dict[ - str, Optional[CustomPromptManagement] - ] = {} + self.prompt_id_to_custom_prompt: Dict[str, Optional[CustomPromptManagement]] = {} """ Guardrail id to CustomGuardrail object mapping """ @@ -139,12 +135,8 @@ class InMemoryPromptRegistry: if initializer: custom_prompt_callback = initializer(litellm_params, prompt) if not isinstance(custom_prompt_callback, CustomPromptManagement): - raise ValueError( - f"CustomPromptManagement is required, got {type(custom_prompt_callback)}" - ) - litellm.logging_callback_manager.add_litellm_callback( - custom_prompt_callback - ) # type: ignore + raise ValueError(f"CustomPromptManagement is required, got {type(custom_prompt_callback)}") + litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback) # type: ignore else: raise ValueError(f"Unsupported prompt: {prompt_integration}") @@ -168,9 +160,7 @@ class InMemoryPromptRegistry: """ return self.IN_MEMORY_PROMPTS.get(prompt_id) - def get_prompt_callback_by_id( - self, prompt_id: str - ) -> Optional[CustomPromptManagement]: + def get_prompt_callback_by_id(self, prompt_id: str) -> Optional[CustomPromptManagement]: """ Get a prompt callback by its ID from memory """ @@ -189,9 +179,7 @@ class InMemoryPromptRegistry: from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id prompts_to_delete = [ - pid - for pid in self.IN_MEMORY_PROMPTS.keys() - if get_base_prompt_id(prompt_id=pid) == base_prompt_id + pid for pid in self.IN_MEMORY_PROMPTS.keys() if get_base_prompt_id(prompt_id=pid) == base_prompt_id ] for pid in prompts_to_delete: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index a8efbadb253..74ec0cc8700 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -82,9 +82,7 @@ def append_query_params(url: Optional[str], params: dict) -> str: if not isinstance(url, str) or url == "": # Preserve previous startup behavior when DATABASE_URL is absent. # Returning an empty string avoids urlparse type errors in test/dev flows. - verbose_proxy_logger.warning( - "append_query_params received empty or non-string URL, returning empty string" - ) + verbose_proxy_logger.warning("append_query_params received empty or non-string URL, returning empty string") return "" parsed_url = urlparse.urlparse(url) parsed_query = urlparse.parse_qs(parsed_url.query) @@ -114,9 +112,7 @@ class ProxyInitializationHelpers: test: Union[bool, str], ): request_model = model or "gpt-3.5-turbo" - click.echo( - f"\nLiteLLM: Making a test ChatCompletions request to your proxy. Model={request_model}" - ) + click.echo(f"\nLiteLLM: Making a test ChatCompletions request to your proxy. Model={request_model}") import openai api_base = f"http://{host}:{port}" @@ -138,9 +134,7 @@ class ProxyInitializationHelpers: ) click.echo(f"\nLiteLLM: response from proxy {response}") - print( - f"\n LiteLLM: Making a test ChatCompletions + streaming r equest to proxy. Model={request_model}" - ) + print(f"\n LiteLLM: Making a test ChatCompletions + streaming r equest to proxy. Model={request_model}") stream_response = client.chat.completions.create( model=request_model, @@ -192,10 +186,7 @@ class ProxyInitializationHelpers: if keepalive_timeout is not None: uvicorn_args["timeout_keep_alive"] = keepalive_timeout if timeout_worker_healthcheck is not None: - if ( - "timeout_worker_healthcheck" - in inspect.signature(uvicorn.Config.__init__).parameters - ): + if "timeout_worker_healthcheck" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["timeout_worker_healthcheck"] = timeout_worker_healthcheck else: print( @@ -224,10 +215,7 @@ class ProxyInitializationHelpers: "has no effect without --max_requests_before_restart\033[0m\n" ) return - if ( - "limit_max_requests_jitter" - in inspect.signature(uvicorn.Config.__init__).parameters - ): + if "limit_max_requests_jitter" in inspect.signature(uvicorn.Config.__init__).parameters: uvicorn_args["limit_max_requests_jitter"] = jitter else: print( @@ -309,9 +297,7 @@ class ProxyInitializationHelpers: uvicorn_args.update(ProxyInitializationHelpers._get_reload_options(config_path)) os.environ["LITELLM_DEV_ENV_HOT_RELOAD"] = "True" env_path = os.path.join(os.getcwd(), ".env") - ProxyInitializationHelpers._patch_statreload_extra_paths( - [config_path, env_path] - ) + ProxyInitializationHelpers._patch_statreload_extra_paths([config_path, env_path]) verbose_proxy_logger.warning( "LiteLLM --reload: worker processes re-read .env with override, so .env " "values win over shell-exported environment variables. Unset a key in .env " @@ -335,9 +321,7 @@ class ProxyInitializationHelpers: from hypercorn.asyncio import serve from hypercorn.config import Config - print( - f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Hypercorn\033[0m\n" - ) + print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Hypercorn\033[0m\n") config = Config() config.bind = [f"{host}:{port}"] @@ -373,18 +357,14 @@ class ProxyInitializationHelpers: from granian import Granian from granian.constants import Interfaces - print( - f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n" - ) + print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n") if max_requests_before_restart is not None: print( "\033[1;33mLiteLLM: --max_requests_before_restart is not supported by Granian " "(Granian uses workers_lifetime in seconds, not a per-request limit).\033[0m\n" ) if ciphers is not None: - print( - "\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n" - ) + print("\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n") kwargs: dict[str, Any] = { "target": "litellm.proxy.proxy_server:app", @@ -403,9 +383,7 @@ class ProxyInitializationHelpers: kwargs["ssl_cert"] = Path(ssl_certfile_path) kwargs["ssl_key"] = Path(ssl_keyfile_path) elif ssl_certfile_path is not None or ssl_keyfile_path is not None: - raise click.ClickException( - "Both --ssl_certfile_path and --ssl_keyfile_path are required for SSL." - ) + raise click.ClickException("Both --ssl_certfile_path and --ssl_keyfile_path are required for SSL.") Granian(**kwargs).serve() @@ -435,9 +413,7 @@ class ProxyInitializationHelpers: self.application = app # FastAPI app super().__init__() - _endpoint_str = ( - f"curl --location 'http://0.0.0.0:{port}/chat/completions' \\" - ) + _endpoint_str = f"curl --location 'http://0.0.0.0:{port}/chat/completions' \\" curl_command = ( _endpoint_str + """ @@ -458,15 +434,9 @@ class ProxyInitializationHelpers: print( '\033[1;34mLiteLLM: Test your local proxy with: "litellm --test" This runs an openai.ChatCompletion request to your proxy [In a new terminal tab]\033[0m\n' ) - print( - f"\033[1;34mLiteLLM: Curl Command Test for your local proxy\n {curl_command} \033[0m\n" - ) - print( - "\033[1;34mDocs: https://docs.litellm.ai/docs/simple_proxy\033[0m\n" - ) - print( - f"\033[1;34mSee all Router/Swagger docs on http://0.0.0.0:{port} \033[0m\n" - ) + print(f"\033[1;34mLiteLLM: Curl Command Test for your local proxy\n {curl_command} \033[0m\n") + print("\033[1;34mDocs: https://docs.litellm.ai/docs/simple_proxy\033[0m\n") + print(f"\033[1;34mSee all Router/Swagger docs on http://0.0.0.0:{port} \033[0m\n") def load_config(self): # note: This Loads the gunicorn config - has nothing to do with LiteLLM Proxy config @@ -486,9 +456,7 @@ class ProxyInitializationHelpers: # gunicorn app function return self.application - print( - f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n" - ) + print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n") gunicorn_options = { "bind": f"{host}:{port}", "workers": num_workers, # default is 1 @@ -509,9 +477,7 @@ class ProxyInitializationHelpers: "has no effect without --max_requests_before_restart\033[0m\n" ) else: - gunicorn_options["max_requests_jitter"] = ( - max_requests_before_restart_jitter - ) + gunicorn_options["max_requests_jitter"] = max_requests_before_restart_jitter # Clean up prometheus .db files when a worker exits (prevents ghost gauge values) if os.environ.get("PROMETHEUS_MULTIPROC_DIR"): @@ -588,15 +554,11 @@ class ProxyInitializationHelpers: from litellm.proxy.prometheus_cleanup import wipe_directory - multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get( - "prometheus_multiproc_dir" - ) + multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get("prometheus_multiproc_dir") auto_created = not multiproc_dir if not multiproc_dir: - multiproc_dir = os.path.join( - tempfile.gettempdir(), "litellm_prometheus_multiproc" - ) + multiproc_dir = os.path.join(tempfile.gettempdir(), "litellm_prometheus_multiproc") os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir os.makedirs(multiproc_dir, exist_ok=True) @@ -607,9 +569,7 @@ class ProxyInitializationHelpers: @click.command() @click.argument("cli_args", nargs=-1) -@click.option( - "--host", default="0.0.0.0", help="Host for the server to listen on.", envvar="HOST" -) +@click.option("--host", default="0.0.0.0", help="Host for the server to listen on.", envvar="HOST") @click.option("--port", default=4000, help="Port to bind the server to.", envvar="PORT") @click.option( "--num_workers", @@ -637,17 +597,13 @@ class ProxyInitializationHelpers: default=litellm.AZURE_DEFAULT_API_VERSION, help="For azure - pass in the api version.", ) -@click.option( - "--model", "-m", default=None, help="The model name to pass to litellm expects" -) +@click.option("--model", "-m", default=None, help="The model name to pass to litellm expects") @click.option( "--alias", default=None, help='The alias for the model - use this to give a litellm model name (e.g. "huggingface/codellama/CodeLlama-7b-Instruct-hf") a more user-friendly name ("codellama")', ) -@click.option( - "--add_key", default=None, help="The model name to pass to litellm expects" -) +@click.option("--add_key", default=None, help="The model name to pass to litellm expects") @click.option("--headers", default=None, help="headers for the API call") @click.option("--save", is_flag=True, type=bool, help="Save the model-specific config") @click.option( @@ -673,12 +629,8 @@ class ProxyInitializationHelpers: type=bool, help="To use celery workers for async endpoints", ) -@click.option( - "--temperature", default=None, type=float, help="Set temperature for the model" -) -@click.option( - "--max_tokens", default=None, type=int, help="Set max tokens for the model" -) +@click.option("--temperature", default=None, type=float, help="Set temperature for the model") +@click.option("--max_tokens", default=None, type=int, help="Set max tokens for the model") @click.option( "--request_timeout", default=None, @@ -926,9 +878,7 @@ def run_server( authenticator = XAIOAuthAuthenticator() auth_data = authenticator.login() - click.echo( - f"xAI OAuth login successful. Credentials saved to {authenticator.auth_file}." - ) + click.echo(f"xAI OAuth login successful. Credentials saved to {authenticator.auth_file}.") if auth_data.get("expires_at"): click.echo(f"Access token expires at {auth_data['expires_at']}.") return @@ -957,9 +907,7 @@ def run_server( save_worker_config, ) except ModuleNotFoundError as e: - raise ModuleNotFoundError( - f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`" - ) + raise ModuleNotFoundError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`") except ImportError as e: if "litellm[proxy]" in str(e): # user is missing a proxy dependency, ask them to pip install litellm[proxy] @@ -1018,9 +966,7 @@ def run_server( try: import uvicorn except Exception: - raise ImportError( - "uvicorn, gunicorn needs to be imported. Run - `pip install 'litellm[proxy]'`" - ) + raise ImportError("uvicorn, gunicorn needs to be imported. Run - `pip install 'litellm[proxy]'`") db_connection_pool_limit = 100 # Starts optional due to config fallback checks; guaranteed non-None before use. @@ -1046,9 +992,7 @@ def run_server( db_name = os.getenv("DATABASE_NAME") db_schema = os.getenv("DATABASE_SCHEMA") - token = generate_iam_auth_token( - db_host=db_host, db_port=db_port, db_user=db_user - ) + token = generate_iam_auth_token(db_host=db_host, db_port=db_port, db_user=db_user) # print(f"token: {token}") _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" @@ -1062,10 +1006,7 @@ def run_server( from litellm.secret_managers.aws_secret_manager import decrypt_env_var - if ( - os.getenv("USE_AWS_KMS", None) is not None - and os.getenv("USE_AWS_KMS") == "True" - ): + if os.getenv("USE_AWS_KMS", None) is not None and os.getenv("USE_AWS_KMS") == "True": ## V2 IMPLEMENTATION OF AWS KMS - USER WANTS TO DECRYPT MULTIPLE KEYS IN THEIR ENV new_env_var = decrypt_env_var() @@ -1083,9 +1024,7 @@ def run_server( import asyncio except Exception: - raise ImportError( - "yaml needs to be imported. Run - `pip install 'litellm[proxy]'`" - ) + raise ImportError("yaml needs to be imported. Run - `pip install 'litellm[proxy]'`") proxy_config = ProxyConfig() _config = asyncio.run(proxy_config.get_config(config_file_path=config)) @@ -1107,21 +1046,15 @@ def run_server( if general_settings is None: general_settings = {} ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ### - key_management_settings = general_settings.get( - "key_management_settings", None - ) + key_management_settings = general_settings.get("key_management_settings", None) if key_management_settings is not None: import litellm - litellm._key_management_settings = KeyManagementSettings( - **key_management_settings - ) + litellm._key_management_settings = KeyManagementSettings(**key_management_settings) if general_settings: ### LOAD SECRET MANAGER ### - key_management_system = general_settings.get( - "key_management_system", None - ) + key_management_system = general_settings.get("key_management_system", None) proxy_config.initialize_secret_manager( key_management_system=key_management_system, config_file_path=config ) @@ -1139,29 +1072,19 @@ def run_server( ) db_connection_timeout = general_settings.get("database_connection_timeout") if db_connection_timeout is None: - db_connection_timeout = general_settings.get( - "database_connection_pool_timeout" - ) + db_connection_timeout = general_settings.get("database_connection_pool_timeout") if db_connection_timeout is None: - db_connection_timeout = ( - LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value - ) + db_connection_timeout = LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value db_connect_timeout = general_settings.get("database_connect_timeout") db_socket_timeout = general_settings.get("database_socket_timeout") - _disable_prepared_statements = general_settings.get( - "database_disable_prepared_statements", False - ) + _disable_prepared_statements = general_settings.get("database_disable_prepared_statements", False) if isinstance(_disable_prepared_statements, str): from litellm.secret_managers.main import str_to_bool - db_disable_prepared_statements = ( - str_to_bool(_disable_prepared_statements) is True - ) + db_disable_prepared_statements = str_to_bool(_disable_prepared_statements) is True else: db_disable_prepared_statements = bool(_disable_prepared_statements) - db_extra_connection_params = general_settings.get( - "database_extra_connection_params" - ) + db_extra_connection_params = general_settings.get("database_extra_connection_params") if database_url and database_url.startswith("os.environ/"): original_dir = os.getcwd() # set the working directory to where this script is @@ -1187,17 +1110,10 @@ def run_server( # Set default values for connection pool settings when no config is used if config is None: - db_connection_pool_limit = ( - LiteLLMDatabaseConnectionPool.database_connection_pool_limit.value - ) - db_connection_timeout = ( - LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value - ) + db_connection_pool_limit = LiteLLMDatabaseConnectionPool.database_connection_pool_limit.value + db_connection_timeout = LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value - if ( - os.getenv("DATABASE_URL", None) is not None - or os.getenv("DIRECT_URL", None) is not None - ): + if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None: from litellm.proxy.db.db_url_settings import ( unsupported_db_scheme, unsupported_db_scheme_message, @@ -1210,9 +1126,7 @@ def run_server( _bad_scheme = unsupported_db_scheme(_candidate_url) if _bad_scheme is not None: print( - f"\033[1;31mLiteLLM Proxy: " - f"{unsupported_db_scheme_message(_db_env, _bad_scheme)}" - "\033[0m", + f"\033[1;31mLiteLLM Proxy: {unsupported_db_scheme_message(_db_env, _bad_scheme)}\033[0m", file=sys.stderr, flush=True, ) @@ -1237,9 +1151,7 @@ def run_server( os.environ["DATABASE_URL"] = modified_url if os.getenv("DIRECT_URL", None) is not None: database_url = os.getenv("DIRECT_URL") - modified_url = append_query_params( - database_url, connection_url_params - ) + modified_url = append_query_params(database_url, connection_url_params) os.environ["DIRECT_URL"] = modified_url subprocess.run(["prisma"], capture_output=True) is_prisma_runnable = True @@ -1253,12 +1165,7 @@ def run_server( should_update_prisma_schema, ) - if ( - should_update_prisma_schema( - general_settings.get("disable_prisma_schema_update") - ) - is False - ): + if should_update_prisma_schema(general_settings.get("disable_prisma_schema_update")) is False: check_prisma_schema_diff(db_url=None) else: if not use_v2_migration_resolver: @@ -1279,8 +1186,7 @@ def run_server( # v1 never raises here, so this only fires when the # operator opted into v2. print( - "\033[1;31mLiteLLM Proxy: Database migration cannot proceed. " - f"{e}\033[0m", + f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {e}\033[0m", file=sys.stderr, flush=True, ) @@ -1329,9 +1235,7 @@ def run_server( port=port, log_config=log_config, keepalive_timeout=keepalive_timeout, - timeout_worker_healthcheck=( - timeout_worker_healthcheck if running_uvicorn else None - ), + timeout_worker_healthcheck=(timeout_worker_healthcheck if running_uvicorn else None), ) # Optional: recycle uvicorn workers after N requests if max_requests_before_restart is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b2a85ddc978..adbba821bf3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -152,9 +152,7 @@ warnings.filterwarnings("default", category=UserWarning) messages: list = [] -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path - for litellm local dev +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path - for litellm local dev try: import logging @@ -190,8 +188,7 @@ def generate_feedback_box(): print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa: T201 print("\033[1;37m" + "# {:^59} #\033[0m".format(message)) # noqa: T201 print( # noqa: T201 - "\033[1;37m" - + "# {:^59} #\033[0m".format("https://github.com/BerriAI/litellm/issues/new") + "\033[1;37m" + "# {:^59} #\033[0m".format("https://github.com/BerriAI/litellm/issues/new") ) print("\033[1;37m" + "#" + " " * box_width + "#\033[0m") # noqa: T201 print("\033[1;37m" + "#" + "-" * box_width + "#\033[0m") # noqa: T201 @@ -631,12 +628,8 @@ except ImportError: server_root_path = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional["EnterpriseLicenseData"] = ( - _license_check.airgapped_license_data -) -global_max_parallel_request_retries_env: Optional[str] = os.getenv( - "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" -) +premium_user_data: Optional["EnterpriseLicenseData"] = _license_check.airgapped_license_data +global_max_parallel_request_retries_env: Optional[str] = os.getenv("LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES") proxy_state = ProxyState() SENSITIVE_DATA_MASKER = SensitiveDataMasker() if global_max_parallel_request_retries_env is None: @@ -650,9 +643,7 @@ global_max_parallel_request_retry_timeout_env: Optional[str] = os.getenv( if global_max_parallel_request_retry_timeout_env is None: global_max_parallel_request_retry_timeout: float = 60.0 else: - global_max_parallel_request_retry_timeout = float( - global_max_parallel_request_retry_timeout_env - ) + global_max_parallel_request_retry_timeout = float(global_max_parallel_request_retry_timeout_env) ui_link = f"{server_root_path}/ui" fallback_login_link = f"{server_root_path}/fallback/login" @@ -662,7 +653,9 @@ ui_message += "\n\n💸 [```LiteLLM Model Cost Map```](https://models.litellm.ai ui_message += f"\n\n🔎 [```LiteLLM Model Hub```]({model_hub_link}). See available models on the proxy. [**Docs**](https://docs.litellm.ai/docs/proxy/ai_hub)" -custom_swagger_message = "[**Customize Swagger Docs**](https://docs.litellm.ai/docs/proxy/enterprise#swagger-docs---custom-routes--branding)" +custom_swagger_message = ( + "[**Customize Swagger Docs**](https://docs.litellm.ai/docs/proxy/enterprise#swagger-docs---custom-routes--branding)" +) ### CUSTOM BRANDING [ENTERPRISE FEATURE] ### _title = os.getenv("DOCS_TITLE", "LiteLLM API") if premium_user else "LiteLLM API" @@ -711,12 +704,7 @@ def cleanup_router_config_variables(): async def proxy_shutdown_event(): - global \ - prisma_client, \ - master_key, \ - user_custom_auth, \ - user_custom_key_generate, \ - user_custom_key_update + global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") if prisma_client: verbose_proxy_logger.debug("Disconnecting from Prisma") @@ -778,9 +766,7 @@ async def _initialize_shared_aiohttp_session(): ) return session except Exception as e: - verbose_proxy_logger.warning( - f"Failed to create shared aiohttp session: {e}. Continuing without session reuse." - ) + verbose_proxy_logger.warning(f"Failed to create shared aiohttp session: {e}. Continuing without session reuse.") return None @@ -825,20 +811,14 @@ async def proxy_startup_event(app: FastAPI): await _hook_fn() else: _hook_fn() - verbose_proxy_logger.info( - "Worker startup hook '%s' executed successfully", _hook_spec - ) + verbose_proxy_logger.info("Worker startup hook '%s' executed successfully", _hook_spec) except Exception as e: - verbose_proxy_logger.error( - "Worker startup hook '%s' failed: %s", _hook_spec, e - ) + verbose_proxy_logger.error("Worker startup hook '%s' failed: %s", _hook_spec, e) raise ## CHECK PREMIUM USER verbose_proxy_logger.debug( - "litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format( - premium_user - ) + "litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format(premium_user) ) if premium_user is False: premium_user = _license_check.is_premium() @@ -851,16 +831,12 @@ async def proxy_startup_event(app: FastAPI): verbose_proxy_logger.debug("worker_config: %s", worker_config) # check if it's a valid file path if env_config_yaml is not None: - if os.path.isfile(env_config_yaml) and proxy_config.is_yaml( - config_file_path=env_config_yaml - ): + if os.path.isfile(env_config_yaml) and proxy_config.is_yaml(config_file_path=env_config_yaml): ( llm_router, llm_model_list, general_settings, - ) = await proxy_config.load_config( - router=llm_router, config_file_path=env_config_yaml - ) + ) = await proxy_config.load_config(router=llm_router, config_file_path=env_config_yaml) elif worker_config is not None: if ( isinstance(worker_config, str) @@ -871,19 +847,13 @@ async def proxy_startup_event(app: FastAPI): llm_router, llm_model_list, general_settings, - ) = await proxy_config.load_config( - router=llm_router, config_file_path=worker_config - ) - elif os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None and isinstance( - worker_config, str - ): + ) = await proxy_config.load_config(router=llm_router, config_file_path=worker_config) + elif os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None and isinstance(worker_config, str): ( llm_router, llm_model_list, general_settings, - ) = await proxy_config.load_config( - router=llm_router, config_file_path=worker_config - ) + ) = await proxy_config.load_config(router=llm_router, config_file_path=worker_config) elif isinstance(worker_config, dict): await initialize(**worker_config) else: @@ -916,10 +886,8 @@ async def proxy_startup_event(app: FastAPI): ## when the proxy cache backend is not Redis ## transaction_buffer_redis_cache = redis_usage_cache if transaction_buffer_redis_cache is None: - transaction_buffer_redis_cache = ( - ProxyStartupEvent._get_transaction_buffer_redis_cache( - general_settings=general_settings - ) + transaction_buffer_redis_cache = ProxyStartupEvent._get_transaction_buffer_redis_cache( + general_settings=general_settings ) ProxyStartupEvent._initialize_startup_logging( @@ -951,11 +919,7 @@ async def proxy_startup_event(app: FastAPI): publish_global_otel_v2_provider, ) - registered = ( - open_telemetry_logger - if isinstance(open_telemetry_logger, OpenTelemetryV2) - else None - ) + registered = open_telemetry_logger if isinstance(open_telemetry_logger, OpenTelemetryV2) else None publish_global_otel_v2_provider( _in_memory_loggers, # any-ok: pre-existing untyped List[Any] global _otel_trace.set_tracer_provider, @@ -976,9 +940,7 @@ async def proxy_startup_event(app: FastAPI): verbose_proxy_logger.debug("About to initialize semantic tool filter") _config = proxy_config.get_config_state() _litellm_settings = _config.get("litellm_settings", {}) - verbose_proxy_logger.debug( - f"litellm_settings keys = {list(_litellm_settings.keys())}" - ) + verbose_proxy_logger.debug(f"litellm_settings keys = {list(_litellm_settings.keys())}") await ProxyStartupEvent._initialize_semantic_tool_filter( llm_router=llm_router, litellm_settings=_litellm_settings, @@ -999,9 +961,7 @@ async def proxy_startup_event(app: FastAPI): verbose_proxy_logger.debug("prisma_client: %s", prisma_client) if prisma_client is not None and litellm.max_budget > 0: - ProxyStartupEvent._add_proxy_budget_to_db( - litellm_proxy_budget_name=litellm_proxy_admin_name - ) + ProxyStartupEvent._add_proxy_budget_to_db(litellm_proxy_budget_name=litellm_proxy_admin_name) asyncio.create_task( ProxyStartupEvent._warm_global_spend_cache( litellm_proxy_admin_name=litellm_proxy_admin_name, @@ -1028,9 +988,7 @@ async def proxy_startup_event(app: FastAPI): # Start background health checks AFTER models are loaded and index is built if use_background_health_checks: - asyncio.create_task( - _run_background_health_check() - ) # start the background health check coroutine. + asyncio.create_task(_run_background_health_check()) # start the background health check coroutine. # Start adaptive-router queue flusher unconditionally — adaptive routers # may be added later via `/config/reload`, and the flusher is a no-op when @@ -1080,9 +1038,7 @@ async def proxy_startup_event(app: FastAPI): verbose_proxy_logger.error(f"Error stopping token refresh task: {e}") # Shutdown event - stop Prisma DB health watchdog task - if prisma_client is not None and hasattr( - prisma_client, "stop_db_health_watchdog_task" - ): + if prisma_client is not None and hasattr(prisma_client, "stop_db_health_watchdog_task"): try: await prisma_client.stop_db_health_watchdog_task() except Exception as e: @@ -1170,9 +1126,7 @@ def ensure_unique_openapi_operation_ids( if not isinstance(operation_id, str): continue operation_entries.append((method, operation, operation_id)) - operation_id_counts[operation_id] = ( - operation_id_counts.get(operation_id, 0) + 1 - ) + operation_id_counts[operation_id] = operation_id_counts.get(operation_id, 0) + 1 used_operation_ids = set(reserved_operation_ids or set()) seen_operation_ids: Set[str] = set() @@ -1190,10 +1144,7 @@ def ensure_unique_openapi_operation_ids( base_operation_id = _strip_operation_id_method_suffix(operation_id) new_operation_id = f"{base_operation_id}_{method}" suffix = 2 - while ( - new_operation_id in used_operation_ids - or new_operation_id in seen_operation_ids - ): + while new_operation_id in used_operation_ids or new_operation_id in seen_operation_ids: new_operation_id = f"{base_operation_id}_{method}_{suffix}" suffix += 1 operation["operationId"] = new_operation_id @@ -1237,9 +1188,7 @@ vertex_live_passthrough_vertex_base = VertexBase() from fastapi.routing import APIWebSocketRoute -def _inject_websocket_stubs_into_openapi_schema( - openapi_schema: dict, websocket_routes: list -) -> dict: +def _inject_websocket_stubs_into_openapi_schema(openapi_schema: dict, websocket_routes: list) -> dict: """ Add a synthetic GET stub for each WebSocket route so it appears in Swagger UI. @@ -1301,15 +1250,11 @@ def get_openapi_schema(): ) # Find all WebSocket routes - websocket_routes = [ - route for route in app.routes if isinstance(route, APIWebSocketRoute) - ] + websocket_routes = [route for route in app.routes if isinstance(route, APIWebSocketRoute)] # Add a synthetic GET stub for each so they render in Swagger UI, # without clobbering existing HTTP operations on the same path. - openapi_schema = _inject_websocket_stubs_into_openapi_schema( - openapi_schema, websocket_routes - ) + openapi_schema = _inject_websocket_stubs_into_openapi_schema(openapi_schema, websocket_routes) # Add LLM API request schema bodies for documentation from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec @@ -1387,9 +1332,7 @@ async def openai_exception_handler(request: Request, exc: ProxyException): ) -def _close_dangling_otel_server_span( - request: Request, status_code: int, exc: Optional[Exception] = None -) -> None: +def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Optional[Exception] = None) -> None: parent_otel_span = getattr(request.state, "parent_otel_span", None) if parent_otel_span is None: return @@ -1409,29 +1352,19 @@ def _close_dangling_otel_server_span( try: from opentelemetry.trace import Status, StatusCode - open_telemetry_logger.set_response_status_code_attribute( - parent_otel_span, status_code - ) + open_telemetry_logger.set_response_status_code_attribute(parent_otel_span, status_code) if status_code >= 400: - open_telemetry_logger.record_error_attributes_on_span( - parent_otel_span, exc, status_code - ) - parent_otel_span.set_status( - Status(StatusCode.ERROR if status_code >= 400 else StatusCode.OK) - ) + open_telemetry_logger.record_error_attributes_on_span(parent_otel_span, exc, status_code) + parent_otel_span.set_status(Status(StatusCode.ERROR if status_code >= 400 else StatusCode.OK)) parent_otel_span.end() except Exception as e: - verbose_proxy_logger.debug( - "Error closing dangling OTEL SERVER span: %s", str(e) - ) + verbose_proxy_logger.debug("Error closing dangling OTEL SERVER span: %s", str(e)) finally: request.state.parent_otel_span = None @app.exception_handler(RequestValidationError) -async def otel_request_validation_exception_handler( - request: Request, exc: RequestValidationError -): +async def otel_request_validation_exception_handler(request: Request, exc: RequestValidationError): _close_dangling_otel_server_span(request, 422, exc=exc) return JSONResponse( status_code=422, @@ -1443,9 +1376,7 @@ async def otel_request_validation_exception_handler( async def otel_unhandled_exception_handler(request: Request, exc: Exception): if isinstance(exc, (ProxyException, HTTPException, RequestValidationError)): raise exc - verbose_proxy_logger.exception( - "Unhandled exception in request: %s", type(exc).__name__ - ) + verbose_proxy_logger.exception("Unhandled exception in request: %s", type(exc).__name__) _close_dangling_otel_server_span(request, 500, exc=exc) return JSONResponse( status_code=500, @@ -1477,11 +1408,7 @@ def _get_cors_config( Returns: Tuple[List[str], bool]: (origins, allow_credentials) """ - _origins_raw = ( - cors_origins_env - if cors_origins_env is not None - else os.getenv("LITELLM_CORS_ORIGINS") - ) + _origins_raw = cors_origins_env if cors_origins_env is not None else os.getenv("LITELLM_CORS_ORIGINS") if _origins_raw is None or _origins_raw.strip() == "": computed_origins = ["*"] else: @@ -1494,9 +1421,7 @@ def _get_cors_config( # (e.g. for non-browser clients that relied on the Access-Control-Allow-Credentials # header being present regardless of origin). _credentials_raw = ( - cors_credentials_env - if cors_credentials_env is not None - else os.getenv("LITELLM_CORS_ALLOW_CREDENTIALS") + cors_credentials_env if cors_credentials_env is not None else os.getenv("LITELLM_CORS_ALLOW_CREDENTIALS") ) if _credentials_raw is not None: computed_credentials = _credentials_raw.strip().lower() == "true" @@ -1586,17 +1511,13 @@ try: ) return True except (PermissionError, OSError) as e: - verbose_proxy_logger.debug( - f"Could not scan {ui_dir} for restructuring detection: {e}" - ) + verbose_proxy_logger.debug(f"Could not scan {ui_dir} for restructuring detection: {e}") return False # No restructured routes found return False - def _try_populate_ui_directory( - source_path: str, target_path: str - ) -> tuple[bool, str]: + def _try_populate_ui_directory(source_path: str, target_path: str) -> tuple[bool, str]: """ Attempt to populate target UI directory from source. @@ -1634,8 +1555,7 @@ try: # Validate packaged UI before proceeding if not _validate_ui_directory(packaged_ui_path): verbose_proxy_logger.error( - f"Packaged UI at {packaged_ui_path} is invalid or incomplete. " - f"UI may not function correctly." + f"Packaged UI at {packaged_ui_path} is invalid or incomplete. UI may not function correctly." ) # Decision tree for UI path selection: @@ -1667,13 +1587,9 @@ try: # Case 4: Runtime UI missing - try to populate else: - verbose_proxy_logger.info( - f"UI not found at {runtime_ui_path}. Attempting to populate from packaged UI." - ) + verbose_proxy_logger.info(f"UI not found at {runtime_ui_path}. Attempting to populate from packaged UI.") - success, error = _try_populate_ui_directory( - packaged_ui_path, runtime_ui_path - ) + success, error = _try_populate_ui_directory(packaged_ui_path, runtime_ui_path) if success: # Case 4a: Population succeeded @@ -1694,9 +1610,7 @@ try: # Validate final UI path if not _validate_ui_directory(ui_path): - verbose_proxy_logger.error( - f"Selected UI path {ui_path} is invalid or incomplete. UI may not work correctly." - ) + verbose_proxy_logger.error(f"Selected UI path {ui_path} is invalid or incomplete. UI may not work correctly.") # Only modify files if a custom server root path is set AND filesystem is writable if server_root_path and server_root_path != "/": @@ -1801,9 +1715,7 @@ try: is_writable = os.access(ui_path, os.W_OK) if is_pre_restructured: - verbose_proxy_logger.info( - f"Skipping UI restructuring: {ui_path} is already pre-restructured" - ) + verbose_proxy_logger.info(f"Skipping UI restructuring: {ui_path} is already pre-restructured") elif not is_writable: verbose_proxy_logger.warning( f"Cannot restructure UI at {ui_path}: path is not writable. " @@ -1814,13 +1726,9 @@ try: _restructure_ui_html_files(ui_path) verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}") except PermissionError as e: - verbose_proxy_logger.exception( - f"Permission error while restructuring UI directory {ui_path}: {e}" - ) + verbose_proxy_logger.exception(f"Permission error while restructuring UI directory {ui_path}: {e}") except Exception as e: - verbose_proxy_logger.exception( - f"Error while restructuring UI directory {ui_path}: {e}" - ) + verbose_proxy_logger.exception(f"Error while restructuring UI directory {ui_path}: {e}") except Exception: pass @@ -1920,8 +1828,7 @@ def mount_swagger_ui(): body = response.body.decode("utf-8") body = body.replace( "const ui = SwaggerUIBundle({", - _lazy_plugin_js - + 'const ui = SwaggerUIBundle({plugins:[LazyLoadPlugin],tagsSorter:"alpha",', + _lazy_plugin_js + 'const ui = SwaggerUIBundle({plugins:[LazyLoadPlugin],tagsSorter:"alpha",', 1, ) return HTMLResponse(content=body) @@ -1965,26 +1872,16 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional["ClientSession"] = ( - None # Global shared session for connection reuse -) +shared_aiohttp_session: Optional["ClientSession"] = None # Global shared session for connection reuse user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) -spend_counter_cache = DualCache( - default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value -) -model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( - dual_cache=user_api_key_cache -) +spend_counter_cache = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) +model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[RedisCache] = ( - None # redis cache used for tracking spend, tpm/rpm limits -) +redis_usage_cache: Optional[RedisCache] = None # redis cache used for tracking spend, tpm/rpm limits polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[ - str -] = [] # Models that should use native provider background mode instead of polling +native_background_mode: List[str] = [] # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -2019,9 +1916,7 @@ prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = N store_model_in_db: bool = False open_telemetry_logger: Optional[OpenTelemetry] = None ### INITIALIZE GLOBAL LOGGING OBJECT ### -proxy_logging_obj: ProxyLogging = ProxyLogging( - user_api_key_cache=user_api_key_cache, premium_user=premium_user -) +proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user) ### REDIS QUEUE ### async_result = None celery_app_conn = None @@ -2060,11 +1955,7 @@ def _resolve_pydantic_type(typ) -> List: typs = [] if origin is Union: # Check if it's a Union (like Optional) for arg in get_args(typ): - if ( - arg is not None - and not isinstance(arg, type(None)) - and "NoneType" not in str(arg) - ): + if arg is not None and not isinstance(arg, type(None)) and "NoneType" not in str(arg): typs.append(arg) elif isinstance(typ, type) and isinstance(typ, BaseModel): return [typ] @@ -2083,9 +1974,7 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): KVUri = os.getenv("AZURE_KEY_VAULT_URI", None) if KVUri is None: - raise Exception( - "Error when loading keys from Azure Key Vault: AZURE_KEY_VAULT_URI is not set." - ) + raise Exception("Error when loading keys from Azure Key Vault: AZURE_KEY_VAULT_URI is not set.") credential = DefaultAzureCredential() @@ -2106,9 +1995,7 @@ def cost_tracking(): global prisma_client if prisma_client is not None: litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger()) - litellm.logging_callback_manager.add_litellm_async_success_callback( - _ProxyDBLogger() - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(_ProxyDBLogger()) # Bounds authoritative DB re-reads when enforcing a budget against a @@ -2175,9 +2062,7 @@ async def get_current_spend( and cached in-process for a few seconds, so a persistently stale counter drives at most one read per counter per window rather than one per request. """ - current, verified = await _read_spend_counter_estimate( - counter_key=counter_key, fallback_spend=fallback_spend - ) + current, verified = await _read_spend_counter_estimate(counter_key=counter_key, fallback_spend=fallback_spend) if fallback_authoritative: verified = True @@ -2201,9 +2086,7 @@ async def get_current_spend( if authoritative is not None: verified = True if authoritative > current: - await _repair_stale_spend_counter( - counter_key=counter_key, db_spend=authoritative - ) + await _repair_stale_spend_counter(counter_key=counter_key, db_spend=authoritative) return authoritative elif fallback_spend > current: # end-user / tag counters have no DB row; fallback_spend is the @@ -2244,9 +2127,7 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_max( - key=counter_key, value=db_spend - ) + await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend) except Exception: verbose_proxy_logger.debug( "Unable to repair stale spend counter %s in Redis", @@ -2271,9 +2152,7 @@ async def reseed_spend_counter_from_db(counter_key: str) -> None: Counters with no DB row (window/end-user/tag) are left untouched rather than deleted, so enforcement keeps reading whatever value they hold. """ - db_spend = await SpendCounterReseed.from_db( - prisma_client=prisma_client, counter_key=counter_key - ) + db_spend = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) if db_spend is None: return await _repair_stale_spend_counter(counter_key=counter_key, db_spend=db_spend) @@ -2290,9 +2169,7 @@ async def _authoritative_floor_spend( if cached is not None: return float(cached) - db_spend = await SpendCounterReseed.from_db( - prisma_client=prisma_client, counter_key=counter_key - ) + db_spend = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) if ( db_spend is None and window_entity_type is not None @@ -2316,9 +2193,7 @@ async def _authoritative_floor_spend( return db_spend -async def _read_spend_counter_estimate( - counter_key: str, fallback_spend: float -) -> tuple[float, bool]: +async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]: """Return (spend, authoritative). ``authoritative`` is True when the value came from Redis or a fresh DB read (cross-pod truth), False when it came from the per-pod in-memory copy or the caller's fallback. Only the @@ -2396,11 +2271,7 @@ async def increment_spend_counters( # if a raw key somehow arrives, hash it; otherwise use as-is to # avoid double-hashing (budget checks read valid_token.token which # is single-hashed). - hashed_token = ( - hash_token(token=token) - if isinstance(token, str) and token.startswith("sk-") - else token - ) + hashed_token = hash_token(token=token) if isinstance(token, str) and token.startswith("sk-") else token key_counter_key = f"spend:key:{hashed_token}" if key_counter_key not in reserved_counter_keys: await _init_and_increment_spend_counter( @@ -2419,11 +2290,7 @@ async def increment_spend_counters( key_budget_limits = json.loads(key_budget_limits) if isinstance(key_budget_limits, list): for window in key_budget_limits: - duration = ( - window["budget_duration"] - if isinstance(window, dict) - else window.budget_duration - ) + duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration key_window_counter = f"spend:key:{hashed_token}:window:{duration}" if key_window_counter not in reserved_counter_keys: from litellm.proxy.spend_tracking.budget_reservation import ( @@ -2457,11 +2324,7 @@ async def increment_spend_counters( team_budget_limits = json.loads(team_budget_limits) if isinstance(team_budget_limits, list): for window in team_budget_limits: - duration = ( - window["budget_duration"] - if isinstance(window, dict) - else window.budget_duration - ) + duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration team_window_counter = f"spend:team:{team_id}:window:{duration}" if team_window_counter not in reserved_counter_keys: from litellm.proxy.spend_tracking.budget_reservation import ( @@ -2523,9 +2386,7 @@ async def _reconcile_budget_reservation_for_counter_update( reconcile_budget_reservation, ) - reserved_counter_keys = get_reserved_counter_keys( - budget_reservation=budget_reservation - ) + reserved_counter_keys = get_reserved_counter_keys(budget_reservation=budget_reservation) try: await reconcile_budget_reservation( budget_reservation=budget_reservation, @@ -2538,9 +2399,7 @@ async def _reconcile_budget_reservation_for_counter_update( exc_info=True, ) try: - await invalidate_budget_reservation_counters( - budget_reservation=budget_reservation - ) + await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: verbose_proxy_logger.exception( "Failed to invalidate reserved counters after reservation reconciliation failed" @@ -2679,21 +2538,15 @@ async def _ensure_spend_counter_initialized( ) if db_spend is None: # DB unavailable - fall back to in-process cache (may be stale). - base_spend = await _get_source_cache_base_spend( - source_cache_key=source_cache_key - ) + base_spend = await _get_source_cache_base_spend(source_cache_key=source_cache_key) if base_spend > 0: - await _increment_spend_counter_cache( - counter_key=counter_key, increment=base_spend - ) + await _increment_spend_counter_cache(counter_key=counter_key, increment=base_spend) async def _get_source_cache_base_spend( source_cache_key: Union[str, List[str]], ) -> float: - source_cache_keys = ( - [source_cache_key] if isinstance(source_cache_key, str) else source_cache_key - ) + source_cache_keys = [source_cache_key] if isinstance(source_cache_key, str) else source_cache_key for cache_key in source_cache_keys: source = await user_api_key_cache.async_get_cache(key=cache_key) if source is None: @@ -2816,12 +2669,8 @@ async def update_cache( else: hashed_token = token verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token) - existing_spend_obj = await user_api_key_cache.async_get_cache( - key=hashed_token, model_type=UserAPIKeyAuth - ) - verbose_proxy_logger.debug( - f"_update_key_cache: existing_spend_obj={existing_spend_obj}" - ) + existing_spend_obj = await user_api_key_cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth) + verbose_proxy_logger.debug(f"_update_key_cache: existing_spend_obj={existing_spend_obj}") if existing_spend_obj is None: return @@ -2865,23 +2714,15 @@ async def update_cache( ) # set cooldown on alert - if ( - existing_spend_obj is not None - and getattr(existing_spend_obj, "team_spend", None) is not None - ): + if existing_spend_obj is not None and getattr(existing_spend_obj, "team_spend", None) is not None: existing_team_spend = existing_spend_obj.team_spend or 0 # Calculate the new cost by adding the existing cost and response_cost existing_spend_obj.team_spend = existing_team_spend + response_cost - if ( - existing_spend_obj is not None - and getattr(existing_spend_obj, "team_member_spend", None) is not None - ): + if existing_spend_obj is not None and getattr(existing_spend_obj, "team_member_spend", None) is not None: existing_team_member_spend = existing_spend_obj.team_member_spend or 0 # Calculate the new cost by adding the existing cost and response_cost - existing_spend_obj.team_member_spend = ( - existing_team_member_spend + response_cost - ) + existing_spend_obj.team_member_spend = existing_team_member_spend + response_cost # Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns # BaseModel values into dicts for Redis (same Codec path as async_set_cache). @@ -2903,9 +2744,7 @@ async def update_cache( if cached_user is None: # do nothing if there is no cache value return - existing_spend_obj = CacheCodec.deserialize( - cached_user, LiteLLM_UserTable - ) + existing_spend_obj = CacheCodec.deserialize(cached_user, LiteLLM_UserTable) if existing_spend_obj is None: return verbose_proxy_logger.debug( @@ -2920,9 +2759,7 @@ async def update_cache( values_to_update_in_cache.append( ( _id, - CacheCodec.serialize( - existing_spend_obj, model_type=LiteLLM_UserTable - ), + CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_UserTable), ) ) ## UPDATE GLOBAL PROXY ## @@ -2934,9 +2771,7 @@ async def update_cache( return elif response_cost is not None and global_proxy_spend is not None: increment = global_proxy_spend + response_cost - values_to_update_in_cache.append( - ("{}:spend".format(litellm_proxy_admin_name), increment) - ) + values_to_update_in_cache.append(("{}:spend".format(litellm_proxy_admin_name), increment)) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update user spend in cache. " @@ -2961,9 +2796,7 @@ async def update_cache( # if user does not exist in LiteLLM_UserTable, create a new user # do nothing if end-user not in api key cache return - existing_spend_obj = CacheCodec.deserialize( - cached_end_user, LiteLLM_EndUserTable - ) + existing_spend_obj = CacheCodec.deserialize(cached_end_user, LiteLLM_EndUserTable) if existing_spend_obj is None: return verbose_proxy_logger.debug( @@ -2978,9 +2811,7 @@ async def update_cache( values_to_update_in_cache.append( ( _id, - CacheCodec.serialize( - existing_spend_obj, model_type=LiteLLM_EndUserTable - ), + CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_EndUserTable), ) ) except Exception as e: @@ -3005,8 +2836,8 @@ async def update_cache( if cached_team is None: # do nothing if team not in api key cache return - existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = ( - CacheCodec.deserialize(cached_team, LiteLLM_TeamTableCachedObj) + existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = CacheCodec.deserialize( + cached_team, LiteLLM_TeamTableCachedObj ) if existing_spend_obj is None: return @@ -3022,9 +2853,7 @@ async def update_cache( values_to_update_in_cache.append( ( _id, - CacheCodec.serialize( - existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj - ), + CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj), ) ) except Exception as e: @@ -3074,9 +2903,7 @@ async def update_cache( values_to_update_in_cache.append( ( cache_key, - CacheCodec.serialize( - existing_tag_obj, model_type=LiteLLM_TagTable - ), + CacheCodec.serialize(existing_tag_obj, model_type=LiteLLM_TagTable), ) ) except Exception as e: @@ -3151,9 +2978,7 @@ def _rss_mb_for_log() -> str: def _is_unexpected_keyword_argument_type_error(exc: BaseException) -> bool: """True when ``exc`` is a TypeError from passing a kwarg the callee does not accept.""" - return isinstance(exc, TypeError) and ( - "unexpected keyword argument" in str(exc).lower() - ) + return isinstance(exc, TypeError) and ("unexpected keyword argument" in str(exc).lower()) async def _run_direct_health_check_with_instrumentation( @@ -3205,11 +3030,7 @@ def _schedule_background_health_check_db_save( _save_background_health_checks_to_db, ) - checked_by = ( - shared_health_manager.pod_id - if shared_health_manager is not None - else "background_health_check" - ) + checked_by = shared_health_manager.pod_id if shared_health_manager is not None else "background_health_check" start_time = time_module.time() asyncio.create_task( _save_background_health_checks_to_db( @@ -3263,9 +3084,7 @@ def _write_health_state_to_router_cache( _effective_unhealthy = unhealthy_endpoints if llm_router.health_check_ignore_transient_errors: _effective_unhealthy = [ - ep - for ep in unhealthy_endpoints - if _get_endpoint_exception_status(ep, _exceptions) not in (429, 408) + ep for ep in unhealthy_endpoints if _get_endpoint_exception_status(ep, _exceptions) not in (429, 408) ] states = build_deployment_health_states( @@ -3311,9 +3130,7 @@ def _write_health_state_to_router_cache( ) except Exception as e: - verbose_proxy_logger.warning( - "Failed to write health state to router cache: %s", str(e) - ) + verbose_proxy_logger.warning("Failed to write health state to router cache: %s", str(e)) _ADAPTIVE_ROUTER_FLUSH_INTERVAL_SECONDS = 10 @@ -3360,11 +3177,7 @@ async def _run_background_health_check(): global redis_usage_cache, prisma_client global background_health_check_loop_active, background_health_check_cycle_seq - if ( - health_check_interval is None - or not isinstance(health_check_interval, int) - or health_check_interval <= 0 - ): + if health_check_interval is None or not isinstance(health_check_interval, int) or health_check_interval <= 0: return if background_health_check_loop_active: @@ -3410,17 +3223,11 @@ async def _run_background_health_check(): # filter out models that have disabled background health checks _llm_model_list = [ - m - for m in _llm_model_list - if not m.get("model_info", {}).get("disable_background_health_check", False) + m for m in _llm_model_list if not m.get("model_info", {}).get("disable_background_health_check", False) ] model_count_enabled = len(_llm_model_list) expected_peak_in_flight = model_count_enabled - if ( - isinstance(health_check_concurrency, int) - and health_check_concurrency > 0 - and model_count_enabled > 0 - ): + if isinstance(health_check_concurrency, int) and health_check_concurrency > 0 and model_count_enabled > 0: expected_peak_in_flight = min(model_count_enabled, health_check_concurrency) verbose_proxy_logger.debug( @@ -3444,9 +3251,7 @@ async def _run_background_health_check(): # Use shared health check if available, otherwise fall back to direct health check # Convert health_check_details to bool for perform_shared_health_check (defaults to True if None) - details_bool = ( - health_check_details if health_check_details is not None else True - ) + details_bool = health_check_details if health_check_details is not None else True _hc_filter = health_check_filter_kwargs_from_general_settings(general_settings) if shared_health_manager is not None: @@ -3523,9 +3328,7 @@ async def _run_background_health_check(): ) # Write health state to router cache for health-check-driven routing - _write_health_state_to_router_cache( - healthy_endpoints, unhealthy_endpoints, _exceptions_by_model_id - ) + _write_health_state_to_router_cache(healthy_endpoints, unhealthy_endpoints, _exceptions_by_model_id) await asyncio.sleep(health_check_interval) @@ -3563,9 +3366,7 @@ _DB_OVERLAY_REMOTE_MODULE_LIST_FIELDS: Dict[str, Tuple[str, ...]] = { def _is_remote_module_url(value: Any) -> bool: - return isinstance(value, str) and ( - value.startswith("s3://") or value.startswith("gcs://") - ) + return isinstance(value, str) and (value.startswith("s3://") or value.startswith("gcs://")) def _scrub_guardrail_inner(inner: Dict[str, Any]) -> None: @@ -3576,15 +3377,13 @@ def _scrub_guardrail_inner(inner: Dict[str, Any]) -> None: cleaned = [c for c in cbs if not _is_remote_module_url(c)] if len(cleaned) != len(cbs): verbose_proxy_logger.warning( - "Refused %d remote-URL entries from DB-overlay " - "litellm_settings.guardrails[...].callbacks", + "Refused %d remote-URL entries from DB-overlay litellm_settings.guardrails[...].callbacks", len(cbs) - len(cleaned), ) inner["callbacks"] = cleaned if _is_remote_module_url(inner.get("guardrail")): verbose_proxy_logger.warning( - "Refused remote-URL guardrail module from DB-overlay " - "litellm_settings.guardrails[...].guardrail: %r", + "Refused remote-URL guardrail module from DB-overlay litellm_settings.guardrails[...].guardrail: %r", inner.get("guardrail"), ) inner["guardrail"] = None @@ -3634,12 +3433,9 @@ def _scrub_db_overlay_remote_module_loads(section: str, db_value: Any) -> Any: cpm = sanitized.get("custom_provider_map") if isinstance(cpm, list): for item in cpm: - if isinstance(item, dict) and _is_remote_module_url( - item.get("custom_handler") - ): + if isinstance(item, dict) and _is_remote_module_url(item.get("custom_handler")): verbose_proxy_logger.warning( - "Refused remote-URL custom_handler from DB-overlay " - "litellm_settings.custom_provider_map: %r", + "Refused remote-URL custom_handler from DB-overlay litellm_settings.custom_provider_map: %r", item.get("custom_handler"), ) item["custom_handler"] = None @@ -3669,8 +3465,7 @@ def _scrub_db_overlay_remote_module_loads(section: str, db_value: Any) -> Any: jwt = sanitized.get("litellm_jwtauth") if isinstance(jwt, dict) and _is_remote_module_url(jwt.get("custom_validate")): verbose_proxy_logger.warning( - "Refused remote-URL custom_validate from DB-overlay " - "general_settings.litellm_jwtauth: %r", + "Refused remote-URL custom_validate from DB-overlay general_settings.litellm_jwtauth: %r", jwt.get("custom_validate"), ) jwt["custom_validate"] = None @@ -3682,9 +3477,7 @@ def _scrub_db_overlay_remote_module_loads(section: str, db_value: Any) -> Any: pte = sanitized.get("pass_through_endpoints") if isinstance(pte, list): for entry in pte: - if isinstance(entry, dict) and _is_remote_module_url( - entry.get("target") - ): + if isinstance(entry, dict) and _is_remote_module_url(entry.get("target")): verbose_proxy_logger.warning( "Refused remote-URL target from DB-overlay " "general_settings.pass_through_endpoints " @@ -3724,9 +3517,7 @@ class ProxyConfig: except Exception as e: raise Exception(f"Error loading yaml file {file_path}: {str(e)}") - async def _get_config_from_file( - self, config_file_path: Optional[str] = None - ) -> dict: + async def _get_config_from_file(self, config_file_path: Optional[str] = None) -> dict: """ Given a config file path, load the config from the file. Args: @@ -3757,9 +3548,7 @@ class ProxyConfig: if config is None: raise Exception("Config cannot be None or Empty.") # Process includes - config = self._process_includes( - config=config, base_dir=os.path.dirname(os.path.abspath(file_path or "")) - ) + config = self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) # verbose_proxy_logger.debug(f"loaded config={json.dumps(config, indent=4)}") return config @@ -3813,8 +3602,7 @@ class ProxyConfig: """ if prisma_client is not None and ( - general_settings.get("store_model_in_db", False) is True - or store_model_in_db + general_settings.get("store_model_in_db", False) is True or store_model_in_db ): # if using - db for config - models are in ModelTable @@ -3825,14 +3613,9 @@ class ProxyConfig: # _encrypt_env_variables_for_db is idempotent — a caller that # already encrypted the values (or re-submitted ciphertext read # back from the DB) will not get a stacked second layer. - if ( - "environment_variables" in config_to_save - and config_to_save["environment_variables"] - ): - config_to_save["environment_variables"] = ( - self._encrypt_env_variables_for_db( - environment_variables=config_to_save["environment_variables"] - ) + if "environment_variables" in config_to_save and config_to_save["environment_variables"]: + config_to_save["environment_variables"] = self._encrypt_env_variables_for_db( + environment_variables=config_to_save["environment_variables"] ) config_to_save.pop("model_list", None) @@ -3859,22 +3642,16 @@ class ProxyConfig: dict: Processed configuration dictionary. """ if depth > max_depth: - verbose_proxy_logger.warning( - f"Maximum recursion depth ({max_depth}) reached while processing config." - ) + verbose_proxy_logger.warning(f"Maximum recursion depth ({max_depth}) reached while processing config.") return config for key, value in config.items(): if isinstance(value, dict): - config[key] = self._check_for_os_environ_vars( - config=value, depth=depth + 1, max_depth=max_depth - ) + config[key] = self._check_for_os_environ_vars(config=value, depth=depth + 1, max_depth=max_depth) elif isinstance(value, list): for item in value: if isinstance(item, dict): - item = self._check_for_os_environ_vars( - config=item, depth=depth + 1, max_depth=max_depth - ) + item = self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth) # if the value is a string and starts with "os.environ/" - then it's an environment variable elif isinstance(value, str) and value.startswith("os.environ/"): config[key] = get_secret(value) @@ -3884,9 +3661,7 @@ class ProxyConfig: team_config: dict = {} for team in all_teams_config: if "team_id" not in team: - raise Exception( - f"team_id missing from team: {SENSITIVE_DATA_MASKER.mask_dict(team)}" - ) + raise Exception(f"team_id missing from team: {SENSITIVE_DATA_MASKER.mask_dict(team)}") if team_id == team["team_id"]: team_config = team break @@ -3909,9 +3684,7 @@ class ProxyConfig: all_teams_config = litellm_settings.get("default_team_settings", None) if all_teams_config is None: return {} - team_config = self._get_team_config( - team_id=team_id, all_teams_config=all_teams_config - ) + team_config = self._get_team_config(team_id=team_id, all_teams_config=all_teams_config) return team_config def _init_cache( @@ -3930,9 +3703,7 @@ class ProxyConfig: litellm.cache = Cache(**cache_params) - if litellm.cache is not None and isinstance( - litellm.cache.cache, (RedisCache, RedisClusterCache) - ): + if litellm.cache is not None and isinstance(litellm.cache.cache, (RedisCache, RedisClusterCache)): ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache spend_counter_cache.attach_redis_cache( @@ -3962,9 +3733,7 @@ class ProxyConfig: # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. elif litellm_config_cache.redis_cache is None: - verbose_proxy_logger.info( - "litellm_config_cache: no Redis configured; cluster-wide cache sharing disabled." - ) + verbose_proxy_logger.info("litellm_config_cache: no Redis configured; cluster-wide cache sharing disabled.") def switch_on_llm_response_caching(self): """ @@ -3975,15 +3744,9 @@ class ProxyConfig: global llm_router import litellm - if ( - llm_router is not None - and litellm.cache is not None - and llm_router.cache_responses is not True - ): + if llm_router is not None and litellm.cache is not None and llm_router.cache_responses is not True: llm_router.cache_responses = True - verbose_proxy_logger.debug( - "Set router.cache_responses=True after initializing cache" - ) + verbose_proxy_logger.debug("Set router.cache_responses=True after initializing cache") async def get_config(self, config_file_path: Optional[str] = None) -> dict: """ @@ -4007,17 +3770,11 @@ class ProxyConfig: bucket_name = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") object_key = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") bucket_type = os.environ.get("LITELLM_CONFIG_BUCKET_TYPE") - verbose_proxy_logger.debug( - "bucket_name: %s, object_key: %s", bucket_name, object_key - ) + verbose_proxy_logger.debug("bucket_name: %s, object_key: %s", bucket_name, object_key) if bucket_type == "gcs": - config = await get_config_file_contents_from_gcs( - bucket_name=bucket_name, object_key=object_key - ) + config = await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=object_key) else: - config = get_file_contents_from_s3( - bucket_name=bucket_name, object_key=object_key - ) + config = get_file_contents_from_s3(bucket_name=bucket_name, object_key=object_key) if config is None: raise Exception("Unable to load config from given source.") @@ -4103,9 +3860,7 @@ class ProxyConfig: for search_tool in search_tools_raw: # Display loaded search tool search_tool_name = search_tool.get("search_tool_name", "") - search_provider = search_tool.get("litellm_params", {}).get( - "search_provider", "" - ) + search_provider = search_tool.get("litellm_params", {}).get("search_provider", "") print( # noqa: T201 f"\033[32m {search_tool_name} ({search_provider})\033[0m" ) @@ -4122,14 +3877,10 @@ class ProxyConfig: # Cast to SearchToolTypedDict for type safety try: - search_tool_typed: SearchToolTypedDict = SearchToolTypedDict( - **search_tool - ) # type: ignore + search_tool_typed: SearchToolTypedDict = SearchToolTypedDict(**search_tool) # type: ignore search_tools_parsed.append(search_tool_typed) except Exception as e: - verbose_proxy_logger.error( - f"Error parsing search tool {search_tool_name}: {str(e)}" - ) + verbose_proxy_logger.error(f"Error parsing search tool {search_tool_name}: {str(e)}") continue return search_tools_parsed if search_tools_parsed else None @@ -4160,9 +3911,7 @@ class ProxyConfig: if environment_variables: for key, value in environment_variables.items(): if key in self._BLOCKED_ENV_KEYS: - verbose_proxy_logger.warning( - "Skipping blocked environment variable key: %s", key - ) + verbose_proxy_logger.warning("Skipping blocked environment variable key: %s", key) continue ######################################################### # handles this scenario: @@ -4172,9 +3921,7 @@ class ProxyConfig: # ``` ######################################################### if isinstance(value, str) and value.startswith("os.environ/"): - resolved_secret_string: Optional[str] = get_secret_str( - secret_name=value - ) + resolved_secret_string: Optional[str] = get_secret_str(secret_name=value) if resolved_secret_string is not None: os.environ[key] = resolved_secret_string else: @@ -4193,9 +3940,7 @@ class ProxyConfig: premium_user = _license_check.is_premium() return - async def load_config( - self, router: Optional[litellm.Router], config_file_path: str - ): + async def load_config(self, router: Optional[litellm.Router], config_file_path: str): """ Load config values into proxy global state """ @@ -4261,9 +4006,7 @@ class ProxyConfig: verbose_proxy_logger.debug("passed cache type=%s", cache_type) - if ( - cache_type == "redis" or cache_type == "redis-semantic" - ) and len(cache_params.keys()) == 0: + if (cache_type == "redis" or cache_type == "redis-semantic") and len(cache_params.keys()) == 0: cache_host = get_secret("REDIS_HOST", None) cache_port = get_secret("REDIS_PORT", None) cache_password = None @@ -4317,15 +4060,10 @@ class ProxyConfig: ## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = `, _redis.py checks for REDIS specific environment variables self._init_cache( cache_params=cache_params, - enable_redis_auth_cache=litellm_settings.get( - "enable_redis_auth_cache", False - ) - is True, + enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True, ) if litellm.cache is not None: - verbose_proxy_logger.debug( - f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}" - ) + verbose_proxy_logger.debug(f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}") elif key == "cache" and value is False: pass elif key == "guardrails": @@ -4366,9 +4104,7 @@ class ProxyConfig: elif key == "priority_reservation_settings": from litellm.types.utils import PriorityReservationSettings - litellm.priority_reservation_settings = PriorityReservationSettings( - **value - ) + litellm.priority_reservation_settings = PriorityReservationSettings(**value) elif key == "callbacks": initialize_callbacks_on_proxy( value=value, @@ -4384,12 +4120,8 @@ class ProxyConfig: litellm.model_group_settings = ModelGroupSettings(**value) elif key == "post_call_rules": - litellm.post_call_rules = [ - get_instance_fn(value=value, config_file_path=config_file_path) - ] - verbose_proxy_logger.debug( - f"litellm.post_call_rules: {litellm.post_call_rules}" - ) + litellm.post_call_rules = [get_instance_fn(value=value, config_file_path=config_file_path)] + verbose_proxy_logger.debug(f"litellm.post_call_rules: {litellm.post_call_rules}") elif key == "max_budget": litellm.max_budget = float(value) elif key == "max_internal_user_budget": @@ -4397,9 +4129,7 @@ class ProxyConfig: elif key == "default_max_internal_user_budget": litellm.default_max_internal_user_budget = float(value) if litellm.max_internal_user_budget is None: - litellm.max_internal_user_budget = ( - litellm.default_max_internal_user_budget - ) + litellm.max_internal_user_budget = litellm.default_max_internal_user_budget elif key == "custom_provider_map": from litellm.utils import custom_llm_setup @@ -4430,18 +4160,14 @@ class ProxyConfig: ) # these are litellm callbacks - "langfuse", "sentry", "wandb" else: - litellm.logging_callback_manager.add_litellm_success_callback( - callback - ) + litellm.logging_callback_manager.add_litellm_success_callback(callback) if "prometheus" in callback: from litellm.integrations.prometheus import ( PrometheusLogger, ) if PrometheusLogger is not None: - verbose_proxy_logger.debug( - "mounting metrics endpoint" - ) + verbose_proxy_logger.debug("mounting metrics endpoint") PrometheusLogger._mount_metrics_endpoint() print( # noqa: T201 f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}" @@ -4461,9 +4187,7 @@ class ProxyConfig: ) # these are litellm callbacks - "langfuse", "sentry", "wandb" else: - litellm.logging_callback_manager.add_litellm_failure_callback( - callback - ) + litellm.logging_callback_manager.add_litellm_failure_callback(callback) print( # noqa: T201 f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}" ) @@ -4486,9 +4210,7 @@ class ProxyConfig: else: litellm.audit_log_callbacks.append(callback) - _store_audit_logs = litellm_settings.get( - "store_audit_logs", litellm.store_audit_logs - ) + _store_audit_logs = litellm_settings.get("store_audit_logs", litellm.store_audit_logs) if _store_audit_logs: print( # noqa: T201 f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}" @@ -4504,25 +4226,16 @@ class ProxyConfig: pass elif key == "responses": # Initialize global polling via cache settings - global \ - polling_via_cache_enabled, \ - native_background_mode, \ - polling_cache_ttl + global polling_via_cache_enabled, native_background_mode, polling_cache_ttl background_mode = value.get("background_mode", {}) - polling_via_cache_enabled = background_mode.get( - "polling_via_cache", False - ) - native_background_mode = background_mode.get( - "native_background_mode", [] - ) + polling_via_cache_enabled = background_mode.get("polling_via_cache", False) + native_background_mode = background_mode.get("native_background_mode", []) polling_cache_ttl = background_mode.get("ttl", 3600) verbose_proxy_logger.debug( f"{blue_color_code} Initialized polling via cache: enabled={polling_via_cache_enabled}, native_background_mode={native_background_mode}, ttl={polling_cache_ttl}{reset_color_code}" ) elif key == "default_team_settings": - for idx, team_setting in enumerate( - value - ): # run through pydantic validation + for idx, team_setting in enumerate(value): # run through pydantic validation try: TeamDefaultSettings(**team_setting) except Exception: @@ -4533,32 +4246,22 @@ class ProxyConfig: raise Exception( f"team_id missing from default_team_settings at index={idx}\npassed in value={type(team_setting)}" ) - verbose_proxy_logger.debug( - f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}" - ) + verbose_proxy_logger.debug(f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}") setattr(litellm, key, value) elif key == "upperbound_key_generate_params": if value is not None and isinstance(value, dict): for _k, _v in value.items(): if isinstance(_v, str) and _v.startswith("os.environ/"): value[_k] = get_secret(_v) - litellm.upperbound_key_generate_params = ( - LiteLLM_UpperboundKeyGenerateParams(**value) - ) + litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(**value) else: - raise Exception( - f"Invalid value set for upperbound_key_generate_params - value={value}" - ) + raise Exception(f"Invalid value set for upperbound_key_generate_params - value={value}") elif key == "json_logs" and value is True: litellm.json_logs = True litellm._turn_on_json() - verbose_proxy_logger.debug( - f"{blue_color_code} Enabled JSON logging via config{reset_color_code}" - ) + verbose_proxy_logger.debug(f"{blue_color_code} Enabled JSON logging via config{reset_color_code}") else: - verbose_proxy_logger.debug( - f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}" - ) + verbose_proxy_logger.debug(f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}") setattr(litellm, key, value) if key == "request_timeout": litellm.request_timeout_explicitly_set = True @@ -4572,11 +4275,7 @@ class ProxyConfig: ) reset_audit_log_callback_cache() - _in_memory_loggers[:] = [ - cb - for cb in _in_memory_loggers - if not isinstance(cb, S3V2Logger) - ] + _in_memory_loggers[:] = [cb for cb in _in_memory_loggers if not isinstance(cb, S3V2Logger)] ## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging general_settings = config.get("general_settings", {}) @@ -4587,13 +4286,9 @@ class ProxyConfig: _hc_ignore_transient = False if general_settings: ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ### - key_management_settings = general_settings.get( - "key_management_settings", None - ) + key_management_settings = general_settings.get("key_management_settings", None) if key_management_settings is not None: - litellm._key_management_settings = KeyManagementSettings( - **key_management_settings - ) + litellm._key_management_settings = KeyManagementSettings(**key_management_settings) ### LOAD SECRET MANAGER ### key_management_system = general_settings.get("key_management_system", None) @@ -4618,9 +4313,7 @@ class ProxyConfig: database_url = get_secret(database_url) verbose_proxy_logger.debug("RETRIEVED DB URL: %s", database_url) ### MASTER KEY ### - master_key = general_settings.get( - "master_key", get_secret("LITELLM_MASTER_KEY", None) - ) + master_key = general_settings.get("master_key", get_secret("LITELLM_MASTER_KEY", None)) if master_key and master_key.startswith("os.environ/"): master_key = get_secret(master_key) # type: ignore @@ -4632,9 +4325,7 @@ class ProxyConfig: "LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use." ) ### USER API KEY CACHE TTL (in-memory + Redis when Redis auth sharing is enabled) ### - user_api_key_cache_ttl = general_settings.get( - "user_api_key_cache_ttl", None - ) + user_api_key_cache_ttl = general_settings.get("user_api_key_cache_ttl", None) if user_api_key_cache_ttl is not None: ttl = float(user_api_key_cache_ttl) # Mirror TTL on Redis as well when ``litellm_settings.enable_redis_auth_cache`` @@ -4686,37 +4377,25 @@ class ProxyConfig: ## pass filepath custom_auth = general_settings.get("custom_auth", None) if custom_auth is not None: - user_custom_auth = get_instance_fn( - value=custom_auth, config_file_path=config_file_path - ) + user_custom_auth = get_instance_fn(value=custom_auth, config_file_path=config_file_path) warn_once_if_custom_auth_skips_common_checks( custom_auth_configured=custom_auth is not None, - run_common_checks=bool( - general_settings.get("custom_auth_run_common_checks", False) - ), + run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)), ) custom_key_generate = general_settings.get("custom_key_generate", None) if custom_key_generate is not None: - user_custom_key_generate = get_instance_fn( - value=custom_key_generate, config_file_path=config_file_path - ) + user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path) custom_key_update = general_settings.get("custom_key_update", None) if custom_key_update is not None: - user_custom_key_update = get_instance_fn( - value=custom_key_update, config_file_path=config_file_path - ) + user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path) custom_sso = general_settings.get("custom_sso", None) if custom_sso is not None: - user_custom_sso = get_instance_fn( - value=custom_sso, config_file_path=config_file_path - ) + user_custom_sso = get_instance_fn(value=custom_sso, config_file_path=config_file_path) - custom_ui_sso_sign_in_handler = general_settings.get( - "custom_ui_sso_sign_in_handler", None - ) + custom_ui_sso_sign_in_handler = general_settings.get("custom_ui_sso_sign_in_handler", None) if custom_ui_sso_sign_in_handler is not None: user_custom_ui_sso_sign_in_handler = get_instance_fn( value=custom_ui_sso_sign_in_handler, @@ -4728,18 +4407,14 @@ class ProxyConfig: ## pass through endpoints if general_settings.get("pass_through_endpoints", None) is not None: - config_passthrough_endpoints = general_settings[ - "pass_through_endpoints" - ] + config_passthrough_endpoints = general_settings["pass_through_endpoints"] await initialize_pass_through_endpoints( pass_through_endpoints=general_settings["pass_through_endpoints"], config_file_path=config_file_path, ) ## ADMIN UI ACCESS ## - ui_access_mode = general_settings.get( - "ui_access_mode", "all" - ) # can be either ["admin_only" or "all"] + ui_access_mode = general_settings.get("ui_access_mode", "all") # can be either ["admin_only" or "all"] ### ALLOWED IP ### allowed_ips = general_settings.get("allowed_ips", None) if allowed_ips is not None and premium_user is False: @@ -4758,52 +4433,28 @@ class ProxyConfig: "proxy_batch_polling_interval", proxy_batch_polling_interval ) ## BATCH WRITER ## - proxy_batch_write_at = general_settings.get( - "proxy_batch_write_at", proxy_batch_write_at - ) + proxy_batch_write_at = general_settings.get("proxy_batch_write_at", proxy_batch_write_at) ## DISABLE SPEND LOGS ## - gives a perf improvement - disable_spend_logs = general_settings.get( - "disable_spend_logs", disable_spend_logs - ) + disable_spend_logs = general_settings.get("disable_spend_logs", disable_spend_logs) ### BACKGROUND HEALTH CHECKS ### # Enable background health checks - use_background_health_checks = general_settings.get( - "background_health_checks", False - ) + use_background_health_checks = general_settings.get("background_health_checks", False) # Enable shared health check state across pods (requires Redis) - use_shared_health_check = general_settings.get( - "use_shared_health_check", False - ) - health_check_interval = general_settings.get( - "health_check_interval", DEFAULT_HEALTH_CHECK_INTERVAL - ) - health_check_concurrency = general_settings.get( - "health_check_concurrency", None - ) + use_shared_health_check = general_settings.get("use_shared_health_check", False) + health_check_interval = general_settings.get("health_check_interval", DEFAULT_HEALTH_CHECK_INTERVAL) + health_check_concurrency = general_settings.get("health_check_concurrency", None) health_check_details = general_settings.get("health_check_details", True) ### INTERACTIONS API SCHEMA ### - _use_legacy_interactions_schema = general_settings.get( - "use_legacy_interactions_schema" - ) + _use_legacy_interactions_schema = general_settings.get("use_legacy_interactions_schema") if _use_legacy_interactions_schema is not None: if isinstance(_use_legacy_interactions_schema, str): - litellm.use_legacy_interactions_schema = ( - _use_legacy_interactions_schema.lower() == "true" - ) + litellm.use_legacy_interactions_schema = _use_legacy_interactions_schema.lower() == "true" else: - litellm.use_legacy_interactions_schema = bool( - _use_legacy_interactions_schema - ) + litellm.use_legacy_interactions_schema = bool(_use_legacy_interactions_schema) # Health-check-driven routing (opt-in, passes through to Router later) - _enable_hc_routing = general_settings.get( - "enable_health_check_routing", False - ) - _hc_staleness = general_settings.get( - "health_check_staleness_threshold", None - ) - _hc_ignore_transient = general_settings.get( - "health_check_ignore_transient_errors", False - ) + _enable_hc_routing = general_settings.get("enable_health_check_routing", False) + _hc_staleness = general_settings.get("health_check_staleness_threshold", None) + _hc_ignore_transient = general_settings.get("health_check_ignore_transient_errors", False) verbose_proxy_logger.info( "background_health_check_config enabled=%s shared=%s interval_seconds=%s max_concurrency=%s details=%s health_check_routing=%s", use_background_health_checks, @@ -4818,19 +4469,12 @@ class ProxyConfig: rbac_role_permissions = general_settings.get("role_permissions", None) if rbac_role_permissions is not None: general_settings["role_permissions"] = [ # validate role permissions - RoleBasedPermissions(**role_permission) - for role_permission in rbac_role_permissions + RoleBasedPermissions(**role_permission) for role_permission in rbac_role_permissions ] ## check if user has set a premium feature in general_settings - if ( - general_settings.get("enforced_params") is not None - and premium_user is not True - ): - raise ValueError( - "Trying to use `enforced_params`" - + CommonProxyErrors.not_premium_user.value - ) + if general_settings.get("enforced_params") is not None and premium_user is not True: + raise ValueError("Trying to use `enforced_params`" + CommonProxyErrors.not_premium_user.value) # check if litellm_license in general_settings if "litellm_license" in general_settings: @@ -4838,8 +4482,7 @@ class ProxyConfig: premium_user = _license_check.is_premium() router_params: dict = { - "cache_responses": litellm.cache - is not None, # cache if user passed in cache values + "cache_responses": litellm.cache is not None, # cache if user passed in cache values } # Health-check-driven routing params (from general_settings) if _enable_hc_routing: @@ -4878,9 +4521,7 @@ class ProxyConfig: assistants_config = AssistantsTypedDict(**assistant_settings) # type: ignore ## SEARCH TOOLS SETTINGS - search_tools: Optional[List[SearchToolTypedDict]] = self.parse_search_tools( - config - ) + search_tools: Optional[List[SearchToolTypedDict]] = self.parse_search_tools(config) ## SANDBOX TOOLS SETTINGS from litellm.sandbox.sandbox_tools import register_sandbox_tools @@ -4897,9 +4538,7 @@ class ProxyConfig: ## default config for vertex ai routes default_vertex_config = config.get("default_vertex_config", None) - passthrough_endpoint_router.set_default_vertex_config( - config=default_vertex_config - ) + passthrough_endpoint_router.set_default_vertex_config(config=default_vertex_config) ## ROUTER SETTINGS (e.g. routing_strategy, ...) router_settings = config.get("router_settings", None) @@ -4911,9 +4550,7 @@ class ProxyConfig: "search_tools", } - available_args = [ - x for x in litellm.Router.get_valid_args() if x not in exclude_args - ] + available_args = [x for x in litellm.Router.get_valid_args() if x not in exclude_args] for k, v in router_settings.items(): if k in available_args: @@ -4972,15 +4609,11 @@ class ProxyConfig: litellm.credential_list = credential_list_dict ## NON-LLM CONFIGS eg. MCP tools, vector stores, etc. - await self._init_non_llm_configs( - config=config, config_file_path=config_file_path - ) + await self._init_non_llm_configs(config=config, config_file_path=config_file_path) return router, router.get_model_list(), general_settings - async def _init_non_llm_configs( - self, config: dict, config_file_path: Optional[str] = None - ): + async def _init_non_llm_configs(self, config: dict, config_file_path: Optional[str] = None): """ Initialize non-LLM configs eg. MCP tools, vector stores, etc. """ @@ -4991,9 +4624,7 @@ class ProxyConfig: global_mcp_tool_registry, ) - global_mcp_tool_registry.load_tools_from_config( - mcp_tools_config, config_file_path=config_file_path - ) + global_mcp_tool_registry.load_tools_from_config(mcp_tools_config, config_file_path=config_file_path) ## AGENTS agent_config = config.get("agent_list", None) @@ -5014,9 +4645,7 @@ class ProxyConfig: litellm_settings = config.get("litellm_settings", {}) mcp_aliases = litellm_settings.get("mcp_aliases", None) - await global_mcp_server_manager.load_servers_from_config( - mcp_servers_config, mcp_aliases - ) + await global_mcp_server_manager.load_servers_from_config(mcp_servers_config, mcp_aliases) ## VECTOR STORES vector_store_registry_config = config.get("vector_store_registry", None) @@ -5027,16 +4656,12 @@ class ProxyConfig: litellm.vector_store_registry = VectorStoreRegistry() # Load vector stores from config - litellm.vector_store_registry.load_vector_stores_from_config( - vector_store_registry_config - ) + litellm.vector_store_registry.load_vector_stores_from_config(vector_store_registry_config) ## WORKER REGISTRY (Control Plane) worker_registry_config = config.get("worker_registry", None) if worker_registry_config: - self.worker_registry = [ - WorkerRegistryEntry(**e) for e in worker_registry_config - ] + self.worker_registry = [WorkerRegistryEntry(**e) for e in worker_registry_config] else: self.worker_registry = [] @@ -5069,9 +4694,7 @@ class ProxyConfig: policy_attachments_config = config.get("policy_attachments", None) - verbose_proxy_logger.info( - f"Policy engine: found {len(policies_config)} policies in config" - ) + verbose_proxy_logger.info(f"Policy engine: found {len(policies_config)} policies in config") # Initialize policies await init_policies( @@ -5094,9 +4717,7 @@ class ProxyConfig: # Ensure proxy_logging_obj.alerting is set for all alerting types _alerting_value = general_settings.get("alerting", None) - verbose_proxy_logger.debug( - f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}" - ) + verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}") proxy_logging_obj.update_values( alerting=_alerting_value, alerting_threshold=general_settings.get("alerting_threshold", 600), @@ -5118,9 +4739,7 @@ class ProxyConfig: logging_integration=_alert, internal_usage_cache=None, llm_router=None, - custom_logger_init_args={ - "alerting_args": general_settings.get("alerting_args", None) - }, + custom_logger_init_args={"alerting_args": general_settings.get("alerting_args", None)}, ) if _logger is not None: litellm.logging_callback_manager.add_litellm_callback(_logger) @@ -5154,9 +4773,7 @@ class ProxyConfig: ) elif key_management_system == KeyManagementSystem.AWS_KMS.value: load_aws_kms(use_aws_kms=True) - elif ( - key_management_system == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value - ): + elif key_management_system == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: from litellm.secret_managers.google_secret_manager import ( GoogleSecretManager, ) @@ -5286,15 +4903,11 @@ class ProxyConfig: deleted_deployments += 1 return deleted_deployments - def _resolve_db_litellm_param( - self, key: str, value: object, resolve_env_refs: bool = True - ) -> object: + def _resolve_db_litellm_param(self, key: str, value: object, resolve_env_refs: bool = True) -> object: if not isinstance(value, str): return value - decrypted_value = decrypt_value_helper( - value=value, key=key, return_original_value=True - ) + decrypted_value = decrypt_value_helper(value=value, key=key, return_original_value=True) if ( resolve_env_refs and key in _DB_LITELLM_PARAM_ENV_REF_KEYS @@ -5335,9 +4948,7 @@ class ProxyConfig: f"Invalid model added to proxy db. Invalid litellm params. litellm_params={_litellm_params}" ) continue # skip to next model - _model_info = self.get_model_info_with_id( - model=m, db_model=True - ) ## 👈 FLAG = True for db_models + _model_info = self.get_model_info_with_id(model=m, db_model=True) ## 👈 FLAG = True for db_models added = llm_router.upsert_deployment( deployment=Deployment( @@ -5416,9 +5027,7 @@ class ProxyConfig: if llm_router is None and master_key is not None: verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") - _model_list: list = self.decrypt_model_list_from_db( - new_models=models_list - ) + _model_list: list = self.decrypt_model_list_from_db(new_models=models_list) # Only create router if we have models or search_tools to route # Router can function with model_list=[] if search_tools are configured if len(_model_list) > 0 or search_tools: @@ -5443,9 +5052,7 @@ class ProxyConfig: self._add_deployment(db_models=models_list) except Exception as e: - verbose_proxy_logger.exception( - f"Error adding/deleting model to llm_router: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error adding/deleting model to llm_router: {str(e)}") if llm_router is not None: llm_model_list = llm_router.get_model_list() @@ -5523,17 +5130,13 @@ class ProxyConfig: existing_callbacks=litellm.callbacks, ) - def _encrypt_env_variables( - self, environment_variables: dict, new_encryption_key: Optional[str] = None - ) -> dict: + def _encrypt_env_variables(self, environment_variables: dict, new_encryption_key: Optional[str] = None) -> dict: """ Encrypts a dictionary of environment variables and returns them. """ encrypted_env_vars = {} for k, v in environment_variables.items(): - encrypted_value = encrypt_value_helper( - value=v, new_encryption_key=new_encryption_key - ) + encrypted_value = encrypt_value_helper(value=v, new_encryption_key=new_encryption_key) encrypted_env_vars[k] = encrypted_value return encrypted_env_vars @@ -5550,16 +5153,12 @@ class ProxyConfig: decrypted_env_vars = {} for k, v in environment_variables.items(): try: - decrypted_value = decrypt_value_helper( - value=v, key=k, return_original_value=return_original_value - ) + decrypted_value = decrypt_value_helper(value=v, key=k, return_original_value=return_original_value) if decrypted_value is not None: os.environ[k] = decrypted_value decrypted_env_vars[k] = decrypted_value except Exception as e: - verbose_proxy_logger.error( - "Error setting env variable: %s - %s", k, str(e) - ) + verbose_proxy_logger.error("Error setting env variable: %s - %s", k, str(e)) return decrypted_env_vars def _decrypt_db_variables(self, variables_dict: dict) -> dict: @@ -5568,9 +5167,7 @@ class ProxyConfig: """ decrypted_variables = {} for k, v in variables_dict.items(): - decrypted_value = decrypt_value_helper( - value=v, key=k, return_original_value=True - ) + decrypted_value = decrypt_value_helper(value=v, key=k, return_original_value=True) decrypted_variables[k] = decrypted_value return decrypted_variables @@ -5647,9 +5244,7 @@ class ProxyConfig: # user_api_key_dict is already the cached/authenticated key object — # no DB call needed. if user_api_key_dict is not None: - key_settings = self._parse_router_settings_value( - getattr(user_api_key_dict, "router_settings", None) - ) + key_settings = self._parse_router_settings_value(getattr(user_api_key_dict, "router_settings", None)) if key_settings is not None: return key_settings @@ -5664,9 +5259,7 @@ class ProxyConfig: user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - team_settings = self._parse_router_settings_value( - getattr(team_obj, "router_settings", None) - ) + team_settings = self._parse_router_settings_value(getattr(team_obj, "router_settings", None)) if team_settings is not None: return team_settings except Exception: @@ -5705,16 +5298,10 @@ class ProxyConfig: ): from litellm.utils import _update_dictionary - combined_router_settings = _update_dictionary( - config_router_settings, db_router_settings.param_value - ) - elif config_router_settings is not None and isinstance( - config_router_settings, dict - ): + combined_router_settings = _update_dictionary(config_router_settings, db_router_settings.param_value) + elif config_router_settings is not None and isinstance(config_router_settings, dict): combined_router_settings = config_router_settings - elif db_router_settings is not None and isinstance( - db_router_settings.param_value, dict - ): + elif db_router_settings is not None and isinstance(db_router_settings.param_value, dict): combined_router_settings = db_router_settings.param_value if combined_router_settings: @@ -5747,9 +5334,7 @@ class ProxyConfig: _merged_alerting = list(_yaml_alerting.union(_db_alerting)) # Preserve order: YAML values first, then DB values _merged_alerting = list(general_settings["alerting"]) + [ - item - for item in _general_settings["alerting"] - if item not in general_settings["alerting"] + item for item in _general_settings["alerting"] if item not in general_settings["alerting"] ] verbose_proxy_logger.debug( f"Merging alerting values: YAML={general_settings['alerting']}, DB={_general_settings['alerting']}, Merged={_merged_alerting}" @@ -5780,13 +5365,8 @@ class ProxyConfig: alert_types=general_settings["alert_types"], llm_router=llm_router ) - if ( - _general_settings is not None - and "alert_to_webhook_url" in _general_settings - ): - general_settings["alert_to_webhook_url"] = _general_settings[ - "alert_to_webhook_url" - ] + if _general_settings is not None and "alert_to_webhook_url" in _general_settings: + general_settings["alert_to_webhook_url"] = _general_settings["alert_to_webhook_url"] proxy_logging_obj.slack_alerting_instance.update_values( alert_to_webhook_url=general_settings["alert_to_webhook_url"], llm_router=llm_router, @@ -5836,22 +5416,16 @@ class ProxyConfig: replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info( - f"Spend log cleanup rescheduled with cron: {cleanup_cron}" - ) + verbose_proxy_logger.info(f"Spend log cleanup rescheduled with cron: {cleanup_cron}") except ValueError: - verbose_proxy_logger.error( - f"Invalid maximum_spend_logs_cleanup_cron value: {cleanup_cron}" - ) + verbose_proxy_logger.error(f"Invalid maximum_spend_logs_cleanup_cron value: {cleanup_cron}") else: # Interval-based scheduling (existing behavior) from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, ) - retention_interval = general_settings.get( - "maximum_spend_logs_retention_interval", "1d" - ) + retention_interval = general_settings.get("maximum_spend_logs_retention_interval", "1d") try: interval_seconds = duration_in_seconds(retention_interval) scheduler.add_job( @@ -5863,13 +5437,9 @@ class ProxyConfig: replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info( - f"Spend log cleanup rescheduled with interval: {retention_interval}" - ) + verbose_proxy_logger.info(f"Spend log cleanup rescheduled with interval: {retention_interval}") except ValueError: - verbose_proxy_logger.error( - "Invalid maximum_spend_logs_retention_interval value" - ) + verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") async def _update_general_settings(self, db_general_settings: Optional[Json]): """ @@ -5881,14 +5451,10 @@ class ProxyConfig: _general_settings = dict(db_general_settings) ## MAX PARALLEL REQUESTS ## if "max_parallel_requests" in _general_settings: - general_settings["max_parallel_requests"] = _general_settings[ - "max_parallel_requests" - ] + general_settings["max_parallel_requests"] = _general_settings["max_parallel_requests"] if "global_max_parallel_requests" in _general_settings: - general_settings["global_max_parallel_requests"] = _general_settings[ - "global_max_parallel_requests" - ] + general_settings["global_max_parallel_requests"] = _general_settings["global_max_parallel_requests"] ## ALERTING ARGS ## if "alerting_args" in _general_settings: @@ -5899,12 +5465,8 @@ class ProxyConfig: ## PASS-THROUGH ENDPOINTS ## if "pass_through_endpoints" in _general_settings: - general_settings["pass_through_endpoints"] = _general_settings[ - "pass_through_endpoints" - ] - await initialize_pass_through_endpoints( - pass_through_endpoints=general_settings["pass_through_endpoints"] - ) + general_settings["pass_through_endpoints"] = _general_settings["pass_through_endpoints"] + await initialize_pass_through_endpoints(pass_through_endpoints=general_settings["pass_through_endpoints"]) ## UI ACCESS MODE ## if "ui_access_mode" in _general_settings: @@ -5920,9 +5482,7 @@ class ProxyConfig: general_settings["store_prompts_in_spend_logs"] = value elif isinstance(value, str): # Case-insensitive string comparison - general_settings["store_prompts_in_spend_logs"] = ( - value.lower() == "true" - ) + general_settings["store_prompts_in_spend_logs"] = value.lower() == "true" else: # For other types, convert to bool general_settings["store_prompts_in_spend_logs"] = bool(value) @@ -5997,14 +5557,10 @@ class ProxyConfig: # DB-sourced ``s3://`` value would otherwise reach # ``_load_instance_from_remote_storage`` without going through # the runtime gate. - db_param_value = _scrub_db_overlay_remote_module_loads( - section=param_name, db_value=db_param_value - ) + db_param_value = _scrub_db_overlay_remote_module_loads(section=param_name, db_value=db_param_value) if param_name == "environment_variables": - decrypted_env_vars = self._decrypt_and_set_db_env_variables( - db_param_value, return_original_value=True - ) + decrypted_env_vars = self._decrypt_and_set_db_env_variables(db_param_value, return_original_value=True) # Normalize keys when loading from DB so services expecting uppercase # (e.g. Datadog) can read them even if stored in lowercase. merged_env_vars: dict = {} @@ -6014,15 +5570,11 @@ class ProxyConfig: merged_env_vars[upper_key] = value os.environ[upper_key] = value - current_config.setdefault("environment_variables", {}).update( - merged_env_vars - ) + current_config.setdefault("environment_variables", {}).update(merged_env_vars) return current_config elif param_name == "litellm_settings" and isinstance(db_param_value, dict): for key, value in db_param_value.items(): - if ( - key in LITELLM_SETTINGS_SAFE_DB_OVERRIDES - ): # params that are safe to override with db values + if key in LITELLM_SETTINGS_SAFE_DB_OVERRIDES: # params that are safe to override with db values setattr(litellm, key, value) # If param doesn't exist in config, add it @@ -6032,9 +5584,7 @@ class ProxyConfig: return current_config # For dictionary values, update only non-none values - if isinstance(current_config[param_name], dict) and isinstance( - db_param_value, dict - ): + if isinstance(current_config[param_name], dict) and isinstance(db_param_value, dict): _deep_merge_dicts(current_config[param_name], db_param_value) else: # Non-dict or mismatched types: DB value replaces config (unchanged behavior) @@ -6049,9 +5599,7 @@ class ProxyConfig: store_model_in_db: Optional[bool], ): if store_model_in_db is not True: - verbose_proxy_logger.info( - "'store_model_in_db' is not True, skipping db updates" - ) + verbose_proxy_logger.info("'store_model_in_db' is not True, skipping db updates") return config _tasks = [] @@ -6071,9 +5619,7 @@ class ProxyConfig: param_name = getattr(response, "param_name", None) param_value = getattr(response, "param_value", None) - verbose_proxy_logger.debug( - f"param_name={param_name}, param_value={param_value}" - ) + verbose_proxy_logger.debug(f"param_name={param_name}, param_value={param_value}") if param_name is not None and param_value is not None: config = self._update_config_fields( @@ -6084,9 +5630,7 @@ class ProxyConfig: return config - def _should_load_db_object( - self, object_type: Union[str, SupportedDBObjectType] - ) -> bool: + def _should_load_db_object(self, object_type: Union[str, SupportedDBObjectType]) -> bool: """ Check if an object type should be loaded from the database based on general_settings.supported_db_objects. @@ -6132,9 +5676,7 @@ class ProxyConfig: return new_models except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format( - str(e) - ) + "litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format(str(e)) ) return None @@ -6169,13 +5711,9 @@ class ProxyConfig: new_models = await self._get_models_from_db(prisma_client=prisma_client) # update llm router - await self._update_llm_router( - new_models=new_models, proxy_logging_obj=proxy_logging_obj - ) + await self._update_llm_router(new_models=new_models, proxy_logging_obj=proxy_logging_obj) - db_general_settings = await get_config_param( - prisma_client, "general_settings" - ) + db_general_settings = await get_config_param(prisma_client, "general_settings") # update general settings if db_general_settings is not None: @@ -6188,9 +5726,7 @@ class ProxyConfig: except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - {}".format(str(e)) ) async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): @@ -6233,9 +5769,7 @@ class ProxyConfig: await self._check_and_reload_model_cost_map(prisma_client=prisma_client) if self._should_load_db_object(object_type="anthropic_beta_headers"): - await self._check_and_reload_anthropic_beta_headers( - prisma_client=prisma_client - ) + await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) if self._should_load_db_object(object_type="sso_settings"): await self._init_sso_settings_in_db(prisma_client=prisma_client) @@ -6244,17 +5778,13 @@ class ProxyConfig: CacheSettingsManager, ) - await CacheSettingsManager.init_cache_settings_in_db( - prisma_client=prisma_client, proxy_config=self - ) + await CacheSettingsManager.init_cache_settings_in_db(prisma_client=prisma_client, proxy_config=self) if self._should_load_db_object(object_type="semantic_filter_settings"): await self._init_semantic_filter_settings_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="config_overrides"): - await self._init_hashicorp_vault_config_override( - prisma_client=prisma_client - ) + await self._init_hashicorp_vault_config_override(prisma_client=prisma_client) async def _init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient): """ @@ -6276,9 +5806,7 @@ class ProxyConfig: if isinstance(litellm_settings, str): litellm_settings = json.loads(litellm_settings) - mcp_semantic_filter_config = litellm_settings.get( - "mcp_semantic_tool_filter", None - ) + mcp_semantic_filter_config = litellm_settings.get("mcp_semantic_tool_filter", None) if mcp_semantic_filter_config is None: return @@ -6287,18 +5815,11 @@ class ProxyConfig: if hasattr(self, "_last_semantic_filter_config"): if self._last_semantic_filter_config == mcp_semantic_filter_config: # If hook is missing or router isn't built yet, reinitialize anyway - active_hooks = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - SemanticToolFilterHook - ) - ) + active_hooks = litellm.logging_callback_manager.get_custom_loggers_for_type(SemanticToolFilterHook) if active_hooks: for active_hook in active_hooks: if isinstance(active_hook, SemanticToolFilterHook): - if ( - active_hook.filter is not None - and active_hook.filter.tool_router is not None - ): + if active_hook.filter is not None and active_hook.filter.tool_router is not None: verbose_proxy_logger.debug( "Semantic filter settings unchanged, skipping reinitialization" ) @@ -6308,9 +5829,7 @@ class ProxyConfig: ) # Remove old hooks using logging callback manager - litellm.logging_callback_manager.remove_callbacks_by_type( - litellm.callbacks, SemanticToolFilterHook - ) + litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, SemanticToolFilterHook) # Initialize new hook if enabled if mcp_semantic_filter_config.get("enabled", False): @@ -6321,9 +5840,7 @@ class ProxyConfig: ) if hook: litellm.logging_callback_manager.add_litellm_callback(hook) - verbose_proxy_logger.info( - "MCP Semantic Filter reinitialized from DB" - ) + verbose_proxy_logger.info("MCP Semantic Filter reinitialized from DB") else: verbose_proxy_logger.info("MCP Semantic Filter disabled") @@ -6331,9 +5848,7 @@ class ProxyConfig: self._last_semantic_filter_config = mcp_semantic_filter_config.copy() except Exception as e: - verbose_proxy_logger.exception( - f"Error initializing semantic filter settings from DB: {e}" - ) + verbose_proxy_logger.exception(f"Error initializing semantic filter settings from DB: {e}") async def _init_sso_settings_in_db(self, prisma_client: PrismaClient): """ @@ -6343,27 +5858,18 @@ class ProxyConfig: try: sso_settings = await call_with_db_reconnect_retry( prisma_client, - lambda: SSOConfigRepository(prisma_client).table.find_unique( - where={"id": "sso_config"} - ), + lambda: SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}), reason="init_sso_settings_in_db_lookup_failure", ) if sso_settings is not None: sso_settings.sso_settings.pop("role_mappings", None) sso_settings.sso_settings.pop("team_mappings", None) sso_settings.sso_settings.pop("ui_access_mode", None) - uppercase_sso_settings = { - key.upper(): value - for key, value in sso_settings.sso_settings.items() - } - self._decrypt_and_set_db_env_variables( - environment_variables=uppercase_sso_settings - ) + uppercase_sso_settings = {key.upper(): value for key, value in sso_settings.sso_settings.items()} + self._decrypt_and_set_db_env_variables(environment_variables=uppercase_sso_settings) except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_sso_settings_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_sso_settings_in_db - {}".format(str(e)) ) async def _init_hashicorp_vault_config_override(self, prisma_client: PrismaClient): @@ -6430,9 +5936,7 @@ class ProxyConfig: """ try: # Get model cost map reload configuration from database - config_record = await get_config_param( - prisma_client, "model_cost_map_reload_config" - ) + config_record = await get_config_param(prisma_client, "model_cost_map_reload_config") if config_record is None or config_record.param_value is None: return # No configuration found, skip reload @@ -6451,21 +5955,15 @@ class ProxyConfig: if force_reload: should_reload = True - verbose_proxy_logger.info( - "Model cost map reload triggered by force reload flag" - ) + verbose_proxy_logger.info("Model cost map reload triggered by force reload flag") elif interval_hours is not None: # Use pod's in-memory last reload time global last_model_cost_map_reload if last_model_cost_map_reload is not None: try: - last_reload_time = datetime.fromisoformat( - last_model_cost_map_reload - ) + last_reload_time = datetime.fromisoformat(last_model_cost_map_reload) time_since_last_reload = current_time - last_reload_time - hours_since_last_reload = ( - time_since_last_reload.total_seconds() / 3600 - ) + hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 if hours_since_last_reload >= interval_hours: should_reload = True @@ -6473,17 +5971,13 @@ class ProxyConfig: f"Model cost map reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}" ) except Exception as e: - verbose_proxy_logger.warning( - f"Error parsing last reload time: {e}" - ) + verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") # If we can't parse the last reload time, reload anyway should_reload = True else: # No last reload time recorded, reload now should_reload = True - verbose_proxy_logger.info( - "Model cost map reload triggered - no previous reload time recorded" - ) + verbose_proxy_logger.info("Model cost map reload triggered - no previous reload time recorded") if should_reload: # Perform the reload @@ -6533,22 +6027,16 @@ class ProxyConfig: ) except Exception as e: - verbose_proxy_logger.exception( - f"Error in _check_and_reload_model_cost_map: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {str(e)}") - async def _check_and_reload_anthropic_beta_headers( - self, prisma_client: PrismaClient - ): + async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient): """ Check if anthropic beta headers config needs to be reloaded based on database configuration. This function runs every 10 seconds as part of _init_non_llm_objects_in_db. """ try: # Get anthropic beta headers reload configuration from database - config_record = await get_config_param( - prisma_client, "anthropic_beta_headers_reload_config" - ) + config_record = await get_config_param(prisma_client, "anthropic_beta_headers_reload_config") if config_record is None or config_record.param_value is None: return # No configuration found, skip reload @@ -6567,21 +6055,15 @@ class ProxyConfig: if force_reload: should_reload = True - verbose_proxy_logger.info( - "Anthropic beta headers reload triggered by force reload flag" - ) + verbose_proxy_logger.info("Anthropic beta headers reload triggered by force reload flag") elif interval_hours is not None: # Use pod's in-memory last reload time global last_anthropic_beta_headers_reload if last_anthropic_beta_headers_reload is not None: try: - last_reload_time = datetime.fromisoformat( - last_anthropic_beta_headers_reload - ) + last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload) time_since_last_reload = current_time - last_reload_time - hours_since_last_reload = ( - time_since_last_reload.total_seconds() / 3600 - ) + hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 if hours_since_last_reload >= interval_hours: should_reload = True @@ -6589,9 +6071,7 @@ class ProxyConfig: f"Anthropic beta headers reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}" ) except Exception as e: - verbose_proxy_logger.warning( - f"Error parsing last reload time: {e}" - ) + verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") # If we can't parse the last reload time, reload anyway should_reload = True else: @@ -6638,19 +6118,13 @@ class ProxyConfig: await invalidate_config_param("anthropic_beta_headers_reload_config") # Count providers in config - provider_count = sum( - 1 - for k in new_config.keys() - if k != "provider_aliases" and k != "description" - ) + provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") verbose_proxy_logger.info( f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}" ) except Exception as e: - verbose_proxy_logger.exception( - f"Error in _check_and_reload_anthropic_beta_headers: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error in _check_and_reload_anthropic_beta_headers: {str(e)}") def _get_prompt_spec_for_db_prompt(self, db_prompt): """ @@ -6680,9 +6154,7 @@ class ProxyConfig: IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec) except Exception as e: verbose_proxy_logger.debug( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - {}".format(str(e)) ) async def _init_guardrails_in_db(self, prisma_client: PrismaClient): @@ -6693,14 +6165,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( + guardrails_in_db: List[Guardrail] = await GuardrailRegistry.get_all_guardrails_from_db( prisma_client=prisma_client ) - verbose_proxy_logger.debug( - "guardrails from the DB %s", str(guardrails_in_db) - ) + verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db)) db_guardrail_ids: set = set() for guardrail in guardrails_in_db: guardrail_id = guardrail.get("guardrail_id") @@ -6712,14 +6180,10 @@ class ProxyConfig: # Drop in-memory DB-backed entries whose row was deleted on another # pod. Config-loaded entries are never touched. - IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails( - db_guardrail_ids=db_guardrail_ids - ) + IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids) except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - {}".format(str(e)) ) async def _init_policies_in_db(self, prisma_client: PrismaClient): @@ -6740,18 +6204,12 @@ class ProxyConfig: await policy_registry.sync_policies_from_db(prisma_client=prisma_client) # Sync attachments from DB to in-memory registry - await attachment_registry.sync_attachments_from_db( - prisma_client=prisma_client - ) + await attachment_registry.sync_attachments_from_db(prisma_client=prisma_client) - verbose_proxy_logger.debug( - "Successfully synced policies and attachments from DB" - ) + verbose_proxy_logger.debug("Successfully synced policies and attachments from DB") except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_policies_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_policies_in_db - {}".format(str(e)) ) async def _init_tool_policy_in_db(self, prisma_client: PrismaClient): @@ -6767,9 +6225,7 @@ class ProxyConfig: verbose_proxy_logger.debug("Successfully synced tool policy from DB") except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_tool_policy_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_tool_policy_in_db - {}".format(str(e)) ) async def _init_vector_stores_in_db(self, prisma_client: PrismaClient): @@ -6777,26 +6233,18 @@ class ProxyConfig: try: # read vector stores from db table - vector_stores = await VectorStoreRegistry._get_vector_stores_from_db( - prisma_client=prisma_client - ) + vector_stores = await VectorStoreRegistry._get_vector_stores_from_db(prisma_client=prisma_client) if len(vector_stores) <= 0: return if litellm.vector_store_registry is None: - litellm.vector_store_registry = VectorStoreRegistry( - vector_stores=vector_stores - ) + litellm.vector_store_registry = VectorStoreRegistry(vector_stores=vector_stores) else: for vector_store in vector_stores: - litellm.vector_store_registry.add_vector_store_to_registry( - vector_store=vector_store - ) + litellm.vector_store_registry.add_vector_store_to_registry(vector_store=vector_store) except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {}".format(str(e)) ) async def _init_vector_store_indexes_in_db(self, prisma_client: PrismaClient): @@ -6804,10 +6252,8 @@ class ProxyConfig: try: # read vector stores from db table - vector_store_indexes = ( - await VectorStoreIndexRegistry._get_vector_store_indexes_from_db( - prisma_client=prisma_client - ) + vector_store_indexes = await VectorStoreIndexRegistry._get_vector_store_indexes_from_db( + prisma_client=prisma_client ) if len(vector_store_indexes) <= 0: @@ -6819,23 +6265,17 @@ class ProxyConfig: ) else: for vector_store_index in vector_store_indexes: - litellm.vector_store_index_registry.upsert_vector_store_index( - vector_store_index=vector_store_index - ) + litellm.vector_store_index_registry.upsert_vector_store_index(vector_store_index=vector_store_index) except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_vector_stores_in_db - {}".format(str(e)) ) async def _init_mcp_servers_in_db(self): from litellm.proxy._experimental.mcp_server.utils import is_mcp_available if not is_mcp_available(): - verbose_proxy_logger.debug( - "MCP module not available, skipping MCP server initialization" - ) + verbose_proxy_logger.debug("MCP module not available, skipping MCP server initialization") return from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -6846,9 +6286,7 @@ class ProxyConfig: await global_mcp_server_manager.reload_servers_from_database() except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format(str(e)) ) async def _init_agents_in_db(self, prisma_client: PrismaClient): @@ -6857,17 +6295,11 @@ class ProxyConfig: ) try: - db_agents = await AGENT_REGISTRY.get_all_agents_from_db( - prisma_client=prisma_client - ) - AGENT_REGISTRY.load_agents_from_db_and_config( - db_agents=db_agents, agent_config=config_agents - ) + db_agents = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client) + AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents, agent_config=config_agents) except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - {}".format(str(e)) ) async def _init_search_tools_in_db(self, prisma_client: PrismaClient): @@ -6883,13 +6315,9 @@ class ProxyConfig: from litellm.router_utils.search_api_router import SearchAPIRouter try: - search_tools = await SearchToolRegistry.get_all_search_tools_from_db( - prisma_client=prisma_client - ) + search_tools = await SearchToolRegistry.get_all_search_tools_from_db(prisma_client=prisma_client) - verbose_proxy_logger.info( - f"Loading {len(search_tools)} search tool(s) from database into router" - ) + verbose_proxy_logger.info(f"Loading {len(search_tools)} search tool(s) from database into router") # Only update router if there are tools in the database # This prevents overwriting config-loaded tools with an empty list @@ -6899,9 +6327,7 @@ class ProxyConfig: await SearchAPIRouter.update_router_search_tools( router_instance=llm_router, search_tools=search_tools ) - verbose_proxy_logger.info( - f"Successfully loaded {len(search_tools)} search tool(s) into router" - ) + verbose_proxy_logger.info(f"Successfully loaded {len(search_tools)} search tool(s) into router") else: verbose_proxy_logger.debug( "Router not initialized yet, search tools will be added when router is created" @@ -6913,9 +6339,7 @@ class ProxyConfig: except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - {}".format(str(e)) ) async def _init_pass_through_endpoints_in_db(self): @@ -6954,9 +6378,7 @@ class ProxyConfig: ## DELETE ## idx_to_delete = [] for idx, credential in enumerate(litellm.credential_list): - if credential.credential_name not in [ - cred.credential_name for cred in combined_list - ]: + if credential.credential_name not in [cred.credential_name for cred in combined_list]: idx_to_delete.append(idx) for idx in sorted(idx_to_delete, reverse=True): litellm.credential_list.pop(idx) @@ -6965,17 +6387,11 @@ class ProxyConfig: try: credentials = await CredentialsRepository(prisma_client).find_all() credentials = [self.decrypt_credentials(cred) for cred in credentials] - await self.delete_credentials( - credentials - ) # delete credentials that are not in the all-up list - CredentialAccessor.upsert_credentials( - credentials - ) # upsert credentials that are in the all-up list + await self.delete_credentials(credentials) # delete credentials that are not in the all-up list + CredentialAccessor.upsert_credentials(credentials) # upsert credentials that are in the all-up list except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy_server.py::get_credentials() - Error getting credentials from DB - {}".format( - str(e) - ) + "litellm.proxy_server.py::get_credentials() - Error getting credentials from DB - {}".format(str(e)) ) return [] @@ -7073,12 +6489,8 @@ async def initialize( # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS verbose_logger.setLevel(level=logging.INFO) # set package log to info - verbose_router_logger.setLevel( - level=logging.INFO - ) # set router logs to info - verbose_proxy_logger.setLevel( - level=logging.INFO - ) # set proxy logs to info + verbose_router_logger.setLevel(level=logging.INFO) # set router logs to info + verbose_proxy_logger.setLevel(level=logging.INFO) # set proxy logs to info elif litellm_log_setting.upper() == "DEBUG": import logging @@ -7089,12 +6501,8 @@ async def initialize( ) verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug - verbose_router_logger.setLevel( - level=logging.DEBUG - ) # set router logs to debug - verbose_proxy_logger.setLevel( - level=logging.DEBUG - ) # set proxy logs to debug + verbose_router_logger.setLevel(level=logging.DEBUG) # set router logs to debug + verbose_proxy_logger.setLevel(level=logging.DEBUG) # set proxy logs to debug dynamic_config = {"general": {}, user_model: {}} if config: ( @@ -7109,9 +6517,7 @@ async def initialize( user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ["AZURE_API_VERSION"] = ( - api_version # set this for azure - litellm can read this from the env - ) + os.environ["AZURE_API_VERSION"] = api_version # set this for azure - litellm can read this from the env if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -7147,9 +6553,7 @@ def data_generator(response): yield f"data: {json.dumps(chunk)}\n\n" -async def async_assistants_data_generator( - response, user_api_key_dict: UserAPIKeyAuth, request_data: dict -): +async def async_assistants_data_generator(response, user_api_key_dict: UserAPIKeyAuth, request_data: dict): verbose_proxy_logger.debug("inside generator") try: time.time() @@ -7174,9 +6578,7 @@ async def async_assistants_data_generator( yield f"data: {done_message}\n\n" except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.async_assistants_data_generator(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.async_assistants_data_generator(): Exception occured - {}".format(str(e)) ) await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -7239,9 +6641,7 @@ def _get_streaming_fallback_metadata( if not isinstance(additional_headers, dict): return False, None, [] - if not _is_positive_int_like( - additional_headers.get("x-litellm-attempted-fallbacks") - ): + if not _is_positive_int_like(additional_headers.get("x-litellm-attempted-fallbacks")): return False, None, [] fallback_model = additional_headers.get("x-litellm-model-group") @@ -7281,11 +6681,7 @@ def _restamp_streaming_chunk_model( fallback_was_attempted: bool = False, fallback_model_from_metadata: str | None = None, ) -> tuple[Any, bool]: - target_model = ( - fallback_model_from_metadata - if fallback_was_attempted - else requested_model_from_client - ) + target_model = fallback_model_from_metadata if fallback_was_attempted else requested_model_from_client # Always return the client-requested model name (not provider-prefixed internal identifiers) # on streaming chunks. # On fallback, use the public OpenAI-compatible model name. This keeps @@ -7298,9 +6694,7 @@ def _restamp_streaming_chunk_model( return chunk, model_mismatch_logged # For Azure Model Router, preserve the actual model used in each chunk - if not fallback_was_attempted and _is_azure_model_router_request( - requested_model_from_client - ): + if not fallback_was_attempted and _is_azure_model_router_request(requested_model_from_client): return chunk, model_mismatch_logged # For fastest_response batch completions, preserve the winning model's name @@ -7308,9 +6702,7 @@ def _restamp_streaming_chunk_model( if not fallback_was_attempted and request_data.get("fastest_response", False): return chunk, model_mismatch_logged - downstream_model = ( - chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None) - ) + downstream_model = chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None) if downstream_model == target_model: return chunk, model_mismatch_logged @@ -7363,10 +6755,7 @@ def _fast_serialize_simple_model_response_stream( return None choice = choices[0] - if ( - getattr(choice, "logprobs", None) is not None - or getattr(choice, "enhancements", None) is not None - ): + if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None: return None delta = getattr(choice, "delta", None) @@ -7384,9 +6773,7 @@ def _fast_serialize_simple_model_response_stream( "provider_specific_fields", "refusal", ) - if any( - getattr(delta, field, None) is not None for field in unsupported_delta_fields - ): + if any(getattr(delta, field, None) is not None for field in unsupported_delta_fields): return None delta_dict: dict = {} @@ -7467,9 +6854,7 @@ _MAX_RAW_SSE_BUFFER_CHARS = 8 * 1024 * 1024 def _pop_complete_sse_frame(buffer: str) -> tuple[str | None, str]: delimiter_positions = [ - (position, delimiter) - for delimiter in _SSE_FRAME_DELIMITERS - if (position := buffer.find(delimiter)) != -1 + (position, delimiter) for delimiter in _SSE_FRAME_DELIMITERS if (position := buffer.find(delimiter)) != -1 ] if not delimiter_positions: return None, buffer @@ -7490,9 +6875,7 @@ async def async_data_generator( client_disconnected = False try: error_message: Optional[str] = None - requested_model_from_client = _get_client_requested_model_for_streaming( - request_data=request_data - ) + requested_model_from_client = _get_client_requested_model_for_streaming(request_data=request_data) ( fallback_was_attempted, fallback_model_from_metadata, @@ -7581,16 +6964,10 @@ async def async_data_generator( break yield frame if len(raw_sse_buffer) > _MAX_RAW_SSE_BUFFER_CHARS: - raise ValueError( - "Raw SSE stream exceeded maximum buffered size without a frame delimiter" - ) + raise ValueError("Raw SSE stream exceeded maximum buffered size without a frame delimiter") raw_passthrough = True elif chunk.startswith(("data:", "event:", ":")): - yield ( - chunk - if chunk.endswith(_SSE_FRAME_DELIMITERS) - else chunk + "\n\n" - ) + yield (chunk if chunk.endswith(_SSE_FRAME_DELIMITERS) else chunk + "\n\n") raw_passthrough = True elif isinstance(chunk, str) and is_raw_sse_stream: raw_sse_buffer += chunk @@ -7600,9 +6977,7 @@ async def async_data_generator( break yield frame if len(raw_sse_buffer) > _MAX_RAW_SSE_BUFFER_CHARS: - raise ValueError( - "Raw SSE stream exceeded maximum buffered size without a frame delimiter" - ) + raise ValueError("Raw SSE stream exceeded maximum buffered size without a frame delimiter") raw_passthrough = True elif isinstance(chunk, str) and chunk.startswith("data: "): error_message = chunk @@ -7629,11 +7004,7 @@ async def async_data_generator( ProxyLogging._fire_deferred_stream_logging(request_data) if raw_sse_buffer: - yield ( - raw_sse_buffer - if raw_sse_buffer.endswith(_SSE_FRAME_DELIMITERS) - else raw_sse_buffer + "\n\n" - ) + yield (raw_sse_buffer if raw_sse_buffer.endswith(_SSE_FRAME_DELIMITERS) else raw_sse_buffer + "\n\n") if error_message is not None: yield error_message @@ -7649,16 +7020,12 @@ async def async_data_generator( # disconnect, so it fires reliably regardless of needs_iterator_wrap # (a nested iterator hook would only see GeneratorExit on GC). if not stream_completed: - proxy_logging_obj._release_max_parallel_requests_on_disconnect( - user_api_key_dict - ) + proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict) client_disconnected = True raise except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format(str(e)) ) await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -7743,10 +7110,7 @@ def giveup(e): and "Max parallel request limit reached" in e.message ) - if ( - general_settings.get("disable_retry_on_max_parallel_request_limit_error") - is True - ): + if general_settings.get("disable_retry_on_max_parallel_request_limit_error") is True: return True # giveup if queuing max parallel request limits is disabled if result: @@ -7766,9 +7130,7 @@ class ProxyStartupEvent: ## COST TRACKING ## cost_tracking() - proxy_logging_obj.startup_event( - llm_router=llm_router, redis_usage_cache=redis_usage_cache - ) + proxy_logging_obj.startup_event(llm_router=llm_router, redis_usage_cache=redis_usage_cache) @staticmethod def _validate_redis_transaction_buffer_config( @@ -7781,8 +7143,8 @@ class ProxyStartupEvent: """ from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( - general_settings.get("use_redis_transaction_buffer", False) + _use_redis_transaction_buffer: Optional[Union[bool, str]] = general_settings.get( + "use_redis_transaction_buffer", False ) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -7817,9 +7179,7 @@ class ProxyStartupEvent: from litellm._redis import _redis_kwargs_from_environment from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: bool | str | None = general_settings.get( - "use_redis_transaction_buffer", False - ) + _use_redis_transaction_buffer: bool | str | None = general_settings.get("use_redis_transaction_buffer", False) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -7841,18 +7201,11 @@ class ProxyStartupEvent: """Initialize MCP semantic tool filter if configured""" from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook - mcp_semantic_filter_config = litellm_settings.get( - "mcp_semantic_tool_filter", None - ) + mcp_semantic_filter_config = litellm_settings.get("mcp_semantic_tool_filter", None) # Only proceed if the feature is configured and enabled - if not mcp_semantic_filter_config or not mcp_semantic_filter_config.get( - "enabled", False - ): - verbose_proxy_logger.debug( - "Semantic tool filter not configured or not enabled, " - "skipping initialization" - ) + if not mcp_semantic_filter_config or not mcp_semantic_filter_config.get("enabled", False): + verbose_proxy_logger.debug("Semantic tool filter not configured or not enabled, skipping initialization") return verbose_proxy_logger.debug( @@ -7869,9 +7222,7 @@ class ProxyStartupEvent: litellm.logging_callback_manager.add_litellm_callback(hook) else: # Only warn if the feature was configured but failed to initialize - verbose_proxy_logger.warning( - "Semantic tool filter hook was configured but failed to initialize" - ) + verbose_proxy_logger.warning("Semantic tool filter hook was configured but failed to initialize") @classmethod def _initialize_jwt_auth( @@ -7906,18 +7257,12 @@ class ProxyStartupEvent: def _add_proxy_budget_to_db(cls, litellm_proxy_budget_name: str): """Adds a global proxy budget to db""" if litellm.budget_duration is None: - raise Exception( - "budget_duration not set on Proxy. budget_duration is required to use max_budget." - ) + raise Exception("budget_duration not set on Proxy. budget_duration is required to use max_budget.") - asyncio.create_task( - cls._upsert_proxy_budget_with_reset_at_backfill(litellm_proxy_budget_name) - ) + asyncio.create_task(cls._upsert_proxy_budget_with_reset_at_backfill(litellm_proxy_budget_name)) @classmethod - async def _upsert_proxy_budget_with_reset_at_backfill( - cls, litellm_proxy_budget_name: str - ) -> None: + async def _upsert_proxy_budget_with_reset_at_backfill(cls, litellm_proxy_budget_name: str) -> None: """ Upsert the proxy admin user row with the configured max_budget / budget_duration, then backfill budget_reset_at if currently NULL. @@ -7957,16 +7302,10 @@ class ProxyStartupEvent: "user_id": litellm_proxy_budget_name, "budget_reset_at": None, }, - data={ - "budget_reset_at": get_budget_reset_time( - budget_duration=litellm.budget_duration - ) - }, + data={"budget_reset_at": get_budget_reset_time(budget_duration=litellm.budget_duration)}, ) except Exception as e: - verbose_proxy_logger.warning( - "Failed to backfill budget_reset_at on proxy admin row: %s", e - ) + verbose_proxy_logger.warning("Failed to backfill budget_reset_at on proxy admin row: %s", e) @classmethod async def _warm_global_spend_cache( @@ -7984,9 +7323,7 @@ class ProxyStartupEvent: prisma_client=prisma_client, ) except Exception as e: - verbose_proxy_logger.debug( - "Global spend cache warm-up at startup skipped or failed: %s", e - ) + verbose_proxy_logger.debug("Global spend cache warm-up at startup skipped or failed: %s", e) @classmethod async def _update_default_team_member_budget(cls): @@ -8021,17 +7358,11 @@ class ProxyStartupEvent: if prisma_client is None: return - db_record = await UISettingsRepository(prisma_client).table.find_unique( - where={"id": "ui_settings"} - ) + db_record = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) if db_record and db_record.ui_settings: raw = db_record.ui_settings ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw) - flags_to_sync = { - k: ui_settings[k] - for k in _RUNTIME_GENERAL_SETTINGS_FLAGS - if k in ui_settings - } + flags_to_sync = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings} if flags_to_sync: general_settings.update(flags_to_sync) verbose_proxy_logger.info( @@ -8039,9 +7370,7 @@ class ProxyStartupEvent: list(flags_to_sync.keys()), ) except Exception as e: - verbose_proxy_logger.debug( - "UI settings sync on startup skipped or failed: %s", e - ) + verbose_proxy_logger.debug("UI settings sync on startup skipped or failed: %s", e) @classmethod async def initialize_scheduled_background_jobs( @@ -8128,9 +7457,7 @@ class ProxyStartupEvent: ### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ### ## Reduces QPS as there are more tags for a single request - tag_spend_update_interval = int( - batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER - ) + tag_spend_update_interval = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER) from litellm.proxy.utils import update_daily_tag_spend scheduler.add_job( @@ -8161,9 +7488,7 @@ class ProxyStartupEvent: ) ### ADD NEW MODELS ### - store_model_in_db = ( - get_secret_bool("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db - ) + store_model_in_db = get_secret_bool("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db # If store_model_in_db is still False, check DB for override. # This breaks the chicken-and-egg where DB has store_model_in_db=True @@ -8173,21 +7498,13 @@ class ProxyStartupEvent: _db_gs_record = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) - if _db_gs_record is not None and isinstance( - _db_gs_record.param_value, dict - ): + if _db_gs_record is not None and isinstance(_db_gs_record.param_value, dict): _db_val = _db_gs_record.param_value.get("store_model_in_db") - if _db_val is True or ( - isinstance(_db_val, str) and _db_val.lower() == "true" - ): + if _db_val is True or (isinstance(_db_val, str) and _db_val.lower() == "true"): store_model_in_db = True - verbose_proxy_logger.info( - "store_model_in_db=True loaded from DB, overriding config/env" - ) + verbose_proxy_logger.info("store_model_in_db=True loaded from DB, overriding config/env") except Exception as e: - verbose_proxy_logger.debug( - "Failed to check DB for store_model_in_db: %s", str(e) - ) + verbose_proxy_logger.debug("Failed to check DB for store_model_in_db: %s", str(e)) if store_model_in_db is True: # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum @@ -8204,9 +7521,7 @@ class ProxyStartupEvent: ) # this will load all existing models on proxy startup - await proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) + await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) ### GET STORED CREDENTIALS ### scheduler.add_job( @@ -8248,18 +7563,12 @@ class ProxyStartupEvent: replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info( - f"Spend log cleanup scheduled with cron: {cleanup_cron}" - ) + verbose_proxy_logger.info(f"Spend log cleanup scheduled with cron: {cleanup_cron}") except ValueError: - verbose_proxy_logger.error( - f"Invalid maximum_spend_logs_cleanup_cron value: {cleanup_cron}" - ) + verbose_proxy_logger.error(f"Invalid maximum_spend_logs_cleanup_cron value: {cleanup_cron}") else: # Interval-based scheduling (existing behavior) - retention_interval = general_settings.get( - "maximum_spend_logs_retention_interval", "1d" - ) + retention_interval = general_settings.get("maximum_spend_logs_retention_interval", "1d") try: interval_seconds = duration_in_seconds(retention_interval) scheduler.add_job( @@ -8272,9 +7581,7 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) except ValueError: - verbose_proxy_logger.error( - "Invalid maximum_spend_logs_retention_interval value" - ) + verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") ### CHECK BATCH COST ### if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: try: @@ -8290,8 +7597,7 @@ class ProxyStartupEvent: scheduler.add_job( check_batch_cost_job.check_batch_cost, "interval", - seconds=proxy_batch_polling_interval - + random.randint(0, 30), # Add small random offset + seconds=proxy_batch_polling_interval + random.randint(0, 30), # Add small random offset # REMOVED jitter parameter - major cause of memory leak id="check_batch_cost_job", replace_existing=True, @@ -8321,21 +7627,16 @@ class ProxyStartupEvent: scheduler.add_job( check_responses_cost_job.check_responses_cost, "interval", - seconds=proxy_batch_polling_interval - + random.randint(0, 30), # Add small random offset + seconds=proxy_batch_polling_interval + random.randint(0, 30), # Add small random offset # REMOVED jitter parameter - major cause of memory leak id="check_responses_cost_job", replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info( - "Responses cost check job scheduled successfully" - ) + verbose_proxy_logger.info("Responses cost check job scheduled successfully") except Exception as e: - verbose_proxy_logger.debug( - f"Failed to setup responses cost checking: {e}" - ) + verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}") verbose_proxy_logger.debug( "Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..." ) @@ -8353,9 +7654,7 @@ class ProxyStartupEvent: ) @classmethod - async def _initialize_spend_tracking_background_jobs( - cls, scheduler: AsyncIOScheduler - ): + async def _initialize_spend_tracking_background_jobs(cls, scheduler: AsyncIOScheduler): """ Initialize the spend tracking and other background jobs 1. CloudZero Background Job @@ -8409,13 +7708,9 @@ class ProxyStartupEvent: integration_token=db_settings.get("integration_token"), base_url=db_settings.get("base_url"), ) - litellm.logging_callback_manager.add_litellm_callback( - vantage_logger - ) + litellm.logging_callback_manager.add_litellm_callback(vantage_logger) except Exception as e: - verbose_proxy_logger.warning( - "Failed to register VantageLogger from DB settings: %s", e - ) + verbose_proxy_logger.warning("Failed to register VantageLogger from DB settings: %s", e) await VantageLogger.init_vantage_background_job(scheduler=scheduler) ######################################################## @@ -8454,9 +7749,7 @@ class ProxyStartupEvent: # Get prisma_client and proxy_logging_obj from global scope if prisma_client is not None: # Reuse the PodLockManager from db_spend_update_writer - pod_lock_manager = ( - proxy_logging_obj.db_spend_update_writer.pod_lock_manager - ) + pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager key_rotation_manager = KeyRotationManager( prisma_client, pod_lock_manager=pod_lock_manager, @@ -8471,24 +7764,16 @@ class ProxyStartupEvent: id="key_rotation_job", ) else: - verbose_proxy_logger.warning( - "Key rotation enabled but prisma_client not available" - ) + verbose_proxy_logger.warning("Key rotation enabled but prisma_client not available") except Exception as e: verbose_proxy_logger.warning(f"Failed to setup key rotation job: {e}") else: - verbose_proxy_logger.debug( - "Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)" - ) + verbose_proxy_logger.debug("Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)") - await cls._initialize_expired_ui_session_key_cleanup_background_job( - scheduler=scheduler - ) + await cls._initialize_expired_ui_session_key_cleanup_background_job(scheduler=scheduler) @classmethod - async def _initialize_expired_ui_session_key_cleanup_background_job( - cls, scheduler: AsyncIOScheduler - ): + async def _initialize_expired_ui_session_key_cleanup_background_job(cls, scheduler: AsyncIOScheduler): """ Initialize the expired UI session key cleanup background job. """ @@ -8508,10 +7793,7 @@ class ProxyStartupEvent: expired_ui_session_key_cleanup_enabled: Optional[bool] = str_to_bool( LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED ) - verbose_proxy_logger.debug( - "expired_ui_session_key_cleanup_enabled: " - f"{expired_ui_session_key_cleanup_enabled}" - ) + verbose_proxy_logger.debug(f"expired_ui_session_key_cleanup_enabled: {expired_ui_session_key_cleanup_enabled}") if expired_ui_session_key_cleanup_enabled is True: try: @@ -8520,15 +7802,11 @@ class ProxyStartupEvent: ) if prisma_client is not None: - pod_lock_manager = ( - proxy_logging_obj.db_spend_update_writer.pod_lock_manager - ) - expired_ui_session_key_cleanup_manager = ( - ExpiredUISessionKeyCleanupManager( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - pod_lock_manager=pod_lock_manager, - ) + pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager + expired_ui_session_key_cleanup_manager = ExpiredUISessionKeyCleanupManager( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + pod_lock_manager=pod_lock_manager, ) verbose_proxy_logger.debug( "Expired UI session key cleanup background job scheduled " @@ -8545,13 +7823,10 @@ class ProxyStartupEvent: ) else: verbose_proxy_logger.warning( - "Expired UI session key cleanup enabled but prisma_client " - "not available" + "Expired UI session key cleanup enabled but prisma_client not available" ) except Exception as e: - verbose_proxy_logger.warning( - f"Failed to setup expired UI session key cleanup job: {e}" - ) + verbose_proxy_logger.warning(f"Failed to setup expired UI session key cleanup job: {e}") else: verbose_proxy_logger.debug( "Expired UI session key cleanup disabled (set " @@ -8573,22 +7848,17 @@ class ProxyStartupEvent: and prisma_client is not None ): print("Alerting: Initializing Weekly/Monthly Spend Reports") # noqa: T201 - spend_report_frequency: str = ( - general_settings.get("spend_report_frequency", "7d") or "7d" - ) + spend_report_frequency: str = general_settings.get("spend_report_frequency", "7d") or "7d" days = int(spend_report_frequency[:-1]) if spend_report_frequency[-1].lower() != "d": - raise ValueError( - "spend_report_frequency must be specified in days, e.g., '1d', '7d'" - ) + raise ValueError("spend_report_frequency must be specified in days, e.g., '1d', '7d'") scheduler.add_job( proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report, "interval", days=days, - next_run_time=datetime.now() - + timedelta(seconds=10 + random.randint(0, 300)), + next_run_time=datetime.now() + timedelta(seconds=10 + random.randint(0, 300)), args=[spend_report_frequency], id="weekly_spend_report_job", replace_existing=True, @@ -8632,9 +7902,7 @@ class ProxyStartupEvent: prisma_client: Optional[PrismaClient] = None if database_url is not None: try: - prisma_client = PrismaClient( - database_url=database_url, proxy_logging_obj=proxy_logging_obj - ) + prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) except Exception as e: raise e @@ -8642,23 +7910,15 @@ class ProxyStartupEvent: await prisma_client.connect() except Exception as e: if "P3018" in str(e) or "P3009" in str(e): - verbose_proxy_logger.debug( - "CRITICAL: DATABASE MIGRATION FAILED" - ) - verbose_proxy_logger.debug( - "Your database is in a 'dirty' state." - ) - verbose_proxy_logger.debug( - "FIX: Run 'prisma migrate resolve --applied '" - ) + verbose_proxy_logger.debug("CRITICAL: DATABASE MIGRATION FAILED") + verbose_proxy_logger.debug("Your database is in a 'dirty' state.") + verbose_proxy_logger.debug("FIX: Run 'prisma migrate resolve --applied '") raise e ## Start RDS IAM token refresh background task if enabled ## # This proactively refreshes IAM tokens before they expire, # preventing the 15-minute connection failure bug (#16220) - if hasattr(prisma_client, "db") and hasattr( - prisma_client.db, "start_token_refresh_task" - ): + if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): await prisma_client.db.start_token_refresh_task() ## Add necessary views to proxy ## @@ -8671,10 +7931,7 @@ class ProxyStartupEvent: ) # set the spend logs row count in proxy state. Don't block execution # run a health check to ensure the DB is ready - if ( - get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) - is not True - ): + if get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) is not True: await prisma_client.health_check() if hasattr(prisma_client, "start_db_health_watchdog_task"): @@ -8777,10 +8034,7 @@ class ProxyStartupEvent: # pyroscope-io expects sample_rate as an integer configure_kwargs["sample_rate"] = int(float(sample_rate_env)) except (ValueError, TypeError): - raise ValueError( - "PYROSCOPE_SAMPLE_RATE must be a number, got: " - f"{sample_rate_env!r}" - ) + raise ValueError(f"PYROSCOPE_SAMPLE_RATE must be a number, got: {sample_rate_env!r}") pyroscope.configure(**configure_kwargs) msg = ( f"LiteLLM: Pyroscope profiling started (app_name={app_name}, server_address={server_address}). " @@ -8797,9 +8051,7 @@ class ProxyStartupEvent: #### API ENDPOINTS #### -@router.get( - "/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] -) +@router.get("/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]) @router.get( "/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] ) # if project requires model list @@ -8836,13 +8088,7 @@ async def model_list( Hiding is presentation-only: a hidden model can still be called directly. """ - global \ - llm_model_list, \ - general_settings, \ - llm_router, \ - prisma_client, \ - user_api_key_cache, \ - proxy_logging_obj + global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj settings = cast(dict[str, object], general_settings) # any-ok: legacy settings @@ -8872,9 +8118,7 @@ async def model_list( ) # Compute once — used in both branches below to hide paused models from the listing. - blocked_names = ( - llm_router.get_fully_blocked_model_names() if llm_router is not None else set() - ) + blocked_names = llm_router.get_fully_blocked_model_names() if llm_router is not None else set() # Opt-in: also hide models whose deployments are all unhealthy per background # health checks. Empty when health state is unavailable or stale (fail open). @@ -8901,9 +8145,7 @@ async def model_list( # Include model access groups if requested if include_model_access_groups: - proxy_model_list = list( - set(proxy_model_list + list(model_access_groups.keys())) - ) + proxy_model_list = list(set(proxy_model_list + list(model_access_groups.keys()))) # Get complete model list including wildcard routes if requested from litellm.proxy.auth.model_checks import get_complete_model_list @@ -8929,9 +8171,7 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - for response_id, lookup_id in TeamModelNameTranslator.listing_entries( - all_models, llm_router, settings - ): + for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings): model_info = create_model_info_response( model_id=lookup_id, provider="openai", @@ -8971,9 +8211,7 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - for response_id, lookup_id in TeamModelNameTranslator.listing_entries( - all_models, llm_router, settings - ): + for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings): model_info = create_model_info_response( model_id=lookup_id, provider="openai", @@ -9019,13 +8257,7 @@ async def model_info( scoping, health filtering, paused deployments) drives both endpoints; the listing's public id must resolve to the same internal deployment here. """ - global \ - llm_model_list, \ - general_settings, \ - llm_router, \ - prisma_client, \ - user_api_key_cache, \ - proxy_logging_obj + global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj settings = cast(dict[str, object], general_settings) # any-ok: legacy settings @@ -9051,9 +8283,7 @@ async def model_info( # Mirror /v1/models' visibility filter so first-occurrence resolution # cannot land on a deployment the listing had hidden. - blocked_names = ( - llm_router.get_fully_blocked_model_names() if llm_router is not None else set() - ) + blocked_names = llm_router.get_fully_blocked_model_names() if llm_router is not None else set() unhealthy_names: set[str] = set() if healthy_only and llm_router is not None: unhealthy_names = await llm_router.async_get_fully_unhealthy_model_names() @@ -9061,9 +8291,7 @@ async def model_info( if hidden_names: all_models = [m for m in all_models if m not in hidden_names] - internal_to_public = TeamModelNameTranslator.build_internal_to_public_map( - llm_router, settings - ) + internal_to_public = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) resolved_model_id = TeamModelNameTranslator.resolve_public_name( model_id=model_id, available_models=all_models, @@ -9157,32 +8385,15 @@ async def chat_completion( # extra_body); otherwise `data["metadata"][k] = v` below raises # TypeError on a string value and 500s the request. data["metadata"] = {} - if ( - hasattr(user_api_key_dict, "user_id") - and user_api_key_dict.user_id is not None - ): + if hasattr(user_api_key_dict, "user_id") and user_api_key_dict.user_id is not None: data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - if ( - hasattr(user_api_key_dict, "team_id") - and user_api_key_dict.team_id is not None - ): + if hasattr(user_api_key_dict, "team_id") and user_api_key_dict.team_id is not None: data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id - if ( - hasattr(user_api_key_dict, "org_id") - and user_api_key_dict.org_id is not None - ): + if hasattr(user_api_key_dict, "org_id") and user_api_key_dict.org_id is not None: data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id - if ( - hasattr(user_api_key_dict, "organization_alias") - and user_api_key_dict.organization_alias is not None - ): - data["metadata"]["user_api_key_org_alias"] = ( - user_api_key_dict.organization_alias - ) - if ( - hasattr(user_api_key_dict, "agent_id") - and user_api_key_dict.agent_id is not None - ): + if hasattr(user_api_key_dict, "organization_alias") and user_api_key_dict.organization_alias is not None: + data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias + if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None: data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) @@ -9223,9 +8434,7 @@ async def chat_completion( _chat_response.choices[0].finish_reason = "content_filter" # type: ignore if data.get("stream", None) is not None and data["stream"] is True: - _iterator = litellm.utils.ModelResponseIterator( - model_response=_chat_response, convert_to_delta=True - ) + _iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True) _streaming_response = litellm.CustomStreamWrapper( completion_stream=_iterator, model=e.model, @@ -9258,9 +8467,7 @@ async def chat_completion( _chat_response.choices[0].message.content = e.message # type: ignore if data.get("stream", None) is not None and data["stream"] is True: - _iterator = litellm.utils.ModelResponseIterator( - model_response=_chat_response, convert_to_delta=True - ) + _iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True) _streaming_response = litellm.CustomStreamWrapper( completion_stream=_iterator, model=data.get("model", ""), @@ -9277,11 +8484,7 @@ async def chat_completion( return StreamingResponse( selected_data_generator, media_type="text/event-stream", - status_code=( - e.status_code - if hasattr(e, "status_code") - else status.HTTP_400_BAD_REQUEST - ), + status_code=(e.status_code if hasattr(e, "status_code") else status.HTTP_400_BAD_REQUEST), ) _usage = litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) _chat_response.usage = _usage # type: ignore @@ -9294,12 +8497,8 @@ async def chat_completion( ) -@router.post( - "/v1/completions", dependencies=[Depends(user_api_key_auth)], tags=["completions"] -) -@router.post( - "/completions", dependencies=[Depends(user_api_key_auth)], tags=["completions"] -) +@router.post("/v1/completions", dependencies=[Depends(user_api_key_auth)], tags=["completions"]) +@router.post("/completions", dependencies=[Depends(user_api_key_auth)], tags=["completions"]) @router.post( "/engines/{model:path}/completions", dependencies=[Depends(user_api_key_auth)], @@ -9341,32 +8540,15 @@ async def completion( if user_api_key_dict is not None: if data.get("metadata") is None: data["metadata"] = {} - if ( - hasattr(user_api_key_dict, "user_id") - and user_api_key_dict.user_id is not None - ): + if hasattr(user_api_key_dict, "user_id") and user_api_key_dict.user_id is not None: data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - if ( - hasattr(user_api_key_dict, "team_id") - and user_api_key_dict.team_id is not None - ): + if hasattr(user_api_key_dict, "team_id") and user_api_key_dict.team_id is not None: data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id - if ( - hasattr(user_api_key_dict, "org_id") - and user_api_key_dict.org_id is not None - ): + if hasattr(user_api_key_dict, "org_id") and user_api_key_dict.org_id is not None: data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id - if ( - hasattr(user_api_key_dict, "organization_alias") - and user_api_key_dict.organization_alias is not None - ): - data["metadata"]["user_api_key_org_alias"] = ( - user_api_key_dict.organization_alias - ) - if ( - hasattr(user_api_key_dict, "agent_id") - and user_api_key_dict.agent_id is not None - ): + if hasattr(user_api_key_dict, "organization_alias") and user_api_key_dict.organization_alias is not None: + data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias + if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None: data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) return await base_llm_response_processor.base_process_llm_request( @@ -9408,9 +8590,7 @@ async def completion( ) # Set usage attribute dynamically (ModelResponse accepts usage in __init__ but it's not in type definition) setattr(_text_response, "usage", _usage) - _iterator = litellm.utils.ModelResponseIterator( - model_response=_text_response, convert_to_delta=True - ) + _iterator = litellm.utils.ModelResponseIterator(model_response=_text_response, convert_to_delta=True) _streaming_response = litellm.TextCompletionStreamWrapper( completion_stream=_iterator, model=e.model, @@ -9455,9 +8635,7 @@ async def completion( ) _chat_response.usage = _usage # type: ignore _chat_response.choices[0].message.content = e.message # type: ignore - _iterator = litellm.utils.ModelResponseIterator( - model_response=_chat_response, convert_to_delta=True - ) + _iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True) _streaming_response = litellm.TextCompletionStreamWrapper( completion_stream=_iterator, model=_data.get("model", ""), @@ -9474,11 +8652,7 @@ async def completion( selected_data_generator, media_type="text/event-stream", headers={}, - status_code=( - e.status_code - if hasattr(e, "status_code") - else status.HTTP_400_BAD_REQUEST - ), + status_code=(e.status_code if hasattr(e, "status_code") else status.HTTP_400_BAD_REQUEST), ) else: _response = litellm.TextCompletionResponse() @@ -9488,11 +8662,7 @@ async def completion( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.completion(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - {}".format(str(e))) error_msg = f"{str(e)}" raise ProxyException( message=getattr(e, "message", error_msg), @@ -9569,58 +8739,34 @@ async def embeddings( # check if provider accept list of tokens as input - e.g. for langchain integration if llm_router is not None and data.get("model") in router_model_names: # Use router's O(1) lookup instead of O(N) iteration through llm_model_list - deployment = llm_router.get_deployment_by_model_group_name( - model_group_name=data["model"] - ) + deployment = llm_router.get_deployment_by_model_group_name(model_group_name=data["model"]) if deployment is not None: litellm_params = deployment.get("litellm_params", {}) or {} litellm_model = litellm_params.get("model", "") # Check if this provider supports token arrays - supports_token_arrays = ( - litellm_model in litellm.open_ai_embedding_models - or any( - litellm_model.startswith(provider) - for provider in LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS - ) + supports_token_arrays = litellm_model in litellm.open_ai_embedding_models or any( + litellm_model.startswith(provider) + for provider in LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS ) if not supports_token_arrays: # non-openai/azure embedding model called with token input - decode tokens input_list = [] for i in data["input"]: - input_list.append( - litellm.decode(model="gpt-3.5-turbo", tokens=i) - ) + input_list.append(litellm.decode(model="gpt-3.5-turbo", tokens=i)) data["input"] = input_list if user_api_key_dict is not None: if data.get("metadata") is None: data["metadata"] = {} - if ( - hasattr(user_api_key_dict, "user_id") - and user_api_key_dict.user_id is not None - ): + if hasattr(user_api_key_dict, "user_id") and user_api_key_dict.user_id is not None: data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - if ( - hasattr(user_api_key_dict, "team_id") - and user_api_key_dict.team_id is not None - ): + if hasattr(user_api_key_dict, "team_id") and user_api_key_dict.team_id is not None: data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id - if ( - hasattr(user_api_key_dict, "org_id") - and user_api_key_dict.org_id is not None - ): + if hasattr(user_api_key_dict, "org_id") and user_api_key_dict.org_id is not None: data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id - if ( - hasattr(user_api_key_dict, "organization_alias") - and user_api_key_dict.organization_alias is not None - ): - data["metadata"]["user_api_key_org_alias"] = ( - user_api_key_dict.organization_alias - ) - if ( - hasattr(user_api_key_dict, "agent_id") - and user_api_key_dict.agent_id is not None - ): + if hasattr(user_api_key_dict, "organization_alias") and user_api_key_dict.organization_alias is not None: + data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias + if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None: data["metadata"]["agent_id"] = user_api_key_dict.agent_id # Use unified request processor (same as chat/completions and responses) @@ -9728,9 +8874,7 @@ async def moderations( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -9758,9 +8902,7 @@ async def moderations( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.moderations(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.moderations(): Exception occured - {}".format(str(e)) ) if isinstance(e, HTTPException): raise ProxyException( @@ -9850,9 +8992,7 @@ async def audio_speech( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -9895,9 +9035,7 @@ async def audio_speech( if "gemini" in request_model_lower and ( "tts" in request_model_lower or "preview-tts" in request_model_lower ): - media_type = ( - "audio/wav" # Gemini TTS returns WAV format after conversion - ) + media_type = "audio/wav" # Gemini TTS returns WAV format after conversion return StreamingResponse( _audio_speech_chunk_generator(response), # type: ignore[arg-type] @@ -9911,11 +9049,7 @@ async def audio_speech( original_exception=e, request_data=data, ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.audio_speech(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.audio_speech(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) raise e @@ -10014,9 +9148,7 @@ async def audio_transcriptions( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10060,9 +9192,7 @@ async def audio_transcriptions( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.audio_transcription(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.audio_transcription(): Exception occured - {}".format(str(e)) ) if isinstance(e, HTTPException): raise ProxyException( @@ -10128,9 +9258,7 @@ async def vertex_ai_live_passthrough_endpoint( @lru_cache(maxsize=_REALTIME_BODY_CACHE_SIZE) -def _realtime_query_params_template( - model: Optional[str], intent: Optional[str] -) -> Tuple[Tuple[str, str], ...]: +def _realtime_query_params_template(model: Optional[str], intent: Optional[str]) -> Tuple[Tuple[str, str], ...]: """ Build a hashable representation of the realtime query params so we can cache the repetitive model/intent combinations. @@ -10148,12 +9276,8 @@ def _realtime_query_params_template( @app.websocket("/realtime") async def realtime_websocket_endpoint( websocket: WebSocket, - model: Optional[str] = fastapi.Query( - None, description="The model to use for the websocket connection." - ), - intent: Optional[str] = fastapi.Query( - None, description="The intent of the websocket connection." - ), + model: Optional[str] = fastapi.Query(None, description="The model to use for the websocket connection."), + intent: Optional[str] = fastapi.Query(None, description="The intent of the websocket connection."), guardrails: Optional[str] = fastapi.Query( None, description="Comma-separated list of guardrail names to apply to this request.", @@ -10161,9 +9285,7 @@ async def realtime_websocket_endpoint( user_api_key_dict=Depends(user_api_key_auth_websocket), ): requested_protocols = [ - p.strip() - for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") - if p.strip() + p.strip() for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") if p.strip() ] accept_kwargs: dict = {} if requested_protocols: @@ -10190,9 +9312,7 @@ async def realtime_websocket_endpoint( await websocket.accept(**accept_kwargs) # Only use explicit parameters, not all query params - query_params = cast( - RealtimeQueryParams, dict(_realtime_query_params_template(model, intent)) - ) + query_params = cast(RealtimeQueryParams, dict(_realtime_query_params_template(model, intent))) data: Dict[str, Any] = { "model": route_model, @@ -10327,16 +9447,12 @@ async def get_assistants( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.aget_assistants(**data) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10363,11 +9479,7 @@ async def get_assistants( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.get_assistants(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.get_assistants(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10426,16 +9538,12 @@ async def create_assistant( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.acreate_assistants(**data) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10463,9 +9571,7 @@ async def create_assistant( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.error( - "litellm.proxy.proxy_server.create_assistant(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.create_assistant(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): @@ -10523,16 +9629,12 @@ async def delete_assistant( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.adelete_assistant(assistant_id=assistant_id, **data) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10560,9 +9662,7 @@ async def delete_assistant( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) verbose_proxy_logger.error( - "litellm.proxy.proxy_server.delete_assistant(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.delete_assistant(): Exception occured - {}".format(str(e)) ) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): @@ -10620,16 +9720,12 @@ async def create_threads( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.acreate_thread(**data) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10656,11 +9752,7 @@ async def create_threads( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.create_threads(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.create_threads(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10715,16 +9807,12 @@ async def get_thread( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.aget_thread(thread_id=thread_id, **data) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10751,11 +9839,7 @@ async def get_thread( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.get_thread(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.get_thread(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10814,16 +9898,12 @@ async def add_messages( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.a_add_message(thread_id=thread_id, **data) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10850,11 +9930,7 @@ async def add_messages( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.add_messages(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.add_messages(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -10909,16 +9985,12 @@ async def get_messages( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.aget_messages(thread_id=thread_id, **data) ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -10945,11 +10017,7 @@ async def get_messages( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.get_messages(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.get_messages(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -11006,14 +10074,10 @@ async def run_thread( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.arun_thread(thread_id=thread_id, **data) - if ( - "stream" in data and data["stream"] is True - ): # use generate_responses to stream responses + if "stream" in data and data["stream"] is True: # use generate_responses to stream responses return await create_response( generator=async_assistants_data_generator( user_api_key_dict=user_api_key_dict, @@ -11027,9 +10091,7 @@ async def run_thread( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -11056,11 +10118,7 @@ async def run_thread( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.run_thread(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.run_thread(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -11122,9 +10180,7 @@ def _get_provider_token_counter( # Use existing LiteLLM logic to determine provider model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=full_model, - custom_llm_provider=deployment.get("litellm_params", {}).get( - "custom_llm_provider" - ), + custom_llm_provider=deployment.get("litellm_params", {}).get("custom_llm_provider"), api_base=deployment.get("litellm_params", {}).get("api_base"), api_key=deployment.get("litellm_params", {}).get("api_key"), ) @@ -11171,9 +10227,7 @@ async def _try_provider_token_count( system: Optional[str] = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" - if not provider_counter.should_use_token_counting_api( - custom_llm_provider=custom_llm_provider - ): + if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): return None try: result = await provider_counter.count_tokens( @@ -11239,9 +10293,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) # Validate request ######################################################### if prompt is None and messages is None and contents is None: - raise HTTPException( - status_code=400, detail="prompt or messages or contents must be provided" - ) + raise HTTPException(status_code=400, detail="prompt or messages or contents must be provided") deployment: Optional[Dict[str, Any]] = None litellm_model_name = None @@ -11275,9 +10327,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) custom_llm_provider: Optional[str] = None if call_endpoint is True and deployment is not None: # Auto-route to the correct provider based on model - provider_counter, _model, custom_llm_provider = _get_provider_token_counter( - deployment, model_to_use - ) + provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use) if _model is not None: model_to_use = _model @@ -11312,9 +10362,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) Optional[CustomHuggingfaceTokenizer], model_info.get("custom_tokenizer", None), ) - _tokenizer_used = litellm.utils._select_tokenizer( - model=model_to_use, custom_tokenizer=custom_tokenizer - ) + _tokenizer_used = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer) tokenizer_used = str(_tokenizer_used["type"]) total_tokens = token_counter( @@ -11356,9 +10404,7 @@ async def supported_openai_params(model: str): ) } except Exception: - raise HTTPException( - status_code=400, detail={"error": "Could not map model={}".format(model)} - ) + raise HTTPException(status_code=400, detail={"error": "Could not map model={}".format(model)}) @router.post( @@ -11405,18 +10451,14 @@ async def _check_if_model_is_user_added( id = model.get("model_info", {}).get("id", None) if id is None: continue - db_model = await ModelRepository(prisma_client).table.find_unique( - where={"model_id": id} - ) + db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id}) if db_model is not None: if db_model.created_by == user_api_key_dict.user_id: filtered_models.append(model) return filtered_models -def _check_if_model_is_team_model( - models: List[DeploymentTypedDict], user_row: LiteLLM_UserTable -) -> List[Dict]: +def _check_if_model_is_team_model(models: List[DeploymentTypedDict], user_row: LiteLLM_UserTable) -> List[Dict]: """ Check if model is a team model @@ -11511,16 +10553,12 @@ def _add_team_models_to_all_models( team_models.setdefault(model_id, set()).add(team_object.team_id) else: for model_name in team_object.models: - _models = llm_router.get_model_list( - model_name=model_name, team_id=team_object.team_id - ) + _models = llm_router.get_model_list(model_name=model_name, team_id=team_object.team_id) if _models is not None: for model in _models: model_id = model.get("model_info", {}).get("id", None) if model_id is not None: - team_models.setdefault(model_id, set()).add( - team_object.team_id - ) + team_models.setdefault(model_id, set()).add(team_object.team_id) return team_models @@ -11549,10 +10587,7 @@ async def _add_access_group_models_to_team_models( # Skip teams with empty models list — they already have access to everything # (handled by _add_team_models_to_all_models) - if ( - not team_object.models - or SpecialModelNames.all_proxy_models.value in team_object.models - ): + if not team_object.models or SpecialModelNames.all_proxy_models.value in team_object.models: continue eligible_teams.append(team_object) @@ -11576,9 +10611,7 @@ async def _add_access_group_models_to_team_models( model_names.update(ag_model_map.get(ag_id, [])) for model_name in model_names: - deployments = llm_router.get_model_list( - model_name=model_name, team_id=team_object.team_id - ) + deployments = llm_router.get_model_list(model_name=model_name, team_id=team_object.team_id) if deployments is not None: for deployment in deployments: model_id = deployment.get("model_info", {}).get("id", None) @@ -11605,19 +10638,11 @@ async def get_all_team_models( if user_teams == "*": team_db_objects = await TeamRepository(prisma_client).table.find_many() - team_db_objects_typed = [ - LiteLLM_TeamTable(**team_db_object.model_dump()) - for team_db_object in team_db_objects - ] + team_db_objects_typed = [LiteLLM_TeamTable(**team_db_object.model_dump()) for team_db_object in team_db_objects] else: - team_db_objects = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": user_teams}} - ) + team_db_objects = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_teams}}) - team_db_objects_typed = [ - LiteLLM_TeamTable(**team_db_object.model_dump()) - for team_db_object in team_db_objects - ] + team_db_objects_typed = [LiteLLM_TeamTable(**team_db_object.model_dump()) for team_db_object in team_db_objects] team_models = _add_team_models_to_all_models( team_db_objects_typed=team_db_objects_typed, @@ -11685,9 +10710,7 @@ async def _populate_team_access_on_models( direct_access_models: List[str] = [] if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: user_teams = "*" - direct_access_models = llm_router.get_model_ids( - exclude_team_models=True - ) # has access to all models + direct_access_models = llm_router.get_model_ids(exclude_team_models=True) # has access to all models elif user_api_key_dict.user_id is not None: user_db_object = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} @@ -11717,9 +10740,7 @@ async def _populate_team_access_on_models( else: can_use_model = True if can_use_model: - _model["model_info"]["access_via_team_ids"] = team_models.get( - model_id, [] - ) + _model["model_info"]["access_via_team_ids"] = team_models.get(model_id, []) direct_access_model_ids = set(direct_access_models) for _model in all_models: @@ -11768,10 +10789,7 @@ def _enrich_model_info_with_litellm_data( if debug is True: _openai_client = "None" if llm_router is not None: - _openai_client = ( - llm_router._get_client(deployment=model, kwargs={}, client_type="async") - or "None" - ) + _openai_client = llm_router._get_client(deployment=model, kwargs={}, client_type="async") or "None" else: _openai_client = "llm_router_is_None" openai_client = str(_openai_client) @@ -11800,9 +10818,7 @@ def _enrich_model_info_with_litellm_data( if len(split_model) > 0: litellm_model = split_model[-1] try: - litellm_model_info = litellm.get_model_info( - model=litellm_model, custom_llm_provider=split_model[0] - ) + litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0]) except Exception: litellm_model_info = {} for k, v in litellm_model_info.items(): @@ -11811,9 +10827,7 @@ def _enrich_model_info_with_litellm_data( model["model_info"] = model_info # don't return the api key / vertex credentials # don't return the llm credentials - model = remove_sensitive_info_from_deployment( - model, excluded_keys={"litellm_credential_name"} - ) + model = remove_sensitive_info_from_deployment(model, excluded_keys={"litellm_credential_name"}) return model @@ -11838,20 +10852,15 @@ async def _get_caller_byok_team_scope( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ): return None - key_team_scope: set[str] = ( - {user_api_key_dict.team_id} if user_api_key_dict.team_id else set() - ) + key_team_scope: set[str] = {user_api_key_dict.team_id} if user_api_key_dict.team_id else set() user_id = user_api_key_dict.user_id if user_id is None: return key_team_scope try: - user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) except Exception: verbose_proxy_logger.exception( - "Failed to look up caller teams while scoping BYOK search; " - "defaulting to key team scope only." + "Failed to look up caller teams while scoping BYOK search; defaulting to key team scope only." ) return key_team_scope if user_row is None: @@ -11859,9 +10868,7 @@ async def _get_caller_byok_team_scope( return key_team_scope | set(user_row.teams or []) -def _byok_row_outside_caller_teams( - model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]] -) -> bool: +def _byok_row_outside_caller_teams(model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]]) -> bool: """Whether a team BYOK row belongs to a team the caller is not a member of. `team_id` is only set on team BYOK rows; non-team rows fall through @@ -11908,9 +10915,7 @@ async def _fetch_db_models_for_search( filter for `team_public_model_name` instead and keep the DB cost bounded by `search`. """ - db_where_condition: Dict[str, Any] = { - "model_name": {"contains": search_lower, "mode": "insensitive"} - } + db_where_condition: Dict[str, Any] = {"model_name": {"contains": search_lower, "mode": "insensitive"}} if db_model_ids_in_router: db_where_condition["model_id"] = {"not": {"in": list(db_model_ids_in_router)}} @@ -11922,9 +10927,7 @@ async def _fetch_db_models_for_search( else: take_limit = max(0, page * size - router_models_count) - db_models_total_count = await ModelRepository(prisma_client).table.count( - where=db_where_condition - ) + db_models_total_count = await ModelRepository(prisma_client).table.count(where=db_where_condition) db_models_raw: list = [] if take_limit > 0: @@ -11938,9 +10941,7 @@ async def _fetch_db_models_for_search( matching_db_rows = [ m for m in db_models_raw - if not is_byok_outside_caller_teams( - m.model_info if isinstance(m.model_info, dict) else {} - ) + if not is_byok_outside_caller_teams(m.model_info if isinstance(m.model_info, dict) else {}) ] decrypted: List[Dict[str, Any]] = [] @@ -12005,9 +11006,7 @@ async def _apply_search_filter_to_models( # name shown in the UI is searchable. if search_lower in (m.get("model_name") or "").lower(): return True - team_public_model_name = (m.get("model_info") or {}).get( - "team_public_model_name" - ) or "" + team_public_model_name = (m.get("model_info") or {}).get("team_public_model_name") or "" return search_lower in team_public_model_name.lower() # Filter models in router by search term, dropping BYOK rows that @@ -12017,8 +11016,7 @@ async def _apply_search_filter_to_models( filtered_router_models = [ m for m in all_models - if _model_matches_search(m) - and not _is_byok_outside_caller_teams(m.get("model_info") or {}) + if _model_matches_search(m) and not _is_byok_outside_caller_teams(m.get("model_info") or {}) ] # Separate filtered models into config vs db models, and track db model IDs @@ -12056,9 +11054,7 @@ async def _apply_search_filter_to_models( ) search_total_count = router_models_count + db_models_total_count except Exception as e: - verbose_proxy_logger.exception( - f"Error querying database models with search: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error querying database models with search: {str(e)}") search_total_count = router_models_count else: search_total_count = router_models_count @@ -12237,9 +11233,7 @@ def _paginate_models_response( } -def _team_models_resolve_to_names( - team_models: List[str], access_groups: Dict[str, Any] -) -> List[str]: +def _team_models_resolve_to_names(team_models: List[str], access_groups: Dict[str, Any]) -> List[str]: """Expand team model entries (including access group names) to concrete model names.""" resolved: List[str] = [] for name in team_models: @@ -12250,14 +11244,10 @@ def _team_models_resolve_to_names( return resolved -async def _load_team_object_for_model_filter( - team_id: str, prisma_client: PrismaClient -) -> Optional[LiteLLM_TeamTable]: +async def _load_team_object_for_model_filter(team_id: str, prisma_client: PrismaClient) -> Optional[LiteLLM_TeamTable]: """Load team row from DB; returns None if missing or on error.""" try: - team_db_object = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + team_db_object = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if team_db_object is None: verbose_proxy_logger.warning(f"Team {team_id} not found in database") return None @@ -12277,10 +11267,7 @@ async def _gather_team_accessible_model_ids( team_accessible_model_ids: Set[str] = set() access_groups = llm_router.get_model_access_groups() if llm_router else {} - if ( - not team_object.models - or SpecialModelNames.all_proxy_models.value in team_object.models - ): + if not team_object.models or SpecialModelNames.all_proxy_models.value in team_object.models: model_list = llm_router.get_model_list() if llm_router else [] if model_list is not None: for model in model_list: @@ -12299,11 +11286,7 @@ async def _gather_team_accessible_model_ids( resolved_model_names.add(model_name) for model_name in resolved_model_names: - _models = ( - llm_router.get_model_list(model_name=model_name, team_id=team_id) - if llm_router - else [] - ) + _models = llm_router.get_model_list(model_name=model_name, team_id=team_id) if llm_router else [] if _models is not None: for model in _models: model_id = model.get("model_info", {}).get("id", None) @@ -12311,13 +11294,8 @@ async def _gather_team_accessible_model_ids( team_accessible_model_ids.add(model_id) try: - if ( - team_object.models - and SpecialModelNames.all_proxy_models.value not in team_object.models - ): - _resolved_names = _team_models_resolve_to_names( - team_object.models, access_groups - ) + if team_object.models and SpecialModelNames.all_proxy_models.value not in team_object.models: + _resolved_names = _team_models_resolve_to_names(team_object.models, access_groups) db_models = await ModelRepository(prisma_client).table.find_many( where={"model_name": {"in": _resolved_names}} ) @@ -12325,9 +11303,7 @@ async def _gather_team_accessible_model_ids( if db_model.model_id: team_accessible_model_ids.add(db_model.model_id) except Exception as e: - verbose_proxy_logger.debug( - f"Error querying database models for team {team_id}: {str(e)}" - ) + verbose_proxy_logger.debug(f"Error querying database models for team {team_id}: {str(e)}") return team_accessible_model_ids @@ -12357,13 +11333,9 @@ async def _authorize_team_id_query( detail={"error": "Not authorized to view this team's models"}, ) try: - user_row = await UserRepository(prisma_client).table.find_unique( - where={"user_id": user_id} - ) + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) except Exception: - verbose_proxy_logger.exception( - "Failed to look up caller teams while authorizing teamId filter" - ) + verbose_proxy_logger.exception("Failed to look up caller teams while authorizing teamId filter") raise HTTPException( status_code=403, detail={"error": "Not authorized to view this team's models"}, @@ -12411,9 +11383,7 @@ async def _filter_models_by_team_id( if team_object is None: return [] - team_accessible_model_ids = await _gather_team_accessible_model_ids( - team_object, team_id, prisma_client, llm_router - ) + team_accessible_model_ids = await _gather_team_accessible_model_ids(team_object, team_id, prisma_client, llm_router) # When filtering by a specific team we want exactly the models that team # can use: its BYOK rows and the deployments resolved from team.models / @@ -12464,18 +11434,14 @@ async def _find_model_by_id( # If not found in config, search in database if found_model is None: try: - db_model = await ModelRepository(prisma_client).table.find_unique( - where={"model_id": model_id} - ) + db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}) if db_model: # Convert database model to router format decrypted_models = proxy_config.decrypt_model_list_from_db([db_model]) if decrypted_models: found_model = decrypted_models[0] except Exception as e: - verbose_proxy_logger.exception( - f"Error querying database for modelId {model_id}: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error querying database for modelId {model_id}: {str(e)}") # If model found, verify search filter if provided if found_model is not None: @@ -12499,24 +11465,16 @@ async def _find_model_by_id( ) async def model_info_v2( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - model: Optional[str] = fastapi.Query( - None, description="Specify the model name (optional)" - ), - user_models_only: Optional[bool] = fastapi.Query( - False, description="Only return models added by this user" - ), + model: Optional[str] = fastapi.Query(None, description="Specify the model name (optional)"), + user_models_only: Optional[bool] = fastapi.Query(False, description="Only return models added by this user"), include_team_models: Optional[bool] = fastapi.Query( False, description="Return all models across all teams user is in." ), debug: Optional[bool] = False, page: int = Query(1, description="Page number", ge=1), size: int = Query(50, description="Page size", ge=1), - search: Optional[str] = fastapi.Query( - None, description="Search model names (case-insensitive partial match)" - ), - modelId: Optional[str] = fastapi.Query( - None, description="Search for a specific model by its unique ID" - ), + search: Optional[str] = fastapi.Query(None, description="Search model names (case-insensitive partial match)"), + modelId: Optional[str] = fastapi.Query(None, description="Search for a specific model by its unique ID"), teamId: Optional[str] = fastapi.Query( None, description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", @@ -12575,12 +11533,7 @@ async def model_info_v2( } ``` """ - global \ - llm_model_list, \ - general_settings, \ - user_config_file_path, \ - proxy_config, \ - llm_router + global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router # Return empty data array when no models are configured (graceful handling for fresh installs) if llm_router is None or not llm_router.model_list: @@ -12791,9 +11744,7 @@ async def model_streaming_metrics( """ _all_api_bases = set() - db_response = await prisma_client.db.query_raw( - sql_query, _selected_model_group, startTime, endTime - ) + db_response = await prisma_client.db.query_raw(sql_query, _selected_model_group, startTime, endTime) _daily_entries: dict = {} # {"Jun 23": {"model1": 0.002, "model2": 0.003}} if db_response is not None: for model_data in db_response: @@ -12996,8 +11947,7 @@ async def model_metrics_slow_responses( endTime = endTime or datetime.now() alerting_threshold = ( - proxy_logging_obj.slack_alerting_instance.alerting_threshold - or DEFAULT_SLACK_ALERTING_THRESHOLD + proxy_logging_obj.slack_alerting_instance.alerting_threshold or DEFAULT_SLACK_ALERTING_THRESHOLD ) alerting_threshold = int(alerting_threshold) @@ -13107,9 +12057,7 @@ async def model_metrics_exceptions( ORDER BY total_exceptions DESC LIMIT 200; """ - db_response = await prisma_client.db.query_raw( - sql_query, startTime, endTime, _selected_model_group, api_key - ) + db_response = await prisma_client.db.query_raw(sql_query, startTime, endTime, _selected_model_group, api_key) response: List[dict] = [] exception_types = set() @@ -13141,9 +12089,7 @@ async def model_metrics_exceptions( return {"data": response, "exception_types": list(exception_types)} -def _deployment_matches_allowed_model_names( - model: Dict[str, Any], allowed_model_names: Set[str] -) -> bool: +def _deployment_matches_allowed_model_names(model: Dict[str, Any], allowed_model_names: Set[str]) -> bool: """Match a router deployment against allowed public model names. Team-scoped rows store an internal routing key in ``model_name``; callers @@ -13156,10 +12102,7 @@ def _deployment_matches_allowed_model_names( if not isinstance(model_info, dict): return False team_public_model_name = model_info.get("team_public_model_name") - return ( - isinstance(team_public_model_name, str) - and team_public_model_name in allowed_model_names - ) + return isinstance(team_public_model_name, str) and team_public_model_name in allowed_model_names def _get_v1_model_info_allowed_model_names( @@ -13200,11 +12143,7 @@ def _filter_v1_model_info_deployments( ) -> List[dict]: if allowed_model_names is None: return all_models - return [ - model - for model in all_models - if _deployment_matches_allowed_model_names(model, allowed_model_names) - ] + return [model for model in all_models if _deployment_matches_allowed_model_names(model, allowed_model_names)] def _translate_model_name_for_response(model: dict) -> dict: @@ -13260,9 +12199,7 @@ def _get_proxy_model_info(model: dict) -> dict: if len(split_model) > 0: litellm_model = split_model[-1] try: - litellm_model_info = litellm.get_model_info( - model=litellm_model, custom_llm_provider=split_model[0] - ) + litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0]) except Exception: litellm_model_info = {} for k, v in litellm_model_info.items(): @@ -13270,9 +12207,7 @@ def _get_proxy_model_info(model: dict) -> dict: model_info[k] = v model["model_info"] = model_info # don't return the llm credentials - model = remove_sensitive_info_from_deployment( - deployment_dict=model, excluded_keys={"litellm_credential_name"} - ) + model = remove_sensitive_info_from_deployment(deployment_dict=model, excluded_keys={"litellm_credential_name"}) return _translate_model_name_for_response(model) @@ -13336,13 +12271,7 @@ async def model_info_v1( ``` """ - global \ - llm_model_list, \ - general_settings, \ - user_config_file_path, \ - proxy_config, \ - llm_router, \ - user_model + global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router, user_model # Unit tests call this handler directly; FastAPI normally resolves Query defaults. if not isinstance(include_team_models, bool): @@ -13386,9 +12315,7 @@ async def model_info_v1( }, ) - if prisma_client is None and ( - include_team_models or (teamId is not None and teamId.strip()) - ): + if prisma_client is None and (include_team_models or (teamId is not None and teamId.strip())): raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, @@ -13400,13 +12327,9 @@ async def model_info_v1( if deployment_info is None: raise HTTPException( status_code=400, - detail={ - "error": f"Model id = {litellm_model_id} not found on litellm proxy" - }, + detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, ) - _deployment_info_dict = _get_proxy_model_info( - model=deployment_info.model_dump(exclude_none=True) - ) + _deployment_info_dict = _get_proxy_model_info(model=deployment_info.model_dump(exclude_none=True)) single_model_list: List[dict] = [_deployment_info_dict] if prisma_client is not None: single_model_list = await _populate_team_access_on_models( @@ -13457,9 +12380,7 @@ async def model_info_v1( all_models = [ model for model in all_models - if not _byok_row_outside_caller_teams( - model.get("model_info") or {}, allowed_team_ids - ) + if not _byok_row_outside_caller_teams(model.get("model_info") or {}, allowed_team_ids) ] if prisma_client is not None: @@ -13474,9 +12395,7 @@ async def model_info_v1( all_models = _filter_models_to_user_accessible(all_models) all_models = [ - _translate_model_name_for_response( - _enrich_model_info_with_litellm_data(model=model, llm_router=llm_router) - ) + _translate_model_name_for_response(_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router)) for model in all_models ] @@ -13686,12 +12605,7 @@ async def model_group_info( } ``` """ - global \ - llm_model_list, \ - general_settings, \ - user_config_file_path, \ - proxy_config, \ - llm_router + global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router # Return empty data array when no models are configured (graceful handling for fresh installs) if llm_model_list is None or llm_router is None or not llm_model_list: @@ -13832,9 +12746,7 @@ async def alerting_settings( is_slack_enabled = False - if general_settings.get("alerting") and isinstance( - general_settings["alerting"], list - ): + if general_settings.get("alerting") and isinstance(general_settings["alerting"], list): if "slack" in general_settings["alerting"]: is_slack_enabled = True @@ -13864,9 +12776,7 @@ async def alerting_settings( field_value=_slack_alerting_args_dict.get(field_name, None), stored_in_db=_stored_in_db, field_default_value=field_info.default, - premium_field=( - True if field_name == "region_outage_alert_ttl" else False - ), + premium_field=(True if field_name == "region_outage_alert_ttl" else False), ) return_val.append(_response_obj) return return_val @@ -13925,20 +12835,12 @@ async def async_queue_request( data["metadata"]["user_api_key"] = user_api_key_dict.api_key data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata _headers = _safe_get_request_headers(request).copy() - _headers.pop( - "authorization", None - ) # do not store the original `sk-..` api key in the db + _headers.pop("authorization", None) # do not store the original `sk-..` api key in the db data["metadata"]["headers"] = _headers - data["metadata"]["user_api_key_alias"] = getattr( - user_api_key_dict, "key_alias", None - ) + data["metadata"]["user_api_key_alias"] = getattr(user_api_key_dict, "key_alias", None) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - data["metadata"]["user_api_key_team_id"] = getattr( - user_api_key_dict, "team_id", None - ) - data["metadata"]["user_api_key_object_permission_id"] = getattr( - user_api_key_dict, "object_permission_id", None - ) + data["metadata"]["user_api_key_team_id"] = getattr(user_api_key_dict, "team_id", None) + data["metadata"]["user_api_key_object_permission_id"] = getattr(user_api_key_dict, "object_permission_id", None) data["metadata"]["user_api_key_team_object_permission_id"] = getattr( user_api_key_dict, "team_object_permission_id", None ) @@ -13956,15 +12858,11 @@ async def async_queue_request( data["api_base"] = user_api_base if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} - ) + raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) response = await llm_router.schedule_acompletion(**data) - if ( - "stream" in data and data["stream"] is True - ): # use generate_responses to stream responses + if "stream" in data and data["stream"] is True: # use generate_responses to stream responses return StreamingResponse( async_data_generator( user_api_key_dict=user_api_key_dict, @@ -14029,9 +12927,7 @@ async def fallback_login(request: Request): ) -@router.post( - "/login", include_in_schema=False -) # hidden since this is a helper for UI sso login +@router.post("/login", include_in_schema=False) # hidden since this is a helper for UI sso login async def login(request: Request): global premium_user, general_settings, master_key from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object @@ -14079,9 +12975,7 @@ async def login(request: Request): return redirect_response -@router.post( - "/v2/login", include_in_schema=False -) # hidden helper for UI logins via API +@router.post("/v2/login", include_in_schema=False) # hidden helper for UI logins via API async def login_v2(request: Request): global premium_user, general_settings, master_key from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object @@ -14130,11 +13024,7 @@ async def login_v2(request: Request): json_response.set_cookie(key="token", value=jwt_token) return json_response except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.login_v2(): Exception occurred - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.login_v2(): Exception occurred - {}".format(str(e))) if isinstance(e, ProxyException): raise e elif isinstance(e, HTTPException): @@ -14208,24 +13098,16 @@ async def login_v3(request: Request): cache_key = f"login_code:{code}" cache_value = {"token": jwt_token, "redirect_url": litellm_dashboard_ui} if redis_usage_cache is not None: - await redis_usage_cache.async_set_cache( - key=cache_key, value=cache_value, ttl=60 - ) + await redis_usage_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) else: - await user_api_key_cache.async_set_cache( - key=cache_key, value=cache_value, ttl=60 - ) + await user_api_key_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) return JSONResponse( content={"code": code, "expires_in": 60}, status_code=status.HTTP_200_OK, ) except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format(str(e))) if isinstance(e, ProxyException): raise e elif isinstance(e, HTTPException): @@ -14245,9 +13127,7 @@ async def login_v3(request: Request): ) -@router.post( - "/v3/login/exchange", include_in_schema=False -) # exchange single-use opaque code for JWT +@router.post("/v3/login/exchange", include_in_schema=False) # exchange single-use opaque code for JWT async def login_v3_exchange(request: Request): try: if not general_settings.get("control_plane_url"): @@ -14301,9 +13181,7 @@ async def login_v3_exchange(request: Request): raise except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format(str(e)) ) raise ProxyException( message=str(e), @@ -14339,21 +13217,15 @@ async def onboarding(invite_link: str, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique( - where={"id": invite_link} - ) + invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": invite_link}) if invite_obj is None: - raise HTTPException( - status_code=401, detail={"error": "Invitation link does not exist in db."} - ) + raise HTTPException(status_code=401, detail={"error": "Invitation link does not exist in db."}) #### CHECK IF EXPIRED # Extract the date part from both datetime objects utc_now_date = litellm.utils.get_utc_datetime().date() expires_at_date = invite_obj.expires_at.date() if expires_at_date < utc_now_date: - raise HTTPException( - status_code=401, detail={"error": "Invitation link has expired."} - ) + raise HTTPException(status_code=401, detail={"error": "Invitation link has expired."}) #### CHECK IF ALREADY USED if invite_obj.is_accepted is True or invite_obj.accepted_at is not None: @@ -14363,14 +13235,10 @@ async def onboarding(invite_link: str, request: Request): ) ### GET USER OBJECT ### - user_obj = await UserRepository(prisma_client).table.find_unique( - where={"user_id": invite_obj.user_id} - ) + user_obj = await UserRepository(prisma_client).table.find_unique(where={"user_id": invite_obj.user_id}) if user_obj is None: - raise HTTPException( - status_code=401, detail={"error": "User does not exist in db."} - ) + raise HTTPException(status_code=401, detail={"error": "User does not exist in db."}) litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): @@ -14390,9 +13258,7 @@ async def onboarding(invite_link: str, request: Request): master_key, algorithm="HS256", ) - disabled_non_admin_personal_key_creation = ( - get_disabled_non_admin_personal_key_creation() - ) + disabled_non_admin_personal_key_creation = get_disabled_non_admin_personal_key_creation() returned_ui_token_object = ReturnedUITokenObject( user_id=user_obj.user_id, @@ -14401,9 +13267,7 @@ async def onboarding(invite_link: str, request: Request): user_role=user_obj.user_role, login_method="username_password", premium_user=premium_user, - auth_header_name=general_settings.get( - "litellm_key_header_name", "Authorization" - ), + auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"), disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, server_root_path=get_server_root_path(), ) @@ -14478,9 +13342,7 @@ async def _rollback_onboarding_invite_claim( }, ) except Exception: - verbose_proxy_logger.exception( - "Failed to roll back onboarding invitation after session key mint failed." - ) + verbose_proxy_logger.exception("Failed to roll back onboarding invitation after session key mint failed.") async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: @@ -14506,9 +13368,7 @@ async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: from litellm.types.proxy.ui_sso import ReturnedUITokenObject - disabled_non_admin_personal_key_creation = ( - get_disabled_non_admin_personal_key_creation() - ) + disabled_non_admin_personal_key_creation = get_disabled_non_admin_personal_key_creation() returned_ui_token_object = ReturnedUITokenObject( user_id=user_obj.user_id, key=key, @@ -14516,9 +13376,7 @@ async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: user_role=user_obj.user_role, login_method="username_password", premium_user=premium_user, - auth_header_name=general_settings.get( - "litellm_key_header_name", "Authorization" - ), + auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"), disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, server_root_path=get_server_root_path(), ) @@ -14551,21 +13409,15 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique( - where={"id": data.invitation_link} - ) + invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": data.invitation_link}) if invite_obj is None: - raise HTTPException( - status_code=401, detail={"error": "Invitation link does not exist in db."} - ) + raise HTTPException(status_code=401, detail={"error": "Invitation link does not exist in db."}) #### CHECK IF EXPIRED # Extract the date part from both datetime objects utc_now_date = litellm.utils.get_utc_datetime().date() expires_at_date = invite_obj.expires_at.date() if expires_at_date < utc_now_date: - raise HTTPException( - status_code=401, detail={"error": "Invitation link has expired."} - ) + raise HTTPException(status_code=401, detail={"error": "Invitation link has expired."}) #### CHECK IF ALREADY USED if invite_obj.is_accepted is True or invite_obj.accepted_at is not None: @@ -14619,9 +13471,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): ) if user_obj is None: - raise HTTPException( - status_code=401, detail={"error": "User does not exist in db."} - ) + raise HTTPException(status_code=401, detail={"error": "User does not exist in db."}) #### MARK LINK AS USED current_time = litellm.utils.get_utc_datetime() @@ -14648,9 +13498,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): raise e raise HTTPException( status_code=500, - detail={ - "error": "Failed to create onboarding session. Please retry the invitation link." - }, + detail={"error": "Failed to create onboarding session. Please retry the invitation link."}, ) from e litellm_dashboard_ui = get_custom_url(str(request.base_url)) @@ -14713,11 +13561,7 @@ async def get_image(): assets_dir = current_dir # Determine default logo path - default_logo = ( - os.path.join(assets_dir, "logo.jpg") - if assets_dir != current_dir - else default_site_logo - ) + default_logo = os.path.join(assets_dir, "logo.jpg") if assets_dir != current_dir else default_site_logo if assets_dir != current_dir and not os.path.exists(default_logo): default_logo = default_site_logo @@ -14734,8 +13578,7 @@ async def get_image(): safe_logo_path, media_type = safe_logo return FileResponse(safe_logo_path, media_type=media_type) verbose_proxy_logger.warning( - "UI_LOGO_PATH %r is not a supported image file or does not exist, " - "falling back to default logo", + "UI_LOGO_PATH %r is not a supported image file or does not exist, falling back to default logo", logo_path, ) logo_path = default_logo @@ -14778,8 +13621,7 @@ async def get_favicon(): safe_favicon_path, media_type = safe_favicon return FileResponse(safe_favicon_path, media_type=media_type) verbose_proxy_logger.warning( - "LITELLM_FAVICON_URL %r is not a supported image file or does not " - "exist, falling back to default favicon", + "LITELLM_FAVICON_URL %r is not a supported image file or does not exist, falling back to default favicon", favicon_url, ) if os.path.exists(default_favicon): @@ -14797,9 +13639,7 @@ async def get_favicon(): response_model=InvitationModel, include_in_schema=False, ) -async def new_invitation( - data: InvitationNew, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth) -): +async def new_invitation(data: InvitationNew, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): """ Allow admin to create invite links, to onboard new users to Admin UI. @@ -14825,14 +13665,11 @@ async def new_invitation( ) # Allow proxy admins and org/team admins (admin status from DB via get_user_object) - has_access = ( - user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - or await _user_has_admin_privileges( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + has_access = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN or await _user_has_admin_privileges( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) if not has_access: raise HTTPException( @@ -14857,9 +13694,7 @@ async def new_invitation( if not can_invite: raise HTTPException( status_code=400, - detail={ - "error": "You can only create invitations for users in your organization or team." - }, + detail={"error": "You can only create invitations for users in your organization or team."}, ) response = await create_invitation_for_user( @@ -14878,9 +13713,7 @@ async def new_invitation( response_model=InvitationModel, include_in_schema=False, ) -async def invitation_info( - invitation_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth) -): +async def invitation_info(invitation_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): """ Allow admin to create invite links, to onboard new users to Admin UI. @@ -14911,9 +13744,7 @@ async def invitation_info( }, ) - response = await InvitationLinkRepository(prisma_client).table.find_unique( - where={"id": invitation_id} - ) + response = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": invitation_id}) if response is None: raise HTTPException( @@ -14957,11 +13788,7 @@ async def invitation_update( if user_api_key_dict.user_id is None: raise HTTPException( status_code=500, - detail={ - "error": "Unable to identify user id. Received={}".format( - user_api_key_dict.user_id - ) - }, + detail={"error": "Unable to identify user id. Received={}".format(user_api_key_dict.user_id)}, ) current_time = litellm.utils.get_utc_datetime() @@ -15036,9 +13863,7 @@ async def invitation_delete( # Org admins can only delete invitations they created if is_other_admin and not is_proxy_admin: - invitation = await InvitationLinkRepository(prisma_client).table.find_unique( - where={"id": data.invitation_id} - ) + invitation = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": data.invitation_id}) if invitation is None: raise HTTPException( status_code=400, @@ -15047,14 +13872,10 @@ async def invitation_delete( if invitation.created_by != user_api_key_dict.user_id: raise HTTPException( status_code=403, - detail={ - "error": "Organization admins can only delete invitations they created." - }, + detail={"error": "Organization admins can only delete invitations they created."}, ) - response = await InvitationLinkRepository(prisma_client).table.delete( - where={"id": data.invitation_id} - ) + response = await InvitationLinkRepository(prisma_client).table.delete(where={"id": data.invitation_id}) if response is None: raise HTTPException( @@ -15083,27 +13904,16 @@ async def update_config( untouched — this endpoint never persists pre-existing YAML values to DB as a side effect of an unrelated update. """ - global \ - llm_router, \ - llm_model_list, \ - general_settings, \ - proxy_config, \ - proxy_logging_obj, \ - master_key, \ - prisma_client + global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client try: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=403, detail="Only proxy admins can update config" - ) + raise HTTPException(status_code=403, detail="Only proxy admins can update config") if prisma_client is None: raise Exception("No DB Connected") async def _read_section(param_name: str) -> dict: - row = await ConfigRepository(prisma_client).table.find_first( - where={"param_name": param_name} - ) + row = await ConfigRepository(prisma_client).table.find_first(where={"param_name": param_name}) if row is None or row.param_value is None: return {} return dict(row.param_value) @@ -15130,10 +13940,7 @@ async def update_config( if k == "alert_to_webhook_url": if "alerting" not in existing: existing["alerting"] = ["slack"] - elif ( - isinstance(existing["alerting"], list) - and "slack" not in existing["alerting"] - ): + elif isinstance(existing["alerting"], list) and "slack" not in existing["alerting"]: existing["alerting"].append("slack") existing[k] = v await _upsert_section("general_settings", existing) @@ -15146,9 +13953,7 @@ async def update_config( if config_info.environment_variables is not None: existing = await _read_section("environment_variables") existing.update( - proxy_config._encrypt_env_variables_for_db( - environment_variables=config_info.environment_variables - ) + proxy_config._encrypt_env_variables_for_db(environment_variables=config_info.environment_variables) ) await _upsert_section("environment_variables", existing) @@ -15166,9 +13971,7 @@ async def update_config( incoming_cb = updated_litellm_settings.get("success_callback") if isinstance(incoming_cb, list): - updated_litellm_settings["success_callback"] = normalize_callback_names( - incoming_cb - ) + updated_litellm_settings["success_callback"] = normalize_callback_names(incoming_cb) merged = {**existing, **updated_litellm_settings} @@ -15180,9 +13983,7 @@ async def update_config( # different code path may still hold mixed-case names, # which would otherwise dedup-miss against the lowercase # incoming entries. - merged["success_callback"] = list( - set(normalize_callback_names(existing_cb) + incoming_cb) - ) + merged["success_callback"] = list(set(normalize_callback_names(existing_cb) + incoming_cb)) else: merged["success_callback"] = list(set(incoming_cb)) @@ -15194,17 +13995,11 @@ async def update_config( updates = config_info.router_settings.dict(exclude_none=True) await _upsert_section("router_settings", {**existing, **updates}) - await proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) + await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) return {"message": "Config updated successfully"} except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.update_config(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.update_config(): Exception occured - {}".format(str(e))) verbose_proxy_logger.debug(traceback.format_exc()) if isinstance(e, HTTPException): raise ProxyException( @@ -15334,9 +14129,7 @@ async def update_config_general_settings( field_value = data.field_value if data.field_name == "plugins": - field_value = _preserve_redacted_plugin_keys( - field_value, general_settings.get("plugins") - ) + field_value = _preserve_redacted_plugin_keys(field_value, general_settings.get("plugins")) general_settings[data.field_name] = field_value @@ -15373,10 +14166,7 @@ _EXTRA_SECRET_GENERAL_SETTINGS_FIELDS = frozenset( def _is_secret_general_setting_field(field_name: str) -> bool: - return ( - field_name in _EXTRA_SECRET_GENERAL_SETTINGS_FIELDS - or SENSITIVE_DATA_MASKER.is_sensitive_key(field_name) - ) + return field_name in _EXTRA_SECRET_GENERAL_SETTINGS_FIELDS or SENSITIVE_DATA_MASKER.is_sensitive_key(field_name) # Matches the cap on _redact_sensitive_litellm_params (the closest analog in the @@ -15396,11 +14186,7 @@ def _redact_secret_values_in_obj(value: JsonValue, depth: int = 0) -> JsonValue: return "REDACTED" if isinstance(value, dict): return { - key: ( - "REDACTED" - if _is_secret_general_setting_field(key) - else _redact_secret_values_in_obj(sub, depth + 1) - ) + key: ("REDACTED" if _is_secret_general_setting_field(key) else _redact_secret_values_in_obj(sub, depth + 1)) for key, sub in value.items() } if isinstance(value, list): @@ -15408,9 +14194,7 @@ def _redact_secret_values_in_obj(value: JsonValue, depth: int = 0) -> JsonValue: return value -def _redact_general_setting_value( - field_name: str, value: JsonValue, is_full_admin: bool -) -> JsonValue: +def _redact_general_setting_value(field_name: str, value: JsonValue, is_full_admin: bool) -> JsonValue: if is_full_admin: return value if _is_secret_general_setting_field(field_name): @@ -15479,11 +14263,7 @@ async def get_config_general_settings( ) if field_name == "plugins" and isinstance(field_value, list): field_value = [ - ( - {k: ("***" if k == "plugin_key" else v) for k, v in p.items()} - if isinstance(p, dict) - else p - ) + ({k: ("***" if k == "plugin_key" else v) for k, v in p.items()} if isinstance(p, dict) else p) for p in field_value ] return ConfigFieldInfo(field_name=field_name, field_value=field_value) @@ -15598,13 +14378,8 @@ async def get_config_list( sub_field, sub_field_info, ) in pydantic_class.model_fields.items(): - if ( - hasattr(sub_field_info, "description") - and sub_field_info.description is not None - ): - nested_fields[ - idx - ].field_description = sub_field_info.description + if hasattr(sub_field_info, "description") and sub_field_info.description is not None: + nested_fields[idx].field_description = sub_field_info.description idx += 1 _stored_in_db = None @@ -15645,9 +14420,7 @@ async def get_config_list( field_name=field_name, field_type=allowed_args[field_name]["type"], field_description=field_info.description or "", - field_value=_redact_general_setting_value( - field_name, _field_value, is_full_admin - ), + field_value=_redact_general_setting_value(field_name, _field_value, is_full_admin), stored_in_db=_stored_in_db, field_default_value=field_info.default, nested_fields=nested_fields, @@ -15768,9 +14541,7 @@ async def delete_callback( if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) try: @@ -15785,24 +14556,18 @@ async def delete_callback( if callback_name not in success_callbacks: raise HTTPException( status_code=404, - detail={ - "error": f"Callback '{callback_name}' not found in active configuration" - }, + detail={"error": f"Callback '{callback_name}' not found in active configuration"}, ) # Remove callback from success_callback list success_callbacks.remove(callback_name) - config.setdefault("litellm_settings", {})["success_callback"] = ( - success_callbacks - ) + config.setdefault("litellm_settings", {})["success_callback"] = success_callbacks # Save the updated configuration await proxy_config.save_config(new_config=config) # Restart the proxy to apply changes - await proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) + await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) return { "message": f"Successfully deleted callback: {callback_name}", @@ -15814,9 +14579,7 @@ async def delete_callback( except HTTPException: raise except Exception as e: - verbose_proxy_logger.error( - f"litellm.proxy.proxy_server.delete_callback(): Exception occurred - {str(e)}" - ) + verbose_proxy_logger.error(f"litellm.proxy.proxy_server.delete_callback(): Exception occurred - {str(e)}") verbose_proxy_logger.debug(traceback.format_exc()) raise ProxyException( message="Error deleting callback: " + str(e), @@ -15838,13 +14601,7 @@ async def get_config(): # return the callbacks and the env variables for the callback """ - global \ - llm_router, \ - llm_model_list, \ - general_settings, \ - proxy_config, \ - proxy_logging_obj, \ - master_key + global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key try: all_available_callbacks = AllCallbacks() @@ -15867,9 +14624,7 @@ async def get_config(): _success_callbacks = normalize_callback(_success_callbacks) _failure_callbacks = normalize_callback(_failure_callbacks) - _success_and_failure_callbacks = normalize_callback( - _success_and_failure_callbacks - ) + _success_and_failure_callbacks = normalize_callback(_success_and_failure_callbacks) _data_to_return = [] """ @@ -15888,21 +14643,13 @@ async def get_config(): """ for _callback in _success_callbacks: - _data_to_return.append( - process_callback(_callback, "success", environment_variables) - ) + _data_to_return.append(process_callback(_callback, "success", environment_variables)) for _callback in _failure_callbacks: - _data_to_return.append( - process_callback(_callback, "failure", environment_variables) - ) + _data_to_return.append(process_callback(_callback, "failure", environment_variables)) for _callback in _success_and_failure_callbacks: - _data_to_return.append( - process_callback( - _callback, "success_and_failure", environment_variables - ) - ) + _data_to_return.append(process_callback(_callback, "success_and_failure", environment_variables)) # Check if slack alerting is on _alerting = _general_settings.get("alerting", []) @@ -15912,24 +14659,14 @@ async def get_config(): "SLACK_WEBHOOK_URL", ] _slack_env_vars = { - _var: ( - value - if (value := environment_variables.get(_var)) is not None - else os.getenv(_var) - ) + _var: (value if (value := environment_variables.get(_var)) is not None else os.getenv(_var)) for _var in _slack_vars } - _slack_env_vars = mask_sensitive_keys( - _slack_env_vars, _ALERTING_SENSITIVE_VARS - ) + _slack_env_vars = mask_sensitive_keys(_slack_env_vars, _ALERTING_SENSITIVE_VARS) _alerting_types = proxy_logging_obj.slack_alerting_instance.alert_types - _all_alert_types = ( - proxy_logging_obj.slack_alerting_instance._all_possible_alert_types() - ) - _alerts_to_webhook = ( - proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url - ) + _all_alert_types = proxy_logging_obj.slack_alerting_instance._all_possible_alert_types() + _alerts_to_webhook = proxy_logging_obj.slack_alerting_instance.alert_to_webhook_url alerting_data.append( { "name": "slack", @@ -15949,9 +14686,7 @@ async def get_config(): "EMAIL_LOGO_URL", "EMAIL_SUPPORT_CONTACT", ] - _email_env_vars = { - _var: environment_variables.get(_var) for _var in _email_vars - } + _email_env_vars = {_var: environment_variables.get(_var) for _var in _email_vars} _email_env_vars = mask_sensitive_keys(_email_env_vars, _ALERTING_SENSITIVE_VARS) alerting_data.append( @@ -15974,11 +14709,7 @@ async def get_config(): "available_callbacks": all_available_callbacks, } except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.get_config(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.get_config(): Exception occured - {}".format(str(e))) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "detail", f"Authentication Error({str(e)})"), @@ -16047,9 +14778,7 @@ async def reload_model_cost_map( try: global prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail="Database connection not available" - ) + raise HTTPException(status_code=500, detail="Database connection not available") # Immediately reload the model cost map in the current pod from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map @@ -16081,23 +14810,15 @@ async def reload_model_cost_map( data={ "create": { "param_name": "model_cost_map_reload_config", - "param_value": safe_dumps( - {"interval_hours": None, "force_reload": True} - ), - }, - "update": { - "param_value": safe_dumps( - {"interval_hours": existing_interval, "force_reload": True} - ) + "param_value": safe_dumps({"interval_hours": None, "force_reload": True}), }, + "update": {"param_value": safe_dumps({"interval_hours": existing_interval, "force_reload": True})}, }, ) await invalidate_config_param("model_cost_map_reload_config") models_count = len(new_model_cost_map) if new_model_cost_map else 0 - verbose_proxy_logger.info( - f"Model cost map reloaded successfully in current pod. Models count: {models_count}" - ) + verbose_proxy_logger.info(f"Model cost map reloaded successfully in current pod. Models count: {models_count}") return { "message": f"Price data reloaded successfully! {models_count} models updated.", @@ -16107,9 +14828,7 @@ async def reload_model_cost_map( } except Exception as e: verbose_proxy_logger.exception(f"Failed to reload model cost map: {str(e)}") - raise HTTPException( - status_code=500, detail=f"Failed to reload model cost map: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Failed to reload model cost map: {str(e)}") @router.post( @@ -16141,9 +14860,7 @@ async def schedule_model_cost_map_reload( try: global prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail="Database connection not available" - ) + raise HTTPException(status_code=500, detail="Database connection not available") # Update database with new reload configuration await ConfigRepository(prisma_client).table.upsert( @@ -16151,22 +14868,14 @@ async def schedule_model_cost_map_reload( data={ "create": { "param_name": "model_cost_map_reload_config", - "param_value": safe_dumps( - {"interval_hours": hours, "force_reload": False} - ), - }, - "update": { - "param_value": safe_dumps( - {"interval_hours": hours, "force_reload": False} - ) + "param_value": safe_dumps({"interval_hours": hours, "force_reload": False}), }, + "update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})}, }, ) await invalidate_config_param("model_cost_map_reload_config") - verbose_proxy_logger.info( - f"Model cost map reload scheduled for every {hours} hours" - ) + verbose_proxy_logger.info(f"Model cost map reload scheduled for every {hours} hours") return { "message": f"Model cost map reload scheduled for every {hours} hours", @@ -16175,9 +14884,7 @@ async def schedule_model_cost_map_reload( "timestamp": datetime.utcnow().isoformat(), } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to schedule model cost map reload: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to schedule model cost map reload: {str(e)}") raise HTTPException( status_code=500, detail=f"Failed to schedule model cost map reload: {str(e)}", @@ -16208,14 +14915,10 @@ async def cancel_model_cost_map_reload( try: global prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail="Database connection not available" - ) + raise HTTPException(status_code=500, detail="Database connection not available") # Remove reload configuration from database - await ConfigRepository(prisma_client).table.delete( - where={"param_name": "model_cost_map_reload_config"} - ) + await ConfigRepository(prisma_client).table.delete(where={"param_name": "model_cost_map_reload_config"}) await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info("Model cost map reload schedule cancelled") @@ -16226,12 +14929,8 @@ async def cancel_model_cost_map_reload( "timestamp": datetime.utcnow().isoformat(), } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to cancel model cost map reload: {str(e)}" - ) - raise HTTPException( - status_code=500, detail=f"Failed to cancel model cost map reload: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to cancel model cost map reload: {str(e)}") + raise HTTPException(status_code=500, detail=f"Failed to cancel model cost map reload: {str(e)}") @router.get( @@ -16258,9 +14957,7 @@ async def get_model_cost_map_reload_status( try: global prisma_client, last_model_cost_map_reload - verbose_proxy_logger.info( - f"Checking model cost map reload status. Last reload: {last_model_cost_map_reload}" - ) + verbose_proxy_logger.info(f"Checking model cost map reload status. Last reload: {last_model_cost_map_reload}") if prisma_client is None: verbose_proxy_logger.info("No database connection, returning not scheduled") @@ -16308,9 +15005,7 @@ async def get_model_cost_map_reload_status( hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 if hours_since_last_reload < interval_hours: - next_run = ( - last_reload_time + timedelta(hours=interval_hours) - ).isoformat() + next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat() except Exception as e: verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") @@ -16321,9 +15016,7 @@ async def get_model_cost_map_reload_status( "next_run": next_run, } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to get model cost map reload status: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to get model cost map reload status: {str(e)}") raise HTTPException( status_code=500, detail=f"Failed to get model cost map reload status: {str(e)}", @@ -16371,9 +15064,7 @@ async def get_model_cost_map_source( "model_count": model_count, } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to get model cost map source info: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to get model cost map source info: {str(e)}") raise HTTPException( status_code=500, detail=f"Failed to get model cost map source info: {str(e)}", @@ -16408,9 +15099,7 @@ async def reload_anthropic_beta_headers( try: global prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail="Database connection not available" - ) + raise HTTPException(status_code=500, detail="Database connection not available") # Immediately reload the beta headers config in the current pod from litellm.anthropic_beta_headers_manager import reload_beta_headers_config @@ -16428,31 +15117,21 @@ async def reload_anthropic_beta_headers( ) existing_beta_interval = None if existing_beta_config and existing_beta_config.param_value: - existing_beta_interval = existing_beta_config.param_value.get( - "interval_hours" - ) + existing_beta_interval = existing_beta_config.param_value.get("interval_hours") await ConfigRepository(prisma_client).table.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ "create": { "param_name": "anthropic_beta_headers_reload_config", - "param_value": safe_dumps( - {"interval_hours": None, "force_reload": True} - ), - }, - "update": { - "param_value": safe_dumps( - {"interval_hours": existing_beta_interval, "force_reload": True} - ) + "param_value": safe_dumps({"interval_hours": None, "force_reload": True}), }, + "update": {"param_value": safe_dumps({"interval_hours": existing_beta_interval, "force_reload": True})}, }, ) await invalidate_config_param("anthropic_beta_headers_reload_config") - provider_count = sum( - 1 for k in new_config.keys() if k not in ["provider_aliases", "description"] - ) + provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"]) verbose_proxy_logger.info( f"Anthropic beta headers config reloaded successfully in current pod. Providers: {provider_count}" ) @@ -16464,12 +15143,8 @@ async def reload_anthropic_beta_headers( "timestamp": current_time.isoformat(), } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to reload anthropic beta headers: {str(e)}" - ) - raise HTTPException( - status_code=500, detail=f"Failed to reload anthropic beta headers: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to reload anthropic beta headers: {str(e)}") + raise HTTPException(status_code=500, detail=f"Failed to reload anthropic beta headers: {str(e)}") @router.post( @@ -16501,9 +15176,7 @@ async def schedule_anthropic_beta_headers_reload( try: global prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail="Database connection not available" - ) + raise HTTPException(status_code=500, detail="Database connection not available") # Update database with new reload configuration await ConfigRepository(prisma_client).table.upsert( @@ -16511,22 +15184,14 @@ async def schedule_anthropic_beta_headers_reload( data={ "create": { "param_name": "anthropic_beta_headers_reload_config", - "param_value": safe_dumps( - {"interval_hours": hours, "force_reload": False} - ), - }, - "update": { - "param_value": safe_dumps( - {"interval_hours": hours, "force_reload": False} - ) + "param_value": safe_dumps({"interval_hours": hours, "force_reload": False}), }, + "update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})}, }, ) await invalidate_config_param("anthropic_beta_headers_reload_config") - verbose_proxy_logger.info( - f"Anthropic beta headers reload scheduled for every {hours} hours" - ) + verbose_proxy_logger.info(f"Anthropic beta headers reload scheduled for every {hours} hours") return { "message": f"Anthropic beta headers reload scheduled for every {hours} hours", @@ -16535,9 +15200,7 @@ async def schedule_anthropic_beta_headers_reload( "timestamp": datetime.utcnow().isoformat(), } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to schedule anthropic beta headers reload: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to schedule anthropic beta headers reload: {str(e)}") raise HTTPException( status_code=500, detail=f"Failed to schedule anthropic beta headers reload: {str(e)}", @@ -16568,14 +15231,10 @@ async def cancel_anthropic_beta_headers_reload( try: global prisma_client if prisma_client is None: - raise HTTPException( - status_code=500, detail="Database connection not available" - ) + raise HTTPException(status_code=500, detail="Database connection not available") # Remove reload configuration from database - await ConfigRepository(prisma_client).table.delete( - where={"param_name": "anthropic_beta_headers_reload_config"} - ) + await ConfigRepository(prisma_client).table.delete(where={"param_name": "anthropic_beta_headers_reload_config"}) await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled") @@ -16586,9 +15245,7 @@ async def cancel_anthropic_beta_headers_reload( "timestamp": datetime.utcnow().isoformat(), } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to cancel anthropic beta headers reload: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to cancel anthropic beta headers reload: {str(e)}") raise HTTPException( status_code=500, detail=f"Failed to cancel anthropic beta headers reload: {str(e)}", @@ -16638,9 +15295,7 @@ async def get_anthropic_beta_headers_reload_status( ) if config_record is None or config_record.param_value is None: - verbose_proxy_logger.info( - "No anthropic beta headers reload configuration found" - ) + verbose_proxy_logger.info("No anthropic beta headers reload configuration found") return { "scheduled": False, "interval_hours": None, @@ -16666,16 +15321,12 @@ async def get_anthropic_beta_headers_reload_status( # Use pod's in-memory last reload time if last_anthropic_beta_headers_reload is not None: try: - last_reload_time = datetime.fromisoformat( - last_anthropic_beta_headers_reload - ) + last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload) time_since_last_reload = current_time - last_reload_time hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 if hours_since_last_reload < interval_hours: - next_run = ( - last_reload_time + timedelta(hours=interval_hours) - ).isoformat() + next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat() except Exception as e: verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") @@ -16686,9 +15337,7 @@ async def get_anthropic_beta_headers_reload_status( "next_run": next_run, } except Exception as e: - verbose_proxy_logger.exception( - f"Failed to get anthropic beta headers reload status: {str(e)}" - ) + verbose_proxy_logger.exception(f"Failed to get anthropic beta headers reload status: {str(e)}") raise HTTPException( status_code=500, detail=f"Failed to get anthropic beta headers reload status: {str(e)}", @@ -16727,9 +15376,7 @@ async def get_adaptive_router_state( status_code=404, detail={"error": "No adaptive_router is configured on this proxy."}, ) - snapshots = [ - await ar.get_state_snapshot() for ar in llm_router.adaptive_routers.values() - ] + snapshots = [await ar.get_state_snapshot() for ar in llm_router.adaptive_routers.values()] return {"routers": snapshots} @@ -16838,9 +15485,7 @@ app.add_middleware( ) -async def _stream_mcp_asgi_response( - handle_fn, scope: dict, receive -) -> "StreamingResponse": +async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse": """ Call an ASGI MCP handler and return a StreamingResponse so SSE/streaming works. @@ -16855,9 +15500,7 @@ async def _stream_mcp_asgi_response( async def bridging_send(message): if message["type"] == "http.response.start": if not headers_ready.done(): - headers_ready.set_result( - (message.get("status", 200), message.get("headers", [])) - ) + headers_ready.set_result((message.get("status", 200), message.get("headers", []))) elif message["type"] == "http.response.body": chunk = message.get("body", b"") if chunk: @@ -16886,14 +15529,10 @@ async def _stream_mcp_asgi_response( handler_task.add_done_callback(_ensure_eof) try: - status, raw_headers = await asyncio.wait_for( - asyncio.shield(headers_ready), timeout=30.0 - ) + status, raw_headers = await asyncio.wait_for(asyncio.shield(headers_ready), timeout=30.0) except asyncio.TimeoutError: handler_task.cancel() - raise HTTPException( - status_code=504, detail="MCP handler did not respond in time" - ) + raise HTTPException(status_code=504, detail="MCP handler did not respond in time") headers_dict = {k.decode("latin-1"): v.decode("latin-1") for k, v in raw_headers} @@ -16950,9 +15589,7 @@ async def toolset_mcp_route(toolset_name: str, request: Request): if prisma_client is None: raise HTTPException(status_code=503, detail="Database not available") - toolset = await global_mcp_server_manager.get_toolset_by_name_cached( - prisma_client, toolset_name - ) + toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, toolset_name) if toolset is None: raise HTTPException( status_code=404, @@ -16964,18 +15601,14 @@ async def toolset_mcp_route(toolset_name: str, request: Request): token = _mcp_active_toolset_id.set(toolset.toolset_id) try: - return await _stream_mcp_asgi_response( - handle_streamable_http_mcp, scope, request.receive - ) + return await _stream_mcp_asgi_response(handle_streamable_http_mcp, scope, request.receive) finally: _mcp_active_toolset_id.reset(token) except HTTPException as e: raise e except Exception as e: - verbose_proxy_logger.exception( - "Error handling toolset MCP route for %s: %s", toolset_name, str(e) - ) + verbose_proxy_logger.exception("Error handling toolset MCP route for %s: %s", toolset_name, str(e)) raise HTTPException(status_code=500, detail="Internal server error") @@ -16989,14 +15622,10 @@ async def _mcp_forward_as_path(path_segment: str, request: Request): # Preserve the public request path for OAuth challenge URL selection. scope["_original_path"] = scope.get("path", "") scope["path"] = f"/mcp/{path_segment}" - return await _stream_mcp_asgi_response( - handle_streamable_http_mcp, scope, request.receive - ) + return await _stream_mcp_asgi_response(handle_streamable_http_mcp, scope, request.receive) -async def _resolve_mcp_csv_tokens( - csv_segment: str, client_ip: Optional[str] -) -> List[str]: +async def _resolve_mcp_csv_tokens(csv_segment: str, client_ip: Optional[str]) -> List[str]: """Validate a comma-separated ``/{name1,name2,...}/mcp`` segment. For each token, check (in order) whether it is a registered MCP server @@ -17068,11 +15697,7 @@ async def _is_mcp_access_group_cached(name: str) -> bool: await user_api_key_cache.async_set_cache( key=cache_key, value=result, - ttl=( - get_management_object_ttl(user_api_key_cache) - if result - else DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL - ), + ttl=(get_management_object_ttl(user_api_key_cache) if result else DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL), ) return result @@ -17100,9 +15725,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): client_ip = IPAddressUtils.get_mcp_client_ip(request) # 1. Registered MCP server alias - if global_mcp_server_manager.get_mcp_server_by_name( - mcp_server_name, client_ip=client_ip - ): + if global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip): return await _mcp_forward_as_path(mcp_server_name, request) # 2. Comma-separated list — validate every token resolves to a known @@ -17115,8 +15738,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): raise HTTPException( status_code=404, detail=( - f"No MCP server, toolset, or access group in " - f"'{mcp_server_name}' resolved to a known target" + f"No MCP server, toolset, or access group in '{mcp_server_name}' resolved to a known target" ), ) return await _mcp_forward_as_path(",".join(resolved_tokens), request) @@ -17128,18 +15750,14 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): handle_streamable_http_mcp, ) - toolset = await global_mcp_server_manager.get_toolset_by_name_cached( - prisma_client, mcp_server_name - ) + toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, mcp_server_name) if toolset is not None: scope = dict(request.scope) scope["_original_path"] = scope.get("path", "") scope["path"] = "/mcp" token = _mcp_active_toolset_id.set(toolset.toolset_id) try: - return await _stream_mcp_asgi_response( - handle_streamable_http_mcp, scope, request.receive - ) + return await _stream_mcp_asgi_response(handle_streamable_http_mcp, scope, request.receive) finally: _mcp_active_toolset_id.reset(token) @@ -17155,7 +15773,5 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): except HTTPException as e: raise e except Exception as e: - verbose_proxy_logger.exception( - "Error handling dynamic MCP route for %s: %s", mcp_server_name, str(e) - ) + verbose_proxy_logger.exception("Error handling dynamic MCP route for %s: %s", mcp_server_name, str(e)) raise HTTPException(status_code=500, detail="Internal server error") diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 78467c4b2e7..91332480d75 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -134,19 +134,13 @@ def _build_endpoints(raw: Dict[str, Any]) -> List[Dict[str, Any]]: for slug, pd in providers.items() if pd.get("endpoints", {}).get(key) ] - result.append( - {"key": key, "label": label, "endpoint": path, "providers": supporting} - ) + result.append({"key": key, "label": label, "endpoint": path, "providers": supporting}) return result def _load_endpoints() -> List[Dict[str, Any]]: - raw = json.loads( - files("litellm") - .joinpath("provider_endpoints_support_backup.json") - .read_text(encoding="utf-8") - ) + raw = json.loads(files("litellm").joinpath("provider_endpoints_support_backup.json").read_text(encoding="utf-8")) return _build_endpoints(raw) @@ -170,9 +164,7 @@ async def public_model_hub(): ) if llm_router is None: - raise HTTPException( - status_code=400, detail=CommonProxyErrors.no_llm_router.value - ) + raise HTTPException(status_code=400, detail=CommonProxyErrors.no_llm_router.value) model_groups: List[ModelGroupInfoProxy] = [] if litellm.public_model_groups is not None: @@ -267,9 +259,7 @@ async def public_skill_hub(): try: prisma_client = await _get_prisma_client() - plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( - where={"enabled": True} - ) + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many(where={"enabled": True}) items = [] for plugin in plugins: raw = plugin.manifest_json or {} @@ -393,9 +383,7 @@ async def get_litellm_blog_posts(): try: posts_data = get_blog_posts(url=litellm.blog_posts_url) except Exception as e: - verbose_logger.warning( - "LiteLLM: get_litellm_blog_posts endpoint fallback triggered: %s", str(e) - ) + verbose_logger.warning("LiteLLM: get_litellm_blog_posts endpoint fallback triggered: %s", str(e)) posts_data = GetBlogPosts.load_local_blog_posts() posts = [BlogPost(**p) for p in posts_data[:5]] @@ -462,9 +450,7 @@ async def get_agent_fields() -> List[AgentCreateInfo]: field_copy["include_in_litellm_params"] = True inherited_fields.append(field_copy) # Append provider credential fields after agent's own fields - agent["credential_fields"] = ( - agent.get("credential_fields", []) + inherited_fields - ) + agent["credential_fields"] = agent.get("credential_fields", []) + inherited_fields # Remove the inherit field from response (not needed by frontend) agent.pop("inherit_credentials_from_provider", None) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 7ff54ac4c5a..f7f6adaa8a2 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -35,9 +35,7 @@ router = APIRouter() def _raise_vector_store_scan_depth_exceeded() -> None: raise HTTPException( status_code=400, - detail={ - "error": f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while scanning vector_store_id values" - }, + detail={"error": f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while scanning vector_store_id values"}, ) @@ -73,9 +71,7 @@ def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]: if not isinstance(value, str) or not value: raise HTTPException( status_code=400, - detail={ - "error": "vector_store_id must be a non-empty string" - }, + detail={"error": "vector_store_id must be a non-empty string"}, ) vector_store_ids.add(value) continue @@ -193,15 +189,11 @@ async def _save_vector_store_to_db_from_rag_ingest( elif hasattr(response, "vector_store_id"): vector_store_id = response.vector_store_id else: - verbose_proxy_logger.warning( - f"Unable to extract vector_store_id from response type: {type(response)}" - ) + verbose_proxy_logger.warning(f"Unable to extract vector_store_id from response type: {type(response)}") return if vector_store_id is None or not isinstance(vector_store_id, str): - verbose_proxy_logger.warning( - "Vector store ID is None or not a string, skipping database save" - ) + verbose_proxy_logger.warning("Vector store ID is None or not a string, skipping database save") return vector_store_config = ingest_options.get("vector_store", {}) @@ -210,9 +202,7 @@ async def _save_vector_store_to_db_from_rag_ingest( # Extract litellm_vector_store_params for custom name and description litellm_vector_store_params = ingest_options.get("litellm_vector_store_params", {}) custom_vector_store_name = litellm_vector_store_params.get("vector_store_name") - custom_vector_store_description = litellm_vector_store_params.get( - "vector_store_description" - ) + custom_vector_store_description = litellm_vector_store_params.get("vector_store_description") # Extract provider-specific params from vector_store_config to save as litellm_params # This ensures params like aws_region_name, embedding_model, etc. are available for search @@ -231,26 +221,20 @@ async def _save_vector_store_to_db_from_rag_ingest( try: # Check if vector store already exists in database - existing_vector_store = await ManagedVectorStoresRepository( - prisma_client - ).table.find_unique(where={"vector_store_id": vector_store_id}) + existing_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + where={"vector_store_id": vector_store_id} + ) # Only create if it doesn't exist if existing_vector_store is None: - verbose_proxy_logger.info( - f"Saving newly created vector store {vector_store_id} to database" - ) + verbose_proxy_logger.info(f"Saving newly created vector store {vector_store_id} to database") # Initialize metadata with first file initial_metadata = {"ingested_files": [file_entry]} # Use custom name if provided, otherwise default - vector_store_name = ( - custom_vector_store_name or f"RAG Vector Store - {vector_store_id[:8]}" - ) - vector_store_description = ( - custom_vector_store_description or "Created via RAG ingest endpoint" - ) + vector_store_name = custom_vector_store_name or f"RAG Vector Store - {vector_store_id[:8]}" + vector_store_description = custom_vector_store_description or "Created via RAG ingest endpoint" await create_vector_store_in_db( vector_store_id=vector_store_id, @@ -259,20 +243,14 @@ async def _save_vector_store_to_db_from_rag_ingest( vector_store_name=vector_store_name, vector_store_description=vector_store_description, vector_store_metadata=initial_metadata, - litellm_params=( - provider_specific_params if provider_specific_params else None - ), + litellm_params=(provider_specific_params if provider_specific_params else None), team_id=user_api_key_dict.team_id, user_id=user_api_key_dict.user_id, ) - verbose_proxy_logger.info( - f"Vector store {vector_store_id} saved to database successfully" - ) + verbose_proxy_logger.info(f"Vector store {vector_store_id} saved to database successfully") else: - verbose_proxy_logger.info( - f"Vector store {vector_store_id} already exists, appending file to metadata" - ) + verbose_proxy_logger.info(f"Vector store {vector_store_id} already exists, appending file to metadata") # Update existing vector store with new file existing_metadata = existing_vector_store.vector_store_metadata or {} @@ -298,16 +276,12 @@ async def _save_vector_store_to_db_from_rag_ingest( ) except Exception as db_error: # Log the error but don't fail the request since ingestion succeeded - verbose_proxy_logger.exception( - f"Failed to save vector store {vector_store_id} to database: {db_error}" - ) + verbose_proxy_logger.exception(f"Failed to save vector store {vector_store_id} to database: {db_error}") async def parse_rag_ingest_request( request: Request, -) -> Tuple[ - Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str] -]: +) -> Tuple[Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str]]: """ Parse RAG ingest request. @@ -378,9 +352,7 @@ async def parse_rag_ingest_request( if "vector_store" not in ingest_options: raise HTTPException( status_code=400, - detail={ - "error": "ingest_options must contain 'vector_store' configuration" - }, + detail={"error": "ingest_options must contain 'vector_store' configuration"}, ) # Credential fields must come from server configuration, not user requests. @@ -481,16 +453,12 @@ async def rag_ingest( try: # Parse request - ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request( - request - ) + ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(request) # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only - if ( - user_api_key_dict.user_role - == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - and not ingest_options.get("vector_store", {}).get("vector_store_id") - ): + if user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value and not ingest_options.get( + "vector_store", {} + ).get("vector_store_id"): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail={ @@ -692,9 +660,7 @@ async def rag_query( proxy_config=proxy_config, ) - verbose_proxy_logger.debug( - f"RAG Query - model: {model}, retrieval_config: {retrieval_config}" - ) + verbose_proxy_logger.debug(f"RAG Query - model: {model}, retrieval_config: {retrieval_config}") # Call query response = await litellm.aquery( diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index 409ce6f50f2..5a443deba83 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -118,12 +118,8 @@ async def _prepare_client_secret_session( llm_model_list: Optional[list], llm_router: Any, ) -> tuple[str, Optional[dict], str]: - session_type = _coerce_realtime_session_type( - req.session.type if req.session else None - ) - session_data: Optional[dict] = ( - req.session.model_dump(exclude_none=True) if req.session else None - ) + session_type = _coerce_realtime_session_type(req.session.type if req.session else None) + session_data: Optional[dict] = req.session.model_dump(exclude_none=True) if req.session else None if session_data is not None: session_data["type"] = session_type @@ -138,9 +134,7 @@ async def _prepare_client_secret_session( ) return model, session_data, session_type - transcription_model_candidates = _transcription_model_candidates_from_session( - session_data or {} - ) + transcription_model_candidates = _transcription_model_candidates_from_session(session_data or {}) if not transcription_model_candidates: _append_model_candidate(transcription_model_candidates, session_model) _append_model_candidate(transcription_model_candidates, req.model) @@ -282,9 +276,7 @@ async def create_realtime_client_secret( call_type="acreate_realtime_client_secret", ) - verbose_proxy_logger.debug( - "WebRTC: /v1/realtime/client_secrets (model=%s)", model - ) + verbose_proxy_logger.debug("WebRTC: /v1/realtime/client_secrets (model=%s)", model) llm_call = await route_request( data=data, @@ -422,16 +414,10 @@ async def proxy_realtime_calls( ) openai_ephemeral_key = decoded_payload.get("ephemeral_key", "") - model = ( - decoded_payload.get("model_id") - or request.query_params.get("model") - or _DEFAULT_REALTIME_MODEL - ) + model = decoded_payload.get("model_id") or request.query_params.get("model") or _DEFAULT_REALTIME_MODEL user_id = decoded_payload.get("user_id") or None team_id = decoded_payload.get("team_id") or None - session_type = _coerce_realtime_session_type( - decoded_payload.get("session_type") - ) + session_type = _coerce_realtime_session_type(decoded_payload.get("session_type")) else: # Backward compatibility: older tokens contained only encrypted upstream key. openai_ephemeral_key = decrypted_token_value @@ -594,9 +580,7 @@ async def create_realtime_transcription_session( call_type="acreate_realtime_transcription_session", ) - verbose_proxy_logger.debug( - "Realtime: /v1/realtime/transcription_sessions (model=%s)", model - ) + verbose_proxy_logger.debug("Realtime: /v1/realtime/transcription_sessions (model=%s)", model) llm_call = await route_request( data=data, diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index ba9046b3c20..a1b8d2a4821 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -63,9 +63,7 @@ async def rerank( ) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="rerank" - ) + data = await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=data, call_type="rerank") ## ROUTE TO CORRECT ENDPOINT ## llm_call = await route_request( @@ -78,9 +76,7 @@ async def rerank( ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) ### RESPONSE HEADERS ### @@ -107,9 +103,7 @@ async def rerank( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.rerank(): Exception occured - {}".format(str(e)) - ) + verbose_proxy_logger.error("litellm.proxy.proxy_server.rerank(): Exception occured - {}".format(str(e))) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 023f903194b..05c36406f36 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -115,9 +115,7 @@ async def responses_api( ResponsePollingHandler, ) - verbose_proxy_logger.info( - f"Starting background response with polling for model={data.get('model')}" - ) + verbose_proxy_logger.info(f"Starting background response with polling for model={data.get('model')}") # Run pre-call checks (rate limits, guardrails, budget) BEFORE creating # polling ID. This ensures rate-limited requests get a synchronous 429 @@ -236,9 +234,7 @@ async def responses_api( verbose_proxy_logger.warning( f"No model_id found in response hidden params for response {response.id}, skipping managed object storage" ) - raise Exception( - "No model_id found in response hidden params" - ) + raise Exception("No model_id found in response hidden params") # Store in managed objects table await managed_files_obj.store_unified_object_id( unified_object_id=response.id, @@ -364,12 +360,10 @@ async def cursor_chat_completions( if isinstance(response, BaseResponsesAPIStreamingIterator): # Transform Responses API iterator to chat completion iterator # Cast to AsyncIterator[str] since BaseResponsesAPIStreamingIterator implements __aiter__/__anext__ - completion_stream = ( - responses_api_bridge.transformation_handler.get_model_response_iterator( - streaming_response=cast(AsyncIterator[str], response), - sync_stream=False, - json_mode=False, - ) + completion_stream = responses_api_bridge.transformation_handler.get_model_response_iterator( + streaming_response=cast(AsyncIterator[str], response), + sync_stream=False, + json_mode=False, ) # Wrap in CustomStreamWrapper to get the async generator logging_obj = request_data.get("litellm_logging_obj") @@ -415,20 +409,18 @@ async def cursor_chat_completions( # Transform non-streaming Responses API response to chat completions format if isinstance(response, ResponsesAPIResponse): logging_obj = processor.data.get("litellm_logging_obj") - transformed_response = ( - responses_api_bridge.transformation_handler.transform_response( - model=processor.data.get("model", ""), - raw_response=response, - model_response=ModelResponse(), - logging_obj=cast(Any, logging_obj), - request_data=processor.data, - messages=processor.data.get("input", []), - optional_params={}, - litellm_params={}, - encoding=None, - api_key=None, - json_mode=None, - ) + transformed_response = responses_api_bridge.transformation_handler.transform_response( + model=processor.data.get("model", ""), + raw_response=response, + model_response=ModelResponse(), + logging_obj=cast(Any, logging_obj), + request_data=processor.data, + messages=processor.data.get("input", []), + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=None, ) return transformed_response @@ -620,9 +612,7 @@ async def delete_response( state = await polling_handler.get_state(response_id) if not state: - raise HTTPException( - status_code=404, detail=f"Polling response {response_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Polling response {response_id} not found") # Delete from cache success = await polling_handler.delete_polling(response_id) @@ -630,9 +620,7 @@ async def delete_response( if success: return DeleteResponseResult(id=response_id, object="response", deleted=True) else: - raise HTTPException( - status_code=500, detail="Failed to delete polling response" - ) + raise HTTPException(status_code=500, detail="Failed to delete polling response") # Normal provider response flow data = await _read_request_body(request=request) @@ -885,9 +873,7 @@ async def cancel_response( state = await polling_handler.get_state(response_id) if not state: - raise HTTPException( - status_code=404, detail=f"Polling response {response_id} not found" - ) + raise HTTPException(status_code=404, detail=f"Polling response {response_id} not found") # Cancel the polling response (sets status to "cancelled") success = await polling_handler.cancel_polling(response_id) @@ -899,9 +885,7 @@ async def cancel_response( # Return the whole state directly (now with status="cancelled") return updated_state else: - raise HTTPException( - status_code=500, detail="Failed to cancel polling response" - ) + raise HTTPException(status_code=500, detail="Failed to cancel polling response") # Normal provider response flow data = await _read_request_body(request=request) @@ -950,9 +934,7 @@ async def _read_ws_model_from_first_frame( except WebSocketDisconnect: return None except Exception: - verbose_proxy_logger.exception( - "Responses WebSocket error reading first message" - ) + verbose_proxy_logger.exception("Responses WebSocket error reading first message") await websocket.close(code=1011, reason="Internal server error") return None @@ -973,10 +955,7 @@ async def _read_ws_model_from_first_frame( await websocket.close(code=1008, reason="Invalid JSON in first message") return None - if ( - not isinstance(first_event, dict) - or first_event.get("type") != "response.create" - ): + if not isinstance(first_event, dict) or first_event.get("type") != "response.create": await websocket.send_text( json.dumps( { @@ -1019,9 +998,7 @@ def _extract_model_from_first_ws_event(first_event: Any) -> Optional[str]: if not isinstance(first_event, dict): return None nested = first_event.get("response") - return ( - nested.get("model") if isinstance(nested, dict) else None - ) or first_event.get("model") + return (nested.get("model") if isinstance(nested, dict) else None) or first_event.get("model") async def _enforce_responses_ws_first_frame_model_auth( @@ -1049,9 +1026,7 @@ async def _enforce_responses_ws_first_frame_model_auth( or general_settings.get("enable_oauth2_proxy_auth", False) ): return - if user_custom_auth is not None and not general_settings.get( - "custom_auth_run_common_checks", False - ): + if user_custom_auth is not None and not general_settings.get("custom_auth_run_common_checks", False): return await _enforce_key_and_fallback_model_access( valid_token=user_api_key_dict, @@ -1073,9 +1048,7 @@ async def _enforce_responses_ws_first_frame_model_auth( @router.websocket("/responses") async def responses_websocket_endpoint( websocket: WebSocket, - model: Optional[str] = fastapi.Query( - None, description="The model to use for the responses WebSocket session." - ), + model: Optional[str] = fastapi.Query(None, description="The model to use for the responses WebSocket session."), user_api_key_dict=Depends(user_api_key_auth_websocket), ): """ @@ -1107,9 +1080,7 @@ async def responses_websocket_endpoint( # Accept the WebSocket handshake. Key was already validated by the Depends # above; we can safely accept regardless of whether ?model= was supplied. requested_protocols = [ - p.strip() - for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") - if p.strip() + p.strip() for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") if p.strip() ] accept_kwargs: dict = {} if requested_protocols: diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index aedfb26326a..67d8c4021d8 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -133,9 +133,7 @@ async def background_streaming_task( nonlocal state_dirty, last_update_time current_time = asyncio.get_event_loop().time() - if state_dirty and ( - force or (current_time - last_update_time) >= UPDATE_INTERVAL - ): + if state_dirty and (force or (current_time - last_update_time) >= UPDATE_INTERVAL): # Convert output_items dict to list for update output_list = list(output_items.values()) await polling_handler.update_state( @@ -207,12 +205,8 @@ async def background_streaming_task( content_list = output_items[item_id]["content"] if content_index < len(content_list): # Update existing content part with accumulated text - if isinstance( - content_list[content_index], dict - ): - content_list[content_index]["text"] = ( - accumulated_text[key] - ) + if isinstance(content_list[content_index], dict): + content_list[content_index]["text"] = accumulated_text[key] state_dirty = True elif event_type == "response.content_part.done": @@ -263,10 +257,7 @@ async def background_streaming_task( ) # Extract error for failed and incomplete responses - if ( - event_type == "response.failed" - or event_type == "response.incomplete" - ): + if event_type == "response.failed" or event_type == "response.incomplete": terminal_error = response_data.get("error") # Core response fields @@ -280,22 +271,14 @@ async def background_streaming_task( instructions_data = response_data.get("instructions") temperature_data = response_data.get("temperature") top_p_data = response_data.get("top_p") - max_output_tokens_data = response_data.get( - "max_output_tokens" - ) - previous_response_id_data = response_data.get( - "previous_response_id" - ) + max_output_tokens_data = response_data.get("max_output_tokens") + previous_response_id_data = response_data.get("previous_response_id") text_data = response_data.get("text") truncation_data = response_data.get("truncation") - parallel_tool_calls_data = response_data.get( - "parallel_tool_calls" - ) + parallel_tool_calls_data = response_data.get("parallel_tool_calls") user_data = response_data.get("user") store_data = response_data.get("store") - incomplete_details_data = response_data.get( - "incomplete_details" - ) + incomplete_details_data = response_data.get("incomplete_details") # Also update output from final response if available if "output" in response_data: @@ -310,9 +293,7 @@ async def background_streaming_task( await flush_state_if_needed() except json.JSONDecodeError as e: - verbose_proxy_logger.warning( - f"Failed to parse streaming chunk: {e}" - ) + verbose_proxy_logger.warning(f"Failed to parse streaming chunk: {e}") pass # Final flush to ensure all accumulated state is saved @@ -348,9 +329,7 @@ async def background_streaming_task( ) except Exception as e: - verbose_proxy_logger.error( - f"Error in background streaming task for {polling_id}: {str(e)}" - ) + verbose_proxy_logger.error(f"Error in background streaming task for {polling_id}: {str(e)}") import traceback verbose_proxy_logger.error(traceback.format_exc()) diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 739df3ce673..44f3cfb32e4 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -77,9 +77,7 @@ class ResponsePollingHandler: value=response.model_dump_json(), # Pydantic v2 method ttl=self.ttl, ) - verbose_proxy_logger.debug( - f"Created initial polling state for {polling_id} with TTL={self.ttl}s" - ) + verbose_proxy_logger.debug(f"Created initial polling state for {polling_id} with TTL={self.ttl}s") return response @@ -143,9 +141,7 @@ class ResponsePollingHandler: # Get current state cached_state = await self.redis_cache.async_get_cache(cache_key) if not cached_state: - verbose_proxy_logger.warning( - f"No cached state found for polling_id: {polling_id}" - ) + verbose_proxy_logger.warning(f"No cached state found for polling_id: {polling_id}") return # Parse existing ResponsesAPIResponse from cache @@ -258,9 +254,7 @@ def should_use_polling_for_request( redis_cache, # RedisCache or None model: str, llm_router, # Router instance or None - native_background_mode: Optional[ - List[str] - ] = None, # List of models that should use native background mode + native_background_mode: Optional[List[str]] = None, # List of models that should use native background mode ) -> bool: """ Determine if polling via cache should be used for a request. @@ -283,9 +277,7 @@ def should_use_polling_for_request( # Check if model is in native_background_mode list - these use native provider background mode if native_background_mode and model in native_background_mode: - verbose_proxy_logger.debug( - f"Model {model} is in native_background_mode list, skipping polling via cache" - ) + verbose_proxy_logger.debug(f"Model {model} is in native_background_mode list, skipping polling via cache") return False # "all" enables polling for all providers @@ -319,13 +311,9 @@ def should_use_polling_for_request( # If ANY deployment's provider matches, enable polling if dep_provider and dep_provider in polling_via_cache_enabled: - verbose_proxy_logger.debug( - f"Polling enabled for model={model}, provider={dep_provider}" - ) + verbose_proxy_logger.debug(f"Polling enabled for model={model}, provider={dep_provider}") return True except Exception as e: - verbose_proxy_logger.debug( - f"Could not resolve provider for model {model}: {e}" - ) + verbose_proxy_logger.debug(f"Could not resolve provider for model {model}: {e}") return False diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index bbd8b75fdd6..63c6baf9ebc 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -47,16 +47,12 @@ def _is_a2a_agent_model(model_name: Any) -> bool: return isinstance(model_name, str) and model_name.startswith("a2a/") -def _raise_if_model_fully_blocked( - llm_router: LitellmRouter, model_name: Any, team_id: Optional[str] -) -> None: +def _raise_if_model_fully_blocked(llm_router: LitellmRouter, model_name: Any, team_id: Optional[str]) -> None: if not isinstance(model_name, str) or not model_name: return if not isinstance(llm_router, litellm.Router): return - deployments = ( - llm_router.get_model_list(model_name=model_name, team_id=team_id) or [] - ) + deployments = llm_router.get_model_list(model_name=model_name, team_id=team_id) or [] if llm_router._are_all_deployments_blocked(deployments): raise litellm.PermissionDeniedError( message="Model is blocked", @@ -64,9 +60,7 @@ def _raise_if_model_fully_blocked( llm_provider="", response=httpx.Response( status_code=403, - request=httpx.Request( - method="POST", url="https://github.com/BerriAI/litellm" - ), + request=httpx.Request(method="POST", url="https://github.com/BerriAI/litellm"), ), ) @@ -154,11 +148,7 @@ def get_team_id_from_data(data: dict) -> Optional[str]: """ Get the team id from the data's metadata or litellm_metadata params. """ - if ( - "metadata" in data - and data["metadata"] is not None - and "user_api_key_team_id" in data["metadata"] - ): + if "metadata" in data and data["metadata"] is not None and "user_api_key_team_id" in data["metadata"]: return data["metadata"].get("user_api_key_team_id") elif ( "litellm_metadata" in data @@ -205,9 +195,7 @@ async def add_shared_session_to_data(data: dict) -> None: if session is not None and not session.closed: data["shared_session"] = session - verbose_proxy_logger.info( - f"SESSION REUSE: Attached shared aiohttp session to request (ID: {id(session)})" - ) + verbose_proxy_logger.info(f"SESSION REUSE: Attached shared aiohttp session to request (ID: {id(session)})") elif session is not None and session.closed: # Session was created at startup but has since closed — recreate it # Use lock to prevent concurrent recreation (avoids session/connector leak) @@ -230,13 +218,9 @@ async def add_shared_session_to_data(data: dict) -> None: "SESSION REUSE: Shared aiohttp session is None after re-check, recreating..." ) try: - new_session = ( - await proxy_server._initialize_shared_aiohttp_session() - ) + new_session = await proxy_server._initialize_shared_aiohttp_session() except Exception: - verbose_proxy_logger.exception( - "SESSION REUSE: Exception during shared session recreation" - ) + verbose_proxy_logger.exception("SESSION REUSE: Exception during shared session recreation") new_session = None if new_session is not None: proxy_server.shared_aiohttp_session = new_session @@ -246,9 +230,7 @@ async def add_shared_session_to_data(data: dict) -> None: "SESSION REUSE: Failed to recreate shared session, continuing without session reuse" ) else: - verbose_proxy_logger.info( - "SESSION REUSE: No shared session available for this request" - ) + verbose_proxy_logger.info("SESSION REUSE: No shared session available for this request") except Exception: # Continue without session reuse — this outer handler covers import failures # and other unexpected errors to avoid breaking the request path. @@ -438,9 +420,7 @@ async def route_request( else: return getattr(litellm, f"{route_type}")(**data) elif llm_router is not None: - _raise_if_model_fully_blocked( - llm_router=llm_router, model_name=data.get("model"), team_id=team_id - ) + _raise_if_model_fully_blocked(llm_router=llm_router, model_name=data.get("model"), team_id=team_id) # Evals API: always route to litellm directly (not through router) # But extract model credentials if a model is provided if route_type in [ @@ -464,22 +444,16 @@ async def route_request( if model and llm_router: try: # Try to get deployment credentials for this model - deployment_creds = llm_router.get_deployment_credentials( - model_id=model - ) + deployment_creds = llm_router.get_deployment_credentials(model_id=model) if not deployment_creds: # Try by model group name - deployment = llm_router.get_deployment_by_model_group_name( - model_group_name=model - ) + deployment = llm_router.get_deployment_by_model_group_name(model_group_name=model) if ( deployment and deployment.litellm_params and not llm_router._is_deployment_blocked(deployment) ): - deployment_creds = deployment.litellm_params.model_dump( - exclude_none=True - ) + deployment_creds = deployment.litellm_params.model_dump(exclude_none=True) # If we found credentials, merge them into data (but don't override user-provided values) if deployment_creds: @@ -541,24 +515,15 @@ async def route_request( # These endpoints don't need a model, use custom_llm_provider directly return getattr(litellm, f"{route_type}")(**data) - team_model_name = ( - llm_router.map_team_model(data["model"], team_id) - if team_id is not None - else None - ) + team_model_name = llm_router.map_team_model(data["model"], team_id) if team_id is not None else None if team_model_name is not None: data["model"] = team_model_name return getattr(llm_router, f"{route_type}")(**data) - elif data["model"] in router_model_names or llm_router.has_model_id( - data["model"] - ): + elif data["model"] in router_model_names or llm_router.has_model_id(data["model"]): return getattr(llm_router, f"{route_type}")(**data) - elif ( - llm_router.model_group_alias is not None - and data["model"] in llm_router.model_group_alias - ): + elif llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias: return getattr(llm_router, f"{route_type}")(**data) elif data["model"] not in router_model_names: @@ -566,16 +531,11 @@ async def route_request( # Priority: 1. Exact model_name match, 2. Wildcard match, 3. deployment_names match if llm_router.router_general_settings.pass_through_all_models: return getattr(litellm, f"{route_type}")(**data) - elif ( - llm_router.default_deployment is not None - or len(llm_router.pattern_router.patterns) > 0 - ): + elif llm_router.default_deployment is not None or len(llm_router.pattern_router.patterns) > 0: return getattr(llm_router, f"{route_type}")(**data) elif data["model"] in llm_router.deployment_names: # Only match deployment_names if no wildcard matched - return getattr(llm_router, f"{route_type}")( - **data, specific_deployment=True - ) + return getattr(llm_router, f"{route_type}")(**data, specific_deployment=True) elif route_type in [ "amoderation", "aget_responses", @@ -628,9 +588,7 @@ async def route_request( route_a2a_agent_request, ) - result = await route_a2a_agent_request( - data, route_type, user_api_key_dict=user_api_key_dict - ) + result = await route_a2a_agent_request(data, route_type, user_api_key_dict=user_api_key_dict) if result is not None: return result # Fall through to raise exception below if result is None diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index 15ed858b988..7b07d259f48 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -170,9 +170,7 @@ async def search( team_object=team_object, ) except Exception as e: - verbose_proxy_logger.error( - f"Search tool authorization failed for {search_tool_name_value}: {str(e)}" - ) + verbose_proxy_logger.error(f"Search tool authorization failed for {search_tool_name_value}: {str(e)}") raise if llm_router is not None and hasattr(llm_router, "search_tools"): @@ -183,16 +181,12 @@ async def search( ) matching_tools = [ - tool - for tool in llm_router.search_tools - if tool.get("search_tool_name") == search_tool_name_value + tool for tool in llm_router.search_tools if tool.get("search_tool_name") == search_tool_name_value ] if matching_tools: search_tool = matching_tools[0] - search_provider = search_tool.get("litellm_params", {}).get( - "search_provider" - ) + search_provider = search_tool.get("litellm_params", {}).get("search_provider") if search_provider: data["custom_llm_provider"] = search_provider @@ -293,9 +287,7 @@ async def list_search_tools( for tool in llm_router.search_tools: tool_info = { "search_tool_name": tool.get("search_tool_name"), - "search_provider": tool.get("litellm_params", {}).get( - "search_provider" - ), + "search_provider": tool.get("litellm_params", {}).get("search_provider"), } # Add description if available diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 5642fcd10c3..848623bead8 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -147,9 +147,7 @@ async def list_search_tools( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - search_tools_from_db = await SEARCH_TOOL_REGISTRY.get_all_search_tools_from_db( - prisma_client=prisma_client - ) + search_tools_from_db = await SEARCH_TOOL_REGISTRY.get_all_search_tools_from_db(prisma_client=prisma_client) db_tool_names = {tool.get("search_tool_name") for tool in search_tools_from_db} @@ -163,9 +161,7 @@ async def list_search_tools( if parsed_tools: config_search_tools = parsed_tools except Exception as e: - verbose_proxy_logger.debug( - f"Could not get config-defined search tools: {e}" - ) + verbose_proxy_logger.debug(f"Could not get config-defined search tools: {e}") for config_search_tool in config_search_tools: tool_name = config_search_tool.get("search_tool_name") @@ -183,9 +179,7 @@ async def list_search_tools( search_tool_id=None, search_tool_name=tool_name, litellm_params=masked_litellm_params_dict, - search_tool_info=( - dict(config_tool_info) if config_tool_info else None - ), + search_tool_info=(dict(config_tool_info) if config_tool_info else None), created_at=None, updated_at=None, is_from_config=True, @@ -193,9 +187,7 @@ async def list_search_tools( ) search_tool_configs = [ - tool - for tool in search_tool_configs - if tool.get("search_tool_name") not in db_tool_names + tool for tool in search_tool_configs if tool.get("search_tool_name") not in db_tool_names ] for db_search_tool in search_tools_from_db: @@ -212,19 +204,13 @@ async def list_search_tools( search_tool_name=db_search_tool.get("search_tool_name", ""), litellm_params=masked_litellm_params_dict, search_tool_info=db_search_tool.get("search_tool_info"), - created_at=_convert_datetime_to_str( - db_search_tool.get("created_at") - ), - updated_at=_convert_datetime_to_str( - db_search_tool.get("updated_at") - ), + created_at=_convert_datetime_to_str(db_search_tool.get("created_at")), + updated_at=_convert_datetime_to_str(db_search_tool.get("updated_at")), is_from_config=False, ) ) - visible_search_tools = await _filter_visible_search_tools( - search_tool_configs, user_api_key_dict - ) + visible_search_tools = await _filter_visible_search_tools(search_tool_configs, user_api_key_dict) return ListSearchToolsResponse(search_tools=visible_search_tools) except Exception as e: @@ -432,8 +418,7 @@ async def delete_search_tool(search_tool_id: str): ) verbose_proxy_logger.debug( - "Successfully deleted search tool from database. " - "Router will be updated by the cron job." + "Successfully deleted search tool from database. Router will be updated by the cron job." ) return result @@ -574,13 +559,9 @@ async def test_search_tool_connection(request: TestSearchToolConnectionRequest): api_base = litellm_params.get("api_base") if not search_provider: - raise HTTPException( - status_code=400, detail="search_provider is required in litellm_params" - ) + raise HTTPException(status_code=400, detail="search_provider is required in litellm_params") - verbose_proxy_logger.debug( - f"Testing connection to search provider: {search_provider}" - ) + verbose_proxy_logger.debug(f"Testing connection to search provider: {search_provider}") # Make a simple test search query with max_results=1 to minimize cost test_query = "test" @@ -593,26 +574,20 @@ async def test_search_tool_connection(request: TestSearchToolConnectionRequest): timeout=10.0, # 10 second timeout for test ) - verbose_proxy_logger.debug( - f"Successfully tested connection to {search_provider} search provider" - ) + verbose_proxy_logger.debug(f"Successfully tested connection to {search_provider} search provider") return { "status": "success", "message": f"Successfully connected to {search_provider} search provider", "test_query": test_query, - "results_count": ( - len(response.results) if response and response.results else 0 - ), + "results_count": (len(response.results) if response and response.results else 0), } except Exception as e: error_message = str(e) error_type = type(e).__name__ - verbose_proxy_logger.exception( - f"Failed to connect to search provider: {error_message}" - ) + verbose_proxy_logger.exception(f"Failed to connect to search provider: {error_message}") # Return error details in a structured format return { @@ -664,9 +639,7 @@ async def get_available_search_providers(): for provider in SearchProviders: try: # Get the config class for this provider - config = ProviderConfigManager.get_provider_search_config( - provider=provider - ) + config = ProviderConfigManager.get_provider_search_config(provider=provider) if config is not None: # Get the UI-friendly name from the config class @@ -679,9 +652,7 @@ async def get_available_search_providers(): } ) except Exception as e: - verbose_proxy_logger.debug( - f"Could not get config for search provider {provider.value}: {e}" - ) + verbose_proxy_logger.debug(f"Could not get config for search provider {provider.value}: {e}") continue return {"providers": available_providers} diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index 2ec2533211b..5934d686b5d 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -44,9 +44,7 @@ class SearchToolRegistry: ########### DB management helpers for search tools ######## ########################################################### - async def add_search_tool_to_db( - self, search_tool: SearchTool, prisma_client: PrismaClient - ): + async def add_search_tool_to_db(self, search_tool: SearchTool, prisma_client: PrismaClient): """ Add a search tool to the database. @@ -59,15 +57,11 @@ class SearchToolRegistry: """ try: search_tool_name = search_tool.get("search_tool_name") - litellm_params: str = safe_dumps( - dict(search_tool.get("litellm_params", {})) - ) + litellm_params: str = safe_dumps(dict(search_tool.get("litellm_params", {}))) search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Create search tool in DB - created_search_tool = await SearchToolsRepository( - prisma_client - ).table.create( + created_search_tool = await SearchToolsRepository(prisma_client).table.create( data={ "search_tool_name": search_tool_name, "litellm_params": litellm_params, @@ -88,9 +82,7 @@ class SearchToolRegistry: verbose_proxy_logger.exception(f"Error adding search tool to DB: {str(e)}") raise Exception(f"Error adding search tool to DB: {str(e)}") - async def delete_search_tool_from_db( - self, search_tool_id: str, prisma_client: PrismaClient - ): + async def delete_search_tool_from_db(self, search_tool_id: str, prisma_client: PrismaClient): """ Delete a search tool from the database. @@ -103,31 +95,25 @@ class SearchToolRegistry: """ try: # Get search tool before deletion for response - existing_tool = await SearchToolsRepository( - prisma_client - ).table.find_unique(where={"search_tool_id": search_tool_id}) + existing_tool = await SearchToolsRepository(prisma_client).table.find_unique( + where={"search_tool_id": search_tool_id} + ) if not existing_tool: raise Exception(f"Search tool with ID {search_tool_id} not found") # Delete from DB - await SearchToolsRepository(prisma_client).table.delete( - where={"search_tool_id": search_tool_id} - ) + await SearchToolsRepository(prisma_client).table.delete(where={"search_tool_id": search_tool_id}) return { "message": f"Search tool {search_tool_id} deleted successfully", "search_tool_name": existing_tool.search_tool_name, } except Exception as e: - verbose_proxy_logger.exception( - f"Error deleting search tool from DB: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error deleting search tool from DB: {str(e)}") raise Exception(f"Error deleting search tool from DB: {str(e)}") - async def update_search_tool_in_db( - self, search_tool_id: str, search_tool: SearchTool, prisma_client: PrismaClient - ): + async def update_search_tool_in_db(self, search_tool_id: str, search_tool: SearchTool, prisma_client: PrismaClient): """ Update a search tool in the database. @@ -141,15 +127,11 @@ class SearchToolRegistry: """ try: search_tool_name = search_tool.get("search_tool_name") - litellm_params: str = safe_dumps( - dict(search_tool.get("litellm_params", {})) - ) + litellm_params: str = safe_dumps(dict(search_tool.get("litellm_params", {}))) search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Update in DB - updated_search_tool = await SearchToolsRepository( - prisma_client - ).table.update( + updated_search_tool = await SearchToolsRepository(prisma_client).table.update( where={"search_tool_id": search_tool_id}, data={ "search_tool_name": search_tool_name, @@ -162,9 +144,7 @@ class SearchToolRegistry: # Convert to dict with ISO formatted datetimes return self._convert_prisma_to_dict(updated_search_tool) except Exception as e: - verbose_proxy_logger.exception( - f"Error updating search tool in DB: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error updating search tool in DB: {str(e)}") raise Exception(f"Error updating search tool in DB: {str(e)}") @staticmethod @@ -192,16 +172,12 @@ class SearchToolRegistry: search_tools: List[SearchTool] = [] for search_tool in search_tools_from_db: # Convert Prisma result to dict with ISO formatted datetimes - search_tool_dict = SearchToolRegistry._convert_prisma_to_dict( - search_tool - ) + search_tool_dict = SearchToolRegistry._convert_prisma_to_dict(search_tool) search_tools.append(SearchTool(**search_tool_dict)) # type: ignore return search_tools except Exception as e: - verbose_proxy_logger.exception( - f"Error getting search tools from DB: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error getting search tools from DB: {str(e)}") raise Exception(f"Error getting search tools from DB: {str(e)}") async def get_search_tool_by_id_from_db( @@ -229,9 +205,7 @@ class SearchToolRegistry: search_tool_dict = self._convert_prisma_to_dict(search_tool) return SearchTool(**search_tool_dict) # type: ignore except Exception as e: - verbose_proxy_logger.exception( - f"Error getting search tool from DB: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error getting search tool from DB: {str(e)}") raise Exception(f"Error getting search tool from DB: {str(e)}") async def get_search_tool_by_name_from_db( @@ -259,7 +233,5 @@ class SearchToolRegistry: search_tool_dict = self._convert_prisma_to_dict(search_tool) return SearchTool(**search_tool_dict) # type: ignore except Exception as e: - verbose_proxy_logger.exception( - f"Error getting search tool from DB: {str(e)}" - ) + verbose_proxy_logger.exception(f"Error getting search tool from DB: {str(e)}") raise Exception(f"Error getting search tool from DB: {str(e)}") diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 9cfd636c308..ca6c2e86789 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -50,9 +50,7 @@ def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set: return set() entries = budget_reservation.get("entries") or [] return { - entry["counter_key"] - for entry in entries - if isinstance(entry, dict) and entry.get("counter_key") is not None + entry["counter_key"] for entry in entries if isinstance(entry, dict) and entry.get("counter_key") is not None } @@ -132,9 +130,7 @@ async def reserve_budget_for_request( cached_spend = await _get_current_counter_value(counter=counter) current_spend = cached_spend + reservation_cost if current_spend > counter.max_budget: - remaining_before_reservation = counter.max_budget - ( - current_spend - reservation_cost - ) + remaining_before_reservation = counter.max_budget - (current_spend - reservation_cost) if remaining_before_reservation > 1e-12: await _resize_applied_reservation( entries=applied_entries, @@ -163,9 +159,7 @@ async def reserve_budget_for_request( if not applied_entries: return None - input_cost = estimate_request_input_cost( - request_body=request_body, route=route, llm_router=llm_router - ) + input_cost = estimate_request_input_cost(request_body=request_body, route=route, llm_router=llm_router) return { "reserved_cost": reservation_cost, "entries": applied_entries, @@ -227,9 +221,7 @@ async def release_budget_reservation_on_cancel( incurred_cost = float(budget_reservation.get("input_cost") or 0.0) try: await asyncio.shield( - reconcile_budget_reservation( - budget_reservation=budget_reservation, actual_cost=incurred_cost - ) + reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) ) except (asyncio.CancelledError, Exception): pass @@ -450,20 +442,11 @@ async def _get_team_member_budget_counter( user_object: Optional[LiteLLM_UserTable], user_api_key_cache: DualCache, ) -> Optional[_BudgetCounter]: - if ( - team_object is None - or team_object.team_id is None - or user_object is None - or valid_token.user_id is None - ): + if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None - membership_cache_key = ( - f"team_membership:{valid_token.user_id}:{team_object.team_id}" - ) - cached_team_membership = await user_api_key_cache.async_get_cache( - key=membership_cache_key - ) + membership_cache_key = f"team_membership:{valid_token.user_id}:{team_object.team_id}" + cached_team_membership = await user_api_key_cache.async_get_cache(key=membership_cache_key) team_membership: Optional[LiteLLM_TeamMembership] = None if isinstance(cached_team_membership, LiteLLM_TeamMembership): team_membership = cached_team_membership @@ -484,11 +467,7 @@ async def _get_team_member_budget_counter( if team_member_budget is None or team_member_budget <= 0: return None - team_member_spend = ( - cast(LiteLLM_TeamMembership, team_membership).spend - if team_membership is not None - else 0.0 - ) + team_member_spend = cast(LiteLLM_TeamMembership, team_membership).spend if team_membership is not None else 0.0 return _BudgetCounter( counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", source_cache_key=membership_cache_key, @@ -609,9 +588,7 @@ async def _reserve_counter( counter_key=counter.counter_key, source_cache_key=counter.source_cache_key, ) - elif ( - counter.spend_log_entity_id is not None and counter.window_start is not None - ): + elif counter.spend_log_entity_id is not None and counter.window_start is not None: initialized = await _ensure_window_spend_counter_initialized( counter_key=counter.counter_key, entity_type=counter.entity_type, @@ -723,9 +700,7 @@ async def _set_reserved_entry_actual_cost( # persisted, so the DB floor would discard it. Keep the original # fail-closed behavior (raise -> reserve_budget_for_request releases and # denies) rather than admitting against an inconsistent counter. - raise RuntimeError( - f"Cannot resize budget reservation against inconsistent counter {counter_key}" - ) + raise RuntimeError(f"Cannot resize budget reservation against inconsistent counter {counter_key}") entry["applied_adjustment"] = target_adjustment @@ -760,9 +735,7 @@ async def _release_applied_entries_best_effort( ) except Exception: counter_key = entry.get("counter_key") - verbose_proxy_logger.exception( - "Failed to release partial budget reservation during exception cleanup" - ) + verbose_proxy_logger.exception("Failed to release partial budget reservation during exception cleanup") if counter_key is None: continue try: @@ -1048,18 +1021,14 @@ def _estimate_input_tokens( if "input" in request_body: return _count_text_tokens(model=model, text=request_body.get("input")) if "query" in request_body or "documents" in request_body: - query_tokens = _count_text_tokens( - model=model, text=request_body.get("query") - ) + query_tokens = _count_text_tokens(model=model, text=request_body.get("query")) document_tokens = _count_text_tokens( model=model, text=request_body.get("documents"), ) return query_tokens + document_tokens except Exception: - verbose_proxy_logger.debug( - "Unable to count input tokens for budget reservation", exc_info=True - ) + verbose_proxy_logger.debug("Unable to count input tokens for budget reservation", exc_info=True) max_input_tokens = _to_int(model_info.get("max_input_tokens")) if max_input_tokens is not None: @@ -1093,10 +1062,7 @@ def _estimate_output_tokens( # the reservation up to remaining team headroom and pin the counter # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. - model_ceiling = ( - _to_int(model_info.get("max_output_tokens")) - or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - ) + model_ceiling = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None: requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK return min(requested, model_ceiling) diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index 71f4a8af111..827814a9716 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -98,15 +98,11 @@ async def _get_cloudzero_settings(): # Decrypt the API key encrypted_api_key = settings.get("api_key") if encrypted_api_key: - decrypted_api_key = decrypt_value_helper( - encrypted_api_key, key="cloudzero_api_key", exception_type="error" - ) + decrypted_api_key = decrypt_value_helper(encrypted_api_key, key="cloudzero_api_key", exception_type="error") if decrypted_api_key is None: raise HTTPException( status_code=500, - detail={ - "error": "Failed to decrypt CloudZero API key. Check your salt key configuration." - }, + detail={"error": "Failed to decrypt CloudZero API key. Check your salt key configuration."}, ) settings["api_key"] = decrypted_api_key @@ -214,21 +210,11 @@ async def update_cloudzero_settings( current_settings = await _get_cloudzero_settings() # Update only provided fields - updated_api_key = ( - request.api_key - if request.api_key is not None - else current_settings["api_key"] - ) + updated_api_key = request.api_key if request.api_key is not None else current_settings["api_key"] updated_connection_id = ( - request.connection_id - if request.connection_id is not None - else current_settings["connection_id"] - ) - updated_timezone = ( - request.timezone - if request.timezone is not None - else current_settings["timezone"] + request.connection_id if request.connection_id is not None else current_settings["connection_id"] ) + updated_timezone = request.timezone if request.timezone is not None else current_settings["timezone"] # Store updated settings using the setter method with encryption await _set_cloudzero_settings( @@ -239,9 +225,7 @@ async def update_cloudzero_settings( verbose_proxy_logger.info("CloudZero settings updated successfully") - return CloudZeroInitResponse( - message="CloudZero settings updated successfully", status="success" - ) + return CloudZeroInitResponse(message="CloudZero settings updated successfully", status="success") except HTTPException as e: if e.status_code == 400: @@ -377,9 +361,7 @@ async def init_cloudzero_settings( verbose_proxy_logger.info("CloudZero settings initialized successfully") - return CloudZeroInitResponse( - message="CloudZero settings initialized successfully", status="success" - ) + return CloudZeroInitResponse(message="CloudZero settings initialized successfully", status="success") except Exception as e: verbose_proxy_logger.error(f"Error initializing CloudZero settings: {str(e)}") @@ -440,9 +422,7 @@ async def cloudzero_dry_run_export( ) except Exception as e: - verbose_proxy_logger.error( - f"Error performing CloudZero dry run export: {str(e)}" - ) + verbose_proxy_logger.error(f"Error performing CloudZero dry run export: {str(e)}") raise HTTPException( status_code=500, detail={"error": f"Failed to perform CloudZero dry run export: {str(e)}"}, @@ -561,15 +541,11 @@ async def delete_cloudzero_settings( # Delete only the CloudZero settings entry # This uses a specific where clause to target only the cloudzero_settings row - await ConfigRepository(prisma_client).table.delete( - where={"param_name": "cloudzero_settings"} - ) + await ConfigRepository(prisma_client).table.delete(where={"param_name": "cloudzero_settings"}) verbose_proxy_logger.info("CloudZero settings deleted successfully") - return CloudZeroInitResponse( - message="CloudZero settings deleted successfully", status="success" - ) + return CloudZeroInitResponse(message="CloudZero settings deleted successfully", status="success") except HTTPException as e: raise e diff --git a/litellm/proxy/spend_tracking/cold_storage_handler.py b/litellm/proxy/spend_tracking/cold_storage_handler.py index 01c698c7122..9511e0d8565 100644 --- a/litellm/proxy/spend_tracking/cold_storage_handler.py +++ b/litellm/proxy/spend_tracking/cold_storage_handler.py @@ -35,9 +35,7 @@ class ColdStorageHandler: Returns: Optional[dict]: The proxy server request dict or None if not found """ - custom_logger = ( - self._injected_cold_storage_logger or self._resolve_cold_storage_logger() - ) + custom_logger = self._injected_cold_storage_logger or self._resolve_cold_storage_logger() if custom_logger is None: return None @@ -49,17 +47,13 @@ class ColdStorageHandler: custom_logger_name = self._select_custom_logger_for_cold_storage() if custom_logger_name is None: return None - return ( - litellm.logging_callback_manager.get_active_custom_logger_for_callback_name( - custom_logger_name - ) - ) + return litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(custom_logger_name) def _select_custom_logger_for_cold_storage( self, ) -> _custom_logger_compatible_callbacks_literal | None: - cold_storage_custom_logger: ( - _custom_logger_compatible_callbacks_literal | None - ) = litellm.cold_storage_custom_logger + cold_storage_custom_logger: _custom_logger_compatible_callbacks_literal | None = ( + litellm.cold_storage_custom_logger + ) return cold_storage_custom_logger diff --git a/litellm/proxy/spend_tracking/spend_log_error_logger.py b/litellm/proxy/spend_tracking/spend_log_error_logger.py index cfe647b6600..e2987481c8b 100644 --- a/litellm/proxy/spend_tracking/spend_log_error_logger.py +++ b/litellm/proxy/spend_tracking/spend_log_error_logger.py @@ -78,8 +78,6 @@ def spend_log_error( return if exc is not None: - verbose_proxy_logger.error( - message, *args, exc_info=(type(exc), exc, exc.__traceback__) - ) + verbose_proxy_logger.error(message, *args, exc_info=(type(exc), exc, exc.__traceback__)) else: verbose_proxy_logger.error(message, *args, exc_info=True) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 3374877572b..5624afcfa5e 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -164,14 +164,10 @@ async def spend_user_fn( user_id = caller_user_id if user_id is not None: - user_info = await prisma_client.get_data( - table_name="user", query_type="find_unique", user_id=user_id - ) + user_info = await prisma_client.get_data(table_name="user", query_type="find_unique", user_id=user_id) result = [user_info] else: - user_info = await prisma_client.get_data( - table_name="user", query_type="find_all" - ) + user_info = await prisma_client.get_data(table_name="user", query_type="find_all") result = user_info _strip_password_from_users(result) @@ -237,9 +233,7 @@ async def view_spend_tags( FROM "LiteLLM_SpendLogs" GROUP BY individual_request_tag; """ - response = await get_spend_by_tags( - start_date=start_date, end_date=end_date, prisma_client=prisma_client - ) + response = await get_spend_by_tags(start_date=start_date, end_date=end_date, prisma_client=prisma_client) return response except Exception as e: @@ -283,9 +277,7 @@ async def get_global_activity_internal_user( AND "user" = $3 GROUP BY date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date, end_date, user_id - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date, end_date, user_id) return db_response @@ -338,9 +330,7 @@ async def get_global_activity( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -355,9 +345,7 @@ async def get_global_activity( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY ): - db_response = await get_global_activity_internal_user( - user_api_key_dict, start_date_obj, end_date_obj - ) + db_response = await get_global_activity_internal_user(user_api_key_dict, start_date_obj, end_date_obj) else: sql_query = """ SELECT @@ -369,9 +357,7 @@ async def get_global_activity( AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -430,9 +416,7 @@ async def get_global_activity_model_internal_user( AND "user" = $3 GROUP BY model_group, date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date, end_date, user_id - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date, end_date, user_id) return db_response @@ -508,9 +492,7 @@ async def get_global_activity_model( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -525,9 +507,7 @@ async def get_global_activity_model( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY ): - db_response = await get_global_activity_model_internal_user( - user_api_key_dict, start_date_obj, end_date_obj - ) + db_response = await get_global_activity_model_internal_user(user_api_key_dict, start_date_obj, end_date_obj) else: sql_query = """ SELECT @@ -540,9 +520,7 @@ async def get_global_activity_model( AND "startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') GROUP BY model_group, date_trunc('day', "startTime") """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -661,9 +639,7 @@ async def get_global_activity_exceptions_per_deployment( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -692,9 +668,7 @@ async def get_global_activity_exceptions_per_deployment( ORDER BY date; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj, model_group - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, model_group) if db_response is None: return [] @@ -711,9 +685,7 @@ async def get_global_activity_exceptions_per_deployment( row["date"] = _date_obj.strftime("%b %d") model_ui_data[_model]["daily_data"].append(row) - model_ui_data[_model]["sum_num_rate_limit_exceptions"] += row.get( - "num_rate_limit_exceptions", 0 - ) + model_ui_data[_model]["sum_num_rate_limit_exceptions"] += row.get("num_rate_limit_exceptions", 0) # sort mode ui data by sum_api_requests -> get top 10 models model_ui_data = dict( @@ -732,9 +704,7 @@ async def get_global_activity_exceptions_per_deployment( { "api_base": model, "daily_data": _sort_daily_data, - "sum_num_rate_limit_exceptions": data[ - "sum_num_rate_limit_exceptions" - ], + "sum_num_rate_limit_exceptions": data["sum_num_rate_limit_exceptions"], } ) @@ -794,9 +764,7 @@ async def get_global_activity_exceptions( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import prisma_client @@ -823,9 +791,7 @@ async def get_global_activity_exceptions( ORDER BY date; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj, model_group - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, model_group) if db_response is None: return [] @@ -902,9 +868,7 @@ async def get_global_spend_provider( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import llm_router, prisma_client @@ -921,9 +885,7 @@ async def get_global_spend_provider( ): user_id = user_api_key_dict.user_id if user_id is None: - raise HTTPException( - status_code=400, detail={"error": "No user_id found"} - ) + raise HTTPException(status_code=400, detail={"error": "No user_id found"}) sql_query = """ SELECT @@ -936,9 +898,7 @@ async def get_global_spend_provider( AND "user" = $3 GROUP BY model_id """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj, user_id - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, user_id) else: sql_query = """ SELECT @@ -950,9 +910,7 @@ async def get_global_spend_provider( AND length(model_id) > 0 GROUP BY model_id """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1066,9 +1024,7 @@ async def get_global_spend_report( detail={"error": "Please provide start_date and end_date"}, ) - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) from litellm.proxy.proxy_server import premium_user, prisma_client @@ -1081,13 +1037,9 @@ async def get_global_spend_report( if premium_user is not True: verbose_proxy_logger.debug("accessing /spend/report but not a premium user") - raise ValueError( - "/spend/report endpoint " + CommonProxyErrors.not_premium_user.value - ) + raise ValueError("/spend/report endpoint " + CommonProxyErrors.not_premium_user.value) if api_key is not None: - verbose_proxy_logger.debug( - "Getting /spend for api_key: [set=%s]", api_key is not None - ) + verbose_proxy_logger.debug("Getting /spend for api_key: [set=%s]", api_key is not None) if api_key.startswith("sk-"): api_key = hash_token(token=api_key) sql_query = """ @@ -1126,17 +1078,13 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj, api_key - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, api_key) if db_response is None: return [] return db_response elif internal_user_id is not None: - verbose_proxy_logger.debug( - "Getting /spend for internal_user_id: %s", internal_user_id - ) + verbose_proxy_logger.debug("Getting /spend for internal_user_id: %s", internal_user_id) sql_query = """ WITH SpendByModelApiKey AS ( SELECT @@ -1173,9 +1121,7 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj, internal_user_id - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj, internal_user_id) if db_response is None: return [] @@ -1242,9 +1188,7 @@ async def get_global_spend_report( group_by_day; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1303,9 +1247,7 @@ async def get_global_spend_report( group_by_day; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1346,9 +1288,7 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1512,9 +1452,7 @@ async def _get_spend_report_for_time_range( # text, which forces Postgres to parse `::timestamptz` using the DB # session timezone and drifts the window by the offset even with the # AT TIME ZONE 'UTC' wrap below. - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) try: @@ -1534,9 +1472,7 @@ async def _get_spend_report_for_time_range( ORDER BY total_spend DESC; """ - response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + response = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) # get spend per tag for today sql_query = """ @@ -1550,15 +1486,11 @@ async def _get_spend_report_for_time_range( ORDER BY total_spend DESC; """ - spend_per_tag = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + spend_per_tag = await prisma_client.db.query_raw(sql_query, start_date_obj, end_date_obj) return response, spend_per_tag except Exception as e: - verbose_proxy_logger.error( - "Exception in _get_daily_spend_reports {}".format(str(e)) - ) + verbose_proxy_logger.error("Exception in _get_daily_spend_reports {}".format(str(e))) @router.post( @@ -1651,10 +1583,7 @@ async def calculate_spend(request: SpendCalculateRequest): _model_in_llm_router = None cost_per_token: CostPerToken | None = None if llm_router is not None: - if ( - llm_router.model_group_alias is not None - and request.model in llm_router.model_group_alias - ): + if llm_router.model_group_alias is not None and request.model in llm_router.model_group_alias: # lookup alias in llm_router _model_group_name = llm_router.model_group_alias[request.model] for model in llm_router.model_list: @@ -1681,10 +1610,7 @@ async def calculate_spend(request: SpendCalculateRequest): _litellm_model_name = _litellm_params.get("model") input_cost_per_token = _litellm_params.get("input_cost_per_token") output_cost_per_token = _litellm_params.get("output_cost_per_token") - if ( - input_cost_per_token is not None - or output_cost_per_token is not None - ): + if input_cost_per_token is not None or output_cost_per_token is not None: cost_per_token = CostPerToken( input_cost_per_token=input_cost_per_token, output_cost_per_token=output_cost_per_token, @@ -1774,12 +1700,8 @@ async def ui_view_spend_logs( default=None, description="Time till which to view key spend", ), - page: int = fastapi.Query( - default=1, description="Page number for pagination", ge=1 - ), - page_size: int = fastapi.Query( - default=50, description="Number of items per page", ge=1, le=100 - ), + page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), + page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=100), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), status_filter: str | None = fastapi.Query( default=None, description="Filter logs by status (e.g., success, failure)" @@ -1789,18 +1711,10 @@ async def ui_view_spend_logs( default=None, description="Filter logs by model ID (litellm model deployment id)", ), - model_group: str | None = fastapi.Query( - default=None, description="Filter logs by model group" - ), - key_alias: str | None = fastapi.Query( - default=None, description="Filter logs by key alias" - ), - end_user: str | None = fastapi.Query( - default=None, description="Filter logs by end user" - ), - error_code: str | None = fastapi.Query( - default=None, description="Filter logs by error code (e.g., '404', '500')" - ), + model_group: str | None = fastapi.Query(default=None, description="Filter logs by model group"), + key_alias: str | None = fastapi.Query(default=None, description="Filter logs by key alias"), + end_user: str | None = fastapi.Query(default=None, description="Filter logs by end user"), + error_code: str | None = fastapi.Query(default=None, description="Filter logs by error code (e.g., '404', '500')"), error_message: str | None = fastapi.Query( default=None, description="Filter logs by error message (partial string match)" ), @@ -1882,11 +1796,7 @@ async def ui_view_spend_logs( return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc) except ValueError: continue - expected = ( - "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" - if is_v2 - else "'YYYY-MM-DD HH:MM:SS'" - ) + expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'" raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid date format: {date_str}. Expected: {expected}", @@ -1983,22 +1893,16 @@ async def ui_view_spend_logs( if not can_view_team: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Not authorized to view team spend for team_id={}".format( - team_id - ) - }, + detail={"error": "Not authorized to view team spend for team_id={}".format(team_id)}, ) where_conditions["team_id"] = team_id where_conditions.pop("user", None) else: if _can_user_view_spend_log(user_api_key_dict=user_api_key_dict): try: - permitted_team_ids = ( - await _get_permitted_team_ids_for_spend_logs( - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - ) + permitted_team_ids = await _get_permitted_team_ids_for_spend_logs( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, ) except Exception: permitted_team_ids = [] @@ -2028,14 +1932,10 @@ async def ui_view_spend_logs( # Date range (always present). Wrap the param side with # `AT TIME ZONE 'UTC'` so comparison against the plain `timestamp` # column does not depend on the DB session timezone (see #22529). - sql_conditions.append( - f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')" - ) + sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')") sql_params.append(start_date_obj) p += 1 - sql_conditions.append( - f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')" - ) + sql_conditions.append(f"\"startTime\" <= (${p}::timestamptz AT TIME ZONE 'UTC')") sql_params.append(end_date_obj) p += 1 @@ -2089,15 +1989,11 @@ async def ui_view_spend_logs( sql_params.append(f"%{key_alias}%") p += 1 if error_code is not None: - sql_conditions.append( - f"metadata->'error_information'->>'error_code' = ${p}" - ) + sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}") sql_params.append(error_code) p += 1 if error_message is not None: - sql_conditions.append( - f"metadata->'error_information'->>'error_message' LIKE ${p}" - ) + sql_conditions.append(f"metadata->'error_information'->>'error_message' LIKE ${p}") sql_params.append(f"%{error_message}%") p += 1 @@ -2326,13 +2222,9 @@ async def ui_view_request_response_for_request_id( start_date_obj: datetime | None = None end_date_obj: datetime | None = None if start_date is not None: - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc) if end_date is not None: - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S").replace( - tzinfo=timezone.utc - ) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc) for custom_logger in custom_loggers: payload = await custom_logger.get_request_response_payload( @@ -2360,9 +2252,7 @@ async def ui_view_request_response_for_request_id( """ db_result = await prisma_client.db.query_raw(sql_query, request_id) if db_result and len(db_result) > 0: - resolved = await _resolve_request_response_payload( - db_result[0], cold_storage_handler=ColdStorageHandler() - ) + resolved = await _resolve_request_response_payload(db_result[0], cold_storage_handler=ColdStorageHandler()) return resolved._asdict() return None @@ -2465,12 +2355,8 @@ async def view_spend_logs( and isinstance(end_date, str) ): # Convert the date strings to datetime objects - start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) - end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace( - tzinfo=timezone.utc - ) + start_date_obj = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) + end_date_obj = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc) # Convert to ISO format strings for Prisma start_date_iso = start_date_obj.isoformat() @@ -2514,34 +2400,24 @@ async def view_spend_logs( }, ) - if ( - isinstance(response, list) - and len(response) > 0 - and isinstance(response[0], dict) - ): + if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict): result: dict = {} for record in response: - dt_object = datetime.strptime( - str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ" - ) # type: ignore + dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ") # type: ignore date = dt_object.date() if date not in result: result[date] = {"users": {}, "models": {}} api_key = record["api_key"] # type: ignore user_id = record["user"] # type: ignore model = record["model"] # type: ignore - result[date]["spend"] = result[date].get("spend", 0) + record.get( - "_sum", {} - ).get("spend", 0) - result[date][api_key] = result[date].get(api_key, 0) + record.get( - "_sum", {} - ).get("spend", 0) - result[date]["users"][user_id] = result[date]["users"].get( - user_id, 0 - ) + record.get("_sum", {}).get("spend", 0) - result[date]["models"][model] = result[date]["models"].get( - model, 0 - ) + record.get("_sum", {}).get("spend", 0) + result[date]["spend"] = result[date].get("spend", 0) + record.get("_sum", {}).get("spend", 0) + result[date][api_key] = result[date].get(api_key, 0) + record.get("_sum", {}).get("spend", 0) + result[date]["users"][user_id] = result[date]["users"].get(user_id, 0) + record.get("_sum", {}).get( + "spend", 0 + ) + result[date]["models"][model] = result[date]["models"].get(model, 0) + record.get("_sum", {}).get( + "spend", 0 + ) return_list = [] final_date = None for k, v in sorted(result.items()): @@ -2581,9 +2457,7 @@ async def view_spend_logs( scoped_filter["user"] = user_id if not scoped_filter: - spend_logs = await prisma_client.get_data( - table_name="spend", query_type="find_all" - ) + spend_logs = await prisma_client.get_data(table_name="spend", query_type="find_all") return spend_logs data = await SpendLogsRepository(prisma_client).table.find_many( @@ -2638,9 +2512,7 @@ async def global_spend_reset(): code=status.HTTP_401_UNAUTHORIZED, ) - await VerificationTokenRepository(prisma_client).table.update_many( - data={"spend": 0.0}, where={} - ) + await VerificationTokenRepository(prisma_client).table.update_many(data={"spend": 0.0}, where={}) await TeamRepository(prisma_client).table.update_many(data={"spend": 0.0}, where={}) return { @@ -2721,9 +2593,7 @@ async def global_spend_refresh(): } except Exception as e: - verbose_proxy_logger.exception( - "Failed to refresh materialized view - {}".format(str(e)) - ) + verbose_proxy_logger.exception("Failed to refresh materialized view - {}".format(str(e))) return { "message": "Failed to refresh materialized view", "status": "failure", @@ -2809,9 +2679,7 @@ async def global_spend_logs( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY ): - response = await global_spend_for_internal_user( - api_key=api_key, user_api_key_dict=user_api_key_dict - ) + response = await global_spend_for_internal_user(api_key=api_key, user_api_key_dict=user_api_key_dict) return response @@ -2907,9 +2775,7 @@ async def global_spend(): ) -async def global_spend_key_internal_user( - user_api_key_dict: UserAPIKeyAuth, limit: int = 10 -): +async def global_spend_key_internal_user(user_api_key_dict: UserAPIKeyAuth, limit: int = 10): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -2977,9 +2843,7 @@ async def global_spend_keys( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY ): - response = await global_spend_key_internal_user( - user_api_key_dict=user_api_key_dict - ) + response = await global_spend_key_internal_user(user_api_key_dict=user_api_key_dict) return response if prisma_client is None: @@ -2996,9 +2860,7 @@ async def global_spend_keys( sql_query = """SELECT * FROM "Last30dKeysBySpend" LIMIT $1 ;""" response = await prisma_client.db.query_raw(sql_query, limit) except ValueError as e: - raise HTTPException( - status_code=422, detail={"error": f"Invalid limit: {limit}, error: {e}"} - ) from e + raise HTTPException(status_code=422, detail={"error": f"Invalid limit: {limit}, error: {e}"}) from e return response @@ -3069,18 +2931,14 @@ async def global_spend_per_team(): total_spend_per_team_ui = [] # order the elements in total_spend_per_team by spend - total_spend_per_team = dict( - sorted(total_spend_per_team.items(), key=lambda item: item[1], reverse=True) - ) + total_spend_per_team = dict(sorted(total_spend_per_team.items(), key=lambda item: item[1], reverse=True)) for team_id in total_spend_per_team: # only add first 10 elements to total_spend_per_team_ui if len(total_spend_per_team_ui) >= 10: break if team_id is None: team_id = "Unassigned" - total_spend_per_team_ui.append( - {"team_id": team_id, "total_spend": total_spend_per_team[team_id]} - ) + total_spend_per_team_ui.append({"team_id": team_id, "total_spend": total_spend_per_team[team_id]}) # sort spend_by_date by it's key (which is a date) @@ -3174,16 +3032,12 @@ GROUP BY end_user ORDER BY total_spend DESC LIMIT 100 """ - response = await prisma_client.db.query_raw( - sql_query, startTime, endTime, selected_api_key - ) + response = await prisma_client.db.query_raw(sql_query, startTime, endTime, selected_api_key) return response -async def global_spend_models_internal_user( - user_api_key_dict: UserAPIKeyAuth, limit: int = 10 -): +async def global_spend_models_internal_user(user_api_key_dict: UserAPIKeyAuth, limit: int = 10): from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -3238,9 +3092,7 @@ async def global_spend_models( user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY ): - response = await global_spend_models_internal_user( - user_api_key_dict=user_api_key_dict, limit=limit - ) + response = await global_spend_models_internal_user(user_api_key_dict=user_api_key_dict, limit=limit) return response if prisma_client is None: @@ -3310,9 +3162,7 @@ async def provider_budgets() -> ProviderBudgetResponse: try: if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": "No llm_router found"} - ) + raise HTTPException(status_code=500, detail={"error": "No llm_router found"}) provider_budget_config = llm_router.provider_budget_config if provider_budget_config is None: @@ -3326,14 +3176,8 @@ async def provider_budgets() -> ProviderBudgetResponse: provider_budget_response_dict: Dict[str, ProviderBudgetResponseObject] = {} for _provider, _budget_info in provider_budget_config.items(): - _provider_spend = ( - await router_budget_logger._get_current_provider_spend(_provider) or 0.0 - ) - _provider_budget_ttl = ( - await router_budget_logger._get_current_provider_budget_reset_at( - _provider - ) - ) + _provider_spend = await router_budget_logger._get_current_provider_spend(_provider) or 0.0 + _provider_budget_ttl = await router_budget_logger._get_current_provider_budget_reset_at(_provider) provider_budget_response_object = ProviderBudgetResponseObject( budget_limit=_budget_info.max_budget, time_period=_budget_info.budget_duration, @@ -3343,15 +3187,11 @@ async def provider_budgets() -> ProviderBudgetResponse: provider_budget_response_dict[_provider] = provider_budget_response_object return ProviderBudgetResponse(providers=provider_budget_response_dict) except Exception as e: - verbose_proxy_logger.exception( - "/provider/budgets: Exception occured - {}".format(str(e)) - ) + verbose_proxy_logger.exception("/provider/budgets: Exception occured - {}".format(str(e))) raise handle_exception_on_proxy(e) -async def get_spend_by_tags( - prisma_client: PrismaClient, start_date=None, end_date=None -): +async def get_spend_by_tags(prisma_client: PrismaClient, start_date=None, end_date=None): response = await prisma_client.db.query_raw(""" SELECT jsonb_array_elements_text(request_tags) AS individual_request_tag, @@ -3508,9 +3348,7 @@ async def ui_view_session_spend_logs( skip = (page - 1) * page_size # Get total count for pagination metadata - total_records = await SpendLogsRepository(prisma_client).table.count( - where=where_conditions - ) + total_records = await SpendLogsRepository(prisma_client).table.count(where=where_conditions) # Query with raw SQL to exclude heavy columns (messages, response, proxy_server_request) sql_query = """ @@ -3527,9 +3365,7 @@ async def ui_view_session_spend_logs( ORDER BY "startTime" DESC LIMIT $2 OFFSET $3 """ - result = await prisma_client.db.query_raw( - sql_query, session_id, page_size, skip - ) + result = await prisma_client.db.query_raw(sql_query, session_id, page_size, skip) total_pages = (total_records + page_size - 1) // page_size @@ -3591,17 +3427,9 @@ async def _build_ui_spend_logs_response( if enrich_session_counts: session_ids = list( { - ( - row.get("session_id") - if isinstance(row, dict) - else getattr(row, "session_id", None) - ) + (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None)) for row in data - if ( - row.get("session_id") - if isinstance(row, dict) - else getattr(row, "session_id", None) - ) + if (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None)) } ) if session_ids: @@ -3614,11 +3442,7 @@ async def _build_ui_spend_logs_response( where={"session_id": {"in": session_ids}}, count={"session_id": True}, ) - count_map = { - r["session_id"]: r["_count"]["session_id"] - for r in counts - if r.get("session_id") - } + count_map = {r["session_id"]: r["_count"]["session_id"] for r in counts if r.get("session_id")} if enrich_session_counts: enriched: List[dict] = [] @@ -3696,9 +3520,7 @@ async def _can_team_member_view_log( if team_id is None: return False - team_row = await TeamRepository(prisma_client).table.find_unique( - where={"team_id": team_id} - ) + team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if team_row is None: return False team_obj = LiteLLM_TeamTable(**team_row.model_dump()) @@ -3759,11 +3581,7 @@ async def _assert_user_can_view_request_id( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "Not authorized to view spend log for request_id={}".format( - request_id - ) - }, + detail={"error": "Not authorized to view spend log for request_id={}".format(request_id)}, ) @@ -3793,9 +3611,7 @@ async def _get_permitted_team_ids_for_spend_logs( if user_obj is None or not user_obj.teams: return [] - team_rows = await TeamRepository(prisma_client).table.find_many( - where={"team_id": {"in": user_obj.teams}} - ) + team_rows = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_obj.teams}}) permitted: List[str] = [] for team_row in team_rows: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index f6e0303e349..22f97feecd9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -70,9 +70,7 @@ def _get_spend_logs_metadata( applied_guardrails: Optional[List[str]] = None, batch_models: Optional[List[str]] = None, mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = None, - vector_store_request_metadata: Optional[ - List[StandardLoggingVectorStoreRequest] - ] = None, + vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = None, guardrail_information: Optional[List[StandardLoggingGuardrailInformation]] = None, usage_object: Optional[dict] = None, model_map_information: Optional[StandardLoggingModelInformation] = None, @@ -113,8 +111,7 @@ def _get_spend_logs_metadata( litellm_call_id=litellm_call_id, ) verbose_proxy_logger.debug( - "getting payload for SpendLogs, available keys in metadata: " - + str(list(metadata.keys())) + "getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys())) ) # Filter the metadata dictionary to include only the specified keys @@ -126,8 +123,8 @@ def _get_spend_logs_metadata( clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata - clean_metadata["vector_store_request_metadata"] = ( - _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) + clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload( + vector_store_request_metadata ) clean_metadata["guardrail_information"] = guardrail_information clean_metadata["usage_object"] = usage_object @@ -163,16 +160,12 @@ def generate_hash_from_response(response_obj: Any) -> str: return hashlib.md5(str(response_obj).encode()).hexdigest() -def get_spend_logs_id( - call_type: str, response_obj: dict, kwargs: dict -) -> Optional[str]: +def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> Optional[str]: if call_type == "aretrieve_batch" or call_type == "acreate_file": # Generate a hash from the response object id: Optional[str] = generate_hash_from_response(response_obj) else: - id = cast(Optional[str], response_obj.get("id")) or cast( - Optional[str], kwargs.get("litellm_call_id") - ) + id = cast(Optional[str], response_obj.get("id")) or cast(Optional[str], kwargs.get("litellm_call_id")) return id @@ -274,9 +267,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs usage = _combined_usage.model_dump() id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs) - standard_logging_payload = cast( - Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None) - ) + standard_logging_payload = cast(Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None)) end_user_id = get_end_user_id_for_cost_tracking(litellm_params) @@ -286,12 +277,8 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs standard_logging_completion_tokens: int = 0 standard_logging_total_tokens: int = 0 if standard_logging_payload is not None: - standard_logging_prompt_tokens = standard_logging_payload.get( - "prompt_tokens", 0 - ) - standard_logging_completion_tokens = standard_logging_payload.get( - "completion_tokens", 0 - ) + standard_logging_prompt_tokens = standard_logging_payload.get("prompt_tokens", 0) + standard_logging_completion_tokens = standard_logging_payload.get("completion_tokens", 0) standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0) if api_key is not None and isinstance(api_key, str): if api_key.startswith("sk-"): @@ -301,24 +288,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if ( standard_logging_payload is not None ): # [TODO] migrate completely to sl payload. currently missing pass-through endpoint data - api_key = ( - api_key - or standard_logging_payload["metadata"].get("user_api_key_hash") - or "" - ) - end_user_id = end_user_id or standard_logging_payload["metadata"].get( - "user_api_key_end_user_id" - ) + api_key = api_key or standard_logging_payload["metadata"].get("user_api_key_hash") or "" + end_user_id = end_user_id or standard_logging_payload["metadata"].get("user_api_key_end_user_id") # BUG FIX: Don't overwrite api_key when standard_logging_payload is None # The api_key was already extracted from metadata (line 243) and hashed (lines 256-259) - request_tags = ( - safe_dumps(metadata.get("tags", [])) - if isinstance(metadata.get("tags", []), list) - else "[]" - ) + request_tags = safe_dumps(metadata.get("tags", [])) if isinstance(metadata.get("tags", []), list) else "[]" if ( - standard_logging_payload is not None - and standard_logging_payload.get("request_tags") is not None + standard_logging_payload is not None and standard_logging_payload.get("request_tags") is not None ): # use 'tags' from standard logging payload instead request_tags = safe_dumps(standard_logging_payload["request_tags"]) @@ -350,9 +326,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs else None ), vector_store_request_metadata=( - standard_logging_payload["metadata"].get( - "vector_store_request_metadata", None - ) + standard_logging_payload["metadata"].get("vector_store_request_metadata", None) if standard_logging_payload is not None else None ), @@ -362,18 +336,12 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs else None ), model_map_information=( - standard_logging_payload["model_map_information"] - if standard_logging_payload is not None - else None + standard_logging_payload["model_map_information"] if standard_logging_payload is not None else None ), guardrail_information=( standard_logging_payload.get("guardrail_information", None) if standard_logging_payload is not None - else ( - metadata.get("standard_logging_guardrail_information", None) - if metadata is not None - else None - ) + else (metadata.get("standard_logging_guardrail_information", None) if metadata is not None else None) ), cold_storage_object_key=( standard_logging_payload["metadata"].get("cold_storage_object_key", None) @@ -382,9 +350,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs ), litellm_overhead_time_ms=litellm_overhead_time_ms, cost_breakdown=( - standard_logging_payload.get("cost_breakdown", None) - if standard_logging_payload is not None - else None + standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None else None ), litellm_call_id=cast( Optional[str], @@ -411,13 +377,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs id = f"{id}_cache_hit{time.time()}" # SpendLogs does not allow duplicate request_id mcp_namespaced_tool_name = None - mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = clean_metadata.get( - "mcp_tool_call_metadata" - ) + mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = clean_metadata.get("mcp_tool_call_metadata") if mcp_tool_call_metadata is not None: - mcp_namespaced_tool_name = mcp_tool_call_metadata.get( - "namespaced_tool_name", None - ) + mcp_namespaced_tool_name = mcp_tool_call_metadata.get("namespaced_tool_name", None) # Extract agent_id for A2A requests (set directly on model_call_details) agent_id: Optional[str] = kwargs.get("agent_id") or metadata.get("agent_id") @@ -443,9 +405,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs spend=kwargs.get("response_cost", 0), total_tokens=usage.get("total_tokens", standard_logging_total_tokens), prompt_tokens=usage.get("prompt_tokens", standard_logging_prompt_tokens), - completion_tokens=usage.get( - "completion_tokens", standard_logging_completion_tokens - ), + completion_tokens=usage.get("completion_tokens", standard_logging_completion_tokens), request_tags=request_tags, end_user=end_user_id or "", api_base=litellm_params.get("api_base", ""), @@ -458,9 +418,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs messages=_get_messages_for_spend_logs_payload( standard_logging_payload=standard_logging_payload, metadata=metadata ), - response=_get_response_for_spend_logs_payload( - payload=standard_logging_payload, kwargs=kwargs - ), + response=_get_response_for_spend_logs_payload(payload=standard_logging_payload, kwargs=kwargs), proxy_server_request=_get_proxy_server_request_for_spend_logs_payload( metadata=metadata, litellm_params=litellm_params, kwargs=kwargs ), @@ -502,10 +460,7 @@ def _get_session_id_for_spend_log( """ from litellm._uuid import uuid - if ( - standard_logging_payload is not None - and standard_logging_payload.get("trace_id") is not None - ): + if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: return str(standard_logging_payload.get("trace_id")) # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs @@ -598,9 +553,7 @@ async def get_spend_by_team_and_customer( group_by_day; """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date, end_date, team_id, customer_id - ) + db_response = await prisma_client.db.query_raw(sql_query, start_date, end_date, team_id, customer_id) if db_response is None: return [] @@ -657,9 +610,7 @@ def _sanitize_request_body_for_spend_logs_payload( def _sanitize_value(value: Any) -> Any: if isinstance(value, dict): - return _sanitize_request_body_for_spend_logs_payload( - value, visited, max_string_length_prompt_in_db - ) + return _sanitize_request_body_for_spend_logs_payload(value, visited, max_string_length_prompt_in_db) elif isinstance(value, list): return [_sanitize_value(item) for item in value] elif isinstance(value, str): @@ -696,11 +647,7 @@ def _sanitize_request_body_for_spend_logs_payload( return value return value - return { - k: _sanitize_value(v) - for k, v in request_body.items() - if k not in _SENSITIVE_REQUEST_BODY_KEYS - } + return {k: _sanitize_value(v) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS} # Quoted-key form: ``"input"`` / ``'messages'`` / ``"prompt"`` followed by @@ -880,9 +827,7 @@ def _sanitize_error_information_for_spend_logs( return cast(StandardLoggingPayloadErrorInformation, sanitized) -def _convert_to_json_serializable_dict( - obj: Any, visited: Optional[set] = None, max_depth: int = 20 -) -> Any: +def _convert_to_json_serializable_dict(obj: Any, visited: Optional[set] = None, max_depth: int = 20) -> Any: """ Convert object to JSON-serializable dict, handling Pydantic models safely. @@ -921,20 +866,12 @@ def _convert_to_json_serializable_dict( # Recursively process the dumped dict return _convert_to_json_serializable_dict(result, visited, max_depth - 1) elif isinstance(obj, dict): - return { - k: _convert_to_json_serializable_dict(v, visited, max_depth - 1) - for k, v in obj.items() - } + return {k: _convert_to_json_serializable_dict(v, visited, max_depth - 1) for k, v in obj.items()} elif isinstance(obj, list): - return [ - _convert_to_json_serializable_dict(item, visited, max_depth - 1) - for item in obj - ] + return [_convert_to_json_serializable_dict(item, visited, max_depth - 1) for item in obj] elif hasattr(obj, "__dict__"): # Handle objects with __dict__ attribute - return _convert_to_json_serializable_dict( - obj.__dict__, visited, max_depth - 1 - ) + return _convert_to_json_serializable_dict(obj.__dict__, visited, max_depth - 1) else: # Primitives (str, int, float, bool, None) pass through return obj @@ -955,9 +892,7 @@ def _get_proxy_server_request_for_spend_logs_payload( If turn_off_message_logging is enabled, redact messages in the request body. """ if _should_store_prompts_and_responses_in_spend_logs(): - _proxy_server_request = cast( - Optional[dict], litellm_params.get("proxy_server_request", {}) - ) + _proxy_server_request = cast(Optional[dict], litellm_params.get("proxy_server_request", {})) if _proxy_server_request is not None: _request_body = _proxy_server_request.get("body", {}) or {} @@ -977,9 +912,7 @@ def _get_proxy_server_request_for_spend_logs_payload( # Build model_call_details dict to check redaction settings model_call_details = { "litellm_params": litellm_params, - "standard_callback_dynamic_params": kwargs.get( - "standard_callback_dynamic_params" - ), + "standard_callback_dynamic_params": kwargs.get("standard_callback_dynamic_params"), } # If redaction is enabled, convert to serializable dict before redacting @@ -1012,8 +945,7 @@ def _get_vector_store_request_for_spend_logs_payload( return None for vector_store_request in vector_store_request_metadata: vector_store_search_response: VectorStoreSearchResponse = ( - vector_store_request.get("vector_store_search_response") - or VectorStoreSearchResponse() + vector_store_request.get("vector_store_search_response") or VectorStoreSearchResponse() ) response_data = vector_store_search_response.get("data", []) or [] for response_item in response_data: @@ -1050,21 +982,15 @@ def _get_response_for_spend_logs_payload( litellm_params = kwargs.get("litellm_params", {}) model_call_details = { "litellm_params": litellm_params, - "standard_callback_dynamic_params": kwargs.get( - "standard_callback_dynamic_params" - ), + "standard_callback_dynamic_params": kwargs.get("standard_callback_dynamic_params"), } # If redaction is enabled, convert to serializable dict before redacting if should_redact_message_logging(model_call_details=model_call_details): response_obj = _convert_to_json_serializable_dict(response_obj) - response_obj = perform_redaction( - model_call_details={}, result=response_obj - ) + response_obj = perform_redaction(model_call_details={}, result=response_obj) - sanitized_wrapper = _sanitize_request_body_for_spend_logs_payload( - {"response": response_obj} - ) + sanitized_wrapper = _sanitize_request_body_for_spend_logs_payload({"response": response_obj}) sanitized_response = sanitized_wrapper.get("response", response_obj) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 1dde31b54cb..52d46e0d267 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -34,9 +34,7 @@ def _get_registered_vantage_logger(): """Return the VantageLogger already registered in litellm.callbacks, if any.""" from litellm.integrations.vantage.vantage_logger import VantageLogger - vantage_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=VantageLogger - ) + vantage_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=VantageLogger) if vantage_loggers: return vantage_loggers[0] return None @@ -98,15 +96,11 @@ async def _get_vantage_settings(): encrypted_api_key = settings.get("api_key") if encrypted_api_key: - decrypted_api_key = decrypt_value_helper( - encrypted_api_key, key="vantage_api_key", exception_type="error" - ) + decrypted_api_key = decrypt_value_helper(encrypted_api_key, key="vantage_api_key", exception_type="error") if decrypted_api_key is None: raise HTTPException( status_code=500, - detail={ - "error": "Failed to decrypt Vantage API key. Check your salt key configuration." - }, + detail={"error": "Failed to decrypt Vantage API key. Check your salt key configuration."}, ) settings["api_key"] = decrypted_api_key @@ -120,9 +114,7 @@ async def _get_vantage_settings(): if decrypted_integration_token is None: raise HTTPException( status_code=500, - detail={ - "error": "Failed to decrypt Vantage integration token. Check your salt key configuration." - }, + detail={"error": "Failed to decrypt Vantage integration token. Check your salt key configuration."}, ) settings["integration_token"] = decrypted_integration_token @@ -215,16 +207,10 @@ async def update_vantage_settings( if not current_settings: raise HTTPException( status_code=404, - detail={ - "error": "Vantage settings not found. Please initialize settings first using /vantage/init" - }, + detail={"error": "Vantage settings not found. Please initialize settings first using /vantage/init"}, ) - updated_api_key = ( - request.api_key - if request.api_key is not None - else current_settings.get("api_key", "") - ) + updated_api_key = request.api_key if request.api_key is not None else current_settings.get("api_key", "") updated_token = ( request.integration_token if request.integration_token is not None @@ -244,9 +230,7 @@ async def update_vantage_settings( verbose_proxy_logger.info("Vantage settings updated successfully") - return VantageInitResponse( - message="Vantage settings updated successfully", status="success" - ) + return VantageInitResponse(message="Vantage settings updated successfully", status="success") except HTTPException: raise @@ -335,9 +319,7 @@ async def init_vantage_settings( verbose_proxy_logger.info("Vantage settings initialized successfully") - return VantageInitResponse( - message="Vantage settings initialized successfully", status="success" - ) + return VantageInitResponse(message="Vantage settings initialized successfully", status="success") except HTTPException: raise @@ -393,26 +375,14 @@ async def vantage_dry_run_export( def _to_json_safe_dicts(frame: pl.DataFrame) -> list: """Cast Decimal columns to Float64 so .to_dicts() produces JSON-serializable float values instead of decimal.Decimal.""" - decimal_cols = [ - col - for col, dtype in zip(frame.columns, frame.dtypes) - if isinstance(dtype, pl.Decimal) - ] + decimal_cols = [col for col, dtype in zip(frame.columns, frame.dtypes) if isinstance(dtype, pl.Decimal)] if decimal_cols: - frame = frame.with_columns( - [pl.col(c).cast(pl.Float64) for c in decimal_cols] - ) + frame = frame.with_columns([pl.col(c).cast(pl.Float64) for c in decimal_cols]) return frame.to_dicts() - usage_sample = ( - _to_json_safe_dicts(data.head(min(50, len(data)))) - if not data.is_empty() - else [] - ) + usage_sample = _to_json_safe_dicts(data.head(min(50, len(data)))) if not data.is_empty() else [] normalized_sample = ( - _to_json_safe_dicts(normalized.head(min(50, len(normalized)))) - if not normalized.is_empty() - else [] + _to_json_safe_dicts(normalized.head(min(50, len(normalized)))) if not normalized.is_empty() else [] ) # Use the same pre-transform column names as @@ -564,22 +534,16 @@ async def delete_vantage_settings( detail={"error": "Vantage settings not found"}, ) - await ConfigRepository(prisma_client).table.delete( - where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} - ) + await ConfigRepository(prisma_client).table.delete(where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}) # Deregister in-memory VantageLogger so the scheduler stops firing from litellm.integrations.vantage.vantage_logger import VantageLogger - litellm.logging_callback_manager.remove_callbacks_by_type( - litellm.callbacks, VantageLogger - ) + litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, VantageLogger) verbose_proxy_logger.info("Vantage settings deleted successfully") - return VantageInitResponse( - message="Vantage settings deleted successfully", status="success" - ) + return VantageInitResponse(message="Vantage settings deleted successfully", status="success") except HTTPException: raise diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py index 8bdaf3fd9e8..e2206541fc6 100644 --- a/litellm/proxy/types_utils/utils.py +++ b/litellm/proxy/types_utils/utils.py @@ -46,14 +46,10 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: spec = importlib.util.spec_from_file_location(module_name, module_file_path) # type: ignore if spec is None: - raise ImportError( - f"Could not find a module specification for {module_file_path}" - ) + raise ImportError(f"Could not find a module specification for {module_file_path}") module = importlib.util.module_from_spec(spec) # type: ignore if spec.loader is None: - raise ImportError( - f"Could not find a module loader for {module_file_path}" - ) + raise ImportError(f"Could not find a module loader for {module_file_path}") spec.loader.exec_module(module) # type: ignore else: # Dynamically import the module @@ -66,18 +62,14 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: except ImportError as e: # Re-raise the exception with a user-friendly message if instance_name and module_name: - raise ImportError( - f"Could not import {instance_name} from {module_name}" - ) from e + raise ImportError(f"Could not import {instance_name} from {module_name}") from e else: raise e except Exception as e: raise e -def _load_instance_from_remote_storage( - remote_url: str, config_file_path: Optional[str] = None -) -> Any: +def _load_instance_from_remote_storage(remote_url: str, config_file_path: Optional[str] = None) -> Any: """ Load custom logger instance from S3 or GCS URL. @@ -130,9 +122,7 @@ def _load_instance_from_remote_storage( # Split by last dot to separate module from instance module_parts = path_and_module.split(".") if len(module_parts) < 2: - raise ValueError( - f"Invalid module specification in {remote_url}. Expected: path/to/module.instance_name" - ) + raise ValueError(f"Invalid module specification in {remote_url}. Expected: path/to/module.instance_name") instance_name = module_parts[-1] module_path = ".".join(module_parts[:-1]) @@ -164,14 +154,10 @@ def _load_instance_from_remote_storage( local_file_path=local_file_path, ) else: # gcs - success = asyncio.run( - _download_gcs_file_wrapper(bucket_name, object_key, local_file_path) - ) + success = asyncio.run(_download_gcs_file_wrapper(bucket_name, object_key, local_file_path)) if not success: - raise ImportError( - f"Failed to download {object_key} from {storage_type} bucket {bucket_name}" - ) + raise ImportError(f"Failed to download {object_key} from {storage_type} bucket {bucket_name}") # Load the module from the downloaded file using the actual module name spec = importlib.util.spec_from_file_location(module_path, local_file_path) @@ -188,33 +174,23 @@ def _load_instance_from_remote_storage( try: os.remove(local_file_path) except Exception as cleanup_error: - verbose_proxy_logger.warning( - f"Could not clean up temporary file {local_file_path}: {cleanup_error}" - ) + verbose_proxy_logger.warning(f"Could not clean up temporary file {local_file_path}: {cleanup_error}") - verbose_proxy_logger.info( - f"Successfully loaded custom logger from {remote_url}" - ) + verbose_proxy_logger.info(f"Successfully loaded custom logger from {remote_url}") return instance except Exception as e: - raise ImportError( - f"Failed to load custom logger from {remote_url}: {str(e)}" - ) from e + raise ImportError(f"Failed to load custom logger from {remote_url}: {str(e)}") from e -async def _download_gcs_file_wrapper( - bucket_name: str, object_key: str, local_file_path: str -) -> bool: +async def _download_gcs_file_wrapper(bucket_name: str, object_key: str, local_file_path: str) -> bool: """Wrapper for GCS download to handle async properly""" try: from litellm.proxy.common_utils.load_config_utils import ( download_python_file_from_gcs, ) - return await download_python_file_from_gcs( - bucket_name, object_key, local_file_path - ) + return await download_python_file_from_gcs(bucket_name, object_key, local_file_path) except Exception as e: from litellm._logging import verbose_proxy_logger @@ -232,8 +208,6 @@ def validate_custom_validate_return_type( return_type = hints.get("return") if return_type != Literal[True]: - raise TypeError( - f"Custom validator must be annotated to return Literal[True], got {return_type}" - ) + raise TypeError(f"Custom validator must be annotated to return Literal[True], got {return_type}") return fn diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 674f7efd837..e4f68a1e9db 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -343,9 +343,7 @@ async def add_allowed_ip(ip_address: IPAddress): if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) # Load existing config @@ -524,9 +522,7 @@ async def get_default_team_settings(): ) -async def update_default_team_member_budget( - teams: List[NewUserRequestTeam], user_api_key_dict: UserAPIKeyAuth -): +async def update_default_team_member_budget(teams: List[NewUserRequestTeam], user_api_key_dict: UserAPIKeyAuth): """ 1. Update the max member budget for the team """ @@ -554,9 +550,7 @@ async def update_default_team_member_budget( async def _update_litellm_setting( - settings: Union[ - DefaultInternalUserParams, DefaultTeamSSOParams, MCPSemanticFilterSettings - ], + settings: Union[DefaultInternalUserParams, DefaultTeamSSOParams, MCPSemanticFilterSettings], settings_key: str, success_message: str, ): @@ -573,9 +567,7 @@ async def _update_litellm_setting( if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) in_memory_var = settings.model_dump(exclude_none=True) @@ -617,9 +609,7 @@ async def update_internal_user_settings( Update the default internal user parameters for SSO users. These settings will be applied to new users who sign in via SSO. """ - if settings.teams is not None and all( - isinstance(team, NewUserRequestTeam) for team in settings.teams - ): + if settings.teams is not None and all(isinstance(team, NewUserRequestTeam) for team in settings.teams): await update_default_team_member_budget( settings.teams, user_api_key_dict=user_api_key_dict, # type: ignore @@ -670,9 +660,7 @@ async def get_sso_settings(): ) # Get SSO config from dedicated table - sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( - where={"id": "sso_config"} - ) + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) # Initialize with defaults sso_settings_dict = {} @@ -709,29 +697,15 @@ async def get_sso_settings(): sso_config = SSOConfig( google_client_id=decrypted_sso_settings_dict.get("google_client_id", None), - google_client_secret=decrypted_sso_settings_dict.get( - "google_client_secret", None - ), - microsoft_client_id=decrypted_sso_settings_dict.get( - "microsoft_client_id", None - ), - microsoft_client_secret=decrypted_sso_settings_dict.get( - "microsoft_client_secret", None - ), + google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None), + microsoft_client_id=decrypted_sso_settings_dict.get("microsoft_client_id", None), + microsoft_client_secret=decrypted_sso_settings_dict.get("microsoft_client_secret", None), microsoft_tenant=decrypted_sso_settings_dict.get("microsoft_tenant", None), generic_client_id=decrypted_sso_settings_dict.get("generic_client_id", None), - generic_client_secret=decrypted_sso_settings_dict.get( - "generic_client_secret", None - ), - generic_authorization_endpoint=decrypted_sso_settings_dict.get( - "generic_authorization_endpoint", None - ), - generic_token_endpoint=decrypted_sso_settings_dict.get( - "generic_token_endpoint", None - ), - generic_userinfo_endpoint=decrypted_sso_settings_dict.get( - "generic_userinfo_endpoint", None - ), + generic_client_secret=decrypted_sso_settings_dict.get("generic_client_secret", None), + generic_authorization_endpoint=decrypted_sso_settings_dict.get("generic_authorization_endpoint", None), + generic_token_endpoint=decrypted_sso_settings_dict.get("generic_token_endpoint", None), + generic_userinfo_endpoint=decrypted_sso_settings_dict.get("generic_userinfo_endpoint", None), proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None), user_email=decrypted_sso_settings_dict.get("user_email"), ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"), @@ -794,9 +768,7 @@ async def update_sso_settings(sso_config: SSOConfig): if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) # Update environment variables @@ -836,9 +808,7 @@ async def update_sso_settings(sso_config: SSOConfig): # Clear environment variable if value is null/empty os.environ.pop(env_var_name, None) - encrypted_sso_data = proxy_config._encrypt_env_variables( - environment_variables=sso_data - ) + encrypted_sso_data = proxy_config._encrypt_env_variables(environment_variables=sso_data) # Save to dedicated SSO table await SSOConfigRepository(prisma_client).table.upsert( @@ -872,9 +842,7 @@ async def update_sso_settings(sso_config: SSOConfig): env_vars_to_remove = set(env_var_mapping.values()) filtered_env_vars = { - key: value - for key, value in environment_variables.items() - if key not in env_vars_to_remove + key: value for key, value in environment_variables.items() if key not in env_vars_to_remove } await ConfigRepository(prisma_client).table.update( @@ -964,9 +932,7 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) # Load existing config @@ -1028,10 +994,7 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): # Handle environment variable encryption if needed stored_config = config.copy() - if ( - "environment_variables" in stored_config - and len(stored_config["environment_variables"]) > 0 - ): + if "environment_variables" in stored_config and len(stored_config["environment_variables"]) > 0: # Only encrypt if there are environment variables to encrypt stored_config["environment_variables"] = proxy_config._encrypt_env_variables( environment_variables=stored_config["environment_variables"] @@ -1099,13 +1062,9 @@ async def update_mcp_semantic_filter_settings( from litellm.proxy.proxy_server import prisma_client, proxy_config if prisma_client is not None: - await proxy_config._init_semantic_filter_settings_in_db( - prisma_client=prisma_client - ) + await proxy_config._init_semantic_filter_settings_in_db(prisma_client=prisma_client) except Exception as e: - verbose_proxy_logger.warning( - f"Failed to reinitialize MCP semantic filter settings immediately: {e}" - ) + verbose_proxy_logger.warning(f"Failed to reinitialize MCP semantic filter settings immediately: {e}") return result @@ -1132,23 +1091,17 @@ async def get_ui_settings_cached() -> Dict[str, Any]: if prisma_client is None: return {} - db_record = await UISettingsRepository(prisma_client).table.find_unique( - where={"id": "ui_settings"} - ) + db_record = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) ui_settings: Dict[str, Any] = {} if db_record and db_record.ui_settings: raw = db_record.ui_settings ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw) # Sanitize - ui_settings = { - k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS - } + ui_settings = {k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS} # 3. Populate cache with TTL - await user_api_key_cache.async_set_cache( - key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL - ) + await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) return ui_settings @@ -1173,9 +1126,7 @@ async def get_ui_settings(): ui_settings: Dict[str, Any] = {} - db_record = await UISettingsRepository(prisma_client).table.find_unique( - where={"id": "ui_settings"} - ) + db_record = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) if db_record and db_record.ui_settings: ui_settings_json = db_record.ui_settings @@ -1185,15 +1136,11 @@ async def get_ui_settings(): ui_settings = dict(ui_settings_json) # Sanitize any unexpected keys from persisted config before returning - ui_settings = { - k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS - } + ui_settings = {k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS} # Sync runtime flags into general_settings so the proxy picks them up # at runtime (covers server restart scenarios). - _flags_to_sync = { - k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings - } + _flags_to_sync = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings} if _flags_to_sync: from litellm.proxy.proxy_server import general_settings @@ -1202,9 +1149,7 @@ async def get_ui_settings(): # Refresh DualCache so other code paths (e.g. /user/filter/ui) see fresh values from litellm.proxy.proxy_server import user_api_key_cache - await user_api_key_cache.async_set_cache( - key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL - ) + await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) # Build config-like object for schema helper config: Dict[str, Any] = {"litellm_settings": {"ui_settings": ui_settings}} @@ -1232,9 +1177,7 @@ async def update_ui_settings( from litellm.proxy.proxy_server import prisma_client, store_model_in_db if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: - raise HTTPException( - status_code=403, detail="Only proxy admins can update UI settings." - ) + raise HTTPException(status_code=403, detail="Only proxy admins can update UI settings.") if prisma_client is None: raise HTTPException( @@ -1245,9 +1188,7 @@ async def update_ui_settings( if store_model_in_db is not True: raise HTTPException( status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, + detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) # Validate against the same effective class GET advertises, so @@ -1263,10 +1204,7 @@ async def update_ui_settings( # Reject enterprise-only settings up front so the caller gets a clear # signal instead of a silent drop. - blocked_enterprise_keys = sorted( - (settings_dict.keys() & _ENTERPRISE_ONLY_UI_SETTINGS) - - ALLOWED_UI_SETTINGS_FIELDS - ) + blocked_enterprise_keys = sorted((settings_dict.keys() & _ENTERPRISE_ONLY_UI_SETTINGS) - ALLOWED_UI_SETTINGS_FIELDS) if blocked_enterprise_keys: raise HTTPException( status_code=403, @@ -1279,16 +1217,12 @@ async def update_ui_settings( ) # Enforce allowlist and drop anything unexpected - incoming = { - k: v for k, v in settings_dict.items() if k in ALLOWED_UI_SETTINGS_FIELDS - } + incoming = {k: v for k, v in settings_dict.items() if k in ALLOWED_UI_SETTINGS_FIELDS} # Merge with existing persisted settings so a partial PATCH doesn't # overwrite fields the caller didn't send. existing: dict = {} - db_existing = await UISettingsRepository(prisma_client).table.find_unique( - where={"id": "ui_settings"} - ) + db_existing = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) if db_existing and db_existing.ui_settings: raw = db_existing.ui_settings existing = json.loads(raw) if isinstance(raw, str) else dict(raw) @@ -1310,9 +1244,7 @@ async def update_ui_settings( # Sync runtime flags to general_settings so the proxy picks them up # at runtime (general_settings is checked in pre-call utils). - _flags_to_sync = { - k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings - } + _flags_to_sync = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings} if _flags_to_sync: from litellm.proxy.proxy_server import general_settings @@ -1321,12 +1253,8 @@ async def update_ui_settings( # Invalidate + set DualCache so subsequent reads see the new values immediately from litellm.proxy.proxy_server import user_api_key_cache - sanitized = { - k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS - } - await user_api_key_cache.async_set_cache( - key=UI_SETTINGS_CACHE_KEY, value=sanitized, ttl=UI_SETTINGS_CACHE_TTL - ) + sanitized = {k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS} + await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=sanitized, ttl=UI_SETTINGS_CACHE_TTL) return { "message": "UI settings updated successfully", @@ -1361,9 +1289,7 @@ async def upload_logo(file: UploadFile = File(...)): # Validate file size (max 5MB) file_content = await file.read() if len(file_content) > 5 * 1024 * 1024: # 5MB - raise HTTPException( - status_code=400, detail="File size too large. Maximum size is 5MB." - ) + raise HTTPException(status_code=400, detail="File size too large. Maximum size is 5MB.") # Create uploads directory if it doesn't exist current_dir = os.path.dirname(os.path.abspath(__file__)) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 51cf30ee454..154a17bc4db 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -71,9 +71,7 @@ except ImportError: try: import backoff except ImportError: - raise ImportError( - "backoff is not installed. Please install it via 'pip install backoff'" - ) + raise ImportError("backoff is not installed. Please install it via 'pip install backoff'") from fastapi import HTTPException, status @@ -340,9 +338,7 @@ def _accepts_litellm_call_info(cb: CustomLogger) -> bool: return _CALLBACK_ACCEPTS_CALL_INFO[key] -def _enrich_http_exception_with_guardrail_context( - exc: BaseException, callback: Any -) -> None: +def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: Any) -> None: """ If `exc` is an HTTPException with a dict `detail`, mutate it in place to add `guardrail_name` and `guardrail_mode` taken from the callback instance. @@ -418,9 +414,7 @@ class ProxyLogging: self.internal_usage_cache: InternalUsageCache = InternalUsageCache( dual_cache=DualCache(default_in_memory_ttl=1) # ping redis cache every 1s ) - self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler( - self.internal_usage_cache - ) + self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache) self.max_budget_limiter = _PROXY_MaxBudgetLimiter() self.cache_control_check = _PROXY_CacheControlCheck() self.alerting: Optional[List] = None @@ -473,16 +467,13 @@ class ProxyLogging: and not self.daily_report_started ): asyncio.create_task( - self.slack_alerting_instance._run_scheduled_daily_report( - llm_router=llm_router - ) + self.slack_alerting_instance._run_scheduled_daily_report(llm_router=llm_router) ) # RUN DAILY REPORT (if scheduled) self.daily_report_started = True if ( self.slack_alerting_instance is not None - and AlertType.llm_requests_hanging - in self.slack_alerting_instance.alert_types + and AlertType.llm_requests_hanging in self.slack_alerting_instance.alert_types and not self.hanging_requests_check_started ): asyncio.create_task( @@ -534,9 +525,7 @@ class ProxyLogging: or "outage_alerts" in self.alert_types or "region_outage_alerts" in self.alert_types ): - litellm.logging_callback_manager.add_litellm_callback( - self.slack_alerting_instance - ) # type: ignore + litellm.logging_callback_manager.add_litellm_callback(self.slack_alerting_instance) # type: ignore litellm.logging_callback_manager.add_litellm_success_callback( self.slack_alerting_instance.response_taking_too_long_callback ) @@ -608,16 +597,10 @@ class ProxyLogging: if isinstance(callback, CustomLogger): litellm.logging_callback_manager.add_litellm_success_callback(callback) litellm.logging_callback_manager.add_litellm_failure_callback(callback) - litellm.logging_callback_manager.add_litellm_async_success_callback( - callback - ) - litellm.logging_callback_manager.add_litellm_async_failure_callback( - callback - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - async def update_request_status( - self, litellm_call_id: str, status: Literal["success", "fail"] - ): + async def update_request_status(self, litellm_call_id: str, status: Literal["success", "fail"]): # only use this if slack alerting is being used if self.alerting is None: return @@ -658,13 +641,9 @@ class ProxyLogging: from litellm.types.llms.openai import ChatCompletionUserMessage # Create a synthetic message that represents the tool call - tool_call_content = ( - f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}" - ) + tool_call_content = f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}" - synthetic_message = ChatCompletionUserMessage( - role="user", content=tool_call_content - ) + synthetic_message = ChatCompletionUserMessage(role="user", content=tool_call_content) # Create synthetic LLM data that guardrails can process synthetic_data = { @@ -686,9 +665,7 @@ class ProxyLogging: return synthetic_data - def _convert_llm_result_to_mcp_response( - self, llm_result, request_obj - ) -> Optional[Any]: + def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> Optional[Any]: """ Convert LLM guardrail result back to MCP response format. """ @@ -707,20 +684,12 @@ class ProxyLogging: modified_messages = llm_result.get("messages") if modified_messages: # Check if content was blocked/modified - original_content = ( - f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}" - ) - new_content = ( - modified_messages[0].get("content", "") if modified_messages else "" - ) + original_content = f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}" + new_content = modified_messages[0].get("content", "") if modified_messages else "" if new_content != original_content: # Content was modified - could be masking, redaction, or blocking - if ( - not new_content - or "blocked" in new_content.lower() - or "violation" in new_content.lower() - ): + if not new_content or "blocked" in new_content.lower() or "violation" in new_content.lower(): # Content was blocked completely return MCPPreCallResponseObject( should_proceed=False, @@ -731,11 +700,7 @@ class ProxyLogging: # Content was masked/redacted - extract the modified arguments try: # Try to parse the modified arguments from the masked content - modified_args = ( - self._extract_modified_arguments_from_content( - new_content, request_obj - ) - ) + modified_args = self._extract_modified_arguments_from_content(new_content, request_obj) if modified_args is not None: # Return the masked/redacted arguments for the MCP call to use return MCPPreCallResponseObject( @@ -750,31 +715,23 @@ class ProxyLogging: ) return None except Exception as e: - verbose_proxy_logger.error( - f"Error parsing modified arguments: {e}" - ) + verbose_proxy_logger.error(f"Error parsing modified arguments: {e}") # Fallback: allow original call return None # If result is a string, it's likely an error message if isinstance(llm_result, str): - return MCPPreCallResponseObject( - should_proceed=False, error_message=llm_result, modified_arguments=None - ) + return MCPPreCallResponseObject(should_proceed=False, error_message=llm_result, modified_arguments=None) return None - def _extract_modified_arguments_from_content( - self, masked_content: str, request_obj - ) -> Optional[dict]: + def _extract_modified_arguments_from_content(self, masked_content: str, request_obj) -> Optional[dict]: """ Extract modified/masked arguments from the guardrail response content. """ import json - verbose_proxy_logger.debug( - f"Extracting modified args from content: {masked_content}" - ) + verbose_proxy_logger.debug(f"Extracting modified args from content: {masked_content}") try: # The format should be: "Tool: \nArguments: " @@ -790,32 +747,22 @@ class ProxyLogging: # Try to parse as JSON first try: modified_args = json.loads(args_text) - verbose_proxy_logger.debug( - f"Successfully parsed JSON args: {modified_args}" - ) + verbose_proxy_logger.debug(f"Successfully parsed JSON args: {modified_args}") return modified_args except json.JSONDecodeError as e: # If JSON parsing fails, try to extract key-value pairs manually - verbose_proxy_logger.debug( - f"Failed to parse JSON arguments: {args_text}, error: {e}" - ) - return self._parse_arguments_manually( - args_text, request_obj.arguments - ) + verbose_proxy_logger.debug(f"Failed to parse JSON arguments: {args_text}, error: {e}") + return self._parse_arguments_manually(args_text, request_obj.arguments) # If we can't find the Arguments: line, return None - verbose_proxy_logger.warning( - "Could not find 'Arguments:' line in masked content" - ) + verbose_proxy_logger.warning("Could not find 'Arguments:' line in masked content") return None except Exception as e: verbose_proxy_logger.error(f"Error extracting modified arguments: {e}") return None - def _parse_arguments_manually( - self, args_text: str, original_args: dict - ) -> Optional[dict]: + def _parse_arguments_manually(self, args_text: str, original_args: dict) -> Optional[dict]: """ Try to manually parse arguments when JSON parsing fails. This is a fallback for cases where the guardrail modifies the format. @@ -831,9 +778,7 @@ class ProxyLogging: for key, original_value in original_args.items(): if isinstance(original_value, str): # Look for the key in the masked content and try to extract its value - pattern = ( - rf"['\"]?{re.escape(key)}['\"]?\s*:\s*['\"]?([^,'\"]*)['\"]?" - ) + pattern = rf"['\"]?{re.escape(key)}['\"]?\s*:\s*['\"]?([^,'\"]*)['\"]?" match = re.search(pattern, args_text, re.IGNORECASE) if match: new_value = match.group(1).strip() @@ -846,29 +791,21 @@ class ProxyLogging: verbose_proxy_logger.error(f"Error in manual argument parsing: {e}") return None - def _convert_llm_result_to_mcp_during_response( - self, llm_result, request_obj - ) -> Optional[Any]: + def _convert_llm_result_to_mcp_during_response(self, llm_result, request_obj) -> Optional[Any]: """ Convert LLM guardrail result back to MCP during call response format. """ # If result is an exception, it means the guardrail wants to stop execution if isinstance(llm_result, Exception): - return MCPDuringCallResponseObject( - should_continue=False, error_message=str(llm_result) - ) + return MCPDuringCallResponseObject(should_continue=False, error_message=str(llm_result)) # If result is a dict with modified messages, check for content filtering if isinstance(llm_result, dict): modified_messages = llm_result.get("messages") if modified_messages: # Check if content was blocked/modified - original_content = ( - f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}" - ) - new_content = ( - modified_messages[0].get("content", "") if modified_messages else "" - ) + original_content = f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}" + new_content = modified_messages[0].get("content", "") if modified_messages else "" if new_content != original_content: # Content was modified, could be masking or blocking @@ -887,15 +824,11 @@ class ProxyLogging: # If result is a string, it's likely an error message if isinstance(llm_result, str): - return MCPDuringCallResponseObject( - should_continue=False, error_message=llm_result - ) + return MCPDuringCallResponseObject(should_continue=False, error_message=llm_result) return None - def get_combined_callback_list( - self, dynamic_success_callbacks: Optional[List], global_callbacks: List - ) -> List: + def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: if dynamic_success_callbacks is None: return list(global_callbacks) return list(set(dynamic_success_callbacks + global_callbacks)) @@ -914,25 +847,20 @@ class ProxyLogging: """ result = { "should_proceed": response.should_proceed, - "modified_arguments": response.modified_arguments - or original_request.arguments, + "modified_arguments": response.modified_arguments or original_request.arguments, "error_message": response.error_message, "hidden_params": response.hidden_params, } return result - def _create_mcp_request_object_from_kwargs( - self, kwargs: dict - ) -> "MCPPreCallRequestObject": + def _create_mcp_request_object_from_kwargs(self, kwargs: dict) -> "MCPPreCallRequestObject": """ Helper function to create MCPPreCallRequestObject from kwargs for standard pre_call_hook. """ from litellm.types.llms.base import HiddenParams from litellm.types.mcp import MCPPreCallRequestObject - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( - kwargs.get("user_api_key_auth") - ) + user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(kwargs.get("user_api_key_auth")) return MCPPreCallRequestObject( tool_name=kwargs.get("name", ""), @@ -942,9 +870,7 @@ class ProxyLogging: hidden_params=HiddenParams(), ) - def _convert_mcp_hook_response_to_kwargs( - self, response_data: Optional[dict], original_kwargs: dict - ) -> dict: + def _convert_mcp_hook_response_to_kwargs(self, response_data: Optional[dict], original_kwargs: dict) -> dict: """ Helper function to convert pre_call_hook response back to kwargs for MCP usage. @@ -1004,11 +930,7 @@ class ProxyLogging: if llm_router is None or not hasattr(llm_router, "guardrail_list"): return False - matching = [ - g - for g in llm_router.guardrail_list - if g.get("guardrail_name") == guardrail_name - ] + matching = [g for g in llm_router.guardrail_list if g.get("guardrail_name") == guardrail_name] return len(matching) > 1 async def _execute_guardrail_hook( @@ -1037,8 +959,7 @@ class ProxyLogging: # Use unified_guardrail if callback has apply_guardrail method has_apply_guardrail = "apply_guardrail" in type(callback).__dict__ use_unified = has_apply_guardrail and not ( - hook_type == "during_call" - and getattr(callback, "use_native_during_call_hook", False) + hook_type == "during_call" and getattr(callback, "use_native_during_call_hook", False) ) if use_unified: data["guardrail_to_apply"] = callback @@ -1096,9 +1017,7 @@ class ProxyLogging: raise ValueError("Router not initialized") # Select guardrail using router's load balancing - selected_guardrail = llm_router.get_available_guardrail( - guardrail_name=guardrail_name - ) + selected_guardrail = llm_router.get_available_guardrail(guardrail_name=guardrail_name) callback = selected_guardrail.get("callback") if callback is None: @@ -1138,10 +1057,7 @@ class ProxyLogging: from litellm.types.guardrails import GuardrailEventHooks # Determine the event type based on call type - if ( - event_type is GuardrailEventHooks.pre_call - and call_type == CallTypes.call_mcp_tool.value - ): + if event_type is GuardrailEventHooks.pre_call and call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.pre_mcp_call # Check if the guardrail should run for this request @@ -1158,9 +1074,7 @@ class ProxyLogging: try: # Check if load balancing should be used - if guardrail_name and self._should_use_guardrail_load_balancing( - guardrail_name - ): + if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): response = await self._execute_guardrail_with_load_balancing( guardrail_name=guardrail_name, hook_type="pre_call", @@ -1180,9 +1094,7 @@ class ProxyLogging: # Process the response if one was returned if response is not None: - data = await self.process_pre_call_hook_response( - response=response, data=data, call_type=call_type - ) + data = await self.process_pre_call_hook_response(response=response, data=data, call_type=call_type) callback.mark_pre_call_hook_ran(data) @@ -1202,9 +1114,7 @@ class ProxyLogging: # Get guardrail name for metrics (fallback if not set) metrics_guardrail_name = ( - guardrail_name - or getattr(callback, "guardrail_name", callback.__class__.__name__) - or "unknown" + guardrail_name or getattr(callback, "guardrail_name", callback.__class__.__name__) or "unknown" ) self._emit_guardrail_metrics( @@ -1240,13 +1150,9 @@ class ProxyLogging: all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, ) else: - lookup_prompt_id = construct_versioned_prompt_id( - prompt_id=prompt_id, version=prompt_version - ) + lookup_prompt_id = construct_versioned_prompt_id(prompt_id=prompt_id, version=prompt_version) - custom_logger = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id( - lookup_prompt_id - ) + custom_logger = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(lookup_prompt_id) prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(lookup_prompt_id) litellm_prompt_id: Optional[str] = None if prompt_spec is not None: @@ -1296,28 +1202,17 @@ class ProxyLogging: if guardrails_in_metadata and isinstance(guardrails_in_metadata, list): applied_guardrails = [] - if ( - isinstance(metadata_standard, dict) - and "applied_guardrails" in metadata_standard - ): + if isinstance(metadata_standard, dict) and "applied_guardrails" in metadata_standard: applied_guardrails = metadata_standard.get("applied_guardrails", []) - elif ( - isinstance(metadata_litellm, dict) - and "applied_guardrails" in metadata_litellm - ): + elif isinstance(metadata_litellm, dict) and "applied_guardrails" in metadata_litellm: applied_guardrails = metadata_litellm.get("applied_guardrails", []) if not isinstance(applied_guardrails, list): applied_guardrails = [] for guardrail_name in guardrails_in_metadata: - if ( - isinstance(guardrail_name, str) - and guardrail_name not in applied_guardrails - ): - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=guardrail_name - ) + if isinstance(guardrail_name, str) and guardrail_name not in applied_guardrails: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( self, @@ -1378,18 +1273,12 @@ class ProxyLogging: if result.terminal_action == "block": original_exception = result.original_exception - if original_exception is not None and not _exception_changes_request_flow( - original_exception - ): + if original_exception is not None and not _exception_changes_request_flow(original_exception): blocking_step = result.step_results[-1] if result.step_results else None if blocking_step is not None: - callback = PipelineExecutor.find_guardrail_callback( - blocking_step.guardrail_name - ) + callback = PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) if callback is not None: - _enrich_http_exception_with_guardrail_context( - original_exception, callback - ) + _enrich_http_exception_with_guardrail_context(original_exception, callback) raise original_exception step_results_serializable = [ @@ -1414,8 +1303,7 @@ class ProxyLogging: if result.terminal_action == "modify_response": raise ModifyResponseException( - message=result.modify_response_message - or "Response modified by pipeline", + message=result.modify_response_message or "Response modified by pipeline", model=data.get("model", "unknown"), request_data=data, guardrail_name=f"pipeline:{policy_name}", @@ -1464,9 +1352,7 @@ class ProxyLogging: if data is None: return None - litellm_logging_obj = cast( - Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None) - ) + litellm_logging_obj = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None)) prompt_id = data.get("prompt_id", None) prompt_version = data.get("prompt_version", None) @@ -1515,10 +1401,7 @@ class ProxyLogging: try: if isinstance(_callback, CustomGuardrail) and data is not None: # Skip guardrails managed by a pipeline - if ( - _callback.guardrail_name - and _callback.guardrail_name in pipeline_managed - ): + if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed: continue result = await self._process_guardrail_callback( @@ -1536,8 +1419,7 @@ class ProxyLogging: _callback is not None and isinstance(_callback, CustomLogger) and "async_pre_call_hook" in vars(_callback.__class__) - and _callback.__class__.async_pre_call_hook - != CustomLogger.async_pre_call_hook + and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook ): if call_type == "call_mcp_tool" and user_api_key_dict is None: continue @@ -1576,18 +1458,14 @@ class ProxyLogging: ) if deferred_route_exc is not None and data is not None: - data = await self._handle_sensitive_data_route_exception( - deferred_route_exc, data, user_api_key_dict - ) + data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict) if data is not None: self._process_guardrail_metadata(data) return data except SensitiveDataRouteException as e: - data = await self._handle_sensitive_data_route_exception( - e, data, user_api_key_dict - ) + data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) if data is not None: self._process_guardrail_metadata(data) return data @@ -1665,18 +1543,14 @@ class ProxyLogging: break @staticmethod - async def _run_guardrail_with_metrics( - callback: Any, coro: Awaitable[Any], hook_type: str - ) -> Any: + async def _run_guardrail_with_metrics(callback: Any, coro: Awaitable[Any], hook_type: str) -> Any: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and enriching any raised HTTPException with the originating callback's `guardrail_name`/`guardrail_mode` before re-raising. """ - guardrail_name = ( - getattr(callback, "guardrail_name", None) or type(callback).__name__ - ) + guardrail_name = getattr(callback, "guardrail_name", None) or type(callback).__name__ start_time = time.perf_counter() status = "success" error_type: Optional[str] = None @@ -1721,9 +1595,7 @@ class ProxyLogging: # Cache for callback-capability detection. Keyed on a signature of # litellm.callbacks (length + each item's id) so we recompute when the # callback list mutates (add/remove) without iterating every request. - _callback_capabilities_cache: ClassVar[ - Dict[Tuple[int, Tuple[int, ...]], "_CallbackCapabilities"] - ] = {} + _callback_capabilities_cache: ClassVar[Dict[Tuple[int, Tuple[int, ...]], "_CallbackCapabilities"]] = {} @staticmethod def _callback_capabilities() -> "_CallbackCapabilities": @@ -1793,9 +1665,9 @@ class ProxyLogging: "async_post_call_streaming_hook", base_streaming_hook, ) - if getattr( - cls_streaming_hook, "__func__", cls_streaming_hook - ) is not getattr(base_streaming_hook, "__func__", base_streaming_hook): + if getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr( + base_streaming_hook, "__func__", base_streaming_hook + ): has_streaming_chunk_override = True if "async_pre_call_hook" in cls_attrs: has_pre_call_override = True @@ -1824,11 +1696,7 @@ class ProxyLogging: @staticmethod def has_streaming_callbacks() -> bool: caps = ProxyLogging._callback_capabilities() - return ( - caps.has_iterator_override - or caps.has_streaming_chunk_override - or caps.has_guardrail - ) + return caps.has_iterator_override or caps.has_streaming_chunk_override or caps.has_guardrail @staticmethod def has_streaming_chunk_hook_overrides() -> bool: @@ -1881,9 +1749,7 @@ class ProxyLogging: ################################################################ # V1 implementation - backwards compatibility - if callback.event_hook is None and hasattr( - callback, "moderation_check" - ): + if callback.event_hook is None and hasattr(callback, "moderation_check"): if callback.moderation_check == "pre_call": # type: ignore return else: @@ -1894,16 +1760,11 @@ class ProxyLogging: if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.during_mcp_call - if ( - callback.should_run_guardrail(data=data, event_type=event_type) - is not True - ): + if callback.should_run_guardrail(data=data, event_type=event_type) is not True: continue # Convert user_api_key_dict to proper format for async_moderation_hook if call_type == CallTypes.call_mcp_tool.value: - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( - user_api_key_dict - ) + user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict) else: user_api_key_auth_dict = user_api_key_dict # Add task to list for parallel execution @@ -1978,9 +1839,7 @@ class ProxyLogging: # Note: user_info is a CallInfo that can represent user/team/org level info. For team budgets, # alert_emails is populated from team_object.metadata.soft_budget_alerting_emails (see auth_checks.py) is_soft_budget_with_alert_emails = ( - type == "soft_budget" - and user_info.alert_emails is not None - and len(user_info.alert_emails) > 0 + type == "soft_budget" and user_info.alert_emails is not None and len(user_info.alert_emails) > 0 ) if self.alerting is None and not is_soft_budget_with_alert_emails: @@ -1997,9 +1856,7 @@ class ProxyLogging: # Call email_logging_instance if: # 1. "email" is in alerting config, OR # 2. It's a soft_budget alert with team-specific alert_emails (bypasses global alerting config) - should_send_email = ( - self.alerting is not None and "email" in self.alerting - ) or is_soft_budget_with_alert_emails + should_send_email = (self.alerting is not None and "email" in self.alerting) or is_soft_budget_with_alert_emails if should_send_email and self.email_logging_instance is not None: await self.email_logging_instance.budget_alerts( @@ -2036,9 +1893,7 @@ class ProxyLogging: # Get the current timestamp current_time = datetime.now().strftime("%H:%M:%S") _proxy_base_url = os.getenv("PROXY_BASE_URL", None) - formatted_message = ( - f"Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" - ) + formatted_message = f"Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" if _proxy_base_url is not None: formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`" @@ -2072,9 +1927,7 @@ class ProxyLogging: else: raise Exception("Missing SENTRY_DSN from environment") - async def failure_handler( - self, original_exception, duration: float, call_type: str, traceback_str="" - ): + async def failure_handler(self, original_exception, duration: float, call_type: str, traceback_str=""): """ Log failed db read/writes @@ -2147,9 +2000,7 @@ class ProxyLogging: """ ### ALERTING ### - await self.update_request_status( - litellm_call_id=request_data.get("litellm_call_id", ""), status="fail" - ) + await self.update_request_status(litellm_call_id=request_data.get("litellm_call_id", ""), status="fail") if AlertType.llm_exceptions in self.alert_types and not isinstance( original_exception, (HTTPException, ProxyException) ): @@ -2229,10 +2080,7 @@ class ProxyLogging: traceback_str=traceback_str, ) # If callback returned an HTTPException, use it (first one wins) - if ( - isinstance(hook_result, HTTPException) - and transformed_exception is None - ): + if isinstance(hook_result, HTTPException) and transformed_exception is None: transformed_exception = hook_result except HTTPException as e: # If callback raised an HTTPException, use it (first one wins) @@ -2244,9 +2092,7 @@ class ProxyLogging: f"[Non-Blocking] Error in async_post_call_failure_hook callback: {e}" ) except Exception as e: - verbose_proxy_logger.exception( - f"[Non-Blocking] Error setting up post_call_failure_hook callback: {e}" - ) + verbose_proxy_logger.exception(f"[Non-Blocking] Error setting up post_call_failure_hook callback: {e}") return transformed_exception @@ -2276,9 +2122,7 @@ class ProxyLogging: ######################################################### if route is None: return False - if not ( - RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route) - ): + if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False return isinstance(original_exception, (HTTPException, ProxyException)) or ( @@ -2297,17 +2141,13 @@ class ProxyLogging: Is triggered when self._is_proxy_only_error() returns True """ - litellm_logging_obj: Optional[Logging] = request_data.get( - "litellm_logging_obj", None - ) + litellm_logging_obj: Optional[Logging] = request_data.get("litellm_logging_obj", None) if litellm_logging_obj is None: from litellm._uuid import uuid request_data["litellm_call_id"] = str(uuid.uuid4()) - user_api_key_logged_metadata = ( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + user_api_key_logged_metadata = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_dict ) litellm_logging_obj, data = litellm.utils.function_setup( @@ -2341,9 +2181,7 @@ class ProxyLogging: input: Union[list, str, dict] = "" normalized_call_type: Optional[str] = None - if "messages" in request_data and isinstance( - request_data["messages"], list - ): + if "messages" in request_data and isinstance(request_data["messages"], list): input = request_data["messages"] litellm_logging_obj.model_call_details["messages"] = input if litellm_logging_obj.call_type != CallTypes.pass_through.value: @@ -2360,9 +2198,7 @@ class ProxyLogging: normalized_call_type = CallTypes.aembedding.value if normalized_call_type is not None: litellm_logging_obj.call_type = normalized_call_type - litellm_logging_obj.model_call_details["call_type"] = ( - normalized_call_type - ) + litellm_logging_obj.model_call_details["call_type"] = normalized_call_type # Pass-through endpoints are logged via the callback loop's # async_post_call_failure_hook — skip pre_call and failure handlers. if litellm_logging_obj.call_type == CallTypes.pass_through.value: @@ -2373,9 +2209,7 @@ class ProxyLogging: # fabricate an LLM-call span for a call that did not happen (and, since # this runs inside the live ``auth`` phase span, would otherwise nest it # under auth). The marker tells them to skip span creation. - litellm_logging_obj.model_call_details[ - LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL - ] = True + litellm_logging_obj.model_call_details[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True litellm_logging_obj.pre_call( input=input, api_key="", @@ -2436,9 +2270,7 @@ class ProxyLogging: ############################################################################# # Merge model-level guardrails before checking which guardrails to run - guardrail_data = _check_and_merge_model_level_guardrails( - data=data, llm_router=llm_router - ) + guardrail_data = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) for callback in guardrail_callbacks: # Main - V2 Guardrails implementation @@ -2517,9 +2349,7 @@ class ProxyLogging: try: # Build litellm_call_info — normalized routing metadata for callbacks - litellm_call_info = self._build_litellm_call_info( - data=data, response=response - ) + litellm_call_info = self._build_litellm_call_info(data=data, response=response) for callback in litellm.callbacks: _callback: Optional[CustomLogger] = None @@ -2550,9 +2380,7 @@ class ProxyLogging: if result is not None: merged_headers.update(result) except Exception as e: - verbose_proxy_logger.exception( - "Error in post_call_response_headers_hook: %s", str(e) - ) + verbose_proxy_logger.exception("Error in post_call_response_headers_hook: %s", str(e)) return merged_headers @staticmethod @@ -2585,9 +2413,7 @@ class ProxyLogging: async def async_post_call_streaming_hook( self, data: dict, - response: Union[ - ModelResponse, EmbeddingResponse, ImageResponse, ModelResponseStream - ], + response: Union[ModelResponse, EmbeddingResponse, ImageResponse, ModelResponseStream], user_api_key_dict: UserAPIKeyAuth, str_so_far: Optional[str] = None, ): @@ -2631,10 +2457,8 @@ class ProxyLogging: ## CHECK FOR MODEL-LEVEL GUARDRAILS (cached per-request) if not _guardrail_data_computed: - _cached_guardrail_data = ( - _check_and_merge_model_level_guardrails( - data=data, llm_router=llm_router - ) + _cached_guardrail_data = _check_and_merge_model_level_guardrails( + data=data, llm_router=llm_router ) _guardrail_data_computed = True @@ -2657,11 +2481,9 @@ class ProxyLogging: complete_response = str_so_far + response_str else: complete_response = response_str - callback_response = ( - await _callback.async_post_call_streaming_hook( - user_api_key_dict=user_api_key_dict, - response=complete_response, - ) + callback_response = await _callback.async_post_call_streaming_hook( + user_api_key_dict=user_api_key_dict, + response=complete_response, ) if callback_response is not None: response = callback_response @@ -2697,18 +2519,14 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router # Merge model-level guardrails before checking which guardrails to run - request_data = _check_and_merge_model_level_guardrails( - data=request_data, llm_router=llm_router - ) + request_data = _check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router) current_response = response for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): if ( - resolved_callback.should_run_guardrail( - data=request_data, event_type=GuardrailEventHooks.post_call - ) + resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True ): continue @@ -2763,9 +2581,7 @@ class ProxyLogging: logging_obj._deferred_stream_complete_args = None asyncio.create_task(_deferred_cb(*_args)) - def _release_max_parallel_requests_on_disconnect( - self, user_api_key_dict: UserAPIKeyAuth - ) -> None: + def _release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None: """ Release the api-key max_parallel_requests slot when a streaming response is cancelled mid-flight (client disconnect). Neither the @@ -2785,11 +2601,7 @@ class ProxyLogging: if not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3): return try: - asyncio.create_task( - limiter.async_release_max_parallel_requests_on_disconnect( - user_api_key_dict - ) - ) + asyncio.create_task(limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)) except RuntimeError: # No running event loop (e.g. interpreter/loop shutdown); the # counter's window TTL will reclaim the slot. @@ -2808,13 +2620,8 @@ class ProxyLogging: This handles checking for if a request is hanging for too long """ ## ALERTING ### - if ( - self.slack_alerting_instance - and self.slack_alerting_instance.alerting is not None - ): - asyncio.create_task( - self.slack_alerting_instance.response_taking_too_long(request_data=data) - ) + if self.slack_alerting_instance and self.slack_alerting_instance.alerting is not None: + asyncio.create_task(self.slack_alerting_instance.response_taking_too_long(request_data=data)) ### DB CONNECTOR ### @@ -2892,9 +2699,7 @@ async def _lookup_deprecated_key( # DualCache for LiteLLM_Config param_name reads. # Redis layer is attached in proxy_server._init_cache. -LITELLM_CONFIG_CACHE_TTL_SECONDS: int = int( - os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60") -) +LITELLM_CONFIG_CACHE_TTL_SECONDS: int = int(os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60")) _CONFIG_CACHE_MISS: str = "__litellm_config_param_miss__" litellm_config_cache: DualCache = DualCache( @@ -2936,13 +2741,9 @@ async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any] if cached is not None: return _unpack_config_row(cached) - row = await prisma_client.get_generic_data( - key="param_name", value=param_name, table_name="config" - ) + row = await prisma_client.get_generic_data(key="param_name", value=param_name, table_name="config") cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS - await litellm_config_cache.async_set_cache( - cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS - ) + await litellm_config_cache.async_set_cache(cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS) return row @@ -2968,9 +2769,7 @@ async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> by_name = {row.param_name: row for row in rows} for name in param_names: row = by_name.get(name) - cache_value: Any = ( - _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS - ) + cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS await litellm_config_cache.async_set_cache( _config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS ) @@ -2988,26 +2787,16 @@ class PrismaClient: ): ## init logging object self.proxy_logging_obj = proxy_logging_obj - self.iam_token_db_auth: Optional[bool] = str_to_bool( - os.getenv("IAM_TOKEN_DB_AUTH") - ) + self.iam_token_db_auth: Optional[bool] = str_to_bool(os.getenv("IAM_TOKEN_DB_AUTH")) verbose_proxy_logger.debug("Creating Prisma Client..") try: from prisma import Prisma # type: ignore except Exception as e: verbose_proxy_logger.error(f"Failed to import Prisma client: {e}") - verbose_proxy_logger.error( - "This usually means 'prisma generate' hasn't been run yet." - ) - verbose_proxy_logger.error( - "Please run 'prisma generate' to generate the Prisma client." - ) - raise Exception( - "Unable to find Prisma binaries. Please run 'prisma generate' first." - ) - iam_flag = ( - self.iam_token_db_auth if self.iam_token_db_auth is not None else False - ) + verbose_proxy_logger.error("This usually means 'prisma generate' hasn't been run yet.") + verbose_proxy_logger.error("Please run 'prisma generate' to generate the Prisma client.") + raise Exception("Unable to find Prisma binaries. Please run 'prisma generate' first.") + iam_flag = self.iam_token_db_auth if self.iam_token_db_auth is not None else False # When read-replica routing is on, tag log lines with [writer]/[reader] # so the two wrappers' interleaved IAM refresh logs can be told apart. # Single-DB deployments get an empty prefix (logs unchanged). @@ -3038,9 +2827,7 @@ class PrismaClient: # the same cadence as the writer. We parse the static endpoint # pieces (host/port/user/db) once from the reader URL — only # the IAM token rotates after that. - reader_iam_endpoint = ( - parse_iam_endpoint_from_url(read_replica_url) if iam_flag else None - ) + reader_iam_endpoint = parse_iam_endpoint_from_url(read_replica_url) if iam_flag else None # Mint a fresh IAM token for the reader BEFORE constructing the # Prisma client. Mirrors what `proxy_cli.py` already does for # the writer (proxy_cli.py:812-832) — without this, the reader @@ -3061,9 +2848,7 @@ class PrismaClient: ) read_replica_url = reader_iam_endpoint.build_url(reader_token) os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url - reader_kwargs: Dict[str, Any] = { - "datasource": {"url": read_replica_url} - } + reader_kwargs: Dict[str, Any] = {"datasource": {"url": read_replica_url}} if http_client is not None: reader_prisma = Prisma(http=http_client, **reader_kwargs) else: @@ -3076,9 +2861,7 @@ class PrismaClient: recreate_uses_datasource=True, log_prefix="[reader]", ) - self.db = RoutingPrismaWrapper( - writer=writer_wrapper, reader=reader_wrapper - ) + self.db = RoutingPrismaWrapper(writer=writer_wrapper, reader=reader_wrapper) verbose_proxy_logger.info( "PrismaClient: read-replica routing enabled via DATABASE_URL_READ_REPLICA" + (" (with IAM token auto-refresh)" if iam_flag else "") @@ -3103,9 +2886,7 @@ class PrismaClient: self._db_reconnect_lock = asyncio.Lock() self._db_health_watchdog_task: Optional[asyncio.Task] = None self._db_last_reconnect_attempt_ts: float = 0.0 - self._db_reconnect_cooldown_seconds: int = max( - 1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15")) - ) + self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15"))) self._db_health_watchdog_interval_seconds: int = max( 5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "30")) ) @@ -3127,9 +2908,7 @@ class PrismaClient: float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")), ) self._consecutive_reconnect_failures: int = 0 - self._reconnect_escalation_threshold: int = max( - 1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3")) - ) + self._reconnect_escalation_threshold: int = max(1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3"))) self._engine_pidfd: int = -1 self._engine_pid: int = 0 self._watching_engine: bool = False @@ -3144,9 +2923,7 @@ class PrismaClient: return self.db.writer return self.db - def get_request_status( - self, payload: Union[dict, SpendLogsPayload] - ) -> Literal["success", "failure"]: + def get_request_status(self, payload: Union[dict, SpendLogsPayload]) -> Literal["success", "failure"]: """ Determine if a request was successful or failed based on payload metadata. @@ -3158,22 +2935,14 @@ class PrismaClient: """ try: # Get metadata and convert to dict if it's a JSON string - payload_metadata: Union[Dict, SpendLogsMetadata, str] = payload.get( - "metadata", {} - ) + payload_metadata: Union[Dict, SpendLogsMetadata, str] = payload.get("metadata", {}) if isinstance(payload_metadata, str): - payload_metadata_json: Union[Dict, SpendLogsMetadata] = cast( - Dict, json.loads(payload_metadata) - ) + payload_metadata_json: Union[Dict, SpendLogsMetadata] = cast(Dict, json.loads(payload_metadata)) else: payload_metadata_json = payload_metadata # Check status in metadata dict - return ( - "failure" - if payload_metadata_json.get("status") == "failure" - else "success" - ) + return "failure" if payload_metadata_json.get("status") == "failure" else "success" except (json.JSONDecodeError, AttributeError): # Default to success if metadata parsing fails @@ -3265,9 +3034,7 @@ class PrismaClient: LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; """) - verbose_proxy_logger.info( - "LiteLLM_VerificationTokenView Created in DB!" - ) + verbose_proxy_logger.info("LiteLLM_VerificationTokenView Created in DB!") else: should_create_views = await should_create_missing_views(db=self.db) if should_create_views: @@ -3275,9 +3042,7 @@ class PrismaClient: else: # don't block execution if these views are missing # Convert lists to sets for efficient difference calculation - ret_view_names_set = ( - set(ret[0]["view_names"]) if ret[0]["view_names"] else set() - ) + ret_view_names_set = set(ret[0]["view_names"]) if ret[0]["view_names"] else set() expected_views_set = set(expected_views) # Find missing views missing_views = expected_views_set - ret_view_names_set @@ -3360,9 +3125,7 @@ class PrismaClient: raise e - async def _query_first_with_cached_plan_fallback( - self, sql_query: str, *args - ) -> Optional[dict]: + async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> Optional[dict]: """ Execute a query, recovering once from PostgreSQL's "cached plan must not change result type" error. @@ -3441,9 +3204,7 @@ class PrismaClient: expires: Optional[datetime] = None, reset_at: Optional[datetime] = None, offset: Optional[int] = None, # pagination, what row number to start from - limit: Optional[ - int - ] = None, # pagination, number of rows to getch when find_all==True + limit: Optional[int] = None, # pagination, number of rows to getch when find_all==True parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, budget_id_list: Optional[List[str]] = None, @@ -3454,33 +3215,25 @@ class PrismaClient: hashed_token: Optional[str] = None try: response: Any = None - if (token is not None and table_name is None) or ( - table_name is not None and table_name == "key" - ): + if (token is not None and table_name is None) or (table_name is not None and table_name == "key"): # check if plain text or hash if token is not None: if isinstance(token, str): hashed_token = _hash_token_if_needed(token=token) - verbose_proxy_logger.debug( - f"PrismaClient: find_unique for token: {hashed_token}" - ) + verbose_proxy_logger.debug(f"PrismaClient: find_unique for token: {hashed_token}") if query_type == "find_unique" and hashed_token is not None: if token is None: raise HTTPException( status_code=400, detail={"error": f"No token passed in. Token={token}"}, ) - response = await VerificationTokenRepository( - self - ).table.find_unique( + response = await VerificationTokenRepository(self).table.find_unique( where={"token": hashed_token}, # type: ignore include={"litellm_budget_table": True}, ) if response is not None: # for prisma we need to cast the expires time to str - if response.expires is not None and isinstance( - response.expires, datetime - ): + if response.expires is not None and isinstance(response.expires, datetime): response.expires = response.expires.isoformat() else: # Token does not exist. @@ -3507,11 +3260,7 @@ class PrismaClient: for r in response: if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() - elif ( - query_type == "find_all" - and expires is not None - and reset_at is not None - ): + elif query_type == "find_all" and expires is not None and reset_at is not None: response = await VerificationTokenRepository(self).table.find_many( where={ # type: ignore "OR": [ @@ -3555,9 +3304,7 @@ class PrismaClient: status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication Error: invalid user key - token does not exist", ) - elif (user_id is not None and table_name is None) or ( - table_name is not None and table_name == "user" - ): + elif (user_id is not None and table_name is None) or (table_name is not None and table_name == "user"): if query_type == "find_unique": if key_val is None: key_val = {"user_id": user_id} @@ -3578,9 +3325,7 @@ class PrismaClient: } ) elif query_type == "find_all" and user_id_list is not None: - response = await UserRepository(self).table.find_many( - where={"user_id": {"in": user_id_list}} - ) + response = await UserRepository(self).table.find_many(where={"user_id": {"in": user_id_list}}) elif query_type == "find_all": if expires is not None: response = await UserRepository(self).table.find_many( # type: ignore @@ -3610,9 +3355,7 @@ class PrismaClient: response = await self.db.query_raw(sql_query, limit, offset) return response elif table_name == "spend": - verbose_proxy_logger.debug( - "PrismaClient: get_data: table_name == 'spend'" - ) + verbose_proxy_logger.debug("PrismaClient: get_data: table_name == 'spend'") if key_val is not None: if query_type == "find_unique": response = await SpendLogsRepository(self).table.find_unique( # type: ignore @@ -3675,19 +3418,13 @@ class PrismaClient: include={"litellm_budget_table": True}, ) elif query_type == "find_all" and team_id_list is not None: - response = await TeamRepository(self).table.find_many( - where={"team_id": {"in": team_id_list}} - ) + response = await TeamRepository(self).table.find_many(where={"team_id": {"in": team_id_list}}) elif query_type == "find_all" and team_id_list is None: - response = await TeamRepository(self).table.find_many( - take=MAX_TEAM_LIST_LIMIT - ) + response = await TeamRepository(self).table.find_many(take=MAX_TEAM_LIST_LIMIT) return response elif table_name == "user_notification": if query_type == "find_unique": - response = await UserNotificationsRepository( - self - ).table.find_unique( # type: ignore + response = await UserNotificationsRepository(self).table.find_unique( # type: ignore where={"user_id": user_id} # type: ignore ) elif query_type == "find_all": @@ -3698,9 +3435,7 @@ class PrismaClient: if token is not None: if isinstance(token, str): hashed_token = _hash_token_if_needed(token=token) - verbose_proxy_logger.debug( - f"PrismaClient: find_unique for token: {hashed_token}" - ) + verbose_proxy_logger.debug(f"PrismaClient: find_unique for token: {hashed_token}") if query_type == "find_unique": if token is None: raise HTTPException( @@ -3752,20 +3487,12 @@ class PrismaClient: WHERE v.token = $1 """ - response = await self._query_first_with_cached_plan_fallback( - sql_query, hashed_token - ) + response = await self._query_first_with_cached_plan_fallback(sql_query, hashed_token) # If not found in main table, check deprecated keys (grace period) # check_deprecated=False on the recursive call prevents unbounded chaining - if ( - response is None - and hashed_token is not None - and check_deprecated - ): - active_token_id = await _lookup_deprecated_key( - db=self.db, hashed_token=hashed_token - ) + if response is None and hashed_token is not None and check_deprecated: + active_token_id = await _lookup_deprecated_key(db=self.db, hashed_token=hashed_token) if active_token_id: # The recursive call returns a finished # LiteLLM_VerificationTokenView; the dict @@ -3779,9 +3506,7 @@ class PrismaClient: check_deprecated=False, ) if deprecated_response is not None: - verbose_proxy_logger.debug( - "Deprecated key used during grace period" - ) + verbose_proxy_logger.debug("Deprecated key used during grace period") return deprecated_response if response is not None: @@ -3791,10 +3516,7 @@ class PrismaClient: response["team_blocked"] = False team_member: Optional[Member] = None - if ( - response["team_members_with_roles"] is not None - and response["user_id"] is not None - ): + if response["team_members_with_roles"] is not None and response["user_id"] is not None: ## find the team member corresponding to user id """ [ @@ -3811,24 +3533,20 @@ class PrismaClient: ] """ for tm in response["team_members_with_roles"]: - if tm.get("user_id") is not None and response[ - "user_id" - ] == tm.get("user_id"): + if tm.get("user_id") is not None and response["user_id"] == tm.get("user_id"): team_member = Member(**tm) response["team_member"] = team_member - response = LiteLLM_VerificationTokenView( - **response, last_refreshed_at=time.time() - ) + response = LiteLLM_VerificationTokenView(**response, last_refreshed_at=time.time()) # for prisma we need to cast the expires time to str - if response.expires is not None and isinstance( - response.expires, datetime - ): + if response.expires is not None and isinstance(response.expires, datetime): response.expires = response.expires.isoformat() return response except Exception as e: import traceback - prisma_query_info = f"LiteLLM Prisma Client Exception: Error with `get_data`. Args passed in: {args_passed_in}" + prisma_query_info = ( + f"LiteLLM Prisma Client Exception: Error with `get_data`. Args passed in: {args_passed_in}" + ) error_msg = prisma_query_info + str(e) print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() @@ -3848,13 +3566,9 @@ class PrismaClient: def jsonify_team_object(self, db_data: dict): db_data = self.jsonify_object(data=db_data) - if db_data.get("members_with_roles", None) is not None and isinstance( - db_data["members_with_roles"], list - ): + if db_data.get("members_with_roles", None) is not None and isinstance(db_data["members_with_roles"], list): db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"]) - if db_data.get("budget_limits", None) is not None and isinstance( - db_data["budget_limits"], list - ): + if db_data.get("budget_limits", None) is not None and isinstance(db_data["budget_limits"], list): db_data["budget_limits"] = json.dumps(db_data["budget_limits"]) return db_data @@ -3869,9 +3583,7 @@ class PrismaClient: async def insert_data( self, data: dict, - table_name: Literal[ - "user", "key", "config", "spend", "team", "user_notification" - ], + table_name: Literal["user", "key", "config", "spend", "team", "user_notification"], ): """ Add a key to the database. If it already exists, do nothing. @@ -3888,12 +3600,8 @@ class PrismaClient: # Strip them so the DB stores NULL via the column's nullable constraint. if db_data.get("budget_limits") is None: db_data.pop("budget_limits", None) - print_verbose( - "PrismaClient: Before upsert into litellm_verificationtoken" - ) - new_verification_token = await VerificationTokenRepository( - self - ).table.upsert( # type: ignore + print_verbose("PrismaClient: Before upsert into litellm_verificationtoken") + new_verification_token = await VerificationTokenRepository(self).table.upsert( # type: ignore where={ "token": hashed_token, }, @@ -3980,9 +3688,7 @@ class PrismaClient: return new_spend_row elif table_name == "user_notification": db_data = self.jsonify_object(data=data) - new_user_notification_row = await UserNotificationsRepository( - self - ).table.upsert( # type: ignore + new_user_notification_row = await UserNotificationsRepository(self).table.upsert( # type: ignore where={"request_id": data["request_id"]}, data={ "create": {**db_data}, # type: ignore @@ -4026,18 +3732,14 @@ class PrismaClient: user_id: Optional[str] = None, team_id: Optional[str] = None, query_type: Literal["update", "update_many"] = "update", - table_name: Optional[ - Literal["user", "key", "config", "spend", "team", "enduser", "budget"] - ] = None, + table_name: Optional[Literal["user", "key", "config", "spend", "team", "enduser", "budget"]] = None, update_key_values: Optional[dict] = None, update_key_values_custom_query: Optional[dict] = None, ): """ Update existing data """ - verbose_proxy_logger.debug( - f"PrismaClient: update_data, table_name: {table_name}" - ) + verbose_proxy_logger.debug(f"PrismaClient: update_data, table_name: {table_name}") start_time = time.time() try: db_data = self.jsonify_object(data=data) @@ -4052,11 +3754,7 @@ class PrismaClient: where={"token": token}, # type: ignore data={**db_data}, # type: ignore ) - verbose_proxy_logger.debug( - "\033[91m" - + f"DB Token Table update succeeded {response}" - + "\033[0m" - ) + verbose_proxy_logger.debug("\033[91m" + f"DB Token Table update succeeded {response}" + "\033[0m") _data: dict = {} if response is not None: try: @@ -4064,11 +3762,7 @@ class PrismaClient: except Exception: _data = response.dict() return {"token": token, "data": _data} - elif ( - user_id is not None - or (table_name is not None and table_name == "user") - and query_type == "update" - ): + elif user_id is not None or (table_name is not None and table_name == "user") and query_type == "update": """ If data['spend'] + data['user'], update the user table with spend info as well """ @@ -4089,16 +3783,10 @@ class PrismaClient: }, ) verbose_proxy_logger.info( - "\033[91m" - + f"DB User Table - update succeeded {update_user_row}" - + "\033[0m" + "\033[91m" + f"DB User Table - update succeeded {update_user_row}" + "\033[0m" ) return {"user_id": user_id, "data": update_user_row} - elif ( - team_id is not None - or (table_name is not None and table_name == "team") - and query_type == "update" - ): + elif team_id is not None or (table_name is not None and table_name == "team") and query_type == "update": """ If data['spend'] + data['user'], update the user table with spend info as well """ @@ -4108,18 +3796,12 @@ class PrismaClient: update_key_values = db_data if "team_id" not in db_data and team_id is not None: db_data["team_id"] = team_id - if "members_with_roles" in db_data and isinstance( - db_data["members_with_roles"], list - ): - db_data["members_with_roles"] = json.dumps( - db_data["members_with_roles"] - ) + if "members_with_roles" in db_data and isinstance(db_data["members_with_roles"], list): + db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"]) if "members_with_roles" in update_key_values and isinstance( update_key_values["members_with_roles"], list ): - update_key_values["members_with_roles"] = json.dumps( - update_key_values["members_with_roles"] - ) + update_key_values["members_with_roles"] = json.dumps(update_key_values["members_with_roles"]) update_team_row = await TeamRepository(self).table.upsert( where={"team_id": team_id}, # type: ignore data={ @@ -4130,9 +3812,7 @@ class PrismaClient: }, ) verbose_proxy_logger.info( - "\033[91m" - + f"DB Team Table - update succeeded {update_team_row}" - + "\033[0m" + "\033[91m" + f"DB Team Table - update succeeded {update_team_row}" + "\033[0m" ) return {"team_id": team_id, "data": update_team_row} elif ( @@ -4151,9 +3831,7 @@ class PrismaClient: if t.token.startswith("sk-"): # type: ignore t.token = self.hash_token(token=t.token) # type: ignore try: - data_json = self.jsonify_object( - data=t.model_dump(exclude_none=True) - ) + data_json = self.jsonify_object(data=t.model_dump(exclude_none=True)) except Exception: data_json = self.jsonify_object(data=t.dict(exclude_none=True)) batcher.litellm_verificationtoken.update( @@ -4161,9 +3839,7 @@ class PrismaClient: data={**data_json}, # type: ignore ) await batcher.commit() - print_verbose( - "\033[91m" + "DB Token Table update succeeded" + "\033[0m" - ) + print_verbose("\033[91m" + "DB Token Table update succeeded" + "\033[0m") elif ( table_name is not None and table_name == "user" @@ -4177,9 +3853,7 @@ class PrismaClient: batcher = self.db.batch_() for idx, user in enumerate(data_list): try: - data_json = self.jsonify_object( - data=user.model_dump(exclude_none=True) - ) + data_json = self.jsonify_object(data=user.model_dump(exclude_none=True)) except Exception: data_json = self.jsonify_object(data=user.dict()) batcher.litellm_usertable.upsert( @@ -4192,9 +3866,7 @@ class PrismaClient: }, ) await batcher.commit() - verbose_proxy_logger.info( - "\033[91m" + "DB User Table Batch update succeeded" + "\033[0m" - ) + verbose_proxy_logger.info("\033[91m" + "DB User Table Batch update succeeded" + "\033[0m") elif ( table_name is not None and table_name == "enduser" @@ -4208,9 +3880,7 @@ class PrismaClient: batcher = self.db.batch_() for enduser in data_list: try: - data_json = self.jsonify_object( - data=enduser.model_dump(exclude_none=True) - ) + data_json = self.jsonify_object(data=enduser.model_dump(exclude_none=True)) except Exception: data_json = self.jsonify_object(data=enduser.dict()) batcher.litellm_endusertable.upsert( @@ -4223,9 +3893,7 @@ class PrismaClient: }, ) await batcher.commit() - verbose_proxy_logger.info( - "\033[91m" + "DB End User Table Batch update succeeded" + "\033[0m" - ) + verbose_proxy_logger.info("\033[91m" + "DB End User Table Batch update succeeded" + "\033[0m") elif ( table_name is not None and table_name == "budget" @@ -4239,9 +3907,7 @@ class PrismaClient: batcher = self.db.batch_() for budget in data_list: try: - data_json = self.jsonify_object( - data=budget.model_dump(exclude_none=True) - ) + data_json = self.jsonify_object(data=budget.model_dump(exclude_none=True)) except Exception: data_json = self.jsonify_object(data=budget.dict()) batcher.litellm_budgettable.upsert( @@ -4254,9 +3920,7 @@ class PrismaClient: }, ) await batcher.commit() - verbose_proxy_logger.info( - "\033[91m" + "DB Budget Table Batch update succeeded" + "\033[0m" - ) + verbose_proxy_logger.info("\033[91m" + "DB Budget Table Batch update succeeded" + "\033[0m") elif ( table_name is not None and table_name == "team" @@ -4268,13 +3932,9 @@ class PrismaClient: batcher = self.db.batch_() for idx, team in enumerate(data_list): try: - data_json = self.jsonify_team_object( - db_data=team.model_dump(exclude_none=True) - ) + data_json = self.jsonify_team_object(db_data=team.model_dump(exclude_none=True)) except Exception: - data_json = self.jsonify_object( - data=team.dict(exclude_none=True) - ) + data_json = self.jsonify_object(data=team.dict(exclude_none=True)) batcher.litellm_teamtable.upsert( where={"team_id": team.team_id}, # type: ignore data={ @@ -4285,9 +3945,7 @@ class PrismaClient: }, ) await batcher.commit() - verbose_proxy_logger.info( - "\033[91m" + "DB Team Table Batch update succeeded" + "\033[0m" - ) + verbose_proxy_logger.info("\033[91m" + "DB Team Table Batch update succeeded" + "\033[0m") except Exception as e: import traceback @@ -4339,38 +3997,22 @@ class PrismaClient: hashed_tokens.append(hashed_token) filter_query: dict = {} if user_id is not None: - filter_query = { - "AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}] - } + filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]} else: filter_query = {"token": {"in": hashed_tokens}} - deleted_tokens = await VerificationTokenRepository( - self - ).table.delete_many( + deleted_tokens = await VerificationTokenRepository(self).table.delete_many( where=filter_query # type: ignore ) verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens) return {"deleted_keys": deleted_tokens} - elif ( - table_name == "team" - and team_id_list is not None - and isinstance(team_id_list, List) - ): + elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, List): # admin only endpoint -> `/team/delete` - await TeamRepository(self).table.delete_many( - where={"team_id": {"in": team_id_list}} - ) + await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) return {"deleted_teams": team_id_list} - elif ( - table_name == "key" - and team_id_list is not None - and isinstance(team_id_list, List) - ): + elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, List): # admin only endpoint -> `/team/delete` - await VerificationTokenRepository(self).table.delete_many( - where={"team_id": {"in": team_id_list}} - ) + await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) except Exception as e: import traceback @@ -4400,13 +4042,9 @@ class PrismaClient: async def connect(self): start_time = time.time() try: - verbose_proxy_logger.debug( - "PrismaClient: connect() called Attempting to Connect to DB" - ) + verbose_proxy_logger.debug("PrismaClient: connect() called Attempting to Connect to DB") if self.db.is_connected() is False: - verbose_proxy_logger.debug( - "PrismaClient: DB not connected, Attempting to Connect to DB" - ) + verbose_proxy_logger.debug("PrismaClient: DB not connected, Attempting to Connect to DB") await self.db.connect() except Exception as e: import traceback @@ -4530,8 +4168,7 @@ class PrismaClient: ) if self._consume_expected_death(pid): verbose_proxy_logger.info( - "PID %s death was planned (engine already replaced); " - "not reconnecting.", + "PID %s death was planned (engine already replaced); not reconnecting.", pid, ) self._cleanup_engine_watcher() @@ -4755,17 +4392,11 @@ class PrismaClient: waitpid thread nor pidfd are available. """ - if ( - self._watching_engine - or self._engine_pidfd >= 0 - or self._engine_wait_thread is not None - ): + if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: return pid = self._get_engine_pid() if pid == 0: - verbose_proxy_logger.debug( - "Could not find prisma-query-engine PID; engine death detection unavailable." - ) + verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") return self._engine_pid = pid self._engine_confirmed_dead = False @@ -4812,9 +4443,7 @@ class PrismaClient: self._cleanup_engine_watcher() asyncio.create_task(self._start_engine_watcher()) - async def _run_reconnect_cycle( - self, timeout_seconds: Optional[float] = None - ) -> None: + async def _run_reconnect_cycle(self, timeout_seconds: Optional[float] = None) -> None: """ Run a reconnect cycle with a single overall timeout budget. @@ -4827,9 +4456,7 @@ class PrismaClient: subprocess.Popen.wait() inside prisma-client-py (see issue #26191). """ effective_timeout = ( - timeout_seconds - if timeout_seconds is not None - else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds ) # Snapshot the writer's engine generation BEFORE any await. Both @@ -4843,9 +4470,7 @@ class PrismaClient: # otherwise slip in and bump the very generation the closure then reads. expected_generation = getattr(self.writer_db, "_engine_generation", None) - engine_is_dead = self._engine_confirmed_dead or ( - self._engine_pid > 0 and not self._is_engine_alive() - ) + engine_is_dead = self._engine_confirmed_dead or (self._engine_pid > 0 and not self._is_engine_alive()) if engine_is_dead: dead_pid = self._engine_pid @@ -4859,9 +4484,7 @@ class PrismaClient: async def _do_heavy_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") if not db_url: - verbose_proxy_logger.error( - "DATABASE_URL not set; cannot recreate Prisma client." - ) + verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") raise RuntimeError("DATABASE_URL not set") # Forward the entry-snapshot generation. The engine was # confirmed dead, but a concurrent IAM refresh may have already @@ -4870,9 +4493,7 @@ class PrismaClient: # direct path there is no SELECT 1 probe here, so the generation # guard is the only thing standing between a crash-reconnect and # a refresh that raced it. - await self.db.recreate_prisma_client( - db_url, expected_generation=expected_generation - ) + await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation) await self._start_engine_watcher() await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) @@ -4883,16 +4504,12 @@ class PrismaClient: # demoting to the lightweight path. self._engine_confirmed_dead = False else: - verbose_proxy_logger.debug( - "Performing Prisma DB reconnect (engine alive or unknown)." - ) + verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).") async def _do_direct_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") if not db_url: - verbose_proxy_logger.error( - "DATABASE_URL not set; cannot reconnect Prisma client." - ) + verbose_proxy_logger.error("DATABASE_URL not set; cannot reconnect Prisma client.") raise RuntimeError("DATABASE_URL not set") # Probe the writer BEFORE recreating. A concurrent IAM token # refresh may have just replaced the engine (issue #29176); if @@ -4921,9 +4538,7 @@ class PrismaClient: # ends up killing the engine anyway, we do it non-blockingly # via `_kill_engine_process` inside `recreate_prisma_client`. self._cleanup_engine_watcher() - await self.db.recreate_prisma_client( - db_url, expected_generation=expected_generation - ) + await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation) await self._start_engine_watcher() # Smoke-test the writer specifically; query_raw on the routing # wrapper sends to the reader, which would not validate the @@ -4939,11 +4554,7 @@ class PrismaClient: timeout_seconds: Optional[float], ) -> bool: now = time.time() - if ( - force is False - and now - self._db_last_reconnect_attempt_ts - < self._db_reconnect_cooldown_seconds - ): + if force is False and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds: verbose_proxy_logger.debug( "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", reason, @@ -4963,18 +4574,14 @@ class PrismaClient: ) self._engine_confirmed_dead = True - verbose_proxy_logger.warning( - "Attempting Prisma DB reconnect. reason=%s", reason - ) + verbose_proxy_logger.warning("Attempting Prisma DB reconnect. reason=%s", reason) reconnect_succeeded = False try: await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) reconnect_succeeded = True self._consecutive_reconnect_failures = 0 - verbose_proxy_logger.info( - "Prisma DB reconnect succeeded. reason=%s", reason - ) + verbose_proxy_logger.info("Prisma DB reconnect succeeded. reason=%s", reason) except Exception as reconnect_err: self._consecutive_reconnect_failures += 1 verbose_proxy_logger.error( @@ -5002,11 +4609,7 @@ class PrismaClient: bool: True if reconnection succeeded, else False. """ now = time.time() - if ( - force is False - and now - self._db_last_reconnect_attempt_ts - < self._db_reconnect_cooldown_seconds - ): + if force is False and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds: verbose_proxy_logger.debug( "Skipping DB reconnect attempt due to cooldown. reason=%s", reason, @@ -5015,9 +4618,7 @@ class PrismaClient: if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await self._attempt_reconnect_inside_lock( - force, reason, timeout_seconds - ) + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) lock_acquired_by_timeout_task = False @@ -5066,9 +4667,7 @@ class PrismaClient: return False try: - return await self._attempt_reconnect_inside_lock( - force, reason, timeout_seconds - ) + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) finally: self._db_reconnect_lock.release() @@ -5078,9 +4677,7 @@ class PrismaClient: - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. """ if self._db_health_watchdog_enabled is not True: - verbose_proxy_logger.debug( - "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" - ) + verbose_proxy_logger.debug("Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED") return if self._db_health_watchdog_task is not None: return @@ -5089,9 +4686,7 @@ class PrismaClient: # mistaken for a crash (issue #29176). Set on the writer wrapper since # the watcher tracks the writer engine. self.writer_db.on_engine_replaced = self._handle_writer_engine_replaced - self._db_health_watchdog_task = asyncio.create_task( - self._db_health_watchdog_loop() - ) + self._db_health_watchdog_task = asyncio.create_task(self._db_health_watchdog_loop()) verbose_proxy_logger.info( "Started Prisma DB health watchdog (interval=%ss, reconnect_cooldown=%ss, probe_timeout=%ss, reconnect_timeout=%ss)", self._db_health_watchdog_interval_seconds, @@ -5125,17 +4720,13 @@ class PrismaClient: except asyncio.CancelledError: break except Exception as e: - if isinstance( - e, asyncio.TimeoutError - ) or PrismaDBExceptionHandler.is_database_connection_error(e): + if isinstance(e, asyncio.TimeoutError) or PrismaDBExceptionHandler.is_database_connection_error(e): await self.attempt_db_reconnect( reason="db_health_watchdog_connection_error", timeout_seconds=self._db_watchdog_reconnect_timeout_seconds, ) else: - verbose_proxy_logger.debug( - "Prisma DB health watchdog observed non-DB error: %s", e - ) + verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e) @backoff.on_exception( backoff.expo, @@ -5198,9 +4789,7 @@ class PrismaClient: try: return await _fetch_row_count() except Exception as e: - verbose_proxy_logger.error( - f"Error getting LiteLLM_SpendLogs row count: {e}" - ) + verbose_proxy_logger.error(f"Error getting LiteLLM_SpendLogs row count: {e}") return 0 @backoff.on_exception( @@ -5225,23 +4814,15 @@ class PrismaClient: ) # Health Check Database Methods - def _validate_response_time( - self, response_time_ms: Optional[float] - ) -> Optional[float]: + def _validate_response_time(self, response_time_ms: Optional[float]) -> Optional[float]: """Validate and clean response time value""" if response_time_ms is None: return None try: value = float(response_time_ms) - return ( - value - if value == value and value not in (float("inf"), float("-inf")) - else None - ) + return value if value == value and value not in (float("inf"), float("-inf")) else None except (ValueError, TypeError): - verbose_proxy_logger.warning( - f"Invalid response_time_ms value: {response_time_ms}" - ) + verbose_proxy_logger.warning(f"Invalid response_time_ms value: {response_time_ms}") return None def _clean_details(self, details: Optional[dict]) -> Optional[dict]: @@ -5286,19 +4867,13 @@ class PrismaClient: } # Add only non-None optional fields - health_check_data.update( - {k: v for k, v in optional_fields.items() if v is not None} - ) + health_check_data.update({k: v for k, v in optional_fields.items() if v is not None}) verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}") - return await HealthCheckRepository(self).table.create( - data=health_check_data - ) + return await HealthCheckRepository(self).table.create(data=health_check_data) except Exception as e: - verbose_proxy_logger.error( - f"Error saving health check result for model {model_name}: {e}" - ) + verbose_proxy_logger.error(f"Error saving health check result for model {model_name}: {e}") return None async def get_health_check_history( @@ -5364,13 +4939,9 @@ async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient): user_row = await db.get_data(user_id=user_id) if user_row is not None: print_verbose(f"User Row: {user_row}, type = {type(user_row)}") - if hasattr(user_row, "model_dump_json") and callable( - getattr(user_row, "model_dump_json") - ): + if hasattr(user_row, "model_dump_json") and callable(getattr(user_row, "model_dump_json")): cache_value = user_row.model_dump_json() - cache.set_cache( - key=cache_key, value=cache_value, ttl=600 - ) # store for 10 minutes + cache.set_cache(key=cache_key, value=cache_value, ttl=600) # store for 10 minutes return @@ -5385,9 +4956,7 @@ def _should_use_smtp_ssl(smtp_port: int) -> bool: def _create_smtp_connection(smtp_host: str, smtp_port: int) -> smtplib.SMTP: if _should_use_smtp_ssl(smtp_port=smtp_port): - return smtplib.SMTP_SSL( - host=smtp_host, port=smtp_port, context=ssl.create_default_context() - ) + return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context()) return smtplib.SMTP(host=smtp_host, port=smtp_port) @@ -5425,9 +4994,7 @@ async def send_email( email_message["From"] = sender_email email_message["To"] = receiver_email email_message["Subject"] = subject - verbose_proxy_logger.debug( - "sending email from %s to %s", sender_email, receiver_email - ) + verbose_proxy_logger.debug("sending email from %s to %s", sender_email, receiver_email) if smtp_host is None: raise ValueError("Trying to use SMTP, but SMTP_HOST is not set") @@ -5459,9 +5026,7 @@ async def send_email( ) except Exception as e: - verbose_proxy_logger.exception( - "An error occurred while sending the email:" + str(e) - ) + verbose_proxy_logger.exception("An error occurred while sending the email:" + str(e)) def hash_token(token: str): @@ -5494,17 +5059,13 @@ def verify_password(password: str, stored: str) -> bool: try: raw = base64.b64decode(stored[7:]) salt, dk = raw[:16], raw[16:] - dk2 = hashlib.scrypt( - password.encode(), salt=salt, n=16384, r=8, p=1, dklen=32 - ) + dk2 = hashlib.scrypt(password.encode(), salt=salt, n=16384, r=8, p=1, dklen=32) return secrets.compare_digest(dk, dk2) except Exception: return False # SHA256 fallback (not vulnerable to pass-the-hash: checks sha256(input) == stored) if len(stored) == 64 and all(c in "0123456789abcdef" for c in stored): - return secrets.compare_digest( - hashlib.sha256(password.encode()).hexdigest().encode(), stored.encode() - ) + return secrets.compare_digest(hashlib.sha256(password.encode()).hexdigest().encode(), stored.encode()) return False @@ -5522,11 +5083,7 @@ async def migrate_passwords_to_scrypt_async(prisma_client) -> str: return len(s) == 64 and all(c in "0123456789abcdef" for c in s) plaintext_users = [ - u - for u in all_with_pw - if u.password - and not u.password.startswith("scrypt:") - and not _is_sha256_hex(u.password) + u for u in all_with_pw if u.password and not u.password.startswith("scrypt:") and not _is_sha256_hex(u.password) ] if not plaintext_users: return "No plaintext passwords found" @@ -5562,14 +5119,10 @@ class ProxyUpdateSpend: for i in range(n_retry_times + 1): start_time = time.time() try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: + async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction: async with transaction.batch_() as batcher: # Sort by end_user_id for consistent lock ordering across pods to prevent deadlocks. - for end_user_id, response_cost in sorted( - end_user_list_transactions.items() - ): + for end_user_id, response_cost in sorted(end_user_list_transactions.items()): if litellm.max_end_user_budget is not None: pass batcher.litellm_endusertable.upsert( @@ -5593,9 +5146,7 @@ class ProxyUpdateSpend: # Optionally, sleep for a bit before retrying await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) + _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) @staticmethod async def update_spend_logs( @@ -5606,20 +5157,14 @@ class ProxyUpdateSpend: logs_to_process: Optional[List[Dict[str, Any]]] = None, ): BATCH_SIZE = 1000 # Preferred size of each batch to write to the database - MAX_LOGS_PER_INTERVAL = ( - 10000 # Maximum number of logs to flush in a single interval - ) + MAX_LOGS_PER_INTERVAL = 10000 # Maximum number of logs to flush in a single interval popped_batch = False if logs_to_process is None: # Atomically read and remove logs to process (protected by lock) async with prisma_client._spend_log_transactions_lock: - logs_to_process = prisma_client.spend_log_transactions[ - :MAX_LOGS_PER_INTERVAL - ] + logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] # Remove the logs we're about to process - prisma_client.spend_log_transactions = ( - prisma_client.spend_log_transactions[len(logs_to_process) :] - ) + prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[len(logs_to_process) :] popped_batch = True if len(logs_to_process) > 0: verbose_proxy_logger.info( @@ -5631,11 +5176,7 @@ class ProxyUpdateSpend: for i in range(n_retry_times + 1): try: base_url = os.getenv("SPEND_LOGS_URL", None) - if ( - len(logs_to_process) > 0 - and base_url is not None - and db_writer_client is not None - ): + if len(logs_to_process) > 0 and base_url is not None and db_writer_client is not None: if not base_url.endswith("/"): base_url += "/" verbose_proxy_logger.debug("base_url: {}".format(base_url)) @@ -5652,16 +5193,11 @@ class ProxyUpdateSpend: else: for j in range(0, len(logs_to_process), BATCH_SIZE): batch = logs_to_process[j : j + BATCH_SIZE] - batch_with_dates = [ - prisma_client.jsonify_object({**entry}) - for entry in batch - ] + batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch] await SpendLogsRepository(prisma_client).table.create_many( data=batch_with_dates, skip_duplicates=True ) - verbose_proxy_logger.debug( - f"Flushed {len(batch)} logs to the DB." - ) + verbose_proxy_logger.debug(f"Flushed {len(batch)} logs to the DB.") # Explicitly clear batch memory del batch, batch_with_dates @@ -5676,8 +5212,7 @@ class ProxyUpdateSpend: if i is None: i = 0 verbose_proxy_logger.warning( - "Spend tracking - DB connection error writing spend logs, " - "retry %d/%d. logs_count=%d, error=%s", + "Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s", i + 1, n_retry_times, len(logs_to_process), @@ -5689,9 +5224,7 @@ class ProxyUpdateSpend: except Exception as e: # Logs already removed from queue at start - don't put them back # This matches the original behavior where logs are removed even on error - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) + _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) finally: # Clean up logs_to_process only if we popped it (caller-owned otherwise) if popped_batch: @@ -5786,12 +5319,10 @@ async def update_daily_tag_spend( proxy_logging_obj=proxy_logging_obj, ) else: - await ( - proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db( - prisma_client=prisma_client, - n_retry_times=n_retry_times, - proxy_logging_obj=proxy_logging_obj, - ) + await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, ) except Exception as e: # NOTE: keep this as a plain ``error`` (no traceback) to match the @@ -5824,9 +5355,7 @@ async def update_spend_logs_job( async with prisma_client._spend_log_transactions_lock: logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] - prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ - len(logs_to_process) : - ] + prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[len(logs_to_process) :] await ProxyUpdateSpend.update_spend_logs( n_retry_times=n_retry_times, @@ -5914,9 +5443,7 @@ async def _monitor_spend_logs_queue( f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff" ) # Exponential backoff when below threshold but still processing - current_interval = min( - current_interval * backoff_multiplier, max_backoff - ) + current_interval = min(current_interval * backoff_multiplier, max_backoff) await update_spend_logs_job( prisma_client=prisma_client, @@ -5925,9 +5452,7 @@ async def _monitor_spend_logs_queue( ) else: # Exponential backoff when no logs to process - current_interval = min( - current_interval * backoff_multiplier, max_backoff - ) + current_interval = min(current_interval * backoff_multiplier, max_backoff) await asyncio.sleep(current_interval) except Exception as e: @@ -5937,9 +5462,7 @@ async def _monitor_spend_logs_queue( await asyncio.sleep(current_interval) -def _raise_failed_update_spend_exception( - e: Exception, start_time: float, proxy_logging_obj: ProxyLogging -): +def _raise_failed_update_spend_exception(e: Exception, start_time: float, proxy_logging_obj: ProxyLogging): """ Raise an exception for failed update spend logs @@ -5948,9 +5471,7 @@ def _raise_failed_update_spend_exception( """ import traceback - error_msg = ( - f"[Non-Blocking]LiteLLM Prisma Client Exception - update spend logs: {str(e)}" - ) + error_msg = f"[Non-Blocking]LiteLLM Prisma Client Exception - update spend logs: {str(e)}" error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() _duration = end_time - start_time @@ -5971,9 +5492,7 @@ def _get_month_end_date(today: date) -> date: return date(today.year, today.month + 1, 1) - timedelta(days=1) -def _is_projected_spend_over_limit( - current_spend: float, soft_budget_limit: Optional[float] -): +def _is_projected_spend_over_limit(current_spend: float, soft_budget_limit: Optional[float]): if soft_budget_limit is None: # If there's no limit, we can't exceed it. return False @@ -6000,9 +5519,7 @@ def _is_projected_spend_over_limit( return False -def _get_projected_spend_over_limit( - current_spend: float, soft_budget_limit: Optional[float] -) -> Optional[tuple]: +def _get_projected_spend_over_limit(current_spend: float, soft_budget_limit: Optional[float]) -> Optional[tuple]: if soft_budget_limit is None: return None @@ -6052,9 +5569,7 @@ def _to_ns(dt): return int(dt.timestamp() * 1e9) -def _check_and_merge_model_level_guardrails( - data: dict, llm_router: Optional[Router] -) -> dict: +def _check_and_merge_model_level_guardrails(data: dict, llm_router: Optional[Router]) -> dict: """ Check if the model has guardrails defined and merge them with existing guardrails in the request data. @@ -6111,9 +5626,7 @@ def _merge_guardrails_with_existing(data: dict, model_level_guardrails: Any) -> # Ensure model_level_guardrails is a list if not isinstance(model_level_guardrails, list): - model_level_guardrails = ( - [model_level_guardrails] if model_level_guardrails else [] - ) + model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] # Combine existing and model-level guardrails metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) @@ -6226,7 +5739,9 @@ def _premium_user_check(feature: Optional[str] = None): if feature: detail_msg = f"This feature is only available for LiteLLM Enterprise users: {feature}. {CommonProxyErrors.not_premium_user.value}" else: - detail_msg = f"This feature is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}" + detail_msg = ( + f"This feature is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}" + ) if not premium_user: raise HTTPException( @@ -6380,14 +5895,14 @@ def construct_database_url_from_env_vars() -> Optional[str]: if database_host and database_username and database_name: # Handle the problem of special character escaping in the database URL database_username_enc = urllib.parse.quote_plus(database_username) - database_password_enc = ( - urllib.parse.quote_plus(database_password) if database_password else "" - ) + database_password_enc = urllib.parse.quote_plus(database_password) if database_password else "" database_name_enc = urllib.parse.quote_plus(database_name) # Construct DATABASE_URL from the provided variables if database_password: - database_url = f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" + database_url = ( + f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" + ) else: database_url = f"postgresql://{database_username_enc}@{database_host}/{database_name_enc}" @@ -6466,9 +5981,7 @@ async def get_available_models_for_user( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await validate_membership( - user_api_key_dict=user_api_key_dict, team_table=team_object - ) + await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) team_models = team_object.models team_models = get_team_models( @@ -6592,9 +6105,7 @@ def validate_model_access( if model_id not in available_models: raise HTTPException( status_code=404, - detail="The model `{}` does not exist or is not accessible".format( - model_id - ), + detail="The model `{}` does not exist or is not accessible".format(model_id), ) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 9c2d297bfa7..4e7890e6ed8 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -42,9 +42,9 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( Raises: HTTPException: If user doesn't have access to the vector store """ - vector_store_to_run: Optional[ - LiteLLM_ManagedVectorStore - ] = await get_litellm_managed_vector_store(vector_store_id=vector_store_id) + vector_store_to_run: Optional[LiteLLM_ManagedVectorStore] = await get_litellm_managed_vector_store( + vector_store_id=vector_store_id + ) if vector_store_to_run is not None: if user_api_key_dict is not None: await assert_user_can_access_vector_store( @@ -56,9 +56,7 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider") if "litellm_credential_name" in vector_store_to_run: - data["litellm_credential_name"] = vector_store_to_run.get( - "litellm_credential_name" - ) + data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name") if "litellm_params" in vector_store_to_run: litellm_params = vector_store_to_run.get("litellm_params", {}) or {} @@ -222,9 +220,7 @@ async def vector_store_create( ) # Get managed vector stores hook - managed_vector_stores: Any = proxy_logging_obj.get_proxy_hook( - "managed_vector_stores" - ) + managed_vector_stores: Any = proxy_logging_obj.get_proxy_hook("managed_vector_stores") if managed_vector_stores is None: raise HTTPException( status_code=500, @@ -277,12 +273,8 @@ async def vector_store_create( ) -@router.get( - "/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] -) -@router.get( - "/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] -) +@router.get("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.get("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) async def vector_store_retrieve( request: Request, vector_store_id: str, @@ -414,12 +406,8 @@ async def vector_store_list( ) -@router.post( - "/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] -) -@router.post( - "/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] -) +@router.post("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.post("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) async def vector_store_update( request: Request, vector_store_id: str, @@ -484,12 +472,8 @@ async def vector_store_update( ) -@router.delete( - "/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] -) -@router.delete( - "/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] -) +@router.delete("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.delete("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) async def vector_store_delete( request: Request, vector_store_id: str, @@ -589,9 +573,9 @@ async def index_create( detail=CommonProxyErrors.db_not_connected_error.value, ) ## 1. check if index already exists - existing_index = await ManagedVectorStoreIndexRepository( - prisma_client - ).table.find_unique(where={"index_name": index_create_request.index_name}) + existing_index = await ManagedVectorStoreIndexRepository(prisma_client).table.find_unique( + where={"index_name": index_create_request.index_name} + ) ## 2. set created_by and updated_by @@ -605,8 +589,6 @@ async def index_create( index_data = index_create_request.model_dump(exclude_none=True) index_data["created_by"] = user_api_key_dict.user_id index_data["updated_by"] = user_api_key_dict.user_id - new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create( - data=jsonify_object(index_data) - ) + new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create(data=jsonify_object(index_data)) return new_index.model_dump() diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 032a3302fdc..84054f1398d 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -141,9 +141,7 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router( - embedding_model: str, llm_router -) -> Optional[Dict[str, Any]]: +def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> Optional[Dict[str, Any]]: """ Resolve embedding config from router's config-defined models. @@ -173,9 +171,7 @@ def _resolve_embedding_config_from_router( for model_name in model_name_candidates: try: # Try to get deployment by model group name (model_name in config) - deployment = llm_router.get_deployment_by_model_group_name( - model_group_name=model_name - ) + deployment = llm_router.get_deployment_by_model_group_name(model_group_name=model_name) if deployment is not None and deployment.litellm_params is not None: litellm_params = deployment.litellm_params @@ -215,17 +211,13 @@ def _resolve_embedding_config_from_router( ) return embedding_config except Exception as e: - verbose_proxy_logger.debug( - f"Error resolving embedding config from router for model {model_name}: {str(e)}" - ) + verbose_proxy_logger.debug(f"Error resolving embedding config from router for model {model_name}: {str(e)}") continue return None -async def _resolve_embedding_config_from_db( - embedding_model: str, prisma_client -) -> Optional[Dict[str, Any]]: +async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) -> Optional[Dict[str, Any]]: """ Resolve embedding config from database model configuration. @@ -254,9 +246,7 @@ async def _resolve_embedding_config_from_db( # Try to find model in database for model_name in model_name_candidates: try: - db_model = await ModelRepository(prisma_client).table.find_first( - where={"model_name": model_name} - ) + db_model = await ModelRepository(prisma_client).table.find_first(where={"model_name": model_name}) if db_model and db_model.litellm_params: # Extract litellm_params (could be dict or JSON string) @@ -271,9 +261,7 @@ async def _resolve_embedding_config_from_db( for k, v in model_params.items(): if isinstance(v, str): # Decrypt value - returns original value if decryption fails or no key is set - decrypted_value = decrypt_value_helper( - value=v, key=k, return_original_value=True - ) + decrypted_value = decrypt_value_helper(value=v, key=k, return_original_value=True) decrypted_params[k] = decrypted_value else: decrypted_params[k] = v @@ -311,17 +299,13 @@ async def _resolve_embedding_config_from_db( ) return embedding_config except Exception as e: - verbose_proxy_logger.debug( - f"Error resolving embedding config for model {model_name}: {str(e)}" - ) + verbose_proxy_logger.debug(f"Error resolving embedding config for model {model_name}: {str(e)}") continue return None -async def _resolve_embedding_config( - embedding_model: str, prisma_client, llm_router=None -) -> Optional[Dict[str, Any]]: +async def _resolve_embedding_config(embedding_model: str, prisma_client, llm_router=None) -> Optional[Dict[str, Any]]: """ Resolve embedding config from either router (config-defined) or database models. @@ -358,13 +342,9 @@ async def _resolve_embedding_config( # First try to resolve from router (config-defined models) if llm_router is not None: - router_config = _resolve_embedding_config_from_router( - embedding_model=embedding_model, llm_router=llm_router - ) + router_config = _resolve_embedding_config_from_router(embedding_model=embedding_model, llm_router=llm_router) if router_config: - verbose_proxy_logger.debug( - f"Resolved embedding config from router for model {embedding_model}" - ) + verbose_proxy_logger.debug(f"Resolved embedding config from router for model {embedding_model}") cache.set_cache(embedding_model, router_config) return router_config @@ -374,9 +354,7 @@ async def _resolve_embedding_config( embedding_model=embedding_model, prisma_client=prisma_client ) if db_config: - verbose_proxy_logger.debug( - f"Resolved embedding config from database for model {embedding_model}" - ) + verbose_proxy_logger.debug(f"Resolved embedding config from database for model {embedding_model}") cache.set_cache(embedding_model, db_config) return db_config @@ -402,9 +380,7 @@ async def _check_vector_store_access( - key-level and team-level ``object_permission.vector_stores`` allowlists - team_id match between key and store """ - return await can_user_access_vector_store( - vector_store=vector_store, user_api_key_dict=user_api_key_dict - ) + return await can_user_access_vector_store(vector_store=vector_store, user_api_key_dict=user_api_key_dict) async def create_vector_store_in_db( @@ -439,9 +415,9 @@ async def create_vector_store_in_db( raise HTTPException(status_code=500, detail="Database not connected") # Check if vector store already exists - existing_vector_store = await ManagedVectorStoresRepository( - prisma_client - ).table.find_unique(where={"vector_store_id": vector_store_id}) + existing_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + where={"vector_store_id": vector_store_id} + ) if existing_vector_store is not None: raise HTTPException( status_code=400, @@ -478,32 +454,22 @@ async def create_vector_store_in_db( # at request-handling time so the cleartext config exists only in # per-request memory and never reaches the database. if litellm_params: - litellm_params_dict = GenericLiteLLMParams(**litellm_params).model_dump( - exclude_none=True - ) + litellm_params_dict = GenericLiteLLMParams(**litellm_params).model_dump(exclude_none=True) data_to_create["litellm_params"] = safe_dumps(litellm_params_dict) else: # Provide empty dict if no litellm_params provided data_to_create["litellm_params"] = safe_dumps({}) # Create in database - _new_vector_store = await ManagedVectorStoresRepository(prisma_client).table.create( - data=data_to_create - ) + _new_vector_store = await ManagedVectorStoresRepository(prisma_client).table.create(data=data_to_create) - new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( - **_new_vector_store.model_dump() - ) + new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore(**_new_vector_store.model_dump()) # Add vector store to registry if litellm.vector_store_registry is not None: - litellm.vector_store_registry.add_vector_store_to_registry( - vector_store=new_vector_store - ) + litellm.vector_store_registry.add_vector_store_to_registry(vector_store=new_vector_store) - verbose_proxy_logger.info( - f"Vector store {vector_store_id} created in database successfully" - ) + verbose_proxy_logger.info(f"Vector store {vector_store_id} created in database successfully") return new_vector_store @@ -568,9 +534,7 @@ async def new_vector_store( # cleartext value persisted by an earlier proxy version doesn't # come back in the response. response_vs = LiteLLM_ManagedVectorStore(**new_vector_store) - response_vs["litellm_params"] = _redact_sensitive_litellm_params( - new_vector_store.get("litellm_params") - ) + response_vs["litellm_params"] = _redact_sensitive_litellm_params(new_vector_store.get("litellm_params")) return { "status": "success", @@ -617,9 +581,7 @@ async def list_vector_stores( try: # Get vector stores from database first (source of truth) - vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( - prisma_client=prisma_client - ) + vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db(prisma_client=prisma_client) # Build map from database vector stores for vector_store in vector_stores_from_db: @@ -630,9 +592,7 @@ async def list_vector_stores( # Process in-memory vector stores if litellm.vector_store_registry is not None: - in_memory_vector_stores = copy.deepcopy( - litellm.vector_store_registry.vector_stores - ) + in_memory_vector_stores = copy.deepcopy(litellm.vector_store_registry.vector_stores) vector_stores_to_delete_from_memory: List[str] = [] @@ -654,12 +614,8 @@ async def list_vector_stores( # Synchronize in-memory registry with database # 1. Remove deleted vector stores from memory for vs_id in vector_stores_to_delete_from_memory: - litellm.vector_store_registry.delete_vector_store_from_registry( - vector_store_id=vs_id - ) - verbose_proxy_logger.debug( - f"Removed deleted vector store {vs_id} from in-memory registry" - ) + litellm.vector_store_registry.delete_vector_store_from_registry(vector_store_id=vs_id) + verbose_proxy_logger.debug(f"Removed deleted vector store {vs_id} from in-memory registry") # 2. Update in-memory registry with database versions (for updates) for vector_store in vector_stores_from_db: @@ -674,9 +630,7 @@ async def list_vector_stores( 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") - ) + redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) accessible_vector_stores.append(redacted) total_count = len(accessible_vector_stores) @@ -725,14 +679,12 @@ async def delete_vector_store( memory_vector_store_exists = False vector_store_to_check = None - existing_vector_store = await ManagedVectorStoresRepository( - prisma_client - ).table.find_unique(where={"vector_store_id": data.vector_store_id}) + existing_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + where={"vector_store_id": data.vector_store_id} + ) if existing_vector_store is not None: db_vector_store_exists = True - vector_store_to_check = LiteLLM_ManagedVectorStore( - **existing_vector_store.model_dump() - ) + vector_store_to_check = LiteLLM_ManagedVectorStore(**existing_vector_store.model_dump()) # Check in-memory registry if litellm.vector_store_registry is not None: @@ -752,9 +704,7 @@ async def delete_vector_store( ) # Check access control - if vector_store_to_check and not await _check_vector_store_access( - vector_store_to_check, user_api_key_dict - ): + if vector_store_to_check and not await _check_vector_store_access(vector_store_to_check, user_api_key_dict): raise HTTPException( status_code=403, detail="Access denied: You do not have permission to delete this vector store", @@ -768,9 +718,7 @@ async def delete_vector_store( # Delete from in-memory registry if exists if memory_vector_store_exists and litellm.vector_store_registry is not None: - litellm.vector_store_registry.delete_vector_store_from_registry( - vector_store_id=data.vector_store_id - ) + litellm.vector_store_registry.delete_vector_store_from_registry(vector_store_id=data.vector_store_id) return { "status": "success", @@ -808,9 +756,7 @@ async def get_vector_store_info( ) if vector_store is not None: # Check access control - if not await _check_vector_store_access( - vector_store, user_api_key_dict - ): + if not await _check_vector_store_access(vector_store, user_api_key_dict): raise HTTPException( status_code=403, detail="Access denied: You do not have permission to access this vector store", @@ -828,17 +774,12 @@ async def get_vector_store_info( vector_store_id=vector_store.get("vector_store_id") or "", custom_llm_provider=vector_store.get("custom_llm_provider") or "", vector_store_name=vector_store.get("vector_store_name") or None, - vector_store_description=vector_store.get( - "vector_store_description" - ) - or None, + vector_store_description=vector_store.get("vector_store_description") or None, vector_store_metadata=parsed_metadata, created_at=vector_store.get("created_at") or None, updated_at=vector_store.get("updated_at") or None, litellm_credential_name=vector_store.get("litellm_credential_name"), - litellm_params=_redact_sensitive_litellm_params( - vector_store.get("litellm_params") - ), + litellm_params=_redact_sensitive_litellm_params(vector_store.get("litellm_params")), team_id=vector_store.get("team_id") or None, user_id=vector_store.get("user_id") or None, ) @@ -851,9 +792,7 @@ async def get_vector_store_info( ) vector_store_dict = dict(vector_store_typed) if "litellm_params" in vector_store_dict: - vector_store_dict["litellm_params"] = _redact_sensitive_litellm_params( - vector_store_dict["litellm_params"] - ) + vector_store_dict["litellm_params"] = _redact_sensitive_litellm_params(vector_store_dict["litellm_params"]) return {"vector_store": vector_store_dict} except HTTPException: # Preserve 403/404 from the access-control / not-found checks above; @@ -900,9 +839,7 @@ async def update_vector_store( # Handle metadata serialization if update_data.get("vector_store_metadata") is not None: - update_data["vector_store_metadata"] = safe_dumps( - update_data["vector_store_metadata"] - ) + update_data["vector_store_metadata"] = safe_dumps(update_data["vector_store_metadata"]) # Handle litellm_params if provided. As with the create path, the # embedding-config auto-resolve previously persisted cleartext @@ -913,9 +850,7 @@ async def update_vector_store( # ``litellm_embedding_model`` reference. if "litellm_params" in update_data: _input_litellm_params: dict = update_data.get("litellm_params", {}) or {} - litellm_params_dict = GenericLiteLLMParams( - **_input_litellm_params - ).model_dump(exclude_none=True) + litellm_params_dict = GenericLiteLLMParams(**_input_litellm_params).model_dump(exclude_none=True) update_data["litellm_params"] = safe_dumps(litellm_params_dict) # Update in database @@ -941,9 +876,7 @@ async def update_vector_store( # credentials) back to the caller — even when the caller only # changed unrelated fields like ``vector_store_description``. response_vs = LiteLLM_ManagedVectorStore(**updated_vs) - response_vs["litellm_params"] = _redact_sensitive_litellm_params( - updated_vs.get("litellm_params") - ) + response_vs["litellm_params"] = _redact_sensitive_litellm_params(updated_vs.get("litellm_params")) return { "status": "success", "message": f"Vector store {vector_store_id} updated successfully", diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index d4afc547031..05149a0f6c0 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -48,10 +48,7 @@ def assert_proxy_admin_for_vector_store_index_management( return raise HTTPException( status_code=403, - detail=( - f"Only proxy admins can {operation} vector store indexes. " - "Contact your LiteLLM administrator." - ), + detail=(f"Only proxy admins can {operation} vector store indexes. Contact your LiteLLM administrator."), ) @@ -171,26 +168,17 @@ async def can_user_access_vector_store( key_object_permission = user_api_key_dict.object_permission if key_object_permission is None: - key_object_permission = await _get_object_permission_for_id( - user_api_key_dict.object_permission_id - ) + key_object_permission = await _get_object_permission_for_id(user_api_key_dict.object_permission_id) if _object_permission_allows_vector_store(key_object_permission, vector_store_id): return True - team_object_permission: Optional[LiteLLM_ObjectPermissionTable] = ( - user_api_key_dict.team_object_permission - ) + team_object_permission: Optional[LiteLLM_ObjectPermissionTable] = user_api_key_dict.team_object_permission if team_object_permission is None: - team_object_permission = await _get_object_permission_for_id( - user_api_key_dict.team_object_permission_id - ) + team_object_permission = await _get_object_permission_for_id(user_api_key_dict.team_object_permission_id) 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 - ): + if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store_team_id: return True return False @@ -247,9 +235,7 @@ async def get_litellm_managed_vector_store( ) if not rows: return None - return _normalize_litellm_params( - LiteLLM_ManagedVectorStore(**rows[0].model_dump()) - ) + return _normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from shared cache: %s", @@ -282,9 +268,7 @@ async def assert_user_can_access_vector_store_id( Unknown ids are treated as provider-native ids and are not rejected here. """ - vector_store = await get_litellm_managed_vector_store( - vector_store_id=vector_store_id - ) + vector_store = await get_litellm_managed_vector_store(vector_store_id=vector_store_id) if vector_store is not None: await assert_user_can_access_vector_store( vector_store=vector_store, @@ -348,10 +332,7 @@ def check_vector_store_permission( if index_config.get("index_name") == index_name: index_permissions = index_config.get("index_permissions", []) - if ( - isinstance(index_permissions, list) - and permission in index_permissions - ): + if isinstance(index_permissions, list) and permission in index_permissions: return True return False @@ -379,15 +360,11 @@ def is_allowed_to_call_vector_store_endpoint( key_metadata = user_api_key_dict.metadata team_metadata = user_api_key_dict.team_metadata - provider_config = ProviderConfigManager.get_provider_vector_stores_config( - provider=provider - ) + provider_config = ProviderConfigManager.get_provider_vector_stores_config(provider=provider) if provider_config is None: return None - provider_vector_store_endpoints = ( - provider_config.get_vector_store_endpoints_by_type() - ) + provider_vector_store_endpoints = provider_config.get_vector_store_endpoints_by_type() # Inline import — auth_utils participates in a proxy import cycle. from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 @@ -413,17 +390,13 @@ def is_allowed_to_call_vector_store_endpoint( # Determine the permission type based on the request permission_type = None for endpoint in provider_vector_store_endpoints["read"]: - if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request_route - ): + if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route): permission_type = "read" break if permission_type is None: for endpoint in provider_vector_store_endpoints["write"]: - if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request_route - ): + if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route): permission_type = "write" break @@ -469,15 +442,11 @@ def is_allowed_to_call_vector_store_files_endpoint( key_metadata = user_api_key_dict.metadata team_metadata = user_api_key_dict.team_metadata - provider_config = ProviderConfigManager.get_provider_vector_store_files_config( - provider=provider - ) + provider_config = ProviderConfigManager.get_provider_vector_store_files_config(provider=provider) if provider_config is None: return None - provider_vector_store_endpoints = ( - provider_config.get_vector_store_file_endpoints_by_type() - ) + provider_vector_store_endpoints = provider_config.get_vector_store_file_endpoints_by_type() # Inline import — auth_utils participates in a proxy import cycle. from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 @@ -486,17 +455,13 @@ def is_allowed_to_call_vector_store_files_endpoint( permission_type: Optional[str] = None for endpoint in provider_vector_store_endpoints.get("read", ()): - if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request_route - ): + if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route): permission_type = "read" break if permission_type is None: for endpoint in provider_vector_store_endpoints.get("write", ()): - if request.method == endpoint[0] and _does_endpoint_match( - endpoint[1], request_route - ): + if request.method == endpoint[0] and _does_endpoint_match(endpoint[1], request_route): permission_type = "write" break diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index f6ceae39779..890db2f73a4 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -100,9 +100,7 @@ def _update_request_data_with_managed_file_id( # Get credentials for the model if llm_router: - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=routing_model - ) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=routing_model) if credentials: prepare_data_with_credentials( data=data, @@ -117,9 +115,7 @@ def _update_request_data_with_managed_file_id( # If we extracted the provider file ID but no routing, still use it if llm_output_file_id: data["file_id"] = llm_output_file_id - verbose_logger.debug( - f"Replaced unified file ID with provider file ID: {llm_output_file_id}" - ) + verbose_logger.debug(f"Replaced unified file ID with provider file ID: {llm_output_file_id}") return data, file_id # Return original managed file ID return data, file_id if decoded_id else None @@ -148,11 +144,7 @@ def _update_request_data_with_managed_file_id( verbose_logger.debug( f"Routing vector store file operation using model: {model_used}" - + ( - f", file_id: {file_id} -> {original_file_id}" - if original_file_id - else "" - ) + + (f", file_id: {file_id} -> {original_file_id}" if original_file_id else "") ) return data, file_id # Return original file ID for response replacement @@ -231,13 +223,9 @@ async def _update_request_data_with_model_routing_hint( if data.get("api_key") is not None or data.get("api_base") is not None: return data - user_controlled_model_hint = request.query_params.get( - "model" - ) or request.headers.get("x-litellm-model") + user_controlled_model_hint = request.query_params.get("model") or request.headers.get("x-litellm-model") model_hint = data.get("model") or user_controlled_model_hint - should_authorize_model_hint = ( - isinstance(model_hint, str) and model_hint == user_controlled_model_hint - ) + should_authorize_model_hint = isinstance(model_hint, str) and model_hint == user_controlled_model_hint should_route = False credentials = None @@ -249,9 +237,7 @@ async def _update_request_data_with_model_routing_hint( llm_router=llm_router, user_api_key_dict=user_api_key_dict, ) - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=model_hint - ) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_hint) should_route = credentials is not None else: if isinstance(model_hint, str) and should_authorize_model_hint: @@ -299,9 +285,7 @@ async def _update_request_data_with_model_routing_hint( openai_credentials = None for model_name in model_names_to_check: - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=model_name - ) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_name) if credentials is None: continue @@ -396,9 +380,7 @@ def _update_request_data_with_litellm_managed_vector_store_registry( if routing_model: data["model"] = routing_model - verbose_logger.info( - f"Routing vector store files operation to model: {routing_model}" - ) + verbose_logger.info(f"Routing vector store files operation to model: {routing_model}") # Replace unified vector store ID with provider resource ID if provider_resource_id: @@ -411,11 +393,7 @@ def _update_request_data_with_litellm_managed_vector_store_registry( # Legacy path: Check vector store registry for non-managed vector stores. vector_store_to_run = managed_vector_store - if ( - vector_store_to_run is None - and should_lookup_registry - and litellm.vector_store_registry is not None - ): + if vector_store_to_run is None and should_lookup_registry and litellm.vector_store_registry is not None: vector_store_to_run = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( vector_store_id=vector_store_id ) @@ -424,9 +402,7 @@ def _update_request_data_with_litellm_managed_vector_store_registry( if "custom_llm_provider" in vector_store_to_run: data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider") if "litellm_credential_name" in vector_store_to_run: - data["litellm_credential_name"] = vector_store_to_run.get( - "litellm_credential_name" - ) + data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name") if "litellm_params" in vector_store_to_run: litellm_params = vector_store_to_run.get("litellm_params", {}) or {} data.update(litellm_params) @@ -468,9 +444,7 @@ def _maybe_check_permissions( return metadata = user_api_key_dict.metadata or {} team_metadata = user_api_key_dict.team_metadata or {} - if not metadata.get("allowed_vector_store_indexes") and not team_metadata.get( - "allowed_vector_store_indexes" - ): + if not metadata.get("allowed_vector_store_indexes") and not team_metadata.get("allowed_vector_store_indexes"): return is_allowed_to_call_vector_store_files_endpoint( provider=provider, diff --git a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py index b47f6a747db..79a188817bf 100644 --- a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py @@ -52,9 +52,7 @@ def _decode_to_convergence(value: str) -> str: def _normalize_langfuse_base_url(base_target_url: str) -> str: - if not ( - base_target_url.startswith("http://") or base_target_url.startswith("https://") - ): + if not (base_target_url.startswith("http://") or base_target_url.startswith("https://")): # Existing behavior allows host-only Langfuse settings. base_target_url = "http://" + base_target_url @@ -113,17 +111,13 @@ def _get_langfuse_proxy_credentials( if not dynamic_langfuse_public_key or not dynamic_langfuse_secret_key: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "Dynamic Langfuse hosts must include dynamic Langfuse credentials" - }, + detail={"error": "Dynamic Langfuse hosts must include dynamic Langfuse credentials"}, ) return dynamic_langfuse_public_key, dynamic_langfuse_secret_key return ( - dynamic_langfuse_public_key - or litellm.utils.get_secret(secret_name="LANGFUSE_PUBLIC_KEY"), - dynamic_langfuse_secret_key - or litellm.utils.get_secret(secret_name="LANGFUSE_SECRET_KEY"), + dynamic_langfuse_public_key or litellm.utils.get_secret(secret_name="LANGFUSE_PUBLIC_KEY"), + dynamic_langfuse_secret_key or litellm.utils.get_secret(secret_name="LANGFUSE_SECRET_KEY"), ) @@ -179,23 +173,16 @@ async def langfuse_proxy_route( decoded_str = decoded_bytes.decode("utf-8") api_key = decoded_str.split(":")[1] # assume api key is passed in as secret key - user_api_key_dict = await user_api_key_auth( - request=request, api_key="Bearer {}".format(api_key) - ) + user_api_key_dict = await user_api_key_auth(request=request, api_key="Bearer {}".format(api_key)) - callback_settings_obj: Optional[TeamCallbackMetadata] = ( - _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=proxy_config - ) + callback_settings_obj: Optional[TeamCallbackMetadata] = _get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, proxy_config=proxy_config ) dynamic_langfuse_public_key: Optional[str] = None dynamic_langfuse_secret_key: Optional[str] = None dynamic_langfuse_host: Optional[str] = None - if ( - callback_settings_obj is not None - and callback_settings_obj.callback_vars is not None - ): + if callback_settings_obj is not None and callback_settings_obj.callback_vars is not None: for k, v in callback_settings_obj.callback_vars.items(): if k == "langfuse_public_key": dynamic_langfuse_public_key = v @@ -206,9 +193,7 @@ async def langfuse_proxy_route( dynamic_host_supplied = dynamic_langfuse_host is not None base_target_url: str = ( - dynamic_langfuse_host - or os.getenv("LANGFUSE_HOST", _DEFAULT_LANGFUSE_HOST) - or _DEFAULT_LANGFUSE_HOST + dynamic_langfuse_host or os.getenv("LANGFUSE_HOST", _DEFAULT_LANGFUSE_HOST) or _DEFAULT_LANGFUSE_HOST ) langfuse_public_key, langfuse_secret_key = _get_langfuse_proxy_credentials( dynamic_host_supplied=dynamic_host_supplied, @@ -221,9 +206,9 @@ async def langfuse_proxy_route( dynamic_host_supplied=dynamic_host_supplied, ) - langfuse_combined_key = "Basic " + b64encode( - f"{langfuse_public_key}:{langfuse_secret_key}".encode("utf-8") - ).decode("ascii") + langfuse_combined_key = "Basic " + b64encode(f"{langfuse_public_key}:{langfuse_secret_key}".encode("utf-8")).decode( + "ascii" + ) target_headers["Authorization"] = langfuse_combined_key ## CREATE PASS-THROUGH diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 8d1c8059dca..5809967cff5 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -264,9 +264,7 @@ async def video_status( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id( - model_id_from_decoded - ) + resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) if resolved_model: data["model"] = resolved_model @@ -364,9 +362,7 @@ async def video_content( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id( - model_id_from_decoded - ) + resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) if resolved_model: data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing @@ -396,9 +392,7 @@ async def video_content( return Response( content=video_bytes, media_type="video/mp4", - headers={ - "Content-Disposition": f"attachment; filename=video_{video_id}.mp4" - }, + headers={"Content-Disposition": f"attachment; filename=video_{video_id}.mp4"}, ) except Exception as e: raise await processor._handle_llm_api_exception( @@ -478,9 +472,7 @@ async def video_remix( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id( - model_id_from_decoded - ) + resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) if resolved_model: data["model"] = resolved_model @@ -566,9 +558,7 @@ async def video_create_character( if video_file: data["video"] = video_file[0] - target_model_name = extract_model_from_target_model_names( - data.get("target_model_names") - ) + target_model_name = extract_model_from_target_model_names(data.get("target_model_names")) if target_model_name and not data.get("model"): data["model"] = target_model_name @@ -602,11 +592,7 @@ async def video_create_character( ) if target_model_name: hidden_params = getattr(response, "_hidden_params", {}) or {} - provider_for_encoding = ( - hidden_params.get("custom_llm_provider") - or custom_llm_provider - or "openai" - ) + provider_for_encoding = hidden_params.get("custom_llm_provider") or custom_llm_provider or "openai" model_id_for_encoding = hidden_params.get("model_id") or data.get("model") response = encode_character_id_in_response( response=response, @@ -687,9 +673,7 @@ async def video_get_character( data["custom_llm_provider"] = custom_llm_provider if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id( - model_id_from_decoded - ) + resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) if resolved_model: data["model"] = resolved_model @@ -798,9 +782,7 @@ async def video_edit( data["custom_llm_provider"] = custom_llm_provider if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id( - model_id_from_decoded - ) + resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) if resolved_model: data["model"] = resolved_model @@ -900,9 +882,7 @@ async def video_extension( data["custom_llm_provider"] = custom_llm_provider if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id( - model_id_from_decoded - ) + resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) if resolved_model: data["model"] = resolved_model diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 689fe4a371b..412a0e87d88 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -7,9 +7,7 @@ from litellm.types.videos.utils import encode_character_id_with_provider def extract_model_from_target_model_names(target_model_names: Any) -> Optional[str]: if isinstance(target_model_names, str): - target_model_names = [ - m.strip() for m in target_model_names.split(",") if m.strip() - ] + target_model_names = [m.strip() for m in target_model_names.split(",") if m.strip()] elif not isinstance(target_model_names, list): return None return target_model_names[0] if target_model_names else None @@ -37,9 +35,7 @@ def get_custom_provider_from_data(data: Dict[str, Any]) -> Optional[str]: return None -def encode_character_id_in_response( - response: Any, custom_llm_provider: str, model_id: Optional[str] -) -> Any: +def encode_character_id_in_response(response: Any, custom_llm_provider: str, model_id: Optional[str]) -> Any: if isinstance(response, dict) and response.get("id"): response["id"] = encode_character_id_with_provider( character_id=response["id"], diff --git a/litellm/proxy_auth/credentials.py b/litellm/proxy_auth/credentials.py index 103b0088d80..5383e17e793 100644 --- a/litellm/proxy_auth/credentials.py +++ b/litellm/proxy_auth/credentials.py @@ -103,8 +103,7 @@ class AzureADCredential: self._initialized = True except ImportError: raise ImportError( - "azure-identity is required for AzureADCredential. " - "Install it with: pip install azure-identity" + "azure-identity is required for AzureADCredential. Install it with: pip install azure-identity" ) result = self._credential.get_token(scope) diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index b49c86b3c8b..527b76e42ff 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -60,9 +60,7 @@ class BaseRAGIngestion(ABC): ingest_options.get("chunking_strategy") or {"type": "auto"}, ) self.embedding_config = ingest_options.get("embedding") - self.vector_store_config: Dict[str, Any] = cast( - Dict[str, Any], ingest_options.get("vector_store") or {} - ) + self.vector_store_config: Dict[str, Any] = cast(Dict[str, Any], ingest_options.get("vector_store") or {}) self.ingest_name = ingest_options.get("name") # Load credentials from litellm_credential_name if provided in vector_store config @@ -82,9 +80,7 @@ class BaseRAGIngestion(ABC): credential_name = self.vector_store_config.get("litellm_credential_name") if credential_name and litellm.credential_list: - credential_values = CredentialAccessor.get_credential_values( - credential_name - ) + credential_values = CredentialAccessor.get_credential_values(credential_name) if not credential_values: return for key, value in credential_values.items(): @@ -129,9 +125,7 @@ class BaseRAGIngestion(ABC): response.raise_for_status() file_content = response.content filename = file_url.split("/")[-1] or "document" - content_type = response.headers.get( - "content-type", "application/octet-stream" - ) + content_type = response.headers.get("content-type", "application/octet-stream") return filename, file_content, content_type, None if file_id: diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index 24452cea213..e4d636b7d30 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -92,29 +92,20 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): BaseAWSLLM.__init__(self) # Use vector_store_id as unified param (maps to knowledge_base_id) - self.knowledge_base_id = self.vector_store_config.get( - "vector_store_id" - ) or self.vector_store_config.get("knowledge_base_id") + self.knowledge_base_id = self.vector_store_config.get("vector_store_id") or self.vector_store_config.get( + "knowledge_base_id" + ) # Optional config self._data_source_id = self.vector_store_config.get("data_source_id") self._s3_bucket = self.vector_store_config.get("s3_bucket") self._s3_prefix: Optional[str] = ( - str(self.vector_store_config.get("s3_prefix")) - if self.vector_store_config.get("s3_prefix") - else None - ) - self.embedding_model = ( - self.vector_store_config.get("embedding_model") - or "amazon.titan-embed-text-v2:0" + str(self.vector_store_config.get("s3_prefix")) if self.vector_store_config.get("s3_prefix") else None ) + self.embedding_model = self.vector_store_config.get("embedding_model") or "amazon.titan-embed-text-v2:0" - self.wait_for_ingestion = self.vector_store_config.get( - "wait_for_ingestion", False - ) - self.ingestion_timeout: int = _get_int( - self.vector_store_config.get("ingestion_timeout"), 300 - ) + self.wait_for_ingestion = self.vector_store_config.get("wait_for_ingestion", False) + self.ingestion_timeout: int = _get_int(self.vector_store_config.get("ingestion_timeout"), 300) # Get AWS region using BaseAWSLLM method _aws_region = self.vector_store_config.get("aws_region_name") @@ -147,16 +138,12 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): def _auto_detect_config(self): """Auto-detect data source ID and S3 bucket from existing Knowledge Base.""" - verbose_logger.debug( - f"Auto-detecting data source and S3 bucket for KB={self.knowledge_base_id}" - ) + verbose_logger.debug(f"Auto-detecting data source and S3 bucket for KB={self.knowledge_base_id}") bedrock_agent = self._get_boto3_client("bedrock-agent") # List data sources for this KB - ds_response = bedrock_agent.list_data_sources( - knowledgeBaseId=self.knowledge_base_id - ) + ds_response = bedrock_agent.list_data_sources(knowledgeBaseId=self.knowledge_base_id) data_sources = ds_response.get("dataSourceSummaries", []) if not data_sources: @@ -178,11 +165,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): dataSourceId=self.data_source_id, ) - s3_config = ( - ds_details.get("dataSource", {}) - .get("dataSourceConfiguration", {}) - .get("s3Configuration", {}) - ) + s3_config = ds_details.get("dataSource", {}).get("dataSourceConfiguration", {}).get("s3Configuration", {}) bucket_arn = s3_config.get("bucketArn", "") if bucket_arn: @@ -220,22 +203,16 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): self.s3_bucket = self._s3_bucket or self._create_s3_bucket(unique_id) # Step 2: Create OpenSearch Serverless collection - collection_name, collection_arn = await self._create_opensearch_collection( - unique_id, account_id, caller_arn - ) + collection_name, collection_arn = await self._create_opensearch_collection(unique_id, account_id, caller_arn) # Step 3: Create OpenSearch index await self._create_opensearch_index(collection_name) # Step 4: Create IAM role for Bedrock - role_arn = await self._create_bedrock_role( - unique_id, account_id, collection_arn - ) + role_arn = await self._create_bedrock_role(unique_id, account_id, collection_arn) # Step 5: Create Knowledge Base - self.knowledge_base_id = await self._create_knowledge_base( - kb_name, role_arn, collection_arn - ) + self.knowledge_base_id = await self._create_knowledge_base(kb_name, role_arn, collection_arn) # Step 6: Create Data Source self.data_source_id = self._create_data_source(kb_name) @@ -254,9 +231,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): create_params: Dict[str, Any] = {"Bucket": bucket_name} if self.aws_region_name != "us-east-1": - create_params["CreateBucketConfiguration"] = { - "LocationConstraint": self.aws_region_name - } + create_params["CreateBucketConfiguration"] = {"LocationConstraint": self.aws_region_name} s3.create_bucket(**create_params) self._created_resources["s3_bucket"] = bucket_name @@ -264,16 +239,12 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): verbose_logger.info(f"Created S3 bucket: {bucket_name}") return bucket_name - async def _create_opensearch_collection( - self, unique_id: str, account_id: str, caller_arn: str - ) -> Tuple[str, str]: + async def _create_opensearch_collection(self, unique_id: str, account_id: str, caller_arn: str) -> Tuple[str, str]: """Create OpenSearch Serverless collection for vector storage.""" oss = self._get_boto3_client("opensearchserverless") collection_name = f"litellm-kb-{unique_id}" - verbose_logger.debug( - f"Creating OpenSearch Serverless collection: {collection_name}" - ) + verbose_logger.debug(f"Creating OpenSearch Serverless collection: {collection_name}") # Create encryption policy oss.create_security_policy( @@ -319,9 +290,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # This ensures the credentials being used have access to the collection # Normalize the caller ARN (convert assumed-role ARN to IAM role ARN if needed) normalized_caller_arn = _normalize_principal_arn(caller_arn, account_id) - verbose_logger.debug( - f"Caller ARN: {caller_arn}, Normalized: {normalized_caller_arn}" - ) + verbose_logger.debug(f"Caller ARN: {caller_arn}, Normalized: {normalized_caller_arn}") principals = [f"arn:aws:iam::{account_id}:root", normalized_caller_arn] # Deduplicate in case caller is root @@ -387,15 +356,9 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Get credentials for signing credentials = self.get_credentials( - aws_access_key_id=_get_str_or_none( - self.vector_store_config.get("aws_access_key_id") - ), - aws_secret_access_key=_get_str_or_none( - self.vector_store_config.get("aws_secret_access_key") - ), - aws_session_token=_get_str_or_none( - self.vector_store_config.get("aws_session_token") - ), + aws_access_key_id=_get_str_or_none(self.vector_store_config.get("aws_access_key_id")), + aws_secret_access_key=_get_str_or_none(self.vector_store_config.get("aws_secret_access_key")), + aws_session_token=_get_str_or_none(self.vector_store_config.get("aws_session_token")), aws_region_name=self.aws_region_name, ) @@ -454,10 +417,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): except Exception as e: last_error = e error_str = str(e) - if ( - "authorization_exception" in error_str.lower() - or "security_exception" in error_str.lower() - ): + if "authorization_exception" in error_str.lower() or "security_exception" in error_str.lower(): verbose_logger.warning( f"OpenSearch index creation attempt {attempt + 1}/{max_retries} failed due to authorization. " f"Waiting {retry_delay}s for policy propagation..." @@ -473,9 +433,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): f"Data access policy may not have propagated. Last error: {last_error}" ) - async def _create_bedrock_role( - self, unique_id: str, account_id: str, collection_arn: str - ) -> str: + async def _create_bedrock_role(self, unique_id: str, account_id: str, collection_arn: str) -> str: """Create IAM role for Bedrock KB.""" iam = self._get_boto3_client("iam") role_name = f"litellm-bedrock-kb-{unique_id}" @@ -513,9 +471,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): { "Effect": "Allow", "Action": ["bedrock:InvokeModel"], - "Resource": [ - f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}" - ], + "Resource": [f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}"], }, { "Effect": "Allow", @@ -545,9 +501,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): verbose_logger.info(f"Created IAM role: {role_arn}") return role_arn - async def _create_knowledge_base( - self, kb_name: str, role_arn: str, collection_arn: str - ) -> str: + async def _create_knowledge_base(self, kb_name: str, role_arn: str, collection_arn: str) -> str: """Create Bedrock Knowledge Base.""" bedrock_agent = self._get_boto3_client("bedrock-agent") @@ -620,40 +574,20 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): try: import boto3 except ImportError: - raise ImportError( - "boto3 is required for Bedrock ingestion. Install with: pip install boto3" - ) + raise ImportError("boto3 is required for Bedrock ingestion. Install with: pip install boto3") # Get credentials using BaseAWSLLM's get_credentials method credentials = self.get_credentials( - aws_access_key_id=_get_str_or_none( - self.vector_store_config.get("aws_access_key_id") - ), - aws_secret_access_key=_get_str_or_none( - self.vector_store_config.get("aws_secret_access_key") - ), - aws_session_token=_get_str_or_none( - self.vector_store_config.get("aws_session_token") - ), + aws_access_key_id=_get_str_or_none(self.vector_store_config.get("aws_access_key_id")), + aws_secret_access_key=_get_str_or_none(self.vector_store_config.get("aws_secret_access_key")), + aws_session_token=_get_str_or_none(self.vector_store_config.get("aws_session_token")), aws_region_name=self.aws_region_name, - aws_session_name=_get_str_or_none( - self.vector_store_config.get("aws_session_name") - ), - aws_profile_name=_get_str_or_none( - self.vector_store_config.get("aws_profile_name") - ), - aws_role_name=_get_str_or_none( - self.vector_store_config.get("aws_role_name") - ), - aws_web_identity_token=_get_str_or_none( - self.vector_store_config.get("aws_web_identity_token") - ), - aws_sts_endpoint=_get_str_or_none( - self.vector_store_config.get("aws_sts_endpoint") - ), - aws_external_id=_get_str_or_none( - self.vector_store_config.get("aws_external_id") - ), + aws_session_name=_get_str_or_none(self.vector_store_config.get("aws_session_name")), + aws_profile_name=_get_str_or_none(self.vector_store_config.get("aws_profile_name")), + aws_role_name=_get_str_or_none(self.vector_store_config.get("aws_role_name")), + aws_web_identity_token=_get_str_or_none(self.vector_store_config.get("aws_web_identity_token")), + aws_sts_endpoint=_get_str_or_none(self.vector_store_config.get("aws_sts_endpoint")), + aws_external_id=_get_str_or_none(self.vector_store_config.get("aws_external_id")), ) # Create session with credentials @@ -711,9 +645,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): await self._ensure_config_initialized() if not file_content or not filename: - verbose_logger.warning( - "No file content or filename provided for Bedrock ingestion" - ) + verbose_logger.warning("No file content or filename provided for Bedrock ingestion") return _get_str_or_none(self.knowledge_base_id), None # Step 1: Upload file to S3 @@ -732,9 +664,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Step 2: Start ingestion job bedrock_agent = self._get_boto3_client("bedrock-agent") - verbose_logger.debug( - f"Starting ingestion job for KB={self.knowledge_base_id}, DS={self.data_source_id}" - ) + verbose_logger.debug(f"Starting ingestion job for KB={self.knowledge_base_id}, DS={self.data_source_id}") ingestion_response = bedrock_agent.start_ingestion_job( knowledgeBaseId=self.knowledge_base_id, dataSourceId=self.data_source_id, @@ -763,9 +693,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) break elif status == "FAILED": - failure_reasons = job_status["ingestionJob"].get( - "failureReasons", [] - ) + failure_reasons = job_status["ingestionJob"].get("failureReasons", []) verbose_logger.error(f"Ingestion failed: {failure_reasons}") break elif status in ("STARTING", "IN_PROGRESS"): diff --git a/litellm/rag/ingestion/file_parsers/pdf_parser.py b/litellm/rag/ingestion/file_parsers/pdf_parser.py index cb42dfd5d8b..a992e957f37 100644 --- a/litellm/rag/ingestion/file_parsers/pdf_parser.py +++ b/litellm/rag/ingestion/file_parsers/pdf_parser.py @@ -37,9 +37,7 @@ def extract_text_from_pdf(file_content: bytes) -> Optional[str]: if text_parts: extracted_text = "\n\n".join(text_parts) - verbose_logger.debug( - f"Extracted {len(extracted_text)} characters from PDF using pypdf" - ) + verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using pypdf") return extracted_text except ImportError: @@ -60,15 +58,11 @@ def extract_text_from_pdf(file_content: bytes) -> Optional[str]: if text_parts: extracted_text = "\n\n".join(text_parts) - verbose_logger.debug( - f"Extracted {len(extracted_text)} characters from PDF using PyPDF2" - ) + verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using PyPDF2") return extracted_text except ImportError: - verbose_logger.debug( - "PyPDF2 not available, PDF extraction requires OCR or pypdf/PyPDF2 library" - ) + verbose_logger.debug("PyPDF2 not available, PDF extraction requires OCR or pypdf/PyPDF2 library") except Exception as e: verbose_logger.debug(f"PDF text extraction failed: {e}") diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index dd0fa94bc91..3f1d46bbb11 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -86,19 +86,11 @@ class GeminiRAGIngestion(BaseRAGIngestion): vector_store_config = cast(Dict[str, Any], self.vector_store_config) # Get API credentials - api_key = ( - cast(Optional[str], vector_store_config.get("api_key")) - or GeminiModelInfo.get_api_key() - ) - api_base = ( - cast(Optional[str], vector_store_config.get("api_base")) - or GeminiModelInfo.get_api_base() - ) + api_key = cast(Optional[str], vector_store_config.get("api_key")) or GeminiModelInfo.get_api_key() + api_base = cast(Optional[str], vector_store_config.get("api_base")) or GeminiModelInfo.get_api_base() if not api_key: - raise ValueError( - "GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search" - ) + raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search") if not api_base: raise ValueError("GEMINI_API_BASE is required") @@ -245,12 +237,8 @@ class GeminiRAGIngestion(BaseRAGIngestion): if white_space_config: request_body["chunkingConfig"] = { "whiteSpaceConfig": { - "maxTokensPerChunk": white_space_config.get( - "max_tokens_per_chunk", 800 - ), - "maxOverlapTokens": white_space_config.get( - "max_overlap_tokens", 400 - ), + "maxTokensPerChunk": white_space_config.get("max_tokens_per_chunk", 800), + "maxOverlapTokens": white_space_config.get("max_overlap_tokens", 400), } } @@ -334,9 +322,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): try: response_data = response.json() # The response should contain the document name or file reference - file_id = response_data.get("name", "") or response_data.get( - "file", {} - ).get("name", "") + file_id = response_data.get("name", "") or response_data.get("file", {}).get("name", "") verbose_logger.debug(f"Upload complete. File ID: {file_id}") return file_id except Exception as e: diff --git a/litellm/rag/ingestion/openai_ingestion.py b/litellm/rag/ingestion/openai_ingestion.py index 61fe7e17ea3..ca5575a7e30 100644 --- a/litellm/rag/ingestion/openai_ingestion.py +++ b/litellm/rag/ingestion/openai_ingestion.py @@ -87,15 +87,11 @@ class OpenAIRAGIngestion(BaseRAGIngestion): api_base = self.vector_store_config.get("api_base") if existing_file_id and not vector_store_id: - raise ValueError( - "vector_store_id is required when ingesting an existing file_id" - ) + raise ValueError("vector_store_id is required when ingesting an existing file_id") # Create vector store if not provided if not vector_store_id: - expires_after = ( - {"anchor": "last_active_at", "days": ttl_days} if ttl_days else None - ) + expires_after = {"anchor": "last_active_at", "days": ttl_days} if ttl_days else None create_response = await vector_store_acreate( name=self.ingest_name or "litellm-rag-ingest", custom_llm_provider="openai", diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 0a5defce962..3ec623657ee 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -68,9 +68,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Extract config self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] self.index_name = self.vector_store_config.get("index_name") - self.distance_metric = self.vector_store_config.get( - "distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC - ) + self.distance_metric = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) self.non_filterable_metadata_keys = self.vector_store_config.get( "non_filterable_metadata_keys", S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, @@ -86,9 +84,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) # Create httpx client (similar to s3_v2.py) - ssl_verify = self._get_ssl_verify( - ssl_verify=self.vector_store_config.get("ssl_verify") - ) + ssl_verify = self._get_ssl_verify(ssl_verify=self.vector_store_config.get("ssl_verify")) self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.RAG, params={"ssl_verify": ssl_verify} if ssl_verify is not None else None, @@ -109,27 +105,19 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): try: model_name = self.embedding_config["model"] - verbose_logger.debug( - f"Auto-detecting dimension by making test embedding request to {model_name}" - ) + verbose_logger.debug(f"Auto-detecting dimension by making test embedding request to {model_name}") # Make a test embedding request test_input = "test" if self.router: - response = await self.router.aembedding( - model=model_name, input=[test_input] - ) + response = await self.router.aembedding(model=model_name, input=[test_input]) else: - response = await litellm.aembedding( - model=model_name, input=[test_input] - ) + response = await litellm.aembedding(model=model_name, input=[test_input]) # Get dimension from the response if response.data and len(response.data) > 0: dimension = len(response.data[0]["embedding"]) - verbose_logger.debug( - f"Auto-detected dimension {dimension} for embedding model {model_name}" - ) + verbose_logger.debug(f"Auto-detected dimension {dimension} for embedding model {model_name}") return dimension except Exception as e: verbose_logger.warning( @@ -188,9 +176,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError( - "Missing botocore to call S3 Vectors. Run 'pip install boto3'." - ) + raise ImportError("Missing botocore to call S3 Vectors. Run 'pip install boto3'.") # Get AWS credentials using BaseAWSLLM's get_credentials method credentials = self.get_credentials( @@ -201,9 +187,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): aws_session_name=self.vector_store_config.get("aws_session_name"), aws_profile_name=self.vector_store_config.get("aws_profile_name"), aws_role_name=self.vector_store_config.get("aws_role_name"), - aws_web_identity_token=self.vector_store_config.get( - "aws_web_identity_token" - ), + aws_web_identity_token=self.vector_store_config.get("aws_web_identity_token"), aws_sts_endpoint=self.vector_store_config.get("aws_sts_endpoint"), aws_external_id=self.vector_store_config.get("aws_external_id"), ) @@ -240,13 +224,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Make the request using specific method (pattern from s3_v2.py) method_upper = method.upper() if method_upper == "PUT": - response = await self.async_httpx_client.put( - url, data=data, headers=signed_headers - ) + response = await self.async_httpx_client.put(url, data=data, headers=signed_headers) elif method_upper == "POST": - response = await self.async_httpx_client.post( - url, data=data, headers=signed_headers - ) + response = await self.async_httpx_client.post(url, data=data, headers=signed_headers) elif method_upper == "GET": response = await self.async_httpx_client.get(url, headers=signed_headers) else: @@ -256,9 +236,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): async def _ensure_vector_bucket_exists(self): """Create vector bucket if it doesn't exist using GetVectorBucket and CreateVectorBucket APIs.""" - verbose_logger.debug( - f"Ensuring S3 vector bucket exists: {self.vector_bucket_name}" - ) + verbose_logger.debug(f"Ensuring S3 vector bucket exists: {self.vector_bucket_name}") # Validate bucket name (AWS S3 naming rules) if len(self.vector_bucket_name) < 3: @@ -279,40 +257,28 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): get_body = safe_dumps({"vectorBucketName": self.vector_bucket_name}) try: - response = await self._sign_and_execute_request( - "POST", get_url, data=get_body - ) + response = await self._sign_and_execute_request("POST", get_url, data=get_body) if response.status_code == 200: verbose_logger.debug(f"Vector bucket {self.vector_bucket_name} exists") return except Exception as e: - verbose_logger.debug( - f"Bucket check failed (may not exist): {e}, attempting to create" - ) + verbose_logger.debug(f"Bucket check failed (may not exist): {e}, attempting to create") # Create vector bucket using CreateVectorBucket API try: verbose_logger.debug(f"Creating vector bucket: {self.vector_bucket_name}") - create_url = ( - f"https://s3vectors.{self.aws_region_name}.api.aws/CreateVectorBucket" - ) + create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateVectorBucket" create_body = safe_dumps({"vectorBucketName": self.vector_bucket_name}) - response = await self._sign_and_execute_request( - "POST", create_url, data=create_body - ) + response = await self._sign_and_execute_request("POST", create_url, data=create_body) if response.status_code in (200, 201): verbose_logger.info(f"Created vector bucket: {self.vector_bucket_name}") elif response.status_code == 409: # Bucket already exists (ConflictException) - verbose_logger.debug( - f"Vector bucket {self.vector_bucket_name} already exists" - ) + verbose_logger.debug(f"Vector bucket {self.vector_bucket_name} already exists") else: - verbose_logger.error( - f"CreateVectorBucket failed: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"CreateVectorBucket failed: {response.status_code} - {response.text}") response.raise_for_status() except Exception as e: verbose_logger.exception(f"Error creating vector bucket: {e}") @@ -320,27 +286,19 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): async def _ensure_vector_index_exists(self): """Create vector index if it doesn't exist using GetIndex and CreateIndex APIs.""" - verbose_logger.debug( - f"Ensuring vector index exists: {self.vector_bucket_name}/{self.index_name}" - ) + verbose_logger.debug(f"Ensuring vector index exists: {self.vector_bucket_name}/{self.index_name}") # Try to get index info using GetIndex API get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetIndex" - get_body = safe_dumps( - {"vectorBucketName": self.vector_bucket_name, "indexName": self.index_name} - ) + get_body = safe_dumps({"vectorBucketName": self.vector_bucket_name, "indexName": self.index_name}) try: - response = await self._sign_and_execute_request( - "POST", get_url, data=get_body - ) + response = await self._sign_and_execute_request("POST", get_url, data=get_body) if response.status_code == 200: verbose_logger.debug(f"Vector index {self.index_name} exists") return except Exception as e: - verbose_logger.debug( - f"Index check failed (may not exist): {e}, attempting to create" - ) + verbose_logger.debug(f"Index check failed (may not exist): {e}, attempting to create") # Create vector index using CreateIndex API try: @@ -358,23 +316,17 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): } if self.non_filterable_metadata_keys: - index_config["metadataConfiguration"] = { - "nonFilterableMetadataKeys": self.non_filterable_metadata_keys - } + index_config["metadataConfiguration"] = {"nonFilterableMetadataKeys": self.non_filterable_metadata_keys} create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateIndex" - response = await self._sign_and_execute_request( - "POST", create_url, data=safe_dumps(index_config) - ) + response = await self._sign_and_execute_request("POST", create_url, data=safe_dumps(index_config)) if response.status_code in (200, 201): verbose_logger.info(f"Created vector index: {self.index_name}") elif response.status_code == 409: verbose_logger.debug(f"Vector index {self.index_name} already exists") else: - verbose_logger.error( - f"CreateIndex failed: {response.status_code} - {response.text}" - ) + verbose_logger.error(f"CreateIndex failed: {response.status_code} - {response.text}") response.raise_for_status() except Exception as e: verbose_logger.exception(f"Error creating vector index: {e}") @@ -387,9 +339,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): Args: vectors: List of vector objects with keys: "key", "data", "metadata" """ - verbose_logger.debug( - f"Storing {len(vectors)} vectors in {self.vector_bucket_name}/{self.index_name}" - ) + verbose_logger.debug(f"Storing {len(vectors)} vectors in {self.vector_bucket_name}/{self.index_name}") url = f"https://s3vectors.{self.aws_region_name}.api.aws/PutVectors" @@ -401,18 +351,12 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): } try: - response = await self._sign_and_execute_request( - "POST", url, data=safe_dumps(request_body) - ) + response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body)) if response.status_code in (200, 201): - verbose_logger.info( - f"Successfully stored {len(vectors)} vectors in index {self.index_name}" - ) + verbose_logger.info(f"Successfully stored {len(vectors)} vectors in index {self.index_name}") else: - verbose_logger.error( - f"PutVectors failed with status {response.status_code}: {response.text}" - ) + verbose_logger.error(f"PutVectors failed with status {response.status_code}: {response.text}") response.raise_for_status() except Exception as e: verbose_logger.exception(f"Error storing vectors: {e}") @@ -432,28 +376,20 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Use embedding config from ingest_options or default if not self.embedding_config: - verbose_logger.warning( - "No embedding config provided, using default text-embedding-3-small" - ) + verbose_logger.warning("No embedding config provided, using default text-embedding-3-small") self.embedding_config = {"model": "text-embedding-3-small"} embedding_model = self.embedding_config.get("model", "text-embedding-3-small") - verbose_logger.debug( - f"Generating embeddings for {len(chunks)} chunks using {embedding_model}" - ) + verbose_logger.debug(f"Generating embeddings for {len(chunks)} chunks using {embedding_model}") # Convert to list to ensure type compatibility input_chunks: List[str] = list(chunks) if self.router: - response = await self.router.aembedding( - model=embedding_model, input=input_chunks - ) + response = await self.router.aembedding(model=embedding_model, input=input_chunks) else: - response = await litellm.aembedding( - model=embedding_model, input=input_chunks - ) + response = await litellm.aembedding(model=embedding_model, input=input_chunks) return [item["embedding"] for item in response.data] @@ -528,9 +464,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): vector_store_id = f"{self.vector_bucket_name}:{self.index_name}" return vector_store_id, filename - async def query_vector_store( - self, vector_store_id: str, query: str, top_k: int = 5 - ) -> Optional[Dict[str, Any]]: + async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> Optional[Dict[str, Any]]: """ Query S3 Vectors using QueryVectors API. @@ -566,15 +500,11 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): } try: - response = await self._sign_and_execute_request( - "POST", url, data=safe_dumps(request_body) - ) + response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body)) if response.status_code == 200: results = response.json() - verbose_logger.debug( - f"Query returned {len(results.get('vectors', []))} results" - ) + verbose_logger.debug(f"Query returned {len(results.get('vectors', []))} results") # Check if query terms appear in results if results.get("vectors"): @@ -587,9 +517,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Return results even if exact match not found return results else: - verbose_logger.error( - f"QueryVectors failed with status {response.status_code}: {response.text}" - ) + verbose_logger.error(f"QueryVectors failed with status {response.status_code}: {response.text}") return None except Exception as e: verbose_logger.exception(f"Error querying vectors: {e}") diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index b7bd87d1f6c..34cd1a88a61 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -95,10 +95,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): Tuple of (rag_corpus_id, file_id) """ if not self.project_id: - raise ValueError( - "vertex_project is required for Vertex AI RAG ingestion. " - "Set it in vector_store config." - ) + raise ValueError("vertex_project is required for Vertex AI RAG ingestion. Set it in vector_store config.") # Get or create RAG corpus rag_corpus_id = self.vector_store_config.get("vector_store_id") @@ -148,10 +145,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Construct URL using vertex base URL helper base_url = get_vertex_base_url(self.location) - url = ( - f"{base_url}/v1beta1/" - f"projects/{self.project_id}/locations/{self.location}/ragCorpora" - ) + url = f"{base_url}/v1beta1/projects/{self.project_id}/locations/{self.location}/ragCorpora" # Build request body with camelCase keys (Vertex AI API format) request_body: Dict[str, Any] = { @@ -197,9 +191,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): raise Exception(error_msg) response_data = response.json() - verbose_logger.debug( - f"Create corpus response: {json.dumps(response_data, indent=2)}" - ) + verbose_logger.debug(f"Create corpus response: {json.dumps(response_data, indent=2)}") # The response is a long-running operation # Check if it's already done or if we need to poll @@ -278,13 +270,9 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): if corpus_name: return corpus_name else: - raise Exception( - f"No corpus name in operation response: {operation_data}" - ) + raise Exception(f"No corpus name in operation response: {operation_data}") - verbose_logger.debug( - f"Operation not done yet, attempt {attempt + 1}/{max_retries}" - ) + verbose_logger.debug(f"Operation not done yet, attempt {attempt + 1}/{max_retries}") await asyncio.sleep(retry_delay) raise Exception(f"Operation timed out after {max_retries} attempts") @@ -345,9 +333,9 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): "rag_file_chunking_config": {"fixed_length_chunking": {}} } - chunking_config = metadata["upload_rag_file_config"][ - "rag_file_transformation_config" - ]["rag_file_chunking_config"]["fixed_length_chunking"] + chunking_config = metadata["upload_rag_file_config"]["rag_file_transformation_config"][ + "rag_file_chunking_config" + ]["fixed_length_chunking"] if chunk_size: chunking_config["chunk_size"] = chunk_size @@ -426,9 +414,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): url = f"{base_url}/v1beta1/{rag_corpus_id}/ragFiles:import" # Build request body with camelCase keys (Vertex AI API format) - request_body: Dict[str, Any] = { - "importRagFilesConfig": {"gcsSource": {"uris": gcs_uris}} - } + request_body: Dict[str, Any] = {"importRagFilesConfig": {"gcsSource": {"uris": gcs_uris}}} # Add chunking configuration if provided chunking_strategy = self.chunking_strategy @@ -443,13 +429,9 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): } # Add max embedding requests per minute if specified - max_embedding_qpm = self.vector_store_config.get( - "max_embedding_requests_per_min" - ) + max_embedding_qpm = self.vector_store_config.get("max_embedding_requests_per_min") if max_embedding_qpm: - request_body["importRagFilesConfig"]["maxEmbeddingRequestsPerMin"] = ( - max_embedding_qpm - ) + request_body["importRagFilesConfig"]["maxEmbeddingRequestsPerMin"] = max_embedding_qpm verbose_logger.debug(f"Importing files from GCS: {url}") verbose_logger.debug(f"Request body: {json.dumps(request_body, indent=2)}") diff --git a/litellm/rag/main.py b/litellm/rag/main.py index e3d354b6c33..6b5f087f902 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -71,10 +71,7 @@ def get_ingestion_class(provider: str) -> Type[BaseRAGIngestion]: ingestion_class = INGESTION_REGISTRY.get(provider) if ingestion_class is None: supported = ", ".join(INGESTION_REGISTRY.keys()) - raise ValueError( - f"Provider '{provider}' is not supported for RAG ingestion. " - f"Supported providers: {supported}" - ) + raise ValueError(f"Provider '{provider}' is not supported for RAG ingestion. Supported providers: {supported}") return ingestion_class @@ -184,9 +181,7 @@ async def aingest( except Exception as e: raise litellm.exception_type( model=None, - custom_llm_provider=ingest_options.get("vector_store", {}).get( - "custom_llm_provider" - ), + custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"), original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, @@ -235,9 +230,7 @@ async def _execute_query_pipeline( documents=documents, top_n=rerank.get("top_n", 5), ) - context_chunks = RAGQuery.get_top_chunks_from_rerank( - search_response, rerank_response - ) + context_chunks = RAGQuery.get_top_chunks_from_rerank(search_response, rerank_response) # 4. Build context message and call completion context_message = RAGQuery.build_context_message(context_chunks) @@ -432,9 +425,7 @@ def ingest( except Exception as e: raise litellm.exception_type( model=None, - custom_llm_provider=ingest_options.get("vector_store", {}).get( - "custom_llm_provider" - ), + custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"), original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index ebbf209e815..65cd8a3572e 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -30,11 +30,7 @@ class RAGQuery: elif isinstance(content, list) and len(content) > 0: # Handle list of content items, extract text from first text item for item in content: - if ( - isinstance(item, dict) - and item.get("type") == "text" - and "text" in item - ): + if isinstance(item, dict) and item.get("type") == "text" and "text" in item: return item["text"] return None @@ -48,9 +44,7 @@ class RAGQuery: for chunk in context_chunks: if isinstance(chunk, dict): - result_content: Optional[List[VectorStoreResultContent]] = chunk.get( - "content" - ) + result_content: Optional[List[VectorStoreResultContent]] = chunk.get("content") if result_content: for content_item in result_content: content_text: Optional[str] = content_item.get("text") @@ -80,9 +74,7 @@ class RAGQuery: message = getattr(choice, "message", None) if message is not None: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(message, "provider_specific_fields", None) or {} - ) + provider_fields = getattr(message, "provider_specific_fields", None) or {} # Add search results provider_fields["search_results"] = search_results @@ -107,9 +99,7 @@ class RAGQuery: return documents @staticmethod - def get_top_chunks_from_rerank( - search_response: Any, rerank_response: Any - ) -> List[Any]: + def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> List[Any]: """Get the original search results corresponding to the top reranked results.""" top_chunks = [] original_results = search_response.get("data", []) diff --git a/litellm/rag/text_splitters/recursive_character_text_splitter.py b/litellm/rag/text_splitters/recursive_character_text_splitter.py index edf6b84312b..48041172369 100644 --- a/litellm/rag/text_splitters/recursive_character_text_splitter.py +++ b/litellm/rag/text_splitters/recursive_character_text_splitter.py @@ -31,18 +31,13 @@ class RecursiveCharacterTextSplitter: """Split text into chunks.""" return self._split_text(text, self.separators) - def _split_text( - self, text: str, separators: List[str], depth: int = 0 - ) -> List[str]: + def _split_text(self, text: str, separators: List[str], depth: int = 0) -> List[str]: """Recursively split text using separators.""" from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH if depth > DEFAULT_MAX_RECURSE_DEPTH: # Max depth reached, return text as-is split into chunk_size pieces - return [ - text[i : i + self.chunk_size] - for i in range(0, len(text), self.chunk_size) - ] + return [text[i : i + self.chunk_size] for i in range(0, len(text), self.chunk_size)] final_chunks: List[str] = [] @@ -109,9 +104,7 @@ class RecursiveCharacterTextSplitter: chunks.append(chunk_text) # Handle overlap - while ( - current_length > self.chunk_overlap and len(current_chunk) > 1 - ): + while current_length > self.chunk_overlap and len(current_chunk) > 1: removed = current_chunk.pop(0) current_length -= len(removed) + len(separator) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index f6f0a92def0..a6fec4729ad 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -43,11 +43,7 @@ base_llm_http_handler = BaseLLMHTTPHandler() def _build_litellm_metadata(kwargs: dict) -> dict: """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" metadata: dict = {**(kwargs.get("litellm_metadata") or {})} - guardrails = ( - (kwargs.get("metadata") or {}).get("guardrails") - or kwargs.get("guardrails") - or [] - ) + guardrails = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or [] if guardrails: metadata["guardrails"] = guardrails return metadata @@ -87,11 +83,7 @@ def _get_realtime_http_provider_config( # Fallback for providers without a dedicated HTTP config (treated as OpenAI-compatible). resolved_api_base = raw_api_base or litellm.api_base or "https://api.openai.com" resolved_api_key = ( - raw_api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - or "" + raw_api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") or "" ) return provider_config, resolved_api_base.rstrip("/"), resolved_api_key @@ -110,11 +102,7 @@ async def acreate_realtime_client_secret( session=RealtimeSessionConfig(**session) if session else None, expires_after=RealtimeExpiresAfter(**expires_after) if expires_after else None, ) - model_name = ( - (req.session.model if req.session is not None else None) - or req.model - or "gpt-4o-realtime-preview" - ) + model_name = (req.session.model if req.session is not None else None) or req.model or "gpt-4o-realtime-preview" litellm_logging_obj: LiteLLMLogging = kwargs.get("litellm_logging_obj") # type: ignore litellm_params = GenericLiteLLMParams(**kwargs) @@ -358,19 +346,9 @@ async def _arealtime( query_params=query_params, ) elif _custom_llm_provider == "azure": - api_base = ( - dynamic_api_base - or litellm_params.api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) + 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") - ) + api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_API_KEY") api_version = api_version or litellm_params.api_version or "2024-10-01-preview" @@ -379,10 +357,7 @@ async def _arealtime( or litellm_params.get("realtime_protocol") or os.environ.get("LITELLM_AZURE_REALTIME_PROTOCOL") ) - if ( - realtime_protocol is None - and (query_params or {}).get("intent") == "transcription" - ): + if realtime_protocol is None and (query_params or {}).get("intent") == "transcription": realtime_protocol = "GA" realtime_protocol = realtime_protocol or "beta" await azure_realtime.async_realtime( @@ -401,19 +376,9 @@ async def _arealtime( litellm_metadata=_build_litellm_metadata(kwargs), ) elif _custom_llm_provider == "openai": - api_base = ( - dynamic_api_base - or litellm_params.api_base - or litellm.api_base - or "https://api.openai.com/" - ) + api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or "https://api.openai.com/" # set API KEY - api_key = ( - dynamic_api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") - ) + api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") await openai_realtime.async_realtime( model=model, @@ -462,15 +427,10 @@ async def _arealtime( ) elif _custom_llm_provider == "xai": api_base = ( - dynamic_api_base - or litellm_params.api_base - or get_secret_str("XAI_API_BASE") - or "https://api.x.ai/v1" + dynamic_api_base or litellm_params.api_base or get_secret_str("XAI_API_BASE") or "https://api.x.ai/v1" ) # set API KEY - api_key = XAIModelInfo.get_api_key( - dynamic_api_key, legacy_generic_before_env=True - ) + api_key = XAIModelInfo.get_api_key(dynamic_api_key, legacy_generic_before_env=True) await xai_realtime.async_realtime( model=model, @@ -503,9 +463,7 @@ async def _arealtime( or get_secret_str("VERTEXAI_LOCATION") ) - resolved_location = vertex_llm_base.get_vertex_region( - vertex_region=vertex_location, model=model - ) + resolved_location = vertex_llm_base.get_vertex_region(vertex_region=vertex_location, model=model) ( access_token, @@ -580,14 +538,10 @@ async def _realtime_health_check( query_params={"model": model}, ) elif custom_llm_provider == "xai": - url = xai_realtime._construct_url( - api_base=api_base or "https://api.x.ai/v1", query_params={"model": model} - ) + url = xai_realtime._construct_url(api_base=api_base or "https://api.x.ai/v1", query_params={"model": model}) elif custom_llm_provider == "vertex_ai": vertex_location = litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") - resolved_location = vertex_llm_base.get_vertex_region( - vertex_region=vertex_location, model=model - ) + resolved_location = vertex_llm_base.get_vertex_region(vertex_region=vertex_location, model=model) ( access_token, resolved_project, @@ -603,9 +557,7 @@ async def _realtime_health_check( ) url = vertex_realtime_config.get_complete_url(api_base=api_base, model=model) ssl_context = get_shared_realtime_ssl_context() - headers = vertex_realtime_config.validate_environment( - headers={}, model=model, api_key=None - ) + headers = vertex_realtime_config.validate_environment(headers={}, model=model, api_key=None) async with websockets.connect( # type: ignore url, additional_headers=headers, diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py index a25620c7b4d..40aeb6df3de 100644 --- a/litellm/repositories/base_repository.py +++ b/litellm/repositories/base_repository.py @@ -29,9 +29,7 @@ class BaseRepository(ABC, Generic[T]): @property def prisma_client(self) -> Any: if self._prisma_client is None: - raise RuntimeError( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client @property @@ -95,9 +93,7 @@ class BaseRepository(ABC, Generic[T]): assert model is not None return model - async def update( - self, id_value: str, data: Dict[str, Any], id_field: str = "id" - ) -> Optional[T]: + async def update(self, id_value: str, data: Dict[str, Any], id_field: str = "id") -> Optional[T]: """Update an existing record.""" record = await self.table.update(where={id_field: id_value}, data=data) return self._to_model(record) diff --git a/litellm/repositories/budget_repository.py b/litellm/repositories/budget_repository.py index 5947701fb4e..aa2676fa72b 100644 --- a/litellm/repositories/budget_repository.py +++ b/litellm/repositories/budget_repository.py @@ -19,9 +19,7 @@ class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): def model_class(self) -> Type[LiteLLM_BudgetTable]: return LiteLLM_BudgetTable - async def find_by_id( - self, budget_id: str, id_field: str = "budget_id" - ) -> Optional[LiteLLM_BudgetTable]: + async def find_by_id(self, budget_id: str, id_field: str = "budget_id") -> Optional[LiteLLM_BudgetTable]: return await super().find_by_id(budget_id, id_field) async def create_budget( diff --git a/litellm/repositories/config_repository.py b/litellm/repositories/config_repository.py index eba7ebe26ca..5af78bf1a6c 100644 --- a/litellm/repositories/config_repository.py +++ b/litellm/repositories/config_repository.py @@ -40,9 +40,7 @@ class ConfigRepository: @property def prisma_client(self) -> Any: if self._prisma_client is None: - raise RuntimeError( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client @property @@ -61,9 +59,7 @@ class ConfigRepository: async def set_param(self, param_name: str, param_value: Any) -> ConfigParam: """Set a config parameter in the database.""" - value_json = ( - json.dumps(param_value) if not isinstance(param_value, str) else param_value - ) + value_json = json.dumps(param_value) if not isinstance(param_value, str) else param_value await self.table.upsert( where={"param_name": param_name}, data={ @@ -111,9 +107,7 @@ class ConfigRepository: else: d[k] = v - def _decrypt_env_variables( - self, env_vars: Dict[str, Any], return_original_value: bool = True - ) -> Dict[str, str]: + def _decrypt_env_variables(self, env_vars: Dict[str, Any], return_original_value: bool = True) -> Dict[str, str]: """Decrypt environment variables from database.""" decrypted: Dict[str, str] = {} for key, value in env_vars.items(): @@ -152,25 +146,19 @@ class ConfigRepository: ) -> dict: """Update config fields with DB values, handling the merge strategy.""" if param_name == "environment_variables": - decrypted_env_vars = self._decrypt_env_variables( - db_param_value, return_original_value=True - ) + decrypted_env_vars = self._decrypt_env_variables(db_param_value, return_original_value=True) merged_env_vars = self._normalize_env_variable_keys(decrypted_env_vars) for env_key, value in merged_env_vars.items(): os.environ[env_key] = value - current_config.setdefault("environment_variables", {}).update( - merged_env_vars - ) + current_config.setdefault("environment_variables", {}).update(merged_env_vars) return current_config if param_name not in current_config: current_config[param_name] = db_param_value return current_config - if isinstance(current_config[param_name], dict) and isinstance( - db_param_value, dict - ): + if isinstance(current_config[param_name], dict) and isinstance(db_param_value, dict): self._deep_merge_dicts(current_config[param_name], db_param_value) else: current_config[param_name] = db_param_value @@ -196,9 +184,7 @@ class ConfigRepository: The merged configuration with DB overrides applied """ if store_model_in_db is not True: - verbose_proxy_logger.info( - "'store_model_in_db' is not True, skipping db config reconciliation" - ) + verbose_proxy_logger.info("'store_model_in_db' is not True, skipping db config reconciliation") return yaml_config tasks = [self.get_param(k) for k in self.CONFIG_PARAMS] @@ -211,9 +197,7 @@ class ConfigRepository: param_name = response.param_name param_value = response.param_value - verbose_proxy_logger.debug( - f"param_name={param_name}, param_value={param_value}" - ) + verbose_proxy_logger.debug(f"param_name={param_name}, param_value={param_value}") if param_name is not None and param_value is not None: config = self._update_config_fields( diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index dd53c753307..b5a315d233c 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -20,9 +20,7 @@ class CredentialsRepository: @property def prisma_client(self) -> Any: if self._prisma_client is None: - raise RuntimeError( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client @property @@ -47,15 +45,11 @@ class CredentialsRepository: return await self.table.create(data=data) async def find_by_name(self, credential_name: str) -> Optional[CredentialItem]: - record = await self.table.find_unique( - where={"credential_name": credential_name} - ) + record = await self.table.find_unique(where={"credential_name": credential_name}) return self._to_model(record) async def update_by_name(self, credential_name: str, data: Dict[str, Any]) -> Any: - return await self.table.update( - where={"credential_name": credential_name}, data=data - ) + return await self.table.update(where={"credential_name": credential_name}, data=data) async def delete_by_name(self, credential_name: str) -> Any: return await self.table.delete(where={"credential_name": credential_name}) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 893cf342d71..0da51519964 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -33,9 +33,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): encrypted = {} for key, value in litellm_params.items(): if isinstance(value, str): - encrypted[key] = encrypt_value_helper( - value, new_encryption_key=self._encryption_key - ) + encrypted[key] = encrypt_value_helper(value, new_encryption_key=self._encryption_key) else: encrypted[key] = value return encrypted @@ -65,15 +63,11 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): data["model_info"] = json.loads(data["model_info"]) if data.get("litellm_params"): - data["litellm_params"] = self._decrypt_litellm_params( - data["litellm_params"] - ) + data["litellm_params"] = self._decrypt_litellm_params(data["litellm_params"]) return LiteLLM_ProxyModelTable(**data) - async def find_by_id( - self, model_id: str, id_field: str = "model_id" - ) -> Optional[LiteLLM_ProxyModelTable]: + async def find_by_id(self, model_id: str, id_field: str = "model_id") -> Optional[LiteLLM_ProxyModelTable]: return await super().find_by_id(model_id, id_field) async def find_by_name(self, model_name: str) -> List[LiteLLM_ProxyModelTable]: @@ -158,14 +152,10 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): """Delete a model.""" return await self.delete(model_id, id_field="model_id") - async def block_model( - self, model_id: str, updated_by: str - ) -> Optional[LiteLLM_ProxyModelTable]: + async def block_model(self, model_id: str, updated_by: str) -> Optional[LiteLLM_ProxyModelTable]: """Block a model.""" return await self.update_model(model_id, updated_by, blocked=True) - async def unblock_model( - self, model_id: str, updated_by: str - ) -> Optional[LiteLLM_ProxyModelTable]: + async def unblock_model(self, model_id: str, updated_by: str) -> Optional[LiteLLM_ProxyModelTable]: """Unblock a model.""" return await self.update_model(model_id, updated_by, blocked=False) diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py index f4d9a8bb90a..063a17cfe92 100644 --- a/litellm/repositories/object_permission_repository.py +++ b/litellm/repositories/object_permission_repository.py @@ -99,12 +99,8 @@ class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): if search_tools is not None: data["search_tools"] = search_tools - return await self.update( - object_permission_id, data, id_field="object_permission_id" - ) + return await self.update(object_permission_id, data, id_field="object_permission_id") - async def delete_permission( - self, object_permission_id: str - ) -> Optional[LiteLLM_ObjectPermissionTable]: + async def delete_permission(self, object_permission_id: str) -> Optional[LiteLLM_ObjectPermissionTable]: """Delete an object permission record.""" return await self.delete(object_permission_id, id_field="object_permission_id") diff --git a/litellm/repositories/organization_repository.py b/litellm/repositories/organization_repository.py index 2d25a43e836..d5f8c990001 100644 --- a/litellm/repositories/organization_repository.py +++ b/litellm/repositories/organization_repository.py @@ -24,13 +24,9 @@ class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): ) -> Optional[LiteLLM_OrganizationTable]: return await super().find_by_id(organization_id, id_field) - async def find_by_alias( - self, organization_alias: str - ) -> Optional[LiteLLM_OrganizationTable]: + async def find_by_alias(self, organization_alias: str) -> Optional[LiteLLM_OrganizationTable]: """Find an organization by alias.""" - records = await self.table.find_many( - where={"organization_alias": organization_alias} - ) + records = await self.table.find_many(where={"organization_alias": organization_alias}) if records: return self._to_model(records[0]) return None @@ -88,16 +84,10 @@ class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): return await self.update(organization_id, data, id_field="organization_id") - async def delete_organization( - self, organization_id: str - ) -> Optional[LiteLLM_OrganizationTable]: + async def delete_organization(self, organization_id: str) -> Optional[LiteLLM_OrganizationTable]: """Delete an organization.""" return await self.delete(organization_id, id_field="organization_id") - async def update_spend( - self, organization_id: str, spend: float - ) -> Optional[LiteLLM_OrganizationTable]: + async def update_spend(self, organization_id: str, spend: float) -> Optional[LiteLLM_OrganizationTable]: """Update organization spend.""" - return await self.update( - organization_id, {"spend": spend}, id_field="organization_id" - ) + return await self.update(organization_id, {"spend": spend}, id_field="organization_id") diff --git a/litellm/repositories/project_repository.py b/litellm/repositories/project_repository.py index 86567dd05fb..86faaf2e13c 100644 --- a/litellm/repositories/project_repository.py +++ b/litellm/repositories/project_repository.py @@ -19,9 +19,7 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): def model_class(self) -> Type[LiteLLM_ProjectTable]: return LiteLLM_ProjectTable - async def find_by_id( - self, project_id: str, id_field: str = "project_id" - ) -> Optional[LiteLLM_ProjectTable]: + async def find_by_id(self, project_id: str, id_field: str = "project_id") -> Optional[LiteLLM_ProjectTable]: return await super().find_by_id(project_id, id_field) async def find_by_alias(self, project_alias: str) -> Optional[LiteLLM_ProjectTable]: @@ -122,8 +120,6 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): """Delete a project.""" return await self.delete(project_id, id_field="project_id") - async def update_spend( - self, project_id: str, spend: float - ) -> Optional[LiteLLM_ProjectTable]: + async def update_spend(self, project_id: str, spend: float) -> Optional[LiteLLM_ProjectTable]: """Update project spend.""" return await self.update(project_id, {"spend": spend}, id_field="project_id") diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 47ea11c0592..7ce4607e1ca 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -21,9 +21,7 @@ class PrismaTableRepository: @property def prisma_client(self) -> Any: if self._prisma_client is None: - raise RuntimeError( - "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" - ) + raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") return self._prisma_client @property diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index 2ae6647060c..3227aa812ca 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -46,9 +46,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return LiteLLM_TeamTable(**data) - async def find_by_id( - self, team_id: str, id_field: str = "team_id" - ) -> Optional[LiteLLM_TeamTable]: + async def find_by_id(self, team_id: str, id_field: str = "team_id") -> Optional[LiteLLM_TeamTable]: return await super().find_by_id(team_id, id_field) async def find_by_alias(self, team_alias: str) -> Optional[LiteLLM_TeamTable]: @@ -58,9 +56,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return self._to_model(records[0]) return None - async def find_by_organization_id( - self, organization_id: str - ) -> List[LiteLLM_TeamTable]: + async def find_by_organization_id(self, organization_id: str) -> List[LiteLLM_TeamTable]: """Find all teams belonging to an organization.""" records = await self.table.find_many(where={"organization_id": organization_id}) return self._to_model_list(records) @@ -219,9 +215,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): data["admins"] = team.admins data["members"] = team.members if team.members_with_roles: - data["members_with_roles"] = json.dumps( - [m.model_dump() for m in team.members_with_roles] - ) + data["members_with_roles"] = json.dumps([m.model_dump() for m in team.members_with_roles]) if team.metadata: data["metadata"] = json.dumps(team.metadata) if team.max_budget is not None: @@ -255,15 +249,11 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): data["allow_team_guardrail_config"] = team.allow_team_guardrail_config return data - async def update_spend( - self, team_id: str, spend: float - ) -> Optional[LiteLLM_TeamTable]: + async def update_spend(self, team_id: str, spend: float) -> Optional[LiteLLM_TeamTable]: """Update team spend.""" return await self.update(team_id, {"spend": spend}, id_field="team_id") - async def add_member( - self, team_id: str, user_id: str - ) -> Optional[LiteLLM_TeamTable]: + async def add_member(self, team_id: str, user_id: str) -> Optional[LiteLLM_TeamTable]: """Add a member to a team using atomic array push operation.""" if not await self.exists(team_id, id_field="team_id"): return None @@ -274,9 +264,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) return self._to_model(record) - async def remove_member( - self, team_id: str, user_id: str - ) -> Optional[LiteLLM_TeamTable]: + async def remove_member(self, team_id: str, user_id: str) -> Optional[LiteLLM_TeamTable]: """Remove a member from a team. Note: Prisma doesn't support atomic array removal, so we use a @@ -290,9 +278,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): members = [m for m in team.members if m != user_id] return await self.update(team_id, {"members": members}, id_field="team_id") - async def add_admin( - self, team_id: str, user_id: str - ) -> Optional[LiteLLM_TeamTable]: + async def add_admin(self, team_id: str, user_id: str) -> Optional[LiteLLM_TeamTable]: """Add an admin to a team using atomic array push operation.""" if not await self.exists(team_id, id_field="team_id"): return None @@ -303,9 +289,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) return self._to_model(record) - async def remove_admin( - self, team_id: str, user_id: str - ) -> Optional[LiteLLM_TeamTable]: + async def remove_admin(self, team_id: str, user_id: str) -> Optional[LiteLLM_TeamTable]: """Remove an admin from a team. Note: Prisma doesn't support atomic array removal, so we use a @@ -319,9 +303,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): admins = [a for a in team.admins if a != user_id] return await self.update(team_id, {"admins": admins}, id_field="team_id") - async def add_models( - self, team_id: str, models: List[str] - ) -> Optional[LiteLLM_TeamTable]: + async def add_models(self, team_id: str, models: List[str]) -> Optional[LiteLLM_TeamTable]: """Add models to a team's allowed models list using atomic array push.""" if not await self.exists(team_id, id_field="team_id"): return None @@ -332,9 +314,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): ) return self._to_model(record) - async def remove_models( - self, team_id: str, models: List[str] - ) -> Optional[LiteLLM_TeamTable]: + async def remove_models(self, team_id: str, models: List[str]) -> Optional[LiteLLM_TeamTable]: """Remove models from a team's allowed models list. Note: Prisma doesn't support atomic array removal, so we use a @@ -346,6 +326,4 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return None current_models = [m for m in team.models if m not in models] - return await self.update( - team_id, {"models": current_models}, id_field="team_id" - ) + return await self.update(team_id, {"models": current_models}, id_field="team_id") diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 4d28b58f0ab..2697f15a6c0 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -34,9 +34,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): return LiteLLM_UserTable(**data) - async def find_by_id( - self, user_id: str, id_field: str = "user_id" - ) -> Optional[LiteLLM_UserTable]: + async def find_by_id(self, user_id: str, id_field: str = "user_id") -> Optional[LiteLLM_UserTable]: return await super().find_by_id(user_id, id_field) async def find_by_email(self, user_email: str) -> Optional[LiteLLM_UserTable]: @@ -51,9 +49,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): record = await self.table.find_unique(where={"sso_user_id": sso_user_id}) return self._to_model(record) - async def find_by_organization_id( - self, organization_id: str - ) -> List[LiteLLM_UserTable]: + async def find_by_organization_id(self, organization_id: str) -> List[LiteLLM_UserTable]: """Find all users in an organization.""" records = await self.table.find_many(where={"organization_id": organization_id}) return self._to_model_list(records) @@ -193,15 +189,11 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): """Delete a user.""" return await self.delete(user_id, id_field="user_id") - async def update_spend( - self, user_id: str, spend: float - ) -> Optional[LiteLLM_UserTable]: + async def update_spend(self, user_id: str, spend: float) -> Optional[LiteLLM_UserTable]: """Update user spend.""" return await self.update(user_id, {"spend": spend}, id_field="user_id") - async def add_to_team( - self, user_id: str, team_id: str - ) -> Optional[LiteLLM_UserTable]: + async def add_to_team(self, user_id: str, team_id: str) -> Optional[LiteLLM_UserTable]: """Add a user to a team using atomic array push operation.""" if not await self.exists(user_id, id_field="user_id"): return None @@ -212,9 +204,7 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): ) return self._to_model(record) - async def remove_from_team( - self, user_id: str, team_id: str - ) -> Optional[LiteLLM_UserTable]: + async def remove_from_team(self, user_id: str, team_id: str) -> Optional[LiteLLM_UserTable]: """Remove a user from a team. Note: Prisma doesn't support atomic array removal, so we use a diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index 56c3e0714aa..3ea5f32629b 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -54,14 +54,10 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): return LiteLLM_VerificationToken(**data) - async def find_by_id( - self, token: str, id_field: str = "token" - ) -> Optional[LiteLLM_VerificationToken]: + async def find_by_id(self, token: str, id_field: str = "token") -> Optional[LiteLLM_VerificationToken]: return await super().find_by_id(token, id_field) - async def find_by_alias( - self, key_alias: str - ) -> Optional[LiteLLM_VerificationToken]: + async def find_by_alias(self, key_alias: str) -> Optional[LiteLLM_VerificationToken]: """Find a token by key alias.""" records = await self.table.find_many(where={"key_alias": key_alias}) if records: @@ -78,9 +74,7 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): records = await self.table.find_many(where={"team_id": team_id}) return self._to_model_list(records) - async def find_by_project_id( - self, project_id: str - ) -> List[LiteLLM_VerificationToken]: + async def find_by_project_id(self, project_id: str) -> List[LiteLLM_VerificationToken]: """Find all tokens belonging to a project.""" records = await self.table.find_many(where={"project_id": project_id}) return self._to_model_list(records) @@ -342,32 +336,22 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): data[field] = json.dumps(data[field]) return data - async def update_spend( - self, token: str, spend: float - ) -> Optional[LiteLLM_VerificationToken]: + async def update_spend(self, token: str, spend: float) -> Optional[LiteLLM_VerificationToken]: """Update token spend.""" return await self.update(token, {"spend": spend}, id_field="token") - async def update_last_active( - self, token: str - ) -> Optional[LiteLLM_VerificationToken]: + async def update_last_active(self, token: str) -> Optional[LiteLLM_VerificationToken]: """Update the last_active timestamp.""" - return await self.update( - token, {"last_active": datetime.utcnow()}, id_field="token" - ) + return await self.update(token, {"last_active": datetime.utcnow()}, id_field="token") - async def block_token( - self, token: str, updated_by: Optional[str] = None - ) -> Optional[LiteLLM_VerificationToken]: + async def block_token(self, token: str, updated_by: Optional[str] = None) -> Optional[LiteLLM_VerificationToken]: """Block a token.""" data: Dict[str, Any] = {"blocked": True} if updated_by is not None: data["updated_by"] = updated_by return await self.update(token, data, id_field="token") - async def unblock_token( - self, token: str, updated_by: Optional[str] = None - ) -> Optional[LiteLLM_VerificationToken]: + async def unblock_token(self, token: str, updated_by: Optional[str] = None) -> Optional[LiteLLM_VerificationToken]: """Unblock a token.""" data: Dict[str, Any] = {"blocked": False} if updated_by is not None: diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 35d40423bca..9320c7fae8a 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -31,10 +31,7 @@ async def arerank( query: str, documents: List[Union[str, Dict[str, Any]]], custom_llm_provider: ( - Literal[ - "cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx" - ] - | None + Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx"] | None ) = None, top_n: int | None = None, rank_fields: List[str] | None = None, @@ -125,9 +122,7 @@ def rerank( "max_chunks_per_doc": max_chunks_per_doc, "max_tokens_per_doc": max_tokens_per_doc, } - present_version_params = [ - k for k, v in unique_version_params.items() if v is not None - ] + present_version_params = [k for k, v in unique_version_params.items() if v is not None] ( model, @@ -141,13 +136,11 @@ def rerank( api_key=optional_params.api_key, ) - rerank_provider_config: BaseRerankConfig = ( - ProviderConfigManager.get_provider_rerank_config( - model=model, - provider=litellm.LlmProviders(_custom_llm_provider), - api_base=optional_params.api_base, - present_version_params=present_version_params, - ) + rerank_provider_config: BaseRerankConfig = ProviderConfigManager.get_provider_rerank_config( + model=model, + provider=litellm.LlmProviders(_custom_llm_provider), + api_base=optional_params.api_base, + present_version_params=present_version_params, ) optional_rerank_params: Dict = get_optional_rerank_params( @@ -195,9 +188,7 @@ def rerank( or _custom_llm_provider == litellm.LlmProviders.LITELLM_PROXY ): # Implement Cohere rerank logic - api_key: str | None = ( - dynamic_api_key or optional_params.api_key or litellm.api_key - ) + api_key: str | None = dynamic_api_key or optional_params.api_key or litellm.api_key api_base: str | None = ( dynamic_api_base @@ -208,9 +199,7 @@ def rerank( ) if api_base is None: - raise Exception( - "Invalid api base. api_base=None. Set in call or via `COHERE_API_BASE` env var." - ) + raise Exception("Invalid api base. api_base=None. Set in call or via `COHERE_API_BASE` env var.") response = base_llm_http_handler.rerank( model=model, custom_llm_provider=_custom_llm_provider, @@ -253,16 +242,11 @@ def rerank( api_key = dynamic_api_key or optional_params.api_key or litellm.api_key api_base = ( - dynamic_api_base - or optional_params.api_base - or litellm.api_base - or get_secret_str("INFINITY_API_BASE") + dynamic_api_base or optional_params.api_base or litellm.api_base or get_secret_str("INFINITY_API_BASE") ) if api_base is None: - raise Exception( - "Invalid api base. api_base=None. Set in call or via `INFINITY_API_BASE` env var." - ) + raise Exception("Invalid api base. api_base=None. Set in call or via `INFINITY_API_BASE` env var.") response = base_llm_http_handler.rerank( model=model, @@ -290,9 +274,7 @@ def rerank( ) if api_key is None: - raise ValueError( - "TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment" - ) + raise ValueError("TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment") response = together_rerank.rerank( model=model, @@ -307,15 +289,10 @@ def rerank( ) elif _custom_llm_provider == litellm.LlmProviders.JINA_AI: if dynamic_api_key is None: - raise ValueError( - "Jina AI API key is required, please set 'JINA_AI_API_KEY' in your environment" - ) + raise ValueError("Jina AI API key is required, please set 'JINA_AI_API_KEY' in your environment") api_base = ( - dynamic_api_base - or optional_params.api_base - or litellm.api_base - or get_secret("BEDROCK_API_BASE") # type: ignore + dynamic_api_base or optional_params.api_base or litellm.api_base or get_secret("BEDROCK_API_BASE") # type: ignore ) response = base_llm_http_handler.rerank( @@ -335,9 +312,7 @@ def rerank( ) elif _custom_llm_provider == litellm.LlmProviders.NVIDIA_NIM: if dynamic_api_key is None: - raise ValueError( - "Nvidia NIM API key is required, please set 'NVIDIA_NIM_API_KEY' in your environment" - ) + raise ValueError("Nvidia NIM API key is required, please set 'NVIDIA_NIM_API_KEY' in your environment") # Note: For rerank, the base URL is different from chat/embeddings # Rerank uses ai.api.nvidia.com instead of integrate.api.nvidia.com @@ -364,10 +339,7 @@ def rerank( ) elif _custom_llm_provider == litellm.LlmProviders.BEDROCK: api_base = ( - dynamic_api_base - or optional_params.api_base - or litellm.api_base - or get_secret("BEDROCK_API_BASE") # type: ignore + dynamic_api_base or optional_params.api_base or litellm.api_base or get_secret("BEDROCK_API_BASE") # type: ignore ) # Merge headers and extra_headers if both are provided @@ -394,17 +366,9 @@ def rerank( ) elif _custom_llm_provider == litellm.LlmProviders.HOSTED_VLLM: # Implement Hosted VLLM rerank logic - api_key = ( - dynamic_api_key - or optional_params.api_key - or get_secret_str("HOSTED_VLLM_API_KEY") - ) + api_key = dynamic_api_key or optional_params.api_key or get_secret_str("HOSTED_VLLM_API_KEY") - api_base = ( - dynamic_api_base - or optional_params.api_base - or get_secret_str("HOSTED_VLLM_API_BASE") - ) + api_base = dynamic_api_base or optional_params.api_base or get_secret_str("HOSTED_VLLM_API_BASE") if api_base is None: raise ValueError( @@ -428,17 +392,9 @@ def rerank( ) elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA: - api_key = ( - dynamic_api_key - or optional_params.api_key - or get_secret_str("DEEPINFRA_API_KEY") - ) + api_key = dynamic_api_key or optional_params.api_key or get_secret_str("DEEPINFRA_API_KEY") - api_base = ( - dynamic_api_base - or optional_params.api_base - or get_secret_str("DEEPINFRA_API_BASE") - ) + api_base = dynamic_api_base or optional_params.api_base or get_secret_str("DEEPINFRA_API_BASE") if api_base is None: raise ValueError( @@ -470,11 +426,7 @@ def rerank( or get_secret_str("FIREWORKS_AI_TOKEN") ) - api_base = ( - dynamic_api_base - or optional_params.api_base - or get_secret_str("FIREWORKS_AI_API_BASE") - ) + api_base = dynamic_api_base or optional_params.api_base or get_secret_str("FIREWORKS_AI_API_BASE") response = base_llm_http_handler.rerank( model=model, @@ -499,11 +451,7 @@ def rerank( or get_secret_str("VOYAGE_AI_API_KEY") ) - api_base = ( - dynamic_api_base - or optional_params.api_base - or get_secret_str("VOYAGE_API_BASE") - ) + api_base = dynamic_api_base or optional_params.api_base or get_secret_str("VOYAGE_API_BASE") response = base_llm_http_handler.rerank( model=model, @@ -586,6 +534,4 @@ def rerank( return response except Exception as e: verbose_logger.error(f"Error in rerank: {str(e)}") - raise exception_type( - model=model, custom_llm_provider=custom_llm_provider, original_exception=e - ) + raise exception_type(model=model, custom_llm_provider=custom_llm_provider, original_exception=e) diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index a85582fb3b4..6e27ad66f54 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -152,11 +152,7 @@ async def _run_vector_searches( vector_store_id=vs_id, query=query, ) - results_data = ( - response.get("data") - if isinstance(response, dict) - else getattr(response, "data", None) - ) + results_data = response.get("data") if isinstance(response, dict) else getattr(response, "data", None) if results_data: all_results.extend(results_data) except Exception as exc: @@ -195,10 +191,7 @@ def _format_search_results_as_tool_output( file_id = _get_field(result, "file_id") filename = _get_field(result, "filename") content_items = _get_field(result, "content") or [] - text_chunks = [ - c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") - for c in content_items - ] + text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] text = " ".join(t for t in text_chunks if t) header = f"[Result {i}" @@ -230,10 +223,7 @@ def _build_search_results_for_include( for result in results: file_id = _get_field(result, "file_id") or "" content_items = _get_field(result, "content") or [] - text_chunks = [ - c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") - for c in content_items - ] + text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] text = " ".join(t for t in text_chunks if t) formatted.append( { @@ -326,27 +316,13 @@ def _build_message_output( def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str: """Pull the assistant's text from the provider's response.""" for item in response.output: - item_type = ( - item.get("type") if isinstance(item, dict) else getattr(item, "type", None) - ) + item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None) if item_type == "message": - content = ( - item.get("content") - if isinstance(item, dict) - else getattr(item, "content", []) - ) + content = item.get("content") if isinstance(item, dict) else getattr(item, "content", []) for block in content or []: - block_type = ( - block.get("type") - if isinstance(block, dict) - else getattr(block, "type", None) - ) + block_type = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) if block_type == "output_text": - raw = ( - block.get("text") - if isinstance(block, dict) - else getattr(block, "text", "") - ) + raw = block.get("text") if isinstance(block, dict) else getattr(block, "text", "") return str(raw) if raw is not None else "" return "" @@ -373,9 +349,7 @@ def _synthesize_responses_api_response( created_at=getattr(original_response, "created_at", int(time.time())), status="completed", model=getattr(original_response, "model", ""), - output=cast( - List[Union[ResponseOutputItem, Dict[str, Any]]], synthesized_output - ), + output=cast(List[Union[ResponseOutputItem, Dict[str, Any]]], synthesized_output), usage=getattr(original_response, "usage", None), error=None, ) @@ -389,9 +363,7 @@ def _synthesize_responses_api_response( else getattr(first_hidden, "response_cost", None) ) if first_cost is not None: - current_cost = ( - hidden.get("response_cost") if isinstance(hidden, dict) else 0 - ) + current_cost = hidden.get("response_cost") if isinstance(hidden, dict) else 0 hidden["response_cost"] = (current_cost or 0) + first_cost synthesized._hidden_params = hidden return synthesized @@ -402,9 +374,7 @@ def _synthesize_responses_api_response( # --------------------------------------------------------------------------- -async def _call_aresponses( - input, model, tools, **kwargs -): # pragma: no cover – thin wrapper for patching in tests +async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover – thin wrapper for patching in tests from litellm.responses.main import aresponses return await aresponses(input=input, model=model, tools=tools, **kwargs) @@ -420,8 +390,7 @@ def _prepare_emulated_file_search_call( updated_kwargs = kwargs if original_stream: verbose_logger.debug( - "Streaming is not yet supported for emulated file_search. " - "Disabling stream for this request." + "Streaming is not yet supported for emulated file_search. Disabling stream for this request." ) updated_kwargs = {**kwargs, "stream": False} @@ -431,16 +400,10 @@ def _prepare_emulated_file_search_call( def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> Tuple[str, str]: """Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item.""" if isinstance(tool_call, dict): - call_id = str( - tool_call.get("call_id") or tool_call.get("id") or fallback_call_id - ) + call_id = str(tool_call.get("call_id") or tool_call.get("id") or fallback_call_id) raw_args = tool_call.get("arguments") or "{}" else: - raw_call_id = ( - getattr(tool_call, "call_id", None) - or getattr(tool_call, "id", None) - or fallback_call_id - ) + raw_call_id = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None) or fallback_call_id call_id = str(raw_call_id) raw_args = getattr(tool_call, "arguments", "{}") or "{}" return call_id, raw_args @@ -470,9 +433,7 @@ async def _execute_file_search_tool_calls( all_results: List[VectorStoreSearchResult] = [] for tool_call in file_search_calls: - call_id, raw_args = _extract_tool_call_fields( - tool_call, fallback_call_id=file_search_call_id - ) + call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) try: args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args @@ -514,9 +475,7 @@ def _build_follow_up_input( Serializes Pydantic model instances to plain dicts so the transformation layer can call .get(). """ original_input_items = ( - list(input) - if isinstance(input, (list, tuple)) - else [{"role": "user", "content": str(input)}] + list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}] ) first_response_output_items: List[Any] = [] for _item in first_response.output: diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 1187640bf0c..8531cc55238 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -34,18 +34,18 @@ class LiteLLMCompletionTransformationHandler: ) -> Union[ ResponsesAPIResponse, BaseResponsesAPIStreamingIterator, - Coroutine[ - Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] - ], + Coroutine[Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]], ]: - litellm_completion_request: dict = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( - model=model, - input=input, - responses_api_request=responses_api_request, - custom_llm_provider=custom_llm_provider, - stream=stream, - extra_headers=extra_headers, - **kwargs, + litellm_completion_request: dict = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=input, + responses_api_request=responses_api_request, + custom_llm_provider=custom_llm_provider, + stream=stream, + extra_headers=extra_headers, + **kwargs, + ) ) if _is_async: @@ -61,17 +61,17 @@ class LiteLLMCompletionTransformationHandler: completion_args.update(litellm_completion_request) completion_args["_skip_responses_api_bridge"] = True - litellm_completion_response: Union[ - ModelResponse, litellm.CustomStreamWrapper - ] = litellm.completion( + litellm_completion_response: Union[ModelResponse, litellm.CustomStreamWrapper] = litellm.completion( **completion_args, ) if isinstance(litellm_completion_response, ModelResponse): - responses_api_response: ResponsesAPIResponse = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - chat_completion_response=litellm_completion_response, - request_input=input, - responses_api_request=responses_api_request, + responses_api_response: ResponsesAPIResponse = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + chat_completion_response=litellm_completion_response, + request_input=input, + responses_api_request=responses_api_request, + ) ) return responses_api_response @@ -85,9 +85,7 @@ class LiteLLMCompletionTransformationHandler: custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), ) - raise ValueError( - f"Unexpected response type: {type(litellm_completion_response)}" - ) + raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") async def async_response_api_handler( self, @@ -96,9 +94,7 @@ class LiteLLMCompletionTransformationHandler: responses_api_request: ResponsesAPIOptionalRequestParams, **kwargs, ) -> Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]: - previous_response_id: Optional[str] = responses_api_request.get( - "previous_response_id" - ) + previous_response_id: Optional[str] = responses_api_request.get("previous_response_id") if previous_response_id: litellm_completion_request = await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( previous_response_id=previous_response_id, @@ -110,17 +106,17 @@ class LiteLLMCompletionTransformationHandler: acompletion_args.update(litellm_completion_request) acompletion_args["_skip_responses_api_bridge"] = True - litellm_completion_response: Union[ - ModelResponse, litellm.CustomStreamWrapper - ] = await litellm.acompletion( + litellm_completion_response: Union[ModelResponse, litellm.CustomStreamWrapper] = await litellm.acompletion( **acompletion_args, ) if isinstance(litellm_completion_response, ModelResponse): - responses_api_response: ResponsesAPIResponse = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - chat_completion_response=litellm_completion_response, - request_input=request_input, - responses_api_request=responses_api_request, + responses_api_response: ResponsesAPIResponse = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + chat_completion_response=litellm_completion_response, + request_input=request_input, + responses_api_request=responses_api_request, + ) ) return responses_api_response @@ -131,11 +127,7 @@ class LiteLLMCompletionTransformationHandler: litellm_custom_stream_wrapper=litellm_completion_response, request_input=request_input, responses_api_request=responses_api_request, - custom_llm_provider=litellm_completion_request.get( - "custom_llm_provider" - ), + custom_llm_provider=litellm_completion_request.get("custom_llm_provider"), litellm_metadata=kwargs.get("litellm_metadata", {}), ) - raise ValueError( - f"Unexpected response type: {type(litellm_completion_response)}" - ) + raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 45ab16b0d4a..68637ea97b3 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -40,17 +40,11 @@ class ResponsesSessionHandler: ChatCompletionSession, ) - verbose_proxy_logger.debug( - "inside get_chat_completion_message_history_for_previous_response_id" - ) + verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") all_spend_logs: List[ SpendLogsPayload - ] = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( - previous_response_id - ) - verbose_proxy_logger.debug( - "found %s spend logs for this response id", len(all_spend_logs) - ) + ] = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(previous_response_id) + verbose_proxy_logger.debug("found %s spend logs for this response id", len(all_spend_logs)) litellm_session_id: Optional[str] = None if len(all_spend_logs) > 0: @@ -66,9 +60,11 @@ class ResponsesSessionHandler: ] ] = [] for spend_log in all_spend_logs: - chat_completion_message_history = await ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( - spend_log=spend_log, - chat_completion_message_history=chat_completion_message_history, + chat_completion_message_history = ( + await ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( + spend_log=spend_log, + chat_completion_message_history=chat_completion_message_history, + ) ) verbose_proxy_logger.debug( @@ -100,10 +96,8 @@ class ResponsesSessionHandler: LiteLLMCompletionResponsesConfig, ) - proxy_server_request_dict = ( - await ResponsesSessionHandler.get_proxy_server_request_from_spend_log( - spend_log=spend_log, - ) + proxy_server_request_dict = await ResponsesSessionHandler.get_proxy_server_request_from_spend_log( + spend_log=spend_log, ) response_input_param: Optional[Union[str, ResponseInputParam]] = None _messages: Optional[Union[str, ResponseInputParam]] = None @@ -142,11 +136,7 @@ class ResponsesSessionHandler: # Add Output messages for this Spend Log ############################################################ _response_output = spend_log.get("response", "{}") - if ( - isinstance(_response_output, dict) - and _response_output - and _response_output != {} - ): + if isinstance(_response_output, dict) and _response_output and _response_output != {}: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response = ModelResponse(**_response_output) for choice in model_response.choices: @@ -161,9 +151,7 @@ class ResponsesSessionHandler: """ Get the parsed proxy server request from the spend log """ - proxy_server_request: Union[str, dict] = ( - spend_log.get("proxy_server_request") or "{}" - ) + proxy_server_request: Union[str, dict] = spend_log.get("proxy_server_request") or "{}" proxy_server_request_dict: Optional[dict] = None if isinstance(proxy_server_request, dict): proxy_server_request_dict = proxy_server_request @@ -173,20 +161,16 @@ class ResponsesSessionHandler: ############################################################ # Check if user has setup cold storage for session handling ############################################################ - if ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_server_request_dict - ): + if ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_server_request_dict): # Try to get cold storage object key from spend log metadata _proxy_server_request_dict: Optional[dict] = None - cold_storage_object_key = ( - ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log( - spend_log - ) - ) + cold_storage_object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log) if cold_storage_object_key: # Use the object key directly from metadata - _proxy_server_request_dict = await ResponsesSessionHandler.get_proxy_server_request_from_cold_storage_with_object_key( - object_key=cold_storage_object_key, + _proxy_server_request_dict = ( + await ResponsesSessionHandler.get_proxy_server_request_from_cold_storage_with_object_key( + object_key=cold_storage_object_key, + ) ) if _proxy_server_request_dict: proxy_server_request_dict = _proxy_server_request_dict @@ -215,9 +199,7 @@ class ResponsesSessionHandler: return metadata_str.get("cold_storage_object_key") return None except (json.JSONDecodeError, TypeError, AttributeError): - verbose_proxy_logger.debug( - "Failed to parse metadata from spend log to extract cold storage object key" - ) + verbose_proxy_logger.debug("Failed to parse metadata from spend log to extract cold storage object key") return None @staticmethod @@ -233,12 +215,12 @@ class ResponsesSessionHandler: Returns: Optional[dict]: The proxy server request dict or None if not found """ - verbose_proxy_logger.debug( - "inside get_proxy_server_request_from_cold_storage_with_object_key..." - ) + verbose_proxy_logger.debug("inside get_proxy_server_request_from_cold_storage_with_object_key...") - proxy_server_request_dict = await COLD_STORAGE_HANDLER.get_proxy_server_request_from_cold_storage_with_object_key( - object_key=object_key, + proxy_server_request_dict = ( + await COLD_STORAGE_HANDLER.get_proxy_server_request_from_cold_storage_with_object_key( + object_key=object_key, + ) ) return proxy_server_request_dict @@ -281,14 +263,8 @@ class ResponsesSessionHandler: verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) - decoded_response_id = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - previous_response_id - ) - ) - previous_response_id = decoded_response_id.get( - "response_id", previous_response_id - ) + decoded_response_id = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) + previous_response_id = decoded_response_id.get("response_id", previous_response_id) if prisma_client is None: return [] diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 582144e6cb8..b1198780bac 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -59,13 +59,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): litellm_metadata: Optional[dict] = None, ): self.model: str = model - self.litellm_custom_stream_wrapper: litellm.CustomStreamWrapper = ( - litellm_custom_stream_wrapper - ) + self.litellm_custom_stream_wrapper: litellm.CustomStreamWrapper = litellm_custom_stream_wrapper self.request_input: Union[str, ResponseInputParam] = request_input - self.responses_api_request: ResponsesAPIOptionalRequestParams = ( - responses_api_request - ) + self.responses_api_request: ResponsesAPIOptionalRequestParams = responses_api_request self.custom_llm_provider: Optional[str] = custom_llm_provider self.litellm_metadata: Optional[dict] = litellm_metadata or {} # Store lightweight dict snapshots for stream_chunk_builder to reduce @@ -81,9 +77,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.sent_output_content_part_done_event: bool = False self.sent_output_item_done_event: bool = False self.sent_annotation_events: bool = False - self.litellm_model_response: Optional[ - Union[ModelResponse, TextCompletionResponse] - ] = None + self.litellm_model_response: Optional[Union[ModelResponse, TextCompletionResponse]] = None self.final_text: str = "" self._cached_item_id: Optional[str] = None self._cached_response_id: Optional[str] = None @@ -92,9 +86,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_args_by_call_id: dict[str, str] = {} self._tool_call_id_by_index: dict[int, str] = {} self._ambiguous_tool_call_indexes: set[int] = set() - self._next_tool_output_index: int = ( - 1 # output_index=0 reserved for the message item - ) + self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item self._final_tool_events_queued: bool = False self._sequence_number: int = 0 self._cached_reasoning_item_id: Optional[str] = None @@ -119,11 +111,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return idx def _normalize_tool_call_index(self, tool_call: object) -> Optional[int]: - idx_raw = ( - tool_call.get("index") - if isinstance(tool_call, dict) - else getattr(tool_call, "index", None) - ) + idx_raw = tool_call.get("index") if isinstance(tool_call, dict) else getattr(tool_call, "index", None) if idx_raw is None: return None try: @@ -140,12 +128,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if hasattr(delta, "thinking_blocks") and delta.thinking_blocks: return False - return ( - delta.content - or delta.function_call - or delta.tool_calls - or chunk.choices[0].finish_reason is not None - ) + return delta.content or delta.function_call or delta.tool_calls or chunk.choices[0].finish_reason is not None def _queue_tool_call_delta_events(self, tool_calls: object) -> None: """ @@ -163,9 +146,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): for tc in tool_calls: tc_index = self._normalize_tool_call_index(tc) - call_id_raw = ( - tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) - ) + call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) call_id = "" if call_id_raw: @@ -187,11 +168,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not call_id: continue - fn = ( - tc.get("function") - if isinstance(tc, dict) - else getattr(tc, "function", None) - ) + fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) fn_name = "" fn_args_delta = "" if isinstance(fn, dict): @@ -232,21 +209,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): for i in range(0, len(fn_args_delta), chunk_size): delta_chunk = fn_args_delta[i : i + chunk_size] self._sequence_number += 1 - delta_event: BaseLiteLLMOpenAIResponseObject = ( - FunctionCallArgumentsDeltaEvent( - type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, - item_id=call_id, - output_index=output_index, - delta=delta_chunk, - ) + delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent( + type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + item_id=call_id, + output_index=output_index, + delta=delta_chunk, ) # Add sequence_number as extra field (BaseLiteLLMOpenAIResponseObject allows extra fields) delta_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(delta_event) - def _queue_final_tool_call_done_events( - self, litellm_complete_object: ModelResponse - ) -> None: + def _queue_final_tool_call_done_events(self, litellm_complete_object: ModelResponse) -> None: """ Ensure tool calls that were not streamed as deltas still get emitted before response.completed. """ @@ -264,19 +237,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return for tc in tool_calls: - call_id_raw = ( - tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) - ) + call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) if not call_id_raw: continue call_id = str(call_id_raw) output_index = self._get_or_assign_tool_output_index(call_id) - fn = ( - tc.get("function") - if isinstance(tc, dict) - else getattr(tc, "function", None) - ) + fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) fn_name = "" fn_args = "" if isinstance(fn, dict): @@ -382,17 +349,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "store": True, } if "temperature" in self.responses_api_request: - response_created_event_data["temperature"] = self.responses_api_request[ - "temperature" - ] + response_created_event_data["temperature"] = self.responses_api_request["temperature"] if "text" in self.responses_api_request: response_created_event_data["text"] = self.responses_api_request["text"] if "tool_choice" in self.responses_api_request: # Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format response_created_event_data["tool_choice"] = ( - LiteLLMCompletionResponsesConfig._transform_tool_choice( - self.responses_api_request["tool_choice"] - ) + LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"]) or "auto" ) else: @@ -406,15 +369,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: response_created_event_data["top_p"] = 1.0 if "truncation" in self.responses_api_request: - response_created_event_data["truncation"] = self.responses_api_request[ - "truncation" - ] + response_created_event_data["truncation"] = self.responses_api_request["truncation"] if "user" in self.responses_api_request: response_created_event_data["user"] = self.responses_api_request["user"] if "metadata" in self.responses_api_request: - response_created_event_data["metadata"] = self.responses_api_request[ - "metadata" - ] + response_created_event_data["metadata"] = self.responses_api_request["metadata"] return response_created_event_data def create_response_created_event(self) -> ResponseCreatedEvent: @@ -473,9 +432,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): item_id=self._cached_item_id, output_index=0, content_index=0, - part=BaseLiteLLMOpenAIResponseObject( - **{"type": "output_text", "text": "", "annotations": []} - ), + part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}), ) event.__dict__["sequence_number"] = self._sequence_number return event @@ -501,10 +458,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ), ) if response is not None and self._accumulated_provider_specific_fields: - if ( - not hasattr(response, "_hidden_params") - or response._hidden_params is None - ): + if not hasattr(response, "_hidden_params") or response._hidden_params is None: response._hidden_params = {} response._hidden_params.setdefault("provider_specific_fields", {}).update( self._accumulated_provider_specific_fields @@ -522,11 +476,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk_dict = chunk.model_dump() hidden_params = 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 - ) + chunk_dict["_hidden_params"] = dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params return chunk_dict def create_reasoning_summary_text_done_event( @@ -593,9 +543,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ), ) - def create_output_text_done_event( - self, litellm_complete_object: ModelResponse - ) -> OutputTextDoneEvent: + def create_output_text_done_event(self, litellm_complete_object: ModelResponse) -> OutputTextDoneEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" @@ -608,20 +556,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): or "", ) - def create_output_content_part_done_event( - self, litellm_complete_object: ModelResponse - ) -> ContentPartDoneEvent: + def create_output_content_part_done_event(self, litellm_complete_object: ModelResponse) -> ContentPartDoneEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" text = getattr(litellm_complete_object.choices[0].message, "content", "") or "" # type: ignore - reasoning_content = ( - getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") - or "" - ) # type: ignore - annotations = getattr( - litellm_complete_object.choices[0].message, "annotations", None - ) # type: ignore + reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" # type: ignore + annotations = getattr(litellm_complete_object.choices[0].message, "annotations", None) # type: ignore part: Optional[PART_UNION_TYPES] = None if reasoning_content: @@ -631,8 +572,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) else: - response_annotations = LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( - annotations=annotations + response_annotations = ( + LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( + annotations=annotations + ) ) part = ContentPartDonePartOutputText( type="output_text", @@ -649,19 +592,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): part=part, ) - def create_output_item_done_event( - self, litellm_complete_object: ModelResponse - ) -> OutputItemDoneEvent: + def create_output_item_done_event(self, litellm_complete_object: ModelResponse) -> OutputItemDoneEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" text = self.litellm_model_response.choices[0].message.content or "" # type: ignore - annotations = getattr( - self.litellm_model_response.choices[0].message, "annotations", None - ) # type: ignore + annotations = getattr(self.litellm_model_response.choices[0].message, "annotations", None) # type: ignore - response_annotations = LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( - annotations=annotations + response_annotations = ( + LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( + annotations=annotations + ) ) return OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, @@ -762,12 +703,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return True return False - def common_done_event_logic( - self, sync_mode: bool = True - ) -> BaseLiteLLMOpenAIResponseObject: - if not self.litellm_model_response or isinstance( - self.litellm_model_response, TextCompletionResponse - ): + def common_done_event_logic(self, sync_mode: bool = True) -> BaseLiteLLMOpenAIResponseObject: + if not self.litellm_model_response or isinstance(self.litellm_model_response, TextCompletionResponse): self.litellm_model_response = self.create_litellm_model_response() if self.litellm_model_response: # If tool calls exist, emit tool events before finishing/response.completed. @@ -786,9 +723,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopAsyncIteration self.finished = self.is_stream_finished() - response_completed_event = self._emit_response_completed_event( - self.litellm_model_response - ) + response_completed_event = self._emit_response_completed_event(self.litellm_model_response) if response_completed_event: return response_completed_event else: @@ -909,53 +844,37 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Incrementally accumulate reasoning content instead of # calling stream_chunk_builder on every chunk (O(n²)) delta = chunk.choices[0].delta if chunk.choices else None - if ( - delta - and hasattr(delta, "reasoning_content") - and delta.reasoning_content - ): - self._accumulated_reasoning_content_parts.append( - delta.reasoning_content - ) + if delta and hasattr(delta, "reasoning_content") and delta.reasoning_content: + self._accumulated_reasoning_content_parts.append(delta.reasoning_content) if self._is_reasoning_end(chunk): - reasoning_content = "".join( - self._accumulated_reasoning_content_parts - ) + reasoning_content = "".join(self._accumulated_reasoning_content_parts) # Ensure we have a valid reasoning_item_id reasoning_item_id = ( - self._reasoning_item_id - or self._cached_reasoning_item_id - or f"rs_{uuid.uuid4()}" + self._reasoning_item_id or self._cached_reasoning_item_id or f"rs_{uuid.uuid4()}" ) # Create text.done event first with its own sequence number self._sequence_number += 1 - text_done_event = ( - self.create_reasoning_summary_text_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) + text_done_event = self.create_reasoning_summary_text_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, ) # Create part.done event second with its own sequence number self._sequence_number += 1 - part_done_event = ( - self.create_reasoning_summary_part_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) + part_done_event = self.create_reasoning_summary_part_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, ) self._sequence_number += 1 - reasoning_output_item_done_event = ( - self.create_reasoning_output_item_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) + reasoning_output_item_done_event = self.create_reasoning_output_item_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, ) self._pending_response_events.extend( [ @@ -967,11 +886,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._reasoning_done_emitted = True self._reasoning_active = False - response_api_chunk = ( - self._transform_chat_completion_chunk_to_response_api_chunk( - chunk - ) - ) + response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: self._pending_response_events.append(response_api_chunk) @@ -1028,18 +943,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # _ensure_output_item_for_chunk queues events on the same chunk. # This mirrors the async path (see __anext__). self.collected_chat_completion_chunks.append( - self._snapshot_chunk_for_stream_chunk_builder( - cast(ModelResponseStream, chunk) - ) + self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk)) ) # Emit any just-queued output_item event if self._pending_response_events: return self._pending_response_events.pop(0) - response_api_chunk = ( - self._transform_chat_completion_chunk_to_response_api_chunk( - chunk - ) - ) + response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: return response_api_chunk # Otherwise, loop to next chunk @@ -1078,9 +987,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._pending_annotation_events = [] for idx, annotation in enumerate(response_annotations): annotation_dict = ( - annotation.model_dump() - if hasattr(annotation, "model_dump") - else dict(annotation) + annotation.model_dump() if hasattr(annotation, "model_dump") else dict(annotation) ) event = OutputTextAnnotationAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, @@ -1122,11 +1029,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 3: Handle tool call deltas (if any) -> queue events and emit them # For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior - if ( - chunk.choices - and hasattr(chunk.choices[0].delta, "tool_calls") - and chunk.choices[0].delta.tool_calls - ): + if chunk.choices and hasattr(chunk.choices[0].delta, "tool_calls") and chunk.choices[0].delta.tool_calls: self._queue_tool_call_delta_events(chunk.choices[0].delta.tool_calls) # Return one pending tool event at a time if self._pending_tool_events: @@ -1134,10 +1037,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 4: If we have pending annotation events, emit the next one # This happens when the current chunk has no text/reasoning content - if ( - hasattr(self, "_pending_annotation_events") - and self._pending_annotation_events - ): + if hasattr(self, "_pending_annotation_events") and self._pending_annotation_events: event = self._pending_annotation_events.pop(0) return event @@ -1147,9 +1047,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return None - def _get_delta_string_from_streaming_choices( - self, choices: List[StreamingChoices] - ) -> str: + def _get_delta_string_from_streaming_choices(self, choices: List[StreamingChoices]) -> str: """ Get the delta string from the streaming choices @@ -1161,30 +1059,25 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chat_completion_delta: ChatCompletionDelta = choice.delta return chat_completion_delta.content or "" - def _emit_response_completed_event( - self, litellm_model_response: ModelResponse - ) -> Optional[ResponseCompletedEvent]: + def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> Optional[ResponseCompletedEvent]: if litellm_model_response: # Add cost to usage object if include_cost_in_streaming_usage is True - if ( - litellm.include_cost_in_streaming_usage - and self.litellm_logging_obj is not None - ): + if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None: usage = getattr(litellm_model_response, "usage", None) if usage is not None: setattr( usage, "cost", - self.litellm_logging_obj._response_cost_calculator( - result=litellm_model_response - ), + self.litellm_logging_obj._response_cost_calculator(result=litellm_model_response), ) # Transform the response - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input=self.request_input, - chat_completion_response=litellm_model_response, - responses_api_request=self.responses_api_request, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input=self.request_input, + chat_completion_response=litellm_model_response, + responses_api_request=self.responses_api_request, + ) ) # Use the cached response ID to ensure consistency across all events diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0bd78e59819..d2a7edd21ef 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -134,9 +134,7 @@ class LiteLLMCompletionResponsesConfig: tool_choice_type = tool_choice.get("type") # If it has a function with name, it's standard OpenAI format - pass through - if tool_choice.get("function") and tool_choice.get("function", {}).get( - "name" - ): + if tool_choice.get("function") and tool_choice.get("function", {}).get("name"): return tool_choice # Handle Cursor IDE dict formats without function name @@ -180,9 +178,7 @@ class LiteLLMCompletionResponsesConfig: response_format = None text_param = responses_api_request.get("text") if text_param: - response_format = LiteLLMCompletionResponsesConfig._transform_text_format_to_response_format( - text_param - ) + response_format = LiteLLMCompletionResponsesConfig._transform_text_format_to_response_format(text_param) # Extract reasoning_effort from reasoning parameter reasoning_effort: Optional[Union[Reasoning, str]] = None @@ -235,16 +231,12 @@ class LiteLLMCompletionResponsesConfig: "include_usage": True, } litellm_completion_request["stream_options"] = stream_options - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( - "litellm_logging_obj" - ) + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") if litellm_logging_obj: litellm_logging_obj.stream_options = stream_options # only pass non-None values - litellm_completion_request = { - k: v for k, v in litellm_completion_request.items() if v is not None - } + litellm_completion_request = {k: v for k, v in litellm_completion_request.items() if v is not None} return litellm_completion_request @staticmethod @@ -295,12 +287,12 @@ class LiteLLMCompletionResponsesConfig: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages """ - chat_completion_session = ChatCompletionSession( - messages=[], litellm_session_id=None - ) + chat_completion_session = ChatCompletionSession(messages=[], litellm_session_id=None) if previous_response_id: - chat_completion_session = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - previous_response_id=previous_response_id + chat_completion_session = ( + await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + previous_response_id=previous_response_id + ) ) _messages = litellm_completion_request.get("messages") or [] session_messages = chat_completion_session.get("messages") or [] @@ -348,15 +340,11 @@ class LiteLLMCompletionResponsesConfig: f"Original request: previous_response_id={previous_response_id}" ), model=litellm_completion_request.get("model", ""), - llm_provider=litellm_completion_request.get( - "custom_llm_provider", "" - ), + llm_provider=litellm_completion_request.get("custom_llm_provider", ""), ) litellm_completion_request["messages"] = combined_messages - litellm_completion_request["litellm_trace_id"] = chat_completion_session.get( - "litellm_session_id" - ) + litellm_completion_request["litellm_trace_id"] = chat_completion_session.get("litellm_session_id") return litellm_completion_request @staticmethod @@ -387,13 +375,13 @@ class LiteLLMCompletionResponsesConfig: elif isinstance(input, list): existing_tool_call_ids: Set[str] = set() for _input in input: - chat_completion_messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( - input_item=_input + chat_completion_messages = ( + LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=_input + ) ) - if LiteLLMCompletionResponsesConfig._is_input_item_function_call( - input_item=_input - ): + if LiteLLMCompletionResponsesConfig._is_input_item_function_call(input_item=_input): call_id_raw = _input.get("call_id") or _input.get("id") or "" if call_id_raw: existing_tool_call_ids.add(str(call_id_raw)) @@ -410,16 +398,12 @@ class LiteLLMCompletionResponsesConfig: if messages: last_msg = messages[-1] last_role = ( - last_msg.get("role") - if isinstance(last_msg, dict) - else getattr(last_msg, "role", None) + last_msg.get("role") if isinstance(last_msg, dict) else getattr(last_msg, "role", None) ) if last_role == "assistant": for new_msg in chat_completion_messages: new_role = ( - new_msg.get("role") - if isinstance(new_msg, dict) - else getattr(new_msg, "role", None) + new_msg.get("role") if isinstance(new_msg, dict) else getattr(new_msg, "role", None) ) if new_role == "assistant": _raw_tcs = ( @@ -427,13 +411,9 @@ class LiteLLMCompletionResponsesConfig: if isinstance(new_msg, dict) else getattr(new_msg, "tool_calls", None) ) - new_tcs: list = ( - _raw_tcs if isinstance(_raw_tcs, list) else [] - ) + new_tcs: list = _raw_tcs if isinstance(_raw_tcs, list) else [] for tc in new_tcs: - LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant( - last_msg, tc - ) + LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant(last_msg, tc) continue ######################################################### @@ -441,9 +421,7 @@ class LiteLLMCompletionResponsesConfig: # preserving the ordering of tool call outputs. Some models require the tool # result to immediately follow the assistant tool call. ######################################################### - if LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( - input_item=_input - ): + if LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(input_item=_input): if not chat_completion_messages: continue @@ -458,9 +436,7 @@ class LiteLLMCompletionResponsesConfig: # Drop assistant tool_calls wrappers if we already have this call_id if role == "assistant": tool_calls: Any = ( - m.get("tool_calls") - if isinstance(m, dict) - else getattr(m, "tool_calls", None) + m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None) ) call_id = "" if ( @@ -574,9 +550,7 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _find_previous_assistant_idx( - messages: List[Any], current_idx: int - ) -> Optional[int]: + def _find_previous_assistant_idx(messages: List[Any], current_idx: int) -> Optional[int]: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": @@ -584,20 +558,14 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _recover_tool_call_id_from_assistant( - assistant_message: Any, message: Any - ) -> str: + def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str: """Try to recover empty tool_call_id from assistant message's tool_calls.""" tool_calls_raw = ( assistant_message.get("tool_calls") if isinstance(assistant_message, dict) else getattr(assistant_message, "tool_calls", None) ) - if ( - tool_calls_raw - and isinstance(tool_calls_raw, list) - and len(tool_calls_raw) > 0 - ): + if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0: first_tool_call = tool_calls_raw[0] if isinstance(first_tool_call, dict): tool_call_id_raw = first_tool_call.get("id", "") @@ -619,9 +587,7 @@ class LiteLLMCompletionResponsesConfig: return [] if isinstance(tool_calls_raw, list): return tool_calls_raw - if hasattr(tool_calls_raw, "__iter__") and not isinstance( - tool_calls_raw, (str, bytes) - ): + if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)): return list(tool_calls_raw) return [] @@ -639,9 +605,7 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _reconstruct_tool_call_from_tools( - tool_call_id: str, tools: List[Any] - ) -> Optional[Dict[str, Any]]: + def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: List[Any]) -> Optional[Dict[str, Any]]: """Reconstruct a minimal tool_call definition from tools list.""" for tool in tools: if isinstance(tool, dict): @@ -683,33 +647,17 @@ class LiteLLMCompletionResponsesConfig: tool_use_definition: Dict[str, Any], tool_call_id: str, index: int ) -> ChatCompletionToolCallChunk: """Create a ChatCompletionToolCallChunk from tool_use_definition.""" - function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "function" - ) - function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - function_raw, "name" - ) - function_arguments_raw = ( - LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - function_raw, "arguments" - ) - ) + function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function") + function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") + function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments") function: Dict[str, Any] = { "name": function_name_raw or "", "arguments": function_arguments_raw or "{}", } - tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "id" - ) - tool_use_id: str = ( - str(tool_use_id_raw) if tool_use_id_raw is not None else str(tool_call_id) - ) - tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "type" - ) - tool_use_type: str = ( - str(tool_use_type_raw) if tool_use_type_raw is not None else "function" - ) + tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id") + tool_use_id: str = str(tool_use_id_raw) if tool_use_id_raw is not None else str(tool_call_id) + tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type") + tool_use_type: str = str(tool_use_type_raw) if tool_use_type_raw is not None else "function" return ChatCompletionToolCallChunk( id=tool_use_id, type=cast(Literal["function"], tool_use_type), @@ -721,9 +669,7 @@ class LiteLLMCompletionResponsesConfig: ) @staticmethod - def _normalize_tool_use_definition( - tool_use_definition: Any, tool_call_id: str - ) -> Optional[Dict[str, Any]]: + def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> Optional[Dict[str, Any]]: """ Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. """ @@ -733,26 +679,12 @@ class LiteLLMCompletionResponsesConfig: if isinstance(tool_use_definition, dict): normalized_definition: Dict[str, Any] = dict(tool_use_definition) else: - tool_use_id_raw = ( - LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "id" - ) - ) - tool_use_type_raw = ( - LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "type" - ) - ) - function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "function" - ) + tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id") + tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type") + function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function") # Object does not expose the expected tool_call fields. - if ( - tool_use_id_raw is None - and tool_use_type_raw is None - and function_raw is None - ): + if tool_use_id_raw is None and tool_use_type_raw is None and function_raw is None: return None normalized_definition = { @@ -763,15 +695,9 @@ class LiteLLMCompletionResponsesConfig: function_raw = normalized_definition.get("function") if function_raw is not None and not isinstance(function_raw, dict): - function_name_raw = ( - LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - function_raw, "name" - ) - ) - function_arguments_raw = ( - LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - function_raw, "arguments" - ) + function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") + function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "arguments" ) if function_name_raw is not None or function_arguments_raw is not None: normalized_definition["function"] = { @@ -784,9 +710,7 @@ class LiteLLMCompletionResponsesConfig: return normalized_definition @staticmethod - def _add_tool_call_to_assistant( - assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk - ) -> None: + def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None: """Add a tool_call to an assistant message.""" if isinstance(assistant_message, dict): prev_assistant_dict = cast(Dict[str, Any], assistant_message) @@ -854,9 +778,7 @@ class LiteLLMCompletionResponsesConfig: # Count non-tool messages to avoid removing all messages # This prevents empty messages list when using previous_response_id without a database - non_tool_messages_count = sum( - 1 for msg in fixed_messages if msg.get("role") != "tool" - ) + non_tool_messages_count = sum(1 for msg in fixed_messages if msg.get("role") != "tool") for i, message in enumerate(fixed_messages): # Only process tool messages - check role first to narrow the type @@ -866,19 +788,11 @@ class LiteLLMCompletionResponsesConfig: # At this point, we know it's a tool message, so it should have tool_call_id # Use get() with default to safely access tool_call_id tool_call_id_raw = ( - message.get("tool_call_id") - if isinstance(message, dict) - else getattr(message, "tool_call_id", None) - ) - tool_call_id: str = ( - str(tool_call_id_raw) if tool_call_id_raw is not None else "" + message.get("tool_call_id") if isinstance(message, dict) else getattr(message, "tool_call_id", None) ) + tool_call_id: str = str(tool_call_id_raw) if tool_call_id_raw is not None else "" - prev_assistant_idx = ( - LiteLLMCompletionResponsesConfig._find_previous_assistant_idx( - fixed_messages, i - ) - ) + prev_assistant_idx = LiteLLMCompletionResponsesConfig._find_previous_assistant_idx(fixed_messages, i) # Try to recover empty tool_call_id from previous assistant message if not tool_call_id and prev_assistant_idx is not None: @@ -912,13 +826,9 @@ class LiteLLMCompletionResponsesConfig: # not just those that had an empty tool_call_id initially if prev_assistant_idx is not None and tool_call_id: prev_assistant = fixed_messages[prev_assistant_idx] - tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list( - prev_assistant - ) + tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list(prev_assistant) - if not LiteLLMCompletionResponsesConfig._check_tool_call_exists( - tool_calls, tool_call_id - ): + if not LiteLLMCompletionResponsesConfig._check_tool_call_exists(tool_calls, tool_call_id): _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) if not _tool_use_definition and tools: @@ -926,23 +836,17 @@ class LiteLLMCompletionResponsesConfig: tool_call_id, tools ) - normalized_tool_use_definition = ( - LiteLLMCompletionResponsesConfig._normalize_tool_use_definition( - _tool_use_definition, tool_call_id - ) + normalized_tool_use_definition = LiteLLMCompletionResponsesConfig._normalize_tool_use_definition( + _tool_use_definition, tool_call_id ) if normalized_tool_use_definition: - tool_call_chunk = ( - LiteLLMCompletionResponsesConfig._create_tool_call_chunk( - normalized_tool_use_definition, - tool_call_id, - len(tool_calls), - ) - ) - LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant( - prev_assistant, tool_call_chunk + tool_call_chunk = LiteLLMCompletionResponsesConfig._create_tool_call_chunk( + normalized_tool_use_definition, + tool_call_id, + len(tool_calls), ) + LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant(prev_assistant, tool_call_chunk) # Remove messages with empty tool_call_id that couldn't be fixed for idx in reversed(messages_to_remove): @@ -977,8 +881,10 @@ class LiteLLMCompletionResponsesConfig: """ if LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(input_item): # handle executed tool call results - return LiteLLMCompletionResponsesConfig._transform_responses_api_tool_call_output_to_chat_completion_message( - tool_call_output=input_item + return ( + LiteLLMCompletionResponsesConfig._transform_responses_api_tool_call_output_to_chat_completion_message( + tool_call_output=input_item + ) ) elif LiteLLMCompletionResponsesConfig._is_input_item_function_call(input_item): # handle function call input items @@ -1078,9 +984,7 @@ class LiteLLMCompletionResponsesConfig: if isinstance(image_url_val, dict): url = image_url_val.get("url") if isinstance(url, str) and url: - normalized_blocks.append( - {"type": "image_url", "image_url": {"url": url}} - ) + normalized_blocks.append({"type": "image_url", "image_url": {"url": url}}) elif isinstance(image_url_val, str) and image_url_val: normalized_blocks.append( { @@ -1113,9 +1017,7 @@ class LiteLLMCompletionResponsesConfig: tool_output_message = ChatCompletionToolMessage( role="tool", - content=_normalize_function_call_output_to_tool_content( - tool_call_output.get("output") - ), + content=_normalize_function_call_output_to_tool_content(tool_call_output.get("output")), tool_call_id=str(call_id), ) @@ -1150,9 +1052,7 @@ class LiteLLMCompletionResponsesConfig: function: dict = _tool_use_definition.get("function") or {} tool_call_chunk = ChatCompletionToolCallChunk( id=_tool_use_definition.get("id") or "", - type=cast( - Literal["function"], _tool_use_definition.get("type") or "function" - ), + type=cast(Literal["function"], _tool_use_definition.get("type") or "function"), function=ChatCompletionToolCallFunctionChunk( name=function.get("name") or "", arguments=str(function.get("arguments") or ""), @@ -1282,15 +1182,11 @@ class LiteLLMCompletionResponsesConfig: elif isinstance(item, dict): if item.get("type") == "input_file": content_list.append( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - item - ) + LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) ) elif item.get("type") == "input_image": image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - item - ) + LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) ) if "cache_control" in item: image_block["cache_control"] = item["cache_control"] @@ -1367,17 +1263,12 @@ class LiteLLMCompletionResponsesConfig: """ if tools is None: return [], None - chat_completion_tools: List[ - Union[ChatCompletionToolParam, OpenAIMcpServerTool] - ] = [] + chat_completion_tools: List[Union[ChatCompletionToolParam, OpenAIMcpServerTool]] = [] web_search_options: Optional[OpenAIWebSearchOptions] = None for tool in tools: if tool.get("type") == "mcp": chat_completion_tools.append(cast(OpenAIMcpServerTool, tool)) - elif ( - tool.get("type") == "web_search_preview" - or tool.get("type") == "web_search" - ): + elif tool.get("type") == "web_search_preview" or tool.get("type") == "web_search": _search_context_size: Literal["low", "medium", "high"] = cast( Literal["low", "medium", "high"], tool.get("search_context_size") ) @@ -1409,25 +1300,17 @@ class LiteLLMCompletionResponsesConfig: if tool.get("defer_loading"): chat_completion_tool["defer_loading"] = tool.get("defer_loading") # type: ignore if tool.get("allowed_callers"): - chat_completion_tool["allowed_callers"] = tool.get( - "allowed_callers" - ) # type: ignore + chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") # type: ignore if tool.get("input_examples"): chat_completion_tool["input_examples"] = tool.get("input_examples") # type: ignore - chat_completion_tools.append( - cast(ChatCompletionToolParam, chat_completion_tool) - ) + chat_completion_tools.append(cast(ChatCompletionToolParam, chat_completion_tool)) else: - chat_completion_tools.append( - cast(Union[ChatCompletionToolParam, OpenAIMcpServerTool], tool) - ) + chat_completion_tools.append(cast(Union[ChatCompletionToolParam, OpenAIMcpServerTool], tool)) return chat_completion_tools, web_search_options @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( - chat_completion_tools: Optional[ - List[Union[ChatCompletionToolParam, OpenAIMcpServerTool]] - ], + chat_completion_tools: Optional[List[Union[ChatCompletionToolParam, OpenAIMcpServerTool]]], ) -> List[Dict[str, Any]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to @@ -1490,9 +1373,7 @@ class LiteLLMCompletionResponsesConfig: if tool.type == "function": function_definition = tool.function provider_specific_fields: Optional[Dict] = None - if hasattr(tool, "provider_specific_fields") and getattr( - tool, "provider_specific_fields", None - ): + if hasattr(tool, "provider_specific_fields") and getattr(tool, "provider_specific_fields", None): provider_specific_fields = getattr(tool, "provider_specific_fields") if not isinstance(provider_specific_fields, dict): provider_specific_fields = ( @@ -1500,12 +1381,10 @@ class LiteLLMCompletionResponsesConfig: if hasattr(provider_specific_fields, "__dict__") else {} ) - elif hasattr( - function_definition, "provider_specific_fields" - ) and getattr(function_definition, "provider_specific_fields", None): - provider_specific_fields = getattr( - function_definition, "provider_specific_fields" - ) + elif hasattr(function_definition, "provider_specific_fields") and getattr( + function_definition, "provider_specific_fields", None + ): + provider_specific_fields = getattr(function_definition, "provider_specific_fields") if not isinstance(provider_specific_fields, dict): provider_specific_fields = ( dict(provider_specific_fields) # type: ignore @@ -1562,9 +1441,7 @@ class LiteLLMCompletionResponsesConfig: return "completed" @staticmethod - def _tool_call_id_from_responses_item( - item_id: Optional[str], call_id: Optional[str] - ) -> str: + def _tool_call_id_from_responses_item(item_id: Optional[str], call_id: Optional[str]) -> str: """Bedrock Mantle returns a non-unique, index-based ``call_id`` (``call_0``, ``call_1``, ... that resets every response) alongside a unique ``id`` (``fc_...``). ``call_id`` is the canonical Responses API correlation key, so @@ -1591,14 +1468,10 @@ class LiteLLMCompletionResponsesConfig: Dictionary in ChatCompletionToolCallChunk format """ # Extract provider_specific_fields if present - provider_specific_fields = getattr( - tool_call_item, "provider_specific_fields", None - ) + provider_specific_fields = getattr(tool_call_item, "provider_specific_fields", None) if provider_specific_fields and not isinstance(provider_specific_fields, dict): provider_specific_fields = ( - dict(provider_specific_fields) - if hasattr(provider_specific_fields, "__dict__") - else {} + dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} ) elif hasattr(tool_call_item, "get") and callable(tool_call_item.get): # type: ignore provider_fields = tool_call_item.get("provider_specific_fields") # type: ignore @@ -1692,28 +1565,20 @@ class LiteLLMCompletionResponsesConfig: model=chat_completion_response.model, object="response", error=getattr(chat_completion_response, "error", None), - incomplete_details=getattr( - chat_completion_response, "incomplete_details", None - ), + incomplete_details=getattr(chat_completion_response, "incomplete_details", None), instructions=getattr(chat_completion_response, "instructions", None), metadata=getattr(chat_completion_response, "metadata", {}), output=LiteLLMCompletionResponsesConfig._transform_chat_completion_choices_to_responses_output( chat_completion_response=chat_completion_response, choices=getattr(chat_completion_response, "choices", []), ), - parallel_tool_calls=getattr( - chat_completion_response, "parallel_tool_calls", False - ), + parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False), temperature=getattr(chat_completion_response, "temperature", 0), tool_choice=getattr(chat_completion_response, "tool_choice", "auto"), tools=getattr(chat_completion_response, "tools", []), top_p=getattr(chat_completion_response, "top_p", None), - max_output_tokens=getattr( - chat_completion_response, "max_output_tokens", None - ), - previous_response_id=getattr( - chat_completion_response, "previous_response_id", None - ), + max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None), + previous_response_id=getattr(chat_completion_response, "previous_response_id", None), reasoning=None, status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( finish_reason @@ -1725,14 +1590,10 @@ class LiteLLMCompletionResponsesConfig: ), user=getattr(chat_completion_response, "user", None), ) - responses_api_response._hidden_params = getattr( - chat_completion_response, "_hidden_params", {} - ) + responses_api_response._hidden_params = getattr(chat_completion_response, "_hidden_params", {}) # Surface provider-specific fields (generic passthrough from any provider) - provider_fields = responses_api_response._hidden_params.get( - "provider_specific_fields" - ) + provider_fields = responses_api_response._hidden_params.get("provider_specific_fields") if provider_fields: setattr(responses_api_response, "provider_specific_fields", provider_fields) @@ -1762,14 +1623,10 @@ class LiteLLMCompletionResponsesConfig: ] = [] responses_output.extend( - LiteLLMCompletionResponsesConfig._extract_reasoning_output_items( - chat_completion_response, choices - ) + LiteLLMCompletionResponsesConfig._extract_reasoning_output_items(chat_completion_response, choices) ) responses_output.extend( - LiteLLMCompletionResponsesConfig._extract_message_output_items( - chat_completion_response, choices - ) + LiteLLMCompletionResponsesConfig._extract_message_output_items(chat_completion_response, choices) ) responses_output.extend( LiteLLMCompletionResponsesConfig.transform_chat_completion_tools_to_responses_tools( @@ -1780,11 +1637,7 @@ class LiteLLMCompletionResponsesConfig: # Convert server-side tool results (e.g. Anthropic code execution) # into code_interpreter_call output items, replacing the corresponding # function_call items so the output matches OpenAI's native shape. - tool_result_items = ( - LiteLLMCompletionResponsesConfig._extract_tool_result_output_items( - chat_completion_response - ) - ) + tool_result_items = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(chat_completion_response) if tool_result_items: result_by_id = {item.id: item for item in tool_result_items} replaced_ids = set(result_by_id.keys()) @@ -1894,11 +1747,7 @@ class LiteLLMCompletionResponsesConfig: for idx, image_item in enumerate(images): # Extract base64 from data URL image_url = image_item.get("image_url", {}).get("url", "") - base64_data = ( - LiteLLMCompletionResponsesConfig._extract_base64_from_data_url( - image_url - ) - ) + base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: image_generation_items.append( @@ -1963,9 +1812,7 @@ class LiteLLMCompletionResponsesConfig: chat_completion_response: ModelResponse, choices: List[Choices], ) -> List[Union[GenericResponseOutputItem, OutputImageGenerationCall]]: - message_output_items: List[ - Union[GenericResponseOutputItem, OutputImageGenerationCall] - ] = [] + message_output_items: List[Union[GenericResponseOutputItem, OutputImageGenerationCall]] = [] for choice in choices: # Check if message has images (image generation) if hasattr(choice.message, "images") and choice.message.images: @@ -2054,8 +1901,10 @@ class LiteLLMCompletionResponsesConfig: message: Message, ) -> OutputText: annotations = getattr(message, "annotations", None) - transformed_annotations = LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( - annotations=annotations + transformed_annotations = ( + LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( + annotations=annotations + ) ) return OutputText( @@ -2068,9 +1917,7 @@ class LiteLLMCompletionResponsesConfig: def _transform_chat_completion_annotations_to_response_output_annotations( annotations: Optional[List[ChatCompletionAnnotation]], ) -> List[GenericResponseOutputItemContentAnnotation]: - response_output_annotations: List[ - GenericResponseOutputItemContentAnnotation - ] = [] + response_output_annotations: List[GenericResponseOutputItemContentAnnotation] = [] if annotations is None: return response_output_annotations @@ -2118,71 +1965,41 @@ class LiteLLMCompletionResponsesConfig: setattr(response_usage, "cost", usage.cost) # Translate prompt_tokens_details to input_tokens_details - if ( - hasattr(usage, "prompt_tokens_details") - and usage.prompt_tokens_details is not None - ): + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: prompt_details = usage.prompt_tokens_details input_details_dict: Dict[str, int] = {} - if ( - hasattr(prompt_details, "cached_tokens") - and prompt_details.cached_tokens is not None - ): + if hasattr(prompt_details, "cached_tokens") and prompt_details.cached_tokens is not None: input_details_dict["cached_tokens"] = prompt_details.cached_tokens else: input_details_dict["cached_tokens"] = 0 - if ( - hasattr(prompt_details, "text_tokens") - and prompt_details.text_tokens is not None - ): + if hasattr(prompt_details, "text_tokens") and prompt_details.text_tokens is not None: input_details_dict["text_tokens"] = prompt_details.text_tokens - if ( - hasattr(prompt_details, "audio_tokens") - and prompt_details.audio_tokens is not None - ): + if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None: input_details_dict["audio_tokens"] = prompt_details.audio_tokens if input_details_dict: - response_usage.input_tokens_details = InputTokensDetails( - **input_details_dict - ) + response_usage.input_tokens_details = InputTokensDetails(**input_details_dict) # Translate completion_tokens_details to output_tokens_details - if ( - hasattr(usage, "completion_tokens_details") - and usage.completion_tokens_details is not None - ): + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details is not None: completion_details = usage.completion_tokens_details output_details_dict: Dict[str, int] = {} - if ( - hasattr(completion_details, "reasoning_tokens") - and completion_details.reasoning_tokens is not None - ): - output_details_dict["reasoning_tokens"] = ( - completion_details.reasoning_tokens - ) + if hasattr(completion_details, "reasoning_tokens") and completion_details.reasoning_tokens is not None: + output_details_dict["reasoning_tokens"] = completion_details.reasoning_tokens else: output_details_dict["reasoning_tokens"] = 0 - if ( - hasattr(completion_details, "text_tokens") - and completion_details.text_tokens is not None - ): + if hasattr(completion_details, "text_tokens") and completion_details.text_tokens is not None: output_details_dict["text_tokens"] = completion_details.text_tokens - if ( - hasattr(completion_details, "image_tokens") - and completion_details.image_tokens is not None - ): + if hasattr(completion_details, "image_tokens") and completion_details.image_tokens is not None: output_details_dict["image_tokens"] = completion_details.image_tokens if output_details_dict: - response_usage.output_tokens_details = OutputTokensDetails( - **output_details_dict - ) + response_usage.output_tokens_details = OutputTokensDetails(**output_details_dict) return response_usage diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 3edcbd430f1..e8fe51ed484 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -188,9 +188,7 @@ async def aresponses_api_with_mcp( # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform) # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata) - user_api_key_auth = kwargs.get("user_api_key_auth") or kwargs.get( - "litellm_metadata", {} - ).get("user_api_key_auth") + user_api_key_auth = kwargs.get("user_api_key_auth") or kwargs.get("litellm_metadata", {}).get("user_api_key_auth") # Extract MCP auth headers from request (for dynamic auth when fetching tools) mcp_auth_header: Optional[str] = None @@ -202,9 +200,7 @@ async def aresponses_api_with_mcp( mcp_server_auth_headers, _, _, - ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( - secret_fields=secret_fields, tools=tools - ) + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request(secret_fields=secret_fields, tools=tools) # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods ( @@ -217,9 +213,7 @@ async def aresponses_api_with_mcp( mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, ) - openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( - original_mcp_tools - ) + openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(original_mcp_tools) # Combine with other tools all_tools = openai_tools + other_tools if (openai_tools or other_tools) else None @@ -279,9 +273,7 @@ async def aresponses_api_with_mcp( ) # Determine if we should auto-execute tools - should_auto_execute = bool( - mcp_tools_with_litellm_proxy - ) and LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( + should_auto_execute = bool(mcp_tools_with_litellm_proxy) and LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy ) @@ -308,14 +300,10 @@ async def aresponses_api_with_mcp( # If auto-execute tools is True, then we need to execute the tool calls ######################################################### if should_auto_execute and isinstance(response, ResponsesAPIResponse): # type: ignore - tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response( - response=response - ) + tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response(response=response) if tool_calls: - user_api_key_auth = kwargs.get("litellm_metadata", {}).get( - "user_api_key_auth" - ) + user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth") # Extract MCP auth headers from the request to pass to MCP server secret_fields = kwargs.get("secret_fields") @@ -347,19 +335,15 @@ async def aresponses_api_with_mcp( ) # Prepare parameters for follow-up call (restores original stream setting) - follow_up_call_params = ( - LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( - call_params=call_params, original_stream_setting=stream or False - ) + follow_up_call_params = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + call_params=call_params, original_stream_setting=stream or False ) # Create tool execution events for streaming if needed tool_execution_events = [] if stream: - tool_execution_events = ( - LiteLLM_Proxy_MCP_Handler._create_tool_execution_events( - tool_calls=tool_calls, tool_results=tool_results - ) + tool_execution_events = LiteLLM_Proxy_MCP_Handler._create_tool_execution_events( + tool_calls=tool_calls, tool_results=tool_results ) final_response = await LiteLLM_Proxy_MCP_Handler._make_follow_up_call( @@ -374,10 +358,7 @@ async def aresponses_api_with_mcp( if ( stream and tool_execution_events - and ( - hasattr(final_response, "__aiter__") - or hasattr(final_response, "__iter__") - ) + and (hasattr(final_response, "__aiter__") or hasattr(final_response, "__iter__")) ): from litellm.responses.mcp.mcp_streaming_iterator import ( MCPEnhancedStreamingIterator, @@ -402,12 +383,10 @@ async def aresponses_api_with_mcp( mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, ) - final_response = ( - LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response( - response=final_response, - mcp_tools_fetched=mcp_tools_for_output, - tool_results=tool_results, - ) + final_response = LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response( + response=final_response, + mcp_tools_fetched=mcp_tools_for_output, + tool_results=tool_results, ) return final_response @@ -458,9 +437,7 @@ async def aresponses( kwargs["aresponses"] = True # Convert text_format to text parameter if provided - text = ResponsesAPIRequestUtils.convert_text_format_to_text_param( - text_format=text_format, text=text - ) + text = ResponsesAPIRequestUtils.convert_text_format_to_text_param(text_format=text_format, text=text) if text is not None: # Update local_vars to include the converted text parameter local_vars["text"] = text @@ -488,13 +465,9 @@ async def aresponses( if isinstance( litellm_logging_obj, LiteLLMLoggingObj - ) and litellm_logging_obj.should_run_prompt_management_hooks( - prompt_id=prompt_id, non_default_params=kwargs - ): + ) and litellm_logging_obj.should_run_prompt_management_hooks(prompt_id=prompt_id, non_default_params=kwargs): if isinstance(input, str): - client_input: List[AllMessageValues] = [ - {"role": "user", "content": input} - ] + client_input: List[AllMessageValues] = [{"role": "user", "content": input}] else: client_input = [ item # type: ignore[misc] @@ -573,9 +546,7 @@ async def aresponses( response._hidden_params["custom_llm_provider"] = custom_llm_provider if response is None: - raise ValueError( - f"Got an unexpected None response from the Responses API: {response}" - ) + raise ValueError(f"Got an unexpected None response from the Responses API: {response}") return response except Exception as e: @@ -615,9 +586,7 @@ def _apply_prompt_management_to_responses_call( if isinstance(item, dict) and "role" in item ] - if isinstance( - litellm_logging_obj, LiteLLMLoggingObj - ) and litellm_logging_obj.should_run_prompt_management_hooks( + if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and litellm_logging_obj.should_run_prompt_management_hooks( prompt_id=prompt_id, non_default_params=kwargs ): ( @@ -700,11 +669,7 @@ def _apply_managed_file_id_mapping( local_vars: Dict[str, Any], ) -> tuple[Union[str, ResponseInputParam], Optional[Iterable[ToolParam]]]: model_file_id_mapping = kwargs.get("model_file_id_mapping") - model_info_id = ( - kwargs.get("model_info", {}).get("id") - if isinstance(kwargs.get("model_info"), dict) - else None - ) + model_info_id = kwargs.get("model_info", {}).get("id") if isinstance(kwargs.get("model_info"), dict) else None input = cast( Union[str, ResponseInputParam], @@ -877,19 +842,13 @@ def _responses_try_dispatch_emulated_file_search( "custom_llm_provider": custom_llm_provider, **( { - **( - {"use_chat_completions_api": True} - if use_chat_completions_api - else {} - ), + **({"use_chat_completions_api": True} if use_chat_completions_api else {}), **{k: v for k, v in kwargs.items() if k not in _internal_skip}, } ), } if _is_async: - return aresponses_with_emulated_file_search( - input=input, model=model, tools=tools, **emulated_kwargs - ) + return aresponses_with_emulated_file_search(input=input, model=model, tools=tools, **emulated_kwargs) return run_async_function( aresponses_with_emulated_file_search, input=input, @@ -948,9 +907,7 @@ def responses( use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs) # Convert text_format to text parameter if provided - text = ResponsesAPIRequestUtils.convert_text_format_to_text_param( - text_format=text_format, text=text - ) + text = ResponsesAPIRequestUtils.convert_text_format_to_text_param(text_format=text_format, text=text) if text is not None: # Update local_vars to include the converted text parameter local_vars["text"] = text @@ -961,21 +918,13 @@ def responses( ######################################################### # MOCK RESPONSE LOGIC ######################################################### - if litellm_params.mock_response and isinstance( - litellm_params.mock_response, str - ): - return mock_responses_api_response( - mock_response=litellm_params.mock_response - ) + if litellm_params.mock_response and isinstance(litellm_params.mock_response, str): + return mock_responses_api_response(mock_response=litellm_params.mock_response) - _stripped_model, _from_chat_completions_prefix = ( - _normalize_openai_chat_completions_responses_model(model) - ) + _stripped_model, _from_chat_completions_prefix = _normalize_openai_chat_completions_responses_model(model) model = _stripped_model local_vars["model"] = model - use_chat_completions_api = ( - use_chat_completions_api or _from_chat_completions_prefix - ) + use_chat_completions_api = use_chat_completions_api or _from_chat_completions_prefix model, custom_llm_provider = _resolve_model_provider_for_responses( model=model, @@ -1002,9 +951,7 @@ def responses( ######################################################### # Update input and tools with provider-specific file IDs if managed files are used ######################################################### - input, tools = _apply_managed_file_id_mapping( - input=input, tools=tools, kwargs=kwargs, local_vars=local_vars - ) + input, tools = _apply_managed_file_id_mapping(input=input, tools=tools, kwargs=kwargs, local_vars=local_vars) ######################################################### # Native MCP Responses API @@ -1046,27 +993,21 @@ def responses( if custom_llm_provider is None: responses_api_provider_config = None else: - responses_api_provider_config = ( - ProviderConfigManager.get_provider_responses_api_config( - model=model, - provider=custom_llm_provider, - ) + responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=custom_llm_provider, ) local_vars.update(kwargs) # Map reasoning_effort (from litellm_params/proxy config) to reasoning when not set if reasoning is None and "reasoning_effort" in local_vars: - _mapped = LiteLLMResponsesTransformationHandler()._map_reasoning_effort( - local_vars.pop("reasoning_effort") - ) + _mapped = LiteLLMResponsesTransformationHandler()._map_reasoning_effort(local_vars.pop("reasoning_effort")) if _mapped is not None: reasoning = _mapped local_vars["reasoning"] = _mapped # Get ResponsesAPIOptionalRequestParams with only valid parameters response_api_optional_params: ResponsesAPIOptionalRequestParams = ( - ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - local_vars - ) + ResponsesAPIRequestUtils.get_requested_response_api_optional_param(local_vars) ) _file_search_dispatch = _responses_try_dispatch_emulated_file_search( @@ -1122,13 +1063,11 @@ def responses( ) # Get optional parameters for the responses API - responses_api_request_params: Dict = ( - ResponsesAPIRequestUtils.get_optional_params_responses_api( - model=model, - responses_api_provider_config=responses_api_provider_config, - response_api_optional_params=response_api_optional_params, - allowed_openai_params=allowed_openai_params, - ) + responses_api_request_params: Dict = ResponsesAPIRequestUtils.get_optional_params_responses_api( + model=model, + responses_api_provider_config=responses_api_provider_config, + response_api_optional_params=response_api_optional_params, + allowed_openai_params=allowed_openai_params, ) litellm_logging_obj.update_from_kwargs( @@ -1141,22 +1080,14 @@ def responses( "aresponses": _is_async, "litellm_call_id": litellm_call_id, "model_info": kwargs.get("model_info"), - "data_residency": infer_openai_data_residency( - custom_llm_provider, litellm_params.api_base - ), - "metadata": ( - kwargs["litellm_metadata"] - if "litellm_metadata" in kwargs - else kwargs.get("metadata") - ), + "data_residency": infer_openai_data_residency(custom_llm_provider, litellm_params.api_base), + "metadata": (kwargs["litellm_metadata"] if "litellm_metadata" in kwargs else kwargs.get("metadata")), }, custom_llm_provider=custom_llm_provider, ) # Decode any litellm-encoded encrypted-content item IDs back to their original IDs - input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - input - ) + input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(input) # Call the handler with _is_async flag instead of directly calling the async handler if custom_llm_provider is None: @@ -1229,15 +1160,11 @@ async def adelete_responses( kwargs["adelete_responses"] = True # get custom llm provider from response_id - decoded_response_id: DecodedResponseId = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id, - ) + decoded_response_id: DecodedResponseId = ResponsesAPIRequestUtils._decode_responses_api_response_id( + response_id=response_id, ) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider func = partial( delete_responses, @@ -1298,15 +1225,11 @@ def delete_responses( litellm_params = GenericLiteLLMParams(**kwargs) # get custom llm provider from response_id - decoded_response_id: DecodedResponseId = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id, - ) + decoded_response_id: DecodedResponseId = ResponsesAPIRequestUtils._decode_responses_api_response_id( + response_id=response_id, ) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider if custom_llm_provider is None: raise ValueError("custom_llm_provider is required but passed as None") @@ -1320,9 +1243,7 @@ def delete_responses( ) if responses_api_provider_config is None: - raise ValueError( - f"DELETE responses is not supported for {custom_llm_provider}" - ) + raise ValueError(f"DELETE responses is not supported for {custom_llm_provider}") local_vars.update(kwargs) @@ -1396,15 +1317,11 @@ async def aget_responses( kwargs["aget_responses"] = True # get custom llm provider from response_id - decoded_response_id: DecodedResponseId = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id, - ) + decoded_response_id: DecodedResponseId = ResponsesAPIRequestUtils._decode_responses_api_response_id( + response_id=response_id, ) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider func = partial( get_responses, @@ -1479,15 +1396,11 @@ def get_responses( litellm_params = GenericLiteLLMParams(**kwargs) # get custom llm provider from response_id - decoded_response_id: DecodedResponseId = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id, - ) + decoded_response_id: DecodedResponseId = ResponsesAPIRequestUtils._decode_responses_api_response_id( + response_id=response_id, ) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider if custom_llm_provider is None: raise ValueError("custom_llm_provider is required but passed as None") @@ -1501,9 +1414,7 @@ def get_responses( ) if responses_api_provider_config is None: - raise ValueError( - f"GET responses is not supported for {custom_llm_provider}" - ) + raise ValueError(f"GET responses is not supported for {custom_llm_provider}") local_vars.update(kwargs) @@ -1573,15 +1484,9 @@ async def alist_input_items( loop = asyncio.get_event_loop() kwargs["alist_input_items"] = True - decoded_response_id = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id - ) - ) + decoded_response_id = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id=response_id) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider func = partial( list_input_items, @@ -1638,15 +1543,9 @@ def list_input_items( litellm_params = GenericLiteLLMParams(**kwargs) - decoded_response_id = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id - ) - ) + decoded_response_id = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id=response_id) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider if custom_llm_provider is None: raise ValueError("custom_llm_provider is required but passed as None") @@ -1659,9 +1558,7 @@ def list_input_items( ) if responses_api_provider_config is None: - raise ValueError( - f"list_input_items is not supported for {custom_llm_provider}" - ) + raise ValueError(f"list_input_items is not supported for {custom_llm_provider}") local_vars.update(kwargs) @@ -1727,15 +1624,11 @@ async def acancel_responses( kwargs["acancel_responses"] = True # get custom llm provider from response_id - decoded_response_id: DecodedResponseId = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id, - ) + decoded_response_id: DecodedResponseId = ResponsesAPIRequestUtils._decode_responses_api_response_id( + response_id=response_id, ) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider func = partial( cancel_responses, @@ -1796,15 +1689,11 @@ def cancel_responses( litellm_params = GenericLiteLLMParams(**kwargs) # get custom llm provider from response_id - decoded_response_id: DecodedResponseId = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id=response_id, - ) + decoded_response_id: DecodedResponseId = ResponsesAPIRequestUtils._decode_responses_api_response_id( + response_id=response_id, ) response_id = decoded_response_id.get("response_id") or response_id - custom_llm_provider = ( - decoded_response_id.get("custom_llm_provider") or custom_llm_provider - ) + custom_llm_provider = decoded_response_id.get("custom_llm_provider") or custom_llm_provider if custom_llm_provider is None: raise ValueError("custom_llm_provider is required but passed as None") @@ -1818,9 +1707,7 @@ def cancel_responses( ) if responses_api_provider_config is None: - raise ValueError( - f"CANCEL responses is not supported for {custom_llm_provider}" - ) + raise ValueError(f"CANCEL responses is not supported for {custom_llm_provider}") local_vars.update(kwargs) @@ -1992,27 +1879,21 @@ def compact_responses( ) if responses_api_provider_config is None: - raise ValueError( - f"COMPACT responses is not supported for {custom_llm_provider}" - ) + raise ValueError(f"COMPACT responses is not supported for {custom_llm_provider}") local_vars.update(kwargs) # Build optional params for compact endpoint response_api_optional_params: ResponsesAPIOptionalRequestParams = ( - ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - local_vars - ) + ResponsesAPIRequestUtils.get_requested_response_api_optional_param(local_vars) ) # Get optional parameters for the responses API - responses_api_request_params: Dict = ( - ResponsesAPIRequestUtils.get_optional_params_responses_api( - model=model, - responses_api_provider_config=responses_api_provider_config, - response_api_optional_params=response_api_optional_params, - allowed_openai_params=None, - ) + responses_api_request_params: Dict = ResponsesAPIRequestUtils.get_optional_params_responses_api( + model=model, + responses_api_provider_config=responses_api_provider_config, + response_api_optional_params=response_api_optional_params, + allowed_openai_params=None, ) # Pre Call logging @@ -2023,18 +1904,14 @@ def compact_responses( litellm_params={ **responses_api_request_params, "litellm_call_id": litellm_call_id, - "data_residency": infer_openai_data_residency( - custom_llm_provider, litellm_params.api_base - ), + "data_residency": infer_openai_data_residency(custom_llm_provider, litellm_params.api_base), }, custom_llm_provider=custom_llm_provider, ) # Decode any litellm-encoded encrypted-content item IDs back to their original IDs # before forwarding to the upstream provider. - input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input( - input - ) + input = ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(input) # Call the handler with _is_async flag instead of directly calling the async handler response = base_llm_http_handler.compact_response_api_handler( @@ -2079,11 +1956,7 @@ def compact_responses( def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: metadata: dict = {**(kwargs.get("litellm_metadata") or {})} - guardrails = ( - (kwargs.get("metadata") or {}).get("guardrails") - or kwargs.get("guardrails") - or [] - ) + guardrails = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or [] if guardrails: metadata["guardrails"] = guardrails return metadata @@ -2139,16 +2012,12 @@ async def _aresponses_websocket( responses_api_provider_config: Optional[BaseResponsesAPIConfig] = None if _custom_llm_provider is not None: - responses_api_provider_config = ( - ProviderConfigManager.get_provider_responses_api_config( - model=model, - provider=litellm.LlmProviders(_custom_llm_provider), - ) + responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=litellm.LlmProviders(_custom_llm_provider), ) - resolved_api_base = ( - dynamic_api_base or litellm_params.api_base or litellm.api_base or None - ) + resolved_api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or None resolved_api_key = ( dynamic_api_key or litellm_params.api_key diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index acb7487f430..10ff67f68d5 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -114,9 +114,7 @@ async def acompletion_with_mcp( ) # Extract user_api_key_auth from metadata or kwargs - user_api_key_auth = kwargs.get("user_api_key_auth") or ( - (kwargs.get("metadata", {}) or {}).get("user_api_key_auth") - ) + user_api_key_auth = kwargs.get("user_api_key_auth") or ((kwargs.get("metadata", {}) or {}).get("user_api_key_auth")) # Extract MCP auth headers before fetching tools (needed for dynamic auth) ( @@ -246,9 +244,7 @@ async def acompletion_with_mcp( async def __aiter__(self): return self - def _add_mcp_list_tools_to_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( StreamingChoices, @@ -260,19 +256,10 @@ async def acompletion_with_mcp( if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - existing_fields = ( - getattr(choice.delta, "provider_specific_fields", None) - or {} - ) - provider_fields = dict( - existing_fields - ) # Create a copy to avoid mutating the original + existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {} + provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = self.openai_tools @@ -283,9 +270,7 @@ async def acompletion_with_mcp( return chunk - def _add_mcp_tool_metadata_to_final_chunk( - self, chunk: ModelResponseStream - ) -> ModelResponseStream: + def _add_mcp_tool_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_tool_calls and mcp_call_results to the final chunk.""" from litellm.types.utils import ( StreamingChoices, @@ -294,25 +279,15 @@ async def acompletion_with_mcp( if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if ( - isinstance(choice, StreamingChoices) - and hasattr(choice, "delta") - and choice.delta - ): + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict # Access the attribute directly to handle Pydantic model attributes correctly existing_fields = {} if hasattr(choice.delta, "provider_specific_fields"): - attr_value = getattr( - choice.delta, "provider_specific_fields", None - ) + attr_value = getattr(choice.delta, "provider_specific_fields", None) if attr_value is not None: # Create a copy to avoid mutating the original - existing_fields = ( - dict(attr_value) - if isinstance(attr_value, dict) - else {} - ) + existing_fields = dict(attr_value) if isinstance(attr_value, dict) else {} provider_fields = existing_fields @@ -375,9 +350,7 @@ async def acompletion_with_mcp( # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] - final_chunk = self._add_mcp_tool_metadata_to_final_chunk( - final_chunk - ) + final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) # If we have tool results, prepare follow-up call if self.tool_results and self.complete_response: await self._prepare_follow_up_call() @@ -414,9 +387,7 @@ async def acompletion_with_mcp( ): from litellm._logging import verbose_logger - verbose_logger.warning( - "Follow-up stream was not created despite having tool results" - ) + verbose_logger.warning("Follow-up stream was not created despite having tool results") raise StopAsyncIteration @@ -445,18 +416,16 @@ async def acompletion_with_mcp( if self.tool_calls: # Execute tool calls - self.tool_results = ( - await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map=self.tool_server_map, - tool_calls=self.tool_calls, - user_api_key_auth=self.user_api_key_auth, - mcp_auth_header=self.mcp_auth_header, - mcp_server_auth_headers=self.mcp_server_auth_headers, - oauth2_headers=self.oauth2_headers, - raw_headers=self.raw_headers, - litellm_call_id=self.litellm_call_id, - litellm_trace_id=self.litellm_trace_id, - ) + self.tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=self.tool_server_map, + tool_calls=self.tool_calls, + user_api_key_auth=self.user_api_key_auth, + mcp_auth_header=self.mcp_auth_header, + mcp_server_auth_headers=self.mcp_server_auth_headers, + oauth2_headers=self.oauth2_headers, + raw_headers=self.raw_headers, + litellm_call_id=self.litellm_call_id, + litellm_trace_id=self.litellm_trace_id, ) async def _prepare_follow_up_call(self): @@ -468,12 +437,10 @@ async def acompletion_with_mcp( return # Create follow-up messages with tool results - follow_up_messages = ( - LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=self.messages, - response=self.complete_response, - tool_results=self.tool_results, - ) + follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=self.messages, + response=self.complete_response, + tool_results=self.tool_results, ) # Make follow-up call with streaming @@ -529,14 +496,10 @@ async def acompletion_with_mcp( completion_stream=None, model=getattr(original_wrapper, "model", "unknown"), logging_obj=getattr(original_wrapper, "logging_obj", None), - custom_llm_provider=getattr( - original_wrapper, "custom_llm_provider", None - ), + custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), stream_options=getattr(original_wrapper, "stream_options", None), make_call=getattr(original_wrapper, "make_call", None), - _response_headers=getattr( - original_wrapper, "_response_headers", None - ), + _response_headers=getattr(original_wrapper, "_response_headers", None), ) self._original_wrapper = original_wrapper self._custom_iterator = custom_iterator @@ -560,9 +523,7 @@ async def acompletion_with_mcp( except RuntimeError: self._sync_loop = asyncio.new_event_loop() asyncio.set_event_loop(self._sync_loop) - self._sync_iterator = _SyncIteratorWrapper( - self._custom_iterator, self._sync_loop - ) + self._sync_iterator = _SyncIteratorWrapper(self._custom_iterator, self._sync_loop) return self._sync_iterator def __next__(self): @@ -615,9 +576,7 @@ async def acompletion_with_mcp( return initial_response # Extract tool calls from response - tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( - response=initial_response - ) + tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(response=initial_response) if not tool_calls: _add_mcp_metadata_to_response( diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index df5de205d45..e969208d1d9 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -69,13 +69,9 @@ class LiteLLM_Proxy_MCP_Handler: for tool in tools: if isinstance(tool, dict) and tool.get("type") == "mcp": server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith( - LITELLM_PROXY_MCP_SERVER_URL - ): + if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): return True - if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match( - server_url - ): + if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url): return True return False @@ -96,9 +92,7 @@ class LiteLLM_Proxy_MCP_Handler: for tool in tools: if isinstance(tool, dict) and tool.get("type") == "mcp": server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith( - LITELLM_PROXY_MCP_SERVER_URL - ): + if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): mcp_tools_with_litellm_proxy.append(tool) elif isinstance(server_url, str): # Also intercept URLs like http://localhost:4000/mcp/atlassian_test @@ -133,22 +127,16 @@ class LiteLLM_Proxy_MCP_Handler: global_mcp_server_manager, ) - tool_permissions = ( - await global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=resolved_toolset_ids - ) - ) - all_server_ids = list( - set(tool_permissions.keys()) | set(resolved_mcp_servers) + tool_permissions = await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=resolved_toolset_ids ) + all_server_ids = list(set(tool_permissions.keys()) | set(resolved_mcp_servers)) existing_op = user_api_key_auth.object_permission if existing_op is not None: merged_tool_perms = dict(existing_op.mcp_tool_permissions or {}) for server_id, tool_names in tool_permissions.items(): existing_tools = merged_tool_perms.get(server_id, []) - merged_tool_perms[server_id] = list( - set(existing_tools) | set(tool_names) - ) + merged_tool_perms[server_id] = list(set(existing_tools) | set(tool_names)) updated_op = existing_op.model_copy( update={ "mcp_servers": all_server_ids, @@ -162,9 +150,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_servers=all_server_ids, mcp_tool_permissions=tool_permissions, ) - return user_api_key_auth.model_copy( - update={"object_permission": updated_op} - ) + return user_api_key_auth.model_copy(update={"object_permission": updated_op}) except Exception as _e: verbose_logger.debug(f"Could not apply toolset permissions: {_e}") return user_api_key_auth @@ -202,12 +188,8 @@ class LiteLLM_Proxy_MCP_Handler: if mcp_tools_with_litellm_proxy: for _tool in mcp_tools_with_litellm_proxy: # if user specifies servers as server_url: litellm_proxy/mcp/zapier,github then return zapier,github - server_url = ( - _tool.get("server_url", "") if isinstance(_tool, dict) else "" - ) - if isinstance(server_url, str) and server_url.startswith( - LITELLM_PROXY_MCP_SERVER_URL_PREFIX - ): + server_url = _tool.get("server_url", "") if isinstance(_tool, dict) else "" + if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX): mcp_servers.append(server_url.split("/")[-1]) # Resolve toolset names: collect all toolset IDs first, then apply their @@ -221,11 +203,7 @@ class LiteLLM_Proxy_MCP_Handler: from litellm.proxy.proxy_server import prisma_client if prisma_client is not None: - toolset = ( - await global_mcp_server_manager.get_toolset_by_name_cached( - prisma_client, name - ) - ) + toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) if toolset is not None: # Access control: only allow if the key explicitly grants this toolset. if user_api_key_auth is not None: @@ -236,20 +214,11 @@ class LiteLLM_Proxy_MCP_Handler: is_admin = _user_has_admin_view(user_api_key_auth) if not is_admin: op = user_api_key_auth.object_permission - granted = ( - getattr(op, "mcp_toolsets", None) - if op - else None - ) + granted = getattr(op, "mcp_toolsets", None) if op else None # None means no grants configured → deny (consistent with # fetch_mcp_toolsets which returns [] for unconfigured keys) - if ( - granted is None - or toolset.toolset_id not in granted - ): - verbose_logger.debug( - f"Key does not have access to toolset '{name}', skipping." - ) + if granted is None or toolset.toolset_id not in granted: + verbose_logger.debug(f"Key does not have access to toolset '{name}', skipping.") continue resolved_toolset_ids.append(toolset.toolset_id) # Don't add to resolved_mcp_servers — toolset scope @@ -261,21 +230,17 @@ class LiteLLM_Proxy_MCP_Handler: # Apply all resolved toolsets at once (union), avoiding permission overwrite. if resolved_toolset_ids and user_api_key_auth is not None: - user_api_key_auth = ( - await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( - resolved_toolset_ids=resolved_toolset_ids, - resolved_mcp_servers=resolved_mcp_servers, - user_api_key_auth=user_api_key_auth, - ) + user_api_key_auth = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( + resolved_toolset_ids=resolved_toolset_ids, + resolved_mcp_servers=resolved_mcp_servers, + user_api_key_auth=user_api_key_auth, ) # When toolsets were resolved we updated object_permission.mcp_servers to the # full union (toolset server IDs + direct server names). Passing a name-based # filter here would exclude those toolset server IDs (which are UUIDs, not # names), so use None and let the auth object's mcp_servers do the filtering. - effective_server_filter = ( - None if resolved_toolset_ids else (resolved_mcp_servers or None) - ) + effective_server_filter = None if resolved_toolset_ids else (resolved_mcp_servers or None) tools = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, @@ -287,9 +252,7 @@ class LiteLLM_Proxy_MCP_Handler: litellm_trace_id=litellm_trace_id, ) - allowed_mcp_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) - ) + allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) allowed_mcp_servers = global_mcp_server_manager.get_mcp_servers_from_ids( # type: ignore[attr-defined] allowed_mcp_server_ids ) @@ -304,9 +267,7 @@ class LiteLLM_Proxy_MCP_Handler: if server is None: continue server_name = ( - getattr(server, "server_name", None) - or getattr(server, "alias", None) - or getattr(server, "name", None) + getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None) ) if isinstance(server_name, str): server_names.append(server_name) @@ -343,9 +304,7 @@ class LiteLLM_Proxy_MCP_Handler: if len(allowed_mcp_servers) == 1: tool_server_map[tool_name] = allowed_mcp_servers[0] else: - _, tool_server_map[tool_name] = split_server_prefix_from_name( - tool_name - ) + _, tool_server_map[tool_name] = split_server_prefix_from_name(tool_name) return deduplicated_tools, tool_server_map @@ -414,9 +373,7 @@ class LiteLLM_Proxy_MCP_Handler: litellm_trace_id=litellm_trace_id, ) - openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( - deduplicated_mcp_tools - ) + openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(deduplicated_mcp_tools) return openai_tools, tool_server_map @@ -457,20 +414,16 @@ class LiteLLM_Proxy_MCP_Handler: ) # Step 2: Filter tools based on allowed_tools parameter - filtered_mcp_tools = ( - LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( - mcp_tools=mcp_tools_fetched, - mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - ) + filtered_mcp_tools = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( + mcp_tools=mcp_tools_fetched, + mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, ) # Step 3: Deduplicate tools after filtering ( deduplicated_mcp_tools, tool_server_map, - ) = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( - filtered_mcp_tools, allowed_mcp_servers - ) + ) = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(filtered_mcp_tools, allowed_mcp_servers) return deduplicated_mcp_tools, tool_server_map @@ -520,15 +473,9 @@ class LiteLLM_Proxy_MCP_Handler: tool_calls: List[Any] = [] for output_item in response.output: # Check if this is a function call output item - if ( - isinstance(output_item, dict) - and output_item.get("type") == "function_call" - ): + if isinstance(output_item, dict) and output_item.get("type") == "function_call": tool_calls.append(output_item) - elif ( - hasattr(output_item, "type") - and getattr(output_item, "type") == "function_call" - ): + elif hasattr(output_item, "type") and getattr(output_item, "type") == "function_call": # Handle pydantic model case tool_calls.append(output_item) @@ -552,9 +499,7 @@ class LiteLLM_Proxy_MCP_Handler: else: tool_calls.append(tool_call) except Exception: - verbose_logger.exception( - "Failed to extract tool calls from chat completion response" - ) + verbose_logger.exception("Failed to extract tool calls from chat completion response") return tool_calls @@ -575,9 +520,7 @@ class LiteLLM_Proxy_MCP_Handler: tool_name = tool_call.get("name") tool_arguments = tool_call.get("arguments") else: - tool_call_id = getattr(tool_call, "call_id", None) or getattr( - tool_call, "id", None - ) + tool_call_id = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None) function_obj = getattr(tool_call, "function", None) if function_obj is not None: @@ -682,9 +625,7 @@ class LiteLLM_Proxy_MCP_Handler: verbose_logger.warning(f"Tool call missing name: {tool_call}") continue - parsed_arguments = LiteLLM_Proxy_MCP_Handler._parse_tool_arguments( - tool_arguments - ) + parsed_arguments = LiteLLM_Proxy_MCP_Handler._parse_tool_arguments(tool_arguments) # Import here to avoid circular import from litellm.proxy.proxy_server import proxy_logging_obj @@ -693,14 +634,8 @@ class LiteLLM_Proxy_MCP_Handler: # Remove the server name prefix if the tool name includes it. sanitized_tool_name = tool_name - unprefixed_name, prefixed_server_name = split_server_prefix_from_name( - tool_name - ) - if ( - prefixed_server_name - and prefixed_server_name == server_name - and unprefixed_name - ): + unprefixed_name, prefixed_server_name = split_server_prefix_from_name(tool_name) + if prefixed_server_name and prefixed_server_name == server_name and unprefixed_name: sanitized_tool_name = unprefixed_name start_time = datetime.now() @@ -743,9 +678,9 @@ class LiteLLM_Proxy_MCP_Handler: if user_api_key: logging_request_data["metadata"]["user_api_key"] = user_api_key - user_identifier = getattr( - user_api_key_auth, "end_user_id", None - ) or getattr(user_api_key_auth, "user_id", None) + user_identifier = getattr(user_api_key_auth, "end_user_id", None) or getattr( + user_api_key_auth, "user_id", None + ) if user_identifier: logging_request_data["user"] = user_identifier @@ -775,38 +710,28 @@ class LiteLLM_Proxy_MCP_Handler: api_key="", ) except Exception: - verbose_logger.exception( - "Failed to run pre_call for MCP tool logging" - ) + verbose_logger.exception("Failed to run pre_call for MCP tool logging") standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = { "name": sanitized_tool_name, "arguments": parsed_arguments, "namespaced_tool_name": tool_name, } - mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name( - tool_name - ) + mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(tool_name) if mcp_server: mcp_info = mcp_server.mcp_info or {} standard_logging_mcp_tool_call["mcp_server_name"] = ( - mcp_info.get("server_name") - or getattr(mcp_server, "server_name", None) - or server_name + mcp_info.get("server_name") or getattr(mcp_server, "server_name", None) or server_name ) logo_url = mcp_info.get("logo_url") if logo_url: standard_logging_mcp_tool_call["mcp_server_logo_url"] = logo_url cost_info = mcp_info.get("mcp_server_cost_info") if cost_info: - standard_logging_mcp_tool_call["mcp_server_cost_info"] = ( - cost_info - ) + standard_logging_mcp_tool_call["mcp_server_cost_info"] = cost_info if litellm_logging_obj: - litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( - standard_logging_mcp_tool_call - ) + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call litellm_logging_obj.model = f"MCP: {tool_name}" litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value @@ -838,9 +763,7 @@ class LiteLLM_Proxy_MCP_Handler: end_time=end_time, ) except Exception: - verbose_logger.exception( - "Failed to log MCP tool call success for %s", tool_name - ) + verbose_logger.exception("Failed to log MCP tool call success for %s", tool_name) # Format result for inclusion in response result_text = LiteLLM_Proxy_MCP_Handler._parse_mcp_result(result) @@ -859,9 +782,7 @@ class LiteLLM_Proxy_MCP_Handler: request_data=logging_request_data, error=e, ) - verbose_logger.error( - f"BlockedPiiEntityError in MCP tool call: {str(e)}" - ) + verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}") error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {str(e)}" tool_results.append( { @@ -877,10 +798,10 @@ class LiteLLM_Proxy_MCP_Handler: request_data=logging_request_data, error=e, ) - verbose_logger.error( - f"GuardrailRaisedException in MCP tool call: {str(e)}" + verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}") + error_message = ( + f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {str(e)}" ) - error_message = f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {str(e)}" tool_results.append( { "tool_call_id": tool_call_id, @@ -933,9 +854,7 @@ class LiteLLM_Proxy_MCP_Handler: from litellm.utils import convert_list_message_to_dict - follow_up_messages: List[Any] = convert_list_message_to_dict( - deepcopy(original_messages) - ) + follow_up_messages: List[Any] = convert_list_message_to_dict(deepcopy(original_messages)) if not follow_up_messages: follow_up_messages = [] @@ -943,9 +862,7 @@ class LiteLLM_Proxy_MCP_Handler: message_to_append: Optional[dict] = None try: first_choice = response.choices[0] - if isinstance(first_choice, Choices) and getattr( - first_choice, "message", None - ): + if isinstance(first_choice, Choices) and getattr(first_choice, "message", None): message_to_append = first_choice.message.model_dump(exclude_none=True) # Ensure tool_calls have arguments field (required by OpenAI API) if message_to_append.get("tool_calls"): @@ -983,9 +900,7 @@ class LiteLLM_Proxy_MCP_Handler: # Add original user input if available to maintain conversation context if original_input: if isinstance(original_input, str): - follow_up_input.append( - {"type": "message", "role": "user", "content": original_input} - ) + follow_up_input.append({"type": "message", "role": "user", "content": original_input}) elif isinstance(original_input, list): follow_up_input.extend(original_input) else: @@ -1173,9 +1088,7 @@ class LiteLLM_Proxy_MCP_Handler: return request_params @staticmethod - def _create_tool_execution_events( - tool_calls: List[Any], tool_results: List[Dict[str, Any]] - ) -> List[Any]: + def _create_tool_execution_events(tool_calls: List[Any], tool_results: List[Dict[str, Any]]) -> List[Any]: """ Create MCP tool execution events for streaming. @@ -1223,9 +1136,7 @@ class LiteLLM_Proxy_MCP_Handler: return tool_execution_events @staticmethod - def _prepare_initial_call_params( - call_params: Dict[str, Any], should_auto_execute: bool - ) -> Dict[str, Any]: + def _prepare_initial_call_params(call_params: Dict[str, Any], should_auto_execute: bool) -> Dict[str, Any]: """ Prepare call parameters for the initial LLM call. @@ -1241,9 +1152,7 @@ class LiteLLM_Proxy_MCP_Handler: return initial_params @staticmethod - def _prepare_follow_up_call_params( - call_params: Dict[str, Any], original_stream_setting: bool - ) -> Dict[str, Any]: + def _prepare_follow_up_call_params(call_params: Dict[str, Any], original_stream_setting: bool) -> Dict[str, Any]: """ Prepare call parameters for the follow-up LLM call after tool execution. diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 839e0232b04..a961271e3f0 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -41,9 +41,7 @@ async def create_mcp_list_tools_events( for tool in mcp_tools_with_litellm_proxy: if isinstance(tool, dict) and "server_url" in tool: server_url = tool.get("server_url") - if isinstance(server_url, str) and server_url.startswith( - "litellm_proxy/mcp/" - ): + if isinstance(server_url, str) and server_url.startswith("litellm_proxy/mcp/"): server_name = server_url.split("/")[-1] mcp_servers.append(server_name) @@ -88,9 +86,7 @@ async def create_mcp_list_tools_events( first_tool = mcp_tools_with_litellm_proxy[0] if isinstance(first_tool, dict): server_label_value = first_tool.get("server_label", "") - server_label = ( - str(server_label_value) if server_label_value is not None else "" - ) + server_label = str(server_label_value) if server_label_value is not None else "" # Format tools for OpenAI output_item.done format formatted_tools = [] @@ -269,7 +265,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.should_auto_execute = self._should_auto_execute_tools() # Streaming state management - self.phase = "initial_response" # initial_response -> mcp_discovery -> tool_execution -> follow_up_response -> finished + self.phase = ( + "initial_response" # initial_response -> mcp_discovery -> tool_execution -> follow_up_response -> finished + ) self.finished = False # Event queues and generation flags @@ -278,15 +276,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): ) self.tool_execution_events: List[ResponsesAPIStreamingResponse] = [] self.mcp_discovery_generated = True # Events are already generated - self.mcp_events = ( - mcp_events # Store the initial MCP events for backward compatibility - ) + self.mcp_events = mcp_events # Store the initial MCP events for backward compatibility self.tool_server_map = tool_server_map # Iterator references - self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = ( - base_iterator # Will be created when needed - ) + self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = base_iterator # Will be created when needed self.follow_up_iterator: Optional[Any] = None # Response collection for tool execution @@ -295,9 +289,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Set up model metadata (will be updated when we get the real iterator) self.model = self.original_request_params.get("model", "unknown") self.litellm_metadata = {} - self.custom_llm_provider = self.original_request_params.get( - "custom_llm_provider", None - ) + self.custom_llm_provider = self.original_request_params.get("custom_llm_provider", None) self.litellm_call_id = self.original_request_params.get("litellm_call_id") self.litellm_trace_id = self.original_request_params.get("litellm_trace_id") @@ -334,15 +326,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if raw_headers_from_request: headers_obj = Headers(raw_headers_from_request) - self.mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( - headers_obj - ) - self.mcp_server_auth_headers = ( - MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) - ) - self.oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers( - headers_obj - ) + self.mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers_obj) + self.mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) + self.oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers_obj) # Also check if headers are provided in tools array (from request body) tools = self.original_request_params.get("tools") @@ -353,10 +339,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if tool_headers and isinstance(tool_headers, dict): # Merge tool headers into mcp_server_auth_headers headers_obj_from_tool = Headers(tool_headers) - tool_mcp_server_auth_headers = ( - MCPRequestHandler._get_mcp_server_auth_headers_from_headers( - headers_obj_from_tool - ) + tool_mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers( + headers_obj_from_tool ) if tool_mcp_server_auth_headers: @@ -369,9 +353,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): ) in tool_mcp_server_auth_headers.items(): if server_alias not in self.mcp_server_auth_headers: self.mcp_server_auth_headers[server_alias] = {} - self.mcp_server_auth_headers[server_alias].update( - headers_dict - ) + self.mcp_server_auth_headers[server_alias].update(headers_dict) # Also merge raw headers if self.raw_headers is None: @@ -384,9 +366,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): LiteLLM_Proxy_MCP_Handler, ) - return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( - self.mcp_tools_with_litellm_proxy - ) + return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(self.mcp_tools_with_litellm_proxy) def __aiter__(self): return self @@ -489,9 +469,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): response_obj = getattr(chunk, "response", None) if response_obj and hasattr(response_obj, "id"): self._cached_response_id = response_obj.id - verbose_logger.debug( - f"Cached response ID: {self._cached_response_id}" - ) + verbose_logger.debug(f"Cached response ID: {self._cached_response_id}") # After emitting response.output_item.added, transition to MCP discovery if not self.initial_events_emitted and hasattr(chunk, "type"): @@ -519,9 +497,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): raise else: # base_iterator is not async iterable (likely a ResponsesAPIResponse) - if self.should_auto_execute and isinstance( - self.base_iterator, ResponsesAPIResponse - ): + if self.should_auto_execute and isinstance(self.base_iterator, ResponsesAPIResponse): self.collected_response = self.base_iterator self.phase = "tool_execution" await self._generate_tool_execution_events() @@ -534,9 +510,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): """Check if this chunk indicates the response is completed""" from litellm.types.llms.openai import ResponsesAPIStreamEvents - return ( - getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED - ) + return getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: """ @@ -552,9 +526,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): response_obj = getattr(chunk, "response", None) if response_obj and hasattr(response_obj, "id"): if response_obj.id != self._cached_response_id: - verbose_logger.debug( - f"Updating response ID from {response_obj.id} to {self._cached_response_id}" - ) + verbose_logger.debug(f"Updating response ID from {response_obj.id} to {self._cached_response_id}") response_obj.id = self._cached_response_id # If auto-execution is enabled, check for completed responses @@ -582,15 +554,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Use the pre-fetched all_tools from original_request_params (no re-processing needed) params_for_llm = {} for key, value in params.items(): - params_for_llm[key] = ( - value # Copy all params as-is since tools are already processed - ) + params_for_llm[key] = value # Copy all params as-is since tools are already processed - tools_count = ( - len(params_for_llm.get("tools", [])) - if params_for_llm.get("tools") - else 0 - ) + tools_count = len(params_for_llm.get("tools", [])) if params_for_llm.get("tools") else 0 verbose_logger.debug(f"Making LLM call with {tools_count} tools") response = await aresponses(**params_for_llm) @@ -600,12 +566,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Copy metadata from the real iterator self.model = getattr(response, "model", self.model) self.litellm_metadata = getattr(response, "litellm_metadata", {}) - self.custom_llm_provider = getattr( - response, "custom_llm_provider", self.custom_llm_provider - ) - verbose_logger.debug( - f"Created base iterator: {type(self.base_iterator)}" - ) + self.custom_llm_provider = getattr(response, "custom_llm_provider", self.custom_llm_provider) + verbose_logger.debug(f"Created base iterator: {type(self.base_iterator)}") else: # Non-streaming response - this shouldn't happen but handle it verbose_logger.warning(f"Got non-streaming response: {type(response)}") @@ -632,11 +594,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): try: # Extract tool calls from the response if self.collected_response is not None: - tool_calls = ( - LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response( - self.collected_response - ) - ) # type: ignore[arg-type] + tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_response(self.collected_response) # type: ignore[arg-type] else: tool_calls = [] if not tool_calls: diff --git a/litellm/responses/sse_output_recovery.py b/litellm/responses/sse_output_recovery.py index 5c18770a611..1546b9414ea 100644 --- a/litellm/responses/sse_output_recovery.py +++ b/litellm/responses/sse_output_recovery.py @@ -26,14 +26,8 @@ def parse_sse_json_chunk(chunk: str) -> Optional[Dict[str, Any]]: # Import locally to avoid a circular import with the streaming handler. from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - stripped_chunk = ( - CustomStreamWrapper._strip_sse_data_from_chunk(chunk.strip()) or "" - ).strip() - if ( - not stripped_chunk - or stripped_chunk == STREAM_SSE_DONE_STRING - or stripped_chunk.startswith("event:") - ): + stripped_chunk = (CustomStreamWrapper._strip_sse_data_from_chunk(chunk.strip()) or "").strip() + if not stripped_chunk or stripped_chunk == STREAM_SSE_DONE_STRING or stripped_chunk.startswith("event:"): return None try: parsed_chunk = json.loads(stripped_chunk) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 1f699e451dc..6df544dee3e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -89,13 +89,9 @@ class BaseResponsesAPIStreamingIterator: # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py _api_base = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get( - "litellm_params", {} - ), - ) - _model_info: Dict = ( - litellm_metadata.get("model_info", {}) if litellm_metadata else {} + optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), ) + _model_info: Dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} self._hidden_params = { "model_id": _model_info.get("id", None), "api_base": _api_base, @@ -139,24 +135,18 @@ class BaseResponsesAPIStreamingIterator: # Format as ResponsesAPIStreamingResponse if isinstance(parsed_chunk, dict): if self.responses_api_provider_config is None: - raise ValueError( - "responses_api_provider_config is required to process live streaming chunks" - ) - openai_responses_api_chunk = ( - self.responses_api_provider_config.transform_streaming_response( - model=self.model, - parsed_chunk=parsed_chunk, - logging_obj=self.logging_obj, - ) + raise ValueError("responses_api_provider_config is required to process live streaming chunks") + openai_responses_api_chunk = self.responses_api_provider_config.transform_streaming_response( + model=self.model, + parsed_chunk=parsed_chunk, + logging_obj=self.logging_obj, ) # Only when the SSE JSON carries a response body (delta events do not). # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a # truthy child Mock for any attribute, which breaks tests and is wrong on stream. if "response" in parsed_chunk: - response_object = getattr( - openai_responses_api_chunk, "response", None - ) + response_object = getattr(openai_responses_api_chunk, "response", None) if response_object is not None: response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=response_object, @@ -168,9 +158,7 @@ class BaseResponsesAPIStreamingIterator: # Encode container_id on streaming events so proxy/UI follow-ups route correctly _event_type = getattr(openai_responses_api_chunk, "type", None) _stream_model_id = ( - self.litellm_metadata.get("model_info", {}).get("id") - if self.litellm_metadata - else None + self.litellm_metadata.get("model_info", {}).get("id") if self.litellm_metadata else None ) if _event_type in ( ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -183,12 +171,8 @@ class BaseResponsesAPIStreamingIterator: custom_llm_provider=self.custom_llm_provider, model_id=_stream_model_id, ) - elif ( - _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED - ): - _annotation = getattr( - openai_responses_api_chunk, "annotation", None - ) + elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED: + _annotation = getattr(openai_responses_api_chunk, "annotation", None) if _annotation is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_annotation, @@ -212,9 +196,7 @@ class BaseResponsesAPIStreamingIterator: ) # Wrap encrypted_content in streaming events (output_item.added, output_item.done) - if self.litellm_metadata and self.litellm_metadata.get( - "encrypted_content_affinity_enabled" - ): + if self.litellm_metadata and self.litellm_metadata.get("encrypted_content_affinity_enabled"): openai_types = _get_openai_response_types() event_type = getattr(openai_responses_api_chunk, "type", None) if event_type in ( @@ -226,9 +208,7 @@ class BaseResponsesAPIStreamingIterator: encrypted_content = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): model_id = ( - self.litellm_metadata.get("model_info", {}).get( - "id" - ) + self.litellm_metadata.get("model_info", {}).get("id") if self.litellm_metadata else None ) @@ -248,23 +228,14 @@ class BaseResponsesAPIStreamingIterator: ): self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True - if ( - litellm.include_cost_in_streaming_usage - and self.logging_obj is not None - ): - response_obj: Optional[Any] = getattr( - openai_responses_api_chunk, "response", None - ) + if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: + response_obj: Optional[Any] = getattr(openai_responses_api_chunk, "response", None) if response_obj: - usage_obj: Optional[Any] = getattr( - response_obj, "usage", None - ) + usage_obj: Optional[Any] = getattr(response_obj, "usage", None) if usage_obj is not None: try: - cost: Optional[float] = ( - self.logging_obj._response_cost_calculator( - result=response_obj - ) + cost: Optional[float] = self.logging_obj._response_cost_calculator( + result=response_obj ) if cost is not None: setattr(usage_obj, "cost", cost) @@ -272,10 +243,7 @@ class BaseResponsesAPIStreamingIterator: # Best-effort usage cost annotation should not break stream replay. pass - if ( - _chunk_type - == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED - ): + if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() else: self._handle_logging_completed_response() @@ -306,13 +274,9 @@ class BaseResponsesAPIStreamingIterator: # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) logging_response = self.completed_response - if self.completed_response is not None and hasattr( - self.completed_response, "model_dump" - ): + if self.completed_response is not None and hasattr(self.completed_response, "model_dump"): try: - logging_response = type(self.completed_response).model_validate( - self.completed_response.model_dump() - ) + logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump()) except Exception: # Fallback to original if serialization fails pass @@ -358,11 +322,7 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj = ( - getattr(self.completed_response, "response", None) - if self.completed_response - else None - ) + response_obj = getattr(self.completed_response, "response", None) if self.completed_response else None error_info = getattr(response_obj, "error", None) if response_obj else None error_message = "Response failed" if isinstance(error_info, dict): @@ -393,10 +353,7 @@ class BaseResponsesAPIStreamingIterator: completed_response = self.completed_response openai_types = _get_openai_response_types() - if ( - getattr(completed_response, "type", None) - != openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED - ): + if getattr(completed_response, "type", None) != openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: return response_obj = self._get_completed_response_object() @@ -408,10 +365,7 @@ class BaseResponsesAPIStreamingIterator: return request_kwargs = getattr(caching_handler, "request_kwargs", None) - if ( - not isinstance(request_kwargs, dict) - or request_kwargs.get("stream") is not True - ): + if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: return request_kwargs = request_kwargs.copy() preset_cache_key = getattr(caching_handler, "preset_cache_key", None) @@ -471,15 +425,11 @@ class BaseResponsesAPIStreamingIterator: typed_call_type = None if typed_call_type is None: try: - typed_call_type = CallTypes( - getattr(self.logging_obj, "call_type", None) - ) + typed_call_type = CallTypes(getattr(self.logging_obj, "call_type", None)) except Exception: typed_call_type = None - request_data = self.request_data or getattr( - self.logging_obj, "model_call_details", {} - ) + request_data = self.request_data or getattr(self.logging_obj, "model_call_details", {}) callbacks = getattr(litellm, "callbacks", None) or [] hooks_ran = False for callback in callbacks: @@ -521,9 +471,9 @@ class BaseResponsesAPIStreamingIterator: pass if "litellm_params" not in request_payload: try: - request_payload["litellm_params"] = getattr( - self.logging_obj, "model_call_details", {} - ).get("litellm_params", {}) + request_payload["litellm_params"] = getattr(self.logging_obj, "model_call_details", {}).get( + "litellm_params", {} + ) except Exception: request_payload["litellm_params"] = {} @@ -818,10 +768,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): evt = self._events[self._idx] self._idx += 1 openai_types = _get_openai_response_types() - if ( - getattr(evt, "type", None) - == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED - ): + if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=True) return evt @@ -835,10 +782,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): evt = self._events[self._idx] self._idx += 1 openai_types = _get_openai_response_types() - if ( - getattr(evt, "type", None) - == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED - ): + if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=False) return evt @@ -891,10 +835,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): evt = self._events[self._idx] self._idx += 1 openai_types = _get_openai_response_types() - if ( - getattr(evt, "type", None) - == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED - ): + if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=True) return evt @@ -908,10 +849,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): evt = self._events[self._idx] self._idx += 1 openai_types = _get_openai_response_types() - if ( - getattr(evt, "type", None) - == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED - ): + if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED: self.completed_response = evt self._log_completed_response(is_async=False) return evt @@ -938,12 +876,8 @@ def _build_response_status_event( update={"status": "in_progress", "output": []}, ) if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED: - return openai_types.ResponseCreatedEvent( - type=event_type, response=in_progress_response - ) - return openai_types.ResponseInProgressEvent( - type=event_type, response=in_progress_response - ) + return openai_types.ResponseCreatedEvent(type=event_type, response=in_progress_response) + return openai_types.ResponseInProgressEvent(type=event_type, response=in_progress_response) def _build_content_part_done_event( @@ -1012,9 +946,7 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - for annotation_index, annotation in enumerate( - part_payload.get("annotations", []) or [] - ): + for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []): events.append( openai_types.OutputTextAnnotationAddedEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, @@ -1068,27 +1000,19 @@ def _build_synthetic_response_events( usage_obj: Optional[Any] = getattr(transformed, "usage", None) if usage_obj is not None: try: - cost: Optional[float] = logging_obj._response_cost_calculator( - result=transformed - ) + cost: Optional[float] = logging_obj._response_cost_calculator(result=transformed) if cost is not None: setattr(usage_obj, "cost", cost) except Exception: pass events: List[Any] = [ - _build_response_status_event( - openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed - ), - _build_response_status_event( - openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed - ), + _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), + _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), ] sequence_number = 0 - for output_index, output_item in enumerate( - getattr(transformed, "output", []) or [] - ): + for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): output_item_payload = _dump_response_object(output_item) item_id = str(output_item_payload.get("id") or transformed.id) item_type = output_item_payload.get("type") @@ -1097,16 +1021,12 @@ def _build_synthetic_response_events( openai_types.OutputItemAddedEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, - item=openai_types.BaseLiteLLMOpenAIResponseObject( - **output_item_payload - ), + item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload), ) ) if item_type == "message": - for content_index, part in enumerate( - output_item_payload.get("content", []) or [] - ): + for content_index, part in enumerate(output_item_payload.get("content", []) or []): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( @@ -1114,9 +1034,7 @@ def _build_synthetic_response_events( item_id=item_id, output_index=output_index, content_index=content_index, - part=openai_types.BaseLiteLLMOpenAIResponseObject( - **part_payload - ), + part=openai_types.BaseLiteLLMOpenAIResponseObject(**part_payload), ) ) _add_text_like_part_events( @@ -1155,9 +1073,7 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate( - output_item_payload.get("summary", []) or [] - ): + for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1189,9 +1105,7 @@ def _build_synthetic_response_events( output_index=output_index, sequence_number=sequence_number, summary_index=summary_index, - part=openai_types.BaseLiteLLMOpenAIResponseObject( - **summary_payload - ), + part=openai_types.BaseLiteLLMOpenAIResponseObject(**summary_payload), ) ) @@ -1201,9 +1115,7 @@ def _build_synthetic_response_events( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, sequence_number=sequence_number, - item=openai_types.BaseLiteLLMOpenAIResponseObject( - **output_item_payload - ), + item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload), ) ) @@ -1231,9 +1143,7 @@ RESPONSES_WS_LOGGED_EVENT_TYPES = [ "error", ] -RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES = frozenset( - {"input_text", "output_text", "text"} -) +RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES = frozenset({"input_text", "output_text", "text"}) class ResponsesWebSocketStreaming: @@ -1317,20 +1227,13 @@ class ResponsesWebSocketStreaming: if item.get("type") == "message" and item.get("role") == "user": content = item.get("content", []) if isinstance(content, str): - self.input_messages.append( - {"role": "user", "content": content} - ) + self.input_messages.append({"role": "user", "content": content}) elif isinstance(content, list): for c in content: - if ( - isinstance(c, dict) - and c.get("type") == "input_text" - ): + if isinstance(c, dict) and c.get("type") == "input_text": text = c.get("text", "") if text: - self.input_messages.append( - {"role": "user", "content": text} - ) + self.input_messages.append({"role": "user", "content": text}) except (json.JSONDecodeError, AttributeError, TypeError): pass @@ -1377,10 +1280,7 @@ class ResponsesWebSocketStreaming: _evt_type = json.loads(response_str).get("type") except (json.JSONDecodeError, TypeError): _evt_type = None - if ( - _evt_type in self._DELTA_EVENT_TYPES - or _evt_type in self._OUTPUT_DONE_EVENT_TYPES - ): + if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES: continue unmasked_str = self._unmask_response_event(response_str) @@ -1459,19 +1359,13 @@ class ResponsesWebSocketStreaming: modified = model_modified for cb in self.guardrail_callbacks: - presidio_config = cb.get_presidio_settings_from_request_data( - self.request_data - ) + presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) # response.create carries client text in two shapes: # flat: {"type": "response.create", "input": ..., "instructions": ...} # nested: {"type": "response.create", "response": {"input": ..., "instructions": ...}} # Mask "input" and "instructions" in both shapes so PII is never # forwarded unmasked regardless of where the client places it. - nested_response = ( - msg_obj.get("response") - if isinstance(msg_obj.get("response"), dict) - else None - ) + nested_response = msg_obj.get("response") if isinstance(msg_obj.get("response"), dict) else None text_containers: list[tuple[dict, str]] = [] for container in (msg_obj, nested_response): if container is None: @@ -1511,8 +1405,7 @@ class ResponsesWebSocketStreaming: for block in value: if ( isinstance(block, dict) - and block.get("type") - in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES + and block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES and isinstance(block.get("text"), str) ): block["text"] = await cb.check_pii( @@ -1567,9 +1460,7 @@ class ResponsesWebSocketStreaming: if not self.guardrail_callbacks: return response_str - pii_tokens: Dict[str, str] = (self.request_data.get("metadata") or {}).get( - "pii_tokens", {} - ) + pii_tokens: Dict[str, str] = (self.request_data.get("metadata") or {}).get("pii_tokens", {}) if not pii_tokens: return response_str @@ -1638,9 +1529,7 @@ class ResponsesWebSocketStreaming: modified = False for cb in self.output_guardrail_callbacks: - presidio_config = cb.get_presidio_settings_from_request_data( - self.request_data - ) + presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) response_obj = evt_obj.get("response") or {} if not isinstance(response_obj, dict): continue @@ -1791,9 +1680,9 @@ class ManagedResponsesWebSocketHandler: self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.litellm_metadata: Dict[str, Any] = litellm_metadata or {} - self.model_group: Optional[str] = self.litellm_metadata.get( - "model_group" - ) or self.litellm_metadata.get("deployment_model_name") + self.model_group: Optional[str] = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( + "deployment_model_name" + ) self.api_key = api_key self.api_base = api_base self.timeout = timeout @@ -1801,9 +1690,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message # Carry through safe pass-through kwargs (e.g. extra_headers) - self.extra_kwargs: Dict[str, Any] = { - k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS - } + self.extra_kwargs: Dict[str, Any] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. # Keyed by the DECODED (pre-encoding) response ID from response.completed. # This avoids the async DB-write race condition where spend logs haven't @@ -1826,17 +1713,13 @@ class ManagedResponsesWebSocketHandler: return json.dumps(chunk, default=str) return json.dumps(str(chunk)) except Exception as exc: - verbose_logger.debug( - "ManagedResponsesWS: failed to serialize chunk: %s", exc - ) + verbose_logger.debug("ManagedResponsesWS: failed to serialize chunk: %s", exc) return None async def _send_error(self, message: str, error_type: str = "server_error") -> None: try: await self.websocket.send_text( - json.dumps( - {"type": "error", "error": {"type": error_type, "message": message}} - ) + json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) ) except Exception: pass @@ -1848,9 +1731,7 @@ class ManagedResponsesWebSocketHandler: The key is the *decoded* response ID (the raw provider response ID before LiteLLM base64-encodes it into the ``resp_...`` format). """ - decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id( - previous_response_id - ) + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) raw_id = decoded.get("response_id", previous_response_id) return list(self._session_history.get(raw_id, [])) @@ -1870,9 +1751,7 @@ class ManagedResponsesWebSocketHandler: Returns *None* if the event doesn't contain a usable ID. """ resp_obj = completed_event.get("response", {}) - encoded_id: Optional[str] = ( - resp_obj.get("id") if isinstance(resp_obj, dict) else None - ) + encoded_id: Optional[str] = resp_obj.get("id") if isinstance(resp_obj, dict) else None if not encoded_id: return None decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id) @@ -1942,9 +1821,7 @@ class ManagedResponsesWebSocketHandler: try: msg_obj = json.loads(raw_message) except json.JSONDecodeError: - await self._send_error( - "Invalid JSON in response.create event", "invalid_request_error" - ) + await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None if msg_obj.get("type") != "response.create": # Silently ignore non-response.create messages (e.g. warmup pings) @@ -1963,9 +1840,7 @@ class ManagedResponsesWebSocketHandler: """Return True for synthetic warmup IDs that only exist on this connection.""" if not response_id: return False - decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id( - response_id - ) + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id) raw_id = decoded.get("response_id", response_id) return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX) @@ -2025,9 +1900,7 @@ class ManagedResponsesWebSocketHandler: """ nested = msg_obj.get("response") response_params: Dict[str, Any] = ( - nested - if isinstance(nested, dict) and nested - else {k: v for k, v in msg_obj.items() if k != "type"} + nested if isinstance(nested, dict) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { param: response_params[param] @@ -2090,9 +1963,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials( - self, call_kwargs: Dict[str, Any], model: Optional[str] = None - ) -> None: + def _inject_credentials(self, call_kwargs: Dict[str, Any], model: Optional[str] = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key @@ -2113,9 +1984,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _update_proxy_request(call_kwargs: Dict[str, Any], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" - proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get( - "proxy_server_request" - ) or {} + proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get("proxy_server_request") or {} if not isinstance(proxy_server_request, dict): return body = dict(proxy_server_request.get("body") or {}) @@ -2132,9 +2001,7 @@ class ManagedResponsesWebSocketHandler: call_kwargs.setdefault("litellm_params", {}) call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request - async def _stream_and_forward( - self, model: str, call_kwargs: Dict[str, Any] - ) -> Optional[Dict[str, Any]]: + async def _stream_and_forward(self, model: str, call_kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]: """ Stream ``litellm.aresponses`` and forward every chunk over the WebSocket. @@ -2148,9 +2015,7 @@ class ManagedResponsesWebSocketHandler: if chunk is None: continue # Read type from the object before serializing to avoid double JSON parse - chunk_type = getattr(chunk, "type", None) or ( - chunk.get("type") if isinstance(chunk, dict) else None - ) + chunk_type = getattr(chunk, "type", None) or (chunk.get("type") if isinstance(chunk, dict) else None) serialized = self._serialize_chunk(chunk) if serialized is None: continue @@ -2162,9 +2027,7 @@ class ManagedResponsesWebSocketHandler: try: await self.websocket.send_text(serialized) except Exception as send_exc: - verbose_logger.debug( - "ManagedResponsesWS: error sending chunk to client: %s", send_exc - ) + verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc) return completed_event # Client disconnected return completed_event @@ -2225,9 +2088,7 @@ class ManagedResponsesWebSocketHandler: try: await self._send_warmup_ack(msg_obj) except Exception as exc: - verbose_logger.debug( - "ManagedResponsesWS: error sending warmup ack: %s", exc - ) + verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc) return call_kwargs = self._build_base_call_kwargs(msg_obj) @@ -2243,33 +2104,21 @@ class ManagedResponsesWebSocketHandler: else: model = requested_model - previous_response_id: Optional[str] = call_kwargs.pop( - "previous_response_id", None - ) + previous_response_id: Optional[str] = call_kwargs.pop("previous_response_id", None) current_messages = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history - prior_history = ( - self._get_history_messages(previous_response_id) - if previous_response_id - else [] - ) + prior_history = self._get_history_messages(previous_response_id) if previous_response_id else [] - self._apply_history( - call_kwargs, previous_response_id, current_messages, prior_history - ) + self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history) self._inject_credentials(call_kwargs, model=model) - self._update_proxy_request( - call_kwargs, requested_model or self.model_group or model - ) + self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) try: completed_event = await self._stream_and_forward(model, call_kwargs) except Exception as exc: - verbose_logger.exception( - "ManagedResponsesWS: error processing response.create: %s", exc - ) + verbose_logger.exception("ManagedResponsesWS: error processing response.create: %s", exc) await self._send_error(str(exc)) return @@ -2292,9 +2141,7 @@ class ManagedResponsesWebSocketHandler: try: message = await self.websocket.receive_text() except Exception as exc: - verbose_logger.debug( - "ManagedResponsesWS: client disconnected: %s", exc - ) + verbose_logger.debug("ManagedResponsesWS: client disconnected: %s", exc) break await self._process_response_create(message) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 60badb57d2a..cff113dc3e5 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -51,9 +51,7 @@ class ResponsesAPIRequestUtils: if k not in supported_params: unsupported_params[k] = non_default_params[k] if unsupported_params: - if litellm.drop_params is True or ( - drop_params is not None and drop_params is True - ): + if litellm.drop_params is True or (drop_params is not None and drop_params is True): pass else: raise litellm.UnsupportedParamsError( @@ -83,9 +81,7 @@ class ResponsesAPIRequestUtils: # Remove None values and internal parameters # Get supported parameters for the model - supported_params = responses_api_provider_config.get_supported_openai_params( - model - ) + supported_params = responses_api_provider_config.get_supported_openai_params(model) non_default_params = cast(Dict, response_api_optional_params) # Check for unsupported parameters @@ -133,21 +129,21 @@ class ResponsesAPIRequestUtils: special_params = params.pop("kwargs", {}) additional_drop_params = params.pop("additional_drop_params", None) - non_default_params = ( - PreProcessNonDefaultParams.base_pre_process_non_default_params( - passed_params=params, - special_params=special_params, - custom_llm_provider=custom_llm_provider, - additional_drop_params=additional_drop_params, - default_param_values={k: None for k in valid_keys}, - additional_endpoint_specific_params=["input"], - ) + non_default_params = PreProcessNonDefaultParams.base_pre_process_non_default_params( + passed_params=params, + special_params=special_params, + custom_llm_provider=custom_llm_provider, + additional_drop_params=additional_drop_params, + default_param_values={k: None for k in valid_keys}, + additional_endpoint_specific_params=["input"], ) # decode previous_response_id if it's a litellm encoded id if "previous_response_id" in non_default_params: - decoded_previous_response_id = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id( - non_default_params["previous_response_id"] + decoded_previous_response_id = ( + ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id( + non_default_params["previous_response_id"] + ) ) non_default_params["previous_response_id"] = decoded_previous_response_id @@ -220,20 +216,16 @@ class ResponsesAPIRequestUtils: responses_api_response.id = updated_id if litellm_metadata.get("encrypted_content_affinity_enabled"): - responses_api_response = ( - ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( - response=responses_api_response, - model_id=model_id, - ) + responses_api_response = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response( + response=responses_api_response, + model_id=model_id, ) # Encode container IDs in the response output - responses_api_response = ( - ResponsesAPIRequestUtils._update_container_ids_in_response( - responses_api_response=responses_api_response, - custom_llm_provider=custom_llm_provider, - litellm_metadata=litellm_metadata, - ) + responses_api_response = ResponsesAPIRequestUtils._update_container_ids_in_response( + responses_api_response=responses_api_response, + custom_llm_provider=custom_llm_provider, + litellm_metadata=litellm_metadata, ) return responses_api_response @@ -275,9 +267,7 @@ class ResponsesAPIRequestUtils: return None @staticmethod - def _wrap_encrypted_content_with_model_id( - encrypted_content: str, model_id: str - ) -> str: + def _wrap_encrypted_content_with_model_id(encrypted_content: str, model_id: str) -> str: """Wrap encrypted_content with model_id metadata for affinity routing. When Codex or other clients send items with encrypted_content but no ID, @@ -316,9 +306,7 @@ class ResponsesAPIRequestUtils: if missing: metadata_b64 += "=" * (4 - missing) - decoded_metadata = base64.b64decode(metadata_b64.encode("utf-8")).decode( - "utf-8" - ) + decoded_metadata = base64.b64decode(metadata_b64.encode("utf-8")).decode("utf-8") model_id = decoded_metadata.replace("model_id:", "") return model_id, original_content except Exception: @@ -356,16 +344,12 @@ class ResponsesAPIRequestUtils: if encrypted_content and isinstance(encrypted_content, str): # Always wrap encrypted_content with model_id for redundancy - item["encrypted_content"] = ( - ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - encrypted_content, model_id - ) + item["encrypted_content"] = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id ) # Also encode the ID if present if item_id and isinstance(item_id, str): - item["id"] = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, item_id - ) + item["id"] = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, item_id) else: item_id = getattr(item, "id", None) encrypted_content = getattr(item, "encrypted_content", None) @@ -381,9 +365,7 @@ class ResponsesAPIRequestUtils: # Also encode the ID if present if item_id and isinstance(item_id, str): try: - item.id = ResponsesAPIRequestUtils._build_encrypted_item_id( - model_id, item_id - ) + item.id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, item_id) except AttributeError: pass @@ -407,9 +389,7 @@ class ResponsesAPIRequestUtils: if isinstance(item, dict): item_id = item.get("id") if item_id and isinstance(item_id, str): - decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id( - item_id - ) + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) if decoded: item["id"] = decoded["item_id"] @@ -418,9 +398,7 @@ class ResponsesAPIRequestUtils: ( _, unwrapped, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - encrypted_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) if unwrapped != encrypted_content: item["encrypted_content"] = unwrapped @@ -433,12 +411,10 @@ class ResponsesAPIRequestUtils: response_id: str, ) -> str: """Build the responses_api_response_id""" - assembled_id: str = str( - SpecialEnums.LITELLM_MANAGED_RESPONSE_COMPLETE_STR.value - ).format(custom_llm_provider, model_id, response_id) - base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode( - "utf-8" + assembled_id: str = str(SpecialEnums.LITELLM_MANAGED_RESPONSE_COMPLETE_STR.value).format( + custom_llm_provider, model_id, response_id ) + base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode("utf-8") return f"resp_{base64_encoded_id}" @staticmethod @@ -470,16 +446,12 @@ class ResponsesAPIRequestUtils: custom_llm_provider = None model_id = None - if ( - len(parts) >= 3 - ): # Full format with custom_llm_provider, model_id, and response_id + if len(parts) >= 3: # Full format with custom_llm_provider, model_id, and response_id custom_llm_provider_part = parts[0] model_id_part = parts[1] response_part = parts[2] - custom_llm_provider = custom_llm_provider_part.replace( - "litellm:custom_llm_provider:", "" - ) + custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "") model_id = model_id_part.replace("model_id:", "") decoded_response_id = response_part.replace("response_id:", "") else: @@ -503,9 +475,7 @@ class ResponsesAPIRequestUtils: """Get the model_id from the response_id""" if response_id is None: return None - decoded_response_id = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id) - ) + decoded_response_id = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id) return decoded_response_id.get("model_id") or None @staticmethod @@ -525,11 +495,7 @@ class ResponsesAPIRequestUtils: Returns: The original previous_response_id """ - decoded_response_id = ( - ResponsesAPIRequestUtils._decode_responses_api_response_id( - previous_response_id - ) - ) + decoded_response_id = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) return decoded_response_id.get("response_id", previous_response_id) @staticmethod @@ -546,9 +512,7 @@ class ResponsesAPIRequestUtils: provider_part = "" if custom_llm_provider is None else custom_llm_provider model_part = "" if model_id is None else model_id assembled_id = f"litellm:custom_llm_provider:{provider_part};model_id:{model_part};container_id:{container_id}" - base64_encoded_id = base64.b64encode(assembled_id.encode("utf-8")).decode( - "utf-8" - ) + base64_encoded_id = base64.b64encode(assembled_id.encode("utf-8")).decode("utf-8") return f"cntr_{base64_encoded_id}" @staticmethod @@ -746,9 +710,7 @@ class ResponsesAPIRequestUtils: if not annotations or not isinstance(annotations, list): return for ann in annotations: - ResponsesAPIRequestUtils._collect_container_ids_from_output_item( - ann, collected - ) + ResponsesAPIRequestUtils._collect_container_ids_from_output_item(ann, collected) @staticmethod def _collect_container_ids_from_message_content( @@ -801,9 +763,7 @@ class ResponsesAPIRequestUtils: nested_obj = getattr(item, "code_interpreter_call", None) if nested_obj is not None: - ResponsesAPIRequestUtils._collect_container_ids_from_output_item( - nested_obj, collected - ) + ResponsesAPIRequestUtils._collect_container_ids_from_output_item(nested_obj, collected) if getattr(item, "type", None) == "message": ResponsesAPIRequestUtils._collect_container_ids_from_message_content( @@ -825,9 +785,7 @@ class ResponsesAPIRequestUtils: collected: set[str] = set() if output: for item in output: - ResponsesAPIRequestUtils._collect_container_ids_from_output_item( - item, collected - ) + ResponsesAPIRequestUtils._collect_container_ids_from_output_item(item, collected) return list(collected) @staticmethod @@ -929,15 +887,9 @@ class ResponsesAPIRequestUtils: if raw_headers_from_request: headers_obj = Headers(raw_headers_from_request) - mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( - headers_obj - ) - mcp_server_auth_headers = ( - MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) - ) - oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers( - headers_obj - ) + mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers_obj) + mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) + oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers_obj) if tools: for tool in tools: @@ -947,10 +899,8 @@ class ResponsesAPIRequestUtils: # Merge tool headers into mcp_server_auth_headers # Extract server-specific headers from tool.headers headers_obj_from_tool = Headers(tool_headers) - tool_mcp_server_auth_headers = ( - MCPRequestHandler._get_mcp_server_auth_headers_from_headers( - headers_obj_from_tool - ) + tool_mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers( + headers_obj_from_tool ) if tool_mcp_server_auth_headers: if mcp_server_auth_headers is None: @@ -962,9 +912,7 @@ class ResponsesAPIRequestUtils: ) in tool_mcp_server_auth_headers.items(): if server_alias not in mcp_server_auth_headers: mcp_server_auth_headers[server_alias] = {} - mcp_server_auth_headers[server_alias].update( - headers_dict - ) + mcp_server_auth_headers[server_alias].update(headers_dict) # Also merge raw headers (non-prefixed headers from tool.headers) if raw_headers_from_request is None: raw_headers_from_request = {} @@ -1008,18 +956,10 @@ class ResponseAPILoggingUtils: if isinstance(usage_input, dict): usage_input = dict(usage_input) # shallow copy; avoid mutating caller # Realtime *_token_details → *_tokens_details when unset. - if ( - usage_input.get("input_tokens_details") is None - and "input_token_details" in usage_input - ): + if usage_input.get("input_tokens_details") is None and "input_token_details" in usage_input: usage_input["input_tokens_details"] = usage_input["input_token_details"] - if ( - usage_input.get("output_tokens_details") is None - and "output_token_details" in usage_input - ): - usage_input["output_tokens_details"] = usage_input[ - "output_token_details" - ] + if usage_input.get("output_tokens_details") is None and "output_token_details" in usage_input: + usage_input["output_tokens_details"] = usage_input["output_token_details"] total_tokens = usage_input.get("total_tokens") if total_tokens is None: input_tokens = usage_input.get("input_tokens") @@ -1035,33 +975,19 @@ class ResponseAPILoggingUtils: prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None if response_api_usage.input_tokens_details: if isinstance(response_api_usage.input_tokens_details, dict): - prompt_tokens_details = PromptTokensDetailsWrapper( - **response_api_usage.input_tokens_details - ) + prompt_tokens_details = PromptTokensDetailsWrapper(**response_api_usage.input_tokens_details) else: prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=getattr( - response_api_usage.input_tokens_details, "cached_tokens", None - ), - audio_tokens=getattr( - response_api_usage.input_tokens_details, "audio_tokens", None - ), - text_tokens=getattr( - response_api_usage.input_tokens_details, "text_tokens", None - ), - image_tokens=getattr( - response_api_usage.input_tokens_details, "image_tokens", None - ), + cached_tokens=getattr(response_api_usage.input_tokens_details, "cached_tokens", None), + audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), + text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), + image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), ) completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None - output_tokens_details = getattr( - response_api_usage, "output_tokens_details", None - ) + output_tokens_details = getattr(response_api_usage, "output_tokens_details", None) if output_tokens_details: completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=getattr( - output_tokens_details, "reasoning_tokens", None - ), + reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None), image_tokens=getattr(output_tokens_details, "image_tokens", None), text_tokens=getattr(output_tokens_details, "text_tokens", None), audio_tokens=getattr(output_tokens_details, "audio_tokens", None), diff --git a/litellm/router.py b/litellm/router.py index 4aae07bd8bd..1d23e838563 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -247,9 +247,7 @@ class Router: def __init__( self, - model_list: Optional[ - Union[List[DeploymentTypedDict], List[Dict[str, Any]]] - ] = None, + model_list: Optional[Union[List[DeploymentTypedDict], List[Dict[str, Any]]]] = None, ## ASSISTANTS API ## assistants_config: Optional[AssistantsTypedDict] = None, ## SEARCH API ## @@ -264,54 +262,34 @@ class Router: redis_db: Optional[int] = None, cache_responses: Optional[bool] = False, cache_kwargs: dict = {}, # additional kwargs to pass to RedisCache (see caching.py) - caching_groups: Optional[ - List[tuple] - ] = None, # if you want to cache across model groups + caching_groups: Optional[List[tuple]] = None, # if you want to cache across model groups client_ttl: int = 3600, # ttl for cached clients - will re-initialize after this time in seconds ## SCHEDULER ## polling_interval: Optional[float] = None, default_priority: Optional[int] = None, ## RELIABILITY ## num_retries: Optional[int] = None, - max_fallbacks: Optional[ - int - ] = None, # max fallbacks to try before exiting the call. Defaults to 5. + max_fallbacks: Optional[int] = None, # max fallbacks to try before exiting the call. Defaults to 5. timeout: Optional[float] = None, stream_timeout: Optional[float] = None, - default_litellm_params: Optional[ - dict - ] = None, # default params for Router.chat.completion.create + default_litellm_params: Optional[dict] = None, # default params for Router.chat.completion.create default_max_parallel_requests: Optional[int] = None, set_verbose: bool = False, debug_level: Literal["DEBUG", "INFO"] = "INFO", - default_fallbacks: Optional[ - List[str] - ] = None, # generic fallbacks, works across all deployments + default_fallbacks: Optional[List[str]] = None, # generic fallbacks, works across all deployments fallbacks: List = [], context_window_fallbacks: List = [], content_policy_fallbacks: List = [], - model_group_alias: Optional[ - Dict[str, Union[str, RouterModelGroupAliasItem]] - ] = {}, + model_group_alias: Optional[Dict[str, Union[str, RouterModelGroupAliasItem]]] = {}, enable_pre_call_checks: bool = False, enable_tag_filtering: bool = False, tag_filtering_match_any: bool = True, retry_after: int = 0, # min time to wait before retrying a failed request - retry_policy: Optional[ - Union[RetryPolicy, dict] - ] = None, # set custom retries for different exceptions - model_group_retry_policy: Dict[ - str, RetryPolicy - ] = {}, # set custom retry policies based on model group - allowed_fails: Optional[ - int - ] = None, # Number of times a deployment can failbefore being added to cooldown - allowed_fails_policy: Optional[ - AllowedFailsPolicy - ] = None, # set custom allowed fails policy - cooldown_time: Optional[ - float - ] = None, # (seconds) time to cooldown a deployment after failure + retry_policy: Optional[Union[RetryPolicy, dict]] = None, # set custom retries for different exceptions + model_group_retry_policy: Dict[str, RetryPolicy] = {}, # set custom retry policies based on model group + allowed_fails: Optional[int] = None, # Number of times a deployment can failbefore being added to cooldown + allowed_fails_policy: Optional[AllowedFailsPolicy] = None, # set custom allowed fails policy + cooldown_time: Optional[float] = None, # (seconds) time to cooldown a deployment after failure disable_cooldowns: Optional[bool] = None, routing_strategy: Literal[ "simple-shuffle", @@ -327,9 +305,7 @@ class Router: routing_groups: Optional[List[Union[RoutingGroup, dict]]] = None, provider_budget_config: Optional[GenericBudgetConfigType] = None, alerting_config: Optional[AlertingConfig] = None, - router_general_settings: Optional[ - RouterGeneralSettings - ] = RouterGeneralSettings(), + router_general_settings: Optional[RouterGeneralSettings] = RouterGeneralSettings(), deployment_affinity_ttl_seconds: int = 3600, model_group_affinity_config: Optional[Dict[str, List[str]]] = None, ignore_invalid_deployments: bool = False, @@ -425,9 +401,7 @@ class Router: verbose_router_logger.setLevel(logging.INFO) elif debug_level == "DEBUG": verbose_router_logger.setLevel(logging.DEBUG) - self.router_general_settings: RouterGeneralSettings = ( - router_general_settings or RouterGeneralSettings() - ) + self.router_general_settings: RouterGeneralSettings = router_general_settings or RouterGeneralSettings() self.assistants_config = assistants_config self.search_tools = search_tools or [] @@ -435,9 +409,7 @@ class Router: self.deployment_names: List = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### - cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = ( - "local" # default to an in-memory cache - ) + cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = "local" # default to an in-memory cache redis_cache = None cache_config: Dict[str, Any] = {} @@ -477,17 +449,15 @@ class Router: ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. ### SCHEDULER ### - self.scheduler = Scheduler( - polling_interval=polling_interval, redis_cache=redis_cache - ) + self.scheduler = Scheduler(polling_interval=polling_interval, redis_cache=redis_cache) self.default_priority = default_priority - self.default_deployment = None # use this to track the users default deployment, when they want to use model = * + self.default_deployment = ( + None # use this to track the users default deployment, when they want to use model = * + ) self.default_max_parallel_requests = default_max_parallel_requests self.provider_default_deployment_ids: List[str] = [] self.pattern_router = PatternMatchRouter() - self.team_pattern_routers: Dict[ - str, PatternMatchRouter - ] = {} # {"TEAM_ID": PatternMatchRouter} + self.team_pattern_routers: Dict[str, PatternMatchRouter] = {} # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} self.complexity_routers: Dict[str, "ComplexityRouter"] = {} self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {} @@ -532,9 +502,7 @@ class Router: else: self.allowed_fails = litellm.allowed_fails self.cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS - self.cooldown_cache = CooldownCache( - cache=self.cache, default_cooldown_time=self.cooldown_time - ) + self.cooldown_cache = CooldownCache(cache=self.cache, default_cooldown_time=self.cooldown_time) self.disable_cooldowns = disable_cooldowns self.enable_health_check_routing = enable_health_check_routing self.enable_weighted_failover = enable_weighted_failover @@ -542,9 +510,7 @@ class Router: _staleness = health_check_staleness_threshold or ( DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER ) - self.health_state_cache = DeploymentHealthCache( - cache=self.cache, staleness_threshold=float(_staleness) - ) + self.health_state_cache = DeploymentHealthCache(cache=self.cache, staleness_threshold=float(_staleness)) self.failed_calls = InMemoryCache() # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown if num_retries is not None: @@ -566,16 +532,12 @@ class Router: # Per-attempt request_timeout, independent of router_settings.timeout. # Only stored when a router timeout is also set, since otherwise # request_timeout already flows through self.timeout above. - self.request_timeout = ( - get_configured_request_timeout() if timeout is not None else None - ) + self.request_timeout = get_configured_request_timeout() if timeout is not None else None self.stream_timeout = stream_timeout self.retry_after = retry_after self.routing_strategy = self._normalize_strategy(routing_strategy) - self._routing_groups_input: Optional[List[Union[RoutingGroup, dict]]] = ( - routing_groups - ) + self._routing_groups_input: Optional[List[Union[RoutingGroup, dict]]] = routing_groups ## SETTING FALLBACKS ## ### validate if it's set + in correct format @@ -592,24 +554,14 @@ class Router: else: self.fallbacks = [{"*": _fallbacks}] - self.context_window_fallbacks = ( - context_window_fallbacks or litellm.context_window_fallbacks - ) + self.context_window_fallbacks = context_window_fallbacks or litellm.context_window_fallbacks - _content_policy_fallbacks = ( - content_policy_fallbacks or litellm.content_policy_fallbacks - ) + _content_policy_fallbacks = content_policy_fallbacks or litellm.content_policy_fallbacks self.validate_fallbacks(fallback_param=_content_policy_fallbacks) self.content_policy_fallbacks = _content_policy_fallbacks - self.total_calls: defaultdict = defaultdict( - int - ) # dict to store total calls made to each model - self.fail_calls: defaultdict = defaultdict( - int - ) # dict to store fail_calls made to each model - self.success_calls: defaultdict = defaultdict( - int - ) # dict to store success_calls made to each model + self.total_calls: defaultdict = defaultdict(int) # dict to store total calls made to each model + self.fail_calls: defaultdict = defaultdict(int) # dict to store fail_calls made to each model + self.success_calls: defaultdict = defaultdict(int) # dict to store success_calls made to each model self.previous_models: List = [] # list to store failed calls (passed in as metadata to next call) # make Router.chat.completions.create compatible for openai.chat.completions.create @@ -620,9 +572,7 @@ class Router: self.default_litellm_params = default_litellm_params self.default_litellm_params.setdefault("timeout", timeout) self.default_litellm_params.setdefault("max_retries", 0) - self.default_litellm_params.setdefault("metadata", {}).update( - {"caching_groups": caching_groups} - ) + self.default_litellm_params.setdefault("metadata", {}).update({"caching_groups": caching_groups}) self.deployment_stats: dict = {} # used for debugging load balancing """ @@ -653,17 +603,11 @@ class Router: self.access_groups = None ## USAGE TRACKING ## if isinstance(litellm._async_success_callback, list): - litellm.logging_callback_manager.add_litellm_async_success_callback( - self.deployment_callback_on_success - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) else: - litellm.logging_callback_manager.add_litellm_async_success_callback( - self.deployment_callback_on_success - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) if isinstance(litellm.success_callback, list): - litellm.logging_callback_manager.add_litellm_success_callback( - self.sync_deployment_callback_on_success - ) + litellm.logging_callback_manager.add_litellm_success_callback(self.sync_deployment_callback_on_success) else: litellm.success_callback = [self.sync_deployment_callback_on_success] if isinstance(litellm._async_failure_callback, list): @@ -671,14 +615,10 @@ class Router: self.async_deployment_callback_on_failure ) else: - litellm._async_failure_callback = [ - self.async_deployment_callback_on_failure - ] + litellm._async_failure_callback = [self.async_deployment_callback_on_failure] ## COOLDOWNS ## if isinstance(litellm.failure_callback, list): - litellm.logging_callback_manager.add_litellm_failure_callback( - self.deployment_callback_on_failure - ) + litellm.logging_callback_manager.add_litellm_failure_callback(self.deployment_callback_on_failure) else: litellm.failure_callback = [self.deployment_callback_on_failure] self.routing_strategy_args = routing_strategy_args @@ -705,12 +645,8 @@ class Router: ) ) - self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = ( - model_group_retry_policy - ) - self.model_group_affinity_config: Optional[Dict[str, List[str]]] = ( - model_group_affinity_config - ) + self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = model_group_retry_policy + self.model_group_affinity_config: Optional[Dict[str, List[str]]] = model_group_affinity_config self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None if allowed_fails_policy is not None: @@ -735,8 +671,7 @@ class Router: # enabled, we still need the DeploymentAffinityCheck callback (with global # flags all False) so per-group config can activate affinity per model group. if self.model_group_affinity_config and not any( - isinstance(cb, DeploymentAffinityCheck) - for cb in (self.optional_callbacks or []) + isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or []) ): if self.optional_callbacks is None: self.optional_callbacks = [] @@ -783,27 +718,13 @@ class Router: Pseudo-destructor to be invoked to clean up global data structures when router is no longer used. For now, unhook router's callbacks from all lists """ - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm._async_success_callback, self - ) - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.success_callback, self - ) - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm._async_failure_callback, self - ) - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.failure_callback, self - ) - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.input_callback, self - ) - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.service_callback, self - ) - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.callbacks, self - ) + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_success_callback, self) + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.success_callback, self) + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_failure_callback, self) + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.failure_callback, self) + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.input_callback, self) + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.service_callback, self) + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, self) # Remove ForwardClientSideHeadersByModelGroup if it exists if self.optional_callbacks is not None: @@ -865,19 +786,12 @@ class Router: return strategy.value return strategy - def _validate_routing_strategy( - self, routing_strategy: Union[RoutingStrategy, str, None] - ) -> None: + def _validate_routing_strategy(self, routing_strategy: Union[RoutingStrategy, str, None]) -> None: # See: https://github.com/BerriAI/litellm/issues/11330 - valid_strategy_strings = ["simple-shuffle", "lar1"] + [ - s.value for s in RoutingStrategy - ] + valid_strategy_strings = ["simple-shuffle", "lar1"] + [s.value for s in RoutingStrategy] if routing_strategy is None: return - is_valid_string = ( - isinstance(routing_strategy, str) - and routing_strategy in valid_strategy_strings - ) + is_valid_string = isinstance(routing_strategy, str) and routing_strategy in valid_strategy_strings is_valid_enum = isinstance(routing_strategy, RoutingStrategy) if not is_valid_string and not is_valid_enum: raise ValueError( @@ -928,11 +842,7 @@ class Router: routing_args={}, ) - if ( - selector is not None - and register_callbacks - and isinstance(litellm.callbacks, list) - ): + if selector is not None and register_callbacks and isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(selector) # type: ignore return selector @@ -948,26 +858,17 @@ class Router: if not selector_ids: return if isinstance(litellm.callbacks, list): - litellm.callbacks = [ - c for c in litellm.callbacks if id(c) not in selector_ids - ] + litellm.callbacks = [c for c in litellm.callbacks if id(c) not in selector_ids] if isinstance(litellm.input_callback, list): - litellm.input_callback = [ - c for c in litellm.input_callback if id(c) not in selector_ids - ] + litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] - def routing_strategy_init( - self, routing_strategy: Union[RoutingStrategy, str], routing_strategy_args: dict - ): + def routing_strategy_init(self, routing_strategy: Union[RoutingStrategy, str], routing_strategy_args: dict): verbose_router_logger.info(f"Routing strategy: {routing_strategy}") self._validate_routing_strategy(routing_strategy) self._reset_custom_routing_strategy() self._unregister_router_selectors( - [ - getattr(self, attr, None) - for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values() - ] + [getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()] ) self.leastbusy_logger: Optional[LeastBusyLoggingHandler] = None @@ -980,9 +881,7 @@ class Router: strategy=routing_strategy, routing_strategy_args=routing_strategy_args, ) - attr = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get( - self._normalize_strategy(routing_strategy) or "" - ) + attr = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(self._normalize_strategy(routing_strategy) or "") # TODO: legacy `self._logger` attributes are read directly by # `get_settings()` and external callers. Fold the default group into # `self._group_selectors["default"]` and drop these attribute writes — @@ -1004,11 +903,7 @@ class Router: attributes set up in `routing_strategy_init`. """ self._unregister_router_selectors( - [ - sel - for selectors in getattr(self, "_group_selectors", {}).values() - for sel in selectors.values() - ] + [sel for selectors in getattr(self, "_group_selectors", {}).values() for sel in selectors.values()] ) self._routing_groups: Dict[str, RoutingGroup] = {} @@ -1018,9 +913,7 @@ class Router: if not groups_input: return - known_model_names = { - m.get("model_name") for m in (self.model_list or []) if m.get("model_name") - } + known_model_names = {m.get("model_name") for m in (self.model_list or []) if m.get("model_name")} seen_group_names: set = set() for raw in groups_input: @@ -1029,9 +922,7 @@ class Router: if not group.group_name: raise ValueError("routing_groups: group_name must be non-empty.") if group.group_name == "default": - raise ValueError( - "routing_groups: 'default' is reserved for the implicit fallback group." - ) + raise ValueError("routing_groups: 'default' is reserved for the implicit fallback group.") if group.group_name in seen_group_names: raise ValueError( f"routing_groups: group names must be unique, duplicate group_name '{group.group_name}'." @@ -1086,17 +977,13 @@ class Router: strategy = self._normalize_strategy(self.routing_strategy) attr = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") selector = getattr(self, attr, None) if attr is not None else None - verbose_router_logger.debug( - "routing_group=default model=%s strategy=%s", model, strategy - ) + verbose_router_logger.debug("routing_group=default model=%s strategy=%s", model, strategy) return strategy, selector group = self._routing_groups[group_name] strategy = self._normalize_strategy(group.routing_strategy) selector = self._group_selectors.get(group_name, {}).get(strategy or "") - verbose_router_logger.debug( - "routing_group=%s model=%s strategy=%s", group_name, model, strategy - ) + verbose_router_logger.debug("routing_group=%s model=%s strategy=%s", group_name, model, strategy) return strategy, selector async def _select_deployment_async( @@ -1210,49 +1097,23 @@ class Router: def _initialize_core_endpoints(self): """Helper to initialize core router endpoints.""" - self.amoderation = self.factory_function( - litellm.amoderation, call_type="moderation" - ) - self.aanthropic_messages = self.factory_function( - litellm.anthropic_messages, call_type="anthropic_messages" - ) - self.anthropic_messages = self.factory_function( - litellm.anthropic_messages, call_type="anthropic_messages" - ) - self.agenerate_content = self.factory_function( - litellm.agenerate_content, call_type="agenerate_content" - ) + self.amoderation = self.factory_function(litellm.amoderation, call_type="moderation") + self.aanthropic_messages = self.factory_function(litellm.anthropic_messages, call_type="anthropic_messages") + self.anthropic_messages = self.factory_function(litellm.anthropic_messages, call_type="anthropic_messages") + self.agenerate_content = self.factory_function(litellm.agenerate_content, call_type="agenerate_content") self.aadapter_generate_content = self.factory_function( litellm.aadapter_generate_content, call_type="aadapter_generate_content" ) - self.aresponses = self.factory_function( - litellm.aresponses, call_type="aresponses" - ) - self.afile_delete = self.factory_function( - litellm.afile_delete, call_type="afile_delete" - ) - self.afile_content = self.factory_function( - litellm.afile_content, call_type="afile_content" - ) + self.aresponses = self.factory_function(litellm.aresponses, call_type="aresponses") + self.afile_delete = self.factory_function(litellm.afile_delete, call_type="afile_delete") + self.afile_content = self.factory_function(litellm.afile_content, call_type="afile_content") self.responses = self.factory_function(litellm.responses, call_type="responses") - self.aget_responses = self.factory_function( - litellm.aget_responses, call_type="aget_responses" - ) - self.acancel_responses = self.factory_function( - litellm.acancel_responses, call_type="acancel_responses" - ) - self.acompact_responses = self.factory_function( - litellm.acompact_responses, call_type="acompact_responses" - ) - self.adelete_responses = self.factory_function( - litellm.adelete_responses, call_type="adelete_responses" - ) - self.alist_input_items = self.factory_function( - litellm.alist_input_items, call_type="alist_input_items" - ) - self._arealtime = self.factory_function( - litellm._arealtime, call_type="_arealtime" - ) + self.aget_responses = self.factory_function(litellm.aget_responses, call_type="aget_responses") + self.acancel_responses = self.factory_function(litellm.acancel_responses, call_type="acancel_responses") + self.acompact_responses = self.factory_function(litellm.acompact_responses, call_type="acompact_responses") + self.adelete_responses = self.factory_function(litellm.adelete_responses, call_type="adelete_responses") + self.alist_input_items = self.factory_function(litellm.alist_input_items, call_type="alist_input_items") + self._arealtime = self.factory_function(litellm._arealtime, call_type="_arealtime") self._aresponses_websocket = self.factory_function( litellm._aresponses_websocket, call_type="_aresponses_websocket" ) @@ -1268,12 +1129,8 @@ class Router: self.aretrieve_fine_tuning_job = self.factory_function( litellm.aretrieve_fine_tuning_job, call_type="aretrieve_fine_tuning_job" ) - self.afile_list = self.factory_function( - litellm.afile_list, call_type="alist_files" - ) - self.aimage_edit = self.factory_function( - litellm.aimage_edit, call_type="aimage_edit" - ) + self.afile_list = self.factory_function(litellm.afile_list, call_type="alist_files") + self.aimage_edit = self.factory_function(litellm.aimage_edit, call_type="aimage_edit") self.allm_passthrough_route = self.factory_function( litellm.allm_passthrough_route, call_type="allm_passthrough_route" ) @@ -1296,39 +1153,17 @@ class Router: update, ) - self.avector_store_search = self.factory_function( - asearch, call_type="avector_store_search" - ) - self.vector_store_search = self.factory_function( - search, call_type="vector_store_search" - ) - self.vector_store_create = self.factory_function( - create, call_type="vector_store_create" - ) - self.avector_store_retrieve = self.factory_function( - aretrieve, call_type="avector_store_retrieve" - ) - self.vector_store_retrieve = self.factory_function( - retrieve, call_type="vector_store_retrieve" - ) - self.avector_store_list = self.factory_function( - alist, call_type="avector_store_list" - ) - self.vector_store_list = self.factory_function( - list, call_type="vector_store_list" - ) - self.avector_store_update = self.factory_function( - aupdate, call_type="avector_store_update" - ) - self.vector_store_update = self.factory_function( - update, call_type="vector_store_update" - ) - self.avector_store_delete = self.factory_function( - adelete, call_type="avector_store_delete" - ) - self.vector_store_delete = self.factory_function( - delete, call_type="vector_store_delete" - ) + self.avector_store_search = self.factory_function(asearch, call_type="avector_store_search") + self.vector_store_search = self.factory_function(search, call_type="vector_store_search") + self.vector_store_create = self.factory_function(create, call_type="vector_store_create") + self.avector_store_retrieve = self.factory_function(aretrieve, call_type="avector_store_retrieve") + self.vector_store_retrieve = self.factory_function(retrieve, call_type="vector_store_retrieve") + self.avector_store_list = self.factory_function(alist, call_type="avector_store_list") + self.vector_store_list = self.factory_function(list, call_type="vector_store_list") + self.avector_store_update = self.factory_function(aupdate, call_type="avector_store_update") + self.vector_store_update = self.factory_function(update, call_type="vector_store_update") + self.avector_store_delete = self.factory_function(adelete, call_type="avector_store_delete") + self.vector_store_delete = self.factory_function(delete, call_type="vector_store_delete") def _initialize_vector_store_file_endpoints(self): """Initialize vector store file endpoints.""" @@ -1411,12 +1246,8 @@ class Router: generate_content_stream, ) - self.agenerate_content = self.factory_function( - agenerate_content, call_type="agenerate_content" - ) - self.generate_content = self.factory_function( - generate_content, call_type="generate_content" - ) + self.agenerate_content = self.factory_function(agenerate_content, call_type="agenerate_content") + self.generate_content = self.factory_function(generate_content, call_type="generate_content") self.agenerate_content_stream = self.factory_function( agenerate_content_stream, call_type="agenerate_content_stream" ) @@ -1459,50 +1290,26 @@ class Router: video_status, ) - self.avideo_generation = self.factory_function( - avideo_generation, call_type="avideo_generation" - ) - self.video_generation = self.factory_function( - video_generation, call_type="video_generation" - ) + self.avideo_generation = self.factory_function(avideo_generation, call_type="avideo_generation") + self.video_generation = self.factory_function(video_generation, call_type="video_generation") self.avideo_list = self.factory_function(avideo_list, call_type="avideo_list") self.video_list = self.factory_function(video_list, call_type="video_list") - self.avideo_status = self.factory_function( - avideo_status, call_type="avideo_status" - ) - self.video_status = self.factory_function( - video_status, call_type="video_status" - ) - self.avideo_content = self.factory_function( - avideo_content, call_type="avideo_content" - ) - self.video_content = self.factory_function( - video_content, call_type="video_content" - ) - self.avideo_remix = self.factory_function( - avideo_remix, call_type="avideo_remix" - ) + self.avideo_status = self.factory_function(avideo_status, call_type="avideo_status") + self.video_status = self.factory_function(video_status, call_type="video_status") + self.avideo_content = self.factory_function(avideo_content, call_type="avideo_content") + self.video_content = self.factory_function(video_content, call_type="video_content") + self.avideo_remix = self.factory_function(avideo_remix, call_type="avideo_remix") self.video_remix = self.factory_function(video_remix, call_type="video_remix") self.avideo_create_character = self.factory_function( avideo_create_character, call_type="avideo_create_character" ) - self.video_create_character = self.factory_function( - video_create_character, call_type="video_create_character" - ) - self.avideo_get_character = self.factory_function( - avideo_get_character, call_type="avideo_get_character" - ) - self.video_get_character = self.factory_function( - video_get_character, call_type="video_get_character" - ) + self.video_create_character = self.factory_function(video_create_character, call_type="video_create_character") + self.avideo_get_character = self.factory_function(avideo_get_character, call_type="avideo_get_character") + self.video_get_character = self.factory_function(video_get_character, call_type="video_get_character") self.avideo_edit = self.factory_function(avideo_edit, call_type="avideo_edit") self.video_edit = self.factory_function(video_edit, call_type="video_edit") - self.avideo_extension = self.factory_function( - avideo_extension, call_type="avideo_extension" - ) - self.video_extension = self.factory_function( - video_extension, call_type="video_extension" - ) + self.avideo_extension = self.factory_function(avideo_extension, call_type="avideo_extension") + self.video_extension = self.factory_function(video_extension, call_type="video_extension") def _initialize_container_endpoints(self): """Initialize container endpoints.""" @@ -1520,30 +1327,14 @@ class Router: _generated_endpoints as container_file_endpoints, ) - self.acreate_container = self.factory_function( - acreate_container, call_type="acreate_container" - ) - self.create_container = self.factory_function( - create_container, call_type="create_container" - ) - self.alist_containers = self.factory_function( - alist_containers, call_type="alist_containers" - ) - self.list_containers = self.factory_function( - list_containers, call_type="list_containers" - ) - self.aretrieve_container = self.factory_function( - aretrieve_container, call_type="aretrieve_container" - ) - self.retrieve_container = self.factory_function( - retrieve_container, call_type="retrieve_container" - ) - self.adelete_container = self.factory_function( - adelete_container, call_type="adelete_container" - ) - self.delete_container = self.factory_function( - delete_container, call_type="delete_container" - ) + self.acreate_container = self.factory_function(acreate_container, call_type="acreate_container") + self.create_container = self.factory_function(create_container, call_type="create_container") + self.alist_containers = self.factory_function(alist_containers, call_type="alist_containers") + self.list_containers = self.factory_function(list_containers, call_type="list_containers") + self.aretrieve_container = self.factory_function(aretrieve_container, call_type="aretrieve_container") + self.retrieve_container = self.factory_function(retrieve_container, call_type="retrieve_container") + self.adelete_container = self.factory_function(adelete_container, call_type="adelete_container") + self.delete_container = self.factory_function(delete_container, call_type="delete_container") # Auto-register JSON-generated container file endpoints for name, func in container_file_endpoints.items(): @@ -1551,18 +1342,10 @@ class Router: def _initialize_skills_endpoints(self): """Initialize Anthropic Skills API endpoints.""" - self.acreate_skill = self.factory_function( - litellm.acreate_skill, call_type="acreate_skill" - ) - self.alist_skills = self.factory_function( - litellm.alist_skills, call_type="alist_skills" - ) - self.aget_skill = self.factory_function( - litellm.aget_skill, call_type="aget_skill" - ) - self.adelete_skill = self.factory_function( - litellm.adelete_skill, call_type="adelete_skill" - ) + self.acreate_skill = self.factory_function(litellm.acreate_skill, call_type="acreate_skill") + self.alist_skills = self.factory_function(litellm.alist_skills, call_type="alist_skills") + self.aget_skill = self.factory_function(litellm.aget_skill, call_type="aget_skill") + self.adelete_skill = self.factory_function(litellm.adelete_skill, call_type="adelete_skill") def _initialize_interactions_endpoints(self): """Initialize Google Interactions API endpoints.""" @@ -1575,30 +1358,14 @@ class Router: from litellm.interactions import delete as delete_interaction from litellm.interactions import get as get_interaction - self.acreate_interaction = self.factory_function( - acreate_interaction, call_type="acreate_interaction" - ) - self.create_interaction = self.factory_function( - create_interaction, call_type="create_interaction" - ) - self.aget_interaction = self.factory_function( - aget_interaction, call_type="aget_interaction" - ) - self.get_interaction = self.factory_function( - get_interaction, call_type="get_interaction" - ) - self.adelete_interaction = self.factory_function( - adelete_interaction, call_type="adelete_interaction" - ) - self.delete_interaction = self.factory_function( - delete_interaction, call_type="delete_interaction" - ) - self.acancel_interaction = self.factory_function( - acancel_interaction, call_type="acancel_interaction" - ) - self.cancel_interaction = self.factory_function( - cancel_interaction, call_type="cancel_interaction" - ) + self.acreate_interaction = self.factory_function(acreate_interaction, call_type="acreate_interaction") + self.create_interaction = self.factory_function(create_interaction, call_type="create_interaction") + self.aget_interaction = self.factory_function(aget_interaction, call_type="aget_interaction") + self.get_interaction = self.factory_function(get_interaction, call_type="get_interaction") + self.adelete_interaction = self.factory_function(adelete_interaction, call_type="adelete_interaction") + self.delete_interaction = self.factory_function(delete_interaction, call_type="delete_interaction") + self.acancel_interaction = self.factory_function(acancel_interaction, call_type="acancel_interaction") + self.cancel_interaction = self.factory_function(cancel_interaction, call_type="cancel_interaction") def _initialize_managed_agents_endpoints(self): """Initialize Google Managed Agents API endpoints (v1beta/agents).""" @@ -1613,30 +1380,16 @@ class Router: from litellm.interactions.agents import list as list_agents from litellm.interactions.agents import list_versions as list_agent_versions - self.acreate_agent = self.factory_function( - acreate_agent, call_type="acreate_agent" - ) - self.create_agent = self.factory_function( - create_agent, call_type="create_agent" - ) - self.alist_agents = self.factory_function( - alist_agents, call_type="alist_agents" - ) + self.acreate_agent = self.factory_function(acreate_agent, call_type="acreate_agent") + self.create_agent = self.factory_function(create_agent, call_type="create_agent") + self.alist_agents = self.factory_function(alist_agents, call_type="alist_agents") self.list_agents = self.factory_function(list_agents, call_type="list_agents") self.aget_agent = self.factory_function(aget_agent, call_type="aget_agent") self.get_agent = self.factory_function(get_agent, call_type="get_agent") - self.adelete_agent = self.factory_function( - adelete_agent, call_type="adelete_agent" - ) - self.delete_agent = self.factory_function( - delete_agent, call_type="delete_agent" - ) - self.alist_agent_versions = self.factory_function( - alist_agent_versions, call_type="alist_agent_versions" - ) - self.list_agent_versions = self.factory_function( - list_agent_versions, call_type="list_agent_versions" - ) + self.adelete_agent = self.factory_function(adelete_agent, call_type="adelete_agent") + self.delete_agent = self.factory_function(delete_agent, call_type="delete_agent") + self.alist_agent_versions = self.factory_function(alist_agent_versions, call_type="alist_agent_versions") + self.list_agent_versions = self.factory_function(list_agent_versions, call_type="list_agent_versions") def _initialize_specialized_endpoints(self): """Helper to initialize specialized router endpoints (vector store, OCR, search, video, container, skills, interactions).""" @@ -1670,9 +1423,7 @@ class Router: f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." ) - def _add_encrypted_content_affinity_check( - self, enable_global_affinity: bool - ) -> None: + def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1685,20 +1436,13 @@ class Router: return callback_list.remove(callback_to_move) insert_index = next( - ( - idx - for idx, callback in enumerate(callback_list) - if isinstance(callback, DeploymentAffinityCheck) - ), + (idx for idx, callback in enumerate(callback_list) if isinstance(callback, DeploymentAffinityCheck)), len(callback_list), ) callback_list.insert(insert_index, callback_to_move) - if ( - enable_global_affinity - or EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - self.model_group_affinity_config - ) + if enable_global_affinity or EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + self.model_group_affinity_config ): if self.optional_callbacks is None: self.optional_callbacks = [] @@ -1712,12 +1456,9 @@ class Router: if existing_ec_callback is not None: existing_ec_callback.router = self existing_ec_callback.enable_global_affinity = ( - existing_ec_callback.enable_global_affinity - or enable_global_affinity - ) - existing_ec_callback.model_group_affinity_config = ( - self.model_group_affinity_config or {} + existing_ec_callback.enable_global_affinity or enable_global_affinity ) + existing_ec_callback.model_group_affinity_config = self.model_group_affinity_config or {} ec_callback = existing_ec_callback else: ec_callback = EncryptedContentAffinityCheck( @@ -1731,9 +1472,7 @@ class Router: _move_before_deployment_affinity(self.optional_callbacks, ec_callback) _move_before_deployment_affinity(litellm.callbacks, ec_callback) - def add_optional_pre_call_checks( - self, optional_pre_call_checks: Optional[OptionalPreCallChecks] - ): + def add_optional_pre_call_checks(self, optional_pre_call_checks: Optional[OptionalPreCallChecks]): if optional_pre_call_checks is None: return @@ -1741,15 +1480,9 @@ class Router: # Unified deployment affinity (session stickiness) # --------------------------------------------------------------------- enable_user_key_affinity = "deployment_affinity" in optional_pre_call_checks - enable_responses_api_affinity = ( - "responses_api_deployment_check" in optional_pre_call_checks - ) + enable_responses_api_affinity = "responses_api_deployment_check" in optional_pre_call_checks enable_session_id_affinity = "session_affinity" in optional_pre_call_checks - if ( - enable_user_key_affinity - or enable_responses_api_affinity - or enable_session_id_affinity - ): + if enable_user_key_affinity or enable_responses_api_affinity or enable_session_id_affinity: if self.optional_callbacks is None: self.optional_callbacks = [] @@ -1761,24 +1494,17 @@ class Router: if existing_affinity_callback is not None: existing_affinity_callback.enable_user_key_affinity = ( - existing_affinity_callback.enable_user_key_affinity - or enable_user_key_affinity + existing_affinity_callback.enable_user_key_affinity or enable_user_key_affinity ) existing_affinity_callback.enable_responses_api_affinity = ( - existing_affinity_callback.enable_responses_api_affinity - or enable_responses_api_affinity + existing_affinity_callback.enable_responses_api_affinity or enable_responses_api_affinity ) existing_affinity_callback.enable_session_id_affinity = ( - existing_affinity_callback.enable_session_id_affinity - or enable_session_id_affinity - ) - existing_affinity_callback.ttl_seconds = ( - self.deployment_affinity_ttl_seconds + existing_affinity_callback.enable_session_id_affinity or enable_session_id_affinity ) + existing_affinity_callback.ttl_seconds = self.deployment_affinity_ttl_seconds if self.model_group_affinity_config: - existing_affinity_callback.model_group_affinity_config = ( - self.model_group_affinity_config - ) + existing_affinity_callback.model_group_affinity_config = self.model_group_affinity_config else: affinity_callback = DeploymentAffinityCheck( cache=self.cache, @@ -1795,9 +1521,7 @@ class Router: # Encrypted content affinity # --------------------------------------------------------------------- self._add_encrypted_content_affinity_check( - enable_global_affinity=( - "encrypted_content_affinity" in optional_pre_call_checks - ) + enable_global_affinity=("encrypted_content_affinity" in optional_pre_call_checks) ) # --------------------------------------------------------------------- @@ -1852,9 +1576,7 @@ class Router: return _deployment_copy except Exception as e: - verbose_router_logger.debug( - f"Error occurred while printing deployment - {str(e)}" - ) + verbose_router_logger.debug(f"Error occurred while printing deployment - {str(e)}") raise e ### COMPLETION, EMBEDDING, IMG GENERATION FUNCTIONS @@ -1917,9 +1639,7 @@ class Router: self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) kwargs.pop("silent_model", None) # Ensure it's not in kwargs either model_name = litellm_params["model"] - potential_model_client = self._get_client( - deployment=deployment, kwargs=kwargs - ) + potential_model_client = self._get_client(deployment=deployment, kwargs=kwargs) # check if provided keys == client keys # dynamic_api_key = kwargs.get("api_key", None) if ( @@ -1944,15 +1664,11 @@ class Router: **kwargs, } response = litellm.completion(**input_kwargs) - verbose_router_logger.info( - f"litellm.completion(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.completion(model={model_name})\033[32m 200 OK\033[0m") ## CHECK CONTENT FILTER ERROR ## if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error( - model=model, response=response, kwargs=kwargs - ) + _should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs) if _should_raise: raise litellm.ContentPolicyViolationError( message="Response output was blocked.", @@ -1971,9 +1687,7 @@ class Router: return response except Exception as e: - verbose_router_logger.info( - f"litellm.completion(model={model_name})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.completion(model={model_name})\033[31m Exception {str(e)}\033[0m") # Set per-deployment num_retries on exception for retry logic if deployment is not None: self._set_deployment_num_retries_on_exception(e, deployment) @@ -2000,10 +1714,7 @@ class Router: # model_group / is_silent_experiment on the silent dict doesn't corrupt # the primary call's metadata. original_metadata = kwargs.get("metadata") - if ( - original_metadata is not None - and silent_kwargs.get("metadata") is original_metadata - ): + if original_metadata is not None and silent_kwargs.get("metadata") is original_metadata: silent_kwargs["metadata"] = dict(original_metadata) if "metadata" not in silent_kwargs: @@ -2028,9 +1739,7 @@ class Router: return silent_kwargs - def _silent_experiment_completion( - self, silent_model: str, messages: List[Any], **kwargs - ): + def _silent_experiment_completion(self, silent_model: str, messages: List[Any], **kwargs): """ Run a silent experiment in the background (thread). """ @@ -2041,9 +1750,7 @@ class Router: messages = copy.deepcopy(messages) - verbose_router_logger.info( - f"Starting silent experiment for model {silent_model}" - ) + verbose_router_logger.info(f"Starting silent experiment for model {silent_model}") silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) @@ -2075,9 +1782,7 @@ class Router: finally: loop.close() except Exception as e: - verbose_router_logger.error( - f"Silent experiment failed for model {silent_model}: {str(e)}" - ) + verbose_router_logger.error(f"Silent experiment failed for model {silent_model}: {str(e)}") # fmt: off @@ -2171,9 +1876,7 @@ class Router: and complete_response_object_usage.usage is not None # type: ignore ): usage_objects.append(complete_response_object_usage) - combined_usage = BaseTokenUsageProcessor.combine_usage_objects( - usage_objects=usage_objects - ) + combined_usage = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) setattr(fallback_item, "usage", combined_usage) @staticmethod @@ -2247,9 +1950,7 @@ class Router: except MidStreamFallbackError as e: from litellm.main import stream_chunk_builder - complete_response_object = stream_chunk_builder( - chunks=model_response.chunks - ) + complete_response_object = stream_chunk_builder(chunks=model_response.chunks) complete_response_object_usage = cast( Optional[Usage], getattr(complete_response_object, "usage", None), @@ -2257,9 +1958,7 @@ class Router: try: # Use the router's fallback system model_group = cast(str, initial_kwargs.get("model")) - fallbacks: Optional[List] = initial_kwargs.get( - "fallbacks", self.fallbacks - ) + fallbacks: Optional[List] = initial_kwargs.get("fallbacks", self.fallbacks) context_window_fallbacks: Optional[List] = initial_kwargs.get( "context_window_fallbacks", self.context_window_fallbacks ) @@ -2285,43 +1984,30 @@ class Router: "prefix": True, }, ] - self._update_kwargs_before_fallbacks( - model=model_group, kwargs=initial_kwargs - ) - fallback_response = ( - await self.async_function_with_fallbacks_common_utils( - e=e, - disable_fallbacks=False, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get( - "include_fallback_errors", False - ) - is True, - ) + self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) + fallback_response = await self.async_function_with_fallbacks_common_utils( + e=e, + disable_fallbacks=False, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + model_group=model_group, + args=(), + kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, ) # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): - prepared_fallback_hidden_params = ( - Router._prepare_fallback_hidden_params(fallback_response) - ) + prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) async for fallback_item in fallback_response: # type: ignore - Router._apply_fallback_hidden_params_to_item( - fallback_item, prepared_fallback_hidden_params - ) + Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - self._combine_fallback_usage( - fallback_item, complete_response_object_usage - ) + self._combine_fallback_usage(fallback_item, complete_response_object_usage) yield fallback_item else: # If fallback returns a non-streaming response, yield None @@ -2329,9 +2015,7 @@ class Router: except Exception as fallback_error: # If fallback also fails, log and re-raise original error - verbose_router_logger.error( - f"Fallback also failed: {fallback_error}" - ) + verbose_router_logger.error(f"Fallback also failed: {fallback_error}") # No fallback handled the mid-stream error, so surface the # real provider exception (e.g. RateLimitError) instead of # leaking the internal MidStreamFallbackError to the client @@ -2355,9 +2039,7 @@ class Router: "stream_with_fallbacks: error closing model_response: %s", e, ) - if fallback_response is not None and hasattr( - fallback_response, "aclose" - ): + if fallback_response is not None and hasattr(fallback_response, "aclose"): try: await fallback_response.aclose() except BaseException as e: @@ -2421,10 +2103,7 @@ class Router: # real Usage instances always have them. prompt = int(getattr(chat, "prompt_tokens", 0) or 0) completion = int(getattr(chat, "completion_tokens", 0) or 0) - total = int( - getattr(chat, "total_tokens", prompt + completion) - or (prompt + completion) - ) + total = int(getattr(chat, "total_tokens", prompt + completion) or (prompt + completion)) return ResponseAPIUsage( input_tokens=prompt, output_tokens=completion, @@ -2636,9 +2315,7 @@ class Router: getattr(source_iterator, "litellm_logging_obj", None), ) self.finished = False - self.responses_api_provider_config = getattr( - source_iterator, "responses_api_provider_config", None - ) + self.responses_api_provider_config = getattr(source_iterator, "responses_api_provider_config", None) self.completed_response = None self.start_time = getattr(source_iterator, "start_time", datetime.now()) self._failure_handled = False @@ -2647,19 +2324,13 @@ class Router: self._completed_response_cache_hit = None self._persist_completed_response_before_logging = True self._stream_created_time = time.time() - self.litellm_metadata = getattr( - source_iterator, "litellm_metadata", None - ) - self.custom_llm_provider = getattr( - source_iterator, "custom_llm_provider", None - ) + self.litellm_metadata = getattr(source_iterator, "litellm_metadata", None) + self.custom_llm_provider = getattr(source_iterator, "custom_llm_provider", None) self.request_data = getattr(source_iterator, "request_data", {}) or {} self.call_type = getattr(source_iterator, "call_type", None) # Preserve hidden params so response headers (model_id, # api_base, additional_headers) keep flowing. - self._hidden_params = dict( - getattr(source_iterator, "_hidden_params", None) or {} - ) + self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) def __aiter__(self): return self @@ -2676,10 +2347,7 @@ class Router: # records nothing on streaming /v1/responses calls — every # follow-up /v1/containers//files call then 403s for # the very key that created the container (#30210). - if ( - self.completed_response is None - and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES - ): + if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: self.completed_response = chunk return chunk @@ -2696,9 +2364,7 @@ class Router: partial_usage = Router._extract_partial_responses_usage(source_iterator) try: model_group = cast(str, initial_kwargs.get("model")) - fallbacks: Optional[List] = initial_kwargs.get( - "fallbacks", self.fallbacks - ) + fallbacks: Optional[List] = initial_kwargs.get("fallbacks", self.fallbacks) context_window_fallbacks: Optional[List] = initial_kwargs.get( "context_window_fallbacks", self.context_window_fallbacks ) @@ -2710,20 +2376,16 @@ class Router: # _ageneric_api_call_with_fallbacks_helper. # original_generic_function is preserved by the caller so # the helper knows what underlying API to invoke per attempt. - initial_kwargs["original_function"] = ( - self._ageneric_api_call_with_fallbacks_helper - ) + initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper if e.is_pre_first_chunk or not e.generated_content: # No content generated before the error — retry with the # original input. Adding a continuation prompt would # waste tokens and confuse the model. pass else: - initial_kwargs["input"] = ( - Router._build_responses_continuation_input( - initial_kwargs.get("input"), - e.generated_content, - ) + initial_kwargs["input"] = Router._build_responses_continuation_input( + initial_kwargs.get("input"), + e.generated_content, ) # The Responses-API path stores observability metadata # under "litellm_metadata" (not the default "metadata") — @@ -2735,42 +2397,29 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) - fallback_response = ( - await self.async_function_with_fallbacks_common_utils( - e=e, - disable_fallbacks=False, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get( - "include_fallback_errors", False - ) - is True, - ) + fallback_response = await self.async_function_with_fallbacks_common_utils( + e=e, + disable_fallbacks=False, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + model_group=model_group, + args=(), + kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, ) if hasattr(fallback_response, "__aiter__"): - prepared_fallback_hidden_params = ( - Router._prepare_fallback_hidden_params(fallback_response) - ) + prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) async for fallback_item in fallback_response: # type: ignore - Router._apply_fallback_hidden_params_to_item( - fallback_item, prepared_fallback_hidden_params - ) + Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if partial_usage is not None: - Router._combine_responses_fallback_usage( - fallback_item, partial_usage - ) + Router._combine_responses_fallback_usage(fallback_item, partial_usage) yield fallback_item else: yield fallback_response except Exception as fallback_error: - verbose_router_logger.error( - f"Responses streaming fallback also failed: {fallback_error}" - ) + verbose_router_logger.error(f"Responses streaming fallback also failed: {fallback_error}") if ( isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None @@ -2787,9 +2436,7 @@ class Router: "stream_with_fallbacks(aresponses): error closing source: %s", exc, ) - if fallback_response is not None and hasattr( - fallback_response, "aclose" - ): + if fallback_response is not None and hasattr(fallback_response, "aclose"): try: await fallback_response.aclose() except BaseException as exc: @@ -2843,18 +2490,14 @@ class Router: except MidStreamFallbackError as e: from litellm.main import stream_chunk_builder - complete_response_object = stream_chunk_builder( - chunks=model_response.chunks - ) + complete_response_object = stream_chunk_builder(chunks=model_response.chunks) complete_response_object_usage = cast( Optional[Usage], getattr(complete_response_object, "usage", None), ) try: model_group = cast(str, initial_kwargs.get("model")) - fallbacks: Optional[List] = initial_kwargs.get( - "fallbacks", router_self.fallbacks - ) + fallbacks: Optional[List] = initial_kwargs.get("fallbacks", router_self.fallbacks) context_window_fallbacks: Optional[List] = initial_kwargs.get( "context_window_fallbacks", router_self.context_window_fallbacks, @@ -2878,9 +2521,7 @@ class Router: "prefix": True, }, ] - router_self._update_kwargs_before_fallbacks( - model=model_group, kwargs=initial_kwargs - ) + router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = router_self.function_with_fallbacks( **initial_kwargs, fallbacks=fallbacks, @@ -2889,29 +2530,21 @@ class Router: ) if hasattr(fallback_response, "__iter__"): - prepared_fallback_hidden_params = ( - Router._prepare_fallback_hidden_params(fallback_response) - ) + prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response) for fallback_item in fallback_response: - Router._apply_fallback_hidden_params_to_item( - fallback_item, prepared_fallback_hidden_params - ) + Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) if ( fallback_item and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - router_self._combine_fallback_usage( - fallback_item, complete_response_object_usage - ) + router_self._combine_fallback_usage(fallback_item, complete_response_object_usage) yield fallback_item else: yield None except Exception as fallback_error: - verbose_router_logger.error( - f"Fallback also failed: {fallback_error}" - ) + verbose_router_logger.error(f"Fallback also failed: {fallback_error}") if ( isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None @@ -2927,9 +2560,7 @@ class Router: "stream_with_fallbacks: error closing model_response: %s", close_err, ) - if fallback_response is not None and hasattr( - fallback_response, "close" - ): + if fallback_response is not None and hasattr(fallback_response, "close"): try: fallback_response.close() except BaseException as close_err: @@ -2940,9 +2571,7 @@ class Router: return SyncFallbackStreamWrapper(stream_with_fallbacks()) - async def _silent_experiment_acompletion( - self, silent_model: str, messages: List[Any], **kwargs - ): + async def _silent_experiment_acompletion(self, silent_model: str, messages: List[Any], **kwargs): """ Run a silent experiment in the background. """ @@ -2953,9 +2582,7 @@ class Router: messages = copy.deepcopy(messages) - verbose_router_logger.info( - f"Starting silent experiment for model {silent_model}" - ) + verbose_router_logger.info(f"Starting silent experiment for model {silent_model}") silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) # Override model_group to correctly attribute metrics to the silent model @@ -2968,9 +2595,7 @@ class Router: **silent_kwargs, ) except Exception as e: - verbose_router_logger.error( - f"Silent experiment failed for model {silent_model}: {str(e)}" - ) + verbose_router_logger.error(f"Silent experiment failed for model {silent_model}: {str(e)}") async def _acompletion( self, model: str, messages: List[Dict[str, str]], **kwargs @@ -3016,9 +2641,7 @@ class Router: # debug how often this deployment picked - self._track_deployment_metrics( - deployment=deployment, parent_otel_span=parent_otel_span - ) + self._track_deployment_metrics(deployment=deployment, parent_otel_span=parent_otel_span) # Check for silent model experiment # Make a local copy of litellm_params to avoid mutating the Router's state @@ -3059,18 +2682,14 @@ class Router: _response = litellm.acompletion(**input_kwargs) - logging_obj: Optional[LiteLLMLogging] = kwargs.get( - "litellm_logging_obj", None - ) + logging_obj: Optional[LiteLLMLogging] = kwargs.get("litellm_logging_obj", None) rpm_semaphore = self._get_client( deployment=deployment, kwargs=kwargs, client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -3093,9 +2712,7 @@ class Router: ## CHECK CONTENT FILTER ERROR ## if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error( - model=model, response=response, kwargs=kwargs - ) + _should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs) if _should_raise: raise litellm.ContentPolicyViolationError( message="Response output was blocked.", @@ -3104,9 +2721,7 @@ class Router: ) self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.acompletion(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.acompletion(model={model_name})\033[32m 200 OK\033[0m") # debug how often this deployment picked self._track_deployment_metrics( deployment=deployment, @@ -3123,12 +2738,10 @@ class Router: return response except litellm.Timeout as e: - deployment_request_timeout_param = _timeout_debug_deployment_dict.get( - "litellm_params", {} - ).get("request_timeout", None) - deployment_timeout_param = _timeout_debug_deployment_dict.get( - "litellm_params", {} - ).get("timeout", None) + deployment_request_timeout_param = _timeout_debug_deployment_dict.get("litellm_params", {}).get( + "request_timeout", None + ) + deployment_timeout_param = _timeout_debug_deployment_dict.get("litellm_params", {}).get("timeout", None) if litellm.expose_router_debug_in_errors: e.message += f"\n\nDeployment Info: request_timeout: {deployment_request_timeout_param}\ntimeout: {deployment_timeout_param}" # Set per-deployment num_retries on exception for retry logic @@ -3137,9 +2750,7 @@ class Router: self._set_failed_deployment_id_on_exception(e, deployment) raise e except Exception as e: - verbose_router_logger.info( - f"litellm.acompletion(model={model_name})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.acompletion(model={model_name})\033[31m Exception {str(e)}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 # Set per-deployment num_retries on exception for retry logic @@ -3169,9 +2780,7 @@ class Router: if _req_num_retries is not None: kwargs["num_retries"] = _req_num_retries else: - kwargs["num_retries"] = ( - self.num_retries if self.num_retries is not None else 0 - ) + kwargs["num_retries"] = self.num_retries if self.num_retries is not None else 0 kwargs.setdefault("litellm_trace_id", str(uuid.uuid4())) model_group_alias: Optional[str] = None if self._get_model_from_alias(model=model): @@ -3180,9 +2789,7 @@ class Router: {"model_group": model, "model_group_alias": model_group_alias} ) - def _set_deployment_num_retries_on_exception( - self, exception: Exception, deployment: dict - ) -> None: + def _set_deployment_num_retries_on_exception(self, exception: Exception, deployment: dict) -> None: """ Set num_retries from deployment litellm_params on the exception. @@ -3201,9 +2808,7 @@ class Router: except (ValueError, TypeError): pass # Skip if value can't be converted to int - def _set_failed_deployment_id_on_exception( - self, exception: Exception, deployment: dict - ) -> None: + def _set_failed_deployment_id_on_exception(self, exception: Exception, deployment: dict) -> None: """ Stamp the failed deployment's `model_info.id` on the exception so the fallback layer can exclude it from subsequent re-picks within the same @@ -3251,17 +2856,13 @@ class Router: """ model_info = deployment.get("model_info", {}).copy() litellm_params = deployment["litellm_params"].copy() - dynamic_litellm_params = get_dynamic_litellm_params( - litellm_params=litellm_params, request_kwargs=kwargs - ) + dynamic_litellm_params = get_dynamic_litellm_params(litellm_params=litellm_params, request_kwargs=kwargs) # Use deployment model_name as model_group for generating model_id metadata_variable_name = _get_router_metadata_variable_name( function_name=function_name, ) model_group = kwargs.get(metadata_variable_name, {}).get("model_group") - _model_id = self._generate_model_id( - model_group=model_group, litellm_params=dynamic_litellm_params - ) + _model_id = self._generate_model_id(model_group=model_group, litellm_params=dynamic_litellm_params) original_model_id = model_info.get("id") model_info["id"] = _model_id model_info["original_model_id"] = original_model_id @@ -3270,9 +2871,7 @@ class Router: litellm_params=LiteLLM_Params(**dynamic_litellm_params), model_info=model_info, ) - self.upsert_deployment( - deployment=deployment_pydantic_obj - ) # add new deployment to router + self.upsert_deployment(deployment=deployment_pydantic_obj) # add new deployment to router return deployment_pydantic_obj @staticmethod @@ -3345,9 +2944,7 @@ class Router: kwargs[metadata_variable_name]["tags"] = merged_tags ## CREDENTIAL NAME AS TAG - credential_name = deployment.get("litellm_params", {}).get( - "litellm_credential_name" - ) + credential_name = deployment.get("litellm_params", {}).get("litellm_credential_name") if credential_name: credential_tag = f"Credential: {credential_name}" existing_tags = kwargs[metadata_variable_name].get("tags") or [] @@ -3363,9 +2960,7 @@ class Router: ) _router_timeout = ( - float(self._explicit_timeout) - if isinstance(self._explicit_timeout, (int, float)) - else None + float(self._explicit_timeout) if isinstance(self._explicit_timeout, (int, float)) else None ) kwargs["timeout"] = resolve_llm_passthrough_timeout( kwargs=kwargs, @@ -3373,13 +2968,9 @@ class Router: router_timeout=_router_timeout, ) else: - kwargs["timeout"] = self._get_timeout( - kwargs=kwargs, data=deployment["litellm_params"] - ) + kwargs["timeout"] = self._get_timeout(kwargs=kwargs, data=deployment["litellm_params"]) - self._update_kwargs_with_default_litellm_params( - kwargs=kwargs, metadata_variable_name=metadata_variable_name - ) + self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) def _get_async_openai_model_client(self, deployment: dict, kwargs: dict): """ @@ -3390,9 +2981,7 @@ class Router: If dynamic api key is provided: Do not re-use the client. Pass model_client=None. The OpenAI/ AzureOpenAI client will be recreated in the handler for the llm provider """ - potential_model_client = self._get_client( - deployment=deployment, kwargs=kwargs, client_type="async" - ) + potential_model_client = self._get_client(deployment=deployment, kwargs=kwargs, client_type="async") # check if provided keys == client keys # dynamic_api_key = kwargs.get("api_key", None) @@ -3407,33 +2996,23 @@ class Router: return model_client - def _get_stream_timeout( - self, kwargs: dict, data: dict - ) -> Optional[Union[float, int]]: + def _get_stream_timeout(self, kwargs: dict, data: dict) -> Optional[Union[float, int]]: """Helper to get stream timeout from kwargs or deployment params""" return ( kwargs.get("stream_timeout", None) # the params dynamically set by user - or data.get( - "stream_timeout", None - ) # timeout set on litellm_params for this deployment + or data.get("stream_timeout", None) # timeout set on litellm_params for this deployment or self.stream_timeout # timeout set on router or self.request_timeout # litellm_settings.request_timeout (per-attempt) or self.default_litellm_params.get("stream_timeout", None) ) - def _get_non_stream_timeout( - self, kwargs: dict, data: dict - ) -> Optional[Union[float, int]]: + def _get_non_stream_timeout(self, kwargs: dict, data: dict) -> Optional[Union[float, int]]: """Helper to get non-stream timeout from kwargs or deployment params""" timeout = ( kwargs.get("timeout", None) # the params dynamically set by user or kwargs.get("request_timeout", None) # the params dynamically set by user - or data.get( - "timeout", None - ) # timeout set on litellm_params for this deployment - or data.get( - "request_timeout", None - ) # timeout set on litellm_params for this deployment + or data.get("timeout", None) # timeout set on litellm_params for this deployment + or data.get("request_timeout", None) # timeout set on litellm_params for this deployment or self.request_timeout # litellm_settings.request_timeout (per-attempt) or self.timeout # timeout set on router (router_settings.timeout) or self.default_litellm_params.get("timeout", None) @@ -3487,9 +3066,7 @@ class Router: """ ############## Helpers for async completion ################## - async def _async_completion_no_exceptions( - model: str, messages: List[AllMessageValues], **kwargs - ): + async def _async_completion_no_exceptions(model: str, messages: List[AllMessageValues], **kwargs): """ Wrapper around self.async_completion that catches exceptions and returns them as a result """ @@ -3521,11 +3098,7 @@ class Router: _tasks = [] for model in models: # add each task but if the task fails - _tasks.append( - _async_completion_no_exceptions( - model=model, messages=messages, **kwargs - ) - ) # type: ignore + _tasks.append(_async_completion_no_exceptions(model=model, messages=messages, **kwargs)) # type: ignore response = await asyncio.gather(*_tasks) return response elif isinstance(messages, list) and all(isinstance(m, list) for m in messages): @@ -3572,9 +3145,7 @@ class Router: ) """ - async def _async_completion_no_exceptions( - model: str, messages: List[AllMessageValues], **kwargs - ): + async def _async_completion_no_exceptions(model: str, messages: List[AllMessageValues], **kwargs): """ Wrapper around self.async_completion that catches exceptions and returns them as a result """ @@ -3586,11 +3157,7 @@ class Router: _tasks = [] for message_request in messages: # add each task but if the task fails - _tasks.append( - _async_completion_no_exceptions( - model=model, messages=message_request, **kwargs - ) - ) + _tasks.append(_async_completion_no_exceptions(model=model, messages=message_request, **kwargs)) response = await asyncio.gather(*_tasks) return response @@ -3634,14 +3201,10 @@ class Router: Wrapper around self.acompletion that catches exceptions and returns them as a result """ try: - result = await self.acompletion( - model=model, messages=messages, stream=stream, **kwargs - ) # type: ignore + result = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs) # type: ignore return result except asyncio.CancelledError: - verbose_router_logger.debug( - "Received 'task.cancel'. Cancelling call w/ model={}.".format(model) - ) + verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model={}.".format(model)) raise except Exception as e: return e @@ -3653,9 +3216,7 @@ class Router: try: result = await task if isinstance(result, (ModelResponse, CustomStreamWrapper)): - verbose_router_logger.debug( - "Received successful response. Cancelling other LLM API calls." - ) + verbose_router_logger.debug("Received successful response. Cancelling other LLM API calls.") # If a desired response is received, cancel all other pending tasks for t in pending_tasks: t.cancel() @@ -3672,9 +3233,7 @@ class Router: for model in models: task = asyncio.create_task( - _async_completion_no_exceptions( - model=model, messages=messages, stream=stream, **kwargs - ) + _async_completion_no_exceptions(model=model, messages=messages, stream=stream, **kwargs) ) pending_tasks.append(task) @@ -3756,22 +3315,16 @@ class Router: if make_request: try: - _response = await self.acompletion( - model=model, messages=messages, stream=stream, **kwargs - ) + _response = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs) _response._hidden_params.setdefault("additional_headers", {}) - _response._hidden_params["additional_headers"].update( - {"x-litellm-request-prioritization-used": True} - ) + _response._hidden_params["additional_headers"].update({"x-litellm-request-prioritization-used": True}) return _response except Exception as e: setattr(e, "priority", priority) raise e else: # Clean up the request from the scheduler queue also before raising the timeout exception - await self.scheduler.remove_request( - request_id=item.request_id, model_name=item.model_name - ) + await self.scheduler.remove_request(request_id=item.request_id, model_name=item.model_name) raise litellm.Timeout( message="Request timed out while polling queue", model=model, @@ -3834,9 +3387,7 @@ class Router: raise e else: # Clean up the request from the scheduler queue also before raising the timeout exception - await self.scheduler.remove_request( - request_id=item.request_id, model_name=item.model_name - ) + await self.scheduler.remove_request(request_id=item.request_id, model_name=item.model_name) raise litellm.Timeout( message="Request timed out while polling queue", model=model, @@ -3878,36 +3429,24 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), ) - self._update_kwargs_with_deployment( - deployment=prompt_management_deployment, kwargs=kwargs - ) + self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) data = prompt_management_deployment["litellm_params"].copy() litellm_model = data.get("model", None) # litellm_agent/ prefix only strips the model name, no prompt_id needed - is_litellm_agent_model = isinstance( - litellm_model, str - ) and litellm_model.startswith("litellm_agent/") + is_litellm_agent_model = isinstance(litellm_model, str) and litellm_model.startswith("litellm_agent/") - prompt_id = kwargs.get("prompt_id") or prompt_management_deployment[ - "litellm_params" - ].get("prompt_id", None) - prompt_variables = kwargs.get( - "prompt_variables" - ) or prompt_management_deployment["litellm_params"].get( + prompt_id = kwargs.get("prompt_id") or prompt_management_deployment["litellm_params"].get("prompt_id", None) + prompt_variables = kwargs.get("prompt_variables") or prompt_management_deployment["litellm_params"].get( "prompt_variables", None ) - prompt_label = kwargs.get("prompt_label", None) or prompt_management_deployment[ - "litellm_params" - ].get("prompt_label", None) + prompt_label = kwargs.get("prompt_label", None) or prompt_management_deployment["litellm_params"].get( + "prompt_label", None + ) - if not is_litellm_agent_model and ( - prompt_id is None or not isinstance(prompt_id, str) - ): - raise ValueError( - f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}" - ) + if not is_litellm_agent_model and (prompt_id is None or not isinstance(prompt_id, str)): + raise ValueError(f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}") if prompt_variables is not None and not isinstance(prompt_variables, dict): raise ValueError( f"Prompt variables is set but not a dictionary. Got={prompt_variables}, type={type(prompt_variables)}" @@ -3935,9 +3474,7 @@ class Router: "prompt_label", "prompt_version", } - filtered_data = { - k: v for k, v in data.items() if k not in prompt_management_params - } + filtered_data = {k: v for k, v in data.items() if k not in prompt_management_params} kwargs = {**filtered_data, **kwargs, **optional_params} kwargs["model"] = model @@ -3970,9 +3507,7 @@ class Router: def _image_generation(self, prompt: str, model: str, **kwargs): model_name = "" try: - verbose_router_logger.debug( - f"Inside _image_generation()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _image_generation()- model: {model}; kwargs: {kwargs}") deployment = self.get_available_deployment( model=model, messages=[{"role": "user", "content": "prompt"}], @@ -4001,9 +3536,7 @@ class Router: } ) self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.image_generation(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.image_generation(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: verbose_router_logger.info( @@ -4037,9 +3570,7 @@ class Router: async def _aimage_generation(self, prompt: str, model: str, **kwargs): model_name = model try: - verbose_router_logger.debug( - f"Inside _image_generation()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _image_generation()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -4075,9 +3606,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -4094,9 +3623,7 @@ class Router: response = await response self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.aimage_generation(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.aimage_generation(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: verbose_router_logger.info( @@ -4151,9 +3678,7 @@ class Router: async def _atranscription(self, file: FileTypes, model: str, **kwargs): model_name = model try: - verbose_router_logger.debug( - f"Inside _atranscription()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _atranscription()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -4187,9 +3712,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -4206,14 +3729,10 @@ class Router: response = await response self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.atranscription(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.atranscription(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"litellm.atranscription(model={model_name})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.atranscription(model={model_name})\033[31m Exception {str(e)}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -4273,9 +3792,7 @@ class Router: async def _aspeech(self, model: str, input: str, voice: str, **kwargs): model_name = model try: - verbose_router_logger.debug( - f"Inside _aspeech()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _aspeech()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -4309,9 +3826,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -4328,14 +3843,10 @@ class Router: response = await response self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.aspeech(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.aspeech(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"litellm.aspeech(model={model_name})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.aspeech(model={model_name})\033[31m Exception {str(e)}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -4364,9 +3875,7 @@ class Router: async def _arerank(self, model: str, **kwargs): model_name = None try: - verbose_router_logger.debug( - f"Inside _rerank()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _rerank()- model: {model}; kwargs: {kwargs}") deployment = await self.async_get_available_deployment( model=model, specific_deployment=kwargs.pop("specific_deployment", None), @@ -4392,14 +3901,10 @@ class Router: ) self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.arerank(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.arerank(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"litellm.arerank(model={model_name})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.arerank(model={model_name})\033[31m Exception {str(e)}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -4429,17 +3934,13 @@ class Router: data = deployment["litellm_params"].copy() for k, v in self.default_litellm_params.items(): - if ( - k not in kwargs - ): # prioritize model-specific params > default router params + if k not in kwargs: # prioritize model-specific params > default router params kwargs[k] = v elif k == "metadata": kwargs[k].update(v) # call via litellm.completion() - return litellm.text_completion( - **{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs} - ) # type: ignore + return litellm.text_completion(**{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs}) # type: ignore except Exception as e: raise e @@ -4482,9 +3983,7 @@ class Router: async def _atext_completion(self, model: str, prompt: str, **kwargs): try: - verbose_router_logger.debug( - f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -4519,9 +4018,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -4538,14 +4035,10 @@ class Router: response = await response self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.atext_completion(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.atext_completion(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"litellm.atext_completion(model={model})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.atext_completion(model={model})\033[31m Exception {str(e)}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -4581,9 +4074,7 @@ class Router: async def _aadapter_completion(self, adapter_id: str, model: str, **kwargs): try: - verbose_router_logger.debug( - f"Inside _aadapter_completion()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _aadapter_completion()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -4618,9 +4109,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -4637,14 +4126,10 @@ class Router: response = await response # type: ignore self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.aadapter_completion(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.aadapter_completion(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"litellm.aadapter_completion(model={model})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.aadapter_completion(model={model})\033[31m Exception {str(e)}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -4662,9 +4147,7 @@ class Router: **kwargs, ) - async def _asearch_with_fallbacks_helper( - self, model: str, original_generic_function: Callable, **kwargs - ): + async def _asearch_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): """ Helper function for search API calls - selects a search tool and calls the original function. Called by async_function_with_fallbacks for each retry attempt. @@ -4703,9 +4186,7 @@ class Router: kwargs=kwargs, metadata_variable_name="litellm_metadata", ) - verbose_router_logger.debug( - f"Inside aguardrail() - guardrail_name: {guardrail_name}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside aguardrail() - guardrail_name: {guardrail_name}; kwargs: {kwargs}") response = await self.async_function_with_fallbacks(**kwargs) return response @@ -4753,9 +4234,7 @@ class Router: """ from litellm.router_strategy.simple_shuffle import simple_shuffle - healthy_deployments = [ - g for g in self.guardrail_list if g.get("guardrail_name") == guardrail_name - ] + healthy_deployments = [g for g in self.guardrail_list if g.get("guardrail_name") == guardrail_name] if not healthy_deployments: raise ValueError(f"No guardrail found with name: {guardrail_name}") @@ -4773,9 +4252,7 @@ class Router: ), ) - async def _ageneric_api_call_with_fallbacks( - self, model: str, original_function: Callable, **kwargs - ): + async def _ageneric_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): """ Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router """ @@ -4783,12 +4260,8 @@ class Router: kwargs["model"] = model kwargs["original_generic_function"] = original_function kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper - self._update_kwargs_before_fallbacks( - model=model, kwargs=kwargs, metadata_variable_name="litellm_metadata" - ) - verbose_router_logger.debug( - f"Inside ageneric_api_call_with_fallbacks() - model: {model}; kwargs: {kwargs}" - ) + self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs, metadata_variable_name="litellm_metadata") + verbose_router_logger.debug(f"Inside ageneric_api_call_with_fallbacks() - model: {model}; kwargs: {kwargs}") response = await self.async_function_with_fallbacks(**kwargs) return response @@ -4830,14 +4303,10 @@ class Router: # If get_llm_provider fails, fall back to using model_name as-is replacement_model_name = model_name - kwargs["endpoint"] = kwargs["endpoint"].replace( - model, replacement_model_name - ) + kwargs["endpoint"] = kwargs["endpoint"].replace(model, replacement_model_name) return kwargs - async def _ageneric_api_call_with_fallbacks_helper( - self, model: str, original_generic_function: Callable, **kwargs - ): + async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): """ Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router """ @@ -4858,9 +4327,7 @@ class Router: return await original_generic_function(model=model, **kwargs) raise e - self._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=function_name - ) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=function_name) data = deployment["litellm_params"].copy() model_name = data["model"] @@ -4877,9 +4344,7 @@ class Router: model=data["model"], custom_llm_provider=custom_llm_provider, ) - custom_llm_provider = ( - custom_llm_provider or inferred_custom_llm_provider - ) + custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider except Exception: custom_llm_provider = None @@ -4901,9 +4366,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -4920,9 +4383,7 @@ class Router: response = await response # type: ignore self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"ageneric_api_call_with_fallbacks(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"ageneric_api_call_with_fallbacks(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: @@ -4968,29 +4429,21 @@ class Router: # helper knows which underlying API to call on fallback. fallback_kwargs: Dict[str, Any] = kwargs.copy() if isinstance(fallback_kwargs.get("litellm_metadata"), dict): - fallback_kwargs["litellm_metadata"] = safe_deep_copy( - fallback_kwargs["litellm_metadata"] - ) + fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) if isinstance(fallback_kwargs.get("metadata"), dict): fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"]) fallback_kwargs["original_generic_function"] = original_function - response = await self._ageneric_api_call_with_fallbacks( - original_function=original_function, **kwargs - ) + response = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) - if kwargs.get("stream") and isinstance( - response, BaseResponsesAPIStreamingIterator - ): + if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator): return await self._aresponses_streaming_iterator( response=response, initial_kwargs=fallback_kwargs, ) return response - def _generic_api_call_with_fallbacks( - self, model: str, original_function: Callable, **kwargs - ): + def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): """ Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router Args: @@ -5001,9 +4454,7 @@ class Router: The response from the handler function """ handler_name = original_function.__name__ - metadata_variable_name = _get_router_metadata_variable_name( - function_name="generic_api_call" - ) + metadata_variable_name = _get_router_metadata_variable_name(function_name="generic_api_call") try: verbose_router_logger.debug( f"Inside _generic_api_call() - handler: {handler_name}, model: {model}; kwargs: {kwargs}" @@ -5019,9 +4470,7 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) - self._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") data = deployment["litellm_params"].copy() model_name = data["model"] @@ -5043,9 +4492,7 @@ class Router: model=data["model"], custom_llm_provider=custom_llm_provider, ) - custom_llm_provider = ( - custom_llm_provider or inferred_custom_llm_provider - ) + custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider except Exception: custom_llm_provider = None @@ -5059,14 +4506,10 @@ class Router: ) self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"{handler_name}(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"{handler_name}(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"{handler_name}(model={model})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"{handler_name}(model={model})\033[31m Exception {str(e)}\033[0m") if model is not None: self.fail_calls[model] += 1 raise e @@ -5091,9 +4534,7 @@ class Router: def _embedding(self, input: Union[str, List], model: str, **kwargs): model_name = None try: - verbose_router_logger.debug( - f"Inside embedding()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside embedding()- model: {model}; kwargs: {kwargs}") deployment = self.get_available_deployment( model=model, input=input, @@ -5103,9 +4544,7 @@ class Router: data = deployment["litellm_params"].copy() model_name = data["model"] - potential_model_client = self._get_client( - deployment=deployment, kwargs=kwargs, client_type="sync" - ) + potential_model_client = self._get_client(deployment=deployment, kwargs=kwargs, client_type="sync") # check if provided keys == client keys # dynamic_api_key = kwargs.get("api_key", None) if ( @@ -5132,14 +4571,10 @@ class Router: } ) self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.embedding(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.embedding(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"litellm.embedding(model={model_name})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.embedding(model={model_name})\033[31m Exception {str(e)}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -5172,9 +4607,7 @@ class Router: async def _aembedding(self, input: Union[str, List], model: str, **kwargs): model_name = None try: - verbose_router_logger.debug( - f"Inside _aembedding()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _aembedding()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -5208,9 +4641,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -5227,14 +4658,10 @@ class Router: response = await response self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.aembedding(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.aembedding(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: - verbose_router_logger.info( - f"litellm.aembedding(model={model_name})\033[31m Exception {str(e)}\033[0m" - ) + verbose_router_logger.info(f"litellm.aembedding(model={model_name})\033[31m Exception {str(e)}\033[0m") if model_name is not None: self.fail_calls[model_name] += 1 raise e @@ -5272,9 +4699,7 @@ class Router: try: from litellm.router_utils.common_utils import add_model_file_id_mappings - verbose_router_logger.debug( - f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) healthy_deployments = await self.async_get_healthy_deployments( model=model, @@ -5310,17 +4735,13 @@ class Router: custom_llm_provider=custom_llm_provider, ) # Preserve explicitly stored provider, fallback to inferred - custom_llm_provider = ( - custom_llm_provider or inferred_custom_llm_provider - ) + custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## purpose = cast(Optional[OpenAIFilesPurpose], kwargs.get("purpose")) file = cast(Optional[FileTypes], kwargs.get("file")) if not file or not purpose: - raise Exception( - "file and file_purpose are required for create_file" - ) + raise Exception("file and file_purpose are required for create_file") replace_model_in_jsonl_bool = should_replace_model_in_jsonl( purpose=purpose, @@ -5335,9 +4756,7 @@ class Router: if ( "gcs_bucket_name" in data ): # TODO: Remove this once we have a better way to handle GCS bucket name: Problem is that we need to pass the gcs_bucket_name to the router for the create_file call but it doesn't show up there - kwargs_copy.setdefault("litellm_metadata", {})[ - "gcs_bucket_name" - ] = data["gcs_bucket_name"] + kwargs_copy.setdefault("litellm_metadata", {})["gcs_bucket_name"] = data["gcs_bucket_name"] response = litellm.acreate_file( **{ **data, @@ -5354,9 +4773,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -5373,9 +4790,7 @@ class Router: response = await response # type: ignore self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m") return response @@ -5396,9 +4811,7 @@ class Router: healthy_deployments=healthy_deployments, responses=responses ) returned_response = cast(OpenAIFileObject, responses[0]) - returned_response._hidden_params["model_file_id_mapping"] = ( - model_file_id_mapping - ) + returned_response._hidden_params["model_file_id_mapping"] = model_file_id_mapping return returned_response except Exception as e: verbose_router_logger.exception( @@ -5430,9 +4843,7 @@ class Router: from litellm.vector_stores.main import acreate # Use the factory function to handle the call - factory_fn = self.factory_function( - acreate, call_type="avector_store_create" - ) + factory_fn = self.factory_function(acreate, call_type="avector_store_create") return await factory_fn(**kwargs) from litellm.vector_stores import acreate as avector_store_create_sdk @@ -5482,9 +4893,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: await self.async_routing_strategy_pre_call_checks( deployment=deployment, parent_otel_span=parent_otel_span @@ -5497,9 +4906,7 @@ class Router: response = await response self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.avector_store_create(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.avector_store_create(model={model_name})\033[32m 200 OK\033[0m") return response except Exception as e: @@ -5530,9 +4937,7 @@ class Router: kwargs["model"] = model kwargs["original_function"] = self._acreate_batch kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - metadata_variable_name = _get_router_metadata_variable_name( - function_name="_acreate_batch" - ) + metadata_variable_name = _get_router_metadata_variable_name(function_name="_acreate_batch") self._update_kwargs_before_fallbacks( model=model, kwargs=kwargs, @@ -5558,9 +4963,7 @@ class Router: **kwargs, ) -> LiteLLMBatch: try: - verbose_router_logger.debug( - f"Inside _acreate_batch()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _acreate_batch()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -5571,9 +4974,7 @@ class Router: data = deployment["litellm_params"].copy() model_name = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="_acreate_batch" - ) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="_acreate_batch") model_client = self._get_async_openai_model_client( deployment=deployment, @@ -5605,9 +5006,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -5624,9 +5023,7 @@ class Router: response = await response # type: ignore self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.acreate_batch(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.acreate_batch(model={model_name})\033[32m 200 OK\033[0m") return response # type: ignore except Exception as e: @@ -5674,9 +5071,7 @@ class Router: data = model_name["litellm_params"].copy() custom_llm_provider = data.get("custom_llm_provider") if model is None: - raise Exception( - f"Model not found in litellm_params for deployment: {model_name}" - ) + raise Exception(f"Model not found in litellm_params for deployment: {model_name}") # Update kwargs with the current model name or any other model-specific adjustments ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## if not custom_llm_provider: @@ -5712,18 +5107,11 @@ class Router: and len(filtered_model_list) > 0 ): results = await asyncio.gather( - *[ - try_retrieve_batch(cast(DeploymentTypedDict, model)) - for model in filtered_model_list - ], + *[try_retrieve_batch(cast(DeploymentTypedDict, model)) for model in filtered_model_list], return_exceptions=True, ) - elif filtered_model_list is not None and isinstance( - filtered_model_list, dict - ): - results = await try_retrieve_batch( - cast(DeploymentTypedDict, filtered_model_list) - ) + elif filtered_model_list is not None and isinstance(filtered_model_list, dict): + results = await try_retrieve_batch(cast(DeploymentTypedDict, filtered_model_list)) else: raise Exception("No healthy deployments found.") @@ -5741,11 +5129,7 @@ class Router: raise receieved_exceptions[0] # Raising the first exception encountered # If no exceptions were encountered, raise a generic exception - raise Exception( - "Unable to find batch in any model. Received errors - {}".format( - receieved_exceptions - ) - ) + raise Exception("Unable to find batch in any model. Received errors - {}".format(receieved_exceptions)) except Exception as e: asyncio.create_task( send_llm_exception_alert( @@ -5769,9 +5153,7 @@ class Router: kwargs["model"] = model kwargs["original_function"] = self._acancel_batch kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - metadata_variable_name = _get_router_metadata_variable_name( - function_name="_acancel_batch" - ) + metadata_variable_name = _get_router_metadata_variable_name(function_name="_acancel_batch") self._update_kwargs_before_fallbacks( model=model, kwargs=kwargs, @@ -5797,9 +5179,7 @@ class Router: **kwargs, ) -> LiteLLMBatch: try: - verbose_router_logger.debug( - f"Inside _acancel_batch()- model: {model}; kwargs: {kwargs}" - ) + verbose_router_logger.debug(f"Inside _acancel_batch()- model: {model}; kwargs: {kwargs}") parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) deployment = await self.async_get_available_deployment( model=model, @@ -5817,9 +5197,7 @@ class Router: data.update(resolved_credentials) data.pop("litellm_credential_name", None) model_name = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="_acancel_batch" - ) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="_acancel_batch") model_client = self._get_async_openai_model_client( deployment=deployment, @@ -5851,9 +5229,7 @@ class Router: client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): + if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): async with rpm_semaphore: """ - Check rpm limits before making the call @@ -5870,9 +5246,7 @@ class Router: response = await response # type: ignore self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.acancel_batch(model={model_name})\033[32m 200 OK\033[0m" - ) + verbose_router_logger.info(f"litellm.acancel_batch(model={model_name})\033[32m 200 OK\033[0m") return response # type: ignore except Exception as e: @@ -5899,16 +5273,12 @@ class Router: async def try_retrieve_batch(model: DeploymentTypedDict): try: # Update kwargs with the current model name or any other model-specific adjustments - return await litellm.alist_batches( - **{**model["litellm_params"], **kwargs} - ) + return await litellm.alist_batches(**{**model["litellm_params"], **kwargs}) except Exception: return None # Check all models in parallel - results = await asyncio.gather( - *[try_retrieve_batch(model) for model in filtered_model_list] - ) + results = await asyncio.gather(*[try_retrieve_batch(model) for model in filtered_model_list]) final_results: Dict = { "object": "list", @@ -6108,9 +5478,7 @@ class Router: client: Optional[Any] = None, **kwargs, ): - return self._generic_api_call_with_fallbacks( - original_function=original_function, **kwargs - ) + return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) return sync_wrapper @@ -6129,9 +5497,7 @@ class Router: if custom_llm_provider and "custom_llm_provider" not in kwargs: kwargs["custom_llm_provider"] = custom_llm_provider if kwargs.get("model"): - return self._generic_api_call_with_fallbacks( - original_function=original_function, **kwargs - ) + return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) return original_function(**kwargs) return vector_store_sync_wrapper @@ -6415,9 +5781,7 @@ class Router: """ from litellm.responses.utils import ResponsesAPIRequestUtils - model_id = ResponsesAPIRequestUtils.get_model_id_from_response_id( - kwargs.get("response_id") - ) + model_id = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) if model_id is not None: kwargs["model"] = model_id return await self._ageneric_api_call_with_fallbacks( @@ -6532,11 +5896,7 @@ class Router: # that fails with RouterRateLimitError whenever the "remaining" entries # are all in cooldown — the inner async_get_healthy_deployments call # would find an empty list and raise immediately. - cooldown_ids = set( - await _async_get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=None - ) - ) + cooldown_ids = set(await _async_get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=None)) remaining = (all_ids - cooldown_ids) - excluded if not remaining: return None @@ -6616,17 +5976,10 @@ class Router: e, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError), ) - _request_team_id: Optional[str] = (kwargs.get("metadata", {}) or {}).get( - "user_api_key_team_id" - ) + _request_team_id: Optional[str] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") # Use wildcard-aware lookup so order-based fallback also works for model # groups resolved via pattern routing (e.g. `openai/*` -> `openai/gpt-4.1-mini`). - all_deployments = ( - self.get_model_list( - model_name=original_model_group, team_id=_request_team_id - ) - or [] - ) + all_deployments = self.get_model_list(model_name=original_model_group, team_id=_request_team_id) or [] _order_set: set = { litellm.utils._get_deployment_order(d) for d in all_deployments @@ -6636,14 +5989,10 @@ class Router: if len(order_values) > 1 and not _skip_order_fallback: # Determine which order levels have already been tried current_target = kwargs.get("_target_order") - skip_up_to = ( - current_target if current_target is not None else order_values[0] - ) + skip_up_to = current_target if current_target is not None else order_values[0] # Build order-based fallback entries (skip already-tried levels) order_fallback_entries: List = [ - {"model": original_model_group, "_target_order": o} - for o in order_values - if o > skip_up_to + {"model": original_model_group, "_target_order": o} for o in order_values if o > skip_up_to ] # Get external fallbacks — handle both standard and non-standard formats external_fallback_group: Optional[List] = None @@ -6661,9 +6010,7 @@ class Router: external_fallback_group = fallbacks[generic_idx]["*"] # Combined list: order fallbacks first, then external - combined_fallbacks = order_fallback_entries + ( - external_fallback_group or [] - ) + combined_fallbacks = order_fallback_entries + (external_fallback_group or []) if combined_fallbacks: input_kwargs.update( @@ -6679,11 +6026,7 @@ class Router: return response # Weighted intra-group failover (simple-shuffle only); see _maybe_run_weighted_failover. - if ( - self.enable_weighted_failover - and not _skip_order_fallback - and original_model_group is not None - ): + if self.enable_weighted_failover and not _skip_order_fallback and original_model_group is not None: response = await self._maybe_run_weighted_failover( exception=e, original_model_group=original_model_group, @@ -6699,9 +6042,7 @@ class Router: verbose_router_logger.info("Trying to fallback b/w models") # check if client-side fallbacks are used (e.g. fallbacks = ["gpt-3.5-turbo", "claude-3-haiku"] or fallbacks=[{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hey, how's it going?"}]}] - is_non_standard_fallback_format = _check_non_standard_fallback_format( - fallbacks=fallbacks - ) + is_non_standard_fallback_format = _check_non_standard_fallback_format(fallbacks=fallbacks) if is_non_standard_fallback_format: input_kwargs.update( @@ -6804,10 +6145,7 @@ class Router: verbose_router_logger.info( f"No fallback model group found for original model_group={model_group}. Fallbacks={fallbacks}" ) - if ( - hasattr(original_exception, "message") - and litellm.expose_router_debug_in_errors - ): + if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={fallbacks}" # type: ignore raise original_exception @@ -6838,22 +6176,15 @@ class Router: ) fallback_failure_exception_str = str(new_exception) - if ( - hasattr(original_exception, "message") - and litellm.expose_router_debug_in_errors - ): + if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: # add the available fallbacks to the exception - original_exception.message += ( - ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore - model_group, - fallback_model_group, - ) + original_exception.message += ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore + model_group, + fallback_model_group, ) if len(fallback_failure_exception_str) > 0: original_exception.message += ( # type: ignore - "\nError doing the fallback: {}".format( - fallback_failure_exception_str - ) + "\nError doing the fallback: {}".format(fallback_failure_exception_str) ) raise original_exception @@ -6868,12 +6199,8 @@ class Router: include_fallback_errors = kwargs.get("include_fallback_errors", False) is True disable_fallbacks: Optional[bool] = kwargs.pop("disable_fallbacks", False) fallbacks: Optional[List] = kwargs.get("fallbacks", self.fallbacks) - context_window_fallbacks: Optional[List] = kwargs.get( - "context_window_fallbacks", self.context_window_fallbacks - ) - content_policy_fallbacks: Optional[List] = kwargs.get( - "content_policy_fallbacks", self.content_policy_fallbacks - ) + context_window_fallbacks: Optional[List] = kwargs.get("context_window_fallbacks", self.context_window_fallbacks) + content_policy_fallbacks: Optional[List] = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) mock_timeout = kwargs.pop("mock_timeout", None) @@ -6887,9 +6214,7 @@ class Router: ) if mock_timeout is not None: - response = await self.async_function_with_retries( - *args, **kwargs, mock_timeout=mock_timeout - ) + response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) else: response = await self.async_function_with_retries(*args, **kwargs) if verbose_router_logger.isEnabledFor(logging.DEBUG): @@ -6965,16 +6290,10 @@ class Router: original_function = kwargs.pop("original_function") fallbacks = kwargs.pop("fallbacks", self.fallbacks) parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) - context_window_fallbacks = kwargs.pop( - "context_window_fallbacks", self.context_window_fallbacks - ) - content_policy_fallbacks = kwargs.pop( - "content_policy_fallbacks", self.content_policy_fallbacks - ) + context_window_fallbacks = kwargs.pop("context_window_fallbacks", self.context_window_fallbacks) + content_policy_fallbacks = kwargs.pop("content_policy_fallbacks", self.content_policy_fallbacks) # Support per-request model_group_retry_policy override (from key/team settings) - model_group_retry_policy = kwargs.pop( - "model_group_retry_policy", self.model_group_retry_policy - ) + model_group_retry_policy = kwargs.pop("model_group_retry_policy", self.model_group_retry_policy) model_group: Optional[str] = kwargs.get("model") num_retries = kwargs.pop("num_retries", None) if num_retries is None: @@ -6994,27 +6313,19 @@ class Router: ) ## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking _metadata["attempted_retries"] = 0 - _metadata["max_retries"] = ( - num_retries # Updated after overrides in exception handler - ) + _metadata["max_retries"] = num_retries # Updated after overrides in exception handler try: - self._handle_mock_testing_rate_limit_error( - model_group=model_group, kwargs=kwargs - ) + self._handle_mock_testing_rate_limit_error(model_group=model_group, kwargs=kwargs) # if the function call is successful, no exception will be raised and we'll break out of the loop response = await self.make_call(original_function, *args, **kwargs) - response = add_retry_headers_to_response( - response=response, attempted_retries=0, max_retries=None - ) + response = add_retry_headers_to_response(response=response, attempted_retries=0, max_retries=None) return response except Exception as e: current_attempt = None original_exception = e deployment_num_retries = getattr(e, "num_retries", None) - if deployment_num_retries is not None and isinstance( - deployment_num_retries, int - ): + if deployment_num_retries is not None and isinstance(deployment_num_retries, int): num_retries = deployment_num_retries """ Retry Logic @@ -7034,9 +6345,7 @@ class Router: # get num_retries from retry policy # Use the model_group captured at the start of the function, or get it from metadata # kwargs.get("model") at this point is the deployment model, not the model_group - _model_group_for_retry_policy = ( - model_group or _metadata.get("model_group") or kwargs.get("model") - ) + _model_group_for_retry_policy = model_group or _metadata.get("model_group") or kwargs.get("model") # Use per-request model_group_retry_policy if provided, otherwise use self _retry_policy_retries = _get_num_retries_from_retry_policy( exception=original_exception, @@ -7068,9 +6377,7 @@ class Router: else: raise - verbose_router_logger.debug( - f"Retrying request with num_retries: {num_retries}" - ) + verbose_router_logger.debug(f"Retrying request with num_retries: {num_retries}") # decides how long to sleep before retry retry_after = self._time_to_sleep_before_retry( e=original_exception, @@ -7089,9 +6396,7 @@ class Router: _metadata["max_retries"] = num_retries # if the function call is successful, no exception will be raised and we'll break out of the loop response = await self.make_call(original_function, *args, **kwargs) - if coroutine_checker.is_async_callable( - response - ): # async errors are often returned as coroutines + if coroutine_checker.is_async_callable(response): # async errors are often returned as coroutines response = await response response = add_retry_headers_to_response( @@ -7154,9 +6459,7 @@ class Router: # is current_attempt + 1, which equals num_retries when all retries are exhausted. # We've already verified num_retries > 0 before entering the loop, so current_attempt # will always be set (never None) when we reach this point. - actual_retries_attempted = ( - current_attempt + 1 if current_attempt is not None else num_retries - ) + actual_retries_attempted = current_attempt + 1 if current_attempt is not None else num_retries setattr(original_exception, "num_retries", actual_retries_attempted) raise original_exception @@ -7167,42 +6470,29 @@ class Router: """ model_group = kwargs.get("model") response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable( - response - ): + if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): response = await response ## PROCESS RESPONSE HEADERS - response = await self.set_response_headers( - response=response, model_group=model_group, request_kwargs=kwargs - ) + response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) return response - def _handle_mock_testing_rate_limit_error( - self, kwargs: dict, model_group: Optional[str] = None - ): + def _handle_mock_testing_rate_limit_error(self, kwargs: dict, model_group: Optional[str] = None): """ Helper function to raise a mock litellm.RateLimitError error for testing purposes. Raises: litellm.RateLimitError error when `mock_testing_rate_limit_error=True` passed in request params """ - mock_testing_rate_limit_error: Optional[bool] = kwargs.pop( - "mock_testing_rate_limit_error", None - ) + mock_testing_rate_limit_error: Optional[bool] = kwargs.pop("mock_testing_rate_limit_error", None) available_models = self.get_model_list(model_name=model_group) num_retries: Optional[int] = None if available_models is not None and len(available_models) == 1: - num_retries = cast( - Optional[int], available_models[0]["litellm_params"].get("num_retries") - ) + num_retries = cast(Optional[int], available_models[0]["litellm_params"].get("num_retries")) - if ( - mock_testing_rate_limit_error is not None - and mock_testing_rate_limit_error is True - ): + if mock_testing_rate_limit_error is not None and mock_testing_rate_limit_error is True: verbose_router_logger.info( f"litellm.router.py::_mock_rate_limit_error() - Raising mock RateLimitError for model={model_group}" ) @@ -7239,16 +6529,10 @@ class Router: _num_all_deployments = len(all_deployments) ### CHECK IF RATE LIMIT / CONTEXT WINDOW ERROR / CONTENT POLICY VIOLATION ERROR w/ fallbacks available / Bad Request Error - if ( - isinstance(error, litellm.ContextWindowExceededError) - and context_window_fallbacks is not None - ): + if isinstance(error, litellm.ContextWindowExceededError) and context_window_fallbacks is not None: raise error - if ( - isinstance(error, litellm.ContentPolicyViolationError) - and content_policy_fallbacks is not None - ): + if isinstance(error, litellm.ContentPolicyViolationError) and content_policy_fallbacks is not None: raise error status_code = getattr(error, "status_code", None) @@ -7273,9 +6557,7 @@ class Router: - if other deployments available -> retry - else -> raise error """ - if ( - _num_all_deployments <= 1 - ): # if there is only 1 deployment for this model group then don't retry + if _num_all_deployments <= 1: # if there is only 1 deployment for this model group then don't retry raise error # then raise error # Do not retry if there are no healthy deployments @@ -7350,11 +6632,7 @@ class Router: ## base case - single deployment if all_deployments is not None and len(all_deployments) == 1: pass - elif ( - healthy_deployments is not None - and isinstance(healthy_deployments, list) - and len(healthy_deployments) > 0 - ): + elif healthy_deployments is not None and isinstance(healthy_deployments, list) and len(healthy_deployments) > 0: return 0 response_headers: Optional[httpx.Headers] = None @@ -7398,9 +6676,7 @@ class Router: # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object is None") if kwargs["litellm_params"].get("metadata") is None: @@ -7410,9 +6686,7 @@ class Router: "deployment", None ) # stable name - works for wildcard routes as well # Get model_group and id from kwargs like the sync version does - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) model_info = kwargs["litellm_params"].get("model_info", {}) or {} id = model_info.get("id", None) if model_group is None or id is None: @@ -7466,13 +6740,9 @@ class Router: # Setup values # ------------ dt = get_utc_datetime() - current_minute = dt.strftime( - "%H-%M" - ) # use the same timezone regardless of system clock + current_minute = dt.strftime("%H-%M") # use the same timezone regardless of system clock - tpm_key = RouterCacheEnum.TPM.value.format( - id=id, current_minute=current_minute, model=deployment_name - ) + tpm_key = RouterCacheEnum.TPM.value.format(id=id, current_minute=current_minute, model=deployment_name) # ------------ # Update usage # ------------ @@ -7489,9 +6759,7 @@ class Router: ) ## RPM - rpm_key = RouterCacheEnum.RPM.value.format( - id=id, current_minute=current_minute, model=deployment_name - ) + rpm_key = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name) pipeline_operations.append( RedisPipelineIncrementOperation( key=rpm_key, @@ -7509,9 +6777,7 @@ class Router: except Exception as e: verbose_router_logger.debug( - "litellm.router.Router::deployment_callback_on_success(): Exception occured - {}".format( - str(e) - ) + "litellm.router.Router::deployment_callback_on_success(): Exception occured - {}".format(str(e)) ) pass @@ -7632,11 +6898,7 @@ class Router: """ Update RPM usage for a deployment """ - deployment_name = kwargs[ - "litellm_params" - ][ - "metadata" - ].get( + deployment_name = kwargs["litellm_params"]["metadata"].get( "deployment", None ) # handles wildcard routes - by giving the original name sent to `litellm.completion` model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) @@ -7649,14 +6911,10 @@ class Router: parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) dt = get_utc_datetime() - current_minute = dt.strftime( - "%H-%M" - ) # use the same timezone regardless of system clock + current_minute = dt.strftime("%H-%M") # use the same timezone regardless of system clock ## RPM - rpm_key = RouterCacheEnum.RPM.value.format( - id=id, current_minute=current_minute, model=deployment_name - ) + rpm_key = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name) await self.cache.async_increment_cache( key=rpm_key, value=1, @@ -7664,9 +6922,7 @@ class Router: ttl=RoutingArgs.ttl.value, ) - def _get_metadata_variable_name_from_kwargs( - self, kwargs: dict - ) -> Literal["metadata", "litellm_metadata"]: + def _get_metadata_variable_name_from_kwargs(self, kwargs: dict) -> Literal["metadata", "litellm_metadata"]: """ Helper to return what the "metadata" field should be called in the request data @@ -7685,9 +6941,7 @@ class Router: When a retry or fallback happens, log the details of the just failed model call - similar to Sentry breadcrumbing """ try: - _metadata_var = ( - "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - ) + _metadata_var = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" # Log failed model as the previous model previous_model = { "exception_type": type(e).__name__, @@ -7715,9 +6969,7 @@ class Router: except Exception as e: raise e - def _update_usage( - self, deployment_id: str, parent_otel_span: Optional[Span] - ) -> int: + def _update_usage(self, deployment_id: str, parent_otel_span: Optional[Span]) -> int: """ Update deployment rpm for that minute @@ -7726,19 +6978,13 @@ class Router: """ rpm_key = deployment_id - request_count = self.cache.get_cache( - key=rpm_key, parent_otel_span=parent_otel_span, local_only=True - ) + request_count = self.cache.get_cache(key=rpm_key, parent_otel_span=parent_otel_span, local_only=True) if request_count is None: request_count = 1 - self.cache.set_cache( - key=rpm_key, value=request_count, local_only=True, ttl=60 - ) # only store for 60s + self.cache.set_cache(key=rpm_key, value=request_count, local_only=True, ttl=60) # only store for 60s else: request_count += 1 - self.cache.set_cache( - key=rpm_key, value=request_count, local_only=True - ) # don't change existing ttl + self.cache.set_cache(key=rpm_key, value=request_count, local_only=True) # don't change existing ttl return request_count @@ -7751,9 +6997,7 @@ class Router: return True return False - def _should_raise_content_policy_error( - self, model: str, response: ModelResponse, kwargs: dict - ) -> bool: + def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: """ Determines if a content policy error should be raised. @@ -7765,9 +7009,7 @@ class Router: if response.choices[0].finish_reason != "content_filter": return False - content_policy_fallbacks = kwargs.get( - "content_policy_fallbacks", self.content_policy_fallbacks - ) + content_policy_fallbacks = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) ### ONLY RAISE ERROR IF CP FALLBACK AVAILABLE ### if content_policy_fallbacks is not None: @@ -7804,9 +7046,7 @@ class Router: litellm_router_instance=self, parent_otel_span=parent_otel_span ) unhealthy_set = set(unhealthy_deployments) - healthy_deployments: list = [ - d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set - ] + healthy_deployments: list = [d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set] healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments @@ -7836,9 +7076,7 @@ class Router: # Convert to set for O(1) lookup instead of O(n) unhealthy_deployments_set = set(unhealthy_deployments) healthy_deployments: list = [ - d - for d in _all_deployments - if d["model_info"]["id"] not in unhealthy_deployments_set + d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_deployments_set ] healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments @@ -7944,14 +7182,12 @@ class Router: for _callback in litellm.callbacks: if isinstance(_callback, CustomLogger): try: - returned_healthy_deployments = ( - await _callback.async_filter_deployments( - model=model, - healthy_deployments=returned_healthy_deployments, - messages=messages, - request_kwargs=request_kwargs, - parent_otel_span=parent_otel_span, - ) + returned_healthy_deployments = await _callback.async_filter_deployments( + model=model, + healthy_deployments=returned_healthy_deployments, + messages=messages, + request_kwargs=request_kwargs, + parent_otel_span=parent_otel_span, ) except Exception as e: ## LOG FAILURE EVENT @@ -8024,9 +7260,7 @@ class Router: if all(model_info.get(f) is not None for f in cache_fields): return try: - backend_info = litellm.get_model_info( - model=backend_model, custom_llm_provider=custom_llm_provider - ) + backend_info = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider) except Exception: return for field in cache_fields: @@ -8082,9 +7316,7 @@ class Router: ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes _model_name = deployment.litellm_params.model if deployment.litellm_params.custom_llm_provider is not None: - _model_name = ( - deployment.litellm_params.custom_llm_provider + "/" + _model_name - ) + _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name # For the shared backend key, strip custom pricing fields so that # one deployment's pricing overrides don't pollute another @@ -8092,12 +7324,8 @@ class Router: # Each deployment's full pricing is already stored under its # unique model_id above. _custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys() - _shared_model_info = { - k: v for k, v in _model_info.items() if k not in _custom_pricing_fields - } - _existing_shared_mode = ( - cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {} - ).get("mode") + _shared_model_info = {k: v for k, v in _model_info.items() if k not in _custom_pricing_fields} + _existing_shared_mode = (cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {}).get("mode") _deployment_mode = _shared_model_info.get("mode") # Keep the built-in bridge mode stable for shared backend keys. # Multiple aliases can point at the same provider/model backend, @@ -8107,12 +7335,8 @@ class Router: # (e.g. chat -> responses) and unrelated mode changes still apply, # and so a missing deployment mode does not silently clear the # existing shared backend mode. - _is_responses_to_chat_downgrade = ( - _existing_shared_mode == "responses" and _deployment_mode == "chat" - ) - _would_clear_existing_mode = ( - _existing_shared_mode is not None and _deployment_mode is None - ) + _is_responses_to_chat_downgrade = _existing_shared_mode == "responses" and _deployment_mode == "chat" + _would_clear_existing_mode = _existing_shared_mode is not None and _deployment_mode is None if _is_responses_to_chat_downgrade or _would_clear_existing_mode: if _deployment_mode is not None: verbose_router_logger.warning( @@ -8134,10 +7358,7 @@ class Router: litellm.register_model(model_cost=_backend_alias_cost) ## Check if LLM Deployment is allowed for this deployment - if ( - self.deployment_is_active_for_environment(deployment=deployment) - is not True - ): + if self.deployment_is_active_for_environment(deployment=deployment) is not True: verbose_router_logger.warning( f"Ignoring deployment {deployment.model_name} as it is not active for environment {deployment.model_info['supported_environments']}" ) @@ -8158,9 +7379,7 @@ class Router: model = deployment.to_json(exclude_none=True) - self._add_model_to_list_and_index_map( - model=model, model_id=deployment.model_info.id - ) + self._add_model_to_list_and_index_map(model=model, model_id=deployment.model_info.id) return deployment except Exception as e: if self.ignore_invalid_deployments: @@ -8197,26 +7416,20 @@ class Router: """ from litellm.router_strategy.auto_router.auto_router import AutoRouter - auto_router_config_path: Optional[str] = ( - deployment.litellm_params.auto_router_config_path - ) + auto_router_config_path: Optional[str] = deployment.litellm_params.auto_router_config_path auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config if auto_router_config_path is None and auto_router_config is None: raise ValueError( "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" ) - default_model: Optional[str] = ( - deployment.litellm_params.auto_router_default_model - ) + default_model: Optional[str] = deployment.litellm_params.auto_router_default_model if default_model is None: raise ValueError( "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" ) - embedding_model: Optional[str] = ( - deployment.litellm_params.auto_router_embedding_model - ) + embedding_model: Optional[str] = deployment.litellm_params.auto_router_embedding_model if embedding_model is None: raise ValueError( "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" @@ -8259,13 +7472,9 @@ class Router: ComplexityRouter, ) - complexity_router_config: Optional[dict] = ( - deployment.litellm_params.complexity_router_config - ) + complexity_router_config: Optional[dict] = deployment.litellm_params.complexity_router_config - default_model: Optional[str] = ( - deployment.litellm_params.complexity_router_default_model - ) + default_model: Optional[str] = deployment.litellm_params.complexity_router_default_model # If no default model specified, try to get from config tiers if default_model is None and complexity_router_config: @@ -8315,38 +7524,22 @@ class Router: litellm._async_success_callback, litellm._async_failure_callback, ): - litellm.logging_callback_manager.remove_callbacks_by_type( - _cb_list, AdaptiveRouterPostCallHook - ) + litellm.logging_callback_manager.remove_callbacks_by_type(_cb_list, AdaptiveRouterPostCallHook) for entry in self.model_list or []: - lp = ( - entry.get("litellm_params") - if isinstance(entry, dict) - else entry.litellm_params - ) - lp_model = ( - (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None - ) + lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params + lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None if not (lp_model and lp_model.startswith("auto_router/adaptive_router")): continue - model_name = ( - entry.get("model_name") if isinstance(entry, dict) else entry.model_name - ) + model_name = entry.get("model_name") if isinstance(entry, dict) else entry.model_name if not model_name or not lp: continue if model_name in self.adaptive_routers: continue deployment = Deployment( model_name=model_name, - litellm_params=( - lp if not isinstance(lp, dict) else LiteLLM_Params(**lp) - ), - model_info=( - entry.get("model_info") - if isinstance(entry, dict) - else entry.model_info - ), + litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)), + model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info), ) self.init_adaptive_router_deployment(deployment=deployment) @@ -8375,9 +7568,7 @@ class Router: raw_config = deployment.litellm_params.adaptive_router_config if raw_config is None: - raise ValueError( - "adaptive_router_config is required for adaptive-router deployments." - ) + raise ValueError("adaptive_router_config is required for adaptive-router deployments.") config = AdaptiveRouterConfig(**raw_config) @@ -8391,18 +7582,14 @@ class Router: continue d = (self.model_list or [])[indices[0]] mi = d.get("model_info") if isinstance(d, dict) else d.model_info - mi_dict: Dict[str, Any] = ( - mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) - ) + mi_dict: Dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) prefs_raw = mi_dict.get("adaptive_router_preferences") if prefs_raw is not None: model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) # `input_cost_per_token` is a LiteLLM_Params field per types/router.py. lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params - lp_dict: Dict[str, Any] = ( - lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) - ) + lp_dict: Dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) cost = lp_dict.get("input_cost_per_token") if cost is not None: model_to_cost[name] = float(cost) @@ -8452,13 +7639,9 @@ class Router: QualityRouter, ) - quality_router_config: Optional[dict] = ( - deployment.litellm_params.quality_router_config - ) + quality_router_config: Optional[dict] = deployment.litellm_params.quality_router_config - default_model: Optional[str] = ( - deployment.litellm_params.quality_router_default_model - ) + default_model: Optional[str] = deployment.litellm_params.quality_router_default_model if default_model is None and quality_router_config: default_model = quality_router_config.get("default_model") @@ -8501,9 +7684,7 @@ class Router: return True litellm_environment = get_secret_str(secret_name="LITELLM_ENVIRONMENT") if litellm_environment is None: - raise ValueError( - "Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env" - ) + raise ValueError("Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env") if litellm_environment not in VALID_LITELLM_ENVIRONMENTS: raise ValueError( @@ -8570,9 +7751,7 @@ class Router: _model_info=_model_info, ) - verbose_router_logger.debug( - f"\nInitialized Model List {self.get_model_names()}" - ) + verbose_router_logger.debug(f"\nInitialized Model List {self.get_model_names()}") self.model_names = {m["model_name"] for m in model_list} # Note: model_name_to_deployment_indices is already built incrementally @@ -8611,15 +7790,12 @@ class Router: api_base, ) = litellm.get_llm_provider( model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.get( - "custom_llm_provider", None - ), + custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), ) # done reading model["litellm_params"] # Check if provider is supported: either in enum or JSON-configured - if ( - custom_llm_provider not in litellm.provider_list - and not JSONProviderRegistry.exists(custom_llm_provider) + if custom_llm_provider not in litellm.provider_list and not JSONProviderRegistry.exists( + custom_llm_provider ): raise Exception(f"Unsupported provider - {custom_llm_provider}") @@ -8628,16 +7804,10 @@ class Router: ############ Users can either pass tpm/rpm as a litellm_param or a router param ########### # for get_available_deployment, we use the litellm_param["rpm"] # in this snippet we also set rpm to be a litellm_param - if ( - deployment.litellm_params.rpm is None - and getattr(deployment, "rpm", None) is not None - ): + if deployment.litellm_params.rpm is None and getattr(deployment, "rpm", None) is not None: deployment.litellm_params.rpm = getattr(deployment, "rpm") - if ( - deployment.litellm_params.tpm is None - and getattr(deployment, "tpm", None) is not None - ): + if deployment.litellm_params.tpm is None and getattr(deployment, "tpm", None) is not None: deployment.litellm_params.tpm = getattr(deployment, "tpm") # Check if user is trying to use model_name == "*" @@ -8653,19 +7823,13 @@ class Router: if "*" in deployment.model_name: # store this as a regex pattern - all deployments matching this pattern will be sent to this deployment # Store deployment.model_name as a regex pattern - self.pattern_router.add_pattern( - deployment.model_name, deployment.to_json(exclude_none=True) - ) + self.pattern_router.add_pattern(deployment.model_name, deployment.to_json(exclude_none=True)) if deployment.model_info.id: self.provider_default_deployment_ids.append(deployment.model_info.id) _team_id = deployment.model_info.get("team_id") _team_public_model_name = deployment.model_info.get("team_public_model_name") - if ( - _team_id is not None - and _team_public_model_name is not None - and "*" in _team_public_model_name - ): + if _team_id is not None and _team_public_model_name is not None and "*" in _team_public_model_name: if _team_id not in self.team_pattern_routers: self.team_pattern_routers[_team_id] = PatternMatchRouter() self.team_pattern_routers[_team_id].add_pattern( @@ -8704,9 +7868,7 @@ class Router: ######################################################### # Check if this is a complexity-router deployment ######################################################### - if self._is_complexity_router_deployment( - litellm_params=deployment.litellm_params - ): + if self._is_complexity_router_deployment(litellm_params=deployment.litellm_params): self.init_complexity_router_deployment(deployment=deployment) # NOTE: adaptive-router deployments are deferred to the end of @@ -8721,9 +7883,7 @@ class Router: return deployment - def _initialize_deployment_for_pass_through( - self, deployment: Deployment, custom_llm_provider: str, model: str - ): + def _initialize_deployment_for_pass_through(self, deployment: Deployment, custom_llm_provider: str, model: str): """ Optional: Initialize deployment for pass-through endpoints if `deployment.litellm_params.use_in_pass_through` is True @@ -8742,17 +7902,10 @@ class Router: credential_values = {} if custom_llm_provider == "vertex_ai": - vertex_project = ( - credential_values.get("vertex_project") - or deployment.litellm_params.vertex_project - ) - vertex_location = ( - credential_values.get("vertex_location") - or deployment.litellm_params.vertex_location - ) + vertex_project = credential_values.get("vertex_project") or deployment.litellm_params.vertex_project + vertex_location = credential_values.get("vertex_location") or deployment.litellm_params.vertex_location vertex_credentials = ( - credential_values.get("vertex_credentials") - or deployment.litellm_params.vertex_credentials + credential_values.get("vertex_credentials") or deployment.litellm_params.vertex_credentials ) if vertex_project is None or vertex_location is None: @@ -8765,14 +7918,8 @@ class Router: vertex_credentials=vertex_credentials, ) else: - api_base = ( - credential_values.get("api_base") - or deployment.litellm_params.api_base - ) - api_key = ( - credential_values.get("api_key") - or deployment.litellm_params.api_key - ) + api_base = credential_values.get("api_base") or deployment.litellm_params.api_base + api_key = credential_values.get("api_key") or deployment.litellm_params.api_key if api_key is None: verbose_router_logger.debug( "Skipping pass-through credential setup for deployment model=%s, custom_llm_provider=%s; no api_key set. Providers like bedrock resolve credentials at request time.", @@ -8834,9 +7981,7 @@ class Router: ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes _model_name = deployment.litellm_params.model if deployment.litellm_params.custom_llm_provider is not None: - _model_name = ( - deployment.litellm_params.custom_llm_provider + "/" + _model_name - ) + _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name # For the shared backend key, strip custom pricing fields so that # one deployment's pricing overrides don't pollute another @@ -8844,9 +7989,7 @@ class Router: # Each deployment's full pricing is already stored under its # unique model_id above (when present). _custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys() - _shared_model_info = { - k: v for k, v in _model_info_dict.items() if k not in _custom_pricing_fields - } + _shared_model_info = {k: v for k, v in _model_info_dict.items() if k not in _custom_pricing_fields} _backend_alias_cost = {_model_name: _shared_model_info} if "responses/" in _model_name: _stripped_model_name = _model_name.replace("responses/", "") @@ -8854,16 +7997,12 @@ class Router: litellm.register_model(model_cost=_backend_alias_cost) # add to model names - self._add_model_to_list_and_index_map( - model=_deployment, model_id=deployment.model_info.id - ) + self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id) self.model_names.add(deployment.model_name) self._sync_deployment_budget_config(deployment=deployment) return deployment - def _update_deployment_indices_after_removal( - self, model_id: str, removal_idx: int - ) -> None: + def _update_deployment_indices_after_removal(self, model_id: str, removal_idx: int) -> None: """ Helper method to update deployment indices after a deployment has been removed from model_list. @@ -8930,9 +8069,7 @@ class Router: - idx: int - the index in model_list """ team_id = (model.get("model_info") or {}).get("team_id") - team_public_model_name = (model.get("model_info") or {}).get( - "team_public_model_name" - ) + team_public_model_name = (model.get("model_info") or {}).get("team_public_model_name") if team_id and team_public_model_name: key = (team_id, team_public_model_name) if key not in self.team_model_to_deployment_indices: @@ -8940,9 +8077,7 @@ class Router: if idx not in self.team_model_to_deployment_indices[key]: self.team_model_to_deployment_indices[key].append(idx) - def _add_model_to_list_and_index_map( - self, model: dict, model_id: Optional[str] = None - ) -> None: + def _add_model_to_list_and_index_map(self, model: dict, model_id: Optional[str] = None) -> None: """ Helper method to add a model to the model_list and update both indices. @@ -8984,9 +8119,7 @@ class Router: # check if deployment already exists _deployment_model_id = deployment.model_info.id or "" - _deployment_on_router: Optional[Deployment] = self.get_deployment( - model_id=_deployment_model_id - ) + _deployment_on_router: Optional[Deployment] = self.get_deployment(model_id=_deployment_model_id) if _deployment_on_router is not None: # deployment with this model_id exists on the router if ( @@ -9009,9 +8142,7 @@ class Router: self.model_list.pop(removal_idx) self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal( - model_id=deployment_id, removal_idx=removal_idx - ) + self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx) # if the model_id is not in router self.add_deployment(deployment=deployment) @@ -9044,9 +8175,7 @@ class Router: item = self.model_list.pop(deployment_idx) self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal( - model_id=id, removal_idx=deployment_idx - ) + self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) _budget_limiter = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) @@ -9095,15 +8224,11 @@ class Router: return if _budget_limiter is None: - self.add_optional_pre_call_checks( - optional_pre_call_checks=["router_budget_limiting"] - ) + self.add_optional_pre_call_checks(optional_pre_call_checks=["router_budget_limiting"]) _budget_limiter = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: - _budget_limiter.register_deployment_budget( - deployment=deployment.to_json(exclude_none=True) - ) + _budget_limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True)) def get_deployment(self, model_id: str) -> Optional[Deployment]: """ @@ -9135,13 +8260,11 @@ class Router: deployment = self.get_deployment(model_id=model_id) if deployment is None or self._is_deployment_blocked(deployment): return None - return CredentialLiteLLMParams( - **deployment.litellm_params.model_dump(exclude_none=True) - ).model_dump(exclude_none=True) + return CredentialLiteLLMParams(**deployment.litellm_params.model_dump(exclude_none=True)).model_dump( + exclude_none=True + ) - def get_deployment_by_model_group_name( - self, model_group_name: str - ) -> Optional[Deployment]: + def get_deployment_by_model_group_name(self, model_group_name: str) -> Optional[Deployment]: """ Returns -> Deployment or None @@ -9163,9 +8286,7 @@ class Router: raise Exception("Model Name invalid - {}".format(type(model))) return None - def get_deployment_credentials_with_provider( - self, model_id: str - ) -> Optional[Dict[str, Any]]: + def get_deployment_credentials_with_provider(self, model_id: str) -> Optional[Dict[str, Any]]: """ Get API credentials and provider info from a model name in model_list. Useful for passthrough endpoints (files, batches, etc.) that need credentials. @@ -9191,9 +8312,7 @@ class Router: # If not found, try by model_group_name if deployment is None: - deployment = self.get_deployment_by_model_group_name( - model_group_name=model_id - ) + deployment = self.get_deployment_by_model_group_name(model_group_name=model_id) # If still not found, check for wildcard pattern matches if deployment is None: @@ -9210,9 +8329,9 @@ class Router: return None # Get basic credentials - credentials = CredentialLiteLLMParams( - **deployment.litellm_params.model_dump(exclude_none=True) - ).model_dump(exclude_none=True) + credentials = CredentialLiteLLMParams(**deployment.litellm_params.model_dump(exclude_none=True)).model_dump( + exclude_none=True + ) # Resolve litellm_credential_name to actual credentials if deployment.litellm_params.litellm_credential_name is not None: @@ -9229,14 +8348,10 @@ class Router: # Add custom_llm_provider if deployment.litellm_params.custom_llm_provider: - credentials["custom_llm_provider"] = ( - deployment.litellm_params.custom_llm_provider - ) + credentials["custom_llm_provider"] = deployment.litellm_params.custom_llm_provider elif "/" in deployment.litellm_params.model: # Extract provider from "provider/model" format - credentials["custom_llm_provider"] = deployment.litellm_params.model.split( - "/" - )[0] + credentials["custom_llm_provider"] = deployment.litellm_params.model.split("/")[0] else: credentials["custom_llm_provider"] = "openai" # default @@ -9252,9 +8367,7 @@ class Router: pass @overload - def get_router_model_info( - self, deployment: None, received_model_name: str, id: str - ) -> ModelMapInfo: + def get_router_model_info(self, deployment: None, received_model_name: str, id: str) -> ModelMapInfo: pass def get_router_model_info( @@ -9287,9 +8400,7 @@ class Router: ## GET BASE MODEL base_model = (deployment.get("model_info") or {}).get("base_model", None) if base_model is None: - base_model = (deployment.get("litellm_params") or {}).get( - "base_model", None - ) + base_model = (deployment.get("litellm_params") or {}).get("base_model", None) model = base_model @@ -9327,9 +8438,7 @@ class Router: if (potential_model.get("model_info") or {}).get("id") == ( deployment.get("model_info") or {} ).get("id"): - model = ( - potential_model.get("litellm_params") or {} - ).get("model") + model = (potential_model.get("litellm_params") or {}).get("model") break except Exception: pass @@ -9383,9 +8492,7 @@ class Router: model_name = model_info["model_name"] return self.get_model_list(model_name=model_name) - def get_deployment_model_info( - self, model_id: str, model_name: str - ) -> Optional[ModelInfo]: + def get_deployment_model_info(self, model_id: str, model_name: str) -> Optional[ModelInfo]: """ For a given model id, return the model info @@ -9447,9 +8554,7 @@ class Router: return model_info - def _set_model_group_info( - self, model_group: str, user_facing_model_group_name: str - ) -> Optional[ModelGroupInfo]: + def _set_model_group_info(self, model_group: str, user_facing_model_group_name: str) -> Optional[ModelGroupInfo]: """ For a given model group name, return the combined model info @@ -9467,14 +8572,9 @@ class Router: return None for model in model_list: is_match = False - if ( - "model_name" in model and model["model_name"] == model_group - ): # exact match + if "model_name" in model and model["model_name"] == model_group: # exact match is_match = True - elif ( - "model_name" in model - and self.pattern_router.route(model_group) is not None - ): # wildcard model + elif "model_name" in model and self.pattern_router.route(model_group) is not None: # wildcard model is_match = True if not is_match: @@ -9482,9 +8582,7 @@ class Router: # model in model group found # litellm_params = LiteLLM_Params(**model["litellm_params"]) # type: ignore # get configurable clientside auth params - configurable_clientside_auth_params = ( - litellm_params.configurable_clientside_auth_params - ) + configurable_clientside_auth_params = litellm_params.configurable_clientside_auth_params # Cache nested dict access to avoid repeated temporary dict allocations model_litellm_params = model.get("litellm_params", {}) @@ -9512,9 +8610,7 @@ class Router: try: model_id = model_info_dict.get("id", None) if model_id is not None: - model_info = self.get_deployment_model_info( - model_id=model_id, model_name=litellm_params.model - ) + model_info = self.get_deployment_model_info(model_id=model_id, model_name=litellm_params.model) else: model_info = None except Exception: @@ -9528,9 +8624,7 @@ class Router: custom_llm_provider=litellm_params.custom_llm_provider, ) except litellm.exceptions.BadRequestError as e: - verbose_router_logger.error( - "litellm.router.py::get_model_group_info() - {}".format(str(e)) - ) + verbose_router_logger.error("litellm.router.py::get_model_group_info() - {}".format(str(e))) if model_info is None: supported_openai_params = litellm.get_supported_openai_params( @@ -9581,8 +8675,7 @@ class Router: and model_info["max_input_tokens"] is not None and ( model_group_info.max_input_tokens is None - or model_info["max_input_tokens"] - > model_group_info.max_input_tokens + or model_info["max_input_tokens"] > model_group_info.max_input_tokens ) ): model_group_info.max_input_tokens = model_info["max_input_tokens"] @@ -9591,36 +8684,27 @@ class Router: and model_info["max_output_tokens"] is not None and ( model_group_info.max_output_tokens is None - or model_info["max_output_tokens"] - > model_group_info.max_output_tokens + or model_info["max_output_tokens"] > model_group_info.max_output_tokens ) ): model_group_info.max_output_tokens = model_info["max_output_tokens"] if model_info.get("input_cost_per_token", None) is not None and ( model_group_info.input_cost_per_token is None - or (model_info["input_cost_per_token"] or 0.0) - > (model_group_info.input_cost_per_token or 0.0) + or (model_info["input_cost_per_token"] or 0.0) > (model_group_info.input_cost_per_token or 0.0) ): - model_group_info.input_cost_per_token = model_info[ - "input_cost_per_token" - ] + model_group_info.input_cost_per_token = model_info["input_cost_per_token"] if model_info.get("output_cost_per_token", None) is not None and ( model_group_info.output_cost_per_token is None - or (model_info["output_cost_per_token"] or 0.0) - > (model_group_info.output_cost_per_token or 0.0) + or (model_info["output_cost_per_token"] or 0.0) > (model_group_info.output_cost_per_token or 0.0) ): - model_group_info.output_cost_per_token = model_info[ - "output_cost_per_token" - ] + model_group_info.output_cost_per_token = model_info["output_cost_per_token"] if ( - model_info.get("supports_parallel_function_calling", None) - is not None + model_info.get("supports_parallel_function_calling", None) is not None and model_info["supports_parallel_function_calling"] is True # type: ignore ): model_group_info.supports_parallel_function_calling = True if ( - model_info.get("supports_vision", None) is not None - and model_info["supports_vision"] is True # type: ignore + model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True # type: ignore ): model_group_info.supports_vision = True if ( @@ -9640,17 +8724,14 @@ class Router: model_group_info.supports_url_context = True if ( - model_info.get("supports_reasoning", None) is not None - and model_info["supports_reasoning"] is True # type: ignore + model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True # type: ignore ): model_group_info.supports_reasoning = True if ( model_info.get("supported_openai_params", None) is not None and model_info["supported_openai_params"] is not None ): - model_group_info.supported_openai_params = model_info[ - "supported_openai_params" - ] + model_group_info.supported_openai_params = model_info["supported_openai_params"] if model_info.get("tpm", None) is not None and _deployment_tpm is None: _deployment_tpm = model_info.get("tpm") if model_info.get("rpm", None) is not None and _deployment_rpm is None: @@ -9675,9 +8756,7 @@ class Router: ## UPDATE WITH CONFIGURABLE CLIENTSIDE AUTH PARAMS FOR MODEL GROUP if configurable_clientside_auth_params is not None: - model_group_info.configurable_clientside_auth_params = ( - configurable_clientside_auth_params - ) + model_group_info.configurable_clientside_auth_params = configurable_clientside_auth_params return model_group_info @@ -9708,13 +8787,9 @@ class Router: ) ## Check if actual model - return self._set_model_group_info( - model_group=model_group, user_facing_model_group_name=model_group - ) + return self._set_model_group_info(model_group=model_group, user_facing_model_group_name=model_group) - async def get_model_group_usage( - self, model_group: str - ) -> Tuple[Optional[int], Optional[int]]: + async def get_model_group_usage(self, model_group: str) -> Tuple[Optional[int], Optional[int]]: """ Returns current tpm/rpm usage for model group @@ -9725,9 +8800,7 @@ class Router: - usage: Tuple[tpm, rpm] """ dt = get_utc_datetime() - current_minute = dt.strftime( - "%H-%M" - ) # use the same timezone regardless of system clock + current_minute = dt.strftime("%H-%M") # use the same timezone regardless of system clock tpm_keys: List[str] = [] rpm_keys: List[str] = [] @@ -9758,9 +8831,7 @@ class Router: ) combined_tpm_rpm_keys = tpm_keys + rpm_keys - combined_tpm_rpm_values = await self.cache.async_batch_get_cache( - keys=combined_tpm_rpm_keys - ) + combined_tpm_rpm_values = await self.cache.async_batch_get_cache(keys=combined_tpm_rpm_keys) if combined_tpm_rpm_values is None: return None, None @@ -9786,9 +8857,7 @@ class Router: return tpm_usage, rpm_usage @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) - def _cached_get_model_group_info( - self, model_group: str - ) -> Optional[ModelGroupInfo]: + def _cached_get_model_group_info(self, model_group: str) -> Optional[ModelGroupInfo]: """ Cached version of get_model_group_info, uses @lru_cache wrapper @@ -9816,14 +8885,10 @@ class Router: returned_dict = {} if tpm_limit is not None: - returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - ( - current_tpm or 0 - ) + returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0) returned_dict["x-ratelimit-limit-tokens"] = tpm_limit if rpm_limit is not None: - returned_dict["x-ratelimit-remaining-requests"] = rpm_limit - ( - current_rpm or 0 - ) + returned_dict["x-ratelimit-remaining-requests"] = rpm_limit - (current_rpm or 0) returned_dict["x-ratelimit-limit-requests"] = rpm_limit return returned_dict @@ -9858,50 +8923,30 @@ class Router: # Lift QualityRouter routing decision into response headers for # transparency. The decision is stashed in request_kwargs.metadata # by QualityRouter.async_pre_routing_hook. - metadata = ( - (request_kwargs.get("metadata") or {}) - if isinstance(request_kwargs, dict) - else {} - ) - decision = ( - metadata.get("quality_router_decision") - if isinstance(metadata, dict) - else None - ) + metadata = (request_kwargs.get("metadata") or {}) if isinstance(request_kwargs, dict) else {} + decision = metadata.get("quality_router_decision") if isinstance(metadata, dict) else None if isinstance(decision, dict): # Only emit headers for fields that have a meaningful value. # `complexity_tier` and `matched_keyword` are mutually exclusive # (the keyword path short-circuits classification), so each # request emits one or the other but not both. if decision.get("routed_model") is not None: - additional_headers["x-litellm-quality-router-model"] = str( - decision["routed_model"] - ) + additional_headers["x-litellm-quality-router-model"] = str(decision["routed_model"]) if decision.get("quality_tier") is not None: - additional_headers["x-litellm-quality-router-tier"] = str( - decision["quality_tier"] - ) + additional_headers["x-litellm-quality-router-tier"] = str(decision["quality_tier"]) if decision.get("routed_via") is not None: - additional_headers["x-litellm-quality-router-via"] = str( - decision["routed_via"] - ) + additional_headers["x-litellm-quality-router-via"] = str(decision["routed_via"]) if decision.get("matched_keyword") is not None: - additional_headers["x-litellm-quality-router-keyword"] = str( - decision["matched_keyword"] - ) + additional_headers["x-litellm-quality-router-keyword"] = str(decision["matched_keyword"]) if decision.get("complexity_tier") is not None: - additional_headers["x-litellm-quality-router-complexity"] = str( - decision["complexity_tier"] - ) + additional_headers["x-litellm-quality-router-complexity"] = str(decision["complexity_tier"]) if ( "x-ratelimit-remaining-tokens" not in additional_headers and "x-ratelimit-remaining-requests" not in additional_headers and model_group is not None ): - remaining_usage = await self.get_remaining_model_group_usage( - model_group - ) + remaining_usage = await self.get_remaining_model_group_usage(model_group) # get_remaining_model_group_usage reads the router's TPM/RPM # counter, which is incremented post-response by @@ -9923,9 +8968,7 @@ class Router: for header, value in remaining_usage.items(): if value is not None: - additional_headers[header] = value - in_flight_delta.get( - header, 0 - ) + additional_headers[header] = value - in_flight_delta.get(header, 0) return response def _build_model_name_index(self, model_list: list) -> None: @@ -9974,9 +9017,7 @@ class Router: self._add_model_to_list_and_index_map(model=model, model_id=model_id) - def get_model_ids( - self, model_name: Optional[str] = None, exclude_team_models: bool = False - ) -> List[str]: + def get_model_ids(self, model_name: Optional[str] = None, exclude_team_models: bool = False) -> List[str]: """ if 'model_name' is none, returns all. @@ -10023,9 +9064,7 @@ class Router: """ return candidate_id in self.model_id_to_deployment_index_map - def resolve_model_name_from_model_id( - self, model_id: Optional[str] - ) -> Optional[str]: + def resolve_model_name_from_model_id(self, model_id: Optional[str]) -> Optional[str]: """ Resolve model_name from model_id. @@ -10077,9 +9116,7 @@ class Router: # No match found return None - def map_team_model( - self, team_model_name: Optional[str], team_id: str - ) -> Optional[str]: + def map_team_model(self, team_model_name: Optional[str], team_id: str) -> Optional[str]: """ Check if team_model_name resolves to team-specific deployments. @@ -10105,26 +9142,21 @@ class Router: # No model was specified (e.g. vector store endpoints). # Return the deployment's public model name so the router # can route to it and inject the BYOK API key. - return model.get("model_info", {}).get( - "team_public_model_name" - ) or model.get("model_name") + return model.get("model_info", {}).get("team_public_model_name") or model.get("model_name") return team_model_name # No team-scoped deployment found; wildcard/pattern routes are # handled downstream by the pattern_router in _common_checks_available_deployment. return None - def should_include_deployment( - self, model_name: str, model: dict, team_id: Optional[str] = None - ) -> bool: + def should_include_deployment(self, model_name: str, model: dict, team_id: Optional[str] = None) -> bool: """ Get the team-specific model name if team_id matches the deployment. """ if ( team_id is not None and (model.get("model_info") or {}).get("team_id") == team_id - and model_name - == (model.get("model_info") or {}).get("team_public_model_name") + and model_name == (model.get("model_info") or {}).get("team_public_model_name") ): return True elif model_name is not None and model["model_name"] == model_name: @@ -10171,9 +9203,7 @@ class Router: # O(k) where k = team deployments for this model_name (typically 1-10) for idx in indices: model = self.model_list[idx] - if not self.should_include_deployment( - model_name=model_name, model=model, team_id=team_id - ): + if not self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): continue if model_alias is not None: alias_model = model.copy() @@ -10191,9 +9221,7 @@ class Router: # O(k) where k = deployments for this model_name (typically 1-10) for idx in indices: model = self.model_list[idx] - if self.should_include_deployment( - model_name=model_name, model=model, team_id=team_id - ): + if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): if model_alias is not None: # Optimized: Use shallow copy since we only modify top-level model_name # This is much faster than deepcopy for nested dict structures @@ -10207,9 +9235,7 @@ class Router: # check if model_name matches any team_public_model_name # O(n) scan but only when team_id lookup fails for idx, model in enumerate(self.model_list): - if self.should_include_deployment( - model_name=model_name, model=model, team_id=team_id - ): + if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): if model_alias is not None: # Optimized: Use shallow copy since we only modify top-level model_name alias_model = model.copy() @@ -10233,9 +9259,7 @@ class Router: for deployment in deployments: model_info = deployment.get("model_info") if self._is_team_specific_model(model_info): - team_model_name = self._get_team_specific_model( - deployment=deployment, team_id=team_id - ) + team_model_name = self._get_team_specific_model(deployment=deployment, team_id=team_id) if team_model_name: model_names.append(team_model_name) else: @@ -10262,17 +9286,14 @@ class Router: blocked_by_name[name] = blocked_by_name[name] and is_blocked else: blocked_by_name[name] = is_blocked - return { - name for name, fully_blocked in blocked_by_name.items() if fully_blocked - } + return {name for name, fully_blocked in blocked_by_name.items() if fully_blocked} @staticmethod def _are_all_deployments_blocked( deployments: List[DeploymentTypedDict], ) -> bool: return len(deployments) > 0 and all( - (deployment.get("model_info") or {}).get("blocked") is True - for deployment in deployments + (deployment.get("model_info") or {}).get("blocked") is True for deployment in deployments ) def _is_model_fully_blocked(self, model: str) -> bool: @@ -10309,9 +9330,7 @@ class Router: """ if self.allowed_fails_policy is not None: return set() - unhealthy_ids = ( - await self.health_state_cache.async_get_unhealthy_deployment_ids() - ) + unhealthy_ids = await self.health_state_cache.async_get_unhealthy_deployment_ids() if not unhealthy_ids: return set() deployments = self.get_model_list() or [] @@ -10330,15 +9349,9 @@ class Router: unhealthy_by_name[name] = unhealthy_by_name[name] and is_unhealthy else: unhealthy_by_name[name] = is_unhealthy - return { - name - for name, fully_unhealthy in unhealthy_by_name.items() - if fully_unhealthy - } + return {name for name, fully_unhealthy in unhealthy_by_name.items() if fully_unhealthy} - def _get_team_specific_model( - self, deployment: DeploymentTypedDict, team_id: Optional[str] = None - ) -> Optional[str]: + def _get_team_specific_model(self, deployment: DeploymentTypedDict, team_id: Optional[str] = None) -> Optional[str]: """ Get the team-specific model name if team_id matches the deployment. @@ -10369,9 +9382,7 @@ class Router: """ return bool(model_info and model_info.get("team_id")) - def get_model_list_from_model_alias( - self, model_name: Optional[str] = None - ) -> List[DeploymentTypedDict]: + def get_model_list_from_model_alias(self, model_name: Optional[str] = None) -> List[DeploymentTypedDict]: """ Helper function to get model list from model alias. @@ -10399,11 +9410,7 @@ class Router: else: continue - returned_models.extend( - self._get_all_deployments( - model_name=_router_model_name, model_alias=model_alias - ) - ) + returned_models.extend(self._get_all_deployments(model_name=_router_model_name, model_alias=model_alias)) return returned_models @@ -10420,22 +9427,16 @@ class Router: returned_models: List[DeploymentTypedDict] = [] if model_name is not None: - returned_models.extend( - self._get_all_deployments(model_name=model_name, team_id=team_id) - ) + returned_models.extend(self._get_all_deployments(model_name=model_name, team_id=team_id)) - returned_models.extend( - self.get_model_list_from_model_alias(model_name=model_name) - ) + returned_models.extend(self.get_model_list_from_model_alias(model_name=model_name)) if len(returned_models) == 0: # check if wildcard route potential_wildcard_models = self.pattern_router.route(model_name) or [] ## check for team-specific wildcard models if team_id is not None and team_id in self.team_pattern_routers: - potential_team_only_wildcard_models = ( - self.team_pattern_routers[team_id].route(model_name) or [] - ) + potential_team_only_wildcard_models = self.team_pattern_routers[team_id].route(model_name) or [] potential_wildcard_models.extend(potential_team_only_wildcard_models) if model_name is not None and potential_wildcard_models is not None: @@ -10483,9 +9484,7 @@ class Router: - team_id: Optional[str] - the team id, to resolve team-specific models """ # Check if this is the no-args hot path (cacheable) - _use_cache = ( - model_name is None and model_access_group is None and team_id is None - ) + _use_cache = model_name is None and model_access_group is None and team_id is None # Return cached result for the no-args hot path if _use_cache and self._access_groups_cache is not None: @@ -10516,16 +9515,12 @@ class Router: return access_groups - def _is_model_access_group_for_wildcard_route( - self, model_access_group: str - ) -> bool: + def _is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: """ Return True if model access group is a wildcard route """ # GET ACCESS GROUPS - access_groups = self.get_model_access_groups( - model_access_group=model_access_group - ) + access_groups = self.get_model_access_groups(model_access_group=model_access_group) if len(access_groups) == 0: return False @@ -10573,9 +9568,7 @@ class Router: ): _settings_to_return[var] = self.lowestlatency_logger.routing_args.json() - _settings_to_return["routing_groups"] = [ - group.model_dump() for group in self._routing_groups.values() - ] + _settings_to_return["routing_groups"] = [group.model_dump() for group in self._routing_groups.values()] return _settings_to_return def update_settings(self, **kwargs): @@ -10637,9 +9630,7 @@ class Router: else: self.routing_strategy_init( routing_strategy=value, - routing_strategy_args=kwargs.get( - "routing_strategy_args", {} - ), + routing_strategy_args=kwargs.get("routing_strategy_args", {}), ) rebuild_routing_groups = True elif var == "routing_strategy_args": @@ -10648,10 +9639,7 @@ class Router: else: verbose_router_logger.debug("Setting {} is not allowed".format(var)) - if ( - relink_lar1_from_args - and self._normalize_strategy(self.routing_strategy) == "lar1" - ): + if relink_lar1_from_args and self._normalize_strategy(self.routing_strategy) == "lar1": from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy apply_lar1_routing_strategy(self, self.routing_strategy_args) @@ -10676,42 +9664,28 @@ class Router: parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs(kwargs) if client_type == "max_parallel_requests": cache_key = "{}_max_parallel_requests_client".format(model_id) - client = self.cache.get_cache( - key=cache_key, local_only=True, parent_otel_span=parent_otel_span - ) + client = self.cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) if client is None: - InitalizeCachedClient.set_max_parallel_requests_client( - litellm_router_instance=self, model=deployment - ) - client = self.cache.get_cache( - key=cache_key, local_only=True, parent_otel_span=parent_otel_span - ) + InitalizeCachedClient.set_max_parallel_requests_client(litellm_router_instance=self, model=deployment) + client = self.cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) return client elif client_type == "async": if kwargs.get("stream") is True: cache_key = f"{model_id}_stream_async_client" - client = self.cache.get_cache( - key=cache_key, local_only=True, parent_otel_span=parent_otel_span - ) + client = self.cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) return client else: cache_key = f"{model_id}_async_client" - client = self.cache.get_cache( - key=cache_key, local_only=True, parent_otel_span=parent_otel_span - ) + client = self.cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) return client else: if kwargs.get("stream") is True: cache_key = f"{model_id}_stream_client" - client = self.cache.get_cache( - key=cache_key, parent_otel_span=parent_otel_span - ) + client = self.cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) return client else: cache_key = f"{model_id}_client" - client = self.cache.get_cache( - key=cache_key, parent_otel_span=parent_otel_span - ) + client = self.cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) return client def _pre_call_checks( @@ -10730,9 +9704,7 @@ class Router: - [TODO] function call and model doesn't support function calling """ - verbose_router_logger.debug( - f"Starting Pre-call checks for deployments in model={model}" - ) + verbose_router_logger.debug(f"Starting Pre-call checks for deployments in model={model}") # Optimized: Use list() shallow copy instead of deepcopy # We only pop from the list, not modify deployment dicts - 100x+ faster on hot path (every request) @@ -10755,10 +9727,7 @@ class Router: current_minute = dt.strftime("%H-%M") rpm_key = f"{model}:rpm:{current_minute}" model_group_cache = ( - self.cache.get_cache( - key=rpm_key, local_only=True, parent_otel_span=parent_otel_span - ) - or {} + self.cache.get_cache(key=rpm_key, local_only=True, parent_otel_span=parent_otel_span) or {} ) # check the in-memory cache used by lowest_latency and usage-based routing. Only check the local cache. for idx, deployment in enumerate(_returned_deployments): # Cache nested dict access to avoid repeated temporary dict allocations @@ -10771,16 +9740,10 @@ class Router: base_model = _model_info.get("base_model", None) if base_model is None: base_model = _litellm_params.get("base_model", None) - model_info = self.get_router_model_info( - deployment=deployment, received_model_name=model - ) + model_info = self.get_router_model_info(deployment=deployment, received_model_name=model) _deployment_model = base_model or _litellm_params.get("model", None) - max_input_tokens = ( - model_info.get("max_input_tokens") - if isinstance(model_info, dict) - else None - ) + max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None if isinstance(max_input_tokens, int): if input_tokens is None: try: @@ -10795,12 +9758,10 @@ class Router: if input_tokens > max_input_tokens: invalid_model_indices.add(idx) _context_window_error = True - _potential_error_str += ( - "Model={}, Max Input Tokens={}, Got={}".format( - _deployment_model, - max_input_tokens, - input_tokens, - ) + _potential_error_str += "Model={}, Max Input Tokens={}, Got={}".format( + _deployment_model, + max_input_tokens, + input_tokens, ) continue except Exception as e: @@ -10810,39 +9771,22 @@ class Router: ## RPM CHECK ## ### get local router cache ### current_request_cache_local = ( - self.cache.get_cache( - key=model_id, local_only=True, parent_otel_span=parent_otel_span - ) - or 0 + self.cache.get_cache(key=model_id, local_only=True, parent_otel_span=parent_otel_span) or 0 ) ### get usage based cache ### - if ( - isinstance(model_group_cache, dict) - and self.routing_strategy != "usage-based-routing-v2" - ): + if isinstance(model_group_cache, dict) and self.routing_strategy != "usage-based-routing-v2": model_group_cache[model_id] = model_group_cache.get(model_id, 0) - current_request = max( - current_request_cache_local, model_group_cache[model_id] - ) + current_request = max(current_request_cache_local, model_group_cache[model_id]) - if ( - isinstance(_litellm_params, dict) - and _litellm_params.get("rpm", None) is not None - ): - if ( - isinstance(_litellm_params["rpm"], int) - and _litellm_params["rpm"] <= current_request - ): + if isinstance(_litellm_params, dict) and _litellm_params.get("rpm", None) is not None: + if isinstance(_litellm_params["rpm"], int) and _litellm_params["rpm"] <= current_request: invalid_model_indices.add(idx) _rate_limit_error = True continue ## REGION CHECK ## - if ( - request_kwargs is not None - and request_kwargs.get("allowed_model_region") is not None - ): + if request_kwargs is not None and request_kwargs.get("allowed_model_region") is not None: allowed_model_region = request_kwargs.get("allowed_model_region") if allowed_model_region is not None: @@ -10876,17 +9820,13 @@ class Router: continue else: # check the non-default openai params in request kwargs - non_default_params = litellm.utils.get_non_default_params( - passed_params=request_kwargs - ) + non_default_params = litellm.utils.get_non_default_params(passed_params=request_kwargs) special_params = ["response_format"] # check if all params are supported for k, v in non_default_params.items(): if k not in supported_openai_params and k in special_params: # if not -> invalid model - verbose_router_logger.debug( - f"INVALID MODEL INDEX @ REQUEST KWARG FILTERING, k={k}" - ) + verbose_router_logger.debug(f"INVALID MODEL INDEX @ REQUEST KWARG FILTERING, k={k}") invalid_model_indices.add(idx) if len(invalid_model_indices) == len(_returned_deployments): @@ -10911,11 +9851,7 @@ class Router: ) if len(invalid_model_indices) > 0: # Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops) - _returned_deployments = [ - d - for i, d in enumerate(_returned_deployments) - if i not in invalid_model_indices - ] + _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in invalid_model_indices] return _returned_deployments @@ -10957,9 +9893,7 @@ class Router: # This intentionally takes priority over team pattern routers below, # so that named team deployments shadow wildcard/pattern routes. if request_team_id is not None: - team_deployments = self._get_all_deployments( - model_name=model, team_id=request_team_id - ) + team_deployments = self._get_all_deployments(model_name=model, team_id=request_team_id) if team_deployments: return model, team_deployments @@ -10971,9 +9905,7 @@ class Router: return model, pattern_deployments if request_team_id is not None and request_team_id in self.team_pattern_routers: - pattern_deployments = self.team_pattern_routers[ - request_team_id - ].get_deployments_by_pattern( + pattern_deployments = self.team_pattern_routers[request_team_id].get_deployments_by_pattern( model=model, ) if pattern_deployments: @@ -10982,9 +9914,7 @@ class Router: if self.default_deployment is not None: # Shallow copy with nested litellm_params copy (100x+ faster than deepcopy) updated_deployment = self.default_deployment.copy() - updated_deployment["litellm_params"] = self.default_deployment[ - "litellm_params" - ].copy() + updated_deployment["litellm_params"] = self.default_deployment["litellm_params"].copy() updated_deployment["litellm_params"]["model"] = model return model, updated_deployment @@ -11013,9 +9943,7 @@ class Router: if request_kwargs is not None: metadata = request_kwargs.get("metadata") or {} litellm_metadata = request_kwargs.get("litellm_metadata") or {} - request_team_id = metadata.get( - "user_api_key_team_id" - ) or litellm_metadata.get("user_api_key_team_id") + request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") # check if aliases set on litellm model alias map if specific_deployment is True: return model, self._get_deployment_by_litellm_model(model=model) @@ -11032,17 +9960,13 @@ class Router: if _model_from_alias is not None: model = _model_from_alias - early = self._try_early_resolve_deployments_for_model_not_in_names( - model=model, request_team_id=request_team_id - ) + early = self._try_early_resolve_deployments_for_model_not_in_names(model=model, request_team_id=request_team_id) if early is not None: return early ## get healthy deployments ### get all deployments - healthy_deployments = self._get_all_deployments( - model_name=model, team_id=request_team_id - ) + healthy_deployments = self._get_all_deployments(model_name=model, team_id=request_team_id) _pre_model_access_group_filter_len = len(healthy_deployments) healthy_deployments = self._filter_deployments_by_model_access_groups( model=model, @@ -11059,9 +9983,7 @@ class Router: # Do not fall back when access-group filtering removed every candidate; # _get_deployment_by_litellm_model does not re-apply that filter. if _pre_model_access_group_filter_len == 0: - _litellm_model_deployments = self._get_deployment_by_litellm_model( - model=model - ) + _litellm_model_deployments = self._get_deployment_by_litellm_model(model=model) healthy_deployments = self._filter_deployments_by_model_access_groups( model=model, healthy_deployments=_litellm_model_deployments, @@ -11073,26 +9995,18 @@ class Router: # being emptied: prevent default-model fallback from bypassing the # restriction (the fallback model may have no access_groups and # would short-circuit the filter). - if ( - len(_litellm_model_deployments) > 0 - and len(healthy_deployments) == 0 - ): + if len(_litellm_model_deployments) > 0 and len(healthy_deployments) == 0: _access_group_filter_emptied_candidates = True if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug( - f"initial list of deployments: {healthy_deployments}" - ) + verbose_router_logger.debug(f"initial list of deployments: {healthy_deployments}") if len(healthy_deployments) == 0: # Check for default fallbacks if no deployments are found for the requested model # Do not fall back to another model when access-group filtering removed every # candidate for the requested name: re-filtering the fallback model can be a # no-op when it has no access_groups, incorrectly serving a different model. - if ( - self._has_default_fallbacks() - and not _access_group_filter_emptied_candidates - ): + if self._has_default_fallbacks() and not _access_group_filter_emptied_candidates: fallback_model = self._get_first_default_fallback() if fallback_model: verbose_router_logger.info( @@ -11100,24 +10014,18 @@ class Router: ) # Re-assign model to the fallback and try to get deployments again model = fallback_model - healthy_deployments = self._get_all_deployments( - model_name=model, team_id=request_team_id - ) - healthy_deployments = ( - self._filter_deployments_by_model_access_groups( - model=model, - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - request_team_id=request_team_id, - ) + healthy_deployments = self._get_all_deployments(model_name=model, team_id=request_team_id) + healthy_deployments = self._filter_deployments_by_model_access_groups( + model=model, + healthy_deployments=healthy_deployments, + request_kwargs=request_kwargs, + request_team_id=request_team_id, ) # If still no deployments after checking for fallbacks, raise an error if len(healthy_deployments) == 0: if self.get_model_list(model_name=model) is None: - message = f"You passed in model={model}. There is no 'model_name' with this string".format( - model - ) + message = f"You passed in model={model}. There is no 'model_name' with this string".format(model) else: message = f"You passed in model={model}. There are no healthy deployments for this model".format( model @@ -11156,9 +10064,7 @@ class Router: metadata = request_kwargs.get("metadata") or {} litellm_metadata = request_kwargs.get("litellm_metadata") or {} - user_api_key_auth = metadata.get("user_api_key_auth") or litellm_metadata.get( - "user_api_key_auth" - ) + user_api_key_auth = metadata.get("user_api_key_auth") or litellm_metadata.get("user_api_key_auth") if user_api_key_auth is None: return healthy_deployments @@ -11170,16 +10076,10 @@ class Router: # If caller has direct model/wildcard/all-proxy access, do not constrain # deployment choice by access group. - if ( - model in allowed_models - or "*" in allowed_models - or "all-proxy-models" in allowed_models - ): + if model in allowed_models or "*" in allowed_models or "all-proxy-models" in allowed_models: return healthy_deployments - access_groups_for_model = self.get_model_access_groups( - model_name=model, team_id=request_team_id - ) + access_groups_for_model = self.get_model_access_groups(model_name=model, team_id=request_team_id) if len(access_groups_for_model) == 0: return healthy_deployments @@ -11192,9 +10092,7 @@ class Router: filtered_deployments = [] for deployment in healthy_deployments: deployment_model_info = deployment.get("model_info") or {} - deployment_access_groups = set( - deployment_model_info.get("access_groups", []) or [] - ) + deployment_access_groups = set(deployment_model_info.get("access_groups", []) or []) if deployment_access_groups & allowed_access_groups: filtered_deployments.append(deployment) @@ -11234,9 +10132,7 @@ class Router: ) if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug( - f"healthy_deployments after team filter: {healthy_deployments}" - ) + verbose_router_logger.debug(f"healthy_deployments after team filter: {healthy_deployments}") healthy_deployments = filter_web_search_deployments( healthy_deployments=healthy_deployments, @@ -11244,9 +10140,7 @@ class Router: ) if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug( - f"healthy_deployments after web search filter: {healthy_deployments}" - ) + verbose_router_logger.debug(f"healthy_deployments after web search filter: {healthy_deployments}") if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -11258,11 +10152,9 @@ class Router: return healthy_deployments # Health-check-based filtering (before cooldown) - healthy_deployments = ( - await self._async_filter_health_check_unhealthy_deployments( - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - ) + healthy_deployments = await self._async_filter_health_check_unhealthy_deployments( + healthy_deployments=healthy_deployments, + parent_otel_span=parent_otel_span, ) cooldown_deployments = await _async_get_cooldown_deployments( @@ -11278,11 +10170,7 @@ class Router: # Safety net: only bypass cooldown filter when health-check routing is # driving cooldown (i.e. allowed_fails_policy is set). Without a policy, # cooldowns are from real request failures and must not be bypassed. - if ( - not healthy_deployments - and self.enable_health_check_routing - and self.allowed_fails_policy is not None - ): + if not healthy_deployments and self.enable_health_check_routing and self.allowed_fails_policy is not None: verbose_router_logger.warning( "All deployments in cooldown via health-check routing, bypassing cooldown filter" ) @@ -11293,9 +10181,7 @@ class Router: healthy_deployments = await self.async_callback_filter_deployments( model=model, healthy_deployments=healthy_deployments, - messages=( - cast(List[AllMessageValues], messages) if messages is not None else None - ), + messages=(cast(List[AllMessageValues], messages) if messages is not None else None), request_kwargs=request_kwargs, parent_otel_span=parent_otel_span, ) @@ -11313,9 +10199,7 @@ class Router: model=model, request_kwargs=request_kwargs, healthy_deployments=healthy_deployments, - metadata_variable_name=self._get_metadata_variable_name_from_kwargs( - request_kwargs - ), + metadata_variable_name=self._get_metadata_variable_name_from_kwargs(request_kwargs), ) ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) @@ -11327,9 +10211,7 @@ class Router: ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in ## this request via weighted-failover. Always honored, regardless of the ## router-level flag, so a stale exclusion key on kwargs cannot escape. - _excluded_deployment_ids = (request_kwargs or {}).pop( - "_excluded_deployment_ids", None - ) + _excluded_deployment_ids = (request_kwargs or {}).pop("_excluded_deployment_ids", None) healthy_deployments = litellm.utils._get_excluded_filtered_deployments( cast(List[Dict], healthy_deployments), excluded_deployment_ids=_excluded_deployment_ids, @@ -11408,10 +10290,7 @@ class Router: return healthy_deployments # When encrypted content affinity pins to a specific deployment, - if ( - request_kwargs.get("_encrypted_content_affinity_pinned") - and len(healthy_deployments) == 1 - ): + if request_kwargs.get("_encrypted_content_affinity_pinned") and len(healthy_deployments) == 1: return healthy_deployments[0] start_time = time.time() @@ -11530,9 +10409,7 @@ class Router: ) # 4. Filter deployments that support pass-through - pass_through_deployments = self._filter_pass_through_deployments( - healthy_deployments=healthy_deployments - ) + pass_through_deployments = self._filter_pass_through_deployments(healthy_deployments=healthy_deployments) if len(pass_through_deployments) == 0: raise litellm.BadRequestError( @@ -11695,9 +10572,7 @@ class Router: ) return healthy_deployments - parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( - request_kwargs - ) + parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs(request_kwargs) # Health-check-based filtering (before cooldown) healthy_deployments = self._filter_health_check_unhealthy_deployments( @@ -11713,11 +10588,7 @@ class Router: healthy_deployments=healthy_deployments, cooldown_deployments=cooldown_deployments, ) - if ( - not healthy_deployments - and self.enable_health_check_routing - and self.allowed_fails_policy is not None - ): + if not healthy_deployments and self.enable_health_check_routing and self.allowed_fails_policy is not None: verbose_router_logger.warning( "All deployments in cooldown via health-check routing, bypassing cooldown filter" ) @@ -11743,9 +10614,7 @@ class Router: ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in ## this request via weighted-failover. See async counterpart in ## async_get_healthy_deployments for details. - _excluded_deployment_ids = (request_kwargs or {}).pop( - "_excluded_deployment_ids", None - ) + _excluded_deployment_ids = (request_kwargs or {}).pop("_excluded_deployment_ids", None) healthy_deployments = litellm.utils._get_excluded_filtered_deployments( healthy_deployments, excluded_deployment_ids=_excluded_deployment_ids, @@ -11756,9 +10625,7 @@ class Router: _cooldown_time = self.cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = _get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) + _cooldown_list = _get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, @@ -11786,16 +10653,12 @@ class Router: ) if deployment is None: - verbose_router_logger.info( - f"get_available_deployment for model: {model}, No deployment available" - ) + verbose_router_logger.info(f"get_available_deployment for model: {model}, No deployment available") model_ids = self.get_model_ids(model_name=model) _cooldown_time = self.cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = _get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) + _cooldown_list = _get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, @@ -11862,9 +10725,7 @@ class Router: ) # 3. Filter deployments that support pass-through - pass_through_deployments = self._filter_pass_through_deployments( - healthy_deployments=healthy_deployments - ) + pass_through_deployments = self._filter_pass_through_deployments(healthy_deployments=healthy_deployments) if len(pass_through_deployments) == 0: # No deployments support pass-through @@ -11875,9 +10736,7 @@ class Router: ) # 4. Apply health-check and cooldown filtering - parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( - request_kwargs - ) + parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs(request_kwargs) pass_through_deployments = self._filter_health_check_unhealthy_deployments( healthy_deployments=pass_through_deployments, parent_otel_span=parent_otel_span, @@ -11889,9 +10748,7 @@ class Router: healthy_deployments=pass_through_deployments, cooldown_deployments=cooldown_deployments, ) - pass_through_deployments = self._filter_blocked_deployments( - pass_through_deployments - ) + pass_through_deployments = self._filter_blocked_deployments(pass_through_deployments) # 5. Apply pre-call checks (if enabled) if self.enable_pre_call_checks and messages is not None: @@ -11907,9 +10764,7 @@ class Router: _cooldown_time = self.cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = _get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) + _cooldown_list = _get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, @@ -11943,9 +10798,7 @@ class Router: _cooldown_time = self.cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span ) - _cooldown_list = _get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) + _cooldown_list = _get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) raise RouterRateLimitError( model=model, cooldown_time=_cooldown_time, @@ -11975,15 +10828,9 @@ class Router: verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") # Convert to set for O(1) lookup and use list comprehension for O(n) filtering cooldown_set = set(cooldown_deployments) - return [ - deployment - for deployment in healthy_deployments - if deployment["model_info"]["id"] not in cooldown_set - ] + return [deployment for deployment in healthy_deployments if deployment["model_info"]["id"] not in cooldown_set] - def _filter_blocked_deployments( - self, healthy_deployments: List[Dict] - ) -> List[Dict]: + def _filter_blocked_deployments(self, healthy_deployments: List[Dict]) -> List[Dict]: """ Filters out deployments that an admin has paused via `LiteLLM_ProxyModelTable.blocked`. @@ -12031,22 +10878,16 @@ class Router: if self.allowed_fails_policy is not None: return healthy_deployments - unhealthy_ids = ( - await self.health_state_cache.async_get_unhealthy_deployment_ids( - parent_otel_span=parent_otel_span - ) + unhealthy_ids = await self.health_state_cache.async_get_unhealthy_deployment_ids( + parent_otel_span=parent_otel_span ) if not unhealthy_ids: return healthy_deployments - filtered = [ - d for d in healthy_deployments if d["model_info"]["id"] not in unhealthy_ids - ] + filtered = [d for d in healthy_deployments if d["model_info"]["id"] not in unhealthy_ids] if not filtered: - verbose_router_logger.warning( - "All deployments marked unhealthy by health checks, bypassing health filter" - ) + verbose_router_logger.warning("All deployments marked unhealthy by health checks, bypassing health filter") return healthy_deployments return filtered @@ -12063,27 +10904,19 @@ class Router: if self.allowed_fails_policy is not None: return healthy_deployments - unhealthy_ids = self.health_state_cache.get_unhealthy_deployment_ids( - parent_otel_span=parent_otel_span - ) + unhealthy_ids = self.health_state_cache.get_unhealthy_deployment_ids(parent_otel_span=parent_otel_span) if not unhealthy_ids: return healthy_deployments - filtered = [ - d for d in healthy_deployments if d["model_info"]["id"] not in unhealthy_ids - ] + filtered = [d for d in healthy_deployments if d["model_info"]["id"] not in unhealthy_ids] if not filtered: - verbose_router_logger.warning( - "All deployments marked unhealthy by health checks, bypassing health filter" - ) + verbose_router_logger.warning("All deployments marked unhealthy by health checks, bypassing health filter") return healthy_deployments return filtered - def _filter_pass_through_deployments( - self, healthy_deployments: List[Dict] - ) -> List[Dict]: + def _filter_pass_through_deployments(self, healthy_deployments: List[Dict]) -> List[Dict]: """ Filter out deployments configured with use_in_pass_through=True @@ -12103,15 +10936,11 @@ class Router: if deployment.get("litellm_params", {}).get("use_in_pass_through", False) ] - verbose_router_logger.debug( - f"Found {len(pass_through_deployments)} deployments with pass-through enabled" - ) + verbose_router_logger.debug(f"Found {len(pass_through_deployments)} deployments with pass-through enabled") return pass_through_deployments - def _track_deployment_metrics( - self, deployment, parent_otel_span: Optional[Span], response=None - ): + def _track_deployment_metrics(self, deployment, parent_otel_span: Optional[Span], response=None): """ Tracks successful requests rpm usage. """ @@ -12120,15 +10949,11 @@ class Router: if response is None: # update self.deployment_stats if model_id is not None: - self._update_usage( - model_id, parent_otel_span - ) # update in-memory cache for tracking + self._update_usage(model_id, parent_otel_span) # update in-memory cache for tracking except Exception as e: verbose_router_logger.error(f"Error in _track_deployment_metrics: {str(e)}") - def get_num_retries_from_retry_policy( - self, exception: Exception, model_group: Optional[str] = None - ): + def get_num_retries_from_retry_policy(self, exception: Exception, model_group: Optional[str] = None): return _get_num_retries_from_retry_policy( exception=exception, model_group=model_group, @@ -12155,10 +10980,7 @@ class Router: and allowed_fails_policy.AuthenticationErrorAllowedFails is not None ): return allowed_fails_policy.AuthenticationErrorAllowedFails - if ( - isinstance(exception, litellm.Timeout) - and allowed_fails_policy.TimeoutErrorAllowedFails is not None - ): + if isinstance(exception, litellm.Timeout) and allowed_fails_policy.TimeoutErrorAllowedFails is not None: return allowed_fails_policy.TimeoutErrorAllowedFails if ( isinstance(exception, litellm.RateLimitError) @@ -12196,13 +11018,9 @@ class Router: litellm.logging_callback_manager.add_litellm_success_callback( _slack_alerting_logger.response_taking_too_long_callback ) - verbose_router_logger.info( - "\033[94m\nInitialized Alerting for litellm.Router\033[0m\n" - ) + verbose_router_logger.info("\033[94m\nInitialized Alerting for litellm.Router\033[0m\n") - def set_custom_routing_strategy( - self, CustomRoutingStrategy: CustomRoutingStrategyBase - ): + def set_custom_routing_strategy(self, CustomRoutingStrategy: CustomRoutingStrategyBase): """ Sets get_available_deployment and async_get_available_deployment on an instanced of litellm.Router diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 4856d7ff4cd..69d6a019e68 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -126,9 +126,7 @@ class AdaptiveRouter: continue if row.model_name not in self.config.available_models: continue - self._cells[(rt, row.model_name)] = BanditCell( - alpha=row.alpha, beta=row.beta - ) + self._cells[(rt, row.model_name)] = BanditCell(alpha=row.alpha, beta=row.beta) loaded += 1 verbose_router_logger.info( "AdaptiveRouter[%s]: loaded %d cells from DB", @@ -165,15 +163,11 @@ class AdaptiveRouter: attribution is enforced post-call via the owner cache (see `claim_or_check_owner`). """ - user_text = ( - get_last_user_message(cast(List[AllMessageValues], messages or [])) or "" - ) + user_text = get_last_user_message(cast(List[AllMessageValues], messages or [])) or "" request_type = classify_prompt(user_text) min_quality_tier = self._extract_min_quality_tier(request_kwargs) - chosen_model = await self.pick_model( - request_type=request_type, min_quality_tier=min_quality_tier - ) + chosen_model = await self.pick_model(request_type=request_type, min_quality_tier=min_quality_tier) verbose_router_logger.debug( "AdaptiveRouter[%s]: classified=%s -> chose %s", self.router_name, @@ -201,10 +195,7 @@ class AdaptiveRouter: """Thompson-sample across eligible models. Stateless per-turn.""" eligible = self._eligible_models(min_quality_tier) if not eligible: - raise ValueError( - f"AdaptiveRouter[{self.router_name}]: no models meet " - f"min_quality_tier={min_quality_tier}" - ) + raise ValueError(f"AdaptiveRouter[{self.router_name}]: no models meet min_quality_tier={min_quality_tier}") cells = {m: self._cells[(request_type, m)] for m in eligible} costs = {m: self.model_to_cost.get(m, 0.0) for m in eligible} @@ -255,9 +246,7 @@ class AdaptiveRouter: async def get_state_snapshot(self) -> Dict[str, Any]: """In-memory snapshot for the introspection endpoint. Cheap; no DB hit.""" cells = [] - for (rt, model), cell in sorted( - self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1]) - ): + for (rt, model), cell in sorted(self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])): total = cell.alpha + cell.beta cells.append( { @@ -327,8 +316,7 @@ class AdaptiveRouter: return [ m for m in self.config.available_models - if (self.model_to_prefs.get(m) or _default_prefs()).quality_tier - >= min_quality_tier + if (self.model_to_prefs.get(m) or _default_prefs()).quality_tier >= min_quality_tier ] # ---- Session state --------------------------------------------------- @@ -378,9 +366,7 @@ class AdaptiveRouter: """Apply one turn, push session snapshot + bandit deltas to the queue.""" state = self.get_or_create_session_state(session_id, model_name, request_type) delta = apply_turn(state, turn) - verbose_router_logger.debug( - "AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta - ) + verbose_router_logger.debug("AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta) # Strip the raw conversation content before persisting. The # last_user/assistant_content and tool_call_history fields are only @@ -396,9 +382,7 @@ class AdaptiveRouter: "pending_tool_calls", ): snapshot.pop(sensitive, None) - await self.queue.add_session_state( - session_id, self.router_name, model_name, snapshot - ) + await self.queue.add_session_state(session_id, self.router_name, model_name, snapshot) d_alpha, d_beta = self._compute_bandit_delta(delta) verbose_router_logger.debug( @@ -414,9 +398,7 @@ class AdaptiveRouter: # back to the session's original type so closing pleasantries don't # misattribute the reward. attribution_type = ( - request_type - if request_type != RequestType.GENERAL - else RequestType(state.classified_type) + request_type if request_type != RequestType.GENERAL else RequestType(state.classified_type) ) cell_key = (attribution_type, model_name) self._cells[cell_key] = apply_delta(self._cells[cell_key], d_alpha, d_beta) @@ -443,13 +425,5 @@ class AdaptiveRouter: - exhaustion -> 0 (uptime issue, tracked separately later) """ d_alpha = float(delta.satisfaction) - d_beta = ( - float( - delta.misalignment - + delta.stagnation - + delta.disengagement - + delta.failure - ) - + 0.5 * delta.loop - ) + d_beta = float(delta.misalignment + delta.stagnation + delta.disengagement + delta.failure) + 0.5 * delta.loop return d_alpha, d_beta diff --git a/litellm/router_strategy/adaptive_router/bandit.py b/litellm/router_strategy/adaptive_router/bandit.py index 1ab96f0e952..4c914e00826 100644 --- a/litellm/router_strategy/adaptive_router/bandit.py +++ b/litellm/router_strategy/adaptive_router/bandit.py @@ -42,9 +42,7 @@ class BanditCell: return max(0, int(self.alpha + self.beta - COLD_START_MASS)) -def initial_cell( - prefs: AdaptiveRouterPreferences, request_type: RequestType -) -> BanditCell: +def initial_cell(prefs: AdaptiveRouterPreferences, request_type: RequestType) -> BanditCell: """ Cold-start prior for a (model, request_type) cell. @@ -54,10 +52,7 @@ def initial_cell( """ if prefs.quality_tier not in BASE_TIER_WEIGHT: valid = sorted(BASE_TIER_WEIGHT) - raise ValueError( - f"quality_tier={prefs.quality_tier} is not supported; " - f"valid tiers are {valid}" - ) + raise ValueError(f"quality_tier={prefs.quality_tier} is not supported; valid tiers are {valid}") base = BASE_TIER_WEIGHT[prefs.quality_tier] bonus = STRENGTH_BONUS if request_type in prefs.strengths else 0.0 mean = min(0.95, base + bonus) diff --git a/litellm/router_strategy/adaptive_router/classifier.py b/litellm/router_strategy/adaptive_router/classifier.py index 0434dfdb63f..1cd66fc7242 100644 --- a/litellm/router_strategy/adaptive_router/classifier.py +++ b/litellm/router_strategy/adaptive_router/classifier.py @@ -102,9 +102,7 @@ _RULES: List[Tuple[Pattern[str], RequestType]] = [ RequestType.WRITING, ), ( - re.compile( - r"^\s*(who|what|when|where|which)\s+(?:is|was|were|are)\b", re.IGNORECASE - ), + re.compile(r"^\s*(who|what|when|where|which)\s+(?:is|was|were|are)\b", re.IGNORECASE), RequestType.FACTUAL_LOOKUP, ), ( diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 99fe5e26f7f..c3e3f8ca74a 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -68,10 +68,7 @@ def _resolve_session_key(kwargs: Dict[str, Any]) -> Optional[str]: # crediting the bandit for conversations that are too short to signal. return None - identity = ":".join( - str(metadata.get(f) or "") if isinstance(metadata, dict) else "" - for f in _IDENTITY_FIELDS - ) + identity = ":".join(str(metadata.get(f) or "") if isinstance(metadata, dict) else "" for f in _IDENTITY_FIELDS) anchor = messages[:SIGNAL_GATE_MIN_MESSAGES] payload = ( identity @@ -146,9 +143,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple: return None, [] msg = choices[0] - msg = getattr(msg, "message", None) or ( - msg.get("message") if isinstance(msg, dict) else None - ) + msg = getattr(msg, "message", None) or (msg.get("message") if isinstance(msg, dict) else None) if msg is None: return None, [] @@ -196,11 +191,7 @@ class AdaptiveRouterPostCallHook(CustomLogger): the header is included for both paths. """ metadata = data.get("metadata") or {} - chosen = ( - metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) - if isinstance(metadata, dict) - else None - ) + chosen = metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) if isinstance(metadata, dict) else None if not chosen: return None return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} @@ -238,17 +229,11 @@ class AdaptiveRouterPostCallHook(CustomLogger): # The pre-routing hook stashes the logical pick under this key. litellm_params = kwargs.get("litellm_params") or {} metadata = litellm_params.get("metadata") or {} - current_model = ( - metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) - if isinstance(metadata, dict) - else None - ) + current_model = metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY) if isinstance(metadata, dict) else None if not current_model: return - if not self.adaptive_router.claim_or_check_owner( - session_key, current_model - ): + if not self.adaptive_router.claim_or_check_owner(session_key, current_model): # A different model owns this conversation — skip attribution. return @@ -259,9 +244,7 @@ class AdaptiveRouterPostCallHook(CustomLogger): request_type = classify_prompt(user_text or "") turn = Turn( user_content=user_text, - assistant_content=( - assistant_text if isinstance(assistant_text, str) else None - ), + assistant_content=(assistant_text if isinstance(assistant_text, str) else None), tool_calls=tool_calls, tool_results=tool_results, response_status=response_status, @@ -273,6 +256,4 @@ class AdaptiveRouterPostCallHook(CustomLogger): turn=turn, ) except Exception as e: - verbose_router_logger.exception( - "AdaptiveRouterPostCallHook: failed to record turn: %s", e - ) + verbose_router_logger.exception("AdaptiveRouterPostCallHook: failed to record turn: %s", e) diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index 5e33a64d27f..2fd1d24fbbe 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -115,9 +115,7 @@ def _jaccard(a: Set[str], b: Set[str]) -> float: _DISENGAGEMENT_PATTERNS = [ - re.compile( - r"\b(forget it|never mind|give up|talk to (?:a )?human|cancel)\b", re.IGNORECASE - ), + re.compile(r"\b(forget it|never mind|give up|talk to (?:a )?human|cancel)\b", re.IGNORECASE), re.compile(r"\b(this (?:isn'?t|is not) working|stop|abort)\b", re.IGNORECASE), re.compile(r"\bi'?ll do it (?:myself|manually)\b", re.IGNORECASE), ] @@ -211,9 +209,7 @@ _EXHAUSTION_KEYWORDS = ( ) -def _detect_exhaustion( - status: Optional[int], tool_results: List[Dict[str, Any]] -) -> bool: +def _detect_exhaustion(status: Optional[int], tool_results: List[Dict[str, Any]]) -> bool: if status is not None and status in _EXHAUSTION_STATUSES: return True for r in tool_results: @@ -246,10 +242,7 @@ def apply_turn(state: SessionState, turn: Turn) -> SignalDelta: # after MIN_TURNS_FOR_CLEAN_CREDIT turns of context. Early "thanks" # on turn 1-2 is noise, not a validated quality signal. current_turn_index = state.turn_count + 1 - if ( - not state.clean_credit_awarded - and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT - ): + if not state.clean_credit_awarded and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT: delta.satisfaction = 1 state.clean_credit_awarded = True if _detect_failure(turn.tool_results): diff --git a/litellm/router_strategy/adaptive_router/update_queue.py b/litellm/router_strategy/adaptive_router/update_queue.py index 1d87feddd84..505d243202c 100644 --- a/litellm/router_strategy/adaptive_router/update_queue.py +++ b/litellm/router_strategy/adaptive_router/update_queue.py @@ -136,9 +136,7 @@ class AdaptiveRouterUpdateQueue: "update": { "alpha": {"increment": payload["delta_alpha"]}, "beta": {"increment": payload["delta_beta"]}, - "total_samples": { - "increment": int(payload["samples_added"]) - }, + "total_samples": {"increment": int(payload["samples_added"])}, }, }, ) @@ -174,9 +172,7 @@ class AdaptiveRouterUpdateQueue: # writes to fields that are part of the @@id. asdict(state) # always carries them, so build a separate update dict. update_payload = { - k: v - for k, v in payload.items() - if k not in ("session_id", "router_name", "model_name") + k: v for k, v in payload.items() if k not in ("session_id", "router_name", "model_name") } await AdaptiveRouterSessionRepository(prisma_client).table.upsert( where={ diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index 58b2c5a3912..dc24623f505 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -144,9 +144,7 @@ class AutoRouter(CustomLogger): ) message_content = self._extract_text_from_messages(messages) - route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer( - text=message_content - ) + route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer(text=message_content) verbose_router_logger.debug(f"route_choice: {route_choice}") if isinstance(route_choice, RouteChoice): model = route_choice.name or self.default_model diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index 1fe22eafdf5..7e163ba16a6 100644 --- a/litellm/router_strategy/auto_router/litellm_encoder.py +++ b/litellm/router_strategy/auto_router/litellm_encoder.py @@ -18,11 +18,7 @@ def litellm_to_list(embeds: litellm.EmbeddingResponse) -> list[list[float]]: :param embeds: The LiteLLM embedding response. :return: A list of embeddings. """ - if ( - not embeds - or not isinstance(embeds, litellm.EmbeddingResponse) - or not embeds.data - ): + if not embeds or not isinstance(embeds, litellm.EmbeddingResponse) or not embeds.data: raise ValueError("No embeddings found in LiteLLM embedding response.") return [x["embedding"] for x in embeds.data] @@ -90,50 +86,34 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): if self.litellm_router_instance is None: raise ValueError("litellm_router_instance is not set") try: - embeds = self.litellm_router_instance.embedding( - input=docs, model=self.model_name, **kwargs - ) + embeds = self.litellm_router_instance.embedding(input=docs, model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: - raise ValueError( - f"{self.type.capitalize()} API call failed. Error: {e}" - ) from e + raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e def encode_documents(self, docs: list[str], **kwargs) -> list[list[float]]: if self.litellm_router_instance is None: raise ValueError("litellm_router_instance is not set") try: - embeds = self.litellm_router_instance.embedding( - input=docs, model=self.model_name, **kwargs - ) + embeds = self.litellm_router_instance.embedding(input=docs, model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: - raise ValueError( - f"{self.type.capitalize()} API call failed. Error: {e}" - ) from e + raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e async def aencode_queries(self, docs: list[str], **kwargs) -> list[list[float]]: if self.litellm_router_instance is None: raise ValueError("litellm_router_instance is not set") try: - embeds = await self.litellm_router_instance.aembedding( - input=docs, model=self.model_name, **kwargs - ) + embeds = await self.litellm_router_instance.aembedding(input=docs, model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: - raise ValueError( - f"{self.type.capitalize()} API call failed. Error: {e}" - ) from e + raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e async def aencode_documents(self, docs: list[str], **kwargs) -> list[list[float]]: if self.litellm_router_instance is None: raise ValueError("litellm_router_instance is not set") try: - embeds = await self.litellm_router_instance.aembedding( - input=docs, model=self.model_name, **kwargs - ) + embeds = await self.litellm_router_instance.aembedding(input=docs, model=self.model_name, **kwargs) return litellm_to_list(embeds) except Exception as e: - raise ValueError( - f"{self.type.capitalize()} API call failed. Error: {e}" - ) from e + raise ValueError(f"{self.type.capitalize()} API call failed. Error: {e}") from e diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 74451729c73..7cc83b0feb1 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -25,9 +25,7 @@ class BaseRoutingStrategy(ABC): if should_batch_redis_writes: self.setup_sync_task(default_sync_interval) - self.in_memory_keys_to_update: set[str] = ( - set() - ) # Set with max size of 1000 keys + self.in_memory_keys_to_update: set[str] = set() # Set with max size of 1000 keys def setup_sync_task(self, default_sync_interval: Optional[Union[int, float]]): """Setup the sync task in a way that's compatible with FastAPI""" @@ -38,9 +36,7 @@ class BaseRoutingStrategy(ABC): asyncio.set_event_loop(loop) self._sync_task = loop.create_task( - self.periodic_sync_in_memory_spend_with_redis( - default_sync_interval=default_sync_interval - ) + self.periodic_sync_in_memory_spend_with_redis(default_sync_interval=default_sync_interval) ) async def cleanup(self): @@ -60,15 +56,11 @@ class BaseRoutingStrategy(ABC): """ results = [] for key, value in increment_list: - result = await self._increment_value_in_current_window( - key=key, value=value, ttl=ttl - ) + result = await self._increment_value_in_current_window(key=key, value=value, ttl=ttl) results.append(result) return results - async def _increment_value_in_current_window( - self, key: str, value: Union[int, float], ttl: int - ): + async def _increment_value_in_current_window(self, key: str, value: Union[int, float], ttl: int): """ Increment spend within existing budget window @@ -92,9 +84,7 @@ class BaseRoutingStrategy(ABC): self.add_to_in_memory_keys_to_update(key=key) return result - async def periodic_sync_in_memory_spend_with_redis( - self, default_sync_interval: Optional[Union[int, float]] - ): + async def periodic_sync_in_memory_spend_with_redis(self, default_sync_interval: Optional[Union[int, float]]): """ Handler that triggers sync_in_memory_spend_with_redis every DEFAULT_REDIS_SYNC_INTERVAL seconds @@ -133,9 +123,7 @@ class BaseRoutingStrategy(ABC): for idx, op in enumerate(self.redis_increment_operation_queue): if op["key"] in compressed_ops: # Add to existing increment - compressed_ops[op["key"]]["increment_value"] += op[ - "increment_value" - ] + compressed_ops[op["key"]]["increment_value"] += op["increment_value"] else: compressed_ops[op["key"]] = op @@ -144,31 +132,22 @@ class BaseRoutingStrategy(ABC): # Convert back to list compressed_queue = list(compressed_ops.values()) - increment_result = ( - await self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=compressed_queue, - ) + increment_result = await self.dual_cache.redis_cache.async_increment_pipeline( + increment_list=compressed_queue, ) self.redis_increment_operation_queue = [ - op - for idx, op in enumerate(self.redis_increment_operation_queue) - if idx not in ops_to_remove + op for idx, op in enumerate(self.redis_increment_operation_queue) if idx not in ops_to_remove ] if increment_result is not None: - return_result = { - key["key"]: op - for key, op in zip(compressed_queue, increment_result) - } + return_result = {key["key"]: op for key, op in zip(compressed_queue, increment_result)} else: return_result = {} return return_result except Exception as e: - verbose_router_logger.error( - f"Error syncing in-memory cache with Redis: {str(e)}" - ) + verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}") self.redis_increment_operation_queue = [] def add_to_in_memory_keys_to_update(self, key: str): @@ -211,17 +190,15 @@ class BaseRoutingStrategy(ABC): return # 2. Fetch all current provider spend from Redis to update in-memory cache - cache_keys = self.get_in_memory_keys_to_update() # if no pattern OR redis cache does not support scan_iter, use in-memory keys + cache_keys = ( + self.get_in_memory_keys_to_update() + ) # if no pattern OR redis cache does not support scan_iter, use in-memory keys cache_keys_list = list(cache_keys) # 1. Snapshot in-memory before in_memory_before_dict = {} - in_memory_before = ( - await self.dual_cache.in_memory_cache.async_batch_get_cache( - keys=cache_keys_list - ) - ) + in_memory_before = await self.dual_cache.in_memory_cache.async_batch_get_cache(keys=cache_keys_list) for k, v in zip(cache_keys_list, in_memory_before): in_memory_before_dict[k] = float(v or 0) @@ -234,9 +211,7 @@ class BaseRoutingStrategy(ABC): for key in cache_keys_list: redis_val = float(redis_values.get(key, 0) or 0) before = float(in_memory_before_dict.get(key, 0) or 0) - after = float( - await self.dual_cache.in_memory_cache.async_get_cache(key=key) or 0 - ) + after = float(await self.dual_cache.in_memory_cache.async_get_cache(key=key) or 0) delta = after - before if after <= redis_val: merged = redis_val + delta @@ -249,11 +224,7 @@ class BaseRoutingStrategy(ABC): # import os # os._exit(1) # raise Exception(f"Redis is behind in-memory cache for key: {key}. This should not happen, since we should be updating redis with in-memory cache.") - await self.dual_cache.in_memory_cache.async_set_cache( - key=key, value=merged - ) + await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=merged) except Exception as e: - verbose_router_logger.exception( - f"Error syncing in-memory cache with Redis: {str(e)}" - ) + verbose_router_logger.exception(f"Error syncing in-memory cache with Redis: {str(e)}") diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 0bb69ca0319..067f38ab11c 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -101,9 +101,7 @@ class RouterBudgetLimiting(CustomLogger): self.dual_cache = dual_cache self.redis_increment_operation_queue: List[RedisPipelineIncrementOperation] = [] asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) - self.provider_budget_config: Optional[GenericBudgetConfigType] = ( - provider_budget_config - ) + self.provider_budget_config: Optional[GenericBudgetConfigType] = provider_budget_config self.deployment_budget_config: Optional[GenericBudgetConfigType] = None self.tag_budget_config: Optional[GenericBudgetConfigType] = None self._init_provider_budgets() @@ -177,9 +175,7 @@ class RouterBudgetLimiting(CustomLogger): potential_deployments=potential_deployments, request_tags=_get_tags_from_request_kwargs( request_kwargs=request_kwargs, - metadata_variable_name=get_metadata_variable_name_from_kwargs( - request_kwargs or {} - ), + metadata_variable_name=get_metadata_variable_name_from_kwargs(request_kwargs or {}), ), ) @@ -228,9 +224,7 @@ class RouterBudgetLimiting(CustomLogger): config = provider_configs[provider] if config.max_budget is None: continue - current_spend = spend_map.get( - f"provider_spend:{provider}:{config.budget_duration}", 0.0 - ) + current_spend = spend_map.get(f"provider_spend:{provider}:{config.budget_duration}", 0.0) self._track_provider_remaining_budget_prometheus( provider=provider, spend=current_spend, @@ -251,9 +245,7 @@ class RouterBudgetLimiting(CustomLogger): model_id = deployment.get("model_info", {}).get("id") if model_id in deployment_configs: config = deployment_configs[model_id] - current_spend = spend_map.get( - f"deployment_spend:{model_id}:{config.budget_duration}", 0.0 - ) + current_spend = spend_map.get(f"deployment_spend:{model_id}:{config.budget_duration}", 0.0) if config.max_budget and current_spend >= config.max_budget: debug_msg = f"Exceeded budget for deployment model_name: {_model_name}, litellm_params.model: {_litellm_model_name}, model_id: {model_id}: {current_spend} >= {config.budget_duration}" verbose_router_logger.debug(debug_msg) @@ -269,10 +261,7 @@ class RouterBudgetLimiting(CustomLogger): f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}", 0.0, ) - if ( - _tag_budget_config.max_budget - and _tag_spend >= _tag_budget_config.max_budget - ): + if _tag_budget_config.max_budget and _tag_spend >= _tag_budget_config.max_budget: debug_msg = f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, tag_budget_limit={_tag_budget_config.max_budget}" verbose_router_logger.debug(debug_msg) deployment_above_budget_info += f"{debug_msg}\n" @@ -313,9 +302,7 @@ class RouterBudgetLimiting(CustomLogger): if self.tag_budget_config: _request_tags = _get_tags_from_request_kwargs( request_kwargs=request_kwargs, - metadata_variable_name=get_metadata_variable_name_from_kwargs( - request_kwargs or {} - ), + metadata_variable_name=get_metadata_variable_name_from_kwargs(request_kwargs or {}), ) for deployment in healthy_deployments: @@ -325,14 +312,9 @@ class RouterBudgetLimiting(CustomLogger): deployment_providers.append(provider) if provider is not None: budget_config = self._get_budget_config_for_provider(provider) - if ( - budget_config is not None - and budget_config.budget_duration is not None - ): + if budget_config is not None and budget_config.budget_duration is not None: provider_configs[provider] = budget_config - cache_keys.append( - f"provider_spend:{provider}:{budget_config.budget_duration}" - ) + cache_keys.append(f"provider_spend:{provider}:{budget_config.budget_duration}") # Check deployment budgets if self.deployment_budget_config: @@ -341,17 +323,13 @@ class RouterBudgetLimiting(CustomLogger): budget_config = self._get_budget_config_for_deployment(model_id) if budget_config is not None: deployment_configs[model_id] = budget_config - cache_keys.append( - f"deployment_spend:{model_id}:{budget_config.budget_duration}" - ) + cache_keys.append(f"deployment_spend:{model_id}:{budget_config.budget_duration}") # Check tag budgets (outside loop — tags are per-request, not per-deployment) for _tag in _request_tags: _tag_budget_config = self._get_budget_config_for_tag(_tag) if _tag_budget_config: - cache_keys.append( - f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}" - ) + cache_keys.append(f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}") return ( cache_keys, provider_configs, @@ -359,9 +337,7 @@ class RouterBudgetLimiting(CustomLogger): deployment_providers, ) - async def _get_or_set_budget_start_time( - self, start_time_key: str, current_time: float, ttl_seconds: int - ) -> float: + async def _get_or_set_budget_start_time(self, start_time_key: str, current_time: float, ttl_seconds: int) -> float: """ Checks if the key = `provider_budget_start_time:{provider}` exists in cache. @@ -370,9 +346,7 @@ class RouterBudgetLimiting(CustomLogger): """ budget_start = await self.dual_cache.async_get_cache(start_time_key) if budget_start is None: - await self.dual_cache.async_set_cache( - key=start_time_key, value=current_time, ttl=ttl_seconds - ) + await self.dual_cache.async_set_cache(key=start_time_key, value=current_time, ttl=ttl_seconds) return current_time return float(budget_start) @@ -396,17 +370,11 @@ class RouterBudgetLimiting(CustomLogger): - stores key: `provider_budget_start_time:{provider}`, value: current_time. This stores the start time of the new budget window """ - await self.dual_cache.async_set_cache( - key=spend_key, value=response_cost, ttl=ttl_seconds - ) - await self.dual_cache.async_set_cache( - key=start_time_key, value=current_time, ttl=ttl_seconds - ) + await self.dual_cache.async_set_cache(key=spend_key, value=response_cost, ttl=ttl_seconds) + await self.dual_cache.async_set_cache(key=start_time_key, value=current_time, ttl=ttl_seconds) return current_time - async def _increment_spend_in_current_window( - self, spend_key: str, response_cost: float, ttl: int - ): + async def _increment_spend_in_current_window(self, spend_key: str, response_cost: float, ttl: int): """ Increment spend within existing budget window @@ -433,26 +401,20 @@ class RouterBudgetLimiting(CustomLogger): # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: raise ValueError("standard_logging_payload is required") response_cost: float = standard_logging_payload.get("response_cost", 0) model_id: str = str(standard_logging_payload.get("model_id", "")) - custom_llm_provider: str = kwargs.get("litellm_params", {}).get( - "custom_llm_provider", None - ) + custom_llm_provider: str = kwargs.get("litellm_params", {}).get("custom_llm_provider", None) if custom_llm_provider is None: raise ValueError("custom_llm_provider is required") budget_config = self._get_budget_config_for_provider(custom_llm_provider) if budget_config: # increment spend for provider - spend_key = ( - f"provider_spend:{custom_llm_provider}:{budget_config.budget_duration}" - ) + spend_key = f"provider_spend:{custom_llm_provider}:{budget_config.budget_duration}" start_time_key = f"provider_budget_start_time:{custom_llm_provider}" await self._increment_spend_for_key( budget_config=budget_config, @@ -481,9 +443,7 @@ class RouterBudgetLimiting(CustomLogger): for _tag in request_tags: _tag_budget_config = self._get_budget_config_for_tag(_tag) if _tag_budget_config: - _tag_spend_key = ( - f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}" - ) + _tag_spend_key = f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}" _tag_start_time_key = f"tag_budget_start_time:{_tag}" await self._increment_spend_for_key( budget_config=_tag_budget_config, @@ -539,9 +499,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=spend_key, response_cost=response_cost, ttl=ttl_for_increment ) - verbose_router_logger.debug( - f"Incremented spend for {spend_key} by {response_cost}" - ) + verbose_router_logger.debug(f"Incremented spend for {spend_key} by {response_cost}") async def periodic_sync_in_memory_spend_with_redis(self): """ @@ -587,9 +545,7 @@ class RouterBudgetLimiting(CustomLogger): self.redis_increment_operation_queue = [] except Exception as e: - verbose_router_logger.error( - f"Error syncing in-memory cache with Redis: {str(e)}" - ) + verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}") async def _sync_in_memory_spend_with_redis(self): """ @@ -619,17 +575,13 @@ class RouterBudgetLimiting(CustomLogger): for provider, config in self.provider_budget_config.items(): if config is None: continue - cache_keys.append( - f"provider_spend:{provider}:{config.budget_duration}" - ) + cache_keys.append(f"provider_spend:{provider}:{config.budget_duration}") if self.deployment_budget_config is not None: for model_id, config in self.deployment_budget_config.items(): if config is None: continue - cache_keys.append( - f"deployment_spend:{model_id}:{config.budget_duration}" - ) + cache_keys.append(f"deployment_spend:{model_id}:{config.budget_duration}") if self.tag_budget_config is not None: for tag, config in self.tag_budget_config.items(): @@ -638,25 +590,17 @@ class RouterBudgetLimiting(CustomLogger): cache_keys.append(f"tag_spend:{tag}:{config.budget_duration}") # Batch fetch current spend values from Redis - redis_values = await self.dual_cache.redis_cache.async_batch_get_cache( - key_list=cache_keys - ) + redis_values = await self.dual_cache.redis_cache.async_batch_get_cache(key_list=cache_keys) # Update in-memory cache with Redis values if isinstance(redis_values, dict): # Check if redis_values is a dictionary for key, value in redis_values.items(): if value is not None: - await self.dual_cache.in_memory_cache.async_set_cache( - key=key, value=float(value) - ) - verbose_router_logger.debug( - f"Updated in-memory cache for {key}: {value}" - ) + await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=float(value)) + verbose_router_logger.debug(f"Updated in-memory cache for {key}: {value}") except Exception as e: - verbose_router_logger.error( - f"Error syncing in-memory cache with Redis: {str(e)}" - ) + verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}") def _get_budget_config_for_deployment( self, @@ -666,9 +610,7 @@ class RouterBudgetLimiting(CustomLogger): return None return self.deployment_budget_config.get(model_id, None) - def _get_budget_config_for_provider( - self, provider: str - ) -> Optional[GenericBudgetInfo]: + def _get_budget_config_for_provider(self, provider: str) -> Optional[GenericBudgetInfo]: if self.provider_budget_config is None: return None return self.provider_budget_config.get(provider, None) @@ -687,9 +629,7 @@ class RouterBudgetLimiting(CustomLogger): provider_resolution_params: Any = deployment_litellm_params elif isinstance(deployment_litellm_params, dict): model = deployment_litellm_params.get("model") or "" - provider_resolution_params = _LiteLLMParamsDictView( - deployment_litellm_params - ) + provider_resolution_params = _LiteLLMParamsDictView(deployment_litellm_params) else: model = "" provider_resolution_params = _LiteLLMParamsDictView({}) @@ -699,15 +639,11 @@ class RouterBudgetLimiting(CustomLogger): litellm_params=provider_resolution_params, ) except Exception: - verbose_router_logger.error( - f"Error getting LLM provider for deployment: {deployment}" - ) + verbose_router_logger.error(f"Error getting LLM provider for deployment: {deployment}") return None return custom_llm_provider - def _track_provider_remaining_budget_prometheus( - self, provider: str, spend: float, budget_limit: float - ): + def _track_provider_remaining_budget_prometheus(self, provider: str, spend: float, budget_limit: float): """ Optional helper - emit provider remaining budget metric to Prometheus @@ -748,9 +684,7 @@ class RouterBudgetLimiting(CustomLogger): current_spend = await self.dual_cache.async_get_cache(spend_key) return float(current_spend) if current_spend is not None else 0.0 - async def _get_current_provider_budget_reset_at( - self, provider: str - ) -> Optional[str]: + async def _get_current_provider_budget_reset_at(self, provider: str) -> Optional[str]: budget_config = self._get_budget_config_for_provider(provider) if budget_config is None: return None @@ -766,9 +700,7 @@ class RouterBudgetLimiting(CustomLogger): return (datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds)).isoformat() - async def _init_provider_budget_in_cache( - self, provider: str, budget_config: GenericBudgetInfo - ): + async def _init_provider_budget_in_cache(self, provider: str, budget_config: GenericBudgetInfo): """ Initialize provider budget in cache by storing the following keys if they don't exist: - provider_spend:{provider}:{budget_config.time_period} - stores the current spend @@ -785,22 +717,16 @@ class RouterBudgetLimiting(CustomLogger): budget_start = await self.dual_cache.async_get_cache(start_time_key) if budget_start is None: budget_start = datetime.now(timezone.utc).timestamp() - await self.dual_cache.async_set_cache( - key=start_time_key, value=budget_start, ttl=ttl_seconds - ) + await self.dual_cache.async_set_cache(key=start_time_key, value=budget_start, ttl=ttl_seconds) _spend_key = await self.dual_cache.async_get_cache(spend_key) if _spend_key is None: - await self.dual_cache.async_set_cache( - key=spend_key, value=0.0, ttl=ttl_seconds - ) + await self.dual_cache.async_set_cache(key=spend_key, value=0.0, ttl=ttl_seconds) @staticmethod def should_init_router_budget_limiter( provider_budget_config: Optional[dict], - model_list: Optional[ - Union[List[DeploymentTypedDict], List[Dict[str, Any]]] - ] = None, + model_list: Optional[Union[List[DeploymentTypedDict], List[Dict[str, Any]]]] = None, ): """ Returns `True` if the router budget routing settings are set and RouterBudgetLimiting should be initialized @@ -821,10 +747,7 @@ class RouterBudgetLimiting(CustomLogger): for _model in model_list: _litellm_params = _model.get("litellm_params", {}) - if ( - _litellm_params.get("max_budget") - or _litellm_params.get("budget_duration") is not None - ): + if _litellm_params.get("max_budget") or _litellm_params.get("budget_duration") is not None: return True return False @@ -849,9 +772,7 @@ class RouterBudgetLimiting(CustomLogger): ) ) - verbose_router_logger.debug( - f"Initalized Provider budget config: {self.provider_budget_config}" - ) + verbose_router_logger.debug(f"Initalized Provider budget config: {self.provider_budget_config}") def _init_deployment_budgets( self, @@ -869,11 +790,7 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug( f"Init Deployment Budget: max_budget: {_max_budget}, budget_duration: {_budget_duration}, model_id: {_model_id}" ) - if ( - _max_budget is not None - and _budget_duration is not None - and _model_id is not None - ): + if _max_budget is not None and _budget_duration is not None and _model_id is not None: _budget_config = GenericBudgetInfo( time_period=_budget_duration, budget_limit=_max_budget, @@ -882,9 +799,7 @@ class RouterBudgetLimiting(CustomLogger): self.deployment_budget_config = {} self.deployment_budget_config[_model_id] = _budget_config - verbose_router_logger.debug( - f"Initialized Deployment Budget Config: {self.deployment_budget_config}" - ) + verbose_router_logger.debug(f"Initialized Deployment Budget Config: {self.deployment_budget_config}") def register_deployment_budget( self, @@ -908,9 +823,7 @@ class RouterBudgetLimiting(CustomLogger): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user if premium_user is not True: - raise ValueError( - f"Tag budgets are an Enterprise only feature, {CommonProxyErrors.not_premium_user}" - ) + raise ValueError(f"Tag budgets are an Enterprise only feature, {CommonProxyErrors.not_premium_user}") if self.tag_budget_config is None: self.tag_budget_config = {} @@ -924,6 +837,4 @@ class RouterBudgetLimiting(CustomLogger): ) self.tag_budget_config[_tag] = _generic_budget_config - verbose_router_logger.debug( - f"Initialized Tag Budget Config: {self.tag_budget_config}" - ) + verbose_router_logger.debug(f"Initialized Tag Budget Config: {self.tag_budget_config}") diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index aa3bcef6392..9a1d845dd65 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -89,12 +89,8 @@ class ComplexityRouter(CustomLogger): # Build effective keyword lists (use config overrides or defaults) self.code_keywords = self.config.code_keywords or DEFAULT_CODE_KEYWORDS - self.reasoning_keywords = ( - self.config.reasoning_keywords or DEFAULT_REASONING_KEYWORDS - ) - self.technical_keywords = ( - self.config.technical_keywords or DEFAULT_TECHNICAL_KEYWORDS - ) + self.reasoning_keywords = self.config.reasoning_keywords or DEFAULT_REASONING_KEYWORDS + self.technical_keywords = self.config.technical_keywords or DEFAULT_TECHNICAL_KEYWORDS self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS # Pre-compile regex patterns for efficiency @@ -106,9 +102,7 @@ class ComplexityRouter(CustomLogger): re.compile(r"[a-z]\)\s", re.IGNORECASE), ] - verbose_router_logger.debug( - f"ComplexityRouter initialized for {model_name} with tiers: {self.config.tiers}" - ) + verbose_router_logger.debug(f"ComplexityRouter initialized for {model_name} with tiers: {self.config.tiers}") def _estimate_tokens(self, text: str) -> int: """ @@ -124,13 +118,9 @@ class ComplexityRouter(CustomLogger): complex_threshold = thresholds.get("complex", 400) if estimated_tokens < simple_threshold: - return DimensionScore( - "tokenCount", -1.0, f"short ({estimated_tokens} tokens)" - ) + return DimensionScore("tokenCount", -1.0, f"short ({estimated_tokens} tokens)") if estimated_tokens > complex_threshold: - return DimensionScore( - "tokenCount", 1.0, f"long ({estimated_tokens} tokens)" - ) + return DimensionScore("tokenCount", 1.0, f"long ({estimated_tokens} tokens)") return DimensionScore("tokenCount", 0, None) def _keyword_matches(self, text: str, keyword: str) -> bool: @@ -174,16 +164,12 @@ class ComplexityRouter(CustomLogger): if match_count >= high_threshold: return ( - DimensionScore( - name, score_high, f"{signal_label} ({', '.join(matches[:3])})" - ), + DimensionScore(name, score_high, f"{signal_label} ({', '.join(matches[:3])})"), match_count, ) if match_count >= low_threshold: return ( - DimensionScore( - name, score_low, f"{signal_label} ({', '.join(matches[:3])})" - ), + DimensionScore(name, score_low, f"{signal_label} ({', '.join(matches[:3])})"), match_count, ) return DimensionScore(name, score_none, None), match_count @@ -202,9 +188,7 @@ class ComplexityRouter(CustomLogger): return DimensionScore("questionComplexity", 0.5, f"{count} questions") return DimensionScore("questionComplexity", 0, None) - def classify( - self, prompt: str, system_prompt: Optional[str] = None - ) -> Tuple[ComplexityTier, float, List[str]]: + def classify(self, prompt: str, system_prompt: Optional[str] = None) -> Tuple[ComplexityTier, float, List[str]]: """ Classify a prompt by complexity. @@ -328,9 +312,7 @@ class ComplexityRouter(CustomLogger): if medium_model: return medium_model - raise ValueError( - f"No model configured for tier {tier_key} and no default_model set" - ) + raise ValueError(f"No model configured for tier {tier_key} and no default_model set") def _resolve_messages( self, @@ -356,9 +338,7 @@ class ComplexityRouter(CustomLogger): call_type: Optional[CallTypes] = None # 1. Try route-based inference from proxy metadata - route = request_kwargs.get("litellm_metadata", {}).get( - "user_api_key_request_route" - ) + route = request_kwargs.get("litellm_metadata", {}).get("user_api_key_request_route") if route: call_types_list = get_call_types_for_route(route) if call_types_list: @@ -396,9 +376,7 @@ class ComplexityRouter(CustomLogger): content = msg.get("content") or "" if isinstance(content, list): text_parts = [ - part.get("text", "") - for part in content - if isinstance(part, dict) and part.get("type") == "text" + part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text" ] content = " ".join(text_parts).strip() if isinstance(content, str) and content: @@ -441,25 +419,18 @@ class ComplexityRouter(CustomLogger): resolved_messages = self._resolve_messages(messages, request_kwargs) if not resolved_messages: - verbose_router_logger.debug( - "ComplexityRouter: No messages could be resolved, skipping routing" - ) + verbose_router_logger.debug("ComplexityRouter: No messages could be resolved, skipping routing") return None # Determine whether the original request used messages directly has_original_messages = messages is not None and len(messages) > 0 - user_message, system_prompt = self._extract_user_message_and_system_prompt( - resolved_messages - ) + user_message, system_prompt = self._extract_user_message_and_system_prompt(resolved_messages) if user_message is None: - verbose_router_logger.debug( - "ComplexityRouter: No user message found, routing to default model" - ) + verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model") return PreRoutingHookResponse( - model=self.config.default_model - or self.get_model_for_tier(ComplexityTier.MEDIUM), + model=self.config.default_model or self.get_model_for_tier(ComplexityTier.MEDIUM), messages=messages if has_original_messages else None, ) @@ -467,8 +438,7 @@ class ComplexityRouter(CustomLogger): routed_model = self.get_model_for_tier(tier) verbose_router_logger.info( - f"ComplexityRouter: tier={tier.value}, score={score:.3f}, " - f"signals={signals}, routed_model={routed_model}" + f"ComplexityRouter: tier={tier.value}, score={score:.3f}, signals={signals}, routed_model={routed_model}" ) return PreRoutingHookResponse( diff --git a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py index 7c1c8f2907d..b46071ccee8 100644 --- a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py +++ b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py @@ -15,9 +15,7 @@ from dataclasses import dataclass from typing import List, Optional, Tuple from unittest.mock import MagicMock -sys.path.insert( - 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) -) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityTier @@ -260,9 +258,7 @@ def run_eval() -> Tuple[int, int, List[dict]]: # Check if pass is_exact_match = tier == case.expected_tier - is_acceptable = ( - case.acceptable_tiers is not None and tier in case.acceptable_tiers - ) + is_acceptable = case.acceptable_tiers is not None and tier in case.acceptable_tiers is_pass = is_exact_match or is_acceptable if is_pass: @@ -274,28 +270,18 @@ def run_eval() -> Tuple[int, int, List[dict]]: { "case": i, "description": case.description, - "prompt": ( - case.prompt[:80] + "..." - if len(case.prompt) > 80 - else case.prompt - ), + "prompt": (case.prompt[:80] + "..." if len(case.prompt) > 80 else case.prompt), "expected": case.expected_tier.value, "actual": tier.value, "score": round(score, 3), "signals": signals, - "acceptable": ( - [t.value for t in case.acceptable_tiers] - if case.acceptable_tiers - else None - ), + "acceptable": ([t.value for t in case.acceptable_tiers] if case.acceptable_tiers else None), } ) # Print result print(f"[{i:2d}] {status} | {case.description}") - print( - f" Expected: {case.expected_tier.value:10s} | Got: {tier.value:10s} | Score: {score:+.3f}" - ) + print(f" Expected: {case.expected_tier.value:10s} | Got: {tier.value:10s} | Score: {score:+.3f}") if signals: print(f" Signals: {', '.join(signals)}") if not is_pass: @@ -312,9 +298,7 @@ def run_eval() -> Tuple[int, int, List[dict]]: print("-" * 70) for f in failures: print(f"Case {f['case']}: {f['description']}") - print( - f" Expected: {f['expected']}, Got: {f['actual']} (score: {f['score']})" - ) + print(f" Expected: {f['expected']}, Got: {f['actual']} (score: {f['score']})") print(f" Signals: {f['signals']}") if f["acceptable"]: print(f" Acceptable: {f['acceptable']}") diff --git a/litellm/router_strategy/lar1_routing.py b/litellm/router_strategy/lar1_routing.py index 31aa5a96bad..b53165f0969 100644 --- a/litellm/router_strategy/lar1_routing.py +++ b/litellm/router_strategy/lar1_routing.py @@ -35,15 +35,9 @@ def lar1_thresholds_from_args( ) -> dict[str, float]: args = routing_strategy_args or {} return { - "low": _coerce_threshold( - args.get("confidence_threshold_low"), DEFAULT_THRESHOLDS["low"] - ), - "medium": _coerce_threshold( - args.get("confidence_threshold_medium"), DEFAULT_THRESHOLDS["medium"] - ), - "high": _coerce_threshold( - args.get("confidence_threshold_high"), DEFAULT_THRESHOLDS["high"] - ), + "low": _coerce_threshold(args.get("confidence_threshold_low"), DEFAULT_THRESHOLDS["low"]), + "medium": _coerce_threshold(args.get("confidence_threshold_medium"), DEFAULT_THRESHOLDS["medium"]), + "high": _coerce_threshold(args.get("confidence_threshold_high"), DEFAULT_THRESHOLDS["high"]), } @@ -66,8 +60,7 @@ def _normalize_thresholds(thresholds: Optional[dict[str, float]]) -> dict[str, f high = merged["high"] if not (0 < low < medium < high < 1): raise ValueError( - "LAR-1 thresholds must satisfy 0 < low < medium < high < 1, " - f"got low={low}, medium={medium}, high={high}" + f"LAR-1 thresholds must satisfy 0 < low < medium < high < 1, got low={low}, medium={medium}, high={high}" ) return merged @@ -75,16 +68,12 @@ def _normalize_thresholds(thresholds: Optional[dict[str, float]]) -> dict[str, f def _parse_lar1_metadata(request_kwargs: dict) -> LAR1Metadata: lar1_raw = request_kwargs.get("metadata", {}).get("lar1", {}) if not isinstance(lar1_raw, dict): - verbose_router_logger.warning( - f"[LAR-1] Invalid lar1 metadata type: {type(lar1_raw).__name__}. Using defaults" - ) + verbose_router_logger.warning(f"[LAR-1] Invalid lar1 metadata type: {type(lar1_raw).__name__}. Using defaults") return LAR1Metadata() try: return LAR1Metadata.model_validate(lar1_raw) except ValidationError as exc: - verbose_router_logger.warning( - f"[LAR-1] Invalid lar1 metadata: {exc}. Using defaults" - ) + verbose_router_logger.warning(f"[LAR-1] Invalid lar1 metadata: {exc}. Using defaults") return LAR1Metadata() @@ -138,8 +127,7 @@ class LAR1RoutingStrategy(CustomRoutingStrategyBase): else: actual_type = selected.get("model_info", {}).get("type", "unknown") verbose_router_logger.warning( - f"[LAR-1] No deployment for type '{target}', " - f"fallback to deployment type '{actual_type}'" + f"[LAR-1] No deployment for type '{target}', fallback to deployment type '{actual_type}'" ) return selected @@ -187,6 +175,5 @@ class LAR1RoutingStrategy(CustomRoutingStrategyBase): def get_available_deployment(self, *args, **kwargs): raise NotImplementedError( - "LAR-1 routing only supports async routing. " - "Enable async_only_mode on the router or use acompletion." + "LAR-1 routing only supports async routing. Enable async_only_mode on the router or use acompletion." ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index e1614388379..819fedde991 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -31,9 +31,7 @@ class LeastBusyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: return @@ -42,14 +40,10 @@ class LeastBusyLoggingHandler(CustomLogger): request_count_api_key = f"{model_group}_request_count" # update cache - request_count_dict = ( - self.router_cache.get_cache(key=request_count_api_key) or {} - ) + request_count_dict = self.router_cache.get_cache(key=request_count_api_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + 1 - self.router_cache.set_cache( - key=request_count_api_key, value=request_count_dict - ) + self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) except Exception: pass @@ -58,9 +52,7 @@ class LeastBusyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: @@ -70,16 +62,12 @@ class LeastBusyLoggingHandler(CustomLogger): request_count_api_key = f"{model_group}_request_count" # decrement count in cache - request_count_dict = ( - self.router_cache.get_cache(key=request_count_api_key) or {} - ) + request_count_dict = self.router_cache.get_cache(key=request_count_api_key) or {} request_count_value: Optional[int] = request_count_dict.get(id, 0) if request_count_value is None: return request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache( - key=request_count_api_key, value=request_count_dict - ) + self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) ### TESTING ### if self.test_flag: @@ -92,9 +80,7 @@ class LeastBusyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: return @@ -103,16 +89,12 @@ class LeastBusyLoggingHandler(CustomLogger): request_count_api_key = f"{model_group}_request_count" # decrement count in cache - request_count_dict = ( - self.router_cache.get_cache(key=request_count_api_key) or {} - ) + request_count_dict = self.router_cache.get_cache(key=request_count_api_key) or {} request_count_value: Optional[int] = request_count_dict.get(id, 0) if request_count_value is None: return request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache( - key=request_count_api_key, value=request_count_dict - ) + self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) ### TESTING ### if self.test_flag: @@ -125,9 +107,7 @@ class LeastBusyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: @@ -137,17 +117,12 @@ class LeastBusyLoggingHandler(CustomLogger): request_count_api_key = f"{model_group}_request_count" # decrement count in cache - request_count_dict = ( - await self.router_cache.async_get_cache(key=request_count_api_key) - or {} - ) + request_count_dict = await self.router_cache.async_get_cache(key=request_count_api_key) or {} request_count_value: Optional[int] = request_count_dict.get(id, 0) if request_count_value is None: return request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache( - key=request_count_api_key, value=request_count_dict - ) + await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) ### TESTING ### if self.test_flag: @@ -160,9 +135,7 @@ class LeastBusyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: return @@ -171,17 +144,12 @@ class LeastBusyLoggingHandler(CustomLogger): request_count_api_key = f"{model_group}_request_count" # decrement count in cache - request_count_dict = ( - await self.router_cache.async_get_cache(key=request_count_api_key) - or {} - ) + request_count_dict = await self.router_cache.async_get_cache(key=request_count_api_key) or {} request_count_value: Optional[int] = request_count_dict.get(id, 0) if request_count_value is None: return request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache( - key=request_count_api_key, value=request_count_dict - ) + await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) ### TESTING ### if self.test_flag: @@ -234,16 +202,12 @@ class LeastBusyLoggingHandler(CustomLogger): all_deployments=all_deployments, ) - async def async_get_available_deployments( - self, model_group: str, healthy_deployments: list - ): + async def async_get_available_deployments(self, model_group: str, healthy_deployments: list): """ Async helper to get deployments using least busy strategy """ request_count_api_key = f"{model_group}_request_count" - all_deployments = ( - await self.router_cache.async_get_cache(key=request_count_api_key) or {} - ) + all_deployments = await self.router_cache.async_get_cache(key=request_count_api_key) or {} return self._get_available_deployments( healthy_deployments=healthy_deployments, all_deployments=all_deployments, diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 3f641d4f0fb..67bdbfbe0ac 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -26,9 +26,7 @@ class LowestCostLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: @@ -85,9 +83,7 @@ class LowestCostLoggingHandler(CustomLogger): ) ## RPM - request_count_dict[id][precise_minute]["rpm"] = ( - request_count_dict[id][precise_minute].get("rpm", 0) + 1 - ) + request_count_dict[id][precise_minute]["rpm"] = request_count_dict[id][precise_minute].get("rpm", 0) + 1 self.router_cache.set_cache(key=cost_key, value=request_count_dict) @@ -96,9 +92,7 @@ class LowestCostLoggingHandler(CustomLogger): self.logged_success += 1 except Exception as e: verbose_logger.exception( - "litellm.router_strategy.lowest_cost.py::log_success_event(): Exception occured - {}".format( - str(e) - ) + "litellm.router_strategy.lowest_cost.py::log_success_event(): Exception occured - {}".format(str(e)) ) pass @@ -110,9 +104,7 @@ class LowestCostLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: @@ -156,9 +148,7 @@ class LowestCostLoggingHandler(CustomLogger): # Update usage # ------------ - request_count_dict = ( - await self.router_cache.async_get_cache(key=cost_key) or {} - ) + request_count_dict = await self.router_cache.async_get_cache(key=cost_key) or {} if id not in request_count_dict: request_count_dict[id] = {} @@ -171,9 +161,7 @@ class LowestCostLoggingHandler(CustomLogger): ) ## RPM - request_count_dict[id][precise_minute]["rpm"] = ( - request_count_dict[id][precise_minute].get("rpm", 0) + 1 - ) + request_count_dict[id][precise_minute]["rpm"] = request_count_dict[id][precise_minute].get("rpm", 0) + 1 await self.router_cache.async_set_cache( key=cost_key, value=request_count_dict @@ -261,32 +249,22 @@ class LowestCostLoggingHandler(CustomLogger): or float("inf") ) item_litellm_model_name = _deployment.get("litellm_params", {}).get("model") - item_litellm_model_cost_map = litellm.model_cost.get( - item_litellm_model_name, {} - ) + item_litellm_model_cost_map = litellm.model_cost.get(item_litellm_model_name, {}) # check if user provided input_cost_per_token and output_cost_per_token in litellm_params item_input_cost = None item_output_cost = None if _deployment.get("litellm_params", {}).get("input_cost_per_token", None): - item_input_cost = _deployment.get("litellm_params", {}).get( - "input_cost_per_token" - ) + item_input_cost = _deployment.get("litellm_params", {}).get("input_cost_per_token") if _deployment.get("litellm_params", {}).get("output_cost_per_token", None): - item_output_cost = _deployment.get("litellm_params", {}).get( - "output_cost_per_token" - ) + item_output_cost = _deployment.get("litellm_params", {}).get("output_cost_per_token") if item_input_cost is None: - item_input_cost = item_litellm_model_cost_map.get( - "input_cost_per_token", 5.0 - ) + item_input_cost = item_litellm_model_cost_map.get("input_cost_per_token", 5.0) if item_output_cost is None: - item_output_cost = item_litellm_model_cost_map.get( - "output_cost_per_token", 5.0 - ) + item_output_cost = item_litellm_model_cost_map.get("output_cost_per_token", 5.0) # if litellm["model"] is not in model_cost map -> use item_cost = $10 @@ -304,9 +282,7 @@ class LowestCostLoggingHandler(CustomLogger): # -------------- # # We use _cost_per_deployment to log to langfuse, slack - this is not used to make a decision on routing # this helps a user to debug why the router picked a specfic deployment # - _deployment_api_base = _deployment.get("litellm_params", {}).get( - "api_base", "" - ) + _deployment_api_base = _deployment.get("litellm_params", {}).get("api_base", "") if _deployment_api_base is not None: _cost_per_deployment[_deployment_api_base] = item_cost # -------------- # @@ -314,8 +290,7 @@ class LowestCostLoggingHandler(CustomLogger): # -------------- # if ( - item_tpm + input_tokens > _deployment_tpm - or item_rpm + 1 > _deployment_rpm + item_tpm + input_tokens > _deployment_tpm or item_rpm + 1 > _deployment_rpm ): # if user passed in tpm / rpm in the model_list continue else: diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 3f2db97e2bf..23476fe7dcc 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -44,9 +44,7 @@ class LowestLatencyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get(metadata_field) is None: pass else: - model_group = kwargs["litellm_params"][metadata_field].get( - "model_group", None - ) + model_group = kwargs["litellm_params"][metadata_field].get("model_group", None) id = (kwargs["litellm_params"].get("model_info") or {}).get("id", None) if model_group is None or id is None: @@ -79,9 +77,7 @@ class LowestLatencyLoggingHandler(CustomLogger): if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = ( - kwargs.get("completion_start_time", end_time) - start_time - ) + time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time final_value: Union[float, timedelta] = response_ms time_to_first_token: Optional[float] = None @@ -99,9 +95,7 @@ class LowestLatencyLoggingHandler(CustomLogger): else: response_seconds = response_ms - final_value = safe_divide_seconds( - response_seconds, completion_tokens - ) + final_value = safe_divide_seconds(response_seconds, completion_tokens) if final_value is not None: final_value = float(final_value) else: @@ -109,39 +103,27 @@ class LowestLatencyLoggingHandler(CustomLogger): if time_to_first_token_response_time is not None: if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = ( - time_to_first_token_response_time.total_seconds() - ) + ttft_seconds = time_to_first_token_response_time.total_seconds() else: ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds( - ttft_seconds, completion_tokens - ) + time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) # ------------ # Update usage # ------------ parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) request_count_dict = ( - self.router_cache.get_cache( - key=latency_key, parent_otel_span=parent_otel_span - ) - or {} + self.router_cache.get_cache(key=latency_key, parent_otel_span=parent_otel_span) or {} ) if id not in request_count_dict: request_count_dict[id] = {} ## Latency - if ( - len(request_count_dict[id].get("latency", [])) - < self.routing_args.max_latency_list_size - ): + if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: request_count_dict[id].setdefault("latency", []).append(final_value) else: - request_count_dict[id]["latency"] = request_count_dict[id][ - "latency" - ][1:] + [final_value] + request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [final_value] ## Time to first token if time_to_first_token is not None: @@ -149,14 +131,11 @@ class LowestLatencyLoggingHandler(CustomLogger): len(request_count_dict[id].get("time_to_first_token", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault( - "time_to_first_token", [] - ).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = ( - request_count_dict[id]["time_to_first_token"][1:] - + [time_to_first_token] - ) + request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ + 1: + ] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -167,9 +146,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ) ## RPM - request_count_dict[id][precise_minute]["rpm"] = ( - request_count_dict[id][precise_minute].get("rpm", 0) + 1 - ) + request_count_dict[id][precise_minute]["rpm"] = request_count_dict[id][precise_minute].get("rpm", 0) + 1 self.router_cache.set_cache( key=latency_key, value=request_count_dict, ttl=self.routing_args.ttl @@ -197,13 +174,9 @@ class LowestLatencyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get(metadata_field) is None: pass else: - model_group = kwargs["litellm_params"][metadata_field].get( - "model_group", None - ) + model_group = kwargs["litellm_params"][metadata_field].get("model_group", None) - id = (kwargs["litellm_params"].get("model_info") or {}).get( - "id", None - ) + id = (kwargs["litellm_params"].get("model_info") or {}).get("id", None) if model_group is None or id is None: return elif isinstance(id, int): @@ -223,23 +196,16 @@ class LowestLatencyLoggingHandler(CustomLogger): } """ latency_key = f"{model_group}_map" - request_count_dict = ( - await self.router_cache.async_get_cache(key=latency_key) or {} - ) + request_count_dict = await self.router_cache.async_get_cache(key=latency_key) or {} if id not in request_count_dict: request_count_dict[id] = {} ## Latency - give 1000s penalty for failing - if ( - len(request_count_dict[id].get("latency", [])) - < self.routing_args.max_latency_list_size - ): + if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: request_count_dict[id].setdefault("latency", []).append(1000.0) else: - request_count_dict[id]["latency"] = request_count_dict[id][ - "latency" - ][1:] + [1000.0] + request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [1000.0] await self.router_cache.async_set_cache( key=latency_key, @@ -266,9 +232,7 @@ class LowestLatencyLoggingHandler(CustomLogger): if kwargs["litellm_params"].get(metadata_field) is None: pass else: - model_group = kwargs["litellm_params"][metadata_field].get( - "model_group", None - ) + model_group = kwargs["litellm_params"][metadata_field].get("model_group", None) id = (kwargs["litellm_params"].get("model_info") or {}).get("id", None) if model_group is None or id is None: @@ -301,9 +265,7 @@ class LowestLatencyLoggingHandler(CustomLogger): time_to_first_token_response_time = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = ( - kwargs.get("completion_start_time", end_time) - start_time - ) + time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time final_value: Union[float, timedelta] = response_ms total_tokens = 0 @@ -321,9 +283,7 @@ class LowestLatencyLoggingHandler(CustomLogger): else: response_seconds = response_ms - final_value = safe_divide_seconds( - response_seconds, completion_tokens - ) + final_value = safe_divide_seconds(response_seconds, completion_tokens) if final_value is not None: final_value = float(final_value) else: @@ -331,14 +291,10 @@ class LowestLatencyLoggingHandler(CustomLogger): if time_to_first_token_response_time is not None: if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = ( - time_to_first_token_response_time.total_seconds() - ) + ttft_seconds = time_to_first_token_response_time.total_seconds() else: ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds( - ttft_seconds, completion_tokens - ) + time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) # ------------ # Update usage # ------------ @@ -356,15 +312,10 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - if ( - len(request_count_dict[id].get("latency", [])) - < self.routing_args.max_latency_list_size - ): + if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: request_count_dict[id].setdefault("latency", []).append(final_value) else: - request_count_dict[id]["latency"] = request_count_dict[id][ - "latency" - ][1:] + [final_value] + request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [final_value] ## Time to first token if time_to_first_token is not None: @@ -372,14 +323,11 @@ class LowestLatencyLoggingHandler(CustomLogger): len(request_count_dict[id].get("time_to_first_token", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault( - "time_to_first_token", [] - ).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = ( - request_count_dict[id]["time_to_first_token"][1:] - + [time_to_first_token] - ) + request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ + 1: + ] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -390,9 +338,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ) ## RPM - request_count_dict[id][precise_minute]["rpm"] = ( - request_count_dict[id][precise_minute].get("rpm", 0) + 1 - ) + request_count_dict[id][precise_minute]["rpm"] = request_count_dict[id][precise_minute].get("rpm", 0) + 1 await self.router_cache.async_set_cache( key=latency_key, value=request_count_dict, ttl=self.routing_args.ttl @@ -510,9 +456,7 @@ class LowestLatencyLoggingHandler(CustomLogger): # -------------- # # We use _latency_per_deployment to log to langfuse, slack - this is not used to make a decision on routing # this helps a user to debug why the router picked a specfic deployment # - _deployment_api_base = _deployment.get("litellm_params", {}).get( - "api_base", "" - ) + _deployment_api_base = _deployment.get("litellm_params", {}).get("api_base", "") if _deployment_api_base is not None: _latency_per_deployment[_deployment_api_base] = item_latency # -------------- # @@ -520,8 +464,7 @@ class LowestLatencyLoggingHandler(CustomLogger): # -------------- # if ( - item_tpm + input_tokens > _deployment_tpm - or item_rpm + 1 > _deployment_rpm + item_tpm + input_tokens > _deployment_tpm or item_rpm + 1 > _deployment_rpm ): # if user passed in tpm / rpm in the model_list continue else: @@ -539,18 +482,14 @@ class LowestLatencyLoggingHandler(CustomLogger): # Find deployments within buffer of lowest latency buffer = self.routing_args.lowest_latency_buffer * lowest_latency - valid_deployments = [ - x for x in sorted_deployments if x[1] <= lowest_latency + buffer - ] + valid_deployments = [x for x in sorted_deployments if x[1] <= lowest_latency + buffer] # Pick a random deployment from valid deployments random_valid_deployment = random.choice(valid_deployments) deployment = random_valid_deployment[0] metadata_field = self._select_metadata_field(request_kwargs) if request_kwargs is not None and metadata_field in request_kwargs: - request_kwargs[metadata_field]["_latency_per_deployment"] = ( - _latency_per_deployment - ) + request_kwargs[metadata_field]["_latency_per_deployment"] = _latency_per_deployment return deployment async def async_get_available_deployments( @@ -564,14 +503,9 @@ class LowestLatencyLoggingHandler(CustomLogger): # get list of potential deployments latency_key = f"{model_group}_map" - parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( - request_kwargs - ) + parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs(request_kwargs) request_count_dict = ( - await self.router_cache.async_get_cache( - key=latency_key, parent_otel_span=parent_otel_span - ) - or {} + await self.router_cache.async_get_cache(key=latency_key, parent_otel_span=parent_otel_span) or {} ) return self._get_available_deployments( @@ -597,15 +531,8 @@ class LowestLatencyLoggingHandler(CustomLogger): # get list of potential deployments latency_key = f"{model_group}_map" - parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( - request_kwargs - ) - request_count_dict = ( - self.router_cache.get_cache( - key=latency_key, parent_otel_span=parent_otel_span - ) - or {} - ) + parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs(request_kwargs) + request_count_dict = self.router_cache.get_cache(key=latency_key, parent_otel_span=parent_otel_span) or {} return self._get_available_deployments( model_group, diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index f807ba7232a..89ad7526f20 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -36,9 +36,7 @@ class LowestTPMLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: @@ -63,17 +61,13 @@ class LowestTPMLoggingHandler(CustomLogger): request_count_dict = self.router_cache.get_cache(key=tpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens - self.router_cache.set_cache( - key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl - ) + self.router_cache.set_cache(key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl) ## RPM request_count_dict = self.router_cache.get_cache(key=rpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + 1 - self.router_cache.set_cache( - key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl - ) + self.router_cache.set_cache(key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl) ### TESTING ### if self.test_flag: @@ -97,9 +91,7 @@ class LowestTPMLoggingHandler(CustomLogger): if kwargs["litellm_params"].get("metadata") is None: pass else: - model_group = kwargs["litellm_params"]["metadata"].get( - "model_group", None - ) + model_group = kwargs["litellm_params"]["metadata"].get("model_group", None) model_info = kwargs["litellm_params"].get("model_info") id = None @@ -127,9 +119,7 @@ class LowestTPMLoggingHandler(CustomLogger): # update cache ## TPM - request_count_dict = ( - await self.router_cache.async_get_cache(key=tpm_key) or {} - ) + request_count_dict = await self.router_cache.async_get_cache(key=tpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens await self.router_cache.async_set_cache( @@ -137,9 +127,7 @@ class LowestTPMLoggingHandler(CustomLogger): ) ## RPM - request_count_dict = ( - await self.router_cache.async_get_cache(key=rpm_key) or {} - ) + request_count_dict = await self.router_cache.async_get_cache(key=rpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + 1 await self.router_cache.async_set_cache( @@ -179,9 +167,7 @@ class LowestTPMLoggingHandler(CustomLogger): tpm_dict = self.router_cache.get_cache(key=tpm_key) rpm_dict = self.router_cache.get_cache(key=rpm_key) - verbose_router_logger.debug( - f"tpm_key={tpm_key}, tpm_dict: {tpm_dict}, rpm_dict: {rpm_dict}" - ) + verbose_router_logger.debug(f"tpm_key={tpm_key}, tpm_dict: {tpm_dict}, rpm_dict: {rpm_dict}") try: input_tokens = token_counter(messages=messages, text=input) except Exception: @@ -238,9 +224,7 @@ class LowestTPMLoggingHandler(CustomLogger): if item_tpm + input_tokens > _deployment_tpm: continue - elif (rpm_dict is not None and item in rpm_dict) and ( - rpm_dict[item] + 1 >= _deployment_rpm - ): + elif (rpm_dict is not None and item in rpm_dict) and (rpm_dict[item] + 1 >= _deployment_rpm): continue elif item_tpm < lowest_tpm: lowest_tpm = item_tpm diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 09c2084c095..d63621a7d77 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -76,9 +76,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): deployment_name = deployment.get("litellm_params", {}).get("model") rpm_key = f"{model_id}:{deployment_name}:rpm:{current_minute}" - local_result = self.router_cache.get_cache( - key=rpm_key, local_only=True - ) # check local result first + local_result = self.router_cache.get_cache(key=rpm_key, local_only=True) # check local result first deployment_rpm = None if deployment_rpm is None: @@ -115,14 +113,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): else: # if local result below limit, check redis ## prevent unnecessary redis checks - result = self.router_cache.increment_cache( - key=rpm_key, value=1, ttl=self.routing_args.ttl - ) + result = self.router_cache.increment_cache(key=rpm_key, value=1, ttl=self.routing_args.ttl) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( - message="Deployment over defined rpm limit={}. current usage={}".format( - deployment_rpm, result - ), + message="Deployment over defined rpm limit={}. current usage={}".format(deployment_rpm, result), llm_provider="", model=deployment.get("litellm_params", {}).get("model"), response=httpx.Response( @@ -144,9 +138,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): raise e return deployment # don't fail calls if eg. redis fails to connect - async def async_pre_call_check( - self, deployment: Dict, parent_otel_span: Optional[Span] - ) -> Optional[Dict]: + async def async_pre_call_check(self, deployment: Dict, parent_otel_span: Optional[Span]) -> Optional[Dict]: """ Pre-call check + update model rpm - Used inside semaphore @@ -205,14 +197,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) else: # if local result below limit, check redis ## prevent unnecessary redis checks - result = await self._increment_value_in_current_window( - key=rpm_key, value=1, ttl=self.routing_args.ttl - ) + result = await self._increment_value_in_current_window(key=rpm_key, value=1, ttl=self.routing_args.ttl) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( - message="Deployment over defined rpm limit={}. current usage={}".format( - deployment_rpm, result - ), + message="Deployment over defined rpm limit={}. current usage={}".format(deployment_rpm, result), llm_provider="", model=deployment.get("litellm_params", {}).get("model"), response=httpx.Response( @@ -241,9 +229,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): """ Update TPM/RPM usage on success """ - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") model_group = standard_logging_object.get("model_group") @@ -260,9 +246,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): # Setup values # ------------ dt = get_utc_datetime() - current_minute = dt.strftime( - "%H-%M" - ) # use the same timezone regardless of system clock + current_minute = dt.strftime("%H-%M") # use the same timezone regardless of system clock tpm_key = f"{id}:{model}:tpm:{current_minute}" # ------------ @@ -271,17 +255,13 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): # update cache ## TPM - self.router_cache.increment_cache( - key=tpm_key, value=total_tokens, ttl=self.routing_args.ttl - ) + self.router_cache.increment_cache(key=tpm_key, value=total_tokens, ttl=self.routing_args.ttl) ### TESTING ### if self.test_flag: self.logged_success += 1 except Exception as e: verbose_logger.exception( - "litellm.proxy.hooks.lowest_tpm_rpm_v2.py::log_success_event(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.hooks.lowest_tpm_rpm_v2.py::log_success_event(): Exception occured - {}".format(str(e)) ) pass @@ -290,9 +270,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): """ Update TPM usage on success """ - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") model_group = standard_logging_object.get("model_group") @@ -307,9 +285,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): # Setup values # ------------ dt = get_utc_datetime() - current_minute = dt.strftime( - "%H-%M" - ) # use the same timezone regardless of system clock + current_minute = dt.strftime("%H-%M") # use the same timezone regardless of system clock tpm_key = f"{id}:{model}:tpm:{current_minute}" # ------------ @@ -346,8 +322,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): lowest_tpm = float("inf") potential_deployments = [] # if multiple deployments have the same low value deployment_lookup = { - deployment.get("model_info", {}).get("id"): deployment - for deployment in healthy_deployments + deployment.get("model_info", {}).get("id"): deployment for deployment in healthy_deployments } for item, item_tpm in all_deployments.items(): ## get the item from model list @@ -525,13 +500,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if _deployment_tpm is None: _deployment_tpm = _deployment.get("tpm", None) if _deployment_tpm is None: - _deployment_tpm = _deployment.get("litellm_params", {}).get( - "tpm", None - ) + _deployment_tpm = _deployment.get("litellm_params", {}).get("tpm", None) if _deployment_tpm is None: - _deployment_tpm = _deployment.get("model_info", {}).get( - "tpm", None - ) + _deployment_tpm = _deployment.get("model_info", {}).get("tpm", None) if _deployment_tpm is None: _deployment_tpm = float("inf") @@ -543,13 +514,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if _deployment_rpm is None: _deployment_rpm = _deployment.get("rpm", None) if _deployment_rpm is None: - _deployment_rpm = _deployment.get("litellm_params", {}).get( - "rpm", None - ) + _deployment_rpm = _deployment.get("litellm_params", {}).get("rpm", None) if _deployment_rpm is None: - _deployment_rpm = _deployment.get("model_info", {}).get( - "rpm", None - ) + _deployment_rpm = _deployment.get("model_info", {}).get("rpm", None) if _deployment_rpm is None: _deployment_rpm = float("inf") @@ -641,13 +608,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if _deployment_tpm is None: _deployment_tpm = _deployment.get("tpm", None) if _deployment_tpm is None: - _deployment_tpm = _deployment.get("litellm_params", {}).get( - "tpm", None - ) + _deployment_tpm = _deployment.get("litellm_params", {}).get("tpm", None) if _deployment_tpm is None: - _deployment_tpm = _deployment.get("model_info", {}).get( - "tpm", None - ) + _deployment_tpm = _deployment.get("model_info", {}).get("tpm", None) if _deployment_tpm is None: _deployment_tpm = float("inf") @@ -659,13 +622,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if _deployment_rpm is None: _deployment_rpm = _deployment.get("rpm", None) if _deployment_rpm is None: - _deployment_rpm = _deployment.get("litellm_params", {}).get( - "rpm", None - ) + _deployment_rpm = _deployment.get("litellm_params", {}).get("rpm", None) if _deployment_rpm is None: - _deployment_rpm = _deployment.get("model_info", {}).get( - "rpm", None - ) + _deployment_rpm = _deployment.get("model_info", {}).get("rpm", None) if _deployment_rpm is None: _deployment_rpm = float("inf") diff --git a/litellm/router_strategy/quality_router/quality_router.py b/litellm/router_strategy/quality_router/quality_router.py index a79b4384f5e..15c26f2c278 100644 --- a/litellm/router_strategy/quality_router/quality_router.py +++ b/litellm/router_strategy/quality_router/quality_router.py @@ -186,17 +186,10 @@ class QualityRouter(CustomLogger): else: # A Pydantic object of some other shape — coerce via its dict. prefs = RoutingPreferences( - **( - raw_prefs.model_dump() - if hasattr(raw_prefs, "model_dump") - else dict(raw_prefs) - ) + **(raw_prefs.model_dump() if hasattr(raw_prefs, "model_dump") else dict(raw_prefs)) ) except Exception as e: - raise ValueError( - f"QualityRouter: model '{name}' has invalid " - f"litellm_routing_preferences: {e}" - ) from e + raise ValueError(f"QualityRouter: model '{name}' has invalid litellm_routing_preferences: {e}") from e tier_int = int(prefs.quality_tier) tier_to_models.setdefault(tier_int, []).append(name) @@ -305,10 +298,7 @@ class QualityRouter(CustomLogger): if self.config.default_model: return self.config.default_model - raise ValueError( - f"QualityRouter: no model available for quality tier {tier} and " - f"no default_model configured" - ) + raise ValueError(f"QualityRouter: no model available for quality tier {tier} and no default_model configured") def _stash_decision( self, @@ -338,9 +328,7 @@ class QualityRouter(CustomLogger): from litellm.types.router import PreRoutingHookResponse if messages is None or len(messages) == 0: - verbose_router_logger.debug( - "QualityRouter: No messages provided, skipping routing" - ) + verbose_router_logger.debug("QualityRouter: No messages provided, skipping routing") return None # Extract last user message and last system prompt — same rules as @@ -353,9 +341,7 @@ class QualityRouter(CustomLogger): content = msg.get("content") or "" if isinstance(content, list): text_parts = [ - part.get("text", "") - for part in content - if isinstance(part, dict) and part.get("type") == "text" + part.get("text", "") for part in content if isinstance(part, dict) and part.get("type") == "text" ] content = " ".join(text_parts).strip() if isinstance(content, str) and content: @@ -365,13 +351,9 @@ class QualityRouter(CustomLogger): system_prompt = content if user_message is None: - verbose_router_logger.debug( - "QualityRouter: No user message found, routing to default model" - ) + verbose_router_logger.debug("QualityRouter: No user message found, routing to default model") if not self.config.default_model: - raise ValueError( - "QualityRouter: no user message and no default_model configured" - ) + raise ValueError("QualityRouter: no user message and no default_model configured") return PreRoutingHookResponse( model=self.config.default_model, messages=messages, @@ -404,14 +386,8 @@ class QualityRouter(CustomLogger): ) # No keyword match → complexity classification flow. - complexity_tier, score, signals = self._scorer.classify( - user_message, system_prompt - ) - complexity_name = ( - complexity_tier.value - if hasattr(complexity_tier, "value") - else str(complexity_tier) - ) + complexity_tier, score, signals = self._scorer.classify(user_message, system_prompt) + complexity_name = complexity_tier.value if hasattr(complexity_tier, "value") else str(complexity_tier) quality_tier = self.config.complexity_to_quality.get(complexity_name) if quality_tier is None: diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index f78acbfbd04..d3349ee29ce 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -43,9 +43,7 @@ def simple_shuffle( for weight_by in ["weight", "rpm", "tpm"]: weight = healthy_deployments[0].get("litellm_params").get(weight_by, None) if weight is not None: - weights = [ - m["litellm_params"].get(weight_by, 0) for m in healthy_deployments - ] + weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] verbose_router_logger.debug(f"\nweight {weights}") total_weight = sum(weights) if total_weight <= 0: diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 07143af38a2..65e76ba909b 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -43,9 +43,7 @@ def _is_valid_deployment_tag_regex( return None -def is_valid_deployment_tag( - deployment_tags: List[str], request_tags: List[str], match_any: bool = True -) -> bool: +def is_valid_deployment_tag(deployment_tags: List[str], request_tags: List[str], match_any: bool = True) -> bool: """ Check if a tag is valid, the matching can be either any or all based on `match_any` flag """ @@ -109,9 +107,7 @@ def _match_deployment( deployment_has_plain_tags = deployment_tags is not None and len(deployment_tags) > 0 strict_tag_check_failed = not match_any and deployment_has_plain_tags if deployment_tag_regex and header_strings and not strict_tag_check_failed: - regex_match = _is_valid_deployment_tag_regex( - deployment_tag_regex, header_strings - ) + regex_match = _is_valid_deployment_tag_regex(deployment_tag_regex, header_strings) if regex_match is not None: return {"matched_via": "tag_regex", "matched_value": regex_match} @@ -141,21 +137,15 @@ async def get_deployments_for_tag( return healthy_deployments if healthy_deployments is None: - verbose_logger.debug( - "get_deployments_for_tag: healthy_deployments is None returning healthy_deployments" - ) + verbose_logger.debug("get_deployments_for_tag: healthy_deployments is None returning healthy_deployments") return healthy_deployments # Tag filtering applies only when there is at least one deployment to evaluate. if isinstance(healthy_deployments, list) and len(healthy_deployments) == 0: - verbose_logger.debug( - "get_deployments_for_tag: empty candidate set; skipping tag filter" - ) + verbose_logger.debug("get_deployments_for_tag: empty candidate set; skipping tag filter") return healthy_deployments - verbose_logger.debug( - "request metadata: %s", request_kwargs.get(metadata_variable_name) - ) + verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name)) if metadata_variable_name in request_kwargs: metadata = request_kwargs[metadata_variable_name] request_tags = metadata.get("tags") @@ -174,12 +164,8 @@ async def get_deployments_for_tag( # behaviour for operators who use plain tags: a request that carries a # User-Agent (all proxy requests do) but targets deployments with no # tag_regex will continue to use the original tag-only code path. - has_regex_deployments = any( - d.get("litellm_params", {}).get("tag_regex") for d in healthy_deployments - ) - has_tag_filter = bool(request_tags) or ( - bool(header_strings) and has_regex_deployments - ) + has_regex_deployments = any(d.get("litellm_params", {}).get("tag_regex") for d in healthy_deployments) + has_tag_filter = bool(request_tags) or (bool(header_strings) and has_regex_deployments) if has_tag_filter: verbose_logger.debug( "get_deployments_for_tag routing: request_tags=%s user_agent=%s", @@ -225,11 +211,7 @@ async def get_deployments_for_tag( f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model={model} and tags={request_tags}" ) - return ( - new_healthy_deployments - if len(new_healthy_deployments) > 0 - else default_deployments - ) + return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments # for Untagged requests use default deployments if set _default_deployments_with_tags = [] diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index ddec753d362..b09c5eada79 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -8,9 +8,7 @@ from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose class InMemoryFile(io.BytesIO): - def __init__( - self, content: bytes, name: str, content_type: str = "application/jsonl" - ): + def __init__(self, content: bytes, name: str, content_type: str = "application/jsonl"): super().__init__(content) self.name = name self.content_type = content_type @@ -55,9 +53,7 @@ def parse_jsonl_with_embedded_newlines(content: str) -> List[dict]: json_object = json.loads(buffer.strip()) json_objects.append(json_object) except json.JSONDecodeError as e: - verbose_logger.error( - f"error parsing final buffer: {buffer[:100]}..., error: {e}" - ) + verbose_logger.error(f"error parsing final buffer: {buffer[:100]}..., error: {e}") raise e return json_objects @@ -109,17 +105,11 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File # the model rewrite is actually applied to tuple-wrapped upload handles; # otherwise a restricted body.model would survive and bypass the batch # model allowlist (which validates the upload target alias). - output = InMemoryFile( - b"", name="modified_file.jsonl", content_type="application/jsonl" - ) + output = InMemoryFile(b"", name="modified_file.jsonl", content_type="application/jsonl") wrote_any = False buffer = "" for raw_line in line_iter: # type: ignore[attr-defined] - buffer += ( - raw_line.decode("utf-8") - if isinstance(raw_line, (bytes, bytearray)) - else raw_line - ) + buffer += raw_line.decode("utf-8") if isinstance(raw_line, (bytes, bytearray)) else raw_line stripped = buffer.strip() if not stripped: buffer = "" @@ -128,13 +118,9 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File json_object = json.loads(stripped) except json.JSONDecodeError: continue # object not complete yet; keep accumulating - if isinstance(json_object, dict) and isinstance( - json_object.get("body"), dict - ): + if isinstance(json_object, dict) and isinstance(json_object.get("body"), dict): json_object["body"]["model"] = new_model_name - output.write( - (("\n" if wrote_any else "") + json.dumps(json_object)).encode("utf-8") - ) + output.write((("\n" if wrote_any else "") + json.dumps(json_object)).encode("utf-8")) wrote_any = True buffer = "" @@ -143,9 +129,7 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File # that followed it). Returning the partial `output` would silently # drop those rows; return the unchanged original so the provider # rejects the batch loudly instead of accepting a truncated one. - verbose_logger.error( - f"error parsing trailing batch content: {buffer[:100]}..." - ) + verbose_logger.error(f"error parsing trailing batch content: {buffer[:100]}...") if hasattr(source, "seek"): try: source.seek(0) # type: ignore[attr-defined] @@ -182,9 +166,7 @@ def _get_router_metadata_variable_name(function_name: Optional[str]) -> str: "_ageneric_api_call_with_fallbacks", ] ) - if function_name and any( - method in function_name for method in ROUTER_METHODS_USING_LITELLM_METADATA - ): + if function_name and any(method in function_name for method in ROUTER_METHODS_USING_LITELLM_METADATA): return "litellm_metadata" else: return "metadata" diff --git a/litellm/router_utils/client_initalization_utils.py b/litellm/router_utils/client_initalization_utils.py index e24d237853e..f192a3d2a8f 100644 --- a/litellm/router_utils/client_initalization_utils.py +++ b/litellm/router_utils/client_initalization_utils.py @@ -13,9 +13,7 @@ else: class InitalizeCachedClient: @staticmethod - def set_max_parallel_requests_client( - litellm_router_instance: LitellmRouter, model: dict - ): + def set_max_parallel_requests_client(litellm_router_instance: LitellmRouter, model: dict): litellm_params = model.get("litellm_params", {}) model_id = model["model_info"]["id"] rpm = litellm_params.get("rpm", None) diff --git a/litellm/router_utils/clientside_credential_handler.py b/litellm/router_utils/clientside_credential_handler.py index 45ade81b2dd..e992ef63658 100644 --- a/litellm/router_utils/clientside_credential_handler.py +++ b/litellm/router_utils/clientside_credential_handler.py @@ -28,11 +28,7 @@ def _admin_config_fields_to_clear_on_base_override() -> List[str]: """ from litellm.types.router import CredentialLiteLLMParams - typed_fields = [ - f - for f in CredentialLiteLLMParams.model_fields - if f not in clientside_credential_keys - ] + typed_fields = [f for f in CredentialLiteLLMParams.model_fields if f not in clientside_credential_keys] kwargs_only_fields = [ # Caller-supplied via **kwargs, not declared on CredentialLiteLLMParams. "organization", @@ -60,9 +56,7 @@ def _admin_config_fields_to_clear_on_base_override() -> List[str]: return typed_fields + kwargs_only_fields -_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE = ( - _admin_config_fields_to_clear_on_base_override() -) +_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE = _admin_config_fields_to_clear_on_base_override() def is_clientside_credential(request_kwargs: dict) -> bool: diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index f6da26ccd7f..18bce2348f6 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -14,9 +14,7 @@ def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: Hash of the credential params, used for mapping the file id to the right model """ sensitive_params = CredentialLiteLLMParams(**litellm_params) - return hashlib.sha256( - json.dumps(sensitive_params.model_dump()).encode() - ).hexdigest() + return hashlib.sha256(json.dumps(sensitive_params.model_dump()).encode()).hexdigest() def add_model_file_id_mappings( @@ -37,9 +35,7 @@ def add_model_file_id_mappings( """ model_file_id_mapping: Dict[str, str] = {} deployments_list: List[Dict] = ( - healthy_deployments - if isinstance(healthy_deployments, list) - else [healthy_deployments] + healthy_deployments if isinstance(healthy_deployments, list) else [healthy_deployments] ) for deployment, response in zip(deployments_list, responses): model_id = deployment.get("model_info", {}).get("id") @@ -62,9 +58,7 @@ def filter_team_based_models( metadata = request_kwargs.get("metadata") or {} litellm_metadata = request_kwargs.get("litellm_metadata") or {} - request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get( - "user_api_key_team_id" - ) + request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") ids_to_remove = set() if isinstance(healthy_deployments, dict): return healthy_deployments @@ -129,9 +123,7 @@ def filter_web_search_deployments( return healthy_deployments # Filter out deployments that don't support web search - final_deployments = [ - d for d in healthy_deployments if _deployment_supports_web_search(d) - ] + final_deployments = [d for d in healthy_deployments if _deployment_supports_web_search(d)] if len(healthy_deployments) > 0 and len(final_deployments) == 0: verbose_logger.warning("No deployments support web search for request") return final_deployments diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index dcfa44381c1..3f3284b315c 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -50,9 +50,7 @@ class CooldownCache: # Store the cooldown information for the deployment separately cooldown_data = CooldownCacheValue( - exception_received=self.exception_masker._mask_value( - str(original_exception) - ), + exception_received=self.exception_masker._mask_value(str(original_exception)), status_code=str(exception_status), timestamp=current_time, cooldown_time=cooldown_time, @@ -60,11 +58,7 @@ class CooldownCache: return cooldown_key, cooldown_data except Exception as e: - verbose_logger.error( - "CooldownCache::_common_add_cooldown_logic - Exception occurred - {}".format( - str(e) - ) - ) + verbose_logger.error("CooldownCache::_common_add_cooldown_logic - Exception occurred - {}".format(str(e))) raise e def add_deployment_to_cooldown( @@ -98,11 +92,7 @@ class CooldownCache: ttl=_cooldown_time, ) except Exception as e: - verbose_logger.error( - "CooldownCache::add_deployment_to_cooldown - Exception occurred - {}".format( - str(e) - ) - ) + verbose_logger.error("CooldownCache::add_deployment_to_cooldown - Exception occurred - {}".format(str(e))) raise e @staticmethod @@ -114,18 +104,14 @@ class CooldownCache: self, model_ids: List[str], parent_otel_span: Optional[Span] ) -> List[Tuple[str, CooldownCacheValue]]: # Generate the keys for the deployments - keys = [ - CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids - ] + keys = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] # Retrieve the values for the keys using mget ## more likely to be none if no models ratelimited. So just check redis every 1s ## each redis call adds ~100ms latency. ## check in memory cache first - results = await self.cache.async_batch_get_cache( - keys=keys, parent_otel_span=parent_otel_span - ) + results = await self.cache.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) active_cooldowns: List[Tuple[str, CooldownCacheValue]] = [] if results is None or all(v is None for v in results): @@ -143,14 +129,9 @@ class CooldownCache: self, model_ids: List[str], parent_otel_span: Optional[Span] ) -> List[Tuple[str, CooldownCacheValue]]: # Generate the keys for the deployments - keys = [ - CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids - ] + keys = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] # Retrieve the values for the keys using mget - results = ( - self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) - or [] - ) + results = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] active_cooldowns = [] # Process the results @@ -161,19 +142,14 @@ class CooldownCache: return active_cooldowns - def get_min_cooldown( - self, model_ids: List[str], parent_otel_span: Optional[Span] - ) -> float: + def get_min_cooldown(self, model_ids: List[str], parent_otel_span: Optional[Span]) -> float: """Return min cooldown time required for a group of model id's.""" # Generate the keys for the deployments keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results = ( - self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) - or [] - ) + results = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] min_cooldown_time: Optional[float] = None # Process the results diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 343328dacf3..313037a6364 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -41,10 +41,7 @@ async def router_cooldown_event_callback( temp_litellm_params = copy.deepcopy(_litellm_params) temp_litellm_params = dict(temp_litellm_params) _model_name = _deployment.get("model_name", None) or "" - _api_base = ( - litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params) - or "" - ) + _api_base = litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params) or "" model_info = _deployment["model_info"] model_id = model_info.id @@ -59,9 +56,7 @@ async def router_cooldown_event_callback( pass # get the prometheus logger from in memory loggers - prometheusLogger: Optional[PrometheusLogger] = ( - _get_prometheus_logger_from_callbacks() - ) + prometheusLogger: Optional[PrometheusLogger] = _get_prometheus_logger_from_callbacks() if prometheusLogger is not None: prometheusLogger.set_deployment_complete_outage( diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 81bfac2ad19..2bc2ed998ca 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -55,9 +55,7 @@ def _is_cooldown_required( """ try: ignored_strings = ["APIConnectionError"] - if ( - exception_str is not None - ): # don't cooldown on litellm api connection errors errors + if exception_str is not None: # don't cooldown on litellm api connection errors errors for ignored_string in ignored_strings: if ignored_string in exception_str: return False @@ -113,10 +111,7 @@ def _should_run_cooldown_logic( - deployment is in litellm_router_instance.provider_default_deployment_ids - exception_status is not one that should be immediately retried (e.g. 401) """ - if ( - deployment is None - or litellm_router_instance.get_model_group(id=deployment) is None - ): + if deployment is None or litellm_router_instance.get_model_group(id=deployment) is None: verbose_router_logger.debug( "Should Not Run Cooldown Logic: deployment id is none or model group can't be found." ) @@ -125,18 +120,12 @@ def _should_run_cooldown_logic( ######################################################### # If time_to_cooldown is 0 or 0.0000000, don't run cooldown logic ######################################################### - if time_to_cooldown is not None and math.isclose( - a=time_to_cooldown, b=0.0, abs_tol=1e-9 - ): - verbose_router_logger.debug( - "Should Not Run Cooldown Logic: time_to_cooldown is effectively 0" - ) + if time_to_cooldown is not None and math.isclose(a=time_to_cooldown, b=0.0, abs_tol=1e-9): + verbose_router_logger.debug("Should Not Run Cooldown Logic: time_to_cooldown is effectively 0") return False if litellm_router_instance.disable_cooldowns: - verbose_router_logger.debug( - "Should Not Run Cooldown Logic: disable_cooldowns is True" - ) + verbose_router_logger.debug("Should Not Run Cooldown Logic: disable_cooldowns is True") return False if deployment is None: @@ -149,15 +138,11 @@ def _should_run_cooldown_logic( exception_status=exception_status, exception_str=str(original_exception), ): - verbose_router_logger.debug( - "Should Not Run Cooldown Logic: _is_cooldown_required returned False" - ) + verbose_router_logger.debug("Should Not Run Cooldown Logic: _is_cooldown_required returned False") return False if deployment in litellm_router_instance.provider_default_deployment_ids: - verbose_router_logger.debug( - "Should Not Run Cooldown Logic: deployment is in provider_default_deployment_ids" - ) + verbose_router_logger.debug("Should Not Run Cooldown Logic: deployment is in provider_default_deployment_ids") return False return True @@ -194,10 +179,7 @@ def _should_cooldown_deployment( is_single_deployment_model_group = True if ( litellm_router_instance.allowed_fails_policy is None - and _is_allowed_fails_set_on_router( - litellm_router_instance=litellm_router_instance - ) - is False + and _is_allowed_fails_set_on_router(litellm_router_instance=litellm_router_instance) is False ): num_successes_this_minute = get_deployment_successes_for_current_minute( litellm_router_instance=litellm_router_instance, deployment_id=deployment @@ -209,9 +191,7 @@ def _should_cooldown_deployment( total_requests_this_minute = num_successes_this_minute + num_fails_this_minute percent_fails = 0.0 if total_requests_this_minute > 0: - percent_fails = num_fails_this_minute / ( - num_successes_this_minute + num_fails_this_minute - ) + percent_fails = num_fails_this_minute / (num_successes_this_minute + num_fails_this_minute) verbose_router_logger.debug( "percent fails for deployment = %s, percent fails = %s, num successes = %s, num fails = %s", deployment, @@ -223,11 +203,7 @@ def _should_cooldown_deployment( exception_status_int = cast_exception_status_to_int(exception_status) if exception_status_int == 429 and not is_single_deployment_model_group: return True - elif ( - percent_fails == 1.0 - and total_requests_this_minute - >= SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD - ): + elif percent_fails == 1.0 and total_requests_this_minute >= SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: # Cooldown if all requests failed and we have reasonable traffic return True elif ( @@ -238,12 +214,7 @@ def _should_cooldown_deployment( # Only apply error rate cooldown when we have enough requests to make the percentage meaningful return True - elif ( - litellm._should_retry( - status_code=cast_exception_status_to_int(exception_status) - ) - is False - ): + elif litellm._should_retry(status_code=cast_exception_status_to_int(exception_status)) is False: return True return False @@ -328,11 +299,9 @@ async def _async_get_cooldown_deployments( Async implementation of '_get_cooldown_deployments' """ model_ids = litellm_router_instance.get_model_ids() - cooldown_models = ( - await litellm_router_instance.cooldown_cache.async_get_active_cooldowns( - model_ids=model_ids, - parent_otel_span=parent_otel_span, - ) + cooldown_models = await litellm_router_instance.cooldown_cache.async_get_active_cooldowns( + model_ids=model_ids, + parent_otel_span=parent_otel_span, ) cached_value_deployment_ids = [] @@ -356,19 +325,15 @@ async def _async_get_cooldown_deployments_with_debug_info( Async implementation of '_get_cooldown_deployments' """ model_ids = litellm_router_instance.get_model_ids() - cooldown_models = ( - await litellm_router_instance.cooldown_cache.async_get_active_cooldowns( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) + cooldown_models = await litellm_router_instance.cooldown_cache.async_get_active_cooldowns( + model_ids=model_ids, parent_otel_span=parent_otel_span ) verbose_router_logger.debug(f"retrieve cooldown models: {cooldown_models}") return cooldown_models -def _get_cooldown_deployments( - litellm_router_instance: LitellmRouter, parent_otel_span: Optional[Span] -) -> List[str]: +def _get_cooldown_deployments(litellm_router_instance: LitellmRouter, parent_otel_span: Optional[Span]) -> List[str]: """ Get the list of models being cooled down for this minute """ @@ -413,9 +378,7 @@ def should_cooldown_based_on_allowed_fails_policy( ) or litellm_router_instance.allowed_fails ) - cooldown_time = ( - litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS - ) + cooldown_time = litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS current_fails = litellm_router_instance.failed_calls.get_cache(key=deployment) or 0 updated_fails = current_fails + 1 @@ -423,9 +386,7 @@ def should_cooldown_based_on_allowed_fails_policy( if updated_fails > allowed_fails: return True else: - litellm_router_instance.failed_calls.set_cache( - key=deployment, value=updated_fails, ttl=cooldown_time - ) + litellm_router_instance.failed_calls.set_cache(key=deployment, value=updated_fails, ttl=cooldown_time) return False diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 891d80d785a..17d29ec0da8 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -43,9 +43,7 @@ def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: return False -def get_fallback_model_group( - fallbacks: List[Any], model_group: str -) -> Tuple[Optional[List[str]], Optional[int]]: +def get_fallback_model_group(fallbacks: List[Any], model_group: str) -> Tuple[Optional[List[str]], Optional[int]]: """ Returns: - fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"] @@ -141,16 +139,12 @@ async def run_async_fallback( kwargs["max_fallbacks"] = max_fallbacks if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors - response = await litellm_router.async_function_with_fallbacks( - *args, **kwargs - ) + response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( response=response, attempted_fallbacks=fallback_depth, - fallback_errors=( - list(fallback_errors) if include_fallback_errors else None - ), + fallback_errors=(list(fallback_errors) if include_fallback_errors else None), ) # callback for successfull_fallback_event(): await log_success_fallback_event( @@ -170,9 +164,7 @@ async def run_async_fallback( raise error_from_fallbacks -async def log_success_fallback_event( - original_model_group: str, kwargs: dict, original_exception: Exception -): +async def log_success_fallback_event(original_model_group: str, kwargs: dict, original_exception: Exception): """ Log a successful fallback event to all registered callbacks. @@ -187,9 +179,7 @@ async def log_success_fallback_event( Errors during logging are caught and reported but do not interrupt the process. """ # Get deduplicated CustomLogger instances from all callback lists - custom_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( - CustomLogger - ) + custom_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type(CustomLogger) for _callback_custom_logger in custom_loggers: try: @@ -199,14 +189,10 @@ async def log_success_fallback_event( original_exception=original_exception, ) except Exception as e: - verbose_router_logger.error( - f"Error in log_success_fallback_event: {str(e)}" - ) + verbose_router_logger.error(f"Error in log_success_fallback_event: {str(e)}") -async def log_failure_fallback_event( - original_model_group: str, kwargs: dict, original_exception: Exception -): +async def log_failure_fallback_event(original_model_group: str, kwargs: dict, original_exception: Exception): """ Log a failed fallback event to all registered callbacks. @@ -221,9 +207,7 @@ async def log_failure_fallback_event( Errors during logging are caught and reported but do not interrupt the process. """ # Get deduplicated CustomLogger instances from all callback lists - custom_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( - CustomLogger - ) + custom_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type(CustomLogger) for _callback_custom_logger in custom_loggers: try: @@ -233,9 +217,7 @@ async def log_failure_fallback_event( original_exception=original_exception, ) except Exception as e: - verbose_router_logger.error( - f"Error in log_failure_fallback_event: {str(e)}" - ) + verbose_router_logger.error(f"Error in log_failure_fallback_event: {str(e)}") def _check_non_standard_fallback_format(fallbacks: Optional[List[Any]]) -> bool: @@ -264,7 +246,5 @@ def _check_non_standard_fallback_format(fallbacks: Optional[List[Any]]) -> bool: return False -def run_non_standard_fallback_format( - fallbacks: Union[List[str], List[Dict[str, Any]]], model_group: str -): +def run_non_standard_fallback_format(fallbacks: Union[List[str], List[Dict[str, Any]]], model_group: str): pass diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/router_utils/get_retry_from_policy.py index 162d6428f85..314917d3f56 100644 --- a/litellm/router_utils/get_retry_from_policy.py +++ b/litellm/router_utils/get_retry_from_policy.py @@ -31,11 +31,7 @@ def get_num_retries_from_retry_policy( """ # if we can find the exception then in the retry policy -> return the number of retries - if ( - model_group_retry_policy is not None - and model_group is not None - and model_group in model_group_retry_policy - ): + if model_group_retry_policy is not None and model_group is not None and model_group in model_group_retry_policy: retry_policy = model_group_retry_policy.get(model_group, None) # type: ignore if retry_policy is None: @@ -43,27 +39,18 @@ def get_num_retries_from_retry_policy( if isinstance(retry_policy, dict): retry_policy = RetryPolicy(**retry_policy) - if ( - isinstance(exception, AuthenticationError) - and retry_policy.AuthenticationErrorRetries is not None - ): + if isinstance(exception, AuthenticationError) and retry_policy.AuthenticationErrorRetries is not None: return retry_policy.AuthenticationErrorRetries if isinstance(exception, Timeout) and retry_policy.TimeoutErrorRetries is not None: return retry_policy.TimeoutErrorRetries - if ( - isinstance(exception, RateLimitError) - and retry_policy.RateLimitErrorRetries is not None - ): + if isinstance(exception, RateLimitError) and retry_policy.RateLimitErrorRetries is not None: return retry_policy.RateLimitErrorRetries if ( isinstance(exception, ContentPolicyViolationError) and retry_policy.ContentPolicyViolationErrorRetries is not None ): return retry_policy.ContentPolicyViolationErrorRetries - if ( - isinstance(exception, BadRequestError) - and retry_policy.BadRequestErrorRetries is not None - ): + if isinstance(exception, BadRequestError) and retry_policy.BadRequestErrorRetries is not None: return retry_policy.BadRequestErrorRetries diff --git a/litellm/router_utils/handle_error.py b/litellm/router_utils/handle_error.py index c23e6ce473a..05bde3d50d1 100644 --- a/litellm/router_utils/handle_error.py +++ b/litellm/router_utils/handle_error.py @@ -74,9 +74,7 @@ async def async_raise_no_deployment_exception( """ Raises a RouterRateLimitError if no deployment is found for the given model. """ - verbose_router_logger.info( - f"get_available_deployment for model: {model}, No deployment available" - ) + verbose_router_logger.info(f"get_available_deployment for model: {model}, No deployment available") model_ids = litellm_router_instance.get_model_ids(model_name=model) _cooldown_time = litellm_router_instance.cooldown_cache.get_min_cooldown( model_ids=model_ids, parent_otel_span=parent_otel_span diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 65b064f19d2..aeb2ba945f5 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -42,9 +42,7 @@ class DeploymentHealthCache: self.cache = cache self.staleness_threshold = staleness_threshold - def set_deployment_health_states( - self, states: Dict[str, DeploymentHealthStateValue] - ) -> None: + def set_deployment_health_states(self, states: Dict[str, DeploymentHealthStateValue]) -> None: """Bulk-write all deployment health states as a single cache entry.""" try: self.cache.set_cache( @@ -71,9 +69,7 @@ class DeploymentHealthCache: and (now - state.get("timestamp", 0)) < self.staleness_threshold } - async def async_get_unhealthy_deployment_ids( - self, parent_otel_span: Optional[Span] = None - ) -> Set[str]: + async def async_get_unhealthy_deployment_ids(self, parent_otel_span: Optional[Span] = None) -> Set[str]: """Return set of deployment IDs currently marked unhealthy and not stale.""" try: raw = await self.cache.async_get_cache(key=self.CACHE_KEY) @@ -85,9 +81,7 @@ class DeploymentHealthCache: ) return set() - def get_unhealthy_deployment_ids( - self, parent_otel_span: Optional[Span] = None - ) -> Set[str]: + def get_unhealthy_deployment_ids(self, parent_otel_span: Optional[Span] = None) -> Set[str]: """Sync version: return set of deployment IDs currently marked unhealthy and not stale.""" try: raw = self.cache.get_cache(key=self.CACHE_KEY) diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 48f85a83411..c08f8e95cf4 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -26,9 +26,7 @@ class PatternUtils: complexity_chars = ["*", "+", "?", "\\", "^", "$", "|", "(", ")"] ret_val = ( len(pattern), # Longer patterns more specific - sum( - pattern.count(char) for char in complexity_chars - ), # More regex complexity + sum(pattern.count(char) for char in complexity_chars), # More regex complexity ) return ret_val @@ -99,27 +97,19 @@ class PatternMatchRouter: # return f"^{regex}$" return re.escape(pattern).replace(r"\*", "(.*)") - def _return_pattern_matched_deployments( - self, matched_pattern: Match, deployments: List[Dict] - ) -> List[Dict]: + def _return_pattern_matched_deployments(self, matched_pattern: Match, deployments: List[Dict]) -> List[Dict]: new_deployments = [] for deployment in deployments: new_deployment = copy.deepcopy(deployment) - new_deployment["litellm_params"]["model"] = ( - PatternMatchRouter.set_deployment_model_name( - matched_pattern=matched_pattern, - litellm_deployment_litellm_model=deployment["litellm_params"][ - "model" - ], - ) + new_deployment["litellm_params"]["model"] = PatternMatchRouter.set_deployment_model_name( + matched_pattern=matched_pattern, + litellm_deployment_litellm_model=deployment["litellm_params"]["model"], ) new_deployments.append(new_deployment) return new_deployments - def route( - self, request: Optional[str], filtered_model_names: Optional[List[str]] = None - ) -> Optional[List[Dict]]: + def route(self, request: Optional[str], filtered_model_names: Optional[List[str]] = None) -> Optional[List[Dict]]: """ Route a requested model to the corresponding llm deployments based on the regex pattern @@ -139,15 +129,10 @@ class PatternMatchRouter: sorted_patterns = PatternUtils.sorted_patterns(self.patterns) regex_filtered_model_names = ( - [self._pattern_to_regex(m) for m in filtered_model_names] - if filtered_model_names is not None - else [] + [self._pattern_to_regex(m) for m in filtered_model_names] if filtered_model_names is not None else [] ) for pattern, llm_deployments in sorted_patterns: - if ( - filtered_model_names is not None - and pattern not in regex_filtered_model_names - ): + if filtered_model_names is not None and pattern not in regex_filtered_model_names: continue pattern_match = re.match(pattern, request) if pattern_match: @@ -201,20 +186,14 @@ class PatternMatchRouter: dynamic_segments = matched_pattern.groups() if len(dynamic_segments) > wildcard_count: - return ( - matched_pattern.string - ) # default to the user input, if unable to map based on wildcards. + return matched_pattern.string # default to the user input, if unable to map based on wildcards. # Replace the corresponding wildcards in the litellm model pattern with extracted segments for segment in dynamic_segments: - litellm_deployment_litellm_model = litellm_deployment_litellm_model.replace( - "*", segment, 1 - ) + litellm_deployment_litellm_model = litellm_deployment_litellm_model.replace("*", segment, 1) return litellm_deployment_litellm_model - def get_pattern( - self, model: str, custom_llm_provider: Optional[str] = None - ) -> Optional[List[Dict]]: + def get_pattern(self, model: str, custom_llm_provider: Optional[str] = None) -> Optional[List[Dict]]: """ Check if a pattern exists for the given model and custom llm provider @@ -238,9 +217,7 @@ class PatternMatchRouter: pass return self.route(model) or self.route(f"{custom_llm_provider}/{model}") - def get_deployments_by_pattern( - self, model: str, custom_llm_provider: Optional[str] = None - ) -> List[Dict]: + def get_deployments_by_pattern(self, model: str, custom_llm_provider: Optional[str] = None) -> List[Dict]: """ Get the deployments by pattern diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index d3e7e2ffa34..c3f58935ef0 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -62,9 +62,7 @@ class DeploymentAffinityCheck(CustomLogger): self.enable_user_key_affinity = enable_user_key_affinity self.enable_responses_api_affinity = enable_responses_api_affinity self.enable_session_id_affinity = enable_session_id_affinity - self.model_group_affinity_config: Dict[str, List[str]] = ( - model_group_affinity_config or {} - ) + self.model_group_affinity_config: Dict[str, List[str]] = model_group_affinity_config or {} for group, flags in self.model_group_affinity_config.items(): unknown = set(flags) - self.VALID_FLAGS if unknown: @@ -179,11 +177,7 @@ class DeploymentAffinityCheck(CustomLogger): return base_model litellm_model_name = litellm_params.get("model") if isinstance(litellm_model_name, str) and litellm_model_name: - return ( - DeploymentAffinityCheck._get_model_map_key_from_litellm_model_name( - litellm_model_name - ) - ) + return DeploymentAffinityCheck._get_model_map_key_from_litellm_model_name(litellm_model_name) return None @@ -270,9 +264,7 @@ class DeploymentAffinityCheck(CustomLogger): """ # 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 - ) + user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(metadata=metadata) if user_key is not None: return user_key @@ -281,25 +273,19 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _get_session_id_from_request_kwargs(request_kwargs: dict) -> Optional[str]: for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): - session_id = DeploymentAffinityCheck._get_session_id_from_metadata_dict( - metadata=metadata - ) + session_id = DeploymentAffinityCheck._get_session_id_from_metadata_dict(metadata=metadata) if session_id is not None: return session_id return None @staticmethod - def _find_deployment_by_model_id( - healthy_deployments: List[dict], model_id: str - ) -> Optional[dict]: + def _find_deployment_by_model_id(healthy_deployments: List[dict], model_id: str) -> Optional[dict]: for deployment in healthy_deployments: model_info = deployment.get("model_info") if not isinstance(model_info, dict): continue deployment_model_id = model_info.get("id") - if deployment_model_id is not None and str(deployment_model_id) == str( - model_id - ): + if deployment_model_id is not None and str(deployment_model_id) == str(model_id): return deployment return None @@ -329,11 +315,7 @@ class DeploymentAffinityCheck(CustomLogger): if enable_responses_api: previous_response_id = request_kwargs.get("previous_response_id") if previous_response_id is not None: - responses_model_id = ( - ResponsesAPIRequestUtils.get_model_id_from_response_id( - str(previous_response_id) - ) - ) + responses_model_id = ResponsesAPIRequestUtils.get_model_id_from_response_id(str(previous_response_id)) if responses_model_id is not None: deployment = self._find_deployment_by_model_id( healthy_deployments=typed_healthy_deployments, @@ -354,22 +336,16 @@ class DeploymentAffinityCheck(CustomLogger): # 2) Session-id -> deployment affinity if enable_session_id: - session_id = self._get_session_id_from_request_kwargs( - request_kwargs=request_kwargs - ) + session_id = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs) if session_id is not None: session_cache_key = self.get_session_affinity_cache_key( model_group=stable_model_map_key, session_id=session_id ) - session_cache_result = await self.cache.async_get_cache( - key=session_cache_key - ) + session_cache_result = await self.cache.async_get_cache(key=session_cache_key) session_model_id: Optional[str] = None if isinstance(session_cache_result, dict): - session_model_id = cast( - Optional[str], session_cache_result.get("model_id") - ) + session_model_id = cast(Optional[str], session_cache_result.get("model_id")) elif isinstance(session_cache_result, str): session_model_id = session_cache_result @@ -399,9 +375,7 @@ class DeploymentAffinityCheck(CustomLogger): if user_key is None: return typed_healthy_deployments - cache_key = self.get_affinity_cache_key( - model_group=stable_model_map_key, user_key=user_key - ) + cache_key = self.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key) cache_result = await self.cache.async_get_cache(key=cache_key) model_id: Optional[str] = None @@ -449,10 +423,7 @@ class DeploymentAffinityCheck(CustomLogger): deployment_model_name: Optional[str] = None for metadata in metadata_dicts: maybe_deployment_model_name = metadata.get("deployment_model_name") - if ( - isinstance(maybe_deployment_model_name, str) - and maybe_deployment_model_name - ): + if isinstance(maybe_deployment_model_name, str) and maybe_deployment_model_name: deployment_model_name = maybe_deployment_model_name break @@ -501,16 +472,12 @@ class DeploymentAffinityCheck(CustomLogger): model_id = model_info.get("id") if not model_id: - verbose_router_logger.warning( - "DeploymentAffinityCheck: model_id missing; skipping affinity cache update." - ) + verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.") return None if user_key is not None: try: - cache_key = self.get_affinity_cache_key( - model_group=deployment_model_name, user_key=user_key - ) + cache_key = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key) await self.cache.async_set_cache( cache_key, DeploymentAffinityCacheValue(model_id=str(model_id)), diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index cab2041964f..9788f5f2299 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -77,9 +77,7 @@ class EncryptedContentAffinityCheck(CustomLogger): super().__init__() self.router = router self.enable_global_affinity = enable_global_affinity - self.model_group_affinity_config: Dict[str, List[str]] = ( - model_group_affinity_config or {} - ) + self.model_group_affinity_config: Dict[str, List[str]] = model_group_affinity_config or {} # ------------------------------------------------------------------ # Helpers @@ -92,10 +90,7 @@ class EncryptedContentAffinityCheck(CustomLogger): if not model_group_affinity_config: return False - return any( - "encrypted_content_affinity" in checks - for checks in model_group_affinity_config.values() - ) + return any("encrypted_content_affinity" in checks for checks in model_group_affinity_config.values()) def _is_enabled_for_model_group(self, model_group: str) -> bool: group_checks = self.model_group_affinity_config.get(model_group) @@ -137,26 +132,20 @@ class EncryptedContentAffinityCheck(CustomLogger): ( model_id, _, - ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - encrypted_content - ) + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) if model_id: return model_id return None @staticmethod - def _find_deployment_by_model_id( - healthy_deployments: List[dict], model_id: str - ) -> Optional[dict]: + def _find_deployment_by_model_id(healthy_deployments: List[dict], model_id: str) -> Optional[dict]: for deployment in healthy_deployments: model_info = deployment.get("model_info") if not isinstance(model_info, dict): continue deployment_model_id = model_info.get("id") - if deployment_model_id is not None and str(deployment_model_id) == str( - model_id - ): + if deployment_model_id is not None and str(deployment_model_id) == str(model_id): return deployment return None @@ -204,15 +193,11 @@ class EncryptedContentAffinityCheck(CustomLogger): originating = self.router.get_deployment(model_id=model_id) if originating is None: return [], None - boundary = self._encryption_boundary_key( - originating.litellm_params.model_dump(exclude_none=True) - ) + boundary = self._encryption_boundary_key(originating.litellm_params.model_dump(exclude_none=True)) if boundary is None: return [], originating matches = [ - d - for d in healthy_deployments - if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary + d for d in healthy_deployments if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary ] return matches, originating @@ -252,9 +237,7 @@ class EncryptedContentAffinityCheck(CustomLogger): # _get_metadata_variable_name_from_kwargs would pick "litellm_metadata" # over "metadata" where tags are actually stored. if "litellm_metadata" in request_kwargs: - request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = ( - True - ) + request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True request_input = request_kwargs.get("input") model_id = self._extract_model_id_from_input(request_input) @@ -279,11 +262,9 @@ class EncryptedContentAffinityCheck(CustomLogger): return [deployment] # Follow-up switched model_name (LIT-2531): pin by Azure resource instead. - boundary_matches, originating = ( - self._find_deployments_on_same_encryption_boundary( - healthy_deployments=typed_healthy_deployments, - model_id=model_id, - ) + boundary_matches, originating = self._find_deployments_on_same_encryption_boundary( + healthy_deployments=typed_healthy_deployments, + model_id=model_id, ) if boundary_matches: verbose_router_logger.debug( @@ -327,9 +308,7 @@ class EncryptedContentAffinityCheck(CustomLogger): llm_provider="", ) - cooldown = await self._get_origin_cooldown( - model_id=model_id, parent_otel_span=parent_otel_span - ) + cooldown = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span) if cooldown is not None and str(cooldown.get("status_code")) == "429": retry_after = self._cooldown_seconds_remaining(cooldown) @@ -384,9 +363,5 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _cooldown_seconds_remaining(cooldown: CooldownCacheValue) -> int: - remaining = ( - float(cooldown.get("timestamp", 0.0)) - + float(cooldown.get("cooldown_time", 0.0)) - - time.time() - ) + remaining = float(cooldown.get("timestamp", 0.0)) + float(cooldown.get("cooldown_time", 0.0)) - time.time() return max(1, int(remaining)) diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index 836f9858744..0c6450b191f 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -46,9 +46,7 @@ class ModelRateLimitingCheck(CustomLogger): def __init__(self, dual_cache: DualCache): self.dual_cache = dual_cache - def _get_deployment_limits( - self, deployment: Dict - ) -> tuple[Optional[int], Optional[int]]: + def _get_deployment_limits(self, deployment: Dict) -> tuple[Optional[int], Optional[int]]: """ Extract TPM and RPM limits from a deployment configuration. @@ -131,9 +129,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check RPM limit (atomic increment-first to avoid race conditions) if rpm_limit is not None: - current_rpm = self.dual_cache.increment_cache( - key=rpm_key, value=1, ttl=RoutingArgs.ttl - ) + current_rpm = self.dual_cache.increment_cache(key=rpm_key, value=1, ttl=RoutingArgs.ttl) if current_rpm is not None and current_rpm > rpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", @@ -155,15 +151,11 @@ class ModelRateLimitingCheck(CustomLogger): except litellm.RateLimitError: raise except Exception as e: - verbose_router_logger.debug( - f"Error in ModelRateLimitingCheck.pre_call_check: {str(e)}" - ) + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.pre_call_check: {str(e)}") # Don't fail the request if rate limit check fails return deployment - async def async_pre_call_check( - self, deployment: Dict, parent_otel_span: Optional[Span] = None - ) -> Optional[Dict]: + async def async_pre_call_check(self, deployment: Dict, parent_otel_span: Optional[Span] = None) -> Optional[Dict]: """ Async pre-call check for model rate limits. @@ -187,9 +179,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check TPM limit if tpm_limit is not None: # First check local cache - current_tpm = await self.dual_cache.async_get_cache( - key=tpm_key, local_only=True - ) + current_tpm = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True) if current_tpm is not None and current_tpm >= tpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", @@ -237,9 +227,7 @@ class ModelRateLimitingCheck(CustomLogger): except litellm.RateLimitError: raise except Exception as e: - verbose_router_logger.debug( - f"Error in ModelRateLimitingCheck.async_pre_call_check: {str(e)}" - ) + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_pre_call_check: {str(e)}") # Don't fail the request if rate limit check fails return deployment @@ -251,9 +239,7 @@ class ModelRateLimitingCheck(CustomLogger): Always tracks tokens - the pre-call check handles enforcement. """ try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_object is None: return @@ -262,9 +248,7 @@ class ModelRateLimitingCheck(CustomLogger): return total_tokens = standard_logging_object.get("total_tokens", 0) - model = standard_logging_object.get("hidden_params", {}).get( - "litellm_model_name" - ) + model = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") verbose_router_logger.debug( f"[TPM TRACKING] model_id={model_id}, total_tokens={total_tokens}, model={model}" @@ -277,9 +261,7 @@ class ModelRateLimitingCheck(CustomLogger): current_minute = dt.strftime("%H-%M") tpm_key = f"{model_id}:{model}:tpm:{current_minute}" - verbose_router_logger.debug( - f"[TPM TRACKING] Incrementing {tpm_key} by {total_tokens}" - ) + verbose_router_logger.debug(f"[TPM TRACKING] Incrementing {tpm_key} by {total_tokens}") await self.dual_cache.async_increment_cache( key=tpm_key, @@ -288,9 +270,7 @@ class ModelRateLimitingCheck(CustomLogger): ) except Exception as e: - verbose_router_logger.debug( - f"Error in ModelRateLimitingCheck.async_log_success_event: {str(e)}" - ) + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.async_log_success_event: {str(e)}") def log_success_event(self, kwargs, response_obj, start_time, end_time): """ @@ -298,9 +278,7 @@ class ModelRateLimitingCheck(CustomLogger): Always tracks tokens - the pre-call check handles enforcement. """ try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") if standard_logging_object is None: return @@ -309,9 +287,7 @@ class ModelRateLimitingCheck(CustomLogger): return total_tokens = standard_logging_object.get("total_tokens", 0) - model = standard_logging_object.get("hidden_params", {}).get( - "litellm_model_name" - ) + model = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") if not model or not total_tokens: return @@ -327,6 +303,4 @@ class ModelRateLimitingCheck(CustomLogger): ) except Exception as e: - verbose_router_logger.debug( - f"Error in ModelRateLimitingCheck.log_success_event: {str(e)}" - ) + verbose_router_logger.debug(f"Error in ModelRateLimitingCheck.log_success_event: {str(e)}") diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index e9c4b69d8ef..581783980fd 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -49,9 +49,7 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: return diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index eab342e5402..faf632f3b40 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -41,9 +41,7 @@ class PromptCachingCache: return obj.dict() elif isinstance(obj, dict): # If the object is a dictionary, serialize it with sorted keys - return json.dumps( - obj, sort_keys=True, separators=(",", ":") - ) # Standardize serialization + return json.dumps(obj, sort_keys=True, separators=(",", ":")) # Standardize serialization elif isinstance(obj, list): # Serialize lists by ensuring each element is handled properly @@ -160,9 +158,7 @@ class PromptCachingCache: # Use serialize_object for consistent and stable serialization data_to_hash = {} if cacheable_messages is not None: - serialized_messages = PromptCachingCache.serialize_object( - cacheable_messages - ) + serialized_messages = PromptCachingCache.serialize_object(cacheable_messages) data_to_hash["messages"] = serialized_messages if tools is not None: serialized_tools = PromptCachingCache.serialize_object(tools) @@ -193,9 +189,7 @@ class PromptCachingCache: if cache_key is None: return None - self.cache.set_cache( - cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300 - ) + self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300) return None async def async_add_model_id( diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index 9bcbcd83653..260bb4cc03a 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -56,9 +56,7 @@ class SearchAPIRouter: try: from litellm.types.router import SearchToolTypedDict - verbose_router_logger.debug( - f"Adding {len(search_tools)} search tools to router" - ) + verbose_router_logger.debug(f"Adding {len(search_tools)} search tools to router") # Convert search tools to the format expected by the router router_search_tools: list = [] @@ -75,14 +73,10 @@ class SearchAPIRouter: # Update the router's search_tools list router_instance.search_tools = router_search_tools - verbose_router_logger.info( - f"Successfully updated router with {len(router_search_tools)} search tool(s)" - ) + verbose_router_logger.info(f"Successfully updated router with {len(router_search_tools)} search tool(s)") except Exception as e: - verbose_router_logger.exception( - f"Error updating router with search tools: {str(e)}" - ) + verbose_router_logger.exception(f"Error updating router with search tools: {str(e)}") raise e @staticmethod @@ -104,15 +98,11 @@ class SearchAPIRouter: ValueError: If no matching search tools are found """ matching_tools = [ - tool - for tool in router_instance.search_tools - if tool.get("search_tool_name") == search_tool_name + tool for tool in router_instance.search_tools if tool.get("search_tool_name") == search_tool_name ] if not matching_tools: - raise ValueError( - f"Search tool '{search_tool_name}' not found in router.search_tools" - ) + raise ValueError(f"Search tool '{search_tool_name}' not found in router.search_tools") return matching_tools @@ -138,14 +128,10 @@ class SearchAPIRouter: search_tool_name = kwargs.get("search_tool_name", kwargs.get("model")) if not search_tool_name: - raise ValueError( - "search_tool_name or model parameter is required for search" - ) + raise ValueError("search_tool_name or model parameter is required for search") # Set up kwargs for the fallback system - kwargs["model"] = ( - search_tool_name # Use model field for compatibility with fallback system - ) + kwargs["model"] = search_tool_name # Use model field for compatibility with fallback system kwargs["original_generic_function"] = original_function # Bind router_instance to the helper method using partial kwargs["original_function"] = partial( @@ -160,9 +146,7 @@ class SearchAPIRouter: metadata_variable_name="litellm_metadata", ) - available_search_tool_names = [ - tool.get("search_tool_name") for tool in router_instance.search_tools - ] + available_search_tool_names = [tool.get("search_tool_name") for tool in router_instance.search_tools] verbose_router_logger.debug( f"Inside SearchAPIRouter.async_search_with_fallbacks() - search_tool_name: {search_tool_name}, Available Search Tools: {available_search_tool_names}, kwargs: {kwargs}" ) @@ -221,17 +205,13 @@ class SearchAPIRouter: litellm_params = selected_tool.get("litellm_params", {}) search_provider = litellm_params.get("search_provider") if not search_provider: - raise ValueError( - f"search_provider not found in litellm_params for search tool '{search_tool_name}'" - ) + raise ValueError(f"search_provider not found in litellm_params for search tool '{search_tool_name}'") api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( tool_litellm_params=litellm_params, ) - verbose_router_logger.debug( - f"Selected search tool with provider: {search_provider}" - ) + verbose_router_logger.debug(f"Selected search tool with provider: {search_provider}") # Call the original search function with the provider config response = await original_generic_function( diff --git a/litellm/sandbox/main.py b/litellm/sandbox/main.py index 76d9994c683..c59552fbd9c 100644 --- a/litellm/sandbox/main.py +++ b/litellm/sandbox/main.py @@ -38,9 +38,7 @@ _LITELLM_INTERNAL_KWARGS = { def _get_config(provider: str) -> BaseSandboxConfig: - config = ProviderConfigManager.get_provider_sandbox_config( - SandboxProviders(provider) - ) + config = ProviderConfigManager.get_provider_sandbox_config(SandboxProviders(provider)) if config is None: raise ValueError(f"Code execution is not supported for provider: {provider}") return config @@ -140,15 +138,9 @@ async def acode_interpreter_tool( **forwarded, ) try: - return await config.arun_code( - container=container, code=code, api_key=api_key, **forwarded - ) + return await config.arun_code(container=container, code=code, api_key=api_key, **forwarded) finally: try: - await config.adelete_sandbox( - container=container, api_key=api_key, api_base=api_base, **forwarded - ) + await config.adelete_sandbox(container=container, api_key=api_key, api_base=api_base, **forwarded) except Exception as e: - litellm._logging.verbose_logger.debug( - f"sandbox: failed to delete ephemeral container: {e}" - ) + litellm._logging.verbose_logger.debug(f"sandbox: failed to delete ephemeral container: {e}") diff --git a/litellm/sandbox/sandbox_tools.py b/litellm/sandbox/sandbox_tools.py index 509b316ee21..44665af30a6 100644 --- a/litellm/sandbox/sandbox_tools.py +++ b/litellm/sandbox/sandbox_tools.py @@ -29,16 +29,12 @@ def _iter_valid_tools(tools: list[dict]) -> Iterator[tuple[str, dict]]: continue name = tool.get("sandbox_tool_name") if not name: - verbose_logger.warning( - "sandbox_tools: skipping entry missing 'sandbox_tool_name': %r", tool - ) + verbose_logger.warning("sandbox_tools: skipping entry missing 'sandbox_tool_name': %r", tool) continue params = tool.get("litellm_params") or {} provider = params.get("sandbox_provider") if not provider: - verbose_logger.warning( - "sandbox_tools: skipping entry missing 'sandbox_provider': %r", tool - ) + verbose_logger.warning("sandbox_tools: skipping entry missing 'sandbox_provider': %r", tool) continue yield ( name, diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 5309971eeda..814fd74333a 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -11,9 +11,7 @@ from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL class SchedulerCacheKeys(enum.Enum): queue = "scheduler:queue" - default_in_memory_ttl = ( - DEFAULT_IN_MEMORY_TTL # cache queue in-memory for 5s when redis cache available - ) + default_in_memory_ttl = DEFAULT_IN_MEMORY_TTL # cache queue in-memory for 5s when redis cache available class FlowItem(BaseModel): @@ -38,12 +36,8 @@ class Scheduler: if redis_cache is not None: # if redis-cache available frequently poll that instead of using in-memory. default_in_memory_ttl = SchedulerCacheKeys.default_in_memory_ttl.value - self.cache = DualCache( - redis_cache=redis_cache, default_in_memory_ttl=default_in_memory_ttl - ) - self.polling_interval = ( - polling_interval or DEFAULT_POLLING_INTERVAL - ) # default to 3ms + self.cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=default_in_memory_ttl) + self.polling_interval = polling_interval or DEFAULT_POLLING_INTERVAL # default to 3ms async def add_request(self, request: FlowItem): # We use the priority directly, as lower values indicate higher priority @@ -69,9 +63,7 @@ class Scheduler: """ queue = await self.get_queue(model_name=model_name) if not queue: - raise Exception( - "Incorrectly setup. Queue is invalid. Queue={}".format(queue) - ) + raise Exception("Incorrectly setup. Queue is invalid. Queue={}".format(queue)) # ------------ # Setup values @@ -101,17 +93,13 @@ class Scheduler: filtered_queue = [item for item in queue if item[1] != request_id] heapq.heapify(filtered_queue) # restore heap invariant after filtering await self.save_queue(queue=filtered_queue, model_name=model_name) - print_verbose( - f"Removed request_id: {request_id} from queue for model: {model_name}" - ) + print_verbose(f"Removed request_id: {request_id} from queue for model: {model_name}") async def peek(self, id: str, model_name: str, health_deployments: list) -> bool: """Return if the id is at the top of the queue. Don't pop the value from heap.""" queue = await self.get_queue(model_name=model_name) if not queue: - raise Exception( - "Incorrectly setup. Queue is invalid. Queue={}".format(queue) - ) + raise Exception("Incorrectly setup. Queue is invalid. Queue={}".format(queue)) # ------------ # Setup values diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 841c003dff3..9680446064a 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -33,9 +33,7 @@ def search_provider_cost_per_query( # Check for tiered pricing (e.g., Exa AI based on max_results) tiered_pricing = model_info.get("tiered_pricing") if tiered_pricing and isinstance(tiered_pricing, list): - max_results = (optional_params or {}).get( - "max_results", 10 - ) # default 10 results + max_results = (optional_params or {}).get("max_results", 10) # default 10 results cost_per_query = 0.0 for tier in tiered_pricing: diff --git a/litellm/search/main.py b/litellm/search/main.py index 15a797c8b4e..4e1920f4599 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -144,9 +144,7 @@ async def asearch( response = init_response if response is None: - raise ValueError( - f"Got an unexpected None response from the Search API: {response}" - ) + raise ValueError(f"Got an unexpected None response from the Search API: {response}") return response except Exception as e: @@ -235,18 +233,14 @@ def search( # Validate query parameter if not isinstance(query, (str, list)): - raise ValueError( - f"query must be a string or list of strings, got {type(query)}" - ) + raise ValueError(f"query must be a string or list of strings, got {type(query)}") if isinstance(query, list) and not all(isinstance(q, str) for q in query): raise ValueError("All items in query list must be strings") # Get provider config - search_provider_config: Optional[BaseSearchConfig] = ( - ProviderConfigManager.get_provider_search_config( - provider=SearchProviders(search_provider), - ) + search_provider_config: Optional[BaseSearchConfig] = ProviderConfigManager.get_provider_search_config( + provider=SearchProviders(search_provider), ) if search_provider_config is None: diff --git a/litellm/secret_managers/aws_secret_manager.py b/litellm/secret_managers/aws_secret_manager.py index 60d0a713eff..090e3ebcfa2 100644 --- a/litellm/secret_managers/aws_secret_manager.py +++ b/litellm/secret_managers/aws_secret_manager.py @@ -88,9 +88,7 @@ class AWSKeyManagementService_V2: raise ValueError("kms_client is None") encrypted_value = os.getenv(secret_name, None) if encrypted_value is None: - raise Exception( - "AWS KMS - Encrypted Value of Key={} is None".format(secret_name) - ) + raise Exception("AWS KMS - Encrypted Value of Key={} is None".format(secret_name)) if isinstance(encrypted_value, str) and encrypted_value.startswith("aws_kms/"): encrypted_value = encrypted_value.replace("aws_kms/", "") @@ -130,11 +128,9 @@ def decrypt_env_var() -> Dict[str, Any]: # iterate through env - for `aws_kms/` new_values = {} for k, v in os.environ.items(): - if ( - k is not None - and isinstance(k, str) - and k.lower().startswith("litellm_secret_aws_kms") - ) or (v is not None and isinstance(v, str) and v.startswith("aws_kms/")): + if (k is not None and isinstance(k, str) and k.lower().startswith("litellm_secret_aws_kms")) or ( + v is not None and isinstance(v, str) and v.startswith("aws_kms/") + ): decrypted_value = aws_kms.decrypt_value(secret_name=k) # reset env var k = re.sub("litellm_secret_aws_kms_", "", k, flags=re.IGNORECASE) diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index ef3c821caf1..2b24ea1ce61 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -92,27 +92,13 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): aws_kwargs = {} if key_management_settings is not None: aws_kwargs = { - "aws_region_name": getattr( - key_management_settings, "aws_region_name", None - ), - "aws_role_name": getattr( - key_management_settings, "aws_role_name", None - ), - "aws_session_name": getattr( - key_management_settings, "aws_session_name", None - ), - "aws_external_id": getattr( - key_management_settings, "aws_external_id", None - ), - "aws_profile_name": getattr( - key_management_settings, "aws_profile_name", None - ), - "aws_web_identity_token": getattr( - key_management_settings, "aws_web_identity_token", None - ), - "aws_sts_endpoint": getattr( - key_management_settings, "aws_sts_endpoint", None - ), + "aws_region_name": getattr(key_management_settings, "aws_region_name", None), + "aws_role_name": getattr(key_management_settings, "aws_role_name", None), + "aws_session_name": getattr(key_management_settings, "aws_session_name", None), + "aws_external_id": getattr(key_management_settings, "aws_external_id", None), + "aws_profile_name": getattr(key_management_settings, "aws_profile_name", None), + "aws_web_identity_token": getattr(key_management_settings, "aws_web_identity_token", None), + "aws_sts_endpoint": getattr(key_management_settings, "aws_sts_endpoint", None), "replica_regions": key_management_settings.replica_regions, } # Remove None values @@ -156,9 +142,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): ) try: - response = await async_client.post( - url=endpoint_url, headers=headers, data=body.decode("utf-8") - ) + response = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() return response.json()["SecretString"] except httpx.TimeoutException: @@ -209,9 +193,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): ) try: - response = sync_client.post( - url=endpoint_url, headers=headers, data=body.decode("utf-8") - ) + response = sync_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) return response.json()["SecretString"] except httpx.TimeoutException: raise ValueError("Timeout error occurred") @@ -242,9 +224,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): """ return json.loads(primary_secret_json_str or "{}") - def sync_read_secret_from_primary_secret( - self, secret_name: str, primary_secret_name: str - ) -> Optional[str]: + def sync_read_secret_from_primary_secret(self, secret_name: str, primary_secret_name: str) -> Optional[str]: """ Read a secret from the primary secret """ @@ -252,15 +232,11 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): primary_secret_kv_pairs = self._parse_primary_secret(primary_secret_json_str) return primary_secret_kv_pairs.get(secret_name) - async def async_read_secret_from_primary_secret( - self, secret_name: str, primary_secret_name: str - ) -> Optional[str]: + async def async_read_secret_from_primary_secret(self, secret_name: str, primary_secret_name: str) -> Optional[str]: """ Read a secret from the primary secret """ - primary_secret_json_str = await self.async_read_secret( - secret_name=primary_secret_name - ) + primary_secret_json_str = await self.async_read_secret(secret_name=primary_secret_name) primary_secret_kv_pairs = self._parse_primary_secret(primary_secret_json_str) return primary_secret_kv_pairs.get(secret_name) @@ -405,9 +381,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): ) try: - response = await async_client.post( - url=endpoint_url, headers=headers, data=body.decode("utf-8") - ) + response = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() return response.json() except httpx.HTTPStatusError as err: @@ -459,9 +433,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): ) try: - response = await async_client.post( - url=endpoint_url, headers=headers, data=body.decode("utf-8") - ) + response = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() return response.json() except httpx.HTTPStatusError as err: @@ -541,9 +513,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): ) try: - response = await async_client.post( - url=endpoint_url, headers=headers, data=body.decode("utf-8") - ) + response = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() return response.json() except httpx.HTTPStatusError as err: @@ -579,17 +549,12 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): optional_params["aws_external_id"] = self.aws_external_id if not optional_params.get("aws_profile_name") and self.aws_profile_name: optional_params["aws_profile_name"] = self.aws_profile_name - if ( - not optional_params.get("aws_web_identity_token") - and self.aws_web_identity_token - ): + if not optional_params.get("aws_web_identity_token") and self.aws_web_identity_token: optional_params["aws_web_identity_token"] = self.aws_web_identity_token if not optional_params.get("aws_sts_endpoint") and self.aws_sts_endpoint: optional_params["aws_sts_endpoint"] = self.aws_sts_endpoint - boto3_credentials_info = self._get_boto_credentials_from_optional_params( - optional_params - ) + boto3_credentials_info = self._get_boto_credentials_from_optional_params(optional_params) # Get endpoint _, endpoint_url = self.get_runtime_endpoint( @@ -614,9 +579,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): } # Sign request - request = AWSRequest( - method="POST", url=endpoint_url, data=body, headers=headers - ) + request = AWSRequest(method="POST", url=endpoint_url, data=body, headers=headers) SigV4Auth( boto3_credentials_info.credentials, "secretsmanager", diff --git a/litellm/secret_managers/base_secret_manager.py b/litellm/secret_managers/base_secret_manager.py index 32a244c5ea9..d33d76093c9 100644 --- a/litellm/secret_managers/base_secret_manager.py +++ b/litellm/secret_managers/base_secret_manager.py @@ -174,7 +174,5 @@ class BaseSecretManager(ABC): except httpx.TimeoutException: raise ValueError("Timeout error occurred") except Exception as e: - verbose_logger.exception( - "Error rotating secret in AWS Secrets Manager: %s", str(e) - ) + verbose_logger.exception("Error rotating secret in AWS Secrets Manager: %s", str(e)) raise diff --git a/litellm/secret_managers/custom_secret_manager_loader.py b/litellm/secret_managers/custom_secret_manager_loader.py index 8405740062d..ab8153c9436 100644 --- a/litellm/secret_managers/custom_secret_manager_loader.py +++ b/litellm/secret_managers/custom_secret_manager_loader.py @@ -30,9 +30,7 @@ def load_custom_secret_manager(config_file_path: Optional[str] = None) -> None: """ if not config_file_path: - raise ValueError( - "CustomSecretManagerException - config_file_path is required to load custom secret manager" - ) + raise ValueError("CustomSecretManagerException - config_file_path is required to load custom secret manager") # Get the custom_secret_manager class path from settings if litellm._key_management_settings is None: @@ -40,9 +38,7 @@ def load_custom_secret_manager(config_file_path: Optional[str] = None) -> None: "CustomSecretManagerException - key_management_settings is required with custom_secret_manager field" ) - custom_secret_manager_path = getattr( - litellm._key_management_settings, "custom_secret_manager", None - ) + custom_secret_manager_path = getattr(litellm._key_management_settings, "custom_secret_manager", None) if not custom_secret_manager_path: raise ValueError( @@ -64,9 +60,7 @@ def load_custom_secret_manager(config_file_path: Optional[str] = None) -> None: spec = importlib.util.spec_from_file_location(_class_name, module_file_path) # type: ignore if not spec: - raise ImportError( - f"Could not find a module specification for {module_file_path}" - ) + raise ImportError(f"Could not find a module specification for {module_file_path}") module = importlib.util.module_from_spec(spec) # type: ignore spec.loader.exec_module(module) # type: ignore @@ -74,9 +68,7 @@ def load_custom_secret_manager(config_file_path: Optional[str] = None) -> None: # Validate that it's a CustomSecretManager subclass if not issubclass(_secret_manager_class, CustomSecretManager): - raise TypeError( - f"CustomSecretManagerException - {_class_name} must be a subclass of CustomSecretManager" - ) + raise TypeError(f"CustomSecretManagerException - {_class_name} must be a subclass of CustomSecretManager") # Instantiate the custom secret manager _secret_manager_instance = _secret_manager_class() diff --git a/litellm/secret_managers/cyberark_secret_manager.py b/litellm/secret_managers/cyberark_secret_manager.py index 11b853412fa..faf6224757f 100644 --- a/litellm/secret_managers/cyberark_secret_manager.py +++ b/litellm/secret_managers/cyberark_secret_manager.py @@ -143,9 +143,7 @@ class CyberArkSecretManager(BaseSecretManager): except httpx.HTTPStatusError as e: # Variable might already exist, which is fine if e.response.status_code in [409, 422]: - verbose_logger.debug( - f"Variable {secret_name} already exists or policy conflict (expected)" - ) + verbose_logger.debug(f"Variable {secret_name} already exists or policy conflict (expected)") else: verbose_logger.warning( f"Could not ensure variable exists: {e.response.status_code} - {e.response.text}" @@ -165,9 +163,7 @@ class CyberArkSecretManager(BaseSecretManager): """ # URL-encode the secret name to handle slashes and special characters encoded_name = quote(secret_name, safe="") - return ( - f"{self.conjur_addr}/secrets/{self.conjur_account}/variable/{encoded_name}" - ) + return f"{self.conjur_addr}/secrets/{self.conjur_account}/variable/{encoded_name}" async def async_read_secret( self, @@ -207,13 +203,9 @@ class CyberArkSecretManager(BaseSecretManager): except httpx.HTTPStatusError as e: if e.response.status_code == 404: - verbose_logger.debug( - f"Secret {secret_name} not found in CyberArk Conjur" - ) + verbose_logger.debug(f"Secret {secret_name} not found in CyberArk Conjur") else: - verbose_logger.exception( - f"Error reading secret from CyberArk Conjur: {e}" - ) + verbose_logger.exception(f"Error reading secret from CyberArk Conjur: {e}") return None except Exception as e: verbose_logger.exception(f"Error reading secret from CyberArk Conjur: {e}") @@ -254,13 +246,9 @@ class CyberArkSecretManager(BaseSecretManager): except httpx.HTTPStatusError as e: if e.response.status_code == 404: - verbose_logger.debug( - f"Secret {secret_name} not found in CyberArk Conjur" - ) + verbose_logger.debug(f"Secret {secret_name} not found in CyberArk Conjur") else: - verbose_logger.exception( - f"Error reading secret from CyberArk Conjur: {e}" - ) + verbose_logger.exception(f"Error reading secret from CyberArk Conjur: {e}") return None except Exception as e: verbose_logger.exception(f"Error reading secret from CyberArk Conjur: {e}") @@ -300,9 +288,7 @@ class CyberArkSecretManager(BaseSecretManager): # Now set the secret value url = self.get_url(secret_name) - response = await async_client.post( - url=url, headers=self._get_request_headers(), content=secret_value - ) + response = await async_client.post(url=url, headers=self._get_request_headers(), content=secret_value) response.raise_for_status() # Update cache @@ -337,8 +323,7 @@ class CyberArkSecretManager(BaseSecretManager): dict: Response indicating operation not supported """ verbose_logger.warning( - "CyberArk Conjur does not support direct secret deletion. " - "Secrets must be removed through policy updates." + "CyberArk Conjur does not support direct secret deletion. Secrets must be removed through policy updates." ) # Clear from cache diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index 184d959b964..2c52054964d 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -61,21 +61,14 @@ def get_azure_ad_token_provider( ) if azure_scope is None: - azure_scope = ( - os.environ.get("AZURE_SCOPE") - or "https://cognitiveservices.azure.com/.default" - ) + azure_scope = os.environ.get("AZURE_SCOPE") or "https://cognitiveservices.azure.com/.default" cred: str = ( azure_credential.value if azure_credential - else None - or os.environ.get("AZURE_CREDENTIAL") - or infer_credential_type_from_environment() - ) - verbose_logger.info( - f"For Azure AD Token Provider, choosing credential type: {cred}" + else None or os.environ.get("AZURE_CREDENTIAL") or infer_credential_type_from_environment() ) + verbose_logger.info(f"For Azure AD Token Provider, choosing credential type: {cred}") credential: Optional[ Union[ ClientSecretCredential, diff --git a/litellm/secret_managers/google_kms.py b/litellm/secret_managers/google_kms.py index 18e25abeb24..d22cd0b38b3 100644 --- a/litellm/secret_managers/google_kms.py +++ b/litellm/secret_managers/google_kms.py @@ -17,13 +17,9 @@ from litellm.proxy._types import KeyManagementSystem def validate_environment(): if "GOOGLE_APPLICATION_CREDENTIALS" not in os.environ: - raise ValueError( - "Missing required environment variable - GOOGLE_APPLICATION_CREDENTIALS" - ) + raise ValueError("Missing required environment variable - GOOGLE_APPLICATION_CREDENTIALS") if "GOOGLE_KMS_RESOURCE_NAME" not in os.environ: - raise ValueError( - "Missing required environment variable - GOOGLE_KMS_RESOURCE_NAME" - ) + raise ValueError("Missing required environment variable - GOOGLE_KMS_RESOURCE_NAME") def load_google_kms(use_google_kms: Optional[bool]): diff --git a/litellm/secret_managers/google_secret_manager.py b/litellm/secret_managers/google_secret_manager.py index 2fd35ced6e8..91284d5eb30 100644 --- a/litellm/secret_managers/google_secret_manager.py +++ b/litellm/secret_managers/google_secret_manager.py @@ -37,23 +37,14 @@ class GoogleSecretManager(GCSBucketBase): self.sync_httpx_client = _get_httpx_client() litellm.secret_manager_client = self litellm._key_management_system = KeyManagementSystem.GOOGLE_SECRET_MANAGER - _refresh_interval = os.environ.get( - "GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL", refresh_interval - ) - _refresh_interval = ( - int(_refresh_interval) if _refresh_interval else refresh_interval - ) - self.cache = InMemoryCache( - default_ttl=_refresh_interval - ) # store in memory for 1 day + _refresh_interval = os.environ.get("GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL", refresh_interval) + _refresh_interval = int(_refresh_interval) if _refresh_interval else refresh_interval + self.cache = InMemoryCache(default_ttl=_refresh_interval) # store in memory for 1 day _always_read_secret_manager = os.environ.get( "GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER", ) - if ( - _always_read_secret_manager - and _always_read_secret_manager.lower() == "true" - ): + if _always_read_secret_manager and _always_read_secret_manager.lower() == "true": self.always_read_secret_manager = True else: # by default this should be False, we want to use in memory caching for this. It's a bad idea to fetch from secret manager for all requests @@ -76,9 +67,7 @@ class GoogleSecretManager(GCSBucketBase): if secret_name in self.cache.cache_dict: return cached_secret - _secret_name = ( - f"projects/{self.PROJECT_ID}/secrets/{secret_name}/versions/latest" - ) + _secret_name = f"projects/{self.PROJECT_ID}/secrets/{secret_name}/versions/latest" headers = self.sync_construct_request_headers() url = f"https://secretmanager.googleapis.com/v1/{_secret_name}:access" @@ -86,15 +75,9 @@ class GoogleSecretManager(GCSBucketBase): response = self.sync_httpx_client.get(url=url, headers=headers) if response.status_code != 200: - verbose_logger.error( - "Google Secret Manager retrieval error: %s", str(response.text) - ) - self.cache.set_cache( - secret_name, None - ) # Cache that the secret was not found - raise ValueError( - f"secret {secret_name} not found in Google Secret Manager. Error: {response.text}" - ) + verbose_logger.error("Google Secret Manager retrieval error: %s", str(response.text)) + self.cache.set_cache(secret_name, None) # Cache that the secret was not found + raise ValueError(f"secret {secret_name} not found in Google Secret Manager. Error: {response.text}") verbose_logger.debug( "Google Secret Manager retrieval response status code: %s", @@ -108,9 +91,7 @@ class GoogleSecretManager(GCSBucketBase): # decode the base64 encoded value if _base64_encoded_value is not None: _decoded_value = base64.b64decode(_base64_encoded_value).decode("utf-8") - self.cache.set_cache( - secret_name, _decoded_value - ) # Cache the retrieved secret + self.cache.set_cache(secret_name, _decoded_value) # Cache the retrieved secret return _decoded_value self.cache.set_cache(secret_name, None) # Cache that the secret was not found diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index 8bb3f801a1e..bd1b1097347 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -51,17 +51,9 @@ class HashicorpSecretManager(BaseSecretManager): litellm.secret_manager_client = self litellm._key_management_system = KeyManagementSystem.HASHICORP_VAULT - _refresh_interval = os.environ.get( - "HCP_VAULT_REFRESH_INTERVAL", SECRET_MANAGER_REFRESH_INTERVAL - ) - _refresh_interval = ( - int(_refresh_interval) - if _refresh_interval - else SECRET_MANAGER_REFRESH_INTERVAL - ) - self.cache = InMemoryCache( - default_ttl=_refresh_interval - ) # store in memory for 1 day + _refresh_interval = os.environ.get("HCP_VAULT_REFRESH_INTERVAL", SECRET_MANAGER_REFRESH_INTERVAL) + _refresh_interval = int(_refresh_interval) if _refresh_interval else SECRET_MANAGER_REFRESH_INTERVAL + self.cache = InMemoryCache(default_ttl=_refresh_interval) # store in memory for 1 day def _verify_required_credentials_exist(self) -> None: """ @@ -147,9 +139,7 @@ class HashicorpSecretManager(BaseSecretManager): ) # Cache the token with its lease duration - self.cache.set_cache( - key="hcp_vault_approle_token", value=token, ttl=_lease_duration - ) + self.cache.set_cache(key="hcp_vault_approle_token", value=token, ttl=_lease_duration) return token except Exception as e: raise RuntimeError(f"Could not authenticate to Vault via AppRole: {e}") @@ -204,9 +194,7 @@ class HashicorpSecretManager(BaseSecretManager): token = resp.json()["auth"]["client_token"] _lease_duration = resp.json()["auth"]["lease_duration"] verbose_logger.debug("Successfully obtained Vault token via TLS cert auth.") - self.cache.set_cache( - key="hcp_vault_token", value=token, ttl=_lease_duration - ) + self.cache.set_cache(key="hcp_vault_token", value=token, ttl=_lease_duration) return token except Exception as e: raise RuntimeError(f"Could not authenticate to Vault via TLS cert: {e}") @@ -232,12 +220,8 @@ class HashicorpSecretManager(BaseSecretManager): - With custom mount: http://127.0.0.1:8200/v1/kv/data/mykey - With path prefix: http://127.0.0.1:8200/v1/secret/data/myapp/mykey """ - resolved_namespace = self._sanitize_path_component( - namespace if namespace is not None else self.vault_namespace - ) - resolved_mount = self._sanitize_path_component( - mount_name if mount_name is not None else self.vault_mount_name - ) + resolved_namespace = self._sanitize_path_component(namespace if namespace is not None else self.vault_namespace) + resolved_mount = self._sanitize_path_component(mount_name if mount_name is not None else self.vault_mount_name) if resolved_mount is None: resolved_mount = "secret" resolved_path_prefix = self._sanitize_path_component( @@ -261,18 +245,14 @@ class HashicorpSecretManager(BaseSecretManager): return None return value_str - def _sanitize_path_component( - self, value: Optional[Union[str, int]] - ) -> Optional[str]: + def _sanitize_path_component(self, value: Optional[Union[str, int]]) -> Optional[str]: sanitized_value = self._sanitize_plain_value(value) if sanitized_value is None: return None sanitized_value = sanitized_value.strip("/") return sanitized_value or None - def _extract_secret_manager_settings( - self, optional_params: Optional[dict] - ) -> Dict[str, Any]: + def _extract_secret_manager_settings(self, optional_params: Optional[dict]) -> Dict[str, Any]: if not isinstance(optional_params, dict): return {} @@ -281,9 +261,7 @@ class HashicorpSecretManager(BaseSecretManager): allowed_keys = {"namespace", "mount", "path_prefix", "data"} return {k: source[k] for k in allowed_keys if k in source} - def _build_secret_target( - self, secret_name: str, optional_params: Optional[dict] - ) -> Dict[str, Any]: + def _build_secret_target(self, secret_name: str, optional_params: Optional[dict]) -> Dict[str, Any]: settings = self._extract_secret_manager_settings(optional_params) namespace = settings.get("namespace", self.vault_namespace) @@ -470,9 +448,7 @@ class HashicorpSecretManager(BaseSecretManager): try: # First verify the old secret exists using _build_secret_target - current_target = self._build_secret_target( - current_secret_name, optional_params - ) + current_target = self._build_secret_target(current_secret_name, optional_params) try: response = await async_client.get( url=current_target["url"], @@ -482,9 +458,7 @@ class HashicorpSecretManager(BaseSecretManager): # Secret exists, we can proceed except httpx.HTTPStatusError as e: if e.response.status_code == 404: - verbose_logger.exception( - f"Current secret {current_secret_name} not found" - ) + verbose_logger.exception(f"Current secret {current_secret_name} not found") return { "status": "error", "message": f"Current secret {current_secret_name} not found", @@ -497,9 +471,7 @@ class HashicorpSecretManager(BaseSecretManager): "message": f"HTTP error occurred while checking current secret: {e.response.text if hasattr(e, 'response') else str(e)}", } except Exception as e: - verbose_logger.exception( - f"Error checking current secret existence: {e}" - ) + verbose_logger.exception(f"Error checking current secret existence: {e}") return { "status": "error", "message": f"Error checking current secret: {e}", @@ -516,10 +488,7 @@ class HashicorpSecretManager(BaseSecretManager): ) # Check if async_write_secret returned an error - if ( - isinstance(create_response, dict) - and create_response.get("status") == "error" - ): + if isinstance(create_response, dict) and create_response.get("status") == "error": return create_response # Verify new secret was created successfully using _build_secret_target @@ -533,9 +502,7 @@ class HashicorpSecretManager(BaseSecretManager): json_resp = response.json() # Use data_key from target to get the correct value data_key = new_target["data_key"] - new_secret_value_from_vault = ( - json_resp.get("data", {}).get("data", {}).get(data_key, None) - ) + new_secret_value_from_vault = json_resp.get("data", {}).get("data", {}).get(data_key, None) if new_secret_value_from_vault != new_secret_value: verbose_logger.exception( f"New secret value mismatch. Expected: {new_secret_value}, Got: {new_secret_value_from_vault}" @@ -546,9 +513,7 @@ class HashicorpSecretManager(BaseSecretManager): } except httpx.HTTPStatusError as e: if e.response.status_code == 404: - verbose_logger.exception( - f"Failed to verify new secret {new_secret_name}" - ) + verbose_logger.exception(f"Failed to verify new secret {new_secret_name}") return { "status": "error", "message": f"Failed to verify new secret {new_secret_name}", @@ -577,10 +542,7 @@ class HashicorpSecretManager(BaseSecretManager): timeout=timeout, ) # Check if async_delete_secret returned an error - if ( - isinstance(delete_response, dict) - and delete_response.get("status") == "error" - ): + if isinstance(delete_response, dict) and delete_response.get("status") == "error": # Log the error but don't fail the rotation since new secret was created successfully verbose_logger.warning( f"Failed to delete old secret {current_secret_name} after rotation: {delete_response.get('message')}" @@ -628,9 +590,7 @@ class HashicorpSecretManager(BaseSecretManager): try: target = self._build_secret_target(secret_name, optional_params) - response = await async_client.delete( - url=target["url"], headers=self._get_request_headers() - ) + response = await async_client.delete(url=target["url"], headers=self._get_request_headers()) response.raise_for_status() # Clear the cache for this secret @@ -646,9 +606,7 @@ class HashicorpSecretManager(BaseSecretManager): verbose_logger.exception(f"Error deleting secret from Hashicorp Vault: {e}") return {"status": "error", "message": str(e)} - def _get_secret_value_from_json_response( - self, json_resp: Optional[dict] - ) -> Optional[str]: + def _get_secret_value_from_json_response(self, json_resp: Optional[dict]) -> Optional[str]: """ Get the secret value from the JSON response diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index f4b1d4a1b69..fa55828a249 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -207,10 +207,7 @@ def get_secret( # https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers#using-custom-actions actions_id_token_request_url = os.getenv("ACTIONS_ID_TOKEN_REQUEST_URL") actions_id_token_request_token = os.getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN") - if ( - actions_id_token_request_url is None - or actions_id_token_request_token is None - ): + if actions_id_token_request_url is None or actions_id_token_request_token is None: raise ValueError( "ACTIONS_ID_TOKEN_REQUEST_URL or ACTIONS_ID_TOKEN_REQUEST_TOKEN not found in environment" ) @@ -278,10 +275,7 @@ def get_secret( raise ValueError("Unsupported OIDC provider") try: - if ( - _should_read_secret_from_secret_manager() - and litellm.secret_manager_client is not None - ): + if _should_read_secret_from_secret_manager() and litellm.secret_manager_client is not None: try: client = litellm.secret_manager_client key_manager = "local" @@ -319,9 +313,7 @@ def get_secret( else: secret = os.environ.get(secret_name) secret_value_as_bool = str_to_bool(secret) if secret is not None else None - if secret_value_as_bool is not None and isinstance( - secret_value_as_bool, bool - ): + if secret_value_as_bool is not None and isinstance(secret_value_as_bool, bool): return secret_value_as_bool else: return secret diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index 3a3cf6272dc..e2fb0b900b8 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -49,20 +49,16 @@ def get_secret_from_manager( if ( key_manager == KeyManagementSystem.AZURE_KEY_VAULT.value - or type(client).__module__ + "." + type(client).__name__ - == "azure.keyvault.secrets._client.SecretClient" + or type(client).__module__ + "." + type(client).__name__ == "azure.keyvault.secrets._client.SecretClient" ): # support Azure Secret Client - from azure.keyvault.secrets import SecretClient secret = client.get_secret(secret_name).value elif ( - key_manager == KeyManagementSystem.GOOGLE_KMS.value - or client.__class__.__name__ == "KeyManagementServiceClient" + key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient" ): encrypted_secret: Any = os.getenv(secret_name) if encrypted_secret is None: - raise ValueError( - "Google KMS requires the encrypted secret to be in the environment!" - ) + raise ValueError("Google KMS requires the encrypted secret to be in the environment!") b64_flag = _is_base64(encrypted_secret) if b64_flag is True: # if passed in as encoded b64 string encrypted_secret = base64.b64decode(encrypted_secret) @@ -77,9 +73,7 @@ def get_secret_from_manager( "ciphertext": ciphertext, } ) - secret = response.plaintext.decode( - "utf-8" - ) # assumes the original value was encoded with utf-8 + secret = response.plaintext.decode("utf-8") # assumes the original value was encoded with utf-8 elif key_manager == KeyManagementSystem.AWS_KMS.value: """ @@ -87,9 +81,7 @@ def get_secret_from_manager( """ encrypted_value = os.getenv(secret_name, None) if encrypted_value is None: - raise Exception( - "AWS KMS - Encrypted Value of Key={} is None".format(secret_name) - ) + raise Exception("AWS KMS - Encrypted Value of Key={} is None".format(secret_name)) # Decode the base64 encoded ciphertext ciphertext_blob = base64.b64decode(encrypted_value) @@ -123,13 +115,9 @@ def get_secret_from_manager( elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: try: secret = client.get_secret_from_google_secret_manager(secret_name) - print_verbose( - f"secret from google secret manager: [set={secret is not None}]" - ) + print_verbose(f"secret from google secret manager: [set={secret is not None}]") if secret is None: - raise ValueError( - f"No secret found in Google Secret Manager for {secret_name}" - ) + raise ValueError(f"No secret found in Google Secret Manager for {secret_name}") except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e @@ -138,9 +126,7 @@ def get_secret_from_manager( try: secret = client.sync_read_secret(secret_name=secret_name) if secret is None: - raise ValueError( - f"No secret found in Hashicorp Secret Manager for {secret_name}" - ) + raise ValueError(f"No secret found in Hashicorp Secret Manager for {secret_name}") except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e @@ -149,9 +135,7 @@ def get_secret_from_manager( try: secret = client.sync_read_secret(secret_name=secret_name) if secret is None: - raise ValueError( - f"No secret found in CyberArk Secret Manager for {secret_name}" - ) + raise ValueError(f"No secret found in CyberArk Secret Manager for {secret_name}") except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e @@ -163,16 +147,10 @@ def get_secret_from_manager( if isinstance(client, CustomSecretManager): secret = client.sync_read_secret( secret_name=secret_name, - optional_params=( - key_management_settings.model_dump() - if key_management_settings - else None - ), + optional_params=(key_management_settings.model_dump() if key_management_settings else None), ) if secret is None: - raise ValueError( - f"No secret found in Custom Secret Manager for {secret_name}" - ) + raise ValueError(f"No secret found in Custom Secret Manager for {secret_name}") else: raise ValueError( f"Custom secret manager client must be an instance of CustomSecretManager, got {type(client).__name__}" diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index 2f0cb1233ae..b43590079fa 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -184,11 +184,7 @@ def _styled_input(prompt: str) -> str: def _yaml_escape(value: str) -> str: """Escape a string for safe embedding in a double-quoted YAML scalar.""" return ( - value.replace("\\", "\\\\") - .replace('"', '\\"') - .replace("\n", "\\n") - .replace("\r", "\\r") - .replace("\t", "\\t") + value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t") ) @@ -242,9 +238,7 @@ class SetupWizard: config_path = Path(os.getcwd()) / "litellm_config.yaml" try: - config_path.write_text( - SetupWizard._build_config(providers, env_vars, master_key) - ) + config_path.write_text(SetupWizard._build_config(providers, env_vars, master_key)) except OSError as exc: print(f"\n {bold(_CROSS + ' Could not write config:')} {exc}") print(" Try running from a directory you have write access to.\n") @@ -283,9 +277,7 @@ class SetupWizard: @staticmethod def _read_key() -> str: """Read one keypress from /dev/tty in raw mode.""" - assert ( - termios is not None and tty is not None - ) # only called when _HAS_RAW_TERMINAL + assert termios is not None and tty is not None # only called when _HAS_RAW_TERMINAL with open("/dev/tty", "rb") as tty_fh: fd = tty_fh.fileno() old = termios.tcgetattr(fd) @@ -366,11 +358,7 @@ class SetupWizard: """Number-based fallback when raw terminal input is unavailable.""" print() print(f" {bold('Add your first model')}") - print( - grey( - " Enter numbers separated by commas (e.g. 1,2). Press Enter to confirm." - ) - ) + print(grey(" Enter numbers separated by commas (e.g. 1,2). Press Enter to confirm.")) print() for i, p in enumerate(PROVIDERS, 1): print(f" {grey(str(i) + '.')} {bold(p['name'])} {grey(p['description'])}") @@ -382,11 +370,7 @@ class SetupWizard: print(grey(" Please select at least one provider.")) continue try: - nums = [ - int(x.strip()) - for x in raw.replace(" ", ",").split(",") - if x.strip() - ] + nums = [int(x.strip()) for x in raw.replace(" ", ",").split(",") if x.strip()] valid = sorted({n for n in nums if 1 <= n <= len(PROVIDERS)}) if not valid: print(grey(f" Enter numbers between 1 and {len(PROVIDERS)}.")) @@ -405,44 +389,30 @@ class SetupWizard: print() print(f" {bold('Enter your API keys')}") print(grey(" Keys are stored only in the generated config file.")) - print( - grey( - " Tip: add litellm_config.yaml to .gitignore to avoid committing secrets." - ) - ) + print(grey(" Tip: add litellm_config.yaml to .gitignore to avoid committing secrets.")) print() for p in providers: if p["env_key"] is None: - print( - f" {green(p['name'])}: {grey('no key needed (uses local Ollama)')}" - ) + print(f" {green(p['name'])}: {grey('no key needed (uses local Ollama)')}") continue key = SetupWizard._prompt_key(p) if not key: continue - for extra_key, extra_hint in zip( - p.get("extra_keys", []), p.get("extra_hints", []) - ): + for extra_key, extra_hint in zip(p.get("extra_keys", []), p.get("extra_hints", [])): val = _styled_input(f" {blue('❯')} {extra_key} {grey(extra_hint)}: ") if val: env_vars[extra_key] = val if p.get("needs_api_base"): - api_base = _styled_input( - f" {blue('❯')} Azure endpoint URL {grey(p.get('api_base_hint', ''))}: " - ) + api_base = _styled_input(f" {blue('❯')} Azure endpoint URL {grey(p.get('api_base_hint', ''))}: ") if api_base: env_vars[f"_LITELLM_AZURE_AI_API_BASE_{p['id'].upper()}"] = api_base - deployment = _styled_input( - f" {blue('❯')} Azure deployment name {grey('(e.g. my-gpt4o)')}: " - ) + deployment = _styled_input(f" {blue('❯')} Azure deployment name {grey('(e.g. my-gpt4o)')}: ") if deployment: - env_vars[f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}"] = ( - deployment - ) + env_vars[f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}"] = deployment # Store the key returned by validation — may be a re-entered replacement env_vars[p["env_key"]] = SetupWizard._validate_and_report(p, key) @@ -454,9 +424,7 @@ class SetupWizard: """Prompt for a provider's API key, with skip option. Returns the key or ''.""" hint = grey(provider.get("key_hint", "")) while True: - key = _styled_input( - f" {blue('❯')} {bold(provider['name'])} API key {hint}: " - ) + key = _styled_input(f" {blue('❯')} {bold(provider['name'])} API key {hint}: ") if key: return key print(grey(" Key is required. Leave blank to skip this provider.")) @@ -480,22 +448,15 @@ class SetupWizard: ) valid = check_valid_key(model=test_model, api_key=api_key) if valid: - print( - f" {green(_CHECK)} {bold(provider['name'])} connected successfully" - ) + print(f" {green(_CHECK)} {bold(provider['name'])} connected successfully") return api_key print(f" {_CROSS} {bold(provider['name'])} {grey('— invalid API key')}") - if ( - _styled_input(f" {blue('❯')} Re-enter key? {grey('(y/N)')}: ").lower() - != "y" - ): + if _styled_input(f" {blue('❯')} Re-enter key? {grey('(y/N)')}: ").lower() != "y": return api_key hint = grey(provider.get("key_hint", "")) - new_key = _styled_input( - f" {blue('❯')} {bold(provider['name'])} API key {hint}: " - ) + new_key = _styled_input(f" {blue('❯')} {bold(provider['name'])} API key {hint}: ") if not new_key: return api_key api_key = new_key @@ -539,9 +500,7 @@ class SetupWizard: continue if p["id"] == "azure": - deployment = env_copy.pop( - f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}", "" - ) + deployment = env_copy.pop(f"_LITELLM_AZURE_DEPLOYMENT_{p['id'].upper()}", "") if not deployment: continue # skip Azure entirely if no deployment name was provided models = [f"azure/{deployment}"] @@ -560,15 +519,11 @@ class SetupWizard: if p["env_key"] and p["env_key"] in env_copy: lines.append(f" api_key: os.environ/{p['env_key']}") if p.get("api_base"): - lines.append( - f' api_base: "{_yaml_escape(str(p["api_base"]))}"' - ) + lines.append(f' api_base: "{_yaml_escape(str(p["api_base"]))}"') elif p.get("needs_api_base"): azure_base_key = f"_LITELLM_AZURE_AI_API_BASE_{p['id'].upper()}" if azure_base_key in env_copy: - lines.append( - f' api_base: "{_yaml_escape(env_copy.pop(azure_base_key))}"' - ) + lines.append(f' api_base: "{_yaml_escape(env_copy.pop(azure_base_key))}"') if p.get("api_version"): lines.append(f" api_version: {p['api_version']}") @@ -611,18 +566,12 @@ class SetupWizard: @staticmethod def _offer_start(config_path: Path, port: int, master_key: str) -> None: - start = _styled_input( - f" {blue('❯')} Start the proxy now? {grey('(Y/n)')}: " - ).lower() + start = _styled_input(f" {blue('❯')} Start the proxy now? {grey('(Y/n)')}: ").lower() if start not in ("", "y", "yes"): print() - print( - f" Run {bold(f'litellm --config {config_path}')} whenever you're ready." - ) + print(f" Run {bold(f'litellm --config {config_path}')} whenever you're ready.") print() - print( - grey(f" Quick test once running: curl http://localhost:{port}/health") - ) + print(grey(f" Quick test once running: curl http://localhost:{port}/health")) print() return diff --git a/litellm/skills/main.py b/litellm/skills/main.py index cee811d84fd..00ee89e5e77 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -50,9 +50,7 @@ def _get_skill_request_metadata( return extra_body["metadata"] metadata = kwargs.get("metadata") - if isinstance(metadata, dict) and isinstance( - metadata.get("requester_metadata"), dict - ): + if isinstance(metadata, dict) and isinstance(metadata.get("requester_metadata"), dict): return metadata["requester_metadata"] return None @@ -208,9 +206,7 @@ def create_skill( # Validate environment and get headers headers = extra_headers or {} - headers = skills_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = skills_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request request_body = skills_api_provider_config.transform_create_skill_request( @@ -223,9 +219,7 @@ def create_skill( from litellm.llms.anthropic.common_utils import AnthropicModelInfo api_base = AnthropicModelInfo.get_api_base(litellm_params.api_base) - url = skills_api_provider_config.get_complete_url( - api_base=api_base, endpoint="skills" - ) + url = skills_api_provider_config.get_complete_url(api_base=api_base, endpoint="skills") # Pre-call logging litellm_logging_obj.update_from_kwargs( @@ -403,9 +397,7 @@ def list_skills( # Validate environment and get headers headers = extra_headers or {} - headers = skills_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = skills_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request url, query_params = skills_api_provider_config.transform_list_skills_request( @@ -566,9 +558,7 @@ def get_skill( # Validate environment and get headers headers = extra_headers or {} - headers = skills_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = skills_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Get API base from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -734,9 +724,7 @@ def delete_skill( # Validate environment and get headers headers = extra_headers or {} - headers = skills_api_provider_config.validate_environment( - headers=headers, litellm_params=litellm_params - ) + headers = skills_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Get API base from litellm.llms.anthropic.common_utils import AnthropicModelInfo diff --git a/litellm/timeout.py b/litellm/timeout.py index 0d03a3e45e8..044b2b1c56f 100644 --- a/litellm/timeout.py +++ b/litellm/timeout.py @@ -70,9 +70,7 @@ def timeout(timeout_duration: float = 0.0, exception_to_raise=Timeout): elif "request_timeout" in kwargs and kwargs["request_timeout"] is not None: local_timeout_duration = kwargs["request_timeout"] try: - value = await asyncio.wait_for( - func(*args, **kwargs), timeout=timeout_duration - ) + value = await asyncio.wait_for(func(*args, **kwargs), timeout=timeout_duration) return value except asyncio.TimeoutError: model = args[0] if len(args) > 0 else kwargs["model"] diff --git a/litellm/types/agents.py b/litellm/types/agents.py index f34631b5600..c6347c17fd2 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -354,9 +354,7 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): LiteLLMSendMessageResponse with _hidden_params support """ response_dict = response.model_dump(mode="json", exclude_none=True) - response_dict = _normalize_a2a_jsonrpc_response( - response_dict, request_id=request_id - ) + response_dict = _normalize_a2a_jsonrpc_response(response_dict, request_id=request_id) return cls(**response_dict) @classmethod @@ -375,6 +373,4 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): Returns: LiteLLMSendMessageResponse with _hidden_params support """ - return cls( - **_normalize_a2a_jsonrpc_response(response_dict, request_id=request_id) - ) + return cls(**_normalize_a2a_jsonrpc_response(response_dict, request_id=request_id)) diff --git a/litellm/types/completion.py b/litellm/types/completion.py index a91f6234fad..380662836aa 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -66,9 +66,7 @@ class ChatCompletionContentPartImageParam(TypedDict, total=False): """The type of the content part.""" -ChatCompletionContentPartParam = Union[ - ChatCompletionContentPartTextParam, ChatCompletionContentPartImageParam -] +ChatCompletionContentPartParam = Union[ChatCompletionContentPartTextParam, ChatCompletionContentPartImageParam] class ChatCompletionUserMessageParam(TypedDict, total=False): diff --git a/litellm/types/containers/main.py b/litellm/types/containers/main.py index 0b0bef39e18..0377426bc93 100644 --- a/litellm/types/containers/main.py +++ b/litellm/types/containers/main.py @@ -132,9 +132,7 @@ class ContainerFileObject(BaseModel): """Represents a container file object.""" id: str - object: Literal[ - "container.file", "container_file" - ] # OpenAI returns "container.file" + object: Literal["container.file", "container_file"] # OpenAI returns "container.file" container_id: str bytes: Optional[int] = None # Can be null for some files created_at: int diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index 5c76736ab90..b2e1fb3d46b 100644 --- a/litellm/types/google_genai/main.py +++ b/litellm/types/google_genai/main.py @@ -23,9 +23,7 @@ if TYPE_CHECKING: generationConfig: Optional[Any] tools: Optional[ToolConfigDict] # type: ignore[assignment, valid-type] - class GenerateContentResponse( - GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject - ): # type: ignore[misc, valid-type] + class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc, valid-type] _hidden_params: dict = {} pass diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index e286c6f7e0d..889e029b902 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -381,17 +381,13 @@ class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): default=None, description="Path to a JSON file containing ad-hoc recognizers for Presidio", ) - mock_redacted_text: Optional[dict] = Field( - default=None, description="Mock redacted text for testing" - ) + mock_redacted_text: Optional[dict] = Field(default=None, description="Mock redacted text for testing") class BedrockGuardrailConfigModel(BaseModel): """Configuration parameters for the AWS Bedrock guardrail""" - guardrailIdentifier: Optional[str] = Field( - default=None, description="The ID of your guardrail on Bedrock" - ) + guardrailIdentifier: Optional[str] = Field(default=None, description="The ID of your guardrail on Bedrock") guardrailVersion: Optional[str] = Field( default=None, description="The version of your Bedrock guardrail (e.g., DRAFT or version number)", @@ -400,59 +396,29 @@ class BedrockGuardrailConfigModel(BaseModel): default=False, description="If True, will not raise an exception when the guardrail is blocked. Useful for OpenWebUI where exceptions can end the chat flow.", ) - aws_region_name: Optional[str] = Field( - default=None, description="AWS region where your guardrail is deployed" - ) - aws_access_key_id: Optional[str] = Field( - default=None, description="AWS access key ID for authentication" - ) - aws_secret_access_key: Optional[str] = Field( - default=None, description="AWS secret access key for authentication" - ) - aws_session_token: Optional[str] = Field( - default=None, description="AWS session token for temporary credentials" - ) - aws_session_name: Optional[str] = Field( - default=None, description="Name of the AWS session" - ) - aws_profile_name: Optional[str] = Field( - default=None, description="AWS profile name for credential retrieval" - ) - aws_role_name: Optional[str] = Field( - default=None, description="AWS role name for assuming roles" - ) + aws_region_name: Optional[str] = Field(default=None, description="AWS region where your guardrail is deployed") + aws_access_key_id: Optional[str] = Field(default=None, description="AWS access key ID for authentication") + aws_secret_access_key: Optional[str] = Field(default=None, description="AWS secret access key for authentication") + aws_session_token: Optional[str] = Field(default=None, description="AWS session token for temporary credentials") + aws_session_name: Optional[str] = Field(default=None, description="Name of the AWS session") + aws_profile_name: Optional[str] = Field(default=None, description="AWS profile name for credential retrieval") + aws_role_name: Optional[str] = Field(default=None, description="AWS role name for assuming roles") aws_web_identity_token: Optional[str] = Field( default=None, description="Web identity token for AWS role assumption" ) - aws_sts_endpoint: Optional[str] = Field( - default=None, description="AWS STS endpoint URL" - ) - aws_bedrock_runtime_endpoint: Optional[str] = Field( - default=None, description="AWS Bedrock runtime endpoint URL" - ) + aws_sts_endpoint: Optional[str] = Field(default=None, description="AWS STS endpoint URL") + aws_bedrock_runtime_endpoint: Optional[str] = Field(default=None, description="AWS Bedrock runtime endpoint URL") class LakeraV2GuardrailConfigModel(BaseModel): """Configuration parameters for the Lakera AI v2 guardrail""" - api_key: Optional[str] = Field( - default=None, description="API key for the Lakera AI service" - ) - api_base: Optional[str] = Field( - default=None, description="Base URL for the Lakera AI API" - ) - project_id: Optional[str] = Field( - default=None, description="Project ID for the Lakera AI project" - ) - payload: Optional[bool] = Field( - default=True, description="Whether to include payload in the response" - ) - breakdown: Optional[bool] = Field( - default=True, description="Whether to include breakdown in the response" - ) - metadata: Optional[Dict] = Field( - default=None, description="Additional metadata to include in the request" - ) + api_key: Optional[str] = Field(default=None, description="API key for the Lakera AI service") + api_base: Optional[str] = Field(default=None, description="Base URL for the Lakera AI API") + project_id: Optional[str] = Field(default=None, description="Project ID for the Lakera AI project") + payload: Optional[bool] = Field(default=True, description="Whether to include payload in the response") + breakdown: Optional[bool] = Field(default=True, description="Whether to include breakdown in the response") + metadata: Optional[Dict] = Field(default=None, description="Additional metadata to include in the request") dev_info: Optional[bool] = Field( default=True, description="Whether to include developer information in the response", @@ -466,15 +432,9 @@ class LakeraV2GuardrailConfigModel(BaseModel): class LassoGuardrailConfigModel(BaseModel): """Configuration parameters for the Lasso guardrail""" - lasso_user_id: Optional[str] = Field( - default=None, description="User ID for the Lasso guardrail" - ) - lasso_conversation_id: Optional[str] = Field( - default=None, description="Conversation ID for the Lasso guardrail" - ) - mask: Optional[bool] = Field( - default=False, description="Enable content masking using Lasso classifix API" - ) + lasso_user_id: Optional[str] = Field(default=None, description="User ID for the Lasso guardrail") + lasso_conversation_id: Optional[str] = Field(default=None, description="Conversation ID for the Lasso guardrail") + mask: Optional[bool] = Field(default=False, description="Enable content masking using Lasso classifix API") class PillarGuardrailConfigModel(BaseModel): @@ -548,21 +508,11 @@ class ZscalerAIGuardConfigModel(BaseModel): class JavelinGuardrailConfigModel(BaseModel): """Configuration parameters for the Javelin guardrail""" - guard_name: Optional[str] = Field( - default=None, description="Name of the Javelin guard to use" - ) - api_version: Optional[str] = Field( - default="v1", description="API version for Javelin service" - ) - metadata: Optional[Dict] = Field( - default=None, description="Additional metadata to send with requests" - ) - application: Optional[str] = Field( - default=None, description="Application name for Javelin service" - ) - config: Optional[Dict] = Field( - default=None, description="Additional configuration for the guardrail" - ) + guard_name: Optional[str] = Field(default=None, description="Name of the Javelin guard to use") + api_version: Optional[str] = Field(default="v1", description="API version for Javelin service") + metadata: Optional[Dict] = Field(default=None, description="Additional metadata to send with requests") + application: Optional[str] = Field(default=None, description="Application name for Javelin service") + config: Optional[Dict] = Field(default=None, description="Additional configuration for the guardrail") class ContentFilterAction(str, Enum): @@ -576,9 +526,7 @@ class BlockedWord(BaseModel): """Represents a blocked word with its action and optional description""" keyword: str = Field(description="The keyword to block or mask") - action: ContentFilterAction = Field( - description="Action to take when keyword is detected (BLOCK or MASK)" - ) + action: ContentFilterAction = Field(description="Action to take when keyword is detected (BLOCK or MASK)") description: Optional[str] = Field( default=None, description="Optional description explaining why this keyword is sensitive", @@ -603,9 +551,7 @@ class ContentFilterPattern(BaseModel): default=None, description="Name for this pattern (used in logging and error messages)", ) - action: ContentFilterAction = Field( - description="Action to take when pattern matches (BLOCK or MASK)" - ) + action: ContentFilterAction = Field(description="Action to take when pattern matches (BLOCK or MASK)") class ContentFilterConfigModel(BaseModel): @@ -639,15 +585,9 @@ class ContentFilterConfigModel(BaseModel): ) -class BaseLitellmParams( - ContentFilterConfigModel -): # works for new and patch update guardrails - api_key: Optional[str] = Field( - default=None, description="API key for the guardrail service" - ) - api_base: Optional[str] = Field( - default=None, description="Base URL for the guardrail service API" - ) +class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails + api_key: Optional[str] = Field(default=None, description="API key for the guardrail service") + api_base: Optional[str] = Field(default=None, description="Base URL for the guardrail service API") experimental_use_latest_role_message_only: Optional[bool] = Field( default=False, @@ -686,12 +626,8 @@ class BaseLitellmParams( ) # guardrails ai params - guard_name: Optional[str] = Field( - default=None, description="Name of the guardrail in guardrails.ai" - ) - default_on: Optional[bool] = Field( - default=None, description="Whether the guardrail is enabled by default" - ) + guard_name: Optional[str] = Field(default=None, description="Name of the guardrail in guardrails.ai") + default_on: Optional[bool] = Field(default=None, description="Whether the guardrail is enabled by default") ################## PII control params ################# ######################################################## @@ -705,13 +641,9 @@ class BaseLitellmParams( ) # pangea params - pangea_input_recipe: Optional[str] = Field( - default=None, description="Recipe for input (LLM request)" - ) + pangea_input_recipe: Optional[str] = Field(default=None, description="Recipe for input (LLM request)") - pangea_output_recipe: Optional[str] = Field( - default=None, description="Recipe for output (LLM response)" - ) + pangea_output_recipe: Optional[str] = Field(default=None, description="Recipe for output (LLM response)") model: Optional[str] = Field( default=None, @@ -739,19 +671,13 @@ class BaseLitellmParams( ) # Model Armor params - template_id: Optional[str] = Field( - default=None, description="The ID of your Model Armor template" - ) - location: Optional[str] = Field( - default=None, description="Google Cloud location/region (e.g., us-central1)" - ) + template_id: Optional[str] = Field(default=None, description="The ID of your Model Armor template") + location: Optional[str] = Field(default=None, description="Google Cloud location/region (e.g., us-central1)") credentials: Optional[str] = Field( default=None, description="Path to Google Cloud credentials JSON file or JSON string", ) - api_endpoint: Optional[str] = Field( - default=None, description="Optional custom API endpoint for Model Armor" - ) + api_endpoint: Optional[str] = Field(default=None, description="Optional custom API endpoint for Model Armor") fail_on_error: Optional[bool] = Field( default=True, description=( @@ -847,21 +773,15 @@ class BaseLitellmParams( @model_validator(mode="after") def validate_sensitive_data_routing(self) -> "BaseLitellmParams": if self.on_sensitive_data == "route" and not self.sensitive_data_route_to_model: - raise ValueError( - "sensitive_data_route_to_model must be set when on_sensitive_data='route'" - ) + raise ValueError("sensitive_data_route_to_model must be set when on_sensitive_data='route'") return self model_config = ConfigDict(extra="allow", protected_namespaces=()) class Mode(BaseModel): - tags: Dict[str, Union[str, List[str]]] = Field( - description="Tags for the guardrail mode" - ) - default: Optional[Union[str, List[str]]] = Field( - default=None, description="Default mode when no tags match" - ) + tags: Dict[str, Union[str, List[str]]] = Field(description="Tags for the guardrail mode") + default: Optional[Union[str, List[str]]] = Field(default=None, description="Default mode when no tags match") class LitellmParams( @@ -970,9 +890,7 @@ class GuardrailInfoResponse(BaseModel): guardrail_info: Optional[Dict] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None - guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = ( - GUARDRAIL_DEFINITION_LOCATION.CONFIG - ) + guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.CONFIG def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/litellm/types/integrations/anthropic_cache_control_hook.py b/litellm/types/integrations/anthropic_cache_control_hook.py index 83e5a9e7f01..3e2d0d688ac 100644 --- a/litellm/types/integrations/anthropic_cache_control_hook.py +++ b/litellm/types/integrations/anthropic_cache_control_hook.py @@ -9,9 +9,7 @@ class CacheControlMessageInjectionPoint(TypedDict): """Type for message-level injection points.""" location: Literal["message"] - role: Optional[ - Literal["user", "system", "assistant"] - ] # Optional: target by role (user, system, assistant) + role: Optional[Literal["user", "system", "assistant"]] # Optional: target by role (user, system, assistant) index: Optional[Union[int, str]] # Optional: target by specific index control: Optional[ChatCompletionCachedContent] diff --git a/litellm/types/integrations/pagerduty.py b/litellm/types/integrations/pagerduty.py index 0fa45f219cb..c41a591728c 100644 --- a/litellm/types/integrations/pagerduty.py +++ b/litellm/types/integrations/pagerduty.py @@ -49,7 +49,9 @@ class AlertingConfig(TypedDict, total=False): failure_threshold_window_seconds: int # Window in seconds # Requests hanging threshold - hanging_threshold_seconds: float # Number of seconds of waiting for a response before a request is considered hanging + hanging_threshold_seconds: ( + float # Number of seconds of waiting for a response before a request is considered hanging + ) hanging_threshold_fails: int # Number of requests hanging in a window hanging_threshold_window_seconds: int # Window in seconds diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 52afc889ae6..8f460b79955 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -561,9 +561,7 @@ class PrometheusMetricLabels: litellm_api_key_max_budget_metric = litellm_remaining_api_key_budget_metric - litellm_api_key_budget_remaining_hours_metric = ( - litellm_remaining_api_key_budget_metric - ) + litellm_api_key_budget_remaining_hours_metric = litellm_remaining_api_key_budget_metric litellm_remaining_user_budget_metric = [ UserAPIKeyLabelNames.USER.value, @@ -706,9 +704,7 @@ class PrometheusMetricLabels: litellm_managed_batch_created_total = _batch_user_labels - litellm_managed_file_size_bytes: List[ - str - ] = [] # labels: purpose, file_type, model, api_provider, user (custom) + litellm_managed_file_size_bytes: List[str] = [] # labels: purpose, file_type, model, api_provider, user (custom) litellm_managed_batch_duration_seconds = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -717,9 +713,7 @@ class PrometheusMetricLabels: litellm_managed_file_created_total = _batch_user_labels - litellm_managed_file_deleted_total: List[ - str - ] = [] # only "result" label, added at metric creation + litellm_managed_file_deleted_total: List[str] = [] # only "result" label, added at metric creation litellm_check_batch_cost_jobs_polled: List[str] = [] @@ -739,19 +733,11 @@ class PrometheusMetricLabels: # Add custom metadata labels custom_labels.extend( - [ - _sanitize_prometheus_label_name(metric) - for metric in litellm.custom_prometheus_metadata_labels - ] + [_sanitize_prometheus_label_name(metric) for metric in litellm.custom_prometheus_metadata_labels] ) # Add custom tags labels - custom_labels.extend( - [ - _sanitize_prometheus_label_name(f"tag_{tag}") - for tag in litellm.custom_prometheus_tags - ] - ) + custom_labels.extend([_sanitize_prometheus_label_name(f"tag_{tag}") for tag in litellm.custom_prometheus_tags]) # Conditionally add stream label to litellm_proxy_total_requests_metric if ( @@ -766,18 +752,12 @@ class PrometheusMetricLabels: # historical label set is preserved across upgrade; enable via # ``litellm.prometheus_emit_rate_limit_labels`` once downstream # dashboards include the new labels in their matchers / aggregations. - if ( - label_name == "litellm_proxy_failed_requests_metric" - and litellm.prometheus_emit_rate_limit_labels is True - ): + if label_name == "litellm_proxy_failed_requests_metric" and litellm.prometheus_emit_rate_limit_labels is True: for _rate_limit_label in ( UserAPIKeyLabelNames.RATE_LIMIT_CATEGORY.value, UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, ): - if ( - _rate_limit_label not in default_labels - and _rate_limit_label not in custom_labels - ): + if _rate_limit_label not in default_labels and _rate_limit_label not in custom_labels: custom_labels.append(_rate_limit_label) _user_budget_metrics = { @@ -785,10 +765,7 @@ class PrometheusMetricLabels: "litellm_user_max_budget_metric", "litellm_user_budget_remaining_hours_metric", } - if ( - label_name in _user_budget_metrics - and litellm.prometheus_user_budget_label_include_email_alias is True - ): + if label_name in _user_budget_metrics and litellm.prometheus_user_budget_label_include_email_alias is True: for label in [ UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.USER_ALIAS.value, diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index c7ee4d87825..db01cb9ae12 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -94,9 +94,7 @@ class BedrockKBModelConfiguration(TypedDict, total=False): class BedrockKBRerankingConfiguration(TypedDict, total=False): """Configuration for reranking in vector search.""" - bedrockRerankingConfiguration: Optional[ - Dict[str, Any] - ] # This could be further typed if needed + bedrockRerankingConfiguration: Optional[Dict[str, Any]] # This could be further typed if needed type: Optional[str] diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index a07642073af..793cc02ff17 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -15,9 +15,7 @@ class Annotation(BaseModel): None, description="Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.", ) - end_index: Optional[int] = Field( - None, description="End of the attributed segment, exclusive." - ) + end_index: Optional[int] = Field(None, description="End of the attributed segment, exclusive.") source: Optional[str] = Field( None, description="Source attributed for a portion of the text. Could be a URL, title, or\nother identifier.", @@ -28,16 +26,12 @@ class DocumentContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[str] = None - type: Literal["document"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) + type: Literal["document"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") class FunctionCallContent(BaseModel): name: str = Field(..., description="The name of the tool to call.") - arguments: Dict[str, Any] = Field( - ..., description="The arguments to pass to the function." - ) + arguments: Dict[str, Any] = Field(..., description="The arguments to pass to the function.") type: Literal["function_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) @@ -49,9 +43,7 @@ class Language(Enum): class CodeExecutionCallArguments(BaseModel): - language: Optional[Language] = Field( - None, description="Programming language of the `code`." - ) + language: Optional[Language] = Field(None, description="Programming language of the `code`.") code: Optional[str] = Field(None, description="The code to be executed.") @@ -62,9 +54,7 @@ class UrlContextCallArguments(BaseModel): class McpServerToolCallContent(BaseModel): name: str = Field(..., description="The name of the tool which was called.") server_name: str = Field(..., description="The name of the used MCP server.") - arguments: Dict[str, Any] = Field( - ..., description="The JSON object of arguments for the function." - ) + arguments: Dict[str, Any] = Field(..., description="The JSON object of arguments for the function.") type: Literal["mcp_server_tool_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) @@ -72,25 +62,17 @@ class McpServerToolCallContent(BaseModel): class GoogleSearchCallArguments(BaseModel): - queries: Optional[List[str]] = Field( - None, description="Web search queries for the following-up web search." - ) + queries: Optional[List[str]] = Field(None, description="Web search queries for the following-up web search.") class CodeExecutionResultContent(BaseModel): result: Optional[str] = Field(None, description="The output of the code execution.") - is_error: Optional[bool] = Field( - None, description="Whether the code execution resulted in an error." - ) - signature: Optional[str] = Field( - None, description="A signature hash for backend validation." - ) + is_error: Optional[bool] = Field(None, description="Whether the code execution resulted in an error.") + signature: Optional[str] = Field(None, description="A signature hash for backend validation.") type: Literal["code_execution_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the code execution call block." - ) + call_id: Optional[str] = Field(None, description="ID to match the ID from the code execution call block.") class Status(Enum): @@ -102,9 +84,7 @@ class Status(Enum): class UrlContextResult(BaseModel): url: Optional[str] = Field(None, description="The URL that was fetched.") - status: Optional[Status] = Field( - None, description="The status of the URL retrieval." - ) + status: Optional[Status] = Field(None, description="The status of the URL retrieval.") class GoogleSearchResult(BaseModel): @@ -119,9 +99,7 @@ class GoogleSearchResult(BaseModel): class FileSearchResult(BaseModel): title: Optional[str] = Field(None, description="The title of the search result.") text: Optional[str] = Field(None, description="The text of the search result.") - file_search_store: Optional[str] = Field( - None, description="The name of the file search store." - ) + file_search_store: Optional[str] = Field(None, description="The name of the file search store.") class SpeechConfig(BaseModel): @@ -142,12 +120,8 @@ class DynamicAgentConfig(BaseModel): class Function(BaseModel): name: Optional[str] = Field(None, description="The name of the function.") - description: Optional[str] = Field( - None, description="A description of the function." - ) - parameters: Optional[Any] = Field( - None, description="The JSON Schema for the function's parameters." - ) + description: Optional[str] = Field(None, description="A description of the function.") + parameters: Optional[Any] = Field(None, description="The JSON Schema for the function's parameters.") type: Literal["function"] @@ -165,9 +139,7 @@ class Environment(Enum): class ComputerUse(BaseModel): type: Literal["computer_use"] - environment: Optional[Environment] = Field( - None, description="The environment being operated." - ) + environment: Optional[Environment] = Field(None, description="The environment being operated.") excludedPredefinedFunctions: Optional[List[str]] = Field( None, description="The list of predefined functions that are excluded from the model call.", @@ -179,12 +151,8 @@ class GoogleSearch(BaseModel): class FileSearch(BaseModel): - file_search_store_names: Optional[List[str]] = Field( - None, description="The file search store names to search." - ) - top_k: Optional[int] = Field( - None, description="The number of semantic retrieval chunks to retrieve." - ) + file_search_store_names: Optional[List[str]] = Field(None, description="The file search store names to search.") + top_k: Optional[int] = Field(None, description="The number of semantic retrieval chunks to retrieve.") metadata_filter: Optional[str] = Field( None, description="Metadata filter to apply to the semantic retrieval documents and chunks.", @@ -219,9 +187,7 @@ class InteractionStatusUpdate(BaseModel): class TextDelta(BaseModel): text: Optional[str] = None - type: Literal["text"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) + type: Literal["text"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") annotations: Optional[List[Annotation]] = Field( None, description="Citation information for model-generated content." ) @@ -231,9 +197,7 @@ class DocumentDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[str] = None - type: Literal["document"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) + type: Literal["document"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") class ThoughtSignatureDelta(BaseModel): @@ -252,9 +216,7 @@ class FunctionCallDelta(BaseModel): type: Literal["function_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class CodeExecutionCallDelta(BaseModel): @@ -262,9 +224,7 @@ class CodeExecutionCallDelta(BaseModel): type: Literal["code_execution_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class UrlContextCallDelta(BaseModel): @@ -272,9 +232,7 @@ class UrlContextCallDelta(BaseModel): type: Literal["url_context_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class GoogleSearchCallDelta(BaseModel): @@ -282,9 +240,7 @@ class GoogleSearchCallDelta(BaseModel): type: Literal["google_search_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class McpServerToolCallDelta(BaseModel): @@ -294,9 +250,7 @@ class McpServerToolCallDelta(BaseModel): type: Literal["mcp_server_tool_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class CodeExecutionResultDelta(BaseModel): @@ -306,9 +260,7 @@ class CodeExecutionResultDelta(BaseModel): type: Literal["code_execution_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the function call block." - ) + call_id: Optional[str] = Field(None, description="ID to match the ID from the function call block.") class UrlContextResultDelta(BaseModel): @@ -318,9 +270,7 @@ class UrlContextResultDelta(BaseModel): type: Literal["url_context_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the function call block." - ) + call_id: Optional[str] = Field(None, description="ID to match the ID from the function call block.") class GoogleSearchResultDelta(BaseModel): @@ -330,9 +280,7 @@ class GoogleSearchResultDelta(BaseModel): type: Literal["google_search_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the function call block." - ) + call_id: Optional[str] = Field(None, description="ID to match the ID from the function call block.") class FileSearchResultDelta(BaseModel): @@ -352,9 +300,7 @@ class ContentStop(BaseModel): class Error(BaseModel): - code: Optional[str] = Field( - None, description="A URI that identifies the error type." - ) + code: Optional[str] = Field(None, description="A URI that identifies the error type.") message: Optional[str] = Field(None, description="A human-readable error message.") @@ -410,28 +356,20 @@ class AgentOption(RootModel[str]): class ImageMimeTypeOption(RootModel[str]): - root: str = Field( - ..., description="The mime type of the image.", title="ImageMimeType" - ) + root: str = Field(..., description="The mime type of the image.", title="ImageMimeType") class AudioMimeTypeOption(RootModel[str]): - root: str = Field( - ..., description="The mime type of the audio.", title="AudioMimeType" - ) + root: str = Field(..., description="The mime type of the audio.", title="AudioMimeType") class VideoMimeTypeOption(RootModel[str]): - root: str = Field( - ..., description="The mime type of the video.", title="VideoMimeType" - ) + root: str = Field(..., description="The mime type of the video.", title="VideoMimeType") class TextContent(BaseModel): text: Optional[str] = Field(None, description="The text content.") - type: Literal["text"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) + type: Literal["text"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") annotations: Optional[List[Annotation]] = Field( None, description="Citation information for model-generated content." ) @@ -441,33 +379,23 @@ class ImageContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[ImageMimeTypeOption] = None - type: Literal["image"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) - resolution: Optional[MediaResolution] = Field( - None, description="The resolution of the media." - ) + type: Literal["image"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") + resolution: Optional[MediaResolution] = Field(None, description="The resolution of the media.") class AudioContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[AudioMimeTypeOption] = None - type: Literal["audio"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) + type: Literal["audio"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") class VideoContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[VideoMimeTypeOption] = None - type: Literal["video"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) - resolution: Optional[MediaResolution] = Field( - None, description="The resolution of the media." - ) + type: Literal["video"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") + resolution: Optional[MediaResolution] = Field(None, description="The resolution of the media.") class ThoughtSummary1(RootModel[Union[TextContent, ImageContent]]): @@ -485,33 +413,23 @@ class CodeExecutionCallContent(BaseModel): type: Literal["code_execution_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class UrlContextCallContent(BaseModel): - arguments: Optional[UrlContextCallArguments] = Field( - None, description="The arguments to pass to the URL context." - ) + arguments: Optional[UrlContextCallArguments] = Field(None, description="The arguments to pass to the URL context.") type: Literal["url_context_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class GoogleSearchCallContent(BaseModel): - arguments: Optional[GoogleSearchCallArguments] = Field( - None, description="The arguments to pass to Google Search." - ) + arguments: Optional[GoogleSearchCallArguments] = Field(None, description="The arguments to pass to Google Search.") type: Literal["google_search_call"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: Optional[str] = Field( - None, description="A unique ID for this specific tool call." - ) + id: Optional[str] = Field(None, description="A unique ID for this specific tool call.") class Result(BaseModel): @@ -519,57 +437,33 @@ class Result(BaseModel): class FunctionResultContent(BaseModel): - name: Optional[str] = Field( - None, description="The name of the tool that was called." - ) - is_error: Optional[bool] = Field( - None, description="Whether the tool call resulted in an error." - ) + name: Optional[str] = Field(None, description="The name of the tool that was called.") + is_error: Optional[bool] = Field(None, description="Whether the tool call resulted in an error.") type: Literal["function_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - result: Union[Result, Dict[str, Any], str] = Field( - ..., description="The result of the tool call." - ) - call_id: str = Field( - ..., description="ID to match the ID from the function call block." - ) + result: Union[Result, Dict[str, Any], str] = Field(..., description="The result of the tool call.") + call_id: str = Field(..., description="ID to match the ID from the function call block.") class UrlContextResultContent(BaseModel): - signature: Optional[str] = Field( - None, description="The signature of the URL context result." - ) - result: Optional[List[UrlContextResult]] = Field( - None, description="The results of the URL context." - ) - is_error: Optional[bool] = Field( - None, description="Whether the URL context resulted in an error." - ) + signature: Optional[str] = Field(None, description="The signature of the URL context result.") + result: Optional[List[UrlContextResult]] = Field(None, description="The results of the URL context.") + is_error: Optional[bool] = Field(None, description="Whether the URL context resulted in an error.") type: Literal["url_context_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the url context call block." - ) + call_id: Optional[str] = Field(None, description="ID to match the ID from the url context call block.") class GoogleSearchResultContent(BaseModel): - signature: Optional[str] = Field( - None, description="The signature of the Google Search result." - ) - result: Optional[List[GoogleSearchResult]] = Field( - None, description="The results of the Google Search." - ) - is_error: Optional[bool] = Field( - None, description="Whether the Google Search resulted in an error." - ) + signature: Optional[str] = Field(None, description="The signature of the Google Search result.") + result: Optional[List[GoogleSearchResult]] = Field(None, description="The results of the Google Search.") + is_error: Optional[bool] = Field(None, description="Whether the Google Search resulted in an error.") type: Literal["google_search_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the google search call block." - ) + call_id: Optional[str] = Field(None, description="ID to match the ID from the google search call block.") class McpServerToolResultContent(BaseModel): @@ -577,36 +471,24 @@ class McpServerToolResultContent(BaseModel): None, description="Name of the tool which is called for this specific tool call.", ) - server_name: Optional[str] = Field( - None, description="The name of the used MCP server." - ) + server_name: Optional[str] = Field(None, description="The name of the used MCP server.") type: Literal["mcp_server_tool_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - result: Union[Result, Dict[str, Any], str] = Field( - ..., description="The result of the tool call." - ) - call_id: str = Field( - ..., description="ID to match the ID from the MCP server tool call block." - ) + result: Union[Result, Dict[str, Any], str] = Field(..., description="The result of the tool call.") + call_id: str = Field(..., description="ID to match the ID from the MCP server tool call block.") class FileSearchResultContent(BaseModel): - result: Optional[List[FileSearchResult]] = Field( - None, description="The results of the File Search." - ) + result: Optional[List[FileSearchResult]] = Field(None, description="The results of the File Search.") type: Literal["file_search_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) class AllowedTools(BaseModel): - mode: Optional[ToolChoiceType] = Field( - None, description="The mode of the tool choice." - ) - tools: Optional[List[str]] = Field( - None, description="The names of the allowed tools." - ) + mode: Optional[ToolChoiceType] = Field(None, description="The mode of the tool choice.") + tools: Optional[List[str]] = Field(None, description="The names of the allowed tools.") class DeepResearchAgentConfig(BaseModel): @@ -630,60 +512,42 @@ class McpServer(BaseModel): None, description="Optional: Fields for authentication headers, timeouts, etc., if needed.", ) - allowed_tools: Optional[List[AllowedTools]] = Field( - None, description="The allowed tools." - ) + allowed_tools: Optional[List[AllowedTools]] = Field(None, description="The allowed tools.") class ModalityTokens(BaseModel): - modality: Optional[ResponseModality] = Field( - None, description="The modality associated with the token count." - ) - tokens: Optional[int] = Field( - None, description="Number of tokens for the modality." - ) + modality: Optional[ResponseModality] = Field(None, description="The modality associated with the token count.") + tokens: Optional[int] = Field(None, description="Number of tokens for the modality.") class ImageDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[ImageMimeTypeOption] = None - type: Literal["image"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) - resolution: Optional[MediaResolution] = Field( - None, description="The resolution of the media." - ) + type: Literal["image"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") + resolution: Optional[MediaResolution] = Field(None, description="The resolution of the media.") class AudioDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[AudioMimeTypeOption] = None - type: Literal["audio"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) + type: Literal["audio"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") class VideoDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[VideoMimeTypeOption] = None - type: Literal["video"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) - resolution: Optional[MediaResolution] = Field( - None, description="The resolution of the media." - ) + type: Literal["video"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") + resolution: Optional[MediaResolution] = Field(None, description="The resolution of the media.") class ThoughtSummaryDelta(BaseModel): type: Literal["thought_summary"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - content: Optional[Union[TextContent, ImageContent]] = Field( - None, discriminator="type" - ) + content: Optional[Union[TextContent, ImageContent]] = Field(None, discriminator="type") class FunctionResultDelta(BaseModel): @@ -692,12 +556,8 @@ class FunctionResultDelta(BaseModel): type: Literal["function_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - result: Optional[Union[Result, str]] = Field( - None, description="Tool call result delta." - ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the function call block." - ) + result: Optional[Union[Result, str]] = Field(None, description="Tool call result delta.") + call_id: Optional[str] = Field(None, description="ID to match the ID from the function call block.") class McpServerToolResultDelta(BaseModel): @@ -706,12 +566,8 @@ class McpServerToolResultDelta(BaseModel): type: Literal["mcp_server_tool_result"] = Field( ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - result: Optional[Union[Result, str]] = Field( - None, description="Tool call result delta." - ) - call_id: Optional[str] = Field( - None, description="ID to match the ID from the function call block." - ) + result: Optional[Union[Result, str]] = Field(None, description="Tool call result delta.") + call_id: Optional[str] = Field(None, description="ID to match the ID from the function call block.") class ErrorEvent(BaseModel): @@ -756,24 +612,16 @@ class ThoughtContent(BaseModel): None, description="Signature to match the backend source to be part of the generation.", ) - type: Literal["thought"] = Field( - ..., description="Used as the OpenAPI type discriminator for the content oneof." - ) - summary: Optional[ThoughtSummary] = Field( - None, description="A summary of the thought." - ) + type: Literal["thought"] = Field(..., description="Used as the OpenAPI type discriminator for the content oneof.") + summary: Optional[ThoughtSummary] = Field(None, description="A summary of the thought.") class ToolChoice(RootModel[Union[ToolChoiceType, ToolChoiceConfig]]): - root: Union[ToolChoiceType, ToolChoiceConfig] = Field( - ..., description="The configuration for tool choice." - ) + root: Union[ToolChoiceType, ToolChoiceConfig] = Field(..., description="The configuration for tool choice.") class Usage(BaseModel): - total_input_tokens: Optional[int] = Field( - None, description="Number of tokens in the prompt (context)." - ) + total_input_tokens: Optional[int] = Field(None, description="Number of tokens in the prompt (context).") input_tokens_by_modality: Optional[List[ModalityTokens]] = Field( None, description="A breakdown of input token usage by modality." ) @@ -790,15 +638,11 @@ class Usage(BaseModel): output_tokens_by_modality: Optional[List[ModalityTokens]] = Field( None, description="A breakdown of output token usage by modality." ) - total_tool_use_tokens: Optional[int] = Field( - None, description="Number of tokens present in tool-use prompt(s)." - ) + total_tool_use_tokens: Optional[int] = Field(None, description="Number of tokens present in tool-use prompt(s).") tool_use_tokens_by_modality: Optional[List[ModalityTokens]] = Field( None, description="A breakdown of tool-use token usage by modality." ) - total_reasoning_tokens: Optional[int] = Field( - None, description="Number of tokens of thoughts for thinking models." - ) + total_reasoning_tokens: Optional[int] = Field(None, description="Number of tokens of thoughts for thinking models.") total_tokens: Optional[int] = Field( None, description="Total token count for the interaction request (prompt + responses + other\ninternal tokens).", @@ -885,29 +729,21 @@ class Turn(BaseModel): None, description="The originator of this turn. Must be user for input or model for\nmodel output.", ) - content: Optional[Union[str, List[Content]]] = Field( - None, description="The content of the turn." - ) + content: Optional[Union[str, List[Content]]] = Field(None, description="The content of the turn.") class GenerationConfig(BaseModel): - temperature: Optional[float] = Field( - None, description="Controls the randomness of the output." - ) + temperature: Optional[float] = Field(None, description="Controls the randomness of the output.") top_p: Optional[float] = Field( None, description="The maximum cumulative probability of tokens to consider when sampling.", ) - seed: Optional[int] = Field( - None, description="Seed used in decoding for reproducibility." - ) + seed: Optional[int] = Field(None, description="Seed used in decoding for reproducibility.") stop_sequences: Optional[List[str]] = Field( None, description="A list of character sequences that will stop output interaction.", ) - tool_choice: Optional[ToolChoice] = Field( - None, description="The tool choice for the interaction." - ) + tool_choice: Optional[ToolChoice] = Field(None, description="The tool choice for the interaction.") thinking_level: Optional[ThinkingLevel] = Field( None, description="The level of thought tokens that the model should generate." ) @@ -917,9 +753,7 @@ class GenerationConfig(BaseModel): max_output_tokens: Optional[int] = Field( None, description="The maximum number of tokens to include in the response." ) - speech_config: Optional[List[SpeechConfig]] = Field( - None, description="Configuration for speech interaction." - ) + speech_config: Optional[List[SpeechConfig]] = Field(None, description="Configuration for speech interaction.") class ContentStart(BaseModel): @@ -943,9 +777,7 @@ class Interaction(BaseModel): ..., description="Output only. A unique identifier for the interaction completion.", ) - status: Status1 = Field( - ..., description="Output only. The status of the interaction." - ) + status: Status1 = Field(..., description="Output only. The status of the interaction.") created: Optional[AwareDatetime] = Field( None, description="Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", @@ -954,19 +786,13 @@ class Interaction(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - outputs: Optional[List[Content]] = Field( - None, description="Output only. Responses from the model." - ) - system_instruction: Optional[str] = Field( - None, description="System instruction for the interaction." - ) + outputs: Optional[List[Content]] = Field(None, description="Output only. Responses from the model.") + system_instruction: Optional[str] = Field(None, description="System instruction for the interaction.") tools: Optional[List[Tool]] = Field( None, description="A list of tool declarations the model may call during interaction.", ) - background: Optional[bool] = Field( - None, description="Whether to run the model interaction in the background." - ) + background: Optional[bool] = Field(None, description="Whether to run the model interaction in the background.") object: Literal["interaction"] = Field( "interaction", description="Output only. The object type of the interaction. Always set to `interaction`.", @@ -987,9 +813,7 @@ class Interaction(BaseModel): None, description="The mime type of the response. This is required if response_format is set.", ) - previous_interaction_id: Optional[str] = Field( - None, description="The ID of the previous interaction, if any." - ) + previous_interaction_id: Optional[str] = Field(None, description="The ID of the previous interaction, if any.") input: Optional[Union[str, List[Content], List[Turn], Content]] = Field( None, description="The inputs for the interaction." ) @@ -1003,12 +827,8 @@ class Interaction(BaseModel): class CreateModelInteractionParams(BaseModel): - model: ModelOption = Field( - ..., description="The name of the `Model` used for generating the interaction." - ) - stream: Optional[bool] = Field( - None, description="Input only. Whether the interaction will be streamed." - ) + model: ModelOption = Field(..., description="The name of the `Model` used for generating the interaction.") + stream: Optional[bool] = Field(None, description="Input only. Whether the interaction will be streamed.") store: Optional[bool] = Field( None, description="Input only. Whether to store the response and request for later retrieval.", @@ -1017,9 +837,7 @@ class CreateModelInteractionParams(BaseModel): None, description="Output only. A unique identifier for the interaction completion.", ) - status: Optional[Status3] = Field( - None, description="Output only. The status of the interaction." - ) + status: Optional[Status3] = Field(None, description="Output only. The status of the interaction.") created: Optional[AwareDatetime] = Field( None, description="Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", @@ -1028,19 +846,13 @@ class CreateModelInteractionParams(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - outputs: Optional[List[Content]] = Field( - None, description="Output only. Responses from the model." - ) - system_instruction: Optional[str] = Field( - None, description="System instruction for the interaction." - ) + outputs: Optional[List[Content]] = Field(None, description="Output only. Responses from the model.") + system_instruction: Optional[str] = Field(None, description="System instruction for the interaction.") tools: Optional[List[Tool]] = Field( None, description="A list of tool declarations the model may call during interaction.", ) - background: Optional[bool] = Field( - None, description="Whether to run the model interaction in the background." - ) + background: Optional[bool] = Field(None, description="Whether to run the model interaction in the background.") usage: Optional[Usage] = Field( None, description="Output only. Statistics on the interaction request's token usage.", @@ -1057,12 +869,8 @@ class CreateModelInteractionParams(BaseModel): None, description="The mime type of the response. This is required if response_format is set.", ) - previous_interaction_id: Optional[str] = Field( - None, description="The ID of the previous interaction, if any." - ) - input: Union[str, List[Content], List[Turn], Content] = Field( - ..., description="The inputs for the interaction." - ) + previous_interaction_id: Optional[str] = Field(None, description="The ID of the previous interaction, if any.") + input: Union[str, List[Content], List[Turn], Content] = Field(..., description="The inputs for the interaction.") generation_config: Optional[GenerationConfig] = Field( None, description="Input only. Configuration parameters for the model interaction.", @@ -1070,12 +878,8 @@ class CreateModelInteractionParams(BaseModel): class CreateAgentInteractionParams(BaseModel): - agent: AgentOption = Field( - ..., description="The name of the `Agent` used for generating the interaction." - ) - stream: Optional[bool] = Field( - None, description="Input only. Whether the interaction will be streamed." - ) + agent: AgentOption = Field(..., description="The name of the `Agent` used for generating the interaction.") + stream: Optional[bool] = Field(None, description="Input only. Whether the interaction will be streamed.") store: Optional[bool] = Field( None, description="Input only. Whether to store the response and request for later retrieval.", @@ -1084,9 +888,7 @@ class CreateAgentInteractionParams(BaseModel): None, description="Output only. A unique identifier for the interaction completion.", ) - status: Optional[Status3] = Field( - None, description="Output only. The status of the interaction." - ) + status: Optional[Status3] = Field(None, description="Output only. The status of the interaction.") created: Optional[AwareDatetime] = Field( None, description="Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", @@ -1095,19 +897,13 @@ class CreateAgentInteractionParams(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - outputs: Optional[List[Content]] = Field( - None, description="Output only. Responses from the model." - ) - system_instruction: Optional[str] = Field( - None, description="System instruction for the interaction." - ) + outputs: Optional[List[Content]] = Field(None, description="Output only. Responses from the model.") + system_instruction: Optional[str] = Field(None, description="System instruction for the interaction.") tools: Optional[List[Tool]] = Field( None, description="A list of tool declarations the model may call during interaction.", ) - background: Optional[bool] = Field( - None, description="Whether to run the model interaction in the background." - ) + background: Optional[bool] = Field(None, description="Whether to run the model interaction in the background.") usage: Optional[Usage] = Field( None, description="Output only. Statistics on the interaction request's token usage.", @@ -1124,12 +920,8 @@ class CreateAgentInteractionParams(BaseModel): None, description="The mime type of the response. This is required if response_format is set.", ) - previous_interaction_id: Optional[str] = Field( - None, description="The ID of the previous interaction, if any." - ) - input: Union[str, List[Content], List[Turn], Content] = Field( - ..., description="The inputs for the interaction." - ) + previous_interaction_id: Optional[str] = Field(None, description="The ID of the previous interaction, if any.") + input: Union[str, List[Content], List[Turn], Content] = Field(..., description="The inputs for the interaction.") agent_config: Optional[Union[DynamicAgentConfig, DeepResearchAgentConfig]] = Field( None, description="Configuration for the agent.", discriminator="type" ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 561d3692ee2..aa9f4dccbd1 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -365,9 +365,7 @@ class AnthropicSystemMessageContent(TypedDict, total=False): cache_control: Optional[Union[dict, ChatCompletionCachedContent]] -AllAnthropicMessageValues = Union[ - AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam -] +AllAnthropicMessageValues = Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam] class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): @@ -385,14 +383,10 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): top_p: Optional[float] mcp_servers: Optional[List[AnthropicMcpServerTool]] context_management: Optional[Dict[str, Any]] - container: Optional[ - Dict[str, Any] - ] # Container config with skills for code execution + container: Optional[Dict[str, Any]] # Container config with skills for code execution output_format: Optional[AnthropicOutputSchema] # Structured outputs support speed: Optional[str] # Fast mode support for Opus models - output_config: Optional[ - AnthropicOutputConfig - ] # Configuration for Claude's output behavior + output_config: Optional[AnthropicOutputConfig] # Configuration for Claude's output behavior cache_control: Optional[Dict[str, Any]] # Automatic prompt caching reasoning_effort: Optional[str] @@ -497,9 +491,7 @@ class ContentBlockStartText(TypedDict): content_block: TextBlock -ContentBlockContentBlockDict = Union[ - ToolUseBlock, TextBlock, ChatCompletionThinkingBlock -] +ContentBlockContentBlockDict = Union[ToolUseBlock, TextBlock, ChatCompletionThinkingBlock] ContentBlockStart = Union[ContentBlockStartToolUse, ContentBlockStartText] diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 85a2b3fee7c..e432e25b6ca 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -89,9 +89,7 @@ class AnthropicMessagesResponse(TypedDict, total=False): id: str model: Optional[str] # This represents the Model type from Anthropic role: Optional[Literal["assistant"]] - stop_reason: Optional[ - Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] - ] + stop_reason: Optional[Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"]] stop_sequence: Optional[str] type: Optional[Literal["message"]] usage: Optional[AnthropicUsage] diff --git a/litellm/types/llms/anthropic_tool_search.py b/litellm/types/llms/anthropic_tool_search.py index 7cdaec2e7cc..98eee267492 100644 --- a/litellm/types/llms/anthropic_tool_search.py +++ b/litellm/types/llms/anthropic_tool_search.py @@ -29,6 +29,4 @@ def get_tool_search_beta_header(custom_llm_provider: str) -> str: """ Get the tool search beta header for a given provider. """ - return TOOL_SEARCH_BETA_HEADER_BY_PROVIDER.get( - custom_llm_provider, TOOL_SEARCH_BETA_HEADER_ANTHROPIC - ) + return TOOL_SEARCH_BETA_HEADER_BY_PROVIDER.get(custom_llm_provider, TOOL_SEARCH_BETA_HEADER_ANTHROPIC) diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index fa8c3a93ef3..a1dc08e5f29 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -28,9 +28,7 @@ class ImageBlock(TypedDict): source: SourceBlock -BedrockVideoTypes = Literal[ - "mp4", "mov", "mkv", "webm", "flv", "mpeg", "mpg", "wmv", "3gp" -] +BedrockVideoTypes = Literal["mp4", "mov", "mkv", "webm", "flv", "mpeg", "mpg", "wmv", "3gp"] class VideoBlock(TypedDict): @@ -38,9 +36,7 @@ class VideoBlock(TypedDict): source: SourceBlock -BedrockDocumentTypes = Literal[ - "pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md" -] +BedrockDocumentTypes = Literal["pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md"] class DocumentBlock(TypedDict): @@ -237,13 +233,9 @@ class ConverseResponseBlock(TypedDict, total=False): additionalModelResponseFields: dict metrics: ConverseMetricsBlock output: Required[ConverseResponseOutputBlock] - stopReason: Required[ - str - ] # end_turn | tool_use | max_tokens | stop_sequence | content_filtered + stopReason: Required[str] # end_turn | tool_use | max_tokens | stop_sequence | content_filtered usage: Required[ConverseTokenUsageBlock] - serviceTier: ( - ServiceTierBlock # Optional - only present when serviceTier was sent in request - ) + serviceTier: ServiceTierBlock # Optional - only present when serviceTier was sent in request class ToolJsonSchemaBlock(TypedDict, total=False): @@ -401,9 +393,7 @@ class OutputConfigBlock(TypedDict, total=False): textFormat: OutputFormat -class CommonRequestObject( - TypedDict, total=False -): # common request object across sync + async flows +class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows additionalModelRequestFields: dict additionalModelResponseFieldPaths: List[str] inferenceConfig: InferenceConfig @@ -487,9 +477,7 @@ class ServerSentEvent: return f"ServerSentEvent(event={self.event}, data={self.data}, id={self.id}, retry={self.retry})" -COHERE_EMBEDDING_INPUT_TYPES = Literal[ - "search_document", "search_query", "classification", "clustering", "image" -] +COHERE_EMBEDDING_INPUT_TYPES = Literal["search_document", "search_query", "classification", "clustering", "image"] class CohereEmbeddingRequest(TypedDict, total=False): @@ -636,9 +624,7 @@ NOVA_DETAIL_LEVELS = Literal["STANDARD_IMAGE", "DOCUMENT_IMAGE"] NOVA_EMBEDDING_MODES = Literal["AUDIO_VIDEO_COMBINED", "AUDIO_VIDEO_SEPARATE"] -NOVA_EMBEDDING_TYPES = Literal[ - "TEXT", "IMAGE", "VIDEO", "AUDIO", "AUDIO_VIDEO_COMBINED" -] +NOVA_EMBEDDING_TYPES = Literal["TEXT", "IMAGE", "VIDEO", "AUDIO", "AUDIO_VIDEO_COMBINED"] class NovaSourceS3Location(TypedDict): @@ -757,9 +743,7 @@ class AmazonStability3TextToImageRequest(TypedDict, total=False): """ prompt: str - aspect_ratio: Literal[ - "16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21" - ] + aspect_ratio: Literal["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"] mode: Literal["image-to-image", "text-to-image"] output_format: Literal["JPEG", "PNG"] seed: int @@ -822,9 +806,7 @@ class AmazonNovaCanvasTextToImageParams(TypedDict, total=False): conditionImage: str -class AmazonNovaCanvasTextToImageRequest( - AmazonNovaCanvasRequestBase, TypedDict, total=False -): +class AmazonNovaCanvasTextToImageRequest(AmazonNovaCanvasRequestBase, TypedDict, total=False): """ Request for Amazon Nova Canvas Text to Image API @@ -847,9 +829,7 @@ class AmazonNovaCanvasColorGuidedGenerationParams(TypedDict, total=False): negativeText: str -class AmazonNovaCanvasColorGuidedRequest( - AmazonNovaCanvasRequestBase, TypedDict, total=False -): +class AmazonNovaCanvasColorGuidedRequest(AmazonNovaCanvasRequestBase, TypedDict, total=False): """ Request for Amazon Nova Canvas Color Guided Generation API @@ -882,9 +862,7 @@ class AmazonNovaCanvasInpaintingParams(TypedDict, total=False): negativeText: str -class AmazonNovaCanvasInpaintingRequest( - AmazonNovaCanvasRequestBase, TypedDict, total=False -): +class AmazonNovaCanvasInpaintingRequest(AmazonNovaCanvasRequestBase, TypedDict, total=False): """ Request for Amazon Nova Canvas Inpainting API @@ -1024,9 +1002,7 @@ class BedrockCreateBatchRequest(TypedDict, total=False): tags: Optional[List[dict]] -BedrockBatchJobStatus = Literal[ - "Submitted", "InProgress", "Completed", "Failed", "Stopping", "Stopped" -] +BedrockBatchJobStatus = Literal["Submitted", "InProgress", "Completed", "Failed", "Stopping", "Stopped"] class BedrockCreateBatchResponse(TypedDict): diff --git a/litellm/types/llms/cohere.py b/litellm/types/llms/cohere.py index ea41bacd965..7ac6a945d5c 100644 --- a/litellm/types/llms/cohere.py +++ b/litellm/types/llms/cohere.py @@ -41,9 +41,7 @@ class ChatHistoryChatBot(TypedDict, total=False): tool_calls: List[ToolCallObject] -ChatHistory = List[ - Union[ChatHistorySystem, ChatHistoryChatBot, ChatHistoryUser, ChatHistoryToolResult] -] +ChatHistory = List[Union[ChatHistorySystem, ChatHistoryChatBot, ChatHistoryUser, ChatHistoryToolResult]] class CohereV2ChatResponseMessageToolCallFunction(TypedDict, total=False): diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index c862bf0e4ac..aa0e33f52b5 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -43,9 +43,7 @@ class DatabricksReasoningContent(TypedDict, total=False): citations: Optional[List[Dict[str, Any]]] -AllDatabricksContentListValues = Union[ - DatabricksTextContent, DatabricksReasoningContent -] +AllDatabricksContentListValues = Union[DatabricksTextContent, DatabricksReasoningContent] AllDatabricksContentValues = Union[str, List[AllDatabricksContentListValues]] diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index e24eb4aebb5..80c006787d3 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -103,9 +103,7 @@ class BidiGenerateContentRealtimeInput(TypedDict, total=False): StartOfSpeechSensitivityEnum = Literal[ "START_SENSITIVITY_UNSPECIFIED", "START_SENSITIVITY_HIGH", "START_SENSITIVITY_LOW" ] -EndOfSpeechSensitivityEnum = Literal[ - "END_SENSITIVITY_UNSPECIFIED", "END_SENSITIVITY_HIGH", "END_SENSITIVITY_LOW" -] +EndOfSpeechSensitivityEnum = Literal["END_SENSITIVITY_UNSPECIFIED", "END_SENSITIVITY_HIGH", "END_SENSITIVITY_LOW"] class AutomaticActivityDetection(TypedDict, total=False): diff --git a/litellm/types/llms/oci.py b/litellm/types/llms/oci.py index 621f40aa31d..aa46877fd6b 100644 --- a/litellm/types/llms/oci.py +++ b/litellm/types/llms/oci.py @@ -416,13 +416,9 @@ class OCIEmbedRequest(BaseModel): compartmentId: str servingMode: OCIServingMode inputs: List[str] - inputType: Optional[str] = ( - None # SEARCH_DOCUMENT | SEARCH_QUERY | CLASSIFICATION | CLUSTERING | IMAGE - ) + inputType: Optional[str] = None # SEARCH_DOCUMENT | SEARCH_QUERY | CLASSIFICATION | CLUSTERING | IMAGE truncate: Optional[str] = "END" # NONE | START | END - outputDimensions: Optional[int] = ( - None # cohere.embed-v4.0+; valid: 256, 512, 1024, 1536 - ) + outputDimensions: Optional[int] = None # cohere.embed-v4.0+; valid: 256, 512, 1024, 1536 class OCIEmbedUsage(BaseModel): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 36b6bac90ea..42015f76442 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -686,9 +686,7 @@ class ChatCompletionFileObjectFile(TypedDict, total=False): filename: str format: str detail: str # For video/image resolution control (low, medium, high, ultra_high) - video_metadata: Dict[ - str, Any - ] # For video-specific metadata (fps, start_offset, end_offset) + video_metadata: Dict[str, Any] # For video-specific metadata (fps, start_offset, end_offset) class ChatCompletionFileObject(TypedDict): @@ -751,9 +749,7 @@ class OpenAIChatCompletionAssistantMessage(TypedDict, total=False): class ChatCompletionAssistantMessage(OpenAIChatCompletionAssistantMessage, total=False): cache_control: ChatCompletionCachedContent - thinking_blocks: Optional[ - List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] - ] + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] reasoning_items: Optional[List[ChatCompletionReasoningItem]] @@ -901,9 +897,7 @@ class ChatCompletionToolChoiceObjectParam(TypedDict): ChatCompletionToolChoiceStringValues = Literal["none", "auto", "required"] -ChatCompletionToolChoiceValues = Union[ - ChatCompletionToolChoiceStringValues, ChatCompletionToolChoiceObjectParam -] +ChatCompletionToolChoiceValues = Union[ChatCompletionToolChoiceStringValues, ChatCompletionToolChoiceObjectParam] class ChatCompletionToolParamFunctionChunk(TypedDict, total=False): @@ -968,9 +962,7 @@ class ChatCompletionDeltaChunk(TypedDict, total=False): role: str -ChatCompletionAssistantContentValue = ( - str # keep as var, used in stream_chunk_builder as well -) +ChatCompletionAssistantContentValue = str # keep as var, used in stream_chunk_builder as well class ChatCompletionResponseMessage(TypedDict, total=False): @@ -981,9 +973,7 @@ class ChatCompletionResponseMessage(TypedDict, total=False): function_call: Optional[ChatCompletionToolCallFunctionChunk] provider_specific_fields: Optional[dict] reasoning_content: Optional[str] - thinking_blocks: Optional[ - List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] - ] + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] class ChatCompletionUsageBlock(TypedDict, total=False): @@ -1003,12 +993,8 @@ class OpenAIChatCompletionChunk(ChatCompletionChunk): class Hyperparameters(BaseModel): batch_size: Optional[Union[str, int]] = None # "Number of examples in each batch." - learning_rate_multiplier: Optional[Union[str, float]] = ( - None # Scaling factor for the learning rate - ) - n_epochs: Optional[Union[str, int]] = ( - None # "The number of epochs to train the model for" - ) + learning_rate_multiplier: Optional[Union[str, float]] = None # Scaling factor for the learning rate + n_epochs: Optional[Union[str, int]] = None # "The number of epochs to train the model for" model_config = {"extra": "allow"} @@ -1037,27 +1023,17 @@ class FineTuningJobCreate(BaseModel): model: str # "The name of the model to fine-tune." training_file: str # "The ID of an uploaded file that contains training data." - hyperparameters: Optional[Hyperparameters] = ( - None # "The hyperparameters used for the fine-tuning job." - ) - suffix: Optional[str] = ( - None # "A string of up to 18 characters that will be added to your fine-tuned model name." - ) - validation_file: Optional[str] = ( - None # "The ID of an uploaded file that contains validation data." - ) - integrations: Optional[List[str]] = ( - None # "A list of integrations to enable for your fine-tuning job." - ) + hyperparameters: Optional[Hyperparameters] = None # "The hyperparameters used for the fine-tuning job." + suffix: Optional[str] = None # "A string of up to 18 characters that will be added to your fine-tuned model name." + validation_file: Optional[str] = None # "The ID of an uploaded file that contains validation data." + integrations: Optional[List[str]] = None # "A list of integrations to enable for your fine-tuning job." seed: Optional[int] = None # "The seed controls the reproducibility of the job." class LiteLLMFineTuningJobCreate(FineTuningJobCreate): custom_llm_provider: Optional[Literal["openai", "azure", "vertex_ai"]] = None - model_config = { - "extra": "allow" - } # This allows the model to accept additional fields + model_config = {"extra": "allow"} # This allows the model to accept additional fields AllEmbeddingInputValues = Union[str, List[str], List[int], List[List[int]]] @@ -1194,9 +1170,7 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False): prompt_cache_retention: Optional[str] stream_options: Optional[dict] top_logprobs: Optional[int] - partial_images: Optional[ - int - ] # Number of partial images to generate (1-3) for streaming image generation + partial_images: Optional[int] # Number of partial images to generate (1-3) for streaming image generation context_management: Optional[List[ContextManagementEntry]] """Context management configuration. E.g. [{\"type\": \"compaction\", \"compact_threshold\": 200000}] for server-side compaction (minimum 1000).""" @@ -1254,9 +1228,7 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} -ResponsesAPIStatus = Literal[ - "completed", "failed", "in_progress", "cancelled", "queued", "incomplete" -] +ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"] """ The status of the response generation. One of: completed, failed, in_progress, cancelled, queued, or incomplete. @@ -1287,9 +1259,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): parallel_tool_calls: Optional[bool] = None temperature: Optional[float] = None tool_choice: Optional[ToolChoice] = None - tools: Optional[ - Union[List[Tool], List[ResponseFunctionToolCall], List[Dict[str, Any]]] - ] = None + tools: Optional[Union[List[Tool], List[ResponseFunctionToolCall], List[Dict[str, Any]]]] = None top_p: Optional[float] = None max_output_tokens: Optional[int] = None previous_response_id: Optional[str] = None @@ -1345,12 +1315,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): return serialized return [ ( - { - k: v - for k, v in item.items() - if v is not None - or k not in ("status", "content", "encrypted_content") - } + {k: v for k, v in item.items() if v is not None or k not in ("status", "content", "encrypted_content")} if isinstance(item, dict) and item.get("type") == "reasoning" else item ) diff --git a/litellm/types/llms/openai_evals.py b/litellm/types/llms/openai_evals.py index 431dd34647f..768f6a4e80a 100644 --- a/litellm/types/llms/openai_evals.py +++ b/litellm/types/llms/openai_evals.py @@ -42,9 +42,7 @@ class DataSourceConfigStoredCompletions(TypedDict, total=False): """Optional metadata for filtering stored completions""" -DataSourceConfig = Union[ - DataSourceConfigCustom, DataSourceConfigLogs, DataSourceConfigStoredCompletions -] +DataSourceConfig = Union[DataSourceConfigCustom, DataSourceConfigLogs, DataSourceConfigStoredCompletions] class LLMAsJudgeGraderConfig(TypedDict, total=False): @@ -80,9 +78,7 @@ class CustomGraderConfig(TypedDict, total=False): """ID of the custom grading function""" -GraderConfig = Union[ - LLMAsJudgeGraderConfig, GroundTruthGraderConfig, CustomGraderConfig -] +GraderConfig = Union[LLMAsJudgeGraderConfig, GroundTruthGraderConfig, CustomGraderConfig] class CreateEvalRequest(TypedDict, total=False): @@ -235,9 +231,7 @@ class DataSourceInlineConfig(TypedDict, total=False): """List of inline samples to use for the run""" -RunDataSourceConfig = Union[ - DataSourceDatasetConfig, DataSourceSampleSetConfig, DataSourceInlineConfig -] +RunDataSourceConfig = Union[DataSourceDatasetConfig, DataSourceSampleSetConfig, DataSourceInlineConfig] class CompletionConfig(TypedDict, total=False): diff --git a/litellm/types/llms/stability.py b/litellm/types/llms/stability.py index c439a3b59e3..52666806203 100644 --- a/litellm/types/llms/stability.py +++ b/litellm/types/llms/stability.py @@ -21,9 +21,7 @@ class StabilityImageGenerationRequest(TypedDict, total=False): prompt: str # Required - text prompt for image generation negative_prompt: Optional[str] # What to avoid in the image - aspect_ratio: Optional[ - str - ] # e.g., "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21" + aspect_ratio: Optional[str] # e.g., "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21" seed: Optional[int] # Random seed for reproducibility (0 to 4294967294) output_format: Optional[Literal["jpeg", "png", "webp"]] # Output format model: Optional[str] # Model variant (e.g., "sd3.5-large", "sd3.5-medium") @@ -144,9 +142,7 @@ class StabilityRemoveBackgroundRequest(TypedDict, total=False): """ image: str # Required - Base64-encoded image - output_format: Optional[ - Literal["png", "webp"] - ] # Output format (no jpeg - needs transparency) + output_format: Optional[Literal["png", "webp"]] # Output format (no jpeg - needs transparency) class StabilityControlRequest(TypedDict, total=False): diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index b28fee51284..64a06825773 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -182,9 +182,7 @@ HarmBlockThreshold = Literal[ ] HarmBlockMethod = Literal["HARM_BLOCK_METHOD_UNSPECIFIED", "SEVERITY", "PROBABILITY"] -HarmProbability = Literal[ - "HARM_PROBABILITY_UNSPECIFIED", "NEGLIGIBLE", "LOW", "MEDIUM", "HIGH" -] +HarmProbability = Literal["HARM_PROBABILITY_UNSPECIFIED", "NEGLIGIBLE", "LOW", "MEDIUM", "HIGH"] HarmSeverity = Literal[ "HARM_SEVERITY_UNSPECIFIED", @@ -210,9 +208,7 @@ class GeminiThinkingConfig(TypedDict, total=False): GeminiResponseModalities = Literal["TEXT", "IMAGE", "AUDIO", "VIDEO"] -GeminiImageAspectRatio = Literal[ - "1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", "21:9" -] +GeminiImageAspectRatio = Literal["1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", "21:9"] GeminiImageSize = Literal["1K", "2K", "4K"] @@ -307,9 +303,7 @@ class UsageMetadata(TypedDict, total=False): cacheTokensDetails: List[PromptTokensDetails] thoughtsTokenCount: int responseTokensDetails: List[PromptTokensDetails] - candidatesTokensDetails: List[ - PromptTokensDetails - ] # Alternative key name used in some responses + candidatesTokensDetails: List[PromptTokensDetails] # Alternative key name used in some responses class TokenCountDetailsResponse(TypedDict): diff --git a/litellm/types/llms/watsonx.py b/litellm/types/llms/watsonx.py index 21e58500c6f..5ca419985f2 100644 --- a/litellm/types/llms/watsonx.py +++ b/litellm/types/llms/watsonx.py @@ -55,9 +55,7 @@ class WatsonXAIEndpoint(str, Enum): CHAT = "/ml/v1/text/chat" CHAT_STREAM = "/ml/v1/text/chat_stream" DEPLOYMENT_TEXT_GENERATION = "/ml/v1/deployments/{deployment_id}/text/generation" - DEPLOYMENT_TEXT_GENERATION_STREAM = ( - "/ml/v1/deployments/{deployment_id}/text/generation_stream" - ) + DEPLOYMENT_TEXT_GENERATION_STREAM = "/ml/v1/deployments/{deployment_id}/text/generation_stream" DEPLOYMENT_CHAT = "/ml/v1/deployments/{deployment_id}/text/chat" DEPLOYMENT_CHAT_STREAM = "/ml/v1/deployments/{deployment_id}/text/chat_stream" EMBEDDINGS = "/ml/v1/text/embeddings" diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index 6d8cb63a15c..9bccfed7c14 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -13,14 +13,10 @@ class CacheSettingsField(BaseModel): field_value: Any field_description: str field_default: Any = None - options: Optional[List[str]] = ( - None # For fields with predefined options/enum values - ) + options: Optional[List[str]] = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: Optional[str] = None # Documentation link for the field - redis_type: Optional[str] = ( - None # Which Redis type this field applies to (node, cluster, sentinel) - ) + redis_type: Optional[str] = None # Which Redis type this field applies to (node, cluster, sentinel) # Redis type descriptions diff --git a/litellm/types/management_endpoints/router_settings_endpoints.py b/litellm/types/management_endpoints/router_settings_endpoints.py index 6b18211cb70..929f2fde132 100644 --- a/litellm/types/management_endpoints/router_settings_endpoints.py +++ b/litellm/types/management_endpoints/router_settings_endpoints.py @@ -12,9 +12,7 @@ from pydantic import BaseModel, Field, field_validator class FallbackCreateRequest(BaseModel): """Request model for creating/updating fallbacks""" - model: str = Field( - description="The model name to configure fallbacks for (e.g., 'gpt-3.5-turbo')" - ) + model: str = Field(description="The model name to configure fallbacks for (e.g., 'gpt-3.5-turbo')") fallback_models: List[str] = Field( description="List of fallback model names in order of priority", min_length=1, @@ -75,9 +73,7 @@ class RouterSettingsField(BaseModel): field_value: Any field_description: str field_default: Any = None - options: Optional[List[str]] = ( - None # For fields with predefined options/enum values - ) + options: Optional[List[str]] = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: Optional[str] = None # Documentation link for the field diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 94e4c68f5e2..eeb814bf776 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -42,9 +42,7 @@ class MCPAuth(str, enum.Enum): # MCP Literals MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http, MCPTransport.stdio] -MCPSpecVersionType = Literal[ - MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025 -] +MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025] MCPAuthType = Optional[ Literal[ MCPAuth.none, @@ -214,7 +212,5 @@ class MCPPostCallResponseObject(BaseModel): Pydantic object used for MCP post_call_hook response """ - mcp_tool_call_response: List[ - Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource] - ] + mcp_tool_call_response: List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]] hidden_params: HiddenParams diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 809da6418d7..343ece91355 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -35,12 +35,8 @@ class MCPServer(BaseModel): disallowed_tools: Optional[List[str]] = None tool_name_to_display_name: Optional[Dict[str, str]] = None tool_name_to_description: Optional[Dict[str, str]] = None - allowed_params: Optional[Dict[str, List[str]]] = ( - None # map of tool names to allowed parameter lists - ) - static_headers: Optional[Dict[str, str]] = ( - None # static headers to forward to the MCP server - ) + allowed_params: Optional[Dict[str, List[str]]] = None # map of tool names to allowed parameter lists + static_headers: Optional[Dict[str, str]] = None # static headers to forward to the MCP server # Admin-configured env vars. Each entry is {name, value, scope, description}. # scope=="global" values are interpolated into static_headers using ${NAME}. # scope=="user" values must be supplied per-user. diff --git a/litellm/types/memory_management.py b/litellm/types/memory_management.py index bde54933699..81f655bf668 100644 --- a/litellm/types/memory_management.py +++ b/litellm/types/memory_management.py @@ -23,9 +23,7 @@ class LiteLLM_MemoryRow(BaseModel): class MemoryCreateRequest(BaseModel): key: str = Field(..., description="Memory key (acts as the namespace in the URL).") - value: str = Field( - ..., description="Memory content. Typically markdown/text for LLM context." - ) + value: str = Field(..., description="Memory content. Typically markdown/text for LLM context.") metadata: Optional[Any] = Field( default=None, description="Optional JSON metadata (tags, structured fields).", diff --git a/litellm/types/prompts/init_prompts.py b/litellm/types/prompts/init_prompts.py index eefd3d3dc87..6db714c333c 100644 --- a/litellm/types/prompts/init_prompts.py +++ b/litellm/types/prompts/init_prompts.py @@ -56,10 +56,7 @@ class PromptSpec(BaseModel): if "prompt_info" not in data: data["prompt_info"] = PromptInfo(prompt_type="config") elif "prompt_info" in data: - if ( - isinstance(data["prompt_info"], dict) - and data["prompt_info"].get("prompt_type") is None - ): + if isinstance(data["prompt_info"], dict) and data["prompt_info"].get("prompt_type") is None: data["prompt_info"]["prompt_type"] = "config" super().__init__(**data) @@ -73,9 +70,7 @@ class PromptTemplateBase(BaseModel): class PromptInfoResponse(BaseModel): prompt_spec: PromptSpec raw_prompt_template: Optional[PromptTemplateBase] = None - environments: Optional[List[str]] = ( - None # All environments this prompt is deployed to - ) + environments: Optional[List[str]] = None # All environments this prompt is deployed to class ListPromptsResponse(BaseModel): diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index 9774cbec614..47dd40df694 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -49,12 +49,8 @@ class RegisterPluginRequest(BaseModel): homepage: Optional[str] = Field(None, description="Plugin homepage URL") keywords: Optional[List[str]] = Field(None, description="Search keywords") category: Optional[str] = Field(None, description="Plugin category") - domain: Optional[str] = Field( - None, description="Skill domain (e.g., 'Productivity')" - ) - namespace: Optional[str] = Field( - None, description="Skill namespace within domain (e.g., 'workflows')" - ) + domain: Optional[str] = Field(None, description="Skill domain (e.g., 'Productivity')") + namespace: Optional[str] = Field(None, description="Skill namespace within domain (e.g., 'workflows')") class PluginResponse(BaseModel): @@ -125,6 +121,4 @@ class MarketplaceResponse(BaseModel): name: str = Field(..., description="Marketplace identifier") owner: PluginOwner = Field(..., description="Marketplace owner") - plugins: List[MarketplacePluginEntry] = Field( - default_factory=list, description="Available plugins" - ) + plugins: List[MarketplacePluginEntry] = Field(default_factory=list, description="Available plugins") diff --git a/litellm/types/proxy/cloudzero_endpoints.py b/litellm/types/proxy/cloudzero_endpoints.py index 3ac5795b358..c50c4f53df6 100644 --- a/litellm/types/proxy/cloudzero_endpoints.py +++ b/litellm/types/proxy/cloudzero_endpoints.py @@ -12,12 +12,8 @@ class CloudZeroInitRequest(BaseModel): """Request model for initializing CloudZero settings""" api_key: str = Field(..., description="CloudZero API key for authentication") - connection_id: str = Field( - ..., description="CloudZero connection ID for data submission" - ) - timezone: str = Field( - default="UTC", description="Timezone for date handling (default: UTC)" - ) + connection_id: str = Field(..., description="CloudZero connection ID for data submission") + timezone: str = Field(default="UTC", description="Timezone for date handling (default: UTC)") class CloudZeroInitResponse(BaseModel): @@ -30,19 +26,13 @@ class CloudZeroInitResponse(BaseModel): class CloudZeroExportRequest(BaseModel): """Request model for CloudZero export operations""" - limit: Optional[int] = Field( - None, description="Optional limit on number of records to export" - ) + limit: Optional[int] = Field(None, description="Optional limit on number of records to export") operation: str = Field( default="replace_hourly", description="CloudZero operation type (replace_hourly or sum)", ) - start_time_utc: Optional[datetime] = Field( - None, description="Start time for data export in UTC" - ) - end_time_utc: Optional[datetime] = Field( - None, description="End time for data export in UTC" - ) + start_time_utc: Optional[datetime] = Field(None, description="Start time for data export in UTC") + end_time_utc: Optional[datetime] = Field(None, description="End time for data export in UTC") class CloudZeroExportResponse(BaseModel): @@ -54,20 +44,14 @@ class CloudZeroExportResponse(BaseModel): dry_run_data: Optional[Dict[str, Any]] = Field( None, description="Dry run data including usage data and CBF transformed data" ) - summary: Optional[Dict[str, Any]] = Field( - None, description="Summary statistics for dry run" - ) + summary: Optional[Dict[str, Any]] = Field(None, description="Summary statistics for dry run") class CloudZeroSettingsView(BaseModel): """Response model for viewing CloudZero settings with masked API key""" - api_key_masked: Optional[str] = Field( - None, description="Masked API key showing only first 4 and last 4 characters" - ) - connection_id: Optional[str] = Field( - None, description="CloudZero connection ID for data submission" - ) + api_key_masked: Optional[str] = Field(None, description="Masked API key showing only first 4 and last 4 characters") + connection_id: Optional[str] = Field(None, description="CloudZero connection ID for data submission") timezone: Optional[str] = Field(None, description="Timezone for date handling") status: Optional[str] = Field(None, description="Configuration status") @@ -75,10 +59,6 @@ class CloudZeroSettingsView(BaseModel): class CloudZeroSettingsUpdate(BaseModel): """Request model for updating CloudZero settings""" - api_key: Optional[str] = Field( - None, description="New CloudZero API key for authentication" - ) - connection_id: Optional[str] = Field( - None, description="New CloudZero connection ID for data submission" - ) + api_key: Optional[str] = Field(None, description="New CloudZero API key for authentication") + connection_id: Optional[str] = Field(None, description="New CloudZero connection ID for data submission") timezone: Optional[str] = Field(None, description="New timezone for date handling") diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py index d73b73502b6..83082942390 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py @@ -19,9 +19,7 @@ class AzureTextModerationRequestBodyOptionalParams(TypedDict, total=False): outputType: Literal["FourSeverityLevels", "EightSeverityLevels"] -class AzureTextModerationGuardrailRequestBody( - AzureTextModerationRequestBodyOptionalParams -): +class AzureTextModerationGuardrailRequestBody(AzureTextModerationRequestBodyOptionalParams): """Configuration parameters for the Azure Text Moderation guardrail""" text: Required[str] diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/cisco_ai_defense.py b/litellm/types/proxy/guardrails/guardrail_hooks/cisco_ai_defense.py index f03fc9e1c32..90c4f4161de 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/cisco_ai_defense.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/cisco_ai_defense.py @@ -116,9 +116,7 @@ class CiscoAIDefenseGuardrailConfigModelOptionalParams(BaseModel): ) -class CiscoAIDefenseGuardrailConfigModel( - GuardrailConfigModel[CiscoAIDefenseGuardrailConfigModelOptionalParams] -): +class CiscoAIDefenseGuardrailConfigModel(GuardrailConfigModel[CiscoAIDefenseGuardrailConfigModelOptionalParams]): """Configuration parameters for the Cisco AI Defense guardrail.""" api_key: Optional[str] = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index ba5985935eb..e967e2b1d9a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py @@ -9,9 +9,7 @@ class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel): pass -class CrowdStrikeAIDRGuardrailConfigModel( - GuardrailConfigModel[CrowdStrikeAIDRGuardrailConfigModelOptionalParams] -): +class CrowdStrikeAIDRGuardrailConfigModel(GuardrailConfigModel[CrowdStrikeAIDRGuardrailConfigModelOptionalParams]): api_key: Optional[str] = Field( default=None, description="The CrowdStrike AIDR API key. Reads from CS_AIDR_TOKEN env var if None.", diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 2fe976463c4..28fb482b3af 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -98,9 +98,7 @@ class GenericGuardrailAPIRequest(BaseModel): description="LiteLLM library version running this proxy.", ) additional_provider_specific_params: Optional[Dict[str, Any]] = None - tool_calls: Optional[ - Union[List[ChatCompletionToolCallChunk], List[ChatCompletionMessageToolCall]] - ] = None + tool_calls: Optional[Union[List[ChatCompletionToolCallChunk], List[ChatCompletionMessageToolCall]]] = None model: Optional[str] = None # the model being used for the LLM call diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py b/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py index 24d8cd85a0b..d5e4ae226e2 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py @@ -42,9 +42,7 @@ class GraySwanGuardrailConfigModelOptionalParams(BaseModel): ) -class GraySwanGuardrailConfigModel( - GuardrailConfigModel[GraySwanGuardrailConfigModelOptionalParams] -): +class GraySwanGuardrailConfigModel(GuardrailConfigModel[GraySwanGuardrailConfigModelOptionalParams]): """Configuration parameters for the Gray Swan guardrail.""" api_key: Optional[str] = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/hiddenlayer.py b/litellm/types/proxy/guardrails/guardrail_hooks/hiddenlayer.py index a6c73dc411b..4a0e5a23389 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/hiddenlayer.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/hiddenlayer.py @@ -32,9 +32,7 @@ class HiddenlayerGuardrailConfigModel(GuardrailConfigModel): description="The Hiddenlayer Secret Key for the Hiddenlayer API.. If not provided, the `HIDDENLAYER_CLIENT_SECRET` environment variable is checked.", ) - version: Optional[int] = Field( - default=2, description="Hiddenlayer guardrail version to use." - ) + version: Optional[int] = Field(default=2, description="Hiddenlayer guardrail version to use.") @staticmethod def ui_friendly_name() -> str: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/javelin.py b/litellm/types/proxy/guardrails/guardrail_hooks/javelin.py index ba33e1adc25..fef597a9ef0 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/javelin.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/javelin.py @@ -79,9 +79,7 @@ class JavelinGuardResponse(TypedDict): assessments: List[ Dict[ str, - JavelinPromptInjectionAssessment - | JavelinTrustSafetyAssessment - | JavelinLanguageDetectionAssessment, + JavelinPromptInjectionAssessment | JavelinTrustSafetyAssessment | JavelinLanguageDetectionAssessment, ] ] @@ -89,21 +87,11 @@ class JavelinGuardResponse(TypedDict): class JavelinGuardrailConfigModel(GuardrailConfigModel): """Configuration parameters for the Javelin guardrail""" - guard_name: Optional[str] = Field( - default=None, description="Name of the Javelin guard to use" - ) - api_version: Optional[str] = Field( - default="v1", description="API version for Javelin service" - ) - metadata: Optional[Dict] = Field( - default=None, description="Additional metadata to send with requests" - ) - application: Optional[str] = Field( - default=None, description="Application name for Javelin service" - ) - config: Optional[Dict] = Field( - default=None, description="Configuration parameters for Javelin service" - ) + guard_name: Optional[str] = Field(default=None, description="Name of the Javelin guard to use") + api_version: Optional[str] = Field(default="v1", description="API version for Javelin service") + metadata: Optional[Dict] = Field(default=None, description="Additional metadata to send with requests") + application: Optional[str] = Field(default=None, description="Application name for Javelin service") + config: Optional[Dict] = Field(default=None, description="Configuration parameters for Javelin service") @staticmethod def ui_friendly_name() -> str: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/lasso.py b/litellm/types/proxy/guardrails/guardrail_hooks/lasso.py index 4fe183c181d..3f1ce6f488d 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/lasso.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/lasso.py @@ -16,9 +16,7 @@ class LassoGuardrailConfigModelOptionalParams(BaseModel): ) -class LassoGuardrailConfigModel( - GuardrailConfigModel[LassoGuardrailConfigModelOptionalParams] -): +class LassoGuardrailConfigModel(GuardrailConfigModel[LassoGuardrailConfigModelOptionalParams]): api_key: Optional[str] = Field( default=None, description="The API key for the Lasso guardrail. If not provided, the `LASSO_API_KEY` environment variable is checked.", diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py index 5fa701574c7..628ac0442de 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py @@ -8,22 +8,14 @@ from .base import GuardrailConfigModel class ModelArmorGuardrailConfigModel(GuardrailConfigModel): """Configuration parameters for Google Cloud Model Armor guardrail""" - template_id: Optional[str] = Field( - default=None, description="The ID of your Model Armor template" - ) - project_id: Optional[str] = Field( - default=None, description="Google Cloud project ID" - ) - location: Optional[str] = Field( - default=None, description="Google Cloud location/region (e.g., us-central1)" - ) + template_id: Optional[str] = Field(default=None, description="The ID of your Model Armor template") + project_id: Optional[str] = Field(default=None, description="Google Cloud project ID") + location: Optional[str] = Field(default=None, description="Google Cloud location/region (e.g., us-central1)") credentials: Optional[str] = Field( default=None, description="Path to Google Cloud credentials JSON file or JSON string", ) - api_endpoint: Optional[str] = Field( - default=None, description="Optional custom API endpoint for Model Armor" - ) + api_endpoint: Optional[str] = Field(default=None, description="Optional custom API endpoint for Model Armor") fail_on_error: Optional[bool] = Field( default=True, description="Whether to fail the request if Model Armor encounters an error", diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py index 0fcc0f2309a..1bc674cb2b5 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py @@ -8,11 +8,9 @@ from ..base import GuardrailConfigModel class BaseOpenAIModerationGuardrailConfigModel(GuardrailConfigModel): """Base configuration model for the OpenAI Moderation guardrail""" - model: Optional[Literal["omni-moderation-latest", "text-moderation-latest"]] = ( - Field( - default="omni-moderation-latest", - description="The OpenAI moderation model to use. 'omni-moderation-latest' supports more categorization options and multi-modal inputs. Defaults to 'omni-moderation-latest'.", - ) + model: Optional[Literal["omni-moderation-latest", "text-moderation-latest"]] = Field( + default="omni-moderation-latest", + description="The OpenAI moderation model to use. 'omni-moderation-latest' supports more categorization options and multi-modal inputs. Defaults to 'omni-moderation-latest'.", ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/pangea.py b/litellm/types/proxy/guardrails/guardrail_hooks/pangea.py index 820da727bba..53423103cdd 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/pangea.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/pangea.py @@ -16,9 +16,7 @@ class PangeaGuardrailConfigModelOptionalParams(BaseModel): ) -class PangeaGuardrailConfigModel( - GuardrailConfigModel[PangeaGuardrailConfigModelOptionalParams] -): +class PangeaGuardrailConfigModel(GuardrailConfigModel[PangeaGuardrailConfigModelOptionalParams]): api_key: Optional[str] = Field( default=None, description="The Pangea API key. Reads from PANGEA_API_KEY env var if None.", diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/pillar.py b/litellm/types/proxy/guardrails/guardrail_hooks/pillar.py index 17391c070d0..e4a9cce33be 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/pillar.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/pillar.py @@ -42,9 +42,7 @@ class PillarGuardrailConfigModelOptionalParams(BaseModel): ) -class PillarGuardrailConfigModel( - GuardrailConfigModel[PillarGuardrailConfigModelOptionalParams] -): +class PillarGuardrailConfigModel(GuardrailConfigModel[PillarGuardrailConfigModelOptionalParams]): """Configuration parameters for the Pillar Security guardrail""" api_key: Optional[str] = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py index b47e40196e0..c30fcc9e1be 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -20,9 +20,7 @@ class ToolPermissionRule(BaseModel): default=None, description="Regex pattern applied to the tool type (e.g., function)", ) - decision: Literal["allow", "deny"] = Field( - description="Whether to allow or deny this tool usage" - ) + decision: Literal["allow", "deny"] = Field(description="Whether to allow or deny this tool usage") allowed_param_patterns: Optional[Dict[str, str]] = Field( default=None, description="Optional regex map enforcing nested parameter values using dot/[] paths", @@ -51,9 +49,7 @@ class ToolPermissionRule(BaseModel): @model_validator(mode="after") def _ensure_target_present(self): if self.tool_name is None and self.tool_type is None: - raise ValueError( - "Each rule must specify at least a tool_name or tool_type regex" - ) + raise ValueError("Each rule must specify at least a tool_name or tool_type regex") return self @@ -63,9 +59,7 @@ class ToolResult(BaseModel): """ type: str = Field(default="tool_result", description="Should be 'tool_result'") - tool_use_id: str = Field( - description="ID of the tool use this result corresponds to" - ) + tool_use_id: str = Field(description="ID of the tool use this result corresponds to") content: str = Field(description="Result content") is_error: bool = Field(default=True, description="Whether this is an error result") diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/vigil_guard.py b/litellm/types/proxy/guardrails/guardrail_hooks/vigil_guard.py index 6d41c24eccd..bc580619ee0 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/vigil_guard.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/vigil_guard.py @@ -8,17 +8,11 @@ from .base import GuardrailConfigModel class VigilGuardGuardrailConfigModel(GuardrailConfigModel): api_base: Optional[str] = Field( default=None, - description=( - "Vigil Guard API base URL. " - "Falls back to the VIGIL_GUARD_URL environment variable." - ), + description=("Vigil Guard API base URL. Falls back to the VIGIL_GUARD_URL environment variable."), ) api_key: Optional[str] = Field( default=None, - description=( - "Vigil Guard API key. " - "Falls back to the VIGIL_GUARD_API_KEY environment variable." - ), + description=("Vigil Guard API key. Falls back to the VIGIL_GUARD_API_KEY environment variable."), ) @staticmethod diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py index a0cd280202c..f522f5b470a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py @@ -108,9 +108,7 @@ class ZscalerAIGuardConfigModel(GuardrailConfigModel): # Check for configuration issues assert api_base is not None # always set via env default above is_resolve_policy = api_base.endswith("/resolve-and-execute-policy") - is_execute_policy = ( - api_base.endswith("/execute-policy") and not is_resolve_policy - ) + is_execute_policy = api_base.endswith("/execute-policy") and not is_resolve_policy # Scenario A: execute-policy without policy_id if is_execute_policy and (policy_id is None or policy_id < 1): diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 948401fb8bf..b57d39f4c1a 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -48,35 +48,19 @@ class KeyMetricWithMetadata(MetricBase): class MetricWithMetadata(MetricBase): metadata: Dict[str, Any] = Field(default_factory=dict) # API key breakdown for this metric (e.g., which API keys are using this MCP server) - api_key_breakdown: Dict[str, KeyMetricWithMetadata] = Field( - default_factory=dict - ) # api_key -> {metrics, metadata} + api_key_breakdown: Dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} class BreakdownMetrics(BaseModel): """Breakdown of spend by different dimensions""" - mcp_servers: Dict[str, MetricWithMetadata] = Field( - default_factory=dict - ) # mcp_server -> {metrics, metadata} - models: Dict[str, MetricWithMetadata] = Field( - default_factory=dict - ) # model -> {metrics, metadata} - model_groups: Dict[str, MetricWithMetadata] = Field( - default_factory=dict - ) # model_group -> {metrics, metadata} - providers: Dict[str, MetricWithMetadata] = Field( - default_factory=dict - ) # provider -> {metrics, metadata} - endpoints: Dict[str, MetricWithMetadata] = Field( - default_factory=dict - ) # endpoint -> {metrics, metadata} - api_keys: Dict[str, KeyMetricWithMetadata] = Field( - default_factory=dict - ) # api_key -> {metrics, metadata} - entities: Dict[str, MetricWithMetadata] = Field( - default_factory=dict - ) # entity -> {metrics, metadata} + mcp_servers: Dict[str, MetricWithMetadata] = Field(default_factory=dict) # mcp_server -> {metrics, metadata} + models: Dict[str, MetricWithMetadata] = Field(default_factory=dict) # model -> {metrics, metadata} + model_groups: Dict[str, MetricWithMetadata] = Field(default_factory=dict) # model_group -> {metrics, metadata} + providers: Dict[str, MetricWithMetadata] = Field(default_factory=dict) # provider -> {metrics, metadata} + endpoints: Dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata} + api_keys: Dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} + entities: Dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata} class DailySpendData(BaseModel): diff --git a/litellm/types/proxy/management_endpoints/config_overrides.py b/litellm/types/proxy/management_endpoints/config_overrides.py index 6f5d661f57a..5c2ef519c78 100644 --- a/litellm/types/proxy/management_endpoints/config_overrides.py +++ b/litellm/types/proxy/management_endpoints/config_overrides.py @@ -56,9 +56,5 @@ class ConfigOverrideSettingsResponse(BaseModel): """Response model for config override settings GET endpoints.""" config_type: str = Field(description="The type of config override") - values: Dict[str, Any] = Field( - description="Current configuration values (sensitive fields decrypted)" - ) - field_schema: Dict[str, Any] = Field( - description="Schema information for UI rendering" - ) + values: Dict[str, Any] = Field(description="Current configuration values (sensitive fields decrypted)") + field_schema: Dict[str, Any] = Field(description="Schema information for UI rendering") diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index 6023094a920..18f37a65d44 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -25,13 +25,9 @@ class UserListResponse(BaseModel): class BulkUpdateUserRequest(BaseModel): """Request for bulk user updates""" - users: Optional[List[UpdateUserRequest]] = ( - None # List of specific user update requests - ) + users: Optional[List[UpdateUserRequest]] = None # List of specific user update requests all_users: Optional[bool] = False # Flag to update all users - user_updates: Optional[UpdateUserRequestNoUserIDorEmail] = ( - None # Updates to apply to all users when all_users=True - ) + user_updates: Optional[UpdateUserRequestNoUserIDorEmail] = None # Updates to apply to all users when all_users=True @field_validator("users", "all_users", "user_updates") @classmethod @@ -40,9 +36,7 @@ class BulkUpdateUserRequest(BaseModel): values = info.data if hasattr(info, "data") else {} # After all fields are set, validate the combination - if ( - info.field_name == "user_updates" - ): # This is the last field, do validation here + if info.field_name == "user_updates": # This is the last field, do validation here users = values.get("users") all_users = values.get("all_users", False) user_updates = v @@ -55,9 +49,7 @@ class BulkUpdateUserRequest(BaseModel): # Cannot specify both users list and all_users if users and all_users: - raise ValueError( - "Cannot specify both 'users' and 'all_users=True'. Choose one approach." - ) + raise ValueError("Cannot specify both 'users' and 'all_users=True'. Choose one approach.") return v diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index d214cdb4f5d..4c107a4cc98 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -61,12 +61,8 @@ class KeyUpdateFields(BaseModel): model_tpm_limit: Optional[Dict[str, Any]] = None model_rpm_limit: Optional[Dict[str, Any]] = None max_parallel_requests: Optional[int] = None - rpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] - ] = None - tpm_limit_type: Optional[ - Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] - ] = None + rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]] = None + tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]] = None # Temporary budget grants (auto-expire). `spend` deliberately omitted — bulk-zeroing it bypasses budget enforcement; admin-only via /key/update. temp_budget_increase: Optional[float] = None @@ -83,9 +79,7 @@ class KeyUpdateFields(BaseModel): def validate_temp_budget(self) -> "KeyUpdateFields": if self.temp_budget_increase is not None or self.temp_budget_expiry is not None: if self.temp_budget_increase is None or self.temp_budget_expiry is None: - raise ValueError( - "temp_budget_increase and temp_budget_expiry must be set together" - ) + raise ValueError("temp_budget_increase and temp_budget_expiry must be set together") return self @model_validator(mode="after") @@ -108,11 +102,7 @@ class BulkUpdateTeamKeysRequest(BaseModel): def validate_selection(self) -> "BulkUpdateTeamKeysRequest": has_key_ids = self.key_ids is not None and len(self.key_ids) > 0 if has_key_ids and self.all_keys_in_team: - raise ValueError( - "Provide either `key_ids` or `all_keys_in_team=True`, not both." - ) + raise ValueError("Provide either `key_ids` or `all_keys_in_team=True`, not both.") if not has_key_ids and not self.all_keys_in_team: - raise ValueError( - "Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`." - ) + raise ValueError("Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`.") return self diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index bbbfc0de9f8..db0c75e26ab 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -21,12 +21,8 @@ class UpdateUsefulLinksRequest(BaseModel): class NewModelGroupRequest(BaseModel): access_group: str # The access group name (e.g., "production-models") - model_names: Optional[List[str]] = ( - None # Existing model groups to include - tags ALL deployments for each name - ) - model_ids: Optional[List[str]] = ( - None # Specific deployment IDs to tag (more precise than model_names) - ) + model_names: Optional[List[str]] = None # Existing model groups to include - tags ALL deployments for each name + model_ids: Optional[List[str]] = None # Specific deployment IDs to tag (more precise than model_names) class NewModelGroupResponse(BaseModel): @@ -40,9 +36,7 @@ class UpdateModelGroupRequest(BaseModel): model_names: Optional[List[str]] = ( None # Updated list of model groups to include - tags ALL deployments for each name ) - model_ids: Optional[List[str]] = ( - None # Specific deployment IDs to tag (more precise than model_names) - ) + model_ids: Optional[List[str]] = None # Specific deployment IDs to tag (more precise than model_names) class DeleteModelGroupResponse(BaseModel): diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index 6270d2d0925..3b1ea8f572e 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -11,9 +11,7 @@ from pydantic import ( ) from pydantic_core.core_schema import SerializerFunctionWrapHandler -SCIM_ENTERPRISE_USER_SCHEMA = ( - "urn:ietf:params:scim:schemas:extension:enterprise:2.0:User" -) +SCIM_ENTERPRISE_USER_SCHEMA = "urn:ietf:params:scim:schemas:extension:enterprise:2.0:User" SCIM_ENTERPRISE_METADATA_KEY = "scim_enterprise" @@ -90,9 +88,7 @@ class SCIMUser(SCIMResource): ) @model_serializer(mode="wrap") - def _omit_absent_enterprise( - self, handler: SerializerFunctionWrapHandler - ) -> Dict[str, Any]: + def _omit_absent_enterprise(self, handler: SerializerFunctionWrapHandler) -> Dict[str, Any]: dumped = handler(self) if self.enterprise_user is None: dumped.pop(SCIM_ENTERPRISE_USER_SCHEMA, None) diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 771eb773c0d..7234cc2650f 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -70,9 +70,7 @@ class RoleMappings(LiteLLMPydanticObjectBase): which role to assign the user based on the roles mapping. """ - provider: str = Field( - description="SSO Provider name (e.g., 'google', 'microsoft', 'generic')" - ) + provider: str = Field(description="SSO Provider name (e.g., 'google', 'microsoft', 'generic')") group_claim: str = Field( description="The field name in the SSO token that contains the groups array (e.g., 'groups', 'roles')" ) diff --git a/litellm/types/proxy/policy_engine/pipeline_types.py b/litellm/types/proxy/policy_engine/pipeline_types.py index b7b82bff723..be9ea11e275 100644 --- a/litellm/types/proxy/policy_engine/pipeline_types.py +++ b/litellm/types/proxy/policy_engine/pipeline_types.py @@ -53,9 +53,7 @@ class PipelineStep(BaseModel): if v is None: return None if v not in VALID_PIPELINE_ACTIONS: - raise ValueError( - f"Invalid action '{v}'. Must be one of: {sorted(VALID_PIPELINE_ACTIONS)}" - ) + raise ValueError(f"Invalid action '{v}'. Must be one of: {sorted(VALID_PIPELINE_ACTIONS)}") return v @@ -79,9 +77,7 @@ class GuardrailPipeline(BaseModel): @classmethod def validate_mode(cls, v: str) -> str: if v not in VALID_PIPELINE_MODES: - raise ValueError( - f"Invalid mode '{v}'. Must be one of: {sorted(VALID_PIPELINE_MODES)}" - ) + raise ValueError(f"Invalid mode '{v}'. Must be one of: {sorted(VALID_PIPELINE_MODES)}") return v diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index cb6590688dc..2c7e8d5afc9 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -203,53 +203,29 @@ class PolicyDBResponse(BaseModel): default="production", description="One of: draft, published, production.", ) - parent_version_id: Optional[str] = Field( - default=None, description="Policy ID this version was cloned from." - ) + parent_version_id: Optional[str] = Field(default=None, description="Policy ID this version was cloned from.") is_latest: bool = Field( default=True, description="True if this is the latest version by version_number.", ) - published_at: Optional[datetime] = Field( - default=None, description="When this version was published." - ) - production_at: Optional[datetime] = Field( - default=None, description="When this version was promoted to production." - ) + published_at: Optional[datetime] = Field(default=None, description="When this version was published.") + production_at: Optional[datetime] = Field(default=None, description="When this version was promoted to production.") inherit: Optional[str] = Field(default=None, description="Parent policy name.") description: Optional[str] = Field(default=None, description="Policy description.") - guardrails_add: List[str] = Field( - default_factory=list, description="Guardrails to add." - ) - guardrails_remove: List[str] = Field( - default_factory=list, description="Guardrails to remove." - ) - condition: Optional[Dict[str, Any]] = Field( - default=None, description="Policy condition." - ) - pipeline: Optional[Dict[str, Any]] = Field( - default=None, description="Optional guardrail pipeline." - ) - created_at: Optional[datetime] = Field( - default=None, description="When the policy was created." - ) - updated_at: Optional[datetime] = Field( - default=None, description="When the policy was last updated." - ) - created_by: Optional[str] = Field( - default=None, description="Who created the policy." - ) - updated_by: Optional[str] = Field( - default=None, description="Who last updated the policy." - ) + guardrails_add: List[str] = Field(default_factory=list, description="Guardrails to add.") + guardrails_remove: List[str] = Field(default_factory=list, description="Guardrails to remove.") + condition: Optional[Dict[str, Any]] = Field(default=None, description="Policy condition.") + pipeline: Optional[Dict[str, Any]] = Field(default=None, description="Optional guardrail pipeline.") + created_at: Optional[datetime] = Field(default=None, description="When the policy was created.") + updated_at: Optional[datetime] = Field(default=None, description="When the policy was last updated.") + created_by: Optional[str] = Field(default=None, description="Who created the policy.") + updated_by: Optional[str] = Field(default=None, description="Who last updated the policy.") class PolicyListDBResponse(BaseModel): """Response for listing policies from the database.""" - policies: List[PolicyDBResponse] = Field( - default_factory=list, description="List of policies." - ) + policies: List[PolicyDBResponse] = Field(default_factory=list, description="List of policies.") total_count: int = Field(default=0, description="Total number of policies.") @@ -337,18 +313,10 @@ class PolicyAttachmentDBResponse(BaseModel): keys: List[str] = Field(default_factory=list, description="Key patterns.") models: List[str] = Field(default_factory=list, description="Model patterns.") tags: List[str] = Field(default_factory=list, description="Tag patterns.") - created_at: Optional[datetime] = Field( - default=None, description="When the attachment was created." - ) - updated_at: Optional[datetime] = Field( - default=None, description="When the attachment was last updated." - ) - created_by: Optional[str] = Field( - default=None, description="Who created the attachment." - ) - updated_by: Optional[str] = Field( - default=None, description="Who last updated the attachment." - ) + created_at: Optional[datetime] = Field(default=None, description="When the attachment was created.") + updated_at: Optional[datetime] = Field(default=None, description="When the attachment was last updated.") + created_by: Optional[str] = Field(default=None, description="Who created the attachment.") + updated_by: Optional[str] = Field(default=None, description="Who last updated the attachment.") class PolicyAttachmentListResponse(BaseModel): @@ -379,12 +347,8 @@ class PipelineTestRequest(BaseModel): class PolicyResolveRequest(BaseModel): """Request body for resolving effective policies/guardrails for a context.""" - team_alias: Optional[str] = Field( - default=None, description="Team alias to resolve for." - ) - key_alias: Optional[str] = Field( - default=None, description="Key alias to resolve for." - ) + team_alias: Optional[str] = Field(default=None, description="Team alias to resolve for.") + key_alias: Optional[str] = Field(default=None, description="Key alias to resolve for.") model: Optional[str] = Field(default=None, description="Model name to resolve for.") tags: Optional[List[str]] = Field(default=None, description="Tags to resolve for.") @@ -431,12 +395,8 @@ class AttachmentImpactResponse(BaseModel): default=0, description="Number of teams that would be affected (named + unnamed).", ) - unnamed_keys_count: int = Field( - default=0, description="Number of affected keys without an alias." - ) - unnamed_teams_count: int = Field( - default=0, description="Number of affected teams without an alias." - ) + unnamed_keys_count: int = Field(default=0, description="Number of affected keys without an alias.") + unnamed_teams_count: int = Field(default=0, description="Number of affected teams without an alias.") sample_keys: List[str] = Field( default_factory=list, description="Sample of affected key aliases (up to 10).", diff --git a/litellm/types/proxy/policy_engine/validation_types.py b/litellm/types/proxy/policy_engine/validation_types.py index e079febcc9e..15751c223aa 100644 --- a/litellm/types/proxy/policy_engine/validation_types.py +++ b/litellm/types/proxy/policy_engine/validation_types.py @@ -30,9 +30,7 @@ class PolicyValidationError(BaseModel): """ policy_name: str = Field(description="Name of the policy with the issue.") - error_type: PolicyValidationErrorType = Field( - description="Type of validation error." - ) + error_type: PolicyValidationErrorType = Field(description="Type of validation error.") message: str = Field(description="Human-readable error message.") field: Optional[str] = Field( default=None, diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index 84199e78340..60171f1ad57 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -12,9 +12,7 @@ class VantageInitRequest(BaseModel): """Request model for initializing Vantage settings""" api_key: str = Field(..., description="Vantage API key for authentication") - integration_token: str = Field( - ..., description="Vantage integration token for the cost-import endpoint" - ) + integration_token: str = Field(..., description="Vantage integration token for the cost-import endpoint") base_url: str = Field( default="https://api.vantage.sh", description="Vantage API base URL (default: https://api.vantage.sh)", @@ -42,20 +40,14 @@ class VantageExportRequest(BaseModel): None, description="Optional limit on number of records to export (default: no limit)", ) - start_time_utc: Optional[datetime] = Field( - None, description="Start time for data export in UTC" - ) - end_time_utc: Optional[datetime] = Field( - None, description="End time for data export in UTC" - ) + start_time_utc: Optional[datetime] = Field(None, description="Start time for data export in UTC") + end_time_utc: Optional[datetime] = Field(None, description="End time for data export in UTC") class VantageDryRunRequest(BaseModel): """Request model for Vantage dry-run operations (capped for preview)""" - limit: Optional[int] = Field( - 500, description="Limit on number of records to preview (default: 500)" - ) + limit: Optional[int] = Field(500, description="Limit on number of records to preview (default: 500)") class VantageExportResponse(BaseModel): @@ -66,9 +58,7 @@ class VantageExportResponse(BaseModel): dry_run_data: Optional[Dict[str, Any]] = Field( None, description="Dry run data including usage data and FOCUS transformed data" ) - summary: Optional[Dict[str, Any]] = Field( - None, description="Summary statistics for dry run" - ) + summary: Optional[Dict[str, Any]] = Field(None, description="Summary statistics for dry run") class VantageSettingsView(BaseModel): @@ -89,12 +79,8 @@ class VantageSettingsView(BaseModel): class VantageSettingsUpdate(BaseModel): """Request model for updating Vantage settings""" - api_key: Optional[str] = Field( - None, description="New Vantage API key for authentication" - ) - integration_token: Optional[str] = Field( - None, description="New Vantage integration token" - ) + api_key: Optional[str] = Field(None, description="New Vantage API key for authentication") + integration_token: Optional[str] = Field(None, description="New Vantage integration token") base_url: Optional[str] = Field(None, description="New Vantage API base URL") @field_validator("api_key", "integration_token") diff --git a/litellm/types/rag.py b/litellm/types/rag.py index 29e35d5fe8d..802bf7a6e9e 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -53,9 +53,7 @@ class OpenAIVectorStoreOptions(TypedDict, total=False): ttl_days: Optional[int] # Time-to-live in days for indexed content # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) - litellm_credential_name: Optional[ - str - ] # Credential name to load from litellm.credential_list + litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list api_key: Optional[str] # Direct API key (alternative to litellm_credential_name) api_base: Optional[str] # Direct API base (alternative to litellm_credential_name) @@ -83,21 +81,13 @@ class BedrockVectorStoreOptions(TypedDict, total=False): # Bedrock-specific options s3_bucket: Optional[str] # S3 bucket (auto-created if not provided) s3_prefix: Optional[str] # S3 key prefix (default: "data/") - embedding_model: Optional[ - str - ] # Embedding model (default: amazon.titan-embed-text-v2:0) + embedding_model: Optional[str] # Embedding model (default: amazon.titan-embed-text-v2:0) data_source_id: Optional[str] # For existing KB: override auto-detected DS - wait_for_ingestion: Optional[ - bool - ] # Wait for completion (default: False - returns immediately) - ingestion_timeout: Optional[ - int - ] # Timeout in seconds if wait_for_ingestion=True (default: 300) + wait_for_ingestion: Optional[bool] # Wait for completion (default: False - returns immediately) + ingestion_timeout: Optional[int] # Timeout in seconds if wait_for_ingestion=True (default: 300) # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) - litellm_credential_name: Optional[ - str - ] # Credential name to load from litellm.credential_list + litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list # AWS auth (uses BaseAWSLLM) aws_access_key_id: Optional[str] @@ -129,14 +119,10 @@ class VertexAIVectorStoreOptions(TypedDict, total=False): vector_store_id: str # RAG corpus ID (required for Vertex AI) # GCP config - vertex_project: Optional[ - str - ] # GCP project ID (uses env VERTEXAI_PROJECT if not set) + vertex_project: Optional[str] # GCP project ID (uses env VERTEXAI_PROJECT if not set) vertex_location: Optional[str] # GCP region (default: us-central1) vertex_credentials: Optional[str] # Path to credentials JSON (uses ADC if not set) - gcs_bucket: Optional[ - str - ] # GCS bucket for file uploads (uses env GCS_BUCKET_NAME if not set) + gcs_bucket: Optional[str] # GCS bucket for file uploads (uses env GCS_BUCKET_NAME if not set) # Import settings wait_for_import: Optional[bool] # Wait for import to complete (default: True) @@ -167,18 +153,12 @@ class S3VectorsVectorStoreOptions(TypedDict, total=False): index_name: Optional[str] # Vector index name (auto-creates if not provided) # Index configuration (for auto-creation) - dimension: Optional[ - int - ] # Vector dimension (auto-detected from embedding model, or default: 1024) + dimension: Optional[int] # Vector dimension (auto-detected from embedding model, or default: 1024) distance_metric: Optional[Literal["cosine", "euclidean"]] # Default: cosine - non_filterable_metadata_keys: Optional[ - List[str] - ] # Keys excluded from filtering (e.g., ["source_text"]) + non_filterable_metadata_keys: Optional[List[str]] # Keys excluded from filtering (e.g., ["source_text"]) # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) - litellm_credential_name: Optional[ - str - ] # Credential name to load from litellm.credential_list + litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list # AWS auth (uses BaseAWSLLM) aws_access_key_id: Optional[str] @@ -232,9 +212,7 @@ class RAGIngestOptions(TypedDict, total=False): name: Optional[str] # Optional pipeline name for logging ocr: Optional[RAGIngestOCROptions] # Optional OCR step - chunking_strategy: Optional[ - RAGChunkingStrategy - ] # RecursiveCharacterTextSplitter args + chunking_strategy: Optional[RAGChunkingStrategy] # RecursiveCharacterTextSplitter args embedding: Optional[RAGIngestEmbeddingOptions] # Embedding model config vector_store: RAGIngestVectorStoreOptions # OpenAI or Bedrock config diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 376d6f66603..9b0629156c1 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -64,9 +64,7 @@ class RerankResponseResult(TypedDict, total=False): class RerankResponse(BaseModel): id: Optional[str] = None - results: Optional[List[RerankResponseResult]] = ( - None # Contains index and relevance_score - ) + results: Optional[List[RerankResponseResult]] = None # Contains index and relevance_score meta: Optional[RerankResponseMeta] = None # Contains api_version and billed_units # Define private attributes using PrivateAttr diff --git a/litellm/types/router.py b/litellm/types/router.py index a1c571ed7f7..e74bd4ceeca 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -24,9 +24,7 @@ class ConfigurableClientsideParamsCustomAuth(TypedDict): api_base: str -CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = Optional[ - List[Union[str, ConfigurableClientsideParamsCustomAuth]] -] +CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = Optional[List[Union[str, ConfigurableClientsideParamsCustomAuth]]] class ModelConfig(BaseModel): @@ -107,9 +105,7 @@ class UpdateRouterConfig(BaseModel): class ModelInfo(BaseModel): - id: Optional[ - str - ] # Allow id to be optional on input, but it will always be present as a str in the model instance + id: Optional[str] # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. updated_at: Optional[datetime.datetime] = None updated_by: Optional[str] = None @@ -117,9 +113,7 @@ class ModelInfo(BaseModel): created_at: Optional[datetime.datetime] = None created_by: Optional[str] = None - base_model: Optional[str] = ( - None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking - ) + base_model: Optional[str] = None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking tier: Optional[Literal["free", "paid"]] = None """ @@ -203,9 +197,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): custom_llm_provider: Optional[str] = None tpm: Optional[int] = None rpm: Optional[int] = None - timeout: Optional[Union[float, str, httpx.Timeout]] = ( - None # if str, pass in as os.environ/ - ) + timeout: Optional[Union[float, str, httpx.Timeout]] = None # if str, pass in as os.environ/ stream_timeout: Optional[Union[float, str]] = ( None # timeout when making stream=True calls, if str, pass in as os.environ/ ) @@ -360,7 +352,9 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): stream_timeout: Optional[Union[float, str]] max_retries: Optional[int] organization: Optional[Union[List, str]] # for openai orgs - configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS # for allowing api base switching on finetuned models + configurable_clientside_auth_params: ( + CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS # for allowing api base switching on finetuned models + ) ## DROP PARAMS ## drop_params: Optional[bool] ## RESPONSES API → CHAT COMPLETIONS BRIDGE ## @@ -437,9 +431,7 @@ class Deployment(BaseModel): elif isinstance(model_info, dict): model_info = ModelInfo(**model_info) - for key in ( - SPECIAL_MODEL_INFO_PARAMS - ): # ensures custom pricing info is consistently in 'model_info' + for key in SPECIAL_MODEL_INFO_PARAMS: # ensures custom pricing info is consistently in 'model_info' field = getattr(litellm_params, key, None) if field is not None: setattr(model_info, key, field) @@ -482,12 +474,8 @@ class RouterErrors(enum.Enum): user_defined_ratelimit_error = "Deployment over user-defined ratelimit." no_deployments_available = "No deployments available for selected model" - no_deployments_with_tag_routing = ( - "Not allowed to access model due to tags configuration" - ) - no_deployments_with_provider_budget_routing = ( - "No deployments available - crossed budget" - ) + no_deployments_with_tag_routing = "Not allowed to access model due to tags configuration" + no_deployments_with_provider_budget_routing = "No deployments available - crossed budget" class AllowedFailsPolicy(BaseModel): @@ -700,9 +688,7 @@ class CustomRoutingStrategyBase: class RouterGeneralSettings(BaseModel): - async_only_mode: bool = Field( - default=False - ) # this will only initialize async clients. Good for memory utils + async_only_mode: bool = Field(default=False) # this will only initialize async clients. Good for memory utils pass_through_all_models: bool = Field( default=False ) # if passed a model not llm_router model list, pass through the request to litellm.acompletion/embedding @@ -812,12 +798,8 @@ class MockRouterTestingParams: return cls( mock_testing_fallbacks=extract_bool_param("mock_testing_fallbacks"), - mock_testing_context_fallbacks=extract_bool_param( - "mock_testing_context_fallbacks" - ), - mock_testing_content_policy_fallbacks=extract_bool_param( - "mock_testing_content_policy_fallbacks" - ), + mock_testing_context_fallbacks=extract_bool_param("mock_testing_context_fallbacks"), + mock_testing_content_policy_fallbacks=extract_bool_param("mock_testing_content_policy_fallbacks"), ) @@ -859,9 +841,7 @@ class AdaptiveRouterWeights(BaseModel): def _weights_sum_to_one(cls, v, info): q = info.data.get("quality", 0.7) if abs(q + v - 1.0) > 0.001: - raise ValueError( - f"weights must sum to 1.0, got quality={q} + cost={v} = {q + v}" - ) + raise ValueError(f"weights must sum to 1.0, got quality={q} + cost={v} = {q + v}") return v diff --git a/litellm/types/search.py b/litellm/types/search.py index e94477fe1f6..341e1efe1d9 100644 --- a/litellm/types/search.py +++ b/litellm/types/search.py @@ -63,9 +63,7 @@ class SearchToolInfoResponse(TypedDict, total=False): search_tool_info: Optional[dict] created_at: Optional[str] updated_at: Optional[str] - is_from_config: Optional[ - bool - ] # True if this tool is defined in config file, False if from DB + is_from_config: Optional[bool] # True if this tool is defined in config file, False if from DB class ListSearchToolsResponse(TypedDict): diff --git a/litellm/types/services.py b/litellm/types/services.py index 580d5653dec..e6155f36fbc 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -59,47 +59,21 @@ Metric types to use for each service - Pod Lock Manager only needs a gauge metric """ DEFAULT_SERVICE_CONFIGS = { - ServiceTypes.REDIS.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, - ServiceTypes.DB.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, - ServiceTypes.BATCH_WRITE_TO_DB.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, - ServiceTypes.RESET_BUDGET_JOB.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, - ServiceTypes.LITELLM.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, - ServiceTypes.ROUTER.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, - ServiceTypes.AUTH.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, - ServiceTypes.PROXY_PRE_CALL.value: { - "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM] - }, + ServiceTypes.REDIS.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.DB.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.BATCH_WRITE_TO_DB.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.RESET_BUDGET_JOB.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.LITELLM.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.ROUTER.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.AUTH.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, + ServiceTypes.PROXY_PRE_CALL.value: {"metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]}, # Operational metrics for DB Transaction Queues ServiceTypes.POD_LOCK_MANAGER.value: {"metrics": [ServiceMetrics.GAUGE]}, - ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE.value: { - "metrics": [ServiceMetrics.GAUGE] - }, - ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE.value: { - "metrics": [ServiceMetrics.GAUGE] - }, - ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE.value: { - "metrics": [ServiceMetrics.GAUGE] - }, - ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE.value: { - "metrics": [ServiceMetrics.GAUGE] - }, - ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE.value: { - "metrics": [ServiceMetrics.GAUGE] - }, + ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]}, + ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]}, + ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]}, + ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]}, + ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]}, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE.value: {"metrics": [ServiceMetrics.GAUGE]}, } @@ -126,9 +100,7 @@ class ServiceLoggerPayload(BaseModel): service: ServiceTypes = Field(description="who is this for? - postgres/redis") duration: float = Field(description="How long did the request take?") call_type: str = Field(description="The call of the service, being made") - event_metadata: Optional[dict] = Field( - description="The metadata logged during service success/failure" - ) + event_metadata: Optional[dict] = Field(description="The metadata logged during service success/failure") def to_json(self, **kwargs): try: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2013a5760b4..279e9b15fe7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -151,9 +151,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_max_reasoning_effort: Optional[bool] supports_output_config: Optional[bool] supports_image_size: Optional[bool] - bedrock_output_config_effort_ceiling: Optional[ - Literal["low", "medium", "high", "max", "xhigh"] - ] + bedrock_output_config_effort_ceiling: Optional[Literal["low", "medium", "high", "max", "xhigh"]] class SearchContextCostPerQuery(TypedDict, total=False): @@ -185,19 +183,13 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): max_output_tokens: Required[Optional[int]] input_cost_per_token: Required[Optional[float]] input_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - input_cost_per_token_priority: Optional[ - float - ] # OpenAI priority service tier pricing + input_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing cache_creation_input_token_cost: Optional[float] cache_creation_input_token_cost_above_200k_tokens: Optional[float] cache_creation_input_token_cost_above_1hr: Optional[float] cache_read_input_token_cost: Optional[float] - cache_read_input_token_cost_flex: Optional[ - float - ] # OpenAI flex service tier pricing - cache_read_input_token_cost_priority: Optional[ - float - ] # OpenAI priority service tier pricing + cache_read_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing + cache_read_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing cache_read_input_token_cost_above_200k_tokens: Optional[float] cache_read_input_token_cost_above_200k_tokens_priority: Optional[float] cache_read_input_token_cost_above_272k_tokens: Optional[float] @@ -206,20 +198,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_character: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models - input_cost_per_token_above_200k_tokens: Optional[ - float - ] # only for vertex ai gemini-2.5-pro models + input_cost_per_token_above_200k_tokens: Optional[float] # only for vertex ai gemini-2.5-pro models input_cost_per_token_above_200k_tokens_priority: Optional[float] - input_cost_per_token_above_272k_tokens: Optional[ - float - ] # GPT-5.4/5.4-pro: prompts >272K priced at 2x input + input_cost_per_token_above_272k_tokens: Optional[float] # GPT-5.4/5.4-pro: prompts >272K priced at 2x input input_cost_per_token_above_272k_tokens_priority: Optional[float] - input_cost_per_token_above_512k_tokens: Optional[ - float - ] # MiniMax-M3: prompts >512K priced at 2x input - input_cost_per_character_above_128k_tokens: Optional[ - float - ] # only for vertex ai models + input_cost_per_token_above_512k_tokens: Optional[float] # MiniMax-M3: prompts >512K priced at 2x input + input_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models input_cost_per_query: Optional[float] # only for rerank models input_cost_per_image: Optional[float] # only for vertex ai models input_cost_per_image_token: Optional[float] # for gpt-image-1 and similar models @@ -230,9 +214,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_batches: Optional[float] output_cost_per_token: Required[Optional[float]] output_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - output_cost_per_token_priority: Optional[ - float - ] # OpenAI priority service tier pricing + output_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing regional_processing_uplift_multiplier_eu: Optional[ float ] # OpenAI EU data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%) @@ -241,23 +223,13 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): ] # OpenAI US data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%) output_cost_per_character: Optional[float] # only for vertex ai models output_cost_per_audio_token: Optional[float] - output_cost_per_token_above_128k_tokens: Optional[ - float - ] # only for vertex ai models - output_cost_per_token_above_200k_tokens: Optional[ - float - ] # only for vertex ai gemini-2.5-pro models + output_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models + output_cost_per_token_above_200k_tokens: Optional[float] # only for vertex ai gemini-2.5-pro models output_cost_per_token_above_200k_tokens_priority: Optional[float] - output_cost_per_token_above_272k_tokens: Optional[ - float - ] # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output + output_cost_per_token_above_272k_tokens: Optional[float] # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output output_cost_per_token_above_272k_tokens_priority: Optional[float] - output_cost_per_token_above_512k_tokens: Optional[ - float - ] # MiniMax-M3: prompts >512K priced at 2x output - output_cost_per_character_above_128k_tokens: Optional[ - float - ] # only for vertex ai models + output_cost_per_token_above_512k_tokens: Optional[float] # MiniMax-M3: prompts >512K priced at 2x output + output_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models output_cost_per_image: Optional[float] output_cost_per_image_token: Optional[float] output_vector_size: Optional[int] @@ -271,16 +243,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): ocr_cost_per_page: Optional[float] # for OCR models ocr_cost_per_credit: Optional[float] # for OCR models priced by credit annotation_cost_per_page: Optional[float] # for OCR models - search_context_cost_per_query: Optional[ - SearchContextCostPerQuery - ] # Cost for using web search tool + search_context_cost_per_query: Optional[SearchContextCostPerQuery] # Cost for using web search tool web_search_billing_unit: Optional[ Literal["per_query", "per_prompt"] ] # "per_query" (Gemini 3.x) or "per_prompt" (Gemini 2.x) citation_cost_per_token: Optional[float] # Cost per citation token for Perplexity - tiered_pricing: Optional[ - List[Dict[str, Any]] - ] # Tiered pricing structure for models like Dashscope + tiered_pricing: Optional[List[Dict[str, Any]]] # Tiered pricing structure for models like Dashscope litellm_provider: Required[str] mode: Required[ Literal[ @@ -1013,9 +981,7 @@ class FunctionCall(OpenAIObject): class Function(OpenAIObject): arguments: str - name: Optional[ - str - ] # can be None - openai e.g.: ChoiceDeltaToolCallFunction(arguments='{"', name=None), type=None) + name: Optional[str] # can be None - openai e.g.: ChoiceDeltaToolCallFunction(arguments='{"', name=None), type=None) def __init__( self, @@ -1024,9 +990,7 @@ class Function(OpenAIObject): **params, ): if arguments is None: - if params.get("parameters", None) is not None and isinstance( - params["parameters"], dict - ): + if params.get("parameters", None) is not None and isinstance(params["parameters"], dict): arguments = json.dumps(params["parameters"]) params.pop("parameters") else: @@ -1167,9 +1131,7 @@ ChatCompletionMessage(content='This is a test', role='assistant', function_call= """ -def add_provider_specific_fields( - object: BaseModel, provider_specific_fields: Optional[Dict[str, Any]] -): +def add_provider_specific_fields(object: BaseModel, provider_specific_fields: Optional[Dict[str, Any]]): if not provider_specific_fields: # set if provider_specific_fields is not empty return setattr(object, "provider_specific_fields", provider_specific_fields) @@ -1183,9 +1145,7 @@ class Message(SafeAttributeModel, OpenAIObject): audio: Optional[ChatCompletionAudioResponse] = None images: Optional[List[ImageURLListItem]] = None reasoning_content: Optional[str] = None - thinking_blocks: Optional[ - List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] - ] = None + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None reasoning_items: Optional[List[ChatCompletionReasoningItem]] = None provider_specific_fields: Optional[Dict[str, Any]] = Field(default=None) annotations: Optional[List[ChatCompletionAnnotation]] = None @@ -1200,11 +1160,7 @@ class Message(SafeAttributeModel, OpenAIObject): images: Optional[List[ImageURLListItem]] = None, provider_specific_fields: Optional[Dict[str, Any]] = None, reasoning_content: Optional[str] = None, - thinking_blocks: Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ] = None, + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None, reasoning_items: Optional[List[ChatCompletionReasoningItem]] = None, annotations: Optional[List[ChatCompletionAnnotation]] = None, **params, @@ -1212,16 +1168,10 @@ class Message(SafeAttributeModel, OpenAIObject): init_values: Dict[str, Any] = { "content": content, "role": role or "assistant", # handle null input - "function_call": ( - FunctionCall(**function_call) if function_call is not None else None - ), + "function_call": (FunctionCall(**function_call) if function_call is not None else None), "tool_calls": ( [ - ( - ChatCompletionMessageToolCall(**tool_call) - if isinstance(tool_call, dict) - else tool_call - ) + (ChatCompletionMessageToolCall(**tool_call) if isinstance(tool_call, dict) else tool_call) for tool_call in tool_calls ] if tool_calls is not None and len(tool_calls) > 0 @@ -1307,9 +1257,7 @@ class Message(SafeAttributeModel, OpenAIObject): class Delta(SafeAttributeModel, OpenAIObject): reasoning_content: Optional[str] = None - thinking_blocks: Optional[ - List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] - ] = None + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None reasoning_items: Optional[List[ChatCompletionReasoningItem]] = None provider_specific_fields: Optional[Dict[str, Any]] = Field(default=None) @@ -1322,11 +1270,7 @@ class Delta(SafeAttributeModel, OpenAIObject): audio: Optional[ChatCompletionAudioResponse] = None, images: Optional[List[ImageURLListItem]] = None, reasoning_content: Optional[str] = None, - thinking_blocks: Optional[ - List[ - Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock] - ] - ] = None, + thinking_blocks: Optional[List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]]] = None, reasoning_items: Optional[List[ChatCompletionReasoningItem]] = None, annotations: Optional[List[ChatCompletionAnnotation]] = None, **params, @@ -1440,9 +1384,7 @@ class Choices(SafeAttributeModel, OpenAIObject): mapped = map_finish_reason(finish_reason) params["finish_reason"] = mapped if finish_reason != mapped: - provider_specific_fields = ( - dict(provider_specific_fields) if provider_specific_fields else {} - ) + provider_specific_fields = dict(provider_specific_fields) if provider_specific_fields else {} provider_specific_fields["native_finish_reason"] = finish_reason else: params["finish_reason"] = "stop" @@ -1459,11 +1401,7 @@ class Choices(SafeAttributeModel, OpenAIObject): params["message"] = Message(**message) elif isinstance(message, BaseModel): # Normalize provider/OpenAI SDK message models into LiteLLM's Message type. - dump = ( - message.model_dump() - if hasattr(message, "model_dump") - else message.dict() - ) + dump = message.model_dump() if hasattr(message, "model_dump") else message.dict() params["message"] = Message(**dump) if logprobs is not None: if isinstance(logprobs, dict): @@ -1501,9 +1439,7 @@ class Choices(SafeAttributeModel, OpenAIObject): setattr(self, key, value) -class CompletionTokensDetailsWrapper( - CompletionTokensDetails -): # wrapper for older openai versions +class CompletionTokensDetailsWrapper(CompletionTokensDetails): # wrapper for older openai versions text_tokens: Optional[int] = None """Text tokens generated by the model.""" @@ -1603,12 +1539,8 @@ class Usage(SafeAttributeModel, CompletionUsage): completion_tokens: Optional[int] = None, total_tokens: Optional[int] = None, reasoning_tokens: Optional[int] = None, - prompt_tokens_details: Optional[ - Union[PromptTokensDetailsWrapper, PromptTokensDetails, dict] - ] = None, - completion_tokens_details: Optional[ - Union[CompletionTokensDetailsWrapper, dict] - ] = None, + prompt_tokens_details: Optional[Union[PromptTokensDetailsWrapper, PromptTokensDetails, dict]] = None, + completion_tokens_details: Optional[Union[CompletionTokensDetailsWrapper, dict]] = None, server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, @@ -1619,9 +1551,7 @@ class Usage(SafeAttributeModel, CompletionUsage): # First, handle existing completion_tokens_details if completion_tokens_details: if isinstance(completion_tokens_details, dict): - _completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details - ) + _completion_tokens_details = CompletionTokensDetailsWrapper(**completion_tokens_details) elif isinstance(completion_tokens_details, CompletionTokensDetails): _completion_tokens_details = completion_tokens_details @@ -1637,10 +1567,7 @@ class Usage(SafeAttributeModel, CompletionUsage): # Auto-calculate text_tokens only if provider didn't set it explicitly # Formula: text_tokens = completion_tokens - reasoning_tokens - image_tokens - audio_tokens - if ( - _completion_tokens_details.text_tokens is None - and completion_tokens is not None - ): + if _completion_tokens_details.text_tokens is None and completion_tokens is not None: calculated_text_tokens = completion_tokens - reasoning_tokens # Subtract other modality tokens if present @@ -1658,49 +1585,33 @@ class Usage(SafeAttributeModel, CompletionUsage): # guarantee prompt_token_details is always a PromptTokensDetailsWrapper if prompt_tokens_details: if isinstance(prompt_tokens_details, dict): - _prompt_tokens_details = PromptTokensDetailsWrapper( - **prompt_tokens_details - ) + _prompt_tokens_details = PromptTokensDetailsWrapper(**prompt_tokens_details) elif isinstance(prompt_tokens_details, PromptTokensDetails): - _prompt_tokens_details = PromptTokensDetailsWrapper( - **prompt_tokens_details.model_dump() - ) + _prompt_tokens_details = PromptTokensDetailsWrapper(**prompt_tokens_details.model_dump()) elif isinstance(prompt_tokens_details, PromptTokensDetailsWrapper): _prompt_tokens_details = prompt_tokens_details ## DEEPSEEK MAPPING ## - if "prompt_cache_hit_tokens" in params and isinstance( - params["prompt_cache_hit_tokens"], int - ): + if "prompt_cache_hit_tokens" in params and isinstance(params["prompt_cache_hit_tokens"], int): if _prompt_tokens_details is None: - _prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=params["prompt_cache_hit_tokens"] - ) + _prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=params["prompt_cache_hit_tokens"]) else: _prompt_tokens_details.cached_tokens = params["prompt_cache_hit_tokens"] ## ANTHROPIC MAPPING ## - if "cache_read_input_tokens" in params and isinstance( - params["cache_read_input_tokens"], int - ): + if "cache_read_input_tokens" in params and isinstance(params["cache_read_input_tokens"], int): if _prompt_tokens_details is None: - _prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=params["cache_read_input_tokens"] - ) + _prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=params["cache_read_input_tokens"]) else: _prompt_tokens_details.cached_tokens = params["cache_read_input_tokens"] - if "cache_creation_input_tokens" in params and isinstance( - params["cache_creation_input_tokens"], int - ): + if "cache_creation_input_tokens" in params and isinstance(params["cache_creation_input_tokens"], int): if _prompt_tokens_details is None: _prompt_tokens_details = PromptTokensDetailsWrapper( cache_creation_tokens=params["cache_creation_input_tokens"] ) else: - _prompt_tokens_details.cache_creation_tokens = params[ - "cache_creation_input_tokens" - ] + _prompt_tokens_details.cache_creation_tokens = params["cache_creation_input_tokens"] super().__init__( prompt_tokens=prompt_tokens or 0, @@ -1724,20 +1635,14 @@ class Usage(SafeAttributeModel, CompletionUsage): del self.cost ## ANTHROPIC MAPPING ## - if "cache_creation_input_tokens" in params and isinstance( - params["cache_creation_input_tokens"], int - ): + if "cache_creation_input_tokens" in params and isinstance(params["cache_creation_input_tokens"], int): self._cache_creation_input_tokens = params["cache_creation_input_tokens"] - if "cache_read_input_tokens" in params and isinstance( - params["cache_read_input_tokens"], int - ): + if "cache_read_input_tokens" in params and isinstance(params["cache_read_input_tokens"], int): self._cache_read_input_tokens = params["cache_read_input_tokens"] ## DEEPSEEK MAPPING ## - if "prompt_cache_hit_tokens" in params and isinstance( - params["prompt_cache_hit_tokens"], int - ): + if "prompt_cache_hit_tokens" in params and isinstance(params["prompt_cache_hit_tokens"], int): self._cache_read_input_tokens = params["prompt_cache_hit_tokens"] for k, v in params.items(): @@ -1859,9 +1764,7 @@ class ModelResponseStream(ModelResponseBase): def __init__( self, - choices: Optional[ - Union[List[StreamingChoices], Union[StreamingChoices, dict, BaseModel]] - ] = None, + choices: Optional[Union[List[StreamingChoices], Union[StreamingChoices, dict, BaseModel]]] = None, id: Optional[str] = None, created: Optional[int] = None, provider_specific_fields: Optional[Dict[str, Any]] = None, @@ -1896,9 +1799,7 @@ class ModelResponseStream(ModelResponseBase): kwargs["usage"] = Usage(**kwargs["usage"]) elif isinstance(kwargs["usage"], BaseModel): dump = ( - kwargs["usage"].model_dump() - if hasattr(kwargs["usage"], "model_dump") - else kwargs["usage"].dict() + kwargs["usage"].model_dump() if hasattr(kwargs["usage"], "model_dump") else kwargs["usage"].dict() ) kwargs["usage"] = Usage(**dump) @@ -1958,11 +1859,7 @@ class ModelResponse(ModelResponseBase): elif isinstance(choice, dict): _new_choice = Choices(**choice) # type: ignore elif isinstance(choice, BaseModel): - dump = ( - choice.model_dump() - if hasattr(choice, "model_dump") - else choice.dict() - ) + dump = choice.model_dump() if hasattr(choice, "model_dump") else choice.dict() _new_choice = Choices(**dump) # type: ignore else: _new_choice = choice @@ -1983,9 +1880,7 @@ class ModelResponse(ModelResponseBase): if isinstance(usage, dict): usage = Usage(**usage) elif isinstance(usage, BaseModel): - dump = ( - usage.model_dump() if hasattr(usage, "model_dump") else usage.dict() - ) + dump = usage.model_dump() if hasattr(usage, "model_dump") else usage.dict() usage = Usage(**dump) else: usage = usage @@ -2468,9 +2363,7 @@ class TranscriptionUsageTokensObject(BaseModel): class TranscriptionResponse(OpenAIObject): text: Optional[str] = None - usage: Optional[ - Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject] - ] = None + usage: Optional[Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject]] = None _hidden_params: dict = {} _response_headers: Optional[dict] = None @@ -2690,23 +2583,17 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata): Specific metadata k,v pairs logged to integration for easier cost tracking and prompt management """ - spend_logs_metadata: Optional[ - dict - ] # special param to log k,v pairs to spendlogs for a call + spend_logs_metadata: Optional[dict] # special param to log k,v pairs to spendlogs for a call requester_ip_address: Optional[str] user_agent: Optional[str] requester_metadata: Optional[dict] - requester_custom_headers: Optional[ - Dict[str, str] - ] # Log any custom (`x-`) headers sent by the client to the proxy. + requester_custom_headers: Optional[Dict[str, str]] # Log any custom (`x-`) headers sent by the client to the proxy. prompt_management_metadata: Optional[StandardLoggingPromptManagementMetadata] mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] applied_guardrails: Optional[List[str]] usage_object: Optional[dict] - cold_storage_object_key: Optional[ - str - ] # S3/GCS object key for cold storage retrieval + cold_storage_object_key: Optional[str] # S3/GCS object key for cold storage retrieval team_alias: Optional[str] team_id: Optional[str] @@ -2786,17 +2673,13 @@ class GuardrailMode(TypedDict, total=False): default: Optional[Union[str, List[str]]] -GuardrailStatus = Literal[ - "success", "guardrail_intervened", "guardrail_failed_to_respond", "not_run" -] +GuardrailStatus = Literal["success", "guardrail_intervened", "guardrail_failed_to_respond", "not_run"] class StandardLoggingGuardrailInformation(TypedDict, total=False): guardrail_name: Optional[str] guardrail_provider: Optional[str] - guardrail_mode: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks], GuardrailMode] - ] + guardrail_mode: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], GuardrailMode]] guardrail_request: Optional[dict] guardrail_response: Optional[Union[dict, str, List[dict]]] guardrail_status: GuardrailStatus @@ -2928,14 +2811,10 @@ class CostBreakdown(TypedDict, total=False): input_cost: float # Cost of raw (non-cached) input tokens only cache_read_cost: float # Cost of cache-read tokens (discounted rate) cache_creation_cost: float # Cost of cache-write tokens (premium rate) - output_cost: ( - float # Cost of output/completion tokens (includes reasoning if applicable) - ) + output_cost: float # Cost of output/completion tokens (includes reasoning if applicable) total_cost: float # Total cost (input + output + tool usage) tool_usage_cost: float # Cost of usage of built-in tools - additional_costs: Dict[ - str, float - ] # Free-form additional costs (e.g., {"azure_model_router_flat_cost": 0.00014}) + additional_costs: Dict[str, float] # Free-form additional costs (e.g., {"azure_model_router_flat_cost": 0.00014}) original_cost: float # Cost before discount (optional) discount_percent: float # Discount percentage applied (e.g., 0.05 = 5%) (optional) discount_amount: float # Discount amount in USD (optional) @@ -2981,9 +2860,7 @@ class StandardLoggingPayload(TypedDict): stream: Optional[bool] response_cost: float cost_breakdown: Optional[CostBreakdown] # Detailed cost breakdown - response_cost_failure_debug_info: Optional[ - StandardLoggingModelCostFailureDebugInformation - ] + response_cost_failure_debug_info: Optional[StandardLoggingModelCostFailureDebugInformation] status: StandardLoggingPayloadStatus status_fields: StandardLoggingPayloadStatusFields custom_llm_provider: Optional[str] @@ -3023,9 +2900,7 @@ from typing import AsyncIterator, Iterator class CustomStreamingDecoder: async def aiter_bytes( self, iterator: AsyncIterator[bytes] - ) -> AsyncIterator[ - Optional[Union[GenericStreamingChunk, StreamingChatCompletionChunk]] - ]: + ) -> AsyncIterator[Optional[Union[GenericStreamingChunk, StreamingChatCompletionChunk]]]: raise NotImplementedError def iter_bytes( @@ -3308,9 +3183,7 @@ all_litellm_params = ( class KeyGenerationConfig(TypedDict, total=False): - required_params: List[ - str - ] # specify params that must be present in the key generation request + required_params: List[str] # specify params that must be present in the key generation request class TeamUIKeyGenerationConfig(KeyGenerationConfig): @@ -3507,9 +3380,7 @@ OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = { ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai"] -LIST_BATCHES_SUPPORTED_PROVIDERS: frozenset[str] = frozenset( - get_args(ListBatchesSupportedProvider) -) +LIST_BATCHES_SUPPORTED_PROVIDERS: frozenset[str] = frozenset(get_args(ListBatchesSupportedProvider)) class SearchProviders(str, Enum): @@ -3561,9 +3432,7 @@ class LiteLLMLoggingBaseClass: def pre_call(self, input, api_key, model=None, additional_args={}): pass - def post_call( - self, original_response, input=None, api_key=None, additional_args={} - ): + def post_call(self, original_response, input=None, api_key=None, additional_args={}): pass @@ -3663,9 +3532,7 @@ class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase): _hidden_params: dict = {} @field_serializer("results") - def _serialize_results( - self, results: OpenAIRealtimeStreamList - ) -> List[Dict[str, Any]]: + def _serialize_results(self, results: OpenAIRealtimeStreamList) -> List[Dict[str, Any]]: return [dict(event) for event in results] def __contains__(self, key): @@ -3721,12 +3588,12 @@ class ExtractedFileData(TypedDict): class SpecialEnums(Enum): LITELM_MANAGED_FILE_ID_PREFIX = "litellm_proxy" - LITELLM_MANAGED_FILE_COMPLETE_STR = "litellm_proxy:{};unified_id,{};target_model_names,{};llm_output_file_id,{};llm_output_file_model_id,{}" - - LITELLM_MANAGED_RESPONSE_COMPLETE_STR = ( - "litellm:custom_llm_provider:{};model_id:{};response_id:{}" + LITELLM_MANAGED_FILE_COMPLETE_STR = ( + "litellm_proxy:{};unified_id,{};target_model_names,{};llm_output_file_id,{};llm_output_file_model_id,{}" ) + LITELLM_MANAGED_RESPONSE_COMPLETE_STR = "litellm:custom_llm_provider:{};model_id:{};response_id:{}" + LITELLM_MANAGED_BATCH_COMPLETE_STR = "litellm_proxy;model_id:{};llm_batch_id:{}" LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR = ( @@ -3735,13 +3602,9 @@ class SpecialEnums(Enum): LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR = "litellm_proxy;model_id:{};generic_response_id:{}" # generic implementation of 'managed batches' - used for finetuning and any future work. - LITELLM_MANAGED_VIDEO_COMPLETE_STR = ( - "litellm:custom_llm_provider:{};model_id:{};video_id:{}" - ) + LITELLM_MANAGED_VIDEO_COMPLETE_STR = "litellm:custom_llm_provider:{};model_id:{};video_id:{}" - LITELLM_PASSTHROUGH_MANAGED_ID_COMPLETE_STR = ( - "litellm_proxy:passthrough;provider:{};unified_id,{};raw_id,{}" - ) + LITELLM_PASSTHROUGH_MANAGED_ID_COMPLETE_STR = "litellm_proxy:passthrough;provider:{};unified_id,{};raw_id,{}" class ServiceTier(Enum): diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 6adfbf4fd35..0076e87e2bc 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -90,9 +90,7 @@ class VectorStoreSearchResult(TypedDict, total=False): class VectorStoreSearchResponse(TypedDict, total=False): """Response after searching a vector store""" - object: Literal[ - "vector_store.search_results.page" - ] # Always "vector_store.search_results.page" + object: Literal["vector_store.search_results.page"] # Always "vector_store.search_results.page" search_query: Optional[str] data: Optional[List[VectorStoreSearchResult]] @@ -176,9 +174,7 @@ class VertexSearchEngineExtraBody(VertexSearchDataStoreExtraBody, total=False): class VectorStoreExpirationPolicy(TypedDict, total=False): """The expiration policy for a vector store""" - anchor: Literal[ - "last_active_at" - ] # Anchor timestamp after which the expiration policy applies + anchor: Literal["last_active_at"] # Anchor timestamp after which the expiration policy applies days: int # Number of days after anchor time that the vector store will expire @@ -225,12 +221,8 @@ class VectorStoreCreateOptionalRequestParams(TypedDict, total=False): name: Optional[str] # Name of the vector store file_ids: Optional[List[str]] # List of File IDs that the vector store should use - expires_after: Optional[ - VectorStoreExpirationPolicy - ] # Expiration policy for the vector store - chunking_strategy: Optional[ - VectorStoreChunkingStrategy - ] # Chunking strategy for the files + expires_after: Optional[VectorStoreExpirationPolicy] # Expiration policy for the vector store + chunking_strategy: Optional[VectorStoreChunkingStrategy] # Chunking strategy for the files metadata: Optional[Dict[str, str]] # Set of key-value pairs for metadata @@ -252,9 +244,7 @@ class VectorStoreCreateResponse(TypedDict, total=False): status: Literal["expired", "in_progress", "completed"] # Status of the vector store expires_after: Optional[VectorStoreExpirationPolicy] # Expiration policy expires_at: Optional[int] # Unix timestamp of when the vector store expires - last_active_at: Optional[ - int - ] # Unix timestamp of when the vector store was last active + last_active_at: Optional[int] # Unix timestamp of when the vector store was last active metadata: Optional[Dict[str, str]] # Metadata associated with the vector store diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index ec0277c789a..30b862886bc 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -73,12 +73,8 @@ class VideoCreateOptionalRequestParams(TypedDict, total=False): """ input_reference: Optional[FileTypes] # File reference for input image - image: Optional[ - Any - ] # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object - parameters: Optional[ - Dict[str, Any] - ] # Provider-specific parameters block passed directly to the API + image: Optional[Any] # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object + parameters: Optional[Dict[str, Any]] # Provider-specific parameters block passed directly to the API model: Optional[str] seconds: Optional[str] size: Optional[str] diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index afa368ecdc3..f1b20618a74 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -35,9 +35,7 @@ def _add_base64_padding(value: str) -> str: return value -def encode_video_id_with_provider( - video_id: str, provider: str, model_id: Optional[str] = None -) -> str: +def encode_video_id_with_provider(video_id: str, provider: str, model_id: Optional[str] = None) -> str: """Encode provider and model_id into video_id using base64.""" if not provider or not video_id: return video_id @@ -51,13 +49,9 @@ def encode_video_id_with_provider( return video_id # ID is not encoded (even if it starts with video_), so encode it - assembled_id = str(SpecialEnums.LITELLM_MANAGED_VIDEO_COMPLETE_STR.value).format( - provider, model_id or "", video_id - ) + assembled_id = str(SpecialEnums.LITELLM_MANAGED_VIDEO_COMPLETE_STR.value).format(provider, model_id or "", video_id) - base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode( - "utf-8" - ) + base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode("utf-8") return f"{VIDEO_ID_PREFIX}{base64_encoded_id}" @@ -101,9 +95,7 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: model_id_part = parts[1] video_id_part = parts[2] - custom_llm_provider = custom_llm_provider_part.replace( - "litellm:custom_llm_provider:", "" - ) + custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "") model_id = model_id_part.replace("model_id:", "") decoded_video_id = video_id_part.replace("video_id:", "") @@ -127,9 +119,7 @@ def extract_original_video_id(encoded_video_id: str) -> str: return decoded.get("video_id", encoded_video_id) -def encode_character_id_with_provider( - character_id: str, provider: str, model_id: Optional[str] = None -) -> str: +def encode_character_id_with_provider(character_id: str, provider: str, model_id: Optional[str] = None) -> str: """Encode provider and model_id into character_id using base64.""" if not provider or not character_id: return character_id @@ -139,9 +129,7 @@ def encode_character_id_with_provider( return character_id assembled_id = CHARACTER_ID_TEMPLATE.format(provider, model_id or "", character_id) - base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode( - "utf-8" - ) + base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode("utf-8") return f"{CHARACTER_ID_PREFIX}{base64_encoded_id}" @@ -184,9 +172,7 @@ def decode_character_id_with_provider(encoded_character_id: str) -> DecodedChara model_id_part = parts[1] character_id_part = parts[2] - custom_llm_provider = custom_llm_provider_part.replace( - "litellm:custom_llm_provider:", "" - ) + custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "") model_id = model_id_part.replace("model_id:", "") decoded_character_id = character_id_part.replace("character_id:", "") @@ -196,9 +182,7 @@ def decode_character_id_with_provider(encoded_character_id: str) -> DecodedChara character_id=decoded_character_id, ) except Exception as e: - verbose_logger.debug( - f"Error decoding character_id '{encoded_character_id}': {e}" - ) + verbose_logger.debug(f"Error decoding character_id '{encoded_character_id}': {e}") return DecodedCharacterId( custom_llm_provider=None, model_id=None, diff --git a/litellm/utils.py b/litellm/utils.py index f11c42cbece..b9e76b6525b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -224,9 +224,7 @@ try: ): json_data = json.load(f) except (ImportError, AttributeError, TypeError): - with resources.open_text( - "litellm.litellm_core_utils.tokenizers", "anthropic_tokenizer.json" - ) as f: + with resources.open_text("litellm.litellm_core_utils.tokenizers", "anthropic_tokenizer.json") as f: json_data = json.load(f) # Convert to str (if necessary) @@ -509,9 +507,7 @@ def custom_llm_setup(): litellm._custom_providers.append(custom_llm["provider"]) -def _add_custom_logger_callback_to_specific_event( - callback: str, logging_event: Literal["success", "failure"] -) -> None: +def _add_custom_logger_callback_to_specific_event(callback: str, logging_event: Literal["success", "failure"]) -> None: """ Add a custom logger callback to the specific event """ @@ -533,44 +529,20 @@ def _add_custom_logger_callback_to_specific_event( ) if callback_class: - if ( - logging_event == "success" - and _custom_logger_class_exists_in_success_callbacks(callback_class) - is False - ): - litellm.logging_callback_manager.add_litellm_success_callback( - callback_class - ) - litellm.logging_callback_manager.add_litellm_async_success_callback( - callback_class - ) + if logging_event == "success" and _custom_logger_class_exists_in_success_callbacks(callback_class) is False: + litellm.logging_callback_manager.add_litellm_success_callback(callback_class) + litellm.logging_callback_manager.add_litellm_async_success_callback(callback_class) if callback in litellm.success_callback: - litellm.success_callback.remove( - callback - ) # remove the string from the callback list + litellm.success_callback.remove(callback) # remove the string from the callback list if callback in litellm._async_success_callback: - litellm._async_success_callback.remove( - callback - ) # remove the string from the callback list - elif ( - logging_event == "failure" - and _custom_logger_class_exists_in_failure_callbacks(callback_class) - is False - ): - litellm.logging_callback_manager.add_litellm_failure_callback( - callback_class - ) - litellm.logging_callback_manager.add_litellm_async_failure_callback( - callback_class - ) + litellm._async_success_callback.remove(callback) # remove the string from the callback list + elif logging_event == "failure" and _custom_logger_class_exists_in_failure_callbacks(callback_class) is False: + litellm.logging_callback_manager.add_litellm_failure_callback(callback_class) + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback_class) if callback in litellm.failure_callback: - litellm.failure_callback.remove( - callback - ) # remove the string from the callback list + litellm.failure_callback.remove(callback) # remove the string from the callback list if callback in litellm._async_failure_callback: - litellm._async_failure_callback.remove( - callback - ) # remove the string from the callback list + litellm._async_failure_callback.remove(callback) # remove the string from the callback list def _custom_logger_class_exists_in_success_callbacks( @@ -584,8 +556,7 @@ def _custom_logger_class_exists_in_success_callbacks( Prevents double adding a custom logger callback to the litellm callbacks """ return any( - isinstance(cb, type(callback_class)) - for cb in litellm.success_callback + litellm._async_success_callback + isinstance(cb, type(callback_class)) for cb in litellm.success_callback + litellm._async_success_callback ) @@ -600,8 +571,7 @@ def _custom_logger_class_exists_in_failure_callbacks( Prevents double adding a custom logger callback to the litellm callbacks """ return any( - isinstance(cb, type(callback_class)) - for cb in litellm.failure_callback + litellm._async_failure_callback + isinstance(cb, type(callback_class)) for cb in litellm.failure_callback + litellm._async_failure_callback ) @@ -682,9 +652,7 @@ def _remove_thought_signature_from_id(tool_call_id: str, separator: str) -> str: return tool_call_id -def _process_assistant_message_tool_calls( - msg_copy: dict, thought_signature_separator: str -) -> dict: +def _process_assistant_message_tool_calls(msg_copy: dict, thought_signature_separator: str) -> dict: """ Process assistant message to remove thought signatures from tool call IDs. """ @@ -707,9 +675,7 @@ def _process_assistant_message_tool_calls( # Remove thought signature from ID if present if isinstance(tc_dict.get("id"), str): if thought_signature_separator in tc_dict["id"]: - tc_dict["id"] = _remove_thought_signature_from_id( - tc_dict["id"], thought_signature_separator - ) + tc_dict["id"] = _remove_thought_signature_from_id(tc_dict["id"], thought_signature_separator) new_tool_calls.append(tc_dict) msg_copy["tool_calls"] = new_tool_calls @@ -730,9 +696,7 @@ def _process_tool_message_id(msg_copy: dict, thought_signature_separator: str) - return msg_copy -def _remove_thought_signatures_from_messages( - messages: List, thought_signature_separator: str -) -> List: +def _remove_thought_signatures_from_messages(messages: List, thought_signature_separator: str) -> List: """ Remove thought signatures from tool call IDs in all messages. """ @@ -750,9 +714,7 @@ def _remove_thought_signatures_from_messages( continue # Process assistant messages with tool_calls - msg_dict = _process_assistant_message_tool_calls( - msg_dict, thought_signature_separator - ) + msg_dict = _process_assistant_message_tool_calls(msg_dict, thought_signature_separator) # Process tool messages with tool_call_id msg_dict = _process_tool_message_id(msg_dict, thought_signature_separator) @@ -783,15 +745,11 @@ def function_setup( function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## LAZY LOAD COROUTINE CHECKER ## - get_coroutine_checker_fn = getattr( - sys.modules[__name__], "get_coroutine_checker" - ) + get_coroutine_checker_fn = getattr(sys.modules[__name__], "get_coroutine_checker") coroutine_checker = get_coroutine_checker_fn() ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = ( - kwargs.pop("callbacks", None) - ) + dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = kwargs.pop("callbacks", None) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -804,36 +762,23 @@ def function_setup( llm_router=None, # type: ignore ) if callback is None or any( - isinstance(cb, type(callback)) - for cb in litellm._async_success_callback + isinstance(cb, type(callback)) for cb in litellm._async_success_callback ): # don't double add a callback continue if callback not in litellm.input_callback: litellm.input_callback.append(callback) # type: ignore if callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback( - callback - ) # type: ignore + litellm.logging_callback_manager.add_litellm_success_callback(callback) # type: ignore if callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback( - callback - ) # type: ignore + litellm.logging_callback_manager.add_litellm_failure_callback(callback) # type: ignore if callback not in litellm._async_success_callback: - litellm.logging_callback_manager.add_litellm_async_success_callback( - callback - ) # type: ignore + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) # type: ignore if callback not in litellm._async_failure_callback: - litellm.logging_callback_manager.add_litellm_async_failure_callback( - callback - ) # type: ignore - print_verbose( - f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}" - ) + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) # type: ignore + print_verbose(f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}") if ( - len(litellm.input_callback) > 0 - or len(litellm.success_callback) > 0 - or len(litellm.failure_callback) > 0 + len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0 ) and len( callback_list # type: ignore ) == 0: # type: ignore @@ -861,21 +806,14 @@ def function_setup( removed_async_items = [] for index, callback in enumerate(litellm.success_callback): # type: ignore if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_success_callback( - callback - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) removed_async_items.append(index) elif callback == "dynamodb" or callback == "openmeter": # dynamo is an async callback, it's used for the proxy and needs to be async # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy - litellm.logging_callback_manager.add_litellm_async_success_callback( - callback - ) + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) removed_async_items.append(index) - elif ( - callback in litellm._known_custom_logger_compatible_callbacks - and isinstance(callback, str) - ): + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): _add_custom_logger_callback_to_specific_event(callback, "success") # Pop the async items from success_callback in reverse order to avoid index issues @@ -886,42 +824,23 @@ def function_setup( removed_async_items = [] for index, callback in enumerate(litellm.failure_callback): # type: ignore if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_failure_callback( - callback - ) + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) removed_async_items.append(index) - elif ( - callback in litellm._known_custom_logger_compatible_callbacks - and isinstance(callback, str) - ): + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): _add_custom_logger_callback_to_specific_event(callback, "failure") # Pop the async items from failure_callback in reverse order to avoid index issues for index in reversed(removed_async_items): litellm.failure_callback.pop(index) ### DYNAMIC CALLBACKS ### - dynamic_success_callbacks: Optional[ - List[Union[str, Callable, "CustomLogger"]] - ] = None - dynamic_async_success_callbacks: Optional[ - List[Union[str, Callable, "CustomLogger"]] - ] = None - dynamic_failure_callbacks: Optional[ - List[Union[str, Callable, "CustomLogger"]] - ] = None - dynamic_async_failure_callbacks: Optional[ - List[Union[str, Callable, "CustomLogger"]] - ] = None - if kwargs.get("success_callback", None) is not None and isinstance( - kwargs["success_callback"], list - ): + dynamic_success_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None + dynamic_async_success_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None + dynamic_failure_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None + dynamic_async_failure_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = None + if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list): removed_async_items = [] for index, callback in enumerate(kwargs["success_callback"]): - if ( - coroutine_checker.is_async_callable(callback) - or callback == "dynamodb" - or callback == "s3" - ): + if coroutine_checker.is_async_callable(callback) or callback == "dynamodb" or callback == "s3": if dynamic_async_success_callbacks is not None and isinstance( dynamic_async_success_callbacks, list ): @@ -933,9 +852,7 @@ def function_setup( for index in reversed(removed_async_items): kwargs["success_callback"].pop(index) dynamic_success_callbacks = kwargs.pop("success_callback") - if kwargs.get("failure_callback", None) is not None and isinstance( - kwargs["failure_callback"], list - ): + if kwargs.get("failure_callback", None) is not None and isinstance(kwargs["failure_callback"], list): dynamic_failure_callbacks = kwargs.pop("failure_callback") if add_breadcrumb: @@ -1020,9 +937,7 @@ def function_setup( # Only process if target is NOT a Gemini model if not _is_gemini_model(model, custom_llm_provider): - verbose_logger.debug( - "Removing thought signatures from tool call IDs for non-Gemini model" - ) + verbose_logger.debug("Removing thought signatures from tool call IDs for non-Gemini model") # Process messages to remove thought signatures processed_messages = _remove_thought_signatures_from_messages( @@ -1039,37 +954,18 @@ def function_setup( except Exception as e: # Log the error but don't fail the request - verbose_logger.warning( - f"Error removing thought signatures from tool call IDs: {str(e)}" - ) - elif ( - call_type == CallTypes.embedding.value - or call_type == CallTypes.aembedding.value - ): + verbose_logger.warning(f"Error removing thought signatures from tool call IDs: {str(e)}") + elif call_type == CallTypes.embedding.value or call_type == CallTypes.aembedding.value: messages = args[1] if len(args) > 1 else kwargs.get("input", None) - elif ( - call_type == CallTypes.image_generation.value - or call_type == CallTypes.aimage_generation.value - ): + elif call_type == CallTypes.image_generation.value or call_type == CallTypes.aimage_generation.value: messages = args[0] if len(args) > 0 else kwargs["prompt"] - elif ( - call_type == CallTypes.moderation.value - or call_type == CallTypes.amoderation.value - ): + elif call_type == CallTypes.moderation.value or call_type == CallTypes.amoderation.value: messages = args[1] if len(args) > 1 else kwargs["input"] - elif ( - call_type == CallTypes.atext_completion.value - or call_type == CallTypes.text_completion.value - ): + elif call_type == CallTypes.atext_completion.value or call_type == CallTypes.text_completion.value: messages = args[0] if len(args) > 0 else kwargs["prompt"] - elif ( - call_type == CallTypes.rerank.value or call_type == CallTypes.arerank.value - ): + elif call_type == CallTypes.rerank.value or call_type == CallTypes.arerank.value: messages = kwargs.get("query") - elif ( - call_type == CallTypes.atranscription.value - or call_type == CallTypes.transcription.value - ): + elif call_type == CallTypes.atranscription.value or call_type == CallTypes.transcription.value: _file_obj: FileTypes = args[1] if len(args) > 1 else kwargs["file"] # Lazy import audio_utils.utils only when needed for transcription calls audio_utils = _get_cached_audio_utils() @@ -1079,20 +975,12 @@ def function_setup( else: kwargs["metadata"] = {"file_checksum": file_checksum} messages = file_checksum - elif ( - call_type == CallTypes.aspeech.value or call_type == CallTypes.speech.value - ): + elif call_type == CallTypes.aspeech.value or call_type == CallTypes.speech.value: messages = kwargs.get("input", "speech") - elif ( - call_type == CallTypes.aresponses.value - or call_type == CallTypes.responses.value - ): + elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value: # Handle both 'input' (standard Responses API) and 'messages' (Cursor chat format) messages = ( - args[0] - if len(args) > 0 - else kwargs.get("input") - or kwargs.get("messages", "default-message-value") + args[0] if len(args) > 0 else kwargs.get("input") or kwargs.get("messages", "default-message-value") ) elif ( call_type == CallTypes.generate_content.value @@ -1119,16 +1007,11 @@ def function_setup( config=kwargs.get("config"), ) transformed_messages = transformed.get("messages", []) - messages = ( - get_last_user_message(transformed_messages) - or "default-message-value" - ) + messages = get_last_user_message(transformed_messages) or "default-message-value" else: messages = "default-message-value" except Exception as e: - verbose_logger.debug( - f"Error extracting messages from Google contents: {str(e)}" - ) + verbose_logger.debug(f"Error extracting messages from Google contents: {str(e)}") messages = "default-message-value" else: messages = "default-message-value" @@ -1138,9 +1021,7 @@ def function_setup( call_type=call_type, ): stream = True - get_litellm_logging_class = getattr( - sys.modules[__name__], "get_litellm_logging_class" - ) + get_litellm_logging_class = getattr(sys.modules[__name__], "get_litellm_logging_class") logging_obj = get_litellm_logging_class()( # Victim for object pool model=model, # type: ignore messages=messages, @@ -1162,9 +1043,7 @@ def function_setup( litellm_params: Dict[str, Any] = {"api_base": ""} if "metadata" in kwargs: litellm_params["metadata"] = kwargs["metadata"] - if "litellm_metadata" in kwargs and isinstance( - kwargs["litellm_metadata"], dict - ): + if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): litellm_params["litellm_metadata"] = kwargs["litellm_metadata"].copy() # For endpoints like /v1/messages that use "litellm_metadata" instead # of "metadata" (to avoid conflicting with provider API metadata fields), @@ -1182,9 +1061,7 @@ def function_setup( ) return logging_obj, kwargs except Exception as e: - verbose_logger.exception( - "litellm.utils.py::function_setup() - [Non-Blocking] Error in function_setup" - ) + verbose_logger.exception("litellm.utils.py::function_setup() - [Non-Blocking] Error in function_setup") raise e @@ -1207,9 +1084,7 @@ async def _client_async_logging_helper( from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=logging_obj.async_success_handler( - result=result, start_time=start_time, end_time=end_time - ) + async_coroutine=logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time) ) ################################################ @@ -1222,9 +1097,7 @@ async def _client_async_logging_helper( ) -def _get_wrapper_num_retries( - kwargs: Dict[str, Any], exception: Exception -) -> Tuple[Optional[int], Dict[str, Any]]: +def _get_wrapper_num_retries(kwargs: Dict[str, Any], exception: Exception) -> Tuple[Optional[int], Dict[str, Any]]: """ Get the number of retries from the kwargs and the retry policy. Used for the wrapper functions. @@ -1234,9 +1107,7 @@ def _get_wrapper_num_retries( if num_retries is None: num_retries = litellm.num_retries if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr( - sys.modules[__name__], "get_num_retries_from_retry_policy" - ) + get_num_retries_from_retry_policy = getattr(sys.modules[__name__], "get_num_retries_from_retry_policy") reset_retry_policy = getattr(sys.modules[__name__], "reset_retry_policy") retry_policy_num_retries = get_num_retries_from_retry_policy( exception=exception, @@ -1249,17 +1120,13 @@ def _get_wrapper_num_retries( return num_retries, kwargs -def _get_wrapper_timeout( - kwargs: Dict[str, Any], exception: Exception -) -> Optional[Union[float, int, httpx.Timeout]]: +def _get_wrapper_timeout(kwargs: Dict[str, Any], exception: Exception) -> Optional[Union[float, int, httpx.Timeout]]: """ Get the timeout from the kwargs Used for the wrapper functions. """ - timeout = cast( - Optional[Union[float, int, httpx.Timeout]], kwargs.get("timeout", None) - ) + timeout = cast(Optional[Union[float, int, httpx.Timeout]], kwargs.get("timeout", None)) return timeout @@ -1285,9 +1152,7 @@ async def async_pre_call_deployment_hook(kwargs: Dict[str, Any], call_type: str) CustomLogger = _get_cached_custom_logger() for callback in litellm.callbacks: if isinstance(callback, CustomLogger): - result = await callback.async_pre_call_deployment_hook( - modified_kwargs, typed_call_type - ) + result = await callback.async_pre_call_deployment_hook(modified_kwargs, typed_call_type) if result is not None: modified_kwargs = result @@ -1329,21 +1194,13 @@ def post_call_processing( pass else: call_type = original_function.__name__ - if ( - call_type == CallTypes.completion.value - or call_type == CallTypes.acompletion.value - ): + if call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value: is_coroutine = check_coroutine(original_response) if is_coroutine is True: pass else: - if ( - isinstance(original_response, ModelResponse) - and len(original_response.choices) > 0 - ): - model_response: Optional[str] = original_response.choices[ - 0 - ].message.content # type: ignore + if isinstance(original_response, ModelResponse) and len(original_response.choices) > 0: + model_response: Optional[str] = original_response.choices[0].message.content # type: ignore if model_response is not None: ### POST-CALL RULES ### rules_obj.post_call_rules(input=model_response, model=model) @@ -1364,8 +1221,7 @@ def post_call_processing( if ( optional_params is not None and "response_format" in optional_params - and optional_params["response_format"] - is not None + and optional_params["response_format"] is not None ): json_response_format: Optional[dict] = None if ( @@ -1373,29 +1229,18 @@ def post_call_processing( optional_params["response_format"], dict, ) - and optional_params["response_format"].get( - "json_schema" - ) - is not None + and optional_params["response_format"].get("json_schema") is not None ): - json_response_format = optional_params[ - "response_format" - ] + json_response_format = optional_params["response_format"] elif _parsing._completions.is_basemodel_type( optional_params["response_format"] # type: ignore ): - json_response_format = ( - type_to_response_format_param( - response_format=optional_params[ - "response_format" - ] - ) + json_response_format = type_to_response_format_param( + response_format=optional_params["response_format"] ) if json_response_format is not None: litellm.litellm_core_utils.json_validation_rule.validate_schema( - schema=json_response_format[ - "json_schema" - ]["schema"], + schema=json_response_format["json_schema"]["schema"], response=model_response, ) except TypeError: @@ -1405,28 +1250,18 @@ def post_call_processing( and "response_format" in optional_params and isinstance(optional_params["response_format"], dict) and "type" in optional_params["response_format"] - and optional_params["response_format"]["type"] - == "json_object" - and "response_schema" - in optional_params["response_format"] + and optional_params["response_format"]["type"] == "json_object" + and "response_schema" in optional_params["response_format"] and isinstance( - optional_params["response_format"][ - "response_schema" - ], + optional_params["response_format"]["response_schema"], dict, ) - and "enforce_validation" - in optional_params["response_format"] - and optional_params["response_format"][ - "enforce_validation" - ] - is True + and "enforce_validation" in optional_params["response_format"] + and optional_params["response_format"]["enforce_validation"] is True ): # schema given, json response expected, and validation enforced litellm.litellm_core_utils.json_validation_rule.validate_schema( - schema=optional_params["response_format"][ - "response_schema" - ], + schema=optional_params["response_format"]["response_schema"], response=model_response, ) @@ -1447,9 +1282,7 @@ def client(original_function): # [OPTIONAL] CHECK MAX RETRIES / REQUEST if litellm.num_retries_per_request is not None: # check if previous_models passed in as ['litellm_params']['metadata]['previous_models'] - previous_models = (kwargs.get("metadata") or {}).get( - "previous_models", None - ) + previous_models = (kwargs.get("metadata") or {}).get("previous_models", None) if previous_models is not None: if litellm.num_retries_per_request <= len(previous_models): raise Exception("Max retries per request hit!") @@ -1460,16 +1293,11 @@ def client(original_function): kwargs=kwargs, call_type=call_type, ): - if ( - "complete_response" in kwargs - and kwargs["complete_response"] is True - ): + if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) - return litellm.stream_chunk_builder( - chunks, messages=kwargs.get("messages", None) - ) + return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: return result @@ -1479,9 +1307,7 @@ def client(original_function): print_args_passed_to_litellm(original_function, args, kwargs) start_time = datetime.datetime.now() result = None - logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get( - "litellm_logging_obj", None - ) + logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get("litellm_logging_obj", None) # only set litellm_call_id if its not in kwargs if "litellm_call_id" not in kwargs: @@ -1491,14 +1317,10 @@ def client(original_function): try: if logging_obj is None: - logging_obj, kwargs = function_setup( - original_function.__name__, rules_obj, start_time, *args, **kwargs - ) + logging_obj, kwargs = function_setup(original_function.__name__, rules_obj, start_time, *args, **kwargs) # Type assertion: logging_obj is guaranteed to be non-None after function_setup - assert logging_obj is not None, ( - "logging_obj should not be None after function_setup" - ) + assert logging_obj is not None, "logging_obj should not be None after function_setup" ## LOAD CREDENTIALS load_credentials_from_list(kwargs) @@ -1522,9 +1344,7 @@ def client(original_function): # [OPTIONAL] CHECK MAX RETRIES / REQUEST if litellm.num_retries_per_request is not None: # check if previous_models passed in as ['litellm_params']['metadata]['previous_models'] - previous_models = (kwargs.get("metadata") or {}).get( - "previous_models", None - ) + previous_models = (kwargs.get("metadata") or {}).get("previous_models", None) if previous_models is not None: if litellm.num_retries_per_request <= len(previous_models): raise Exception("Max retries per request hit!") @@ -1537,10 +1357,7 @@ def client(original_function): if ( ( ( - ( - kwargs.get("caching", None) is None - and litellm.cache is not None - ) + (kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True ) and kwargs.get("cache", {}).get("no-cache", False) is not True @@ -1555,16 +1372,14 @@ def client(original_function): ): # allow users to control returning cached responses from the completion function # checking cache verbose_logger.debug("INSIDE CHECKING SYNC CACHE") - caching_handler_response: "CachingHandlerResponse" = ( - _llm_caching_handler._sync_get_cache( - model=model or "", - original_function=original_function, - logging_obj=logging_obj, - start_time=start_time, - call_type=call_type, - kwargs=kwargs, - args=args, - ) + caching_handler_response: "CachingHandlerResponse" = _llm_caching_handler._sync_get_cache( + model=model or "", + original_function=original_function, + logging_obj=logging_obj, + start_time=start_time, + call_type=call_type, + kwargs=kwargs, + args=args, ) if caching_handler_response.cached_result is not None: @@ -1575,8 +1390,7 @@ def client(original_function): if ( kwargs.get("max_tokens", None) is not None and model is not None - and litellm.modify_params - is True # user is okay with params being modified + and litellm.modify_params is True # user is okay with params being modified and ( call_type == CallTypes.acompletion.value or call_type == CallTypes.completion.value @@ -1611,21 +1425,14 @@ def client(original_function): kwargs=kwargs, call_type=call_type, ): - if ( - "complete_response" in kwargs - and kwargs["complete_response"] is True - ): + if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) - return litellm.stream_chunk_builder( - chunks, messages=kwargs.get("messages", None) - ) + return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: # RETURN RESULT - update_response_metadata = getattr( - sys.modules[__name__], "update_response_metadata" - ) + update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") update_response_metadata( result=result, logging_obj=logging_obj, @@ -1678,9 +1485,7 @@ def client(original_function): end_time, ) # RETURN RESULT - update_response_metadata = getattr( - sys.modules[__name__], "update_response_metadata" - ) + update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") update_response_metadata( result=result, logging_obj=logging_obj, @@ -1693,29 +1498,19 @@ def client(original_function): except Exception as e: call_type = original_function.__name__ if call_type == CallTypes.completion.value: - num_retries = ( - kwargs.get("num_retries", None) or litellm.num_retries or None - ) + num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None if kwargs.get("retry_policy", None): get_num_retries_from_retry_policy = getattr( sys.modules[__name__], "get_num_retries_from_retry_policy" ) - reset_retry_policy = getattr( - sys.modules[__name__], "reset_retry_policy" - ) + reset_retry_policy = getattr(sys.modules[__name__], "reset_retry_policy") num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops - litellm.num_retries = ( - None # set retries to None to prevent infinite loops - ) - context_window_fallback_dict = kwargs.get( - "context_window_fallback_dict", {} - ) + kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops + litellm.num_retries = None # set retries to None to prevent infinite loops + context_window_fallback_dict = kwargs.get("context_window_fallback_dict", {}) _is_litellm_router_call = "model_group" in ( kwargs.get("metadata") or {} @@ -1742,26 +1537,18 @@ def client(original_function): kwargs["model"] = context_window_fallback_dict[model] return original_function(*args, **kwargs) elif call_type == CallTypes.responses.value: - num_retries = ( - kwargs.get("num_retries", None) or litellm.num_retries or None - ) + num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None if kwargs.get("retry_policy", None): get_num_retries_from_retry_policy = getattr( sys.modules[__name__], "get_num_retries_from_retry_policy" ) - reset_retry_policy = getattr( - sys.modules[__name__], "reset_retry_policy" - ) + reset_retry_policy = getattr(sys.modules[__name__], "reset_retry_policy") num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops - litellm.num_retries = ( - None # set retries to None to prevent infinite loops - ) + kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops + litellm.num_retries = None # set retries to None to prevent infinite loops _is_litellm_router_call = "model_group" in ( kwargs.get("metadata") or {} @@ -1791,12 +1578,8 @@ def client(original_function): print_args_passed_to_litellm(original_function, args, kwargs) start_time = datetime.datetime.now() result = None - _update_response_metadata = getattr( - sys.modules[__name__], "update_response_metadata" - ) - logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get( - "litellm_logging_obj", None - ) + _update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") + logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get("litellm_logging_obj", None) LLMCachingHandler = _get_cached_llm_caching_handler() _llm_caching_handler: "LLMCachingHandler" = LLMCachingHandler( original_function=original_function, @@ -1815,14 +1598,10 @@ def client(original_function): try: if logging_obj is None: - logging_obj, kwargs = function_setup( - original_function.__name__, rules_obj, start_time, *args, **kwargs - ) + logging_obj, kwargs = function_setup(original_function.__name__, rules_obj, start_time, *args, **kwargs) # Type assertion: logging_obj is guaranteed to be non-None after function_setup - assert logging_obj is not None, ( - "logging_obj should not be None after function_setup" - ) + assert logging_obj is not None, "logging_obj should not be None after function_setup" modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type) if modified_kwargs is not None: @@ -1850,23 +1629,20 @@ def client(original_function): print_verbose( f"ASYNC kwargs[caching]: {kwargs.get('caching', False)}; litellm.cache: {litellm.cache}; kwargs.get('cache'): {kwargs.get('cache', None)}" ) - _caching_handler_response: "Optional[CachingHandlerResponse]" = ( - await _llm_caching_handler._async_get_cache( - model=model or "", - original_function=original_function, - logging_obj=logging_obj, - start_time=start_time, - call_type=call_type, - kwargs=kwargs, - args=args, - ) + _caching_handler_response: "Optional[CachingHandlerResponse]" = await _llm_caching_handler._async_get_cache( + model=model or "", + original_function=original_function, + logging_obj=logging_obj, + start_time=start_time, + call_type=call_type, + kwargs=kwargs, + args=args, ) if _caching_handler_response is not None: if ( _caching_handler_response.cached_result is not None - and _caching_handler_response.final_embedding_cached_response - is None + and _caching_handler_response.final_embedding_cached_response is None ): return _caching_handler_response.cached_result @@ -1877,8 +1653,7 @@ def client(original_function): if ( kwargs.get("max_tokens", None) is not None and model is not None - and litellm.modify_params - is True # user is okay with params being modified + and litellm.modify_params is True # user is okay with params being modified and ( call_type == CallTypes.acompletion.value or call_type == CallTypes.completion.value @@ -1915,16 +1690,11 @@ def client(original_function): kwargs=kwargs, call_type=call_type, ): - if ( - "complete_response" in kwargs - and kwargs["complete_response"] is True - ): + if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) - return litellm.stream_chunk_builder( - chunks, messages=kwargs.get("messages", None) - ) + return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: _update_response_metadata( result=result, @@ -2002,8 +1772,7 @@ def client(original_function): if ( isinstance(result, EmbeddingResponse) and _caching_handler_response is not None - and _caching_handler_response.final_embedding_cached_response - is not None + and _caching_handler_response.final_embedding_cached_response is not None ): return _llm_caching_handler._combine_cached_embedding_response_with_api_result( _caching_handler_response=_caching_handler_response, @@ -2033,18 +1802,14 @@ def client(original_function): except Exception as e: raise e try: - await logging_obj.async_failure_handler( - e, traceback_exception, start_time, end_time - ) + await logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time) except Exception as e: raise e call_type = original_function.__name__ num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) if call_type == CallTypes.acompletion.value: - context_window_fallback_dict = kwargs.get( - "context_window_fallback_dict", {} - ) + context_window_fallback_dict = kwargs.get("context_window_fallback_dict", {}) _is_litellm_router_call = "model_group" in ( kwargs.get("metadata") or {} @@ -2054,14 +1819,10 @@ def client(original_function): num_retries and not _is_litellm_router_call ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: - litellm.num_retries = ( - None # set retries to None to prevent infinite loops - ) + litellm.num_retries = None # set retries to None to prevent infinite loops kwargs["num_retries"] = num_retries kwargs["original_function"] = original_function - if isinstance( - e, openai.RateLimitError - ): # rate limiting specific error + if isinstance(e, openai.RateLimitError): # rate limiting specific error kwargs["retry_strategy"] = "exponential_backoff_retry" elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" @@ -2088,14 +1849,10 @@ def client(original_function): num_retries and not _is_litellm_router_call ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: - litellm.num_retries = ( - None # set retries to None to prevent infinite loops - ) + litellm.num_retries = None # set retries to None to prevent infinite loops kwargs["num_retries"] = num_retries kwargs["original_function"] = original_function - if isinstance( - e, openai.RateLimitError - ): # rate limiting specific error + if isinstance(e, openai.RateLimitError): # rate limiting specific error kwargs["retry_strategy"] = "exponential_backoff_retry" elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" @@ -2103,9 +1860,7 @@ def client(original_function): except Exception: pass - setattr( - e, "num_retries", num_retries - ) ## IMPORTANT: returns the deployment's num_retries to the router + setattr(e, "num_retries", num_retries) ## IMPORTANT: returns the deployment's num_retries to the router timeout = _get_wrapper_timeout(kwargs=kwargs, exception=e) setattr(e, "timeout", timeout) @@ -2178,9 +1933,7 @@ def _is_streaming_request( return call_type in _STREAMING_CALL_TYPES -def _select_tokenizer( - model: str, custom_tokenizer: Optional[CustomHuggingfaceTokenizer] = None -): +def _select_tokenizer(model: str, custom_tokenizer: Optional[CustomHuggingfaceTokenizer] = None): if custom_tokenizer is not None: _tokenizer = create_pretrained_tokenizer( identifier=custom_tokenizer["identifier"], @@ -2214,9 +1967,7 @@ def _return_openai_tokenizer(model: str) -> SelectTokenizerResponse: def _return_huggingface_tokenizer(model: str) -> Optional[SelectTokenizerResponse]: if model in litellm.cohere_models and "command-r" in model: # cohere - cohere_tokenizer = Tokenizer.from_pretrained( - "Xenova/c4ai-command-r-v01-tokenizer" - ) + cohere_tokenizer = Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer") return {"type": "huggingface_tokenizer", "tokenizer": cohere_tokenizer} # anthropic elif model in litellm.anthropic_models and "claude-3" not in model: @@ -2275,38 +2026,28 @@ def decode( tokenizer_json = custom_tokenizer or _select_tokenizer(model=model) if tokenizer_json["type"] == "huggingface_tokenizer": if skip_special_tokens: - tokens = _strip_huggingface_special_token_ids( - tokenizer_json["tokenizer"], tokens - ) - dec = tokenizer_json["tokenizer"].decode( - tokens, skip_special_tokens=skip_special_tokens - ) + tokens = _strip_huggingface_special_token_ids(tokenizer_json["tokenizer"], tokens) + dec = tokenizer_json["tokenizer"].decode(tokens, skip_special_tokens=skip_special_tokens) return dec dec = tokenizer_json["tokenizer"].decode(tokens) return dec -def _strip_huggingface_special_token_ids( - tokenizer: Tokenizer, tokens: List[int] -) -> List[int]: +def _strip_huggingface_special_token_ids(tokenizer: Tokenizer, tokens: List[int]) -> List[int]: try: added_tokens_decoder = tokenizer.get_added_tokens_decoder() except Exception: return tokens special_token_ids = { - token_id - for token_id, added_token in added_tokens_decoder.items() - if getattr(added_token, "special", False) + token_id for token_id, added_token in added_tokens_decoder.items() if getattr(added_token, "special", False) } if not special_token_ids: return tokens return [token for token in tokens if token not in special_token_ids] -def create_pretrained_tokenizer( - identifier: str, revision="main", auth_token: Optional[str] = None -): +def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: Optional[str] = None): """ Creates a tokenizer from an existing file on a HuggingFace repository to be used with `token_counter`. @@ -2326,9 +2067,7 @@ def create_pretrained_tokenizer( auth_token=auth_token, # type: ignore ) except Exception as e: - verbose_logger.error( - f"Error creating pretrained tokenizer: {e}. Defaulting to version without 'auth_token'." - ) + verbose_logger.error(f"Error creating pretrained tokenizer: {e}. Defaulting to version without 'auth_token'.") tokenizer = Tokenizer.from_pretrained(identifier, revision=revision) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} @@ -2481,9 +2220,7 @@ def supports_native_streaming(model: str, custom_llm_provider: Optional[str]) -> model=model, custom_llm_provider=custom_llm_provider ) - model_info = _get_model_info_helper( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) supports_native_streaming = model_info.get("supports_native_streaming", True) if supports_native_streaming is None: supports_native_streaming = True @@ -2495,9 +2232,7 @@ def supports_native_streaming(model: str, custom_llm_provider: Optional[str]) -> return False -def supports_response_schema( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_response_schema(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model + provider supports 'response_schema' as a param. @@ -2513,9 +2248,7 @@ def supports_response_schema( ## GET LLM PROVIDER ## try: get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") - model, custom_llm_provider, _, _ = get_llm_provider( - model=model, custom_llm_provider=custom_llm_provider - ) + model, custom_llm_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) except Exception as e: verbose_logger.debug( f"Model not found or error in checking response schema support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {str(e)}" @@ -2540,9 +2273,7 @@ def supports_response_schema( ) -def supports_parallel_function_calling( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_parallel_function_calling(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model supports parallel tool calls and return a boolean value. """ @@ -2553,9 +2284,7 @@ def supports_parallel_function_calling( ) -def supports_function_calling( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_function_calling(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model supports function calling and return a boolean value. @@ -2580,9 +2309,7 @@ def supports_tool_choice(model: str, custom_llm_provider: Optional[str] = None) """ Check if the given model supports `tool_choice` and return a boolean value. """ - return _supports_factory( - model=model, custom_llm_provider=custom_llm_provider, key="supports_tool_choice" - ) + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_tool_choice") def _supports_provider_info_factory( @@ -2592,9 +2319,7 @@ def _supports_provider_info_factory( Check if the given model supports a provider specific model info and return a boolean value. """ - provider_info = get_provider_info( - model=model, custom_llm_provider=custom_llm_provider - ) + provider_info = get_provider_info(model=model, custom_llm_provider=custom_llm_provider) if provider_info is not None and provider_info.get(key, False) is True: return True @@ -2620,9 +2345,7 @@ def _supports_factory(model: str, custom_llm_provider: Optional[str], key: str) model=model, custom_llm_provider=custom_llm_provider ) - model_info = _get_model_info_helper( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) if model_info.get(key, False) is True: return True @@ -2637,9 +2360,7 @@ def _supports_factory(model: str, custom_llm_provider: Optional[str], key: str) if bare_entry.get(key, False) is True: return True - supported_by_provider = _supports_provider_info_factory( - model, custom_llm_provider, key - ) + supported_by_provider = _supports_provider_info_factory(model, custom_llm_provider, key) if supported_by_provider is not None: return supported_by_provider @@ -2649,18 +2370,14 @@ def _supports_factory(model: str, custom_llm_provider: Optional[str], key: str) f"Model not found or error in checking {key} support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {str(e)}" ) - supported_by_provider = _supports_provider_info_factory( - model, custom_llm_provider, key - ) + supported_by_provider = _supports_provider_info_factory(model, custom_llm_provider, key) if supported_by_provider is not None: return supported_by_provider return False -def _is_explicitly_disabled_factory( - model: str, custom_llm_provider: Optional[str], key: str -) -> bool: +def _is_explicitly_disabled_factory(model: str, custom_llm_provider: Optional[str], key: str) -> bool: """Return True only when the model map explicitly sets *key* to ``False``. This is the opt-out mirror of :func:`_supports_factory`. Where @@ -2677,9 +2394,7 @@ def _is_explicitly_disabled_factory( model, custom_llm_provider, _, _ = litellm.get_llm_provider( model=model, custom_llm_provider=custom_llm_provider ) - model_info = _get_model_info_helper( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) val = model_info.get(key) if val is False: return True @@ -2701,30 +2416,20 @@ def _is_explicitly_disabled_factory( def supports_audio_input(model: str, custom_llm_provider: Optional[str] = None) -> bool: """Check if a given model supports audio input in a chat completion call""" - return _supports_factory( - model=model, custom_llm_provider=custom_llm_provider, key="supports_audio_input" - ) + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_audio_input") def supports_pdf_input(model: str, custom_llm_provider: Optional[str] = None) -> bool: """Check if a given model supports pdf input in a chat completion call""" - return _supports_factory( - model=model, custom_llm_provider=custom_llm_provider, key="supports_pdf_input" - ) + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_pdf_input") -def supports_audio_output( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_audio_output(model: str, custom_llm_provider: Optional[str] = None) -> bool: """Check if a given model supports audio output in a chat completion call""" - return _supports_factory( - model=model, custom_llm_provider=custom_llm_provider, key="supports_audio_input" - ) + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_audio_input") -def supports_prompt_caching( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_prompt_caching(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model supports prompt caching and return a boolean value. @@ -2745,9 +2450,7 @@ def supports_prompt_caching( ) -def supports_computer_use( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_computer_use(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model supports computer use and return a boolean value. @@ -2790,14 +2493,10 @@ def supports_reasoning(model: str, custom_llm_provider: Optional[str] = None) -> """ Check if the given model supports reasoning and return a boolean value. """ - return _supports_factory( - model=model, custom_llm_provider=custom_llm_provider, key="supports_reasoning" - ) + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_reasoning") -def supports_native_structured_output( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_native_structured_output(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model supports native structured outputs and return a boolean value. """ @@ -2808,9 +2507,7 @@ def supports_native_structured_output( ) -def get_supported_regions( - model: str, custom_llm_provider: Optional[str] = None -) -> Optional[List[str]]: +def get_supported_regions(model: str, custom_llm_provider: Optional[str] = None) -> Optional[List[str]]: """ Get a list of supported regions for a given model and provider. @@ -2823,9 +2520,7 @@ def get_supported_regions( model=model, custom_llm_provider=custom_llm_provider ) - model_info = _get_model_info_helper( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info = _get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) # Get the key used in model_cost to look up supported_regions # since ModelInfoBase doesn't include this field @@ -2852,9 +2547,7 @@ def get_supported_regions( return None -def supports_embedding_image_input( - model: str, custom_llm_provider: Optional[str] = None -) -> bool: +def supports_embedding_image_input(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model supports embedding image input and return a boolean value. """ @@ -2921,9 +2614,7 @@ _CACHE_PRICING_FIELDS = ( ) -def _resolve_builtin_model_cost_entry( - key: str, provider: str -) -> Optional[Dict[str, Any]]: +def _resolve_builtin_model_cost_entry(key: str, provider: str) -> Optional[Dict[str, Any]]: """Best-effort lookup of a built-in ``model_cost`` entry for a custom key whose shape ``get_model_info`` cannot resolve (double provider prefixes like ``bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region aliases). @@ -3002,15 +2693,10 @@ def register_model(model_cost: Union[str, dict]): except Exception: existing_model = {} model_cost_key = key - builtin_entry = _resolve_builtin_model_cost_entry( - key=_key_str, provider=provider - ) + builtin_entry = _resolve_builtin_model_cost_entry(key=_key_str, provider=provider) if builtin_entry is not None: for field in _CACHE_PRICING_FIELDS: - if ( - value.get(field) is None - and builtin_entry.get(field) is not None - ): + if value.get(field) is None and builtin_entry.get(field) is not None: existing_model[field] = builtin_entry[field] elif ( value.get("cache_creation_input_token_cost") is None @@ -3051,9 +2737,7 @@ def register_model(model_cost: Union[str, dict]): # Invalidate case-insensitive lookup map since model_cost was modified _invalidate_model_cost_lowercase_map() - verbose_logger.debug( - f"added/updated model={model_cost_key} in litellm.model_cost: {model_cost_key}" - ) + verbose_logger.debug(f"added/updated model={model_cost_key} in litellm.model_cost: {model_cost_key}") # add new model names to provider lists if value.get("litellm_provider") == "openai": if key not in litellm.open_ai_chat_completion_models: @@ -3105,26 +2789,19 @@ def register_model(model_cost: Union[str, dict]): def _should_drop_param(k, additional_drop_params) -> bool: - if ( - additional_drop_params is not None - and isinstance(additional_drop_params, list) - and k in additional_drop_params - ): + if additional_drop_params is not None and isinstance(additional_drop_params, list) and k in additional_drop_params: return True # allow user to drop specific params for a model - e.g. vllm - logit bias return False -def _get_non_default_params( - passed_params: dict, default_params: dict, additional_drop_params: Optional[list] -) -> dict: +def _get_non_default_params(passed_params: dict, default_params: dict, additional_drop_params: Optional[list]) -> dict: non_default_params = {} for k, v in passed_params.items(): if ( k in default_params and v != default_params[k] - and _should_drop_param(k=k, additional_drop_params=additional_drop_params) - is False + and _should_drop_param(k=k, additional_drop_params=additional_drop_params) is False ): non_default_params[k] = v @@ -3162,11 +2839,7 @@ def get_optional_params_transcription( "timestamp_granularities": None, } - non_default_params = { - k: v - for k, v in passed_params.items() - if (k in default_params and v != default_params[k]) - } + non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])} optional_params = {} ## raise exception if non-default value passed for non-openai/azure embedding calls @@ -3175,9 +2848,8 @@ def get_optional_params_transcription( keys = list(non_default_params.keys()) for k in keys: if ( - (drop_params is True or litellm.drop_params is True) - and k not in supported_params - ): # drop the unsupported non-default values + drop_params is True or litellm.drop_params is True + ) and k not in supported_params: # drop the unsupported non-default values non_default_params.pop(k, None) elif k not in supported_params: raise UnsupportedParamsError( @@ -3239,9 +2911,7 @@ def _map_openai_size_to_vertex_ai_aspect_ratio(size: Optional[str]) -> str: "1792x1024": "16:9", # Landscape "1024x1792": "9:16", # Portrait } - return size_to_aspect_ratio.get( - size, "1:1" - ) # Default to square if size not recognized + return size_to_aspect_ratio.get(size, "1:1") # Default to square if size not recognized def get_optional_params_image_gen( @@ -3275,9 +2945,7 @@ def get_optional_params_image_gen( elif k == "hf_model_name" and custom_llm_provider != "sagemaker": continue elif ( - k.startswith("vertex_") - and custom_llm_provider != "vertex_ai" - and custom_llm_provider != "vertex_ai_beta" + k.startswith("vertex_") and custom_llm_provider != "vertex_ai" and custom_llm_provider != "vertex_ai_beta" ): # allow dynamically setting vertex ai init logic continue passed_params[k] = v @@ -3307,9 +2975,8 @@ def get_optional_params_image_gen( keys = list(non_default_params.keys()) for k in keys: if ( - (litellm.drop_params is True or drop_params is True) - and k not in supported_params - ): # drop the unsupported non-default values + litellm.drop_params is True or drop_params is True + ) and k not in supported_params: # drop the unsupported non-default values non_default_params.pop(k, None) passed_params.pop(k, None) elif k not in supported_params: @@ -3320,9 +2987,7 @@ def get_optional_params_image_gen( return non_default_params if provider_config is not None: - supported_params = provider_config.get_supported_openai_params( - model=model or "" - ) + supported_params = provider_config.get_supported_openai_params(model=model or "") _check_valid_arg(supported_params=supported_params) optional_params = provider_config.map_openai_params( non_default_params=non_default_params, @@ -3340,9 +3005,7 @@ def get_optional_params_image_gen( config_class = litellm.BedrockImageGeneration.get_config_class(model=model) supported_params = config_class.get_supported_openai_params(model=model) _check_valid_arg(supported_params=supported_params) - optional_params = config_class.map_openai_params( - non_default_params=non_default_params, optional_params={} - ) + optional_params = config_class.map_openai_params(non_default_params=non_default_params, optional_params={}) elif custom_llm_provider == "vertex_ai": supported_params = ["n", "size"] """ @@ -3354,15 +3017,11 @@ def get_optional_params_image_gen( # Map OpenAI size parameter to Vertex AI aspectRatio if size is not None: - optional_params["aspectRatio"] = _map_openai_size_to_vertex_ai_aspect_ratio( - size - ) + optional_params["aspectRatio"] = _map_openai_size_to_vertex_ai_aspect_ratio(size) openai_params: list[str] = list(default_params.keys()) if provider_config is not None: - supported_params = provider_config.get_supported_openai_params( - model=model or "" - ) + supported_params = provider_config.get_supported_openai_params(model=model or "") openai_params = list(supported_params) optional_params = add_provider_specific_params_to_optional_params( @@ -3374,9 +3033,7 @@ def get_optional_params_image_gen( ) # remove keys with None or empty dict/list values to avoid sending empty payloads optional_params = { - k: v - for k, v in optional_params.items() - if v is not None and (not isinstance(v, (dict, list)) or len(v) > 0) + k: v for k, v in optional_params.items() if v is not None and (not isinstance(v, (dict, list)) or len(v) > 0) } return optional_params @@ -3394,9 +3051,7 @@ def get_optional_params_embeddings( **kwargs, ): # Lazy load get_supported_openai_params - get_supported_openai_params = getattr( - sys.modules[__name__], "get_supported_openai_params" - ) + get_supported_openai_params = getattr(sys.modules[__name__], "get_supported_openai_params") # retrieve all parameters passed to the function passed_params = locals() @@ -3417,9 +3072,7 @@ def get_optional_params_embeddings( if k not in supported_params: unsupported_params[k] = non_default_params[k] if unsupported_params: - if litellm.drop_params is True or ( - drop_params is not None and drop_params is True - ): + if litellm.drop_params is True or (drop_params is not None and drop_params is True): pass else: raise UnsupportedParamsError( @@ -3427,32 +3080,25 @@ def get_optional_params_embeddings( message=f"{custom_llm_provider} does not support parameters: {unsupported_params}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n", ) - non_default_params = ( - PreProcessNonDefaultParams.embedding_pre_process_non_default_params( - passed_params=passed_params, - special_params=special_params, - custom_llm_provider=custom_llm_provider, - additional_drop_params=additional_drop_params, - model=model, - ) + non_default_params = PreProcessNonDefaultParams.embedding_pre_process_non_default_params( + passed_params=passed_params, + special_params=special_params, + custom_llm_provider=custom_llm_provider, + additional_drop_params=additional_drop_params, + model=model, ) provider_config: Optional[BaseEmbeddingConfig] = None optional_params = {} - if ( - custom_llm_provider is not None - and custom_llm_provider in LlmProviders._member_map_.values() - ): + if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_embedding_config( model=model, provider=LlmProviders(custom_llm_provider), ) if provider_config is not None: - supported_params: Optional[list] = provider_config.get_supported_openai_params( - model=model - ) + supported_params: Optional[list] = provider_config.get_supported_openai_params(model=model) _check_valid_arg(supported_params=supported_params) optional_params = provider_config.map_openai_params( non_default_params=non_default_params, @@ -3469,11 +3115,7 @@ def get_optional_params_embeddings( for param in supported_params: if param in OPENAI_EMBEDDING_PARAMS: continue - if ( - param in passed_params - and passed_params[param] is not None - and param not in optional_params - ): + if param in passed_params and passed_params[param] is not None and param not in optional_params: optional_params[param] = passed_params[param] ## raise exception if non-default value passed for non-openai/azure embedding calls elif custom_llm_provider == "openai": @@ -3543,9 +3185,7 @@ def get_optional_params_embeddings( non_default_params=non_default_params, optional_params={}, kwargs=kwargs ) elif custom_llm_provider == "lm_studio": - supported_params = ( - litellm.LmStudioEmbeddingConfig().get_supported_openai_params() - ) + supported_params = litellm.LmStudioEmbeddingConfig().get_supported_openai_params() _check_valid_arg(supported_params=supported_params) optional_params = litellm.LmStudioEmbeddingConfig().map_openai_params( non_default_params=non_default_params, optional_params={} @@ -3572,9 +3212,7 @@ def get_optional_params_embeddings( supported_params = object.get_supported_openai_params() _check_valid_arg(supported_params=supported_params) - optional_params = object.map_openai_params( - non_default_params=non_default_params, optional_params={} - ) + optional_params = object.map_openai_params(non_default_params=non_default_params, optional_params={}) elif custom_llm_provider == "mistral": supported_params = get_supported_openai_params( model=model, @@ -3606,22 +3244,18 @@ def get_optional_params_embeddings( ) _check_valid_arg(supported_params=supported_params) if litellm.VoyageContextualEmbeddingConfig.is_contextualized_embeddings(model): - optional_params = ( - litellm.VoyageContextualEmbeddingConfig().map_openai_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - drop_params=drop_params if drop_params is not None else False, - ) + optional_params = litellm.VoyageContextualEmbeddingConfig().map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=drop_params if drop_params is not None else False, ) elif litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model): - optional_params = ( - litellm.VoyageMultimodalEmbeddingConfig().map_openai_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - drop_params=drop_params if drop_params is not None else False, - ) + optional_params = litellm.VoyageMultimodalEmbeddingConfig().map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=drop_params if drop_params is not None else False, ) else: optional_params = litellm.VoyageEmbeddingConfig().map_openai_params( @@ -3703,9 +3337,7 @@ def get_optional_params_embeddings( if "dimensions" in non_default_params: optional_params["dimensions"] = non_default_params.pop("dimensions") if len(non_default_params.keys()) > 0: - if ( - litellm.drop_params is True or drop_params is True - ): # drop the unsupported non-default values + if litellm.drop_params is True or drop_params is True: # drop the unsupported non-default values keys = list(non_default_params.keys()) for k in keys: non_default_params.pop(k, None) @@ -3720,9 +3352,7 @@ def get_optional_params_embeddings( and custom_llm_provider not in litellm.openai_compatible_providers ): if len(non_default_params.keys()) > 0: - if ( - litellm.drop_params is True or drop_params is True - ): # drop the unsupported non-default values + if litellm.drop_params is True or drop_params is True: # drop the unsupported non-default values keys = list(non_default_params.keys()) for k in keys: non_default_params.pop(k, None) @@ -3827,9 +3457,7 @@ def _remove_json_schema_refs(schema, max_depth=10): return schema -def _remove_unsupported_params( - non_default_params: dict, supported_openai_params: Optional[List[str]] -) -> dict: +def _remove_unsupported_params(non_default_params: dict, supported_openai_params: Optional[List[str]]) -> dict: """ Remove unsupported params from non_default_params """ @@ -3863,9 +3491,7 @@ def filter_out_litellm_params(kwargs: dict) -> dict: >>> # filtered = {"query": "test"} """ - return { - key: value for key, value in kwargs.items() if key not in all_litellm_params - } + return {key: value for key, value in kwargs.items() if key not in all_litellm_params} class PreProcessNonDefaultParams: @@ -3884,8 +3510,7 @@ class PreProcessNonDefaultParams: # bedrock-mantle configs), never as a request body field continue if k.startswith("aws_") and ( - custom_llm_provider != "bedrock" - and not custom_llm_provider.startswith("sagemaker") + custom_llm_provider != "bedrock" and not custom_llm_provider.startswith("sagemaker") ): # allow dynamically setting boto3 init logic continue elif k == "hf_model_name" and custom_llm_provider != "sagemaker": @@ -3912,10 +3537,7 @@ class PreProcessNonDefaultParams: and k not in additional_endpoint_specific_params and k in default_param_values and v != default_param_values[k] - and _should_drop_param( - k=k, additional_drop_params=additional_drop_params - ) - is False + and _should_drop_param(k=k, additional_drop_params=additional_drop_params) is False ) } @@ -3931,15 +3553,13 @@ class PreProcessNonDefaultParams: remove_sensitive_keys: bool = False, add_provider_specific_params: bool = False, ) -> dict: - non_default_params = ( - PreProcessNonDefaultParams.base_pre_process_non_default_params( - passed_params=passed_params, - special_params=special_params, - custom_llm_provider=custom_llm_provider, - additional_drop_params=additional_drop_params, - default_param_values={k: None for k in OPENAI_EMBEDDING_PARAMS}, - additional_endpoint_specific_params=["input"], - ) + non_default_params = PreProcessNonDefaultParams.base_pre_process_non_default_params( + passed_params=passed_params, + special_params=special_params, + custom_llm_provider=custom_llm_provider, + additional_drop_params=additional_drop_params, + default_param_values={k: None for k in OPENAI_EMBEDDING_PARAMS}, + additional_endpoint_specific_params=["input"], ) return non_default_params @@ -3971,10 +3591,8 @@ def pre_process_non_default_params( if "response_format" in non_default_params: if provider_config is not None: - non_default_params["response_format"] = ( - provider_config.get_json_schema_from_pydantic_object( - response_format=non_default_params["response_format"] - ) + non_default_params["response_format"] = provider_config.get_json_schema_from_pydantic_object( + response_format=non_default_params["response_format"] ) else: non_default_params["response_format"] = type_to_response_format_param( @@ -3990,10 +3608,7 @@ def pre_process_non_default_params( parameters = tool_function.get("parameters", None) if parameters is not None: new_parameters = copy.deepcopy(parameters) - if ( - "additionalProperties" in new_parameters - and new_parameters["additionalProperties"] is False - ): + if "additionalProperties" in new_parameters and new_parameters["additionalProperties"] is False: new_parameters.pop("additionalProperties", None) tool_function["parameters"] = new_parameters @@ -4025,9 +3640,7 @@ def remove_sensitive_keys_from_dict(d: dict) -> dict: return d -def pre_process_optional_params( - passed_params: dict, non_default_params: dict, custom_llm_provider: str -) -> dict: +def pre_process_optional_params(passed_params: dict, non_default_params: dict, custom_llm_provider: str) -> dict: """For .completion(), preprocess optional params""" optional_params: Dict = {} @@ -4042,15 +3655,10 @@ def pre_process_optional_params( non_default_params=passed_params, optional_params=optional_params ) elif custom_llm_provider == "bedrock": - optional_params = ( - litellm.AmazonBedrockGlobalConfig().map_special_auth_params( - non_default_params=passed_params, optional_params=optional_params - ) + optional_params = litellm.AmazonBedrockGlobalConfig().map_special_auth_params( + non_default_params=passed_params, optional_params=optional_params ) - elif ( - custom_llm_provider == "vertex_ai" - or custom_llm_provider == "vertex_ai_beta" - ): + elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": optional_params = litellm.VertexAIConfig().map_special_auth_params( non_default_params=passed_params, optional_params=optional_params ) @@ -4060,11 +3668,7 @@ def pre_process_optional_params( ) ## raise exception if function calling passed in for a provider that doesn't support it - if ( - "functions" in non_default_params - or "function_call" in non_default_params - or "tools" in non_default_params - ): + if "functions" in non_default_params or "function_call" in non_default_params or "tools" in non_default_params: if ( custom_llm_provider == "ollama" and custom_llm_provider != "text-completion-openai" @@ -4095,23 +3699,13 @@ def pre_process_optional_params( if custom_llm_provider == "ollama": # ollama actually supports json output optional_params["format"] = "json" - litellm.add_function_to_prompt = ( - True # so that main.py adds the function call to the prompt - ) + litellm.add_function_to_prompt = True # so that main.py adds the function call to the prompt if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("tools") - ) - non_default_params.pop( - "tool_choice", None - ) # causes ollama requests to hang + optional_params["functions_unsupported_model"] = non_default_params.pop("tools") + non_default_params.pop("tool_choice", None) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("functions") - ) - elif ( - litellm.add_function_to_prompt - ): # if user opts to add it to prompt instead + optional_params["functions_unsupported_model"] = non_default_params.pop("functions") + elif litellm.add_function_to_prompt: # if user opts to add it to prompt instead optional_params["functions_unsupported_model"] = non_default_params.pop( "tools", non_default_params.pop("functions", None) ) @@ -4175,9 +3769,7 @@ def get_optional_params( # OpenAI param. passed_params.pop("base_model", None) provider_config: Optional[BaseConfig] = None - if custom_llm_provider is not None and custom_llm_provider in [ - provider.value for provider in LlmProviders - ]: + if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: provider_config = ProviderConfigManager.get_provider_chat_config( model=model, provider=LlmProviders(custom_llm_provider), @@ -4204,15 +3796,9 @@ def get_optional_params( Args: supported_params: List[str] - supported params from the litellm config """ - verbose_logger.info( - f"\nLiteLLM completion() model= {model}; provider = {custom_llm_provider}" - ) - verbose_logger.debug( - f"\nLiteLLM: Params passed to completion() {passed_params}" - ) - verbose_logger.debug( - f"\nLiteLLM: Non-Default params passed to completion() {non_default_params}" - ) + verbose_logger.info(f"\nLiteLLM completion() model= {model}; provider = {custom_llm_provider}") + verbose_logger.debug(f"\nLiteLLM: Params passed to completion() {passed_params}") + verbose_logger.debug(f"\nLiteLLM: Non-Default params passed to completion() {non_default_params}") unsupported_params = {} for k in non_default_params.keys(): if k not in supported_params: @@ -4229,9 +3815,7 @@ def get_optional_params( unsupported_params[k] = non_default_params[k] if unsupported_params: - if litellm.drop_params is True or ( - drop_params is not None and drop_params is True - ): + if litellm.drop_params is True or (drop_params is not None and drop_params is True): for k in unsupported_params.keys(): non_default_params.pop(k, None) else: @@ -4240,16 +3824,12 @@ def get_optional_params( message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.", ) - get_supported_openai_params = getattr( - sys.modules[__name__], "get_supported_openai_params" - ) + get_supported_openai_params = getattr(sys.modules[__name__], "get_supported_openai_params") supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider, base_model=base_model ) if supported_params is None: - supported_params = get_supported_openai_params( - model=model, custom_llm_provider="openai" - ) + supported_params = get_supported_openai_params(model=model, custom_llm_provider="openai") supported_params = supported_params or [] allowed_openai_params = allowed_openai_params or [] @@ -4265,32 +3845,20 @@ def get_optional_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "anthropic_text": optional_params = litellm.AnthropicTextConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) optional_params = litellm.AnthropicTextConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere": @@ -4299,11 +3867,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "triton": optional_params = litellm.TritonConfig().map_openai_params( @@ -4318,55 +3882,35 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "replicate": optional_params = litellm.ReplicateConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "predibase": optional_params = litellm.PredibaseConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "huggingface": optional_params = litellm.HuggingFaceChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "together_ai": optional_params = litellm.TogetherAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "vertex_ai" and ( model in litellm.vertex_chat_models @@ -4380,11 +3924,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "gemini": @@ -4392,96 +3932,58 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) - elif custom_llm_provider == "vertex_ai_beta" or ( - custom_llm_provider == "vertex_ai" and "gemini" in model - ): + elif custom_llm_provider == "vertex_ai_beta" or (custom_llm_provider == "vertex_ai" and "gemini" in model): optional_params = litellm.VertexGeminiConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) - elif litellm.VertexAIAnthropicConfig.is_supported_model( - model=model, custom_llm_provider=custom_llm_provider - ): + elif litellm.VertexAIAnthropicConfig.is_supported_model(model=model, custom_llm_provider=custom_llm_provider): optional_params = litellm.VertexAIAnthropicConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "vertex_ai": if model in litellm.vertex_mistral_models: if "codestral" in model: - optional_params = ( - litellm.CodestralTextCompletionConfig().map_openai_params( - model=model, - non_default_params=non_default_params, - optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), - ) + optional_params = litellm.CodestralTextCompletionConfig().map_openai_params( + model=model, + non_default_params=non_default_params, + optional_params=optional_params, + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) else: optional_params = litellm.MistralConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif model in litellm.vertex_ai_ai21_models: optional_params = litellm.VertexAIAi21Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) else: # use generic openai-like param mapping optional_params = litellm.VertexAILlama3Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "sagemaker": @@ -4490,11 +3992,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "bedrock": BedrockModelInfo = getattr(sys.modules[__name__], "BedrockModelInfo") @@ -4505,61 +4003,36 @@ def get_optional_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif bedrock_route == "openai": optional_params = litellm.AmazonBedrockOpenAIConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif "anthropic" in bedrock_base_model and bedrock_route == "invoke": - if ( - bedrock_base_model - in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names() - ): + if bedrock_base_model in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names(): optional_params = litellm.AmazonAnthropicConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) else: - optional_params = ( - litellm.AmazonAnthropicClaudeConfig().map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), - ) + optional_params = litellm.AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) if bedrock_route == "claude_platform": optional_params = BedrockModelInfo.map_claude_platform_auth_params( @@ -4570,44 +4043,28 @@ def get_optional_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "ollama": optional_params = litellm.OllamaConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "ollama_chat": optional_params = litellm.OllamaChatConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "nlp_cloud": optional_params = litellm.NLPCloudConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "petals": @@ -4615,55 +4072,35 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "deepinfra": optional_params = litellm.DeepInfraConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "perplexity" and provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "mistral" or custom_llm_provider == "codestral": optional_params = litellm.MistralConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "text-completion-codestral": optional_params = litellm.CodestralTextCompletionConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "text-completion-inception": @@ -4671,11 +4108,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "databricks": @@ -4683,33 +4116,21 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "nvidia_nim": optional_params = litellm.NvidiaNimConfig().map_openai_params( model=model, non_default_params=non_default_params, optional_params=optional_params, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "cerebras": optional_params = litellm.CerebrasConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "xai": optional_params = litellm.XAIChatConfig().map_openai_params( @@ -4722,110 +4143,70 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "fireworks_ai": optional_params = litellm.FireworksAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "volcengine": optional_params = litellm.VolcEngineConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "hosted_vllm": optional_params = litellm.HostedVLLMChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "vllm": optional_params = litellm.VLLMConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "groq": optional_params = litellm.GroqChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "bedrock_mantle": optional_params = litellm.BedrockMantleChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "deepseek": optional_params = litellm.DeepSeekChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "openrouter": optional_params = litellm.OpenrouterConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "watsonx": optional_params = litellm.IBMWatsonXChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) # WatsonX-text param check for param in passed_params.keys(): @@ -4838,61 +4219,37 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "openai": optional_params = litellm.OpenAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "nebius": optional_params = litellm.NebiusConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "azure": _azure_detection_model = base_model or model - if litellm.AzureOpenAIO1Config().is_o_series_model( - model=_azure_detection_model - ): + if litellm.AzureOpenAIO1Config().is_o_series_model(model=_azure_detection_model): optional_params = litellm.AzureOpenAIO1Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=_azure_detection_model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) - elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( - model=_azure_detection_model - ): + elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=_azure_detection_model): optional_params = litellm.AzureOpenAIGPT5Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=_azure_detection_model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) else: verbose_logger.debug( @@ -4911,33 +4268,21 @@ def get_optional_params( optional_params=optional_params, model=_azure_detection_model, api_version=api_version, # type: ignore - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif provider_config is not None: optional_params = provider_config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) else: # assume passing in params for openai-like api optional_params = litellm.OpenAILikeChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=( - drop_params - if drop_params is not None and isinstance(drop_params, bool) - else False - ), + drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) # if user passed in non-default kwargs for specific providers/models, pass them along optional_params = add_provider_specific_params_to_optional_params( @@ -4976,18 +4321,9 @@ def add_provider_specific_params_to_optional_params( Add provider specific params to optional_params """ - if ( - custom_llm_provider - in ["openai", "azure", "text-completion-openai"] - + litellm.openai_compatible_providers - ): + if custom_llm_provider in ["openai", "azure", "text-completion-openai"] + litellm.openai_compatible_providers: # for openai, azure we should pass the extra/passed params within `extra_body` https://github.com/openai/openai-python/blob/ac33853ba10d13ac149b1fa3ca6dba7d613065c9/src/openai/resources/models.py#L46 - if ( - _should_drop_param( - k="extra_body", additional_drop_params=additional_drop_params - ) - is False - ): + if _should_drop_param(k="extra_body", additional_drop_params=additional_drop_params) is False: extra_body = dict(passed_params.pop("extra_body", None) or {}) for k in passed_params.keys(): if k not in openai_params and passed_params[k] is not None: @@ -5000,34 +4336,22 @@ def add_provider_specific_params_to_optional_params( } if additional_drop_params is not None: - processed_extra_body = { - k: v - for k, v in initial_extra_body.items() - if k not in additional_drop_params - } + processed_extra_body = {k: v for k, v in initial_extra_body.items() if k not in additional_drop_params} else: processed_extra_body = initial_extra_body - _ensure_extra_body_is_safe = getattr( - sys.modules[__name__], "_ensure_extra_body_is_safe" - ) - optional_params["extra_body"] = _ensure_extra_body_is_safe( - extra_body=processed_extra_body - ) + _ensure_extra_body_is_safe = getattr(sys.modules[__name__], "_ensure_extra_body_is_safe") + optional_params["extra_body"] = _ensure_extra_body_is_safe(extra_body=processed_extra_body) else: for k in passed_params.keys(): if k not in openai_params and passed_params[k] is not None: - if _should_drop_param( - k=k, additional_drop_params=additional_drop_params - ): + if _should_drop_param(k=k, additional_drop_params=additional_drop_params): continue optional_params[k] = passed_params[k] return optional_params -def _apply_openai_param_overrides( - optional_params: dict, non_default_params: dict, allowed_openai_params: list -): +def _apply_openai_param_overrides(optional_params: dict, non_default_params: dict, allowed_openai_params: list): """ If user passes in allowed_openai_params, apply them to optional_params @@ -5122,13 +4446,9 @@ def _get_deployment_order(deployment: Union[Dict, Any]) -> Optional[int]: return order -def _get_order_filtered_deployments( - healthy_deployments: List[Dict], target_order: Optional[int] = None -) -> List: +def _get_order_filtered_deployments(healthy_deployments: List[Dict], target_order: Optional[int] = None) -> List: if target_order is not None: - filtered = [ - d for d in healthy_deployments if _get_deployment_order(d) == target_order - ] + filtered = [d for d in healthy_deployments if _get_deployment_order(d) == target_order] if filtered: return filtered # target_order doesn't match any deployment (e.g., external fallback model) — return all @@ -5136,18 +4456,13 @@ def _get_order_filtered_deployments( # Default: pick min order group _valid_orders: List[int] = [ - o - for deployment in healthy_deployments - for o in [_get_deployment_order(deployment)] - if o is not None + o for deployment in healthy_deployments for o in [_get_deployment_order(deployment)] if o is not None ] min_order: Optional[int] = min(_valid_orders) if _valid_orders else None if min_order is not None: filtered_deployments = [ - deployment - for deployment in healthy_deployments - if _get_deployment_order(deployment) == min_order + deployment for deployment in healthy_deployments if _get_deployment_order(deployment) == min_order ] return filtered_deployments @@ -5174,16 +4489,10 @@ def _get_excluded_filtered_deployments( return healthy_deployments excluded_set = set(excluded_deployment_ids) - return [ - d - for d in healthy_deployments - if (d.get("model_info") or {}).get("id") not in excluded_set - ] + return [d for d in healthy_deployments if (d.get("model_info") or {}).get("id") not in excluded_set] -def _get_model_region( - custom_llm_provider: str, litellm_params: LiteLLM_Params -) -> Optional[str]: +def _get_model_region(custom_llm_provider: str, litellm_params: LiteLLM_Params) -> Optional[str]: """ Return the region for a model, for a given provider """ @@ -5220,14 +4529,10 @@ def _infer_model_region(litellm_params: LiteLLM_Params) -> Optional[AllowedModel model=litellm_params.model, litellm_params=litellm_params ) - model_region = _get_model_region( - custom_llm_provider=custom_llm_provider, litellm_params=litellm_params - ) + model_region = _get_model_region(custom_llm_provider=custom_llm_provider, litellm_params=litellm_params) if model_region is None: - verbose_logger.debug( - "Cannot infer model region for model: {}".format(litellm_params.model) - ) + verbose_logger.debug("Cannot infer model region for model: {}".format(litellm_params.model)) return None if custom_llm_provider == "azure": @@ -5283,9 +4588,7 @@ def _is_region_us(litellm_params: LiteLLM_Params) -> bool: return False -def is_region_allowed( - litellm_params: LiteLLM_Params, allowed_model_region: str -) -> bool: +def is_region_allowed(litellm_params: LiteLLM_Params, allowed_model_region: str) -> bool: """ Return true/false if a deployment is in the EU """ @@ -5294,9 +4597,7 @@ def is_region_allowed( return False -def get_model_region( - litellm_params: LiteLLM_Params, mode: Optional[str] -) -> Optional[str]: +def get_model_region(litellm_params: LiteLLM_Params, mode: Optional[str]) -> Optional[str]: """ Pass the litellm params for an azure model, and get back the region """ @@ -5394,9 +4695,7 @@ def get_api_key(llm_provider: str, dynamic_api_key: Optional[str]): api_key = api_key or litellm.ai21_key or get_secret("AI211_API_KEY") # aleph_alpha elif llm_provider == "aleph_alpha": - api_key = ( - api_key or litellm.aleph_alpha_key or get_secret("ALEPH_ALPHA_API_KEY") - ) + api_key = api_key or litellm.aleph_alpha_key or get_secret("ALEPH_ALPHA_API_KEY") # baseten elif llm_provider == "baseten": api_key = api_key or litellm.baseten_key or get_secret("BASETEN_API_KEY") @@ -5405,9 +4704,7 @@ def get_api_key(llm_provider: str, dynamic_api_key: Optional[str]): api_key = api_key or litellm.cohere_key or get_secret("COHERE_API_KEY") # huggingface elif llm_provider == "huggingface": - api_key = ( - api_key or litellm.huggingface_key or get_secret("HUGGINGFACE_API_KEY") - ) + api_key = api_key or litellm.huggingface_key or get_secret("HUGGINGFACE_API_KEY") # nlp_cloud elif llm_provider == "nlp_cloud": api_key = api_key or litellm.nlp_cloud_key or get_secret("NLP_CLOUD_API_KEY") @@ -5417,10 +4714,7 @@ def get_api_key(llm_provider: str, dynamic_api_key: Optional[str]): # together_ai elif llm_provider == "together_ai": api_key = ( - api_key - or litellm.togetherai_api_key - or get_secret("TOGETHERAI_API_KEY") - or get_secret("TOGETHER_AI_TOKEN") + api_key or litellm.togetherai_api_key or get_secret("TOGETHERAI_API_KEY") or get_secret("TOGETHER_AI_TOKEN") ) # nebius elif llm_provider == "nebius": @@ -5539,9 +4833,7 @@ def _strip_model_name(model: str, custom_llm_provider: Optional[str]) -> str: if custom_llm_provider and custom_llm_provider in ["bedrock", "bedrock_converse"]: stripped_bedrock_model = _get_base_bedrock_model(model_name=model) return stripped_bedrock_model - elif custom_llm_provider and ( - custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini" - ): + elif custom_llm_provider and (custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini"): strip_version = _strip_stable_vertex_version(model_name=model) return strip_version elif custom_llm_provider and (custom_llm_provider == "databricks"): @@ -5688,20 +4980,13 @@ def _check_provider_match(model_info: dict, custom_llm_provider: Optional[str]) so normalising the two cases keeps custom pricing applied consistently. """ if custom_llm_provider and ( - model_info.get("litellm_provider") is not None - and model_info["litellm_provider"] != custom_llm_provider + model_info.get("litellm_provider") is not None and model_info["litellm_provider"] != custom_llm_provider ): - if custom_llm_provider == "vertex_ai" and model_info[ - "litellm_provider" - ].startswith("vertex_ai"): + if custom_llm_provider == "vertex_ai" and model_info["litellm_provider"].startswith("vertex_ai"): return True - elif custom_llm_provider == "fireworks_ai" and model_info[ - "litellm_provider" - ].startswith("fireworks_ai"): + elif custom_llm_provider == "fireworks_ai" and model_info["litellm_provider"].startswith("fireworks_ai"): return True - elif custom_llm_provider.startswith("bedrock") and model_info[ - "litellm_provider" - ].startswith("bedrock"): + elif custom_llm_provider.startswith("bedrock") and model_info["litellm_provider"].startswith("bedrock"): return True elif ( custom_llm_provider == "litellm_proxy" @@ -5746,27 +5031,19 @@ def _get_potential_model_names( except Exception: split_model = model combined_model_name = model - stripped_model_name = _strip_model_name( - model=model, custom_llm_provider=custom_llm_provider - ) + stripped_model_name = _strip_model_name(model=model, custom_llm_provider=custom_llm_provider) combined_stripped_model_name = stripped_model_name - elif ( - custom_llm_provider and model.startswith(custom_llm_provider + "/") + elif custom_llm_provider and model.startswith( + custom_llm_provider + "/" ): # handle case where custom_llm_provider is provided and model starts with custom_llm_provider split_model = model.split("/", 1)[1] combined_model_name = model - stripped_model_name = _strip_model_name( - model=split_model, custom_llm_provider=custom_llm_provider - ) - combined_stripped_model_name = "{}/{}".format( - custom_llm_provider, stripped_model_name - ) + stripped_model_name = _strip_model_name(model=split_model, custom_llm_provider=custom_llm_provider) + combined_stripped_model_name = "{}/{}".format(custom_llm_provider, stripped_model_name) else: split_model = model combined_model_name = "{}/{}".format(custom_llm_provider, model) - stripped_model_name = _strip_model_name( - model=model, custom_llm_provider=custom_llm_provider - ) + stripped_model_name = _strip_model_name(model=model, custom_llm_provider=custom_llm_provider) combined_stripped_model_name = "{}/{}".format( custom_llm_provider, stripped_model_name, @@ -5822,9 +5099,7 @@ def _cached_get_model_info_helper( ) -def get_provider_info( - model: str, custom_llm_provider: Optional[str] -) -> Optional[ProviderSpecificModelInfo]: +def get_provider_info(model: str, custom_llm_provider: Optional[str]) -> Optional[ProviderSpecificModelInfo]: ## PROVIDER-SPECIFIC INFORMATION # if custom_llm_provider == "predibase": # _model_info["supports_response_schema"] = True @@ -5880,19 +5155,13 @@ def _get_model_info_helper( elif model + "@latest" in litellm.vertex_ai_ai21_models: model = model + "@latest" ########################## - potential_model_names = _get_potential_model_names( - model=model, custom_llm_provider=custom_llm_provider - ) + potential_model_names = _get_potential_model_names(model=model, custom_llm_provider=custom_llm_provider) - verbose_logger.debug( - f"checking potential_model_names in litellm.model_cost: {potential_model_names}" - ) + verbose_logger.debug(f"checking potential_model_names in litellm.model_cost: {potential_model_names}") combined_model_name = potential_model_names["combined_model_name"] stripped_model_name = potential_model_names["stripped_model_name"] - combined_stripped_model_name = potential_model_names[ - "combined_stripped_model_name" - ] + combined_stripped_model_name = potential_model_names["combined_stripped_model_name"] split_model = potential_model_names["split_model"] custom_llm_provider = potential_model_names["custom_llm_provider"] model_cost_custom_llm_provider = custom_llm_provider @@ -6010,9 +5279,7 @@ def _get_model_info_helper( raise ValueError( "This model isn't mapped yet. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" ) - _input_cost_per_token: Optional[float] = _model_info.get( - "input_cost_per_token" - ) + _input_cost_per_token: Optional[float] = _model_info.get("input_cost_per_token") if _input_cost_per_token is None: # default value to 0, be noisy about this verbose_logger.debug( @@ -6022,9 +5289,7 @@ def _get_model_info_helper( ) _input_cost_per_token = 0 - _output_cost_per_token: Optional[float] = _model_info.get( - "output_cost_per_token" - ) + _output_cost_per_token: Optional[float] = _model_info.get("output_cost_per_token") if _output_cost_per_token is None: # default value to 0, be noisy about this verbose_logger.debug( @@ -6040,21 +5305,13 @@ def _get_model_info_helper( max_input_tokens=_model_info.get("max_input_tokens", None), max_output_tokens=_model_info.get("max_output_tokens", None), input_cost_per_token=_input_cost_per_token, - input_cost_per_token_flex=_model_info.get( - "input_cost_per_token_flex", None - ), - input_cost_per_token_priority=_model_info.get( - "input_cost_per_token_priority", None - ), - cache_creation_input_token_cost=_model_info.get( - "cache_creation_input_token_cost", None - ), + input_cost_per_token_flex=_model_info.get("input_cost_per_token_flex", None), + input_cost_per_token_priority=_model_info.get("input_cost_per_token_priority", None), + cache_creation_input_token_cost=_model_info.get("cache_creation_input_token_cost", None), cache_creation_input_token_cost_above_200k_tokens=_model_info.get( "cache_creation_input_token_cost_above_200k_tokens", None ), - cache_read_input_token_cost=_model_info.get( - "cache_read_input_token_cost", None - ), + cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( "cache_read_input_token_cost_above_200k_tokens", None ), @@ -6070,79 +5327,43 @@ def _get_model_info_helper( cache_read_input_token_cost_above_512k_tokens=_model_info.get( "cache_read_input_token_cost_above_512k_tokens", None ), - cache_read_input_token_cost_flex=_model_info.get( - "cache_read_input_token_cost_flex", None - ), - cache_read_input_token_cost_priority=_model_info.get( - "cache_read_input_token_cost_priority", None - ), + cache_read_input_token_cost_flex=_model_info.get("cache_read_input_token_cost_flex", None), + cache_read_input_token_cost_priority=_model_info.get("cache_read_input_token_cost_priority", None), cache_creation_input_token_cost_above_1hr=_model_info.get( "cache_creation_input_token_cost_above_1hr", None ), - input_cost_per_character=_model_info.get( - "input_cost_per_character", None - ), - input_cost_per_token_above_128k_tokens=_model_info.get( - "input_cost_per_token_above_128k_tokens", None - ), - input_cost_per_token_above_200k_tokens=_model_info.get( - "input_cost_per_token_above_200k_tokens", None - ), + input_cost_per_character=_model_info.get("input_cost_per_character", None), + input_cost_per_token_above_128k_tokens=_model_info.get("input_cost_per_token_above_128k_tokens", None), + input_cost_per_token_above_200k_tokens=_model_info.get("input_cost_per_token_above_200k_tokens", None), input_cost_per_token_above_200k_tokens_priority=_model_info.get( "input_cost_per_token_above_200k_tokens_priority", None ), - input_cost_per_token_above_272k_tokens=_model_info.get( - "input_cost_per_token_above_272k_tokens", None - ), + input_cost_per_token_above_272k_tokens=_model_info.get("input_cost_per_token_above_272k_tokens", None), input_cost_per_token_above_272k_tokens_priority=_model_info.get( "input_cost_per_token_above_272k_tokens_priority", None ), - input_cost_per_token_above_512k_tokens=_model_info.get( - "input_cost_per_token_above_512k_tokens", None - ), + input_cost_per_token_above_512k_tokens=_model_info.get("input_cost_per_token_above_512k_tokens", None), input_cost_per_query=_model_info.get("input_cost_per_query", None), input_cost_per_second=_model_info.get("input_cost_per_second", None), - input_cost_per_audio_token=_model_info.get( - "input_cost_per_audio_token", None - ), - input_cost_per_image_token=_model_info.get( - "input_cost_per_image_token", None - ), + input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), + input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), input_cost_per_image=_model_info.get("input_cost_per_image", None), - input_cost_per_audio_per_second=_model_info.get( - "input_cost_per_audio_per_second", None - ), - input_cost_per_video_per_second=_model_info.get( - "input_cost_per_video_per_second", None - ), - input_cost_per_token_batches=_model_info.get( - "input_cost_per_token_batches" - ), - output_cost_per_token_batches=_model_info.get( - "output_cost_per_token_batches" - ), + input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None), + input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None), + input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"), + output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"), output_cost_per_token=_output_cost_per_token, - output_cost_per_token_flex=_model_info.get( - "output_cost_per_token_flex", None - ), - output_cost_per_token_priority=_model_info.get( - "output_cost_per_token_priority", None - ), + output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None), + output_cost_per_token_priority=_model_info.get("output_cost_per_token_priority", None), regional_processing_uplift_multiplier_eu=_model_info.get( "regional_processing_uplift_multiplier_eu", None ), regional_processing_uplift_multiplier_us=_model_info.get( "regional_processing_uplift_multiplier_us", None ), - output_cost_per_audio_token=_model_info.get( - "output_cost_per_audio_token", None - ), - output_cost_per_character=_model_info.get( - "output_cost_per_character", None - ), - output_cost_per_reasoning_token=_model_info.get( - "output_cost_per_reasoning_token", None - ), + output_cost_per_audio_token=_model_info.get("output_cost_per_audio_token", None), + output_cost_per_character=_model_info.get("output_cost_per_character", None), + output_cost_per_reasoning_token=_model_info.get("output_cost_per_reasoning_token", None), output_cost_per_token_above_128k_tokens=_model_info.get( "output_cost_per_token_above_128k_tokens", None ), @@ -6165,103 +5386,52 @@ def _get_model_info_helper( "output_cost_per_token_above_512k_tokens", None ), output_cost_per_second=_model_info.get("output_cost_per_second", None), - output_cost_per_second_1080p=_model_info.get( - "output_cost_per_second_1080p", None - ), - output_cost_per_video_per_second=_model_info.get( - "output_cost_per_video_per_second", None - ), + output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), + output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), - output_cost_per_image_token=_model_info.get( - "output_cost_per_image_token", None - ), + output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_vector_size=_model_info.get("output_vector_size", None), - citation_cost_per_token=_model_info.get( - "citation_cost_per_token", None - ), + citation_cost_per_token=_model_info.get("citation_cost_per_token", None), tiered_pricing=_model_info.get("tiered_pricing", None), - litellm_provider=_model_info.get( - "litellm_provider", custom_llm_provider - ), + litellm_provider=_model_info.get("litellm_provider", custom_llm_provider), mode=_model_info.get("mode"), # type: ignore - supports_system_messages=_model_info.get( - "supports_system_messages", None - ), - supports_response_schema=_model_info.get( - "supports_response_schema", None - ), + supports_system_messages=_model_info.get("supports_system_messages", None), + supports_response_schema=_model_info.get("supports_response_schema", None), supports_vision=_model_info.get("supports_vision", None), - supports_function_calling=_model_info.get( - "supports_function_calling", None - ), + supports_function_calling=_model_info.get("supports_function_calling", None), supports_tool_choice=_model_info.get("supports_tool_choice", None), - supports_assistant_prefill=_model_info.get( - "supports_assistant_prefill", None - ), - supports_prompt_caching=_model_info.get( - "supports_prompt_caching", None - ), + supports_assistant_prefill=_model_info.get("supports_assistant_prefill", None), + supports_prompt_caching=_model_info.get("supports_prompt_caching", None), supports_audio_input=_model_info.get("supports_audio_input", None), supports_audio_output=_model_info.get("supports_audio_output", None), supports_pdf_input=_model_info.get("supports_pdf_input", None), - supports_embedding_image_input=_model_info.get( - "supports_embedding_image_input", None - ), - supports_native_streaming=_model_info.get( - "supports_native_streaming", None - ), - supports_native_structured_output=_model_info.get( - "supports_native_structured_output", None - ), + supports_embedding_image_input=_model_info.get("supports_embedding_image_input", None), + supports_native_streaming=_model_info.get("supports_native_streaming", None), + supports_native_structured_output=_model_info.get("supports_native_structured_output", None), supports_web_search=_model_info.get("supports_web_search", None), supports_url_context=_model_info.get("supports_url_context", None), supports_reasoning=_model_info.get("supports_reasoning", None), - supports_adaptive_thinking=_model_info.get( - "supports_adaptive_thinking", None - ), - supports_none_reasoning_effort=_model_info.get( - "supports_none_reasoning_effort", None - ), - supports_minimal_reasoning_effort=_model_info.get( - "supports_minimal_reasoning_effort", None - ), - supports_low_reasoning_effort=_model_info.get( - "supports_low_reasoning_effort", None - ), - supports_xhigh_reasoning_effort=_model_info.get( - "supports_xhigh_reasoning_effort", None - ), - supports_max_reasoning_effort=_model_info.get( - "supports_max_reasoning_effort", None - ), - bedrock_output_config_effort_ceiling=_model_info.get( - "bedrock_output_config_effort_ceiling", None - ), + supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None), + supports_none_reasoning_effort=_model_info.get("supports_none_reasoning_effort", None), + supports_minimal_reasoning_effort=_model_info.get("supports_minimal_reasoning_effort", None), + supports_low_reasoning_effort=_model_info.get("supports_low_reasoning_effort", None), + supports_xhigh_reasoning_effort=_model_info.get("supports_xhigh_reasoning_effort", None), + supports_max_reasoning_effort=_model_info.get("supports_max_reasoning_effort", None), + bedrock_output_config_effort_ceiling=_model_info.get("bedrock_output_config_effort_ceiling", None), supports_computer_use=_model_info.get("supports_computer_use", None), - search_context_cost_per_query=_model_info.get( - "search_context_cost_per_query", None - ), - web_search_billing_unit=_model_info.get( - "web_search_billing_unit", None - ), + search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None), + web_search_billing_unit=_model_info.get("web_search_billing_unit", None), tpm=_model_info.get("tpm", None), rpm=_model_info.get("rpm", None), ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None), ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None), - annotation_cost_per_page=_model_info.get( - "annotation_cost_per_page", None - ), - provider_specific_entry=_model_info.get( - "provider_specific_entry", None - ), + annotation_cost_per_page=_model_info.get("annotation_cost_per_page", None), + provider_specific_entry=_model_info.get("provider_specific_entry", None), uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), ) for cost_key, cost_value in _model_info.items(): - if ( - cost_key not in returned_model_info - and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None - ): + if cost_key not in returned_model_info and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None: returned_model_info[cost_key] = cost_value # type: ignore[literal-required] return returned_model_info except Exception as e: @@ -6279,9 +5449,7 @@ def _build_model_info( api_base: Optional[str] = None, api_key: Optional[str] = None, ) -> ModelInfo: - supported_openai_params = litellm.get_supported_openai_params( - model=model, custom_llm_provider=custom_llm_provider - ) + supported_openai_params = litellm.get_supported_openai_params(model=model, custom_llm_provider=custom_llm_provider) _model_info = _get_model_info_helper( model=model, @@ -6290,9 +5458,7 @@ def _build_model_info( api_key=api_key, ) - provider_info = get_provider_info( - model=model, custom_llm_provider=custom_llm_provider - ) + provider_info = get_provider_info(model=model, custom_llm_provider=custom_llm_provider) if provider_info: for key, value in provider_info.items(): if value is not None: @@ -6310,9 +5476,7 @@ def _cached_get_model_info( custom_llm_provider: Optional[str] = None, api_base: Optional[str] = None, ) -> ModelInfo: - return _build_model_info( - model=model, custom_llm_provider=custom_llm_provider, api_base=api_base - ) + return _build_model_info(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base) def get_model_info( @@ -6494,9 +5658,7 @@ def function_to_dict(input_function) -> dict: # noqa: C901 "enum": param_enum, } - parameters[param_name] = dict( - [(k, v) for k, v in param_dict.items() if isinstance(v, str)] - ) + parameters[param_name] = dict([(k, v) for k, v in param_dict.items() if isinstance(v, str)]) # Check if the parameter has no default value (i.e., it's required) if param.default == param.empty: @@ -6585,10 +5747,7 @@ def get_provider_fields(custom_llm_provider: str) -> List[ProviderField]: def create_proxy_transport_and_mounts(): - proxies = { - key: None if url is None else Proxy(url=url) - for key, url in get_environment_proxies().items() - } + proxies = {key: None if url is None else Proxy(url=url) for key, url in get_environment_proxies().items()} sync_proxy_mounts = {} async_proxy_mounts = {} @@ -6652,21 +5811,12 @@ def validate_environment( else: missing_keys.append("OPENAI_API_KEY") elif custom_llm_provider == "azure": - if ( - "AZURE_API_BASE" in os.environ - and "AZURE_API_VERSION" in os.environ - and "AZURE_API_KEY" in os.environ - ): + if "AZURE_API_BASE" in os.environ and "AZURE_API_VERSION" in os.environ and "AZURE_API_KEY" in os.environ: keys_in_environment = True else: - missing_keys.extend( - ["AZURE_API_BASE", "AZURE_API_VERSION", "AZURE_API_KEY"] - ) + missing_keys.extend(["AZURE_API_BASE", "AZURE_API_VERSION", "AZURE_API_KEY"]) elif custom_llm_provider == "anthropic": - if ( - "ANTHROPIC_API_KEY" in os.environ - or "ANTHROPIC_AUTH_TOKEN" in os.environ - ): + if "ANTHROPIC_API_KEY" in os.environ or "ANTHROPIC_AUTH_TOKEN" in os.environ: keys_in_environment = True else: missing_keys.append("ANTHROPIC_API_KEY") @@ -6731,18 +5881,13 @@ def validate_environment( else: missing_keys.append("NLP_CLOUD_API_KEY") elif custom_llm_provider == "bedrock" or custom_llm_provider == "sagemaker": - if ( - "AWS_ACCESS_KEY_ID" in os.environ - and "AWS_SECRET_ACCESS_KEY" in os.environ - ) or ( + if ("AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ) or ( # IAM role, profile, or web identity auth don't require access keys "AWS_ROLE_ARN" in os.environ or "AWS_PROFILE" in os.environ or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ - or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" - in os.environ # ECS task role - or "AWS_CONTAINER_CREDENTIALS_FULL_URI" - in os.environ # ECS/Fargate full URI credential delivery + or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ # ECS task role + or "AWS_CONTAINER_CREDENTIALS_FULL_URI" in os.environ # ECS/Fargate full URI credential delivery ): keys_in_environment = True else: @@ -6809,18 +5954,12 @@ def validate_environment( keys_in_environment = True else: missing_keys.append("VOLCENGINE_API_KEY") - elif ( - custom_llm_provider == "codestral" - or custom_llm_provider == "text-completion-codestral" - ): + elif custom_llm_provider == "codestral" or custom_llm_provider == "text-completion-codestral": if "CODESTRAL_API_KEY" in os.environ: keys_in_environment = True else: missing_keys.append("CODESTRAL_API_KEY") - elif ( - custom_llm_provider == "inception" - or custom_llm_provider == "text-completion-inception" - ): + elif custom_llm_provider == "inception" or custom_llm_provider == "text-completion-inception": if "INCEPTION_API_KEY" in os.environ: keys_in_environment = True else: @@ -6867,8 +6006,7 @@ def validate_environment( missing_keys.append("FIREWORKS_AI_API_KEY") elif custom_llm_provider == "cloudflare": if "CLOUDFLARE_API_KEY" in os.environ and ( - "CLOUDFLARE_ACCOUNT_ID" in os.environ - or "CLOUDFLARE_API_BASE" in os.environ + "CLOUDFLARE_ACCOUNT_ID" in os.environ or "CLOUDFLARE_API_BASE" in os.environ ): keys_in_environment = True else: @@ -6919,10 +6057,7 @@ def validate_environment( missing_keys.append("OPENAI_API_KEY") ## anthropic elif model in litellm.anthropic_models: - if ( - "ANTHROPIC_API_KEY" in os.environ - or "ANTHROPIC_AUTH_TOKEN" in os.environ - ): + if "ANTHROPIC_API_KEY" in os.environ or "ANTHROPIC_AUTH_TOKEN" in os.environ: keys_in_environment = True else: missing_keys.append("ANTHROPIC_API_KEY") @@ -7062,10 +6197,7 @@ def prompt_token_calculator(model, messages): def valid_model(model): try: # for a given model name, check if the user has the right permissions to access the model - if ( - model in litellm.open_ai_chat_completion_models - or model in litellm.open_ai_text_completion_models - ): + if model in litellm.open_ai_chat_completion_models or model in litellm.open_ai_text_completion_models: openai.models.retrieve(model) else: messages = [{"role": "user", "content": "Hello World"}] @@ -7087,9 +6219,7 @@ def check_valid_key(model: str, api_key: str): """ messages = [{"role": "user", "content": "Hey, how's it going?"}] try: - litellm.completion( - model=model, messages=messages, api_key=api_key, max_tokens=10 - ) + litellm.completion(model=model, messages=messages, api_key=api_key, max_tokens=10) return True except AuthenticationError: return False @@ -7281,9 +6411,7 @@ class TextCompletionStreamWrapper: response["created"] = chunk.get("created", None) response["model"] = chunk.get("model", None) text_choices = TextChoices() - if isinstance( - chunk, Choices - ): # chunk should always be of type StreamingChoices + if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception delta = chunk["choices"][0]["delta"] text_choices["text"] = delta["content"] @@ -7293,17 +6421,12 @@ class TextCompletionStreamWrapper: response["choices"] = [text_choices] # only pass usage when stream_options["include_usage"] is True - if ( - self.stream_options - and self.stream_options.get("include_usage", False) is True - ): + if self.stream_options and self.stream_options.get("include_usage", False) is True: response["usage"] = chunk.get("usage", None) return response except Exception as e: - raise Exception( - f"Error occurred converting to text completion object - chunk: {chunk}; Error: {str(e)}" - ) + raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {str(e)}") def __next__(self): # model_response = ModelResponse(stream=True, model=self.model) @@ -7339,9 +6462,7 @@ class TextCompletionStreamWrapper: raise StopAsyncIteration -def mock_completion_streaming_obj( - model_response, mock_response, model, n: Optional[int] = None -): +def mock_completion_streaming_obj(model_response, mock_response, model, n: Optional[int] = None): if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -7356,9 +6477,7 @@ def mock_completion_streaming_obj( for j in range(n): _streaming_choice = litellm.utils.StreamingChoices( index=j, - delta=litellm.utils.Delta( - role="assistant", content=mock_response[i : i + 3] - ), + delta=litellm.utils.Delta(role="assistant", content=mock_response[i : i + 3]), ) _all_choices.append(_streaming_choice) model_response.choices = _all_choices @@ -7385,9 +6504,7 @@ async def async_mock_completion_streaming_obj( for j in range(n): _streaming_choice = litellm.utils.StreamingChoices( index=j, - delta=litellm.utils.Delta( - role="assistant", content=mock_response[i : i + 3] - ), + delta=litellm.utils.Delta(role="assistant", content=mock_response[i : i + 3]), ) _all_choices.append(_streaming_choice) model_response.choices = _all_choices @@ -7417,13 +6534,9 @@ def process_system_message(system_message, max_tokens, model): system_message_tokens = get_token_count([system_message_event], model) if system_message_tokens > max_tokens: - print_verbose( - "`tokentrimmer`: Warning, system message exceeds token limit. Trimming..." - ) + print_verbose("`tokentrimmer`: Warning, system message exceeds token limit. Trimming...") # shorten system message to fit within max_tokens - new_system_message = shorten_message_to_fit_limit( - system_message_event, max_tokens, model - ) + new_system_message = shorten_message_to_fit_limit(system_message_event, max_tokens, model) system_message_tokens = get_token_count([new_system_message], model) return system_message_event, max_tokens - system_message_tokens @@ -7440,9 +6553,7 @@ def process_messages(messages, max_tokens, model): verbose_logger.debug(f"processing final_messages: {final_messages}") used_tokens = get_token_count(final_messages, model) available_tokens = max_tokens - used_tokens - verbose_logger.debug( - f"used_tokens: {used_tokens}, available_tokens: {available_tokens}" - ) + verbose_logger.debug(f"used_tokens: {used_tokens}, available_tokens: {available_tokens}") if available_tokens <= 3: break @@ -7453,21 +6564,15 @@ def process_messages(messages, max_tokens, model): max_tokens=max_tokens, model=model, ) - verbose_logger.debug( - f"final_messages after attempt_message_addition: {final_messages}" - ) + verbose_logger.debug(f"final_messages after attempt_message_addition: {final_messages}") verbose_logger.debug(f"Final messages: {final_messages}") return final_messages -def attempt_message_addition( - final_messages, message, available_tokens, max_tokens, model -): +def attempt_message_addition(final_messages, message, available_tokens, max_tokens, model): temp_messages = [message] + final_messages temp_message_tokens = get_token_count(messages=temp_messages, model=model) - verbose_logger.debug( - f"temp_message_tokens: {temp_message_tokens}, max_tokens: {max_tokens}" - ) + verbose_logger.debug(f"temp_message_tokens: {temp_message_tokens}, max_tokens: {max_tokens}") if temp_message_tokens <= max_tokens: return temp_messages @@ -7477,9 +6582,7 @@ def attempt_message_addition( # fit updated_message to be within temp_message_tokens - max_tokens (aka the amount temp_message_tokens is greate than max_tokens) updated_message = shorten_message_to_fit_limit(message, available_tokens, model) if can_add_message(updated_message, final_messages, max_tokens, model): - verbose_logger.debug( - "can add message, returning [updated_message] + final_messages" - ) + verbose_logger.debug("can add message, returning [updated_message] + final_messages") return [updated_message] + final_messages else: verbose_logger.debug("cannot add message, returning final_messages") @@ -7496,9 +6599,7 @@ def get_token_count(messages, model): return token_counter(model=model, messages=messages) -def shorten_message_to_fit_limit( - message, tokens_needed, model: Optional[str], raise_error_on_max_limit: bool = False -): +def shorten_message_to_fit_limit(message, tokens_needed, model: Optional[str], raise_error_on_max_limit: bool = False): """ Shorten a message to fit within a token limit by removing characters from the middle. @@ -7523,9 +6624,7 @@ def shorten_message_to_fit_limit( while attempts < MAX_TOKEN_TRIMMING_ATTEMPTS: verbose_logger.debug(f"getting token count for message: {message}") total_tokens = get_token_count([message], model) - verbose_logger.debug( - f"total_tokens: {total_tokens}, tokens_needed: {tokens_needed}" - ) + verbose_logger.debug(f"total_tokens: {total_tokens}, tokens_needed: {tokens_needed}") if total_tokens <= tokens_needed: break @@ -7637,9 +6736,7 @@ def trim_messages( messages = [message for message in messages if message["role"] != "system"] verbose_logger.debug(f"Processed system message: {system_message_event}") - final_messages = process_messages( - messages=messages, max_tokens=max_tokens, model=model - ) + final_messages = process_messages(messages=messages, max_tokens=max_tokens, model=model) verbose_logger.debug(f"Processed messages: {final_messages}") # Add system message to the beginning of the final messages @@ -7649,19 +6746,13 @@ def trim_messages( if len(tool_messages) > 0: final_messages.extend(tool_messages) - verbose_logger.debug( - f"Final messages: {final_messages}, return_response_tokens: {return_response_tokens}" - ) - if ( - return_response_tokens - ): # if user wants token count with new trimmed messages + verbose_logger.debug(f"Final messages: {final_messages}, return_response_tokens: {return_response_tokens}") + if return_response_tokens: # if user wants token count with new trimmed messages response_tokens = max_tokens - get_token_count(final_messages, model) return final_messages, response_tokens return final_messages except Exception as e: # [NON-Blocking, if error occurs just return final_messages - verbose_logger.exception( - "Got exception while token trimming - {}".format(str(e)) - ) + verbose_logger.exception("Got exception while token trimming - {}".format(str(e))) return original_messages @@ -7675,11 +6766,7 @@ class AvailableModelsCache(InMemoryCache): def _get_env_hash(self) -> str: """Create a hash of relevant environment variables""" - env_vars = { - k: v - for k, v in os.environ.items() - if k.startswith(("OPENAI", "ANTHROPIC", "AZURE", "AWS")) - } + env_vars = {k: v for k, v in os.environ.items() if k.startswith(("OPENAI", "ANTHROPIC", "AZURE", "AWS"))} return str(hash(frozenset(env_vars.items()))) def _check_env_changed(self) -> bool: @@ -7754,10 +6841,7 @@ def _infer_valid_provider_from_env_vars( # PROVIDER_API_KEY. Example: OPENAI_API_KEY, COHERE_API_KEY expected_provider_key_1 = f"{env_provider_1.upper()}_API_KEY" expected_provider_key_2 = f"{env_provider_2.upper()}_API_KEY" - if ( - expected_provider_key_1 in environ_keys - or expected_provider_key_2 in environ_keys - ): + if expected_provider_key_1 in environ_keys or expected_provider_key_2 in environ_keys: # key is set valid_providers.append(provider) @@ -7770,9 +6854,7 @@ def _get_valid_models_from_provider_api( litellm_params: Optional[LiteLLM_Params] = None, ) -> List[str]: try: - cached_result = _model_cache.get_cached_model_info( - custom_llm_provider, litellm_params - ) + cached_result = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params) if cached_result is not None: return cached_result @@ -7821,9 +6903,7 @@ def get_valid_models( litellm_params.api_base = api_base ################################# - check_provider_endpoint = ( - check_provider_endpoint or litellm.check_provider_endpoint - ) + check_provider_endpoint = check_provider_endpoint or litellm.check_provider_endpoint # get keys set in .env valid_providers: List[str] = [] @@ -7846,11 +6926,7 @@ def get_valid_models( if provider == "azure": valid_models.append("Azure-LLM") - elif ( - provider_config is not None - and check_provider_endpoint - and provider is not None - ): + elif provider_config is not None and check_provider_endpoint and provider is not None: valid_models.extend( _get_valid_models_from_provider_api( provider_config, @@ -7859,9 +6935,7 @@ def get_valid_models( ) ) else: - models_for_provider = copy.deepcopy( - litellm.models_by_provider.get(provider, []) - ) + models_for_provider = copy.deepcopy(litellm.models_by_provider.get(provider, [])) valid_models.extend(models_for_provider) return valid_models @@ -7875,17 +6949,9 @@ def print_args_passed_to_litellm(original_function, args, kwargs): return try: # we've already printed this for acompletion, don't print for completion - if ( - "acompletion" in kwargs - and kwargs["acompletion"] is True - and original_function.__name__ == "completion" - ): + if "acompletion" in kwargs and kwargs["acompletion"] is True and original_function.__name__ == "completion": return - elif ( - "aembedding" in kwargs - and kwargs["aembedding"] is True - and original_function.__name__ == "embedding" - ): + elif "aembedding" in kwargs and kwargs["aembedding"] is True and original_function.__name__ == "embedding": return elif ( "aimg_generation" in kwargs @@ -7903,17 +6969,11 @@ def print_args_passed_to_litellm(original_function, args, kwargs): "\033[92mRequest to litellm:\033[0m", ) if args and kwargs: - print_verbose( - f"\033[92mlitellm.{original_function.__name__}({args_str}, {kwargs_str})\033[0m" - ) + print_verbose(f"\033[92mlitellm.{original_function.__name__}({args_str}, {kwargs_str})\033[0m") elif args: - print_verbose( - f"\033[92mlitellm.{original_function.__name__}({args_str})\033[0m" - ) + print_verbose(f"\033[92mlitellm.{original_function.__name__}({args_str})\033[0m") elif kwargs: - print_verbose( - f"\033[92mlitellm.{original_function.__name__}({kwargs_str})\033[0m" - ) + print_verbose(f"\033[92mlitellm.{original_function.__name__}({kwargs_str})\033[0m") else: print_verbose(f"\033[92mlitellm.{original_function.__name__}()\033[0m") print_verbose("\n") # new line after @@ -7924,9 +6984,7 @@ def print_args_passed_to_litellm(original_function, args, kwargs): def get_logging_id(start_time, response_obj): try: - response_id = ( - "time-" + start_time.strftime("%H-%M-%S-%f") + "_" + response_obj.get("id") - ) + response_id = "time-" + start_time.strftime("%H-%M-%S-%f") + "_" + response_obj.get("id") return response_id except Exception: return None @@ -7945,9 +7003,7 @@ def _get_base_model_from_metadata(model_call_details=None): _get_base_model_from_litellm_call_metadata = getattr( sys.modules[__name__], "_get_base_model_from_litellm_call_metadata" ) - base_model_from_metadata = _get_base_model_from_litellm_call_metadata( - metadata=metadata - ) + base_model_from_metadata = _get_base_model_from_litellm_call_metadata(metadata=metadata) if base_model_from_metadata is not None: return base_model_from_metadata @@ -7964,12 +7020,8 @@ class ModelResponseIterator: def __init__(self, model_response: ModelResponse, convert_to_delta: bool = False): if convert_to_delta is True: _stream_response = ModelResponseStream() - _stream_response.choices[0].delta.content = model_response.choices[ - 0 - ].message.content # type: ignore - self.model_response: Union[ModelResponse, ModelResponseStream] = ( - _stream_response - ) + _stream_response.choices[0].delta.content = model_response.choices[0].message.content # type: ignore + self.model_response: Union[ModelResponse, ModelResponseStream] = _stream_response else: self.model_response = model_response self.is_done = False @@ -8138,9 +7190,7 @@ def any_assistant_message_has_thinking_blocks( for message in messages: if message.get("role") == "assistant": thinking_blocks = message.get("thinking_blocks") - if thinking_blocks is not None and ( - not hasattr(thinking_blocks, "__len__") or len(thinking_blocks) > 0 - ): + if thinking_blocks is not None and (not hasattr(thinking_blocks, "__len__") or len(thinking_blocks) > 0): return True return False @@ -8174,9 +7224,7 @@ def last_assistant_with_tool_calls_has_no_thinking_blocks( # Check if it has thinking_blocks thinking_blocks = last_assistant_with_tools.get("thinking_blocks") - return thinking_blocks is None or ( - hasattr(thinking_blocks, "__len__") and len(thinking_blocks) == 0 - ) + return thinking_blocks is None or (hasattr(thinking_blocks, "__len__") and len(thinking_blocks) == 0) def add_dummy_tool(custom_llm_provider: str) -> List[ChatCompletionToolParam]: @@ -8225,9 +7273,7 @@ def convert_to_dict(message: Union[BaseModel, dict]) -> dict: elif isinstance(message, dict): return message else: - raise TypeError( - f"Invalid message type: {type(message)}. Expected dict or Pydantic model." - ) + raise TypeError(f"Invalid message type: {type(message)}. Expected dict or Pydantic model.") def convert_list_message_to_dict(messages: List): @@ -8325,9 +7371,7 @@ def validate_chat_completion_user_messages(messages: List[AllMessageValues]): for item in user_content: if isinstance(item, dict): if item.get("type") not in ValidUserMessageContentTypes: - raise Exception( - f"invalid content type={item.get('type')}" - ) + raise Exception(f"invalid content type={item.get('type')}") except Exception as e: if isinstance(e, KeyError): raise Exception( @@ -8362,10 +7406,7 @@ def validate_chat_completion_tool_choice( return tool_choice elif isinstance(tool_choice, dict): # Handle Cursor IDE format: {"type": "auto"} -> return as-is - if ( - tool_choice.get("type") in ["auto", "none", "required"] - and "function" not in tool_choice - ): + if tool_choice.get("type") in ["auto", "none", "required"] and "function" not in tool_choice: return tool_choice # Standard OpenAI format: {"type": "function", "function": {...}} @@ -8392,11 +7433,7 @@ def validate_openai_optional_params( Returns: Validated stop parameter (truncated to 4 elements if needed) """ - if ( - stop is not None - and isinstance(stop, list) - and not litellm.disable_stop_sequence_limit - ): + if stop is not None and isinstance(stop, list) and not litellm.disable_stop_sequence_limit: # Truncate to 4 elements if more are provided as openai only supports up to 4 stop sequences if len(stop) > 4: stop = stop[:4] @@ -8407,9 +7444,7 @@ def validate_openai_optional_params( @lru_cache(maxsize=1) def _get_bundled_model_cost_map() -> Dict[str, Any]: try: - model_cost_path = resources.files("litellm").joinpath( - "model_prices_and_context_window_backup.json" - ) + model_cost_path = resources.files("litellm").joinpath("model_prices_and_context_window_backup.json") return json.loads(model_cost_path.read_text()) except Exception: return {} @@ -8701,15 +7736,11 @@ class ProviderConfigManager: # Handle Azure before the generic map so base_model can be threaded through if provider == LlmProviders.AZURE: - return ProviderConfigManager._get_azure_config( - model=model, base_model=base_model - ) + return ProviderConfigManager._get_azure_config(model=model, base_model=base_model) # Initialize provider config map lazily (avoids circular imports) if ProviderConfigManager._PROVIDER_CONFIG_MAP is None: - ProviderConfigManager._PROVIDER_CONFIG_MAP = ( - ProviderConfigManager._build_provider_config_map() - ) + ProviderConfigManager._PROVIDER_CONFIG_MAP = ProviderConfigManager._build_provider_config_map() # O(1) dictionary lookup — Python classes first (custom overrides take priority) config_entry = ProviderConfigManager._PROVIDER_CONFIG_MAP.get(provider) @@ -8739,9 +7770,7 @@ class ProviderConfigManager: ) -> Optional[BaseEmbeddingConfig]: if ( litellm.LlmProviders.VOYAGE == provider - and litellm.VoyageContextualEmbeddingConfig.is_contextualized_embeddings( - model - ) + and litellm.VoyageContextualEmbeddingConfig.is_contextualized_embeddings(model) ): return litellm.VoyageContextualEmbeddingConfig() elif ( @@ -8765,10 +7794,7 @@ class ProviderConfigManager: from litellm.llms.oci.embed.transformation import OCIEmbedConfig return OCIEmbedConfig() - elif ( - litellm.LlmProviders.COHERE == provider - or litellm.LlmProviders.COHERE_CHAT == provider - ): + elif litellm.LlmProviders.COHERE == provider or litellm.LlmProviders.COHERE_CHAT == provider: from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig return CohereEmbeddingConfig() @@ -8831,10 +7857,7 @@ class ProviderConfigManager: api_base: Optional[str], present_version_params: List[str], ) -> BaseRerankConfig: - if ( - litellm.LlmProviders.COHERE == provider - or litellm.LlmProviders.COHERE_CHAT == provider - ): + if litellm.LlmProviders.COHERE == provider or litellm.LlmProviders.COHERE_CHAT == provider: if should_use_cohere_v1_client(api_base, present_version_params): return litellm.CohereRerankConfig() else: @@ -8878,9 +7901,7 @@ class ProviderConfigManager: model: str, provider: LlmProviders, ) -> Optional[BaseAnthropicMessagesConfig]: - return ProviderConfigManager._get_provider_anthropic_messages_config_cached( - model=model, provider=provider - ) + return ProviderConfigManager._get_provider_anthropic_messages_config_cached(model=model, provider=provider) @staticmethod @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) @@ -9011,9 +8032,7 @@ class ProviderConfigManager: from litellm.llms.openai_like.json_loader import JSONProviderRegistry # Resolve provider string for JSON lookup - provider_str = ( - provider.value if isinstance(provider, LlmProviders) else str(provider) - ) + provider_str = provider.value if isinstance(provider, LlmProviders) else str(provider) # Try to convert to enum for Python class lookup first. # Python classes take priority over JSON (they have custom overrides). @@ -9027,16 +8046,12 @@ class ProviderConfigManager: pass # Check Python classes first (custom overrides take priority) - result = ProviderConfigManager._get_python_responses_api_config( - provider_enum, model - ) + result = ProviderConfigManager._get_python_responses_api_config(provider_enum, model) if result is not None: return result # Fall back to JSON providers (generic OpenAI-compatible) - if JSONProviderRegistry.exists( - provider_str - ) and JSONProviderRegistry.supports_responses_api(provider_str): + if JSONProviderRegistry.exists(provider_str) and JSONProviderRegistry.supports_responses_api(provider_str): provider_config = JSONProviderRegistry.get(provider_str) if provider_config is not None: return create_responses_config_class(provider_config)() @@ -9059,10 +8074,7 @@ class ProviderConfigManager: # Note: GPT models (gpt-3.5, gpt-4, gpt-5, etc.) support temperature parameter # O-series models (o1, o3) do not contain "gpt" and have different parameter restrictions is_gpt_model = model and "gpt" in model.lower() - is_o_series = model and ( - "o_series" in model.lower() - or (supports_reasoning(model) and not is_gpt_model) - ) + is_o_series = model and ("o_series" in model.lower() or (supports_reasoning(model) and not is_gpt_model)) if is_o_series: return litellm.AzureOpenAIOSeriesResponsesAPIConfig() @@ -9115,8 +8127,7 @@ class ProviderConfigManager: if not model or not mantle_supports_responses(model, litellm.model_cost): return None return litellm.BedrockMantleResponsesAPIConfig( - use_openai_path=mantle_base_segment(model, litellm.model_cost) - == "openai/v1" + use_openai_path=mantle_base_segment(model, litellm.model_cost) == "openai/v1" ) return None @@ -9866,17 +8877,12 @@ def get_end_user_id_for_cost_tracking( service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking. """ - get_litellm_metadata_from_kwargs = getattr( - sys.modules[__name__], "get_litellm_metadata_from_kwargs" - ) - _metadata = cast( - dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params)) - ) + get_litellm_metadata_from_kwargs = getattr(sys.modules[__name__], "get_litellm_metadata_from_kwargs") + _metadata = cast(dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params))) end_user_id = cast( Optional[str], - litellm_params.get("user_api_key_end_user_id") - or _metadata.get("user_api_key_end_user_id"), + litellm_params.get("user_api_key_end_user_id") or _metadata.get("user_api_key_end_user_id"), ) if litellm.disable_end_user_cost_tracking: return None @@ -9891,17 +8897,13 @@ def get_end_user_id_for_cost_tracking( return end_user_id -def should_use_cohere_v1_client( - api_base: Optional[str], present_version_params: List[str] -): +def should_use_cohere_v1_client(api_base: Optional[str], present_version_params: List[str]): if not api_base: return False uses_v1_params = ("max_chunks_per_doc" in present_version_params) and ( "max_tokens_per_doc" not in present_version_params ) - return api_base.endswith("/v1/rerank") or ( - uses_v1_params and not api_base.endswith("/v2/rerank") - ) + return api_base.endswith("/v1/rerank") or (uses_v1_params and not api_base.endswith("/v2/rerank")) def is_prompt_caching_valid_prompt( @@ -9918,9 +8920,7 @@ def is_prompt_caching_valid_prompt( try: if messages is None and tools is None: return False - if custom_llm_provider is not None and not model.startswith( - custom_llm_provider - ): + if custom_llm_provider is not None and not model.startswith(custom_llm_provider): model = custom_llm_provider + "/" + model token_count = token_counter( messages=messages, @@ -10001,11 +9001,7 @@ def _add_path_to_api_base(api_base: str, ending_path: str) -> str: def get_standard_openai_params(params: dict) -> dict: - return { - k: v - for k, v in params.items() - if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS and v is not None - } + return {k: v for k, v in params.items() if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS and v is not None} def get_non_default_completion_params(kwargs: dict) -> dict: @@ -10088,9 +9084,7 @@ def add_openai_metadata( return None # Only include non-hidden parameters visible_metadata: Dict[str, str] = { - str(k): v - for k, v in metadata.items() - if k != "hidden_params" and isinstance(v, str) + str(k): v for k, v in metadata.items() if k != "hidden_params" and isinstance(v, str) } # max 16 keys allowed by openai - trim down to 16 @@ -10157,9 +9151,7 @@ def return_raw_request(endpoint: CallTypes, kwargs: dict) -> RawRequestTypedDict except Exception as e: received_exception = str(e) - raw_request_typed_dict = litellm_logging_obj.model_call_details.get( - "raw_request_typed_dict" - ) + raw_request_typed_dict = litellm_logging_obj.model_call_details.get("raw_request_typed_dict") if raw_request_typed_dict: return cast(RawRequestTypedDict, raw_request_typed_dict) else: diff --git a/litellm/vector_store_files/main.py b/litellm/vector_store_files/main.py index 0d4d516d03a..c7580625f82 100644 --- a/litellm/vector_store_files/main.py +++ b/litellm/vector_store_files/main.py @@ -43,11 +43,7 @@ def _prepare_registry_credentials( if litellm.vector_store_registry is None: return try: - registry_credentials = ( - litellm.vector_store_registry.get_credentials_for_vector_store( - vector_store_id - ) - ) + registry_credentials = litellm.vector_store_registry.get_credentials_for_vector_store(vector_store_id) if registry_credentials: kwargs.update(registry_credentials) except Exception: @@ -136,14 +132,10 @@ def create( provider=LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Vector store file create is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store file create is not supported for {custom_llm_provider}") local_vars.update(kwargs) - create_request: VectorStoreFileCreateRequest = ( - VectorStoreFileRequestUtils.get_create_request_params(local_vars) - ) + create_request: VectorStoreFileCreateRequest = VectorStoreFileRequestUtils.get_create_request_params(local_vars) create_request["file_id"] = file_id litellm_logging_obj.update_from_kwargs( @@ -252,9 +244,7 @@ def list( timeout: Optional[Union[float, httpx.Timeout]] = None, custom_llm_provider: Optional[str] = None, **kwargs, -) -> Union[ - VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse] -]: +) -> Union[VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse]]: local_vars = locals() try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore @@ -271,14 +261,10 @@ def list( provider=LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Vector store file list is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store file list is not supported for {custom_llm_provider}") local_vars.update(kwargs) - list_query: VectorStoreFileListQueryParams = ( - VectorStoreFileRequestUtils.get_list_query_params(local_vars) - ) + list_query: VectorStoreFileListQueryParams = VectorStoreFileRequestUtils.get_list_query_params(local_vars) litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -385,9 +371,7 @@ def retrieve( provider=LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Vector store file retrieve is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store file retrieve is not supported for {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -478,9 +462,7 @@ def retrieve_content( timeout: Optional[Union[float, httpx.Timeout]] = None, custom_llm_provider: Optional[str] = None, **kwargs, -) -> Union[ - VectorStoreFileContentResponse, Coroutine[Any, Any, VectorStoreFileContentResponse] -]: +) -> Union[VectorStoreFileContentResponse, Coroutine[Any, Any, VectorStoreFileContentResponse]]: local_vars = locals() try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore @@ -497,9 +479,7 @@ def retrieve_content( provider=LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Vector store file content retrieve is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store file content retrieve is not supported for {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -613,14 +593,10 @@ def update( provider=LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Vector store file update is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store file update is not supported for {custom_llm_provider}") local_vars.update(kwargs) - update_request: VectorStoreFileUpdateRequest = ( - VectorStoreFileRequestUtils.get_update_request_params(local_vars) - ) + update_request: VectorStoreFileUpdateRequest = VectorStoreFileRequestUtils.get_update_request_params(local_vars) update_request["attributes"] = attributes litellm_logging_obj.update_from_kwargs( @@ -715,9 +691,7 @@ def delete( timeout: Optional[Union[float, httpx.Timeout]] = None, custom_llm_provider: Optional[str] = None, **kwargs, -) -> Union[ - VectorStoreFileDeleteResponse, Coroutine[Any, Any, VectorStoreFileDeleteResponse] -]: +) -> Union[VectorStoreFileDeleteResponse, Coroutine[Any, Any, VectorStoreFileDeleteResponse]]: local_vars = locals() try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore @@ -734,9 +708,7 @@ def delete( provider=LlmProviders(custom_llm_provider) ) if provider_config is None: - raise ValueError( - f"Vector store file delete is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store file delete is not supported for {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/vector_store_files/utils.py b/litellm/vector_store_files/utils.py index 1ee5b47e306..0f97af0066f 100644 --- a/litellm/vector_store_files/utils.py +++ b/litellm/vector_store_files/utils.py @@ -13,33 +13,23 @@ class VectorStoreFileRequestUtils: @staticmethod def _filter_params(params: Dict[str, Any], model: Any) -> Dict[str, Any]: valid_keys = get_type_hints(model).keys() - return { - key: value - for key, value in params.items() - if key in valid_keys and value is not None - } + return {key: value for key, value in params.items() if key in valid_keys and value is not None} @staticmethod def get_create_request_params( params: Dict[str, Any], ) -> VectorStoreFileCreateRequest: - filtered = VectorStoreFileRequestUtils._filter_params( - params=params, model=VectorStoreFileCreateRequest - ) + filtered = VectorStoreFileRequestUtils._filter_params(params=params, model=VectorStoreFileCreateRequest) return cast(VectorStoreFileCreateRequest, filtered) @staticmethod def get_list_query_params(params: Dict[str, Any]) -> VectorStoreFileListQueryParams: - filtered = VectorStoreFileRequestUtils._filter_params( - params=params, model=VectorStoreFileListQueryParams - ) + filtered = VectorStoreFileRequestUtils._filter_params(params=params, model=VectorStoreFileListQueryParams) return cast(VectorStoreFileListQueryParams, filtered) @staticmethod def get_update_request_params( params: Dict[str, Any], ) -> VectorStoreFileUpdateRequest: - filtered = VectorStoreFileRequestUtils._filter_params( - params=params, model=VectorStoreFileUpdateRequest - ) + filtered = VectorStoreFileRequestUtils._filter_params(params=params, model=VectorStoreFileUpdateRequest) return cast(VectorStoreFileUpdateRequest, filtered) diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 13f2f27d3fa..f768ee75545 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -190,9 +190,7 @@ def create( litellm_params = GenericLiteLLMParams(**kwargs) ## MOCK RESPONSE LOGIC - if litellm_params.mock_response and isinstance( - litellm_params.mock_response, dict - ): + if litellm_params.mock_response and isinstance(litellm_params.mock_response, dict): return mock_vector_store_create_response( mock_response=VectorStoreCreateResponse(**litellm_params.mock_response) ) @@ -212,25 +210,19 @@ def create( custom_llm_provider = custom_llm_provider # get provider config - using vector store custom logger for now - vector_store_provider_config = ( - ProviderConfigManager.get_provider_vector_stores_config( - provider=litellm.LlmProviders(custom_llm_provider), - api_type=api_type, - ) + vector_store_provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, ) if vector_store_provider_config is None: - raise ValueError( - f"Vector store create is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store create is not supported for {custom_llm_provider}") local_vars.update(kwargs) # Get VectorStoreCreateOptionalRequestParams with only valid parameters vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams = ( - VectorStoreRequestUtils.get_requested_vector_store_create_optional_param( - local_vars - ) + VectorStoreRequestUtils.get_requested_vector_store_create_optional_param(local_vars) ) # Pre Call logging @@ -379,11 +371,7 @@ def search( # pull credentials from registry if available if litellm.vector_store_registry is not None and vector_store_id is not None: try: - registry_credentials = ( - litellm.vector_store_registry.get_credentials_for_vector_store( - vector_store_id - ) - ) + registry_credentials = litellm.vector_store_registry.get_credentials_for_vector_store(vector_store_id) kwargs.update(registry_credentials) except Exception: pass @@ -392,9 +380,7 @@ def search( litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs) ## MOCK RESPONSE LOGIC - if litellm_params.mock_response and isinstance( - litellm_params.mock_response, (str, builtins.list) - ): + if litellm_params.mock_response and isinstance(litellm_params.mock_response, (str, builtins.list)): mock_results = None if isinstance(litellm_params.mock_response, builtins.list): mock_results = litellm_params.mock_response # type: ignore[assignment] @@ -415,17 +401,13 @@ def search( custom_llm_provider = custom_llm_provider # get provider config - using vector store custom logger for now - vector_store_provider_config = ( - ProviderConfigManager.get_provider_vector_stores_config( - provider=litellm.LlmProviders(custom_llm_provider), - api_type=api_type, - ) + vector_store_provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, ) if vector_store_provider_config is None: - raise ValueError( - f"Vector store search is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store search is not supported for {custom_llm_provider}") local_vars.update(kwargs) @@ -572,17 +554,13 @@ def retrieve( api_type = None custom_llm_provider = custom_llm_provider - vector_store_provider_config = ( - ProviderConfigManager.get_provider_vector_stores_config( - provider=litellm.LlmProviders(custom_llm_provider), - api_type=api_type, - ) + vector_store_provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, ) if vector_store_provider_config is None: - raise ValueError( - f"Vector store retrieve is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store retrieve is not supported for {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -720,17 +698,13 @@ def list( api_type = None custom_llm_provider = custom_llm_provider - vector_store_provider_config = ( - ProviderConfigManager.get_provider_vector_stores_config( - provider=litellm.LlmProviders(custom_llm_provider), - api_type=api_type, - ) + vector_store_provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, ) if vector_store_provider_config is None: - raise ValueError( - f"Vector store list is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store list is not supported for {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -876,24 +850,18 @@ def update( api_type = None custom_llm_provider = custom_llm_provider - vector_store_provider_config = ( - ProviderConfigManager.get_provider_vector_stores_config( - provider=litellm.LlmProviders(custom_llm_provider), - api_type=api_type, - ) + vector_store_provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, ) if vector_store_provider_config is None: - raise ValueError( - f"Vector store update is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store update is not supported for {custom_llm_provider}") local_vars.update(kwargs) vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams = ( - VectorStoreRequestUtils.get_requested_vector_store_create_optional_param( - local_vars - ) + VectorStoreRequestUtils.get_requested_vector_store_create_optional_param(local_vars) ) litellm_logging_obj.update_from_kwargs( @@ -1025,17 +993,13 @@ def delete( api_type = None custom_llm_provider = custom_llm_provider - vector_store_provider_config = ( - ProviderConfigManager.get_provider_vector_stores_config( - provider=litellm.LlmProviders(custom_llm_provider), - api_type=api_type, - ) + vector_store_provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, ) if vector_store_provider_config is None: - raise ValueError( - f"Vector store delete is not supported for {custom_llm_provider}" - ) + raise ValueError(f"Vector store delete is not supported for {custom_llm_provider}") litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/vector_stores/utils.py b/litellm/vector_stores/utils.py index 9430c68ccc9..27b8d546af1 100644 --- a/litellm/vector_stores/utils.py +++ b/litellm/vector_stores/utils.py @@ -25,9 +25,7 @@ class VectorStoreRequestUtils: VectorStoreSearchOptionalRequestParams instance with only the valid parameters """ valid_keys = get_type_hints(VectorStoreSearchOptionalRequestParams).keys() - filtered_params = { - k: v for k, v in params.items() if k in valid_keys and v is not None - } + filtered_params = {k: v for k, v in params.items() if k in valid_keys and v is not None} optional_params = vector_store_provider_config.map_openai_params( non_default_params=params, @@ -51,8 +49,6 @@ class VectorStoreRequestUtils: VectorStoreCreateOptionalRequestParams instance with only the valid parameters """ valid_keys = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() - filtered_params = { - k: v for k, v in params.items() if k in valid_keys and v is not None - } + filtered_params = {k: v for k, v in params.items() if k in valid_keys and v is not None} return cast(VectorStoreCreateOptionalRequestParams, filtered_params) diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 94f0483e1cc..5070db1c89e 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -25,12 +25,8 @@ else: class VectorStoreIndexRegistry: - def __init__( - self, vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = [] - ): - self.vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = ( - vector_store_indexes - ) + def __init__(self, vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = []): + self.vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = vector_store_indexes def get_vector_store_indexes(self) -> List[LiteLLM_ManagedVectorStoreIndex]: """ @@ -38,9 +34,7 @@ class VectorStoreIndexRegistry: """ return self.vector_store_indexes - def get_vector_store_index_by_name( - self, vector_store_index_name: str - ) -> Optional[LiteLLM_ManagedVectorStoreIndex]: + def get_vector_store_index_by_name(self, vector_store_index_name: str) -> Optional[LiteLLM_ManagedVectorStoreIndex]: """ Returns the vector store index by name """ @@ -49,17 +43,13 @@ class VectorStoreIndexRegistry: return vector_store_index return None - def upsert_vector_store_index( - self, vector_store_index: LiteLLM_ManagedVectorStoreIndex - ): + def upsert_vector_store_index(self, vector_store_index: LiteLLM_ManagedVectorStoreIndex): """ Adds a vector store index to the registry. If it already exists, it will be updated. """ - for i, _vector_store_index in enumerate[LiteLLM_ManagedVectorStoreIndex]( - self.vector_store_indexes - ): + for i, _vector_store_index in enumerate[LiteLLM_ManagedVectorStoreIndex](self.vector_store_indexes): if _vector_store_index.index_name == vector_store_index.index_name: self.vector_store_indexes[i] = vector_store_index return @@ -69,9 +59,7 @@ class VectorStoreIndexRegistry: """ Deletes a vector store index from the registry """ - self.vector_store_indexes = [ - index for index in self.vector_store_indexes if index != vector_store_index - ] + self.vector_store_indexes = [index for index in self.vector_store_indexes if index != vector_store_index] def is_vector_store_index(self, vector_store_index_name: str) -> bool: """ @@ -95,16 +83,12 @@ class VectorStoreIndexRegistry: """ vector_stores_from_db: List[LiteLLM_ManagedVectorStoreIndex] = [] if prisma_client is not None: - _vector_stores_from_db = await ManagedVectorStoreIndexRepository( - prisma_client - ).table.find_many( + _vector_stores_from_db = await ManagedVectorStoreIndexRepository(prisma_client).table.find_many( order={"created_at": "desc"}, ) for vector_store in _vector_stores_from_db: _dict_vector_store = dict(vector_store) - _litellm_managed_vector_store = LiteLLM_ManagedVectorStoreIndex( - **_dict_vector_store - ) + _litellm_managed_vector_store = LiteLLM_ManagedVectorStoreIndex(**_dict_vector_store) vector_stores_from_db.append(_litellm_managed_vector_store) return vector_stores_from_db @@ -112,9 +96,7 @@ class VectorStoreIndexRegistry: class VectorStoreRegistry: def __init__(self, vector_stores: List[LiteLLM_ManagedVectorStore] = []): self.vector_stores: List[LiteLLM_ManagedVectorStore] = vector_stores - self.vector_store_ids_to_vector_store_map: Dict[ - str, LiteLLM_ManagedVectorStore - ] = {} + self.vector_store_ids_to_vector_store_map: Dict[str, LiteLLM_ManagedVectorStore] = {} def _extract_tool_params(self, tool: Dict) -> VectorStoreToolParams: """ @@ -130,9 +112,7 @@ class VectorStoreRegistry: return VectorStoreToolParams(**kwargs) - def get_vector_store_ids_to_run( - self, non_default_params: Dict, tools: Optional[List[Dict]] = None - ) -> List[str]: + def get_vector_store_ids_to_run(self, non_default_params: Dict, tools: Optional[List[Dict]] = None) -> List[str]: """ Returns the vector store ids to run @@ -146,9 +126,7 @@ class VectorStoreRegistry: vector_store_ids.extend(vector_store_ids_param) # 2. check if vector_store_ids is provided as a tool in the request - vector_store_ids = self._get_vector_store_ids_from_tool_calls( - tools=tools, vector_store_ids=vector_store_ids - ) + vector_store_ids = self._get_vector_store_ids_from_tool_calls(tools=tools, vector_store_ids=vector_store_ids) return list(dict.fromkeys(vector_store_ids)) @@ -184,8 +162,7 @@ class VectorStoreRegistry: # Check if all vector_store_ids are recognized in the registry recognised = all( - any(vs.get("vector_store_id") == vs_id for vs in self.vector_stores) - for vs_id in tool_vector_store_ids + any(vs.get("vector_store_id") == vs_id for vs in self.vector_stores) for vs_id in tool_vector_store_ids ) if recognised: @@ -215,9 +192,7 @@ class VectorStoreRegistry: This will return the first vector store found in the registry. """ - vector_store_ids = self.get_vector_store_ids_to_run( - non_default_params=non_default_params, tools=tools - ) + vector_store_ids = self.get_vector_store_ids_to_run(non_default_params=non_default_params, tools=tools) # check if the vector store ids are in the registry if len(vector_store_ids) <= 0: @@ -248,27 +223,21 @@ class VectorStoreRegistry: This ensures synchronization across multiple instances. """ # First check in-memory registry - vector_store = self.get_litellm_managed_vector_store_from_registry( - vector_store_id - ) + vector_store = self.get_litellm_managed_vector_store_from_registry(vector_store_id) if vector_store is not None: return vector_store # Fall back to database if not found in memory if prisma_client is not None: try: - vector_stores_from_db = await self._get_vector_stores_from_db( - prisma_client=prisma_client - ) + vector_stores_from_db = await self._get_vector_stores_from_db(prisma_client=prisma_client) for db_vector_store in vector_stores_from_db: if db_vector_store.get("vector_store_id") == vector_store_id: # Add to in-memory registry for future use self.add_vector_store_to_registry(vector_store=db_vector_store) return db_vector_store except Exception as e: - verbose_logger.debug( - f"Error fetching vector store from database: {str(e)}" - ) + verbose_logger.debug(f"Error fetching vector store from database: {str(e)}") return None @@ -299,14 +268,10 @@ class VectorStoreRegistry: List of vector stores with tool parameters merged into litellm_params """ # Pop vector_store_ids from params - vector_store_ids: List[str] = ( - non_default_params.pop("vector_store_ids", None) or [] - ) + vector_store_ids: List[str] = non_default_params.pop("vector_store_ids", None) or [] # Extract params from tools and collect IDs - params_by_id = self.get_and_pop_recognised_vector_store_tools( - tools=tools, vector_store_ids=vector_store_ids - ) + params_by_id = self.get_and_pop_recognised_vector_store_tools(tools=tools, vector_store_ids=vector_store_ids) vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = [] @@ -318,9 +283,7 @@ class VectorStoreRegistry: # Merge tool params if they exist if vector_store_id in params_by_id: - existing_params = ( - vector_store_copy.get("litellm_params", {}) or {} - ) + existing_params = vector_store_copy.get("litellm_params", {}) or {} tool_params_dict = params_by_id[vector_store_id].to_dict() # Tool params take precedence over existing params tool_params_dict.update(existing_params) @@ -353,14 +316,10 @@ class VectorStoreRegistry: List of vector stores with tool parameters merged into litellm_params """ # Pop vector_store_ids from params - vector_store_ids: List[str] = ( - non_default_params.pop("vector_store_ids", None) or [] - ) + vector_store_ids: List[str] = non_default_params.pop("vector_store_ids", None) or [] # Extract params from tools and collect IDs - params_by_id = self.get_and_pop_recognised_vector_store_tools( - tools=tools, vector_store_ids=vector_store_ids - ) + params_by_id = self.get_and_pop_recognised_vector_store_tools(tools=tools, vector_store_ids=vector_store_ids) vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = [] @@ -378,35 +337,27 @@ class VectorStoreRegistry: if vector_store is not None and prisma_client is not None: try: # Check if it still exists in database - db_vector_store = await ManagedVectorStoresRepository( - prisma_client - ).table.find_unique(where={"vector_store_id": vector_store_id}) + db_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + where={"vector_store_id": vector_store_id} + ) if db_vector_store is None: # Vector store was deleted from database, remove from cache verbose_logger.debug( f"Vector store {vector_store_id} found in memory but deleted from database, removing from cache" ) - self.delete_vector_store_from_registry( - vector_store_id=vector_store_id - ) + self.delete_vector_store_from_registry(vector_store_id=vector_store_id) vector_store = None except Exception as e: - verbose_logger.debug( - f"Error verifying vector store {vector_store_id} in database: {str(e)}" - ) + verbose_logger.debug(f"Error verifying vector store {vector_store_id} in database: {str(e)}") # Fall back to database if not found in memory (or was deleted) if vector_store is None and prisma_client is not None: try: - vector_store = ( - await self.get_litellm_managed_vector_store_from_registry_or_db( - vector_store_id=vector_store_id, prisma_client=prisma_client - ) + vector_store = await self.get_litellm_managed_vector_store_from_registry_or_db( + vector_store_id=vector_store_id, prisma_client=prisma_client ) except Exception as e: - verbose_logger.debug( - f"Error fetching vector store {vector_store_id} from database: {str(e)}" - ) + verbose_logger.debug(f"Error fetching vector store {vector_store_id} from database: {str(e)}") if vector_store is not None: # Create a copy to avoid modifying the registry @@ -442,13 +393,9 @@ class VectorStoreRegistry: """ for vector_store_config in vector_stores_config: # cast to VectorStoreConfig - litellm_vector_store_config = LiteLLM_VectorStoreConfig( - **vector_store_config - ) + litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config) vector_store_name = litellm_vector_store_config.get("vector_store_name") - vector_store_litellm_params: Dict[str, Any] = ( - litellm_vector_store_config.get("litellm_params") or {} - ) + vector_store_litellm_params: Dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {} vector_store_id = vector_store_litellm_params.get("vector_store_id") if vector_store_id is None: @@ -466,12 +413,8 @@ class VectorStoreRegistry: custom_llm_provider=custom_llm_provider, litellm_params=vector_store_litellm_params, vector_store_name=vector_store_name, - vector_store_description=vector_store_litellm_params.get( - "vector_store_description" - ), - vector_store_metadata=vector_store_litellm_params.get( - "vector_store_metadata" - ), + vector_store_description=vector_store_litellm_params.get("vector_store_description"), + vector_store_metadata=vector_store_litellm_params.get("vector_store_metadata"), created_at=datetime.now(timezone.utc), updated_at=datetime.now(timezone.utc), ) @@ -522,9 +465,7 @@ class VectorStoreRegistry: if vector_store.get("vector_store_id") != vector_store_id ] - def update_vector_store_in_registry( - self, vector_store_id: str, updated_data: LiteLLM_ManagedVectorStore - ): + def update_vector_store_in_registry(self, vector_store_id: str, updated_data: LiteLLM_ManagedVectorStore): """Update or add a vector store in the registry""" for i, vector_store in enumerate(self.vector_stores): if vector_store.get("vector_store_id") == vector_store_id: @@ -545,16 +486,12 @@ class VectorStoreRegistry: """ vector_stores_from_db: List[LiteLLM_ManagedVectorStore] = [] if prisma_client is not None: - _vector_stores_from_db = await ManagedVectorStoresRepository( - prisma_client - ).table.find_many( + _vector_stores_from_db = await ManagedVectorStoresRepository(prisma_client).table.find_many( order={"created_at": "desc"}, ) for vector_store in _vector_stores_from_db: _dict_vector_store = dict(vector_store) - _litellm_managed_vector_store = LiteLLM_ManagedVectorStore( - **_dict_vector_store - ) + _litellm_managed_vector_store = LiteLLM_ManagedVectorStore(**_dict_vector_store) vector_stores_from_db.append(_litellm_managed_vector_store) return vector_stores_from_db diff --git a/litellm/videos/main.py b/litellm/videos/main.py index b087f1e88d8..6d81fec36b3 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -206,33 +206,25 @@ def video_generation( ) # get provider config - video_generation_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_generation_provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_generation_provider_config is None: - raise ValueError( - f"video generation is not supported for {custom_llm_provider}" - ) + raise ValueError(f"video generation is not supported for {custom_llm_provider}") local_vars.update(kwargs) # Get VideoGenerationOptionalRequestParams with only valid parameters video_generation_optional_params: VideoCreateOptionalRequestParams = ( - VideoGenerationRequestUtils.get_requested_video_generation_optional_param( - local_vars - ) + VideoGenerationRequestUtils.get_requested_video_generation_optional_param(local_vars) ) # Get optional parameters for the video generation API - video_generation_request_params: Dict = ( - VideoGenerationRequestUtils.get_optional_params_video_generation( - model=model, - video_generation_provider_config=video_generation_provider_config, - video_generation_optional_params=video_generation_optional_params, - ) + video_generation_request_params: Dict = VideoGenerationRequestUtils.get_optional_params_video_generation( + model=model, + video_generation_provider_config=video_generation_provider_config, + video_generation_optional_params=video_generation_optional_params, ) # Pre Call logging @@ -337,17 +329,13 @@ def video_content( litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_provider_config is None: - raise ValueError( - f"video support download is not supported for {custom_llm_provider}" - ) + raise ValueError(f"video support download is not supported for {custom_llm_provider}") local_vars.update(kwargs) # For video content download, we don't need complex optional parameter handling @@ -613,11 +601,9 @@ def video_remix( litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_remix_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_remix_provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_remix_provider_config is None: @@ -712,9 +698,7 @@ async def avideo_list( # get custom llm provider so we can use this for mapping exceptions if custom_llm_provider is None: - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model="", api_base=local_vars.get("api_base", None) - ) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model="", api_base=local_vars.get("api_base", None)) func = partial( video_list, @@ -832,11 +816,9 @@ def video_list( litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_list_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_list_provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_list_provider_config is None: @@ -1058,11 +1040,9 @@ def video_status( litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_status_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_status_provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_status_provider_config is None: @@ -1201,17 +1181,13 @@ def video_create_character( litellm_params = GenericLiteLLMParams(**kwargs) - provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if provider_config is None: - raise ValueError( - f"video create character is not supported for {custom_llm_provider}" - ) + raise ValueError(f"video create character is not supported for {custom_llm_provider}") local_vars.update(kwargs) request_params: Dict = {"name": name} @@ -1330,17 +1306,13 @@ def video_get_character( litellm_params = GenericLiteLLMParams(**kwargs) - provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if provider_config is None: - raise ValueError( - f"video get character is not supported for {custom_llm_provider}" - ) + raise ValueError(f"video get character is not supported for {custom_llm_provider}") local_vars.update(kwargs) request_params: Dict = {"character_id": character_id} @@ -1462,11 +1434,9 @@ def video_edit( litellm_params = GenericLiteLLMParams(**kwargs) - provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if provider_config is None: @@ -1597,17 +1567,13 @@ def video_extension( litellm_params = GenericLiteLLMParams(**kwargs) - provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + provider_config: Optional[BaseVideoConfig] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if provider_config is None: - raise ValueError( - f"video extension is not supported for {custom_llm_provider}" - ) + raise ValueError(f"video extension is not supported for {custom_llm_provider}") local_vars.update(kwargs) request_params: Dict = { diff --git a/litellm/videos/utils.py b/litellm/videos/utils.py index 06dfa5d1396..42e0d7f4a27 100644 --- a/litellm/videos/utils.py +++ b/litellm/videos/utils.py @@ -69,14 +69,11 @@ class VideoGenerationRequestUtils: base_params_raw = { key: value for key, value in params.items() - if key not in {"kwargs", "extra_body", "prompt", "model"} - and value is not None + if key not in {"kwargs", "extra_body", "prompt", "model"} and value is not None } base_params = filter_out_litellm_params(kwargs=base_params_raw) - cleaned_kwargs = filter_out_litellm_params( - kwargs={k: v for k, v in raw_kwargs.items() if v is not None} - ) + cleaned_kwargs = filter_out_litellm_params(kwargs={k: v for k, v in raw_kwargs.items() if v is not None}) optional_params: Dict[str, Any] = { **base_params, diff --git a/ruff.toml b/ruff.toml index 082a8f83a0c..a09bc663ff1 100644 --- a/ruff.toml +++ b/ruff.toml @@ -12,10 +12,6 @@ lint.external = [ ] line-length = 120 -# `ruff format` (replacing Black) must wrap at 88, the width Black used and the whole -# history is formatted to. The global line-length stays 120 because E501 and the import -# sorter (I001, strict gate) are tuned to it and ruff has no per-formatter line-length, -# so 88 is passed at the `ruff format --line-length 88` call sites (Makefile + CI). format.exclude = ["**/enterprise/**"] # Was the top-level `exclude`. Scoped to lint so `ruff format` still formats these paths